[Feature] Spec V2 DFlash Support (#23000)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -230,10 +230,15 @@ 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."
         )
 
-    server_args.disable_overlap_schedule = True
-    logger.warning(
-        "Overlap scheduler is disabled when using DFLASH speculative decoding (spec v2 is not supported yet)."
-    )
+    if not envs.SGLANG_ENABLE_SPEC_V2.get():
+        # The V1 worker only supports non-overlap scheduling.
+        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/environ.py

@@ -621,6 +621,7 @@ class Envs:
     # Overlap Spec V2
     SGLANG_ENABLE_SPEC_V2 = EnvBool(True)
     SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False)
+    SGLANG_DFLASH_PREFILL_REFILL_TARGET = EnvInt(None)
 
     # Spec Config
     SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True)

python/sglang/srt/managers/overlap_utils.py

@@ -150,19 +150,44 @@ class FutureMap:
     def _lazy_init_forward_buf(self, draft_input: EagleDraftInput):
         self._forward_buf_initialized = True
 
-        topk_p0 = draft_input.topk_p[0]
-        topk_index0 = draft_input.topk_index[0]
-        self.topk_p_buf = torch.empty(
-            (self.req_pool_size, *topk_p0.shape),
-            dtype=topk_p0.dtype,
-            device=self.device,
+        self.need_verified_id = getattr(draft_input, "verified_id", None) is not None
+        self.need_bonus_tokens = getattr(draft_input, "bonus_tokens", None) is not None
+        self.need_topk = self.spec_algo.need_topk()
+        self.need_hidden_states = (
+            spec_need_hidden_states()
+            and getattr(draft_input, "hidden_states", None) is not None
         )
-        self.topk_index_buf = torch.empty(
-            (self.req_pool_size, *topk_index0.shape),
-            dtype=topk_index0.dtype,
-            device=self.device,
-        )
-        if spec_need_hidden_states():
+
+        if self.need_verified_id:
+            verified_id0 = draft_input.verified_id[0]
+            self.verified_id_buf = (
+                torch.full(
+                    (self.req_pool_size, *verified_id0.shape),
+                    -1,
+                    dtype=verified_id0.dtype,
+                    device=self.device,
+                )
+                if _DEBUG_ASSERT
+                else torch.empty(
+                    (self.req_pool_size, *verified_id0.shape),
+                    dtype=verified_id0.dtype,
+                    device=self.device,
+                )
+            )
+        if self.need_topk:
+            topk_p0 = draft_input.topk_p[0]
+            topk_index0 = draft_input.topk_index[0]
+            self.topk_p_buf = torch.empty(
+                (self.req_pool_size, *topk_p0.shape),
+                dtype=topk_p0.dtype,
+                device=self.device,
+            )
+            self.topk_index_buf = torch.empty(
+                (self.req_pool_size, *topk_index0.shape),
+                dtype=topk_index0.dtype,
+                device=self.device,
+            )
+        if self.need_hidden_states:
             hidden_states0 = draft_input.hidden_states[0]
             self.hidden_states_buf = torch.empty(
                 (self.req_pool_size, *hidden_states0.shape),
@@ -178,38 +203,68 @@ class FutureMap:
         if draft_input is None:
             # FIXME(lsyin): only prefill; not compatible with mixed mode
             return
+        if self.spec_algo.is_dflash() and getattr(
+            draft_input, "direct_carry_valid", False
+        ):
+            return
         indices = draft_input.future_indices
+        if indices.shape[0] == 0:
+            return
         # FIXME: indices = batch.req_pool_indices, pinned 2 iters via
         # record_batch_in_overlap; record_stream here is redundant.
         indices.record_stream(torch.get_device_module(self.device).current_stream())
-        hidden_states_buf = (
-            self.hidden_states_buf if spec_need_hidden_states() else None
-        )
-        (
-            draft_input.topk_p,
-            draft_input.topk_index,
-            draft_input.bonus_tokens,
-            hidden_states,
-        ) = gather_spec_extras(
-            indices,
-            self.topk_p_buf,
-            self.topk_index_buf,
-            self.output_tokens_buf,
-            hidden_states_buf,
-        )
-        if hidden_states is not None:
-            draft_input.hidden_states = hidden_states
-        if _DEBUG_ASSERT:
-            _assert_nonneg_and_invalidate(
-                draft_input.bonus_tokens, self.output_tokens_buf, indices
+        if self.need_verified_id:
+            draft_input.verified_id = self.verified_id_buf[indices]
+        if self.need_topk:
+            hidden_states_buf = (
+                self.hidden_states_buf if self.need_hidden_states else None
             )
+            (
+                draft_input.topk_p,
+                draft_input.topk_index,
+                bonus_tokens,
+                hidden_states,
+            ) = gather_spec_extras(
+                indices,
+                self.topk_p_buf,
+                self.topk_index_buf,
+                self.output_tokens_buf,
+                hidden_states_buf,
+            )
+            if self.need_bonus_tokens:
+                draft_input.bonus_tokens = bonus_tokens
+            if hidden_states is not None:
+                draft_input.hidden_states = hidden_states
+        elif self.need_bonus_tokens:
+            draft_input.bonus_tokens = self.output_tokens_buf[indices]
+        if self.need_hidden_states and not self.need_topk:
+            draft_input.hidden_states = self.hidden_states_buf[indices]
+        if _DEBUG_ASSERT:
+            if self.need_verified_id:
+                _assert_nonneg_and_invalidate(
+                    draft_input.verified_id, self.verified_id_buf, indices
+                )
+            if self.need_bonus_tokens:
+                _assert_nonneg_and_invalidate(
+                    draft_input.bonus_tokens, self.output_tokens_buf, indices
+                )
 
     def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None:
-        # seq_lens_cpu may be needed on the host for kernel-launch prep (some backends).
-        # Run this D2H on a standalone stream to avoid chain-blocking forward_n ->
-        # prepare_{n+1}: a sync on the schedule stream would inherit its WAR barrier and
-        # stall the host until forward_n ends.
-        fi = batch.spec_info.future_indices if batch.spec_info is not None else None
+        # Lazy pull from new_seq_lens_buf for spec_v2 (accept_lens not known to
+        # schedule). DFLASH intentionally keeps host-side lengths lagging and
+        # uses its carried KV allocation watermark for planning, so only the GPU
+        # seq_lens is resolved there. Other spec-v2 algorithms still need the CPU
+        # mirror for host planning; use a private D2H stream for those copies.
+        draft_input = batch.spec_info
+        if draft_input is None:
+            return
+        if self.spec_algo.is_dflash() and getattr(
+            draft_input, "direct_carry_valid", False
+        ):
+            batch.seq_lens = draft_input.new_seq_lens
+            return
+
+        fi = draft_input.future_indices
         if fi is None:
             return
         if self.publish_ready is not None:
@@ -220,6 +275,11 @@ class FutureMap:
                 self.publish_ready.wait()
         batch.seq_lens = self.new_seq_lens_buf[fi]
 
+        if self.spec_algo.is_dflash():
+            # DFLASH keeps seq_lens_cpu as the lagging committed host view;
+            # planning/reserved host lengths live on DFlashDraftInputV2.
+            return
+
         if not self.needs_cpu_seq_lens:
             # GPU gather above is kept (SB.seq_lens must advance each verify);
             # skip the .cpu() D2H. Downstream takes the GPU-only path.
@@ -277,14 +337,21 @@ class FutureMap:
         draft_input: EagleDraftInput = payload
         if not self._forward_buf_initialized:
             self._lazy_init_forward_buf(draft_input)
-        self.output_tokens_buf[indices] = draft_input.bonus_tokens.to(
-            self.output_tokens_buf.dtype
-        )
-        self.topk_p_buf[indices] = draft_input.topk_p.to(self.topk_p_buf.dtype)
-        self.topk_index_buf[indices] = draft_input.topk_index.to(
-            self.topk_index_buf.dtype
-        )
-        if spec_need_hidden_states():
+        if self.need_verified_id:
+            self.verified_id_buf[indices] = draft_input.verified_id.to(
+                self.verified_id_buf.dtype
+            )
+        if self.need_bonus_tokens:
+            self.output_tokens_buf[indices] = draft_input.bonus_tokens.to(
+                self.output_tokens_buf.dtype
+            )
+
+        if self.need_topk:
+            self.topk_p_buf[indices] = draft_input.topk_p.to(self.topk_p_buf.dtype)
+            self.topk_index_buf[indices] = draft_input.topk_index.to(
+                self.topk_index_buf.dtype
+            )
+        if self.need_hidden_states:
             self.hidden_states_buf[indices] = draft_input.hidden_states.to(
                 self.hidden_states_buf.dtype
             )

python/sglang/srt/managers/scheduler.py

@@ -234,7 +234,11 @@ from sglang.srt.plugins import load_plugins
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
 from sglang.srt.session.session_controller import SessionController
-from sglang.srt.speculative.dflash_utils import validate_dflash_request
+from sglang.srt.speculative.dflash_utils import (
+    resolve_dflash_prefill_refill_target,
+    should_delay_dflash_prefill_for_batching,
+    validate_dflash_request,
+)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -838,6 +842,11 @@ class Scheduler(
             _,
             _,
         ) = self.tp_worker.get_worker_info()
+        self.dflash_prefill_refill_target = (
+            resolve_dflash_prefill_refill_target(self.max_running_requests)
+            if self.spec_algorithm.is_dflash()
+            else 1
+        )
         if not get_global_server_args().pp_max_micro_batch_size:
             get_global_server_args().pp_max_micro_batch_size = max(
                 self.max_running_requests // self.ps.pp_size, 1
@@ -1412,8 +1421,10 @@ class Scheduler(
         self.schedule_stream = self.device_module.Stream(priority=0)
         if self.device == "cpu":
             self.schedule_stream.synchronize = lambda: None  # No-op for CPU
-        # WAR barrier is CUDA-only; other platforms keep the pre-barrier behavior.
-        self._war_barrier_enabled = is_cuda()
+        # DFLASH fences its shared req_to_token writes with verify_done /
+        # plan-stream deps, so the global WAR barrier only serializes plan
+        # overlap. TODO: generalize this global-barrier enablement policy.
+        self._war_barrier_enabled = is_cuda() and not self.spec_algorithm.is_dflash()
         with self.device_module.StreamContext(self.schedule_stream):
             dispatch_event_loop(self)
 
@@ -2029,13 +2040,12 @@ class Scheduler(
             return
 
         if self.spec_algorithm.is_dflash():
-            error_msg = validate_dflash_request(req)
+            error_msg = validate_dflash_request(req, self.enable_overlap)
             if error_msg is not None:
                 req.set_finish_with_abort(error_msg)
                 self.init_req_max_new_tokens(req)
                 self._add_request_to_queue(req)
                 return
-
         # Handle multimodal inputs
         if recv_req.mm_inputs is not None:
             image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs)
@@ -2544,6 +2554,19 @@ class Scheduler(
         res = min(res, self.req_to_token_pool.available_size())
         return res
 
+    def _should_delay_dflash_prefill_for_batching(self, running_bs: int) -> bool:
+        if not self.spec_algorithm.is_dflash():
+            return False
+        if running_bs <= 0 or self.chunked_req is not None:
+            return False
+
+        return should_delay_dflash_prefill_for_batching(
+            running_bs=running_bs,
+            num_allocatable_reqs=self.get_num_allocatable_reqs(running_bs),
+            max_running_requests=self.max_running_requests,
+            prefill_refill_target=self.dflash_prefill_refill_target,
+        )
+
     def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
         prefill_delayer_single_pass = None
         if self.prefill_delayer:
@@ -2586,6 +2609,8 @@ class Scheduler(
             return None
 
         running_bs = len(self.running_batch.reqs)
+        if self._should_delay_dflash_prefill_for_batching(running_bs):
+            return None
 
         # Ignore the check if self.chunked_req is not None.
         # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0,

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -561,12 +561,17 @@ class SchedulerBatchResultProcessor:
                 continue
 
             if req.finished():
-                # -1 because prepare_for_decode pre-claimed the bonus slot.
-                req.kv_committed_len -= 1
+                if not batch.spec_algorithm.is_dflash():
+                    # EAGLE prepare_for_decode pre-claimed the bonus slot.
+                    req.kv_committed_len -= 1
                 continue
 
-            # -1 because prepare_for_decode pre-claimed the bonus slot.
-            req.kv_committed_len += accept_lens[i] - 1
+            if batch.spec_algorithm.is_dflash():
+                # DFLASH materialized accepted draft tokens plus the bonus token.
+                req.kv_committed_len += accept_lens[i]
+            else:
+                # EAGLE prepare_for_decode pre-claimed the bonus slot.
+                req.kv_committed_len += accept_lens[i] - 1
             req.spec_verify_ct += 1
 
             num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]

python/sglang/srt/mem_cache/memory_pool.py

@@ -31,6 +31,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+import triton
 
 from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
 from sglang.srt.configs.mamba_utils import BaseLinearStateParams
@@ -47,6 +48,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
 from sglang.srt.mem_cache.triton_ops.cache_move import (
     copy_all_layer_kv_cache_tiled,
+    set_kv_buffer_prefix_valid_tiled,
 )
 from sglang.srt.mem_cache.utils import (
     get_mla_kv_buffer_triton,
@@ -143,6 +145,67 @@ def _set_kv_buffer_impl(
         v_cache[indices] = v
 
 
+def _set_kv_buffer_prefix_valid_impl(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    loc_2d: torch.Tensor,
+    commit_lens: torch.Tensor,
+    row_dim: int,
+    store_dtype: torch.dtype,
+) -> None:
+    if k.numel() == 0 or loc_2d.numel() == 0 or commit_lens.numel() == 0:
+        return
+
+    if not k.is_contiguous():
+        k = k.contiguous()
+    if not v.is_contiguous():
+        v = v.contiguous()
+    if not loc_2d.is_contiguous():
+        loc_2d = loc_2d.contiguous()
+    if not commit_lens.is_contiguous():
+        commit_lens = commit_lens.contiguous()
+
+    row_bytes = row_dim * store_dtype.itemsize
+    if row_bytes <= 0:
+        return
+
+    if row_bytes >= 8192:
+        bytes_per_tile = 512
+        num_warps = 8
+    elif row_bytes >= 4096:
+        bytes_per_tile = 256
+        num_warps = 4
+    else:
+        bytes_per_tile = 128
+        num_warps = 4
+
+    grid = (
+        int(loc_2d.shape[0]),
+        int(loc_2d.shape[1]),
+        triton.cdiv(row_bytes, bytes_per_tile),
+    )
+
+    set_kv_buffer_prefix_valid_tiled[grid](
+        k,
+        v,
+        k_cache,
+        v_cache,
+        loc_2d,
+        commit_lens,
+        int(k.stride(0) * k.element_size()),
+        int(v.stride(0) * v.element_size()),
+        int(k_cache.stride(0) * k_cache.element_size()),
+        int(v_cache.stride(0) * v_cache.element_size()),
+        int(loc_2d.shape[1]),
+        ROW_BYTES=row_bytes,
+        BYTES_PER_TILE=bytes_per_tile,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+
+
 class ReqToTokenPool:
     """A memory pool that maps a request to its token locations."""
 
@@ -1284,6 +1347,91 @@ class MHATokenToKVPool(KVCache):
             same_kv_dim=self.same_kv_dim,
         )
 
+    def set_kv_buffer_prefix_valid(
+        self,
+        layer: RadixAttention,
+        loc_2d: torch.Tensor,
+        commit_lens: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        layer_id_override: Optional[int] = None,
+    ):
+        if layer_id_override is not None:
+            layer_id = layer_id_override
+        else:
+            layer_id = layer.layer_id
+
+        if loc_2d.ndim != 2:
+            raise ValueError(f"loc_2d must be rank-2, got shape={tuple(loc_2d.shape)}.")
+        if commit_lens.ndim != 1 or commit_lens.shape[0] != loc_2d.shape[0]:
+            raise ValueError(
+                "commit_lens must match loc_2d batch size: "
+                f"{tuple(commit_lens.shape)=} {tuple(loc_2d.shape)=}."
+            )
+
+        num_rows = int(loc_2d.numel())
+        if cache_k.shape[0] != num_rows or cache_v.shape[0] != num_rows:
+            raise ValueError(
+                "dense KV rows must match loc_2d size: "
+                f"{tuple(cache_k.shape)=} {tuple(cache_v.shape)=} {tuple(loc_2d.shape)=}."
+            )
+
+        if cache_k.dtype != self.dtype:
+            if k_scale is not None:
+                cache_k.div_(k_scale)
+            if v_scale is not None:
+                cache_v.div_(v_scale)
+            cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
+
+        if self.store_dtype != self.dtype:
+            cache_k = cache_k.contiguous().view(self.store_dtype)
+            cache_v = cache_v.contiguous().view(self.store_dtype)
+        else:
+            cache_k = cache_k.contiguous()
+            cache_v = cache_v.contiguous()
+
+        if loc_2d.device != self.k_buffer[0].device:
+            loc_2d = loc_2d.to(device=self.k_buffer[0].device, non_blocking=True)
+        if commit_lens.device != self.k_buffer[0].device:
+            commit_lens = commit_lens.to(
+                device=self.k_buffer[0].device, non_blocking=True
+            )
+        if loc_2d.dtype != torch.int64:
+            loc_2d = loc_2d.to(torch.int64)
+        if commit_lens.dtype != torch.int32:
+            commit_lens = commit_lens.to(torch.int32)
+
+        if not (_is_cuda or _is_hip):
+            row_offsets = torch.arange(loc_2d.shape[1], device=loc_2d.device)
+            valid_mask = row_offsets[None, :] < commit_lens.to(torch.int64)[:, None]
+            valid_idx = torch.nonzero(valid_mask.reshape(-1), as_tuple=False).flatten()
+            if valid_idx.numel() == 0:
+                return
+            self.set_kv_buffer(
+                layer,
+                loc_2d.reshape(-1).index_select(0, valid_idx),
+                cache_k.index_select(0, valid_idx),
+                cache_v.index_select(0, valid_idx),
+                k_scale,
+                v_scale,
+                layer_id_override=layer_id,
+            )
+            return
+
+        _set_kv_buffer_prefix_valid_impl(
+            cache_k,
+            cache_v,
+            self.k_buffer[layer_id - self.start_layer],
+            self.v_buffer[layer_id - self.start_layer],
+            loc_2d,
+            commit_lens,
+            row_dim=self.row_dim,
+            store_dtype=self.store_dtype,
+        )
+
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         # Zero-layer pool (e.g. all-SWA model's full sub-pool) has no buffers.
         if self.layer_num == 0:

python/sglang/srt/mem_cache/triton_ops/cache_move.py

@@ -2,6 +2,53 @@ import triton
 import triton.language as tl
 
 
+@triton.jit
+def set_kv_buffer_prefix_valid_tiled(
+    src_k_ptr,
+    src_v_ptr,
+    dst_k_ptr,
+    dst_v_ptr,
+    loc_2d_ptr,
+    commit_len_ptr,
+    src_k_row_stride,
+    src_v_row_stride,
+    dst_k_row_stride,
+    dst_v_row_stride,
+    block_size,
+    ROW_BYTES: tl.constexpr,
+    BYTES_PER_TILE: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    row = tl.program_id(1)
+    tid = tl.program_id(2)
+
+    commit_len = tl.load(commit_len_ptr + bid)
+    if row >= commit_len:
+        return
+
+    byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE)
+    mask_byte = byte_off < ROW_BYTES
+    tl.multiple_of(byte_off, 16)
+
+    loc = tl.load(loc_2d_ptr + bid * block_size + row)
+    src_row = bid * block_size + row
+
+    src_k_ptr = tl.cast(src_k_ptr, tl.pointer_type(tl.uint8))
+    src_v_ptr = tl.cast(src_v_ptr, tl.pointer_type(tl.uint8))
+    dst_k_ptr = tl.cast(dst_k_ptr, tl.pointer_type(tl.uint8))
+    dst_v_ptr = tl.cast(dst_v_ptr, tl.pointer_type(tl.uint8))
+
+    src_k_row_ptr = src_k_ptr + src_row * src_k_row_stride + byte_off
+    src_v_row_ptr = src_v_ptr + src_row * src_v_row_stride + byte_off
+    dst_k_row_ptr = dst_k_ptr + loc * dst_k_row_stride + byte_off
+    dst_v_row_ptr = dst_v_ptr + loc * dst_v_row_stride + byte_off
+
+    k_val = tl.load(src_k_row_ptr, mask=mask_byte, other=0)
+    v_val = tl.load(src_v_row_ptr, mask=mask_byte, other=0)
+    tl.store(dst_k_row_ptr, k_val, mask=mask_byte)
+    tl.store(dst_v_row_ptr, v_val, mask=mask_byte)
+
+
 @triton.jit
 def copy_all_layer_kv_cache_tiled(
     data_ptrs,

python/sglang/srt/server_args.py

@@ -1781,7 +1781,7 @@ class ServerArgs:
                 if self.speculative_algorithm == "STANDALONE":
                     # standalonedraft model and cuda graphs
                     reserved_mem += 6 * 1024
-                elif self.speculative_algorithm != "NGRAM":
+                elif self.speculative_algorithm not in {"NGRAM", "DFLASH"}:
                     # eagle draft models and cuda graphs
                     reserved_mem += 4 * 1024
 

python/sglang/srt/speculative/dflash_info.py

@@ -1,21 +1,25 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import List, Tuple
+from typing import TYPE_CHECKING, List, 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.layers.sampler import apply_custom_logit_processor
 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
+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,
@@ -23,6 +27,9 @@ from sglang.srt.speculative.dflash_utils import (
 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(
     *,
@@ -161,7 +168,7 @@ class DFlashVerifyInput(SpecInput):
     # Kept for compatibility with attention backends that gate tree metadata by `topk > 1`.
     # DFLASH verify is linear (non-tree), so this is always 1.
     topk: int = 1
-    # Custom attention "allow mask" for TARGET_VERIFY in backends that require it (e.g. triton).
+    # Custom attention "allow mask" for TARGET_VERIFY in backends that require it.
     # Semantics follow SGLang speculative conventions: True means the (q, k) pair is allowed.
     custom_mask: torch.Tensor | None = None
     capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
@@ -251,12 +258,53 @@ class DFlashVerifyInput(SpecInput):
             else torch.empty((0,), dtype=torch.bool, device=batch.device)
         )
 
+    def prepare_for_v2_verify(
+        self,
+        batch: ScheduleBatch,
+        target_worker: "TpModelWorker",
+    ) -> 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
+        metadata or eager attention metadata so the actual forward can run with
+        `skip_attn_backend_init=True`.
+        """
+        batch.input_ids = self.draft_token
+        batch.spec_info = self
+        batch.forward_mode = (
+            ForwardMode.IDLE
+            if batch.forward_mode.is_idle()
+            else ForwardMode.TARGET_VERIFY
+        )
+        batch.capture_hidden_mode = self.capture_hidden_mode
+        verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
+
+        can_run_cuda_graph = bool(
+            target_worker.model_runner.decode_cuda_graph_runner
+            and target_worker.model_runner.decode_cuda_graph_runner.can_run(
+                verify_forward_batch
+            )
+        )
+        if can_run_cuda_graph:
+            target_worker.model_runner.decode_cuda_graph_runner.replay_prepare(
+                verify_forward_batch
+            )
+        elif not batch.forward_mode.is_idle():
+            target_worker.model_runner.attn_backend.init_forward_metadata(
+                verify_forward_batch
+            )
+
+        return verify_forward_batch, can_run_cuda_graph
+
     def generate_attn_arg_prefill(
         self,
         req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
         req_to_token: torch.Tensor,
+        kv_start_idx: Optional[torch.Tensor] = None,
     ):
         device = req_pool_indices.device
         bs = len(req_pool_indices)
@@ -283,7 +331,7 @@ class DFlashVerifyInput(SpecInput):
             req_pool_indices,
             paged_kernel_lens,
             cum_kv_seq_len,
-            None,
+            kv_start_idx,
             kv_indices,
             req_to_token.size(1),
         )
@@ -339,28 +387,11 @@ class DFlashVerifyInput(SpecInput):
                     "DFLASH verify sampling_info size mismatch: "
                     f"len(sampling_info)={len(sampling_info)}, bs={bs}."
                 )
-
-            # Keep speculative verify semantics consistent with normal sampling path.
-            if sampling_info.has_custom_logit_processor:
-                apply_custom_logit_processor(
-                    logits_output.next_token_logits,
-                    sampling_info,
-                    num_tokens_in_batch=self.draft_token_num,
-                )
-
-            if (
-                sampling_info.penalizer_orchestrator.is_required
-                or sampling_info.logit_bias is not None
-            ):
-                linear_penalty = torch.zeros(
-                    (bs, logits_output.next_token_logits.shape[1]),
-                    dtype=torch.float32,
-                    device=device,
-                )
-                sampling_info.apply_logits_bias(linear_penalty)
-                logits_output.next_token_logits.add_(
-                    torch.repeat_interleave(linear_penalty, self.draft_token_num, dim=0)
-                )
+            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 (
@@ -368,10 +399,17 @@ class DFlashVerifyInput(SpecInput):
             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(

python/sglang/srt/speculative/dflash_info_v2.py

@@ -0,0 +1,350 @@
+"""DFLASH spec-v2 overlap scheduling data structures."""
+
+import contextlib
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+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.server_args import get_global_server_args
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
+from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
+from sglang.srt.utils.common import is_pin_memory_available
+
+_OVERLAP_PLAN_STREAMS: dict[str, torch.cuda.Stream] = {}
+
+
+def _get_overlap_plan_stream(
+    device: torch.device | str,
+) -> tuple[Optional[torch.cuda.Stream], contextlib.AbstractContextManager]:
+    """Return an optional plan stream/context for overlap scheduling prep kernels."""
+    if not envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
+        return None, contextlib.nullcontext()
+
+    device_str = str(device)
+    stream = _OVERLAP_PLAN_STREAMS.get(device_str)
+    if stream is None:
+        stream = torch.get_device_module(device_str).Stream()
+        _OVERLAP_PLAN_STREAMS[device_str] = stream
+    return stream, torch.get_device_module(device_str).stream(stream)
+
+
+@dataclass
+class DFlashDraftInputV2(SpecInput):
+    """Draft-side state carried across overlap iterations (spec-v2)."""
+
+    # Legacy Eagle-shaped fields kept only for dataclass compatibility. DFLASH
+    # overlap carries new_seq_lens / verified_id directly in the common
+    # no-shape-change path; FutureMap remains the fallback for filter/merge.
+    topk_p: torch.Tensor
+    topk_index: torch.Tensor
+    verified_id: torch.Tensor
+    new_seq_lens: torch.Tensor
+    hidden_states: torch.Tensor
+    verify_done: Optional[torch.cuda.Event] = None
+    max_top_k: int = 1
+    uniform_top_k_value: Optional[int] = None
+    cur_allocated_seq_lens_cpu: Optional[torch.Tensor] = None
+    planning_seq_lens_cpu: Optional[torch.Tensor] = None
+    planning_seq_lens_sum: Optional[int] = None
+    reserved_seq_lens_cpu: Optional[torch.Tensor] = None
+    reserved_seq_lens_sum: Optional[int] = None
+    direct_carry_valid: bool = True
+    _prepare_committed_kv_lens_cpu_buf: Optional[torch.Tensor] = None
+    _prepare_planning_kv_lens_cpu_buf: Optional[torch.Tensor] = None
+    _prepare_batch_seq_lens_cpu_buf: Optional[torch.Tensor] = None
+    _prepare_cur_kv_lens_cpu_buf: Optional[torch.Tensor] = None
+    _prepare_nxt_kv_lens_cpu_buf: Optional[torch.Tensor] = None
+    _prepare_cur_kv_lens_gpu_buf: Optional[torch.Tensor] = None
+    _prepare_nxt_kv_lens_gpu_buf: Optional[torch.Tensor] = None
+
+    # Filled by scheduler after dispatch.
+    future_indices: Optional[torch.Tensor] = None
+
+    def __post_init__(self):
+        super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT)
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        # Spec v2 draft state itself does not change token accounting.
+        return (1, 1)
+
+    def _ensure_prepare_length_buffers(
+        self, bs: int, device: torch.device | str
+    ) -> None:
+        pin_memory = is_pin_memory_available(device)
+
+        def needs_cpu_alloc(buf: Optional[torch.Tensor]) -> bool:
+            return buf is None or buf.numel() < bs or buf.is_pinned() != pin_memory
+
+        def needs_gpu_alloc(buf: Optional[torch.Tensor]) -> bool:
+            return buf is None or buf.numel() < bs or str(buf.device) != str(device)
+
+        def grown_capacity(buf: Optional[torch.Tensor]) -> int:
+            current = 0 if buf is None else int(buf.numel())
+            return max(bs, 32, current * 2 if current > 0 else 0)
+
+        if needs_cpu_alloc(self._prepare_committed_kv_lens_cpu_buf):
+            capacity = grown_capacity(self._prepare_committed_kv_lens_cpu_buf)
+            self._prepare_committed_kv_lens_cpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
+            )
+            self._prepare_planning_kv_lens_cpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
+            )
+            self._prepare_batch_seq_lens_cpu_buf = torch.empty(
+                (capacity,), dtype=torch.int64, device="cpu"
+            )
+            self._prepare_cur_kv_lens_cpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
+            )
+            self._prepare_nxt_kv_lens_cpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
+            )
+
+        if needs_gpu_alloc(self._prepare_cur_kv_lens_gpu_buf):
+            capacity = grown_capacity(self._prepare_cur_kv_lens_gpu_buf)
+            self._prepare_cur_kv_lens_gpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device=device
+            )
+            self._prepare_nxt_kv_lens_gpu_buf = torch.empty(
+                (capacity,), dtype=torch.int32, device=device
+            )
+
+    @classmethod
+    def create_idle_input(cls, device: torch.device) -> "DFlashDraftInputV2":
+        return cls(
+            topk_p=torch.empty((0, 0), device=device, dtype=torch.float32),
+            topk_index=torch.empty((0, 0), device=device, dtype=torch.int64),
+            verified_id=torch.empty((0,), device=device, dtype=torch.int32),
+            new_seq_lens=torch.empty((0,), device=device, dtype=torch.int64),
+            hidden_states=torch.empty((0, 0), device=device, dtype=torch.float16),
+            verify_done=None,
+        )
+
+    def prepare_for_decode(self, batch: ScheduleBatch):
+        """Allocate headroom in the shared req_to_token pool for the next DFLASH step.
+
+        DFLASH spec-v2 uses overlap scheduling's "over-allocation" approach: we reserve
+        future KV slots ahead of time so the worker can gather `out_cache_loc` directly
+        from `req_to_token` without allocator backup/restore. CPU metadata intentionally
+        lags by one iteration; keep it separate from the reserved upper bound that backs
+        the overallocated mapping.
+        """
+        plan_stream, plan_stream_ctx = _get_overlap_plan_stream(batch.device)
+        if plan_stream is None and self.verify_done is not None:
+            # Ensure previous forward is completed before mutating shared buffers.
+            self.verify_done.synchronize()
+
+        bs = batch.batch_size()
+        if bs == 0:
+            return
+        self._ensure_prepare_length_buffers(bs, batch.device)
+        assert self._prepare_committed_kv_lens_cpu_buf is not None
+        assert self._prepare_planning_kv_lens_cpu_buf is not None
+        assert self._prepare_batch_seq_lens_cpu_buf is not None
+        assert self._prepare_cur_kv_lens_cpu_buf is not None
+        assert self._prepare_nxt_kv_lens_cpu_buf is not None
+        assert self._prepare_cur_kv_lens_gpu_buf is not None
+        assert self._prepare_nxt_kv_lens_gpu_buf is not None
+        committed_kv_lens_cpu_t = self._prepare_committed_kv_lens_cpu_buf[:bs]
+        planning_kv_lens_cpu_t = self._prepare_planning_kv_lens_cpu_buf[:bs]
+        batch_seq_lens_cpu_t = self._prepare_batch_seq_lens_cpu_buf[:bs]
+        cur_kv_lens_cpu_t = self._prepare_cur_kv_lens_cpu_buf[:bs]
+        cur_allocated_seq_lens_cpu = self.cur_allocated_seq_lens_cpu
+
+        # For DFLASH, each decode step needs a fixed-size verify block.
+        block_size = int(get_global_server_args().speculative_num_draft_tokens)
+        if block_size <= 0:
+            raise ValueError(
+                f"DFLASH invalid speculative_num_draft_tokens={block_size}."
+            )
+        page_size = batch.token_to_kv_pool_allocator.page_size
+        nxt_kv_lens_cpu_t = self._prepare_nxt_kv_lens_cpu_buf[:bs]
+        committed_seq_lens_sum = 0
+        planning_seq_lens_sum = 0
+        reserved_seq_lens_sum = 0
+        num_needed_tokens = 0
+        max_top_k = 1
+        uniform_top_k_value = None
+        uniform_top_k = True
+        for i, req in enumerate(batch.reqs):
+            committed_len = int(req.kv_committed_len)
+            if cur_allocated_seq_lens_cpu is not None and i < len(
+                cur_allocated_seq_lens_cpu
+            ):
+                cur_alloc_len = int(cur_allocated_seq_lens_cpu[i])
+            else:
+                cur_alloc_len = int(req.kv_allocated_len)
+            planning_len = committed_len + block_size
+            reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)
+            top_k = int(req.sampling_params.top_k)
+
+            committed_kv_lens_cpu_t[i] = committed_len
+            batch_seq_lens_cpu_t[i] = committed_len
+            cur_kv_lens_cpu_t[i] = cur_alloc_len
+            planning_kv_lens_cpu_t[i] = planning_len
+            nxt_kv_lens_cpu_t[i] = reserved_len
+
+            committed_seq_lens_sum += committed_len
+            planning_seq_lens_sum += planning_len
+            reserved_seq_lens_sum += reserved_len
+            num_needed_tokens += reserved_len - cur_alloc_len
+
+            if top_k > max_top_k:
+                max_top_k = top_k
+            if i == 0:
+                uniform_top_k_value = top_k
+            elif uniform_top_k and top_k != uniform_top_k_value:
+                uniform_top_k = False
+
+        self.max_top_k = max(max_top_k, 1)
+        self.uniform_top_k_value = uniform_top_k_value if uniform_top_k else None
+
+        caller_stream = None
+        if plan_stream is not None:
+            caller_stream = torch.get_device_module(batch.device).current_stream()
+
+        with plan_stream_ctx:
+            if plan_stream is not None and caller_stream is not None:
+                # `batch.seq_lens`, `batch.req_pool_indices`, and related tensors may
+                # have just been rebuilt on the scheduler stream by filter/merge ops.
+                # The plan stream must wait for those writes before reading them.
+                plan_stream.wait_stream(caller_stream)
+
+            if plan_stream is not None and self.verify_done is not None:
+                plan_stream.wait_event(self.verify_done)
+
+            cur_kv_lens = self._prepare_cur_kv_lens_gpu_buf[:bs]
+            nxt_kv_lens = self._prepare_nxt_kv_lens_gpu_buf[:bs]
+            cur_kv_lens.copy_(cur_kv_lens_cpu_t, non_blocking=True)
+            nxt_kv_lens.copy_(nxt_kv_lens_cpu_t, non_blocking=True)
+
+            if num_needed_tokens > 0:
+                if page_size == 1:
+                    out_cache_loc = alloc_token_slots(
+                        batch.tree_cache, num_needed_tokens
+                    )
+                else:
+                    last_loc = get_last_loc(
+                        batch.req_to_token_pool.req_to_token,
+                        batch.req_pool_indices,
+                        cur_kv_lens,
+                    )
+                    out_cache_loc = alloc_paged_token_slots_extend(
+                        batch.tree_cache,
+                        cur_kv_lens,
+                        cur_kv_lens_cpu_t,
+                        nxt_kv_lens,
+                        nxt_kv_lens_cpu_t,
+                        last_loc,
+                        num_needed_tokens,
+                    )
+
+                # Updating req_to_token is a write to a shared tensor: it must not overlap
+                # with the previous batch's forward, which also reads req_to_token.
+                assign_req_to_token_pool_func(
+                    batch.req_pool_indices,
+                    batch.req_to_token_pool.req_to_token,
+                    cur_kv_lens,
+                    nxt_kv_lens,
+                    out_cache_loc,
+                    bs,
+                )
+        if caller_stream is not None:
+            # Enqueue the dependency on the caller's stream, not inside the
+            # plan-stream context, so forward work cannot observe partially
+            # prepared req_to_token / KV allocation state.
+            caller_stream.wait_stream(plan_stream)
+
+        # This request-side high-water mark is what release_kv_cache() uses to
+        # reclaim any DFLASH over-allocation if the request finishes later.
+        for i, req in enumerate(batch.reqs):
+            req.kv_allocated_len = max(req.kv_allocated_len, int(nxt_kv_lens_cpu_t[i]))
+
+        # Preserve the lagging committed CPU view on the batch and carry the
+        # tighter host-side planning bound separately from the full reserved
+        # allocator upper bound. Overlap scheduling only drifts by at most one
+        # DFlash block on the committed prefix lengths.
+        batch.seq_lens_cpu = batch_seq_lens_cpu_t
+        batch.seq_lens_sum = committed_seq_lens_sum
+        self.planning_seq_lens_cpu = planning_kv_lens_cpu_t
+        self.planning_seq_lens_sum = planning_seq_lens_sum
+        self.reserved_seq_lens_cpu = nxt_kv_lens_cpu_t
+        self.reserved_seq_lens_sum = reserved_seq_lens_sum
+
+    def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
+        if self.cur_allocated_seq_lens_cpu is not None:
+            self.cur_allocated_seq_lens_cpu = self.cur_allocated_seq_lens_cpu[
+                new_indices.cpu()
+            ]
+        if self.planning_seq_lens_cpu is not None:
+            self.planning_seq_lens_cpu = self.planning_seq_lens_cpu[new_indices.cpu()]
+            self.planning_seq_lens_sum = int(self.planning_seq_lens_cpu.sum().item())
+        if self.reserved_seq_lens_cpu is not None:
+            self.reserved_seq_lens_cpu = self.reserved_seq_lens_cpu[new_indices.cpu()]
+            self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item())
+
+        if self.future_indices is not None:
+            self.future_indices = self.future_indices[new_indices]
+            self.direct_carry_valid = False
+            return
+
+        self.topk_p = self.topk_p[new_indices]
+        self.topk_index = self.topk_index[new_indices]
+        self.verified_id = self.verified_id[new_indices]
+        self.new_seq_lens = self.new_seq_lens[new_indices]
+        self.hidden_states = self.hidden_states[new_indices]
+
+    def merge_batch(self, spec_info: "DFlashDraftInputV2"):
+        if self.cur_allocated_seq_lens_cpu is not None:
+            assert spec_info.cur_allocated_seq_lens_cpu is not None
+            self.cur_allocated_seq_lens_cpu = torch.cat(
+                [self.cur_allocated_seq_lens_cpu, spec_info.cur_allocated_seq_lens_cpu]
+            )
+        elif spec_info.cur_allocated_seq_lens_cpu is not None:
+            self.cur_allocated_seq_lens_cpu = spec_info.cur_allocated_seq_lens_cpu
+
+        if self.planning_seq_lens_cpu is not None:
+            assert spec_info.planning_seq_lens_cpu is not None
+            self.planning_seq_lens_cpu = torch.cat(
+                [self.planning_seq_lens_cpu, spec_info.planning_seq_lens_cpu]
+            )
+            self.planning_seq_lens_sum = int(self.planning_seq_lens_cpu.sum().item())
+        elif spec_info.planning_seq_lens_cpu is not None:
+            self.planning_seq_lens_cpu = spec_info.planning_seq_lens_cpu
+            self.planning_seq_lens_sum = spec_info.planning_seq_lens_sum
+
+        if self.reserved_seq_lens_cpu is not None:
+            assert spec_info.reserved_seq_lens_cpu is not None
+            self.reserved_seq_lens_cpu = torch.cat(
+                [self.reserved_seq_lens_cpu, spec_info.reserved_seq_lens_cpu]
+            )
+            self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item())
+        elif spec_info.reserved_seq_lens_cpu is not None:
+            self.reserved_seq_lens_cpu = spec_info.reserved_seq_lens_cpu
+            self.reserved_seq_lens_sum = spec_info.reserved_seq_lens_sum
+
+        if self.future_indices is not None:
+            assert spec_info.future_indices is not None
+            self.future_indices = torch.cat(
+                [self.future_indices, spec_info.future_indices]
+            )
+            self.direct_carry_valid = False
+            return
+
+        self.topk_p = torch.cat([self.topk_p, spec_info.topk_p], dim=0)
+        self.topk_index = torch.cat([self.topk_index, spec_info.topk_index], dim=0)
+        self.verified_id = torch.cat([self.verified_id, spec_info.verified_id], dim=0)
+        self.new_seq_lens = torch.cat(
+            [self.new_seq_lens, spec_info.new_seq_lens], dim=0
+        )
+        self.hidden_states = torch.cat(
+            [self.hidden_states, spec_info.hidden_states], dim=0
+        )

python/sglang/srt/speculative/dflash_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass
 from numbers import Integral
 from typing import Any, List, Optional, Tuple
@@ -7,12 +8,16 @@ from typing import Any, List, Optional, Tuple
 import torch
 import torch.nn.functional as F
 
+from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.layers.sampler import apply_custom_logit_processor
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.utils import is_cuda, is_musa
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
 
+logger = logging.getLogger(__name__)
+
 _DFLASH_SAMPLING_VERIFY_AVAILABLE = False
 _DFLASH_CHAIN_VERIFY_BUFFERS: dict[tuple[Optional[int], int], dict[str, Any]] = {}
 _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = frozenset(
@@ -20,6 +25,7 @@ _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = frozenset(
         "FlashInferAttnBackend",
         "FlashInferMLAAttnBackend",
         "FlashAttentionBackend",
+        "TritonAttnBackend",
         "TRTLLMHAAttnBackend",
         "TRTLLMMLABackend",
     }
@@ -49,6 +55,36 @@ def is_dflash_sampling_verify_available() -> bool:
     return _DFLASH_SAMPLING_VERIFY_AVAILABLE
 
 
+def resolve_dflash_prefill_refill_target(max_running_requests: int) -> int:
+    """Choose how many free running-request slots DFlash waits for before refill."""
+    override = envs.SGLANG_DFLASH_PREFILL_REFILL_TARGET.get()
+    if override is not None:
+        return override
+
+    max_running_requests = max(0, int(max_running_requests))
+    if max_running_requests < 8:
+        return 1
+    return min(4, max(2, (max_running_requests + 5) // 6))
+
+
+def should_delay_dflash_prefill_for_batching(
+    *,
+    running_bs: int,
+    num_allocatable_reqs: int,
+    max_running_requests: int,
+    prefill_refill_target: int,
+) -> bool:
+    if running_bs <= 0:
+        return False
+
+    target_prefill_bs = int(prefill_refill_target)
+    if target_prefill_bs <= 1:
+        return False
+
+    target_prefill_bs = min(target_prefill_bs, int(max_running_requests))
+    return int(num_allocatable_reqs) < target_prefill_bs
+
+
 def scale_kv_cell_size_per_token_for_dflash(
     *,
     target_cell_size_per_token: int,
@@ -101,6 +137,95 @@ def resolve_dflash_verify_mask_policy(attn_backend: Any) -> tuple[str, bool]:
     return backend_name, (backend_name not in _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS)
 
 
+def apply_dflash_verify_logits_adjustments(
+    *,
+    next_token_logits: torch.Tensor,
+    sampling_info: Any,
+    draft_token_num: int,
+) -> None:
+    """Apply sampling-time logit adjustments for DFlash verify in place.
+
+    This keeps v1 and v2 verify semantics aligned while letting overlap scheduling
+    use the cheaper precomputed `acc_linear_penalties` path instead of allocating a
+    repeated `[bs * draft_token_num, vocab]` penalty tensor every step.
+    """
+    if sampling_info is None:
+        return
+    if next_token_logits.ndim != 2:
+        raise ValueError(
+            "next_token_logits must be 2D, "
+            f"got shape={tuple(next_token_logits.shape)}."
+        )
+    if draft_token_num <= 0:
+        raise ValueError(f"draft_token_num must be positive, got {draft_token_num}.")
+
+    bs = len(sampling_info)
+    if next_token_logits.shape[0] != bs * draft_token_num:
+        raise ValueError(
+            "next_token_logits row count mismatch for DFlash verify adjustments. "
+            f"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}."
+        )
+
+    if sampling_info.has_custom_logit_processor:
+        apply_custom_logit_processor(
+            next_token_logits,
+            sampling_info,
+            num_tokens_in_batch=draft_token_num,
+        )
+
+    acc_linear_penalties = getattr(sampling_info, "acc_linear_penalties", None)
+    penalizer = getattr(sampling_info, "penalizer_orchestrator", None)
+    vocab_mask = getattr(sampling_info, "vocab_mask", None)
+    logit_bias = getattr(sampling_info, "logit_bias", None)
+
+    logits_3d: Optional[torch.Tensor] = None
+
+    def get_logits_3d() -> torch.Tensor:
+        nonlocal logits_3d
+        if logits_3d is None:
+            logits_3d = next_token_logits.reshape(bs, draft_token_num, -1)
+        return logits_3d
+
+    # Dense fallback only when we need live penalizer application or a vocab mask.
+    # In overlap scheduling the common path is `acc_linear_penalties`, which can be
+    # broadcast over the verify block without materializing a repeated buffer.
+    if (
+        penalizer is not None and penalizer.is_required and acc_linear_penalties is None
+    ) or vocab_mask is not None:
+        linear_penalty = torch.zeros(
+            (bs, next_token_logits.shape[1]),
+            dtype=torch.float32,
+            device=next_token_logits.device,
+        )
+        sampling_info.apply_logits_bias(linear_penalty)
+        get_logits_3d().add_(
+            linear_penalty[:, None, :].to(dtype=next_token_logits.dtype)
+        )
+        return
+
+    if acc_linear_penalties is not None:
+        if (
+            acc_linear_penalties.device != next_token_logits.device
+            or acc_linear_penalties.dtype != next_token_logits.dtype
+        ):
+            acc_linear_penalties = acc_linear_penalties.to(
+                device=next_token_logits.device,
+                dtype=next_token_logits.dtype,
+            )
+        get_logits_3d().add_(acc_linear_penalties[:, None, :])
+
+    if logit_bias is not None:
+        if (
+            logit_bias.device != next_token_logits.device
+            or logit_bias.dtype != next_token_logits.dtype
+        ):
+            logit_bias = logit_bias.to(
+                device=next_token_logits.device,
+                dtype=next_token_logits.dtype,
+            )
+        get_logits_3d().add_(logit_bias[:, None, :])
+
+
 def _get_or_create_chain_verify_buffers(
     *,
     bs: int,
@@ -465,6 +590,8 @@ def compute_dflash_sampling_correct_drafts_and_bonus(
     candidates: torch.Tensor,
     next_token_logits: torch.Tensor,
     sampling_info: Any,
+    max_top_k: Optional[int] = None,
+    uniform_top_k_value: Optional[int] = None,
     threshold_single: Optional[float] = None,
     threshold_acc: Optional[float] = None,
     uniform_samples: Optional[torch.Tensor] = None,
@@ -561,12 +688,19 @@ def compute_dflash_sampling_correct_drafts_and_bonus(
         ).to(dtype=torch.int64)
         vocab_size = int(scaled_logits.shape[-1])
         repeated_top_ks.clamp_(min=1, max=vocab_size)
-        max_top_k = int(repeated_top_ks.max().item())
+        if max_top_k is None:
+            max_top_k = int(repeated_top_ks.max().item())
+        else:
+            max_top_k = int(max_top_k)
+        if max_top_k < 1:
+            max_top_k = 1
+        elif max_top_k > vocab_size:
+            max_top_k = vocab_size
 
         # Sparse exact path for top-k/top-p (top-k-first semantics), then scatter to dense.
         if 0 < max_top_k < vocab_size:
             topk_logits, topk_indices = torch.topk(scaled_logits, k=max_top_k, dim=-1)
-            if not torch.all(repeated_top_ks == max_top_k):
+            if uniform_top_k_value is None or int(uniform_top_k_value) != max_top_k:
                 ranks = torch.arange(max_top_k, device=device, dtype=torch.int64)[
                     None, :
                 ]
@@ -639,10 +773,13 @@ def compute_dflash_sampling_correct_drafts_and_bonus(
     return correct_len, bonus
 
 
-def validate_dflash_request(req: Req) -> Optional[str]:
+def validate_dflash_request(req: Req, enable_overlap: bool) -> Optional[str]:
     if req.return_logprob:
         return "DFLASH speculative decoding does not support return_logprob yet."
 
+    if enable_overlap and req.return_hidden_states:
+        return "DFLASH speculative decoding does not support return_hidden_states yet."
+
     if (
         req.sampling_params.json_schema is not None
         or req.sampling_params.regex is not None

python/sglang/srt/speculative/dflash_worker.py

@@ -115,7 +115,7 @@ class DFlashWorker:
             _fb = "triton" if _torch.version.hip else "flashinfer"
             logger.warning(
                 "DFLASH draft worker does not support 'trtllm_mha' because the "
-                "draft path requires non-causal attention. Falling back to "
+                "draft path requires per-layer DFlash attention. Falling back to "
                 "'%s'.",
                 _fb,
             )
@@ -159,6 +159,8 @@ class DFlashWorker:
         )
         set_global_server_args_for_scheduler(saved_server_args)
         self.draft_model_runner = self.draft_worker.model_runner
+        # Keep the same alias that other spec-v2 workers expose.
+        self.draft_worker.draft_runner = self.draft_model_runner
         self.draft_model = self.draft_model_runner.model
         draft_config = parse_dflash_draft_config(
             draft_hf_config=self.draft_model_runner.model_config.hf_config
@@ -179,6 +181,7 @@ class DFlashWorker:
                     self.block_size,
                     model_block_size,
                 )
+        self.speculative_num_draft_tokens = int(self.block_size)
 
         self._mask_token = draft_config.mask_token
         self._mask_token_id_override = draft_config.mask_token_id
@@ -212,6 +215,9 @@ class DFlashWorker:
         self._draft_block_tokens_buf: Optional[torch.Tensor] = (
             None  # [cap_bs, block_size]
         )
+        self._draft_verify_out_cache_loc_buf: Optional[torch.Tensor] = (
+            None  # [cap_bs, block_size]
+        )
         self._draft_block_end_buf: Optional[torch.Tensor] = None  # [cap_bs]
         self._draft_seq_lens_cpu_buf: Optional[torch.Tensor] = None  # [cap_bs] on CPU
         self._draft_block_spec_info = DFlashVerifyInput(
@@ -224,11 +230,13 @@ class DFlashWorker:
         self._draft_greedy_gathered_max_buf: Optional[torch.Tensor] = None
         self._draft_greedy_gathered_ids_buf: Optional[torch.Tensor] = None
         self._draft_greedy_gather_cap: int = 0
+        self._draft_greedy_local_max_buf: Optional[torch.Tensor] = None
+        self._draft_greedy_local_arg_buf: Optional[torch.Tensor] = None
+        self._draft_greedy_local_cap: int = 0
         self._draft_greedy_best_rank_buf: Optional[torch.Tensor] = None
         self._draft_greedy_rank_index_buf: Optional[torch.Tensor] = None
         self._draft_greedy_selected_ids_buf: Optional[torch.Tensor] = None
         self._draft_greedy_index_cap: int = 0
-
         self._use_fused_kv_materialize = is_cuda()
         self._fused_kv_helper: Optional[object] = None
         if self._use_fused_kv_materialize:
@@ -296,6 +304,8 @@ class DFlashWorker:
                 num_kv_heads=first_attn.num_kv_heads,
                 head_dim=first_attn.head_dim,
                 device=self.device,
+                max_position_hint=self.target_worker.model_runner.model_config.context_len
+                + int(self.block_size),
             )
             if self.tp_rank == 0:
                 logger.info(
@@ -334,6 +344,9 @@ class DFlashWorker:
         self._draft_block_tokens_buf = torch.empty(
             (new_cap, block_size), dtype=torch.long, device=device
         )
+        self._draft_verify_out_cache_loc_buf = torch.empty(
+            (new_cap, block_size), dtype=torch.int64, device=device
+        )
         self._draft_block_end_buf = torch.empty(
             (new_cap,), dtype=torch.int32, device=device
         )
@@ -345,6 +358,11 @@ class DFlashWorker:
         # Delegate anything not implemented yet to the target worker.
         return getattr(self.target_worker, name)
 
+    def on_verify_complete_cpu(
+        self, num_correct_drafts_per_req: list[int], batch_size: int = 0
+    ) -> None:
+        pass
+
     def clear_cache_pool(self):
         # The target worker owns the shared KV allocator/cache. For the compact
         # sliding-window path, the draft req->token view is rebuilt from committed
@@ -562,7 +580,7 @@ class DFlashWorker:
                 "`shard_indices` attributes."
             )
 
-        # --- 2) Draft a non-causal block with the draft model.
+        # --- 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
@@ -570,38 +588,40 @@ class DFlashWorker:
         assert self._draft_block_end_buf is not None
         assert self._draft_seq_lens_cpu_buf is not None
 
-        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))
         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:
@@ -744,18 +764,51 @@ class DFlashWorker:
         def _cast_hs(x: torch.Tensor) -> torch.Tensor:
             return x if x.dtype == weight_dtype else x.to(weight_dtype)
 
+        def _ensure_local_reduce_buffers(
+            chunk_len: int,
+            value_dtype: torch.dtype,
+            device: torch.device,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            if (
+                self._draft_greedy_local_cap < chunk_len
+                or self._draft_greedy_local_max_buf is None
+                or self._draft_greedy_local_arg_buf is None
+                or self._draft_greedy_local_max_buf.dtype != value_dtype
+                or self._draft_greedy_local_max_buf.device != device
+                or self._draft_greedy_local_arg_buf.device != device
+            ):
+                cap = max(int(chunk_size), chunk_len)
+                self._draft_greedy_local_max_buf = torch.empty(
+                    (cap,), dtype=value_dtype, device=device
+                )
+                self._draft_greedy_local_arg_buf = torch.empty(
+                    (cap,), dtype=torch.int64, device=device
+                )
+                self._draft_greedy_local_cap = cap
+            return (
+                self._draft_greedy_local_max_buf[:chunk_len],
+                self._draft_greedy_local_arg_buf[:chunk_len],
+            )
+
         # Fast path (common): single-rank greedy sampling over the base vocab shard.
         # Avoids extra max/id bookkeeping that is only needed for TP sync or added vocab.
+        #
+        # DFLASH draft sampling only materializes a small fixed block of hidden states
+        # each step. On tp=1, splitting those states into many 256-token chunks adds
+        # extra matmul/argmax launches without reducing peak memory meaningfully.
         if tp_size == 1 and num_added == 0:
-            for start in range(0, num_tokens, int(chunk_size)):
-                end = min(num_tokens, start + int(chunk_size))
+            fast_chunk_size = max(int(chunk_size), 1024)
+            for start in range(0, num_tokens, fast_chunk_size):
+                end = min(num_tokens, start + fast_chunk_size)
                 hs = _cast_hs(hidden_states[start:end])
                 if num_org > 0:
                     base_logits = torch.matmul(hs, weight[:num_org].T)
-                    out_tokens[start:end] = (
-                        torch.argmax(base_logits, dim=-1).to(torch.long)
-                        + org_vocab_start
+                    local_max, local_arg = _ensure_local_reduce_buffers(
+                        end - start, base_logits.dtype, hs.device
                     )
+                    torch.max(base_logits, dim=-1, out=(local_max, local_arg))
+                    out_tokens[start:end].copy_(local_arg)
+                    out_tokens[start:end].add_(org_vocab_start)
                 else:
                     out_tokens[start:end] = 0
             return out_tokens
@@ -768,7 +821,10 @@ class DFlashWorker:
             # Base vocab logits.
             if num_org > 0:
                 base_logits = torch.matmul(hs, weight[:num_org].T)
-                local_max, local_arg = torch.max(base_logits, dim=-1)
+                local_max, local_arg = _ensure_local_reduce_buffers(
+                    chunk_len, base_logits.dtype, hs.device
+                )
+                torch.max(base_logits, dim=-1, out=(local_max, local_arg))
             else:
                 local_max = torch.full(
                     (chunk_len,),
@@ -966,11 +1022,13 @@ class DFlashWorker:
                     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",
@@ -978,10 +1036,7 @@ class DFlashWorker:
                     )
                     self._use_fused_kv_materialize = False
                     self._fused_kv_helper = None
-                    self._append_target_hidden_sequential(
-                        ctx_hidden, ctx_positions, ctx_cache_loc
-                    )
-            else:
+            if not wrote_with_fused_kv:
                 self._append_target_hidden_sequential(
                     ctx_hidden, ctx_positions, ctx_cache_loc
                 )
@@ -1011,6 +1066,160 @@ class DFlashWorker:
         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,
+        *,
+        target_hidden: torch.Tensor,
+        cache_loc: torch.Tensor,
+        positions: torch.Tensor,
+        cache_loc_2d: Optional[torch.Tensor] = None,
+        commit_lens: Optional[torch.Tensor] = None,
+    ) -> None:
+        """Materialize target context features into the draft KV cache at explicit slots.
+
+        For the spec-v2 overlap path, callers can pass dense `[bs, block_size]`
+        `cache_loc_2d` plus `commit_lens`; the prefix-valid writer then commits
+        only the live prefix rows without constructing masked/packed index tensors.
+        """
+        if target_hidden is None:
+            raise RuntimeError("DFLASH missing target hidden context features.")
+        if target_hidden.numel() == 0:
+            return
+        if target_hidden.ndim != 2:
+            raise ValueError(
+                "DFLASH target_hidden must be 2D, "
+                f"got shape={tuple(target_hidden.shape)}."
+            )
+
+        if cache_loc.ndim != 1:
+            raise ValueError(
+                f"DFLASH cache_loc must be 1D, got shape={tuple(cache_loc.shape)}."
+            )
+        if positions.ndim != 1:
+            raise ValueError(
+                f"DFLASH positions must be 1D, got shape={tuple(positions.shape)}."
+            )
+        num_tokens = int(target_hidden.shape[0])
+        if int(cache_loc.numel()) != num_tokens:
+            raise ValueError(
+                "DFLASH cache_loc length mismatch: "
+                f"cache_loc={int(cache_loc.numel())}, target_hidden={num_tokens}."
+            )
+        if int(positions.numel()) != num_tokens:
+            raise ValueError(
+                "DFLASH positions length mismatch: "
+                f"positions={int(positions.numel())}, target_hidden={num_tokens}."
+            )
+        if cache_loc_2d is not None:
+            if cache_loc_2d.ndim != 2:
+                raise ValueError(
+                    "DFLASH cache_loc_2d must be 2D, "
+                    f"got shape={tuple(cache_loc_2d.shape)}."
+                )
+            if int(cache_loc_2d.numel()) != num_tokens:
+                raise ValueError(
+                    "DFLASH cache_loc_2d size mismatch: "
+                    f"cache_loc_2d={int(cache_loc_2d.numel())}, target_hidden={num_tokens}."
+                )
+            if commit_lens is None:
+                raise ValueError(
+                    "DFLASH cache_loc_2d requires commit_lens for prefix-valid writes."
+                )
+
+        device = self.model_runner.device
+        if cache_loc.device != device:
+            cache_loc = cache_loc.to(device, non_blocking=True)
+        if positions.device != device:
+            positions = positions.to(device, non_blocking=True)
+        if target_hidden.device != device:
+            target_hidden = target_hidden.to(device, non_blocking=True)
+
+        if cache_loc.dtype != torch.int64:
+            cache_loc = cache_loc.to(torch.int64)
+        if positions.dtype != torch.int64:
+            positions = positions.to(torch.int64)
+        if cache_loc_2d is not None:
+            if cache_loc_2d.device != device:
+                cache_loc_2d = cache_loc_2d.to(device, non_blocking=True)
+            if cache_loc_2d.dtype != torch.int64:
+                cache_loc_2d = cache_loc_2d.to(torch.int64)
+        if commit_lens is not None:
+            if commit_lens.device != device:
+                commit_lens = commit_lens.to(device, non_blocking=True)
+            if commit_lens.dtype != torch.int32:
+                commit_lens = commit_lens.to(torch.int32)
+
+        with torch.inference_mode():
+            ctx_hidden = self.draft_model.project_target_hidden(target_hidden)
+
+            if cache_loc_2d is not None:
+                bs = int(commit_lens.shape[0])
+                if int(cache_loc_2d.shape[0]) != bs:
+                    raise ValueError(
+                        "DFLASH cache_loc_2d batch size mismatch: "
+                        f"cache_loc_2d={tuple(cache_loc_2d.shape)}, commit_lens={tuple(commit_lens.shape)}."
+                    )
+                if bs == 0:
+                    return
+                if self._use_fused_kv_materialize and self._fused_kv_helper is not None:
+                    try:
+                        self._append_target_hidden_fused(
+                            ctx_hidden=ctx_hidden,
+                            ctx_positions=positions,
+                            ctx_cache_loc=cache_loc,
+                            ctx_cache_loc_2d=cache_loc_2d,
+                            commit_lens=commit_lens,
+                        )
+                        return
+                    except Exception as e:
+                        logger.warning(
+                            "DFLASH fused prefix-direct KV append failed; falling back to the per-layer prefix-direct path: %s",
+                            e,
+                        )
+                        self._use_fused_kv_materialize = False
+                        self._fused_kv_helper = None
+
+                for layer in self.draft_model.layers:
+                    attn = layer.self_attn
+                    k, v = attn.kv_proj_only(ctx_hidden)
+                    k = attn.apply_k_norm(k)
+                    k = attn.apply_k_rope(positions, k)
+                    k = k.view(-1, attn.num_kv_heads, attn.head_dim)
+                    v = v.view(-1, attn.num_kv_heads, attn.head_dim)
+
+                    self.draft_model_runner.token_to_kv_pool.set_kv_buffer_prefix_valid(
+                        attn.attn,
+                        cache_loc_2d,
+                        commit_lens,
+                        k,
+                        v,
+                        attn.attn.k_scale,
+                        attn.attn.v_scale,
+                    )
+                return
+
+            if self._use_fused_kv_materialize and self._fused_kv_helper is not None:
+                try:
+                    self._append_target_hidden_fused(
+                        ctx_hidden=ctx_hidden,
+                        ctx_positions=positions,
+                        ctx_cache_loc=cache_loc,
+                    )
+                    return
+                except Exception as e:
+                    logger.warning(
+                        "DFLASH fused KV append-by-loc failed; falling back to sequential path: %s",
+                        e,
+                    )
+                    self._use_fused_kv_materialize = False
+                    self._fused_kv_helper = None
+
+            self._append_target_hidden_sequential(
+                ctx_hidden=ctx_hidden,
+                ctx_positions=positions,
+                ctx_cache_loc=cache_loc,
+            )
+
     def _append_target_hidden_sequential(
         self,
         ctx_hidden: torch.Tensor,
@@ -1041,23 +1250,39 @@ class DFlashWorker:
         ctx_hidden: torch.Tensor,
         ctx_positions: torch.Tensor,
         ctx_cache_loc: torch.Tensor,
+        ctx_cache_loc_2d: Optional[torch.Tensor] = None,
+        commit_lens: Optional[torch.Tensor] = None,
     ) -> None:
         """Fused KV materialization using batched projection + Triton kernel."""
         token_to_kv_pool = self.draft_model_runner.token_to_kv_pool
-        layers = self.draft_model.layers
+        if self._fused_kv_helper is None:
+            raise RuntimeError("DFLASH fused KV helper is not initialized.")
 
         def _write_layer_kv(
-            layer_idx: int, cache_k: torch.Tensor, cache_v: torch.Tensor
+            layer_idx: int,
+            cache_k: torch.Tensor,
+            cache_v: torch.Tensor,
         ) -> None:
-            attn = layers[layer_idx].self_attn.attn
-            token_to_kv_pool.set_kv_buffer(
-                attn,
-                ctx_cache_loc,
-                cache_k,
-                cache_v,
-                attn.k_scale,
-                attn.v_scale,
-            )
+            attn = self.draft_model.layers[layer_idx].self_attn.attn
+            if ctx_cache_loc_2d is not None and commit_lens is not None:
+                token_to_kv_pool.set_kv_buffer_prefix_valid(
+                    attn,
+                    ctx_cache_loc_2d,
+                    commit_lens,
+                    cache_k,
+                    cache_v,
+                    attn.k_scale,
+                    attn.v_scale,
+                )
+            else:
+                token_to_kv_pool.set_kv_buffer(
+                    attn,
+                    ctx_cache_loc,
+                    cache_k,
+                    cache_v,
+                    attn.k_scale,
+                    attn.v_scale,
+                )
 
         self._fused_kv_helper.materialize(
             ctx_hidden=ctx_hidden,

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -0,0 +1,680 @@
+import logging
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler import GenerationBatchResult
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+    compute_position,
+)
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.dflash_info import DFlashVerifyInput
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+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.dflash_worker import DFlashWorker
+from sglang.srt.speculative.eagle_info_v2 import assign_extend_cache_locs_func
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
+from sglang.srt.speculative.triton_ops.dflash_accept_bonus import (
+    _compute_dflash_accept_bonus_triton_unchecked,
+)
+from sglang.srt.speculative.triton_ops.dflash_prepare_block import (
+    _prepare_dflash_draft_block_unchecked,
+)
+from sglang.srt.utils import is_cuda, is_hip
+
+logger = logging.getLogger(__name__)
+
+
+class DFlashWorkerV2(DFlashWorker):
+    """DFLASH speculative decoding worker (spec-v2 overlap scheduling).
+
+    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.
+    """
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: Optional[int],
+        moe_ep_rank: int,
+        attn_cp_rank: int,
+        moe_dp_rank: int,
+        nccl_port: int,
+        target_worker: TpModelWorker,
+    ):
+        super().__init__(
+            server_args=server_args,
+            gpu_id=gpu_id,
+            tp_rank=tp_rank,
+            dp_rank=dp_rank,
+            moe_ep_rank=moe_ep_rank,
+            attn_cp_rank=attn_cp_rank,
+            moe_dp_rank=moe_dp_rank,
+            nccl_port=nccl_port,
+            target_worker=target_worker,
+        )
+        supports_gpu_triton = is_cuda() or is_hip()
+        self._use_triton_prepare_block = supports_gpu_triton
+        self._use_triton_accept_bonus = supports_gpu_triton
+        self._accept_bonus_buffer_cap: int = 0
+        self._accept_bonus_buffer_slot: int = 0
+        self._accept_len_buf: Optional[torch.Tensor] = None
+        self._commit_lens_bufs: List[torch.Tensor] = []
+        self._bonus_id_bufs: List[torch.Tensor] = []
+        self._out_tokens_bufs: List[torch.Tensor] = []
+        self._new_seq_lens_bufs: List[torch.Tensor] = []
+
+    def _ensure_accept_bonus_buffers(self, bs: int) -> None:
+        if self._accept_bonus_buffer_cap >= int(bs):
+            return
+
+        new_cap = max(
+            int(bs),
+            (
+                self._accept_bonus_buffer_cap * 2
+                if self._accept_bonus_buffer_cap > 0
+                else int(bs)
+            ),
+        )
+        device = self.device
+        block_size = int(self.block_size)
+        self._accept_len_buf = torch.empty((new_cap,), dtype=torch.int32, device=device)
+        self._commit_lens_bufs = [
+            torch.empty((new_cap,), dtype=torch.int32, device=device) for _ in range(2)
+        ]
+        self._bonus_id_bufs = [
+            torch.empty((new_cap,), dtype=torch.int32, device=device) for _ in range(2)
+        ]
+        self._out_tokens_bufs = [
+            torch.empty((new_cap, block_size), dtype=torch.int64, device=device)
+            for _ in range(2)
+        ]
+        self._new_seq_lens_bufs = [
+            torch.empty((new_cap,), dtype=torch.int64, device=device) for _ in range(2)
+        ]
+        self._accept_bonus_buffer_cap = new_cap
+
+    def _next_accept_bonus_buffers(self, bs: int) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        self._ensure_accept_bonus_buffers(bs)
+        assert self._accept_len_buf is not None
+        slot = self._accept_bonus_buffer_slot
+        self._accept_bonus_buffer_slot = (slot + 1) % 2
+        return (
+            self._accept_len_buf[:bs],
+            self._commit_lens_bufs[slot][:bs],
+            self._bonus_id_bufs[slot][:bs],
+            self._out_tokens_bufs[slot][:bs],
+            self._new_seq_lens_bufs[slot][:bs],
+        )
+
+    def _validate_phase1_sampling_support(
+        self, model_worker_batch: ScheduleBatch
+    ) -> None:
+        sampling_info = model_worker_batch.sampling_info
+        if sampling_info is None or sampling_info.is_all_greedy:
+            return
+
+        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
+
+    def _make_next_draft_input_prefill(
+        self,
+        *,
+        verified_id: torch.Tensor,
+        seq_lens: torch.Tensor,
+        verify_done: Optional[torch.cuda.Event] = None,
+        cur_allocated_seq_lens_cpu: Optional[torch.Tensor] = None,
+    ) -> DFlashDraftInputV2:
+        bs = int(seq_lens.numel())
+        device = verified_id.device
+        return DFlashDraftInputV2(
+            topk_p=torch.empty((bs, 0), device=device, dtype=torch.float32),
+            topk_index=torch.empty((bs, 0), device=device, dtype=torch.int64),
+            verified_id=verified_id.to(dtype=torch.int32),
+            new_seq_lens=seq_lens.to(dtype=torch.int64),
+            hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16),
+            verify_done=verify_done,
+            cur_allocated_seq_lens_cpu=cur_allocated_seq_lens_cpu,
+        )
+
+    def _make_next_draft_input_decode(
+        self,
+        *,
+        verified_id: torch.Tensor,
+        new_seq_lens: torch.Tensor,
+        verify_done: Optional[torch.cuda.Event] = None,
+        cur_allocated_seq_lens_cpu: Optional[torch.Tensor] = None,
+    ) -> DFlashDraftInputV2:
+        bs = int(new_seq_lens.numel())
+        device = verified_id.device
+        return DFlashDraftInputV2(
+            topk_p=torch.empty((bs, 0), device=device, dtype=torch.float32),
+            topk_index=torch.empty((bs, 0), device=device, dtype=torch.int64),
+            verified_id=verified_id.to(dtype=torch.int32),
+            new_seq_lens=new_seq_lens.to(dtype=torch.int64),
+            hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16),
+            verify_done=verify_done,
+            cur_allocated_seq_lens_cpu=cur_allocated_seq_lens_cpu,
+        )
+
+    def forward_batch_generation(
+        self,
+        model_worker_batch: ScheduleBatch,
+        on_publish=None,
+    ) -> GenerationBatchResult:
+        if getattr(model_worker_batch, "return_logprob", False):
+            raise ValueError(
+                "DFLASH speculative decoding does not support return_logprob yet."
+            )
+        self._validate_phase1_sampling_support(model_worker_batch)
+
+        if (
+            model_worker_batch.forward_mode.is_extend()
+            or model_worker_batch.is_extend_in_batch
+        ):
+            # Target prefill: capture DFlash aux hidden states for prompt tokens.
+            model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
+            batch_output = self.target_worker.forward_batch_generation(
+                model_worker_batch
+            )
+
+            logits_output, next_token_ids = (
+                batch_output.logits_output,
+                batch_output.next_token_ids,
+            )
+            batch_output.new_seq_lens = model_worker_batch.seq_lens
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+
+            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 (
+                model_worker_batch.extend_lens is None
+                or model_worker_batch.prefix_lens is None
+            ):
+                raise RuntimeError(
+                    "DFLASH expected extend_lens / prefix_lens to be populated in extend mode, "
+                    "but got None."
+                )
+
+            # Materialize prompt tokens into the draft KV cache immediately. This is required
+            # for radix cache safety (the scheduler may update radix after prefill returns).
+            device = next_token_ids.device
+            ctx_lens = torch.tensor(
+                model_worker_batch.extend_lens, dtype=torch.int32, device=device
+            )
+            draft_seq_lens = torch.tensor(
+                model_worker_batch.prefix_lens, dtype=torch.int32, device=device
+            )
+
+            if model_worker_batch.out_cache_loc is None:
+                raise RuntimeError(
+                    "DFLASH prefill expected out_cache_loc, but got None."
+                )
+            positions, _ = compute_position(
+                self.model_runner.server_args.attention_backend,
+                draft_seq_lens,
+                ctx_lens,
+                int(sum(model_worker_batch.extend_lens)),
+            )
+            self._append_target_hidden_to_draft_kv_by_loc(
+                target_hidden=logits_output.hidden_states,
+                cache_loc=model_worker_batch.out_cache_loc,
+                positions=positions,
+            )
+
+            # Avoid copying large hidden-state buffers to CPU in overlap scheduling.
+            logits_output.hidden_states = None
+
+            batch_output.next_draft_input = self._make_next_draft_input_prefill(
+                verified_id=next_token_ids,
+                seq_lens=model_worker_batch.seq_lens,
+                cur_allocated_seq_lens_cpu=model_worker_batch.seq_lens_cpu,
+            )
+            verify_done = torch.get_device_module(device).Event()
+            verify_done.record()
+            batch_output.next_draft_input.verify_done = verify_done
+            return batch_output
+
+        # Decode / target-verify stage.
+        if model_worker_batch.spec_info is None:
+            model_worker_batch.spec_info = DFlashDraftInputV2.create_idle_input(
+                device=self.device
+            )
+
+        draft_input = model_worker_batch.spec_info
+        if not isinstance(draft_input, DFlashDraftInputV2):
+            raise RuntimeError(
+                "DFLASH spec-v2 expected DFlashDraftInputV2 state on the running batch."
+            )
+
+        if model_worker_batch.forward_mode.is_idle():
+            empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device)
+            empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device)
+            next_draft_input = self._make_next_draft_input_decode(
+                verified_id=torch.empty((0,), device=self.device, dtype=torch.int32),
+                new_seq_lens=torch.empty((0,), device=self.device, dtype=torch.int64),
+            )
+            if on_publish is not None:
+                on_publish(next_draft_input.new_seq_lens)
+            verify_done = torch.get_device_module(self.device).Event()
+            verify_done.record()
+            next_draft_input.verify_done = verify_done
+            return GenerationBatchResult(
+                logits_output=None,
+                next_token_ids=empty_ids,
+                accept_lens=empty_lens,
+                next_draft_input=next_draft_input,
+                can_run_cuda_graph=False,
+                speculative_num_draft_tokens=int(self.block_size),
+                new_seq_lens=next_draft_input.new_seq_lens,
+            )
+
+        # `seq_lens` is carried over from the previous overlap iteration and may have been
+        # produced on another stream.
+        model_worker_batch.seq_lens.record_stream(
+            torch.get_device_module(self.device).current_stream()
+        )
+
+        bs = len(model_worker_batch.seq_lens)
+        device = self.device
+
+        # --- 1) Draft a fixed block with the draft model.
+        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"):
+            raise RuntimeError(
+                "DFLASH requires the target model to expose `lm_head` with `weight`."
+            )
+
+        block_size = int(self.block_size)
+        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_verify_out_cache_loc_buf is not None
+        assert self._draft_block_end_buf is not None
+        assert self._draft_seq_lens_cpu_buf is not None
+
+        block_ids = self._draft_block_ids_buf[:bs]
+        prefix_lens = model_worker_batch.seq_lens
+        positions_2d = self._draft_block_positions_buf[:bs]
+        verify_out_cache_loc_2d = self._draft_verify_out_cache_loc_buf[:bs]
+        if self._use_triton_prepare_block:
+            try:
+                _prepare_dflash_draft_block_unchecked(
+                    verified_id=draft_input.verified_id.view(-1),
+                    prefix_lens=prefix_lens.view(-1),
+                    req_pool_indices=model_worker_batch.req_pool_indices.view(-1),
+                    req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                    block_ids_out=block_ids,
+                    positions_out=positions_2d,
+                    cache_loc_out=verify_out_cache_loc_2d,
+                    mask_token_id=int(self._mask_token_id),
+                )
+            except Exception as e:
+                self._use_triton_prepare_block = False
+                logger.warning(
+                    "DFLASH Triton prepare_block failed; falling back to eager path: %s",
+                    e,
+                )
+                block_ids.fill_(int(self._mask_token_id))
+                block_ids[:, 0].copy_(draft_input.verified_id)
+                torch.add(
+                    prefix_lens.unsqueeze(1),
+                    self._block_pos_offsets,
+                    out=positions_2d,
+                )
+                end_offset = prefix_lens + block_size
+                verify_out_cache_loc = assign_extend_cache_locs_func(
+                    req_pool_indices=model_worker_batch.req_pool_indices,
+                    req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                    start_offset=prefix_lens,
+                    end_offset=end_offset,
+                    batch_size=bs,
+                    draft_token_num=block_size,
+                    device=device,
+                )
+                verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size))
+        else:
+            block_ids.fill_(int(self._mask_token_id))
+            block_ids[:, 0].copy_(draft_input.verified_id)
+            torch.add(
+                prefix_lens.unsqueeze(1),
+                self._block_pos_offsets,
+                out=positions_2d,
+            )
+            end_offset = prefix_lens + block_size
+            verify_out_cache_loc = assign_extend_cache_locs_func(
+                req_pool_indices=model_worker_batch.req_pool_indices,
+                req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                start_offset=prefix_lens,
+                end_offset=end_offset,
+                batch_size=bs,
+                draft_token_num=block_size,
+                device=device,
+            )
+            verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size))
+
+        noise_embedding = embed_module(block_ids)
+        input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1])
+
+        positions = positions_2d.reshape(-1)
+        verify_out_cache_loc = verify_out_cache_loc_2d.reshape(-1)
+
+        seq_lens_cpu = self._draft_seq_lens_cpu_buf[:bs]
+        if self.use_compact_draft_cache:
+            # Rebuild the draft-local sliding-window view from committed target state.
+            draft_prefix_lens = self._compute_compact_draft_seq_lens(prefix_lens)
+            seq_lens_cpu.copy_(draft_prefix_lens.to(device="cpu", dtype=torch.int32))
+
+            suffix_start = prefix_lens.to(torch.int64) - draft_prefix_lens.to(
+                torch.int64
+            )
+            suffix_cache_loc = self._gather_req_to_token_segments(
+                req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                req_pool_indices=model_worker_batch.req_pool_indices,
+                start=suffix_start,
+                lengths=draft_prefix_lens,
+            )
+            assign_req_to_token_pool_func(
+                model_worker_batch.req_pool_indices,
+                self.draft_model_runner.req_to_token_pool.req_to_token,
+                torch.zeros_like(draft_prefix_lens),
+                draft_prefix_lens,
+                suffix_cache_loc,
+                bs,
+            )
+
+            block_end = self._draft_block_end_buf[:bs]
+            torch.add(draft_prefix_lens, block_size, out=block_end)
+            assign_req_to_token_pool_func(
+                model_worker_batch.req_pool_indices,
+                self.draft_model_runner.req_to_token_pool.req_to_token,
+                draft_prefix_lens,
+                block_end,
+                verify_out_cache_loc,
+                bs,
+            )
+            draft_seq_lens = draft_prefix_lens
+            draft_seq_lens_sum = int(seq_lens_cpu.sum().item())
+        else:
+            # Non-windowed path uses the shared overallocated mapping directly.
+            # Backend planning only needs a safe upper bound for the committed
+            # prefix lengths, not the full allocator reservation length.
+            draft_seq_lens = prefix_lens
+            if draft_input.planning_seq_lens_cpu is not None:
+                seq_lens_cpu.copy_(draft_input.planning_seq_lens_cpu)
+                draft_seq_lens_sum = int(draft_input.planning_seq_lens_sum)
+            elif draft_input.reserved_seq_lens_cpu is not None:
+                seq_lens_cpu.copy_(draft_input.reserved_seq_lens_cpu)
+                draft_seq_lens_sum = int(draft_input.reserved_seq_lens_sum)
+            elif model_worker_batch.seq_lens_cpu is not None:
+                seq_lens_cpu.copy_(model_worker_batch.seq_lens_cpu)
+                draft_seq_lens_sum = (
+                    int(model_worker_batch.seq_lens_sum)
+                    if model_worker_batch.seq_lens_sum is not None
+                    else int(model_worker_batch.seq_lens_cpu.sum())
+                )
+            else:
+                seq_lens_cpu.copy_(prefix_lens.to("cpu", dtype=torch.int32))
+                draft_seq_lens_sum = int(prefix_lens.sum().item())
+
+        forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            batch_size=bs,
+            input_ids=block_ids.flatten(),
+            req_pool_indices=model_worker_batch.req_pool_indices,
+            seq_lens=draft_seq_lens,
+            out_cache_loc=verify_out_cache_loc,
+            seq_lens_sum=draft_seq_lens_sum,
+            seq_lens_cpu=seq_lens_cpu,
+            positions=positions,
+            input_embeds=input_embeds,
+            spec_algorithm=SpeculativeAlgorithm.DFLASH,
+            spec_info=self._draft_block_spec_info,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+        )
+
+        with torch.inference_mode():
+            draft_logits_output = self.draft_model_runner.forward(
+                forward_batch
+            ).logits_output
+
+        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, int(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, int(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)
+
+        # --- 2) Target verify.
+        # TARGET_VERIFY uses standard causal masking; custom masks are unnecessary here.
+        custom_mask = None
+
+        verify_input_ids = draft_tokens.reshape(-1)
+        verify_input = DFlashVerifyInput(
+            draft_token=verify_input_ids,
+            positions=positions,
+            draft_token_num=int(self.block_size),
+            custom_mask=custom_mask,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+        )
+
+        model_worker_batch.out_cache_loc = verify_out_cache_loc
+        sampling_info = model_worker_batch.sampling_info
+
+        need_mamba_verify_commit = hasattr(
+            self.target_worker.model_runner.attn_backend,
+            "update_mamba_state_after_mtp_verify",
+        )
+        seq_lens_pre_verify = (
+            model_worker_batch.seq_lens.clone() if need_mamba_verify_commit else None
+        )
+        seq_lens_cpu_backup = model_worker_batch.seq_lens_cpu
+        seq_lens_sum_backup = model_worker_batch.seq_lens_sum
+        if draft_input.planning_seq_lens_cpu is not None:
+            model_worker_batch.seq_lens_cpu = draft_input.planning_seq_lens_cpu
+            model_worker_batch.seq_lens_sum = int(draft_input.planning_seq_lens_sum)
+        elif draft_input.reserved_seq_lens_cpu is not None:
+            model_worker_batch.seq_lens_cpu = draft_input.reserved_seq_lens_cpu
+            model_worker_batch.seq_lens_sum = int(draft_input.reserved_seq_lens_sum)
+
+        verify_forward_batch, _ = verify_input.prepare_for_v2_verify(
+            model_worker_batch, self.target_worker
+        )
+        model_worker_batch.seq_lens_cpu = seq_lens_cpu_backup
+        model_worker_batch.seq_lens_sum = seq_lens_sum_backup
+
+        target_out = self.target_worker.forward_batch_generation(
+            batch=None,
+            forward_batch=verify_forward_batch,
+            is_verify=True,
+            skip_attn_backend_init=True,
+        )
+        logits_output = target_out.logits_output
+        can_run_cuda_graph = target_out.can_run_cuda_graph
+
+        if sampling_info is not None:
+            apply_dflash_verify_logits_adjustments(
+                next_token_logits=logits_output.next_token_logits,
+                sampling_info=sampling_info,
+                draft_token_num=int(self.block_size),
+            )
+
+        candidates = draft_tokens
+        new_seq_lens = None
+        if (
+            sampling_info is not None
+            and not sampling_info.is_all_greedy
+            and is_dflash_sampling_verify_available()
+        ):
+            accept_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=draft_input.max_top_k,
+                uniform_top_k_value=draft_input.uniform_top_k_value,
+            )
+            commit_lens = accept_len.to(torch.int32) + 1  # [bs]
+            out_tokens = torch.empty(
+                (bs, int(self.block_size)), dtype=torch.int64, device=device
+            )
+            if int(self.block_size) > 1:
+                out_tokens[:, : int(self.block_size) - 1].copy_(candidates[:, 1:])
+            out_tokens[:, int(self.block_size) - 1].fill_(0)
+            out_tokens.scatter_(1, accept_len.to(torch.int64)[:, None], bonus[:, None])
+        else:
+            target_predict = torch.argmax(logits_output.next_token_logits, dim=-1).view(
+                bs, int(self.block_size)
+            )
+            if self._use_triton_accept_bonus:
+                try:
+                    (
+                        accept_len,
+                        commit_lens,
+                        bonus,
+                        out_tokens,
+                        new_seq_lens,
+                    ) = self._next_accept_bonus_buffers(bs)
+                    _compute_dflash_accept_bonus_triton_unchecked(
+                        candidates=candidates,
+                        target_top1=target_predict,
+                        accept_lens_out=accept_len,
+                        commit_lens_out=commit_lens,
+                        bonus_ids_out=bonus,
+                        out_tokens_out=out_tokens,
+                        prefix_lens=prefix_lens,
+                        new_seq_lens_out=new_seq_lens,
+                    )
+                except Exception as e:
+                    self._use_triton_accept_bonus = False
+                    logger.warning(
+                        "DFLASH Triton accept/bonus failed; falling back to eager path: %s",
+                        e,
+                    )
+                    accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
+                        candidates=candidates,
+                        target_predict=target_predict,
+                    )
+                    commit_lens = accept_len.to(torch.int32) + 1  # [bs]
+                    out_tokens = torch.empty(
+                        (bs, int(self.block_size)),
+                        dtype=torch.int64,
+                        device=device,
+                    )
+                    if int(self.block_size) > 1:
+                        out_tokens[:, : int(self.block_size) - 1].copy_(
+                            candidates[:, 1:]
+                        )
+                    out_tokens[:, int(self.block_size) - 1].fill_(0)
+                    out_tokens.scatter_(
+                        1, accept_len.to(torch.int64)[:, None], bonus[:, None]
+                    )
+            else:
+                accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
+                    candidates=candidates,
+                    target_predict=target_predict,
+                )
+                commit_lens = accept_len.to(torch.int32) + 1  # [bs]
+                out_tokens = torch.empty(
+                    (bs, int(self.block_size)), dtype=torch.int64, device=device
+                )
+                if int(self.block_size) > 1:
+                    out_tokens[:, : int(self.block_size) - 1].copy_(candidates[:, 1:])
+                out_tokens[:, int(self.block_size) - 1].fill_(0)
+                out_tokens.scatter_(
+                    1, accept_len.to(torch.int64)[:, None], bonus[:, None]
+                )
+
+        if need_mamba_verify_commit:
+            assert seq_lens_pre_verify is not None
+            self._update_target_mamba_state_after_verify(
+                batch=model_worker_batch,
+                seq_lens_pre_verify=seq_lens_pre_verify,
+                commit_lens=commit_lens,
+            )
+
+        if new_seq_lens is None:
+            new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)
+        if on_publish is not None:
+            on_publish(new_seq_lens)
+
+        # --- 3) Materialize committed verify-input tokens into draft KV cache.
+        hidden = logits_output.hidden_states
+        if hidden is None:
+            raise RuntimeError(
+                "DFLASH verify requires target hidden states, but got None."
+            )
+        hidden = hidden.view(bs, int(self.block_size), -1)
+
+        self._append_target_hidden_to_draft_kv_by_loc(
+            target_hidden=hidden.reshape(-1, hidden.shape[-1]),
+            cache_loc=verify_out_cache_loc,
+            cache_loc_2d=verify_out_cache_loc_2d,
+            positions=positions,
+            commit_lens=commit_lens,
+        )
+
+        # Avoid copying large hidden-state buffers to CPU in overlap scheduling.
+        logits_output.hidden_states = None
+
+        next_draft_input = self._make_next_draft_input_decode(
+            verified_id=bonus,
+            new_seq_lens=new_seq_lens,
+            cur_allocated_seq_lens_cpu=draft_input.reserved_seq_lens_cpu,
+        )
+        verify_done = torch.get_device_module(device).Event()
+        verify_done.record()
+        next_draft_input.verify_done = verify_done
+
+        return GenerationBatchResult(
+            logits_output=logits_output,
+            next_token_ids=out_tokens.reshape(-1),
+            accept_lens=commit_lens,
+            can_run_cuda_graph=can_run_cuda_graph,
+            next_draft_input=next_draft_input,
+            speculative_num_draft_tokens=int(self.block_size),
+            # The non-overlap (sync) scheduler path advances batch.seq_lens
+            # from the result; overlap carries it via next_draft_input instead.
+            new_seq_lens=new_seq_lens,
+        )

python/sglang/srt/speculative/spec_info.py

@@ -152,7 +152,20 @@ class SpeculativeAlgorithm(Enum):
         return None
 
     def supports_spec_v2(self) -> bool:
-        return self.is_eagle() or self.is_standalone() or self.is_ngram()
+        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())
+        )
+
+    def need_topk(self) -> bool:
+        return self.is_eagle() or self.is_standalone()
 
     def get_num_tokens_per_bs_for_target_verify(
         self, num_draft_tokens: int, is_draft_worker: bool
@@ -171,13 +184,14 @@ class SpeculativeAlgorithm(Enum):
             not self.is_none()
         ), "Cannot create worker for NONE speculative algorithm."
 
-        enable_overlap = not server_args.disable_overlap_schedule
-
         if self.is_dflash():
-            if enable_overlap:
-                raise ValueError(
-                    "DFLASH does not support overlap scheduling (spec v2)."
-                )
+            # 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
+
+                return DFlashWorkerV2
             from sglang.srt.speculative.dflash_worker import DFlashWorker
 
             return DFlashWorker

python/sglang/srt/speculative/spec_utils.py

@@ -159,9 +159,9 @@ def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
         server_args = get_global_server_args()
 
     # STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
-    # multi_layer_eagle handles hidden_states internally, not via FutureMap.
+    # multi_layer_eagle and DFLASH don't relay hidden_states through FutureMap.
     # TODO(lsyin): also skip when step == 1.
-    if server_args.speculative_algorithm == "STANDALONE":
+    if server_args.speculative_algorithm in ("STANDALONE", "DFLASH"):
         return False
     return not server_args.enable_multi_layer_eagle
 

python/sglang/srt/speculative/triton_ops/dflash_accept_bonus.py

@@ -0,0 +1,140 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _dflash_accept_bonus_contig_kernel(
+    candidates_ptr,
+    target_top1_ptr,
+    accept_lens_out_ptr,
+    commit_lens_out_ptr,
+    bonus_ids_out_ptr,
+    out_tokens_ptr,
+    prefix_lens_ptr,
+    new_seq_lens_out_ptr,
+    candidates_row_stride,
+    target_row_stride,
+    accept_stride,
+    commit_stride,
+    bonus_stride,
+    out_tokens_row_stride,
+    prefix_lens_stride,
+    new_seq_lens_stride,
+    block_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+    draft_mask = cols < (block_size - 1)
+
+    candidate_row_ptr = candidates_ptr + row * candidates_row_stride
+    target_row_ptr = target_top1_ptr + row * target_row_stride
+    candidate_tail = tl.load(candidate_row_ptr + cols + 1, mask=draft_mask, other=0)
+
+    accept_len = tl.full((), 0, tl.int32)
+    prefix_live = tl.full((), 1, tl.int32)
+    for col in range(BLOCK_SIZE - 1):
+        in_range = col < (block_size - 1)
+        candidate_id = tl.load(candidate_row_ptr + (col + 1), mask=in_range, other=0)
+        target_id = tl.load(target_row_ptr + col, mask=in_range, other=0)
+        match_i32 = (candidate_id == target_id).to(tl.int32)
+        keep = in_range & (prefix_live != 0) & (match_i32 != 0)
+        accept_len += keep.to(tl.int32)
+        prefix_live = tl.where(in_range, prefix_live & match_i32, prefix_live)
+
+    commit_len = accept_len + 1
+    bonus_id = tl.load(target_row_ptr + accept_len.to(tl.int64))
+    new_seq_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride) + commit_len
+
+    tl.store(accept_lens_out_ptr + row * accept_stride, accept_len)
+    tl.store(commit_lens_out_ptr + row * commit_stride, commit_len)
+    tl.store(bonus_ids_out_ptr + row * bonus_stride, bonus_id)
+    tl.store(new_seq_lens_out_ptr + row * new_seq_lens_stride, new_seq_len)
+
+    out_val = tl.where(draft_mask, candidate_tail, 0)
+    out_val = tl.where(cols == accept_len, bonus_id, out_val)
+    tl.store(
+        out_tokens_ptr + row * out_tokens_row_stride + cols, out_val, mask=row_mask
+    )
+
+
+def _pick_num_warps(block_size: int) -> int:
+    if block_size <= 16:
+        return 1
+    if block_size <= 32:
+        return 2
+    if block_size <= 64:
+        return 4
+    return 8
+
+
+def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
+    return x.ndim == 2 and x.is_contiguous()
+
+
+def _compute_dflash_accept_bonus_triton_unchecked(
+    candidates: torch.Tensor,
+    target_top1: torch.Tensor,
+    accept_lens_out: torch.Tensor,
+    commit_lens_out: torch.Tensor,
+    bonus_ids_out: torch.Tensor,
+    out_tokens_out: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    new_seq_lens_out: torch.Tensor,
+) -> None:
+    batch_size, block_size = candidates.shape
+    if batch_size == 0:
+        return
+
+    if not _is_row_major_contiguous_2d(candidates):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous candidates.")
+    if not _is_row_major_contiguous_2d(target_top1):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous target_top1.")
+    if not _is_row_major_contiguous_2d(out_tokens_out):
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous out_tokens_out."
+        )
+    if not accept_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous accept_lens_out."
+        )
+    if not commit_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous commit_lens_out."
+        )
+    if not bonus_ids_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous bonus_ids_out."
+        )
+    if prefix_lens.ndim != 1:
+        raise ValueError("DFLASH Triton accept_bonus requires 1D prefix_lens.")
+    if not new_seq_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous new_seq_lens_out."
+        )
+
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _dflash_accept_bonus_contig_kernel[(batch_size,)](
+        candidates,
+        target_top1,
+        accept_lens_out,
+        commit_lens_out,
+        bonus_ids_out,
+        out_tokens_out,
+        prefix_lens,
+        new_seq_lens_out,
+        candidates.stride(0),
+        target_top1.stride(0),
+        accept_lens_out.stride(0),
+        commit_lens_out.stride(0),
+        bonus_ids_out.stride(0),
+        out_tokens_out.stride(0),
+        prefix_lens.stride(0),
+        new_seq_lens_out.stride(0),
+        block_size,
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )

python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py

@@ -0,0 +1,123 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _prepare_dflash_draft_block_contig_kernel(
+    verified_id_ptr,
+    prefix_lens_ptr,
+    req_pool_indices_ptr,
+    req_to_token_ptr,
+    block_ids_out_ptr,
+    positions_out_ptr,
+    cache_loc_out_ptr,
+    verified_id_stride,
+    prefix_lens_stride,
+    req_pool_indices_stride,
+    req_to_token_row_stride,
+    block_ids_row_stride,
+    positions_row_stride,
+    cache_loc_row_stride,
+    req_to_token_width,
+    block_size,
+    mask_token_id,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+
+    prefix_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride)
+    req_idx = tl.load(req_pool_indices_ptr + row * req_pool_indices_stride)
+    verified_id = tl.load(verified_id_ptr + row * verified_id_stride)
+
+    logical_pos = prefix_len.to(tl.int64) + cols
+    valid = row_mask & (logical_pos < req_to_token_width)
+    req_row_ptr = req_to_token_ptr + req_idx * req_to_token_row_stride
+    slot_ids = tl.load(req_row_ptr + logical_pos, mask=valid, other=0)
+
+    block_ids = tl.full((BLOCK_SIZE,), mask_token_id, tl.int64)
+    block_ids = tl.where(cols == 0, verified_id.to(tl.int64), block_ids)
+    tl.store(
+        block_ids_out_ptr + row * block_ids_row_stride + cols, block_ids, mask=row_mask
+    )
+    tl.store(
+        positions_out_ptr + row * positions_row_stride + cols,
+        logical_pos,
+        mask=row_mask,
+    )
+    tl.store(
+        cache_loc_out_ptr + row * cache_loc_row_stride + cols,
+        slot_ids.to(tl.int64),
+        mask=row_mask,
+    )
+
+
+def _pick_num_warps(block_size: int) -> int:
+    if block_size <= 16:
+        return 1
+    if block_size <= 32:
+        return 2
+    if block_size <= 64:
+        return 4
+    return 8
+
+
+def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
+    return x.ndim == 2 and x.is_contiguous()
+
+
+def _prepare_dflash_draft_block_unchecked(
+    verified_id: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    block_ids_out: torch.Tensor,
+    positions_out: torch.Tensor,
+    cache_loc_out: torch.Tensor,
+    mask_token_id: int,
+) -> None:
+    batch_size = int(verified_id.numel())
+    if batch_size == 0:
+        return
+
+    if req_to_token.ndim != 2 or req_to_token.stride(1) != 1:
+        raise ValueError("DFLASH Triton prepare_block requires row-major req_to_token.")
+    if not _is_row_major_contiguous_2d(block_ids_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous block_ids_out."
+        )
+    if not _is_row_major_contiguous_2d(positions_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous positions_out."
+        )
+    if not _is_row_major_contiguous_2d(cache_loc_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous cache_loc_out."
+        )
+
+    block_size = int(block_ids_out.shape[1])
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _prepare_dflash_draft_block_contig_kernel[(batch_size,)](
+        verified_id,
+        prefix_lens,
+        req_pool_indices,
+        req_to_token,
+        block_ids_out,
+        positions_out,
+        cache_loc_out,
+        verified_id.stride(0),
+        prefix_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        block_ids_out.stride(0),
+        positions_out.stride(0),
+        cache_loc_out.stride(0),
+        int(req_to_token.shape[1]),
+        block_size,
+        int(mask_token_id),
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )

python/sglang/srt/speculative/triton_ops/fused_kv_materialize.py

@@ -13,10 +13,10 @@
 # ==============================================================================
 """Fused Triton kernel for DFlash KV materialization.
 
-Combines: KV projection (cuBLAS) + RMSNorm + RoPE (Triton), then pool-managed KV writes.
+Combines: KV projection + RMSNorm + RoPE, then pool-managed KV writes.
 """
 
-from typing import Callable, List
+from typing import Callable, List, Optional
 
 import torch
 import triton
@@ -24,45 +24,58 @@ import triton.language as tl
 
 
 @triton.jit
-def _fused_norm_rope_kernel(
-    kv_ptr,  # [total_ctx, kv_size * 2]
-    k_norm_weight_ptr,  # [head_dim]
+def _fused_norm_rope_kernel_stacked(
+    kv_ptr,  # [total_ctx, n_layers, kv_size * 2]
+    k_norm_weight_ptr,  # [n_layers, head_dim]
+    eps_ptr,  # [n_layers]
     cos_sin_cache_ptr,  # [max_pos, rotary_dim]
     positions_ptr,  # [total_ctx]
-    k_out_ptr,  # [total_ctx, num_kv_heads, head_dim]
-    v_out_ptr,  # [total_ctx, num_kv_heads, head_dim]
+    k_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
+    v_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
     kv_stride_ctx,
+    kv_stride_layer,
+    k_norm_weight_stride_layer,
     cos_sin_stride_pos,
+    k_out_stride_layer,
     k_out_stride_ctx,
     k_out_stride_head,
+    v_out_stride_layer,
     v_out_stride_ctx,
     v_out_stride_head,
     total_ctx,
+    n_layers: tl.constexpr,
     num_kv_heads: tl.constexpr,
     head_dim: tl.constexpr,
     kv_size: tl.constexpr,
     rotary_dim: tl.constexpr,
     half_rotary_dim: tl.constexpr,
-    eps: tl.constexpr,
     BLOCK_HD: tl.constexpr,
 ):
-    """Fused RMSNorm(K) + RoPE(K) materialization. Grid: (total_ctx, num_kv_heads)."""
+    """Fused RMSNorm(K) + RoPE(K) materialization. Grid: (total_ctx, num_kv_heads, n_layers)."""
     ctx_id = tl.program_id(0)
     head_id = tl.program_id(1)
-    if ctx_id >= total_ctx:
+    layer_id = tl.program_id(2)
+    if ctx_id >= total_ctx or layer_id >= n_layers:
         return
 
-    # Load metadata
     position = tl.load(positions_ptr + ctx_id)
-
-    # Compute base pointers
-    kv_base = kv_ptr + ctx_id * kv_stride_ctx
+    eps = tl.load(eps_ptr + layer_id).to(tl.float32)
+    kv_base = kv_ptr + ctx_id * kv_stride_ctx + layer_id * kv_stride_layer
     k_base = kv_base + head_id * head_dim
     v_base = kv_base + kv_size + head_id * head_dim
-    k_write = k_out_ptr + ctx_id * k_out_stride_ctx + head_id * k_out_stride_head
-    v_write = v_out_ptr + ctx_id * v_out_stride_ctx + head_id * v_out_stride_head
+    k_write = (
+        k_out_ptr
+        + layer_id * k_out_stride_layer
+        + ctx_id * k_out_stride_ctx
+        + head_id * k_out_stride_head
+    )
+    v_write = (
+        v_out_ptr
+        + layer_id * v_out_stride_layer
+        + ctx_id * v_out_stride_ctx
+        + head_id * v_out_stride_head
+    )
 
-    # Load K and V
     offs = tl.arange(0, BLOCK_HD)
     mask_hd = offs < head_dim
     mask_half = offs < half_rotary_dim
@@ -70,36 +83,38 @@ def _fused_norm_rope_kernel(
     k_raw = tl.load(k_base + offs, mask=mask_hd, other=0.0).to(tl.float32)
     v_raw = tl.load(v_base + offs, mask=mask_hd, other=0.0)
 
-    # RMSNorm on K
     inv_rms = tl.rsqrt(tl.sum(k_raw * k_raw) / head_dim + eps)
-    norm_w = tl.load(k_norm_weight_ptr + offs, mask=mask_hd, other=1.0).to(tl.float32)
+    norm_w = tl.load(
+        k_norm_weight_ptr + layer_id * k_norm_weight_stride_layer + offs,
+        mask=mask_hd,
+        other=1.0,
+    ).to(tl.float32)
     k_normed = k_raw * inv_rms * norm_w
 
-    # RoPE (neox style): k_first, k_second -> rotated
     cos_sin_base = cos_sin_cache_ptr + position * cos_sin_stride_pos
     cos_v = tl.load(cos_sin_base + offs, mask=mask_half, other=1.0).to(tl.float32)
     sin_v = tl.load(
         cos_sin_base + half_rotary_dim + offs, mask=mask_half, other=0.0
     ).to(tl.float32)
 
-    # Extract first/second halves of K for rotation
     k_first = tl.where(mask_half, k_normed, 0.0)
     k_second_raw = tl.load(
         k_base + half_rotary_dim + offs, mask=mask_half, other=0.0
     ).to(tl.float32)
     norm_w_second = tl.load(
-        k_norm_weight_ptr + half_rotary_dim + offs, mask=mask_half, other=1.0
+        k_norm_weight_ptr
+        + layer_id * k_norm_weight_stride_layer
+        + half_rotary_dim
+        + offs,
+        mask=mask_half,
+        other=1.0,
     ).to(tl.float32)
     k_second = k_second_raw * inv_rms * norm_w_second
 
-    # Apply rotation
     k_rot_first = k_first * cos_v - k_second * sin_v
     k_rot_second = k_second * cos_v + k_first * sin_v
 
-    # Store V (no transform)
     tl.store(v_write + offs, v_raw, mask=mask_hd)
-
-    # Store K: rotated halves + pass-through
     tl.store(k_write + offs, k_rot_first.to(v_raw.dtype), mask=mask_half)
     tl.store(
         k_write + half_rotary_dim + offs, k_rot_second.to(v_raw.dtype), mask=mask_half
@@ -108,70 +123,117 @@ def _fused_norm_rope_kernel(
     tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass)
 
 
-def _fused_norm_rope(
-    kv: torch.Tensor,  # [total_ctx, kv_size*2]
-    k_norm_weight: torch.Tensor,  # [head_dim]
+def _fused_norm_rope_stacked(
+    kv: torch.Tensor,  # [total_ctx, n_layers, kv_size*2]
+    k_norm_weight: torch.Tensor,  # [n_layers, head_dim]
+    eps: torch.Tensor,  # [n_layers]
     cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]
     positions: torch.Tensor,  # [total_ctx]
     num_kv_heads: int,
     head_dim: int,
     rotary_dim: int,
-    eps: float = 1e-6,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Fused RMSNorm + RoPE materialization for a single layer."""
-    total_ctx = kv.shape[0]
+    """Fused RMSNorm + RoPE materialization for all layers."""
+    if kv.ndim != 3:
+        raise ValueError(
+            "Invalid stacked fused KV projection shape: "
+            f"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2]."
+        )
+
+    total_ctx, n_layers, kv_dim = kv.shape
     if total_ctx == 0:
         empty = torch.empty(
-            (0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
+            (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
         )
         return empty, empty
 
     kv_size = num_kv_heads * head_dim
-    if kv.shape[1] != kv_size * 2:
+    if kv_dim != kv_size * 2:
         raise ValueError(
             "Invalid fused KV projection shape: "
-            f"got {tuple(kv.shape)}, expected second dim {kv_size * 2}."
+            f"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}."
         )
     if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
         raise ValueError(
             "Invalid fused KV rotary/head dim pair: "
             f"rotary_dim={rotary_dim}, head_dim={head_dim}."
         )
+    if k_norm_weight.shape != (n_layers, head_dim):
+        raise ValueError(
+            "Invalid stacked k_norm_weight shape for fused KV materialization: "
+            f"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}."
+        )
+    if eps.shape != (n_layers,):
+        raise ValueError(
+            "Invalid stacked eps shape for fused KV materialization: "
+            f"got {tuple(eps.shape)}, expected {(n_layers,)}."
+        )
 
     half_rotary_dim = rotary_dim // 2
     BLOCK_HD = triton.next_power_of_2(head_dim)
 
-    # Ensure int64 for indexing
     if positions.device != kv.device:
         positions = positions.to(device=kv.device, dtype=torch.int64)
     elif positions.dtype != torch.int64:
         positions = positions.to(torch.int64)
 
-    k_out = torch.empty(
-        (total_ctx, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
-    )
-    v_out = torch.empty_like(k_out)
+    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)
+    if k_out is None:
+        k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device)
+    else:
+        if k_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid k_out shape for fused KV materialization: "
+                f"got {tuple(k_out.shape)}, expected {expected_shape}."
+            )
+        if k_out.device != kv.device or k_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid k_out device/dtype for fused KV materialization: "
+                f"got device={k_out.device}, dtype={k_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
+    if v_out is None:
+        v_out = torch.empty_like(k_out)
+    else:
+        if v_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid v_out shape for fused KV materialization: "
+                f"got {tuple(v_out.shape)}, expected {expected_shape}."
+            )
+        if v_out.device != kv.device or v_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid v_out device/dtype for fused KV materialization: "
+                f"got device={v_out.device}, dtype={v_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
 
-    _fused_norm_rope_kernel[(total_ctx, num_kv_heads)](
+    _fused_norm_rope_kernel_stacked[(total_ctx, num_kv_heads, n_layers)](
         kv,
         k_norm_weight,
+        eps,
         cos_sin_cache,
         positions,
         k_out,
         v_out,
         kv.stride(0),
+        kv.stride(1),
+        k_norm_weight.stride(0),
         cos_sin_cache.stride(0),
         k_out.stride(0),
         k_out.stride(1),
+        k_out.stride(2),
         v_out.stride(0),
         v_out.stride(1),
+        v_out.stride(2),
         total_ctx,
+        n_layers,
         num_kv_heads,
         head_dim,
         kv_size,
         rotary_dim,
         half_rotary_dim,
-        eps,
         BLOCK_HD,
     )
     return k_out, v_out
@@ -180,8 +242,8 @@ def _fused_norm_rope(
 class FusedKVMaterializeHelper:
     """Fused KV materialization helper using batched projection.
 
-    Uses torch.einsum for batched KV projection across all layers,
-    then a Triton kernel for fused RMSNorm + RoPE materialization per layer.
+    Uses a single large GEMM across all layers, then a Triton kernel for fused
+    RMSNorm + RoPE materialization across all layers.
     """
 
     def __init__(
@@ -191,12 +253,15 @@ class FusedKVMaterializeHelper:
         num_kv_heads: int,
         head_dim: int,
         device: torch.device,
+        max_position_hint: Optional[int] = None,
     ):
         self.num_kv_heads = num_kv_heads
         self.head_dim = head_dim
         self.rotary_emb = rotary_emb
         self.n_layers = len(layers)
         self.device = device
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.layer_out_dim = 2 * self.kv_size
 
         self.rotary_dim = int(getattr(rotary_emb, "rotary_dim", head_dim))
         self.is_neox_style = bool(getattr(rotary_emb, "is_neox_style", True))
@@ -209,10 +274,24 @@ class FusedKVMaterializeHelper:
                 f"rotary_dim={self.rotary_dim}, head_dim={self.head_dim}."
             )
 
-        # Pre-extract and stack weights for batched projection.
+        self.max_position_hint = (
+            max(int(max_position_hint) - 1, 0)
+            if max_position_hint is not None
+            else None
+        )
+        self._reserved_rope_cache_len = int(
+            getattr(self.rotary_emb, "cos_sin_cache", torch.empty((0,))).shape[0]
+        )
+        self._mm_out_supported = True
+        self._workspace_capacity = 0
+        self._workspace_dtype: Optional[torch.dtype] = None
+        self._proj_workspace: Optional[torch.Tensor] = None
+        self._k_workspace: Optional[torch.Tensor] = None
+        self._v_workspace: Optional[torch.Tensor] = None
+
         kv_weights = []
-        self.k_norm_weights = []
-        self.eps_values = []
+        k_norm_weights = []
+        eps_values = []
 
         for layer_id, layer in enumerate(layers):
             attn = layer.self_attn
@@ -240,15 +319,72 @@ class FusedKVMaterializeHelper:
                     f"got (rotary_dim={layer_rotary_dim}, neox={layer_is_neox}) at layer {layer_id}."
                 )
 
-            # Extract KV portion of QKV weight
             qkv_w = attn.qkv_proj.weight
             kv_weight = qkv_w[attn.q_size : attn.q_size + 2 * attn.kv_size]
             kv_weights.append(kv_weight)
-            self.k_norm_weights.append(attn.k_norm.weight)
-            self.eps_values.append(attn.k_norm.variance_epsilon)
+            k_norm_weights.append(attn.k_norm.weight)
+            eps_values.append(float(attn.k_norm.variance_epsilon))
 
-        # Stack for batched einsum: [n_layers, kv_size*2, hidden_size]
-        self.batched_kv_weight = torch.stack(kv_weights)
+        flat_kv_weight = torch.stack(kv_weights).reshape(
+            self.n_layers * self.layer_out_dim, -1
+        )
+        self.flat_kv_weight_t = flat_kv_weight.transpose(0, 1).contiguous()
+        self.k_norm_weights = torch.stack(k_norm_weights).contiguous()
+        self.eps_values = torch.tensor(
+            eps_values, dtype=torch.float32, device=self.device
+        )
+
+        if self.max_position_hint is not None:
+            self._ensure_rope_cache(self.max_position_hint)
+
+    def _ensure_rope_cache(self, max_position: int) -> torch.Tensor:
+        if max_position + 1 > self._reserved_rope_cache_len:
+            ensure_cos_sin_cache_length = getattr(
+                self.rotary_emb, "_ensure_cos_sin_cache_length", None
+            )
+            if callable(ensure_cos_sin_cache_length):
+                ensure_cos_sin_cache_length(max_position)
+                self._reserved_rope_cache_len = int(
+                    self.rotary_emb.cos_sin_cache.shape[0]
+                )
+
+        cos_sin_cache = self.rotary_emb.cos_sin_cache
+        if max_position >= int(cos_sin_cache.shape[0]):
+            raise RuntimeError(
+                "RoPE cos/sin cache is too short for fused KV materialization: "
+                f"max_position={max_position}, cache_len={int(cos_sin_cache.shape[0])}."
+            )
+        if cos_sin_cache.device != self.device:
+            cos_sin_cache = cos_sin_cache.to(self.device)
+        return cos_sin_cache
+
+    def _ensure_workspace(self, total_ctx: int, dtype: torch.dtype) -> None:
+        if (
+            self._workspace_capacity >= total_ctx
+            and self._workspace_dtype == dtype
+            and self._proj_workspace is not None
+            and self._k_workspace is not None
+            and self._v_workspace is not None
+        ):
+            return
+
+        new_capacity = max(1, total_ctx)
+        if self._workspace_capacity > 0:
+            new_capacity = max(new_capacity, self._workspace_capacity * 2)
+
+        self._proj_workspace = torch.empty(
+            (new_capacity, self.n_layers * self.layer_out_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._k_workspace = torch.empty(
+            (self.n_layers, new_capacity, self.num_kv_heads, self.head_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._v_workspace = torch.empty_like(self._k_workspace)
+        self._workspace_capacity = new_capacity
+        self._workspace_dtype = dtype
 
     def materialize(
         self,
@@ -269,35 +405,53 @@ class FusedKVMaterializeHelper:
                 f"positions={positions.numel()}, total_ctx={total_ctx}."
             )
 
-        max_position = int(positions.max().item())
-        ensure_cos_sin_cache_length = getattr(
-            self.rotary_emb, "_ensure_cos_sin_cache_length", None
+        if ctx_hidden.device != self.device:
+            ctx_hidden = ctx_hidden.to(self.device, non_blocking=True)
+        if ctx_hidden.dtype != self.flat_kv_weight_t.dtype:
+            ctx_hidden = ctx_hidden.to(self.flat_kv_weight_t.dtype)
+        if positions.device != self.device:
+            positions = positions.to(
+                device=self.device, dtype=torch.int64, non_blocking=True
+            )
+        elif positions.dtype != torch.int64:
+            positions = positions.to(torch.int64)
+
+        max_position = (
+            self.max_position_hint
+            if self.max_position_hint is not None
+            else int(positions.max().item())
         )
-        if callable(ensure_cos_sin_cache_length):
-            ensure_cos_sin_cache_length(max_position)
+        cos_sin_cache = self._ensure_rope_cache(max_position)
 
-        cos_sin_cache = self.rotary_emb.cos_sin_cache
-        if max_position >= int(cos_sin_cache.shape[0]):
-            raise RuntimeError(
-                "RoPE cos/sin cache is too short for fused KV materialization: "
-                f"max_position={max_position}, cache_len={int(cos_sin_cache.shape[0])}."
-            )
-        if cos_sin_cache.device != ctx_hidden.device:
-            cos_sin_cache = cos_sin_cache.to(ctx_hidden.device)
+        self._ensure_workspace(total_ctx, ctx_hidden.dtype)
+        assert self._proj_workspace is not None
+        assert self._k_workspace is not None
+        assert self._v_workspace is not None
 
-        # Batched KV projection: [n_layers, total_ctx, kv_size*2]
-        kv_all = torch.einsum("th,loh->lto", ctx_hidden, self.batched_kv_weight)
+        proj_out_2d = self._proj_workspace[:total_ctx]
+        if self._mm_out_supported:
+            try:
+                torch.mm(ctx_hidden, self.flat_kv_weight_t, out=proj_out_2d)
+            except Exception:
+                self._mm_out_supported = False
+                proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
+        else:
+            proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
 
-        # Per-layer fused norm/RoPE/materialize, then delegate writes to the KV pool.
-        for layer_id in range(self.n_layers):
-            cache_k, cache_v = _fused_norm_rope(
-                kv_all[layer_id],
-                self.k_norm_weights[layer_id],
-                cos_sin_cache,
-                positions,
-                self.num_kv_heads,
-                self.head_dim,
-                self.rotary_dim,
-                self.eps_values[layer_id],
-            )
-            write_layer_kv(layer_id, cache_k, cache_v)
+        proj_out = proj_out_2d.view(total_ctx, self.n_layers, self.layer_out_dim)
+        tmp_k = self._k_workspace[:, :total_ctx]
+        tmp_v = self._v_workspace[:, :total_ctx]
+        cache_k, cache_v = _fused_norm_rope_stacked(
+            proj_out,
+            self.k_norm_weights,
+            self.eps_values,
+            cos_sin_cache,
+            positions,
+            self.num_kv_heads,
+            self.head_dim,
+            self.rotary_dim,
+            k_out=tmp_k,
+            v_out=tmp_v,
+        )
+        for layer_idx in range(self.n_layers):
+            write_layer_kv(layer_idx, cache_k[layer_idx], cache_v[layer_idx])