diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index b53326a6c..1ad52b123 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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, ) ) diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 3e2f93c2f..9488d1bb2 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/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: diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index 60650c8b3..9717414ef 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/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. diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index bef1f6791..51de39c6f 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f552102b8..bd9184408 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/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( diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 4b7c177a1..07028f04d 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/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 ( diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 3c90c034e..ebbfc2b28 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/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) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 812196ffd..7754eb6b5 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4ab9d4e80..4a3135df1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index fb4c6173e..95f3a5e5c 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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( diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 6550ccbe1..ba8d6265a 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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. diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 35b00f40d..2b9470f38 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 798bf9509..2ba072dcd 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index a3ca1ea54..ae30ece64 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index cf1a2da56..6aba4416b 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/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( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 4c055535f..a7a325dad 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/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( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 4b74c65aa..2f7903d24 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index c244afef9..336b6419e 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/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: diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py new file mode 100644 index 000000000..766891f5a --- /dev/null +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -0,0 +1,71 @@ +"""Regression: TBO filter_batch resets the attention plan marker on children. + +filter_batch's completeness guard raises for any non-None ForwardBatch field +missing from the child dict; the plan marker defaults to False (non-None) and +crashed TBO cuda-graph capture until reset. CPU-only. +""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang.srt.batch_overlap.two_batch_overlap as tbo +from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_target_verify_batch(bs: int) -> ForwardBatch: + return ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=bs, + input_ids=torch.zeros(bs, dtype=torch.long), + positions=torch.zeros(bs, dtype=torch.long), + out_cache_loc=torch.zeros(bs, dtype=torch.long), + req_pool_indices=torch.zeros(bs, dtype=torch.long), + seq_lens=torch.ones(bs, dtype=torch.int32), + seq_lens_cpu=torch.ones(bs, dtype=torch.int32), + seq_lens_sum=bs, + spec_info=None, + ) + + +def _filter(batch: ForwardBatch, *, lo: int, hi: int) -> ForwardBatch: + fake_args = SimpleNamespace(moe_dense_tp_size=None, attention_backend="fa3") + with patch.object(tbo, "get_attention_tp_size", lambda: 1), patch.object( + tbo, "get_global_server_args", lambda: fake_args + ): + return TboForwardBatchPreparer.filter_batch( + batch, + start_token_index=lo, + end_token_index=hi, + start_seq_index=lo, + end_seq_index=hi, + out_num_token_non_padded=torch.tensor(hi - lo), + ) + + +class TestTboFilterBatchMarker(CustomTestCase): + def test_filter_batch_resets_plan_marker_on_children(self): + child = _filter(_make_target_verify_batch(8), lo=0, hi=4) + self.assertEqual(child.batch_size, 4) + self.assertFalse(child.forward_metadata_ready) + self.assertIsNone(child.forward_metadata_planned_bs) + self.assertIsNone(child.forward_metadata_planned_num_tokens) + self.assertFalse(child.forward_metadata_replan_equivalent) + + def test_pre_planned_parent_does_not_leak_ready_into_children(self): + parent = _make_target_verify_batch(8) + parent.mark_forward_metadata_ready(replan_equivalent=True) + child = _filter(parent, lo=0, hi=4) + self.assertFalse(child.forward_metadata_ready) + self.assertFalse(child.forward_metadata_replan_equivalent) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_forward_metadata_plan_record.py b/test/registered/unit/model_executor/test_forward_metadata_plan_record.py new file mode 100644 index 000000000..770e8c32c --- /dev/null +++ b/test/registered/unit/model_executor/test_forward_metadata_plan_record.py @@ -0,0 +1,127 @@ +"""Unit tests for the ForwardBatch attention plan marker / plan record. + +Covers the contract behind ``skip_attn_backend_init`` deprecation: + * fresh batches need planning; marked batches don't + * the plan record (planned bs / num tokens) snapshots mark-time shapes + * reshape after marking triggers a re-plan only for sites that opted + into ``replan_equivalent``; re-marking re-records the new shapes + * the deprecated kwarg shim maps explicit values onto the marker + (mapped, not ignored) and warns once per process + +Pure dataclass logic — CPU only. +""" + +import unittest +import warnings + +import torch + +import sglang.srt.model_executor.forward_batch_info as fbi +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_batch(bs: int = 2, num_tokens: int = 2) -> ForwardBatch: + return ForwardBatch( + forward_mode=ForwardMode.DECODE, + batch_size=bs, + input_ids=torch.zeros(num_tokens, dtype=torch.long), + req_pool_indices=torch.zeros(bs, dtype=torch.long), + seq_lens=torch.ones(bs, dtype=torch.long), + out_cache_loc=torch.zeros(bs, dtype=torch.long), + seq_lens_sum=bs, + ) + + +class TestForwardMetadataPlanRecord(CustomTestCase): + def test_fresh_batch_needs_planning(self): + fb = _make_batch() + self.assertTrue(fb.needs_forward_metadata_init()) + self.assertFalse(fb.forward_metadata_ready) + + def test_mark_records_shapes_and_skips_planning(self): + fb = _make_batch(bs=3, num_tokens=7) + fb.mark_forward_metadata_ready() + self.assertFalse(fb.needs_forward_metadata_init()) + self.assertEqual(fb.forward_metadata_planned_bs, 3) + self.assertEqual(fb.forward_metadata_planned_num_tokens, 7) + + def test_reshape_without_opt_in_keeps_skipping(self): + # Wrapper regimes must never auto-re-plan (would clobber per-step metadata). + fb = _make_batch(bs=2) + fb.mark_forward_metadata_ready() + fb.batch_size = 4 # DP padding reshapes the batch + self.assertFalse(fb.needs_forward_metadata_init()) + + def test_reshape_with_opt_in_replans(self): + fb = _make_batch(bs=2, num_tokens=2) + fb.mark_forward_metadata_ready(replan_equivalent=True) + self.assertFalse(fb.needs_forward_metadata_init()) + + fb.batch_size = 4 # bs drift (prepare_mlp_sync_batch decode pad) + self.assertTrue(fb.needs_forward_metadata_init()) + + fb.batch_size = 2 + fb.input_ids = torch.zeros(6, dtype=torch.long) # token drift + self.assertTrue(fb.needs_forward_metadata_init()) + + def test_remark_re_records_padded_shapes(self): + # Per-step loops re-mark each plan; the re-mark must snapshot padded shapes. + fb = _make_batch(bs=2) + fb.mark_forward_metadata_ready(replan_equivalent=True) + fb.batch_size = 4 + self.assertTrue(fb.needs_forward_metadata_init()) + fb.mark_forward_metadata_ready(replan_equivalent=True) + self.assertFalse(fb.needs_forward_metadata_init()) + self.assertEqual(fb.forward_metadata_planned_bs, 4) + + +class TestDeprecatedSkipKwargShim(CustomTestCase): + def setUp(self): + self._saved_warned = fbi._skip_attn_backend_init_warned + fbi._skip_attn_backend_init_warned = False + + def tearDown(self): + fbi._skip_attn_backend_init_warned = self._saved_warned + + def test_none_is_a_silent_no_op(self): + fb = _make_batch() + with warnings.catch_warnings(): + warnings.simplefilter("error") + fb.apply_deprecated_skip_attn_backend_init(None) + self.assertTrue(fb.needs_forward_metadata_init()) + + def test_true_maps_onto_marker_and_warns(self): + # Mapped, not ignored: a no-op would silently re-plan multi-step metadata. + fb = _make_batch() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fb.apply_deprecated_skip_attn_backend_init(True) + self.assertFalse(fb.needs_forward_metadata_init()) + self.assertEqual(len(caught), 1) + self.assertTrue(issubclass(caught[0].category, DeprecationWarning)) + + def test_false_warns_but_does_not_mark(self): + fb = _make_batch() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fb.apply_deprecated_skip_attn_backend_init(False) + self.assertTrue(fb.needs_forward_metadata_init()) + self.assertEqual(len(caught), 1) + + def test_warns_once_per_process(self): + # Hot-loop guard: per-forward callers must not pay warnings.warn repeatedly. + fb = _make_batch() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fb.apply_deprecated_skip_attn_backend_init(True) + fb.apply_deprecated_skip_attn_backend_init(True) + _make_batch().apply_deprecated_skip_attn_backend_init(False) + self.assertEqual(len(caught), 1) + + +if __name__ == "__main__": + unittest.main()