Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
0ed5987
fix: unified radix cache tests write markers into the pool, not a copy
alphabetc1 Aug 26, 2026
c12432c
fix: charge the Mamba slot a prefill admission reserved but never deb…
alphabetc1 Aug 26, 2026
b0647bf
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Aug 26, 2026
a2be592
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Aug 28, 2026
22442b2
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Aug 28, 2026
f0c9be3
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Aug 31, 2026
63166bc
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 1, 2026
7eff811
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 1, 2026
46a5b29
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 1, 2026
88d9b39
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 1, 2026
71961ad
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 1, 2026
c0156fe
refactor: trim the comments on the mamba gap reserve and marker writes
alphabetc1 Sep 1, 2026
1e12165
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
aa390d5
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
8c4366e
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
1f3bdca
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
44168ae
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
77c4178
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 2, 2026
b11ce87
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 3, 2026
d1af0ce
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 3, 2026
a015104
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 3, 2026
a64052c
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 3, 2026
66314fd
Merge branch 'main' into fix/mamba-gap-reserve-double-read
alphabetc1 Sep 4, 2026
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
9 changes: 6 additions & 3 deletions python/sglang/srt/managers/schedule_policy.py
Original file line number Diff line number Diff line change
Expand Up @@ -1222,7 +1222,10 @@ def add_one_req(
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
Expand Down Expand Up @@ -1369,7 +1372,7 @@ def add_one_req(
CLIP_MAX_NEW_TOKENS,
),
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
mamba_gap_reserve=mamba_gap_reserve,
Comment thread
alphabetc1 marked this conversation as resolved.
)
self._account_prefill_cache_admission(req, prefix_len)
else:
Expand Down Expand Up @@ -1416,7 +1419,7 @@ def add_one_req(
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)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4982,10 +4982,9 @@ def _get_full_kv_pool(self, allocator):
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)
Expand All @@ -5000,9 +4999,9 @@ def _fill_mamba_state(self, req_to_token_pool, indices, marker):
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)
Expand Down
Loading