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

[Mamba] Support speculative decoding with extra_buffer_lazy (#30437)

Browse Source
This commit is contained in:
YAMY
2026-07-21 19:38:03 +08:00
committed by GitHub
parent 1a53f231d6
commit 7c9257529f
5 changed files with 253 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+79 -13
View File
@@ -85,7 +85,6 @@ from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
EvictParams,
MatchPrefixParams,
zero_match_result,
)
@@ -1684,21 +1683,42 @@ class _MambaRadixCacheV2TrackEntry(NamedTuple):
track_seqlen: int
def set_mamba_track_indices_from_reqs(batch):
"""Build mamba_track_indices from req objects (authoritative source)."""
def mamba_lazy_spec_in_window(
req, mamba_track_interval: int, max_draft_tokens: int
) -> bool:
"""Whether a track-interval crossing is reachable by an in-flight verify.
kv_committed_len lags device seq_lens by up to one verify under overlap;
the 2x window absorbs it.
"""
seq_len = req.kv_committed_len
window = 2 * max_draft_tokens
return seq_len // mamba_track_interval != (seq_len + window) // mamba_track_interval
def set_mamba_track_indices_from_reqs(
batch, track_positions: Optional[List[int]] = None
):
"""Build mamba_track_indices from req objects (authoritative source).
track_positions: optional per-req ping-pong position override (the lazy
spec track plan, see mamba_lazy_spec_prepare).
"""
req_to_token_pool = batch.req_to_token_pool
all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[
batch.req_pool_indices
] # (bs, ping_pong_size), int64, on device
# Guard: mamba_next_track_idx may be None for requests that haven't
# gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
# Default to 0 (first ping-pong slot) to avoid TypeError.
if track_positions is None:
# Guard: mamba_next_track_idx may be None for requests that haven't
# gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
# Default to 0 (first ping-pong slot) to avoid TypeError.
track_positions = [
req.mamba_next_track_idx if req.mamba_next_track_idx is not None else 0
for req in batch.reqs
]
idx = (
torch.tensor(
[
req.mamba_next_track_idx if req.mamba_next_track_idx is not None else 0
for req in batch.reqs
],
track_positions,
dtype=torch.int64,
pin_memory=True,
)
@@ -1892,6 +1912,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_indices: torch.Tensor = None # shape: [b], int64
mamba_track_mask: torch.Tensor = None # shape: [b], bool
mamba_track_seqlens: torch.Tensor = None # shape: [b], int64
# Lazy + spec: this iteration's per-req scatter positions
# (see mamba_lazy_spec_prepare).
mamba_lazy_spec_track_positions_cpu: Optional[List[int]] = None # shape: [b]
# Deferred mamba init ops: COW pairs and clear indices (performed on forward stream)
mamba_cow_src_indices: torch.Tensor = None
mamba_cow_dst_indices: torch.Tensor = None
@@ -2752,14 +2775,54 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if envs.SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL.get():
new_slot = None
else:
# No evict-retry: a transient slot is not worth evicting a
# cached checkpoint for; on failure tracking degrades in place.
new_slot = pool.mamba_allocator.alloc(1)
if new_slot is None:
self.tree_cache.evict(EvictParams(num_tokens=0, mamba_num=1))
new_slot = pool.mamba_allocator.alloc(1)
if new_slot is not None:
pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0])
req.mamba_next_track_idx = other_idx
def mamba_lazy_spec_prepare(self, mamba_track_interval: int, max_draft_tokens: int):
"""Lazy-mode spec counterpart of mamba_lazy_prealloc_at_boundary.
A crossing is only *possible* at prepare time (accept length is
unknown), so ensure the pending slot exists for reqs whose next
boundary is reachable, WITHOUT swapping mamba_next_track_idx; the
mask-gated commit writes it only on a real crossing, and
_mamba_lazy_spec_confirm_crossing promotes it afterwards. The per-req
scatter position is recorded on the batch and rides the result-queue
copy (forward isolation restores batch fields).
"""
pool = self.req_to_token_pool
track_positions: List[int] = []
for req in self.reqs:
buf = req.mamba_ping_pong_track_buffer
assert buf is not None
if not mamba_lazy_spec_in_window(
req, mamba_track_interval, max_draft_tokens
):
# No crossing reachable: the scatter mask stays -1, the
# position is never written.
track_positions.append(req.mamba_next_track_idx)
continue
other_idx = 1 - req.mamba_next_track_idx
has_pending = buf[other_idx].item() != -1
if not has_pending:
if envs.SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL.get():
new_slot = None
else:
# No evict-retry: a transient slot is not worth
# evicting a cached checkpoint for.
new_slot = pool.mamba_allocator.alloc(1)
if new_slot is not None:
pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0])
has_pending = True
# On failure the verify scatters in place into the keep slot.
track_positions.append(
other_idx if has_pending else req.mamba_next_track_idx
)
self.mamba_lazy_spec_track_positions_cpu = track_positions
def cumulate_penalty_output_tokens(self):
# Under overlap batch.input_ids is just a placeholder here -- the
# real token is relayed via future_map and resolved at forward
@@ -2910,6 +2973,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.mamba_track_indices = None
self.mamba_track_mask = None
self.mamba_track_seqlens = None
self.mamba_lazy_spec_track_positions_cpu = None
self.mamba_cow_src_indices = None
self.mamba_cow_dst_indices = None
self.mamba_clear_indices = None
@@ -2968,6 +3032,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.mamba_track_indices = None
self.mamba_track_mask = None
self.mamba_track_seqlens = None
self.mamba_lazy_spec_track_positions_cpu = None
if self.return_logprob and other.return_logprob:
self.top_logprobs_nums = self.top_logprobs_nums + other.top_logprobs_nums
self.token_ids_logprobs = self.token_ids_logprobs + other.token_ids_logprobs
@@ -3022,6 +3087,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_indices=self.mamba_track_indices,
mamba_track_mask=self.mamba_track_mask,
mamba_track_seqlens=self.mamba_track_seqlens,
mamba_lazy_spec_track_positions_cpu=self.mamba_lazy_spec_track_positions_cpu,
dp_cooperation_info=self.dp_cooperation_info,
prefill_stats=self.prefill_stats,
fpm_start_time=self.fpm_start_time,
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+63
View File
@@ -21,6 +21,7 @@ from sglang.srt.managers.schedule_batch import (
FINISH_MATCHED_TOKEN,
Req,
ScheduleBatch,
mamba_lazy_spec_in_window,
)
from sglang.srt.mem_cache.common import (
maybe_cache_unfinished_req,
@@ -1012,6 +1013,11 @@ class SchedulerBatchResultProcessor:
req, batch, result, i
)
if lazy and not batch.spec_algorithm.is_none():
# For spec, at_boundary means this step actually crossed an interval.
self._mamba_lazy_spec_update(req, batch, i, at_boundary, track_seqlen)
return
if not at_boundary:
return
@@ -1025,6 +1031,63 @@ class SchedulerBatchResultProcessor:
)
)
def _mamba_lazy_spec_update(
self, req: Req, batch: ScheduleBatch, i: int, crossed: bool, track_seqlen: int
) -> None:
"""Lazy + spec post-processing.
Running req with a confirmed crossing: the commit scattered the
crossing state into the planned position; promote it to keep unless
the plan already pointed at the keep slot (in-place fallback, or
promoted by an earlier overlapped confirmation).
Finishing req: donate the keep slot only if no scatter wrote it
(this step) or may still write it (the in-flight next step).
"""
positions = batch.mamba_lazy_spec_track_positions_cpu
planned_pos = (
positions[i]
if positions is not None and i < len(positions)
else None # filtered/merged snapshot without a plan: be conservative
)
if req.finished():
# Skip the donation if a scatter wrote or may still write the keep slot.
keep_written_by_this_step = (
crossed and planned_pos == req.mamba_next_track_idx
)
server_args = get_server_args()
other_idx = 1 - req.mamba_next_track_idx
# Recompute the in-flight verify's plan (kv_committed_len is
# frozen since its prepare, so the recompute is exact).
keep_may_be_written_in_flight = req.mamba_ping_pong_track_buffer[
other_idx
].item() == -1 and mamba_lazy_spec_in_window(
req,
server_args.mamba_track_interval,
server_args.max_speculative_num_draft_tokens,
)
if (
planned_pos is None
or keep_written_by_this_step
or keep_may_be_written_in_flight
):
req.mamba_lazy_is_insert = False
return
if not crossed or planned_pos is None:
return
if planned_pos != req.mamba_next_track_idx:
# Promote pending -> keep: free the old checkpoint, repoint.
pool = batch.req_to_token_pool
keep_idx = req.mamba_next_track_idx
keep_val = req.mamba_ping_pong_track_buffer[keep_idx]
pool.mamba_allocator.free(keep_val.unsqueeze(0))
pool.set_mamba_ping_pong_slot(req, keep_idx, -1)
req.mamba_next_track_idx = planned_pos
# else: in-place fallback, or promoted by an earlier confirmation —
# keep holds the track_seqlen state either way.
req.mamba_last_track_seqlen = track_seqlen
def _mamba_check_track_boundary(self, req, batch, result, i):
"""Check if this decode step crosses a mamba track interval boundary.
python/sglang/srt/server_args.py
+13 -3
View File
@@ -4816,10 +4816,20 @@ class ServerArgs:
assert (
is_cuda() or is_musa() or is_npu() or is_hip()
), "extra_buffer needs CUDA/MUSA/NPU/ROCm (FLA)."
if view.mamba_radix_cache_strategy == "extra_buffer_lazy":
# The PD-disagg decode pool is not wired for lazy slots.
assert view.disaggregation_mode == "null", (
"extra_buffer_lazy unsupported under PD disaggregation; use "
"--mamba-radix-cache-strategy extra_buffer."
)
if view.speculative_num_draft_tokens is not None:
assert (
view.mamba_radix_cache_strategy != "extra_buffer_lazy"
), "extra_buffer_lazy unsupported with spec."
if view.mamba_radix_cache_strategy == "extra_buffer_lazy":
# The dflash family's verify bypasses prepare_mamba_track_for_verify.
assert view.speculative_algorithm not in ("DFLASH", "DSPARK"), (
f"extra_buffer_lazy unsupported with "
f"{view.speculative_algorithm}; use "
"--mamba-radix-cache-strategy extra_buffer."
)
assert view.mamba_track_interval >= view.speculative_num_draft_tokens
if view.page_size is not None:
assert view.mamba_track_interval % view.page_size == 0
python/sglang/srt/speculative/spec_utils.py
+22 -2
View File
@@ -667,10 +667,23 @@ def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
the mask also keeps a stale extend-time mask from triggering in-forward
tracking during TARGET_VERIFY; tracking is done in
commit_mamba_states_after_verify instead.
Lazy: gather the positions planned by mamba_lazy_spec_prepare. Runs
inside forward isolation, so it must not mutate req/pool state.
"""
if not get_server_args().enable_mamba_extra_buffer():
server_args = get_server_args()
if not server_args.enable_mamba_extra_buffer():
return
set_mamba_track_indices_from_reqs(batch)
track_positions = None
if server_args.enable_mamba_extra_buffer_lazy():
track_positions = batch.mamba_lazy_spec_track_positions_cpu
assert track_positions is not None and len(track_positions) == len(
batch.reqs
), (
"lazy spec verify without a track plan: mamba_lazy_spec_prepare "
"must run in prepare_for_decode for every spec decode iteration"
)
set_mamba_track_indices_from_reqs(batch, track_positions)
batch.mamba_track_mask = None
batch.mamba_track_seqlens = None
@@ -833,6 +846,13 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
"""eagle/ngram share a stateless free function; dflash keeps stateful
prep on its draft input -- the dispatcher routes.
"""
server_args = get_server_args()
if server_args.enable_mamba_extra_buffer_lazy():
# Scheduler phase (outside forward isolation).
batch.mamba_lazy_spec_prepare(
server_args.mamba_track_interval,
server_args.max_speculative_num_draft_tokens,
)
if batch.spec_algorithm.is_dflash_family():
batch.spec_info.prepare_for_decode(batch)
else:
test/registered/models_e2e/test_qwen3_next_models_mtp.py
+76 -1
View File
@@ -1,3 +1,4 @@
import os
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
@@ -6,11 +7,35 @@ from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=290, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=430, stage="base-c", runner_config="4-gpu-h100")
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
def _mtp_args(*, strategy, steps, topk, draft_tokens, track_interval):
return [
"--trust-remote-code",
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
str(steps),
"--speculative-eagle-topk",
str(topk),
"--speculative-num-draft-tokens",
str(draft_tokens),
"--mem-fraction-static",
"0.8",
"--tp",
"4",
"--chunked-prefill-size",
"2048",
"--mamba-scheduler-strategy",
strategy,
"--mamba-track-interval",
str(track_interval),
]
class TestQwen3NextMTPTopk(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
@@ -71,5 +96,55 @@ class TestQwen3NextMTPV2(GSM8KMixin, KLDivergenceMixin, DefaultServerBase):
]
class TestQwen3NextMTPLazyV2(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
# extra_buffer_lazy + MTP: the pending-slot track plan
# (mamba_lazy_spec_prepare / _mamba_lazy_spec_confirm_crossing) replaces
# the per-req second ping-pong slot. Branching mixin exercises donation
# of lazily tracked states.
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.0035
other_args = _mtp_args(
strategy="extra_buffer_lazy",
steps=3,
topk=1,
draft_tokens=4,
track_interval=128,
)
@unittest.skip("Manual-only: forces all lazy spec preallocs to fail")
class TestQwen3NextMTPLazyAllocFail(GSM8KMixin, KLDivergenceMixin, DefaultServerBase):
# Small track interval (= draft tokens) maximizes boundary crossings;
# forced alloc failure exercises the in-place fallback and the
# finished-req skip-insert path on every crossing.
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.0035
other_args = _mtp_args(
strategy="extra_buffer_lazy",
steps=3,
topk=1,
draft_tokens=4,
track_interval=4,
)
@classmethod
def setUpClass(cls):
os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
super().setUpClass()
@classmethod
def tearDownClass(cls):
super().tearDownClass()
os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
if __name__ == "__main__":
unittest.main()