Add the KV-canary install API and forward-path wiring (#26809)

python/sglang/srt/kv_canary/api.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
+from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
+from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+def install_canary(
+    *,
+    server_args: "ServerArgs",
+    model_runner: "ModelRunner",
+) -> Optional[CanaryManager]:
+    config = CanaryConfig.from_env(server_args)
+    if config.mode is CanaryMode.NONE:
+        return None
+
+    assert server_args.disable_piecewise_cuda_graph, (
+        "kv-canary: piecewise cuda graph is not supported by the current "
+        "SingleForwardManager design; pass --disable-piecewise-cuda-graph "
+        "when canary is enabled"
+    )
+
+    device = torch.device(model_runner.device)
+    buffer_groups = attach_canary_buffers(
+        pool=model_runner.token_to_kv_pool,
+        config=config,
+        device=device,
+        kv_token_id_vs_position_offset=0,
+    )
+    launch_capacities = CanaryLaunchCapacities.from_args(
+        server_args=model_runner.server_args,
+        req_to_token_pool_size=model_runner.req_to_token_pool.size,
+        max_seq_len_per_req=model_runner.req_to_token_pool.req_to_token.shape[1],
+        pool_slot_count=model_runner.max_total_num_tokens,
+    )
+    swa_window_size = model_runner.sliding_window_size or 0
+    manager = CanaryManager(
+        config=config,
+        buffer_groups=buffer_groups,
+        device=device,
+        req_to_token_pool=model_runner.req_to_token_pool,
+        launch_capacities=launch_capacities,
+        swa_window_size=swa_window_size,
+    )
+
+    _patch_model_forward(model_runner=model_runner, manager=manager)
+
+    # Single-line summary of every knob that controls canary behavior at boot time.
+    # Disaggregation mode is included so PD logs are unambiguous about which side this is.
+    logger.info(
+        "install_canary: disaggregation_mode=%s config=%s "
+        "launch_capacities=%s n_buffer_groups=%d buffer_group_kinds=%s "
+        "swa_window_size=%d",
+        server_args.disaggregation_mode,
+        config,
+        launch_capacities,
+        len(buffer_groups),
+        [g.kind.name for g in buffer_groups],
+        swa_window_size,
+    )
+    return manager
+
+
+def _patch_model_forward(
+    *, model_runner: "ModelRunner", manager: CanaryManager
+) -> None:
+    def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any:
+        forward_batch = _extract_forward_batch(args, kwargs)
+        assert (
+            forward_batch is not None
+        ), "kv-canary: patched model.forward called without a ForwardBatch"
+
+        canary_pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
+        output = original(*args, **kwargs)
+        manager.post_ops_maybe_inside_graph(forward_batch, canary_pre_ops_output)
+        return output
+
+    wrap_method(model_runner.model, "forward", wrapper=_with_canary_bracketing)
+
+
+def _extract_forward_batch(args, kwargs) -> Optional[ForwardBatch]:
+    if "forward_batch" in kwargs and isinstance(kwargs["forward_batch"], ForwardBatch):
+        return kwargs["forward_batch"]
+    for arg in args:
+        if isinstance(arg, ForwardBatch):
+            return arg
+    return None

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1122,7 +1122,11 @@ class CudaGraphRunner:
 
             self.deepep_adapter.capture(is_extend_in_batch=False)
 
-            canary_ctx = contextlib.nullcontext()
+            canary_ctx = (
+                c.with_active_single_forward_manager(0)
+                if (c := self.model_runner.canary_manager) is not None
+                else contextlib.nullcontext()
+            )
             with canary_ctx:
                 for _ in range(2):
                     self.device_module.synchronize()

python/sglang/srt/model_executor/model_runner.py

@@ -105,6 +105,8 @@ from sglang.srt.eplb.expert_location import (
 )
 from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
 from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
+from sglang.srt.kv_canary.api import install_canary
+from sglang.srt.kv_canary.runner.canary_manager import context_tuple
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.attention_registry import (
     ATTENTION_BACKENDS,
@@ -751,6 +753,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # Init memory pool and attention backends
         self.init_memory_pool(pre_model_load_memory)
 
+        # Must be called AFTER init_memory_pool so the pool object exists for
+        # canary to monkey-patch, and BEFORE init_device_graphs so warmup
+        # forwards captured into the graph see the patched pool methods.
+        self.canary_manager = install_canary(
+            server_args=server_args,
+            model_runner=self,
+        )
+
         # Init ngram embedding token table
         self.maybe_init_ngram_embedding()
 
@@ -828,6 +838,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
         self.prealloc_symmetric_memory_pool()
 
+        if self.canary_manager is not None and not self.is_draft_worker:
+            self.canary_manager.mark_init_finished()
+
     def adjust_hybrid_swa_layers_for_pp(self):
         if not self.is_hybrid_swa:
             return
@@ -3204,7 +3217,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
 
-        canary_ctx = contextlib.nullcontext()
+        canary_ctx = (
+            context_tuple(
+                c.with_ops_outside_graph(
+                    single_forward_indices=[0],
+                    maybe_inaccurate_forward_batch=forward_batch,
+                ),
+                c.with_active_single_forward_manager(0),
+            )
+            if not self.is_draft_worker and ((c := self.canary_manager) is not None)
+            else contextlib.nullcontext()
+        )
 
         with (
             canary_ctx,

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -310,6 +310,7 @@ class MockModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -296,6 +296,7 @@ class DSAMockModelRunner(ModelRunner):
         else:
             spec_num_draft_tokens = 0
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -320,6 +320,7 @@ class MockDSV4ModelRunner:
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py

@@ -302,6 +302,7 @@ class DualChunkMockModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -205,6 +205,7 @@ class MockGDNModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -211,6 +211,7 @@ class MockKDAModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -220,6 +220,7 @@ class MockLightningModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -315,6 +315,7 @@ class MockMamba2ModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         # MambaMixer2 asserts the layer_cache is a `SpeculativeState`

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -222,6 +222,7 @@ class MockMLAModelRunner(ModelRunner):
         # does the BF16->FP8 cast on the way in.
         self.kv_cache_dtype = torch.float8_e4m3fn if fp8_kv_cache else dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1