diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 34ec28371..f38d953dc 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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 diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index cea6ff90b..e0a0d7303 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -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, - ) diff --git a/python/sglang/srt/speculative/dflash_worker.py b/python/sglang/srt/speculative/dflash_worker.py index 38d5cf67d..bc3d6cf0f 100644 --- a/python/sglang/srt/speculative/dflash_worker.py +++ b/python/sglang/srt/speculative/dflash_worker.py @@ -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, - ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 4a57375b9..a44b5c637 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -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__( diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index cd1595a2f..1aa761e76 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -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 diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 06123097a..cb040ab99 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -73,9 +73,6 @@ _OWNER_SITES = { ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_committed_len"): 1, ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_allocated_len"): 1, ("speculative/eagle_info.py", "EagleVerifyInput.verify", "spec_verify_ct"): 1, - ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "kv_committed_len"): 1, - ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "kv_allocated_len"): 1, - ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "spec_verify_ct"): 1, # disaggregation decode prealloc ( "disaggregation/decode.py",