[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:
co-authored by
Mohammad Angkad
parent
020703923d
commit
df0dc44931
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user