minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Fix] Vacuous marker writes in the cache tests, and an undebited Mamba admission slot (#36415)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Shuwen Wang
2026-09-04 19:43:17 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 3b678549c3
commit e4adf63275
2 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_policy.py
+6 -3
View File
@@ -1222,7 +1222,10 @@ class PrefillAdder:
total_tokens = cand_extend_input_len + max_new + self.page_size total_tokens = cand_extend_input_len + max_new + self.page_size
# Shared Mamba pool: fold the new mamba state's shared-gap cost into # Shared Mamba pool: fold the new mamba state's shared-gap cost into
# `total_tokens` so both `rem_total_tokens` gates reflect the joint budget. # `total_tokens` so both `rem_total_tokens` gates reflect the joint budget.
total_tokens += self._mamba_gap_budget_for_req(req) # Read before `init_load_back` binds `req.mamba_pool_idx` — after that
# this returns 0, so the debit sites below reuse the value.
mamba_gap_reserve = self._mamba_gap_budget_for_req(req)
total_tokens += mamba_gap_reserve
# adjusting the input_tokens based on host_hit_length and page_size # adjusting the input_tokens based on host_hit_length and page_size
real_input_tokens = cand_extend_input_len - req.host_hit_length real_input_tokens = cand_extend_input_len - req.host_hit_length
@@ -1369,7 +1372,7 @@ class PrefillAdder:
CLIP_MAX_NEW_TOKENS, CLIP_MAX_NEW_TOKENS,
), ),
req.retracted_stain, req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req), mamba_gap_reserve=mamba_gap_reserve,
) )
self._account_prefill_cache_admission(req, prefix_len) self._account_prefill_cache_admission(req, prefix_len)
else: else:
@@ -1416,7 +1419,7 @@ class PrefillAdder:
trunc_len, trunc_len,
0, 0,
req.retracted_stain, req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req), mamba_gap_reserve=mamba_gap_reserve,
) )
self._account_prefill_cache_admission(req, prefix_len) self._account_prefill_cache_admission(req, prefix_len)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+5 -6
View File
@@ -4982,10 +4982,9 @@ class UnifiedRadixCacheSuite:
def _fill_full_kv(self, allocator, indices, marker): def _fill_full_kv(self, allocator, indices, marker):
kv_pool = self._get_full_kv_pool(allocator) kv_pool = self._get_full_kv_pool(allocator)
layer_id = kv_pool.start_layer layer_id = kv_pool.start_layer
k_buf = kv_pool.get_key_buffer(layer_id) # `buf[indices]` is an advanced-index copy — assign, never `fill_()`.
v_buf = kv_pool.get_value_buffer(layer_id) kv_pool.get_key_buffer(layer_id)[indices] = marker
k_buf[indices].fill_(marker) kv_pool.get_value_buffer(layer_id)[indices] = marker + 1
v_buf[indices].fill_(marker + 1)
def _snapshot_full_kv(self, allocator, indices): def _snapshot_full_kv(self, allocator, indices):
kv_pool = self._get_full_kv_pool(allocator) kv_pool = self._get_full_kv_pool(allocator)
@@ -5000,9 +4999,9 @@ class UnifiedRadixCacheSuite:
return return
mamba_indices = indices.reshape(-1) mamba_indices = indices.reshape(-1)
mamba_cache = req_to_token_pool.mamba_pool.mamba_cache mamba_cache = req_to_token_pool.mamba_pool.mamba_cache
mamba_cache.temporal[:, mamba_indices].fill_(marker) mamba_cache.temporal[:, mamba_indices] = marker
for offset, conv_buf in enumerate(mamba_cache.conv, start=1): for offset, conv_buf in enumerate(mamba_cache.conv, start=1):
conv_buf[:, mamba_indices].fill_(marker + offset) conv_buf[:, mamba_indices] = marker + offset
def _snapshot_mamba_state(self, req_to_token_pool, indices): def _snapshot_mamba_state(self, req_to_token_pool, indices):
mamba_indices = indices.reshape(-1) mamba_indices = indices.reshape(-1)