refactor: remove ModelWorkerBatch indirection (#25516)

python/sglang/bench_one_batch.py

@@ -457,8 +457,7 @@ def extend(reqs, model_runner):
     )
     batch.prepare_for_extend()
     _maybe_prepare_mlp_sync_batch(batch, model_runner)
-    model_worker_batch = batch.get_model_worker_batch()
+    forward_batch = ForwardBatch.init_new(batch, model_runner)
-    forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
     logits_output = model_runner.forward(forward_batch).logits_output
     next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits, batch
@@ -469,8 +468,7 @@ def decode(input_token_ids, batch, model_runner):
     batch.output_ids = input_token_ids
     batch.prepare_for_decode()
     _maybe_prepare_mlp_sync_batch(batch, model_runner)
-    model_worker_batch = batch.get_model_worker_batch()
+    forward_batch = ForwardBatch.init_new(batch, model_runner)
-    forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
     logits_output = model_runner.forward(forward_batch).logits_output
     next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits

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

@@ -136,9 +136,8 @@ class SchedulerMlxOverlapMixin:
             self.process_batch_result(pending.batch_copy, result)
 
         def _launch_fresh(batch: "ScheduleBatch") -> MlxPendingJob:
-            mwb = batch.get_model_worker_batch()
             lazy_tokens, prefills, extends, decode, mode = (
-                self.tp_worker.async_forward_batch_generation_mlx(mwb)
+                self.tp_worker.async_forward_batch_generation_mlx(batch)
             )
             return MlxPendingJob(
                 lazy_tokens=lazy_tokens,

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

@@ -23,7 +23,7 @@ from sglang.srt.hardware_backend.mlx.model_runner import (
     MlxPendingExtend,
     MlxPendingPrefill,
 )
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
@@ -94,20 +94,20 @@ class MlxTpModelWorker(TpModelWorker):
 
     def forward_batch_generation(
         self,
-        model_worker_batch: ModelWorkerBatch,
+        batch: Optional[ScheduleBatch],
         forward_batch: Optional[ForwardBatch] = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         is_verify: bool = False,
         skip_attn_backend_init=False,
     ) -> GenerationBatchResult:
         """Override to route through MLX model runner."""
-        if model_worker_batch is not None:
+        if batch is not None:
             self._ensure_mlx_pool_initialized()
-            return self._forward_batch_generation_mlx(model_worker_batch)
+            return self._forward_batch_generation_mlx(batch)
 
         # Fallback to standard path for None batches
         return super().forward_batch_generation(
-            model_worker_batch,
+            batch,
             forward_batch,
             pp_proxy_tensors,
             is_verify,
@@ -125,14 +125,13 @@ class MlxTpModelWorker(TpModelWorker):
             self._mlx_active_rids |= current_rids
 
     def _forward_batch_generation_mlx(
-        self,
+        self, batch: ScheduleBatch
-        model_worker_batch: ModelWorkerBatch,
     ) -> GenerationBatchResult:
         """Run forward pass through the MLX model runner (greedy only)."""
         from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 
-        forward_mode = model_worker_batch.forward_mode
+        forward_mode = batch.forward_mode
-        reqs = model_worker_batch.reqs
+        reqs = batch.reqs
 
         if forward_mode.is_idle():
             return GenerationBatchResult(
@@ -148,9 +147,9 @@ class MlxTpModelWorker(TpModelWorker):
             # Ensure pool is up-to-date before PoolBackedCache reads it
             # for prefix-cached prefills.  Only runs on extend batches.
             self._mlx_runner.flush_all_decode_kv()
-            input_ids_cpu = model_worker_batch.input_ids.cpu().tolist()
+            input_ids_cpu = batch.input_ids.cpu().tolist()
-            out_cache_loc_cpu = model_worker_batch.out_cache_loc.cpu().tolist()
+            out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
-            extend_seq_lens = model_worker_batch.extend_seq_lens
+            extend_seq_lens = batch.extend_lens
 
             offset = 0  # into input_ids_cpu
             slot_offset = 0  # into out_cache_loc_cpu
@@ -226,10 +225,7 @@ class MlxTpModelWorker(TpModelWorker):
             can_run_cuda_graph=False,
         )
 
-    def async_forward_batch_generation_mlx(
+    def async_forward_batch_generation_mlx(self, batch: ScheduleBatch) -> tuple[
-        self,
-        model_worker_batch: ModelWorkerBatch,
-    ) -> tuple[
         Union[mx.array, None],
         list[MlxPendingPrefill],
         list[MlxPendingExtend],
@@ -258,8 +254,8 @@ class MlxTpModelWorker(TpModelWorker):
         """
         self._ensure_mlx_pool_initialized()
 
-        forward_mode = model_worker_batch.forward_mode
+        forward_mode = batch.forward_mode
-        reqs = model_worker_batch.reqs
+        reqs = batch.reqs
 
         if forward_mode.is_idle():
             return None, [], [], None, "idle"
@@ -277,16 +273,13 @@ class MlxTpModelWorker(TpModelWorker):
             # Ensure the pool is up-to-date before any PoolBackedCache
             # reads it for prefix-cached prefills. Mirror the sync path.
             self._mlx_runner.flush_all_decode_kv()
-            return self._async_extend_batch(model_worker_batch)
+            return self._async_extend_batch(batch)
 
         raise ValueError(
             f"MLX async runner does not support forward mode: {forward_mode}"
         )
 
-    def _async_extend_batch(
+    def _async_extend_batch(self, batch: ScheduleBatch) -> tuple[
-        self,
-        model_worker_batch: ModelWorkerBatch,
-    ) -> tuple[
         Union[mx.array, None],
         list[MlxPendingPrefill],
         list[MlxPendingExtend],
@@ -294,10 +287,10 @@ class MlxTpModelWorker(TpModelWorker):
         str,
     ]:
         """Launch each request in an EXTEND batch lazily and kick GPU work."""
-        reqs = model_worker_batch.reqs
+        reqs = batch.reqs
-        input_ids_cpu = model_worker_batch.input_ids.cpu().tolist()
+        input_ids_cpu = batch.input_ids.cpu().tolist()
-        out_cache_loc_cpu = model_worker_batch.out_cache_loc.cpu().tolist()
+        out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
-        extend_seq_lens = model_worker_batch.extend_seq_lens
+        extend_seq_lens = batch.extend_lens
 
         offset = 0
         slot_offset = 0

python/sglang/srt/managers/overlap_utils.py

@@ -9,7 +9,7 @@ from sglang.srt.speculative.spec_utils import spec_need_hidden_states
 from sglang.srt.utils import is_cuda, is_hip
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.scheduler import GenerationBatchResult
     from sglang.srt.speculative.eagle_info import EagleDraftInput
     from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -127,12 +127,12 @@ class FutureMap:
         indices = torch.arange(start, end, dtype=torch.int64, device=self.device)
         return FutureIndices(indices=indices, interval=slice(start, end))
 
-    def resolve_future(self, model_worker_batch: ModelWorkerBatch):
+    def resolve_future(self, batch: ScheduleBatch):
         if self.spec_algo.is_none():
-            _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
+            _resolve_future_token_ids(batch.input_ids, self.token_ids_buf)
         else:
             # TODO(lsyin): write future indices into spec_info.future_indices
-            draft_input: EagleDraftInput = model_worker_batch.spec_info
+            draft_input: EagleDraftInput = batch.spec_info
             if draft_input is None:
                 # FIXME(lsyin): No future exists, only for prefill batch, not compatible with mixed mode
                 return
@@ -140,9 +140,8 @@ class FutureMap:
             # The indices tensor was allocated on the default stream but is
             # used here on the forward stream. Meanwhile, the old spec_info
             # holding this tensor will lose all Python references (replaced at
-            # model_worker_batch.spec_info and batch.spec_info), so the
+            # batch.spec_info), so the caching allocator (torch GC) could
-            # caching allocator (torch GC) could reclaim the memory before
+            # reclaim the memory before the GPU finishes reading it.
-            # the GPU finishes reading it.
             indices.record_stream(torch.get_device_module(self.device).current_stream())
             draft_input.topk_p = self.topk_p_buf[indices]
             draft_input.topk_index = self.topk_index_buf[indices]

python/sglang/srt/managers/schedule_batch.py

@@ -22,17 +22,13 @@ Store information about requests and batches.
 
 The following is the flow of data structures for a batch:
 
-ScheduleBatch -> ModelWorkerBatch -> ForwardBatch
+ScheduleBatch -> ForwardBatch
 
 - ScheduleBatch is managed by `scheduler.py::Scheduler`.
   It contains high-level scheduling data. Most of the data is on the CPU.
-- ModelWorkerBatch is managed by `tp_worker.py::TpModelWorker`.
-  It is a subset of `ScheduleBatch` that only contains data related to the model forward on GPU.
-  It will be transformed from CPU scheduler to GPU model runner.
 - ForwardBatch is managed by `model_runner.py::ModelRunner`.
   It contains low-level tensor data. Most of the data consists of GPU tensors.
-
+  It is constructed directly from a ScheduleBatch by `ForwardBatch.init_new`.
-TODO(lmzheng): ModelWorkerBatch seems a bit redundant and we consider removing it in the future.
 """
 
 import copy
@@ -1430,7 +1426,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     global_num_tokens: Optional[List[int]] = None
     global_num_tokens_for_logprob: Optional[List[int]] = None
     is_extend_in_batch: bool = False
-    all_extend_in_batch: bool = False
+    all_extend_in_batch: bool = False  # plumbing for downstream forks (PR #19639)
     can_run_dp_cuda_graph: bool = False
     tbo_split_seq_index: Optional[int] = None
     global_forward_mode: Optional[ForwardMode] = None
@@ -1470,7 +1466,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     split_prefill_finished: bool = False
     split_forward_count: int = 1
     split_forward_batch: ForwardBatch = None
+
+    # One-shot per-forward overrides; init_new consumes and resets.
     seq_lens_cpu_cache: torch.Tensor = None
+    capture_hidden_mode: Optional[CaptureHiddenMode] = None
+    return_hidden_states_before_norm: bool = False
 
     # Forward-pass metrics
     fpm_start_time: float = 0.0
@@ -2532,89 +2532,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         if self.spec_info:
             self.spec_info.merge_batch(other.spec_info)
 
-    def get_model_worker_batch(
-        self, seq_lens_cpu_cache: Optional[torch.Tensor] = None
-    ) -> ModelWorkerBatch:
-        if self.forward_mode.is_decode_or_idle():
-            extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = None
-        else:
-            extend_seq_lens = self.extend_lens
-            extend_prefix_lens = self.prefix_lens
-            extend_logprob_start_lens = self.extend_logprob_start_lens
-
-        if self.sampling_info:
-            if self.has_grammar:
-                self.sampling_info.grammars = [req.grammar for req in self.reqs]
-            else:
-                self.sampling_info.grammars = None
-
-        seq_lens_cpu = (
-            seq_lens_cpu_cache if seq_lens_cpu_cache is not None else self.seq_lens_cpu
-        )
-
-        return ModelWorkerBatch(
-            forward_mode=self.forward_mode,
-            input_ids=self.input_ids,
-            req_pool_indices=self.req_pool_indices,
-            seq_lens=self.seq_lens,
-            orig_seq_lens=self.orig_seq_lens,
-            out_cache_loc=self.out_cache_loc,
-            seq_lens_cpu=seq_lens_cpu,
-            seq_lens_sum=self.seq_lens_sum,
-            return_logprob=self.return_logprob,
-            top_logprobs_nums=self.top_logprobs_nums,
-            token_ids_logprobs=self.token_ids_logprobs,
-            global_num_tokens=self.global_num_tokens,
-            global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
-            is_extend_in_batch=self.is_extend_in_batch,
-            all_extend_in_batch=self.all_extend_in_batch,
-            can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
-            tbo_split_seq_index=self.tbo_split_seq_index,
-            global_forward_mode=self.global_forward_mode,
-            extend_num_tokens=self.extend_num_tokens,
-            extend_seq_lens=extend_seq_lens,
-            extend_prefix_lens=extend_prefix_lens,
-            extend_logprob_start_lens=extend_logprob_start_lens,
-            multimodal_inputs=self.multimodal_inputs,
-            encoder_cached=self.encoder_cached,
-            encoder_lens=self.encoder_lens,
-            encoder_lens_cpu=self.encoder_lens_cpu,
-            encoder_out_cache_loc=self.encoder_out_cache_loc,
-            lora_ids=[req.lora_id for req in self.reqs],
-            sampling_info=self.sampling_info,
-            input_embeds=self.input_embeds,
-            replace_embeds=self.replace_embeds,
-            replace_positions=self.replace_positions,
-            ne_token_table=self.ne_token_table,
-            token_type_ids=self.token_type_ids,
-            spec_algorithm=self.spec_algorithm,
-            spec_info=self.spec_info,
-            hicache_consumer_index=self.hicache_consumer_index,
-            capture_hidden_mode=(
-                CaptureHiddenMode.FULL
-                if self.return_hidden_states
-                else (
-                    getattr(
-                        self.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL
-                    )
-                    if self.spec_info
-                    else CaptureHiddenMode.NULL
-                )
-            ),
-            extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
-            is_prefill_only=self.is_prefill_only,
-            multi_item_delimiter_indices=self.multi_item_delimiter_indices,
-            dimensions=self.dimensions,
-            return_pooled_hidden_states=self.return_pooled_hidden_states,
-            dllm_block_offsets=[req.dllm_block_offset for req in self.reqs],
-            dllm_config=self.dllm_config,
-            reqs=self.reqs,
-            has_grammar=self.has_grammar,
-            mamba_track_indices=self.mamba_track_indices,
-            mamba_track_mask=self.mamba_track_mask,
-            mamba_track_seqlens=self.mamba_track_seqlens,
-        )
-
     def copy(self):
         # Only contain fields that will be used by process_batch_result.
         # Shallow-copy the reqs list so that in-place mutations (filter_batch,
@@ -2632,8 +2549,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             global_num_tokens=self.global_num_tokens,
             global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
             can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
-            all_extend_in_batch=self.all_extend_in_batch,
             is_extend_in_batch=self.is_extend_in_batch,
+            all_extend_in_batch=self.all_extend_in_batch,
             is_prefill_only=self.is_prefill_only,
             seq_lens_cpu=self.seq_lens_cpu,
             enable_overlap=self.enable_overlap,
@@ -2747,108 +2664,3 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "
             f"#req={(len(self.reqs))})"
         )
-
-
-@dataclasses.dataclass
-class ModelWorkerBatch:
-    # The forward mode
-    forward_mode: ForwardMode
-    # The input ids
-    input_ids: torch.Tensor
-    # The indices of requests in the req_to_token_pool
-    req_pool_indices: torch.Tensor
-    # The sequence length
-    seq_lens: torch.Tensor
-    # The indices of output tokens in the token_to_kv_pool_allocator
-    out_cache_loc: torch.Tensor
-    # The sequence length tensor on CPU
-    seq_lens_cpu: Optional[torch.Tensor]
-    seq_lens_sum: int
-
-    # For logprob
-    return_logprob: bool
-    top_logprobs_nums: Optional[List[int]]
-    token_ids_logprobs: Optional[List[List[int]]]
-
-    # For DP attention
-    global_num_tokens: Optional[List[int]]
-    global_num_tokens_for_logprob: Optional[List[int]]
-    is_extend_in_batch: bool
-    all_extend_in_batch: bool
-    can_run_dp_cuda_graph: bool
-    tbo_split_seq_index: Optional[int]
-    global_forward_mode: Optional[ForwardMode]
-
-    # For extend
-    extend_num_tokens: Optional[int]
-    extend_seq_lens: Optional[List[int]]
-    extend_prefix_lens: Optional[List[int]]
-    extend_logprob_start_lens: Optional[List[int]]
-    extend_input_logprob_token_ids: Optional[torch.Tensor]
-
-    # For multimodal
-    multimodal_inputs: Optional[List[MultimodalInputs]]
-
-    # For encoder-decoder
-    encoder_cached: Optional[List[bool]]
-    encoder_lens: Optional[torch.Tensor]
-    encoder_lens_cpu: Optional[List[int]]
-    encoder_out_cache_loc: Optional[torch.Tensor]
-
-    # For LoRA
-    lora_ids: Optional[List[str]]
-
-    # Sampling info
-    sampling_info: SamplingBatchInfo
-
-    # The original sequence lengths, Qwen-1M related
-    orig_seq_lens: Optional[torch.Tensor] = None
-
-    # The input Embeds
-    input_embeds: Optional[torch.Tensor] = None
-    replace_embeds: Optional[torch.Tensor] = None
-    replace_positions: Optional[torch.Tensor] = None
-
-    # token table for ngram embedding
-    ne_token_table: Optional[torch.Tensor] = None
-
-    # For corss-encoder model
-    token_type_ids: Optional[torch.Tensor] = None
-
-    # Speculative decoding
-    spec_algorithm: SpeculativeAlgorithm = None
-
-    spec_info: Optional[SpecInput] = None
-
-    # If set, the output of the batch contains the hidden states of the run.
-    capture_hidden_mode: CaptureHiddenMode = None
-    hicache_consumer_index: int = -1
-
-    # For matryoshka embeddings
-    dimensions: Optional[list[int]] = None
-
-    # Whether to return pooled hidden states (pre-head transformer output)
-    return_pooled_hidden_states: bool = False
-
-    # Whether this batch is prefill-only (no token generation needed)
-    is_prefill_only: bool = False
-
-    # Pre-computed delimiter indices for multi-item scoring (CPU tensors, one per request)
-    multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None
-
-    # Diffusion LLM
-    dllm_block_offsets: Optional[List[int]] = None
-    dllm_config: Optional[DllmConfig] = None
-
-    # For constrained decoding
-    # FIXME(lsyin): remove this after fully overlap grammar
-    reqs: Optional[List[Req]] = None
-    has_grammar: bool = False
-
-    # For hidden states before normal
-    return_hidden_states_before_norm: bool = False
-
-    # For mamba state tracking
-    mamba_track_indices: Optional[torch.Tensor] = None  # shape: [b], int64
-    mamba_track_mask: Optional[torch.Tensor] = None  # shape: [b], bool
-    mamba_track_seqlens: Optional[torch.Tensor] = None  # shape: [b], int64

python/sglang/srt/managers/scheduler.py

@@ -13,6 +13,7 @@
 # ==============================================================================
 """A scheduler that manages a tensor parallel GPU worker."""
 
+import dataclasses
 import faulthandler
 import logging
 import os
@@ -20,7 +21,7 @@ import signal
 import sys
 import time
 from collections import deque
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, Deque, Dict, List, Optional, Tuple, Union
@@ -159,7 +160,6 @@ from sglang.srt.managers.prefill_delayer import (
 )
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
-    ModelWorkerBatch,
     MultimodalInputs,
     Req,
     ScheduleBatch,
@@ -2992,14 +2992,61 @@ class Scheduler(
         batch.prepare_for_decode()
         return batch
 
-    def record_batch_in_overlap(self, model_worker_batch: ModelWorkerBatch):
+    def record_batch_in_overlap(self, batch: ScheduleBatch):
         # FIXME(lsyin): hacky way to keep a reference to avoid GPU tensors being freed by torch GC
         # NOTE: More Reliable: record all tensors into the forward stream
         # NOTE: - for all future tensors, we shall always read from future map
         #       - for all non-future tensors (produced only by schedule stream),
         #       we shall keep its reference not being release during all the forwarding pass
+        # Snapshot all fields: spec V2 rebinds seq_lens / spec_info mid-forward.
+        attr_snapshot = [
+            getattr(batch, f.name, None) for f in dataclasses.fields(batch)
+        ]
         self.batch_record_ct = (self.batch_record_ct + 1) % 2
-        self.batch_record_buf[self.batch_record_ct] = model_worker_batch
+        # List (not tuple) so that workers can register additional refs via
+        # GenerationBatchResult.extra_keep_alive_refs after forward returns.
+        self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot]
+
+    @contextmanager
+    def _overlap_forward_isolation(self, batch: ScheduleBatch):
+        """Make SB transactional across one overlap forward.
+
+        1. Snapshot SB fields so V2's mid-forward mutations (forward_mode /
+           input_ids / seq_lens / spec_info / ...) can be undone. V1 / non-spec
+           only need sampling_info restored - V1 carries spec_info forward as
+           next-iter draft input.
+        2. Substitute sampling_info with a forward-only copy (orchestrator=None,
+           shares the pre-accumulated penalty buffer) so V2's multiple init_new
+           calls don't double-accumulate penalties.
+        3. Pin (batch, snapshot) into batch_record_buf for 2 iters so GPU
+           tensors in the snapshot survive the caching allocator past the
+           forward stream. Must run AFTER the sampling_info swap so the
+           forward-only copy gets pinned.
+        """
+        # 1. snapshot
+        snapshot_v2_full = batch.is_spec_v2
+        sched_snapshot = (
+            {f.name: getattr(batch, f.name) for f in dataclasses.fields(batch)}
+            if snapshot_v2_full
+            else None
+        )
+        sched_sampling_info = batch.sampling_info
+
+        # 2. sampling_info substitute
+        if sched_sampling_info is not None:
+            batch.sampling_info = sched_sampling_info.copy_for_forward()
+
+        # 3. pin for 2-iter tensor lifetime
+        self.record_batch_in_overlap(batch)
+
+        try:
+            yield
+        finally:
+            if snapshot_v2_full:
+                for name, value in sched_snapshot.items():
+                    setattr(batch, name, value)
+            else:
+                batch.sampling_info = sched_sampling_info
 
     def run_batch(
         self,
@@ -3022,42 +3069,33 @@ class Scheduler(
 
         # Run forward
         if self.is_generation:
-            if self.spec_algorithm.is_none() or self.enable_overlap:
-                # In most cases, we use the model worker batch to run the forward.
-                worker_batch_or_batch = batch.get_model_worker_batch()
-            else:
-                # In speculative decoding v1 (non-overlap) case, we use the batch directly.
-                # TODO(lsyin): delete this branch after unifying the abstraction.
-                worker_batch_or_batch = batch
-
             if self.enable_overlap:
-                model_worker_batch = worker_batch_or_batch
+                with self._overlap_forward_isolation(batch):
-                self.record_batch_in_overlap(model_worker_batch)
+                    bs = len(batch.seq_lens)
+                    future_indices = self.future_map.alloc_future_indices(bs)
 
-                # Sampling info will be modified during forward, so we store a copy.
+                    with self.forward_stream_ctx:
-                model_worker_batch.sampling_info = (
+                        self.forward_stream.wait_stream(self.schedule_stream)
-                    model_worker_batch.sampling_info.copy_for_forward()
+                        self.future_map.resolve_future(batch)
-                )
+                        # FIXME: pp is not compatible with overlap
-                bs = len(model_worker_batch.seq_lens)
+                        batch_result = self.model_worker.forward_batch_generation(batch)
-                future_indices = self.future_map.alloc_future_indices(bs)
+                        # Park any refs the worker wants kept alive 2 iters
-
+                        # (cross-stream tensor lifetime; pinned in the same
-                with self.forward_stream_ctx:
+                        # ring slot as the SB attr snapshot).
-                    self.forward_stream.wait_stream(self.schedule_stream)
+                        if batch_result.extra_keep_alive_refs:
-                    self.future_map.resolve_future(model_worker_batch)
+                            self.batch_record_buf[self.batch_record_ct].extend(
-                    batch_result = self.model_worker.forward_batch_generation(
+                                batch_result.extra_keep_alive_refs
-                        model_worker_batch
+                            )
-                        # here pp is not compatible with overlap
+                        # FIXME(lsyin): maybe move this to forward_batch_generation
-                    )
+                        batch_result.copy_done = self.device_module.Event()
-                    # FIXME(lsyin): maybe move this to forward_batch_generation
+                        if batch_result.delay_sample_func is None:
-                    batch_result.copy_done = self.device_module.Event()
+                            self.future_map.store_to_map(future_indices, batch_result)
-                    if batch_result.delay_sample_func is None:
+                            batch_result.copy_to_cpu(
-                        self.future_map.store_to_map(future_indices, batch_result)
+                                return_logprob=batch.return_logprob,
-                        batch_result.copy_to_cpu(
+                                return_hidden_states=batch.return_hidden_states,
-                            return_logprob=batch.return_logprob,
+                            )
-                            return_hidden_states=batch.return_hidden_states,
+                        else:
-                        )
+                            batch_result.future_indices = future_indices
-                    else:
-                        batch_result.future_indices = future_indices
 
                 # FIXME(lsyin): move this assignment elsewhere
                 future_indices_or_next_token_ids = -future_indices.indices
@@ -3082,7 +3120,7 @@ class Scheduler(
                     else {}
                 )
                 batch_result = self.model_worker.forward_batch_generation(
-                    worker_batch_or_batch, **kwargs
+                    batch, **kwargs
                 )
                 future_indices_or_next_token_ids = batch_result.next_token_ids
                 self.update_cache_from_scheduler(batch, batch_result)
@@ -3109,24 +3147,18 @@ class Scheduler(
 
             ret = batch_result
         else:  # embedding or reward model
-            model_worker_batch = batch.get_model_worker_batch()
-
             if self.enable_overlap:
-                self.record_batch_in_overlap(model_worker_batch)
+                self.record_batch_in_overlap(batch)
                 with self.forward_stream_ctx:
                     self.forward_stream.wait_stream(self.schedule_stream)
-                    pooler_output = self.tp_worker.forward_batch_embedding(
+                    pooler_output = self.tp_worker.forward_batch_embedding(batch)
-                        model_worker_batch
-                    )
                     ret = EmbeddingBatchResult(
                         embeddings=pooler_output.embeddings,
                         pooled_hidden_states=pooler_output.pooled_hidden_states,
                     )
                     ret.copy_to_cpu()
             else:
-                pooler_output = self.tp_worker.forward_batch_embedding(
+                pooler_output = self.tp_worker.forward_batch_embedding(batch)
-                    model_worker_batch
-                )
                 ret = EmbeddingBatchResult(
                     embeddings=pooler_output.embeddings,
                     pooled_hidden_states=pooler_output.pooled_hidden_states,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -639,9 +639,8 @@ class SchedulerPPMixin:
 
                 start = time.perf_counter()
                 batch.prepare_for_extend()
-                model_worker_batch = batch.get_model_worker_batch()
 
-                forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
+                forward_batch = ForwardBatch.init_new(batch, model_runner)
                 set_is_extend_in_batch(batch.forward_mode.is_extend())
 
                 _ = model_runner.forward(

python/sglang/srt/managers/tp_worker.py

@@ -36,7 +36,7 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
 )
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -209,8 +209,8 @@ class BaseTpWorker(ABC):
         )
         return result
 
-    def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
+    def forward_batch_embedding(self, batch: ScheduleBatch):
-        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
+        forward_batch = ForwardBatch.init_new(batch, self.model_runner)
         output = self.model_runner.forward(forward_batch).logits_output
         return output  # Returns EmbeddingPoolerOutput
 
@@ -446,7 +446,7 @@ class TpModelWorker(BaseTpWorker):
 
     def forward_batch_generation(
         self,
-        model_worker_batch: ModelWorkerBatch,
+        batch: Optional[ScheduleBatch],
         forward_batch: Optional[ForwardBatch] = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         is_verify: bool = False,
@@ -455,12 +455,12 @@ class TpModelWorker(BaseTpWorker):
         # 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 model worker batch
+        # Get forward batch from schedule batch
-        if model_worker_batch is not None:
+        if batch is not None:
             # update the consumer index of hicache to the running batch
-            self.set_hicache_consumer(model_worker_batch.hicache_consumer_index)
+            self.set_hicache_consumer(batch.hicache_consumer_index)
 
-            forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
+            forward_batch = ForwardBatch.init_new(batch, self.model_runner)
         else:
             # FIXME(lsyin): unify the interface of forward_batch
             assert forward_batch is not None
@@ -490,7 +490,7 @@ class TpModelWorker(BaseTpWorker):
             if (
                 self.enable_overlap
                 and not self.enable_spec
-                and model_worker_batch.sampling_info.grammars is not None
+                and forward_batch.sampling_info.grammars is not None
             ):
 
                 def sample_batch_func():
@@ -502,7 +502,7 @@ class TpModelWorker(BaseTpWorker):
                 batch_result.delay_sample_func = sample_batch_func
                 return batch_result
 
-            if not model_worker_batch.is_prefill_only:
+            if not forward_batch.is_prefill_only:
                 # For normal requests, sample the next token ids.
                 batch_result.next_token_ids = self.model_runner.sample(
                     logits_output, forward_batch
@@ -511,17 +511,17 @@ class TpModelWorker(BaseTpWorker):
                 # For prefill-only requests, create dummy token IDs on CPU
                 # The size should match the batch size (number of sequences), not total tokens
                 batch_result.next_token_ids = torch.zeros(
-                    len(model_worker_batch.seq_lens),
+                    len(forward_batch.seq_lens),
                     dtype=torch.long,
-                    device=model_worker_batch.input_ids.device,
+                    device=forward_batch.input_ids.device,
                 )
                 if (
-                    model_worker_batch.return_logprob
+                    forward_batch.return_logprob
                     and logits_output.next_token_logits is not None
                 ):
                     # NOTE: Compute logprobs without full sampling
                     self.model_runner.compute_logprobs_only(
-                        logits_output, model_worker_batch
+                        logits_output, forward_batch
                     )
 
             return batch_result
@@ -540,19 +540,17 @@ class TpModelWorker(BaseTpWorker):
 
     def forward_batch_split_prefill(self, batch: ScheduleBatch):
         if batch.split_index == 0:
-            model_worker_batch = batch.get_model_worker_batch()
+            forward_batch = ForwardBatch.init_new(batch, self.model_runner)
-            forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
             batch.split_forward_batch = forward_batch
-            batch.seq_lens_cpu_cache = model_worker_batch.seq_lens_cpu
-        else:
-            model_worker_batch = batch.get_model_worker_batch(batch.seq_lens_cpu_cache)
 
         out = self.model_runner.forward(
             batch.split_forward_batch, split_forward_count=batch.split_forward_count
         )
         logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
         if logits_output:
-            next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+            next_token_ids = self.model_runner.sample(
+                logits_output, batch.split_forward_batch
+            )
         else:
             next_token_ids = None
         batch_result = GenerationBatchResult(

python/sglang/srt/managers/utils.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Union
 
 import torch
 
@@ -48,6 +48,11 @@ class GenerationBatchResult:
     # relay path: forward stream -> next step forward
     next_draft_input: Optional[EagleDraftInput] = None
 
+    # Refs the worker wants scheduler to keep alive for the same 2-iter window
+    # as batch_record_buf. Used for cross-stream tensor lifetime (e.g. a spec
+    # V2 verify ForwardBatch whose tensors must outlive mid-iter SB rebinds).
+    extra_keep_alive_refs: Optional[List[Any]] = None
+
     # Routed experts: pending async D2H for overlap scheduling
     routed_experts_output: Optional[TopkCaptureOutput] = None
     indexer_topk_output: Optional[TopkCaptureOutput] = None

python/sglang/srt/model_executor/forward_batch_info.py

@@ -16,15 +16,13 @@ Store information about a forward batch.
 
 The following is the flow of data structures for a batch:
 
-ScheduleBatch -> ModelWorkerBatch -> ForwardBatch
+ScheduleBatch -> ForwardBatch
 
 - ScheduleBatch is managed by `scheduler.py::Scheduler`.
   It contains high-level scheduling data. Most of the data is on the CPU.
-- ModelWorkerBatch is managed by `tp_worker.py::TpModelWorker`.
-  It is a subset of `ScheduleBatch` that only contains data related to the model forward on GPU.
-  It will be transformed from CPU scheduler to GPU model runner.
 - ForwardBatch is managed by `model_runner.py::ModelRunner`.
   It contains low-level tensor data. Most of the data consists of GPU tensors.
+  It is constructed directly from a ScheduleBatch by `ForwardBatch.init_new`.
 """
 
 from __future__ import annotations
@@ -68,7 +66,7 @@ if TYPE_CHECKING:
     from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
     from sglang.srt.layers.logits_processor import LogitsProcessorOutput
     from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
-    from sglang.srt.managers.schedule_batch import ModelWorkerBatch, MultimodalInputs
+    from sglang.srt.managers.schedule_batch import MultimodalInputs, ScheduleBatch
     from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
     from sglang.srt.model_executor.model_runner import ModelRunner
     from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -390,6 +388,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     dp_local_num_tokens: Optional[torch.Tensor] = None  # cached info at runtime
     global_dp_buffer_len: Optional[int] = None
     is_extend_in_batch: bool = False
+    # Mirrors ScheduleBatch.all_extend_in_batch; kept for downstream forks.
     all_extend_in_batch: bool = False
     can_run_dp_cuda_graph: bool = False
     global_forward_mode: Optional[ForwardMode] = None
@@ -443,9 +442,63 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     @classmethod
     def init_new(
         cls,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         model_runner: ModelRunner,
     ):
+        # Consume one-shot per-forward overrides from SB; reset to defaults so
+        # the next forward on the same SB starts clean. See SB field comment
+        # for the contract.
+        capture_hidden_mode = batch.capture_hidden_mode
+        batch.capture_hidden_mode = None
+        seq_lens_cpu_cache = batch.seq_lens_cpu_cache
+        batch.seq_lens_cpu_cache = None
+        return_hidden_states_before_norm = batch.return_hidden_states_before_norm
+        batch.return_hidden_states_before_norm = False
+
+        # capture_hidden_mode default: derive from SB.return_hidden_states /
+        # spec_info.capture_hidden_mode when caller did not override.
+        if capture_hidden_mode is None:
+            if batch.return_hidden_states:
+                capture_hidden_mode = CaptureHiddenMode.FULL
+            elif batch.spec_info is not None:
+                capture_hidden_mode = getattr(
+                    batch.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL
+                )
+            else:
+                capture_hidden_mode = CaptureHiddenMode.NULL
+
+        # extend-mode-only fields are None on decode/idle
+        if batch.forward_mode.is_decode_or_idle():
+            extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = None
+        else:
+            extend_seq_lens = batch.extend_lens
+            extend_prefix_lens = batch.prefix_lens
+            extend_logprob_start_lens = batch.extend_logprob_start_lens
+
+        # Mirror the grammars-population behavior previously done in
+        # ScheduleBatch.get_model_worker_batch.
+        if batch.sampling_info is not None:
+            if batch.has_grammar:
+                batch.sampling_info.grammars = [req.grammar for req in batch.reqs]
+            else:
+                batch.sampling_info.grammars = None
+
+        # ScheduleBatch.sampling_info is already swapped to the forward-only
+        # copy by Scheduler.run_batch under overlap mode (see save/restore
+        # block there). Use it directly.
+        if seq_lens_cpu_cache is not None:
+            # Stale-cache guard: shape must match current GPU seq_lens. Mismatch
+            # means caller forgot to refresh the override after batch size
+            # changed (e.g. filter/merge_batch); using a stale cache would
+            # propagate wrong CPU mirror to downstream DP / cudagraph logic.
+            assert seq_lens_cpu_cache.shape == batch.seq_lens.shape, (
+                f"seq_lens_cpu_cache shape {seq_lens_cpu_cache.shape} != "
+                f"seq_lens {batch.seq_lens.shape}; stale override on batch?"
+            )
+            seq_lens_cpu = seq_lens_cpu_cache
+        else:
+            seq_lens_cpu = batch.seq_lens_cpu
+
         ret = cls(
             forward_mode=batch.forward_mode,
             batch_size=len(batch.seq_lens),
@@ -462,7 +515,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             encoder_lens_cpu=batch.encoder_lens_cpu,
             encoder_out_cache_loc=batch.encoder_out_cache_loc,
             seq_lens_sum=batch.seq_lens_sum,
-            seq_lens_cpu=batch.seq_lens_cpu,
+            seq_lens_cpu=seq_lens_cpu,
             orig_seq_lens=batch.orig_seq_lens,
             return_logprob=batch.return_logprob,
             top_logprobs_nums=batch.top_logprobs_nums,
@@ -473,22 +526,22 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             global_forward_mode=batch.global_forward_mode,
             is_prefill_only=batch.is_prefill_only,
             multi_item_delimiter_indices=batch.multi_item_delimiter_indices,
-            lora_ids=batch.lora_ids,
+            lora_ids=[req.lora_id for req in batch.reqs],
             sampling_info=batch.sampling_info,
             req_to_token_pool=model_runner.req_to_token_pool,
             token_to_kv_pool=model_runner.token_to_kv_pool,
             attn_backend=model_runner.attn_backend,
             spec_algorithm=batch.spec_algorithm,
             spec_info=batch.spec_info,
-            capture_hidden_mode=batch.capture_hidden_mode,
+            capture_hidden_mode=capture_hidden_mode,
             input_embeds=batch.input_embeds,
             replace_embeds=batch.replace_embeds,
             replace_positions=batch.replace_positions,
             token_type_ids=batch.token_type_ids,
             tbo_split_seq_index=batch.tbo_split_seq_index,
             dimensions=batch.dimensions,
-            return_hidden_states_before_norm=batch.return_hidden_states_before_norm,
             return_pooled_hidden_states=batch.return_pooled_hidden_states,
+            return_hidden_states_before_norm=return_hidden_states_before_norm,
             rids=[req.rid for req in batch.reqs],
         )
         device = model_runner.device
@@ -542,7 +595,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             ret.positions = torch.tensor(
                 [
                     i
-                    for block_offset in batch.dllm_block_offsets
+                    for block_offset in (req.dllm_block_offset for req in batch.reqs)
                     for i in range(block_offset, block_offset + block_size)
                 ],
                 dtype=positions_dtype,
@@ -558,13 +611,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             if ret.positions is None:
                 ret.positions = clamp_position(batch.seq_lens)
         else:
-            assert isinstance(batch.extend_seq_lens, list)
+            assert isinstance(extend_seq_lens, list)
-            assert isinstance(batch.extend_prefix_lens, list)
+            assert isinstance(extend_prefix_lens, list)
-            ret.extend_seq_lens = torch.tensor(
+            ret.extend_seq_lens = torch.tensor(extend_seq_lens, dtype=torch.int32).to(
-                batch.extend_seq_lens, dtype=torch.int32
+                device, non_blocking=True
-            ).to(device, non_blocking=True)
+            )
             ret.extend_prefix_lens = torch.tensor(
-                batch.extend_prefix_lens, dtype=torch.int32
+                extend_prefix_lens, dtype=torch.int32
             ).to(device, non_blocking=True)
             ret.extend_num_tokens = batch.extend_num_tokens
             positions, ret.extend_start_loc = compute_position(
@@ -575,12 +628,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             )
             if ret.positions is None:
                 ret.positions = positions
-            ret.extend_prefix_lens_cpu = batch.extend_prefix_lens
+            ret.extend_prefix_lens_cpu = extend_prefix_lens
-            ret.extend_seq_lens_cpu = batch.extend_seq_lens
+            ret.extend_seq_lens_cpu = extend_seq_lens
-            ret.extend_logprob_start_lens_cpu = batch.extend_logprob_start_lens
+            ret.extend_logprob_start_lens_cpu = extend_logprob_start_lens
 
         if model_runner.use_ngram_embedding:
-            ret._init_ngram_embedding_info(batch, model_runner, device)
+            ret._init_ngram_embedding_info(batch, device)
 
         if model_runner.model_is_mrope:
             if (
@@ -681,9 +734,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             or self.contains_image_inputs()
         )
 
-    def _init_ngram_embedding_info(
+    def _init_ngram_embedding_info(self, batch: ScheduleBatch, device: torch.device):
-        self, batch: ModelWorkerBatch, model_runner: ModelRunner, device: torch.device
-    ):
         if self.forward_mode.is_decode():
             column_starts, req_lens = self.seq_lens - 1, 1
         else:
@@ -697,7 +748,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         )
 
     def compute_spec_mrope_positions(
-        self, model_runner: ModelRunner, batch: ModelWorkerBatch
+        self, model_runner: ModelRunner, batch: ScheduleBatch
     ):
         # TODO support batched deltas
         batch_size = self.seq_lens.shape[0]
@@ -708,7 +759,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             mrope_deltas = []
             extend_lens = []
             for batch_idx in range(batch_size):
-                extend_seq_len = batch.extend_seq_lens[batch_idx]
+                extend_seq_len = batch.extend_lens[batch_idx]
                 extend_lens.append(extend_seq_len)
                 mrope_delta = (
                     torch.zeros(1, dtype=torch.int64)
@@ -761,9 +812,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         mrope_positions = mm_input.mrope_position_delta_repeated_cache + seq_len
         return mrope_positions
 
-    def _compute_mrope_positions(
+    def _compute_mrope_positions(self, model_runner: ModelRunner, batch: ScheduleBatch):
-        self, model_runner: ModelRunner, batch: ModelWorkerBatch
-    ):
         # batch_size * [3 * seq_len]
         batch_size = self.seq_lens_cpu.shape[0]
         mrope_positions_list = [[]] * batch_size
@@ -787,8 +836,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                     mrope_positions_list[batch_idx] = mrope_positions
             elif self.forward_mode.is_extend(include_draft_extend_v2=True):
                 extend_seq_len, extend_prefix_len = (
-                    batch.extend_seq_lens[batch_idx],
+                    batch.extend_lens[batch_idx],
-                    batch.extend_prefix_lens[batch_idx],
+                    batch.prefix_lens[batch_idx],
                 )
                 if (
                     mm_input is None

python/sglang/srt/speculative/dflash_worker.py

@@ -1,12 +1,12 @@
 import logging
 import math
 from copy import deepcopy
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
 from sglang.srt.distributed import get_tp_group
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.mem_cache.common import get_last_loc
@@ -1109,26 +1109,16 @@ class DFlashWorker:
         )
 
     def forward_batch_generation(
-        self,
+        self, batch: ScheduleBatch, **kwargs
-        batch: Union[ScheduleBatch, ModelWorkerBatch],
-        **kwargs,
     ) -> GenerationBatchResult:
         if getattr(batch, "return_logprob", False):
             raise RuntimeError(
                 "Invariant broken: DFLASH batch requested return_logprob, but scheduler should have rejected this request."
             )
 
-        if isinstance(batch, ModelWorkerBatch):
-            # Should not happen for spec-v1 (non-overlap) scheduling, but keep a sane fallback.
-            return self.target_worker.forward_batch_generation(batch, **kwargs)
-
         if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
-            model_worker_batch = batch.get_model_worker_batch()
+            batch.capture_hidden_mode = CaptureHiddenMode.FULL
-            model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
+            batch_result = self.target_worker.forward_batch_generation(batch, **kwargs)
-
-            batch_result = self.target_worker.forward_batch_generation(
-                model_worker_batch, **kwargs
-            )
             logits_output, next_token_ids = (
                 batch_result.logits_output,
                 batch_result.next_token_ids,
@@ -1139,12 +1129,9 @@ class DFlashWorker:
                     "Make sure the target model has DFlash layers-to-capture configured."
                 )
 
-            if (
+            if batch.extend_lens is None or batch.prefix_lens is None:
-                model_worker_batch.extend_seq_lens is None
-                or model_worker_batch.extend_prefix_lens is None
-            ):
                 raise RuntimeError(
-                    "DFLASH expected extend_seq_lens / extend_prefix_lens to be populated in extend mode, but got None."
+                    "DFLASH expected extend_lens / prefix_lens to be populated in extend mode, but got None."
                 )
 
             # Materialize the prompt tokens into the draft KV cache immediately. This is required
@@ -1158,9 +1145,7 @@ class DFlashWorker:
                     return x if x.dtype == torch.int32 else x.to(torch.int32)
                 return torch.tensor(x, dtype=torch.int32, device=device)
 
-            extend_seq_lens = _to_int32_device_tensor(
+            extend_seq_lens = _to_int32_device_tensor(batch.extend_lens)
-                model_worker_batch.extend_seq_lens
-            )
             draft_input = DFlashDraftInput(
                 bonus_tokens=next_token_ids.to(torch.int64),
                 target_hidden=logits_output.hidden_states,
@@ -1168,7 +1153,7 @@ class DFlashWorker:
                 draft_seq_lens=(
                     torch.zeros_like(extend_seq_lens)
                     if self.use_compact_draft_cache
-                    else _to_int32_device_tensor(model_worker_batch.extend_prefix_lens)
+                    else _to_int32_device_tensor(batch.prefix_lens)
                 ),
             )
             self._append_target_hidden_to_draft_kv(batch, draft_input)
@@ -1191,9 +1176,8 @@ class DFlashWorker:
 
         self._prepare_for_speculative_decoding(batch, draft_input)
 
-        model_worker_batch = batch.get_model_worker_batch()
+        assert batch.forward_mode.is_target_verify()
-        assert model_worker_batch.forward_mode.is_target_verify()
+        verify_input = batch.spec_info
-        verify_input = model_worker_batch.spec_info
         assert isinstance(verify_input, DFlashVerifyInput)
         need_mamba_verify_commit = hasattr(
             self.target_worker.model_runner.attn_backend,
@@ -1204,7 +1188,7 @@ class DFlashWorker:
         )
 
         batch_result = self.target_worker.forward_batch_generation(
-            model_worker_batch, is_verify=True, **kwargs
+            batch, is_verify=True, **kwargs
         )
         logits_output, can_run_cuda_graph = (
             batch_result.logits_output,

python/sglang/srt/speculative/eagle_info_v2.py

@@ -14,7 +14,7 @@ from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.utils import get_alloc_len_per_decode
 from sglang.srt.mem_cache.common import (
     alloc_paged_token_slots_extend,
@@ -177,7 +177,7 @@ class EagleDraftInputV2Mixin:
     def prepare_for_v2_draft(
         self: EagleDraftInput,
         req_to_token_pool: ReqToTokenPool,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         cuda_graph_runner: EAGLEDraftCudaGraphRunner,
         draft_model_runner: ModelRunner,
         topk: int,
@@ -211,15 +211,15 @@ class EagleDraftInputV2Mixin:
             if draft_model_runner.spec_algorithm.is_standalone()
             else CaptureHiddenMode.LAST
         )
-        batch.capture_hidden_mode = capture_mode
         self.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
+        batch.capture_hidden_mode = capture_mode
         forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
         can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
         return forward_batch, can_cuda_graph
 
     def prepare_for_extend_to_fill_draft_kvcache(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         predict: torch.Tensor,
         num_draft_tokens: int,
         draft_model_runner: Any,
@@ -233,20 +233,20 @@ class EagleDraftInputV2Mixin:
         batch.seq_lens = batch.seq_lens + num_draft_tokens
         batch.seq_lens_cpu = batch.seq_lens_cpu + num_draft_tokens
         batch.seq_lens_sum += extend_num_tokens
-        batch.extend_seq_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))]
+        batch.extend_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))]
-        batch.extend_prefix_lens = seq_lens_cpu_.tolist()
+        batch.prefix_lens = seq_lens_cpu_.tolist()
         batch.extend_num_tokens = extend_num_tokens
         capture_mode = (
             CaptureHiddenMode.NULL
             if draft_model_runner.spec_algorithm.is_standalone()
             else CaptureHiddenMode.FULL
         )
-        batch.capture_hidden_mode = capture_mode
         batch.forward_mode = (
             ForwardMode.IDLE
             if batch.forward_mode.is_idle()
             else ForwardMode.DRAFT_EXTEND_V2
         )
+        batch.capture_hidden_mode = capture_mode
         forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
         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:
@@ -259,7 +259,7 @@ class EagleVerifyInputV2Mixin:
     def prepare_for_v2_verify(
         self: EagleVerifyInput,
         req_to_token_pool: ReqToTokenPool,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         target_worker: TpModelWorker,
     ):
         if not batch.forward_mode.is_idle():
@@ -332,7 +332,7 @@ class EagleVerifyInputV2Mixin:
 
     def sample(
         self: EagleVerifyInput,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         logits_output: LogitsProcessorOutput,
         vocab_mask: torch.Tensor = None,
     ):

python/sglang/srt/speculative/eagle_worker.py

@@ -583,14 +583,13 @@ class EAGLEWorker(TpModelWorker):
         """
         # Forward with the target model and get hidden states.
         # We need the full hidden states to prefill the KV cache of the draft model.
-        model_worker_batch = batch.get_model_worker_batch()
         capture_mode = (
             CaptureHiddenMode.NULL
             if self.speculative_algorithm.is_standalone()
             else CaptureHiddenMode.FULL
         )
-        model_worker_batch.capture_hidden_mode = capture_mode
+        batch.capture_hidden_mode = capture_mode
-        batch_result = self.target_worker.forward_batch_generation(model_worker_batch)
+        batch_result = self.target_worker.forward_batch_generation(batch)
         logits_output, next_token_ids = (
             batch_result.logits_output,
             batch_result.next_token_ids,
@@ -598,7 +597,7 @@ class EAGLEWorker(TpModelWorker):
         return (
             logits_output,
             next_token_ids,
-            model_worker_batch.seq_lens_cpu,
+            batch.seq_lens_cpu,
             batch_result.can_run_cuda_graph,
         )
 
@@ -776,11 +775,8 @@ class EAGLEWorker(TpModelWorker):
         batch.return_hidden_states = False
 
         # Get forward batch
-        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
-        assert model_worker_batch.capture_hidden_mode == draft_capture_mode
+        assert forward_batch.capture_hidden_mode == draft_capture_mode
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.draft_model_runner
-        )
         can_cuda_graph = self.cuda_graph_runner and self.cuda_graph_runner.can_run(
             forward_batch
         )
@@ -944,11 +940,6 @@ class EAGLEWorker(TpModelWorker):
             else ForwardMode.IDLE
         )
 
-        model_worker_batch = batch.get_model_worker_batch(
-            seq_lens_cpu_cache=spec_info.seq_lens_cpu
-        )
-        assert model_worker_batch.capture_hidden_mode == spec_info.capture_hidden_mode
-
         if batch.has_grammar:
             retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
             retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
@@ -957,8 +948,9 @@ class EAGLEWorker(TpModelWorker):
             ).cpu()
 
         # Forward
+        batch.seq_lens_cpu_cache = spec_info.seq_lens_cpu
         batch_result = self.target_worker.forward_batch_generation(
-            model_worker_batch, is_verify=True
+            batch, is_verify=True
         )
         logits_output, can_run_cuda_graph = (
             batch_result.logits_output,
@@ -1136,12 +1128,8 @@ class EAGLEWorker(TpModelWorker):
             else CaptureHiddenMode.LAST
         )
         batch.spec_info.capture_hidden_mode = capture_mode
-        model_worker_batch = batch.get_model_worker_batch(
+        batch.seq_lens_cpu_cache = seq_lens_cpu
-            seq_lens_cpu_cache=seq_lens_cpu
+        forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
-        )
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.draft_model_runner
-        )
         forward_batch.return_logprob = False
         if mm_input_embeds is not None:
             forward_batch.mm_input_embeds = mm_input_embeds
@@ -1199,14 +1187,11 @@ class EAGLEWorker(TpModelWorker):
 
         batch.return_hidden_states = False
         # Verify-time construction of EagleDraftExtendInput uses the dataclass
-        # default (LAST); the worker overrides here so get_model_worker_batch()
+        # default (LAST); override here so ForwardBatch.init_new picks up the
-        # propagates the correct mode (NULL for STANDALONE).
+        # correct mode (NULL for STANDALONE).
         draft_extend_input.capture_hidden_mode = draft_extend_capture_mode
-        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
-        assert model_worker_batch.capture_hidden_mode == draft_extend_capture_mode
+        assert forward_batch.capture_hidden_mode == draft_extend_capture_mode
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.draft_model_runner
-        )
         if forward_batch.seq_lens_cpu is not None:
             forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()
         else:

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -27,7 +27,7 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
 )
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
@@ -59,6 +59,8 @@ from sglang.srt.speculative.spec_utils import (
     load_token_map,
     maybe_detect_nan,
     maybe_detect_oob,
+    record_stream_each,
+    record_stream_for_v2_verify,
     select_top_k_tokens,
 )
 from sglang.srt.utils.common import (
@@ -335,11 +337,11 @@ class EagleDraftWorker(BaseDraftWorker):
                 f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB.",
             )
 
-    def draft(self, model_worker_batch: ModelWorkerBatch):
+    def draft(self, batch: ScheduleBatch):
-        draft_input: EagleDraftInput = model_worker_batch.spec_info
+        draft_input: EagleDraftInput = batch.spec_info
         forward_batch, can_cuda_graph = draft_input.prepare_for_v2_draft(
             self.req_to_token_pool,
-            model_worker_batch,
+            batch,
             self.cuda_graph_runner,
             self.draft_runner,
             self.topk,
@@ -363,7 +365,7 @@ class EagleDraftWorker(BaseDraftWorker):
                 forward_batch
             )
 
-        if model_worker_batch.forward_mode.is_idle():
+        if batch.forward_mode.is_idle():
             return EagleVerifyInput.create_idle_input(
                 self.topk,
                 self.speculative_num_steps,
@@ -388,8 +390,8 @@ class EagleDraftWorker(BaseDraftWorker):
             parent_list,
             top_scores_index,
             draft_tokens,
-            model_worker_batch.seq_lens,
+            batch.seq_lens,
-            model_worker_batch.seq_lens_sum,
+            batch.seq_lens_sum,
             self.topk,
             self.speculative_num_steps,
             self.speculative_num_draft_tokens,
@@ -512,7 +514,7 @@ class EagleDraftWorker(BaseDraftWorker):
 
     def _draft_extend_for_prefill(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         target_hidden_states: torch.Tensor,
         next_token_ids: torch.Tensor,
         mm_input_embeds: Optional[torch.Tensor] = None,
@@ -528,7 +530,7 @@ class EagleDraftWorker(BaseDraftWorker):
         # Construct input_ids
         if not batch.forward_mode.is_idle():
             pt = 0
-            for i, extend_len in enumerate(batch.extend_seq_lens):
+            for i, extend_len in enumerate(batch.extend_lens):
                 input_ids = batch.input_ids[pt : pt + extend_len]
                 batch.input_ids[pt : pt + extend_len] = torch.cat(
                     (input_ids[1:], next_token_ids[i].reshape(1))
@@ -547,7 +549,15 @@ class EagleDraftWorker(BaseDraftWorker):
 
         batch.spec_info = next_draft_input
 
-        # Run forward
+        # Run forward (LAST mode: only the final hidden state per request,
+        # to feed the next draft step which expects [bs, hidden_dim]).
+        # STANDALONE skips hidden states end-to-end.
+        capture_hidden_mode = (
+            CaptureHiddenMode.NULL
+            if self.speculative_algorithm.is_standalone()
+            else CaptureHiddenMode.LAST
+        )
+        batch.capture_hidden_mode = capture_hidden_mode
         forward_batch = ForwardBatch.init_new(batch, self.draft_runner)
         forward_batch.return_logprob = False
         if mm_input_embeds is not None:
@@ -564,7 +574,7 @@ class EagleDraftWorker(BaseDraftWorker):
         return next_draft_input
 
     def _draft_extend_for_decode(
-        self, batch: ModelWorkerBatch, batch_result: GenerationBatchResult
+        self, batch: ScheduleBatch, batch_result: GenerationBatchResult
     ):
         # Batch 2: Draft extend
         draft_input = EagleDraftInput(
@@ -740,29 +750,18 @@ class EAGLEWorkerV2(BaseSpecWorker):
         # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
         pass
 
-    def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
+    def forward_batch_generation(self, batch: ScheduleBatch):
-        if (
+        if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
-            model_worker_batch.forward_mode.is_extend()
-            or model_worker_batch.is_extend_in_batch
-        ):
             # Target prefill
             target_capture_mode = (
                 CaptureHiddenMode.NULL
                 if self.speculative_algorithm.is_standalone()
                 else CaptureHiddenMode.FULL
             )
-            model_worker_batch.capture_hidden_mode = target_capture_mode
+            batch.capture_hidden_mode = target_capture_mode
-            batch_output = self.target_worker.forward_batch_generation(
+            batch_output = self.target_worker.forward_batch_generation(batch)
-                model_worker_batch
-            )
 
             # Draft prefill
-            draft_capture_mode = (
-                CaptureHiddenMode.NULL
-                if self.speculative_algorithm.is_standalone()
-                else CaptureHiddenMode.LAST
-            )
-            model_worker_batch.capture_hidden_mode = draft_capture_mode
             with (
                 self.draft_worker.draft_tp_context(
                     self.draft_worker.draft_runner.tp_group
@@ -772,7 +771,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
             ):
                 batch_output.next_draft_input = (
                     self.draft_worker._draft_extend_for_prefill(
-                        model_worker_batch,
+                        batch,
                         batch_output.logits_output.hidden_states,
                         batch_output.next_token_ids,
                         batch_output.logits_output.mm_input_embeds,
@@ -780,13 +779,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 )
                 return batch_output
         else:
-            if model_worker_batch.spec_info is None:
+            if batch.spec_info is None:
                 capture_mode = (
                     CaptureHiddenMode.NULL
                     if self.speculative_algorithm.is_standalone()
                     else CaptureHiddenMode.LAST
                 )
-                model_worker_batch.spec_info = EagleDraftInput.create_idle_input(
+                batch.spec_info = EagleDraftInput.create_idle_input(
                     device=self.device,
                     hidden_size=EagleDraftInput.hidden_size_for(self.draft_worker),
                     dtype=EagleDraftInput.dtype_for(self.draft_worker),
@@ -800,9 +799,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
             ):
-                verify_input: EagleVerifyInput = self.draft_worker.draft(
+                verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
-                    model_worker_batch
-                )
             assert verify_input.is_verify_input()
             # Record a CUDA event after draft() GPU work is dispatched.
             # This event will be waited on by plan_stream in verify()
@@ -811,8 +808,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
             if self.plan_stream:
                 self._draft_done_event = torch.get_device_module(self.device).Event()
                 self._draft_done_event.record()
-            model_worker_batch.spec_info = verify_input
+            batch.spec_info = verify_input
-            batch_output = self.verify(model_worker_batch)
+            batch_output = self.verify(batch)
             with (
                 self.draft_worker.draft_tp_context(
                     self.draft_worker.draft_runner.tp_group
@@ -820,9 +817,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
             ):
-                self.draft_worker._draft_extend_for_decode(
+                self.draft_worker._draft_extend_for_decode(batch, batch_output)
-                    model_worker_batch, batch_output
-                )
 
             return batch_output
 
@@ -967,16 +962,11 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 sa.speculative_num_draft_tokens,
             ) = backup
 
-    def verify(self, batch: ModelWorkerBatch):
+    def verify(self, batch: ScheduleBatch):
-        # Since batch.seq_lens is allocated in another stream, we need
+        fwd_stream = torch.get_device_module(self.device).current_stream()
-        # record_stream() to prevent pytorch gc and reuse the gpu memory
-        # while forward_stream is still running.
-        batch.seq_lens.record_stream(
-            torch.get_device_module(self.device).current_stream()
-        )
-
-        # Parse args
         verify_input: EagleVerifyInput = batch.spec_info
+        record_stream_for_v2_verify(batch, verify_input, fwd_stream)
+
         verify_input.num_tokens_per_req = self.speculative_num_steps + 1
         bs = len(batch.seq_lens)
 
@@ -997,6 +987,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 )
             )
 
+        # Cover post-prepare rebinds: draft_token, plan_stream-allocated out_cache_loc.
+        record_stream_each((batch.input_ids, batch.out_cache_loc), fwd_stream)
+
         # Correct some buffers due to the overlap plan
         if self.plan_stream:
             torch.get_device_module(self.device).current_stream().wait_stream(
@@ -1037,7 +1030,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
 
         # Run target verify batch in the main compute stream (GPU compute)
         forward_batch_output = self.target_worker.forward_batch_generation(
-            model_worker_batch=None,
+            batch=None,
             forward_batch=verify_forward_batch,
             is_verify=True,
             skip_attn_backend_init=True,
@@ -1102,13 +1095,17 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 batch, logits_output, predict, accept_index, self.speculative_num_steps
             )
 
-        # Construct the next draft input
         next_draft_input = EagleDraftInput(
             bonus_tokens=bonus_tokens,
             new_seq_lens=new_seq_lens,
             verify_done=verify_done,
         )
 
+        # verify_forward_batch transitively holds verify-time GPU tensors
+        # (draft_token / out_cache_loc / ...) that must outlive the imminent
+        # batch.input_ids rebind in prepare_for_extend_to_fill_draft_kvcache,
+        # until the next iter's verify_done.synchronize() in filter_batch.
+        # Scheduler pins it in batch_record_buf for the 2-iter window.
         return GenerationBatchResult(
             logits_output=logits_output,
             next_token_ids=predict,
@@ -1118,11 +1115,12 @@ class EAGLEWorkerV2(BaseSpecWorker):
             accept_lens=accept_lens,
             routed_experts_output=forward_batch_output.routed_experts_output,
             indexer_topk_output=forward_batch_output.indexer_topk_output,
+            extra_keep_alive_refs=[verify_forward_batch],
         )
 
     def _mamba_verify_update(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         verify_input: EagleVerifyInput,
         accept_lens: torch.Tensor,
         accept_index: torch.Tensor,
@@ -1184,7 +1182,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
 
     def move_accepted_tokens_to_target_kvcache(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         accept_index: torch.Tensor,
         num_correct_drafts: torch.Tensor,
     ):

python/sglang/srt/speculative/frozen_kv_mtp_worker.py

@@ -389,12 +389,8 @@ class FrozenKVMTPWorker(TpModelWorker):
         batch.spec_info = draft_input
 
         try:
-            model_worker_batch = batch.get_model_worker_batch(
+            batch.seq_lens_cpu_cache = seq_lens_cpu
-                seq_lens_cpu_cache=seq_lens_cpu
+            forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
-            )
-            forward_batch = ForwardBatch.init_new(
-                model_worker_batch, self.draft_model_runner
-            )
             forward_batch.return_logprob = False
             if mm_input_embeds is not None:
                 forward_batch.mm_input_embeds = mm_input_embeds
@@ -491,13 +487,12 @@ class FrozenKVMTPWorker(TpModelWorker):
     def forward_target_extend(
         self, batch: ScheduleBatch
     ) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
-        model_worker_batch = batch.get_model_worker_batch()
+        batch.capture_hidden_mode = CaptureHiddenMode.FULL
-        model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
+        batch_result = self.target_worker.forward_batch_generation(batch)
-        batch_result = self.target_worker.forward_batch_generation(model_worker_batch)
         return (
             batch_result.logits_output,
             batch_result.next_token_ids,
-            model_worker_batch.seq_lens_cpu,
+            batch.seq_lens_cpu,
             batch_result.can_run_cuda_graph,
         )
 
@@ -593,11 +588,8 @@ class FrozenKVMTPWorker(TpModelWorker):
         batch.seq_lens_sum = torch.sum(batch.seq_lens).item()
         batch.return_hidden_states = False
 
-        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
-        assert model_worker_batch.capture_hidden_mode == CaptureHiddenMode.LAST
+        assert forward_batch.capture_hidden_mode == CaptureHiddenMode.LAST
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.draft_model_runner
-        )
         self._set_positions(forward_batch)
         self._expand_for_topk_draft(forward_batch)
 
@@ -718,11 +710,6 @@ class FrozenKVMTPWorker(TpModelWorker):
             else ForwardMode.IDLE
         )
 
-        model_worker_batch = batch.get_model_worker_batch(
-            seq_lens_cpu_cache=spec_info.seq_lens_cpu
-        )
-        assert model_worker_batch.capture_hidden_mode == spec_info.capture_hidden_mode
-
         if batch.has_grammar:
             retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
             retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
@@ -730,8 +717,9 @@ class FrozenKVMTPWorker(TpModelWorker):
                 spec_info.retrieve_next_token.shape
             ).cpu()
 
+        batch.seq_lens_cpu_cache = spec_info.seq_lens_cpu
         batch_result = self.target_worker.forward_batch_generation(
-            model_worker_batch, is_verify=True
+            batch, is_verify=True
         )
         logits_output, can_run_cuda_graph = (
             batch_result.logits_output,

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -378,15 +378,14 @@ class MultiLayerEagleWorker(TpModelWorker):
         """
         # Forward with the target model and get hidden states.
         # We need the full hidden states to prefill the KV cache of the draft model.
-        model_worker_batch = batch.get_model_worker_batch()
         capture_mode = (
             CaptureHiddenMode.NULL
             if self.speculative_algorithm.is_standalone()
             else CaptureHiddenMode.FULL
         )
-        model_worker_batch.capture_hidden_mode = capture_mode
+        batch.capture_hidden_mode = capture_mode
-        model_worker_batch.return_hidden_states_before_norm = True
+        batch.return_hidden_states_before_norm = True
-        batch_result = self.target_worker.forward_batch_generation(model_worker_batch)
+        batch_result = self.target_worker.forward_batch_generation(batch)
         logits_output, next_token_ids = (
             batch_result.logits_output,
             batch_result.next_token_ids,
@@ -394,7 +393,7 @@ class MultiLayerEagleWorker(TpModelWorker):
         return (
             logits_output,
             next_token_ids,
-            model_worker_batch.seq_lens_cpu,
+            batch.seq_lens_cpu,
             batch_result.can_run_cuda_graph,
         )
 
@@ -431,11 +430,8 @@ class MultiLayerEagleWorker(TpModelWorker):
         batch.return_hidden_states = False
 
         # Get forward batch
-        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
-        assert model_worker_batch.capture_hidden_mode == draft_capture_mode
+        assert forward_batch.capture_hidden_mode == draft_capture_mode
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.mtp_model_runner(0)
-        )
         forward_batch.can_run_dp_cuda_graph = False
         forward_batch.return_hidden_states_before_norm = True
 
@@ -545,12 +541,6 @@ class MultiLayerEagleWorker(TpModelWorker):
             else ForwardMode.IDLE
         )
 
-        model_worker_batch = batch.get_model_worker_batch(
-            seq_lens_cpu_cache=spec_info.seq_lens_cpu
-        )
-        assert model_worker_batch.capture_hidden_mode == spec_info.capture_hidden_mode
-        model_worker_batch.return_hidden_states_before_norm = True
-
         if batch.has_grammar:
             retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
             retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
@@ -559,8 +549,10 @@ class MultiLayerEagleWorker(TpModelWorker):
             ).cpu()
 
         # Forward
+        batch.seq_lens_cpu_cache = spec_info.seq_lens_cpu
+        batch.return_hidden_states_before_norm = True
         batch_result = self.target_worker.forward_batch_generation(
-            model_worker_batch, is_verify=True
+            batch, is_verify=True
         )
         logits_output, can_run_cuda_graph = (
             batch_result.logits_output,
@@ -682,12 +674,8 @@ class MultiLayerEagleWorker(TpModelWorker):
             else CaptureHiddenMode.LAST
         )
         batch.spec_info.capture_hidden_mode = capture_mode
-        model_worker_batch = batch.get_model_worker_batch(
+        batch.seq_lens_cpu_cache = seq_lens_cpu
-            seq_lens_cpu_cache=seq_lens_cpu
+        forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
-        )
-        forward_batch = ForwardBatch.init_new(
-            model_worker_batch, self.mtp_model_runner(0)
-        )
         forward_batch.return_logprob = False
         forward_batch.return_hidden_states_before_norm = True
         topk_p_list = []
@@ -764,11 +752,9 @@ class MultiLayerEagleWorker(TpModelWorker):
         )
 
         batch.return_hidden_states = False
-        model_worker_batch = batch.get_model_worker_batch()
+        draft_extend_input.capture_hidden_mode = draft_extend_capture_mode
-        assert model_worker_batch.capture_hidden_mode == draft_extend_capture_mode
+        forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
-        forward_batch = ForwardBatch.init_new(
+        assert forward_batch.capture_hidden_mode == draft_extend_capture_mode
-            model_worker_batch, self.mtp_model_runner(0)
-        )
         forward_batch.return_hidden_states_before_norm = True
         if forward_batch.seq_lens_cpu is not None:
             forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -25,7 +25,7 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightFromDiskReqInput,
     UpdateWeightsFromIPCReqInput,
 )
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import (
@@ -50,6 +50,8 @@ from sglang.srt.speculative.spec_utils import (
     draft_tp_context,
     maybe_detect_nan,
     maybe_detect_oob,
+    record_stream_each,
+    record_stream_for_v2_verify,
     select_top_k_tokens,
 )
 from sglang.srt.utils.common import empty_context, fast_topk
@@ -227,11 +229,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
                 forward_batch, batch_result
             )
 
-    def draft(self, model_worker_batch: ModelWorkerBatch):
+    def draft(self, batch: ScheduleBatch):
-        draft_input: EagleDraftInput = model_worker_batch.spec_info
+        draft_input: EagleDraftInput = batch.spec_info
         forward_batch, can_cuda_graph = draft_input.prepare_for_v2_draft(
             self.req_to_token_pool,
-            model_worker_batch,
+            batch,
             self.cuda_graph_runner,
             self.draft_runner_list[0],
             self.topk,
@@ -241,7 +243,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         # Run draft
         parent_list, top_scores_index, draft_tokens = self.draft_forward(forward_batch)
 
-        if model_worker_batch.forward_mode.is_idle():
+        if batch.forward_mode.is_idle():
             return EagleVerifyInput.create_idle_input(
                 self.topk,
                 self.speculative_num_steps,
@@ -265,8 +267,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
             parent_list,
             top_scores_index,
             draft_tokens,
-            model_worker_batch.seq_lens,
+            batch.seq_lens,
-            model_worker_batch.seq_lens_sum,
+            batch.seq_lens_sum,
             self.topk,
             self.speculative_num_steps,
             self.speculative_num_draft_tokens,
@@ -366,7 +368,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
 
     def _draft_extend_for_prefill(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
         target_hidden_states: torch.Tensor,
         next_token_ids: torch.Tensor,
     ):
@@ -390,9 +392,20 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
 
         batch.spec_info = next_draft_input
 
+        # Chain-style MTP needs FULL to get all-token hidden states;
+        # non-chain only needs LAST (the target model's hidden states).
+        # STANDALONE skips hidden states end-to-end.
+        if self.speculative_algorithm.is_standalone():
+            draft_capture_hidden_mode = CaptureHiddenMode.NULL
+        elif self.chain_mtp_hidden_states:
+            draft_capture_hidden_mode = CaptureHiddenMode.FULL
+        else:
+            draft_capture_hidden_mode = CaptureHiddenMode.LAST
+
         # Run forward
+        batch.capture_hidden_mode = draft_capture_hidden_mode
+        batch.return_hidden_states_before_norm = True
         forward_batch = ForwardBatch.init_new(batch, self.draft_runner_list[0])
-        forward_batch.return_hidden_states_before_norm = True
 
         # Construct input_ids
         if not batch.forward_mode.is_idle():
@@ -453,7 +466,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         return next_draft_input
 
     def _draft_extend_for_decode(
-        self, batch: ModelWorkerBatch, batch_result: GenerationBatchResult
+        self, batch: ScheduleBatch, batch_result: GenerationBatchResult
     ):
         # Batch 2: Draft extend
         draft_input = EagleDraftInput(
@@ -656,74 +669,58 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
         # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
         pass
 
-    def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
+    def forward_batch_generation(self, batch: ScheduleBatch):
-        if (
+        if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
-            model_worker_batch.forward_mode.is_extend()
-            or model_worker_batch.is_extend_in_batch
-        ):
             # Target prefill
             target_capture_mode = (
                 CaptureHiddenMode.NULL
                 if self.speculative_algorithm.is_standalone()
                 else CaptureHiddenMode.FULL
             )
-            model_worker_batch.capture_hidden_mode = target_capture_mode
+            batch.capture_hidden_mode = target_capture_mode
-            batch_output = self.target_worker.forward_batch_generation(
+            batch_output = self.target_worker.forward_batch_generation(batch)
-                model_worker_batch
-            )
 
             # Chain-style MTP needs FULL to get all-token hidden states;
             # non-chain only needs LAST (the target model's hidden states).
-            model_worker_batch.capture_hidden_mode = (
-                CaptureHiddenMode.FULL
-                if self.draft_worker.chain_mtp_hidden_states
-                else CaptureHiddenMode.LAST
-            )
             batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill(
-                model_worker_batch,
+                batch,
                 batch_output.logits_output.hidden_states,
                 batch_output.next_token_ids,
             )
             return batch_output
         else:
-            if model_worker_batch.spec_info is None:
+            if batch.spec_info is None:
                 capture_mode = (
                     CaptureHiddenMode.NULL
                     if self.speculative_algorithm.is_standalone()
                     else CaptureHiddenMode.LAST
                 )
-                model_worker_batch.spec_info = EagleDraftInput.create_idle_input(
+                batch.spec_info = EagleDraftInput.create_idle_input(
                     device=self.device,
                     hidden_size=EagleDraftInput.hidden_size_for(self.draft_worker),
                     dtype=EagleDraftInput.dtype_for(self.draft_worker),
                     topk=self.topk * self.speculative_num_steps,
                     capture_hidden_mode=capture_mode,
                 )
-            draft_input: EagleDraftInput = model_worker_batch.spec_info
+            verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
-            verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch)
             assert verify_input.is_verify_input()
             # Record a CUDA event after draft() GPU work is dispatched.
             if self.plan_stream:
                 self._draft_done_event = torch.get_device_module(self.device).Event()
                 self._draft_done_event.record()
-            model_worker_batch.spec_info = verify_input
+            batch.spec_info = verify_input
-            batch_output = self.verify(model_worker_batch)
+            batch_output = self.verify(batch)
-            self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output)
+            self.draft_worker._draft_extend_for_decode(batch, batch_output)
             return batch_output
 
     def verify(
         self,
-        batch: ModelWorkerBatch,
+        batch: ScheduleBatch,
     ):
-        # Since batch.seq_lens is allocated in another stream, we need
+        fwd_stream = torch.get_device_module(self.device).current_stream()
-        # record_stream() to prevent pytorch gc and reuse the gpu memory
-        # while forward_stream is still running.
-        batch.seq_lens.record_stream(
-            torch.get_device_module(self.device).current_stream()
-        )
-
-        # Parse args
         verify_input: EagleVerifyInput = batch.spec_info
+        record_stream_for_v2_verify(batch, verify_input, fwd_stream)
+
         bs = len(batch.seq_lens)
 
         # Batch 1: Target verify
@@ -741,6 +738,9 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
                 )
             )
 
+        # Cover post-prepare rebinds: draft_token, plan_stream-allocated out_cache_loc.
+        record_stream_each((batch.input_ids, batch.out_cache_loc), fwd_stream)
+
         # Correct some buffers due to the overlap plan
         if self.plan_stream:
             torch.get_device_module(self.device).current_stream().wait_stream(
@@ -760,7 +760,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
             )
         # Run target verify batch in the main compute stream
         forward_batch_output = self.target_worker.forward_batch_generation(
-            model_worker_batch=None,
+            batch=None,
             forward_batch=verify_forward_batch,
             is_verify=True,
             skip_attn_backend_init=True,
@@ -795,12 +795,14 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
                 batch, logits_output, predict, accept_index, self.speculative_num_steps
             )
 
-        # Construct the next draft input
         next_draft_input = EagleDraftInput(
             bonus_tokens=bonus_tokens,
             new_seq_lens=new_seq_lens,
             verify_done=verify_done,
         )
+        # verify_forward_batch transitively holds verify-time GPU tensors that
+        # must outlive the imminent batch.input_ids rebind; scheduler pins it
+        # in batch_record_buf via extra_keep_alive_refs. See EAGLEWorkerV2.verify.
         return GenerationBatchResult(
             logits_output=logits_output,
             next_token_ids=predict,
@@ -810,6 +812,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
             accept_lens=accept_lens,
             routed_experts_output=forward_batch_output.routed_experts_output,
             indexer_topk_output=forward_batch_output.indexer_topk_output,
+            extra_keep_alive_refs=[verify_forward_batch],
         )
 
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):

python/sglang/srt/speculative/ngram_worker.py

@@ -267,13 +267,12 @@ class NGRAMWorker:
 
         set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
 
-        model_worker_batch = batch.get_model_worker_batch()
+        spec_info = batch.spec_info
-        spec_info = model_worker_batch.spec_info
         num_correct_drafts = 0
         accept_lens = None
         num_correct_drafts_per_req_cpu = None
 
-        if model_worker_batch.forward_mode.is_target_verify():
+        if batch.forward_mode.is_target_verify():
             if batch.has_grammar:
                 retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
                 retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
@@ -284,14 +283,14 @@ class NGRAMWorker:
             set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
 
             batch_result = self.target_worker.forward_batch_generation(
-                model_worker_batch, is_verify=True
+                batch, is_verify=True
             )
             logits_output, can_run_cuda_graph = (
                 batch_result.logits_output,
                 batch_result.can_run_cuda_graph,
             )
 
-            verify_input: NgramVerifyInput = model_worker_batch.spec_info
+            verify_input: NgramVerifyInput = batch.spec_info
             vocab_mask = None
             if batch.has_grammar:
                 # Generate the logit mask for structured output.
@@ -349,9 +348,7 @@ class NGRAMWorker:
             batch.forward_mode = ForwardMode.DECODE
 
         else:
-            batch_result = self.target_worker.forward_batch_generation(
+            batch_result = self.target_worker.forward_batch_generation(batch)
-                model_worker_batch
-            )
             logits_output, next_token_ids, can_run_cuda_graph = (
                 batch_result.logits_output,
                 batch_result.next_token_ids,

python/sglang/srt/speculative/spec_info.py

@@ -18,7 +18,7 @@ from sglang.srt.speculative.spec_registry import (
 
 if TYPE_CHECKING:
     from sglang.srt.managers.overlap_utils import FutureMap
-    from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.tp_worker import TpModelWorker
     from sglang.srt.server_args import ServerArgs
     from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
@@ -256,11 +256,11 @@ class SpecInput(ABC):
         pass
 
     def get_spec_adjusted_global_num_tokens(
-        self, forward_batch: ModelWorkerBatch
+        self, batch: ScheduleBatch
     ) -> Tuple[List[int], List[int]]:
         c1, c2 = self.get_spec_adjust_token_coefficient()
-        global_num_tokens = [x * c1 for x in forward_batch.global_num_tokens]
+        global_num_tokens = [x * c1 for x in batch.global_num_tokens]
         global_num_tokens_for_logprob = [
-            x * c2 for x in forward_batch.global_num_tokens_for_logprob
+            x * c2 for x in batch.global_num_tokens_for_logprob
         ]
         return global_num_tokens, global_num_tokens_for_logprob

python/sglang/srt/speculative/spec_utils.py

@@ -52,6 +52,53 @@ TREE_SPEC_KERNEL_AVAILABLE = (
 )  # This kernel is only available for CUDA and MUSA now
 
 
+def record_stream_each(tensors, stream):
+    """Call record_stream(stream) on each cuda tensor in `tensors`, skipping
+    non-tensor / non-cuda entries. Tells the caching allocator that the
+    tensors are also used on `stream`, so memory is not recycled while
+    queued work is still in flight after Python refs drop.
+    """
+    for t in tensors:
+        if isinstance(t, torch.Tensor) and t.is_cuda:
+            t.record_stream(stream)
+
+
+def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
+    """Mark pre-prepare SB / verify_input GPU tensors as used on `fwd_stream`.
+
+    Spec V2 mutates SB mid-forward (`prepare_for_v2_verify` rebinds
+    `batch.input_ids` / `out_cache_loc`; `_draft_extend_for_decode` later
+    replaces `batch.input_ids` again). Each rebind drops the only SB Python
+    ref to the old tensor while the verify forward kernel may still be
+    reading its memory on `fwd_stream`; `record_stream` tells the caching
+    allocator to wait for `fwd_stream` before recycling the block.
+
+    Covers pre-prepare tensors only; caller must also `record_stream_each`
+    the post-prepare rebinds (new `batch.input_ids` / `out_cache_loc`).
+    """
+    candidates = [
+        batch.seq_lens,
+        batch.req_pool_indices,
+        batch.input_ids,
+        batch.out_cache_loc,
+    ]
+    if verify_input is not None:
+        candidates.extend(
+            [
+                getattr(verify_input, attr, None)
+                for attr in (
+                    "draft_token",
+                    "custom_mask",
+                    "positions",
+                    "retrieve_index",
+                    "retrieve_next_token",
+                    "retrieve_next_sibling",
+                )
+            ]
+        )
+    record_stream_each(candidates, fwd_stream)
+
+
 def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
     if server_args is None:
         server_args = get_global_server_args()

test/manual/test_forward_split_prefill.py

@@ -121,8 +121,7 @@ class TestForwardSplitPrefill(CustomTestCase):
             batch.prepare_for_extend()
 
         # Create forward batch
-        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(batch, self.model_runner)
-        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
 
         return forward_batch