[bug1] keep full kv when swa skips leaf data (#29351)
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
@@ -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 |
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user