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

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

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,