feat(router): publish GMS placement from workers

Signed-off-by: Maksim Khadkevich <mkhadkevich@nvidia.com>

Author: hutm

components/src/dynamo/router/__main__.py

@@ -203,6 +203,7 @@ async def worker(runtime: DistributedRuntime):
         f"router_track_output_blocks={config.router_track_output_blocks}, "
         f"router_assume_kv_reuse={config.router_assume_kv_reuse}, "
         f"router_track_prefill_tokens={config.router_track_prefill_tokens}, "
+        f"router_gms_decode_transfer={config.router_gms_decode_transfer}, "
         f"router_ttl_secs={config.router_ttl_secs}"
     )
 

components/src/dynamo/sglang/register.py

@@ -291,6 +291,10 @@ async def _get_runtime_config(
         dynamo_args.enable_local_indexer and not is_decode_worker
     )
 
+    if os.environ.get("GMS_SGLANG_DAEMON_SOCKET"):
+        runtime_config.set_gms_placement_enabled()
+        logging.info("Publishing GMS placement capability to discovery")
+
     start_dp_rank, end_dp_rank = model_card_dp_rank_bounds(server_args)
     registered_dp_size = end_dp_rank - start_dp_rank
     runtime_config.data_parallel_start_rank = start_dp_rank

components/src/dynamo/sglang/request_handlers/llm/decode_handler.py

@@ -14,6 +14,10 @@
 from dynamo.common.constants import DisaggregationMode
 from dynamo.common.multimodal.image_loader import ImageLoader
 from dynamo.common.utils.engine_response import normalize_finish_reason
+from dynamo.gms_router_policy import (
+    maybe_fetch_gms_placement,
+    resolve_env_gms_daemon_socket,
+)
 from dynamo.sglang._compat import filter_supported_async_generate_kwargs
 from dynamo.sglang.args import Config
 from dynamo.sglang.publisher import DynamoSglangPublisher
@@ -444,6 +448,13 @@ async def generate(
         if lora_path:
             logging.debug(f"Request {context.id()} will use LoRA adapter: {lora_path}")
 
+        await maybe_fetch_gms_placement(
+            request,
+            resolve_env_gms_daemon_socket("GMS_SGLANG_DAEMON_SOCKET"),
+            logger=logging.getLogger(__name__),
+            request_id=context.id(),
+        )
+
         if self.serving_mode == DisaggregationMode.DECODE:
             # Check if bootstrap_info is pre-computed in the request (from frontend)
             bootstrap_info = request.get("bootstrap_info")

components/src/dynamo/trtllm/request_handlers/handler_base.py

@@ -36,6 +36,10 @@
 from dynamo._core import Client, Context
 from dynamo.common.gms_failover import release_attached_gms_failover_lock
 from dynamo.common.utils.structural_tag import serialize_structural_tag
+from dynamo.gms_router_policy import (
+    maybe_fetch_gms_placement,
+    resolve_env_gms_daemon_socket,
+)
 from dynamo.health_check import HEALTH_CHECK_KEY
 from dynamo.llm.exceptions import EngineShutdown
 from dynamo.logits_processing.examples import HelloWorldLogitsProcessor
@@ -106,12 +110,20 @@ def _collective_rpc(self, method: str, rpc_tags: list[str]) -> None:
         """Call TRT-LLM collective_rpc for KV cache sleep/wake."""
         rpc = getattr(self._engine.llm, "_collective_rpc", None)
         if rpc is None:
-            logger.warning(
-                "TRT-LLM does not expose _collective_rpc; skipping %s", method
+            raise RuntimeError(
+                f"TRT-LLM does not expose _collective_rpc; cannot {method} KV cache"
             )
-            return
 
-        rpc(method, args=(rpc_tags,), kwargs={}, non_block=False)
+        try:
+            rpc(method, args=(rpc_tags,), kwargs={}, non_block=False)
+        except NotImplementedError as exc:
+            message = str(exc)
+            if "MPI collective_rpc only supports model_world_size == 1" in message:
+                raise RuntimeError(
+                    "TRT-LLM MPI executor does not support "
+                    f"collective_rpc({method}) for multi-rank KV cache control"
+                ) from exc
+            raise
 
     @staticmethod
     def _release_gms_weights() -> None:
@@ -1004,6 +1016,16 @@ async def _generate_locally_impl(
         # Normalize OpenAI format to TRT-LLM internal format
         self._normalize_request_format(request)
 
+        await maybe_fetch_gms_placement(
+            request,
+            resolve_env_gms_daemon_socket(
+                "GMS_TRTLLM_DAEMON_SOCKET",
+                default_when_cross_node="/tmp/gms.sock",
+            ),
+            logger=logging.getLogger(__name__),
+            request_id=context.id(),
+        )
+
         # Setup disaggregated params based on PREFILL/DECODE mode
         (
             disaggregated_params,

components/src/dynamo/trtllm/tests/test_trtllm_sleep_wake_handlers.py

@@ -20,7 +20,10 @@
         allow_module_level=True,
     )
 
-from dynamo.trtllm.request_handlers.handler_base import HandlerBase
+from dynamo.trtllm.request_handlers.handler_base import (
+    HandlerBase,
+    TRTLLMEngineQuiesceController,
+)
 
 pytestmark = [
     pytest.mark.unit,
@@ -101,6 +104,33 @@ async def test_mark_request_finished_is_idempotent():
     assert handler._inflight_requests == 0
 
 
+@pytest.mark.asyncio
+async def test_trtllm_quiesce_fails_when_collective_rpc_missing():
+    controller = TRTLLMEngineQuiesceController(SimpleNamespace(llm=SimpleNamespace()))
+
+    with pytest.raises(RuntimeError, match="_collective_rpc"):
+        await controller.quiesce(["kv_cache"])
+
+    assert controller.is_quiesced is False
+
+
+@pytest.mark.asyncio
+async def test_trtllm_quiesce_fails_closed_for_unsupported_multi_rank_rpc():
+    def collective_rpc(*args, **kwargs):
+        raise NotImplementedError(
+            "MPI collective_rpc only supports model_world_size == 1"
+        )
+
+    controller = TRTLLMEngineQuiesceController(
+        SimpleNamespace(llm=SimpleNamespace(_collective_rpc=collective_rpc))
+    )
+
+    with pytest.raises(RuntimeError, match="multi-rank KV cache control"):
+        await controller.quiesce(["kv_cache"])
+
+    assert controller.is_quiesced is False
+
+
 # ---------------------------------------------------------------------------
 # release_memory_occupation
 # ---------------------------------------------------------------------------

components/src/dynamo/trtllm/workers/llm_worker.py

@@ -49,6 +49,7 @@
 )
 from dynamo.common.utils.runtime import parse_endpoint
 from dynamo.common.utils.topology import apply_topology_config
+from dynamo.gms_router_policy import resolve_env_gms_daemon_socket
 from dynamo.llm import (
     KvEventPublisher,
     MediaDecoder,
@@ -634,6 +635,13 @@ async def init_llm_worker(
             config.enable_local_indexer
             and config.disaggregation_mode != DisaggregationMode.DECODE
         )
+        gms_daemon_socket = resolve_env_gms_daemon_socket(
+            "GMS_TRTLLM_DAEMON_SOCKET",
+            default_when_cross_node="/tmp/gms.sock",
+        )
+        if gms_daemon_socket:
+            runtime_config.set_gms_placement_enabled()
+            logging.info("Publishing GMS placement capability to discovery")
         # Set data_parallel_size for attention DP mode
         # This enables the router's scheduler to correctly iterate over all dp_ranks
         # Need to name ADP as `data_parallel_size` for parity with other frameworks

components/src/dynamo/vllm/handlers.py

@@ -57,6 +57,10 @@
 from dynamo.common.utils.input_params import InputParamManager
 from dynamo.common.utils.structural_tag import serialize_structural_tag
 from dynamo.common.utils.time_section import time_and_log_code_section
+from dynamo.gms_router_policy import (
+    maybe_fetch_gms_placement,
+    resolve_vllm_gms_daemon_socket,
+)
 from dynamo.llm import (
     KvEventPublisher,
     ModelInput,
@@ -2303,6 +2307,13 @@ async def _generate_token_mode(self, request, context, request_id):
 
         mm_processor_kwargs = self._get_mm_processor_kwargs(request)
 
+        await maybe_fetch_gms_placement(
+            request,
+            resolve_vllm_gms_daemon_socket(self.engine_client.vllm_config),
+            logger=logger,
+            request_id=request_id,
+        )
+
         multi_modal_data: Dict[str, Any] | None = None
         pre_rendered: Dict[str, Any] | None = None
         if is_decode_only:

components/src/dynamo/vllm/main.py

@@ -91,6 +91,7 @@ def _maybe_isolate_jit_cache_dirs_by_container() -> dict[str, tuple[str | None,
 )
 from dynamo.common.utils.runtime import create_runtime
 from dynamo.common.utils.topology import apply_topology_config
+from dynamo.gms_router_policy import resolve_vllm_gms_daemon_socket
 from dynamo.llm import (
     KvEventPublisher,
     MediaDecoder,
@@ -947,6 +948,10 @@ async def register_vllm_model(
     if stream_interval is not None:
         runtime_config.set_engine_specific("stream_interval", str(stream_interval))
 
+    if resolve_vllm_gms_daemon_socket(vllm_config):
+        runtime_config.set_gms_placement_enabled()
+        logging.info("Publishing GMS placement capability to discovery")
+
     runtime_config.data_parallel_start_rank = dp_range[0]
     runtime_config.data_parallel_size = dp_range[1]
 

lib/llm/src/kv_router/publisher/batching.rs

@@ -96,6 +96,7 @@ impl BatchingState {
                     data: KvCacheEventData::Removed(filtered),
                     dp_rank,
                 },
+                None,
             )
             .await;
             emitted = true;
@@ -112,6 +113,7 @@ impl BatchingState {
                     data: KvCacheEventData::Stored(data),
                     dp_rank,
                 },
+                None,
             )
             .await;
             emitted = true;

lib/llm/src/kv_router/publisher/event_processor.rs

@@ -153,6 +153,7 @@ pub(super) async fn run_event_processor_loop<P: RouterEventSink + Send + Sync +
                                 data: KvCacheEventData::Cleared,
                                 dp_rank: event.dp_rank,
                             },
+                            None,
                         )
                         .await;
                         batching_state.next_publish_id += 1;