minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

refactor: remove ModelWorkerBatch indirection (#25516)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-17 18:05:39 -07:00
committed by GitHub
parent b3803164cb
commit 58ece60703
21 changed files with 430 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/bench_one_batch.py
+2 -4
View File
@@ -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
python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py
+1 -2
View File
@@ -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
+20 -27
View File
@@ -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
python/sglang/srt/managers/overlap_utils.py
+6 -7
View File
@@ -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]
python/sglang/srt/managers/schedule_batch.py
+8 -196
View File
@@ -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
python/sglang/srt/managers/scheduler.py
+80 -48
View File
@@ -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,
python/sglang/srt/managers/scheduler_pp_mixin.py
+1 -2
View File
@@ -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
+18 -20
View File
@@ -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(
python/sglang/srt/managers/utils.py
+6 -1
View File
@@ -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
+80 -31
View File
@@ -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
python/sglang/srt/speculative/dflash_worker.py
+12 -28
View File
@@ -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,
python/sglang/srt/speculative/eagle_info_v2.py
+9 -9
View File
@@ -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,
):
python/sglang/srt/speculative/eagle_worker.py
+13 -28
View File
@@ -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:
python/sglang/srt/speculative/eagle_worker_v2.py
+48 -50
View File
@@ -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,
):
python/sglang/srt/speculative/frozen_kv_mtp_worker.py
+9 -21
View File
@@ -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,
python/sglang/srt/speculative/multi_layer_eagle_worker.py
+14 -28
View File
@@ -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()
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+46 -43
View File
@@ -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):
python/sglang/srt/speculative/ngram_worker.py
+5 -8
View File
@@ -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,
python/sglang/srt/speculative/spec_info.py
+4 -4
View File
@@ -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
+47
View File
@@ -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()