fix(sglang): attach the KVCR linker to the external state agent

components/src/dynamo/sglang/args.py

@@ -44,6 +44,7 @@
 from dynamo.sglang.capacity import sglang_dp_layout
 from dynamo.sglang.elastic_ep_preflight import check_elastic_ep_backend
 from dynamo.sglang.gateway import effective_gateway_workers
+from dynamo.sglang.state_agent import validate_state_agent_worker
 
 configure_dynamo_logging()
 PREFILL_DECODE_DISAGGREGATION_MODE = "pd"
@@ -105,6 +106,7 @@ def __init__(
         # SGLang chose the value or the user did.
         self.attention_backend_from_cli = attention_backend_from_cli
         self.serving_mode = self._set_serving_strategy()
+        validate_state_agent_worker(self)
         if getattr(dynamo_args, "freeze_gc_after_init", False) and (
             self.serving_mode == DisaggregationMode.PREFILL
             or any(

components/src/dynamo/sglang/backend_args.py

@@ -56,6 +56,15 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
 
         g = parser.add_argument_group("Dynamo SGLang Options")
 
+        add_argument(
+            g,
+            flag_name="--kvcr-state-agent-owner",
+            env_var="DYN_SGL_KVCR_STATE_AGENT_OWNER",
+            default=None,
+            arg_type=str,
+            help="Attach an aggregated direct KVCR linker to the external KV state agent using this stable CacheOwnerId.",
+        )
+
         add_negatable_bool_argument(
             g,
             flag_name="--use-sglang-tokenizer",
@@ -216,6 +225,7 @@ class DynamoSGLangConfig(ConfigBase):
 
     use_sglang_tokenizer: bool
     gateway_workers: Optional[int] = None
+    kvcr_state_agent_owner: Optional[str] = None
     # Internal roles derived from the canonical multimodal arguments in args.py.
     multimodal_encode_worker: bool = False
     multimodal_worker: bool = False

components/src/dynamo/sglang/init_llm.py

@@ -34,6 +34,7 @@
 )
 from dynamo.sglang.register import register_model_with_readiness_gate
 from dynamo.sglang.request_handlers import DecodeWorkerHandler, PrefillWorkerHandler
+from dynamo.sglang.state_agent import StateAgentLifecycle
 
 
 async def _warmup_prefill_engine(
@@ -201,7 +202,9 @@ async def init_decode(
         )
 
     body_failed = True
+    state_agent = StateAgentLifecycle()
     try:
+        await state_agent.start(config, generate_endpoint)
         gather_tasks = [
             generate_endpoint.serve_endpoint(
                 handler.generate,
@@ -246,12 +249,15 @@ async def init_decode(
     else:
         body_failed = False
     finally:
-        await finish_worker_teardown(
-            metrics_task,
-            handler.cleanup,
-            run_deferred_handlers,
-            body_failed=body_failed,
-        )
+        try:
+            await state_agent.close()
+        finally:
+            await finish_worker_teardown(
+                metrics_task,
+                handler.cleanup,
+                run_deferred_handlers,
+                body_failed=body_failed,
+            )
 
 
 async def init_prefill(

components/src/dynamo/sglang/kv_hints.py

@@ -104,20 +104,22 @@ def kvcr_hint_endpoints(
     }
 
 
-def publish_kvcr_hint_capabilities(
-    runtime_config, engine, server_args, worker_type
-) -> None:
-    """Opt in only for the supported single-host KVCR linker or HiCache path."""
+def kvcr_hint_sources(
+    server_args: Any, worker_type: Any
+) -> tuple[str, dict[str, str]] | None:
+    """The role and per-DP-rank source-control endpoints this worker advertises.
+
+    None when the worker runs no supported single-host KVCR linker or HiCache
+    path with remote hints enabled.
+    """
     server_args = resolved_server_args(server_args)
     mode = kvcr_mode(server_args)
     if (
-        engine is None
-        or mode is None
+        mode is None
         or getattr(server_args, "nnodes", 1) != 1
         or getattr(server_args, "pp_size", 1) != 1
-        or not filter_supported_async_generate_kwargs(engine, {"kv_hints": None})
     ):
-        return
+        return None
     extra = server_args.hicache_storage_backend_extra_config
     if isinstance(extra, str):
         extra = extra.strip()
@@ -133,15 +135,15 @@ def publish_kvcr_hint_capabilities(
         raise ValueError("KVCR HiCache extra configuration must be an object")
     # Legacy linker revisions have no switch; newer ones can disable hints.
     if extra.get("enable_remote_hint", True) is not True:
-        return
+        return None
     if (
         mode == "linker"
         and extra.get("advertise_host") is None
         and extra.get("enable_remote_hint", False) is not True
     ):
         # The newer control_advertise_host dialect defaults remote hints off.
         # The legacy advertise_host runtime has no switch and enables peers.
-        return
+        return None
     if mode == "linker" and (
         "control_port" not in extra
         or (
@@ -151,14 +153,14 @@ def publish_kvcr_hint_capabilities(
     ):
         # Local-only/default linkers can still serve, but cannot advertise a
         # guessed source-control endpoint to another worker.
-        return
+        return None
     # WorkerType is a PyO3 enum, not enum.Enum: its wire role is __str__,
     # and accessing .value aborts registration after the model has loaded.
     role = str(worker_type)
     if role == "agg":
         role = "aggregated"
     if role not in ("aggregated", "prefill", "decode"):
-        return
+        return None
     # DP attention partitions TP into independent attention/cache groups.
     dp_size = server_args.attn_dp_size
     attn_cp_size = server_args.attn_cp_size
@@ -171,6 +173,21 @@ def publish_kvcr_hint_capabilities(
         pp_size=server_args.pp_size,
         mode=mode,
     )
+    return role, endpoints
+
+
+def publish_kvcr_hint_capabilities(
+    runtime_config, engine, server_args, worker_type
+) -> None:
+    """Opt in only for the supported single-host KVCR linker or HiCache path."""
+    if engine is None or not filter_supported_async_generate_kwargs(
+        engine, {"kv_hints": None}
+    ):
+        return
+    sources = kvcr_hint_sources(server_args, worker_type)
+    if sources is None:
+        return
+    role, endpoints = sources
     for key, value in (
         (KV_HINT_TRANSFER_WORKER_TYPE_RUNTIME_KEY, role),
         (KV_HINT_TRANSFER_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY, endpoints),

components/src/dynamo/sglang/publisher.py

@@ -396,7 +396,13 @@ def init_kv_event_publish(self) -> List[KvEventPublisher]:
             List of KvEventPublisher instances if KV event publishing is enabled,
             empty list otherwise.
         """
-        if self.dynamo_args.use_kv_events and not publishes_kv_events(self.server_args):
+        if getattr(self.dynamo_args, "kvcr_state_agent_owner", None) is not None:
+            # The state agent consumes the raw stream; legacy normalization
+            # rejects KVCR-owned residency events. Keep use_kv_events advertised.
+            pass
+        elif self.dynamo_args.use_kv_events and not publishes_kv_events(
+            self.server_args
+        ):
             logging.info(
                 "Non-leader node (node_rank=%s) shares the leader's single KV "
                 "rank slice; skipping KV event publishing so the router sees "

components/src/dynamo/sglang/register.py

@@ -562,6 +562,8 @@ async def get_runtime_config(
         dynamo_args.enable_local_indexer and not is_decode_worker
     )
     runtime_config.kv_event_publishing_enabled = dynamo_args.use_kv_events
+    if getattr(dynamo_args, "kvcr_state_agent_owner", None) is not None:
+        runtime_config.kv_event_source_mode = "state_agent_v2"
 
     start_dp_rank, end_dp_rank = model_card_dp_rank_bounds(server_args)
     registered_dp_size = end_dp_rank - start_dp_rank

components/src/dynamo/sglang/state_agent.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Opt-in state-agent attachment for an aggregated KVCR linker worker."""
+
+import asyncio
+import json
+import logging
+import re
+from typing import Any
+from urllib.parse import urlparse
+
+from dynamo.common.constants import DisaggregationMode
+from dynamo.llm import KvStateAttachmentOwner, WorkerType
+from dynamo.sglang._compat import resolved_server_args
+from dynamo.sglang.capacity import kv_event_block_size
+from dynamo.sglang.kv_hints import kvcr_hint_sources, kvcr_mode
+
+_CACHE_OWNER_PATTERN = re.compile(
+    r"^[0-9a-f]{32}:[0-9a-f]{32}/[0-9a-f]{16}/[0-9a-f]{32}$"
+)
+_CLEANUP_TIMEOUT_SECS = 30.0
+
+
+class StateAgentLifecycle:
+    """Start one attachment after engine init and close it once on shutdown."""
+
+    def __init__(self) -> None:
+        self._owner: KvStateAttachmentOwner | None = None
+
+    async def start(self, config: Any, endpoint: Any) -> None:
+        descriptor = state_agent_descriptor(config)
+        if descriptor is None:
+            return
+        owner = KvStateAttachmentOwner(endpoint, endpoint.connection_id(), [descriptor])
+        # Cancelling this await drops the native waiter, which withdraws its
+        # own registrations; nothing is left for close() to undo.
+        await owner.start()
+        self._owner = owner
+
+    async def close(self) -> None:
+        owner, self._owner = self._owner, None
+        if owner is None:
+            return
+        try:
+            await asyncio.wait_for(owner.close(), _CLEANUP_TIMEOUT_SECS)
+        except asyncio.TimeoutError:
+            logging.warning(
+                "KV state attachment cleanup timed out; native cleanup continues"
+            )
+
+
+def validate_state_agent_worker(config: Any) -> str | None:
+    """Validate launch intent without forcing SGLang's engine resolution."""
+    dynamo_args = config.dynamo_args
+    owner = getattr(dynamo_args, "kvcr_state_agent_owner", None)
+    if owner is None:
+        return None
+    if not isinstance(owner, str) or not _CACHE_OWNER_PATTERN.fullmatch(owner):
+        raise ValueError("--kvcr-state-agent-owner requires a canonical CacheOwnerId")
+    server_args = resolved_server_args(config.server_args)
+    if config.serving_mode != DisaggregationMode.AGGREGATED:
+        raise ValueError("KVCR state agent supports only aggregated LLM workers")
+    if kvcr_mode(server_args) != "linker" or not dynamo_args.use_kv_events:
+        raise ValueError(
+            "KVCR state agent requires the direct KVCR linker and KV events"
+        )
+    return owner
+
+
+def state_agent_descriptor(config: Any) -> dict[str, Any] | None:
+    """Read final event and linker coordinates after engine initialization."""
+    owner = validate_state_agent_worker(config)
+    if owner is None:
+        return None
+    server_args = resolved_server_args(config.server_args)
+    dynamo_args = config.dynamo_args
+    events = server_args.kv_events_config
+    events = json.loads(events) if isinstance(events, str) else events
+    if not isinstance(events, dict) or events.get("publisher") != "zmq":
+        raise ValueError("KVCR state agent requires a ZMQ event publisher")
+    if events.get("topic", "") != "":
+        raise ValueError("KVCR state agent requires an empty event topic")
+    endpoint = urlparse(events.get("endpoint", ""))
+    if endpoint.scheme != "tcp" or not endpoint.hostname or not endpoint.port:
+        raise ValueError("KVCR state agent requires an explicit TCP event port")
+    # Registration advertises the same source; its resolver already encodes
+    # the linker dialects, and the raw event stream shares the advertise host.
+    sources = kvcr_hint_sources(server_args, WorkerType.Aggregated)
+    if sources is None:
+        raise ValueError("KVCR state agent requires remote linker hints")
+    hint_endpoint = sources[1]["0"]
+    host = urlparse(hint_endpoint).netloc.rsplit(":", 1)[0]
+    return {
+        "cache_owner_id": owner,
+        "global_dp_rank": 0,
+        "kv_state_endpoint": dynamo_args.kv_state_endpoint
+        or f"{dynamo_args.namespace}/{dynamo_args.component}/{dynamo_args.endpoint}",
+        "indexer_domain_id": owner.split("/", 1)[0],
+        "kv_block_size": kv_event_block_size(server_args),
+        "raw_zmq_endpoint": f"tcp://{host}:{endpoint.port}",
+        "raw_topic": "",
+        "image_token_id": None,
+        "video_token_id": None,
+        "router_hint_source": {
+            "source_control_endpoint": hint_endpoint,
+            "worker_type": "aggregated",
+        },
+    }