diff --git a/python/sglang/srt/mem_cache/unified_cache_components/README.md b/python/sglang/srt/mem_cache/unified_cache_components/README.md index 6716f9774..ab8847031 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache_components/README.md @@ -125,9 +125,8 @@ Insert a key-value pair into the tree. - If partially within window: **splits node** at boundary, recovers SWA on the window portion (returns `start_idx`) - If entirely outside window: returns `prefix_len` (no consumption) - Mamba: returns `prefix_len` (no consumption, default behavior) -4. Before creating a new leaf, checks `should_skip_leaf_creation()` per component — any veto aborts leaf creation and frees remaining value -5. Creates leaf via `_add_new_node` (clones value tensor, updates Full leaf-set tracking) -6. Calls `commit_insert_component_data()` per component on the final target node (SWA may trigger a secondary split for window boundary; Mamba sets mamba pool indices and inserts into Mamba LRU) +4. Creates leaf via `_add_new_node` (clones value tensor, updates Full leaf-set tracking). A leaf survives on its Full value alone, so it is materialized even when an auxiliary component holds only a tombstone for the span (e.g. the whole leaf is outside the SWA window) +5. Calls `commit_insert_component_data()` per component on the final target node (SWA may trigger a secondary split for window boundary; Mamba sets mamba pool indices and inserts into Mamba LRU) --- @@ -262,7 +261,6 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d | Hook | Purpose | Called By | Default | |------|---------|-----------|----------| | `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_helper` | returns `prefix_len` | -| `should_skip_leaf_creation()` | Veto leaf creation when the entire new leaf would be a tombstone for this component. SWA: vetoes if `swa_evicted_seqlen ≥ total_prefix_len + key_len`. | `_insert_helper` | `False` | | `recover_after_unevict()` | Rebuild auxiliary component data after `_unevict_node_on_insert()` restores a Full device value from fresh KV indices. SWA uses this to rebuild in-window SWA data. | `_insert_helper` | no-op | | `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_helper` | no-op | diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index d6ef24cd4..ff161d68d 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -207,11 +207,6 @@ class SWAComponent(TreeComponent): # Branch 3: entire value_slice is outside SWA window — not consumed return prefix_len - def should_skip_leaf_creation( - self, total_prefix_len: int, key_len: int, params: InsertParams - ) -> bool: - return params.swa_evicted_seqlen >= total_prefix_len + key_len - def recover_after_unevict( self, node: UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index f83e37f2f..425eea226 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -190,13 +190,6 @@ class TreeComponent(ABC): portion: value_slice[dup_start:consumed_from].""" return prefix_len - def should_skip_leaf_creation( - self, total_prefix_len: int, key_len: int, params: InsertParams - ) -> bool: - """Return True to veto leaf creation when the entire new leaf would - be a tombstone for this component.""" - return False - def recover_after_unevict( self, node: UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index fa26bbee5..5fd3c95ef 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1159,21 +1159,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): child_key = key.child_key(self.page_size) is_new_leaf = False - # Create new leaf for remaining suffix + # Create new leaf for remaining suffix. A leaf survives on its Full + # value alone; auxiliary components (SWA, Mamba) may legitimately hold + # only a tombstone for this span (e.g. the whole leaf is outside the SWA + # window). Materialize it anyway so the Full KV stays cacheable. if len(key): - if any( - comp.should_skip_leaf_creation( - total_prefix_len=total_prefix_length, - key_len=len(key), - params=params, - ) - for comp in self._components_tuple - ): - # TODO: When leaf creation is skipped, We should release all component - # resources here or propagate a flag so that - # cleanup_after_caching_req can free them properly. - self.token_to_kv_pool_allocator.free(value) - return InsertResult(prefix_len=total_prefix_length) target_node = self._add_new_node(node, key, value, priority=priority) is_new_leaf = True else: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index eee7f1d21..47638f54e 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1287,6 +1287,41 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(m.device_indices), len(seq)) tree.sanity_check() + def test_swa_insert_keeps_full_leaf_when_entire_span_is_outside_window(self): + # A leaf survives on its Full value alone: even when the whole span is + # past the SWA window (swa_evicted_seqlen >= total), the Full leaf must + # be materialized (and the Full KV kept) so the prefix stays cacheable. + # Runs across all SWA configs, including page_size > sliding_window_size + # (the dsv4-style edge case). + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("requires SWA without Mamba") + tree, allocator, _ = build_fixture(self.cfg) + + tokens = self._make_seq(1, 2) + value = self._alloc(allocator, len(tokens)) + if value is None: + self.skipTest("insufficient pool for this config") + full_available_before = allocator.full_attn_allocator.available_size() + + tree.insert( + InsertParams( + key=RadixKey(array("q", tokens)), + value=value, + prev_prefix_len=0, + swa_evicted_seqlen=len(tokens), + ) + ) + + self.assertEqual( + allocator.full_attn_allocator.available_size(), full_available_before + ) + node = next(iter(tree.root_node.children.values())) + self.assertTrue( + torch.equal(node.component_data[ComponentType.FULL].value, value) + ) + self.assertIsNone(node.component_data[ComponentType.SWA].value) + tree.sanity_check() + def test_swa_evict_cascades(self): """Evict SWA tokens via swa_num_tokens — cascades to lower-priority components.""" if not self.cfg.has_swa: