diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 2dee4ec64..97a3dceb6 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index b68d59161..248a92939 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5c77b66d4..0bc65b6de 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 39c7ef31b..4671818d7 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -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: diff --git a/test/registered/models_e2e/test_qwen3_next_models_mtp.py b/test/registered/models_e2e/test_qwen3_next_models_mtp.py index f74d1f6af..385d58f5a 100644 --- a/test/registered/models_e2e/test_qwen3_next_models_mtp.py +++ b/test/registered/models_e2e/test_qwen3_next_models_mtp.py @@ -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()