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

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()