[kv canary] Support UnifiedRadixCache in kv-canary and bracket nested model.forward (#30574)

python/sglang/srt/kv_canary/api.py

@@ -102,15 +102,22 @@ def install_canary(
 
 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"
+        with manager.model_forward_bracket_scope() as should_bracket:
+            if not should_bracket:
+                # Nested model.forward calls share the active SingleForwardManager.
+                # Only the outermost call may run kv-canary pre/post ops; otherwise
+                # the phase checker sees a second pre-op before the first post-op.
+                return original(*args, **kwargs)
 
-        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
+            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)
 

python/sglang/srt/kv_canary/radix_cache_walker.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+from sglang.srt.mem_cache.unified_cache_components import (
+    BASE_COMPONENT_TYPE,
+    ComponentType,
+)
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
     from sglang.srt.mem_cache.radix_cache import TreeNode
+    from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode
 
 
 @dataclass(frozen=True, slots=True, kw_only=True)
@@ -34,7 +40,11 @@ def walk_radix_cache_for_canary(
     req. ``swa_resident_only=True`` skips SWA-tombstoned nodes (slots evicted from the SWA
     window)."""
     cache_type = type(radix_cache)
-    if cache_type is not RadixCache and cache_type is not SWARadixCache:
+    if (
+        cache_type is not RadixCache
+        and cache_type is not SWARadixCache
+        and cache_type is not UnifiedRadixCache
+    ):
         raise NotImplementedError(
             f"walk_radix_cache_for_canary does not support {cache_type.__name__}"
         )
@@ -68,7 +78,7 @@ def walk_radix_cache_for_canary(
 
 def _walk_radix_subtree(
     *,
-    node: TreeNode,
+    node: TreeNode | UnifiedTreeNode,
     radix_cache: BasePrefixCache,
     depth: int,
     parent_last_slot: int,
@@ -79,10 +89,7 @@ def _walk_radix_subtree(
     unlocked_only: bool,
     swa_resident_only: bool,
 ) -> None:
-    if isinstance(node.value, torch.Tensor):
-        node_slots = [int(s) for s in node.value.tolist()]
-    else:
-        node_slots = []
+    node_slots = _node_slots_for_canary(node=node, radix_cache=radix_cache)
 
     if unlocked_only:
         emit_slots = not is_root and _node_is_unlocked_for_canary(
@@ -105,7 +112,12 @@ def _walk_radix_subtree(
             prev_slot_buf.append(prev)
         chain_last_slot = slot
 
-    child_depth = depth + len(node_slots)
+    child_depth = depth + _node_len_for_canary(
+        node=node,
+        radix_cache=radix_cache,
+        node_slots=node_slots,
+        is_root=is_root,
+    )
     for child in node.children.values():
         _walk_radix_subtree(
             node=child,
@@ -121,9 +133,40 @@ def _walk_radix_subtree(
         )
 
 
+def _node_slots_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+) -> list[int]:
+    value: Any
+    if type(radix_cache) is UnifiedRadixCache:
+        value = node.component_data[BASE_COMPONENT_TYPE].value
+    else:
+        value = node.value
+
+    if isinstance(value, torch.Tensor):
+        return [int(s) for s in value.tolist()]
+    return []
+
+
+def _node_len_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+    node_slots: list[int],
+    is_root: bool,
+) -> int:
+    if type(radix_cache) is not UnifiedRadixCache:
+        return len(node_slots)
+
+    if is_root or node.key is None:
+        return len(node_slots)
+    return len(node.key)
+
+
 def _node_is_unlocked_for_canary(
     *,
-    node: TreeNode,
+    node: TreeNode | UnifiedTreeNode,
     radix_cache: BasePrefixCache,
 ) -> bool:
     if type(radix_cache) is RadixCache:
@@ -132,6 +175,9 @@ def _node_is_unlocked_for_canary(
     if type(radix_cache) is SWARadixCache:
         return node.full_lock_ref == 0
 
+    if type(radix_cache) is UnifiedRadixCache:
+        return node.component_data[BASE_COMPONENT_TYPE].lock_ref == 0
+
     raise NotImplementedError(
         f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}"
     )
@@ -139,10 +185,15 @@ def _node_is_unlocked_for_canary(
 
 def _node_is_swa_resident_for_canary(
     *,
-    node: TreeNode,
+    node: TreeNode | UnifiedTreeNode,
     radix_cache: BasePrefixCache,
 ) -> bool:
-    if type(radix_cache) is not SWARadixCache:
-        return True
+    if type(radix_cache) is SWARadixCache:
+        return not node.swa_tombstone
 
-    return not node.swa_tombstone
+    if type(radix_cache) is UnifiedRadixCache:
+        if not radix_cache.supports_swa():
+            return True
+        return node.component_data[ComponentType.SWA].value is not None
+
+    return True

python/sglang/srt/kv_canary/runner/canary_manager.py

@@ -61,6 +61,7 @@ class CanaryManager:
         self._swa_allocator: Optional[SWATokenToKVPoolAllocator] = swa_allocator
         self._outer_step_counter: int = 0
         self._active_single_forward_manager_index: Optional[int] = None
+        self._model_forward_bracket_depth: int = 0
 
         self._buffer_groups: tuple[CanaryBufferGroup, ...] = tuple(buffer_groups)
 
@@ -182,6 +183,22 @@ class CanaryManager:
             )
             self._active_single_forward_manager_index = None
 
+    @contextlib.contextmanager
+    def model_forward_bracket_scope(self) -> Iterator[bool]:
+        """Return whether this is the outermost patched ``model.forward`` call.
+
+        Some model implementations enter another patched forward from inside the
+        top-level forward (for example, a vision-language model calling its inner
+        language model). Kv-canary owns one pre/post bracket per active
+        SingleForwardManager; nested brackets would run a second pre-op while the
+        phase checker is already in the first bracket.
+        """
+        self._model_forward_bracket_depth += 1
+        try:
+            yield self._model_forward_bracket_depth == 1
+        finally:
+            self._model_forward_bracket_depth -= 1
+
     def pre_ops_maybe_inside_graph(
         self, forward_batch: ForwardBatch
     ) -> _PreOpsMaybeInsideGraphOutput:

python/sglang/srt/model_executor/runner/base_runner.py

@@ -264,9 +264,19 @@ class BaseRunner(ABC):
         Supplied by warmup() (the decode runner's captured buffers when a graph
         runner exists; a freshly-allocated dummy set in the eager path).
         """
+        mr = self.model_runner
+        canary_run_ctx = (
+            c.with_active_single_forward_manager(0)
+            if (c := mr.canary_manager) is not None
+            else empty_context()
+        )
 
         def forward_fn():
-            self._dummy_run(batch_size=batch_size, buffers=buffers)
+            self._dummy_run(
+                batch_size=batch_size,
+                buffers=buffers,
+                run_ctx=canary_run_ctx,
+            )
 
         run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)