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

[2/N][Mixed] Mixed chunk prefill with spec enabled (#36933)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Yuwei An
2026-08-31 10:48:07 -07:00
committed by GitHub
co-authored by Claude Fable 5
parent 9cf157c252
commit 07d84ebd6d
13 changed files with 320 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+8 -23
View File
@@ -326,16 +326,6 @@ def _handle_dflash(server_args: ServerArgs) -> None:
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
)
if cfg.enable_mixed_chunk:
declare_resolution(
server_args,
"_handle_dflash",
enable_mixed_chunk=False,
)
logger.warning(
"Mixed chunked prefill is disabled because of using dflash speculative decoding."
)
def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
from sglang.srt.speculative.dspark_components.dspark_config import (
@@ -524,16 +514,6 @@ def _handle_dspark(server_args: ServerArgs) -> None:
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
)
if cfg.enable_mixed_chunk:
declare_resolution(
server_args,
"_handle_dspark",
enable_mixed_chunk=False,
)
logger.warning(
"Mixed chunked prefill is disabled because of using dspark speculative decoding."
)
from sglang.srt.speculative.ragged_verify import (
RaggedVerifyMode,
read_ragged_verify_mode,
@@ -688,15 +668,20 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
"speculative decoding."
)
if cfg.enable_mixed_chunk:
# Mixed steps degrade running requests to a plain 1-token decode.
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
algo = SpeculativeAlgorithm.from_string(cfg.speculative_algorithm)
if cfg.enable_mixed_chunk and not algo.supports_mixed_chunk():
declare_resolution(
server_args,
"_handle_eagle_family",
enable_mixed_chunk=False,
)
logger.warning(
"Mixed chunked prefill is disabled because of using "
"eagle speculative decoding."
"Mixed chunked prefill is disabled: %s speculative decoding does "
"not support it.",
cfg.speculative_algorithm,
)
model_arch = model_config_of(server_args).hf_config.architectures[0]
python/sglang/srt/arg_groups/validation_hook.py
+11 -1
View File
@@ -84,9 +84,19 @@ def check_server_args(server_args: Any):
# Check speculative decoding
if cfg.speculative_algorithm is not None:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
# Running requests degrade to a plain 1-token decode inside a
# mixed step; only workers with a verified resume path allow it.
assert (
not cfg.enable_mixed_chunk
), "enable_mixed_chunk is required for speculative decoding"
or SpeculativeAlgorithm.from_string(
cfg.speculative_algorithm
).supports_mixed_chunk()
), (
"enable_mixed_chunk is not supported with "
f"speculative_algorithm={cfg.speculative_algorithm}"
)
# Check chunked prefill
# Skip validation if chunked prefill is disabled (i.e., size <= 0).
python/sglang/srt/managers/overlap_utils.py
+55
View File
@@ -93,6 +93,8 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None:
if batch.prefill_input_ids_cpu is not None:
prefill_gpu = batch.prefill_input_ids_cpu.to(batch.device, non_blocking=True)
if batch.mix_running_indices is not None:
if batch.enable_overlap and not batch.spec_algorithm.is_none():
future_map.resolve_mixed_spec_tails(batch)
decode_gpu = future_map.output_tokens_buf[batch.mix_running_indices]
if _DEBUG_ASSERT:
_assert_nonneg_and_invalidate(
@@ -267,6 +269,8 @@ class FutureMap:
self.needs_cpu_seq_lens = needs_cpu_seq_lens
self.needs_confidence_relay = needs_confidence_relay
self.req_pool_size = req_to_token_pool.req_to_token.shape[0]
# Kept for the mixed-tail late binding (reserved-slot gather).
self.req_to_token = req_to_token_pool.req_to_token
if _DEBUG_ASSERT:
# Poisoned init: every row must be written before its first gather.
@@ -457,6 +461,57 @@ class FutureMap:
draft_input.bonus_tokens, self.output_tokens_buf, indices
)
def stash_bonus_tokens(
self, indices: torch.Tensor, bonus_tokens: torch.Tensor
) -> None:
"""Write only output_tokens_buf rows; for relays carrying no draft
extras (stash() would lazy-init the spec bufs from the payload)."""
self.output_tokens_buf[indices] = bonus_tokens.to(self.output_tokens_buf.dtype)
def resolve_mixed_spec_tails(self, batch: ScheduleBatch) -> None:
"""Late-bind a spec mixed batch's decode tails (overlap): schedule-time
lengths lag the in-flight step's accept count, so rebuild the tail rows
from the published committed lengths behind the publish fence."""
idx = batch.mix_running_indices
n = int(idx.shape[0])
if n == 0:
return
if self.publish_ready is not None:
if _is_hip:
self.publish_ready.synchronize()
else:
self.publish_ready.wait()
fresh = self.new_seq_lens_buf[idx]
seq_lens = batch.seq_lens.clone()
seq_lens[-n:] = fresh + 1
batch.seq_lens = seq_lens
out_cache_loc = batch.out_cache_loc.clone()
out_cache_loc[-n:] = self.req_to_token[idx.long(), fresh.long()].to(
out_cache_loc.dtype
)
batch.out_cache_loc = out_cache_loc
if self.fwd_prepare_d2h_stream is None or self.publish_ready is None:
fresh_cpu = fresh.cpu() # bootstrap / non-CUDA
else:
self.fwd_prepare_d2h_stream.wait_event(self.publish_ready)
with torch.get_device_module(self.device).stream(
self.fwd_prepare_d2h_stream
):
self.new_seq_lens_cpu_pinned.copy_(
self.new_seq_lens_buf, non_blocking=True
)
self.fwd_prepare_d2h_stream.synchronize()
fresh_cpu = self.new_seq_lens_cpu_pinned[batch.mix_running_indices_cpu]
if batch.seq_lens_cpu is not None:
seq_lens_cpu = batch.seq_lens_cpu.clone()
seq_lens_cpu[-n:] = fresh_cpu + 1
batch.seq_lens_cpu = seq_lens_cpu
batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
batch.prefix_lens = batch.prefix_lens[:-n] + [
int(x) for x in fresh_cpu.tolist()
]
def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None:
# Lazy pull from new_seq_lens_buf for spec_v2 (accept_lens not known to
# schedule). The CPU mirror is gated by needs_cpu_seq_lens; backends that
python/sglang/srt/managers/schedule_batch.py
+48 -3
View File
@@ -2156,6 +2156,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Staging consumed by resolve_forward_inputs (prefill H2D / mixed gather).
prefill_input_ids_cpu: Optional[torch.Tensor] = None
mix_running_indices: Optional[torch.Tensor] = None
# CPU twin of mix_running_indices; lets the overlap tail resolve gather
# pinned mirrors without a device sync.
mix_running_indices_cpu: Optional[torch.Tensor] = None
input_embeds: torch.Tensor = None # shape: [b, hidden_size], float32
# Token replacement embeddings and absolute positions (optional).
@@ -2853,13 +2856,55 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Decode tokens of the running portion live in future_map.output_tokens_buf.
self.input_ids = None
self.mix_running_indices = running_batch.req_pool_indices
out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
self.mix_running_indices_cpu = running_batch.req_pool_indices_cpu
if not self.spec_algorithm.is_none():
# Spec keeps no per-step out_cache_loc on the running batch; gather
# each tail's bonus slot at the committed length (rebound under overlap).
tail_base = torch.tensor(
[r.seqlen - 1 for r in running_batch.reqs],
dtype=torch.int64,
device=self.seq_lens.device,
)
running_out_cache_loc = self.req_to_token_pool.req_to_token[
running_batch.req_pool_indices.long(),
tail_base,
].to(self.out_cache_loc.dtype)
# The spec relay is unresolved at schedule time, so merge_batch
# would null seq_lens_cpu; rebuild the tails from request state.
running_seq_lens_cpu = torch.tensor(
[int(r.seqlen) for r in running_batch.reqs], dtype=torch.int64
)
if self.seq_lens_cpu is None:
merged_seq_lens_cpu = running_seq_lens_cpu
else:
merged_seq_lens_cpu = torch.cat(
[self.seq_lens_cpu, running_seq_lens_cpu]
)
else:
# Non-spec: the running batch carries prepared seq_lens_cpu
# (r.seqlen lags it under overlap); merge_batch concats it.
tail_base = None
running_out_cache_loc = running_batch.out_cache_loc
merged_seq_lens_cpu = None
out_cache_loc = torch.cat([self.out_cache_loc, running_out_cache_loc])
self.merge_batch(running_batch)
self.out_cache_loc = out_cache_loc
if merged_seq_lens_cpu is not None:
self.seq_lens_cpu = merged_seq_lens_cpu
if tail_base is not None:
# Spec seq_lens sit at the committed base (bonus token pending);
# this step commits it, so tails carry base + 1 or attention drops the row.
merged = self.seq_lens.clone()
merged[-running_bs:] = tail_base + 1
self.seq_lens = merged
# For overlap scheduler, the output_ids has one step delay
delta = 0 if self.enable_overlap else -1
# For overlap scheduler, the output_ids has one step delay;
# spec tail request state carries no delay in either mode.
if self.spec_algorithm.is_none():
delta = 0 if self.enable_overlap else -1
else:
delta = -1
# NOTE: prefix_indices is what has been cached, but we don't cache each decode step
self.prefix_lens = self.prefix_lens + [
python/sglang/srt/managers/scheduler.py
+14 -3
View File
@@ -3678,9 +3678,20 @@ class Scheduler(
running_batch.prepare_for_decode()
new_batch.mix_with_running(running_batch)
new_batch.decoding_reqs = running_batch.reqs
running_batch = ScheduleBatch(
reqs=[], batch_is_full=running_batch.batch_is_full
)
if not self.enable_overlap and not self.spec_algorithm.is_none():
# Non-overlap spec never writes the relay; stash the
# tails' pending tokens for the mixed input resolve.
last_tokens = torch.tensor(
[r.output_ids[-1] for r in running_batch.reqs],
dtype=torch.int64,
device=self.device,
)
self.future_map.stash_bonus_tokens(
running_batch.req_pool_indices, last_tokens
)
running_batch = ScheduleBatch(
reqs=[], batch_is_full=running_batch.batch_is_full
)
else:
new_batch.decoding_reqs = None
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+9
View File
@@ -328,6 +328,15 @@ class SchedulerBatchResultProcessor:
self._maybe_update_reasoning_tokens(req, next_token_id)
req.update_finish_state()
# A mixed spec tail committed its pending bonus token; advance
# so the next spec prepare_for_decode reserves from the right base.
if (
not req.finished()
and batch.decoding_reqs
and req in batch.decoding_reqs
and not batch.spec_algorithm.is_none()
):
req.kv.kv_committed_len += 1
if req.finished():
self._maybe_collect_routed_experts(req)
self._maybe_collect_indexer_topk(req)
python/sglang/srt/speculative/dflash_worker_v2.py
+3 -2
View File
@@ -1769,7 +1769,8 @@ class DFlashWorkerV2(BaseSpecWorker):
batch_output.next_token_ids,
)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_TARGET, next_token_ids)
batch_output.new_seq_lens = batch.seq_lens
new_seq_lens = batch.seq_lens
batch_output.new_seq_lens = new_seq_lens
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
@@ -1814,7 +1815,7 @@ class DFlashWorkerV2(BaseSpecWorker):
batch_output.next_draft_input = self._make_next_draft_input_prefill(
bonus_tokens=next_token_ids,
seq_lens=batch.seq_lens,
seq_lens=new_seq_lens,
)
return batch_output
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+3 -2
View File
@@ -490,7 +490,8 @@ class DSparkWorkerV2(BaseSpecWorker):
logits_output = batch_output.logits_output
next_token_ids = batch_output.next_token_ids
self._tp_sync.sync(SpecTpSyncSite.DSPARK_TARGET, next_token_ids)
batch_output.new_seq_lens = batch.seq_lens
new_seq_lens = batch.seq_lens
batch_output.new_seq_lens = new_seq_lens
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
@@ -547,7 +548,7 @@ class DSparkWorkerV2(BaseSpecWorker):
batch_output.next_draft_input = make_next_draft_input(
bonus_tokens=next_token_ids,
new_seq_lens=batch.seq_lens,
new_seq_lens=new_seq_lens,
)
return batch_output
python/sglang/srt/speculative/spec_info.py
+13
View File
@@ -129,6 +129,19 @@ class SpeculativeAlgorithm(Enum):
def supports_target_verify_for_draft(self) -> bool:
return self.is_dflash_family()
def supports_mixed_chunk(self) -> bool:
"""Whether mixed chunk prefill may stay enabled with this algorithm.
ngram cannot join as is: its overlap relay skips output_tokens_buf,
which the mixed input resolve reads.
"""
return self in (
SpeculativeAlgorithm.EAGLE,
SpeculativeAlgorithm.EAGLE3,
SpeculativeAlgorithm.DFLASH,
SpeculativeAlgorithm.DSPARK,
)
def supports_ragged_verify(self) -> bool:
"""Whether this algorithm's verify step may carry a RaggedVerifyLayout
(per-request verify lengths); gates the token-bucket-keyed verify
python/sglang/srt/speculative/spec_registry.py
+3
View File
@@ -70,6 +70,9 @@ class CustomSpecAlgo:
def is_eagle(self) -> bool:
return False
def supports_mixed_chunk(self) -> bool:
return False
def is_eagle3(self) -> bool:
return False