feat: Add KV publisher and receiver. Add KV aware routing example.

Signed-off-by: Neelay Shah <neelays@nvidia.com>
Co-authored-by: aflowers <aflowers@nvidia.com>
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>
Co-authored-by: hongkuanz <hongkuanz@nvidia.com>
Co-authored-by: Neelay Shah <neelays@nvidia.com>

runtime/rust/python-wheel/python/triton_distributed_rs/__init__.py → python-wheel/python/triton_distributed_rs/__init__.py

@@ -20,6 +20,7 @@ from typing import Any, AsyncGenerator, Callable, Type
 
 from pydantic import BaseModel, ValidationError
 from triton_distributed_rs._core import DistributedRuntime
+from triton_distributed_rs._core import KvRouter as KvRouter
 
 
 def triton_worker():

runtime/rust/python-wheel/python/triton_distributed_rs/_core.pyi → python-wheel/python/triton_distributed_rs/_core.pyi

-from typing import AsyncGenerator, AsyncIterator, Callable
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import AsyncGenerator, AsyncIterator, Callable, List
 
 class JsonLike:
     """
@@ -54,6 +69,12 @@ class Component:
         """
         ...
 
+    def event_subject(self, name: str) -> str:
+        """
+        Create an event subject
+        """
+        ...
+
 class Endpoint:
     """
     An Endpoint is a single API endpoint
@@ -74,6 +95,12 @@ class Endpoint:
         """
         ...
 
+    async def lease_id(self) -> int:
+        """
+        Return primary lease id. Currently, cannot set a different lease id.
+        """
+        ...
+
 class Client:
     """
     A client capable of calling served instances of an endpoint
@@ -98,3 +125,22 @@ class Client:
         Pick a specific instance of the endpoint
         """
         ...
+
+class KvRouter:
+    """
+    The runtime object for a distributed NOVA applications
+    """
+
+    ...
+
+    def __init__(self, drt: DistributedRuntime, component: Component) -> KvRouter:
+        """
+        Create a `KvRouter` object that is associated with the `component`
+        """
+
+    def schedule(self, token_ids: List[int], lora_id: int) -> str:
+        """
+        Return the worker id that should handle the given token ids,
+        exception will be raised if there is no worker available.
+        """
+        ...

python-wheel/rust/llm.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::*;
+
+pub mod kv;

runtime/rust/Cargo.toml

@@ -74,3 +74,5 @@ rand = { version = "0.8"}
 
 [dev-dependencies]
 assert_matches = "1.5.0"
+env_logger = "0.11"
+rstest = "0.23.0"