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

[Spec] Remove the DFLASH V1 worker path (#27959)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-11 14:37:58 -07:00
committed by GitHub
parent 6e885c844f
commit acdb39edd2
6 changed files with 21 additions and 884 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+6 -8
View File
@@ -230,15 +230,13 @@ 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 not envs.SGLANG_ENABLE_SPEC_V2.get():
# The V1 worker only supports non-overlap scheduling.
# SGLANG_ENABLE_SPEC_V2=False selects the non-overlap (synchronous) spec v2
# path instead of the overlap-scheduled one; both run the V2 worker.
if (
not envs.SGLANG_ENABLE_SPEC_V2.get()
and not server_args.disable_overlap_schedule
):
server_args.disable_overlap_schedule = True
logger.warning(
"Spec v1 is used for DFLASH speculative decoding because "
"SGLANG_ENABLE_SPEC_V2 is off; overlap schedule is disabled."
)
else:
logger.warning("Spec v2 is enabled by default for DFLASH speculative decoding.")
if server_args.enable_mixed_chunk:
server_args.enable_mixed_chunk = False
python/sglang/srt/speculative/dflash_info.py
+4 -395
View File
@@ -1,161 +1,26 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional, Tuple
from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
alloc_token_slots,
get_last_loc,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.speculative.dflash_utils import (
apply_dflash_verify_logits_adjustments,
compute_dflash_correct_drafts_and_bonus,
compute_dflash_sampling_correct_drafts_and_bonus,
is_dflash_sampling_verify_available,
)
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
if TYPE_CHECKING:
from sglang.srt.managers.tp_worker import TpModelWorker
def _compute_paged_keep_slots(
*,
prefix_lens: torch.Tensor,
commit_lens: torch.Tensor,
draft_token_num: int,
page_size: int,
) -> torch.Tensor:
"""Compute how many draft slots per request must remain allocated.
The allocator frees at page granularity for paged mode, so we can only release
full pages from the tail after verify.
"""
if page_size <= 1:
raise ValueError(f"Expected page_size > 1, got {page_size}.")
seq_dtype = prefix_lens.dtype
extended_lens = prefix_lens + int(draft_token_num)
new_lens = prefix_lens + commit_lens.to(seq_dtype)
aligned_new_lens = ((new_lens + page_size - 1) // page_size) * page_size
keep_lens = torch.minimum(aligned_new_lens, extended_lens)
keep_slots = (keep_lens - prefix_lens).to(torch.int64)
keep_slots.clamp_(min=0, max=int(draft_token_num))
return keep_slots
@dataclass
class DFlashDraftInput(SpecInput):
"""Per-batch DFlash draft state for spec-v1 (non-overlap) scheduling.
This object is stored on `ScheduleBatch.spec_info` between decode iterations.
It is NOT sent to model attention backends; the DFlash worker uses it to run
the draft model and to track draft-side cache progress.
When draft windowing is disabled, `draft_seq_lens` matches the committed target
prefix length already materialized in the draft KV cache. When windowing is
enabled, `draft_seq_lens` is the logical resident length in the draft worker's
compact req-to-token mapping. In paged mode this may exceed the requested
window by up to `page_size - 1` so the local page table remains valid. `ctx_lens`
tracks newly committed target tokens that still need draft KV materialization.
"""
# Current token to start the next DFlash block (one per request).
bonus_tokens: torch.Tensor
# Flattened context features for tokens that need to be appended into the draft cache.
# Shape: [sum(ctx_lens), K * hidden_size], where K is the number of target-layer
# hidden-state features concatenated per token (len(dflash_config.target_layer_ids),
# or default K == draft_num_layers for existing checkpoints).
target_hidden: torch.Tensor
# Context lengths per request, used to slice `target_hidden`. Device tensor (int32).
ctx_lens: torch.Tensor
# How many committed tokens are visible to the draft worker per request.
draft_seq_lens: torch.Tensor
def __post_init__(self):
super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT)
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
# Draft state does not change token accounting.
return (1, 1)
def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
old_ctx_lens = self.ctx_lens
old_target_hidden = self.target_hidden
self.bonus_tokens = self.bonus_tokens[new_indices]
self.ctx_lens = old_ctx_lens[new_indices]
self.draft_seq_lens = self.draft_seq_lens[new_indices]
if old_target_hidden is None or old_target_hidden.numel() == 0:
self.target_hidden = old_target_hidden
return
# Rebuild target_hidden for the filtered batch using vectorized indexing.
old_bs = int(old_ctx_lens.shape[0])
offsets = torch.zeros(
(old_bs + 1,), dtype=torch.int64, device=old_ctx_lens.device
)
offsets[1:].copy_(old_ctx_lens.to(torch.int64).cumsum(0))
start = offsets[:-1]
seg_start = start[new_indices]
seg_lens = old_ctx_lens[new_indices].to(torch.int64)
max_len = int(seg_lens.max().item()) if seg_lens.numel() > 0 else 0
if max_len <= 0:
self.target_hidden = old_target_hidden[:0]
return
r = torch.arange(max_len, device=old_ctx_lens.device, dtype=torch.int64)[
None, :
]
pos2d = seg_start[:, None] + r
mask = r < seg_lens[:, None]
flat_pos = pos2d[mask]
self.target_hidden = (
old_target_hidden.index_select(0, flat_pos)
if flat_pos.numel() > 0
else old_target_hidden[:0]
)
def merge_batch(self, spec_info: "DFlashDraftInput"):
self.bonus_tokens = torch.cat(
[self.bonus_tokens, spec_info.bonus_tokens], dim=0
)
self.ctx_lens = torch.cat([self.ctx_lens, spec_info.ctx_lens], dim=0)
self.draft_seq_lens = torch.cat(
[self.draft_seq_lens, spec_info.draft_seq_lens], dim=0
)
if self.target_hidden is None or self.target_hidden.numel() == 0:
self.target_hidden = spec_info.target_hidden
elif (
spec_info.target_hidden is not None and spec_info.target_hidden.numel() > 0
):
self.target_hidden = torch.cat(
[self.target_hidden, spec_info.target_hidden], dim=0
)
@dataclass
class DFlashVerifyInput(SpecInput):
"""Inputs for a target-model verify forward in DFlash (spec-v1).
"""Inputs for a target-model verify forward in DFlash.
The verify forward is run with `ForwardMode.TARGET_VERIFY` so that the target
model returns logits for all tokens in the block, enabling accept-length
@@ -184,80 +49,6 @@ class DFlashVerifyInput(SpecInput):
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
return self.draft_token_num, self.draft_token_num
def prepare_for_verify(
self,
batch: ScheduleBatch,
page_size: int,
*,
build_custom_mask: bool = True,
):
if batch.forward_mode.is_idle():
return
batch.input_ids = self.draft_token
if page_size == 1:
batch.out_cache_loc = alloc_token_slots(
batch.tree_cache, len(batch.input_ids)
)
end_offset = batch.seq_lens + self.draft_token_num
else:
prefix_lens = batch.seq_lens
prefix_lens_cpu = batch.seq_lens_cpu
end_offset = prefix_lens + self.draft_token_num
end_offset_cpu = prefix_lens_cpu + self.draft_token_num
last_loc = get_last_loc(
batch.req_to_token_pool.req_to_token,
batch.req_pool_indices,
prefix_lens,
)
batch.out_cache_loc = alloc_paged_token_slots_extend(
batch.tree_cache,
prefix_lens,
prefix_lens_cpu,
end_offset,
end_offset_cpu,
last_loc,
len(batch.input_ids),
)
bs = batch.batch_size()
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
end_offset,
batch.out_cache_loc,
bs,
)
if not build_custom_mask:
self.custom_mask = None
return
if self.draft_token_num <= 0:
raise ValueError(
f"DFLASH draft_token_num must be positive, got {self.draft_token_num}."
)
mask_chunks: List[torch.Tensor] = []
q_len = int(self.draft_token_num)
q_idx = torch.arange(q_len, device=batch.device, dtype=torch.int32).unsqueeze(1)
for prefix_len in batch.seq_lens_cpu.tolist():
prefix_len_i = int(prefix_len)
kv_len = prefix_len_i + q_len
k_idx = torch.arange(
kv_len, device=batch.device, dtype=torch.int32
).unsqueeze(0)
# Allow attending to the full prefix and to tokens up to (and including) the
# current query position within the verify block (standard causal masking).
allow = k_idx <= (prefix_len_i + q_idx)
mask_chunks.append(allow.flatten())
self.custom_mask = (
torch.cat(mask_chunks, dim=0)
if mask_chunks
else torch.empty((0,), dtype=torch.bool, device=batch.device)
)
def prepare_for_v2_verify(
self,
batch: ScheduleBatch,
@@ -265,9 +56,8 @@ class DFlashVerifyInput(SpecInput):
) -> tuple[ForwardBatch, bool]:
"""Prepare a DFLASH verify forward batch for overlap scheduling.
Unlike spec-v1, the overlap path already computes and stores
`batch.out_cache_loc` before this method is called. This helper only
packages the verify forward and pre-initializes either CUDA-graph replay
The caller computes and stores `batch.out_cache_loc` before this
method is called. This helper only packages the verify forward and pre-initializes either CUDA-graph replay
metadata or eager attention metadata so the actual forward can run with
`skip_attn_backend_init=True`.
"""
@@ -357,184 +147,3 @@ class DFlashVerifyInput(SpecInput):
)
self.custom_mask = mask
return kv_indices, cum_kv_seq_len, qo_indptr, mask
def verify(
self,
*,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
page_size: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[int]]:
"""DFlash verification for greedy and non-greedy sampling.
Returns:
new_bonus_tokens: int64 tensor [bs] (the new current token per request)
commit_lens: int32 tensor [bs] (how many verify-input tokens are committed)
next_target_hidden: tensor [sum(commit_lens), feature_dim]
num_correct_drafts_per_req_cpu: list[int] (accepted draft tokens per request)
"""
if batch.forward_mode.is_idle():
empty = torch.empty((0,), dtype=torch.int64, device=batch.device)
return empty, empty.to(torch.int32), empty, []
bs = batch.batch_size()
device = logits_output.next_token_logits.device
sampling_info = batch.sampling_info
if sampling_info is not None:
if len(sampling_info) != bs:
raise RuntimeError(
"DFLASH verify sampling_info size mismatch: "
f"len(sampling_info)={len(sampling_info)}, bs={bs}."
)
apply_dflash_verify_logits_adjustments(
next_token_logits=logits_output.next_token_logits,
sampling_info=sampling_info,
draft_token_num=self.draft_token_num,
)
candidates = self.draft_token.view(bs, self.draft_token_num)
if (
sampling_info is not None
and not sampling_info.is_all_greedy
and is_dflash_sampling_verify_available()
):
top_ks = [int(req.sampling_params.top_k) for req in batch.reqs]
correct_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus(
candidates=candidates,
next_token_logits=logits_output.next_token_logits,
sampling_info=sampling_info,
max_top_k=max(max(top_ks), 1) if top_ks else 1,
uniform_top_k_value=(
top_ks[0]
if top_ks and all(top_k == top_ks[0] for top_k in top_ks)
else None
),
)
else:
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1).view(
bs, self.draft_token_num
)
correct_len, bonus = compute_dflash_correct_drafts_and_bonus(
candidates=candidates,
target_predict=target_predict,
)
# Single D2H transfer: candidates[1:] + correct_len + bonus
packed = torch.cat(
[candidates[:, 1:], correct_len.unsqueeze(1), bonus.unsqueeze(1)], dim=1
).cpu()
max_acc = self.draft_token_num - 1
num_correct_drafts_per_req_cpu: List[int] = []
commit_lens_cpu: List[int] = []
new_bonus_tokens_list: List[int] = []
for i, req in enumerate(batch.reqs):
acc_len = int(packed[i, max_acc].item())
proposed = packed[i, :acc_len].tolist() + [
int(packed[i, max_acc + 1].item())
]
appended = 0
for token_id in proposed:
token_id = int(token_id)
req.output_ids.append(token_id)
appended += 1
req.update_finish_state()
if req.finished():
break
if req.grammar is not None:
req.grammar.accept_token(token_id)
if req.output_ids:
new_bonus_token = int(req.output_ids[-1])
elif req.origin_input_ids:
# If no token was appended in this verify step, keep the current token unchanged.
new_bonus_token = int(req.origin_input_ids[-1])
else:
raise RuntimeError(
"DFLASH verify cannot determine current token: both output_ids and origin_input_ids are empty."
)
commit_lens_cpu.append(appended)
new_bonus_tokens_list.append(new_bonus_token)
num_correct_drafts_per_req_cpu.append(max(0, appended - 1))
req.spec_verify_ct += 1
req.spec_num_correct_drafts += num_correct_drafts_per_req_cpu[-1]
commit_lens = torch.tensor(commit_lens_cpu, dtype=torch.int32, device=device)
new_bonus_tokens = torch.tensor(
new_bonus_tokens_list, dtype=torch.int64, device=device
)
# Free uncommitted KV cache slots and compact out_cache_loc.
if page_size == 1:
out_cache_loc = batch.out_cache_loc.view(bs, self.draft_token_num)
keep_mask = (
torch.arange(self.draft_token_num, device=device)[None, :]
< commit_lens[:, None]
)
batch.token_to_kv_pool_allocator.free(out_cache_loc[~keep_mask])
batch.out_cache_loc = out_cache_loc[keep_mask]
else:
out_cache_loc = batch.out_cache_loc.view(bs, self.draft_token_num)
row_offsets = torch.arange(self.draft_token_num, device=device)[None, :]
keep_slots = _compute_paged_keep_slots(
prefix_lens=batch.seq_lens,
commit_lens=commit_lens,
draft_token_num=self.draft_token_num,
page_size=page_size,
)
free_mask = row_offsets >= keep_slots[:, None]
batch.token_to_kv_pool_allocator.free(out_cache_loc[free_mask])
keep_mask = row_offsets < commit_lens[:, None]
batch.out_cache_loc = out_cache_loc[keep_mask]
# Update req-level KV cache accounting.
for req, commit_len in zip(batch.reqs, commit_lens_cpu, strict=True):
req.kv_committed_len += commit_len
req.kv_allocated_len = req.kv_committed_len
# Update req_to_token pool mapping for newly committed tokens.
end_offset = batch.seq_lens + commit_lens.to(batch.seq_lens.dtype)
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
end_offset,
batch.out_cache_loc,
bs,
)
# Update batch seq lens.
batch.seq_lens.add_(commit_lens.to(batch.seq_lens.dtype))
batch.seq_lens_cpu.add_(
torch.tensor(commit_lens_cpu, dtype=batch.seq_lens_cpu.dtype)
)
# Keep seq_lens_sum in sync; flashinfer indices updaters rely on this for buffer sizing.
batch.seq_lens_sum += sum(commit_lens_cpu)
# Build next-step context features from the committed verify-input tokens.
hidden = logits_output.hidden_states
if hidden is None:
raise RuntimeError(
"DFLASH verify requires target hidden states, but got None."
)
hidden = hidden.view(bs, self.draft_token_num, -1)
segments: List[torch.Tensor] = []
for i, ln in enumerate(commit_lens_cpu):
if ln > 0:
segments.append(hidden[i, :ln, :])
next_target_hidden = torch.cat(segments, dim=0) if segments else hidden[:0]
# Avoid confusing downstream consumers (spec-v1 decode doesn't use this).
logits_output.hidden_states = None
return (
new_bonus_tokens,
commit_lens,
next_target_hidden,
num_correct_drafts_per_req_cpu,
)
python/sglang/srt/speculative/dflash_worker.py
+3 -459
View File
@@ -7,28 +7,18 @@ import torch
from sglang.srt.distributed import get_tp_group
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
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.server_args import (
ServerArgs,
get_global_server_args,
set_global_server_args_for_scheduler,
)
from sglang.srt.speculative.dflash_info import DFlashDraftInput, DFlashVerifyInput
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
from sglang.srt.speculative.dflash_utils import (
can_dflash_use_fused_qkv_proj,
is_dflash_sampling_verify_available,
parse_dflash_draft_config,
resolve_dflash_verify_mask_policy,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
from sglang.srt.utils import is_cuda, is_npu
_is_npu = is_npu()
@@ -51,7 +41,7 @@ def _get_fused_kv_materialize_helper():
class DFlashWorker:
"""DFlash speculative decoding worker (spec-v1, tp>=1/pp=1)."""
"""Shared DFLASH infrastructure (draft model, draft KV materialization)."""
def __init__(
self,
@@ -540,184 +530,6 @@ class DFlashWorker:
return int(resolved_id)
def _prepare_for_speculative_decoding(
self, batch: ScheduleBatch, draft_input: DFlashDraftInput
):
if batch.forward_mode.is_extend() or batch.forward_mode.is_idle():
return
if batch.has_grammar:
raise RuntimeError(
"Invariant broken: DFLASH batch has grammar constraints, but scheduler should have rejected this request."
)
if batch.sampling_info is not None and not batch.sampling_info.is_all_greedy:
if (
not is_dflash_sampling_verify_available()
and not self._warned_sampling_fallback
and self.tp_rank == 0
):
logger.warning(
"DFLASH non-greedy verification is unavailable on this build/device; "
"falling back to greedy argmax verification."
)
self._warned_sampling_fallback = True
bs = batch.batch_size()
# --- 1) Append any newly committed tokens into the draft KV cache.
self._append_target_hidden_to_draft_kv(batch, draft_input)
target_model = self.target_worker.model_runner.model
embed_module = target_model.get_input_embeddings()
lm_head = getattr(target_model, "lm_head", None)
if (
lm_head is None
or not hasattr(lm_head, "weight")
or not hasattr(lm_head, "shard_indices")
):
raise RuntimeError(
"DFLASH requires the target model to expose a vocab-parallel `lm_head` with `weight` and "
"`shard_indices` attributes."
)
# --- 2) Draft a fixed block with the draft model.
self._ensure_draft_block_buffers(bs)
assert self._draft_block_ids_buf is not None
assert self._draft_block_positions_buf is not None
assert self._draft_block_tokens_buf is not None
assert self._draft_block_end_buf is not None
assert self._draft_seq_lens_cpu_buf is not None
allocator = self.draft_model_runner.token_to_kv_pool_allocator
token_to_kv_pool_state_backup = allocator.backup_state()
try:
block_ids = self._draft_block_ids_buf[:bs]
block_ids.fill_(int(self._mask_token_id))
block_ids[:, 0].copy_(draft_input.bonus_tokens.to(torch.long))
noise_embedding = embed_module(block_ids)
input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1])
# For spec-v1, the draft KV cache is always materialized before drafting the
# next block. `target_prefix_lens` stay absolute for RoPE; `draft_prefix_lens`
# are the logical resident lengths in the draft-local cache.
target_prefix_lens = batch.seq_lens # int32, device
draft_prefix_lens = draft_input.draft_seq_lens
if draft_prefix_lens.dtype != torch.int32:
draft_prefix_lens = draft_prefix_lens.to(torch.int32)
if draft_prefix_lens.device != self.device:
draft_prefix_lens = draft_prefix_lens.to(self.device, non_blocking=True)
positions_2d = self._draft_block_positions_buf[:bs]
torch.add(
target_prefix_lens.unsqueeze(1),
self._block_pos_offsets,
out=positions_2d,
)
positions = positions_2d.reshape(-1)
block_start = draft_prefix_lens
block_end = self._draft_block_end_buf[:bs]
torch.add(block_start, int(self.block_size), out=block_end)
seq_lens_cpu = self._draft_seq_lens_cpu_buf[:bs]
seq_lens_cpu.copy_(draft_prefix_lens.to(device="cpu", dtype=torch.int32))
if self.page_size == 1:
block_cache_loc = allocator.alloc(bs * self.block_size)
else:
block_end_cpu = seq_lens_cpu + int(self.block_size)
last_loc = get_last_loc(
self.draft_model_runner.req_to_token_pool.req_to_token,
batch.req_pool_indices,
block_start,
)
block_cache_loc = allocator.alloc_extend(
block_start,
seq_lens_cpu,
block_end,
block_end_cpu,
last_loc,
bs * self.block_size,
)
if block_cache_loc is None:
raise RuntimeError(
f"DFLASH draft OOM when allocating {bs * self.block_size} block tokens."
)
assign_req_to_token_pool_func(
batch.req_pool_indices,
self.draft_model_runner.req_to_token_pool.req_to_token,
block_start,
block_end,
block_cache_loc,
bs,
)
# Use TARGET_VERIFY mode (cuda-graphable) to run a fixed-size draft block.
# In this mode, `seq_lens` stores the prefix lengths; attention backends
# derive kv_len by adding `draft_token_num`.
draft_spec_info = self._draft_block_spec_info
seq_lens = draft_prefix_lens
seq_lens_sum = int(draft_prefix_lens.sum().item())
forward_batch = ForwardBatch(
forward_mode=ForwardMode.TARGET_VERIFY,
batch_size=bs,
input_ids=block_ids.flatten(),
req_pool_indices=batch.req_pool_indices,
seq_lens=seq_lens,
out_cache_loc=block_cache_loc,
seq_lens_sum=seq_lens_sum,
seq_lens_cpu=seq_lens_cpu,
positions=positions,
input_embeds=input_embeds,
spec_algorithm=SpeculativeAlgorithm.DFLASH,
spec_info=draft_spec_info,
capture_hidden_mode=CaptureHiddenMode.NULL,
)
with torch.inference_mode():
draft_logits_output = self.draft_model_runner.forward(
forward_batch
).logits_output
finally:
# Drop the speculative block from the shared allocator (EAGLE3-style).
allocator.restore_state(token_to_kv_pool_state_backup)
draft_hidden = draft_logits_output.hidden_states
if draft_hidden is None:
raise RuntimeError("DFLASH draft model returned no hidden states.")
draft_hidden = draft_hidden.view(bs, self.block_size, -1)
draft_next = self._greedy_sample_from_vocab_parallel_head(
hidden_states=draft_hidden[:, 1:, :].reshape(-1, draft_hidden.shape[-1]),
lm_head=lm_head,
).view(bs, self.block_size - 1)
draft_tokens = self._draft_block_tokens_buf[:bs]
draft_tokens[:, 0].copy_(block_ids[:, 0])
draft_tokens[:, 1:].copy_(draft_next)
positions = positions_2d.reshape(-1)
verify_input = DFlashVerifyInput(
draft_token=draft_tokens.reshape(-1),
positions=positions,
draft_token_num=self.block_size,
)
_, build_custom_mask = resolve_dflash_verify_mask_policy(
self.model_runner.attn_backend
)
verify_input.prepare_for_verify(
batch,
self.page_size,
build_custom_mask=build_custom_mask,
)
batch.forward_mode = (
ForwardMode.TARGET_VERIFY
if not batch.forward_mode.is_idle()
else ForwardMode.IDLE
)
batch.spec_info = verify_input
batch.return_hidden_states = False
def _greedy_sample_from_vocab_parallel_head(
self,
*,
@@ -928,144 +740,6 @@ class DFlashWorker:
return out_tokens
def _append_target_hidden_to_draft_kv(
self,
batch: ScheduleBatch,
draft_input: DFlashDraftInput,
) -> None:
"""Materialize the target hidden-state features into the draft KV cache.
This must be run before exposing new tokens to radix cache (prefix hits), otherwise
another request could reuse target KV indices without having draft KV values.
"""
bs = batch.batch_size()
device = self.model_runner.device
if draft_input.target_hidden is None:
raise RuntimeError(
"DFLASH draft state missing target_hidden context features."
)
if draft_input.ctx_lens.numel() != bs:
raise RuntimeError(
f"DFLASH ctx_lens length mismatch: got {draft_input.ctx_lens.numel()} for bs={bs}."
)
if draft_input.draft_seq_lens.numel() != bs:
raise RuntimeError(
f"DFLASH draft_seq_lens length mismatch: got {draft_input.draft_seq_lens.numel()} for bs={bs}."
)
total_ctx = int(draft_input.target_hidden.shape[0])
if total_ctx <= 0:
draft_input.ctx_lens = torch.zeros_like(draft_input.ctx_lens)
draft_input.target_hidden = draft_input.target_hidden[:0]
return
target_req_to_token = batch.req_to_token_pool.req_to_token
draft_req_to_token = self.draft_model_runner.req_to_token_pool.req_to_token
req_pool_indices = batch.req_pool_indices
if req_pool_indices.dtype != torch.int64:
req_pool_indices = req_pool_indices.to(torch.int64)
ctx_lens = draft_input.ctx_lens
if ctx_lens.dtype != torch.int32:
ctx_lens = ctx_lens.to(torch.int32)
if ctx_lens.device != device:
ctx_lens = ctx_lens.to(device, non_blocking=True)
ctx_start = batch.seq_lens.to(torch.int64) - ctx_lens.to(torch.int64)
if bs == 1:
# Fast path for single request.
max_ctx = int(total_ctx)
if max_ctx <= self._block_pos_offsets.numel():
r = self._block_pos_offsets[:max_ctx]
else:
r = torch.arange(max_ctx, device=device, dtype=torch.int64)
pos2d = ctx_start[:, None] + r[None, :] # [1, ctx]
cache2d = target_req_to_token[req_pool_indices[:, None], pos2d] # [1, ctx]
ctx_cache_loc = cache2d.reshape(-1).to(torch.int64) # [ctx]
ctx_positions = pos2d.reshape(-1) # [ctx]
else:
# In decode mode, ctx_lens <= block_size so we can skip the .item() sync.
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
max_ctx = int(ctx_lens.max().item())
else:
max_ctx = int(self.block_size)
if max_ctx <= 0:
raise RuntimeError(f"DFLASH invalid max_ctx={max_ctx} for KV append.")
if max_ctx <= self._block_pos_offsets.numel():
r = self._block_pos_offsets[:max_ctx]
else:
r = torch.arange(max_ctx, device=device, dtype=torch.int64)
r = r[None, :] # [1, max_ctx]
pos2d = ctx_start[:, None] + r # [bs, max_ctx]
mask = r < ctx_lens[:, None]
# Batched gather of cache locations and positions.
ctx_cache_loc = self._gather_req_to_token_masked(
req_to_token=target_req_to_token,
req_pool_indices=req_pool_indices,
pos2d=pos2d,
mask=mask,
context="DFLASH target hidden KV append",
) # [sum(ctx_lens)]
ctx_positions = pos2d[mask] # [sum(ctx_lens)]
with torch.inference_mode():
ctx_hidden = self.draft_model.project_target_hidden(
draft_input.target_hidden
) # [sum(ctx), hidden]
if ctx_hidden.shape[0] != ctx_cache_loc.numel():
raise RuntimeError(
f"DFLASH ctx_hidden/cache_loc mismatch: {ctx_hidden.shape[0]} vs {ctx_cache_loc.numel()}."
)
wrote_with_fused_kv = False
if self._use_fused_kv_materialize and self._fused_kv_helper is not None:
try:
self._append_target_hidden_fused(
ctx_hidden, ctx_positions, ctx_cache_loc
)
wrote_with_fused_kv = True
except Exception as e:
logger.warning(
"DFLASH fused KV append failed; falling back to sequential path: %s",
e,
)
self._use_fused_kv_materialize = False
self._fused_kv_helper = None
if not wrote_with_fused_kv:
self._append_target_hidden_sequential(
ctx_hidden, ctx_positions, ctx_cache_loc
)
if self.use_compact_draft_cache:
new_draft_seq_lens = self._compute_compact_draft_seq_lens(batch.seq_lens)
suffix_start = batch.seq_lens.to(torch.int64) - new_draft_seq_lens.to(
torch.int64
)
suffix_cache_loc = self._gather_req_to_token_segments(
req_to_token=target_req_to_token,
req_pool_indices=req_pool_indices,
start=suffix_start,
lengths=new_draft_seq_lens,
)
assign_req_to_token_pool_func(
batch.req_pool_indices,
draft_req_to_token,
torch.zeros_like(new_draft_seq_lens),
new_draft_seq_lens,
suffix_cache_loc,
bs,
)
draft_input.draft_seq_lens = new_draft_seq_lens
else:
draft_input.draft_seq_lens = batch.seq_lens.to(dtype=torch.int32)
draft_input.ctx_lens = torch.zeros_like(ctx_lens)
draft_input.target_hidden = draft_input.target_hidden[:0]
def _append_target_hidden_to_draft_kv_by_loc(
self,
*,
@@ -1335,133 +1009,3 @@ class DFlashWorker:
mamba_steps_to_track=mamba_steps_to_track,
model=self.target_worker.model_runner.model,
)
def forward_batch_generation(
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 batch.forward_mode.is_extend() or batch.is_extend_in_batch:
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,
)
if logits_output.hidden_states is None:
raise RuntimeError(
"DFLASH requires target aux hidden capture for prefill, but got None. "
"Make sure the target model has DFlash layers-to-capture configured."
)
if batch.extend_lens is None or batch.prefix_lens is None:
raise RuntimeError(
"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
# for radix cache support, since the scheduler may update radix after prefill returns.
device = next_token_ids.device
def _to_int32_device_tensor(x, *, device=device):
if isinstance(x, torch.Tensor):
if x.device != device:
x = x.to(device, non_blocking=True)
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(batch.extend_lens)
draft_input = DFlashDraftInput(
bonus_tokens=next_token_ids.to(torch.int64),
target_hidden=logits_output.hidden_states,
ctx_lens=extend_seq_lens,
draft_seq_lens=(
torch.zeros_like(extend_seq_lens)
if self.use_compact_draft_cache
else _to_int32_device_tensor(batch.prefix_lens)
),
)
self._append_target_hidden_to_draft_kv(batch, draft_input)
batch.spec_info = draft_input
return GenerationBatchResult(
logits_output=logits_output,
next_token_ids=next_token_ids,
num_correct_drafts=0,
can_run_cuda_graph=batch_result.can_run_cuda_graph,
)
# Decode / target-verify stage.
draft_input = batch.spec_info
if not isinstance(draft_input, DFlashDraftInput):
raise RuntimeError(
"DFLASH decode requires DFlashDraftInput state on the running batch. "
"This usually means the request did not complete the prefill stage."
)
self._prepare_for_speculative_decoding(batch, draft_input)
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,
"update_mamba_state_after_mtp_verify",
)
seq_lens_pre_verify = (
batch.seq_lens.clone() if need_mamba_verify_commit else None
)
batch_result = self.target_worker.forward_batch_generation(
batch, is_verify=True, **kwargs
)
logits_output, can_run_cuda_graph = (
batch_result.logits_output,
batch_result.can_run_cuda_graph,
)
(
new_bonus_tokens,
commit_lens,
next_target_hidden,
num_correct_drafts_per_req_cpu,
) = verify_input.verify(
batch=batch,
logits_output=logits_output,
page_size=self.page_size,
)
if need_mamba_verify_commit:
assert seq_lens_pre_verify is not None
self._update_target_mamba_state_after_verify(
batch=batch,
seq_lens_pre_verify=seq_lens_pre_verify,
commit_lens=commit_lens,
)
# Update draft state for the next iteration. Also materialize the committed verify tokens
# into the draft KV cache immediately so radix cache entries are safe to reuse.
draft_input.bonus_tokens = new_bonus_tokens
draft_input.target_hidden = next_target_hidden
draft_input.ctx_lens = commit_lens
self._append_target_hidden_to_draft_kv(batch, draft_input)
batch.spec_info = draft_input
batch.forward_mode = ForwardMode.DECODE
num_correct_drafts = sum(num_correct_drafts_per_req_cpu)
if not self._logged_first_verify and self.tp_rank == 0:
logger.info(
"DFLASH verify completed. num_correct_drafts_per_req=%s",
num_correct_drafts_per_req_cpu,
)
self._logged_first_verify = True
return GenerationBatchResult(
logits_output=logits_output,
next_token_ids=new_bonus_tokens,
num_correct_drafts=num_correct_drafts,
num_correct_drafts_per_req_cpu=num_correct_drafts_per_req_cpu,
can_run_cuda_graph=can_run_cuda_graph,
)
python/sglang/srt/speculative/dflash_worker_v2.py
+3 -4
View File
@@ -37,11 +37,10 @@ logger = logging.getLogger(__name__)
class DFlashWorkerV2(DFlashWorker):
"""DFLASH speculative decoding worker (spec-v2 overlap scheduling).
"""DFLASH speculative decoding worker (spec-v2).
This is intentionally implemented as a *separate* worker from the existing
spec-v1 `DFlashWorker` (non-overlap), to keep the v1 path stable and to
minimize risk while bringing up overlap scheduling.
Drives both overlap and non-overlap scheduling, same as EAGLE: the
scheduler runs it synchronously when overlap is disabled.
"""
def __init__(
python/sglang/srt/speculative/spec_info.py
+5 -15
View File
@@ -152,16 +152,11 @@ class SpeculativeAlgorithm(Enum):
return None
def supports_spec_v2(self) -> bool:
from sglang.srt.environ import envs
# DFLASH still ships a V1 worker; SGLANG_ENABLE_SPEC_V2=0 selects it
# and must flip the scheduler schema together with the worker.
# TODO: drop the env gate once the DFLASH V1 worker is removed.
return (
self.is_eagle()
or self.is_standalone()
or self.is_ngram()
or (self.is_dflash() and envs.SGLANG_ENABLE_SPEC_V2.get())
or self.is_dflash()
)
def need_topk(self) -> bool:
@@ -185,16 +180,11 @@ class SpeculativeAlgorithm(Enum):
), "Cannot create worker for NONE speculative algorithm."
if self.is_dflash():
# Keyed off the same env gate as supports_spec_v2() so the worker
# and the scheduler schema always agree. With the gate on, the V2
# worker drives both overlap and non-overlap, same as EAGLE.
if self.supports_spec_v2():
from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
# V2 worker drives both overlap and non-overlap (scheduler runs it
# synchronously when overlap is disabled), same as EAGLE.
from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
return DFlashWorkerV2
from sglang.srt.speculative.dflash_worker import DFlashWorker
return DFlashWorker
return DFlashWorkerV2
if self.is_frozen_kv_mtp():
# V2 worker drives both overlap and non-overlap (scheduler runs it