diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index eea2a0768..5ee20bead 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1996,7 +1996,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--mamba-radix-cache-strategy` - The strategy to use for mamba scheduler. auto currently defaults to no_buffer. 1. no_buffer does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. extra_buffer supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes 2x for non-spec; 1+(1/(2+speculative_num_draft_tokens))x for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. extra_buffer is strictly better for non-KV-cache-bound cases; for KV-cache-bound cases, the tradeoff depends on whether enabling overlap outweighs reduced max running requests. 2b. mamba caching at radix cache branching point is strictly better than non-branch but requires kernel support, currently only extra_buffer supports branching. + The strategy to use for mamba scheduler. auto currently defaults to no_buffer. 1. no_buffer does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. extra_buffer supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes 2x for non-spec; 1+(1/(2+speculative_num_draft_tokens))x for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. extra_buffer is strictly better for non-KV-cache-bound cases; for KV-cache-bound cases, the tradeoff depends on whether enabling overlap outweighs reduced max running requests. 2b. mamba caching at radix cache branching point is strictly better than non-branch but requires kernel support, currently only extra_buffer supports branching. 3. extra_buffer_lazy lowers extra_buffer's slot cost by allocating one track slot per request instead of two; the second slot is allocated on demand at track-interval boundaries (for speculative decoding it is reserved ahead of each verify window and committed only for accepted boundary crossings). Compatible with speculative decoding (EAGLE/NGRAM/DSPARK/DFLASH); not supported under PD disaggregation. `auto` auto, no_buffer, extra_buffer, extra_buffer_lazy diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d7c0c15ab..4f23cc0d8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5772,13 +5772,9 @@ class ServerArgs: "extra_buffer_lazy unsupported under PD disaggregation; use " "--mamba-radix-cache-strategy extra_buffer." ) - algo = (view.speculative_algorithm or "").upper() - # dspark verifies through prepare_mamba_track_for_verify (lazy plan - # wired); dflash bypasses that hook, so it stays unsupported. - assert algo != "DFLASH", ( - f"extra_buffer_lazy unsupported with {view.speculative_algorithm}; " - "use --mamba-radix-cache-strategy extra_buffer." - ) + # eagle/ngram/dspark/dflash all verify through + # prepare_mamba_track_for_verify (lazy plan wired); dflash gained + # the hook in DFlashVerifyInput.prepare_for_verify. if view.speculative_num_draft_tokens is not None: assert view.mamba_track_interval >= view.speculative_num_draft_tokens if view.page_size is not None: diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 6b2cc6ec4..4fd0d1a27 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -62,6 +62,8 @@ class DFlashVerifyInput(SpecInput): metadata or eager attention metadata so the actual forward can run with `skip_attn_backend_init=True`. """ + from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify + batch.input_ids = self.draft_token batch.spec_info = self batch.forward_mode = ( @@ -69,6 +71,12 @@ class DFlashVerifyInput(SpecInput): if batch.forward_mode.is_idle() else ForwardMode.TARGET_VERIFY ) + if not batch.forward_mode.is_idle(): + # Rebuild mamba track indices (lazy: gather the positions planned + # by mamba_lazy_spec_prepare) and clear the stale extend-time mask + # before init_new snapshots them into the verify ForwardBatch. + # Same hook eagle/ngram/dspark run before TARGET_VERIFY. + prepare_mamba_track_for_verify(batch) verify_forward_batch = ForwardBatch.init_new( batch, target_worker.model_runner, diff --git a/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py b/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py new file mode 100644 index 000000000..cee36b559 --- /dev/null +++ b/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py @@ -0,0 +1,134 @@ +"""Unit tests for DFLASH + mamba-radix-cache-strategy extra_buffer_lazy: +server_args validation accepts the pairing, and DFlashVerifyInput.prepare_for_verify +runs prepare_mamba_track_for_verify (the hook eagle/ngram/dspark already run) +before ForwardBatch.init_new snapshots the track fields.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative import dflash_info +from sglang.srt.speculative.dflash_info import DFlashVerifyInput + + +def _lazy_view(**overrides): + view = SimpleNamespace( + mamba_radix_cache_strategy="extra_buffer_lazy", + disaggregation_mode="null", + speculative_algorithm="DFLASH", + speculative_num_draft_tokens=8, + mamba_track_interval=256, + page_size=64, + chunked_prefill_size=None, + ) + for key, value in overrides.items(): + setattr(view, key, value) + return view + + +class TestValidateMambaExtraBufferLazyDflash(CustomTestCase): + """The DFLASH rejection is gone; the neighboring invariants still hold.""" + + def _validate(self, view): + fake_self = SimpleNamespace(mamba_cache_chunk_size=64) + with mock.patch( + "sglang.srt.arg_groups.overrides.supports_mamba_cache_extra_buffer", + return_value=True, + ), mock.patch( + # Keep the test runnable on CPU-only hosts: the platform assert is + # not what is under test here. + "sglang.srt.server_args.is_cuda", + return_value=True, + ): + ServerArgs._validate_mamba_extra_buffer( + fake_self, view, "Qwen3NextForCausalLM" + ) + + def test_dflash_with_extra_buffer_lazy_is_accepted(self): + self._validate(_lazy_view()) + + def test_dspark_still_accepted(self): + self._validate(_lazy_view(speculative_algorithm="DSPARK")) + + def test_pd_disaggregation_still_rejected(self): + with self.assertRaisesRegex(AssertionError, "PD disaggregation"): + self._validate(_lazy_view(disaggregation_mode="decode")) + + def test_track_interval_must_cover_draft_tokens(self): + with self.assertRaises(AssertionError): + self._validate( + _lazy_view(speculative_num_draft_tokens=512, mamba_track_interval=256) + ) + + +class TestDflashVerifyRunsMambaTrackHook(CustomTestCase): + """prepare_for_verify calls prepare_mamba_track_for_verify after the batch + is stamped TARGET_VERIFY and before ForwardBatch.init_new; idle batches + skip the hook.""" + + def _spec_input(self): + return DFlashVerifyInput( + draft_token=torch.tensor([1, 2, 3, 4], dtype=torch.long), + positions=torch.tensor([0, 1, 2, 3], dtype=torch.long), + draft_token_num=4, + ) + + def _run(self, forward_mode): + calls = [] + batch = SimpleNamespace(forward_mode=forward_mode) + attn_backend = SimpleNamespace( + init_forward_metadata=lambda fb: calls.append("init_forward_metadata") + ) + target_worker = SimpleNamespace( + model_runner=SimpleNamespace( + decode_cuda_graph_runner=None, attn_backend=attn_backend + ) + ) + + def fake_hook(hook_batch): + calls.append(("hook", hook_batch.forward_mode)) + + fake_forward_batch = SimpleNamespace() + + def fake_init_new(*args, **kwargs): + calls.append("init_new") + return fake_forward_batch + + with mock.patch( + "sglang.srt.speculative.spec_utils.prepare_mamba_track_for_verify", + side_effect=fake_hook, + ), mock.patch.object( + dflash_info.ForwardBatch, "init_new", side_effect=fake_init_new + ): + out, can_run_cuda_graph = self._spec_input().prepare_for_verify( + batch, target_worker + ) + self.assertIs(out, fake_forward_batch) + self.assertFalse(can_run_cuda_graph) + return calls, batch + + def test_hook_runs_before_init_new_on_verify(self): + calls, batch = self._run(ForwardMode.DECODE) + self.assertEqual( + calls, + [("hook", ForwardMode.TARGET_VERIFY), "init_new", "init_forward_metadata"], + ) + self.assertEqual(batch.forward_mode, ForwardMode.TARGET_VERIFY) + + def test_idle_batch_skips_hook(self): + calls, batch = self._run(ForwardMode.IDLE) + self.assertEqual(calls, ["init_new"]) + self.assertEqual(batch.forward_mode, ForwardMode.IDLE) + + +if __name__ == "__main__": + unittest.main()