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

python/sglang/srt/kv_canary/api.py

@@ -102,6 +102,13 @@ def install_canary(
 
 def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) -> None:
     def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any:
+        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)
+
             forward_batch = _extract_forward_batch(args, kwargs)
             assert (
                 forward_batch is not None

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 = _node_slots_for_canary(node=node, radix_cache=radix_cache)
-        node_slots = [int(s) for s in node.value.tolist()]
-    else:
-        node_slots = []
 
     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:
+    if type(radix_cache) is SWARadixCache:
-        return True
-
         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)
 

test/registered/kv_canary/test_self_unit_radix_walker.py

@@ -6,6 +6,11 @@ import torch
 
 from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode
+from sglang.srt.mem_cache.unified_cache_components import (
+    BASE_COMPONENT_TYPE,
+    ComponentType,
+)
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache, UnifiedTreeNode
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_radix_cache
 from sglang.test.test_utils import CustomTestCase
@@ -124,6 +129,78 @@ class TestSelfUnitRadixWalker(CustomTestCase):
         )
         self.assertEqual(result.slot_indices.tolist(), [3, 4])
 
+    def _make_unified_cache(
+        self, tree_components: tuple[ComponentType, ...]
+    ) -> UnifiedRadixCache:
+        cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
+        cache.tree_components = tree_components
+        cache.components = {ct: None for ct in tree_components}
+        root = UnifiedTreeNode(tree_components)
+        root.component_data[BASE_COMPONENT_TYPE].value = torch.tensor(
+            [], dtype=torch.int32, device=self.device
+        )
+        cache.root_node = root
+        return cache
+
+    def _add_unified_child(
+        self,
+        cache: UnifiedRadixCache,
+        slots: list[int],
+        *,
+        lock_ref: int = 0,
+        swa_value: list[int] | None = None,
+    ) -> UnifiedTreeNode:
+        child = UnifiedTreeNode(cache.tree_components)
+        child.parent = cache.root_node
+        base = child.component_data[BASE_COMPONENT_TYPE]
+        base.value = torch.tensor(slots, dtype=torch.int32, device=self.device)
+        base.lock_ref = lock_ref
+        if swa_value is not None:
+            child.component_data[ComponentType.SWA].value = torch.tensor(
+                swa_value, dtype=torch.int32, device=self.device
+            )
+        cache.root_node.children[child.id] = child
+        return child
+
+    def test_unified_walk_emits_full_component_slots(self):
+        """Verify unified radix walking emits the base (full) component slots."""
+        cache = self._make_unified_cache((ComponentType.FULL,))
+        self._add_unified_child(cache, [10, 20, 30])
+        result = walk_radix_cache_for_canary(radix_cache=cache)
+        self.assertEqual(result.slot_indices.tolist(), [10, 20, 30])
+        self.assertEqual(result.positions.tolist(), [0, 1, 2])
+        self.assertEqual(result.prev_slot_indices.tolist(), [-1, 10, 20])
+
+    def test_unified_walk_unlocked_only_uses_full_lock_ref(self):
+        """Verify unified radix walking honors the base component lock reference."""
+        cache = self._make_unified_cache((ComponentType.FULL,))
+        self._add_unified_child(cache, [1, 2], lock_ref=1)
+        self._add_unified_child(cache, [3, 4])
+        result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True)
+        self.assertEqual(result.slot_indices.tolist(), [3, 4])
+
+    def test_unified_swa_resident_only_skips_evicted_swa_nodes(self):
+        """Verify unified radix walking skips nodes whose SWA storage was evicted."""
+        cache = self._make_unified_cache((ComponentType.FULL, ComponentType.SWA))
+        self._add_unified_child(cache, [1, 2], swa_value=None)
+        self._add_unified_child(cache, [3, 4], swa_value=[3, 4])
+        result = walk_radix_cache_for_canary(
+            radix_cache=cache,
+            swa_resident_only=True,
+        )
+        self.assertEqual(result.slot_indices.tolist(), [3, 4])
+
+    def test_unified_swa_resident_only_noop_without_swa_component(self):
+        """Verify swa_resident_only is a no-op when SWA is not enabled."""
+        cache = self._make_unified_cache((ComponentType.FULL,))
+        self._add_unified_child(cache, [1, 2])
+        self._add_unified_child(cache, [3, 4])
+        result = walk_radix_cache_for_canary(
+            radix_cache=cache,
+            swa_resident_only=True,
+        )
+        self.assertEqual(result.slot_indices.tolist(), [1, 2, 3, 4])
+
 
 if __name__ == "__main__":
     unittest.main()