Replace skip_attn_backend_init with a batch-carried attention plan marker (+ staleness re-plan) (#27193)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -761,6 +761,12 @@ class TboForwardBatchPreparer:
                 token_ids_logprobs=None,
                 next_token_logits_buffer=None,
                 return_hidden_states_before_norm=False,
+                # TBO children start unplanned — planned by the TBO-aware init
+                # flow; a stale parent "ready" would wrongly skip that.
+                forward_metadata_ready=False,
+                forward_metadata_planned_bs=None,
+                forward_metadata_planned_num_tokens=None,
+                forward_metadata_replan_equivalent=False,
             )
         )
 

python/sglang/srt/hardware_backend/mlx/tp_worker.py

@@ -98,7 +98,7 @@ class MlxTpModelWorker(TpModelWorker):
         forward_batch: Optional[ForwardBatch] = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         is_verify: bool = False,
-        skip_attn_backend_init=False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
     ) -> GenerationBatchResult:
         """Override to route through MLX model runner."""
         if batch is not None:

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -180,10 +180,9 @@ class NPUGraphRunner(CudaGraphRunner):
     def replay(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
-        if not skip_attn_backend_init:
+        if forward_batch.needs_forward_metadata_init():
             self.replay_prepare(forward_batch, pp_proxy_tensors)
         else:
             # In speculative decoding, these two fields are still needed.

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -956,9 +956,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             or self.forward_decode_metadata
         )
 
-        # Ensure batch_size is sufficient, the batch size increase due to the padding from the forward batch
-        # FIXME(@rainj-me), refactor the skip_attn_backend_init, init_forward_metadata for attn backends
-        # and padding logic in prepare_mlp_sync_batch to avoid this
+        # Backstop: metadata was built pre-pad (marked) and DP padding then
+        # grew the batch. The marker path deliberately does not re-plan
+        # post-pad (DSA can't rebuild on a padded batch, see #27091), so this
+        # local re-plan catches the size mismatch.
         batch_size = getattr(metadata, "batch_size", None)
         if batch_size is not None and batch_size < forward_batch.batch_size:
             self.init_forward_metadata(forward_batch)
@@ -1058,9 +1059,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 or self.forward_decode_metadata
             )
 
-            # Ensure batch_size is sufficient, the batch size increase due to the padding from the forward batch
-            # FIXME(@rainj-me), refactor the skip_attn_backend_init, init_forward_metadata for attn backends
-            # and padding logic in prepare_mlp_sync_batch to avoid this
+            # Backstop: metadata was built pre-pad (marked) and DP padding
+            # then grew the batch. The marker path deliberately does not
+            # re-plan post-pad (DSA can't rebuild on a padded batch, see
+            # #27091), so this local re-plan catches the size mismatch.
             batch_size = getattr(metadata, "batch_size", None)
             if batch_size is not None and batch_size < forward_batch.batch_size:
                 self.init_forward_metadata(forward_batch)

python/sglang/srt/managers/tp_worker.py

@@ -450,11 +450,8 @@ class TpModelWorker(BaseTpWorker):
         forward_batch: Optional[ForwardBatch] = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         is_verify: bool = False,
-        skip_attn_backend_init=False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
     ) -> GenerationBatchResult:
-        # FIXME(lsyin): maybe remove skip_attn_backend_init in forward_batch_generation,
-        #               which requires preparing replay to always be in this function
-
         # Get forward batch from schedule batch
         if batch is not None:
             # update the consumer index of hicache to the running batch
@@ -465,6 +462,9 @@ class TpModelWorker(BaseTpWorker):
             # FIXME(lsyin): unify the interface of forward_batch
             assert forward_batch is not None
 
+        # Deprecated kwarg: pre-planners mark the batch themselves now.
+        forward_batch.apply_deprecated_skip_attn_backend_init(skip_attn_backend_init)
+
         if self.is_dllm():
             return self._forward_batch_generation_dllm(forward_batch)
 
@@ -472,7 +472,6 @@ class TpModelWorker(BaseTpWorker):
             out = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
-                skip_attn_backend_init=skip_attn_backend_init,
             )
             logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
             batch_result = GenerationBatchResult(
@@ -529,7 +528,6 @@ class TpModelWorker(BaseTpWorker):
             out = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
-                skip_attn_backend_init=skip_attn_backend_init,
             )
             pp_proxy_tensors, can_run_cuda_graph = out.logits_output, out.can_run_graph
             return GenerationBatchResult(

python/sglang/srt/model_executor/cpu_graph_runner.py

@@ -816,7 +816,6 @@ class CPUGraphRunner:
     def replay(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
         assert (

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1163,14 +1163,14 @@ class CudaGraphRunner:
     def replay(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
         self.deepep_adapter.replay()
 
-        if not skip_attn_backend_init:
+        if forward_batch.needs_forward_metadata_init():
             self.replay_prepare(forward_batch, pp_proxy_tensors)
         else:
+            # Pre-planned (plan-stream replay_prepare already ran).
             # In speculative decoding, these two fields are still needed.
             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)

python/sglang/srt/model_executor/forward_batch_info.py

@@ -28,6 +28,7 @@ ScheduleBatch -> ForwardBatch
 from __future__ import annotations
 
 import hashlib
+import warnings
 from dataclasses import dataclass
 from enum import IntEnum, auto
 from functools import total_ordering
@@ -72,6 +73,10 @@ if TYPE_CHECKING:
     from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
     from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
 
+# Warn-once flag for the deprecated skip_attn_backend_init kwarg; see
+# ForwardBatch.apply_deprecated_skip_attn_backend_init.
+_skip_attn_backend_init_warned = False
+
 _is_npu = is_npu()
 
 
@@ -459,6 +464,94 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     req_all_ids_flat: Optional[torch.Tensor] = None
     req_all_ids_lens: Optional[torch.Tensor] = None
 
+    # Attention planning state. True iff attention metadata for this batch has
+    # already been planned outside ModelRunner.forward (multi-step draft
+    # pre-plan, plan-stream replay_prepare, hand-built spec batches), so the
+    # forward path must not plan again. Only such pre-planners may set this —
+    # ModelRunner / graph runners never mark after their own planning. The
+    # marker is only valid for the planning regime (backend set) it was set
+    # under; a fresh batch from init_new always starts unplanned.
+    forward_metadata_ready: bool = False
+    # Shapes the batch had when it was marked (plan record). Lets the
+    # judgment predicate detect staleness when DP padding
+    # (prepare_mlp_sync_batch) reshapes the batch after pre-planning.
+    # Deliberately plain ints — no planner object ref on ForwardBatch
+    # (runtime refs were removed from this dataclass on purpose).
+    forward_metadata_planned_bs: Optional[int] = None
+    forward_metadata_planned_num_tokens: Optional[int] = None
+    # Whether the forward path may re-plan this batch when its shapes no
+    # longer match the plan record. Only mark sites where the forward
+    # path's own init_forward_metadata is equivalent to the pre-plan
+    # (same backend object, no special context) may opt in; multi-step
+    # wrapper plans and view-context plans must keep this False — a
+    # forward-path re-plan would clobber their metadata.
+    forward_metadata_replan_equivalent: bool = False
+
+    def mark_forward_metadata_ready(self, replan_equivalent: bool = False):
+        """Record that attention metadata was pre-planned for this batch.
+
+        Call right next to the out-of-forward planning action
+        (e.g. ``draft_attn_backend.init_forward_metadata(fb)`` or
+        ``graph_runner.replay_prepare(fb)``). Records the batch shapes so
+        staleness is detectable; pass ``replan_equivalent=True`` only when
+        a forward-path re-plan is equivalent to the pre-plan (see field
+        docs).
+        """
+        self.forward_metadata_ready = True
+        self.forward_metadata_planned_bs = self.batch_size
+        self.forward_metadata_planned_num_tokens = (
+            self.input_ids.shape[0] if self.input_ids is not None else 0
+        )
+        self.forward_metadata_replan_equivalent = replan_equivalent
+
+    def needs_forward_metadata_init(self) -> bool:
+        """Single judgment point for whether the forward path must plan.
+
+        A marked batch is treated as stale — and re-planned — when its
+        shapes no longer match the plan record AND the mark site declared
+        the re-plan safe (replan_equivalent). This runs after
+        prepare_mlp_sync_batch in _forward_raw, so the re-plan sees the
+        padded (final) shapes. Sites that cannot opt in (multi-step
+        wrapper plans etc.) keep today's behavior: marked stays skipped,
+        backends' defensive checks remain the backstop.
+        """
+        if not self.forward_metadata_ready:
+            return True
+        if not self.forward_metadata_replan_equivalent:
+            return False
+        num_tokens = self.input_ids.shape[0] if self.input_ids is not None else 0
+        return (
+            self.batch_size != self.forward_metadata_planned_bs
+            or num_tokens != self.forward_metadata_planned_num_tokens
+        )
+
+    def apply_deprecated_skip_attn_backend_init(
+        self, skip_attn_backend_init: Optional[bool]
+    ) -> None:
+        """Map the deprecated ``skip_attn_backend_init`` kwarg onto the marker.
+
+        Mapped, not ignored: callers passing True relied on planning being
+        skipped — ignoring the flag would silently re-plan and corrupt
+        pre-planned multi-step draft metadata. Warns once per process (a
+        module flag, not the warnings filter, so the hot decode loop never
+        pays warnings.warn per forward).
+        """
+        if skip_attn_backend_init is None:
+            return
+        global _skip_attn_backend_init_warned
+        if not _skip_attn_backend_init_warned:
+            _skip_attn_backend_init_warned = True
+            warnings.warn(
+                "skip_attn_backend_init is deprecated and will be removed; "
+                "pre-planners should call "
+                "ForwardBatch.mark_forward_metadata_ready() after planning "
+                "instead. The flag is mapped onto the marker for now.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        if skip_attn_backend_init:
+            self.mark_forward_metadata_ready()
+
     @classmethod
     def init_new(
         cls,

python/sglang/srt/model_executor/model_runner.py

@@ -3074,12 +3074,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def forward_decode(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
         pp_proxy_tensors=None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
         # Set extra arguments
         pdmux_override = False
-        if not skip_attn_backend_init:
+        if forward_batch.needs_forward_metadata_init():
             if hasattr(self.model, "prepare_forward_batch"):
                 # Prepare model-specific attention metadata before planning,
                 # e.g. Moss-VL's prefill cross-attention custom mask.
@@ -3123,7 +3122,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def forward_extend(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
         pp_proxy_tensors=None,
     ) -> Tuple[
         Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool
@@ -3167,7 +3165,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             return (ret, can_run_graph)
 
         # Launch model forward
-        if not skip_attn_backend_init:
+        if forward_batch.needs_forward_metadata_init():
             if hasattr(self.model, "prepare_forward_batch"):
                 # Prepare model-specific attention metadata before planning,
                 # e.g. Moss-VL's prefill cross-attention custom mask.
@@ -3249,11 +3247,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def forward(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool = False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         reinit_attn_backend: bool = False,
         split_forward_count: int = 1,
     ) -> ModelRunnerOutput:
+        # Deprecated kwarg: pre-planners mark the batch themselves now.
+        forward_batch.apply_deprecated_skip_attn_backend_init(skip_attn_backend_init)
+
         self.forward_pass_id += 1
 
         # Try msprob debugger
@@ -3292,7 +3293,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         ):
             output = self._forward_raw(
                 forward_batch,
-                skip_attn_backend_init,
                 pp_proxy_tensors,
                 reinit_attn_backend,
                 split_forward_count,
@@ -3301,7 +3301,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 output = self._maybe_rebalance_after_rank_fault(
                     output,
                     forward_batch,
-                    skip_attn_backend_init,
                     pp_proxy_tensors,
                     reinit_attn_backend,
                     split_forward_count,
@@ -3343,7 +3342,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def _forward_raw(
         self,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool,
         pp_proxy_tensors: Optional[PPProxyTensors],
         reinit_attn_backend: bool = False,
         split_forward_count: int = 1,
@@ -3379,7 +3377,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             if can_run_graph:
                 ret = self.graph_runner.replay(
                     forward_batch,
-                    skip_attn_backend_init=skip_attn_backend_init,
                     pp_proxy_tensors=pp_proxy_tensors,
                 )
                 return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph)
@@ -3415,7 +3412,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             if forward_batch.forward_mode.is_decode():
                 ret = self.forward_decode(
                     forward_batch,
-                    skip_attn_backend_init=skip_attn_backend_init,
                     pp_proxy_tensors=pp_proxy_tensors,
                 )
             elif forward_batch.forward_mode.is_split_prefill():
@@ -3427,7 +3423,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
                 ret, can_run_graph = self.forward_extend(
                     forward_batch,
-                    skip_attn_backend_init=skip_attn_backend_init,
                     pp_proxy_tensors=pp_proxy_tensors,
                 )
             elif forward_batch.forward_mode.is_idle():
@@ -3588,7 +3583,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self,
         output: ModelRunnerOutput,
         forward_batch: ForwardBatch,
-        skip_attn_backend_init: bool,
         pp_proxy_tensors: Optional[PPProxyTensors],
         reinit_attn_backend: bool,
         split_forward_count: int,
@@ -3606,7 +3600,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     break
             output = self._forward_raw(
                 forward_batch,
-                skip_attn_backend_init,
                 pp_proxy_tensors,
                 reinit_attn_backend,
                 split_forward_count,

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -377,6 +377,9 @@ class EAGLEDraftCudaGraphRunner:
             self.draft_attn_backend.init_forward_metadata_out_graph(
                 forward_batch, in_capture=True
             )
+            # The capture batch is planned here (out-of-forward), so the
+            # per-step forwards inside draft_forward must not re-plan.
+            forward_batch.mark_forward_metadata_ready()
             self.deepep_adapter.capture(is_extend_in_batch=False)
             self._capture_init(run_once)
             out = self._capture_graph(

python/sglang/srt/speculative/eagle_info_v2.py

@@ -259,6 +259,12 @@ class EagleDraftInputV2Mixin:
         can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
         if not batch.forward_mode.is_idle() and not can_cuda_graph:
             draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
+            # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer
+            # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch
+            # (the `_batch_size == batch_size` assertion, see #27091); the
+            # marked pre-pad metadata is used as-is, matching the proven
+            # skip_attn_backend_init=True behavior.
+            forward_batch.mark_forward_metadata_ready()
         return forward_batch
 
 
@@ -325,6 +331,7 @@ class EagleVerifyInputV2Mixin:
         )
         if can_run_cuda_graph:
             target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
+            verify_forward_batch.mark_forward_metadata_ready()
         # Non-cuda-graph: defer init to forward_extend, which runs after
         # `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
         # here would use pre-pad shapes and trip DSv4 indexer shape match.

python/sglang/srt/speculative/eagle_worker.py

@@ -783,6 +783,7 @@ class EAGLEWorker(TpModelWorker):
             ):
                 # Skip attention backend init for idle mode or 1-step draft
                 self.draft_attn_backend.init_forward_metadata(forward_batch)
+                forward_batch.mark_forward_metadata_ready()
             # Run forward steps
             parent_list, top_scores_index, draft_tokens = self.draft_forward(
                 forward_batch
@@ -898,7 +899,7 @@ class EAGLEWorker(TpModelWorker):
                 ForwardContext(attn_backend=self.draft_attn_backend.attn_backends[i])
             ):
                 logits_output = self.draft_model_runner.forward(
-                    forward_batch, skip_attn_backend_init=True
+                    forward_batch
                 ).logits_output
             maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}")
             maybe_detect_inf(logits_output.next_token_logits, f"draft_forward step {i}")
@@ -1220,6 +1221,7 @@ class EAGLEWorker(TpModelWorker):
                     or self.draft_model_runner.attn_backend
                 )
                 attn_backend.init_forward_metadata(forward_batch)
+                forward_batch.mark_forward_metadata_ready()
             # Publish the chosen backend via ForwardContext so model code
             # picks it up for this forward (no runner-attr mutation).
             if attn_backend is not None:
@@ -1228,7 +1230,7 @@ class EAGLEWorker(TpModelWorker):
                 ctx_mgr = contextlib.nullcontext()
             with ctx_mgr:
                 logits_output = self.draft_model_runner.forward(
-                    forward_batch, skip_attn_backend_init=True
+                    forward_batch
                 ).logits_output
             # Non-cuda-graph path: compute topk_p / topk_index inline.
             probs = torch.softmax(logits_output.next_token_logits, dim=-1)

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -427,6 +427,7 @@ class EagleDraftWorker(BaseDraftWorker):
                     # Skip attention backend init for 1-step draft,
                     # `draft_forward` only does sample in this case.
                     self.draft_attn_backend.init_forward_metadata(forward_batch)
+                    forward_batch.mark_forward_metadata_ready()
                 parent_list, top_scores_index, draft_tokens = self.draft_forward(
                     forward_batch
                 )
@@ -559,9 +560,7 @@ class EagleDraftWorker(BaseDraftWorker):
             with forward_context(
                 ForwardContext(attn_backend=self.draft_attn_backend.attn_backends[i])
             ), canary_index_ctx:
-                logits_output = self.draft_runner.forward(
-                    forward_batch, skip_attn_backend_init=True
-                ).logits_output
+                logits_output = self.draft_runner.forward(forward_batch).logits_output
             maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}")
             maybe_detect_inf(logits_output.next_token_logits, f"draft_forward step {i}")
             if self.topk == 1 and not _is_hip:
@@ -750,7 +749,7 @@ class EagleDraftWorker(BaseDraftWorker):
                 )
             else:
                 draft_logits_output = self.draft_runner.forward(
-                    forward_batch, skip_attn_backend_init=True
+                    forward_batch
                 ).logits_output
 
         maybe_detect_nan(
@@ -1181,13 +1180,14 @@ class EAGLEWorkerV2(BaseSpecWorker):
             ).cpu()
 
         # Run target verify batch in the main compute stream (GPU compute).
-        # Only skip metadata init when cuda-graph already ran replay_prepare;
-        # the non-cuda-graph path needs forward_extend's init (post-pad).
+        # Metadata init is skipped iff cuda-graph already ran replay_prepare —
+        # prepare_for_v2_verify marked the batch in exactly that case; the
+        # non-cuda-graph path stays unmarked and gets forward_extend's init
+        # (post-pad).
         forward_batch_output = self.target_worker.forward_batch_generation(
             batch=None,
             forward_batch=verify_forward_batch,
             is_verify=True,
-            skip_attn_backend_init=can_run_cuda_graph,
         )
         logits_output = forward_batch_output.logits_output
 

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -280,9 +280,9 @@ class FrozenKVMTPCudaGraphRunner:
             set_is_extend_in_batch(False)
 
             hidden_states_backup = forward_batch.spec_info.hidden_states
-            ret = self.frozen_kv_mtp_worker.draft_forward(
-                forward_batch, skip_attn_backend_init=True
-            )
+            # The capture batch is marked by the capture metadata helper
+            # below, so draft_forward skips its eager plan.
+            ret = self.frozen_kv_mtp_worker.draft_forward(forward_batch)
             forward_batch.spec_info.hidden_states = hidden_states_backup
             return ret
 

python/sglang/srt/speculative/frozen_kv_mtp_worker.py

@@ -283,6 +283,7 @@ class FrozenKVMTPWorker(TpModelWorker):
             forward_batch.seq_lens_sum = torch.sum(forward_batch.seq_lens).item()
         with self._frozen_kv_target_view(forward_batch):
             self.draft_attn_backend.init_forward_metadata(forward_batch)
+        forward_batch.mark_forward_metadata_ready()
 
     def _init_frozen_kv_metadata_capture_cuda_graph(
         self, forward_batch: ForwardBatch
@@ -291,6 +292,7 @@ class FrozenKVMTPWorker(TpModelWorker):
             self.draft_attn_backend.init_forward_metadata_out_graph(
                 forward_batch, in_capture=True
             )
+        forward_batch.mark_forward_metadata_ready()
 
     def _init_frozen_kv_metadata_replay_cuda_graph(
         self, forward_batch: ForwardBatch, bs: int, seq_lens_sum: int
@@ -622,9 +624,7 @@ class FrozenKVMTPWorker(TpModelWorker):
             seq_lens_cpu=batch.seq_lens_cpu,
         )
 
-    def draft_forward(
-        self, forward_batch: ForwardBatch, skip_attn_backend_init: bool = False
-    ):
+    def draft_forward(self, forward_batch: ForwardBatch):
         spec_info = forward_batch.spec_info
         assert isinstance(spec_info, FrozenKVMTPDraftInput)
 
@@ -634,7 +634,7 @@ class FrozenKVMTPWorker(TpModelWorker):
 
         # Seed + recurrent iters share the same `seq_lens - 1` rope position,
         # so one init covers the loop. Must run even at num_steps == 1.
-        if not skip_attn_backend_init:
+        if forward_batch.needs_forward_metadata_init():
             self._init_frozen_kv_metadata(forward_batch)
 
         # Seed iter: assistant forward on (bonus_token, target_h) to produce
@@ -657,9 +657,7 @@ class FrozenKVMTPWorker(TpModelWorker):
             self._target_kv_pool_view(forward_batch),
             forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
         ):
-            seed_output = self.draft_model_runner.forward(
-                forward_batch, skip_attn_backend_init=True
-            ).logits_output
+            seed_output = self.draft_model_runner.forward(forward_batch).logits_output
 
         maybe_detect_nan(
             seed_output.next_token_logits, "frozen_kv_mtp_draft: seed iter"
@@ -703,7 +701,7 @@ class FrozenKVMTPWorker(TpModelWorker):
                 forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
             ):
                 logits_output = self.draft_model_runner.forward(
-                    forward_batch, skip_attn_backend_init=True
+                    forward_batch
                 ).logits_output
 
             maybe_detect_nan(

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -764,10 +764,12 @@ class MultiLayerEagleWorker(TpModelWorker):
                     self.mtp_model_runner(step).attn_backend.init_forward_metadata(
                         forward_batch
                     )
+                    # Planned pre-pad; do NOT opt into post-pad re-plan — a
+                    # DP-padded re-plan breaks DSA's indexer schedule_meta
+                    # (see #27091). Use the marked pre-pad metadata as-is.
+                    forward_batch.mark_forward_metadata_ready()
                 logits_output = (
-                    self.mtp_model_runner(step)
-                    .forward(forward_batch, skip_attn_backend_init=True)
-                    .logits_output
+                    self.mtp_model_runner(step).forward(forward_batch).logits_output
                 )
 
             maybe_detect_nan(

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -516,6 +516,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
                 + batch_result.accept_lens
                 - 1
             )
+            # NOTE: this non-graph path runs the per-step forwards without any
+            # pre-plan (see warning above). Mark the batch so the forward path
+            # keeps skipping metadata init — preserves the pre-existing
+            # behavior; the latent issue is tracked by the warning.
+            forward_batch.mark_forward_metadata_ready()
 
         for step in range(self.speculative_num_steps):
             # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}")
@@ -530,8 +535,10 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
                     draft_logits_output.topk_index,
                 )
             else:
+                # Skip relies on the unconditional mark above (pre-existing
+                # no-pre-plan behavior preserved verbatim).
                 draft_logits_output = self.draft_runner_list[step].forward(
-                    forward_batch, skip_attn_backend_init=True
+                    forward_batch
                 )
                 probs = torch.softmax(
                     draft_logits_output.logits_output.next_token_logits[select_index],
@@ -768,12 +775,16 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
                     else None
                 ),
             )
+        # NOTE: metadata init is skipped here unconditionally, although
+        # prepare_for_v2_verify only plans when cuda-graph replay_prepare ran.
+        # eagle_worker_v2 re-inits the non-graph path instead (post-pad); this
+        # worker has not adopted that fix, so preserve its behavior verbatim.
+        verify_forward_batch.mark_forward_metadata_ready()
         # Run target verify batch in the main compute stream
         forward_batch_output = self.target_worker.forward_batch_generation(
             batch=None,
             forward_batch=verify_forward_batch,
             is_verify=True,
-            skip_attn_backend_init=True,
         )
         logits_output = forward_batch_output.logits_output
 

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py

@@ -390,6 +390,7 @@ def _run_eagle_draft_eager(
         init_eager_metadata(worker, batch, settings)
     else:
         worker.draft_attn_backend.init_forward_metadata(batch)
+    batch.mark_forward_metadata_ready()  # mirror production: pre-plan marks
     return worker.draft_forward(batch)
 
 
@@ -397,7 +398,7 @@ def _run_frozen_kv_mtp_eager(
     worker: _FrozenKVMTPWorkerHarness,
     batch: ForwardBatch,
 ):
-    return worker.draft_forward(batch, skip_attn_backend_init=False)
+    return worker.draft_forward(batch)
 
 
 def _capture_eagle_draft_graph_runner(
@@ -588,8 +589,10 @@ class _DenseEagleDraftForward:
             hidden_size, vocab_size, bias=False, dtype=dtype, device=device
         )
 
-    def __call__(self, forward_batch: ForwardBatch, *, skip_attn_backend_init: bool):
-        del skip_attn_backend_init
+    def __call__(self, forward_batch: ForwardBatch):
+        assert (
+            forward_batch.forward_metadata_ready
+        ), "draft-loop forward reached the runner without a pre-planned batch"
         spec_info = forward_batch.spec_info
         hidden_states = spec_info.hidden_states
         if hidden_states is None:
@@ -625,8 +628,10 @@ class _FrozenKVMTPDenseDraftForward:
             hidden_size, vocab_size, bias=False, dtype=dtype, device=device
         )
 
-    def __call__(self, forward_batch: ForwardBatch, *, skip_attn_backend_init: bool):
-        del skip_attn_backend_init
+    def __call__(self, forward_batch: ForwardBatch):
+        assert (
+            forward_batch.forward_metadata_ready
+        ), "draft-loop forward reached the runner without a pre-planned batch"
         spec_info = forward_batch.spec_info
         hidden_states = spec_info.hidden_states
         if hidden_states is None:
@@ -1012,8 +1017,10 @@ class _MLAEagleDraftForward:
             hidden_size, vocab_size, bias=False, dtype=dtype, device=device
         )
 
-    def __call__(self, forward_batch: ForwardBatch, *, skip_attn_backend_init: bool):
-        del skip_attn_backend_init
+    def __call__(self, forward_batch: ForwardBatch):
+        assert (
+            forward_batch.forward_metadata_ready
+        ), "draft-loop forward reached the runner without a pre-planned batch"
         spec_info = forward_batch.spec_info
         hidden_states = spec_info.hidden_states
         if hidden_states is None:
@@ -1268,8 +1275,10 @@ class _DSV4EagleDraftForward:
             hidden_size, vocab_size, bias=False, dtype=dtype, device=device
         )
 
-    def __call__(self, forward_batch: ForwardBatch, *, skip_attn_backend_init: bool):
-        del skip_attn_backend_init
+    def __call__(self, forward_batch: ForwardBatch):
+        assert (
+            forward_batch.forward_metadata_ready
+        ), "draft-loop forward reached the runner without a pre-planned batch"
         spec_info = forward_batch.spec_info
         hidden_states = spec_info.hidden_states
         if hidden_states is None:
@@ -1568,8 +1577,10 @@ class _DSAEagleDraftForward:
             torch.full_like(indices, -1),
         )
 
-    def __call__(self, forward_batch: ForwardBatch, *, skip_attn_backend_init: bool):
-        del skip_attn_backend_init
+    def __call__(self, forward_batch: ForwardBatch):
+        assert (
+            forward_batch.forward_metadata_ready
+        ), "draft-loop forward reached the runner without a pre-planned batch"
         spec_info = forward_batch.spec_info
         hidden_states = spec_info.hidden_states
         if hidden_states is None: