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

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():
+    # SGLANG_ENABLE_SPEC_V2=False selects the non-overlap (synchronous) spec v2
-        # The V1 worker only supports non-overlap scheduling.
+    # 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

@@ -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
+        The caller computes and stores `batch.out_cache_loc` before this
-        `batch.out_cache_loc` before this method is called. This helper only
+        method is called. This helper only packages the verify forward and pre-initializes either CUDA-graph replay
-        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

@@ -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
-from sglang.srt.model_executor.forward_batch_info import (
-    CaptureHiddenMode,
-    ForwardBatch,
-    ForwardMode,
-)
 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

@@ -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
+    Drives both overlap and non-overlap scheduling, same as EAGLE: the
-    spec-v1 `DFlashWorker` (non-overlap), to keep the v1 path stable and to
+    scheduler runs it synchronously when overlap is disabled.
-    minimize risk while bringing up overlap scheduling.
     """
 
     def __init__(

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
+            # V2 worker drives both overlap and non-overlap (scheduler runs it
-            # and the scheduler schema always agree. With the gate on, the V2
+            # synchronously when overlap is disabled), same as EAGLE.
-            # worker drives both overlap and non-overlap, same as EAGLE.
+            from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
-            if self.supports_spec_v2():
-                from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
 
-                return DFlashWorkerV2
+            return DFlashWorkerV2
-            from sglang.srt.speculative.dflash_worker import DFlashWorker
-
-            return DFlashWorker
 
         if self.is_frozen_kv_mtp():
             # V2 worker drives both overlap and non-overlap (scheduler runs it

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",