[bug1] keep full kv when swa skips leaf data (#29351)

Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Feng Yao
2026-07-01 00:51:18 +08:00
committed by GitHub
co-authored by ispobock
parent c70cc96905
commit 1c2523b719
5 changed files with 41 additions and 30 deletions
@@ -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: