diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 61fa3ba5a..fa854fdef 100644 --- a/python/sglang/bench_one_batch.py +++ b/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(model_worker_batch, model_runner) + forward_batch = ForwardBatch.init_new(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(model_worker_batch, model_runner) + forward_batch = ForwardBatch.init_new(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 diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 2a9785b58..340c0d72e 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/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, diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 8dddf97f2..c4ecb05aa 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/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, - model_worker_batch: ModelWorkerBatch, + self, batch: ScheduleBatch ) -> 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 - reqs = model_worker_batch.reqs + forward_mode = batch.forward_mode + 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() - out_cache_loc_cpu = model_worker_batch.out_cache_loc.cpu().tolist() - extend_seq_lens = model_worker_batch.extend_seq_lens + input_ids_cpu = batch.input_ids.cpu().tolist() + out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist() + 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( - self, - model_worker_batch: ModelWorkerBatch, - ) -> tuple[ + def async_forward_batch_generation_mlx(self, batch: ScheduleBatch) -> 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 - reqs = model_worker_batch.reqs + forward_mode = batch.forward_mode + 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( - self, - model_worker_batch: ModelWorkerBatch, - ) -> tuple[ + def _async_extend_batch(self, batch: ScheduleBatch) -> 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 - input_ids_cpu = model_worker_batch.input_ids.cpu().tolist() - out_cache_loc_cpu = model_worker_batch.out_cache_loc.cpu().tolist() - extend_seq_lens = model_worker_batch.extend_seq_lens + reqs = batch.reqs + input_ids_cpu = batch.input_ids.cpu().tolist() + out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist() + extend_seq_lens = batch.extend_lens offset = 0 slot_offset = 0 diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 6e0a7285a..118ef04b3 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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 - # caching allocator (torch GC) could reclaim the memory before - # the GPU finishes reading it. + # batch.spec_info), so the caching allocator (torch GC) could + # reclaim the memory before 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] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e215da305..0b6d7b3e9 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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. - -TODO(lmzheng): ModelWorkerBatch seems a bit redundant and we consider removing it in the future. + It is constructed directly from a ScheduleBatch by `ForwardBatch.init_new`. """ 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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4091d884a..29b2a40ee 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 - self.record_batch_in_overlap(model_worker_batch) + with self._overlap_forward_isolation(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. - model_worker_batch.sampling_info = ( - model_worker_batch.sampling_info.copy_for_forward() - ) - bs = len(model_worker_batch.seq_lens) - future_indices = self.future_map.alloc_future_indices(bs) - - with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.schedule_stream) - self.future_map.resolve_future(model_worker_batch) - batch_result = self.model_worker.forward_batch_generation( - 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() - if batch_result.delay_sample_func is None: - self.future_map.store_to_map(future_indices, batch_result) - batch_result.copy_to_cpu( - return_logprob=batch.return_logprob, - return_hidden_states=batch.return_hidden_states, - ) - else: - batch_result.future_indices = future_indices + with self.forward_stream_ctx: + self.forward_stream.wait_stream(self.schedule_stream) + self.future_map.resolve_future(batch) + # FIXME: pp is not compatible with overlap + batch_result = self.model_worker.forward_batch_generation(batch) + # Park any refs the worker wants kept alive 2 iters + # (cross-stream tensor lifetime; pinned in the same + # ring slot as the SB attr snapshot). + if batch_result.extra_keep_alive_refs: + self.batch_record_buf[self.batch_record_ct].extend( + batch_result.extra_keep_alive_refs + ) + # FIXME(lsyin): maybe move this to forward_batch_generation + batch_result.copy_done = self.device_module.Event() + if batch_result.delay_sample_func is None: + self.future_map.store_to_map(future_indices, batch_result) + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, + return_hidden_states=batch.return_hidden_states, + ) + 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( - model_worker_batch - ) + pooler_output = self.tp_worker.forward_batch_embedding(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( - model_worker_batch - ) + pooler_output = self.tp_worker.forward_batch_embedding(batch) ret = EmbeddingBatchResult( embeddings=pooler_output.embeddings, pooled_hidden_states=pooler_output.pooled_hidden_states, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 154ad94be..9bf2b5a4d 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 60e105d93..773e61f67 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/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): - forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner) + def forward_batch_embedding(self, batch: ScheduleBatch): + 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 - if model_worker_batch is not None: + # Get forward batch from schedule batch + 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(model_worker_batch, self.model_runner) + forward_batch = ForwardBatch.init_new(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( diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 1a7224a74..4b7879b5e 100644 --- a/python/sglang/srt/managers/utils.py +++ b/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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 1671c5e42..d52414c06 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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(batch.extend_prefix_lens, list) - ret.extend_seq_lens = torch.tensor( - batch.extend_seq_lens, dtype=torch.int32 - ).to(device, non_blocking=True) + assert isinstance(extend_seq_lens, list) + assert isinstance(extend_prefix_lens, list) + ret.extend_seq_lens = torch.tensor(extend_seq_lens, dtype=torch.int32).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_seq_lens_cpu = batch.extend_seq_lens - ret.extend_logprob_start_lens_cpu = batch.extend_logprob_start_lens + ret.extend_prefix_lens_cpu = extend_prefix_lens + ret.extend_seq_lens_cpu = extend_seq_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( - self, batch: ModelWorkerBatch, model_runner: ModelRunner, device: torch.device - ): + def _init_ngram_embedding_info(self, batch: ScheduleBatch, 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( - self, model_runner: ModelRunner, batch: ModelWorkerBatch - ): + def _compute_mrope_positions(self, model_runner: ModelRunner, batch: ScheduleBatch): # 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_prefix_lens[batch_idx], + batch.extend_lens[batch_idx], + batch.prefix_lens[batch_idx], ) if ( mm_input is None diff --git a/python/sglang/srt/speculative/dflash_worker.py b/python/sglang/srt/speculative/dflash_worker.py index b3d688c84..87ddcfe23 100644 --- a/python/sglang/srt/speculative/dflash_worker.py +++ b/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, - batch: Union[ScheduleBatch, ModelWorkerBatch], - **kwargs, + self, batch: ScheduleBatch, **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() - model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL - - batch_result = self.target_worker.forward_batch_generation( - model_worker_batch, **kwargs - ) + batch.capture_hidden_mode = CaptureHiddenMode.FULL + batch_result = self.target_worker.forward_batch_generation(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 ( - model_worker_batch.extend_seq_lens is None - or model_worker_batch.extend_prefix_lens is None - ): + if batch.extend_lens is None or batch.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( - model_worker_batch.extend_seq_lens - ) + extend_seq_lens = _to_int32_device_tensor(batch.extend_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 model_worker_batch.forward_mode.is_target_verify() - verify_input = model_worker_batch.spec_info + assert batch.forward_mode.is_target_verify() + verify_input = 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, diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 4f49fed27..68dcc12e5 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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_prefix_lens = seq_lens_cpu_.tolist() + batch.extend_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))] + 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, ): diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 0a50b9182..b1931214f 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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_result = self.target_worker.forward_batch_generation(model_worker_batch) + batch.capture_hidden_mode = capture_mode + 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() - assert model_worker_batch.capture_hidden_mode == draft_capture_mode - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.draft_model_runner - ) + forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner) + assert forward_batch.capture_hidden_mode == draft_capture_mode 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( - seq_lens_cpu_cache=seq_lens_cpu - ) - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.draft_model_runner - ) + batch.seq_lens_cpu_cache = seq_lens_cpu + forward_batch = ForwardBatch.init_new(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() - # propagates the correct mode (NULL for STANDALONE). + # default (LAST); override here so ForwardBatch.init_new picks up the + # correct mode (NULL for STANDALONE). draft_extend_input.capture_hidden_mode = draft_extend_capture_mode - model_worker_batch = batch.get_model_worker_batch() - assert model_worker_batch.capture_hidden_mode == draft_extend_capture_mode - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.draft_model_runner - ) + forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner) + assert forward_batch.capture_hidden_mode == draft_extend_capture_mode if forward_batch.seq_lens_cpu is not None: forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item() else: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index ef23afd74..8a815eb5e 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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): - draft_input: EagleDraftInput = model_worker_batch.spec_info + def draft(self, batch: ScheduleBatch): + 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, - model_worker_batch.seq_lens_sum, + batch.seq_lens, + 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): - if ( - model_worker_batch.forward_mode.is_extend() - or model_worker_batch.is_extend_in_batch - ): + def forward_batch_generation(self, batch: ScheduleBatch): + if batch.forward_mode.is_extend() or 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_output = self.target_worker.forward_batch_generation( - model_worker_batch - ) + batch.capture_hidden_mode = target_capture_mode + batch_output = self.target_worker.forward_batch_generation(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( - model_worker_batch - ) + verify_input: EagleVerifyInput = self.draft_worker.draft(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_output = self.verify(model_worker_batch) + batch.spec_info = verify_input + 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( - model_worker_batch, batch_output - ) + self.draft_worker._draft_extend_for_decode(batch, batch_output) return batch_output @@ -967,16 +962,11 @@ class EAGLEWorkerV2(BaseSpecWorker): sa.speculative_num_draft_tokens, ) = backup - def verify(self, batch: ModelWorkerBatch): - # Since batch.seq_lens is allocated in another stream, we need - # 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 + def verify(self, batch: ScheduleBatch): + fwd_stream = torch.get_device_module(self.device).current_stream() 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, ): diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index 480dd7144..c52e01942 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/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( - seq_lens_cpu_cache=seq_lens_cpu - ) - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.draft_model_runner - ) + batch.seq_lens_cpu_cache = seq_lens_cpu + forward_batch = ForwardBatch.init_new(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() - model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL - batch_result = self.target_worker.forward_batch_generation(model_worker_batch) + batch.capture_hidden_mode = CaptureHiddenMode.FULL + batch_result = self.target_worker.forward_batch_generation(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() - assert model_worker_batch.capture_hidden_mode == CaptureHiddenMode.LAST - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.draft_model_runner - ) + forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner) + assert forward_batch.capture_hidden_mode == CaptureHiddenMode.LAST 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, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 6217f36f2..ebdeec9d8 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/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 - model_worker_batch.return_hidden_states_before_norm = True - batch_result = self.target_worker.forward_batch_generation(model_worker_batch) + batch.capture_hidden_mode = capture_mode + batch.return_hidden_states_before_norm = True + 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() - assert model_worker_batch.capture_hidden_mode == draft_capture_mode - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.mtp_model_runner(0) - ) + forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0)) + assert forward_batch.capture_hidden_mode == draft_capture_mode 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( - seq_lens_cpu_cache=seq_lens_cpu - ) - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.mtp_model_runner(0) - ) + batch.seq_lens_cpu_cache = seq_lens_cpu + forward_batch = ForwardBatch.init_new(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() - assert model_worker_batch.capture_hidden_mode == draft_extend_capture_mode - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.mtp_model_runner(0) - ) + draft_extend_input.capture_hidden_mode = draft_extend_capture_mode + forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0)) + assert forward_batch.capture_hidden_mode == draft_extend_capture_mode 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() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 4650619e9..8cbf2dc57 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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): - draft_input: EagleDraftInput = model_worker_batch.spec_info + def draft(self, batch: ScheduleBatch): + 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, - model_worker_batch.seq_lens_sum, + batch.seq_lens, + 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): - if ( - model_worker_batch.forward_mode.is_extend() - or model_worker_batch.is_extend_in_batch - ): + def forward_batch_generation(self, batch: ScheduleBatch): + if batch.forward_mode.is_extend() or 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_output = self.target_worker.forward_batch_generation( - model_worker_batch - ) + batch.capture_hidden_mode = target_capture_mode + batch_output = self.target_worker.forward_batch_generation(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(model_worker_batch) + verify_input: EagleVerifyInput = self.draft_worker.draft(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_output = self.verify(model_worker_batch) - self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output) + batch.spec_info = verify_input + batch_output = self.verify(batch) + 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 - # 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 + fwd_stream = torch.get_device_module(self.device).current_stream() 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): diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 49cbe7022..d53522d6b 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/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 = model_worker_batch.spec_info + spec_info = 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( - model_worker_batch - ) + batch_result = self.target_worker.forward_batch_generation(batch) logits_output, next_token_ids, can_run_cuda_graph = ( batch_result.logits_output, batch_result.next_token_ids, diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 0bbf357d5..ddfddf11a 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 64cb54cd6..3a39dcd49 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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() diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index 7c23f4f14..ab83965a6 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/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(model_worker_batch, self.model_runner) + forward_batch = ForwardBatch.init_new(batch, self.model_runner) return forward_batch