[Fix] Forward SWA prealloc reclaim through the DSV4 HiSparse allocator (#40354)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
This commit is contained in:
Mohammad Miadh Angkad
2026-09-19 20:27:51 -07:00
committed by GitHub
co-authored by Mohammad Angkad
parent 020703923d
commit df0dc44931
2 changed files with 27 additions and 0 deletions
@@ -376,6 +376,14 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def swa_available_size(self):
return self.logical_attn_allocator.swa_available_size()
def reclaim_for_prealloc(
self, tree_cache, full_tokens: int, swa_tokens: int
) -> str | None:
# C4 needs no reclaim here: full_available_size prices it into the budget.
return self.logical_attn_allocator.reclaim_for_prealloc(
tree_cache, full_tokens, swa_tokens
)
def free_swa(self, free_indices: torch.Tensor):
self.logical_attn_allocator.free_swa(free_indices)