From 7998fecfd11411832f435301b07ce4e457618898 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 10 Jul 2026 10:58:55 -0700 Subject: [PATCH] [kv canary] Support UnifiedRadixCache in kv-canary and bracket nested model.forward (#30574) --- python/sglang/srt/kv_canary/api.py | 23 ++++-- .../srt/kv_canary/radix_cache_walker.py | 77 +++++++++++++++---- .../srt/kv_canary/runner/canary_manager.py | 17 ++++ .../srt/model_executor/runner/base_runner.py | 12 ++- .../kv_canary/test_self_unit_radix_walker.py | 77 +++++++++++++++++++ 5 files changed, 184 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index dee30c590..fdae24039 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/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) diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py index 516220dea..f70b15869 100644 --- a/python/sglang/srt/kv_canary/radix_cache_walker.py +++ b/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 diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index b6d3ca0f2..c8b73d0a4 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 56420661a..603db6e34 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/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) diff --git a/test/registered/kv_canary/test_self_unit_radix_walker.py b/test/registered/kv_canary/test_self_unit_radix_walker.py index 4caef8b8c..395b7a627 100644 --- a/test/registered/kv_canary/test_self_unit_radix_walker.py +++ b/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()