Separate swa and local attention chunk cache eviction (#15820)

This commit is contained in:
Ke Bao
2025-12-26 09:34:22 +08:00
committed by GitHub
parent 2f66b0671b
commit 7b7e357f61
8 changed files with 50 additions and 29 deletions
+2 -6
View File
@@ -342,9 +342,7 @@ def alloc_for_extend(
# free out-of-window swa tokens
if isinstance(batch.tree_cache, SWAChunkCache):
for req, pre_len in zip(batch.reqs, batch.prefix_lens):
batch.tree_cache.evict_swa(
req, pre_len, batch.model_config.attention_chunk_size
)
batch.tree_cache.evict_swa(req, pre_len)
bs = len(batch.reqs)
prefix_tensors = [r.prefix_indices for r in batch.reqs]
@@ -437,9 +435,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
"""
if isinstance(batch.tree_cache, SWAChunkCache):
for req in batch.reqs:
batch.tree_cache.evict_swa(
req, req.seqlen - 1, batch.model_config.attention_chunk_size
)
batch.tree_cache.evict_swa(req, req.seqlen - 1)
bs = batch.seq_lens.shape[0]