From e4adf63275005f4b3e6f8d74ae4ff203a75d05df Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:43:17 +0800 Subject: [PATCH] [Fix] Vacuous marker writes in the cache tests, and an undebited Mamba admission slot (#36415) Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/srt/managers/schedule_policy.py | 9 ++++++--- .../mem_cache/test_unified_radix_cache_unittest.py | 11 +++++------ 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d08e86093..ade73e45f 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -1222,7 +1222,10 @@ class PrefillAdder: total_tokens = cand_extend_input_len + max_new + self.page_size # 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 += 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 real_input_tokens = cand_extend_input_len - req.host_hit_length @@ -1369,7 +1372,7 @@ class PrefillAdder: CLIP_MAX_NEW_TOKENS, ), 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) else: @@ -1416,7 +1419,7 @@ class PrefillAdder: trunc_len, 0, 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) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 5bc947289..35d92ded4 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -4982,10 +4982,9 @@ class UnifiedRadixCacheSuite: def _fill_full_kv(self, allocator, indices, marker): kv_pool = self._get_full_kv_pool(allocator) layer_id = kv_pool.start_layer - k_buf = kv_pool.get_key_buffer(layer_id) - v_buf = kv_pool.get_value_buffer(layer_id) - k_buf[indices].fill_(marker) - v_buf[indices].fill_(marker + 1) + # `buf[indices]` is an advanced-index copy — assign, never `fill_()`. + kv_pool.get_key_buffer(layer_id)[indices] = marker + kv_pool.get_value_buffer(layer_id)[indices] = marker + 1 def _snapshot_full_kv(self, allocator, indices): kv_pool = self._get_full_kv_pool(allocator) @@ -5000,9 +4999,9 @@ class UnifiedRadixCacheSuite: return mamba_indices = indices.reshape(-1) 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): - conv_buf[:, mamba_indices].fill_(marker + offset) + conv_buf[:, mamba_indices] = marker + offset def _snapshot_mamba_state(self, req_to_token_pool, indices): mamba_indices = indices.reshape(-1)