feat(backends): handle conditional disagg bypass annotation

components/src/dynamo/sglang/backend_args.py

@@ -165,6 +165,13 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
             "(1=minimal, 2=per-request [default], 3=+decode_loop, 4=full).",
         )
 
+        g.add_argument(
+            "--enable-conditional-disagg",
+            action=argparse.BooleanOptionalAction,
+            default=False,
+            help=argparse.SUPPRESS,
+        )
+
 
 class DynamoSGLangConfig(ConfigBase):
     """Configuration for Dynamo SGLang wrapper (SGLang-specific only)."""
@@ -184,6 +191,7 @@ class DynamoSGLangConfig(ConfigBase):
     video_generation_worker: bool
     enable_rl: bool
     frontend_decoding: bool = False
+    enable_conditional_disagg: bool = False
     sglang_trace_level: int
 
     # Extra served names beyond the primary, parsed from --served-model-name.
@@ -196,6 +204,13 @@ def validate(self) -> None:
                 str(self.embedding_transfer_mode)
             )
 
+        if self.enable_conditional_disagg:
+            raise ValueError(
+                "--enable-conditional-disagg is not supported by the SGLang "
+                "backend yet. SGLang decode workers cannot run "
+                "conditional-disagg bypass requests as local prefill+decode."
+            )
+
         if (self.disagg_config is not None) ^ (self.disagg_config_key is not None):
             raise ValueError(
                 "Both 'disagg_config' and 'disagg_config_key' must be provided together."

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

@@ -39,6 +39,17 @@
     "top_k",
     "min_p",
 )
+BYPASS_REMOTE_PREFILL_ANNOTATION = "x-bypass-remote-prefill"
+
+
+def _raise_if_conditional_disagg_bypass(request: Dict[str, Any]) -> None:
+    if BYPASS_REMOTE_PREFILL_ANNOTATION not in (request.get("annotations") or []):
+        return
+    raise RuntimeError(
+        "SGLang backend does not support conditional-disagg bypass yet. "
+        "Disable --router-conditional-disagg or use a backend with "
+        "decode-side local prefill support."
+    )
 
 
 def _nvext_extra_field_requested(request: Dict[str, Any], field: str) -> bool:
@@ -341,6 +352,7 @@ async def generate(
             RuntimeError: If no bootstrap info received from prefill worker.
         """
         logging.debug(f"New Request ID: {context.id()}")
+        _raise_if_conditional_disagg_bypass(request)
         trace_id = context.trace_id
         sampling_params = self._build_sampling_params(request)
         input_param = self._get_input_param(request)

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

@@ -70,6 +70,8 @@
 
 logger = logging.getLogger(__name__)
 
+BYPASS_REMOTE_PREFILL_ANNOTATION = "x-bypass-remote-prefill"
+
 
 class TRTLLMEnginePauseController:
     """Adapts TRT-LLM sleep/wake to the standard pause controller interface.
@@ -703,10 +705,20 @@ def _setup_disaggregated_params_for_mode(
         disaggregated_params = None
         epd_metadata: dict[str, Any] = {}
 
-        # Canary probe: use its pre-built disagg params (skip prefill_result decode
-        # and skip the mode-specific request_type overrides).
-        if request.get(HEALTH_CHECK_KEY) and request.get("disaggregated_params"):
-            return LlmDisaggregatedParams(**request["disaggregated_params"]), None, {}
+        use_request_disagg_params = request.get(HEALTH_CHECK_KEY) or (
+            self.disaggregation_mode == DisaggregationMode.DECODE
+            and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or [])
+        )
+
+        # Canary probes and conditional-disagg bypasses run a full
+        # context+generation request on a disagg-mode worker, so they use
+        # the pre-built params and skip the normal prefill-result handoff.
+        if use_request_disagg_params and request.get("disaggregated_params"):
+            return (
+                LlmDisaggregatedParams(**request["disaggregated_params"]),
+                ep_disaggregated_params,
+                {},
+            )
 
         # PREFILL mode: setup context_only params
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
@@ -984,6 +996,16 @@ async def _generate_locally_impl(
 
         # Normalize OpenAI format to TRT-LLM internal format
         self._normalize_request_format(request)
+        bypass_remote_prefill = (
+            self.disaggregation_mode == DisaggregationMode.DECODE
+            and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or [])
+        )
+        if bypass_remote_prefill:
+            logging.debug(
+                "DECODE: conditional-disagg bypass annotation present; "
+                "running request as AGG (prefill+decode on this worker)."
+            )
+            request["disaggregated_params"] = {"request_type": "context_and_generation"}
 
         # Setup disaggregated params based on PREFILL/DECODE mode
         (
@@ -1013,6 +1035,7 @@ async def _generate_locally_impl(
         if (
             self.disaggregation_mode == DisaggregationMode.DECODE
             and disaggregated_params is None
+            and not bypass_remote_prefill
         ):
             logging.error("DECODE: disaggregated_params is None but required!")
             logging.error(f"DECODE: Request keys: {list(request.keys())}")
@@ -1171,11 +1194,13 @@ async def _generate_locally_impl(
                 cache_salt=cache_salt,
             )
 
-            # In disagg decode mode, wrap abort() to defer until first token
-            # (KV transfer complete).
+            # In disagg decode mode with remote prefill, wrap abort() to defer
+            # until the first token is received (KV transfer complete).
             abort_guard = (
                 _DeferredAbort(generation_result)
                 if self.disaggregation_mode == DisaggregationMode.DECODE
+                and disaggregated_params is not None
+                and not bypass_remote_prefill
                 else None
             )
 

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

@@ -24,7 +24,10 @@
 from dynamo.trtllm.constants import DisaggregationMode
 from dynamo.trtllm.health_check import TrtllmHealthCheckPayload
 from dynamo.trtllm.llm_engine import TrtllmLLMEngine
-from dynamo.trtllm.request_handlers.handler_base import HandlerBase
+from dynamo.trtllm.request_handlers.handler_base import (
+    BYPASS_REMOTE_PREFILL_ANNOTATION,
+    HandlerBase,
+)
 
 pytestmark = [
     pytest.mark.unit,
@@ -717,6 +720,13 @@ def _make_prefill_handler(self, machine_id: int = 42) -> HandlerBase:
         handler.disaggregation_mode = DisaggregationMode.PREFILL
         return handler
 
+    def _make_decode_handler(self) -> HandlerBase:
+        config = MagicMock()
+        config.shutdown_event = None
+        handler = _ConcreteHandler(config)
+        handler.disaggregation_mode = DisaggregationMode.DECODE
+        return handler
+
     def test_disagg_request_id_populated_in_prefill_mode(self):
         """When mode is PREFILL and no ep_disaggregated_params, disagg_request_id is set."""
         handler = self._make_prefill_handler()
@@ -782,6 +792,39 @@ def test_different_machine_ids_produce_different_id_ranges(self):
         )
         assert params_a.disagg_request_id != params_b.disagg_request_id
 
+    def test_decode_conditional_bypass_uses_request_disagg_params(self):
+        """Conditional-disagg bypass runs full context+generation on decode."""
+        handler = self._make_decode_handler()
+        params, _, _ = handler._setup_disaggregated_params_for_mode(
+            request={
+                "annotations": [BYPASS_REMOTE_PREFILL_ANNOTATION],
+                "disaggregated_params": {"request_type": "context_and_generation"},
+            },
+            ep_disaggregated_params=None,
+        )
+        assert params.request_type == "context_and_generation"
+
+    def test_decode_conditional_bypass_preserves_epd_params(self):
+        """Conditional-disagg bypass keeps EPD params for multimodal handling."""
+        handler = self._make_decode_handler()
+        ep_disaggregated_params = MagicMock()
+
+        (
+            params,
+            returned_ep_params,
+            epd_metadata,
+        ) = handler._setup_disaggregated_params_for_mode(
+            request={
+                "annotations": [BYPASS_REMOTE_PREFILL_ANNOTATION],
+                "disaggregated_params": {"request_type": "context_and_generation"},
+            },
+            ep_disaggregated_params=ep_disaggregated_params,
+        )
+
+        assert params.request_type == "context_and_generation"
+        assert returned_ep_params is ep_disaggregated_params
+        assert epd_metadata == {}
+
 
 class TestHealthCheckPriority:
     """Verify generate_locally forwards the correct priority to generate_async.

components/src/dynamo/vllm/args.py

@@ -401,10 +401,13 @@ def create_kv_events_config(
     dynamo_config: Config, engine_config: AsyncEngineArgs
 ) -> Optional[KVEventsConfig]:
     """Create KVEventsConfig for prefix caching if needed."""
-    if dynamo_config.disaggregation_mode == DisaggregationMode.DECODE:
+    if (
+        dynamo_config.disaggregation_mode == DisaggregationMode.DECODE
+        and not dynamo_config.enable_conditional_disagg
+    ):
         logger.info(
             "Decode worker detected (disaggregation_mode=decode): "
-            "kv_events_config disabled (decode workers don't publish KV events)"
+            "kv_events_config disabled (pass --enable-conditional-disagg to opt in)"
         )
         return None
 

components/src/dynamo/vllm/backend_args.py

@@ -82,6 +82,18 @@ def add_arguments(self, parser) -> None:
             "Mark this as a decode worker which does not publish KV events.",
         )
 
+        add_negatable_bool_argument(
+            g,
+            flag_name="--enable-conditional-disagg",
+            env_var="DYN_VLLM_ENABLE_CONDITIONAL_DISAGG",
+            default=False,
+            help=(
+                "Opt this decode worker into conditional-disagg support. Decode "
+                "workers publish KV events so the Dynamo router can see decode-side "
+                "cache, and annotated bypass requests run as local prefill+decode."
+            ),
+        )
+
         add_negatable_bool_argument(
             g,
             flag_name="--use-vllm-tokenizer",
@@ -418,6 +430,7 @@ class DynamoVllmConfig(ConfigBase):
     ]  # None when not provided; resolved to enum in validate()
     is_prefill_worker: bool
     is_decode_worker: bool
+    enable_conditional_disagg: bool
     use_vllm_tokenizer: bool
 
     # Multimodal

components/src/dynamo/vllm/handlers.py

@@ -101,6 +101,11 @@
 configure_dynamo_logging()
 logger = logging.getLogger(__name__)
 
+# Marker set by the Rust conditional-disagg bypass path. When present on a
+# DECODE-mode worker, the request runs as local prefill+decode instead of
+# expecting KV-transfer metadata from an upstream prefill worker.
+BYPASS_REMOTE_PREFILL_ANNOTATION = "x-bypass-remote-prefill"
+
 _GENERATE_REASONING_SUPPORT_CACHE_ATTR = "_dynamo_generate_reasoning_support"
 _DELTA_REQUEST_OUTPUT_KIND = RequestOutputKind.DELTA
 _RL_INIT_WEIGHTS_TIMEOUT_ENV = "DYN_RL_INIT_WEIGHTS_TIMEOUT_S"
@@ -2998,6 +3003,15 @@ async def _generate_token_mode(self, request, context, request_id):
 
         mode = cast(DisaggregationMode, self.config.disaggregation_mode)
         is_decode_only = mode == DisaggregationMode.DECODE
+        if is_decode_only and BYPASS_REMOTE_PREFILL_ANNOTATION in (
+            request.get("annotations") or []
+        ):
+            logger.debug(
+                "DECODE: conditional-disagg bypass annotation present; "
+                "running request as AGG (prefill+decode on this worker)."
+            )
+            is_decode_only = False
+            mode = DisaggregationMode.AGGREGATED
         has_mm_data = request.get("multi_modal_data") is not None
         mixed_embeds: tuple[torch.Tensor, list[int], list[bool]] | None = None
 
@@ -3214,11 +3228,20 @@ async def _generate_text_mode(self, request, context, request_id):
 
         trace_headers = context.trace_headers()
 
+        is_decode_only = self.config.disaggregation_mode == DisaggregationMode.DECODE
+        if is_decode_only and BYPASS_REMOTE_PREFILL_ANNOTATION in (
+            request.get("annotations") or []
+        ):
+            logger.debug(
+                "DECODE: conditional-disagg bypass annotation present; "
+                "running text-mode request as AGG (prefill+decode on this worker)."
+            )
+            is_decode_only = False
+
         # Mirror _generate_token_mode: in disagg decode mode route aborts through
         # the per-request deferred guard so engine_client.abort() never fires in
         # the unsafe pre-first-token window, and the admin abort_request route can
         # reach this request via self._deferred_aborts.
-        is_decode_only = self.config.disaggregation_mode == DisaggregationMode.DECODE
         async with _deferred_abort_guard(
             self.engine_client,
             request_id,

components/src/dynamo/vllm/main.py

@@ -367,9 +367,16 @@ def setup_kv_event_publisher(
     if not config.engine_args.enable_prefix_caching:
         return None
 
-    # Skip KV event publishing for decode workers
-    if config.disaggregation_mode == DisaggregationMode.DECODE:
-        logger.info("Skipping KV event publisher setup for decode worker")
+    # Decode workers do not publish KV events by default. Conditional disagg
+    # opts in because decode-side cache visibility is needed for bypass routing.
+    if (
+        config.disaggregation_mode == DisaggregationMode.DECODE
+        and not config.enable_conditional_disagg
+    ):
+        logger.info(
+            "Skipping KV event publisher setup for decode worker "
+            "(pass --enable-conditional-disagg to opt in)"
+        )
         return None
 
     if config.engine_args.kv_events_config is None: