Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions python/sglang/srt/mem_cache/allocator/hisparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,14 @@ def full_available_size(self):
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)

Expand Down
19 changes: 19 additions & 0 deletions test/registered/unit/mem_cache/test_hisparse_allocator.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,25 @@ def test_forwards_swa_tail_allocation_to_logical_allocator(self):
self.assertEqual(kwargs["extend_num_tokens"], 512)
self.assertEqual(kwargs["swa_tail_len"], 128)

def test_forwards_prealloc_reclaim_to_logical_allocator(self):
"""PD decode preallocation must not crash on the HiSparse composite."""
allocator = object.__new__(DeepSeekV4HiSparseTokenToKVPoolAllocator)
logical_allocator = MagicMock(spec=["reclaim_for_prealloc"])
allocator.logical_attn_allocator = logical_allocator
logical_allocator.reclaim_for_prealloc.return_value = None

tree_cache = object()
self.assertIsNone(allocator.reclaim_for_prealloc(tree_cache, 512, 256))
logical_allocator.reclaim_for_prealloc.assert_called_once_with(
tree_cache, 512, 256
)

logical_allocator.reclaim_for_prealloc.return_value = "SWA eviction short"
self.assertEqual(
allocator.reclaim_for_prealloc(tree_cache, 512, 256),
"SWA eviction short",
)

def test_hisparse_budget_uses_full_logical_capacity_for_swa_tail(self):
from sglang.srt.disaggregation.decode import DecodePreallocQueue

Expand Down
Loading