[Spec] Support mamba-radix-cache-strategy extra_buffer_lazy with DFLASH (#34763)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

docs/docs/advanced_features/server_arguments.mdx

@@ -1996,7 +1996,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--mamba-radix-cache-strategy`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The strategy to use for mamba scheduler. <code>auto</code> currently defaults to <code>no_buffer</code>. 1. <code>no_buffer</code> does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. <code>extra_buffer</code> supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes <code>2x</code> for non-spec; <code>1+(1/(2+speculative_num_draft_tokens))x</code> for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. <code>extra_buffer</code> 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.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The strategy to use for mamba scheduler. <code>auto</code> currently defaults to <code>no_buffer</code>. 1. <code>no_buffer</code> does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. <code>extra_buffer</code> supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes <code>2x</code> for non-spec; <code>1+(1/(2+speculative_num_draft_tokens))x</code> for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. <code>extra_buffer</code> 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. <code>extra_buffer_lazy</code> lowers <code>extra_buffer</code>&apos;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.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`auto`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>no_buffer</code>, <code>extra_buffer</code>, <code>extra_buffer_lazy</code></td>
     </tr>

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:

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,

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()