chore: import upstream snapshot with attribution

python/sglang/srt/speculative/adaptive_runtime_state.py

@@ -0,0 +1,134 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol
+
+from sglang.srt.speculative.adaptive_spec_params import AdaptiveSpeculativeParams
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
+    from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+    from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+        EAGLEDraftCudaGraphRunner,
+    )
+    from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
+        EAGLEDraftExtendCudaGraphRunner,
+    )
+
+
+@dataclass
+class SpecRuntimeState:
+    """A complete set of runtime resources bound to a specific speculative
+    decoding configuration.
+
+    Each decode round runs three stages — draft, verify, extend — and every
+    stage has shape-dependent resources (attention backends and CUDA graphs)
+    that must match the current configuration.  Switching adaptive steps
+    means swapping the entire state atomically.
+    """
+
+    # -- Configuration (determines shapes for all stages) --
+    speculative_num_steps: int
+    speculative_num_draft_tokens: int
+
+    # -- Draft stage: draft model multi-step autoregressive generation --
+    draft_attn_backend: "AttentionBackend | None"
+    cuda_graph_runner: "EAGLEDraftCudaGraphRunner | None"
+
+    # -- Verify stage: target model one-pass tree verification --
+    target_attn_backend: "AttentionBackend"
+    target_graph_runner: "DecodeCudaGraphRunner | CPUGraphRunner | None"
+
+    # -- Extend stage: draft model KV cache catch-up after verify --
+    draft_extend_attn_backend: "AttentionBackend | None"
+    cuda_graph_runner_for_draft_extend: "EAGLEDraftExtendCudaGraphRunner | None"
+
+
+class AdaptiveSpecWorker(Protocol):
+    """Protocol that a worker must implement to use AdaptiveController."""
+
+    speculative_num_steps: int
+
+    def build_adaptive_runtime_state(
+        self,
+        speculative_num_steps: int,
+        speculative_num_draft_tokens: int,
+        cuda_graph_bs: list[int] | None = None,
+    ) -> SpecRuntimeState: ...
+
+    def apply_runtime_state(self, state: SpecRuntimeState) -> None: ...
+
+
+class AdaptiveController:
+    """Facade that owns adaptive decision-making and runtime state switching.
+
+    Works with any worker that implements AdaptiveSpecWorker protocol:
+      - build_adaptive_runtime_state(steps, draft_tokens) → runtime state
+      - apply_runtime_state(state) → apply it to the worker
+
+    The worker only needs to:
+      1. Call register() for the initial state, then init_states()
+         once during startup.
+      2. Call on_verify_complete(num_correct_drafts_per_req) after each decode verify.
+    """
+
+    def __init__(self, worker: AdaptiveSpecWorker, config_path: str | None = None):
+        self.worker = worker
+        self.params = AdaptiveSpeculativeParams(
+            initial_steps=worker.speculative_num_steps,
+            cfg_path=config_path,
+        )
+        self._states: dict[int, SpecRuntimeState] = {}
+
+    @property
+    def candidate_steps(self) -> list[int]:
+        return self.params.candidate_steps
+
+    def register(self, state: SpecRuntimeState, steps: int | None = None) -> None:
+        """Register a pre-built runtime state.
+
+        *steps* defaults to state.speculative_num_steps when not given.
+        """
+        key = steps if steps is not None else state.speculative_num_steps
+        self._states[key] = state
+
+    def init_states(self, cuda_graph_bs: list[int] | None = None) -> None:
+        """Build and register runtime states for all candidate steps."""
+        self.params.set_cuda_graph_bs(cuda_graph_bs)
+
+        for steps in self.candidate_steps:
+            if steps in self._states:
+                continue
+
+            pruned_bs = self.params.cuda_graph_bs_for_step(steps)
+            state = self.worker.build_adaptive_runtime_state(
+                speculative_num_steps=steps,
+                speculative_num_draft_tokens=steps + 1,
+                cuda_graph_bs=pruned_bs,
+            )
+            self._states[steps] = state
+
+        # Start on the initial step.
+        self._activate(self.worker.speculative_num_steps)
+
+    def activate_step_by_batch(self, batch_size: int) -> None:
+        target = self.params.get_steps_for_batch(batch_size)
+        if target != self.worker.speculative_num_steps:
+            self._activate(target)
+
+    def on_verify_complete(
+        self, num_correct_drafts_per_req: list[int], batch_size: int
+    ) -> None:
+        """Feed verify results; switch runtime state if EMA warrants it."""
+        new_step = self.params.on_verify_complete(
+            num_correct_drafts_per_req, batch_size
+        )
+        if new_step is not None:
+            self._activate(new_step)
+
+    def _activate(self, speculative_num_steps: int) -> None:
+        state = self._states.get(speculative_num_steps)
+        if state is None:
+            raise ValueError(
+                f"Missing adaptive runtime state for steps={speculative_num_steps}"
+            )
+        self.worker.apply_runtime_state(state)

python/sglang/srt/speculative/adaptive_spec_params.py

@@ -0,0 +1,345 @@
+"""Adaptive speculative decoding parameters.
+
+Adjusts speculative_num_steps at runtime based on observed acceptance lengths.
+"""
+
+from __future__ import annotations
+
+import bisect
+import json
+import logging
+import math
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from sglang.srt.utils import log_info_on_rank0
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
+    "1": {
+        "candidate_steps": [1, 3, 7],
+        "up_hysteresis": 0.0,
+        "down_hysteresis": -0.25,
+        "ceiling_coeff": 0,
+    },
+    "8": {
+        "candidate_steps": [0, 1, 3],
+        "up_hysteresis": 0.0,
+        "down_hysteresis": 0.0,
+        "ceiling_coeff": 0,
+    },
+    "32": {
+        "candidate_steps": [0, 1],
+        "up_hysteresis": 0.0,
+        "down_hysteresis": 0.0,
+        "ceiling_coeff": 0,
+    },
+    "64": {
+        "candidate_steps": [0],
+        "up_hysteresis": 0.0,
+        "down_hysteresis": 0.0,
+        "ceiling_coeff": 0,
+    },
+}
+
+
+def adaptive_unsupported_reason(server_args: ServerArgs) -> str | None:
+    """Return why adaptive spec cannot run under the given server args, or None if supported."""
+    from sglang.srt.arg_groups.overrides import resolved_view
+
+    if server_args.speculative_algorithm not in ("EAGLE", "EAGLE3"):
+        return (
+            f"speculative_algorithm={server_args.speculative_algorithm} "
+            "(only EAGLE/EAGLE3 are supported)"
+        )
+    if (
+        server_args.speculative_eagle_topk is not None
+        and server_args.speculative_eagle_topk != 1
+    ):
+        return (
+            f"speculative_eagle_topk={server_args.speculative_eagle_topk} "
+            "(only topk=1 is supported)"
+        )
+    if resolved_view(server_args).enable_dp_attention:
+        return (
+            "enable_dp_attention=True is not supported "
+            "(adaptive tier decisions are not synchronized across DP ranks)"
+        )
+    if resolved_view(server_args).enable_multi_layer_eagle:
+        return (
+            "enable_multi_layer_eagle=True is not supported "
+            "(MultiLayerEagleWorkerV2 does not implement adaptive)"
+        )
+    if server_args.enable_two_batch_overlap:
+        return (
+            "enable_two_batch_overlap=True is not supported "
+            "(adaptive state swap would discard the TboAttnBackend wrapper)"
+        )
+    if server_args.enable_pdmux:
+        return (
+            "enable_pdmux=True is not supported "
+            "(adaptive state swap does not update decode_attn_backend_group)"
+        )
+    return None
+
+
+def _load_adaptive_config(
+    cfg_path: str | None,
+) -> tuple[dict, dict[int, dict]]:
+    """Load and validate adaptive config.
+
+    Uses ``DEFAULT_ADAPTIVE_CONFIG`` when *cfg_path* is ``None``.
+    """
+    if cfg_path is not None:
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+    else:
+        cfg = DEFAULT_ADAPTIVE_CONFIG
+
+    bs_entries: dict[int, dict] = {}
+    for key, entry in cfg.items():
+        if not key.isdigit():
+            continue
+
+        steps = entry.get("candidate_steps")
+        if (
+            not isinstance(steps, list)
+            or not steps
+            or not all(isinstance(s, int) and s >= 0 for s in steps)
+        ):
+            raise ValueError(
+                f"BS {key}: candidate_steps must be a list of non-negative ints, "
+                f"got {steps!r}"
+            )
+        bs_entries[int(key)] = entry
+
+    if not bs_entries:
+        raise ValueError(
+            "speculative_adaptive_config must contain at least one integer-string "
+            'BS key, e.g. {"1": {"candidate_steps": [1,3,7]}}. '
+            f"Got keys: {list(cfg.keys())}"
+        )
+    return cfg, bs_entries
+
+
+def resolve_candidate_steps_from_config(
+    cfg_path: str | None = None,
+) -> list[int]:
+    """Union of every BS slot's candidate steps; sizes the runtime buffers."""
+    _, bs_entries = _load_adaptive_config(cfg_path)
+    all_steps: set[int] = set()
+    for entry in bs_entries.values():
+        all_steps.update(entry["candidate_steps"])
+    return sorted(all_steps)
+
+
+class AdaptiveStepSlot:
+    """Tracks acceptance rate via EMA and adapts num_steps accordingly.
+
+    The core idea: if drafts are consistently accepted, try more steps;
+    if drafts are consistently rejected early, reduce steps to avoid waste.
+
+    Formula: target_steps = clamp(round(ema_accept_len) + 1, min_steps, max_steps)
+    - Probes one step beyond observed acceptance
+    - EMA smoothing prevents oscillation
+    - Only updates every `update_interval` batches for stability
+    - num_steps can be selected from different candidate sets on different batch_sizes
+    """
+
+    def __init__(self, initial_steps: int, cfg: dict):
+        candidates = sorted(set(cfg["candidate_steps"]))
+        assert len(candidates) >= 1, "candidate_steps must have at least 1 value"
+        self.candidate_steps = candidates
+
+        self.ema_alpha = cfg.get("ema_alpha", 0.2)
+        self.update_interval = cfg.get("update_interval", 5)
+        self.warmup_batches = cfg.get("warmup_batches", 10)
+        self.down_hysteresis = cfg.get("down_hysteresis", -0.25)
+        self.up_hysteresis = cfg.get("up_hysteresis", 0.0)
+        self.ceiling_coeff = cfg.get("ceiling_coeff", 0)
+
+        if initial_steps in self.candidate_steps:
+            self.current_steps = initial_steps
+        else:
+            self.current_steps = self.candidate_steps[len(self.candidate_steps) // 2]
+
+        # Initialize EMA at current steps - 1 (neutral starting point)
+        self.ema_accept_len = float(self.current_steps - 1)
+        self._batch_count = 0
+
+    def update(self, num_correct_drafts_per_req: list[int]) -> bool:
+        """Update EMA with observed accept lengths. Returns True if params changed.
+
+        Args:
+            num_correct_drafts_per_req: Per-request accepted draft token counts from last verify.
+        """
+        if not num_correct_drafts_per_req:
+            return False
+
+        if self.current_steps > 0:
+            batch_avg = sum(num_correct_drafts_per_req) / len(
+                num_correct_drafts_per_req
+            )
+            self.ema_accept_len = (
+                1 - self.ema_alpha
+            ) * self.ema_accept_len + self.ema_alpha * batch_avg
+
+        self._batch_count += 1
+        if self._batch_count <= self.warmup_batches:
+            return False
+
+        if (self._batch_count - self.warmup_batches) % self.update_interval != 0:
+            return False
+
+        return self._recompute_params()
+
+    def _recompute_params(self) -> bool:
+        """Recompute steps from EMA. Returns True if params changed."""
+        old_steps = self.current_steps
+        current_idx = self.candidate_steps.index(old_steps)
+        old_idx = current_idx
+
+        # Probe the smallest positive step after a zero-step nospec interval.
+        if old_steps == 0:
+            current_idx = min(current_idx + 1, len(self.candidate_steps) - 1)
+            target = self.candidate_steps[current_idx]
+            if target > 0 and self.ema_accept_len < 0:
+                # A slot initialized at steps=0 has no draft acceptance history;
+                # start the first positive-step probe from that step's neutral EMA.
+                self.ema_accept_len = float(target - 1)
+            return self._apply_target_steps(old_steps, target)
+
+        # TODO: Consider limiting step changes to avoid overshooting.
+        while current_idx > 0:
+            prev_step = self.candidate_steps[current_idx - 1]
+            # A zero-step candidate disables drafting. Treat zero accepted drafts
+            # as low enough to reach it when it is the floor candidate.
+            drop_threshold = 0.5 if prev_step == 0 else prev_step - 0.5
+            drop_threshold += self.down_hysteresis
+            if self.ema_accept_len <= drop_threshold:
+                current_idx -= 1
+            else:
+                break
+
+        moved_down = current_idx < old_idx
+        if not moved_down:
+            while current_idx < len(self.candidate_steps) - 1:
+                current_step = self.candidate_steps[current_idx]
+                rise_threshold = current_step - 0.5 + self.up_hysteresis
+                if self.ema_accept_len > rise_threshold:
+                    current_idx += 1
+                else:
+                    break
+
+        target = self.candidate_steps[current_idx]
+        # EMA ceiling: only caps downward — never blocks step-ups, so the
+        # system can explore higher steps and let the EMA catch up.
+        if self.ceiling_coeff > 0:
+            ceiling = max(1, math.ceil(self.ema_accept_len * self.ceiling_coeff))
+            if target > ceiling and target <= old_steps:
+                while current_idx > 0 and self.candidate_steps[current_idx] > ceiling:
+                    current_idx -= 1
+                target = self.candidate_steps[current_idx]
+
+        return self._apply_target_steps(old_steps, target)
+
+    def _apply_target_steps(self, old_steps: int, target: int) -> bool:
+        if target != old_steps:
+            self.current_steps = target
+            log_info_on_rank0(
+                logger,
+                f"Adaptive spec params updated: steps {old_steps} -> {target} "
+                f"(ema_accept_len={self.ema_accept_len:.2f})",
+            )
+            return True
+        return False
+
+
+class AdaptiveSpeculativeParams:
+    """Routes ``batch_size`` to the correct per-BS slot.
+
+    A slot is a per-BS configuration of adaptive step selection.
+    """
+
+    def __init__(
+        self,
+        initial_steps: int,
+        cfg_path: str | None = None,
+    ):
+        cfg, bs_entries = _load_adaptive_config(cfg_path)
+        self._bs_list: list[int] = sorted(bs_entries)
+        self._slots: dict[int, AdaptiveStepSlot] = {}
+        self._cuda_graph_bs: list[int] | None = None
+
+        for bs, entry in sorted(bs_entries.items()):
+            self._slots[bs] = AdaptiveStepSlot(
+                initial_steps=initial_steps,
+                cfg={**cfg, **entry},
+            )
+
+        first_slot = self._slots[self._bs_list[0]]
+        log_info_on_rank0(
+            logger,
+            f"AdaptiveSpeculativeParams initialized: "
+            f"steps={first_slot.current_steps}, "
+            f"candidate_steps={first_slot.candidate_steps}",
+        )
+
+    @cached_property
+    def candidate_steps(self) -> list[int]:
+        """Union of all BS slots' candidate steps."""
+        return sorted({s for p in self._slots.values() for s in p.candidate_steps})
+
+    def set_cuda_graph_bs(self, cuda_graph_bs: list[int] | None) -> None:
+        self._cuda_graph_bs = sorted(cuda_graph_bs) if cuda_graph_bs else None
+
+    def get_steps_for_batch(self, batch_size: int) -> int:
+        return self._route(batch_size).current_steps
+
+    def on_verify_complete(
+        self, num_correct_drafts_per_req: list[int], batch_size: int
+    ) -> int | None:
+        """Feed verify results to the matching BS slot's EMA.
+
+        Returns the new step if a switch is warranted, else ``None``.
+        """
+        params = self._route(batch_size)
+        if params.update(num_correct_drafts_per_req):
+            return params.current_steps
+        return None
+
+    def cuda_graph_bs_for_step(self, step: int) -> list[int] | None:
+        """Return cuda_graph_bs values that can reach *step* at runtime.
+
+        Returns ``None`` when CUDA graphs are disabled (``set_cuda_graph_bs``
+        was never called or was called with ``None``).
+        """
+        if self._cuda_graph_bs is None:
+            return None
+        return [
+            v
+            for v in self._cuda_graph_bs
+            if step in self._slots[self._find_closest_bs(v)].candidate_steps
+        ]
+
+    def _route(self, batch_size: int) -> AdaptiveStepSlot:
+        """Map *batch_size* → pad to CUDA-graph BS → closest slot."""
+        return self._slots[
+            self._find_closest_bs(self._pad_to_cuda_graph_bs(batch_size))
+        ]
+
+    def _pad_to_cuda_graph_bs(self, batch_size: int) -> int:
+        if self._cuda_graph_bs is None:
+            return batch_size
+        idx = bisect.bisect_left(self._cuda_graph_bs, batch_size)
+        return (
+            self._cuda_graph_bs[idx] if idx < len(self._cuda_graph_bs) else batch_size
+        )
+
+    def _find_closest_bs(self, target: int) -> int:
+        idx = bisect.bisect_right(self._bs_list, target) - 1
+        return self._bs_list[max(0, idx)]

python/sglang/srt/speculative/base_spec_worker.py

@@ -0,0 +1,355 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import assign_draft_cache_locs_contiguous_cpu
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+        EAGLEDraftCudaGraphRunner,
+    )
+    from sglang.srt.speculative.eagle_info import (
+        EagleDraftExtendInput,
+        EagleDraftInput,
+    )
+
+
+def duplicate_prefix_tail_to_draft_branches(
+    token_to_kv_pool,
+    rows: torch.Tensor,
+    prefix_base: torch.Tensor,
+    last_page: torch.Tensor,
+    num_new_pages: torch.Tensor,
+    topk: int,
+    page_size: int,
+) -> None:
+    """Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1).
+
+    The draft-decode expand pass reads each branch's own draft page by block id
+    (cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the
+    real prefix tail (branch 0's first page already is it). Mirrors V1 #7725.
+    """
+    if topk <= 1:
+        return
+    bs = rows.shape[0]
+    page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64)
+    branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view(
+        1, topk - 1, 1
+    )
+    # Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch.
+    src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand(
+        bs, topk - 1, page_size
+    )
+    # Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size).
+    tgt_pos = (
+        prefix_base.view(bs, 1, 1)
+        + branches * (num_new_pages.view(bs, 1, 1) * page_size)
+        + page_off.view(1, 1, page_size)
+    )
+    # Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's
+    # own draft slots and must not be overwritten.
+    vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand(
+        bs, topk - 1, page_size
+    )
+    src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    if src_slots.numel() > 0:
+        token_to_kv_pool.move_kv_cache(tgt_slots, src_slots)
+
+
+class EagleDraftWorkerBase(ABC):
+    @abstractmethod
+    def draft():
+        pass
+
+    @abstractmethod
+    def draft_extend():
+        pass
+
+    def alloc_memory_pool(self, **kwargs):
+        pass
+
+    def init_attention_backends(self):
+        """Subclasses wrap this with their context managers (draft_tp_context,
+        speculative_moe_backend_context, etc.) rather than reimplementing it."""
+        self.draft_worker.init_attention_backends()
+        self.init_attention_backend()
+
+    def init_cuda_graphs(self):
+        """Capture draft graphs (decode disabled on the draft TpModelWorker)."""
+        self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False)
+        self._capture_cuda_graphs()
+
+    def prepare_for_draft_extend(
+        self,
+        draft_extend_input: EagleDraftExtendInput,
+        batch: ScheduleBatch,
+        predict: torch.Tensor,
+        num_draft_tokens: int,
+        draft_model_runner: Any,
+        cuda_graph_runner: Any,
+    ):
+        from sglang.srt.model_executor.forward_batch_info import (
+            CaptureHiddenMode,
+            ForwardBatch,
+            ForwardMode,
+        )
+        from sglang.srt.utils.async_probe import maybe_detect_oob
+        from sglang.srt.utils.common import is_npu
+
+        bs = len(batch.seq_lens)
+        extend_num_tokens = bs * num_draft_tokens
+        # When seq_lens_cpu is absent, stay on GPU-only path -- no .tolist()/.cpu().
+        gpu_only = batch.seq_lens_cpu is None
+
+        batch.spec_info = draft_extend_input
+        # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode)
+        # may run this under a plan stream; casting inside the plan stream creates a
+        # cross-stream dependency that can lead to data races and break MTP acceptance.
+        # The caller should cast to int64 before entering the plan stream context.
+        batch.input_ids = predict
+        maybe_detect_oob(
+            batch.input_ids,
+            0,
+            batch.model_config.vocab_size,
+            "v2 prepare_for_draft_extend input_ids",
+        )
+        # init_new requires both list or both Tensor;
+        # gpu_only emits device tensors to skip H2D.
+        if gpu_only:
+            batch.prefix_lens = batch.seq_lens.to(torch.int32)
+            batch.extend_lens = torch.full(
+                (bs,), num_draft_tokens, dtype=torch.int32, device=batch.seq_lens.device
+            )
+        else:
+            batch.prefix_lens = batch.seq_lens_cpu.tolist()
+            batch.extend_lens = [num_draft_tokens] * bs
+        batch.extend_num_tokens = extend_num_tokens
+        capture_mode = (
+            CaptureHiddenMode.NULL
+            if draft_model_runner.spec_algorithm.is_standalone()
+            else CaptureHiddenMode.FULL
+        )
+        batch.forward_mode = (
+            ForwardMode.IDLE
+            if batch.forward_mode.is_idle()
+            else ForwardMode.DRAFT_EXTEND_V2
+        )
+        batch.capture_hidden_mode = capture_mode
+        forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
+        # Forward sees post-write length (draft extend writes num_draft_tokens
+        # slots); mutation stays on forward_batch to preserve SB.seq_lens.
+        forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens
+        if not gpu_only:
+            forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
+            forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum())
+        else:
+            # Supply CPU mirror (extend_seq_lens are all num_draft_tokens) so
+            # backend max() reads from list without a per-iter D2H sync.
+            forward_batch.extend_seq_lens_cpu = [num_draft_tokens] * bs
+        can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+            forward_batch
+        )
+        if not batch.forward_mode.is_idle() and not can_cuda_graph:
+            draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
+            # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer
+            # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch
+            # (the `_batch_size == batch_size` assertion, see #27091); the
+            # marked pre-pad metadata is used as-is, matching the proven
+            # skip_attn_backend_init=True behavior.
+            # On NPU with --disable-cuda-graph, block_table shape won't match
+            # after prepare_mlp_sync_batch padding; defer re-init to
+            # forward_extend (post-pad) instead.
+            if not is_npu() or can_cuda_graph:
+                forward_batch.mark_forward_metadata_ready()
+        return forward_batch
+
+    def prepare_for_draft(
+        self,
+        draft_input: EagleDraftInput,
+        req_to_token_pool: ReqToTokenPool,
+        batch: ScheduleBatch,
+        cuda_graph_runner: EAGLEDraftCudaGraphRunner,
+        draft_model_runner: ModelRunner,
+        topk: int,
+        num_steps: int,
+    ):
+        from sglang.kernels.ops.speculative.cache_locs import (
+            assign_draft_cache_locs_contiguous,
+        )
+        from sglang.srt.model_executor.forward_batch_info import (
+            CaptureHiddenMode,
+            ForwardBatch,
+        )
+
+        if not batch.forward_mode.is_idle():
+            bs = len(batch.seq_lens)
+
+            # Assign cache locations (draft-write targets).
+            page_size = batch.token_to_kv_pool_allocator.page_size
+            if page_size == 1 or topk == 1:
+                batch.out_cache_loc = torch.empty(
+                    (bs * topk * num_steps,),
+                    dtype=torch.int64,
+                    device=batch.device,
+                )
+                if _is_cpu:
+                    assign_draft_cache_locs_contiguous_cpu(
+                        batch.req_pool_indices,
+                        req_to_token_pool.req_to_token,
+                        batch.seq_lens,
+                        batch.out_cache_loc,
+                        req_to_token_pool.req_to_token.shape[1],
+                        topk,
+                        num_steps,
+                    )
+                else:
+                    # FIXME(lsyin): align with the default code path
+                    assign_draft_cache_locs_contiguous[(bs,)](
+                        batch.req_pool_indices,
+                        req_to_token_pool.req_to_token,
+                        batch.seq_lens,
+                        batch.out_cache_loc,
+                        req_to_token_pool.req_to_token.shape[1],
+                        topk,
+                        num_steps,
+                    )
+            else:
+                # page_size > 1 + topk > 1: per-branch page-aligned draft pages.
+                # Reduce out_cache_loc from the page-aligned tree region down to the
+                # dense draft slots (skip each branch's duplicated prefix-tail slots
+                # and trailing padding), matching generate_draft_decode_kv_indices'
+                # paged read formula: prefix_base + t*num_new_pages*page + last_page + s.
+                # base is batch.seq_lens (== KV-ready committed prefix at draft time;
+                # the bonus is the tree root written by verify, not part of [0:seq_lens]).
+                rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()]
+                seq_lens = batch.seq_lens.to(torch.int64)
+                last_page = seq_lens % page_size
+                prefix_base = seq_lens - last_page
+                num_new_pages = (last_page + num_steps + page_size - 1) // page_size
+                topk_ids = torch.arange(
+                    topk, device=rows.device, dtype=torch.int64
+                ).view(1, topk)
+                starts = (
+                    prefix_base.view(bs, 1)
+                    + topk_ids * (num_new_pages.view(bs, 1) * page_size)
+                    + last_page.view(bs, 1)
+                )
+                steps = torch.arange(
+                    num_steps, device=rows.device, dtype=torch.int64
+                ).view(1, 1, num_steps)
+                pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps)
+                batch.out_cache_loc = (
+                    torch.gather(rows, 1, pos).reshape(-1).contiguous()
+                )
+
+                # Each branch's page-aligned region starts with `last_page` hole slots
+                # overlapping the prefix tail page; duplicate the real prefix-tail KV
+                # into them so whole-page reads stay coherent (see helper docstring).
+                duplicate_prefix_tail_to_draft_branches(
+                    draft_model_runner.token_to_kv_pool,
+                    rows,
+                    prefix_base,
+                    last_page,
+                    num_new_pages,
+                    topk,
+                    page_size,
+                )
+
+        # Get a forward batch
+        draft_input.num_tokens_per_req = topk
+        draft_input.num_tokens_for_logprob_per_req = topk
+        capture_mode = (
+            CaptureHiddenMode.NULL
+            if draft_model_runner.spec_algorithm.is_standalone()
+            else CaptureHiddenMode.LAST
+        )
+        draft_input.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
+        batch.capture_hidden_mode = capture_mode
+        forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
+        can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+            forward_batch
+        )
+        return forward_batch, can_cuda_graph
+
+
+class BaseSpecWorker(ABC):
+    @property
+    @abstractmethod
+    def target_worker(self) -> TpModelWorker:
+        pass
+
+    @property
+    @abstractmethod
+    def draft_worker(self) -> EagleDraftWorkerBase:
+        pass
+
+    @property
+    def war_fastpath_runner(self):
+        # The runner that runs the step's LAST shared-buffer-reading phase --
+        # it owns the read-done event the scheduler's WAR barrier waits on.
+        # Default is the target runner; override if the last phase runs
+        # elsewhere (eagle's draft_extend runs on the draft runner).
+        return self.target_worker.model_runner
+
+    @property
+    def spec_v2_attn_backends(self) -> tuple:
+        """Attn backends touched by spec_v2 forward; OR-ed by decide_needs_cpu_seq_lens.
+        Default returns target only; subclasses extend with draft backends."""
+        return (self.target_worker.model_runner.attn_backend,)
+
+    @abstractmethod
+    def clear_cache_pool(self):
+        # TODO: move this abstract method to BaseTpWorker and call through self.model_runner
+        pass
+
+    def alloc_memory_pool(self, **kwargs):
+        pass
+
+    def init_attention_backends(self):
+        pass
+
+    def init_cuda_graphs(self):
+        pass
+
+    def on_verify_complete_cpu(
+        self, num_correct_drafts_per_req: list[int], batch_size: int = 0
+    ) -> None:
+        """Hook called after verify finishes and accept counts are on CPU.
+
+        Default no-op. Adaptive-aware workers override this to feed the
+        controller without forcing a GPU→CPU sync in the worker hot path.
+        """
+        pass
+
+    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        """Hook called by the batch-result processor when a request finishes.
+
+        Default no-op. DSpark overrides this to settle / censor its
+        block-accept estimator state for the finished request.
+        """
+        pass
+
+    def activate_step_by_batch(self, batch_size: int) -> None:
+        """Activate the optimal adaptive step for the current batch size.
+
+        Default no-op. Adaptive-aware workers override this to switch
+        the runtime state before each draft round.
+        """
+        pass

python/sglang/srt/speculative/cpp_ngram/.clang-format

@@ -0,0 +1,15 @@
+BasedOnStyle: Google
+IndentWidth: 2
+ColumnLimit: 120
+AllowShortFunctionsOnASingleLine: Empty
+DerivePointerAlignment: false
+PointerAlignment: Left
+NamespaceIndentation: None
+SortIncludes: true
+AllowShortLoopsOnASingleLine: false
+BinPackParameters: false                 # Prevents packing parameters in declarations
+BinPackArguments: false                  # Prevents packing arguments in function calls
+AlignAfterOpenBracket: AlwaysBreak       # Forces a break after the opening parenthesis
+AlignOperands: Align                     # Aligns arguments vertically
+PenaltyBreakBeforeFirstCallParameter: 1  # Encourages breaking before the first argument
+PenaltyReturnTypeOnItsOwnLine: 100       # Keeps return type with function name

python/sglang/srt/speculative/cpp_ngram/external_corpus.py

@@ -0,0 +1,62 @@
+import json
+from collections.abc import Iterator
+from pathlib import Path
+
+# Must match SuffixAutomaton::kSeparatorToken in suffix_automaton.h.
+SEPARATOR_TOKEN = -(2**31)
+
+# Default chunk size for streaming tokenized documents into the SAM.
+DEFAULT_CHUNK_SIZE = 4096
+
+
+def iter_external_corpus_chunks(
+    path: str, tokenizer, max_tokens: int, chunk_size: int = DEFAULT_CHUNK_SIZE
+) -> Iterator[list[int]]:
+    """Chunk documents and yield fixed-size token chunks from a JSONL corpus file."""
+    corpus_path = Path(path)
+    if not corpus_path.is_file():
+        raise ValueError(f"External ngram corpus path does not exist: {path}")
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required to load an external ngram corpus.")
+    if max_tokens <= 0:
+        raise ValueError("External ngram corpus max tokens must be positive.")
+
+    total_tokens = 0
+    has_previous_doc = False
+    with corpus_path.open("r", encoding="utf-8") as f:
+        for line_no, line in enumerate(f, start=1):
+            if not line.strip():
+                continue
+
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Invalid JSON in external ngram corpus at line {line_no}: {e.msg}"
+                ) from e
+
+            if not isinstance(record, str):
+                raise ValueError(
+                    "Invalid external ngram corpus record at line "
+                    f"{line_no}: expected a JSON string."
+                )
+
+            token_ids = list(tokenizer.encode(record, add_special_tokens=False))
+            if not token_ids:
+                continue
+
+            separator_cost = 1 if has_previous_doc else 0
+            next_total_tokens = total_tokens + separator_cost + len(token_ids)
+            if next_total_tokens > max_tokens:
+                raise ValueError(
+                    "External ngram corpus exceeds the configured token limit "
+                    f"({max_tokens}) at line {line_no} after loading "
+                    f"{total_tokens} tokens."
+                )
+            total_tokens = next_total_tokens
+
+            if has_previous_doc:
+                token_ids = [SEPARATOR_TOKEN] + token_ids
+            for i in range(0, len(token_ids), chunk_size):
+                yield token_ids[i : i + chunk_size]
+            has_previous_doc = True

python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+
+import logging
+from collections.abc import Iterable, Sequence
+from typing import Dict, List, Tuple
+
+import numpy as np
+
+from sglang.jit_kernel.ngram_corpus import get_ngram_corpus_cls
+
+logger = logging.getLogger(__name__)
+
+
+class NgramCorpus:
+    def __init__(
+        self,
+        max_trie_depth=18,
+        min_bfs_breadth=1,
+        max_bfs_breadth=8,
+        draft_token_num=8,
+        match_type="BFS",
+        capacity=1000000,
+        external_sam_budget=0,
+        external_corpus_max_tokens=10000000,
+    ) -> None:
+        cls = get_ngram_corpus_cls()
+        self._obj = cls(
+            capacity=capacity,
+            max_trie_depth=max_trie_depth,
+            min_bfs_breadth=min_bfs_breadth,
+            max_bfs_breadth=max_bfs_breadth,
+            draft_token_num=draft_token_num,
+            match_type=match_type,
+            external_sam_budget=external_sam_budget,
+            external_corpus_max_tokens=external_corpus_max_tokens,
+        )
+        self.draft_token_num = draft_token_num
+        self.external_corpus_max_tokens = external_corpus_max_tokens
+        self._req_id_to_state_id: Dict[str, int] = {}
+        self._next_state_id: int = 0
+        self._corpus_token_counts: Dict[str, int] = {}
+        self._total_loaded_tokens: int = 0
+
+    def _get_state_id(self, req_id: str) -> int:
+        sid = self._req_id_to_state_id.get(req_id)
+        if sid is None:
+            sid = self._next_state_id
+            self._next_state_id += 1
+            self._req_id_to_state_id[req_id] = sid
+        return sid
+
+    def batch_put(self, batch_tokens: List[List[int]]):
+        self._obj.insert(batch_tokens)
+
+    def synchronize(self):
+        self._obj.synchronize()  # type: ignore
+
+    @property
+    def remaining_token_budget(self) -> int:
+        return self.external_corpus_max_tokens - self._total_loaded_tokens
+
+    def load_external_corpus_named(
+        self, corpus_id: str, chunks: Iterable[Sequence[int]]
+    ) -> int:
+        if corpus_id in self._corpus_token_counts:
+            raise ValueError(
+                f"External corpus '{corpus_id}' already exists. Remove it before "
+                f"adding a new corpus with the same id."
+            )
+        # Note(kpham-sgl): remaining_token_budget is stale (e.g if there are removes
+        # during the load), which makes the budget more conservative than it should be.
+        # This is acceptable because otherwise load_external_corpus_named would need to check the budget after each chunk,
+        # which would be inefficient.
+        _, loaded_token_count = self._obj.load_external_corpus_named(
+            corpus_id, chunks, self.remaining_token_budget
+        )
+        return loaded_token_count
+
+    # Commit corpus bookkeeping after successful load. Call only at background thread join.
+    # (or after synchronous load_external_corpus_named returns)
+    def commit_external_corpus_load(
+        self, corpus_id: str, loaded_token_count: int
+    ) -> None:
+        self._corpus_token_counts[corpus_id] = loaded_token_count
+        self._total_loaded_tokens += loaded_token_count
+
+    def remove_external_corpus(self, corpus_id: str) -> None:
+        self._obj.remove_corpus(corpus_id)
+        old_count = self._corpus_token_counts.pop(corpus_id, 0)
+        self._total_loaded_tokens -= old_count
+
+    def list_external_corpora(self) -> Dict[str, int]:
+        return self._obj.list_corpora()
+
+    def reset(self):
+        self._obj.reset()  # type: ignore
+        self._req_id_to_state_id.clear()
+        self._next_state_id = 0
+
+    def batch_get(
+        self,
+        req_ids: List[str],
+        batch_tokens: List[List[int]],
+        total_lens: List[int],
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        state_ids = [self._get_state_id(rid) for rid in req_ids]
+        return self._obj.match_stateful(state_ids, batch_tokens, total_lens)
+
+    def erase_match_state(self, req_ids: List[str]):
+        state_ids = []
+        for rid in req_ids:
+            sid = self._req_id_to_state_id.pop(rid, None)
+            if sid is not None:
+                state_ids.append(sid)
+        if state_ids:
+            self._obj.erase_states(state_ids)
+
+    def leaf_paths_from_mask(
+        self, tokens: List[int], tree_mask: List[List[int]]
+    ) -> List[List[int]]:
+        """
+        Find all leaf paths according to the binary tree_mask (i.e., paths that are not prefixes of any other path).
+
+        Args:
+            mask   : List[List[int]]   # nxn binary matrix
+            tokens : List[int]         # token list corresponding to columns
+
+        Returns:
+            List[List[int]]            # token lists of only the leaf paths, preserving their order of appearance
+        """
+
+        row_sets = [
+            (i, {idx for idx, v in enumerate(row) if v == 1})
+            for i, row in enumerate(tree_mask)
+        ]
+        leaf_sets = []
+        leaf_rows = []
+
+        for i, cur_set in reversed(row_sets):
+            if any(cur_set <= kept for kept in leaf_sets):
+                continue
+            leaf_sets.append(cur_set)
+            leaf_rows.append(i)
+
+        leaf_rows.reverse()
+        result = []
+        for r in leaf_rows:
+            path = [tokens[col] for col in range(len(tokens)) if tree_mask[r][col] == 1]
+            result.append(path)
+
+        return result
+
+    def debug_result(
+        self, decoding_ids: np.ndarray, decoding_masks: np.ndarray, tokenizer=None
+    ):
+        decoding_ids = decoding_ids.reshape(-1, self.draft_token_num)
+        decoding_masks = decoding_masks.reshape(
+            -1, self.draft_token_num, self.draft_token_num
+        )
+        logger.info(f"\n{decoding_ids=}\n{decoding_masks=}")
+        for i in range(decoding_ids.shape[0]):
+            leaf_paths = self.leaf_paths_from_mask(
+                decoding_ids[i].tolist(), decoding_masks[i].tolist()
+            )
+            if tokenizer is None:
+                logger.info(f"draft path {i}: {leaf_paths}")
+            else:
+                logger.info(f"result {i}:")
+                for leaf_path in leaf_paths:
+                    logger.info(
+                        f"draft path {i}: {leaf_path} -> {tokenizer.decode(leaf_path, ensure_ascii=False)}"
+                    )
+
+
+# main function
+if __name__ == "__main__":
+    format = f"%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s"
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format=format,
+        datefmt="%Y-%m-%d %H:%M:%S",
+        force=True,
+    )
+
+    token_ids = [
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+        [1, 2, 3, 44, 55, 66, 77, 88, 99, 100],
+    ]
+    corpus = NgramCorpus(max_trie_depth=12, draft_token_num=8)
+    corpus.batch_put(token_ids)
+
+    corpus.synchronize()
+    queries = [[1, 2, 3], [3, 44], [3, 6, 999]]
+    decoding_ids, decoding_masks = corpus.batch_get(
+        req_ids=[f"query-{i}" for i in range(len(queries))],
+        batch_tokens=queries,
+        total_lens=[len(q) for q in queries],
+    )
+
+    corpus.debug_result(decoding_ids, decoding_masks)

python/sglang/srt/speculative/decoupled_spec_io.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, Optional
+
+
+class DraftMeshMessageType(str, Enum):
+    CONTROL_BATCH = "control_batch"
+    TAIL_STREAM_OUTPUT_BATCH = "tail_stream_output_batch"
+
+
+@dataclass(frozen=True)
+class DraftReqKey:
+    """Request identity on the drafter side.
+
+    The original request_id is only unique within the verifier that owns it.
+    src_verifier_rank keeps the drafter-side request table unambiguous when
+    multiple verifier ranks send work to the same drafter rank.
+    """
+
+    src_verifier_rank: int
+    request_id: str
+
+
+def build_draft_scheduler_rid(draft_key: DraftReqKey) -> str:
+    return f"draft:{int(draft_key.src_verifier_rank)}:{draft_key.request_id}"
+
+
+def parse_draft_scheduler_rid(rid: str) -> DraftReqKey:
+    if rid.startswith("draft:"):
+        encoded = rid[len("draft:") :]
+        rank_text, sep, request_id = encoded.partition(":")
+        if sep and request_id:
+            return DraftReqKey(
+                src_verifier_rank=int(rank_text),
+                request_id=request_id,
+            )
+
+    raise ValueError(f"Invalid decoupled draft scheduler rid: {rid}")
+
+
+@dataclass
+class DraftSync:
+    """Open or re-open a drafter request from a verifier-owned prefix.
+
+    The verifier is the source of truth for committed tokens. DraftSync gives
+    the drafter the prompt and already committed output prefix that it must
+    align to before it can emit draft tail tokens.
+    """
+
+    request_id: str
+    src_verifier_rank: int
+    dst_drafter_rank: int
+    prompt_token_ids: list[int] = field(default_factory=list)
+    committed_outputs: list[int] = field(default_factory=list)
+
+    @property
+    def draft_key(self) -> DraftReqKey:
+        return DraftReqKey(
+            src_verifier_rank=int(self.src_verifier_rank),
+            request_id=self.request_id,
+        )
+
+
+@dataclass
+class VerifyCommit:
+    """
+    Sent from verifier to drafter to commit a portion of the draft outputs.
+
+    committed_tokens is the verifier-committed contiguous output segment:
+    output_ids[
+        pre_verify_committed_len:
+        pre_verify_committed_len + len(committed_tokens)
+    ].
+    Drafter must align its reqs to these committed tokens,
+    and sometimes needs to truncate tokens / reprefill.
+    """
+
+    request_id: str
+    src_verifier_rank: int
+    dst_drafter_rank: int
+    pre_verify_committed_len: int
+    committed_tokens: list[int]
+
+    @property
+    def draft_key(self) -> DraftReqKey:
+        return DraftReqKey(
+            src_verifier_rank=int(self.src_verifier_rank),
+            request_id=self.request_id,
+        )
+
+    def validate_committed_tokens(self) -> None:
+        if not self.committed_tokens:
+            raise ValueError(
+                "VerifyCommit committed_tokens must be non-empty: "
+                f"request_id={self.request_id} "
+                f"pre_verify_committed_len={self.pre_verify_committed_len}"
+            )
+        if int(self.pre_verify_committed_len) < 0:
+            raise ValueError(
+                "VerifyCommit pre_verify_committed_len must be non-negative: "
+                f"request_id={self.request_id} "
+                f"pre_verify_committed_len={self.pre_verify_committed_len}"
+            )
+
+
+@dataclass
+class DraftClose:
+    request_id: str
+    src_verifier_rank: int
+    dst_drafter_rank: int
+    reason: str
+
+    @property
+    def draft_key(self) -> DraftReqKey:
+        return DraftReqKey(
+            src_verifier_rank=int(self.src_verifier_rank),
+            request_id=self.request_id,
+        )
+
+
+@dataclass
+class DraftTailStreamOutput:
+    """
+    Drafter sends one output token to the verifier-side DraftTailBuffer.
+
+    base_committed_len records the verifier prefix length that the drafter used
+    as the base when this token was emitted. The verifier compares it with its
+    stale-base boundary before accepting the token as tail data or as
+    pending-prefix confirmation.
+
+    new_token_pos is the 0-based output token position for new_token. Normal
+    decode streams send the latest generated token.
+    """
+
+    src_drafter_rank: int
+    dst_verifier_rank: int
+    request_id: str
+    base_committed_len: int
+    new_token_pos: int
+    new_token: int
+
+
+@dataclass
+class DraftTailStreamOutputBatch:
+    outputs: list[DraftTailStreamOutput] = field(default_factory=list)
+
+
+@dataclass
+class DraftControlBatch:
+    dst_drafter_rank: int
+    sync_messages: list[DraftSync] = field(default_factory=list)
+    verify_commit_messages: list[VerifyCommit] = field(default_factory=list)
+    close_messages: list[DraftClose] = field(default_factory=list)
+
+
+@dataclass
+class VerifierCommitSegment:
+    """Contiguous VerifyCommit messages coalesced for one drafter request.
+
+    When receiving contiguous VerifyCommit messages for the same draft req,
+    the transport thread(TokenSync thread at drafter side) coalesces them into a single VerifierCommitSegment.
+
+    VerifierCommitSegment represents a contiguous verifier-committed token segment for drafter,
+    and drafter scheduler should align with these segments before emitting tail tokens
+    """
+
+    draft_key: DraftReqKey
+    dst_drafter_rank: int
+    pre_verify_committed_len: int
+    committed_tokens: list[int] = field(default_factory=list)
+
+    @property
+    def end_committed_len(self) -> int:
+        return int(self.pre_verify_committed_len) + len(self.committed_tokens)
+
+    def append_message(self, message: VerifyCommit) -> None:
+        """
+        It runs on TokenSyncThread under _pending_lock. That loop only
+        catches zmq.error.ContextTerminated, so a raise here escapes _run and
+        silently kills the drafter control thread. It then stops applying
+        ALL requests' controls while the verifier keeps pushing.
+
+        TODO: 1. peer-data violations (non-contiguous / invalid len)
+        should quarantine just that request (drop + add to close_keys), not
+        crash the thread. 2. phase 5.c will handle the drafter failure by
+        degrading the verifier into normal autoregressive decoding.
+        """
+        if message.draft_key != self.draft_key:
+            raise RuntimeError(
+                "Verifier commit segment received a commit for a different "
+                f"request: segment_key={self.draft_key} message_key={message.draft_key}"
+            )
+        if int(message.dst_drafter_rank) != int(self.dst_drafter_rank):
+            raise RuntimeError(
+                "Verifier commit segment received a commit for a different "
+                "drafter rank: "
+                f"request_id={message.request_id} "
+                f"segment_drafter_rank={self.dst_drafter_rank} "
+                f"message_drafter_rank={message.dst_drafter_rank}"
+            )
+        message.validate_committed_tokens()
+        pre_verify_committed_len = int(message.pre_verify_committed_len)
+        if pre_verify_committed_len != self.end_committed_len:
+            raise RuntimeError(
+                "Verifier commit segment requires contiguous VerifyCommit "
+                "messages: "
+                f"request_id={message.request_id} "
+                f"expected_pre_verify_committed_len={self.end_committed_len} "
+                f"actual_pre_verify_committed_len={pre_verify_committed_len}"
+            )
+
+        token_ids = [int(token_id) for token_id in message.committed_tokens]
+        self.committed_tokens.extend(token_ids)
+
+    def extract_prefix(self, num_tokens: int) -> VerifierCommitSegment:
+        num_tokens = int(num_tokens)
+        if num_tokens <= 0:
+            raise ValueError(
+                "Verifier commit segment prefix length must be positive: "
+                f"request_id={self.draft_key.request_id} num_tokens={num_tokens}"
+            )
+        if num_tokens > len(self.committed_tokens):
+            raise ValueError(
+                "Verifier commit segment prefix length exceeds segment length: "
+                f"request_id={self.draft_key.request_id} "
+                f"num_tokens={num_tokens} "
+                f"segment_len={len(self.committed_tokens)}"
+            )
+
+        prefix_tokens = [
+            int(token_id) for token_id in self.committed_tokens[:num_tokens]
+        ]
+        remaining_tokens = [
+            int(token_id) for token_id in self.committed_tokens[num_tokens:]
+        ]
+        prefix_segment = VerifierCommitSegment(
+            draft_key=self.draft_key,
+            dst_drafter_rank=int(self.dst_drafter_rank),
+            pre_verify_committed_len=int(self.pre_verify_committed_len),
+            committed_tokens=prefix_tokens,
+        )
+        self.pre_verify_committed_len = int(self.pre_verify_committed_len) + num_tokens
+        self.committed_tokens = remaining_tokens
+        return prefix_segment
+
+
+@dataclass
+class DraftControlInbox:
+    """Drafter-side inbox for verifier control messages.
+
+    The TokenSync thread temporarily stores incoming control messages here.
+    The drafter scheduler extracts and consumes them each time it finishes a decoding step.
+    """
+
+    sync_messages: list[DraftSync] = field(default_factory=list)
+    verifier_commit_segments: dict[DraftReqKey, VerifierCommitSegment] = field(
+        default_factory=dict
+    )
+    close_keys: set[DraftReqKey] = field(default_factory=set)
+
+    def is_empty(self) -> bool:
+        return (
+            not self.sync_messages
+            and not self.verifier_commit_segments
+            and not self.close_keys
+        )
+
+    def pending_control_count(self) -> int:
+        return (
+            len(self.sync_messages)
+            + len(self.verifier_commit_segments)
+            + len(self.close_keys)
+        )
+
+    def add_control_batch_locked(self, batch: DraftControlBatch) -> None:
+        for message in batch.close_messages:
+            self.add_close_key_locked(message.draft_key)
+        for message in batch.sync_messages:
+            if message.draft_key not in self.close_keys:
+                self.sync_messages.append(message)
+        for message in batch.verify_commit_messages:
+            self.add_verify_commit_locked(message)
+
+    def add_close_key_locked(self, draft_key: DraftReqKey) -> None:
+        self.close_keys.add(draft_key)
+        self.verifier_commit_segments.pop(draft_key, None)
+        self.sync_messages = [
+            message for message in self.sync_messages if message.draft_key != draft_key
+        ]
+
+    def add_verify_commit_locked(self, message: VerifyCommit) -> None:
+        if message.draft_key in self.close_keys:
+            return
+        segment = self.verifier_commit_segments.get(message.draft_key)
+        if segment is None:
+            segment = VerifierCommitSegment(
+                draft_key=message.draft_key,
+                dst_drafter_rank=int(message.dst_drafter_rank),
+                pre_verify_committed_len=int(message.pre_verify_committed_len),
+            )
+            segment.append_message(message)
+            self.verifier_commit_segments[message.draft_key] = segment
+            return
+        segment.append_message(message)
+
+    def extract_ready_controls_locked(
+        self,
+        consumable_commit_len: Callable[[VerifierCommitSegment], int],
+    ) -> ReadyDraftControls:
+        ready_controls = ReadyDraftControls()
+
+        if self.close_keys:
+            ready_controls.close_keys = self.close_keys
+            self.close_keys = set()
+
+        if self.sync_messages:
+            ready_controls.sync_messages = self.sync_messages
+            self.sync_messages = []
+
+        for draft_key, segment in list(self.verifier_commit_segments.items()):
+            consumable_len = consumable_commit_len(segment)
+            if consumable_len <= 0:
+                continue
+
+            ready_controls.ready_commit_segments.append(
+                segment.extract_prefix(consumable_len)
+            )
+            if not segment.committed_tokens:
+                self.verifier_commit_segments.pop(draft_key, None)
+
+        return ready_controls
+
+
+@dataclass
+class ReadyDraftControls:
+    sync_messages: list[DraftSync] = field(default_factory=list)
+    close_keys: set[DraftReqKey] = field(default_factory=set)
+    ready_commit_segments: list[VerifierCommitSegment] = field(default_factory=list)
+
+    def is_empty(self) -> bool:
+        return (
+            not self.sync_messages
+            and not self.close_keys
+            and not self.ready_commit_segments
+        )
+
+    def extracted_control_count(self) -> int:
+        return (
+            len(self.sync_messages)
+            + len(self.close_keys)
+            + len(self.ready_commit_segments)
+        )
+
+
+@dataclass
+class DraftMeshMessage:
+    message_type: DraftMeshMessageType
+    control_batch: Optional[DraftControlBatch] = None
+    tail_stream_output_batch: Optional[DraftTailStreamOutputBatch] = None
+
+    @staticmethod
+    def from_control_batch(message: DraftControlBatch) -> DraftMeshMessage:
+        return DraftMeshMessage(
+            message_type=DraftMeshMessageType.CONTROL_BATCH,
+            control_batch=message,
+        )
+
+    @staticmethod
+    def from_tail_stream_output_batch(
+        message: DraftTailStreamOutputBatch,
+    ) -> DraftMeshMessage:
+        return DraftMeshMessage(
+            message_type=DraftMeshMessageType.TAIL_STREAM_OUTPUT_BATCH,
+            tail_stream_output_batch=message,
+        )
+
+
+@dataclass(frozen=True)
+class DecoupledSpecIpcConfig:
+    bind_endpoint: str
+    connect_endpoints: tuple[str, ...]
+    rank: int

python/sglang/srt/speculative/dflash_info.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+@dataclass
+class DFlashVerifyInput(SpecInput):
+    """Inputs for a target-model verify forward in DFlash.
+
+    The verify forward is run with `ForwardMode.TARGET_VERIFY` so that the target
+    model returns logits for all tokens in the block, enabling accept-length
+    computation.
+    """
+
+    draft_token: torch.Tensor
+    positions: torch.Tensor
+    draft_token_num: int
+    # 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.
+    # 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
+
+    # Shape info for padding (e.g., DP attention / CUDA graph).
+    num_tokens_per_req: int = -1
+
+    ragged_verify_layout: Optional[RaggedVerifyLayout] = None
+
+    def __post_init__(self):
+        super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY)
+        if self.num_tokens_per_req == -1:
+            self.num_tokens_per_req = int(self.draft_token_num)
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        return self.draft_token_num, self.draft_token_num
+
+    def prepare_for_verify(
+        self,
+        batch: ScheduleBatch,
+        target_worker: TpModelWorker,
+    ) -> tuple[ForwardBatch, bool]:
+        """Prepare a DFLASH verify forward batch for overlap scheduling.
+
+        The caller computes and stores `batch.out_cache_loc` before this
+        method is called. This helper only packages the verify forward and pre-initializes either CUDA-graph replay
+        metadata or eager attention metadata so the actual forward can run with
+        `skip_attn_backend_init=True`.
+        """
+        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_graph(
+                verify_forward_batch
+            )
+        )
+        if can_run_cuda_graph:
+            target_worker.model_runner.decode_cuda_graph_runner.load_batch(
+                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)
+
+        layout = self.ragged_verify_layout
+
+        if layout is None:
+            qo_indptr = torch.arange(
+                0,
+                (bs + 1) * self.draft_token_num,
+                step=self.draft_token_num,
+                dtype=torch.int32,
+                device=device,
+            )
+            verify_lens = self.draft_token_num
+            kv_indices_extra = self.draft_token_num * bs
+        else:
+            qo_indptr = layout.qo_indptr_device
+            verify_lens = layout.verify_lens
+            kv_indices_extra = layout.total_verify_tokens
+
+        cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=device)
+        paged_kernel_lens = paged_kernel_lens + verify_lens
+        cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum + kv_indices_extra,
+            dtype=torch.int32,
+            device=device,
+        )
+        create_flashinfer_kv_indices_triton[(bs,)](
+            req_to_token,
+            req_pool_indices,
+            paged_kernel_lens,
+            cum_kv_seq_len,
+            kv_start_idx,
+            kv_indices,
+            req_to_token.size(1),
+        )
+        mask = self.custom_mask
+        if mask is not None:
+            mask_numel = (
+                paged_kernel_lens_sum * self.draft_token_num
+                + (self.draft_token_num**2) * bs
+            )
+            if mask.numel() < mask_numel:
+                # FIXME(attn): temporary fix for custom mask padding with cuda graph
+                mask = torch.cat(
+                    [
+                        mask,
+                        torch.full(
+                            (mask_numel - mask.numel(),),
+                            True,
+                            dtype=torch.bool,
+                            device=device,
+                        ),
+                    ],
+                    dim=0,
+                )
+                self.custom_mask = mask
+        return kv_indices, cum_kv_seq_len, qo_indptr, mask

python/sglang/srt/speculative/dflash_info_v2.py

@@ -0,0 +1,288 @@
+"""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.runtime_context import get_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; DFLASH relays via FutureMap so these are unused.
+    topk_p: torch.Tensor
+    topk_index: torch.Tensor
+    bonus_tokens: torch.Tensor
+    new_seq_lens: torch.Tensor
+    hidden_states: torch.Tensor
+    max_top_k: int = 1
+    uniform_top_k_value: Optional[int] = None
+    reserved_seq_lens_cpu: Optional[torch.Tensor] = None
+    reserved_seq_lens_sum: Optional[int] = 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
+
+    verify_token_budget: Optional[int] = 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
+
+        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)
+
+        # The three CPU scratch buffers grow together; capacity is the only
+        # invariant (batch is int64 non-pinned, cur/nxt are int32 pinned).
+        if needs_cpu_alloc(self._prepare_batch_seq_lens_cpu_buf):
+            capacity = grown_capacity(self._prepare_batch_seq_lens_cpu_buf)
+            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),
+            bonus_tokens=torch.empty((0,), device=device, dtype=torch.int64),
+            new_seq_lens=torch.empty((0,), device=device, dtype=torch.int64),
+            hidden_states=torch.empty((0, 0), device=device, dtype=torch.float16),
+        )
+
+    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)
+
+        bs = batch.batch_size()
+        if bs == 0:
+            return
+        self._ensure_prepare_length_buffers(bs, batch.device)
+        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
+        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]
+
+        # For DFLASH, each decode step needs a fixed-size verify block.
+        block_size = int(get_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
+        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)
+            # Read the allocation watermark from the req object like EAGLE.
+            cur_alloc_len = int(req.kv_allocated_len)
+            reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)
+            top_k = int(req.sampling_params.top_k)
+
+            batch_seq_lens_cpu_t[i] = committed_len
+            cur_kv_lens_cpu_t[i] = cur_alloc_len
+            nxt_kv_lens_cpu_t[i] = reserved_len
+
+            committed_seq_lens_sum += committed_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)
+
+            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]))
+
+        # Seed committed; overlap's resolve overwrites it with the published value.
+        batch.seq_lens_cpu = batch_seq_lens_cpu_t
+        batch.seq_lens_sum = committed_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.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]
+            return
+
+        self.topk_p = self.topk_p[new_indices]
+        self.topk_index = self.topk_index[new_indices]
+        self.bonus_tokens = self.bonus_tokens[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.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]
+            )
+            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.bonus_tokens = torch.cat(
+            [self.bonus_tokens, spec_info.bonus_tokens], 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

@@ -0,0 +1,814 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+from numbers import Integral
+from typing import Any, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+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(
+    {
+        "FlashInferAttnBackend",
+        "FlashInferMLAAttnBackend",
+        "FlashAttentionBackend",
+        "TritonAttnBackend",
+        "TRTLLMHAAttnBackend",
+        "TRTLLMMLABackend",
+    }
+)
+
+
+if is_cuda() or is_musa():
+    try:
+        from sgl_kernel import (
+            top_k_renorm_prob,
+            top_p_renorm_prob,
+            tree_speculative_sampling_target_only,
+        )
+
+        _DFLASH_SAMPLING_VERIFY_AVAILABLE = True
+    except Exception:
+        top_k_renorm_prob = None
+        top_p_renorm_prob = None
+        tree_speculative_sampling_target_only = None
+else:
+    top_k_renorm_prob = None
+    top_p_renorm_prob = None
+    tree_speculative_sampling_target_only = None
+
+
+def is_dflash_sampling_verify_available() -> bool:
+    return _DFLASH_SAMPLING_VERIFY_AVAILABLE
+
+
+def scale_kv_cell_size_per_token_for_dflash(
+    *,
+    target_cell_size_per_token: int,
+    target_num_layers: int,
+    draft_num_layers: int,
+    draft_cell_size_per_token: Optional[int] = None,
+) -> int:
+    """Compute bytes/token budget for combined target+draft KV pools (DFLASH).
+
+    DFLASH runs a separate draft runner with its own KV pool. The target runner's
+    token capacity must fit both pools in aggregate.
+
+    Returns:
+        Approximate per-token bytes for (target KV + draft KV), expressed as a
+        scaled version of `target_cell_size_per_token`, unless an explicit
+        `draft_cell_size_per_token` is provided (in which case we sum them).
+    """
+    if target_cell_size_per_token <= 0:
+        raise ValueError(
+            "target_cell_size_per_token must be positive, "
+            f"got {target_cell_size_per_token}."
+        )
+
+    if draft_cell_size_per_token is not None:
+        draft_cell_size_per_token = int(draft_cell_size_per_token)
+        if draft_cell_size_per_token <= 0:
+            raise ValueError(
+                "draft_cell_size_per_token must be positive when provided, "
+                f"got {draft_cell_size_per_token}."
+            )
+        return int(target_cell_size_per_token) + int(draft_cell_size_per_token)
+
+    if target_num_layers <= 0 or draft_num_layers <= 0:
+        return int(target_cell_size_per_token)
+
+    total_layers = int(target_num_layers) + int(draft_num_layers)
+    return (
+        int(target_cell_size_per_token) * int(total_layers) + int(target_num_layers) - 1
+    ) // int(target_num_layers)
+
+
+def resolve_dflash_verify_mask_policy(attn_backend: Any) -> tuple[str, bool]:
+    backend = attn_backend
+    for _ in range(4):
+        full_backend = getattr(backend, "full_attn_backend", None)
+        if full_backend is None:
+            break
+        backend = full_backend
+    backend_name = type(backend).__name__
+    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,
+    draft_token_num: int,
+    device: torch.device,
+) -> tuple[
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]:
+    key = (device.index, int(draft_token_num))
+    cached = _DFLASH_CHAIN_VERIFY_BUFFERS.get(key)
+    cap_bs = 0 if cached is None else int(cached["cap_bs"])
+    if cap_bs < bs:
+        new_cap = max(int(bs), cap_bs * 2 if cap_bs > 0 else int(bs))
+        retrieve_index = torch.arange(
+            new_cap * draft_token_num, dtype=torch.int64, device=device
+        ).view(new_cap, draft_token_num)
+        row_next = torch.arange(
+            1, draft_token_num + 1, dtype=torch.int64, device=device
+        )
+        row_next[-1] = -1
+        retrieve_next_token = row_next.unsqueeze(0).expand(new_cap, -1).clone()
+        retrieve_next_sibling = torch.full(
+            (new_cap, draft_token_num), -1, dtype=torch.int64, device=device
+        )
+        predicts = torch.empty(
+            (new_cap * draft_token_num,), dtype=torch.int32, device=device
+        )
+        accept_index = torch.empty(
+            (new_cap, draft_token_num), dtype=torch.int32, device=device
+        )
+        accept_token_num = torch.empty((new_cap,), dtype=torch.int32, device=device)
+        cached = {
+            "cap_bs": int(new_cap),
+            "retrieve_index": retrieve_index,
+            "retrieve_next_token": retrieve_next_token,
+            "retrieve_next_sibling": retrieve_next_sibling,
+            "predicts": predicts,
+            "accept_index": accept_index,
+            "accept_token_num": accept_token_num,
+        }
+        _DFLASH_CHAIN_VERIFY_BUFFERS[key] = cached
+
+    assert cached is not None
+    retrieve_index = cached["retrieve_index"][:bs]
+    retrieve_next_token = cached["retrieve_next_token"][:bs]
+    retrieve_next_sibling = cached["retrieve_next_sibling"][:bs]
+    predicts = cached["predicts"][: bs * draft_token_num]
+    accept_index = cached["accept_index"][:bs]
+    accept_token_num = cached["accept_token_num"][:bs]
+    return (
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        predicts,
+        accept_index,
+        accept_token_num,
+    )
+
+
+def build_target_layer_ids(num_target_layers: int, num_draft_layers: int) -> List[int]:
+    """Select target layer indices used to build DFlash context features.
+
+    Args:
+        num_target_layers: Number of transformer layers in the runtime target model.
+        num_draft_layers: Number of layers in the DFlash draft model.
+
+    Returns:
+        A list of 0-based target layer indices of length `num_draft_layers`.
+
+    Notes:
+        - DFlash uses hidden states after each selected target layer (HF-style).
+        - SGLang captures "before layer i", so the model hook will typically add +1
+          when mapping to capture points.
+    """
+    if num_target_layers <= 0:
+        raise ValueError(
+            f"num_target_layers must be positive, got {num_target_layers}."
+        )
+    if num_draft_layers <= 0:
+        raise ValueError(f"num_draft_layers must be positive, got {num_draft_layers}.")
+
+    if num_draft_layers == 1:
+        return [num_target_layers // 2]
+
+    start = 1
+    end = num_target_layers - 3
+    if end < start:
+        raise ValueError(
+            "DFlash layer selection requires num_target_layers >= 4. "
+            f"Got num_target_layers={num_target_layers}."
+        )
+
+    span = end - start
+    return [
+        int(round(start + (i * span) / (num_draft_layers - 1)))
+        for i in range(num_draft_layers)
+    ]
+
+
+def get_dflash_layer_types(config: Any) -> Optional[Sequence[str]]:
+    text_config = _get_text_config(config)
+    layer_types = _cfg_get(text_config, "layer_types", _cfg_get(config, "layer_types"))
+    if layer_types is None:
+        return None
+    if isinstance(layer_types, str) or not isinstance(layer_types, Sequence):
+        raise ValueError(
+            "DFLASH config.layer_types must be a sequence of attention type strings."
+        )
+    return layer_types
+
+
+def get_dflash_attention_sliding_window_size(config: Any) -> Optional[int]:
+    layer_types = get_dflash_layer_types(config)
+    if layer_types is None or "sliding_attention" not in layer_types:
+        return None
+
+    text_config = _get_text_config(config)
+    sliding_window = _cfg_get(
+        text_config, "sliding_window", _cfg_get(config, "sliding_window")
+    )
+    if sliding_window is None:
+        raise ValueError(
+            "DFLASH sliding_attention layers require config.sliding_window."
+        )
+
+    # HF sliding windows include the current token; SGLang stores window_left.
+    return int(sliding_window) - 1
+
+
+def _cfg_get(config: Any, key: str, default: Any = None) -> Any:
+    if isinstance(config, dict):
+        return config.get(key, default)
+    return getattr(config, key, default)
+
+
+def _get_text_config(config: Any) -> Any:
+    if config is None:
+        return None
+    if isinstance(config, dict):
+        return config.get("text_config", config)
+    text_config = getattr(config, "text_config", None)
+    if text_config is not None:
+        return text_config
+    get_text_config = getattr(config, "get_text_config", None)
+    if callable(get_text_config):
+        try:
+            resolved = get_text_config()
+            if resolved is not None:
+                return resolved
+        except TypeError:
+            pass
+    return config
+
+
+def _get_dflash_config(config: Any) -> dict:
+    if isinstance(config, dict):
+        cfg = config.get("dflash_config", None)
+    else:
+        cfg = getattr(config, "dflash_config", None)
+    if cfg is None:
+        return {}
+    if isinstance(cfg, dict):
+        return cfg
+
+    try:
+        return dict(cfg)
+    except Exception:
+        return {}
+
+
+def _parse_optional_int(
+    value: Any,
+    *,
+    field_name: str,
+    min_value: Optional[int] = None,
+) -> Optional[int]:
+    if value is None:
+        return None
+    try:
+        parsed = int(value)
+    except Exception as e:
+        raise ValueError(f"Invalid {field_name}={value!r}.") from e
+    if min_value is not None and parsed < int(min_value):
+        comparator = "positive" if int(min_value) == 1 else f">= {int(min_value)}"
+        raise ValueError(f"{field_name} must be {comparator}, got {parsed}.")
+    return parsed
+
+
+@dataclass(frozen=True)
+class DFlashDraftConfig:
+    num_hidden_layers: Optional[int]
+    num_target_layers: Optional[int]
+    block_size: Optional[int]
+    target_layer_ids: Optional[List[int]]
+    mask_token: str
+    mask_token_id: Optional[int]
+
+    def require_num_layers(self) -> int:
+        if self.num_hidden_layers is None:
+            raise ValueError(
+                "DFLASH requires draft num_hidden_layers in config. "
+                "Got config without num_hidden_layers."
+            )
+        return int(self.num_hidden_layers)
+
+    def resolve_block_size(self, *, default: Optional[int] = None) -> Optional[int]:
+        return self.block_size if self.block_size is not None else default
+
+    def resolve_target_layer_ids(
+        self,
+        *,
+        target_num_layers: int,
+        draft_num_layers: Optional[int] = None,
+    ) -> List[int]:
+        target_num_layers = int(target_num_layers)
+        if target_num_layers <= 0:
+            raise ValueError(
+                f"target_num_layers must be positive, got {target_num_layers}."
+            )
+
+        if self.target_layer_ids is None:
+            if draft_num_layers is None:
+                draft_num_layers = self.require_num_layers()
+            return build_target_layer_ids(target_num_layers, int(draft_num_layers))
+
+        resolved = list(self.target_layer_ids)
+        if len(resolved) <= 0:
+            raise ValueError(
+                "DFLASH dflash_config.target_layer_ids must be non-empty. "
+                f"Got len(target_layer_ids)={len(resolved)}."
+            )
+        for idx, val in enumerate(resolved):
+            if val < 0 or val >= target_num_layers:
+                raise ValueError(
+                    "DFLASH target_layer_ids contains an out-of-range layer id. "
+                    f"target_layer_ids[{idx}]={val}, target_num_layers={target_num_layers}."
+                )
+        return resolved
+
+
+def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig:
+    """Parse and validate DFLASH draft config fields from HF config/dict."""
+    dflash_cfg = _get_dflash_config(draft_hf_config)
+    draft_text_config = _get_text_config(draft_hf_config)
+
+    num_hidden_layers = _parse_optional_int(
+        _cfg_get(draft_text_config, "num_hidden_layers", None),
+        field_name="DFLASH draft num_hidden_layers",
+        min_value=1,
+    )
+    raw_num_target_layers = dflash_cfg.get(
+        "num_target_layers",
+        _cfg_get(draft_hf_config, "num_target_layers", None),
+    )
+    num_target_layers = _parse_optional_int(
+        raw_num_target_layers,
+        field_name="DFLASH draft num_target_layers",
+        min_value=1,
+    )
+
+    # Keep support for current checkpoints where block_size is top-level.
+    raw_block_size = dflash_cfg.get(
+        "block_size",
+        _cfg_get(draft_hf_config, "block_size", None),
+    )
+    block_size = _parse_optional_int(
+        raw_block_size,
+        field_name="DFLASH block_size",
+        min_value=1,
+    )
+
+    layer_ids = dflash_cfg.get(
+        "target_layer_ids",
+        _cfg_get(draft_hf_config, "target_layer_ids", None),
+    )
+    parsed_target_layer_ids: Optional[List[int]]
+    if layer_ids is None:
+        parsed_target_layer_ids = None
+    else:
+        if not isinstance(layer_ids, (list, tuple)):
+            raise ValueError(
+                "DFLASH dflash_config.target_layer_ids must be a list of ints, "
+                f"got type={type(layer_ids).__name__}."
+            )
+        parsed_target_layer_ids = [int(x) for x in layer_ids]
+        if len(parsed_target_layer_ids) <= 0:
+            raise ValueError(
+                "DFLASH dflash_config.target_layer_ids must be non-empty. "
+                f"Got len(target_layer_ids)={len(parsed_target_layer_ids)}."
+            )
+
+    mask_token = dflash_cfg.get("mask_token", None)
+    if mask_token is None:
+        mask_token = DEFAULT_DFLASH_MASK_TOKEN
+    if not isinstance(mask_token, str) or not mask_token:
+        raise ValueError(
+            "DFLASH dflash_config.mask_token must be a non-empty string, "
+            f"got {mask_token!r}."
+        )
+
+    mask_token_id = dflash_cfg.get("mask_token_id", None)
+    if mask_token_id is not None:
+        if not isinstance(mask_token_id, Integral) or isinstance(mask_token_id, bool):
+            raise ValueError(
+                "DFLASH dflash_config.mask_token_id must be an integer, "
+                f"got {mask_token_id!r} (type={type(mask_token_id).__name__})."
+            )
+        mask_token_id = int(mask_token_id)
+        if mask_token_id < 0:
+            raise ValueError(
+                "DFLASH dflash_config.mask_token_id must be non-negative, "
+                f"got {mask_token_id}."
+            )
+
+    return DFlashDraftConfig(
+        num_hidden_layers=num_hidden_layers,
+        num_target_layers=num_target_layers,
+        block_size=block_size,
+        target_layer_ids=parsed_target_layer_ids,
+        mask_token=mask_token,
+        mask_token_id=mask_token_id,
+    )
+
+
+def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]:
+    """Validate whether DFlash can slice KV weights from a fused QKV linear layer."""
+    quant_method = getattr(qkv_proj, "quant_method", None)
+    if not isinstance(quant_method, UnquantizedLinearMethod):
+        return (
+            False,
+            "quantized qkv_proj is not supported for this path "
+            f"(quant_method={type(quant_method).__name__})",
+        )
+    if not hasattr(qkv_proj, "weight"):
+        return False, "qkv weight tensor is missing"
+    return True, ""
+
+
+def can_dflash_use_fused_qkv_proj(qkv_proj: Any) -> Tuple[bool, str]:
+    """Validate whether a QKV layer is eligible for DFlash fused KV materialization."""
+    eligible, reason = can_dflash_slice_qkv_weight(qkv_proj)
+    if not eligible:
+        return False, reason
+    if getattr(qkv_proj, "bias", None) is not None:
+        return False, "qkv bias is not supported for fused KV path"
+    return True, ""
+
+
+def compute_dflash_correct_drafts_and_bonus(
+    *,
+    candidates: torch.Tensor,
+    target_predict: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Compute DFlash accept lengths and bonus tokens (greedy verify rule).
+
+    Args:
+        candidates: Token ids proposed by the DFlash draft, including the current token.
+            Shape: [bs, block_size]. candidates[:, 0] is the current token.
+        target_predict: Token ids predicted by the target model for each position in the block.
+            Shape: [bs, block_size]. target_predict[:, t] corresponds to argmax at position t.
+
+    Returns:
+        correct_len: int32 tensor [bs], number of accepted *draft* tokens (excluding current token and bonus token).
+        bonus: int64 tensor [bs], the target-predicted token at index correct_len (the "bonus" token to append).
+
+    Notes:
+        Matches the reference implementation rule:
+          accept while candidates[:, 1:] == target_predict[:, :-1] consecutively.
+    """
+    if candidates.ndim != 2:
+        raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}")
+    if target_predict.shape != candidates.shape:
+        raise ValueError(
+            "target_predict must have the same shape as candidates. "
+            f"candidates.shape={tuple(candidates.shape)}, target_predict.shape={tuple(target_predict.shape)}"
+        )
+
+    bs, block_size = candidates.shape
+    if bs <= 0:
+        raise ValueError(f"batch size must be positive, got {bs}.")
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}.")
+
+    matches = candidates[:, 1:] == target_predict[:, :-1]
+    correct_len = matches.to(torch.int32).cumprod(dim=1).sum(dim=1)
+    bonus = target_predict[torch.arange(bs, device=target_predict.device), correct_len]
+    return correct_len, bonus.to(torch.int64)
+
+
+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,
+    uniform_samples_for_final_sampling: Optional[torch.Tensor] = None,
+    use_sparse_topk: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Compute DFlash accept lengths and bonus tokens for non-greedy sampling.
+
+    This is a chain-specialized variant of speculative target-only verification:
+      - DFlash proposals are linear (topk == 1), so each verify level has at most one candidate.
+      - When a candidate is rejected at a level, the final token is sampled from
+        `relu(q - p)` where `p` has only the rejected candidate mass.
+    """
+    if not _DFLASH_SAMPLING_VERIFY_AVAILABLE:
+        raise RuntimeError(
+            "DFLASH non-greedy verification is unavailable on this build/device."
+        )
+    if candidates.ndim != 2:
+        raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}")
+    if next_token_logits.ndim != 2:
+        raise ValueError(
+            "next_token_logits must be 2D, "
+            f"got shape={tuple(next_token_logits.shape)}."
+        )
+
+    bs, draft_token_num = candidates.shape
+    if bs <= 0:
+        raise ValueError(f"batch size must be positive, got {bs}.")
+    if draft_token_num <= 0:
+        raise ValueError(f"draft_token_num must be positive, got {draft_token_num}.")
+    if next_token_logits.shape[0] != bs * draft_token_num:
+        raise ValueError(
+            "next_token_logits row count mismatch. "
+            f"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}."
+        )
+    if candidates.device != next_token_logits.device:
+        raise ValueError(
+            "candidates and next_token_logits must be on the same device, "
+            f"got {candidates.device} and {next_token_logits.device}."
+        )
+
+    if threshold_single is None:
+        from sglang.srt.runtime_context import get_server_args
+
+        threshold_single = get_server_args().speculative_accept_threshold_single
+    if threshold_acc is None:
+        from sglang.srt.runtime_context import get_server_args
+
+        threshold_acc = get_server_args().speculative_accept_threshold_acc
+    threshold_single = float(threshold_single)
+    threshold_acc = max(float(threshold_acc), 1e-9)
+
+    device = next_token_logits.device
+
+    if uniform_samples is None:
+        uniform_samples = torch.rand(
+            (bs, draft_token_num), dtype=torch.float32, device=device
+        )
+    else:
+        if uniform_samples.shape != (bs, draft_token_num):
+            raise ValueError(
+                "uniform_samples shape mismatch. "
+                f"Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}."
+            )
+        uniform_samples = uniform_samples.to(device=device, dtype=torch.float32)
+
+    if uniform_samples_for_final_sampling is None:
+        uniform_samples_for_final_sampling = torch.rand(
+            (bs,), dtype=torch.float32, device=device
+        )
+    else:
+        if uniform_samples_for_final_sampling.shape != (bs,):
+            raise ValueError(
+                "uniform_samples_for_final_sampling shape mismatch. "
+                f"Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}."
+            )
+        uniform_samples_for_final_sampling = uniform_samples_for_final_sampling.to(
+            device=device,
+            dtype=torch.float32,
+        )
+
+    target_probs = build_dflash_verify_target_probs(
+        next_token_logits=next_token_logits,
+        sampling_info=sampling_info,
+        draft_token_num=draft_token_num,
+        bs=bs,
+        max_top_k=max_top_k,
+        uniform_top_k_value=uniform_top_k_value,
+        use_sparse_topk=use_sparse_topk,
+    )
+    draft_probs = torch.zeros_like(target_probs)
+
+    (
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        predicts,
+        accept_index,
+        accept_token_num,
+    ) = _get_or_create_chain_verify_buffers(
+        bs=bs,
+        draft_token_num=draft_token_num,
+        device=device,
+    )
+    candidates_i64 = (
+        candidates if candidates.dtype == torch.int64 else candidates.to(torch.int64)
+    )
+    tree_speculative_sampling_target_only(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates_i64,
+        retrive_index=retrieve_index,
+        retrive_next_token=retrieve_next_token,
+        retrive_next_sibling=retrieve_next_sibling,
+        uniform_samples=uniform_samples,
+        uniform_samples_for_final_sampling=uniform_samples_for_final_sampling,
+        target_probs=target_probs,
+        draft_probs=draft_probs,
+        threshold_single=threshold_single,
+        threshold_acc=threshold_acc,
+        deterministic=True,
+    )
+
+    correct_len = accept_token_num
+    row_ids = torch.arange(bs, dtype=torch.long, device=device)
+    accept_pos = accept_index[row_ids, correct_len.to(torch.long)].to(torch.long)
+    bonus = predicts[accept_pos].to(torch.int64)
+    return correct_len, bonus
+
+
+def build_dflash_verify_target_probs(
+    *,
+    next_token_logits: torch.Tensor,
+    sampling_info: Any,
+    draft_token_num: int,
+    bs: int,
+    max_top_k: Optional[int] = None,
+    uniform_top_k_value: Optional[int] = None,
+    use_sparse_topk: bool = True,
+) -> torch.Tensor:
+    device = next_token_logits.device
+    need_top_k = bool(getattr(sampling_info, "need_top_k_sampling", True))
+    need_top_p = bool(getattr(sampling_info, "need_top_p_sampling", False))
+    expanded_temperature = torch.repeat_interleave(
+        sampling_info.temperatures, draft_token_num, dim=0
+    )
+    scaled_logits = next_token_logits / expanded_temperature
+    sparse_topk_applied = False
+
+    if use_sparse_topk and need_top_k:
+        repeated_top_ks = torch.repeat_interleave(
+            sampling_info.top_ks, draft_token_num, dim=0
+        ).to(dtype=torch.int64)
+        vocab_size = int(scaled_logits.shape[-1])
+        repeated_top_ks.clamp_(min=1, max=vocab_size)
+        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 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, :
+                ]
+                valid = ranks < repeated_top_ks.unsqueeze(1)
+                topk_logits = topk_logits.masked_fill(~valid, float("-inf"))
+
+            topk_probs = F.softmax(topk_logits, dim=-1)
+            if need_top_p:
+                repeated_top_ps = torch.repeat_interleave(
+                    sampling_info.top_ps, draft_token_num, dim=0
+                )
+                topk_probs = top_p_renorm_prob(topk_probs, repeated_top_ps)
+
+            target_probs = torch.zeros_like(scaled_logits, dtype=topk_probs.dtype)
+            target_probs.scatter_(1, topk_indices, topk_probs)
+            sparse_topk_applied = True
+
+    if not sparse_topk_applied:
+        target_probs = F.softmax(scaled_logits, dim=-1)
+        if need_top_k:
+            target_probs = top_k_renorm_prob(
+                target_probs,
+                torch.repeat_interleave(sampling_info.top_ks, draft_token_num, dim=0),
+            )
+        if need_top_p:
+            target_probs = top_p_renorm_prob(
+                target_probs,
+                torch.repeat_interleave(sampling_info.top_ps, draft_token_num, dim=0),
+            )
+    return target_probs.view(bs, draft_token_num, -1).contiguous()
+
+
+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
+        or req.sampling_params.ebnf is not None
+        or req.sampling_params.structural_tag is not None
+    ):
+        return (
+            "DFLASH speculative decoding does not support "
+            "grammar-constrained decoding yet."
+        )
+
+    return None

python/sglang/srt/speculative/draft_utils.py

@@ -0,0 +1,394 @@
+import logging
+
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.common import (
+    cpu_has_amx_support,
+    is_blackwell,
+    is_cpu,
+    is_hip,
+    is_musa,
+    is_npu,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class DraftBackendFactory:
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        draft_model_runner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        self.server_args = server_args
+        self.draft_model_runner = draft_model_runner
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.draft_attn_backend = server_args.speculative_draft_attention_backend
+
+    def _create_backend(
+        self, backend_name: str, backend_map: dict, error_template: str
+    ):
+        backend_type = (
+            self.draft_attn_backend
+            if self.draft_attn_backend
+            else getattr(self.server_args, backend_name)
+        )
+        if backend_type is None:
+            backend_type = self.server_args.attention_backend
+
+        if backend_type not in backend_map:
+            raise ValueError(error_template.format(backend_type=backend_type))
+
+        return backend_map[backend_type]()
+
+    def create_decode_backend(self):
+        # No multi-step draft backend for steps=0 (nospec) or steps=1.
+        if self.speculative_num_steps <= 1:
+            return None
+
+        backend_map = {
+            "flashinfer": self._create_flashinfer_decode_backend,
+            "triton": self._create_triton_decode_backend,
+            "intel_amx": self._create_intel_amx_decode_backend,
+            "aiter": self._create_aiter_decode_backend,
+            "fa3": self._create_fa3_decode_backend,
+            "hybrid_linear_attn": self._create_hybrid_linear_attn_decode_backend,
+            "flashmla": self._create_flashmla_decode_backend,
+            "trtllm_mha": self._create_trtllm_mha_decode_backend,
+            "trtllm_mla": self._create_trtllm_mla_decode_backend,
+            "cutedsl_mla": self._create_cutedsl_mla_decode_backend,
+            "tokenspeed_mla": self._create_tokenspeed_mla_decode_backend,
+            "dsa": self._create_dsa_decode_backend,
+            "nsa": self._create_dsa_decode_backend,  # Deprecated alias for "dsa"
+            "ascend": self._create_ascend_decode_backend,
+            "fa4": self._create_fa4_decode_backend,
+            "dsv4": self._create_dsv4_decode_backend,
+        }
+
+        return self._create_backend(
+            "decode_attention_backend",
+            backend_map,
+            "EAGLE is not supported in decode attention backend {backend_type}",
+        )
+
+    def create_draft_extend_backend(self):
+        backend_map = {
+            "flashinfer": self._create_flashinfer_prefill_backend,
+            "triton": self._create_triton_prefill_backend,
+            "intel_amx": self._create_intel_amx_prefill_backend,
+            "aiter": self._create_aiter_prefill_backend,
+            "fa3": self._create_fa3_prefill_backend,
+            "hybrid_linear_attn": self._create_hybrid_linear_attn_prefill_backend,
+            "flashmla": self._create_flashmla_prefill_backend,
+            "trtllm_mha": self._create_trtllm_mha_prefill_backend,
+            "trtllm_mla": self._create_trtllm_mla_prefill_backend,
+            # cute-dsl MLA only supports decode; draft-extend falls back to trtllm-gen.
+            "cutedsl_mla": self._create_trtllm_mla_prefill_backend,
+            "tokenspeed_mla": self._create_tokenspeed_mla_prefill_backend,
+            "dsa": self._create_dsa_prefill_backend,
+            "nsa": self._create_dsa_prefill_backend,  # Deprecated alias for "dsa"
+            "ascend": self._create_ascend_prefill_backend,
+            "fa4": self._create_fa4_prefill_backend,
+            "dsv4": self._create_dsv4_prefill_backend,
+        }
+        backend_name = (
+            "decode_attention_backend"
+            if self.server_args.speculative_attention_mode == "decode"
+            else "prefill_attention_backend"
+        )
+        return self._create_backend(
+            backend_name,
+            backend_map,
+            "EAGLE is not supported in attention backend {backend_type}",
+        )
+
+    def _create_dsa_decode_backend(self):
+        from sglang.srt.layers.attention.dsa_backend import (
+            DeepseekSparseAttnMultiStepBackend,
+        )
+
+        return DeepseekSparseAttnMultiStepBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_dsa_prefill_backend(self):
+        from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
+
+        return DeepseekSparseAttnBackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_flashinfer_decode_backend(self):
+        if not self.draft_model_runner.use_mla_backend:
+            from sglang.srt.layers.attention.flashinfer_backend import (
+                FlashInferMultiStepDraftBackend,
+            )
+
+            return FlashInferMultiStepDraftBackend(
+                self.draft_model_runner, self.topk, self.speculative_num_steps
+            )
+        else:
+            from sglang.srt.layers.attention.flashinfer_mla_backend import (
+                FlashInferMLAMultiStepDraftBackend,
+            )
+
+            return FlashInferMLAMultiStepDraftBackend(
+                self.draft_model_runner, self.topk, self.speculative_num_steps
+            )
+
+    def _create_triton_decode_backend(self):
+        from sglang.srt.layers.attention.triton_backend import (
+            TritonMultiStepDraftBackend,
+        )
+
+        return TritonMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_intel_amx_decode_backend(self):
+        from sglang.srt.layers.attention.intel_amx_backend import (
+            IntelAMXMultiStepDraftBackend,
+        )
+
+        return IntelAMXMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_hybrid_linear_attn_decode_backend(self):
+        if is_cpu() and cpu_has_amx_support():
+            return self._create_intel_amx_decode_backend()
+        if is_blackwell():
+            return self._create_triton_decode_backend()
+        return self._create_fa3_decode_backend()
+
+    def _create_hybrid_linear_attn_prefill_backend(self):
+        if is_cpu() and cpu_has_amx_support():
+            return self._create_intel_amx_prefill_backend()
+        if is_blackwell():
+            return self._create_triton_prefill_backend()
+        return self._create_fa3_prefill_backend()
+
+    def _create_aiter_decode_backend(self):
+        from sglang.srt.layers.attention.aiter_backend import AiterMultiStepDraftBackend
+
+        return AiterMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_fa_decode_backend(self, fa_impl_ver: int = 3):
+        if not is_musa():
+            from sglang.srt.layers.attention.flashattention_backend import (
+                FlashAttentionMultiStepBackend,
+            )
+        else:
+            from sglang.srt.hardware_backend.musa.attention.flashattention_backend import (
+                MusaFlashAttentionMultiStepBackend as FlashAttentionMultiStepBackend,
+            )
+
+        return FlashAttentionMultiStepBackend(
+            self.draft_model_runner,
+            self.topk,
+            self.speculative_num_steps,
+            fa_impl_ver=fa_impl_ver,
+        )
+
+    def _create_fa3_decode_backend(self):
+        return self._create_fa_decode_backend(fa_impl_ver=3)
+
+    def _create_fa4_decode_backend(self):
+        return self._create_fa_decode_backend(fa_impl_ver=4)
+
+    def _create_flashmla_decode_backend(self):
+        from sglang.srt.layers.attention.flashmla_backend import (
+            FlashMLAMultiStepDraftBackend,
+        )
+
+        return FlashMLAMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_trtllm_mha_decode_backend(self):
+        from sglang.srt.layers.attention.trtllm_mha_backend import (
+            TRTLLMHAAttnMultiStepDraftBackend,
+        )
+
+        return TRTLLMHAAttnMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_trtllm_mla_decode_backend(self, backend: str = "trtllm-gen"):
+        if not self.draft_model_runner.use_mla_backend:
+            raise ValueError(
+                "trtllm_mla backend requires MLA model (use_mla_backend=True)."
+            )
+
+        from sglang.srt.layers.attention.trtllm_mla_backend import (
+            TRTLLMMLAMultiStepDraftBackend,
+        )
+
+        return TRTLLMMLAMultiStepDraftBackend(
+            self.draft_model_runner,
+            self.topk,
+            self.speculative_num_steps,
+            backend=backend,
+        )
+
+    def _create_cutedsl_mla_decode_backend(self):
+        return self._create_trtllm_mla_decode_backend(backend="cute-dsl")
+
+    def _create_tokenspeed_mla_decode_backend(self):
+        if not self.draft_model_runner.use_mla_backend:
+            raise ValueError(
+                "tokenspeed_mla backend requires MLA model (use_mla_backend=True)."
+            )
+
+        from sglang.srt.layers.attention.tokenspeed_mla_backend import (
+            TokenspeedMLAMultiStepDraftBackend,
+        )
+
+        return TokenspeedMLAMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_ascend_decode_backend(self):
+        from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
+            AscendAttnMultiStepDraftBackend,
+        )
+
+        return AscendAttnMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_dsv4_decode_backend(self):
+        # Decode here is the EAGLE multi-step draft decode path.
+        if is_npu():
+            from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import (
+                DeepseekV4AscendMultiStepDraftBackend,
+            )
+
+            return DeepseekV4AscendMultiStepDraftBackend(
+                self.draft_model_runner, self.topk, self.speculative_num_steps
+            )
+        elif is_hip():
+            from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
+                DeepseekV4MultiStepBackend,
+            )
+        else:
+            from sglang.srt.layers.attention.deepseek_v4_backend import (
+                DeepseekV4MultiStepBackend,
+            )
+
+        return DeepseekV4MultiStepBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_flashinfer_prefill_backend(self):
+        if not self.draft_model_runner.use_mla_backend:
+            from sglang.srt.layers.attention.flashinfer_backend import (
+                FlashInferAttnBackend,
+            )
+
+            return FlashInferAttnBackend(self.draft_model_runner, skip_prefill=False)
+        else:
+            from sglang.srt.layers.attention.flashinfer_mla_backend import (
+                FlashInferMLAAttnBackend,
+            )
+
+            return FlashInferMLAAttnBackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_triton_prefill_backend(self):
+        from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
+
+        return TritonAttnBackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_intel_amx_prefill_backend(self):
+        from sglang.srt.layers.attention.intel_amx_backend import IntelAMXAttnBackend
+
+        return IntelAMXAttnBackend(self.draft_model_runner)
+
+    def _create_aiter_prefill_backend(self):
+        from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
+
+        return AiterAttnBackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_fa_prefill_backend(self, fa_impl_ver: int = 3):
+        if not is_musa():
+            from sglang.srt.layers.attention.flashattention_backend import (
+                FlashAttentionBackend,
+            )
+        else:
+            from sglang.srt.hardware_backend.musa.attention.flashattention_backend import (
+                MusaFlashAttentionBackend as FlashAttentionBackend,
+            )
+        return FlashAttentionBackend(
+            self.draft_model_runner, skip_prefill=False, fa_impl_ver=fa_impl_ver
+        )
+
+    def _create_fa3_prefill_backend(self):
+        return self._create_fa_prefill_backend(fa_impl_ver=3)
+
+    def _create_fa4_prefill_backend(self):
+        return self._create_fa_prefill_backend(fa_impl_ver=4)
+
+    def _create_trtllm_mha_prefill_backend(self):
+        from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
+
+        return TRTLLMHAAttnBackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_trtllm_mla_prefill_backend(self):
+        if not self.draft_model_runner.use_mla_backend:
+            raise ValueError(
+                "trtllm_mla backend requires MLA model (use_mla_backend=True)."
+            )
+
+        from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
+
+        return TRTLLMMLABackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_tokenspeed_mla_prefill_backend(self):
+        if not self.draft_model_runner.use_mla_backend:
+            raise ValueError(
+                "tokenspeed_mla backend requires MLA model (use_mla_backend=True)."
+            )
+
+        from sglang.srt.layers.attention.tokenspeed_mla_backend import (
+            TokenspeedMLABackend,
+        )
+
+        return TokenspeedMLABackend(self.draft_model_runner, skip_prefill=False)
+
+    def _create_ascend_prefill_backend(self):
+        from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
+            AscendAttnBackend,
+        )
+
+        return AscendAttnBackend(self.draft_model_runner)
+
+    def _create_flashmla_prefill_backend(self):
+        logger.warning(
+            "flashmla prefill backend is not yet supported for draft extend."
+        )
+        return None
+
+    def _create_dsv4_prefill_backend(self):
+        # On NPU the "dsv4" backend resolves to the Ascend V4 subclass; its
+        # draft-extend path uses the registered DSV4 prefill backend.
+        if is_npu():
+            from sglang.srt.layers.attention.attention_registry import (
+                ATTENTION_BACKENDS,
+            )
+
+            return ATTENTION_BACKENDS["dsv4"](self.draft_model_runner)
+        elif is_hip():
+            from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
+                DeepseekV4HipRadixBackend,
+            )
+
+            return DeepseekV4HipRadixBackend(
+                self.draft_model_runner, skip_prefill=False
+            )
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            DeepseekV4AttnBackend,
+        )
+
+        return DeepseekV4AttnBackend(self.draft_model_runner, skip_prefill=False)

python/sglang/srt/speculative/draft_worker_common.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import logging
+from copy import deepcopy
+from typing import TYPE_CHECKING, Optional
+
+import msgspec
+import torch
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+from sglang.srt.runtime_context import get_context, get_server_args
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.dflash_info import DFlashVerifyInput
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_DRAFT_BACKENDS = ("flashinfer", "fa3", "fa4", "triton", "ascend")
+
+
+class DraftWorkerBundle(msgspec.Struct, frozen=True):
+    draft_worker: TpModelWorker
+    draft_model_runner: ModelRunner
+    draft_model: torch.nn.Module
+    resolved_attention_backend: str
+
+
+def _resolve_draft_attention_backend_fallback(
+    *, draft_server_args: ServerArgs, algo_label: str
+) -> str:
+    draft_backend = draft_server_args.speculative_draft_attention_backend
+    if draft_backend is None:
+        draft_backend, _ = draft_server_args.get_attention_backends()
+    if draft_backend is None:
+        return "triton" if torch.version.hip else "flashinfer"
+    if draft_backend not in _SUPPORTED_DRAFT_BACKENDS:
+        fallback = "triton" if torch.version.hip else "flashinfer"
+        logger.warning(
+            "%s draft worker only supports attention_backend in %s for now, "
+            "but got %r. Falling back to '%s'.",
+            algo_label,
+            _SUPPORTED_DRAFT_BACKENDS,
+            draft_backend,
+            fallback,
+        )
+        return fallback
+    return draft_backend
+
+
+def build_draft_tp_worker(
+    *,
+    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_model_config: ModelConfig,
+    algo_label: str,
+    attention_backend_override: Optional[str] = None,
+) -> DraftWorkerBundle:
+    draft_server_args = deepcopy(server_args)
+    # An override names a draft-specific backend the caller has already
+    # validated (e.g. a self-drafting architecture); it skips the generic
+    # supported-backend fallback below.
+    draft_backend = attention_backend_override or (
+        _resolve_draft_attention_backend_fallback(
+            draft_server_args=draft_server_args, algo_label=algo_label
+        )
+    )
+    # Post-resolution ServerArgs rejects bare assignment; route the draft-copy
+    # adjustments through the audited mutation point. Keep the resolved value
+    # on speculative_draft_attention_backend: downstream draft-worker logic
+    # keys on that field (backend selection in _get_attention_backend and the
+    # fa4-draft KV dtype override in configure_kv_cache_dtype), so nulling it
+    # would silently skip those paths. context_length keeps the draft aligned
+    # with the target.
+    draft_server_args.override(
+        "draft_worker.build",
+        skip_tokenizer_init=True,
+        speculative_draft_attention_backend=draft_backend,
+        prefill_attention_backend=None,
+        decode_attention_backend=None,
+        attention_backend=draft_backend,
+        context_length=target_model_config.context_len,
+    )
+
+    saved_server_args = get_server_args()
+    try:
+        draft_worker = TpModelWorker(
+            server_args=draft_server_args,
+            gpu_id=gpu_id,
+            tp_rank=tp_rank,
+            moe_ep_rank=moe_ep_rank,
+            pp_rank=0,
+            attn_cp_rank=attn_cp_rank,
+            moe_dp_rank=moe_dp_rank,
+            dp_rank=dp_rank,
+            nccl_port=nccl_port,
+            is_draft_worker=True,
+        )
+    finally:
+        get_context().set_server_args(saved_server_args)
+
+    draft_model_runner = draft_worker.model_runner
+    draft_worker.draft_runner = draft_model_runner
+    return DraftWorkerBundle(
+        draft_worker=draft_worker,
+        draft_model_runner=draft_model_runner,
+        draft_model=draft_model_runner.model,
+        resolved_attention_backend=draft_backend,
+    )
+
+
+def make_draft_input_v2(
+    *,
+    bonus_tokens: torch.Tensor,
+    new_seq_lens: torch.Tensor,
+) -> DFlashDraftInputV2:
+    bs = int(new_seq_lens.numel())
+    device = bonus_tokens.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),
+        bonus_tokens=bonus_tokens.to(dtype=torch.int64),
+        new_seq_lens=new_seq_lens.to(dtype=torch.int64),
+        hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16),
+    )
+
+
+def make_draft_block_spec_info(
+    *,
+    draft_token_num: int,
+    device: torch.device,
+) -> DFlashVerifyInput:
+    return DFlashVerifyInput(
+        draft_token=torch.empty((0,), dtype=torch.long, device=device),
+        positions=torch.empty((0,), dtype=torch.int64, device=device),
+        draft_token_num=int(draft_token_num),
+        custom_mask=None,
+        capture_hidden_mode=CaptureHiddenMode.NULL,
+    )
+
+
+def make_draft_sampler_capture_hook(draft_sampler):
+
+    def capture_hook(runner, out, forward_batch, num_tokens):
+        del runner, num_tokens
+        if not isinstance(out, LogitsProcessorOutput) or out.hidden_states is None:
+            raise RuntimeError(
+                "draft sampler set but the draft forward has no "
+                "hidden_states to capture into the graph."
+            )
+        draft_sampler(out.hidden_states, forward_batch.input_ids)
+
+    return capture_hook
+
+
+def build_block_pos_offsets(*, length: int, device: torch.device) -> torch.Tensor:
+    return torch.arange(int(length), device=device, dtype=torch.int64)

python/sglang/srt/speculative/dspark_components/dspark_block_accept_estimator.py

@@ -0,0 +1,835 @@
+from __future__ import annotations
+
+import json
+import logging
+import math
+from collections import deque
+from pathlib import Path
+from typing import Any, List, Optional, Tuple, Union
+
+import msgspec
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+logger = logging.getLogger(__name__)
+
+_GATHER_ROW_CHUNK = 512
+_STATE_SWEEP_INTERVAL = 1024
+_STATE_EXPIRE_STEPS = 4096
+_FLUSH_EVERY_STEPS = 16
+_PENDING_BUCKET_MIN = 16
+_DEFAULT_ONLINE_WINDOW_STEPS = 256
+
+
+SKIP_STEP_WARNING = (
+    "skipping step: {} (pending blocks of affected requests "
+    "are dropped by the seq-len continuity check)"
+)
+
+
+def block_accept_skip_reason(
+    *,
+    logits_adjustments_are_noop: bool,
+    corrected_logits: Optional[Any],
+) -> Optional[str]:
+    if not logits_adjustments_are_noop:
+        return (
+            "non-noop logits adjustments (penalizer/logit_bias/grammar) "
+            "in batch; cross-step conditioning of the gathered target "
+            "probabilities would be state-dependent"
+        )
+    if corrected_logits is None:
+        return "corrected_logits unavailable (folded draft path)"
+    return None
+
+
+def warn_once(warned_reasons: set, *, reason: str) -> None:
+    if reason not in warned_reasons:
+        warned_reasons.add(reason)
+        logger.warning(
+            "DSPARK block accept estimate recorder: %s (warned once)", reason
+        )
+
+
+def gather_chunked_token_logprobs(
+    *,
+    logits,
+    row_indices,
+    token_indices,
+    per_row_temps,
+    chunk_size: int,
+):
+    """Chunked per-row token logprob gather: logprob of token_indices[i] under
+    logits[row_indices[i]] / per_row_temps[i], computed chunk_size rows at a
+    time to bound the fp32 softmax workspace."""
+    results = []
+    for start in range(0, row_indices.shape[0], chunk_size):
+        end = start + chunk_size
+        rows = logits[row_indices[start:end]].to(torch.float32)
+        rows = rows / per_row_temps[start:end, None]
+        log_norm = torch.logsumexp(rows, dim=-1)
+        token_logits = rows.gather(dim=1, index=token_indices[start:end, None]).squeeze(
+            1
+        )
+        results.append(token_logits - log_norm)
+    return torch.cat(results)
+
+
+def _pending_bucket(count: int) -> int:
+    if count == 0:
+        return 0
+    bucket = _PENDING_BUCKET_MIN
+    while bucket < count:
+        bucket *= 2
+    return bucket
+
+
+class _CeilingSnapshot(msgspec.Struct):
+    window_lo: float
+    window_hi: float
+    window_blocks: int
+    window_horizon: int
+    cumulative_lo: float
+    cumulative_hi: float
+    cumulative_blocks: int
+
+
+class _OnlineCeiling:
+    def __init__(self, *, log_interval: int, window_steps: int) -> None:
+        self._log_interval = log_interval
+        self._window_steps = window_steps
+        self._steps: deque[Tuple[int, float, float, int]] = deque()
+        self._win_lo = 0.0
+        self._win_hi = 0.0
+        self._win_count = 0
+        self._cum_lo = 0.0
+        self._cum_hi = 0.0
+        self._cum_count = 0
+        self._max_forward_ct = 0
+
+    def add(self, *, forward_ct: int, lo: float, hi: float) -> None:
+        self._max_forward_ct = max(self._max_forward_ct, forward_ct)
+        if self._steps and self._steps[-1][0] == forward_ct:
+            fct, slo, shi, c = self._steps[-1]
+            self._steps[-1] = (fct, slo + lo, shi + hi, c + 1)
+        else:
+            self._steps.append((forward_ct, lo, hi, 1))
+        self._win_lo += lo
+        self._win_hi += hi
+        self._win_count += 1
+        self._cum_lo += lo
+        self._cum_hi += hi
+        self._cum_count += 1
+        self._evict(forward_ct=self._max_forward_ct)
+
+    def _evict(self, *, forward_ct: int) -> None:
+        cutoff = forward_ct - self._window_steps
+        while self._steps and self._steps[0][0] <= cutoff:
+            _, slo, shi, c = self._steps.popleft()
+            self._win_lo -= slo
+            self._win_hi -= shi
+            self._win_count -= c
+
+    def estimate(self) -> Optional[_CeilingSnapshot]:
+        if self._cum_count == 0:
+            return None
+        return _CeilingSnapshot(
+            window_lo=self._win_lo / self._win_count,
+            window_hi=self._win_hi / self._win_count,
+            window_blocks=self._win_count,
+            window_horizon=min(self._window_steps, self._max_forward_ct),
+            cumulative_lo=self._cum_lo / self._cum_count,
+            cumulative_hi=self._cum_hi / self._cum_count,
+            cumulative_blocks=self._cum_count,
+        )
+
+    def maybe_log(self, *, forward_ct: int) -> None:
+        if self._log_interval <= 0 or forward_ct % self._log_interval != 0:
+            return
+        snap = self.estimate()
+        if snap is None:
+            return
+        logger.info(
+            "DSpark uncapped-acc-len estimate (forward_ct=%d): "
+            "last %d passes ~%.3f [%.3f, %.3f] w=%.3f (%d blocks) | "
+            "cumulative ~%.3f [%.3f, %.3f] w=%.3f (%d blocks)",
+            forward_ct,
+            snap.window_horizon,
+            0.5 * (snap.window_lo + snap.window_hi),
+            snap.window_lo,
+            snap.window_hi,
+            snap.window_hi - snap.window_lo,
+            snap.window_blocks,
+            0.5 * (snap.cumulative_lo + snap.cumulative_hi),
+            snap.cumulative_lo,
+            snap.cumulative_hi,
+            snap.cumulative_hi - snap.cumulative_lo,
+            snap.cumulative_blocks,
+        )
+
+
+class _PendingBlock(msgspec.Struct):
+    forward_ct: int
+    anchor_pos: int
+    window: int
+    trimmed_tokens: List[int]
+    next_offset: int
+    q_lps: List[float] = []
+    est_prod: float = 1.0
+    est_lo_extra: float = 0.0
+
+
+class _RequestState(msgspec.Struct):
+    expected_seq_len: int = -1
+    last_seen_ct: int = 0
+    pending: List[_PendingBlock] = []
+
+
+class _PendingPlan(msgspec.Struct):
+    rows: List[int]
+    tokens: List[int]
+    slot_lookup: dict[tuple[int, int, int], int]
+
+
+class _SettleBatch(msgspec.Struct):
+    forward_ct: int
+    rids: List[str]
+    row_meta: List[List[int]]
+    drafts: List[List[int]]
+    q_all: List[List[float]]
+    target_diag: List[List[float]]
+    pending_logprobs: List[float]
+    slot_lookup: dict[tuple[int, int, int], int]
+
+    @classmethod
+    def from_bundle(cls, bundle: dict[str, Any]) -> _SettleBatch:
+        return cls(
+            forward_ct=bundle["forward_ct"],
+            rids=bundle["rids"],
+            row_meta=bundle["row_meta"].tolist(),
+            drafts=bundle["draft_tokens"].tolist(),
+            q_all=bundle["q_all"].tolist(),
+            target_diag=bundle["target_diag_logprobs"].tolist(),
+            pending_logprobs=bundle["pending_logprobs"].tolist(),
+            slot_lookup=bundle["pending_slot_lookup"],
+        )
+
+
+class BlockAcceptEstimateRecorder:
+    def __init__(
+        self,
+        *,
+        path: str,
+        gamma: int,
+        device: Union[str, torch.device],
+        online_log_interval: int = 0,
+        online_window_steps: int = 0,
+    ) -> None:
+        self._gamma = gamma
+        self._last_forward_ct = 0
+        if path:
+            self._path: Optional[Path] = Path(path)
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+            self._file = self._path.open("w")
+        else:
+            self._path = None
+            self._file = None
+        self._device = torch.device(device)
+        self._states: dict[str, _RequestState] = {}
+        self._steps_since_flush = 0
+        self._observed_step_ct = 0
+        self._discontinuity_drop_ct = 0
+        self._skipped_step_ct = 0
+        self._warned_skip_reasons: set[str] = set()
+        self._finish_intents: dict[str, bool] = {}
+
+        self._online = _OnlineCeiling(
+            log_interval=online_log_interval,
+            window_steps=(
+                online_window_steps
+                if online_window_steps > 0
+                else (
+                    online_log_interval
+                    if online_log_interval > 0
+                    else _DEFAULT_ONLINE_WINDOW_STEPS
+                )
+            ),
+        )
+
+        self._retained_h2d: List[torch.Tensor] = []
+        self._delayed: Optional[DelayedDeviceHostHandler] = None
+        if self._device.type == "cuda":
+            self._delayed = DelayedDeviceHostHandler(
+                d2h_stream=torch.cuda.Stream(device=self._device)
+            )
+
+        logger.info(
+            "DSPARK block accept estimate recorder enabled: path=%s gamma=%d "
+            "async=%s online_log_interval=%d",
+            path,
+            gamma,
+            self._delayed is not None,
+            online_log_interval,
+        )
+
+    def observe_verify_step(
+        self,
+        *,
+        forward_ct: int,
+        rids: List[str],
+        draft_tokens: torch.Tensor,
+        corrected_logits: Optional[torch.Tensor],
+        draft_temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        target_logits: torch.Tensor,
+        target_temperatures: torch.Tensor,
+        truncated_sampling_mask: Optional[torch.Tensor],
+        logits_adjustments_are_noop: bool,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        bonus: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        layout: Optional[RaggedVerifyLayout],
+    ) -> None:
+        if (
+            self._delayed is not None
+            and torch.cuda.is_available()
+            and torch.cuda.is_current_stream_capturing()
+        ):
+            return
+
+        skip_reason = self._skip_reason(
+            logits_adjustments_are_noop=logits_adjustments_are_noop,
+            corrected_logits=corrected_logits,
+        )
+        if skip_reason is not None:
+            self._skip_step(reason=skip_reason)
+
+        def compute_on_device() -> Optional[dict[str, Any]]:
+            if skip_reason is not None:
+                return None
+            return self._build_device_bundle(
+                forward_ct=forward_ct,
+                rids=rids,
+                draft_tokens=draft_tokens,
+                corrected_logits=corrected_logits,
+                draft_temperatures=draft_temperatures,
+                greedy_mask=greedy_mask,
+                target_logits=target_logits,
+                target_temperatures=target_temperatures,
+                truncated_sampling_mask=truncated_sampling_mask,
+                correct_len=correct_len,
+                cap_trim_lens=cap_trim_lens,
+                bonus=bonus,
+                prefix_lens=prefix_lens,
+                layout=layout,
+            )
+
+        if self._delayed is not None:
+            self._delayed.step(
+                compute_on_device=compute_on_device,
+                postprocess_on_host=self._settle_and_write,
+            )
+        else:
+            bundle = compute_on_device()
+            if bundle is not None:
+                self._settle_and_write(bundle)
+
+    def flush(self) -> None:
+        if self._delayed is not None:
+            self._delayed.step(
+                compute_on_device=lambda: None,
+                postprocess_on_host=self._settle_and_write,
+            )
+            self._apply_all_finish_intents()
+        if self._file is not None:
+            self._file.flush()
+        self._steps_since_flush = 0
+
+    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        if self._delayed is None:
+            self._finalize_request(
+                rid=rid, natural_stop=natural_stop, forward_ct=self._last_forward_ct
+            )
+        else:
+            self._finish_intents[rid] = natural_stop
+
+    def _apply_all_finish_intents(self) -> None:
+        for rid in list(self._finish_intents):
+            self._finalize_request(
+                rid=rid,
+                natural_stop=self._finish_intents.pop(rid),
+                forward_ct=self._last_forward_ct,
+            )
+
+    def _finalize_request(
+        self, *, rid: str, natural_stop: bool, forward_ct: int
+    ) -> None:
+        state = self._states.pop(rid, None)
+        if state is None:
+            return
+        for block in state.pending:
+            if natural_stop:
+                self._finalize_eos_online(block, forward_ct=forward_ct)
+            else:
+                self._finalize_at_end_online(block, forward_ct=forward_ct)
+        if natural_stop and state.pending:
+            self._write_eos_marker(rid=rid, blocks=state.pending)
+
+    def _finalize_eos_online(self, block: _PendingBlock, *, forward_ct: int) -> None:
+        lo = block.window + 1.0 + block.est_lo_extra
+        self._online.add(forward_ct=forward_ct, lo=lo, hi=lo)
+
+    def _write_eos_marker(self, *, rid: str, blocks: List[_PendingBlock]) -> None:
+        if self._file is None:
+            return
+        marker = {"rid": rid, "eos_end": [block.forward_ct for block in blocks]}
+        self._file.write(json.dumps(marker) + "\n")
+
+    def online_estimate(self) -> Optional[_CeilingSnapshot]:
+        return self._online.estimate()
+
+    def estimate_log_suffix(self) -> Optional[str]:
+        snap = self.online_estimate()
+        if snap is None:
+            return None
+        mid = 0.5 * (snap.cumulative_lo + snap.cumulative_hi)
+        return (
+            f"est uncap acc len: {mid:.2f} "
+            f"[{snap.cumulative_lo:.2f}, {snap.cumulative_hi:.2f}]"
+        )
+
+    def drain_pending_online(self) -> None:
+        for state in self._states.values():
+            for block in state.pending:
+                self._finalize_at_end_online(block, forward_ct=self._last_forward_ct)
+            state.pending = []
+
+    def _finalize_walk_online(
+        self, block: _PendingBlock, *, diverged: bool, forward_ct: int
+    ) -> None:
+        base = block.window + 1.0
+        lo = base + block.est_lo_extra
+        if diverged:
+            offset = block.next_offset - 1
+            tail = (
+                block.est_prod * (self._gamma - offset) if offset < self._gamma else 0.0
+            )
+        else:
+            tail = 0.0
+        self._online.add(forward_ct=forward_ct, lo=lo, hi=lo + tail)
+
+    def _finalize_at_end_online(self, block: _PendingBlock, *, forward_ct: int) -> None:
+        base = block.window + 1.0
+        lo = base + block.est_lo_extra
+        tail = block.est_prod * (self._gamma - block.next_offset + 1)
+        self._online.add(forward_ct=forward_ct, lo=lo, hi=lo + tail)
+
+    def _build_device_bundle(
+        self,
+        *,
+        forward_ct: int,
+        rids: List[str],
+        draft_tokens: torch.Tensor,
+        corrected_logits: torch.Tensor,
+        draft_temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        target_logits: torch.Tensor,
+        target_temperatures: torch.Tensor,
+        truncated_sampling_mask: Optional[torch.Tensor],
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        bonus: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        layout: Optional[RaggedVerifyLayout],
+    ) -> dict[str, Any]:
+        gamma = self._gamma
+        rows_per_request = gamma + 1
+        bs = len(rids)
+        device = target_logits.device
+        assert draft_tokens.shape == (bs, gamma)
+        assert corrected_logits.shape[0] == bs and corrected_logits.shape[1] == gamma
+        assert target_logits.shape[0] == bs * rows_per_request
+
+        if truncated_sampling_mask is not None:
+            truncated_mask = truncated_sampling_mask
+        else:
+            truncated_mask = torch.zeros(bs, dtype=torch.bool, device=device)
+        if layout is not None:
+            verify_lens = layout.verify_lens
+        else:
+            verify_lens = torch.full(
+                (bs,), rows_per_request, dtype=torch.int32, device=device
+            )
+
+        draft_temps_full = (
+            draft_temperatures.reshape(bs).to(torch.float32).repeat_interleave(gamma)
+        )
+        target_temps_full = (
+            target_temperatures.reshape(bs)
+            .to(torch.float32)
+            .repeat_interleave(rows_per_request)
+        )
+        draft_flat = draft_tokens.reshape(-1)
+
+        q_all = self._gather_logprobs(
+            logits=corrected_logits.reshape(bs * gamma, -1),
+            row_indices=torch.arange(bs * gamma, device=device),
+            token_indices=draft_flat,
+            temps=draft_temps_full,
+        ).reshape(bs, gamma)
+        target_diag = self._gather_logprobs(
+            logits=target_logits,
+            row_indices=self._diag_rows(bs=bs, rows_per_request=rows_per_request),
+            token_indices=draft_flat,
+            temps=target_temps_full,
+        ).reshape(bs, gamma)
+
+        self._retained_h2d = []
+        plan = self._plan_pending(bs=bs, rows_per_request=rows_per_request, rids=rids)
+        pending_logprobs = self._gather_pending(
+            plan=plan,
+            target_logits=target_logits,
+            target_temps_full=target_temps_full,
+            device=device,
+        )
+
+        return {
+            "forward_ct": int(forward_ct),
+            "rids": list(rids),
+            "row_meta": self._pack_row_meta(
+                correct_len=correct_len,
+                cap_trim_lens=cap_trim_lens,
+                bonus=bonus,
+                prefix_lens=prefix_lens,
+                greedy_mask=greedy_mask,
+                truncated_mask=truncated_mask,
+                verify_lens=verify_lens,
+            ),
+            "draft_tokens": draft_tokens,
+            "q_all": q_all,
+            "target_diag_logprobs": target_diag,
+            "pending_logprobs": pending_logprobs,
+            "pending_slot_lookup": plan.slot_lookup,
+        }
+
+    def _diag_rows(self, *, bs: int, rows_per_request: int) -> torch.Tensor:
+        device = self._device
+        return (
+            (torch.arange(bs, device=device) * rows_per_request)[:, None]
+            + torch.arange(self._gamma, device=device)[None, :]
+        ).reshape(-1)
+
+    def _plan_pending(
+        self, *, bs: int, rows_per_request: int, rids: List[str]
+    ) -> _PendingPlan:
+        gamma = self._gamma
+        rows: List[int] = []
+        tokens: List[int] = []
+        slot_lookup: dict[tuple[int, int, int], int] = {}
+        for b in range(bs):
+            state = self._states.get(rids[b])
+            if state is None or not state.pending or state.expected_seq_len < 0:
+                continue
+            expected_seq_len = state.expected_seq_len
+            for block_idx, block in enumerate(state.pending):
+                offset = block.next_offset
+                while offset <= gamma:
+                    row = block.anchor_pos + offset - expected_seq_len
+                    if row < 0 or row >= rows_per_request:
+                        break
+                    slot_lookup[(b, block_idx, offset)] = len(rows)
+                    rows.append(b * rows_per_request + row)
+                    tokens.append(block.trimmed_tokens[offset - block.window - 1])
+                    offset += 1
+        return _PendingPlan(rows=rows, tokens=tokens, slot_lookup=slot_lookup)
+
+    def _gather_pending(
+        self,
+        *,
+        plan: _PendingPlan,
+        target_logits: torch.Tensor,
+        target_temps_full: torch.Tensor,
+        device: torch.device,
+    ) -> torch.Tensor:
+        bucket = _pending_bucket(len(plan.rows))
+        rows = plan.rows + [0] * (bucket - len(plan.rows))
+        tokens = plan.tokens + [0] * (bucket - len(plan.tokens))
+        return self._gather_logprobs(
+            logits=target_logits,
+            row_indices=self._host_to_device_async(rows, device=device),
+            token_indices=self._host_to_device_async(tokens, device=device),
+            temps=target_temps_full,
+        )
+
+    def _pack_row_meta(
+        self,
+        *,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        bonus: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        truncated_mask: torch.Tensor,
+        verify_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.stack(
+            [
+                correct_len.to(torch.int64),
+                cap_trim_lens.to(torch.int64),
+                bonus.to(torch.int64),
+                prefix_lens.to(torch.int64),
+                greedy_mask.to(torch.int64),
+                truncated_mask.to(torch.int64),
+                verify_lens.to(torch.int64),
+            ],
+            dim=1,
+        )
+
+    def _settle_and_write(self, bundle: dict[str, Any]) -> None:
+        batch = _SettleBatch.from_bundle(bundle)
+        self._last_forward_ct = batch.forward_ct
+        for b in range(len(batch.rids)):
+            self._settle_row(b=b, batch=batch)
+        self._finish_step(forward_ct=batch.forward_ct)
+        self._apply_all_finish_intents()
+
+    def _settle_row(self, *, b: int, batch: _SettleBatch) -> None:
+        forward_ct = batch.forward_ct
+        rid = batch.rids[b]
+        state = self._states.setdefault(rid, _RequestState())
+        state.last_seen_ct = forward_ct
+
+        cl, cap_trim, bonus_token, seq_len, is_greedy, is_truncated, verify_len = (
+            batch.row_meta[b]
+        )
+        window = verify_len - 1
+        assert 0 <= cl <= window <= self._gamma
+
+        self._drop_pending_on_discontinuity(
+            state, seq_len=seq_len, forward_ct=forward_ct
+        )
+        state.expected_seq_len = seq_len + cl + 1
+
+        if is_greedy or is_truncated:
+            if is_truncated and not is_greedy:
+                self._warn_once(
+                    reason="requests with top-k/top-p/min-p sampling are "
+                    "excluded per-row; the estimator only supports "
+                    "pure-temperature sampling (processed target distribution "
+                    "would differ from plain softmax(logits/T))"
+                )
+            state.pending = []
+            return
+
+        record: dict[str, Any] = {
+            "rid": rid,
+            "fct": forward_ct,
+            "w": window,
+            "cl": cl,
+            "ct": cap_trim,
+        }
+        num_old_pending = len(state.pending)
+        if cl == window and window < self._gamma:
+            self._open_block(
+                state,
+                record,
+                drafts_row=batch.drafts[b],
+                q_all_row=batch.q_all[b],
+                window=window,
+                seq_len=seq_len,
+                forward_ct=forward_ct,
+            )
+        else:
+            self._online.add(forward_ct=forward_ct, lo=cl + 1.0, hi=cl + 1.0)
+
+        pending_gathers = self._settle_pending(
+            b=b,
+            batch=batch,
+            state=state,
+            realized=batch.drafts[b][:cl] + [bonus_token],
+            cl=cl,
+            seq_len=seq_len,
+            num_old_pending=num_old_pending,
+        )
+        if pending_gathers:
+            record["pg"] = pending_gathers
+        if self._file is not None:
+            self._file.write(json.dumps(record) + "\n")
+
+    def _open_block(
+        self,
+        state: _RequestState,
+        record: dict[str, Any],
+        *,
+        drafts_row: List[int],
+        q_all_row: List[float],
+        window: int,
+        seq_len: int,
+        forward_ct: int,
+    ) -> None:
+        trimmed_tokens = drafts_row[window : self._gamma]
+        q_lps = q_all_row[window : self._gamma]
+        state.pending.append(
+            _PendingBlock(
+                forward_ct=forward_ct,
+                anchor_pos=seq_len - 1,
+                window=window,
+                trimmed_tokens=trimmed_tokens,
+                next_offset=window + 1,
+                q_lps=q_lps,
+            )
+        )
+        record["trimmed_tokens"] = trimmed_tokens
+        record["q_lp"] = q_lps
+
+    def _settle_pending(
+        self,
+        *,
+        b: int,
+        batch: _SettleBatch,
+        state: _RequestState,
+        realized: List[int],
+        cl: int,
+        seq_len: int,
+        num_old_pending: int,
+    ) -> List[list]:
+        gamma = self._gamma
+        pending_gathers: List[list] = []
+        kept_pending: List[_PendingBlock] = []
+        for block_idx, block in enumerate(state.pending):
+            diverged = False
+            while block.next_offset <= gamma:
+                row = block.anchor_pos + block.next_offset - seq_len
+                assert row >= 0
+                if row > cl:
+                    break
+                token = block.trimmed_tokens[block.next_offset - block.window - 1]
+                if block_idx < num_old_pending:
+                    p_lp = batch.pending_logprobs[
+                        batch.slot_lookup[(b, block_idx, block.next_offset)]
+                    ]
+                else:
+                    p_lp = batch.target_diag[b][row]
+                pending_gathers.append(
+                    [block.forward_ct, block.next_offset, p_lp, token, realized[row]]
+                )
+                self._accumulate_online(block, p_lp=p_lp)
+                block.next_offset += 1
+                if realized[row] != token:
+                    diverged = True
+                    break
+            if not diverged and block.next_offset <= gamma:
+                kept_pending.append(block)
+            else:
+                self._finalize_walk_online(
+                    block, diverged=diverged, forward_ct=batch.forward_ct
+                )
+        state.pending = kept_pending
+        return pending_gathers
+
+    def _accumulate_online(self, block: _PendingBlock, *, p_lp: float) -> None:
+        a = min(1.0, math.exp(p_lp - block.q_lps[block.next_offset - block.window - 1]))
+        block.est_prod *= a
+        block.est_lo_extra += block.est_prod
+
+    def _drop_pending_on_discontinuity(
+        self, state: _RequestState, *, seq_len: int, forward_ct: int
+    ) -> None:
+        if state.expected_seq_len < 0 or seq_len == state.expected_seq_len:
+            return
+        if not state.pending:
+            return
+        self._discontinuity_drop_ct += len(state.pending)
+        for block in state.pending:
+            self._finalize_at_end_online(block, forward_ct=forward_ct)
+        state.pending = []
+
+    def _finish_step(self, *, forward_ct: int) -> None:
+        self._observed_step_ct += 1
+        if self._file is not None:
+            self._steps_since_flush += 1
+            if self._steps_since_flush >= _FLUSH_EVERY_STEPS:
+                self._file.flush()
+                self._steps_since_flush = 0
+        if self._observed_step_ct % _STATE_SWEEP_INTERVAL == 0:
+            self._sweep_states(forward_ct=forward_ct)
+        self._online.maybe_log(forward_ct=forward_ct)
+
+    def _host_to_device_async(
+        self, values: List[int], *, device: torch.device
+    ) -> torch.Tensor:
+        host = torch.tensor(values, dtype=torch.long, pin_memory=device.type == "cuda")
+        self._retained_h2d.append(host)
+        return host.to(device=device, non_blocking=True)
+
+    def _gather_logprobs(
+        self,
+        *,
+        logits: torch.Tensor,
+        row_indices: torch.Tensor,
+        token_indices: torch.Tensor,
+        temps: torch.Tensor,
+    ) -> torch.Tensor:
+        if row_indices.numel() == 0:
+            return torch.zeros(0, dtype=torch.float32, device=logits.device)
+        per_row_temps = temps[row_indices].clamp_min(1e-5)
+        return gather_chunked_token_logprobs(
+            logits=logits,
+            row_indices=row_indices,
+            token_indices=token_indices,
+            per_row_temps=per_row_temps,
+            chunk_size=_GATHER_ROW_CHUNK,
+        )
+
+    def _sweep_states(self, *, forward_ct: int) -> None:
+        expired = [
+            rid
+            for rid, state in self._states.items()
+            if forward_ct - state.last_seen_ct > _STATE_EXPIRE_STEPS
+        ]
+        for rid in expired:
+            for block in self._states[rid].pending:
+                self._finalize_at_end_online(block, forward_ct=forward_ct)
+            del self._states[rid]
+            self._finish_intents.pop(rid, None)
+
+    def _skip_reason(
+        self,
+        *,
+        logits_adjustments_are_noop: bool,
+        corrected_logits: Optional[torch.Tensor],
+    ) -> Optional[str]:
+        return block_accept_skip_reason(
+            logits_adjustments_are_noop=logits_adjustments_are_noop,
+            corrected_logits=corrected_logits,
+        )
+
+    def _skip_step(self, *, reason: str) -> None:
+        self._skipped_step_ct += 1
+        self._warn_once(reason=SKIP_STEP_WARNING.format(reason))
+
+    def _warn_once(self, *, reason: str) -> None:
+        warn_once(self._warned_skip_reasons, reason=reason)
+
+
+def create_block_accept_estimate_recorder(
+    *, gamma: int, device: Union[str, torch.device], tp_rank: int
+) -> Optional[BlockAcceptEstimateRecorder]:
+    if tp_rank != 0:
+        return None
+
+    path = envs.SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH.get()
+    online_log_interval = envs.SGLANG_DSPARK_BLOCK_ACCEPT_ONLINE_INTERVAL.get()
+    if not path and online_log_interval <= 0:
+        return None
+
+    return BlockAcceptEstimateRecorder(
+        path=path,
+        gamma=gamma,
+        device=device,
+        online_log_interval=online_log_interval,
+    )

python/sglang/srt/speculative/dspark_components/dspark_config.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import msgspec
+
+from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DSPARK_GAMMA = 7
+SUPPORTED_DSPARK_MARKOV_HEAD_TYPES = ("vanilla", "gated", "rnn")
+
+# The dsv4 self-drafting checkpoint runs its draft attention on the dedicated
+# DeepSeek-V4 backend instead of the generic draft-backend fallback.
+DSV4_DRAFT_ATTENTION_BACKEND = "dsv4"
+
+
+def draft_is_deepseek_v4(*, server_args: ServerArgs) -> bool:
+    from sglang.srt.configs.model_config import is_deepseek_v4
+    from sglang.srt.utils.hf_transformers_utils import get_config
+
+    draft_model_path = server_args.speculative_draft_model_path
+    if not draft_model_path:
+        return False
+    draft_hf_config = get_config(
+        draft_model_path,
+        trust_remote_code=server_args.trust_remote_code,
+        revision=server_args.speculative_draft_model_revision,
+        model_override_args=json.loads(server_args.json_model_override_args),
+        model_config_parser=server_args.model_config_parser,
+    )
+    return draft_hf_config is not None and is_deepseek_v4(draft_hf_config)
+
+
+def dspark_gamma_from_num_draft_tokens(num_draft_tokens: int) -> int:
+    gamma = int(num_draft_tokens) - 1
+    if gamma < 1:
+        raise ValueError(
+            "DSpark speculative_num_draft_tokens must be >= 2 (= gamma + 1), "
+            f"got {num_draft_tokens}."
+        )
+    return gamma
+
+
+class DSparkDraftConfig(msgspec.Struct, frozen=True):
+    num_hidden_layers: Optional[int]
+    num_target_layers: Optional[int]
+    gamma: Optional[int]
+    target_layer_ids: Optional[List[int]]
+    mask_token: str
+    mask_token_id: Optional[int]
+    markov_rank: int
+    markov_head_type: Optional[str]
+
+    def resolve_gamma(self, *, default: Optional[int] = None) -> Optional[int]:
+        return self.gamma if self.gamma is not None else default
+
+    def require_markov(self) -> bool:
+        return int(self.markov_rank) > 0
+
+
+class DSparkRuntimeConfig(msgspec.Struct, frozen=True):
+    gamma: int
+    verify_num_draft_tokens: int
+    mask_token_id: int
+
+
+def resolve_runtime_config(
+    *,
+    draft_hf_config: Any,
+    speculative_num_draft_tokens: Optional[int],
+    target_vocab_size: int,
+) -> DSparkRuntimeConfig:
+    """Resolve and validate the worker-facing DSpark runtime knobs (gamma,
+    verify window, mask token) from the draft checkpoint config, with
+    server_args.speculative_num_draft_tokens taking precedence for gamma."""
+    draft_config = parse_dspark_draft_config(draft_hf_config=draft_hf_config)
+    if not draft_config.require_markov():
+        raise ValueError(
+            "DSpark draft requires markov_rank > 0; got "
+            f"markov_rank={draft_config.markov_rank}."
+        )
+
+    if speculative_num_draft_tokens is None:
+        gamma = int(draft_config.resolve_gamma(default=None) or 0)
+        if gamma < 1:
+            raise ValueError(
+                "DSpark could not resolve gamma from the draft config and "
+                "speculative_num_draft_tokens is unset."
+            )
+    else:
+        gamma = dspark_gamma_from_num_draft_tokens(int(speculative_num_draft_tokens))
+        config_gamma = draft_config.resolve_gamma(default=None)
+        if config_gamma is not None and int(config_gamma) != gamma:
+            logger.warning(
+                "DSpark gamma mismatch: using gamma=%s (from "
+                "speculative_num_draft_tokens=%s) but draft config block_size=%s.",
+                gamma,
+                speculative_num_draft_tokens,
+                config_gamma,
+            )
+
+    if draft_config.mask_token_id is None:
+        raise ValueError(
+            "DSpark requires mask_token_id to be set in the draft model config."
+        )
+    mask_token_id = int(draft_config.mask_token_id)
+    if mask_token_id >= target_vocab_size:
+        raise ValueError(
+            f"DSpark mask_token_id={mask_token_id} is outside the target "
+            f"vocab size {target_vocab_size}."
+        )
+
+    return DSparkRuntimeConfig(
+        gamma=gamma,
+        verify_num_draft_tokens=gamma + 1,
+        mask_token_id=mask_token_id,
+    )
+
+
+def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]:
+    """Load the draft checkpoint's hf config and read its DSpark gamma
+    (block_size). Raises on config-load failure; callers pick the fallback."""
+    from sglang.srt.utils.hf_transformers_utils import get_config
+
+    draft_hf_config = get_config(
+        server_args.speculative_draft_model_path,
+        trust_remote_code=server_args.trust_remote_code,
+        revision=server_args.speculative_draft_model_revision,
+        model_override_args=json.loads(server_args.json_model_override_args),
+    )
+    return parse_dspark_draft_config(draft_hf_config=draft_hf_config).resolve_gamma(
+        default=None
+    )
+
+
+def checkpoint_bundles_dspark_draft(hf_config: Any) -> bool:
+    """The checkpoint carries a bundled DSpark draft head, marked by the
+    prefixed dspark_* keys on the target hf config. Single source of truth
+    for the bundling convention (draft-path defaulting, draft-arch remap)."""
+    return any(
+        _cfg_get(hf_config, key, None) is not None
+        for key in (
+            "dspark_block_size",
+            "dspark_markov_rank",
+            "dspark_noise_token_id",
+            "dspark_target_layer_ids",
+        )
+    )
+
+
+def _cfg_get(config: Any, key: str, default: Any = None) -> Any:
+    if isinstance(config, dict):
+        return config.get(key, default)
+    return getattr(config, key, default)
+
+
+def _get_text_config(config: Any) -> Any:
+    if config is None:
+        return None
+    if isinstance(config, dict):
+        return config.get("text_config", config)
+    text_config = getattr(config, "text_config", None)
+    if text_config is not None:
+        return text_config
+    return config
+
+
+def _get_dspark_config(config: Any) -> dict:
+    cfg = _cfg_get(config, "dspark_config", None)
+    if cfg is None:
+        return {}
+    if isinstance(cfg, dict):
+        return cfg
+    try:
+        return dict(cfg)
+    except Exception:
+        return {}
+
+
+def parse_dspark_draft_config(*, draft_hf_config: Any) -> DSparkDraftConfig:
+    base = parse_dflash_draft_config(draft_hf_config=draft_hf_config)
+
+    dspark_cfg = _get_dspark_config(draft_hf_config)
+    text_config = _get_text_config(draft_hf_config)
+
+    prefixed_block_size = _cfg_get(draft_hf_config, "dspark_block_size", None)
+    prefixed_markov_rank = _cfg_get(draft_hf_config, "dspark_markov_rank", None)
+    prefixed_markov_head_type = _cfg_get(
+        draft_hf_config, "dspark_markov_head_type", None
+    )
+    prefixed_noise_token_id = _cfg_get(draft_hf_config, "dspark_noise_token_id", None)
+    prefixed_target_layer_ids = _cfg_get(
+        draft_hf_config, "dspark_target_layer_ids", None
+    )
+    uses_prefixed = any(
+        value is not None
+        for value in (
+            prefixed_block_size,
+            prefixed_markov_rank,
+            prefixed_noise_token_id,
+            prefixed_target_layer_ids,
+        )
+    )
+
+    raw_markov_rank = (
+        prefixed_markov_rank
+        if prefixed_markov_rank is not None
+        else dspark_cfg.get(
+            "markov_rank",
+            _cfg_get(
+                text_config, "markov_rank", _cfg_get(draft_hf_config, "markov_rank", 0)
+            ),
+        )
+    )
+    markov_rank = int(raw_markov_rank) if raw_markov_rank is not None else 0
+    if markov_rank < 0:
+        raise ValueError(f"DSpark markov_rank must be >= 0, got {markov_rank}.")
+
+    markov_head_type = (
+        prefixed_markov_head_type
+        if prefixed_markov_head_type is not None
+        else dspark_cfg.get(
+            "markov_head_type",
+            _cfg_get(
+                text_config,
+                "markov_head_type",
+                _cfg_get(draft_hf_config, "markov_head_type", None),
+            ),
+        )
+    )
+    if markov_rank > 0 and markov_head_type is None and not uses_prefixed:
+        raise ValueError(
+            "DSpark requires markov_head_type when markov_rank > 0, got None."
+        )
+    if markov_head_type is not None:
+        markov_head_type = str(markov_head_type).lower()
+        if markov_head_type not in SUPPORTED_DSPARK_MARKOV_HEAD_TYPES:
+            raise ValueError(
+                f"Unsupported DSpark markov_head_type={markov_head_type!r}. "
+                f"Supported: {SUPPORTED_DSPARK_MARKOV_HEAD_TYPES}."
+            )
+
+    raw_mask_token_id = (
+        prefixed_noise_token_id
+        if prefixed_noise_token_id is not None
+        else dspark_cfg.get(
+            "mask_token_id",
+            _cfg_get(
+                text_config,
+                "mask_token_id",
+                _cfg_get(draft_hf_config, "mask_token_id", base.mask_token_id),
+            ),
+        )
+    )
+    mask_token_id = int(raw_mask_token_id) if raw_mask_token_id is not None else None
+    if mask_token_id is not None and mask_token_id < 0:
+        raise ValueError(
+            f"DSpark mask_token_id must be non-negative, got {mask_token_id}."
+        )
+
+    gamma = (
+        int(prefixed_block_size) if prefixed_block_size is not None else base.block_size
+    )
+
+    if prefixed_target_layer_ids is not None:
+        if not isinstance(prefixed_target_layer_ids, (list, tuple)) or not len(
+            prefixed_target_layer_ids
+        ):
+            raise ValueError(
+                "DSpark dspark_target_layer_ids must be a non-empty list of ints, "
+                f"got {prefixed_target_layer_ids!r}."
+            )
+        target_layer_ids: Optional[List[int]] = [
+            int(x) for x in prefixed_target_layer_ids
+        ]
+    else:
+        target_layer_ids = base.target_layer_ids
+
+    return DSparkDraftConfig(
+        num_hidden_layers=base.num_hidden_layers,
+        num_target_layers=base.num_target_layers,
+        gamma=gamma,
+        target_layer_ids=target_layer_ids,
+        mask_token=base.mask_token,
+        mask_token_id=mask_token_id,
+        markov_rank=markov_rank,
+        markov_head_type=markov_head_type,
+    )

python/sglang/srt/speculative/dspark_components/dspark_draft.py

@@ -0,0 +1,421 @@
+from __future__ import annotations
+
+import logging
+from contextlib import nullcontext
+from typing import Optional
+
+import msgspec
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+from sglang.srt.speculative.draft_worker_common import make_draft_input_v2
+from sglang.srt.speculative.dspark_components.dspark_planner import VerifyWindow
+from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import (
+    SampleStepTokens,
+)
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import draft_tp_context
+
+logger = logging.getLogger(__name__)
+
+
+class DraftBlockResult(msgspec.Struct, frozen=True):
+    draft_tokens: torch.Tensor
+    corrected_logits: Optional[torch.Tensor]
+    greedy_mask: torch.Tensor
+    temperatures: torch.Tensor
+
+
+class DraftForwardResult(msgspec.Struct, frozen=True):
+    draft_block_ids: torch.Tensor
+    raw_hidden: torch.Tensor
+    draft_hidden_3d: torch.Tensor
+    can_run_graph: bool
+
+
+class DraftProposal(msgspec.Struct, frozen=True):
+    draft_block_ids: torch.Tensor
+    draft_block: DraftBlockResult
+    draft_hidden: Optional[torch.Tensor]
+    confidence: Optional[torch.Tensor] = None
+    confidence_tap: Optional[torch.Tensor] = None
+    folded: bool = False
+
+
+def greedy_step_sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
+    del step_idx
+    return torch.argmax(step_logits, dim=-1)
+
+
+class DsparkDraftSampler:
+
+    def __init__(self, *, model, gamma, max_bs, device, confidence_fn=None, out=None):
+        self.model = model
+        self.markov_head = model.markov_head
+        self.gamma = int(gamma)
+        if out is not None:
+            assert out.shape == (int(max_bs) * self.gamma,) and out.dtype == torch.int64
+            self.out = out
+        else:
+            self.out = torch.empty(
+                (int(max_bs) * self.gamma,), dtype=torch.int64, device=device
+            )
+        self.confidence_fn = confidence_fn
+        self.confidence_out = (
+            torch.empty((int(max_bs), self.gamma), dtype=torch.float32, device=device)
+            if confidence_fn is not None
+            else None
+        )
+
+    def __call__(self, hidden_states, input_ids):
+        bs = hidden_states.shape[0] // self.gamma
+        base_logits, confidence_tap = self.model.compute_base_logits(hidden_states)
+        base_logits = base_logits.view(bs, self.gamma, -1)
+        anchor = input_ids.view(bs, self.gamma)[:, 0]
+        draft_tokens, _ = self.markov_head.sample_block(
+            base_logits,
+            first_prev_tokens=anchor,
+            hidden_states=hidden_states.view(bs, self.gamma, -1),
+            sampler=greedy_step_sampler,
+        )
+        self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1))
+        if self.confidence_out is not None:
+            confidence = self.confidence_fn(
+                draft_hidden=hidden_states.view(bs, self.gamma, -1),
+                anchor_tokens=anchor,
+                draft_tokens=draft_tokens,
+                confidence_tap=confidence_tap,
+            )
+            self.confidence_out[:bs].copy_(confidence)
+
+
+def maybe_build_draft_sampler(
+    *,
+    draft_model,
+    gamma: int,
+    max_bs: int,
+    device,
+    tp_rank: int,
+    confidence_fn=None,
+    out=None,
+) -> Optional[DsparkDraftSampler]:
+    """Build the graph-folded greedy draft sampler, or return None (with the
+    reason logged) when the draft model cannot support folding and the
+    proposal must stay eager."""
+
+    def _eager(reason):
+        if tp_rank == 0:
+            logger.info("DSpark draft greedy proposal kept eager (reason=%s).", reason)
+        return None
+
+    if gamma <= 0:
+        return _eager("gamma<=0")
+    if not hasattr(draft_model, "compute_base_logits"):
+        return _eager("no compute_base_logits")
+    if getattr(draft_model, "markov_head", None) is None:
+        return _eager("no markov head")
+    if tp_rank == 0:
+        logger.info("DSpark draft greedy proposal folded into the draft cuda graph.")
+    return DsparkDraftSampler(
+        model=draft_model,
+        gamma=gamma,
+        max_bs=max_bs,
+        device=device,
+        confidence_fn=confidence_fn,
+        out=out,
+    )
+
+
+def make_next_draft_input(
+    *,
+    bonus_tokens: torch.Tensor,
+    new_seq_lens: torch.Tensor,
+) -> DFlashDraftInputV2:
+    return make_draft_input_v2(bonus_tokens=bonus_tokens, new_seq_lens=new_seq_lens)
+
+
+def resolve_greedy_mask(
+    *,
+    bs: int,
+    sampling_info,
+    device: torch.device,
+) -> torch.Tensor:
+    if sampling_info is None:
+        return torch.ones(bs, dtype=torch.bool, device=device)
+    return (sampling_info.top_ks <= 1).view(-1)
+
+
+def sample_draft_block(
+    *,
+    base_logits: torch.Tensor,
+    anchor_tokens: torch.Tensor,
+    draft_hidden: torch.Tensor,
+    sampling_info,
+    markov_head,
+    device: torch.device,
+) -> DraftBlockResult:
+    bs = base_logits.shape[0]
+    greedy_mask = resolve_greedy_mask(bs=bs, sampling_info=sampling_info, device=device)
+    any_sampling = sampling_info is not None and not sampling_info.is_all_greedy
+    fast_sampling = envs.SGLANG_DSPARK_FAST_SAMPLING.get()
+
+    if sampling_info is None:
+        temperatures = torch.ones(bs, dtype=torch.float32, device=device)
+    else:
+        temperatures = (
+            sampling_info.temperatures.view(-1).to(torch.float32).clamp_min(1e-5)
+        )
+
+    if not any_sampling:
+
+        def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
+            return torch.argmax(step_logits, dim=-1)
+
+    else:
+
+        def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
+            if fast_sampling:
+                exp_noise = torch.empty(
+                    step_logits.shape, dtype=torch.float32, device=step_logits.device
+                ).exponential_(1)
+                return SampleStepTokens.execute(
+                    step_logits=step_logits,
+                    temperatures=temperatures,
+                    greedy_mask=greedy_mask,
+                    exp_noise=exp_noise,
+                )
+            else:
+                probs = torch.softmax(
+                    step_logits.float() / temperatures[:, None], dim=-1
+                )
+                argmax_tokens = torch.argmax(step_logits, dim=-1)
+                sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
+                return torch.where(greedy_mask, argmax_tokens, sampled_tokens)
+
+    draft_tokens, corrected_logits = markov_head.sample_block(
+        base_logits,
+        first_prev_tokens=anchor_tokens,
+        hidden_states=draft_hidden,
+        sampler=sampler,
+    )
+    return DraftBlockResult(
+        draft_tokens=draft_tokens,
+        corrected_logits=corrected_logits,
+        greedy_mask=greedy_mask,
+        temperatures=temperatures,
+    )
+
+
+class DraftBlockProposer:
+    def __init__(
+        self,
+        *,
+        draft_model,
+        draft_model_runner,
+        gamma: int,
+        mask_token_id: int,
+        draft_block_spec_info,
+        dp_moe_sync: bool = False,
+    ) -> None:
+        self.draft_model = draft_model
+        self.draft_model_runner = draft_model_runner
+        self.gamma = gamma
+        self._mask_token_id = mask_token_id
+        self._draft_block_spec_info = draft_block_spec_info
+        self._draft_sampler = None
+        self._dp_moe_sync = dp_moe_sync
+
+    def attach_draft_sampler(self, draft_sampler) -> None:
+        self._draft_sampler = draft_sampler
+
+    def _base_logits_context(self):
+        if self._dp_moe_sync:
+            return draft_tp_context(get_parallel().attn_tp_group)
+        return nullcontext()
+
+    def propose(
+        self,
+        *,
+        batch: ScheduleBatch,
+        draft_input: DFlashDraftInputV2,
+        verify_window: VerifyWindow,
+        bs: int,
+        device: str,
+        target_model,
+        sampling_info,
+    ) -> DraftProposal:
+        embed_module = target_model.get_input_embeddings()
+        fwd = self._run_forward(
+            batch=batch,
+            draft_input=draft_input,
+            verify_window=verify_window,
+            bs=bs,
+            device=device,
+            embed_module=embed_module,
+        )
+        draft_block_ids = fwd.draft_block_ids
+
+        draft_sampler = self._draft_sampler
+        all_greedy = sampling_info is None or sampling_info.is_all_greedy
+        folded_confidence = None
+        confidence_tap = None
+        folded = False
+        if draft_sampler is not None and fwd.can_run_graph and all_greedy:
+            folded = True
+            if sampling_info is None:
+                temperatures = torch.ones(bs, dtype=torch.float32, device=device)
+            else:
+                temperatures = (
+                    sampling_info.temperatures.view(-1)
+                    .to(torch.float32)
+                    .clamp_min(1e-5)
+                )
+            draft_block = DraftBlockResult(
+                draft_tokens=draft_sampler.out[: bs * self.gamma].view(bs, self.gamma),
+                corrected_logits=None,
+                greedy_mask=resolve_greedy_mask(
+                    bs=bs, sampling_info=sampling_info, device=device
+                ),
+                temperatures=temperatures,
+            )
+            if draft_sampler.confidence_out is not None:
+                folded_confidence = draft_sampler.confidence_out[:bs]
+        else:
+            with self._base_logits_context():
+                base_logits, confidence_tap = self.draft_model.compute_base_logits(
+                    fwd.raw_hidden
+                )
+                base_logits = base_logits.view(bs, self.gamma, -1)
+            draft_block = sample_draft_block(
+                base_logits=base_logits,
+                anchor_tokens=draft_block_ids[:, 0],
+                draft_hidden=fwd.draft_hidden_3d,
+                sampling_info=sampling_info,
+                markov_head=self.draft_model.markov_head,
+                device=device,
+            )
+        return DraftProposal(
+            draft_block_ids=draft_block_ids,
+            draft_block=draft_block,
+            draft_hidden=fwd.draft_hidden_3d,
+            confidence=folded_confidence,
+            confidence_tap=confidence_tap,
+            folded=folded,
+        )
+
+    def run_idle_participation(self, batch: ScheduleBatch) -> None:
+        if not self._dp_moe_sync or batch.global_num_tokens is None:
+            return
+        device = self.draft_model_runner.device
+        empty_long = torch.empty((0,), dtype=torch.int64, device=device)
+        idle_batch = ForwardBatch(
+            forward_mode=ForwardMode.IDLE,
+            batch_size=0,
+            input_ids=empty_long,
+            req_pool_indices=empty_long,
+            seq_lens=empty_long,
+            out_cache_loc=empty_long,
+            seq_lens_sum=0,
+            seq_lens_cpu=torch.empty((0,), dtype=torch.int64),
+            positions=empty_long,
+            spec_algorithm=SpeculativeAlgorithm.DSPARK,
+            spec_info=self._draft_block_spec_info,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+        )
+        self._fill_dp_moe_sync_metadata(idle_batch, batch)
+        with torch.inference_mode():
+            self.draft_model_runner.forward(idle_batch)
+
+    def _run_forward(
+        self,
+        *,
+        batch: ScheduleBatch,
+        draft_input: DFlashDraftInputV2,
+        verify_window: VerifyWindow,
+        bs: int,
+        device: str,
+        embed_module,
+    ) -> DraftForwardResult:
+        gamma = self.gamma
+        prefix_lens = batch.seq_lens
+        positions_2d = verify_window.positions_2d
+        verify_cache_loc_2d = verify_window.verify_cache_loc_2d
+
+        draft_block_ids = torch.full(
+            (bs, gamma), int(self._mask_token_id), dtype=torch.long, device=device
+        )
+        draft_block_ids[:, 0].copy_(draft_input.bonus_tokens.view(-1))
+        draft_positions = positions_2d[:, :gamma].reshape(-1)
+        draft_cache_loc = verify_cache_loc_2d[:, :gamma].reshape(-1)
+
+        draft_owns_embed = hasattr(self.draft_model, "forward_embed")
+        draft_input_embeds: Optional[torch.Tensor] = None
+        if not draft_owns_embed:
+            noise_embedding = embed_module(draft_block_ids)
+            draft_input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1])
+
+        if batch.seq_lens_cpu is not None:
+            draft_seq_lens_cpu = batch.seq_lens_cpu + gamma
+            draft_seq_lens_sum = int(draft_seq_lens_cpu.sum())
+        elif draft_input.reserved_seq_lens_cpu is not None:
+            draft_seq_lens_cpu = draft_input.reserved_seq_lens_cpu
+            draft_seq_lens_sum = int(draft_input.reserved_seq_lens_sum)
+        else:
+            raise RuntimeError("DSpark decode expected batch.seq_lens_cpu, got None")
+
+        draft_forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            batch_size=bs,
+            input_ids=draft_block_ids.flatten(),
+            req_pool_indices=batch.req_pool_indices,
+            seq_lens=prefix_lens,
+            out_cache_loc=draft_cache_loc,
+            seq_lens_sum=draft_seq_lens_sum,
+            seq_lens_cpu=draft_seq_lens_cpu,
+            positions=draft_positions,
+            input_embeds=draft_input_embeds,
+            spec_algorithm=SpeculativeAlgorithm.DSPARK,
+            spec_info=self._draft_block_spec_info,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+        )
+        self._fill_dp_moe_sync_metadata(draft_forward_batch, batch)
+        with torch.inference_mode():
+            draft_out = self.draft_model_runner.forward(draft_forward_batch)
+        logits_output = draft_out.logits_output
+        raw_hidden = logits_output.hidden_states
+        if raw_hidden is None:
+            raise RuntimeError("DSpark draft model returned no hidden states.")
+        draft_hidden_3d = raw_hidden.view(bs, gamma, -1)
+        return DraftForwardResult(
+            draft_block_ids=draft_block_ids,
+            raw_hidden=raw_hidden,
+            draft_hidden_3d=draft_hidden_3d,
+            can_run_graph=draft_out.can_run_graph,
+        )
+
+    def _fill_dp_moe_sync_metadata(
+        self, forward_batch: ForwardBatch, batch: ScheduleBatch
+    ) -> None:
+        if not self._dp_moe_sync or batch.global_num_tokens is None:
+            return
+        gnt, gnt_logprob = (
+            self._draft_block_spec_info.get_spec_adjusted_global_num_tokens(batch)
+        )
+        device = self.draft_model_runner.device
+        forward_batch.global_num_tokens_cpu = gnt
+        forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob
+        forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to(
+            device, non_blocking=True
+        )
+        forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor(
+            gnt_logprob, dtype=torch.int64
+        ).to(device, non_blocking=True)
+        forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph

python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py

@@ -0,0 +1,157 @@
+from typing import Optional
+
+import torch
+
+from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import (
+    BuildCommitInjectLayout,
+)
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+class TargetHiddenKvInjector:
+    def __init__(
+        self,
+        *,
+        draft_model,
+        draft_model_runner,
+        model_runner,
+        device,
+        verify_num_draft_tokens: int,
+        block_pos_offsets: torch.Tensor,
+    ) -> None:
+        self.draft_model = draft_model
+        self.draft_model_runner = draft_model_runner
+        self.model_runner = model_runner
+        self.device = device
+        self.verify_num_draft_tokens = verify_num_draft_tokens
+        self._block_pos_offsets = block_pos_offsets
+
+    def inject_target_hidden(
+        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:
+        if target_hidden is None or target_hidden.numel() == 0:
+            return
+        device = self.model_runner.device
+        cache_loc = cache_loc.to(device=device, dtype=torch.int64, non_blocking=True)
+        positions = positions.to(device=device, dtype=torch.int64, non_blocking=True)
+        target_hidden = target_hidden.to(device=device, non_blocking=True)
+        n_real = positions.shape[0]
+        if target_hidden.shape[0] > n_real:
+            target_hidden = target_hidden[:n_real]
+        if cache_loc_2d is not None:
+            cache_loc_2d = cache_loc_2d.to(
+                device=device, dtype=torch.int64, non_blocking=True
+            )
+        if commit_lens is not None:
+            commit_lens = commit_lens.to(
+                device=device, dtype=torch.int32, non_blocking=True
+            )
+
+        pool = self.draft_model_runner.token_to_kv_pool
+        if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"):
+            self._inject_mla(
+                pool=pool,
+                target_hidden=target_hidden,
+                cache_loc=cache_loc,
+                positions=positions,
+                cache_loc_2d=cache_loc_2d,
+                commit_lens=commit_lens,
+            )
+            return
+
+        with torch.inference_mode():
+            self.draft_model.write_target_hidden_kv(
+                target_hidden=target_hidden,
+                pool=pool,
+                positions=positions,
+                cache_loc=cache_loc,
+                cache_loc_2d=cache_loc_2d,
+                commit_lens=commit_lens,
+            )
+
+    def _inject_mla(
+        self,
+        *,
+        pool,
+        target_hidden: torch.Tensor,
+        cache_loc: torch.Tensor,
+        positions: torch.Tensor,
+        cache_loc_2d: Optional[torch.Tensor],
+        commit_lens: Optional[torch.Tensor],
+    ) -> None:
+        swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32)
+        if commit_lens is not None and cache_loc_2d is not None:
+            bs, verify_len = cache_loc_2d.shape
+            col = torch.arange(verify_len, device=cache_loc.device).view(1, -1)
+            committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
+            swa_loc = torch.where(committed_mask, swa_loc, torch.full_like(swa_loc, -1))
+
+        with torch.inference_mode():
+            self.draft_model.write_target_hidden_kv(
+                main_hidden=target_hidden,
+                swa_loc=swa_loc,
+                positions=positions,
+                pool=pool,
+            )
+
+    def inject_ragged(
+        self,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        hidden_strided: torch.Tensor,
+        commit_lens: torch.Tensor,
+        bs: int,
+    ) -> None:
+        stride = self.verify_num_draft_tokens
+        prefix_lens = batch.seq_lens
+        hidden = hidden_strided.view(bs, stride, -1)
+
+        pool = self.draft_model_runner.token_to_kv_pool
+        if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"):
+            if hidden_strided.numel() == 0:
+                return
+            inject_layout = BuildCommitInjectLayout.execute(
+                req_pool_indices=batch.req_pool_indices,
+                req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                prefix_lens=prefix_lens,
+                block_pos_offsets=self._block_pos_offsets[:stride],
+                full_to_swa_mapping=pool.full_to_swa_index_mapping,
+                commit_lens=commit_lens,
+                stride=stride,
+            )
+            with torch.inference_mode():
+                self.draft_model.write_target_hidden_kv(
+                    main_hidden=hidden.reshape(-1, hidden.shape[-1]),
+                    swa_loc=inject_layout.swa_loc,
+                    positions=inject_layout.positions,
+                    pool=pool,
+                )
+            return
+
+        positions_2d = prefix_lens.unsqueeze(1) + self._block_pos_offsets
+        verify_cache_loc = assign_extend_cache_locs_func(
+            req_pool_indices=batch.req_pool_indices,
+            req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+            start_offset=prefix_lens,
+            end_offset=prefix_lens + stride,
+            batch_size=bs,
+            draft_token_num=stride,
+            device=self.device,
+        )
+        verify_cache_loc_2d = verify_cache_loc.view(bs, stride)
+        self.inject_target_hidden(
+            target_hidden=hidden.reshape(-1, hidden.shape[-1]),
+            cache_loc=verify_cache_loc,
+            cache_loc_2d=verify_cache_loc_2d,
+            positions=positions_2d.reshape(-1),
+            commit_lens=commit_lens,
+        )

python/sglang/srt/speculative/dspark_components/dspark_observability.py

@@ -0,0 +1,961 @@
+from __future__ import annotations
+
+import logging
+import math
+import statistics
+import time
+from collections import deque
+from contextlib import contextmanager, nullcontext
+from enum import Enum
+from typing import Callable, ContextManager, Iterator, Optional, Union
+
+import msgspec
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.sampling.sampling_params import TOP_K_ALL
+from sglang.srt.speculative.dflash_utils import compute_dflash_correct_drafts_and_bonus
+from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import (
+    create_block_accept_estimate_recorder,
+)
+from sglang.srt.speculative.dspark_components.dspark_sts import StsDataRecorder
+from sglang.srt.speculative.dspark_components.dspark_verify import (
+    verify_logits_adjustments_are_noop,
+)
+
+logger = logging.getLogger(__name__)
+
+_NULL_SEGMENT = nullcontext()
+
+ALL_COMPONENTS_TOKEN = "all"
+
+
+class InfoComponent(str, Enum):
+    CORE = "core"
+    STEP_CPU_TIME = "step_cpu_time"
+    STEP_GPU_TIME = "step_gpu_time"
+    DRAFT_GPU_TIME = "draft_gpu_time"
+    TARGET_VERIFY_GPU_TIME = "target_verify_gpu_time"
+    REQS = "reqs"
+
+
+class InfoSegment(str, Enum):
+    STEP = "step"
+    DRAFT = "draft"
+    TARGET_VERIFY = "target_verify"
+
+
+INFO_DUMP_MAX_RECORDS = 200_000
+INFO_DUMP_MAX_STEP_CPU_SECONDS = 1.0
+
+
+def resolve_enabled_components() -> set[InfoComponent]:
+    """Components enabled via env: SGLANG_DSPARK_DEBUG_DUMP tokens, plus the
+    published SPS-profiling switch SGLANG_DSPARK_ENABLE_SPS_RECORD=1, which is
+    an alias for the core,step_cpu_time components the SPS table fit needs."""
+    components = resolve_components(envs.SGLANG_DSPARK_DEBUG_DUMP.get())
+    if envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.get():
+        components |= {InfoComponent.CORE, InfoComponent.STEP_CPU_TIME}
+    return components
+
+
+def resolve_components(raw: tuple[str, ...]) -> set[InfoComponent]:
+    tokens = {token.strip() for token in raw if token.strip()}
+    if not tokens:
+        return set()
+    if ALL_COMPONENTS_TOKEN in tokens:
+        return set(InfoComponent)
+    try:
+        return {InfoComponent(token) for token in tokens}
+    except ValueError as exc:
+        valid = [component.value for component in InfoComponent]
+        raise ValueError(
+            f"Invalid SGLANG_DSPARK_DEBUG_DUMP token in {sorted(tokens)}; "
+            f"valid: {valid} or '{ALL_COMPONENTS_TOKEN}'."
+        ) from exc
+
+
+class ReqDetail(msgspec.Struct, omit_defaults=True):
+    req_pool_index: int
+    prefix_len: int
+    verify_len: int
+    acc_len: int
+    correct_drafts: int
+    cap_trim: int
+    bonus_token: int
+    draft_tokens: list[int]
+    rid: Optional[str] = None
+    confidence: Optional[list[float]] = None
+    survival: Optional[list[float]] = None
+
+
+class DecodeStepRecord(msgspec.Struct, omit_defaults=True):
+    forward_ct: int
+    bs: int = -1
+    mode: str = ""
+    budget: Optional[int] = None
+    lag_steps: Optional[int] = None
+    num_running_reqs: int = -1
+    num_verify_tokens: int = -1
+    verify_tokens_local: int = -1
+    verify_tokens_dp_synced: int = -1
+    verify_tokens_graph_key: int = -1
+    predicted_step_ms: Optional[float] = None
+    predicted_theta: Optional[float] = None
+    step_cpu_ms: Optional[float] = None
+    step_gpu_ms: Optional[float] = None
+    draft_gpu_ms: Optional[float] = None
+    target_verify_gpu_ms: Optional[float] = None
+    reqs: Optional[list[ReqDetail]] = None
+
+
+class DecodeStepObservation(msgspec.Struct):
+    forward_ct: int
+    bs: int
+    mode: str
+    budget: Optional[int]
+    lag_steps: Optional[int]
+    num_verify_tokens: int
+    verify_tokens_local: int
+    verify_tokens_dp_synced: int
+    verify_tokens_graph_key: int
+    predicted_step_ms: Optional[float]
+    predicted_theta: Optional[float]
+    verify_lens: Optional[torch.Tensor]
+    confidence: Optional[torch.Tensor]
+    req_pool_indices: torch.Tensor
+    prefix_lens: torch.Tensor
+    draft_tokens: torch.Tensor
+    bonus_tokens: torch.Tensor
+    correct_len: torch.Tensor
+    cap_trim_lens: torch.Tensor
+    commit_lens: torch.Tensor
+    rids: Optional[list[str]]
+
+
+class _PendingStep(msgspec.Struct):
+    forward_ct: int
+    bs: int
+    mode: str
+    budget: Optional[int]
+    lag_steps: Optional[int]
+    num_verify_tokens: int
+    verify_tokens_local: int
+    verify_tokens_dp_synced: int
+    verify_tokens_graph_key: int
+    predicted_step_ms: Optional[float]
+    predicted_theta: Optional[float]
+    step_cpu_ms: Optional[float]
+    rids: Optional[list[str]]
+    future: Optional[FutureTensors]
+    segment_events: dict[InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event]]
+
+
+class DsparkInfoDumper:
+    def __init__(
+        self,
+        *,
+        components: set[Union[InfoComponent, str]],
+        gamma: int,
+        verify_num_draft_tokens: int,
+        attn_tp_rank: int,
+        device: torch.device,
+        mode_value: str,
+        sps_report_interval: int = 0,
+        max_records: int = INFO_DUMP_MAX_RECORDS,
+        max_step_cpu_seconds: float = INFO_DUMP_MAX_STEP_CPU_SECONDS,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self.gamma = int(gamma)
+        self.verify_num_draft_tokens = int(verify_num_draft_tokens)
+        self.attn_tp_rank = int(attn_tp_rank)
+        self.device = device
+        self.mode_value = mode_value
+        self._clock = clock
+        self._max_step_cpu_seconds = max_step_cpu_seconds
+
+        self._components: set[InfoComponent] = {
+            InfoComponent(component) for component in components
+        }
+        self._sps_report_interval = int(sps_report_interval)
+        if self._sps_report_interval > 0:
+            self._components.add(InfoComponent.STEP_GPU_TIME)
+        # Dedup within an attention-TP group only: records describe the
+        # DP-rank-local batch, so under dp-attention every DP rank must keep
+        # dumping (the SPS profiler reads one payload per DP rank).
+        self.enabled = bool(self._components) and self.attn_tp_rank == 0
+        self._sps_window: list[tuple[float, float]] = []
+        self._sps_mismatched = 0
+
+        self._records: deque[DecodeStepRecord] = deque(maxlen=max_records)
+        self._pending: Optional[_PendingStep] = None
+        self._prev_stamp: Optional[float] = None
+
+        self._d2h_stream: Optional[torch.cuda.Stream] = None
+        if self.enabled and InfoComponent.REQS in self._components:
+            self._d2h_stream = torch.cuda.Stream(device=device)
+
+        self._current_segments: dict[
+            InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event]
+        ] = {}
+        self._open_segments: dict[InfoSegment, torch.cuda.Event] = {}
+
+    def begin_step(self) -> None:
+        if not self.enabled:
+            return
+        self._current_segments = {}
+        self._open_segments = {}
+        if InfoComponent.STEP_GPU_TIME in self._components:
+            self._open_segment(InfoSegment.STEP)
+
+    def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]:
+        if not self.enabled:
+            return _NULL_SEGMENT
+        segment = InfoSegment(name)
+        if not self._segment_enabled(segment):
+            return _NULL_SEGMENT
+        return self._active_segment(segment)
+
+    @contextmanager
+    def _active_segment(self, segment: InfoSegment) -> Iterator[None]:
+        self._open_segment(segment)
+        try:
+            yield
+        finally:
+            self._close_segment(segment)
+
+    def observe_decode_step(self, obs: DecodeStepObservation) -> None:
+        if not self.enabled:
+            return
+        if InfoComponent.STEP_GPU_TIME in self._components:
+            self._close_segment(InfoSegment.STEP)
+
+        now = self._clock()
+        step_cpu_ms = self._step_cpu_ms(now=now)
+        self._drain_pending()
+
+        future = (
+            self._stage_reqs(obs) if InfoComponent.REQS in self._components else None
+        )
+        self._pending = _PendingStep(
+            forward_ct=int(obs.forward_ct),
+            bs=int(obs.bs),
+            mode=obs.mode,
+            budget=None if obs.budget is None else int(obs.budget),
+            lag_steps=None if obs.lag_steps is None else int(obs.lag_steps),
+            num_verify_tokens=int(obs.num_verify_tokens),
+            verify_tokens_local=int(obs.verify_tokens_local),
+            verify_tokens_dp_synced=int(obs.verify_tokens_dp_synced),
+            verify_tokens_graph_key=int(obs.verify_tokens_graph_key),
+            predicted_step_ms=obs.predicted_step_ms,
+            predicted_theta=obs.predicted_theta,
+            step_cpu_ms=step_cpu_ms,
+            rids=obs.rids,
+            future=future,
+            segment_events=self._current_segments,
+        )
+        self._current_segments = {}
+        self._prev_stamp = now
+
+    def note_non_decode_step(self) -> None:
+        if not self.enabled:
+            return
+        self._drain_pending()
+        self._prev_stamp = None
+        self._current_segments = {}
+        self._open_segments = {}
+
+    def flush(self) -> None:
+        if not self.enabled:
+            return
+        self._drain_pending()
+
+    def clear(self) -> None:
+        self._records.clear()
+        self._pending = None
+        self._prev_stamp = None
+        self._current_segments = {}
+        self._open_segments = {}
+        self._sps_window = []
+        self._sps_mismatched = 0
+
+    def dump(self) -> Optional[dict]:
+        if not self.enabled:
+            return None
+        self.flush()
+        return {
+            "mode": self.mode_value,
+            "gamma": self.gamma,
+            "verify_num_draft_tokens": self.verify_num_draft_tokens,
+            "components": sorted(component.value for component in self._components),
+            "records": [msgspec.to_builtins(record) for record in self._records],
+        }
+
+    def _segment_enabled(self, segment: InfoSegment) -> bool:
+        if segment is InfoSegment.STEP:
+            return InfoComponent.STEP_GPU_TIME in self._components
+        if segment is InfoSegment.DRAFT:
+            return InfoComponent.DRAFT_GPU_TIME in self._components
+        if segment is InfoSegment.TARGET_VERIFY:
+            return InfoComponent.TARGET_VERIFY_GPU_TIME in self._components
+        return False
+
+    def _open_segment(self, segment: InfoSegment) -> None:
+        start = torch.cuda.Event(enable_timing=True)
+        start.record()
+        self._open_segments[segment] = start
+
+    def _close_segment(self, segment: InfoSegment) -> None:
+        start = self._open_segments.pop(segment, None)
+        if start is None:
+            return
+        end = torch.cuda.Event(enable_timing=True)
+        end.record()
+        self._current_segments[segment] = (start, end)
+
+    def _stage_reqs(self, obs: DecodeStepObservation) -> Optional[FutureTensors]:
+        tensors: dict[str, torch.Tensor] = {
+            "req_pool_indices": obs.req_pool_indices,
+            "prefix_lens": obs.prefix_lens,
+            "draft_tokens": obs.draft_tokens,
+            "bonus_tokens": obs.bonus_tokens,
+            "correct_len": obs.correct_len,
+            "cap_trim_lens": obs.cap_trim_lens,
+            "commit_lens": obs.commit_lens,
+        }
+        if obs.verify_lens is not None:
+            tensors["verify_lens"] = obs.verify_lens
+        if obs.confidence is not None:
+            tensors["confidence"] = obs.confidence
+        return FutureTensors.device_to_host(tensors, d2h_stream=self._d2h_stream)
+
+    def _drain_pending(self) -> None:
+        pending = self._pending
+        self._pending = None
+        if pending is None:
+            return
+
+        record = DecodeStepRecord(forward_ct=pending.forward_ct)
+        if InfoComponent.CORE in self._components:
+            record.bs = pending.bs
+            record.mode = pending.mode
+            record.budget = pending.budget
+            record.lag_steps = pending.lag_steps
+            record.num_running_reqs = pending.bs
+            record.num_verify_tokens = pending.num_verify_tokens
+            record.verify_tokens_local = pending.verify_tokens_local
+            record.verify_tokens_dp_synced = pending.verify_tokens_dp_synced
+            record.verify_tokens_graph_key = pending.verify_tokens_graph_key
+            record.predicted_step_ms = pending.predicted_step_ms
+            record.predicted_theta = pending.predicted_theta
+        if InfoComponent.STEP_CPU_TIME in self._components:
+            record.step_cpu_ms = pending.step_cpu_ms
+        if InfoComponent.STEP_GPU_TIME in self._components:
+            record.step_gpu_ms = self._segment_ms(pending, InfoSegment.STEP)
+        if InfoComponent.DRAFT_GPU_TIME in self._components:
+            record.draft_gpu_ms = self._segment_ms(pending, InfoSegment.DRAFT)
+        if InfoComponent.TARGET_VERIFY_GPU_TIME in self._components:
+            record.target_verify_gpu_ms = self._segment_ms(
+                pending, InfoSegment.TARGET_VERIFY
+            )
+        if InfoComponent.REQS in self._components and pending.future is not None:
+            record.reqs = self._build_reqs(
+                host=pending.future.wait(), bs=pending.bs, rids=pending.rids
+            )
+        elif pending.future is not None:
+            pending.future.wait()
+
+        self._records.append(record)
+        if self._sps_report_interval > 0:
+            self._report_sps_prediction(pending=pending, step_gpu_ms=record.step_gpu_ms)
+
+    def _report_sps_prediction(
+        self, *, pending: _PendingStep, step_gpu_ms: Optional[float]
+    ) -> None:
+        predicted = pending.predicted_step_ms
+        if predicted is None or step_gpu_ms is None:
+            return
+        matched = (
+            pending.budget is not None
+            and pending.bs + pending.budget == pending.num_verify_tokens
+        )
+        if not matched:
+            self._sps_mismatched += 1
+            return
+        self._sps_window.append((predicted, step_gpu_ms))
+        if len(self._sps_window) < self._sps_report_interval:
+            return
+
+        predictions = [p for p, _ in self._sps_window]
+        actuals = [a for _, a in self._sps_window]
+        abs_err = [abs(p - a) for p, a in self._sps_window]
+        rel_err = [abs(p - a) / a * 100 for p, a in self._sps_window if a > 0]
+        total = len(self._sps_window) + self._sps_mismatched
+        logger.info(
+            "DSpark SPS prediction: n=%d  mean predicted=%.3fms  mean actual=%.3fms  "
+            "MAE=%.3fms  median rel-err=%.1f%%  mean bias(pred-actual)=%+.3fms  "
+            "M_mismatch_rate=%.1f%% (%d/%d)",
+            len(self._sps_window),
+            statistics.fmean(predictions),
+            statistics.fmean(actuals),
+            statistics.fmean(abs_err),
+            statistics.median(rel_err) if rel_err else float("nan"),
+            statistics.fmean([p - a for p, a in self._sps_window]),
+            self._sps_mismatched / total * 100 if total else 0.0,
+            self._sps_mismatched,
+            total,
+        )
+        self._sps_window = []
+        self._sps_mismatched = 0
+
+    def _step_cpu_ms(self, *, now: float) -> Optional[float]:
+        prev = self._prev_stamp
+        if prev is None:
+            return None
+        step_cpu = now - prev
+        if not (0.0 < step_cpu <= self._max_step_cpu_seconds):
+            return None
+        return round(step_cpu * 1000.0, 4)
+
+    def _segment_ms(
+        self, pending: _PendingStep, segment: InfoSegment
+    ) -> Optional[float]:
+        events = pending.segment_events.get(segment)
+        if events is None:
+            return None
+        start, end = events
+        end.synchronize()
+        elapsed_ms = start.elapsed_time(end)
+        if elapsed_ms > self._max_step_cpu_seconds * 1000.0:
+            return None
+        return round(elapsed_ms, 4)
+
+    def _build_reqs(
+        self, *, host: dict, bs: int, rids: Optional[list[str]]
+    ) -> list[ReqDetail]:
+        req_ids = host["req_pool_indices"].tolist()
+        prefixes = host["prefix_lens"].tolist()
+        draft_rows = host["draft_tokens"].tolist()
+        bonus = host["bonus_tokens"].tolist()
+        correct = host["correct_len"].tolist()
+        cap_trim = host["cap_trim_lens"].tolist()
+        commit = host["commit_lens"].tolist()
+        verify_lens = host["verify_lens"].tolist() if "verify_lens" in host else None
+        if "confidence" in host:
+            conf_host = host["confidence"].float()
+            conf_rows = conf_host.tolist()
+            survival_rows = torch.cumprod(conf_host, dim=1).tolist()
+        else:
+            conf_rows = None
+            survival_rows = None
+
+        reqs: list[ReqDetail] = []
+        for row in range(bs):
+            verify_len = (
+                self.verify_num_draft_tokens
+                if verify_lens is None
+                else int(verify_lens[row])
+            )
+            reqs.append(
+                ReqDetail(
+                    rid=None if rids is None else rids[row],
+                    req_pool_index=int(req_ids[row]),
+                    prefix_len=int(prefixes[row]),
+                    verify_len=verify_len,
+                    acc_len=int(commit[row]),
+                    correct_drafts=int(correct[row]),
+                    cap_trim=int(cap_trim[row]),
+                    bonus_token=int(bonus[row]),
+                    draft_tokens=[int(t) for t in draft_rows[row]],
+                    confidence=(
+                        None
+                        if conf_rows is None
+                        else [round(float(p), 4) for p in conf_rows[row]]
+                    ),
+                    survival=(
+                        None
+                        if survival_rows is None
+                        else [round(float(p), 4) for p in survival_rows[row]]
+                    ),
+                )
+            )
+        return reqs
+
+
+EPS_PROB = 1e-8
+
+
+def _format_float(value: float, digits: int = 4) -> str:
+    value = float(value)
+    if math.isnan(value):
+        return "nan"
+    return f"{value:.{digits}f}"
+
+
+class PerPositionConfidenceMetrics:
+
+    def __init__(
+        self,
+        *,
+        gamma: int,
+        device: torch.device,
+        num_coarse_bins: int = 15,
+        num_fine_bins: int = 1024,
+    ) -> None:
+        self.gamma = int(gamma)
+        self.num_coarse_bins = int(num_coarse_bins)
+        self.num_fine_bins = int(num_fine_bins)
+        self.coarse_count = torch.zeros(
+            (self.gamma, self.num_coarse_bins), dtype=torch.float64, device=device
+        )
+        self.coarse_pred = torch.zeros_like(self.coarse_count)
+        self.coarse_target = torch.zeros_like(self.coarse_count)
+        self.fine_pos = torch.zeros(
+            (self.gamma, self.num_fine_bins), dtype=torch.float64, device=device
+        )
+        self.fine_neg = torch.zeros_like(self.fine_pos)
+        self.brier_num = torch.zeros(self.gamma, dtype=torch.float64, device=device)
+
+    def update(self, *, survival: torch.Tensor, prefix_mask: torch.Tensor) -> None:
+        assert survival.shape == prefix_mask.shape
+        assert survival.dim() == 2 and survival.shape[1] == self.gamma
+
+        probs = survival.to(torch.float64).clamp(EPS_PROB, 1.0 - EPS_PROB)
+        targets = prefix_mask.to(torch.float64)
+        bs = probs.shape[0]
+
+        probs_flat = probs.reshape(-1)
+        targets_flat = targets.reshape(-1)
+        weights = torch.ones_like(probs_flat)
+        pos_idx = (
+            torch.arange(self.gamma, device=probs.device)
+            .view(1, -1)
+            .expand(bs, self.gamma)
+            .reshape(-1)
+        )
+
+        coarse_idx = (
+            (probs_flat * self.num_coarse_bins)
+            .long()
+            .clamp_(0, self.num_coarse_bins - 1)
+        )
+        flat_coarse = pos_idx * self.num_coarse_bins + coarse_idx
+        self.coarse_count.view(-1).scatter_add_(0, flat_coarse, weights)
+        self.coarse_pred.view(-1).scatter_add_(0, flat_coarse, probs_flat)
+        self.coarse_target.view(-1).scatter_add_(0, flat_coarse, targets_flat)
+
+        fine_idx = (
+            (probs_flat * self.num_fine_bins).long().clamp_(0, self.num_fine_bins - 1)
+        )
+        flat_fine = pos_idx * self.num_fine_bins + fine_idx
+        self.fine_pos.view(-1).scatter_add_(0, flat_fine, targets_flat)
+        self.fine_neg.view(-1).scatter_add_(0, flat_fine, 1.0 - targets_flat)
+
+        self.brier_num.add_((probs - targets).pow(2).sum(dim=0))
+
+    @staticmethod
+    def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float:
+        total_pos = float(pos_hist.sum())
+        total_neg = float(neg_hist.sum())
+        if total_pos <= 0.0 or total_neg <= 0.0:
+            return float("nan")
+        cum_neg = torch.cumsum(neg_hist, dim=0)
+        cum_neg_before = cum_neg - neg_hist
+        pair = (pos_hist * cum_neg_before).sum() + 0.5 * (pos_hist * neg_hist).sum()
+        return float(pair) / (total_pos * total_neg)
+
+    def compute(self) -> list[dict]:
+        coarse_count = self.coarse_count.cpu()
+        coarse_pred = self.coarse_pred.cpu()
+        coarse_target = self.coarse_target.cpu()
+        fine_pos = self.fine_pos.cpu()
+        fine_neg = self.fine_neg.cpu()
+        brier_num = self.brier_num.cpu()
+
+        out: list[dict] = []
+        for pos in range(self.gamma):
+            weights = coarse_count[pos]
+            total = float(weights.sum())
+            if total <= 1e-12:
+                out.append(
+                    {
+                        "position": pos,
+                        "total_weight": 0.0,
+                        "ece": float("nan"),
+                        "auc": float("nan"),
+                        "brier": float("nan"),
+                        "pred_mean": float("nan"),
+                        "target_mean": float("nan"),
+                        "reliability": [],
+                    }
+                )
+                continue
+
+            denom = weights.clamp_min(1e-12)
+            avg_pred = coarse_pred[pos] / denom
+            avg_target = coarse_target[pos] / denom
+            bin_err = (avg_pred - avg_target).abs()
+            ece = float((bin_err * weights).sum()) / total
+            auc = self._auroc_from_hist(fine_pos[pos], fine_neg[pos])
+            brier = float(brier_num[pos]) / total
+            reliability = []
+            for bin_idx in range(self.num_coarse_bins):
+                weight = float(weights[bin_idx])
+                if weight <= 0.0:
+                    continue
+                reliability.append(
+                    {
+                        "bin": bin_idx,
+                        "range": [
+                            bin_idx / self.num_coarse_bins,
+                            (bin_idx + 1) / self.num_coarse_bins,
+                        ],
+                        "avg_pred": float(avg_pred[bin_idx]),
+                        "avg_target": float(avg_target[bin_idx]),
+                        "weight": weight,
+                    }
+                )
+            out.append(
+                {
+                    "position": pos,
+                    "total_weight": total,
+                    "ece": ece,
+                    "auc": auc,
+                    "brier": brier,
+                    "pred_mean": float(coarse_pred[pos].sum()) / total,
+                    "target_mean": float(coarse_target[pos].sum()) / total,
+                    "reliability": reliability,
+                }
+            )
+        return out
+
+    def format_table(self) -> str:
+        rows = self.compute()
+        header = (
+            f"{'pos':>3} {'count':>12} {'pred':>8} {'target':>8} "
+            f"{'ece':>8} {'auc':>8} {'brier':>8}"
+        )
+        lines = [
+            "DSpark confidence-head per-position calibration "
+            "(cumprod survival vs leading-correct-prefix)",
+            header,
+        ]
+        for row in rows:
+            lines.append(
+                f"{row['position']:>3} {row['total_weight']:>12.0f} "
+                f"{_format_float(row['pred_mean']):>8} "
+                f"{_format_float(row['target_mean']):>8} "
+                f"{_format_float(row['ece']):>8} "
+                f"{_format_float(row['auc']):>8} "
+                f"{_format_float(row['brier']):>8}"
+            )
+        return "\n".join(lines)
+
+
+class ConfidenceMetricsProbe:
+
+    def __init__(
+        self,
+        *,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        tp_rank: int,
+        print_every: int = 256,
+    ) -> None:
+        self.gamma = int(gamma)
+        self.verify_num_draft_tokens = int(verify_num_draft_tokens)
+        self.tp_rank = int(tp_rank)
+        self.print_every = int(print_every)
+        self._metrics: Optional[PerPositionConfidenceMetrics] = None
+        self._step_ct: int = 0
+        self._compact_warned: bool = False
+
+    def maybe_observe(
+        self,
+        *,
+        carries_confidence: bool,
+        is_compact_mode: bool,
+        confidence_raw: Optional[torch.Tensor],
+        verify_ids_2d: torch.Tensor,
+        target_logits: torch.Tensor,
+        bs: int,
+    ) -> None:
+        if not envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.get():
+            return
+        if self.tp_rank != 0:
+            return
+        if not carries_confidence:
+            return
+        if is_compact_mode:
+            if not self._compact_warned:
+                logger.warning(
+                    "SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS is ignored under "
+                    "SGLANG_RAGGED_VERIFY_MODE=compact (padded verify rows corrupt the "
+                    "per-position prefix label); run cap-accept or static to measure it."
+                )
+                self._compact_warned = True
+            return
+        if confidence_raw is None:
+            return
+
+        target_predict = torch.argmax(target_logits, dim=-1).view(
+            bs, self.verify_num_draft_tokens
+        )
+        num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus(
+            candidates=verify_ids_2d,
+            target_predict=target_predict,
+        )
+        positions = torch.arange(self.gamma, device=confidence_raw.device).view(1, -1)
+        prefix_mask = (positions < num_correct_drafts.view(-1, 1)).to(torch.float32)
+        survival = torch.cumprod(torch.sigmoid(confidence_raw.float()), dim=1)
+
+        if self._metrics is None:
+            self._metrics = PerPositionConfidenceMetrics(
+                gamma=self.gamma, device=confidence_raw.device
+            )
+        self._metrics.update(survival=survival, prefix_mask=prefix_mask)
+        self._step_ct += 1
+        if self._step_ct % self.print_every == 0:
+            logger.info("%s", self._metrics.format_table())
+
+
+_STS_COLLECT_FLUSH_EVERY: int = 256
+
+
+class DsparkStepObservers:
+    """Facade over the per-step observability sinks (info dumper, confidence
+    probe, STS collection, block-accept estimator). The worker's decode path
+    makes one call per step; all sink gating and field derivation live here
+    so the hot path stays free of observer plumbing."""
+
+    def __init__(
+        self,
+        *,
+        planner,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        tp_rank: int,
+        device,
+        simulate_acc_len: float,
+    ) -> None:
+        self._planner = planner
+        self._gamma = int(gamma)
+        self._verify_num_draft_tokens = int(verify_num_draft_tokens)
+        self._simulate_acc_len = float(simulate_acc_len)
+
+        self._confidence_probe = ConfidenceMetricsProbe(
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            tp_rank=tp_rank,
+        )
+        self._info_dumper = DsparkInfoDumper(
+            components=resolve_enabled_components(),
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            attn_tp_rank=get_parallel().attn_tp_rank,
+            device=device,
+            mode_value=planner.mode_value,
+            sps_report_interval=envs.SGLANG_DSPARK_LOG_SPS_PRED_INTERVAL.get(),
+        )
+        self._block_accept_recorder = create_block_accept_estimate_recorder(
+            gamma=gamma, device=device, tp_rank=tp_rank
+        )
+        if self._simulate_acc_len > 0 and self._block_accept_recorder is not None:
+            raise ValueError(
+                "SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH cannot be combined with "
+                "SGLANG_SIMULATE_ACC_LEN (simulated correct_len breaks the "
+                "accept-probability bookkeeping of the estimator)."
+            )
+        self._sts_collect_path = envs.SGLANG_DSPARK_STS_COLLECT_PATH.get()
+        self._sts_recorder: Optional[StsDataRecorder] = None
+
+    # --- step lifecycle -------------------------------------------------
+
+    def begin_step(self) -> None:
+        self._info_dumper.begin_step()
+
+    def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]:
+        return self._info_dumper.segment(name)
+
+    def note_prefill_step(self) -> None:
+        self._info_dumper.note_non_decode_step()
+        if self._block_accept_recorder is not None:
+            self._block_accept_recorder.flush()
+
+    def note_idle_decode_step(self) -> None:
+        self._info_dumper.note_non_decode_step()
+
+    # --- scheduler-facing hooks ------------------------------------------
+
+    def dump_info_records(self) -> Optional[dict]:
+        dumped = self._info_dumper.dump()
+        if dumped is None:
+            return None
+        dumped["simulate_acc_len"] = (
+            self._simulate_acc_len if self._simulate_acc_len > 0 else None
+        )
+        return dumped
+
+    def clear_info_records(self) -> None:
+        self._info_dumper.clear()
+
+    def block_accept_estimate_log_suffix(self) -> Optional[str]:
+        if self._block_accept_recorder is None:
+            return None
+        return self._block_accept_recorder.estimate_log_suffix()
+
+    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        if self._block_accept_recorder is None:
+            return
+        self._block_accept_recorder.note_request_finished(
+            rid=rid, natural_stop=natural_stop
+        )
+
+    # --- per-step observation --------------------------------------------
+
+    def observe_verify_step(
+        self,
+        *,
+        forward_ct: int,
+        reqs,
+        bs: int,
+        proposal_folded: bool,
+        verify_ids_2d: torch.Tensor,
+        target_logits: Optional[torch.Tensor],
+        layout,
+        confidence: Optional[torch.Tensor],
+        prefix_lens: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        draft_block,
+        sampling_info,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        bonus: torch.Tensor,
+        commit_lens: torch.Tensor,
+        verify_token_budget: Optional[int],
+        req_pool_indices: torch.Tensor,
+        verify_tier_num_tokens: int,
+        dp_tier_num_tokens: Optional[int],
+    ) -> None:
+        planner = self._planner
+        if not proposal_folded:
+            self._maybe_record_sts_collect(
+                verify_ids_2d=verify_ids_2d,
+                target_logits=target_logits,
+                bs=bs,
+            )
+            self._confidence_probe.maybe_observe(
+                carries_confidence=planner.carries_confidence,
+                is_compact_mode=planner.is_compact_mode,
+                confidence_raw=planner.last_confidence_raw,
+                verify_ids_2d=verify_ids_2d,
+                target_logits=target_logits,
+                bs=bs,
+            )
+        if self._block_accept_recorder is not None and not proposal_folded:
+            self._block_accept_recorder.observe_verify_step(
+                forward_ct=forward_ct,
+                rids=[req.rid for req in reqs],
+                draft_tokens=draft_tokens,
+                corrected_logits=draft_block.corrected_logits,
+                draft_temperatures=draft_block.temperatures,
+                greedy_mask=draft_block.greedy_mask,
+                target_logits=target_logits,
+                target_temperatures=(
+                    sampling_info.temperatures
+                    if sampling_info is not None
+                    else draft_block.temperatures
+                ),
+                truncated_sampling_mask=(
+                    (sampling_info.top_ks != TOP_K_ALL)
+                    | (sampling_info.top_ps != 1.0)
+                    | (sampling_info.min_ps > 0)
+                    if sampling_info is not None
+                    else None
+                ),
+                logits_adjustments_are_noop=verify_logits_adjustments_are_noop(
+                    sampling_info
+                ),
+                correct_len=correct_len,
+                cap_trim_lens=cap_trim_lens,
+                bonus=bonus,
+                prefix_lens=prefix_lens,
+                layout=layout,
+            )
+        if self._info_dumper.enabled:
+            budget_decision = planner.take_budget_decision()
+            predicted_step_ms = (
+                None
+                if budget_decision is None
+                or budget_decision.predicted_step_seconds is None
+                else budget_decision.predicted_step_seconds * 1e3
+            )
+            predicted_theta = (
+                None if budget_decision is None else budget_decision.predicted_theta
+            )
+            num_verify_tokens = (
+                layout.graph_num_tokens
+                if layout is not None
+                else int(verify_ids_2d.numel())
+            )
+            self._info_dumper.observe_decode_step(
+                DecodeStepObservation(
+                    forward_ct=forward_ct,
+                    bs=bs,
+                    mode=planner.mode_value,
+                    budget=verify_token_budget,
+                    lag_steps=planner.lag_steps,
+                    num_verify_tokens=num_verify_tokens,
+                    verify_tokens_local=verify_tier_num_tokens,
+                    verify_tokens_dp_synced=(
+                        -1 if dp_tier_num_tokens is None else int(dp_tier_num_tokens)
+                    ),
+                    verify_tokens_graph_key=num_verify_tokens,
+                    predicted_step_ms=predicted_step_ms,
+                    predicted_theta=predicted_theta,
+                    verify_lens=layout.verify_lens if layout is not None else None,
+                    confidence=confidence,
+                    req_pool_indices=req_pool_indices,
+                    prefix_lens=prefix_lens,
+                    draft_tokens=draft_tokens,
+                    bonus_tokens=bonus,
+                    correct_len=correct_len,
+                    cap_trim_lens=cap_trim_lens,
+                    commit_lens=commit_lens,
+                    rids=[req.rid for req in reqs],
+                )
+            )
+
+    def _maybe_record_sts_collect(
+        self,
+        *,
+        verify_ids_2d: torch.Tensor,
+        target_logits: Optional[torch.Tensor],
+        bs: int,
+    ) -> None:
+        if not self._sts_collect_path:
+            return
+        if not self._planner.carries_confidence:
+            return
+        confidence_raw = self._planner.last_confidence_raw
+        if confidence_raw is None:
+            return
+        if self._sts_recorder is None:
+            self._sts_recorder = StsDataRecorder(
+                path_stem=self._sts_collect_path,
+                gamma=self._gamma,
+                flush_every=_STS_COLLECT_FLUSH_EVERY,
+            )
+        target_predict = torch.argmax(target_logits, dim=-1).view(
+            bs, self._verify_num_draft_tokens
+        )
+        num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus(
+            candidates=verify_ids_2d,
+            target_predict=target_predict,
+        )
+        self._sts_recorder.record(
+            confidence_raw=confidence_raw,
+            num_correct_drafts=num_correct_drafts,
+        )

python/sglang/srt/speculative/dspark_components/dspark_sps.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import bisect
+from typing import Optional
+
+import msgspec
+
+
+def floor_probe_index(edges: list[int], batch_tokens: int) -> int:
+    idx = bisect.bisect_right(edges, batch_tokens) - 1
+    return max(0, min(idx, len(edges) - 1))
+
+
+class SpsCostTable(msgspec.Struct, frozen=True):
+    sample_batch_tokens: list[int]
+    sample_steps_per_sec: list[float]
+    max_batch_tokens: int
+
+    def __post_init__(self) -> None:
+        if not self.sample_batch_tokens:
+            raise ValueError("SpsCostTable requires at least one probe.")
+        if self.sample_batch_tokens != sorted(set(self.sample_batch_tokens)):
+            raise ValueError(
+                "sample_batch_tokens must be strictly increasing (monotone-sorted "
+                f"invariant), got {self.sample_batch_tokens}."
+            )
+        if len(self.sample_batch_tokens) != len(self.sample_steps_per_sec):
+            raise ValueError(
+                "sample_batch_tokens and sample_steps_per_sec must have equal length, "
+                f"got {len(self.sample_batch_tokens)} vs {len(self.sample_steps_per_sec)}."
+            )
+        if self.max_batch_tokens < self.sample_batch_tokens[-1]:
+            raise ValueError(
+                "max_batch_tokens must be >= the largest probe, got "
+                f"{self.max_batch_tokens} < {self.sample_batch_tokens[-1]}."
+            )
+
+    def lookup(self, batch_tokens: int) -> float:
+        return self.sample_steps_per_sec[
+            floor_probe_index(self.sample_batch_tokens, batch_tokens)
+        ]
+
+    def to_json(self) -> str:
+        return msgspec.json.encode(self).decode("utf-8")
+
+    @classmethod
+    def from_json(cls, data: str) -> SpsCostTable:
+        return msgspec.json.decode(data.encode("utf-8"), type=cls)
+
+
+def _interp_clamped(xs: list[int], ys: list[float], x: float) -> float:
+    if x <= xs[0]:
+        return ys[0]
+    if x >= xs[-1]:
+        return ys[-1]
+    hi = bisect.bisect_right(xs, x)
+    lo = hi - 1
+    frac = (x - xs[lo]) / (xs[hi] - xs[lo])
+    return ys[lo] + frac * (ys[hi] - ys[lo])
+
+
+class SpsAdditiveCostTable(msgspec.Struct, frozen=True):
+
+    bias_seconds: float
+    bs_probes: list[int]
+    alpha_seconds: list[float]
+    m_probes: list[int]
+    theta_seconds: list[float]
+
+    def __post_init__(self) -> None:
+        for name, probes, values in (
+            ("bs", self.bs_probes, self.alpha_seconds),
+            ("m", self.m_probes, self.theta_seconds),
+        ):
+            if not probes:
+                raise ValueError(f"SpsAdditiveCostTable requires {name}_probes.")
+            if probes != sorted(set(probes)):
+                raise ValueError(
+                    f"{name}_probes must be strictly increasing, got {probes}."
+                )
+            if len(probes) != len(values):
+                raise ValueError(
+                    f"{name}_probes and its values must have equal length, got "
+                    f"{len(probes)} vs {len(values)}."
+                )
+        if self.bias_seconds <= 0:
+            raise ValueError(f"bias_seconds must be > 0, got {self.bias_seconds}.")
+
+    def step_time(self, *, num_reqs: int, budget: int) -> float:
+        return (
+            self.bias_seconds
+            + _interp_clamped(self.bs_probes, self.alpha_seconds, float(num_reqs))
+            + _interp_clamped(
+                self.m_probes, self.theta_seconds, float(num_reqs + budget)
+            )
+        )
+
+    def to_json(self) -> str:
+        return msgspec.json.encode(self).decode("utf-8")
+
+    @classmethod
+    def from_json(cls, data: str) -> SpsAdditiveCostTable:
+        return msgspec.json.decode(data.encode("utf-8"), type=cls)
+
+
+def profile_sps_table(
+    *,
+    probes: list[tuple[int, float]],
+    max_batch_tokens: Optional[int] = None,
+) -> SpsCostTable:
+    if not probes:
+        raise ValueError("profile_sps_table requires at least one probe.")
+
+    sorted_probes = sorted(probes, key=lambda probe: probe[0])
+
+    sample_batch_tokens: list[int] = []
+    sample_steps_per_sec: list[float] = []
+    for batch_tokens, steps_per_sec in sorted_probes:
+        batch_tokens = int(batch_tokens)
+        if batch_tokens < 1:
+            raise ValueError(
+                f"profile_sps_table requires batch_tokens >= 1, got {batch_tokens}."
+            )
+        if sample_batch_tokens and batch_tokens == sample_batch_tokens[-1]:
+            raise ValueError(
+                "profile_sps_table requires unique batch_tokens per probe; "
+                f"batch_tokens={batch_tokens} appears more than once. Median the "
+                "repeated samples per batch_tokens before calling the assembler."
+            )
+        sample_batch_tokens.append(batch_tokens)
+        sample_steps_per_sec.append(float(steps_per_sec))
+
+    resolved_max = (
+        int(max_batch_tokens)
+        if max_batch_tokens is not None
+        else sample_batch_tokens[-1]
+    )
+    return SpsCostTable(
+        sample_batch_tokens=sample_batch_tokens,
+        sample_steps_per_sec=sample_steps_per_sec,
+        max_batch_tokens=resolved_max,
+    )
+
+
+def load_sps_table_from_path(path: str):
+    with open(path, "r", encoding="utf-8") as f:
+        data = f.read()
+    if '"bias_seconds"' in data:
+        return SpsAdditiveCostTable.from_json(data)
+    return SpsCostTable.from_json(data)
+
+
+def build_uninitialized_sps_table(*, max_batch_tokens: int) -> SpsCostTable:
+    return SpsCostTable(
+        sample_batch_tokens=[1],
+        sample_steps_per_sec=[1.0],
+        max_batch_tokens=max_batch_tokens,
+    )
+
+
+def is_uninitialized_sps_table(table: SpsCostTable | SpsAdditiveCostTable) -> bool:
+    if isinstance(table, SpsAdditiveCostTable):
+        return False
+    return len(table.sample_batch_tokens) <= 1

python/sglang/srt/speculative/dspark_components/dspark_sts.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import msgspec
+import torch
+
+
+class DSparkStsCalibration(msgspec.Struct, frozen=True, omit_defaults=True):
+    temperatures: list[float]
+    dataset: str = ""
+    num_samples: int = 0
+    ece_before: list[float] = []
+    ece_after: list[float] = []
+
+    def __post_init__(self) -> None:
+        if not self.temperatures:
+            raise ValueError("DSparkStsCalibration requires at least one temperature.")
+        for temperature in self.temperatures:
+            if temperature <= 0:
+                raise ValueError(
+                    "DSparkStsCalibration temperatures must all be > 0, got "
+                    f"{self.temperatures}."
+                )
+
+    def to_json(self) -> str:
+        return msgspec.json.encode(self).decode("utf-8")
+
+    @classmethod
+    def from_json(cls, data: str) -> DSparkStsCalibration:
+        return msgspec.json.decode(data.encode("utf-8"), type=cls)
+
+
+def load_sts_calibration_from_path(path: str) -> DSparkStsCalibration:
+    with open(path, "r", encoding="utf-8") as f:
+        return DSparkStsCalibration.from_json(f.read())
+
+
+class StsDataRecorder:
+    def __init__(self, *, path_stem: str, gamma: int, flush_every: int) -> None:
+        self.path_stem = path_stem
+        self.gamma = int(gamma)
+        self.flush_every = int(flush_every)
+        self._logits_buffer: list[torch.Tensor] = []
+        self._prefix_mask_buffer: list[torch.Tensor] = []
+        self._shard_ct = 0
+
+    def record(
+        self, *, confidence_raw: torch.Tensor, num_correct_drafts: torch.Tensor
+    ) -> None:
+        logits = confidence_raw.detach().to(device="cpu", dtype=torch.float32)
+        positions = torch.arange(self.gamma).view(1, -1)
+        counts = (
+            num_correct_drafts.detach().to(device="cpu", dtype=torch.int64).view(-1, 1)
+        )
+        prefix_mask = (positions < counts).to(torch.float32)
+        self._logits_buffer.append(logits)
+        self._prefix_mask_buffer.append(prefix_mask)
+        if len(self._logits_buffer) >= self.flush_every:
+            self.flush()
+
+    def flush(self) -> None:
+        if not self._logits_buffer:
+            return
+        shard_path = Path(f"{self.path_stem}.{self._shard_ct}.pt")
+        shard_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(
+            {
+                "logits": torch.cat(self._logits_buffer, dim=0),
+                "prefix_mask": torch.cat(self._prefix_mask_buffer, dim=0),
+            },
+            shard_path,
+        )
+        self._logits_buffer.clear()
+        self._prefix_mask_buffer.clear()
+        self._shard_ct += 1

python/sglang/srt/speculative/dspark_components/dspark_verify.py

@@ -0,0 +1,716 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
+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
+from sglang.srt.speculative.dspark_components.dspark_draft import DraftBlockResult
+from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
+    TargetHiddenKvInjector,
+)
+from sglang.srt.speculative.dspark_components.dspark_planner import (
+    VerifyWindow,
+    apply_logits_adjustments_strided,
+)
+from sglang.srt.speculative.dspark_components.kernels.dspark_accept import (
+    AcceptGreedy,
+    AcceptSampling,
+    FinalizeAcceptLens,
+    SelectMixedAccept,
+    SoftmaxTemp,
+    accept_greedy_triton,
+    finalize_accept_lens_triton,
+)
+from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import (
+    BuildCommitInjectLayout,
+    BuildOutTokens,
+    BuildRaggedVerifyWindow,
+    RaggedVerifyWindow,
+    ScatterCompactToStrided,
+    scatter_compact_to_strided_into,
+)
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+def verify_logits_adjustments_are_noop(sampling_info) -> bool:
+    if sampling_info is None:
+        return True
+    if sampling_info.has_custom_logit_processor:
+        return False
+    if getattr(sampling_info, "acc_linear_penalties", None) is not None:
+        return False
+    penalizer = getattr(sampling_info, "penalizer_orchestrator", None)
+    if penalizer is not None and penalizer.is_required:
+        return False
+    if getattr(sampling_info, "vocab_mask", None) is not None:
+        return False
+    if getattr(sampling_info, "logit_bias", None) is not None:
+        return False
+    return True
+
+
+class TargetVerifyResult(msgspec.Struct, frozen=True):
+    logits_output: object
+    can_run_cuda_graph: bool
+
+
+class TargetVerifyExecutor:
+    def __init__(
+        self,
+        *,
+        target_worker,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        model_runner,
+        kv_injector: TargetHiddenKvInjector,
+        verify_epilogue=None,
+        simulate_acc_len: float = 0.0,
+    ) -> None:
+        self.target_worker = target_worker
+        self.gamma = int(gamma)
+        self.verify_num_draft_tokens = verify_num_draft_tokens
+        self.model_runner = model_runner
+        self.kv_injector = kv_injector
+        self.verify_epilogue = verify_epilogue
+        self._verify_backend_self_adds_seq_lens_cache: Optional[bool] = None
+        self._simulate_acc_len = float(simulate_acc_len)
+        self._simulated_correct_drafts_buf: Optional[torch.Tensor] = None
+
+    def accept_and_finalize(
+        self,
+        *,
+        folded_accept: bool,
+        bs: int,
+        verify_ids_2d: torch.Tensor,
+        target_logits: Optional[torch.Tensor],
+        draft_block: DraftBlockResult,
+        sampling_info,
+        draft_input: DFlashDraftInputV2,
+        layout: Optional[RaggedVerifyLayout],
+        prefix_lens: torch.Tensor,
+        draft_tokens: torch.Tensor,
+    ) -> AcceptOuts:
+        """Produce the per-request accept outcome after target verify.
+
+        Folded path: the accept/finalize/out-token kernels already ran inside
+        the target-verify cuda graph (DsparkVerifyEpilogue); read its buffers.
+        Eager path: run them here, including the SGLANG_SIMULATE_ACC_LEN
+        override.
+        """
+        if folded_accept:
+            return self.verify_epilogue.read_accept(bs)
+
+        correct_len, bonus, cap_trim_lens = accept_draft_tokens(
+            candidates=verify_ids_2d,
+            target_logits=target_logits,
+            draft_block=draft_block,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=self.gamma,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            cutoff_layout=layout,
+        )
+        if self._simulate_acc_len > 0:
+            correct_len = self._simulated_correct_len(
+                bs=bs, dtype=correct_len.dtype, device=correct_len.device
+            )
+
+        finalized = FinalizeAcceptLens.execute(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=prefix_lens,
+        )
+        out_tokens = BuildOutTokens.execute(
+            draft_tokens=draft_tokens,
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            gamma=self.gamma,
+        )
+        return AcceptOuts(
+            correct_len=correct_len,
+            bonus=bonus,
+            cap_trim_lens=finalized.cap_trim_lens,
+            commit_lens=finalized.commit_lens,
+            new_seq_lens=finalized.new_seq_lens,
+            out_tokens=out_tokens,
+        )
+
+    def _simulated_correct_len(
+        self, *, bs: int, dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        buf = self._simulated_correct_drafts_buf
+        if buf is None or buf.numel() < bs or buf.dtype != dtype:
+            correct_target = int(
+                round(min(max(self._simulate_acc_len - 1.0, 0.0), float(self.gamma)))
+            )
+            buf = torch.full(
+                (max(bs, 512),), correct_target, dtype=dtype, device=device
+            )
+            self._simulated_correct_drafts_buf = buf
+        return buf[:bs]
+
+    def run_idle_participation(
+        self,
+        *,
+        batch: ScheduleBatch,
+        idle_layout: Optional[RaggedVerifyLayout],
+    ) -> None:
+        """Run a dummy target-verify forward so an idle DP rank joins the
+        token-keyed collective ops of the busy ranks' verify step."""
+        device = self.model_runner.device
+        if self.verify_epilogue is not None:
+            self.verify_epilogue.begin_step(None, armed=False)
+        num_dummy_tokens = (
+            idle_layout.graph_num_tokens if idle_layout is not None else 0
+        )
+        verify_input = DFlashVerifyInput(
+            draft_token=torch.zeros(
+                (num_dummy_tokens,), dtype=torch.int64, device=device
+            ),
+            positions=torch.zeros(
+                (num_dummy_tokens,), dtype=torch.int64, device=device
+            ),
+            draft_token_num=self.verify_num_draft_tokens,
+            custom_mask=None,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+            ragged_verify_layout=idle_layout,
+        )
+        batch.out_cache_loc = torch.zeros(
+            (num_dummy_tokens,), dtype=torch.int64, device=device
+        )
+        if idle_layout is not None:
+            num_dummy_slots = int(idle_layout.verify_lens.numel())
+            batch.seq_lens = torch.ones(
+                (num_dummy_slots,), dtype=torch.int64, device=device
+            )
+            batch.req_pool_indices = torch.zeros(
+                (num_dummy_slots,), dtype=torch.int64, device=device
+            )
+            batch.seq_lens_cpu = torch.ones((num_dummy_slots,), dtype=torch.int64)
+            batch.seq_lens_sum = num_dummy_slots
+            batch.forward_mode = ForwardMode.TARGET_VERIFY
+        verify_forward_batch, _ = verify_input.prepare_for_verify(
+            batch, self.target_worker
+        )
+        self.target_worker.forward_batch_generation(
+            batch=None,
+            forward_batch=verify_forward_batch,
+            is_verify=True,
+            skip_attn_backend_init=True,
+        )
+
+    def run_non_compact(
+        self,
+        *,
+        batch: ScheduleBatch,
+        draft_input: DFlashDraftInputV2,
+        verify_ids_2d: torch.Tensor,
+        verify_window: VerifyWindow,
+        sampling_info,
+    ) -> TargetVerifyResult:
+        verify_w = self.verify_num_draft_tokens
+        positions_2d = verify_window.positions_2d
+        verify_cache_loc = verify_window.verify_cache_loc
+
+        verify_input = DFlashVerifyInput(
+            draft_token=verify_ids_2d.reshape(-1),
+            positions=positions_2d.reshape(-1),
+            draft_token_num=verify_w,
+            custom_mask=None,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+        )
+        batch.out_cache_loc = verify_cache_loc
+        seq_lens_cpu_backup = batch.seq_lens_cpu
+        seq_lens_sum_backup = batch.seq_lens_sum
+        if not self._verify_backend_self_adds_seq_lens():
+            if seq_lens_cpu_backup is not None:
+                batch.seq_lens_cpu = seq_lens_cpu_backup + verify_w
+                batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
+            elif draft_input.reserved_seq_lens_cpu is not None:
+                batch.seq_lens_cpu = draft_input.reserved_seq_lens_cpu
+                batch.seq_lens_sum = int(draft_input.reserved_seq_lens_sum)
+
+        result = self._forward_prepared_verify(
+            batch=batch,
+            verify_input=verify_input,
+            seq_lens_cpu_backup=seq_lens_cpu_backup,
+            seq_lens_sum_backup=seq_lens_sum_backup,
+        )
+
+        if sampling_info is not None:
+            apply_dflash_verify_logits_adjustments(
+                next_token_logits=result.logits_output.next_token_logits,
+                sampling_info=sampling_info,
+                draft_token_num=verify_w,
+            )
+
+        return result
+
+    def _forward_prepared_verify(
+        self,
+        *,
+        batch: ScheduleBatch,
+        verify_input: DFlashVerifyInput,
+        seq_lens_cpu_backup,
+        seq_lens_sum_backup,
+    ) -> TargetVerifyResult:
+        verify_forward_batch, _ = verify_input.prepare_for_verify(
+            batch, self.target_worker
+        )
+        batch.seq_lens_cpu = seq_lens_cpu_backup
+        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,
+        )
+        return TargetVerifyResult(
+            logits_output=target_out.logits_output,
+            can_run_cuda_graph=target_out.can_run_cuda_graph,
+        )
+
+    def commit_hidden(
+        self,
+        *,
+        batch: ScheduleBatch,
+        layout: Optional[RaggedVerifyLayout],
+        hidden_strided: Optional[torch.Tensor],
+        verify_window: VerifyWindow,
+        logits_output,
+        commit_lens: torch.Tensor,
+        bs: int,
+        run_compact: bool,
+    ) -> None:
+        if run_compact:
+            self.kv_injector.inject_ragged(
+                batch=batch,
+                layout=layout,
+                hidden_strided=hidden_strided,
+                commit_lens=commit_lens,
+                bs=bs,
+            )
+            return
+        hidden = logits_output.hidden_states
+        if hidden is None:
+            raise RuntimeError("DSpark verify requires target hidden states, got None.")
+        hidden = hidden.view(bs, self.verify_num_draft_tokens, -1)
+        self.kv_injector.inject_target_hidden(
+            target_hidden=hidden.reshape(-1, hidden.shape[-1]),
+            cache_loc=verify_window.verify_cache_loc,
+            cache_loc_2d=verify_window.verify_cache_loc_2d,
+            positions=verify_window.positions_2d.reshape(-1),
+            commit_lens=commit_lens,
+        )
+
+    def _run_ragged(
+        self,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        ragged_window: RaggedVerifyWindow,
+        sampling_info,
+    ) -> TargetVerifyResult:
+        verify_input = DFlashVerifyInput(
+            draft_token=ragged_window.verify_ids,
+            positions=ragged_window.positions,
+            draft_token_num=self.verify_num_draft_tokens,
+            custom_mask=None,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+            ragged_verify_layout=layout,
+        )
+        batch.out_cache_loc = ragged_window.verify_cache_loc
+        seq_lens_cpu_backup = batch.seq_lens_cpu
+        seq_lens_sum_backup = batch.seq_lens_sum
+        if seq_lens_cpu_backup is not None:
+            verify_lens_cpu = (
+                layout.verify_lens_cpu
+                if layout.verify_lens_cpu is not None
+                else layout.verify_lens.cpu().tolist()
+            )
+            batch.seq_lens_cpu = seq_lens_cpu_backup + torch.tensor(
+                verify_lens_cpu, dtype=seq_lens_cpu_backup.dtype
+            )
+            batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
+
+        return self._forward_prepared_verify(
+            batch=batch,
+            verify_input=verify_input,
+            seq_lens_cpu_backup=seq_lens_cpu_backup,
+            seq_lens_sum_backup=seq_lens_sum_backup,
+        )
+
+    def run_compact(
+        self,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        bs: int,
+        device: str,
+        sampling_info,
+        inject_gate: bool = False,
+    ) -> tuple[TargetVerifyResult, torch.Tensor]:
+        ragged_window = BuildRaggedVerifyWindow.execute(
+            batch=batch,
+            layout=layout,
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            model_runner=self.model_runner,
+        )
+        if self.verify_epilogue is not None:
+            self.verify_epilogue.begin_step(layout.verify_lens, armed=inject_gate)
+        target_verify = self._run_ragged(
+            batch=batch,
+            layout=layout,
+            ragged_window=ragged_window,
+            sampling_info=sampling_info,
+        )
+        logits_output = target_verify.logits_output
+
+        stride = self.verify_num_draft_tokens
+        if self.verify_epilogue is not None and target_verify.can_run_cuda_graph:
+            strided_logits = self.verify_epilogue.strided_logits
+            hidden_strided = self.verify_epilogue.strided_hidden
+            assert strided_logits is not None and hidden_strided is not None, (
+                "verify epilogue buffers unwritten after a graph replay -- the "
+                "replayed graph was captured without the epilogue"
+            )
+            strided_logits = strided_logits[: bs * stride]
+            hidden_strided = hidden_strided[: bs * stride]
+        else:
+            compact_logits = logits_output.next_token_logits
+            strided_logits = ScatterCompactToStrided.execute(
+                compact=compact_logits,
+                layout=layout,
+                fill_value=0.0,
+                verify_num_draft_tokens=stride,
+            )
+            compact_hidden = logits_output.hidden_states
+            if compact_hidden is None:
+                raise RuntimeError(
+                    "DSpark verify requires target hidden states, got None."
+                )
+            hidden_strided = ScatterCompactToStrided.execute(
+                compact=compact_hidden,
+                layout=layout,
+                fill_value=0.0,
+                verify_num_draft_tokens=stride,
+            )
+        apply_logits_adjustments_strided(
+            next_token_logits=strided_logits,
+            sampling_info=sampling_info,
+            verify_num_draft_tokens=stride,
+        )
+        logits_output.next_token_logits = strided_logits
+        logits_output.hidden_states = hidden_strided
+        return target_verify, hidden_strided
+
+    def _verify_backend_self_adds_seq_lens(self) -> bool:
+        if self._verify_backend_self_adds_seq_lens_cache is None:
+            backend = self.target_worker.model_runner.attn_backend
+            self._verify_backend_self_adds_seq_lens_cache = hasattr(
+                backend, "make_forward_metadata_from_raw_verify"
+            )
+        return self._verify_backend_self_adds_seq_lens_cache
+
+
+class CommitInjectCtx(msgspec.Struct):
+
+    draft_model: object
+    block_pos_offsets: torch.Tensor
+    resolve_pool: object
+    resolve_req_to_token: object
+
+
+class AcceptOuts(msgspec.Struct):
+    correct_len: torch.Tensor
+    bonus: torch.Tensor
+    cap_trim_lens: torch.Tensor
+    commit_lens: torch.Tensor
+    new_seq_lens: torch.Tensor
+    out_tokens: torch.Tensor
+
+
+class DsparkVerifyEpilogue:
+
+    def __init__(
+        self,
+        *,
+        max_bs: int,
+        verify_num_draft_tokens: int,
+        device,
+        commit_ctx: Optional[CommitInjectCtx] = None,
+    ) -> None:
+        self.max_bs = int(max_bs)
+        self.stride = int(verify_num_draft_tokens)
+        self.gamma = self.stride - 1
+        self.commit_ctx = commit_ctx
+        self.inject_gate_buf = torch.zeros((1,), dtype=torch.int32, device=device)
+        self.verify_lens_buf = torch.zeros(
+            (self.max_bs,), dtype=torch.int64, device=device
+        )
+        self.draft_tokens_buf = torch.zeros(
+            (self.max_bs * self.gamma,), dtype=torch.int64, device=device
+        )
+        self.correct_len_buf = torch.zeros(
+            (self.max_bs,), dtype=torch.int64, device=device
+        )
+        self.bonus_buf = torch.zeros((self.max_bs,), dtype=torch.int64, device=device)
+        self.cap_trim_lens_buf = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device=device
+        )
+        self.commit_lens_buf = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device=device
+        )
+        self.new_seq_lens_buf = torch.zeros(
+            (self.max_bs,), dtype=torch.int64, device=device
+        )
+        self.out_tokens_buf = torch.zeros(
+            (self.max_bs, self.stride), dtype=torch.int64, device=device
+        )
+        self.strided_logits: Optional[torch.Tensor] = None
+        self.strided_hidden: Optional[torch.Tensor] = None
+
+    def capture_hook(self, runner, out, forward_batch, num_tokens) -> None:
+        if runner.model_runner.is_draft_worker or not runner.ragged_verify_mode:
+            return
+        if (
+            not isinstance(out, LogitsProcessorOutput)
+            or out.next_token_logits is None
+            or out.hidden_states is None
+        ):
+            return
+        self(
+            compact_logits=out.next_token_logits,
+            compact_hidden=out.hidden_states,
+            input_ids=forward_batch.input_ids,
+            seq_lens=forward_batch.seq_lens,
+            req_pool_indices=forward_batch.req_pool_indices,
+            bs=forward_batch.batch_size,
+        )
+
+    def begin_step(self, verify_lens, armed: bool) -> None:
+        if verify_lens is None:
+            self.verify_lens_buf.zero_()
+        else:
+            bs = verify_lens.shape[0]
+            self.verify_lens_buf[:bs].copy_(verify_lens)
+            if bs < self.max_bs:
+                self.verify_lens_buf[bs:].zero_()
+        self.inject_gate_buf.fill_(1 if armed else 0)
+
+    def read_accept(self, bs: int) -> AcceptOuts:
+        return AcceptOuts(
+            correct_len=self.correct_len_buf[:bs],
+            bonus=self.bonus_buf[:bs],
+            cap_trim_lens=self.cap_trim_lens_buf[:bs],
+            commit_lens=self.commit_lens_buf[:bs],
+            new_seq_lens=self.new_seq_lens_buf[:bs],
+            out_tokens=self.out_tokens_buf[:bs],
+        )
+
+    @property
+    def folds_commit(self) -> bool:
+        if self.commit_ctx is None:
+            return False
+        pool = self.commit_ctx.resolve_pool()
+        return hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope")
+
+    def _ensure_out(
+        self, buf: Optional[torch.Tensor], compact: torch.Tensor
+    ) -> torch.Tensor:
+        if (
+            buf is not None
+            and buf.dtype == compact.dtype
+            and buf.shape[1] == compact.shape[1]
+        ):
+            return buf
+        assert not torch.cuda.is_current_stream_capturing(), (
+            "DsparkVerifyEpilogue output buffers must be allocated during "
+            "warmup, not inside graph capture (pool memory is unreadable "
+            "post-replay)."
+        )
+        return torch.empty(
+            (self.max_bs * self.stride, compact.shape[1]),
+            dtype=compact.dtype,
+            device=compact.device,
+        )
+
+    def __call__(
+        self,
+        *,
+        compact_logits: torch.Tensor,
+        compact_hidden: torch.Tensor,
+        input_ids: torch.Tensor,
+        seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        bs: int,
+    ) -> None:
+        self.strided_logits = self._ensure_out(self.strided_logits, compact_logits)
+        self.strided_hidden = self._ensure_out(self.strided_hidden, compact_hidden)
+        verify_lens = self.verify_lens_buf[:bs]
+        self._scatter(compact_logits, compact_hidden, verify_lens, bs)
+        commit_lens = self._accept(input_ids, seq_lens, verify_lens, bs)
+        if self.folds_commit:
+            self._commit_inject(
+                commit_lens, verify_lens, seq_lens, req_pool_indices, bs
+            )
+
+    def _scatter(self, compact_logits, compact_hidden, verify_lens, bs: int) -> None:
+        scatter_compact_to_strided_into(
+            compact=compact_logits,
+            verify_lens=verify_lens,
+            out=self.strided_logits[: bs * self.stride],
+            stride=self.stride,
+            fill_value=0.0,
+        )
+        scatter_compact_to_strided_into(
+            compact=compact_hidden,
+            verify_lens=verify_lens,
+            out=self.strided_hidden[: bs * self.stride],
+            stride=self.stride,
+            fill_value=0.0,
+        )
+
+    def _accept(self, input_ids, seq_lens, verify_lens, bs: int) -> torch.Tensor:
+        candidates = torch.zeros(
+            (bs * self.stride, 1), dtype=input_ids.dtype, device=input_ids.device
+        )
+        scatter_compact_to_strided_into(
+            compact=input_ids.view(-1, 1),
+            verify_lens=verify_lens,
+            out=candidates,
+            stride=self.stride,
+            fill_value=0,
+        )
+        correct_len, bonus, cap_trim_lens = accept_greedy_triton(
+            candidates=candidates.view(bs, self.stride),
+            target_logits=self.strided_logits[: bs * self.stride],
+            verify_num_draft_tokens=self.stride,
+            cutoff_verify_lens=verify_lens,
+        )
+        finalized = finalize_accept_lens_triton(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=seq_lens[:bs],
+        )
+        out_tokens = BuildOutTokens.execute(
+            draft_tokens=self.draft_tokens_buf[: bs * self.gamma].view(bs, self.gamma),
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=self.stride,
+            gamma=self.gamma,
+        )
+        self.correct_len_buf[:bs].copy_(correct_len)
+        self.bonus_buf[:bs].copy_(bonus)
+        self.cap_trim_lens_buf[:bs].copy_(cap_trim_lens.to(torch.int32))
+        self.commit_lens_buf[:bs].copy_(finalized.commit_lens)
+        self.new_seq_lens_buf[:bs].copy_(finalized.new_seq_lens)
+        self.out_tokens_buf[:bs].copy_(out_tokens.view(bs, self.stride))
+        return finalized.commit_lens
+
+    def _commit_inject(
+        self, commit_lens, verify_lens, seq_lens, req_pool_indices, bs: int
+    ) -> None:
+        ctx = self.commit_ctx
+        pool = ctx.resolve_pool()
+        gated_commit_lens = (
+            torch.minimum(commit_lens, verify_lens.to(torch.int32))
+            * self.inject_gate_buf
+        )
+        inject_layout = BuildCommitInjectLayout.execute(
+            req_pool_indices=req_pool_indices,
+            req_to_token=ctx.resolve_req_to_token(),
+            prefix_lens=seq_lens[:bs],
+            block_pos_offsets=ctx.block_pos_offsets[: self.stride],
+            full_to_swa_mapping=pool.full_to_swa_index_mapping,
+            commit_lens=gated_commit_lens,
+            stride=self.stride,
+        )
+        with torch.inference_mode():
+            ctx.draft_model.write_target_hidden_kv(
+                main_hidden=self.strided_hidden[: bs * self.stride],
+                swa_loc=inject_layout.swa_loc,
+                positions=inject_layout.positions,
+                pool=pool,
+            )
+
+
+def accept_draft_tokens(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_block: DraftBlockResult,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_layout: Optional[RaggedVerifyLayout] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    greedy_mask = draft_block.greedy_mask
+    cutoff_verify_lens = None if cutoff_layout is None else cutoff_layout.verify_lens
+    all_greedy = sampling_info is None or sampling_info.is_all_greedy
+    if all_greedy:
+        return AcceptGreedy.execute(
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+    bs, gamma_rows, vocab = draft_block.corrected_logits.shape
+    draft_probs = SoftmaxTemp.execute(
+        logits=draft_block.corrected_logits.reshape(bs * gamma_rows, vocab),
+        temperatures=draft_block.temperatures,
+        rows_per_request=gamma_rows,
+    ).view(bs, gamma_rows, vocab)
+    if not sampling_info.is_any_greedy:
+        return AcceptSampling.execute(
+            candidates=candidates,
+            target_logits=target_logits,
+            draft_probs=draft_probs,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+    greedy_len, greedy_bonus, greedy_trim = AcceptGreedy.execute(
+        candidates=candidates,
+        target_logits=target_logits,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    sampling_len, sampling_bonus, sampling_trim = AcceptSampling.execute(
+        candidates=candidates,
+        target_logits=target_logits,
+        draft_probs=draft_probs,
+        sampling_info=sampling_info,
+        draft_input=draft_input,
+        gamma=gamma,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    selected = SelectMixedAccept.execute(
+        greedy_mask=greedy_mask,
+        greedy_len=greedy_len,
+        greedy_bonus=greedy_bonus,
+        greedy_trim=greedy_trim,
+        sampling_len=sampling_len,
+        sampling_bonus=sampling_bonus,
+        sampling_trim=sampling_trim,
+    )
+    return selected.correct_len, selected.bonus, selected.cap_trim_lens

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -0,0 +1,693 @@
+import logging
+from contextlib import nullcontext
+from typing import Optional
+
+import torch
+
+from sglang.srt.environ import envs
+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,
+    compute_position,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+from sglang.srt.speculative.draft_worker_common import (
+    build_block_pos_offsets,
+    build_draft_tp_worker,
+    make_draft_block_spec_info,
+    make_draft_sampler_capture_hook,
+)
+from sglang.srt.speculative.dspark_components.dspark_config import (
+    DSV4_DRAFT_ATTENTION_BACKEND,
+    draft_is_deepseek_v4,
+    resolve_runtime_config,
+)
+from sglang.srt.speculative.dspark_components.dspark_draft import (
+    DraftBlockProposer,
+    make_next_draft_input,
+    maybe_build_draft_sampler,
+)
+from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
+    TargetHiddenKvInjector,
+)
+from sglang.srt.speculative.dspark_components.dspark_observability import (
+    DsparkStepObservers,
+    InfoSegment,
+)
+from sglang.srt.speculative.dspark_components.dspark_planner import (
+    DSparkVerifyPlanner,
+    alloc_verify_window,
+    dp_global_verify_tier_num_tokens,
+    idle_ragged_layout,
+)
+from sglang.srt.speculative.dspark_components.dspark_verify import (
+    CommitInjectCtx,
+    DsparkVerifyEpilogue,
+    TargetVerifyExecutor,
+    verify_logits_adjustments_are_noop,
+)
+from sglang.srt.speculative.spec_utils import draft_tp_context
+from sglang.srt.utils import get_available_gpu_memory, is_cuda
+
+logger = logging.getLogger(__name__)
+
+
+class DSparkWorkerV2(BaseSpecWorker):
+
+    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,
+    ):
+        self.server_args = server_args
+        self.gpu_id = gpu_id
+        self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.moe_ep_rank = moe_ep_rank
+        self.attn_cp_rank = attn_cp_rank
+        self.moe_dp_rank = moe_dp_rank
+        self.nccl_port = nccl_port
+        self._target_worker = target_worker
+        self.model_runner = target_worker.model_runner
+        self.page_size = server_args.page_size
+        self.device = target_worker.device
+
+        self._draft_is_moe = draft_is_deepseek_v4(server_args=server_args)
+        self._draft_dp_context_enabled = (
+            server_args.enable_dp_attention and not self._draft_is_moe
+        )
+        attn_tp_size = server_args.tp_size // max(server_args.dp_size, 1)
+        if server_args.enable_dp_attention and self._draft_is_moe and attn_tp_size > 1:
+            raise ValueError(
+                "DSpark + dp attention with a DeepSeek-V4 (MoE) draft requires "
+                "attn_tp == 1 (set --dp-size == --tp). attn_tp > 1 corrupts the "
+                "MoE-under-DP all-reduce."
+            )
+
+        with self._draft_context():
+            bundle = build_draft_tp_worker(
+                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_model_config=target_worker.model_runner.model_config,
+                algo_label="DSPARK",
+                attention_backend_override=(
+                    DSV4_DRAFT_ATTENTION_BACKEND if self._draft_is_moe else None
+                ),
+            )
+        self._draft_worker = bundle.draft_worker
+        self.draft_model_runner = bundle.draft_model_runner
+        self.draft_model = bundle.draft_model
+        self._draft_sampler = None
+
+        runtime_config = resolve_runtime_config(
+            draft_hf_config=self.draft_model_runner.model_config.hf_config,
+            speculative_num_draft_tokens=server_args.speculative_num_draft_tokens,
+            target_vocab_size=int(
+                self.target_worker.model_runner.model_config.vocab_size
+            ),
+        )
+        self.gamma = runtime_config.gamma
+        self.verify_num_draft_tokens = runtime_config.verify_num_draft_tokens
+        self.speculative_num_draft_tokens = self.verify_num_draft_tokens
+        self._mask_token_id = runtime_config.mask_token_id
+
+        if self.tp_rank == 0:
+            logger.info(
+                "Initialized DSpark draft runner. attention_backend=%s, model=%s, "
+                "gamma=%s, verify_num_draft_tokens=%s, mask_token_id=%s, "
+                "markov_head=%s",
+                bundle.resolved_attention_backend,
+                self.draft_model.__class__.__name__,
+                self.gamma,
+                self.verify_num_draft_tokens,
+                self._mask_token_id,
+                type(self.draft_model.markov_head).__name__,
+            )
+
+        self._block_pos_offsets = build_block_pos_offsets(
+            length=self.verify_num_draft_tokens, device=self.device
+        )
+        self._draft_block_spec_info = make_draft_block_spec_info(
+            draft_token_num=int(self.gamma), device=self.device
+        )
+
+        target_model = self.target_worker.model_runner.model
+        lm_head = getattr(target_model, "lm_head", None)
+        if lm_head is None or not hasattr(lm_head, "weight"):
+            raise RuntimeError(
+                "DSpark requires the target model to expose `lm_head` with `weight`."
+            )
+        self.draft_model.attach_shared_modules(
+            embed_tokens=self._resolve_target_embed_tokens(target_model),
+            lm_head=lm_head,
+        )
+
+        self._verify_planner = DSparkVerifyPlanner(
+            draft_model=self.draft_model,
+            gamma=self.gamma,
+            model_runner=self.model_runner,
+            device=self.device,
+            tp_rank=self.tp_rank,
+            server_args=self.server_args,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+        )
+        if (
+            server_args.enable_dp_attention
+            and not self._draft_is_moe
+            and self._verify_planner.is_compact_mode
+            and not server_args.disable_cuda_graph
+        ):
+            raise ValueError(
+                "DSpark dense-draft compact verify under --enable-dp-attention does not "
+                "yet support cuda graph (idle DP groups cannot join the token-keyed "
+                "compact graph). Re-run with --disable-cuda-graph (eager is lossless), "
+                "or use SGLANG_RAGGED_VERIFY_MODE=static. The dsv4 (MoE) draft supports "
+                "cuda graph under DP."
+            )
+        self._kv_injector = TargetHiddenKvInjector(
+            draft_model=self.draft_model,
+            draft_model_runner=self.draft_model_runner,
+            model_runner=self.model_runner,
+            device=self.device,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            block_pos_offsets=self._block_pos_offsets,
+        )
+        self._proposer = DraftBlockProposer(
+            draft_model=self.draft_model,
+            draft_model_runner=self.draft_model_runner,
+            gamma=self.gamma,
+            mask_token_id=self._mask_token_id,
+            draft_block_spec_info=self._draft_block_spec_info,
+            dp_moe_sync=self._draft_is_moe and server_args.enable_dp_attention,
+        )
+        self._verify_epilogue = None
+        if (
+            self._verify_planner.is_compact_mode
+            and not server_args.disable_cuda_graph
+            and is_cuda()
+        ):
+            self._verify_epilogue = DsparkVerifyEpilogue(
+                max_bs=max(server_args.cuda_graph_config.decode.bs),
+                verify_num_draft_tokens=self.verify_num_draft_tokens,
+                device=self.device,
+                commit_ctx=CommitInjectCtx(
+                    draft_model=self.draft_model,
+                    block_pos_offsets=self._block_pos_offsets,
+                    resolve_pool=lambda: self.draft_model_runner.token_to_kv_pool,
+                    resolve_req_to_token=lambda: (
+                        self.model_runner.req_to_token_pool.req_to_token
+                    ),
+                ),
+            )
+            self.model_runner.capture_tail_hooks.append(
+                self._verify_epilogue.capture_hook
+            )
+
+        self._simulate_acc_len = float(envs.SGLANG_SIMULATE_ACC_LEN.get())
+        if (
+            self._simulate_acc_len > 0
+            and self._simulate_acc_len != 1.0
+            and not self._verify_planner.is_verify_all
+        ):
+            raise ValueError(
+                "SGLANG_SIMULATE_ACC_LEN>1.0 with DSpark requires a verify-all "
+                "schedule (SGLANG_RAGGED_VERIFY_MODE=static, or =compact with the "
+                "uninitialized/flat SPS table): a constant simulated correct_len>0 "
+                "can exceed a trimmed request's verify budget (cap-accept, or "
+                "compact with a profiled SPS table) and break the cutoff/cap "
+                "accounting. SGLANG_SIMULATE_ACC_LEN=1.0 yields correct_len=0 "
+                "(commit is the bonus token only), which stays within every verify "
+                "budget and is safe in any mode. Got mode="
+                f"{self._verify_planner.mode_value!r}, simulate_acc_len="
+                f"{self._simulate_acc_len}."
+            )
+
+        self._verify_executor = TargetVerifyExecutor(
+            target_worker=self.target_worker,
+            gamma=self.gamma,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            model_runner=self.model_runner,
+            kv_injector=self._kv_injector,
+            verify_epilogue=self._verify_epilogue,
+            simulate_acc_len=self._simulate_acc_len,
+        )
+
+        self._forced_budget_frac: Optional[float] = None
+
+        self._observers = DsparkStepObservers(
+            planner=self._verify_planner,
+            gamma=self.gamma,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            tp_rank=self.tp_rank,
+            device=self.device,
+            simulate_acc_len=self._simulate_acc_len,
+        )
+
+    def _resolve_target_embed_tokens(self, target_model):
+        if hasattr(target_model, "get_input_embeddings"):
+            return target_model.get_input_embeddings()
+        return target_model.model.get_input_embeddings()
+
+    @property
+    def carries_confidence(self) -> bool:
+        return self._verify_planner.carries_confidence
+
+    @property
+    def target_worker(self) -> TpModelWorker:
+        return self._target_worker
+
+    @property
+    def draft_worker(self):
+        return self._draft_worker
+
+    @property
+    def spec_v2_attn_backends(self) -> tuple:
+        return (
+            self._target_worker.model_runner.attn_backend,
+            self.draft_model_runner.attn_backend,
+        )
+
+    def __getattr__(self, name):
+        if name == "_target_worker":
+            raise AttributeError(name)
+        return getattr(self.target_worker, name)
+
+    def _draft_context(self):
+        if self._draft_dp_context_enabled:
+            return draft_tp_context(get_parallel().attn_tp_group)
+        return nullcontext()
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        self._draft_worker.alloc_memory_pool(
+            memory_pool_config=memory_pool_config,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+
+    def init_attention_backends(self):
+        with self._draft_context():
+            self._draft_worker.init_attention_backends()
+
+    def init_cuda_graphs(self):
+        capture_decode_cuda_graph = not self.server_args.disable_cuda_graph
+        if is_cuda() and capture_decode_cuda_graph:
+            available_mem = get_available_gpu_memory(self.device, self.gpu_id)
+            if available_mem < 1.0:
+                capture_decode_cuda_graph = False
+                logger.warning(
+                    "Disable DSpark draft cuda graph because only %.2f GB GPU "
+                    "memory is available after target backend initialization.",
+                    available_mem,
+                )
+        with self._draft_context():
+            if capture_decode_cuda_graph:
+                self._draft_sampler = self._maybe_build_draft_sampler()
+                if self._draft_sampler is not None:
+                    self.draft_model_runner.capture_tail_hooks.append(
+                        make_draft_sampler_capture_hook(self._draft_sampler)
+                    )
+                self._proposer.attach_draft_sampler(self._draft_sampler)
+            self._draft_worker.init_cuda_graphs(
+                capture_decode_cuda_graph=capture_decode_cuda_graph
+            )
+
+    def _maybe_build_draft_sampler(self):
+        return maybe_build_draft_sampler(
+            draft_model=self.draft_model,
+            gamma=self.gamma,
+            max_bs=max(self.server_args.cuda_graph_config.decode.bs),
+            device=self.device,
+            tp_rank=self.tp_rank,
+            confidence_fn=(
+                self._verify_planner.compute_confidence_tensor
+                if self._verify_planner.carries_confidence
+                else None
+            ),
+            out=(
+                self._verify_epilogue.draft_tokens_buf
+                if self._verify_epilogue is not None
+                else None
+            ),
+        )
+
+    def clear_cache_pool(self):
+        pass
+
+    def set_dspark_forced_budget_frac(self, frac: Optional[float]) -> None:
+        self._forced_budget_frac = frac
+        self._verify_planner.set_forced_budget_frac(frac)
+
+    def dump_info_records(self) -> Optional[dict]:
+        return self._observers.dump_info_records()
+
+    def clear_info_records(self) -> None:
+        self._observers.clear_info_records()
+
+    def block_accept_estimate_log_suffix(self) -> Optional[str]:
+        return self._observers.block_accept_estimate_log_suffix()
+
+    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        self._observers.note_request_finished(rid=rid, natural_stop=natural_stop)
+
+    def forward_batch_generation(
+        self,
+        batch: ScheduleBatch,
+        on_publish=None,
+    ) -> GenerationBatchResult:
+        if getattr(batch, "return_logprob", False):
+            raise ValueError(
+                "DSpark speculative decoding does not support return_logprob yet."
+            )
+
+        if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
+            self._verify_planner.note_non_decode_step()
+            self._observers.note_prefill_step()
+            return self._forward_prefill(batch, on_publish)
+
+        return self._forward_decode(batch, on_publish)
+
+    def _forward_prefill(
+        self, batch: ScheduleBatch, on_publish
+    ) -> GenerationBatchResult:
+        if batch.forward_mode.is_idle():
+            if self.server_args.enable_dp_attention:
+                batch.capture_hidden_mode = CaptureHiddenMode.FULL
+                self.target_worker.forward_batch_generation(batch)
+            return self._decode_idle_result(on_publish=on_publish)
+
+        batch.capture_hidden_mode = CaptureHiddenMode.FULL
+        batch_output = self.target_worker.forward_batch_generation(batch)
+        logits_output = batch_output.logits_output
+        next_token_ids = batch_output.next_token_ids
+        batch_output.new_seq_lens = 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(
+                "DSpark requires target aux hidden capture for prefill, but got None. "
+                "Make sure the target model has DFlash layers-to-capture configured."
+            )
+        if batch.extend_lens is None or batch.prefix_lens is None:
+            raise RuntimeError(
+                "DSpark expected extend_lens / prefix_lens in extend mode, got None."
+            )
+        if batch.out_cache_loc is None:
+            raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.")
+
+        device = next_token_ids.device
+        ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device)
+        draft_seq_lens = torch.tensor(
+            batch.prefix_lens, dtype=torch.int32, device=device
+        )
+        positions, _ = compute_position(
+            self.model_runner.server_args.attention_backend,
+            draft_seq_lens,
+            ctx_lens,
+            int(sum(batch.extend_lens)),
+        )
+        self._kv_injector.inject_target_hidden(
+            target_hidden=logits_output.hidden_states,
+            cache_loc=batch.out_cache_loc,
+            positions=positions,
+        )
+        logits_output.hidden_states = None
+
+        batch_output.next_draft_input = make_next_draft_input(
+            bonus_tokens=next_token_ids,
+            new_seq_lens=batch.seq_lens,
+        )
+        return batch_output
+
+    def _idle_verify_ragged_layout(self, batch: ScheduleBatch):
+        if batch.global_num_tokens is None or not self._verify_planner.is_compact_mode:
+            return None
+        global_bs = max(batch.global_num_tokens)
+        if global_bs <= 0:
+            return None
+        return idle_ragged_layout(
+            tier_num_reqs=global_bs,
+            dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch),
+            device=self.device,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            model_runner=self.model_runner,
+        )
+
+    def _dp_verify_tier_num_tokens(self, batch: ScheduleBatch) -> Optional[int]:
+        if not (
+            self._draft_is_moe
+            and self.server_args.enable_dp_attention
+            and batch.global_num_tokens is not None
+            and self._verify_planner.is_compact_mode
+        ):
+            return None
+        return dp_global_verify_tier_num_tokens(
+            global_tier_num_tokens=batch.global_spec_verify_tier_num_tokens
+        )
+
+    def _decode_idle_result(
+        self,
+        *,
+        on_publish,
+    ) -> GenerationBatchResult:
+        next_draft_input = make_next_draft_input(
+            bonus_tokens=torch.empty((0,), device=self.device, dtype=torch.int64),
+            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)
+        return GenerationBatchResult(
+            logits_output=None,
+            next_token_ids=torch.empty((0,), dtype=torch.int64, device=self.device),
+            accept_lens=torch.empty((0,), dtype=torch.int32, device=self.device),
+            block_accept_lens=torch.empty((0,), dtype=torch.int32, device=self.device),
+            next_draft_input=next_draft_input,
+            can_run_cuda_graph=False,
+            speculative_num_draft_tokens=int(self.verify_num_draft_tokens),
+            new_seq_lens=next_draft_input.new_seq_lens,
+        )
+
+    def _forward_decode(
+        self, batch: ScheduleBatch, on_publish
+    ) -> GenerationBatchResult:
+        if batch.spec_info is None:
+            batch.spec_info = DFlashDraftInputV2.create_idle_input(device=self.device)
+        draft_input = batch.spec_info
+        if not isinstance(draft_input, DFlashDraftInputV2):
+            raise RuntimeError(
+                "DSpark spec-v2 expected DFlashDraftInputV2 state on the running batch."
+            )
+
+        if batch.forward_mode.is_idle():
+            self._observers.note_idle_decode_step()
+            if self.server_args.enable_dp_attention:
+                if self._draft_is_moe:
+                    self._proposer.run_idle_participation(batch)
+                self._verify_executor.run_idle_participation(
+                    batch=batch, idle_layout=self._idle_verify_ragged_layout(batch)
+                )
+            return self._decode_idle_result(on_publish=on_publish)
+
+        batch.seq_lens.record_stream(
+            torch.get_device_module(self.device).current_stream()
+        )
+        bs = len(batch.seq_lens)
+        device = self.device
+        prefix_lens = batch.seq_lens
+
+        self._observers.begin_step()
+
+        target_model = self.target_worker.model_runner.model
+
+        verify_window = alloc_verify_window(
+            batch=batch,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=self.verify_num_draft_tokens,
+            block_pos_offsets=self._block_pos_offsets,
+            model_runner=self.model_runner,
+        )
+
+        sampling_info = batch.sampling_info
+        with self._draft_context(), self._observers.segment(InfoSegment.DRAFT):
+            proposal = self._proposer.propose(
+                batch=batch,
+                draft_input=draft_input,
+                verify_window=verify_window,
+                bs=bs,
+                device=device,
+                target_model=target_model,
+                sampling_info=sampling_info,
+            )
+        draft_block_ids = proposal.draft_block_ids
+        draft_block = proposal.draft_block
+        draft_tokens = draft_block.draft_tokens
+
+        confidence = proposal.confidence
+        if confidence is None:
+            confidence = self._verify_planner.compute_confidence_tensor(
+                draft_hidden=proposal.draft_hidden,
+                anchor_tokens=draft_block_ids[:, 0],
+                draft_tokens=draft_tokens,
+                confidence_tap=proposal.confidence_tap,
+            )
+
+        verify_token_budget = self._verify_planner.resolve_verify_token_budget(
+            draft_input=draft_input,
+            confidence=confidence,
+            prefix_lens=prefix_lens,
+            req_pool_indices=batch.req_pool_indices,
+        )
+
+        global_num_reqs = (
+            max(batch.global_num_tokens)
+            if self._draft_is_moe
+            and self.server_args.enable_dp_attention
+            and batch.global_num_tokens is not None
+            else None
+        )
+        layout = self._verify_planner.schedule_layout(
+            req_pool_indices=batch.req_pool_indices,
+            prefix_lens=prefix_lens,
+            device=device,
+            confidence=confidence,
+            budget=verify_token_budget,
+            global_num_reqs=global_num_reqs,
+            dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch),
+        )
+        run_compact = self._verify_planner.should_run_compact(layout=layout)
+
+        verify_ids_2d = torch.cat(
+            [draft_block_ids[:, :1], draft_tokens], dim=1
+        ).contiguous()
+
+        fold_eligible = (
+            self._verify_executor.verify_epilogue is not None
+            and proposal.folded
+            and verify_logits_adjustments_are_noop(sampling_info)
+            and self._simulate_acc_len <= 0
+        )
+        with self._observers.segment(InfoSegment.TARGET_VERIFY):
+            if run_compact:
+                target_verify, hidden_strided = self._verify_executor.run_compact(
+                    batch=batch,
+                    layout=layout,
+                    draft_block_ids=draft_block_ids,
+                    draft_tokens=draft_tokens,
+                    bs=bs,
+                    device=device,
+                    sampling_info=sampling_info,
+                    inject_gate=fold_eligible,
+                )
+            else:
+                target_verify = self._verify_executor.run_non_compact(
+                    batch=batch,
+                    draft_input=draft_input,
+                    verify_ids_2d=verify_ids_2d,
+                    verify_window=verify_window,
+                    sampling_info=sampling_info,
+                )
+                hidden_strided = None
+        logits_output = target_verify.logits_output
+        can_run_cuda_graph = target_verify.can_run_cuda_graph
+
+        epilogue = self._verify_executor.verify_epilogue
+        folded_accept = fold_eligible and run_compact and can_run_cuda_graph
+        accept = self._verify_executor.accept_and_finalize(
+            folded_accept=folded_accept,
+            bs=bs,
+            verify_ids_2d=verify_ids_2d,
+            target_logits=logits_output.next_token_logits,
+            draft_block=draft_block,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            layout=layout,
+            prefix_lens=prefix_lens,
+            draft_tokens=draft_tokens,
+        )
+        if on_publish is not None:
+            if confidence is not None:
+                on_publish(accept.new_seq_lens, confidence=confidence)
+            else:
+                on_publish(accept.new_seq_lens)
+
+        folded_commit = folded_accept and epilogue.folds_commit
+        if not folded_commit:
+            self._verify_executor.commit_hidden(
+                batch=batch,
+                layout=layout,
+                hidden_strided=hidden_strided,
+                verify_window=verify_window,
+                logits_output=logits_output,
+                commit_lens=accept.commit_lens,
+                bs=bs,
+                run_compact=run_compact,
+            )
+        logits_output.hidden_states = None
+
+        self._observers.observe_verify_step(
+            forward_ct=int(batch.forward_iter),
+            reqs=batch.reqs,
+            bs=bs,
+            proposal_folded=proposal.folded,
+            verify_ids_2d=verify_ids_2d,
+            target_logits=logits_output.next_token_logits,
+            layout=layout,
+            confidence=confidence,
+            prefix_lens=prefix_lens,
+            draft_tokens=draft_tokens,
+            draft_block=draft_block,
+            sampling_info=sampling_info,
+            correct_len=accept.correct_len,
+            cap_trim_lens=accept.cap_trim_lens,
+            bonus=accept.bonus,
+            commit_lens=accept.commit_lens,
+            verify_token_budget=verify_token_budget,
+            req_pool_indices=batch.req_pool_indices,
+            verify_tier_num_tokens=int(batch.spec_verify_tier_num_tokens),
+            dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch),
+        )
+
+        next_draft_input = make_next_draft_input(
+            bonus_tokens=accept.bonus,
+            new_seq_lens=accept.new_seq_lens,
+        )
+        return GenerationBatchResult(
+            logits_output=logits_output,
+            next_token_ids=accept.out_tokens.reshape(-1),
+            accept_lens=accept.commit_lens,
+            block_accept_lens=accept.commit_lens + accept.cap_trim_lens,
+            cap_lens=(
+                layout.verify_lens.to(torch.int32) if layout is not None else None
+            ),
+            can_run_cuda_graph=can_run_cuda_graph,
+            next_draft_input=next_draft_input,
+            speculative_num_draft_tokens=int(self.verify_num_draft_tokens),
+            new_seq_lens=accept.new_seq_lens,
+        )
+
+    def get_confidence_budget_prepare(self):
+        return self._verify_planner.confidence_budget_prepare()

python/sglang/srt/speculative/dspark_components/kernels/dispatch.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import torch
+
+
+def inputs_on_cuda(*args, **kwargs) -> bool:
+    """Route kernel dispatch by input placement: the first tensor argument
+    decides. CUDA inputs take the fused triton kernel; CPU inputs take the
+    torch reference implementation (triton is CUDA-only, and CPU-side callers
+    such as unit tests exercise the reference path)."""
+    for value in (*args, *kwargs.values()):
+        if isinstance(value, torch.Tensor):
+            return value.is_cuda
+    raise AssertionError("kernel dispatch requires at least one tensor argument")

python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py

@@ -0,0 +1,862 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+from sglang.srt.speculative.dflash_utils import (
+    _get_or_create_chain_verify_buffers,
+    build_dflash_verify_target_probs,
+    compute_dflash_correct_drafts_and_bonus,
+)
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.speculative.reject_sampling import chain_speculative_sampling_triton
+
+
+class AcceptSampling:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        draft_probs: torch.Tensor,
+        sampling_info,
+        draft_input: DFlashDraftInputV2,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_sampling(
+            candidates=candidates,
+            target_logits=target_logits,
+            draft_probs=draft_probs,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        draft_probs: torch.Tensor,
+        sampling_info,
+        draft_input: DFlashDraftInputV2,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_sampling_triton(
+            candidates=candidates,
+            target_logits=target_logits,
+            draft_probs=draft_probs,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+
+def _accept_sampling_core(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    device = candidates.device
+    if not sampling_info.need_top_k_sampling and not sampling_info.need_top_p_sampling:
+        target_probs = SoftmaxTemp.execute(
+            logits=target_logits,
+            temperatures=sampling_info.temperatures,
+            rows_per_request=verify_num_draft_tokens,
+        ).view(bs, verify_num_draft_tokens, -1)
+    else:
+        target_probs = build_dflash_verify_target_probs(
+            next_token_logits=target_logits,
+            sampling_info=sampling_info,
+            draft_token_num=verify_num_draft_tokens,
+            bs=bs,
+            max_top_k=draft_input.max_top_k,
+            uniform_top_k_value=draft_input.uniform_top_k_value,
+        )
+    (
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        predicts,
+        accept_index,
+        accept_token_num,
+    ) = _get_or_create_chain_verify_buffers(
+        bs=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    uniform_samples = torch.rand((bs, gamma), dtype=torch.float32, device=device)
+    uniform_samples_final = torch.rand((bs,), dtype=torch.float32, device=device)
+    chain_speculative_sampling_triton(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        retrive_index=retrieve_index,
+        retrive_next_token=retrieve_next_token,
+        retrive_next_sibling=retrieve_next_sibling,
+        uniform_samples=uniform_samples,
+        uniform_samples_for_final_sampling=uniform_samples_final,
+        target_probs=target_probs,
+        draft_probs=draft_probs,
+        threshold_single=1.0,
+        threshold_acc=1.0,
+        deterministic=True,
+    )
+    correct_len = accept_token_num
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+    else:
+        cap_trim_lens = torch.zeros_like(correct_len)
+    return correct_len, cap_trim_lens, accept_index, predicts
+
+
+def accept_sampling(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    device = candidates.device
+    correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core(
+        candidates=candidates,
+        target_logits=target_logits,
+        draft_probs=draft_probs,
+        sampling_info=sampling_info,
+        draft_input=draft_input,
+        gamma=gamma,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    row_ids = torch.arange(bs, dtype=torch.long, device=device)
+    accept_pos = accept_index[row_ids, correct_len.to(torch.long)].to(torch.long)
+    bonus = predicts[accept_pos].to(torch.int64)
+    return correct_len, bonus, cap_trim_lens
+
+
+@triton.jit
+def _gather_two_level_bonus_kernel(
+    accept_index_ptr,
+    predicts_ptr,
+    correct_len_ptr,
+    out_ptr,
+    cols,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64)
+    accept_pos = tl.load(accept_index_ptr + offs * cols + cl, mask=mask, other=0).to(
+        tl.int64
+    )
+    bonus = tl.load(predicts_ptr + accept_pos, mask=mask, other=0)
+    tl.store(out_ptr + offs, bonus.to(tl.int64), mask=mask)
+
+
+def gather_two_level_bonus_triton(
+    *,
+    accept_index: torch.Tensor,
+    predicts: torch.Tensor,
+    correct_len: torch.Tensor,
+) -> torch.Tensor:
+    bs, cols = accept_index.shape
+    accept_index = accept_index.contiguous()
+    predicts = predicts.contiguous()
+    correct_len = correct_len.contiguous()
+    out = torch.empty(bs, dtype=torch.int64, device=accept_index.device)
+    BLOCK = 256
+    grid = (triton.cdiv(bs, BLOCK),)
+    _gather_two_level_bonus_kernel[grid](
+        accept_index, predicts, correct_len, out, cols, bs, BLOCK=BLOCK
+    )
+    return out
+
+
+def accept_sampling_triton(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core(
+        candidates=candidates,
+        target_logits=target_logits,
+        draft_probs=draft_probs,
+        sampling_info=sampling_info,
+        draft_input=draft_input,
+        gamma=gamma,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    bonus = gather_two_level_bonus_triton(
+        accept_index=accept_index, predicts=predicts, correct_len=correct_len
+    )
+    return correct_len, bonus, cap_trim_lens
+
+
+try:
+    from flashinfer.sampling import softmax as _flashinfer_softmax
+except ImportError:
+    _flashinfer_softmax = None
+
+
+class SoftmaxTemp:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if not inputs_on_cuda(*args, **kwargs):
+            return cls.torch(*args, **kwargs)
+        if _flashinfer_softmax is not None:
+            return cls.flashinfer(*args, **kwargs)
+        return cls.triton(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp_triton(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+    @classmethod
+    def flashinfer(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp_flashinfer(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+
+def softmax_temp(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    num_rows = logits.shape[0]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temp_per_row = torch.repeat_interleave(
+        temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0
+    )
+    scaled = logits.to(torch.float32) / temp_per_row[:, None]
+    return torch.softmax(scaled, dim=-1)
+
+
+@triton.jit
+def _softmax_temp_kernel(
+    logits_ptr,
+    temp_ptr,
+    out_ptr,
+    vocab,
+    rows_per_request,
+    logits_row_stride,
+    BLOCK_V: tl.constexpr,
+):
+    row = tl.program_id(0)
+    temp = tl.load(temp_ptr + row // rows_per_request).to(tl.float32)
+    base = logits_ptr + row.to(tl.int64) * logits_row_stride
+    out_base = out_ptr + row.to(tl.int64) * vocab
+
+    row_max = -float("inf")
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        row_max = tl.maximum(row_max, tl.max(x, axis=0))
+
+    sum_exp = 0.0
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        e = tl.exp(x - row_max)
+        e = tl.where(vmask, e, 0.0)
+        sum_exp += tl.sum(e, axis=0)
+
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        e = tl.exp(x - row_max)
+        tl.store(out_base + offs, e / sum_exp, mask=vmask)
+
+
+def softmax_temp_triton(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    num_rows, vocab = logits.shape[0], logits.shape[-1]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temperatures = temperatures.reshape(bs).to(torch.float32).contiguous()
+    out = torch.empty((num_rows, vocab), dtype=torch.float32, device=logits.device)
+    BLOCK_V = 4096
+    _softmax_temp_kernel[(num_rows,)](
+        logits,
+        temperatures,
+        out,
+        vocab,
+        rows_per_request,
+        logits.stride(0),
+        BLOCK_V=BLOCK_V,
+    )
+    return out
+
+
+def softmax_temp_flashinfer(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    if _flashinfer_softmax is None:
+        raise RuntimeError(
+            "softmax_temp_flashinfer requires flashinfer.sampling.softmax, "
+            "which is unavailable in this environment"
+        )
+    num_rows, vocab = logits.shape[0], logits.shape[-1]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temp_per_row = torch.repeat_interleave(
+        temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0
+    ).contiguous()
+    logits_2d = logits.to(torch.float32).contiguous()
+    return _flashinfer_softmax(logits=logits_2d, temperature=temp_per_row)
+
+
+class MixedAcceptSelectResult(msgspec.Struct):
+    correct_len: torch.Tensor
+    bonus: torch.Tensor
+    cap_trim_lens: torch.Tensor
+
+
+class SelectMixedAccept:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> MixedAcceptSelectResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        greedy_mask: torch.Tensor,
+        greedy_len: torch.Tensor,
+        greedy_bonus: torch.Tensor,
+        greedy_trim: torch.Tensor,
+        sampling_len: torch.Tensor,
+        sampling_bonus: torch.Tensor,
+        sampling_trim: torch.Tensor,
+    ) -> MixedAcceptSelectResult:
+        return select_mixed_accept(
+            greedy_mask=greedy_mask,
+            greedy_len=greedy_len,
+            greedy_bonus=greedy_bonus,
+            greedy_trim=greedy_trim,
+            sampling_len=sampling_len,
+            sampling_bonus=sampling_bonus,
+            sampling_trim=sampling_trim,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        greedy_mask: torch.Tensor,
+        greedy_len: torch.Tensor,
+        greedy_bonus: torch.Tensor,
+        greedy_trim: torch.Tensor,
+        sampling_len: torch.Tensor,
+        sampling_bonus: torch.Tensor,
+        sampling_trim: torch.Tensor,
+    ) -> MixedAcceptSelectResult:
+        return select_mixed_accept_triton(
+            greedy_mask=greedy_mask,
+            greedy_len=greedy_len,
+            greedy_bonus=greedy_bonus,
+            greedy_trim=greedy_trim,
+            sampling_len=sampling_len,
+            sampling_bonus=sampling_bonus,
+            sampling_trim=sampling_trim,
+        )
+
+
+def select_mixed_accept(
+    *,
+    greedy_mask: torch.Tensor,
+    greedy_len: torch.Tensor,
+    greedy_bonus: torch.Tensor,
+    greedy_trim: torch.Tensor,
+    sampling_len: torch.Tensor,
+    sampling_bonus: torch.Tensor,
+    sampling_trim: torch.Tensor,
+) -> MixedAcceptSelectResult:
+    correct_len = torch.where(
+        greedy_mask, greedy_len.to(sampling_len.dtype), sampling_len
+    )
+    bonus = torch.where(greedy_mask, greedy_bonus, sampling_bonus)
+    cap_trim_lens = torch.where(
+        greedy_mask, greedy_trim.to(sampling_trim.dtype), sampling_trim
+    )
+    return MixedAcceptSelectResult(
+        correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens
+    )
+
+
+@triton.jit
+def _mixed_accept_select_kernel(
+    greedy_mask_ptr,
+    greedy_len_ptr,
+    greedy_bonus_ptr,
+    greedy_trim_ptr,
+    sampling_len_ptr,
+    sampling_bonus_ptr,
+    sampling_trim_ptr,
+    correct_len_ptr,
+    bonus_ptr,
+    cap_trim_ptr,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    is_greedy = tl.load(greedy_mask_ptr + offs, mask=mask, other=0) != 0
+
+    g_len = tl.load(greedy_len_ptr + offs, mask=mask, other=0)
+    s_len = tl.load(sampling_len_ptr + offs, mask=mask, other=0)
+    tl.store(correct_len_ptr + offs, tl.where(is_greedy, g_len, s_len), mask=mask)
+
+    g_bonus = tl.load(greedy_bonus_ptr + offs, mask=mask, other=0)
+    s_bonus = tl.load(sampling_bonus_ptr + offs, mask=mask, other=0)
+    tl.store(bonus_ptr + offs, tl.where(is_greedy, g_bonus, s_bonus), mask=mask)
+
+    g_trim = tl.load(greedy_trim_ptr + offs, mask=mask, other=0)
+    s_trim = tl.load(sampling_trim_ptr + offs, mask=mask, other=0)
+    tl.store(cap_trim_ptr + offs, tl.where(is_greedy, g_trim, s_trim), mask=mask)
+
+
+def select_mixed_accept_triton(
+    *,
+    greedy_mask: torch.Tensor,
+    greedy_len: torch.Tensor,
+    greedy_bonus: torch.Tensor,
+    greedy_trim: torch.Tensor,
+    sampling_len: torch.Tensor,
+    sampling_bonus: torch.Tensor,
+    sampling_trim: torch.Tensor,
+) -> MixedAcceptSelectResult:
+    bs = greedy_mask.shape[0]
+    device = greedy_mask.device
+
+    correct_len = torch.empty(bs, dtype=sampling_len.dtype, device=device)
+    bonus = torch.empty(bs, dtype=sampling_bonus.dtype, device=device)
+    cap_trim_lens = torch.empty(bs, dtype=sampling_trim.dtype, device=device)
+    BLOCK = 256
+    _mixed_accept_select_kernel[(triton.cdiv(bs, BLOCK),)](
+        greedy_mask,
+        greedy_len,
+        greedy_bonus,
+        greedy_trim,
+        sampling_len,
+        sampling_bonus,
+        sampling_trim,
+        correct_len,
+        bonus,
+        cap_trim_lens,
+        bs,
+        BLOCK=BLOCK,
+    )
+    return MixedAcceptSelectResult(
+        correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens
+    )
+
+
+class AcceptGreedy:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_greedy(
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_greedy_triton(
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+
+def accept_greedy(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    target_predict = torch.argmax(target_logits, dim=-1).view(
+        bs, verify_num_draft_tokens
+    )
+    correct_len, bonus = compute_dflash_correct_drafts_and_bonus(
+        candidates=candidates,
+        target_predict=target_predict,
+    )
+    cap_trim_lens = torch.zeros_like(correct_len)
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+        row_ids = torch.arange(bs, device=target_predict.device)
+        bonus = target_predict[row_ids, correct_len.to(torch.long)].to(torch.int64)
+    return correct_len, bonus, cap_trim_lens
+
+
+@triton.jit
+def _gather_row_bonus_kernel(
+    table_ptr,
+    idx_ptr,
+    out_ptr,
+    cols,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    idx = tl.load(idx_ptr + offs, mask=mask, other=0).to(tl.int64)
+    val = tl.load(table_ptr + offs * cols + idx, mask=mask, other=0)
+    tl.store(out_ptr + offs, val.to(tl.int64), mask=mask)
+
+
+def gather_row_bonus_triton(*, table: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
+    bs, cols = table.shape
+    table = table.contiguous()
+    idx = idx.contiguous()
+    out = torch.empty(bs, dtype=torch.int64, device=table.device)
+    BLOCK = 256
+    grid = (triton.cdiv(bs, BLOCK),)
+    _gather_row_bonus_kernel[grid](table, idx, out, cols, bs, BLOCK=BLOCK)
+    return out
+
+
+def accept_greedy_triton(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    target_predict = torch.argmax(target_logits, dim=-1).view(
+        bs, verify_num_draft_tokens
+    )
+    correct_len, bonus = compute_dflash_correct_drafts_and_bonus(
+        candidates=candidates,
+        target_predict=target_predict,
+    )
+    cap_trim_lens = torch.zeros_like(correct_len)
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+        bonus = gather_row_bonus_triton(table=target_predict, idx=correct_len)
+    return correct_len, bonus, cap_trim_lens
+
+
+class FinalizeAcceptLensResult(msgspec.Struct):
+    commit_lens: torch.Tensor
+    new_seq_lens: torch.Tensor
+    cap_trim_lens: torch.Tensor
+
+
+class FinalizeAcceptLens:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> FinalizeAcceptLensResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+    ) -> FinalizeAcceptLensResult:
+        return finalize_accept_lens(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=prefix_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+    ) -> FinalizeAcceptLensResult:
+        return finalize_accept_lens_triton(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=prefix_lens,
+        )
+
+
+def finalize_accept_lens(
+    *,
+    correct_len: torch.Tensor,
+    cap_trim_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> FinalizeAcceptLensResult:
+    commit_lens = correct_len.to(torch.int32) + 1
+    new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)
+    return FinalizeAcceptLensResult(
+        commit_lens=commit_lens,
+        new_seq_lens=new_seq_lens,
+        cap_trim_lens=cap_trim_lens.to(torch.int32),
+    )
+
+
+@triton.jit
+def _finalize_accept_lens_kernel(
+    correct_len_ptr,
+    cap_trim_ptr,
+    prefix_lens_ptr,
+    commit_lens_ptr,
+    new_seq_lens_ptr,
+    cap_trim_out_ptr,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    commit = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int32) + 1
+    prefix = tl.load(prefix_lens_ptr + offs, mask=mask, other=0)
+    trim = tl.load(cap_trim_ptr + offs, mask=mask, other=0).to(tl.int32)
+    tl.store(commit_lens_ptr + offs, commit, mask=mask)
+    tl.store(new_seq_lens_ptr + offs, prefix + commit, mask=mask)
+    tl.store(cap_trim_out_ptr + offs, trim, mask=mask)
+
+
+def finalize_accept_lens_triton(
+    *,
+    correct_len: torch.Tensor,
+    cap_trim_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> FinalizeAcceptLensResult:
+    bs = correct_len.shape[0]
+    device = correct_len.device
+
+    commit_lens = torch.empty(bs, dtype=torch.int32, device=device)
+    new_seq_lens = torch.empty(bs, dtype=prefix_lens.dtype, device=device)
+    cap_trim_out = torch.empty(bs, dtype=torch.int32, device=device)
+    BLOCK = 256
+    _finalize_accept_lens_kernel[(triton.cdiv(bs, BLOCK),)](
+        correct_len,
+        cap_trim_lens,
+        prefix_lens,
+        commit_lens,
+        new_seq_lens,
+        cap_trim_out,
+        bs,
+        BLOCK=BLOCK,
+    )
+    return FinalizeAcceptLensResult(
+        commit_lens=commit_lens,
+        new_seq_lens=new_seq_lens,
+        cap_trim_lens=cap_trim_out,
+    )
+
+
+class CapCorrectLen:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        verify_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return cap_correct_len(
+            correct_len=correct_len,
+            verify_lens=verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        verify_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return cap_correct_len_triton(
+            correct_len=correct_len,
+            verify_lens=verify_lens,
+        )
+
+
+def cap_correct_len(
+    *,
+    correct_len: torch.Tensor,
+    verify_lens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ell_r = (verify_lens.to(device=correct_len.device) - 1).to(correct_len.dtype)
+    capped = torch.minimum(correct_len, ell_r)
+    cap_trim_lens = correct_len - capped
+    return capped, cap_trim_lens
+
+
+@triton.jit
+def _cap_correct_len_kernel(
+    correct_len_ptr,
+    verify_lens_ptr,
+    capped_ptr,
+    trim_ptr,
+    n,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64)
+    vl = tl.load(verify_lens_ptr + offs, mask=mask, other=0).to(tl.int64)
+    ell = vl - 1
+    capped = tl.minimum(cl, ell)
+    trim = cl - capped
+    tl.store(capped_ptr + offs, capped, mask=mask)
+    tl.store(trim_ptr + offs, trim, mask=mask)
+
+
+def cap_correct_len_triton(
+    *,
+    correct_len: torch.Tensor,
+    verify_lens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    device = correct_len.device
+    correct_len = correct_len.contiguous()
+    verify_lens = verify_lens.to(device=device).contiguous()
+    n = correct_len.shape[0]
+    capped = torch.empty_like(correct_len)
+    trim = torch.empty_like(correct_len)
+    BLOCK = 1024
+    grid = (triton.cdiv(n, BLOCK),)
+    _cap_correct_len_kernel[grid](
+        correct_len, verify_lens, capped, trim, n, BLOCK=BLOCK
+    )
+    return capped, trim

python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py

@@ -0,0 +1,491 @@
+from __future__ import annotations
+
+from typing import Tuple
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.utils import ceil_align
+
+
+class DsparkWindowGather(msgspec.Struct, frozen=True):
+    num_q: int
+    bs: int
+    context_lens: torch.Tensor
+    req_pool_indices_per_request: torch.Tensor
+    offsets: torch.Tensor
+    invalid: torch.Tensor
+
+
+class ComputeDsparkWindowGather:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> DsparkWindowGather:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        seq_lens_casual: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+    ) -> DsparkWindowGather:
+        return compute_dspark_window_gather(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            block_size=block_size,
+            swa_window=swa_window,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        seq_lens_casual: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+    ) -> DsparkWindowGather:
+        return compute_dspark_window_gather_triton(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            block_size=block_size,
+            swa_window=swa_window,
+        )
+
+
+class BuildDsparkSwaPageIndices:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        req_to_token: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        req_pool_indices_per_request: torch.Tensor,
+        offsets: torch.Tensor,
+        invalid: torch.Tensor,
+        out_loc: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return build_dspark_swa_page_indices(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_per_request=req_pool_indices_per_request,
+            offsets=offsets,
+            invalid=invalid,
+            out_loc=out_loc,
+            context_lens=context_lens,
+            block_size=block_size,
+            swa_window=swa_window,
+            page_index_aligned_size=page_index_aligned_size,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_to_token: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        req_pool_indices_per_request: torch.Tensor,
+        offsets: torch.Tensor,
+        invalid: torch.Tensor,
+        out_loc: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return build_dspark_swa_page_indices_triton(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_per_request=req_pool_indices_per_request,
+            offsets=offsets,
+            out_loc=out_loc,
+            context_lens=context_lens,
+            block_size=block_size,
+            swa_window=swa_window,
+            page_index_aligned_size=page_index_aligned_size,
+        )
+
+
+def compute_dspark_window_gather(
+    *,
+    seq_lens_casual: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+) -> DsparkWindowGather:
+    seq_lens_casual = seq_lens_casual.to(torch.int32)
+    num_q = seq_lens_casual.size(0)
+    assert num_q % block_size == 0, (
+        f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not "
+        f"divisible by block_size={block_size}."
+    )
+    bs = num_q // block_size
+    device = seq_lens_casual.device
+
+    first_token = torch.arange(bs, device=device, dtype=torch.int64) * block_size
+    prefix_lens = (seq_lens_casual[first_token] - 1).to(torch.int32)
+    context_lens = torch.clamp(prefix_lens, max=swa_window).to(torch.int32)
+    req_pool_indices_per_request = req_pool_indices_repeated[first_token]
+
+    offsets = (
+        prefix_lens.to(torch.int64).unsqueeze(1)
+        - swa_window
+        + torch.arange(swa_window, device=device, dtype=torch.int64).unsqueeze(0)
+    )
+    invalid = offsets < 0
+    offsets = offsets.clamp(min=0)
+
+    return DsparkWindowGather(
+        num_q=num_q,
+        bs=bs,
+        context_lens=context_lens,
+        req_pool_indices_per_request=req_pool_indices_per_request,
+        offsets=offsets,
+        invalid=invalid,
+    )
+
+
+@triton.jit
+def _window_gather_kernel(
+    seq_lens_casual_ptr,
+    req_pool_rep_ptr,
+    context_lens_ptr,
+    req_pool_out_ptr,
+    offsets_ptr,
+    invalid_ptr,
+    block_size,
+    swa_window,
+    W_BLOCK: tl.constexpr,
+):
+    i = tl.program_id(0)
+    ft = i * block_size
+    prefix = tl.load(seq_lens_casual_ptr + ft).to(tl.int64) - 1
+    tl.store(context_lens_ptr + i, tl.minimum(prefix, swa_window).to(tl.int32))
+    tl.store(req_pool_out_ptr + i, tl.load(req_pool_rep_ptr + ft))
+    col = tl.arange(0, W_BLOCK)
+    cmask = col < swa_window
+    off = prefix - swa_window + col
+    tl.store(invalid_ptr + i * swa_window + col, off < 0, mask=cmask)
+    tl.store(offsets_ptr + i * swa_window + col, tl.maximum(off, 0), mask=cmask)
+
+
+def compute_dspark_window_gather_triton(
+    *,
+    seq_lens_casual: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+) -> DsparkWindowGather:
+    seq_lens_casual = seq_lens_casual.to(torch.int32).contiguous()
+    num_q = seq_lens_casual.size(0)
+    assert num_q % block_size == 0, (
+        f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not "
+        f"divisible by block_size={block_size}."
+    )
+    bs = num_q // block_size
+    device = seq_lens_casual.device
+    req_pool_indices_repeated = req_pool_indices_repeated.to(device=device).contiguous()
+    context_lens = torch.empty(bs, dtype=torch.int32, device=device)
+    req_pool_out = torch.empty(bs, dtype=req_pool_indices_repeated.dtype, device=device)
+    offsets = torch.empty((bs, swa_window), dtype=torch.int64, device=device)
+    invalid = torch.empty((bs, swa_window), dtype=torch.bool, device=device)
+    W_BLOCK = triton.next_power_of_2(swa_window)
+    _window_gather_kernel[(bs,)](
+        seq_lens_casual,
+        req_pool_indices_repeated,
+        context_lens,
+        req_pool_out,
+        offsets,
+        invalid,
+        block_size,
+        swa_window,
+        W_BLOCK=W_BLOCK,
+    )
+    return DsparkWindowGather(
+        num_q=num_q,
+        bs=bs,
+        context_lens=context_lens,
+        req_pool_indices_per_request=req_pool_out,
+        offsets=offsets,
+        invalid=invalid,
+    )
+
+
+def build_dspark_swa_page_indices(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_per_request: torch.Tensor,
+    offsets: torch.Tensor,
+    invalid: torch.Tensor,
+    out_loc: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if offsets.ndim != 2 or offsets.shape[1] != swa_window:
+        raise ValueError(
+            "offsets must be [bs, swa_window]; "
+            f"got shape={tuple(offsets.shape)} (swa_window={swa_window})."
+        )
+    bs = offsets.shape[0]
+    device = offsets.device
+    context_lens = context_lens.to(device=device, dtype=torch.int32)
+
+    window_full_locs = req_to_token[
+        req_pool_indices_per_request[:, None].to(torch.int64), offsets
+    ]
+    window_full_locs = window_full_locs.masked_fill(invalid, 0)
+    window_swa_locs = full_to_swa_mapping[window_full_locs].to(torch.int32)
+    window_swa_locs = window_swa_locs.masked_fill(invalid, -1)
+
+    block_full_locs = out_loc[: bs * block_size].view(bs, block_size)
+    block_swa_locs = full_to_swa_mapping[block_full_locs].to(torch.int32)
+
+    target_width = ceil_align(swa_window + block_size, page_index_aligned_size)
+
+    swa_page_indices = _compact_dspark_window_then_block(
+        window_swa_locs=window_swa_locs,
+        block_swa_locs=block_swa_locs,
+        context_lens=context_lens,
+        target_width=target_width,
+        block_size=block_size,
+        swa_window=swa_window,
+    )
+
+    swa_page_indices = (
+        swa_page_indices.view(bs, 1, target_width)
+        .expand(bs, block_size, target_width)
+        .reshape(bs * block_size, target_width)
+        .contiguous()
+    )
+    swa_topk_lengths = (
+        (context_lens + block_size)
+        .view(bs, 1)
+        .expand(bs, block_size)
+        .reshape(bs * block_size)
+        .contiguous()
+        .to(torch.int32)
+    )
+    return swa_page_indices, swa_topk_lengths
+
+
+def _compact_dspark_window_then_block(
+    *,
+    window_swa_locs: torch.Tensor,
+    block_swa_locs: torch.Tensor,
+    context_lens: torch.Tensor,
+    target_width: int,
+    block_size: int,
+    swa_window: int,
+) -> torch.Tensor:
+    bs = window_swa_locs.shape[0]
+    device = window_swa_locs.device
+    out = torch.full((bs, target_width), -1, dtype=torch.int32, device=device)
+
+    j = torch.arange(swa_window, device=device, dtype=torch.int32).view(1, -1)
+    shift = (swa_window - context_lens.view(-1, 1)).to(torch.int32)
+    src_col = (shift + j).clamp_(min=0, max=swa_window - 1).to(torch.int64)
+    gathered = torch.gather(window_swa_locs, dim=1, index=src_col)
+    valid = j < context_lens.view(-1, 1)
+    out[:, :swa_window] = torch.where(valid, gathered, -1)
+
+    block_col = context_lens.view(-1, 1) + torch.arange(
+        block_size, device=device, dtype=torch.int32
+    ).view(1, -1)
+    block_rows = torch.arange(bs, device=device).view(-1, 1).expand(-1, block_size)
+    out[block_rows, block_col] = block_swa_locs
+    return out
+
+
+@triton.jit
+def _swa_page_indices_kernel(
+    req_to_token_ptr,
+    full_to_swa_ptr,
+    req_pool_ptr,
+    offsets_ptr,
+    out_loc_ptr,
+    context_lens_ptr,
+    out_ptr,
+    topk_ptr,
+    rt_stride,
+    swa_window,
+    block_size,
+    target_width,
+    TW_BLOCK: tl.constexpr,
+):
+    q = tl.program_id(0)
+    i = q // block_size
+    cl = tl.load(context_lens_ptr + i)
+    rp = tl.load(req_pool_ptr + i).to(tl.int64)
+    k = tl.arange(0, TW_BLOCK)
+    kmask = k < target_width
+    in_window = k < cl
+    src_col = tl.minimum(tl.maximum((swa_window - cl) + k, 0), swa_window - 1)
+    wmask = kmask & in_window
+    off = tl.load(offsets_ptr + i * swa_window + src_col, mask=wmask, other=0).to(
+        tl.int64
+    )
+    win_full = tl.load(req_to_token_ptr + rp * rt_stride + off, mask=wmask, other=0).to(
+        tl.int64
+    )
+    win_swa = tl.load(full_to_swa_ptr + win_full, mask=wmask, other=-1).to(tl.int32)
+
+    in_block = (k >= cl) & (k < cl + block_size)
+    bmask = kmask & in_block
+    bcol = tl.maximum(k - cl, 0)
+    blk_full = tl.load(out_loc_ptr + i * block_size + bcol, mask=bmask, other=0).to(
+        tl.int64
+    )
+    blk_swa = tl.load(full_to_swa_ptr + blk_full, mask=bmask, other=-1).to(tl.int32)
+
+    val = tl.where(in_window, win_swa, tl.where(in_block, blk_swa, -1))
+    tl.store(out_ptr + q * target_width + k, val.to(tl.int32), mask=kmask)
+    tl.store(topk_ptr + q, (cl + block_size).to(tl.int32))
+
+
+def build_dspark_swa_page_indices_triton(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_per_request: torch.Tensor,
+    offsets: torch.Tensor,
+    out_loc: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if offsets.ndim != 2 or offsets.shape[1] != swa_window:
+        raise ValueError(
+            "offsets must be [bs, swa_window]; "
+            f"got shape={tuple(offsets.shape)} (swa_window={swa_window})."
+        )
+    bs = offsets.shape[0]
+    device = offsets.device
+    req_pool = req_pool_indices_per_request.to(device=device).contiguous()
+    offsets = offsets.to(torch.int64).contiguous()
+    out_loc = out_loc[: bs * block_size].contiguous()
+    context_lens = context_lens.to(device=device, dtype=torch.int32).contiguous()
+    rt_stride = req_to_token.stride(0)
+    target_width = ceil_align(swa_window + block_size, page_index_aligned_size)
+    n_q = bs * block_size
+    swa_page_indices = torch.empty(
+        (n_q, target_width), dtype=torch.int32, device=device
+    )
+    swa_topk_lengths = torch.empty(n_q, dtype=torch.int32, device=device)
+    TW_BLOCK = triton.next_power_of_2(target_width)
+    _swa_page_indices_kernel[(n_q,)](
+        req_to_token,
+        full_to_swa_mapping,
+        req_pool,
+        offsets,
+        out_loc,
+        context_lens,
+        swa_page_indices,
+        swa_topk_lengths,
+        rt_stride,
+        swa_window,
+        block_size,
+        target_width,
+        TW_BLOCK=TW_BLOCK,
+    )
+    return swa_page_indices, swa_topk_lengths
+
+
+class BuildBlockSeqLensCausal:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        seq_lens: torch.Tensor,
+        block_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        return build_block_seq_lens_causal(
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        seq_lens: torch.Tensor,
+        block_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        return build_block_seq_lens_causal_triton(
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=device,
+        )
+
+
+def build_block_seq_lens_causal(
+    *,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    prefix = seq_lens.to(torch.int32)
+    steps = torch.arange(1, block_size + 1, device=device, dtype=torch.int32)
+    return (prefix[:, None] + steps[None, :]).reshape(-1)
+
+
+@triton.jit
+def _block_seq_lens_casual_kernel(
+    seq_lens_ptr,
+    out_ptr,
+    block_size,
+    n_out,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_out
+    row = offs // block_size
+    col = offs % block_size
+    prefix = tl.load(seq_lens_ptr + row, mask=mask, other=0)
+    tl.store(out_ptr + offs, (prefix + col + 1).to(tl.int32), mask=mask)
+
+
+def build_block_seq_lens_causal_triton(
+    *,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    seq_lens = seq_lens.to(device=device, dtype=torch.int64).contiguous()
+    n_rows = seq_lens.shape[0]
+    n_out = n_rows * block_size
+    out = torch.empty(n_out, dtype=torch.int32, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(n_out, BLOCK),)
+    _block_seq_lens_casual_kernel[grid](seq_lens, out, block_size, n_out, BLOCK=BLOCK)
+    return out

python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py

@@ -0,0 +1,443 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+
+_BLOCK_V = 1024
+_IDX_SENTINEL = tl.constexpr(2147483647)
+
+
+class SampleStepTokens:
+    @classmethod
+    def execute(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        if step_logits.is_cuda:
+            return cls.triton(
+                step_logits=step_logits,
+                temperatures=temperatures,
+                greedy_mask=greedy_mask,
+                exp_noise=exp_noise,
+            )
+        return cls.torch(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        return sample_step_tokens(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        return sample_step_tokens_triton(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+
+def sample_step_tokens(
+    *,
+    step_logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    greedy_mask: torch.Tensor,
+    exp_noise: torch.Tensor,
+) -> torch.Tensor:
+    probs = torch.softmax(step_logits.float() / temperatures[:, None], dim=-1)
+    noise = torch.where(greedy_mask[:, None], 1.0, exp_noise)
+    return probs.div_(noise).argmax(dim=-1)
+
+
+@triton.jit
+def _online_partial_kernel(
+    logits_ptr,
+    temperatures_ptr,
+    greedy_mask_ptr,
+    exp_noise_ptr,
+    tile_max_ptr,
+    partial_key_ptr,
+    partial_idx_ptr,
+    V,
+    stride_row,
+    n_tiles,
+    BLOCK_V: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    offs = tile * BLOCK_V + tl.arange(0, BLOCK_V)
+    mask = offs < V
+    logits = tl.load(
+        logits_ptr + row * stride_row + offs, mask=mask, other=float("-inf")
+    ).to(tl.float32)
+    temperature = tl.load(temperatures_ptr + row)
+    s = logits / temperature
+    tile_max = tl.max(s, axis=0)
+    greedy = tl.load(greedy_mask_ptr + row) != 0
+    noise = tl.load(exp_noise_ptr + row * V + offs, mask=mask, other=1.0)
+    denom = tl.where(greedy, 1.0, noise)
+    key = tl.exp(s - tile_max) / denom
+    key = tl.where(mask, key, -1.0)
+    tile_best = tl.max(key, axis=0)
+    idx = tl.where(key == tile_best, offs, _IDX_SENTINEL)
+    tl.store(tile_max_ptr + row * n_tiles + tile, tile_max)
+    tl.store(partial_key_ptr + row * n_tiles + tile, tile_best)
+    tl.store(partial_idx_ptr + row * n_tiles + tile, tl.min(idx, axis=0))
+
+
+@triton.jit
+def _online_combine_kernel(
+    tile_max_ptr,
+    partial_key_ptr,
+    partial_idx_ptr,
+    next_tokens_ptr,
+    n_tiles,
+    BLOCK_TILES: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_TILES)
+    mask = offs < n_tiles
+    tile_max = tl.load(
+        tile_max_ptr + row * n_tiles + offs, mask=mask, other=float("-inf")
+    )
+    keys = tl.load(partial_key_ptr + row * n_tiles + offs, mask=mask, other=-1.0)
+    idxs = tl.load(
+        partial_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL
+    )
+    global_max = tl.max(tile_max, axis=0)
+    rescaled = keys * tl.exp(tile_max - global_max)
+    rescaled = tl.where(mask, rescaled, -1.0)
+    best = tl.max(rescaled, axis=0)
+    cand = tl.where(rescaled == best, idxs, _IDX_SENTINEL)
+    tl.store(next_tokens_ptr + row, tl.min(cand, axis=0).to(tl.int64))
+
+
+def sample_step_tokens_triton(
+    *,
+    step_logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    greedy_mask: torch.Tensor,
+    exp_noise: torch.Tensor,
+) -> torch.Tensor:
+    bs, V = step_logits.shape
+    device = step_logits.device
+    assert step_logits.stride(1) == 1, "step_logits rows must be contiguous"
+    stride_row = step_logits.stride(0)
+    temperatures = temperatures.to(torch.float32).contiguous()
+    greedy_mask = greedy_mask.to(torch.int32).contiguous()
+    exp_noise = exp_noise.to(torch.float32).contiguous()
+
+    n_tiles = triton.cdiv(V, _BLOCK_V)
+    block_tiles = triton.next_power_of_2(n_tiles)
+
+    tile_max = torch.empty((bs, n_tiles), dtype=torch.float32, device=device)
+    partial_key = torch.empty((bs, n_tiles), dtype=torch.float32, device=device)
+    partial_idx = torch.empty((bs, n_tiles), dtype=torch.int32, device=device)
+    next_tokens = torch.empty((bs,), dtype=torch.int64, device=device)
+
+    tile_grid = (bs, n_tiles)
+    row_grid = (bs,)
+
+    _online_partial_kernel[tile_grid](
+        step_logits,
+        temperatures,
+        greedy_mask,
+        exp_noise,
+        tile_max,
+        partial_key,
+        partial_idx,
+        V,
+        stride_row,
+        n_tiles,
+        BLOCK_V=_BLOCK_V,
+    )
+    _online_combine_kernel[row_grid](
+        tile_max,
+        partial_key,
+        partial_idx,
+        next_tokens,
+        n_tiles,
+        BLOCK_TILES=block_tiles,
+    )
+    return next_tokens
+
+
+_STACKED_WEIGHT_CACHE: dict[int, _StackedWkvWeight] = {}
+
+
+class CommitKvProj:
+    @classmethod
+    def execute(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        if main_x.is_cuda and _fused_commit_kv_proj_supported(wkv_linears=wkv_linears):
+            return cls.triton(main_x=main_x, wkv_linears=wkv_linears)
+        return cls.torch(main_x=main_x, wkv_linears=wkv_linears)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        return commit_kv_proj(main_x=main_x, wkv_linears=wkv_linears)
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        return commit_kv_proj_fused(main_x=main_x, wkv_linears=wkv_linears)
+
+
+def commit_kv_proj(
+    *,
+    main_x: torch.Tensor,
+    wkv_linears: list[torch.nn.Module],
+) -> list[torch.Tensor]:
+    return [linear(main_x)[0] for linear in wkv_linears]
+
+
+def commit_kv_proj_fused(
+    *,
+    main_x: torch.Tensor,
+    wkv_linears: list[torch.nn.Module],
+) -> list[torch.Tensor]:
+    num_stages = len(wkv_linears)
+    stacked = _stacked_wkv_weight(wkv_linears=wkv_linears)
+
+    if stacked.fp8_scale is not None:
+        quant_method = wkv_linears[0].quant_method
+        kv_all = quant_method.w8a8_block_fp8_linear(
+            input=main_x,
+            weight=stacked.weight,
+            block_size=quant_method.quant_config.weight_block_size,
+            weight_scale=stacked.fp8_scale,
+            input_scale=None,
+            bias=None,
+        )
+    else:
+        kv_all = torch.nn.functional.linear(main_x, stacked.weight)
+
+    head_dim = kv_all.shape[-1] // num_stages
+    return [
+        kv_all[:, i * head_dim : (i + 1) * head_dim].contiguous()
+        for i in range(num_stages)
+    ]
+
+
+class _StackedWkvWeight(msgspec.Struct):
+    weight: torch.Tensor
+    fp8_scale: Optional[torch.Tensor]
+
+
+def _stacked_wkv_weight(*, wkv_linears: list[torch.nn.Module]) -> _StackedWkvWeight:
+    key = id(wkv_linears[0])
+    cached = _STACKED_WEIGHT_CACHE.get(key)
+    if cached is None:
+        cached = _build_stacked_wkv_weight(wkv_linears=wkv_linears)
+        _STACKED_WEIGHT_CACHE[key] = cached
+    return cached
+
+
+def _block_quant_stack_applies(*, wkv_linears: list[torch.nn.Module]) -> bool:
+    quant_method = wkv_linears[0].quant_method
+    block_quant = hasattr(quant_method, "block_quant") and quant_method.block_quant
+    if not (block_quant and hasattr(quant_method, "w8a8_block_fp8_linear")):
+        return False
+    block_out = quant_method.quant_config.weight_block_size[0]
+    return all(
+        linear.weight.dtype == torch.float8_e4m3fn
+        and linear.weight.shape[0] % block_out == 0
+        for linear in wkv_linears
+    )
+
+
+def _dequant_supported(linear: torch.nn.Module) -> bool:
+    """Mirrors the preconditions asserted in _dequant_linear_weight."""
+    weight = linear.weight
+    if weight.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return True
+    if weight.dtype != torch.float8_e4m3fn:
+        return False
+    block = 128
+    out_dim, in_dim = weight.shape
+    expected_scale_shape = (
+        (out_dim + block - 1) // block,
+        (in_dim + block - 1) // block,
+    )
+    return tuple(linear.weight_scale_inv.shape) == expected_scale_shape
+
+
+def _fused_commit_kv_proj_supported(*, wkv_linears: list[torch.nn.Module]) -> bool:
+    """Whether _build_stacked_wkv_weight can handle these weights; unsupported
+    quant schemes fall back to the per-linear torch path in execute()."""
+    if _block_quant_stack_applies(wkv_linears=wkv_linears):
+        return True
+    return all(_dequant_supported(linear) for linear in wkv_linears)
+
+
+def _build_stacked_wkv_weight(
+    *, wkv_linears: list[torch.nn.Module]
+) -> _StackedWkvWeight:
+    if _block_quant_stack_applies(wkv_linears=wkv_linears):
+        weight = torch.cat([linear.weight for linear in wkv_linears], dim=0)
+        if wkv_linears[0].weight_scale_inv.dtype == torch.int32:
+            from sglang.srt.layers.quantization.fp8_utils import (
+                inverse_transform_scale_ue8m0,
+                transform_scale_ue8m0,
+            )
+
+            sf_fp32 = torch.cat(
+                [
+                    inverse_transform_scale_ue8m0(
+                        linear.weight_scale_inv, mn=linear.weight.shape[0]
+                    )
+                    for linear in wkv_linears
+                ],
+                dim=0,
+            )
+            scale = transform_scale_ue8m0(sf_fp32, mn=weight.shape[0])
+            return _StackedWkvWeight(weight=weight, fp8_scale=scale)
+        scale = torch.cat([linear.weight_scale_inv for linear in wkv_linears], dim=0)
+        if scale.dim() >= 2 and scale.stride(-2) != 1:
+            scale = scale.transpose(-2, -1).contiguous().transpose(-2, -1)
+        return _StackedWkvWeight(weight=weight, fp8_scale=scale)
+    weight = torch.cat(
+        [_dequant_linear_weight(linear) for linear in wkv_linears], dim=0
+    )
+    return _StackedWkvWeight(weight=weight, fp8_scale=None)
+
+
+def _dequant_linear_weight(linear: torch.nn.Module) -> torch.Tensor:
+    weight = linear.weight
+    if weight.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return weight.to(torch.bfloat16)
+    assert weight.dtype == torch.float8_e4m3fn, (
+        f"unsupported wkv weight dtype {weight.dtype} for the fused commit kv proj; "
+        f"execute() should have routed this to the torch path "
+        f"(_fused_commit_kv_proj_supported)"
+    )
+    block = 128
+    scale = linear.weight_scale_inv
+    out_dim, in_dim = weight.shape
+    expected_scale_shape = (
+        (out_dim + block - 1) // block,
+        (in_dim + block - 1) // block,
+    )
+    assert tuple(scale.shape) == expected_scale_shape, (
+        f"wkv weight_scale_inv shape {tuple(scale.shape)} does not match the "
+        f"128x128 block grid {expected_scale_shape} for weight {tuple(weight.shape)}; "
+        f"execute() should have routed this to the torch path "
+        f"(_fused_commit_kv_proj_supported)"
+    )
+    scale_full = scale.repeat_interleave(block, dim=0)[:out_dim]
+    scale_full = scale_full.repeat_interleave(block, dim=1)[:, :in_dim]
+    return (weight.to(torch.float32) * scale_full.to(torch.float32)).to(torch.bfloat16)
+
+
+_BLOCK = 1024
+
+
+class BuildStepLocal:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+        return build_step_local(bias=bias, base_local=base_local)
+
+    @classmethod
+    def triton(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+        return build_step_local_triton(bias=bias, base_local=base_local)
+
+
+def build_step_local(*, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+    per_partition = base_local.shape[-1]
+    pad = per_partition - bias.shape[-1]
+    padded = (
+        F.pad(bias.to(torch.float32), (0, pad)) if pad > 0 else bias.to(torch.float32)
+    )
+    return base_local + padded
+
+
+@triton.jit
+def _build_step_local_kernel(
+    bias_ptr,
+    base_ptr,
+    out_ptr,
+    org_width,
+    per_partition,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    offs = tile * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < per_partition
+    base = tl.load(base_ptr + row * per_partition + offs, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    bias = tl.load(
+        bias_ptr + row * org_width + offs, mask=offs < org_width, other=0.0
+    ).to(tl.float32)
+    tl.store(out_ptr + row * per_partition + offs, base + bias, mask=mask)
+
+
+def build_step_local_triton(
+    *, bias: torch.Tensor, base_local: torch.Tensor
+) -> torch.Tensor:
+    bs, per_partition = base_local.shape
+    org_width = bias.shape[-1]
+    base_local = base_local.contiguous()
+    bias = bias.contiguous()
+    out = torch.empty(
+        (bs, per_partition), dtype=torch.float32, device=base_local.device
+    )
+    grid = (bs, triton.cdiv(per_partition, _BLOCK))
+    _build_step_local_kernel[grid](
+        bias, base_local, out, org_width, per_partition, BLOCK=_BLOCK
+    )
+    return out

python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import (
+    inputs_on_cuda,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.dspark_components.dspark_planner import (
+        DSparkScheduleConfig,
+    )
+
+
+class ScheduleVerifyLensTopk:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        confidence: torch.Tensor,
+        budget: int,
+        cfg: DSparkScheduleConfig,
+    ) -> torch.Tensor:
+        return schedule_verify_lens_topk(confidence=confidence, budget=budget, cfg=cfg)
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        confidence: torch.Tensor,
+        budget: int,
+        cfg: DSparkScheduleConfig,
+    ) -> torch.Tensor:
+        return schedule_verify_lens_topk_triton(
+            confidence=confidence, budget=budget, cfg=cfg
+        )
+
+
+def compute_sort_survival(confidence: torch.Tensor) -> torch.Tensor:
+    return torch.cumprod(confidence.to(torch.float32), dim=1)
+
+
+def schedule_verify_lens_topk(
+    *,
+    confidence: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    return schedule_verify_lens_topk_from_survival(
+        survival_probs=compute_sort_survival(confidence), budget=budget, cfg=cfg
+    )
+
+
+def schedule_verify_lens_topk_from_survival(
+    *,
+    survival_probs: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    num_requests, _gamma = survival_probs.shape
+    max_len = cfg.resolved_max_verify_len()
+    device = survival_probs.device
+
+    selected_extra = torch.zeros(num_requests, dtype=torch.int64, device=device)
+    if budget > 0:
+        candidate_window = survival_probs[:, :max_len]
+        num_candidates = candidate_window.numel()
+        if num_candidates > 0:
+            request_index = (
+                torch.arange(num_requests, device=device)
+                .view(num_requests, 1)
+                .expand_as(candidate_window)
+            )
+            position_index = (
+                torch.arange(candidate_window.shape[1], device=device)
+                .view(1, candidate_window.shape[1])
+                .expand_as(candidate_window)
+            )
+            valid = candidate_window >= cfg.survival_eps
+
+            flat_prob = candidate_window.reshape(-1).to(torch.float64)
+            flat_request = request_index.reshape(-1)
+            flat_position = position_index.reshape(-1)
+            flat_valid = valid.reshape(-1)
+
+            order = _value_independent_descending_order(
+                probs=flat_prob,
+                positions=flat_position,
+                requests=flat_request,
+                valid=flat_valid,
+            )
+
+            take = min(int(budget), num_candidates)
+            chosen = order[:take]
+            chosen_requests = flat_request[chosen]
+            chosen_valid = flat_valid[chosen].to(torch.int64)
+            selected_extra.scatter_add_(0, chosen_requests, chosen_valid)
+
+    min_len = torch.full(
+        (num_requests,), cfg.min_verify_len, dtype=torch.int64, device=device
+    )
+    verify_lens = min_len + selected_extra
+    lower_bound = max(cfg.min_verify_len, 1)
+    verify_lens = torch.clamp(verify_lens, min=lower_bound, max=max_len)
+    return verify_lens.to(torch.int32)
+
+
+def _value_independent_descending_order(
+    *,
+    probs: torch.Tensor,
+    positions: torch.Tensor,
+    requests: torch.Tensor,
+    valid: torch.Tensor,
+) -> torch.Tensor:
+    masked_prob = torch.where(valid, probs, torch.full_like(probs, float("-inf")))
+    num_candidates = masked_prob.numel()
+    order = torch.arange(num_candidates, device=probs.device)
+    order = order[torch.argsort(requests[order], stable=True)]
+    order = order[torch.argsort(positions[order], stable=True)]
+    order = order[torch.argsort(-masked_prob[order], stable=True)]
+    return order
+
+
+@triton.jit
+def _schedule_topk_prep_kernel(
+    confidence_ptr,
+    survival_ptr,
+    selected_extra_ptr,
+    gamma,
+    cols,
+    G_P2: tl.constexpr,
+):
+    row = tl.program_id(0)
+    g = tl.arange(0, G_P2)
+    conf = tl.load(
+        confidence_ptr + row.to(tl.int64) * gamma + g, mask=g < gamma, other=1.0
+    ).to(tl.float32)
+    surv = tl.cumprod(conf, axis=0)
+    tl.store(survival_ptr + row.to(tl.int64) * cols + g, surv, mask=g < cols)
+    tl.store(selected_extra_ptr + row, 0)
+
+
+@triton.jit
+def _schedule_topk_finalize_kernel(
+    selected_extra_ptr,
+    out_ptr,
+    min_verify_len,
+    lower_bound,
+    max_len,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    extra = tl.load(selected_extra_ptr + offs, mask=mask, other=0).to(tl.int32)
+    lens = min_verify_len + extra
+    lens = tl.maximum(lens, lower_bound)
+    lens = tl.minimum(lens, max_len)
+    tl.store(out_ptr + offs, lens, mask=mask)
+
+
+@triton.jit
+def _schedule_topk_selected_extra_kernel(
+    survival_ptr,
+    selected_extra_ptr,
+    budget,
+    cols,
+    n,
+    survival_eps,
+    BLOCK_C: tl.constexpr,
+    BLOCK_CP: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    c = pid * BLOCK_C + tl.arange(0, BLOCK_C)
+    cmask = c < n
+    r = c // cols
+    p = c % cols
+    sp = tl.load(survival_ptr + c, mask=cmask, other=0.0)
+    valid_c = sp >= survival_eps
+    mp = tl.where(valid_c, sp, float("-inf"))
+    rank = tl.zeros([BLOCK_C], dtype=tl.int32)
+    for cp0 in range(0, n, BLOCK_CP):
+        cp = cp0 + tl.arange(0, BLOCK_CP)
+        cpmask = cp < n
+        rp = cp // cols
+        pp = cp % cols
+        spp = tl.load(survival_ptr + cp, mask=cpmask, other=0.0)
+        validp = spp >= survival_eps
+        mpp = tl.where(validp, spp, float("-inf"))
+        gt = mpp[None, :] > mp[:, None]
+        eq = mpp[None, :] == mp[:, None]
+        pos_lt = pp[None, :] < p[:, None]
+        pos_eq = pp[None, :] == p[:, None]
+        req_lt = rp[None, :] < r[:, None]
+        before = gt | (eq & (pos_lt | (pos_eq & req_lt)))
+        before = before & cpmask[None, :]
+        rank += tl.sum(before.to(tl.int32), axis=1)
+    selected = valid_c & (rank < budget)
+    tl.atomic_add(selected_extra_ptr + r, selected.to(tl.int32), mask=cmask)
+
+
+def schedule_verify_lens_topk_triton(
+    *,
+    confidence: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    num_requests, gamma = confidence.shape
+    max_len = cfg.resolved_max_verify_len()
+    device = confidence.device
+    cols = min(max_len, gamma)
+    n = num_requests * cols
+
+    selected_extra = torch.empty(num_requests, dtype=torch.int32, device=device)
+    survival = torch.empty((num_requests, cols), dtype=torch.float32, device=device)
+    _schedule_topk_prep_kernel[(num_requests,)](
+        confidence.contiguous(),
+        survival,
+        selected_extra,
+        gamma,
+        cols,
+        G_P2=triton.next_power_of_2(max(gamma, 1)),
+    )
+    if budget > 0 and n > 0:
+        BLOCK_C = 64
+        BLOCK_CP = 256
+        grid = (triton.cdiv(n, BLOCK_C),)
+        _schedule_topk_selected_extra_kernel[grid](
+            survival,
+            selected_extra,
+            int(budget),
+            cols,
+            n,
+            float(cfg.survival_eps),
+            BLOCK_C=BLOCK_C,
+            BLOCK_CP=BLOCK_CP,
+        )
+
+    verify_lens = torch.empty(num_requests, dtype=torch.int32, device=device)
+    BLOCK = 256
+    _schedule_topk_finalize_kernel[(triton.cdiv(num_requests, BLOCK),)](
+        selected_extra,
+        verify_lens,
+        int(cfg.min_verify_len),
+        max(cfg.min_verify_len, 1),
+        int(max_len),
+        num_requests,
+        BLOCK=BLOCK,
+    )
+    return verify_lens

python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py

@@ -0,0 +1,871 @@
+from __future__ import annotations
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+class RaggedVerifyWindow(msgspec.Struct, frozen=True):
+    positions: torch.Tensor
+    verify_cache_loc: torch.Tensor
+    verify_ids: torch.Tensor
+
+
+class BuildRaggedVerifyWindow:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> RaggedVerifyWindow:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        bs: int,
+        device: str,
+        verify_num_draft_tokens: int,
+        model_runner,
+    ) -> RaggedVerifyWindow:
+        return build_ragged_verify_window(
+            batch=batch,
+            layout=layout,
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            model_runner=model_runner,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        bs: int,
+        device: str,
+        verify_num_draft_tokens: int,
+        model_runner,
+    ) -> RaggedVerifyWindow:
+        return build_ragged_verify_window_triton(
+            batch=batch,
+            layout=layout,
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            model_runner=model_runner,
+        )
+
+
+def build_ragged_verify_window(
+    *,
+    batch: ScheduleBatch,
+    layout: RaggedVerifyLayout,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    bs: int,
+    device: str,
+    verify_num_draft_tokens: int,
+    model_runner,
+) -> RaggedVerifyWindow:
+    prefix_lens = batch.seq_lens
+    verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32)
+    padded_total = layout.graph_num_tokens
+
+    req_id, within, valid = compact_row_index(
+        verify_lens=verify_lens, padded_total=padded_total, device=device
+    )
+    safe_req = req_id.clamp(max=bs - 1)
+    positions = torch.where(
+        valid,
+        prefix_lens.to(torch.int64)[safe_req] + within,
+        torch.zeros_like(within),
+    )
+    real_cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=batch.req_pool_indices,
+        req_to_token=model_runner.req_to_token_pool.req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype),
+        batch_size=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    verify_cache_loc = torch.nn.functional.pad(
+        real_cache_loc, (0, padded_total - real_cache_loc.shape[0])
+    )
+    verify_cache_loc = torch.where(
+        valid, verify_cache_loc, torch.zeros_like(verify_cache_loc)
+    )
+
+    verify_ids = compact_verify_ids(
+        draft_block_ids=draft_block_ids,
+        draft_tokens=draft_tokens,
+        layout=layout,
+        device=device,
+    )
+
+    return RaggedVerifyWindow(
+        positions=positions,
+        verify_cache_loc=verify_cache_loc,
+        verify_ids=verify_ids,
+    )
+
+
+@triton.jit
+def _ragged_finalize_kernel(
+    req_ptr,
+    within_ptr,
+    prefix_ptr,
+    cache_ptr,
+    pos_out_ptr,
+    cache_out_ptr,
+    bs,
+    n,
+    real_len,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    req = tl.load(req_ptr + offs, mask=mask, other=0)
+    within = tl.load(within_ptr + offs, mask=mask, other=0)
+    valid = req < bs
+    safe_req = tl.minimum(req, bs - 1)
+    prefix = tl.load(prefix_ptr + safe_req, mask=mask, other=0)
+    pos = tl.where(valid, prefix + within, 0)
+    lmask = mask & (offs < real_len)
+    cl = tl.load(cache_ptr + offs, mask=lmask, other=0)
+    cl = tl.where(valid, cl, 0)
+    tl.store(pos_out_ptr + offs, pos, mask=mask)
+    tl.store(cache_out_ptr + offs, cl, mask=mask)
+
+
+def build_ragged_verify_window_triton(
+    *,
+    batch: ScheduleBatch,
+    layout: RaggedVerifyLayout,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    bs: int,
+    device: str,
+    verify_num_draft_tokens: int,
+    model_runner,
+) -> RaggedVerifyWindow:
+    prefix_lens = batch.seq_lens
+    verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32)
+    padded_total = layout.graph_num_tokens
+
+    req_id, within, _valid = compact_row_index_triton(
+        verify_lens=verify_lens, padded_total=padded_total, device=device
+    )
+    real_cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=batch.req_pool_indices,
+        req_to_token=model_runner.req_to_token_pool.req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype),
+        batch_size=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    prefix_i64 = prefix_lens.to(device=device, dtype=torch.int64).contiguous()
+    positions = torch.empty(padded_total, dtype=torch.int64, device=device)
+    verify_cache_loc = torch.empty(
+        padded_total, dtype=real_cache_loc.dtype, device=device
+    )
+    BLOCK = 256
+    grid = (triton.cdiv(padded_total, BLOCK),)
+    _ragged_finalize_kernel[grid](
+        req_id,
+        within,
+        prefix_i64,
+        real_cache_loc,
+        positions,
+        verify_cache_loc,
+        bs,
+        padded_total,
+        real_cache_loc.shape[0],
+        BLOCK=BLOCK,
+    )
+
+    verify_ids = compact_verify_ids_triton(
+        draft_block_ids=draft_block_ids,
+        draft_tokens=draft_tokens,
+        layout=layout,
+        device=device,
+    )
+    return RaggedVerifyWindow(
+        positions=positions,
+        verify_cache_loc=verify_cache_loc,
+        verify_ids=verify_ids,
+    )
+
+
+_SEARCH_NBITS = 11
+
+
+class CompactRowIndex:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        padded_total: int,
+        device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return compact_row_index(
+            verify_lens=verify_lens,
+            padded_total=padded_total,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        padded_total: int,
+        device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return compact_row_index_triton(
+            verify_lens=verify_lens,
+            padded_total=padded_total,
+            device=device,
+        )
+
+
+class CompactVerifyIds:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        device: str,
+    ) -> torch.Tensor:
+        return compact_verify_ids(
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            layout=layout,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        device: str,
+    ) -> torch.Tensor:
+        return compact_verify_ids_triton(
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            layout=layout,
+            device=device,
+        )
+
+
+def compact_verify_ids(
+    *,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    device: str,
+) -> torch.Tensor:
+    req_id, within, valid = compact_row_index(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    bs = layout.verify_lens.shape[0]
+    safe_req = req_id.clamp(max=bs - 1)
+    anchors = draft_block_ids[:, 0]
+    drafts = draft_tokens[safe_req, (within - 1).clamp_min(0)]
+    verify_ids = torch.where(within == 0, anchors[safe_req], drafts)
+    verify_ids = torch.where(valid, verify_ids, torch.zeros_like(verify_ids))
+    return verify_ids.to(torch.int64)
+
+
+def compact_row_index(
+    *,
+    verify_lens: torch.Tensor,
+    padded_total: int,
+    device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    verify_lens = verify_lens.to(device=device, dtype=torch.int64)
+    bs = int(verify_lens.numel())
+    incl = torch.cumsum(verify_lens, dim=0)
+    start = incl - verify_lens
+    real_total = incl[-1]
+    row = torch.arange(padded_total, device=device, dtype=torch.int64)
+    valid = row < real_total
+    req_id = torch.searchsorted(incl, row, right=True)
+    req_id = torch.where(valid, req_id, torch.full_like(req_id, bs))
+    within = torch.where(
+        valid, row - start[req_id.clamp(max=bs - 1)], torch.zeros_like(row)
+    )
+    return req_id, within, valid
+
+
+@triton.jit
+def _compact_row_index_kernel(
+    incl_ptr,
+    req_out_ptr,
+    within_out_ptr,
+    valid_out_ptr,
+    bs,
+    n,
+    BLOCK: tl.constexpr,
+    NBITS: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    row = offs.to(tl.int64)
+    real_total = tl.load(incl_ptr + (bs - 1))
+    lo = tl.zeros([BLOCK], dtype=tl.int32)
+    hi = tl.full([BLOCK], bs, dtype=tl.int32)
+    for _ in range(NBITS):
+        mid = (lo + hi) // 2
+        active = lo < hi
+        val = tl.load(incl_ptr + tl.minimum(mid, bs - 1), mask=mask, other=0)
+        go_right = val <= row
+        lo = tl.where(active & go_right, mid + 1, lo)
+        hi = tl.where(active & (~go_right), mid, hi)
+    req = lo
+    gidx = tl.maximum(req - 1, 0)
+    start = tl.load(incl_ptr + gidx, mask=mask, other=0)
+    start = tl.where(req > 0, start, 0)
+    valid = row < real_total
+    within = tl.where(valid, row - start, 0)
+    req_final = tl.where(valid, req.to(tl.int64), bs)
+    tl.store(req_out_ptr + offs, req_final, mask=mask)
+    tl.store(within_out_ptr + offs, within, mask=mask)
+    tl.store(valid_out_ptr + offs, valid, mask=mask)
+
+
+def compact_row_index_triton(
+    *,
+    verify_lens: torch.Tensor,
+    padded_total: int,
+    device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    verify_lens = verify_lens.to(device=device, dtype=torch.int64).contiguous()
+    bs = verify_lens.shape[0]
+    incl = torch.cumsum(verify_lens, dim=0).contiguous()
+    req = torch.empty(padded_total, dtype=torch.int64, device=device)
+    within = torch.empty(padded_total, dtype=torch.int64, device=device)
+    valid = torch.empty(padded_total, dtype=torch.bool, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(padded_total, BLOCK),)
+    _compact_row_index_kernel[grid](
+        incl, req, within, valid, bs, padded_total, BLOCK=BLOCK, NBITS=_SEARCH_NBITS
+    )
+    return req, within, valid
+
+
+@triton.jit
+def _compact_verify_ids_gather_kernel(
+    req_ptr,
+    within_ptr,
+    draft_block_ids_ptr,
+    draft_tokens_ptr,
+    out_ptr,
+    bs,
+    gamma,
+    n,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    req = tl.load(req_ptr + offs, mask=mask, other=0)
+    within = tl.load(within_ptr + offs, mask=mask, other=0)
+    valid = req < bs
+    safe_req = tl.minimum(req, bs - 1)
+    anchor = tl.load(draft_block_ids_ptr + safe_req * gamma, mask=mask, other=0)
+    wcol = tl.maximum(within - 1, 0)
+    draft = tl.load(draft_tokens_ptr + safe_req * gamma + wcol, mask=mask, other=0)
+    v = tl.where(within == 0, anchor, draft)
+    v = tl.where(valid, v, 0)
+    tl.store(out_ptr + offs, v.to(tl.int64), mask=mask)
+
+
+def compact_verify_ids_triton(
+    *,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    device: str,
+) -> torch.Tensor:
+    req, within, _valid = compact_row_index_triton(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    bs = layout.verify_lens.shape[0]
+    gamma = draft_tokens.shape[1]
+    draft_block_ids = draft_block_ids.to(device=device, dtype=torch.int64).contiguous()
+    draft_tokens = draft_tokens.to(device=device, dtype=torch.int64).contiguous()
+    n = layout.graph_num_tokens
+    out = torch.empty(n, dtype=torch.int64, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(n, BLOCK),)
+    _compact_verify_ids_gather_kernel[grid](
+        req, within, draft_block_ids, draft_tokens, out, bs, gamma, n, BLOCK=BLOCK
+    )
+    return out
+
+
+class ScatterCompactToStrided:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        compact: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        fill_value: float,
+        verify_num_draft_tokens: int,
+    ) -> torch.Tensor:
+        return scatter_compact_to_strided(
+            compact=compact,
+            layout=layout,
+            fill_value=fill_value,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        compact: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        fill_value: float,
+        verify_num_draft_tokens: int,
+    ) -> torch.Tensor:
+        return scatter_compact_to_strided_triton(
+            compact=compact,
+            layout=layout,
+            fill_value=fill_value,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+        )
+
+
+def scatter_compact_to_strided(
+    *,
+    compact: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    fill_value: float,
+    verify_num_draft_tokens: int,
+) -> torch.Tensor:
+    stride = verify_num_draft_tokens
+    bs = layout.verify_lens.shape[0]
+    dim = compact.shape[1]
+    device = compact.device
+    compact = compact[: layout.graph_num_tokens]
+    strided = torch.full(
+        (bs * stride + 1, dim), fill_value, dtype=compact.dtype, device=device
+    )
+    req_id, within, valid = compact_row_index(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    sink = bs * stride
+    strided_pos = torch.where(
+        valid,
+        req_id.clamp(max=bs - 1) * stride + within,
+        torch.full_like(within, sink),
+    )
+    strided.index_copy_(0, strided_pos, compact)
+    return strided[: bs * stride]
+
+
+@triton.jit
+def _scatter_compact_to_strided_kernel(
+    compact_ptr,
+    verify_lens_ptr,
+    start_ptr,
+    out_ptr,
+    stride,
+    dim,
+    fill_value,
+    BLOCK_D: tl.constexpr,
+):
+    o = tl.program_id(0).to(tl.int64)
+    dblk = tl.program_id(1)
+    i = o // stride
+    w = o % stride
+    vl_i = tl.load(verify_lens_ptr + i)
+    start_i = tl.load(start_ptr + i)
+    d = dblk * BLOCK_D + tl.arange(0, BLOCK_D)
+    dmask = d < dim
+    in_range = w < vl_i
+    src = tl.where(in_range, start_i + w, 0)
+    val = tl.load(compact_ptr + src * dim + d, mask=dmask & in_range, other=0)
+    val = tl.where(in_range, val, fill_value)
+    tl.store(out_ptr + o * dim + d, val, mask=dmask)
+
+
+def scatter_compact_to_strided_into(
+    *,
+    compact: torch.Tensor,
+    verify_lens: torch.Tensor,
+    out: torch.Tensor,
+    stride: int,
+    fill_value: float,
+) -> torch.Tensor:
+    dim = compact.shape[1]
+    fill_value = float(fill_value) if out.dtype.is_floating_point else int(fill_value)
+    compact = compact.contiguous()
+    verify_lens = verify_lens.to(dtype=torch.int64).contiguous()
+    start = (torch.cumsum(verify_lens, dim=0) - verify_lens).contiguous()
+    n_out = out.shape[0]
+    BLOCK_D = 1024
+    grid = (n_out, triton.cdiv(dim, BLOCK_D))
+    _scatter_compact_to_strided_kernel[grid](
+        compact,
+        verify_lens,
+        start,
+        out,
+        stride,
+        dim,
+        fill_value,
+        BLOCK_D=BLOCK_D,
+    )
+    return out
+
+
+def scatter_compact_to_strided_triton(
+    *,
+    compact: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    fill_value: float,
+    verify_num_draft_tokens: int,
+) -> torch.Tensor:
+    stride = verify_num_draft_tokens
+    bs = layout.verify_lens.shape[0]
+    dim = compact.shape[1]
+    device = compact.device
+    out = torch.empty((bs * stride, dim), dtype=compact.dtype, device=device)
+    return scatter_compact_to_strided_into(
+        compact=compact,
+        verify_lens=layout.verify_lens.to(device=device),
+        out=out,
+        stride=stride,
+        fill_value=fill_value,
+    )
+
+
+class CommitInjectLayoutResult(msgspec.Struct):
+    swa_loc: torch.Tensor
+    positions: torch.Tensor
+
+
+class BuildCommitInjectLayout:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> CommitInjectLayoutResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        block_pos_offsets: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        commit_lens: torch.Tensor,
+        stride: int,
+    ) -> CommitInjectLayoutResult:
+        return build_commit_inject_layout(
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            prefix_lens=prefix_lens,
+            block_pos_offsets=block_pos_offsets,
+            full_to_swa_mapping=full_to_swa_mapping,
+            commit_lens=commit_lens,
+            stride=stride,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        block_pos_offsets: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        commit_lens: torch.Tensor,
+        stride: int,
+    ) -> CommitInjectLayoutResult:
+        return build_commit_inject_layout_triton(
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            prefix_lens=prefix_lens,
+            block_pos_offsets=block_pos_offsets,
+            full_to_swa_mapping=full_to_swa_mapping,
+            commit_lens=commit_lens,
+            stride=stride,
+        )
+
+
+def build_commit_inject_layout(
+    *,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_pos_offsets: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    commit_lens: torch.Tensor,
+    stride: int,
+) -> CommitInjectLayoutResult:
+    from sglang.kernels.ops.speculative.cache_locs import (
+        assign_extend_cache_locs_func,
+    )
+
+    bs = req_pool_indices.shape[0]
+    device = req_pool_indices.device
+
+    positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets[:stride]
+    positions = positions_2d.reshape(-1).to(dtype=torch.int64)
+
+    cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + stride,
+        batch_size=bs,
+        draft_token_num=stride,
+        device=device,
+    ).to(dtype=torch.int64)
+    swa_loc = full_to_swa_mapping[cache_loc].to(torch.int32)
+
+    col = torch.arange(stride, device=device).view(1, -1)
+    committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
+    swa_loc = torch.where(committed, swa_loc, torch.full_like(swa_loc, -1))
+
+    return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions)
+
+
+@triton.jit
+def _commit_inject_layout_kernel(
+    req_pool_ptr,
+    req_to_token_ptr,
+    prefix_lens_ptr,
+    block_pos_offsets_ptr,
+    full_to_swa_ptr,
+    commit_lens_ptr,
+    swa_loc_ptr,
+    positions_ptr,
+    rt_stride,
+    stride,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    r = offs // stride
+    c = offs % stride
+
+    prefix = tl.load(prefix_lens_ptr + r, mask=mask, other=0).to(tl.int64)
+    pos_off = tl.load(block_pos_offsets_ptr + c, mask=mask, other=0).to(tl.int64)
+    rp = tl.load(req_pool_ptr + r, mask=mask, other=0).to(tl.int64)
+    full_loc = tl.load(
+        req_to_token_ptr + rp * rt_stride + prefix + pos_off, mask=mask, other=0
+    ).to(tl.int64)
+    swa = tl.load(full_to_swa_ptr + full_loc, mask=mask, other=-1).to(tl.int32)
+
+    commit_len = tl.load(commit_lens_ptr + r, mask=mask, other=0).to(tl.int64)
+    swa = tl.where(c.to(tl.int64) < commit_len, swa, -1)
+
+    tl.store(swa_loc_ptr + offs, swa, mask=mask)
+    tl.store(positions_ptr + offs, prefix + pos_off, mask=mask)
+
+
+def build_commit_inject_layout_triton(
+    *,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_pos_offsets: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    commit_lens: torch.Tensor,
+    stride: int,
+) -> CommitInjectLayoutResult:
+    bs = req_pool_indices.shape[0]
+    n = bs * stride
+    device = req_pool_indices.device
+
+    swa_loc = torch.empty(n, dtype=torch.int32, device=device)
+    positions = torch.empty(n, dtype=torch.int64, device=device)
+    BLOCK = 256
+    _commit_inject_layout_kernel[(triton.cdiv(n, BLOCK),)](
+        req_pool_indices,
+        req_to_token,
+        prefix_lens,
+        block_pos_offsets,
+        full_to_swa_mapping,
+        commit_lens,
+        swa_loc,
+        positions,
+        req_to_token.stride(0),
+        stride,
+        n,
+        BLOCK=BLOCK,
+    )
+    return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions)
+
+
+class BuildOutTokens:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        draft_tokens: torch.Tensor,
+        correct_len: torch.Tensor,
+        bonus: torch.Tensor,
+        verify_num_draft_tokens: int,
+        gamma: int,
+    ) -> torch.Tensor:
+        return build_out_tokens(
+            draft_tokens=draft_tokens,
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            gamma=gamma,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        draft_tokens: torch.Tensor,
+        correct_len: torch.Tensor,
+        bonus: torch.Tensor,
+        verify_num_draft_tokens: int,
+        gamma: int,
+    ) -> torch.Tensor:
+        return build_out_tokens_triton(
+            draft_tokens=draft_tokens,
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            gamma=gamma,
+        )
+
+
+def build_out_tokens(
+    *,
+    draft_tokens: torch.Tensor,
+    correct_len: torch.Tensor,
+    bonus: torch.Tensor,
+    verify_num_draft_tokens: int,
+    gamma: int,
+) -> torch.Tensor:
+    bs = draft_tokens.shape[0]
+    out_tokens = torch.empty(
+        (bs, verify_num_draft_tokens),
+        dtype=torch.int64,
+        device=draft_tokens.device,
+    )
+    out_tokens[:, :gamma].copy_(draft_tokens)
+    out_tokens[:, gamma].fill_(0)
+    out_tokens.scatter_(1, correct_len.to(torch.int64)[:, None], bonus[:, None])
+    return out_tokens
+
+
+@triton.jit
+def _build_out_tokens_kernel(
+    draft_tokens_ptr,
+    correct_len_ptr,
+    bonus_ptr,
+    out_ptr,
+    gamma,
+    T,
+    n_out,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_out
+    b = offs // T
+    k = offs % T
+    cl = tl.load(correct_len_ptr + b, mask=mask, other=0).to(tl.int32)
+    bonus = tl.load(bonus_ptr + b, mask=mask, other=0)
+    draft_mask = mask & (k < gamma)
+    draft = tl.load(draft_tokens_ptr + b * gamma + k, mask=draft_mask, other=0)
+    val = tl.where(k == cl, bonus, tl.where(k < gamma, draft, 0))
+    tl.store(out_ptr + offs, val.to(tl.int64), mask=mask)
+
+
+def build_out_tokens_triton(
+    *,
+    draft_tokens: torch.Tensor,
+    correct_len: torch.Tensor,
+    bonus: torch.Tensor,
+    verify_num_draft_tokens: int,
+    gamma: int,
+) -> torch.Tensor:
+    bs = draft_tokens.shape[0]
+    T = verify_num_draft_tokens
+    device = draft_tokens.device
+    draft_tokens = draft_tokens.to(torch.int64).contiguous()
+    correct_len_i = correct_len.to(torch.int64).contiguous()
+    bonus_i = bonus.to(torch.int64).contiguous()
+    out = torch.empty((bs, T), dtype=torch.int64, device=device)
+    n_out = bs * T
+    BLOCK = 256
+    grid = (triton.cdiv(n_out, BLOCK),)
+    _build_out_tokens_kernel[grid](
+        draft_tokens, correct_len_i, bonus_i, out, gamma, T, n_out, BLOCK=BLOCK
+    )
+    return out

python/sglang/srt/speculative/eagle_disaggregation.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.managers.overlap_utils import RelayPayload
+from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+from sglang.srt.speculative.eagle_info import EagleDraftInput
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.overlap_utils import FutureMap
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.server_args import ServerArgs
+
+
+def build_eagle_disagg_draft_input(
+    batch: ScheduleBatch,
+    server_args: ServerArgs,
+    last_tokens_tensor: torch.Tensor,
+    future_map: FutureMap,
+) -> EagleDraftInput:
+    num_states = server_args.speculative_eagle_topk
+    if server_args.enable_multi_layer_eagle:
+        num_states *= server_args.speculative_num_steps
+
+    topk_p = torch.stack(
+        [
+            torch.as_tensor(
+                req.output_topk_p[:num_states],
+                device=batch.device,
+                dtype=torch.float32,
+            )
+            for req in batch.reqs
+        ],
+        dim=0,
+    )
+    topk_index = torch.stack(
+        [
+            torch.as_tensor(
+                req.output_topk_index[:num_states],
+                device=batch.device,
+                dtype=torch.int64,
+            )
+            for req in batch.reqs
+        ],
+        dim=0,
+    )
+
+    hidden_states = torch.stack(
+        [req.hidden_states_tensor for req in batch.reqs], dim=0
+    ).to(batch.device)
+
+    spec_info = EagleDraftInput(
+        topk_p=topk_p,
+        topk_index=topk_index,
+        hidden_states=hidden_states,
+        bonus_tokens=last_tokens_tensor,
+    )
+    spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
+
+    if batch.enable_overlap:
+        spec_info.future_indices = batch.req_pool_indices
+        # Seed the relay buf with the known seq_lens; publish's chained record
+        # keeps the in-flight forward's fence intact (see FutureMap.publish).
+        future_map.publish(spec_info.future_indices, batch.seq_lens)
+        future_map.stash(
+            spec_info.future_indices, RelayPayload.from_draft_input(spec_info)
+        )
+
+    return spec_info

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -0,0 +1,668 @@
+from __future__ import annotations
+
+import contextlib
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+from sglang.srt.model_executor.runner import (
+    DecodeCudaGraphRunner,
+    DeepEPCudaGraphRunnerAdapter,
+    ShapeKey,
+    _grouped_foreach_copy_,
+    get_batch_sizes_to_capture,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
+from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+from sglang.srt.runtime_context import get_flags
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.speculative.eagle_info import EagleDraftInput
+from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
+from sglang.srt.utils import (
+    require_attn_tp_gather,
+    require_gathered_buffer,
+    require_mlp_sync,
+    require_mlp_tp_gather,
+)
+from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+@dataclass
+class EagleDraftInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    req_pool_indices: torch.Tensor
+    out_cache_loc: torch.Tensor
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    rids_int: Optional[torch.Tensor]
+    bootstrap_room_ids_int: Optional[torch.Tensor]
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    extend_seq_lens: torch.Tensor
+    topk_p: torch.Tensor
+    topk_index: torch.Tensor
+    draft_probs: Optional[torch.Tensor]
+    hidden_states: Optional[torch.Tensor]
+    global_num_tokens_gpu: Optional[torch.Tensor]
+    global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+    dsa_seed_topk: Optional[torch.Tensor] = None
+
+
+class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
+    """EAGLE draft cuda-graph runner.
+
+    Subclasses DecodeCudaGraphRunner to inherit the outer capture
+    loop (capture()), bucket-padding helper (_pad_to_bucket),
+    and the backend-driven capture/replay scaffolding. EAGLE-specific
+    bits — buffer dataclass, dummy ForwardBatch construction in
+    capture_one_shape, replay output unwrap, and can_run_graph — are
+    overridden.
+
+    EAGLE does not call DecodeCudaGraphRunner.__init__ (that init
+    sets up many decode-only fields like SWA/encoder-decoder/MLA-aware
+    state). Instead it sets up its own state directly while making sure
+    the parent's capture() / backend contract is satisfied.
+    """
+
+    def __init__(
+        self,
+        eagle_worker: EagleDraftWorker,
+        *,
+        draft_attn_backend=None,
+        speculative_num_steps: Optional[int] = None,
+    ):
+        # Parse args
+        self.eagle_worker = eagle_worker
+        if not hasattr(eagle_worker, "model_runner"):
+            # V2: EagleDraftWorker
+            self.model_runner = model_runner = eagle_worker.draft_runner
+        else:
+            self.model_runner = model_runner = eagle_worker.model_runner
+
+        # Fields the parent's capture() reads:
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.tp_size = model_runner.tp_size
+        self.dp_size = model_runner.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.enable_torch_compile = get_flags().capture.enable_torch_compile
+        self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
+        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
+        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
+        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
+        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+        self.enable_profile_cuda_graph = (
+            model_runner.server_args.enable_profile_cuda_graph
+        )
+        self.speculative_num_steps = (
+            model_runner.server_args.speculative_num_steps
+            if speculative_num_steps is None
+            else speculative_num_steps
+        )
+        self.topk = model_runner.server_args.speculative_eagle_topk
+        self.draft_attn_backend = draft_attn_backend or model_runner.draft_attn_backend
+
+        # Patch_model in parent's capture() needs an attn_backend reference.
+        # EAGLE doesn't use it (capture_one_shape calls draft_forward instead),
+        # but the field must exist.
+        self.attn_backend = self.draft_attn_backend
+
+        # Disable parent paths that don't apply to EAGLE.
+        self.compile_bs = []  # disables patch_model torch.compile wrapping
+        self.enable_pdmux = False
+        self.record_nolora_graph = False
+        self.is_dllm = False
+
+        self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
+
+        # Capture-time globals required by parent's capture_one_shape signature.
+        self.capture_forward_mode = ForwardMode.DECODE
+        self.capture_hidden_mode = CaptureHiddenMode.LAST
+
+        # Bucket sizes
+        self.capture_bs, _ = get_batch_sizes_to_capture(model_runner)
+        self.num_tokens_per_bs = self.topk
+        self.max_bs = max(self.capture_bs)
+        self.max_num_token = self.max_bs * self.num_tokens_per_bs
+
+        # Attention backend init
+        self.draft_attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
+        self.seq_len_fill_value = self.draft_attn_backend.attn_backends[
+            0
+        ].get_cuda_graph_seq_len_fill_value()
+        self.extend_seq_lens_cpu = [self.seq_len_fill_value] * self.max_bs
+
+        if self.enable_torch_compile:
+            set_torch_compile_config()
+
+        # Static buffers
+        with torch.device(model_runner.device):
+            input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int64)
+            out_cache_loc = torch.zeros(
+                (self.max_num_token * self.speculative_num_steps,),
+                dtype=self._cache_loc_dtype(),
+            )
+            positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
+            rids_int = (
+                torch.zeros((self.max_bs,), dtype=torch.int64)
+                if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get()
+                else None
+            )
+            bootstrap_room_ids_int = (
+                torch.full((self.max_bs,), -1, dtype=torch.int64)
+                if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get()
+                else None
+            )
+            seq_lens = torch.full(
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
+            )
+            extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
+            topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
+            topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)
+            draft_probs = (
+                torch.zeros(
+                    (self.max_bs, self.model_runner.model_config.vocab_size),
+                    dtype=torch.float32,
+                )
+                if self.model_runner.server_args.speculative_use_rejection_sampling
+                else None
+            )
+            _hidden_size, _hidden_dtype = get_draft_recurrent_hidden_state_spec(
+                model_runner
+            )
+            hidden_states = (
+                torch.zeros(
+                    (self.max_bs, _hidden_size),
+                    dtype=_hidden_dtype,
+                )
+                if _hidden_size is not None
+                else None
+            )
+
+            self.temperatures = torch.ones((self.max_bs, 1), dtype=torch.float)
+
+            if self.require_gathered_buffer:
+                if self.require_mlp_tp_gather:
+                    global_num_tokens_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                else:
+                    assert self.require_attn_tp_gather
+                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (1,), dtype=torch.int32
+                    )
+            else:
+                global_num_tokens_gpu = None
+                global_num_tokens_for_logprob_gpu = None
+
+        seq_lens_cpu = torch.full(
+            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
+        )
+
+        dsa_seed_topk = (
+            torch.zeros(
+                (self.max_bs, self.eagle_worker.dsa_index_topk),
+                dtype=torch.int32,
+                device=model_runner.device,
+            )
+            if self.eagle_worker.seed_dsa_topk_from_draft_extend
+            else None
+        )
+
+        self.buffers = EagleDraftInputBuffers(
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            out_cache_loc=out_cache_loc,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            rids_int=rids_int,
+            bootstrap_room_ids_int=bootstrap_room_ids_int,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            extend_seq_lens=extend_seq_lens,
+            topk_p=topk_p,
+            topk_index=topk_index,
+            draft_probs=draft_probs,
+            hidden_states=hidden_states,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            dsa_seed_topk=dsa_seed_topk,
+        )
+        self.buffers.share_buffers()
+
+        self.backend = resolve_decode_backend(self)
+
+        # Capture
+        try:
+            with model_capture_mode():
+                self.capture()
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+            )
+
+    def _replay_graph(self, shape_key, forward_batch):
+        return self.backend.replay(shape_key, forward_batch)
+
+    # -----------------------------------------------------------------
+    # Helpers
+    # -----------------------------------------------------------------
+    def _cache_loc_dtype(self):
+        return torch.int64
+
+    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
+        # EAGLE doesn't use stream_idx / lora variants.
+        return ShapeKey(size=bs)
+
+    # -----------------------------------------------------------------
+    # can_run_graph
+    # -----------------------------------------------------------------
+    def can_run_graph(self, forward_batch: ForwardBatch):
+        if self.require_mlp_tp_gather:
+            cuda_graph_bs = (
+                max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
+                if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
+                else max(forward_batch.global_num_tokens_cpu)
+            )
+        else:
+            cuda_graph_bs = forward_batch.batch_size
+
+        is_bs_supported = (
+            self.backend.can_run(forward_batch, cuda_graph_bs)
+            if self.disable_padding
+            else cuda_graph_bs <= self.max_bs
+        )
+
+        if self.require_mlp_sync:
+            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+
+        return is_bs_supported
+
+    # -----------------------------------------------------------------
+    # Capture (per-shape)
+    # -----------------------------------------------------------------
+    def capture_one_shape(
+        self,
+        size: int,
+        forward: Callable,
+        stream_idx: Optional[int] = None,
+        variant_label: Optional[str] = None,
+    ):
+        num_seqs = size  # EAGLE legacy name
+        buffers = self.buffers
+        num_tokens = num_seqs * self.num_tokens_per_bs
+
+        # Graph inputs
+        req_pool_indices = buffers.req_pool_indices[:num_seqs]
+        seq_lens = buffers.seq_lens[:num_seqs]
+        seq_lens_cpu = buffers.seq_lens_cpu[:num_seqs]
+        extend_seq_lens = buffers.extend_seq_lens[:num_seqs]
+        extend_seq_lens_cpu = self.extend_seq_lens_cpu[:num_seqs]
+        out_cache_loc = buffers.out_cache_loc[: num_tokens * self.speculative_num_steps]
+        positions = buffers.positions[:num_tokens]
+        mrope_positions = buffers.mrope_positions[:, :num_tokens]
+        rids_int = buffers.rids_int[:num_seqs] if buffers.rids_int is not None else None
+        bootstrap_room_ids_int = (
+            buffers.bootstrap_room_ids_int[:num_seqs]
+            if buffers.bootstrap_room_ids_int is not None
+            else None
+        )
+        hidden_states = (
+            buffers.hidden_states[:num_seqs]
+            if buffers.hidden_states is not None
+            else None
+        )
+        topk_p = buffers.topk_p[:num_seqs]
+        topk_index = buffers.topk_index[:num_seqs]
+        draft_probs = (
+            buffers.draft_probs[:num_seqs] if buffers.draft_probs is not None else None
+        )
+
+        if self.require_mlp_tp_gather:
+            global_num_tokens_cpu = [num_tokens] * self.dp_size
+        elif self.require_attn_tp_gather:
+            global_num_tokens_cpu = [num_tokens]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            num_tokens_tensor = torch.tensor(
+                global_num_tokens_cpu,
+                dtype=torch.int32,
+                device=buffers.input_ids.device,
+            )
+            buffers.global_num_tokens_gpu.copy_(num_tokens_tensor)
+            buffers.global_num_tokens_for_logprob_gpu.copy_(num_tokens_tensor)
+            global_num_tokens = buffers.global_num_tokens_gpu
+            global_num_tokens_for_logprob = buffers.global_num_tokens_for_logprob_gpu
+        else:
+            global_dp_buffer_len = None
+            global_num_tokens = None
+            global_num_tokens_for_logprob = None
+
+        capture_mode = (
+            CaptureHiddenMode.NULL
+            if self.model_runner.spec_algorithm.is_standalone()
+            else CaptureHiddenMode.LAST
+        )
+        spec_info = EagleDraftInput(
+            topk_p=topk_p,
+            topk_index=topk_index,
+            draft_probs=draft_probs,
+            hidden_states=hidden_states,
+            capture_hidden_mode=capture_mode,
+        )
+        if self.buffers.dsa_seed_topk is not None:
+            spec_info.dsa_topk_indices = self.buffers.dsa_seed_topk[:num_seqs]
+
+        sampling_info = SamplingBatchInfo(
+            temperatures=self.temperatures[:num_seqs],
+            top_ps=torch.ones((num_seqs,), dtype=torch.float),
+            top_ks=torch.full((num_seqs,), -1, dtype=torch.int32),
+            min_ps=torch.zeros((num_seqs,), dtype=torch.float),
+            is_all_greedy=False,
+            is_any_greedy=False,
+            need_top_p_sampling=False,
+            need_top_k_sampling=False,
+            need_min_p_sampling=False,
+            vocab_size=self.model_runner.model_config.vocab_size,
+        )
+
+        forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.DECODE,
+            batch_size=num_seqs,
+            input_ids=None,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            extend_seq_lens=extend_seq_lens,
+            extend_seq_lens_cpu=extend_seq_lens_cpu,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=seq_lens.sum().item(),
+            return_logprob=False,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            global_num_tokens_gpu=global_num_tokens,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
+            spec_algorithm=self.model_runner.spec_algorithm,
+            spec_info=spec_info,
+            sampling_info=sampling_info,
+            rids_int=rids_int,
+            bootstrap_room_ids_int=bootstrap_room_ids_int,
+            capture_hidden_mode=(
+                spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
+            ),
+        )
+
+        def run_once():
+            self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
+
+            forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(
+                global_dp_buffer_len,
+                num_tokens,
+                forward_batch.dp_padding_mode.is_max_len(),
+                global_num_tokens_cpu,
+            )
+            set_is_extend_in_batch(False)
+
+            output_cache_loc_backup = forward_batch.out_cache_loc
+            hidden_states_backup = forward_batch.spec_info.hidden_states
+
+            ret = self.eagle_worker.draft_forward(forward_batch)
+
+            forward_batch.out_cache_loc = output_cache_loc_backup
+            forward_batch.spec_info.hidden_states = hidden_states_backup
+            forward_batch.positions.sub_(self.eagle_worker.speculative_num_steps - 1)
+            return ret
+
+        with forward_context(ForwardContext(attn_backend=self.draft_attn_backend)):
+            self.draft_attn_backend.init_forward_metadata_out_graph(
+                forward_batch, in_capture=True
+            )
+            # The capture batch is planned here (out-of-forward), so the
+            # per-step forwards inside draft_forward must not re-plan.
+            forward_batch.mark_forward_metadata_ready()
+            self.deepep_adapter.capture(is_extend_in_batch=False)
+            shape_key = self._make_graph_key(num_seqs)
+            post_warmup_hook = getattr(
+                self.draft_attn_backend, "on_after_cuda_graph_warmup", None
+            )
+            maybe_flashinfer_autotune_speculative_draft(
+                self,
+                run_once,
+                post_warmup_hook=post_warmup_hook,
+                skip_logits=False,
+            )
+            self.backend.capture_one(
+                shape_key,
+                run_once,
+                dummies=None,
+                post_warmup_hook=post_warmup_hook,
+            )
+
+    def _postprocess_output_to_raw_bs(self, out, raw_bs):
+        parent_list, top_scores_index, draft_tokens, draft_probs = (
+            t[:raw_bs] if t is not None else None for t in out
+        )
+        return parent_list, top_scores_index, draft_tokens, draft_probs
+
+    # -----------------------------------------------------------------
+    # Replay
+    # -----------------------------------------------------------------
+    def execute(self, forward_batch: ForwardBatch):
+        assert forward_batch.out_cache_loc is not None
+        self.deepep_adapter.replay()
+        buffers = self.buffers
+
+        raw_bs = forward_batch.batch_size
+        raw_num_token = raw_bs * self.num_tokens_per_bs
+
+        # Pad to nearest captured shape
+        if self.require_mlp_tp_gather:
+            max_num_tokens = max(forward_batch.global_num_tokens_cpu)
+            max_batch_size = (
+                max_num_tokens // self.num_tokens_per_bs
+                if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
+                else max_num_tokens
+            )
+            bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
+        else:
+            bs = self._pad_to_bucket(raw_bs, self.capture_bs)
+
+        if bs != raw_bs:
+            buffers.seq_lens.fill_(self.seq_len_fill_value)
+            buffers.out_cache_loc.zero_()
+            buffers.positions.zero_()
+            if buffers.rids_int is not None:
+                buffers.rids_int.zero_()
+            if buffers.bootstrap_room_ids_int is not None:
+                buffers.bootstrap_room_ids_int.fill_(-1)
+            buffers.topk_p.zero_()
+            buffers.topk_index.zero_()
+            if buffers.draft_probs is not None:
+                buffers.draft_probs.zero_()
+            if buffers.hidden_states is not None:
+                buffers.hidden_states.zero_()
+            if buffers.dsa_seed_topk is not None:
+                buffers.dsa_seed_topk.zero_()
+            buffers.req_pool_indices.zero_()
+
+        num_tokens = bs * self.num_tokens_per_bs
+
+        maybe_detect_nan(
+            forward_batch.spec_info.topk_p,
+            "EagleDraftCudaGraphRunner.replay: topk_p",
+        )
+        maybe_detect_oob(
+            forward_batch.spec_info.topk_index,
+            0,
+            self.model_runner.model_config.vocab_size,
+            "EagleDraftCudaGraphRunner.replay: topk_index vs vocab_size="
+            f"{self.model_runner.model_config.vocab_size}",
+        )
+
+        # Common inputs — batch the small per-field device copies into a grouped
+        # foreach copy (one foreach call per dtype pair) to cut launch overhead.
+        # hidden_states is handled separately below (see note), and seq_lens_cpu
+        # is handled further down since it lives on host.
+        copy_dsts = [
+            buffers.seq_lens[:raw_bs],
+            buffers.out_cache_loc[: raw_num_token * self.speculative_num_steps],
+            buffers.positions[:raw_num_token],
+            buffers.topk_p[:raw_bs],
+            buffers.topk_index[:raw_bs],
+            buffers.req_pool_indices[:raw_bs],
+        ]
+        copy_srcs = [
+            forward_batch.seq_lens,
+            forward_batch.out_cache_loc,
+            forward_batch.positions,
+            forward_batch.spec_info.topk_p,
+            forward_batch.spec_info.topk_index,
+            forward_batch.req_pool_indices,
+        ]
+        if buffers.rids_int is not None and forward_batch.rids_int is not None:
+            copy_dsts.append(buffers.rids_int[:raw_bs])
+            copy_srcs.append(forward_batch.rids_int)
+        if (
+            buffers.bootstrap_room_ids_int is not None
+            and forward_batch.bootstrap_room_ids_int is not None
+        ):
+            copy_dsts.append(buffers.bootstrap_room_ids_int[:raw_bs])
+            copy_srcs.append(forward_batch.bootstrap_room_ids_int)
+        _grouped_foreach_copy_(copy_dsts, copy_srcs)
+
+        # hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync
+        # DMA engine; foreach would force the ~3x slower compute-kernel copy.
+        if (
+            buffers.draft_probs is not None
+            and forward_batch.spec_info.draft_probs is not None
+        ):
+            buffers.draft_probs[:raw_bs].copy_(forward_batch.spec_info.draft_probs)
+        if (
+            buffers.hidden_states is not None
+            and forward_batch.spec_info.hidden_states is not None
+        ):
+            buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
+        if buffers.dsa_seed_topk is not None:
+            seed = forward_batch.spec_info.dsa_topk_indices
+            if seed is not None:
+                buffers.dsa_seed_topk[:raw_bs].copy_(seed)
+            else:
+                buffers.dsa_seed_topk[:raw_bs].zero_()
+        # Only rejection sampling reads temperatures (renorm_draft_probs); skip
+        # the copy otherwise to keep the non-RS path free of extra work.
+        if (
+            self.model_runner.server_args.speculative_use_rejection_sampling
+            and forward_batch.sampling_info is not None
+        ):
+            self.temperatures[:raw_bs].copy_(
+                forward_batch.sampling_info.temperatures[:raw_bs]
+            )
+
+        # TODO(ch-wan): support num_token_non_padded
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
+
+        # Save the raw seq_lens_sum; it is restored after replay. While the graph
+        # runs it must reflect the padded fake rows (set below), since draft decode
+        # backends read seq_lens_sum to size/slice kv_indices.
+        raw_seq_lens_sum = forward_batch.seq_lens_sum
+
+        if bs != raw_bs:
+            forward_batch.batch_size = bs
+            forward_batch.seq_lens = buffers.seq_lens[:bs]
+            forward_batch.req_pool_indices = buffers.req_pool_indices[:bs]
+            forward_batch.positions = buffers.positions[:num_tokens]
+            if raw_seq_lens_sum is not None:
+                forward_batch.seq_lens_sum = (
+                    raw_seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
+                )
+            if buffers.rids_int is not None and forward_batch.rids_int is not None:
+                forward_batch.rids_int = buffers.rids_int[:bs]
+            if (
+                buffers.bootstrap_room_ids_int is not None
+                and forward_batch.bootstrap_room_ids_int is not None
+            ):
+                forward_batch.bootstrap_room_ids_int = buffers.bootstrap_room_ids_int[
+                    :bs
+                ]
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+            buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+            forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:bs]
+
+        # forward_batch.batch_size was overwritten to bs above when padding.
+        self.draft_attn_backend.init_forward_metadata_out_graph(forward_batch)
+        self.raw_bs = raw_bs
+        self.bs = bs
+
+        # Replay via backend
+        shape_key = self._make_graph_key(bs)
+        timer_ctx = (
+            self.model_runner.device_timer.wrap(metadata={"category": "eagle_draft"})
+            if self.model_runner.device_timer
+            else contextlib.nullcontext()
+        )
+        with timer_ctx:
+            out = self._replay_graph(shape_key, forward_batch)
+
+        if bs != raw_bs:
+            out = self._postprocess_output_to_raw_bs(out, raw_bs)
+            forward_batch.batch_size = raw_bs
+            forward_batch.positions = buffers.positions[:raw_num_token]
+            forward_batch.seq_lens = buffers.seq_lens[:raw_bs]
+            forward_batch.req_pool_indices = buffers.req_pool_indices[:raw_bs]
+            if buffers.rids_int is not None and forward_batch.rids_int is not None:
+                forward_batch.rids_int = buffers.rids_int[:raw_bs]
+            if (
+                buffers.bootstrap_room_ids_int is not None
+                and forward_batch.bootstrap_room_ids_int is not None
+            ):
+                forward_batch.bootstrap_room_ids_int = buffers.bootstrap_room_ids_int[
+                    :raw_bs
+                ]
+            if forward_batch.seq_lens_cpu is not None:
+                forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:raw_bs]
+            forward_batch.seq_lens_sum = raw_seq_lens_sum
+
+        return out

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -0,0 +1,617 @@
+from __future__ import annotations
+
+import contextlib
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+from sglang.srt.model_executor.runner import (
+    DecodeCudaGraphRunner,
+    DeepEPCudaGraphRunnerAdapter,
+    ShapeKey,
+    _grouped_foreach_copy_,
+    get_batch_sizes_to_capture,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
+from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+from sglang.srt.runtime_context import get_flags
+from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
+from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim
+from sglang.srt.speculative.spec_utils import fast_topk
+from sglang.srt.utils import (
+    is_hip,
+    require_attn_tp_gather,
+    require_gathered_buffer,
+    require_mlp_sync,
+    require_mlp_tp_gather,
+)
+
+_is_hip = is_hip()
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+@dataclass
+class EagleDraftExtendInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    req_pool_indices: torch.Tensor
+    out_cache_loc: torch.Tensor
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    hidden_states: Optional[torch.Tensor]
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    extend_seq_lens: torch.Tensor
+    num_correct_drafts: torch.Tensor
+    num_accept_tokens: torch.Tensor
+    next_token_logits_buffer: torch.Tensor
+    global_num_tokens_gpu: Optional[torch.Tensor]
+    global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+    dsa_seed_topk_capture: Optional[torch.Tensor] = None
+
+
+class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
+    """EAGLE draft-extend cuda-graph runner.
+
+    Subclasses DecodeCudaGraphRunner to inherit the outer capture
+    loop + backend scaffolding. Overrides capture_one_shape,
+    replay, can_run_graph for EAGLE-specific draft-extend semantics.
+    """
+
+    def __init__(
+        self,
+        eagle_worker: EagleDraftWorker,
+        *,
+        draft_extend_attn_backend=None,
+        speculative_num_steps: Optional[int] = None,
+    ):
+        # Parse args
+        self.eagle_worker = eagle_worker
+        self.model_runner = model_runner = eagle_worker.draft_runner
+        self.forward_mode = ForwardMode.DRAFT_EXTEND_V2
+
+        # Fields the parent's capture() reads:
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.tp_size = model_runner.tp_size
+        self.dp_size = model_runner.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.enable_torch_compile = get_flags().capture.enable_torch_compile
+        self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
+        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
+        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
+        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
+        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+        self.enable_profile_cuda_graph = (
+            model_runner.server_args.enable_profile_cuda_graph
+        )
+        self.speculative_num_steps = (
+            model_runner.server_args.speculative_num_steps
+            if speculative_num_steps is None
+            else speculative_num_steps
+        )
+        self.topk = model_runner.server_args.speculative_eagle_topk
+        self.draft_extend_attn_backend = (
+            draft_extend_attn_backend or eagle_worker.draft_extend_attn_backend
+        )
+        self.attn_backend = self.draft_extend_attn_backend
+
+        # Disable parent paths that don't apply.
+        self.compile_bs = []
+        self.enable_pdmux = False
+        self.record_nolora_graph = False
+        self.is_dllm = False
+
+        self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
+
+        self.capture_forward_mode = self.forward_mode
+        self.capture_hidden_mode = CaptureHiddenMode.LAST
+
+        self.capture_bs, _ = get_batch_sizes_to_capture(model_runner)
+        self.padded_static_len = -1
+
+        # Size cuda-graph buffers by num_draft_tokens (full tree width), not
+        # num_steps + 1, or topk > 1 draft-extend overflows them.
+        self.num_tokens_per_bs = model_runner.server_args.speculative_num_draft_tokens
+        self.max_bs = max(self.capture_bs)
+        self.max_num_token = self.max_bs * self.num_tokens_per_bs
+
+        self.draft_extend_attn_backend.init_cuda_graph_state(
+            self.max_bs, self.max_num_token
+        )
+        self.seq_len_fill_value = (
+            self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value()
+        )
+        self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs
+
+        if self.enable_torch_compile:
+            set_torch_compile_config()
+
+        with torch.device(model_runner.device):
+            input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int64)
+            out_cache_loc = torch.ones(
+                (self.max_num_token,), dtype=self._cache_loc_dtype()
+            )
+            positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
+
+            # Width and dtype both come from the draft `model_runner` so the
+            # source stays consistent (the draft dtype matches the target dtype
+            # that produced these hidden states).
+            _hidden_dtype = model_runner.model_config.dtype
+            _hidden_size = (
+                None
+                if self.eagle_worker.speculative_algorithm.is_standalone()
+                else get_draft_input_from_target_hidden_dim(model_runner)
+            )
+            hidden_states = (
+                torch.zeros(
+                    (self.max_num_token, _hidden_size),
+                    dtype=_hidden_dtype,
+                )
+                if _hidden_size is not None
+                else None
+            )
+            self.seq_len_fill_value = (
+                self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value()
+            )
+            seq_lens = torch.full(
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
+            )
+            extend_seq_lens = torch.full(
+                (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32
+            )
+            num_correct_drafts = torch.full(
+                (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32
+            )
+            num_accept_tokens = torch.full(
+                (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32
+            )
+
+            if self.require_gathered_buffer:
+                if self.require_mlp_tp_gather:
+                    global_num_tokens_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                else:
+                    assert self.require_attn_tp_gather
+                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (1,), dtype=torch.int32
+                    )
+            else:
+                global_num_tokens_gpu = None
+                global_num_tokens_for_logprob_gpu = None
+
+            hot_token_id = getattr(self.eagle_worker, "hot_token_id", None)
+            if hasattr(
+                self.model_runner.model_config.hf_config, "draft_vocab_size"
+            ):  # llama_eagle
+                vocab_size = self.model_runner.model_config.hf_config.draft_vocab_size
+            elif hasattr(
+                self.model_runner.model_config.hf_config, "hot_vocab_size"
+            ):  # llama_eagle3
+                vocab_size = self.model_runner.model_config.hf_config.hot_vocab_size
+            elif hot_token_id is not None:
+                # FR-Spec: reduced vocab is injected via a late
+                # json_model_override_args, so hf_config lacks it; size from the head.
+                vocab_size = len(hot_token_id)
+            else:
+                vocab_size = self.model_runner.model_config.vocab_size
+
+            next_token_logits_buffer = (
+                self.model_runner.graph_shared_output.get_logits_buffer(
+                    vocab_size, rows=self.max_bs * self.num_tokens_per_bs
+                )
+            )
+
+        seq_lens_cpu = torch.full(
+            (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
+        )
+
+        dsa_seed_topk_capture = (
+            torch.full(
+                (self.max_num_token, self.eagle_worker.dsa_index_topk),
+                -1,
+                dtype=torch.int32,
+                device=model_runner.device,
+            )
+            if self.eagle_worker.seed_dsa_topk_from_draft_extend
+            else None
+        )
+
+        self.buffers = EagleDraftExtendInputBuffers(
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            out_cache_loc=out_cache_loc,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            hidden_states=hidden_states,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            extend_seq_lens=extend_seq_lens,
+            num_correct_drafts=num_correct_drafts,
+            num_accept_tokens=num_accept_tokens,
+            next_token_logits_buffer=next_token_logits_buffer,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            dsa_seed_topk_capture=dsa_seed_topk_capture,
+        )
+        self.buffers.share_buffers()
+
+        self.backend = resolve_decode_backend(self)
+
+        try:
+            with model_capture_mode():
+                self.capture()
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+            )
+
+    def _replay_graph(self, shape_key, forward_batch):
+        return self.backend.replay(shape_key, forward_batch)
+
+    def _cache_loc_dtype(self):
+        return torch.int64
+
+    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
+        return ShapeKey(size=bs)
+
+    def can_run_graph(self, forward_batch: ForwardBatch):
+        if self.require_mlp_tp_gather:
+            cuda_graph_bs = (
+                max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
+                if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
+                else max(forward_batch.global_num_tokens_cpu)
+            )
+        else:
+            cuda_graph_bs = forward_batch.seq_lens.numel()
+
+        is_bs_supported = (
+            self.backend.can_run(forward_batch, cuda_graph_bs)
+            if self.disable_padding
+            else cuda_graph_bs <= self.max_bs
+        )
+
+        if self.require_mlp_sync:
+            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+
+        return is_bs_supported
+
+    def capture_one_shape(
+        self,
+        size: int,
+        forward: Callable,
+        stream_idx: Optional[int] = None,
+        variant_label: Optional[str] = None,
+    ):
+        bs = size
+        buffers = self.buffers
+        num_tokens = bs * self.num_tokens_per_bs
+
+        # Graph inputs
+        input_ids = buffers.input_ids[:num_tokens]
+        req_pool_indices = buffers.req_pool_indices[:bs]
+        seq_lens = buffers.seq_lens[:bs]
+        seq_lens_cpu = buffers.seq_lens_cpu[:bs]
+        extend_seq_lens = buffers.extend_seq_lens[:bs]
+        extend_seq_lens_cpu = self.extend_seq_lens_cpu[:bs]
+        out_cache_loc = buffers.out_cache_loc[:num_tokens]
+        positions = buffers.positions[:num_tokens]
+        mrope_positions = buffers.mrope_positions[:, :num_tokens]
+        hidden_states = (
+            buffers.hidden_states[:num_tokens]
+            if buffers.hidden_states is not None
+            else None
+        )
+        num_correct_drafts = buffers.num_correct_drafts[:bs]
+        num_accept_tokens = buffers.num_accept_tokens[:bs]
+        next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens]
+
+        # pruned_states = num_tokens (all tokens)
+        num_tokens_for_logprob = num_tokens
+
+        if self.require_mlp_tp_gather:
+            global_num_tokens_cpu = [num_tokens] * self.dp_size
+        elif self.require_attn_tp_gather:
+            global_num_tokens_cpu = [num_tokens]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            buffers.global_num_tokens_gpu.copy_(
+                torch.tensor(
+                    global_num_tokens_cpu,
+                    dtype=torch.int32,
+                    device=buffers.input_ids.device,
+                )
+            )
+            buffers.global_num_tokens_for_logprob_gpu.copy_(
+                torch.tensor(
+                    [num_tokens_for_logprob] * len(global_num_tokens_cpu),
+                    dtype=torch.int32,
+                    device=buffers.input_ids.device,
+                )
+            )
+        else:
+            global_dp_buffer_len = None
+
+        spec_info = EagleDraftExtendInput(
+            hidden_states=hidden_states,
+            num_correct_drafts=num_correct_drafts,
+            num_accept_tokens=num_accept_tokens,
+            # Padded tree width per req; drives the constant qo layout.
+            num_tokens_per_req=self.num_tokens_per_bs,
+        )
+
+        forward_batch = ForwardBatch(
+            forward_mode=self.forward_mode,
+            batch_size=bs,
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            next_token_logits_buffer=next_token_logits_buffer,
+            extend_seq_lens=extend_seq_lens,
+            extend_seq_lens_cpu=extend_seq_lens_cpu,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=seq_lens.sum().item(),
+            return_logprob=False,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            global_num_tokens_gpu=buffers.global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
+            spec_algorithm=self.model_runner.spec_algorithm,
+            spec_info=spec_info,
+            capture_hidden_mode=CaptureHiddenMode.LAST,
+            padded_static_len=self.padded_static_len,
+        )
+
+        if self.buffers.dsa_seed_topk_capture is not None:
+            spec_info.dsa_seed_topk_capture = self.buffers.dsa_seed_topk_capture[
+                :num_tokens
+            ]
+
+        def run_once():
+            self.draft_extend_attn_backend.init_forward_metadata_in_graph(forward_batch)
+
+            # Clean intermediate result cache for DP attention
+            forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(
+                global_dp_buffer_len,
+                num_tokens,
+                forward_batch.dp_padding_mode.is_max_len(),
+                global_num_tokens_cpu,
+            )
+            set_is_extend_in_batch(False)
+
+            output_cache_loc_backup = forward_batch.out_cache_loc
+            hidden_states_backup = forward_batch.spec_info.hidden_states
+
+            ret = self.model_runner.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+            )
+            # ROCm's argmax tie-breaks differently from CUDA's softmax+max
+            # path on FP8 logits, which corrupts MTP draft selection on AMD.
+            # Keep the fastpath CUDA-only.
+            if self.topk == 1 and not _is_hip:
+                ret.topk_index = torch.argmax(
+                    ret.next_token_logits, dim=-1, keepdim=True
+                )
+                ret.topk_p = torch.ones_like(ret.topk_index, dtype=torch.float32)
+            else:
+                probs = torch.softmax(ret.next_token_logits, dim=-1)
+                ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
+
+            forward_batch.out_cache_loc = output_cache_loc_backup
+            forward_batch.spec_info.hidden_states = hidden_states_backup
+            return ret
+
+        with forward_context(
+            ForwardContext(attn_backend=self.draft_extend_attn_backend)
+        ):
+            self.draft_extend_attn_backend.init_forward_metadata_out_graph(
+                forward_batch, in_capture=True
+            )
+            self.deepep_adapter.capture(is_extend_in_batch=True)
+            canary_ctx = (
+                c.with_active_single_forward_manager(0)
+                if (c := self.model_runner.canary_manager) is not None
+                else contextlib.nullcontext()
+            )
+            with canary_ctx:
+                shape_key = self._make_graph_key(bs)
+                post_warmup_hook = getattr(
+                    self.draft_extend_attn_backend,
+                    "on_after_cuda_graph_warmup",
+                    None,
+                )
+                maybe_flashinfer_autotune_speculative_draft(
+                    self,
+                    run_once,
+                    post_warmup_hook=post_warmup_hook,
+                    skip_logits=False,
+                )
+                self.backend.capture_one(
+                    shape_key,
+                    run_once,
+                    dummies=None,
+                    post_warmup_hook=post_warmup_hook,
+                )
+
+    def execute(self, forward_batch: ForwardBatch):
+        assert forward_batch.out_cache_loc is not None
+        self.deepep_adapter.replay()
+        buffers = self.buffers
+
+        raw_bs = forward_batch.batch_size
+        num_tokens = forward_batch.input_ids.shape[0]
+        if self.require_mlp_tp_gather:
+            max_num_tokens = max(forward_batch.global_num_tokens_cpu)
+            max_batch_size = (
+                max_num_tokens // self.num_tokens_per_bs
+                if self.model_runner.spec_algorithm.is_eagle()
+                else max_num_tokens
+            )
+            bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
+        else:
+            bs = self._pad_to_bucket(raw_bs, self.capture_bs)
+
+        if bs * self.num_tokens_per_bs != num_tokens:
+            buffers.seq_lens.fill_(self.seq_len_fill_value)
+            buffers.out_cache_loc.zero_()
+            buffers.positions.zero_()
+            # Pair with seq_lens fill: padded rows must point at reserved
+            # req_pool slot 0 (req_to_token[0, :] is all zeros from init).
+            buffers.req_pool_indices.zero_()
+            buffers.num_correct_drafts.fill_(self.num_tokens_per_bs)
+            buffers.num_accept_tokens.fill_(self.num_tokens_per_bs)
+            buffers.extend_seq_lens.fill_(self.num_tokens_per_bs)
+
+        # Batch the small per-field device copies into a grouped foreach copy
+        # (one foreach call per dtype pair) to cut launch overhead. hidden_states
+        # is handled separately below (see note), and seq_lens_cpu is handled
+        # further down since it lives on host.
+        copy_dsts = [
+            buffers.input_ids[:num_tokens],
+            buffers.seq_lens[:raw_bs],
+            buffers.out_cache_loc[:num_tokens],
+            buffers.positions[:num_tokens],
+            buffers.req_pool_indices[:raw_bs],
+        ]
+        copy_srcs = [
+            forward_batch.input_ids,
+            forward_batch.seq_lens,
+            forward_batch.out_cache_loc,
+            forward_batch.positions,
+            forward_batch.req_pool_indices,
+        ]
+        if forward_batch.extend_seq_lens is not None:
+            copy_dsts.append(buffers.extend_seq_lens[:raw_bs])
+            copy_srcs.append(forward_batch.extend_seq_lens)
+        else:
+            buffers.extend_seq_lens[:raw_bs].fill_(self.num_tokens_per_bs)
+        if forward_batch.spec_info.num_correct_drafts is not None:
+            copy_dsts.append(buffers.num_correct_drafts[:raw_bs])
+            copy_srcs.append(forward_batch.spec_info.num_correct_drafts)
+            copy_dsts.append(buffers.num_accept_tokens[:raw_bs])
+            copy_srcs.append(forward_batch.spec_info.num_accept_tokens)
+        _grouped_foreach_copy_(copy_dsts, copy_srcs)
+
+        # hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync
+        # DMA engine; foreach would force the ~3x slower compute-kernel copy.
+        if (
+            buffers.hidden_states is not None
+            and forward_batch.spec_info.hidden_states is not None
+            and forward_batch.spec_info.hidden_states.shape[1]
+            == buffers.hidden_states.shape[1]
+        ):
+            buffers.hidden_states[:num_tokens].copy_(
+                forward_batch.spec_info.hidden_states
+            )
+
+        # TODO(ch-wan): support num_token_non_padded
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+            buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+        if forward_batch.extend_seq_lens_cpu is not None:
+            self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu
+        else:
+            self.extend_seq_lens_cpu[:raw_bs] = [self.num_tokens_per_bs] * raw_bs
+        if bs > raw_bs:
+            self.extend_seq_lens_cpu[raw_bs:bs] = [self.num_tokens_per_bs] * (
+                bs - raw_bs
+            )
+        forward_batch.spec_info.extend_seq_lens_cpu = list(
+            self.extend_seq_lens_cpu[:bs]
+        )
+        forward_batch.spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs]
+
+        if bs != raw_bs:
+            forward_batch.spec_info.positions = buffers.positions[:num_tokens]
+            forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[:bs]
+            forward_batch.spec_info.num_accept_tokens = buffers.num_accept_tokens[:bs]
+
+        from types import SimpleNamespace
+
+        seq_lens_sum = forward_batch.seq_lens_sum
+        if seq_lens_sum is not None:
+            seq_lens_sum = seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
+        fb_view = SimpleNamespace(
+            batch_size=bs,
+            forward_mode=self.forward_mode,
+            input_ids=getattr(forward_batch, "input_ids", None),
+            req_pool_indices=buffers.req_pool_indices,
+            seq_lens=buffers.seq_lens,
+            seq_lens_sum=seq_lens_sum,
+            seq_lens_cpu=buffers.seq_lens_cpu,
+            encoder_lens=None,
+            out_cache_loc=buffers.out_cache_loc[:num_tokens],
+            out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
+            spec_info=forward_batch.spec_info,
+        )
+        self.draft_extend_attn_backend.init_forward_metadata_out_graph(fb_view)
+
+        # Snapshot built -- the forward is done reading the shared pool. Publish
+        # a read-done event the scheduler's WAR barrier waits on.
+        read_done = self.device_module.Event()
+        read_done.record()
+        self.model_runner.war_fastpath_read_done_event = read_done
+
+        self.raw_bs = raw_bs
+        self.bs = bs
+        shape_key = self._make_graph_key(bs)
+        timer_ctx = (
+            self.model_runner.device_timer.wrap(
+                metadata={"category": "eagle_draft_extend"}
+            )
+            if self.model_runner.device_timer
+            else contextlib.nullcontext()
+        )
+        with timer_ctx:
+            out = self._replay_graph(shape_key, forward_batch)
+
+        out = LogitsProcessorOutput(
+            next_token_logits=out.next_token_logits[:num_tokens],
+            hidden_states=out.hidden_states[:num_tokens],
+        )
+        return out

python/sglang/srt/speculative/eagle_info.py

@@ -0,0 +1,406 @@
+import logging
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+
+from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EagleVerifyInput(SpecInput):
+    draft_token: torch.Tensor
+    custom_mask: torch.Tensor
+    positions: torch.Tensor
+    retrieve_index: torch.Tensor
+    retrieve_next_token: torch.Tensor
+    retrieve_next_sibling: torch.Tensor
+    retrieve_cum_len: torch.Tensor
+    spec_steps: int
+    topk: int
+    draft_token_num: int
+    capture_hidden_mode: CaptureHiddenMode
+    seq_lens_sum: int
+    seq_lens_cpu: torch.Tensor
+    grammar: BaseGrammarObject = None
+    # Stacked per-step draft proposal distribution q, shape (bs, num_steps,
+    # vocab); only set under rejection sampling. Consumed by the verify kernel.
+    draft_probs: torch.Tensor = None
+
+    # Shape info for padding
+    num_tokens_per_req: int = -1  # -1 auto-fills from draft_token_num.
+
+    def __post_init__(self):
+        super().__init__(SpecInputType.EAGLE_VERIFY)
+        if self.num_tokens_per_req < 0:
+            self.num_tokens_per_req = self.draft_token_num
+
+    @property
+    def max_tree_depth(self) -> int:
+        """Longest root-to-leaf chain of the verify tree, incl. the root;
+        bounds the accept_index row width. EAGLE trees are depth-bounded by
+        the draft loop. Algorithms with other tree shapes override this."""
+        return self.spec_steps + 1
+
+    @property
+    def tree_topk(self) -> int:
+        """Branching factor passed to the tree-verify kernels; -1 means an
+        irregular tree (no fixed per-level branching)."""
+        return self.topk
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        return self.draft_token_num, self.draft_token_num
+
+    @classmethod
+    def create_idle_input(
+        cls, topk: int, spec_steps: int, num_verify_tokens: int, device: str
+    ):
+        return cls(
+            draft_token=torch.empty((0,), dtype=torch.long, device=device),
+            custom_mask=torch.full((0,), True, dtype=torch.bool, device=device),
+            positions=torch.empty((0,), dtype=torch.int64, device=device),
+            retrieve_index=torch.full(
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
+            ),
+            retrieve_next_token=torch.full(
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
+            ),
+            retrieve_next_sibling=torch.full(
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
+            ),
+            retrieve_cum_len=None,
+            topk=topk,
+            draft_token_num=num_verify_tokens,
+            spec_steps=spec_steps,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+            seq_lens_sum=0,
+            seq_lens_cpu=torch.empty((0,), dtype=torch.int64),
+        )
+
+    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,
+    ):
+        device = req_pool_indices.device
+        batch_size = len(req_pool_indices)
+        qo_indptr = torch.arange(
+            0,
+            (1 + batch_size) * self.draft_token_num,
+            step=self.draft_token_num,
+            dtype=torch.int32,
+            device=device,
+        )
+        cum_kv_seq_len = torch.zeros(
+            (batch_size + 1,), dtype=torch.int32, device=device
+        )
+
+        paged_kernel_lens = paged_kernel_lens + self.draft_token_num
+        cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum + self.draft_token_num * batch_size,
+            dtype=torch.int32,
+            device=device,
+        )
+        create_flashinfer_kv_indices_triton[(batch_size,)](
+            req_to_token,
+            req_pool_indices,
+            paged_kernel_lens,
+            cum_kv_seq_len,
+            None,
+            kv_indices,
+            req_to_token.size(1),
+        )
+        mask_numel = (
+            paged_kernel_lens_sum * self.draft_token_num
+            + (self.draft_token_num**2) * batch_size
+        )
+        if self.custom_mask.numel() < mask_numel:
+            # FIXME(attn): temporary fix for custom mask padding with cuda graph
+            self.custom_mask = torch.cat(
+                [
+                    self.custom_mask,
+                    torch.full(
+                        (mask_numel - self.custom_mask.numel(),),
+                        True,
+                        dtype=torch.bool,
+                        device=device,
+                    ),
+                ],
+                dim=0,
+            )
+
+        return kv_indices, cum_kv_seq_len, qo_indptr, self.custom_mask
+
+
+@dataclass
+class EagleDraftInput(SpecInput):
+    # For idle stubs use `create_idle_input`, not the bare ctor: `filter_batch`
+    # / `merge_batch` slice / cat `topk_p` / `topk_index` / `hidden_states` /
+    # `bonus_tokens` unconditionally.
+
+    # shape: (b, topk)
+    topk_p: torch.Tensor = None
+    topk_index: torch.Tensor = None
+    # Draft proposal q from draft-extend, only set under rejection sampling:
+    # (b, vocab) single-layer; (b, num_steps, vocab) multi-layer chain.
+    draft_probs: torch.Tensor = None
+    # shape: (b, hidden_size) - one hidden per req, consumed by `draft` forward.
+    # None when the spec algorithm's draft doesn't read hidden_states
+    # (e.g., STANDALONE — vanilla LLM draft).
+    hidden_states: Optional[torch.Tensor] = None
+    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
+
+    # Survives across draft steps: spec_info is shared by reference across the
+    # per-step forwards (each runs on a copied ForwardBatch, dropping writebacks).
+    dsa_topk_indices: Optional[torch.Tensor] = None
+
+    # Per-req bonus token (the "+1" target prediction at end of each accept
+    # chain); the worker copies it here post-extend for next iter's draft.
+    bonus_tokens: torch.Tensor = None
+
+    # shape: (b + 1,)
+    kv_indptr: torch.Tensor = None
+    kv_indices: torch.Tensor = None
+
+    num_tokens_per_req: int = -1
+    num_tokens_for_logprob_per_req: int = -1
+
+    # V2 overlap worker only: req_pool_indices used as buf slot keys.
+    future_indices: Optional[torch.Tensor] = None
+
+    def __post_init__(self):
+        super().__init__(SpecInputType.EAGLE_DRAFT)
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        return self.num_tokens_per_req, self.num_tokens_for_logprob_per_req
+
+    @classmethod
+    def create_idle_input(
+        cls,
+        device: torch.device,
+        hidden_size: Optional[int],
+        dtype: Optional[torch.dtype],
+        topk: int,
+        capture_hidden_mode: CaptureHiddenMode,
+        vocab_size: int = 0,
+    ):
+        return cls(
+            bonus_tokens=torch.empty((0,), device=device, dtype=torch.int32),
+            hidden_states=(
+                torch.empty((0, hidden_size), device=device, dtype=dtype)
+                if hidden_size is not None
+                else None
+            ),
+            topk_p=torch.empty((0, topk), device=device, dtype=torch.float32),
+            topk_index=torch.empty((0, topk), device=device, dtype=torch.int64),
+            draft_probs=(
+                torch.empty((0, vocab_size), device=device, dtype=torch.float32)
+                if get_server_args().speculative_use_rejection_sampling
+                else None
+            ),
+            capture_hidden_mode=capture_hidden_mode,
+        )
+
+    def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
+        if self.future_indices is not None:
+            self.future_indices = self.future_indices[new_indices]
+            return
+
+        strict_check = envs.SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK.get()
+        if has_been_filtered:
+            # in eagle_utils.py:verify, we have already filtered the batch by `unfinished_index`
+            # therefore, we don't need to filter the batch again in scheduler
+            error_msg = f"length of new_indices: {len(new_indices)} != length of topk_p: {len(self.topk_p)}, this should not happen"
+            if len(new_indices) != len(self.topk_p):
+                if strict_check:
+                    raise ValueError(error_msg)
+                else:
+                    logger.warning(error_msg)
+
+            self.topk_p = self.topk_p[: len(new_indices)]
+            self.topk_index = self.topk_index[: len(new_indices)]
+            if self.draft_probs is not None:
+                self.draft_probs = self.draft_probs[: len(new_indices)]
+            if self.hidden_states is not None:
+                self.hidden_states = self.hidden_states[: len(new_indices)]
+            self.bonus_tokens = self.bonus_tokens[: len(new_indices)]
+            if self.dsa_topk_indices is not None:
+                self.dsa_topk_indices = self.dsa_topk_indices[: len(new_indices)]
+        else:
+            # in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index`
+            self.topk_p = self.topk_p[new_indices]
+            self.topk_index = self.topk_index[new_indices]
+            if self.draft_probs is not None:
+                self.draft_probs = self.draft_probs[new_indices]
+            if self.hidden_states is not None:
+                self.hidden_states = self.hidden_states[new_indices]
+            self.bonus_tokens = self.bonus_tokens[new_indices]
+            if self.dsa_topk_indices is not None:
+                self.dsa_topk_indices = self.dsa_topk_indices[new_indices]
+
+    def merge_batch(self, spec_info: "EagleDraftInput"):
+        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]
+            )
+            return
+
+        # Detect idle stub by `topk_index` length (idle inputs have
+        # shape[0] == 0 across all fields). Don't use `hidden_states is None`:
+        # for STANDALONE all non-idle inputs also have None hidden_states.
+        if len(self.topk_index) == 0:
+            self.hidden_states = spec_info.hidden_states
+            self.bonus_tokens = spec_info.bonus_tokens
+            self.topk_p = spec_info.topk_p
+            self.topk_index = spec_info.topk_index
+            self.draft_probs = spec_info.draft_probs
+            self.dsa_topk_indices = spec_info.dsa_topk_indices
+            return
+        if len(spec_info.topk_index) == 0:
+            return
+        if self.hidden_states is not None and spec_info.hidden_states is not None:
+            self.hidden_states = torch.cat(
+                [self.hidden_states, spec_info.hidden_states], axis=0
+            )
+        self.bonus_tokens = torch.cat(
+            [self.bonus_tokens, spec_info.bonus_tokens], axis=0
+        )
+        self.topk_p = torch.cat([self.topk_p, spec_info.topk_p])
+        self.topk_index = torch.cat([self.topk_index, spec_info.topk_index])
+        if self.dsa_topk_indices is not None and spec_info.dsa_topk_indices is not None:
+            self.dsa_topk_indices = torch.cat(
+                [self.dsa_topk_indices, spec_info.dsa_topk_indices]
+            )
+        else:
+            self.dsa_topk_indices = None
+        if self.draft_probs is not None and spec_info.draft_probs is not None:
+            self.draft_probs = torch.cat([self.draft_probs, spec_info.draft_probs])
+
+
+@dataclass
+class EagleDraftExtendInput(SpecInput):
+    """Inputs to the draft-extend forward (the fill-draft-kvcache pass after
+    target prefill / verify).
+
+    Installed on `batch.spec_info` by the worker's `_draft_extend_for_*`
+    (and synthetically by draft-extend cuda-graph capture), then replaced
+    with a fresh `EagleDraftInput` for the next iter's draft.
+    """
+
+    # Target-model hidden states for the draft-extend forward; None when the
+    # draft doesn't read hidden_states (e.g., STANDALONE). Shape: decode
+    # (bs * num_draft_tokens, hidden), prefill (extend_num_tokens, hidden).
+    hidden_states: Optional[torch.Tensor] = None
+
+    # Per-req accept counts. `num_accept_tokens = num_correct_drafts + 1`.
+    # Both kept for cuda-graph buffer indexing.
+    num_correct_drafts: torch.Tensor = None
+    num_accept_tokens: torch.Tensor = None
+    # CPU view, read by attention backends during the extend forward.
+    num_accept_tokens_cpu: List[int] = None
+
+    # Per-req batch-state slices for the draft-extend forward:
+    #   - input_ids:        accept tokens flat over surviving reqs
+    #   - seq_lens / _cpu:  per-req sequence length (post-accept)
+    #   - req_pool_indices: per-req kv-pool slot
+    input_ids: torch.Tensor = None
+    seq_lens: torch.Tensor = None
+    seq_lens_cpu: torch.Tensor = None
+    req_pool_indices: torch.Tensor = None
+
+    #   - positions: shape `[total_accepted]`.
+    #   - bonus_tokens: shape `[bs]`; read post-extend to populate next iter's
+    #     `EagleDraftInput.bonus_tokens`.
+    positions: Optional[torch.Tensor] = None
+    bonus_tokens: Optional[torch.Tensor] = None
+
+    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.LAST
+    num_tokens_per_req: int = -1
+    num_tokens_for_logprob_per_req: int = 1
+
+    dsa_seed_topk_capture: Optional[torch.Tensor] = None
+    dsa_seed_topk_select: Optional[torch.Tensor] = None
+
+    # None for draft-extend's idle batch; attention backends fall back to
+    # rebuilding plain metadata from seq_lens when this is None.
+    kv_indptr: torch.Tensor = None
+
+    def __post_init__(self):
+        super().__init__(SpecInputType.EAGLE_DRAFT_EXTEND)
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        return self.num_tokens_per_req, self.num_tokens_for_logprob_per_req
+
+    @classmethod
+    def create_idle_input(
+        cls,
+        device: torch.device,
+        hidden_size: Optional[int],
+        dtype: Optional[torch.dtype],
+        capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.LAST,
+    ) -> "EagleDraftExtendInput":
+        return cls(
+            hidden_states=(
+                torch.empty((0, hidden_size), device=device, dtype=dtype)
+                if hidden_size is not None
+                else None
+            ),
+            num_correct_drafts=torch.empty((0,), device=device, dtype=torch.int32),
+            num_accept_tokens=torch.empty((0,), device=device, dtype=torch.int32),
+            num_accept_tokens_cpu=[],
+            input_ids=torch.empty((0,), device=device, dtype=torch.long),
+            seq_lens=torch.empty((0,), device=device, dtype=torch.int64),
+            seq_lens_cpu=torch.empty((0,), dtype=torch.int64),
+            req_pool_indices=torch.empty((0,), device=device, dtype=torch.int64),
+            capture_hidden_mode=capture_hidden_mode,
+        )
+
+    def generate_attn_arg_prefill(
+        self,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: Optional[int],
+        req_to_token: torch.Tensor,
+    ):
+        device = req_pool_indices.device
+        bs = self.num_correct_drafts.numel()
+        # Constant num_tokens_per_req qo layout (required for cuda-graph capture).
+        qo_indptr = torch.arange(
+            0,
+            (bs + 1) * self.num_tokens_per_req,
+            step=self.num_tokens_per_req,
+            dtype=torch.int32,
+            device=device,
+        )
+        cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=device)
+        cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+
+        if paged_kernel_lens_sum is None:
+            paged_kernel_lens_sum = cum_kv_seq_len[-1]
+
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum, dtype=torch.int32, device=device
+        )
+
+        create_flashinfer_kv_indices_triton[(bs,)](
+            req_to_token,
+            req_pool_indices,
+            paged_kernel_lens,
+            cum_kv_seq_len,
+            None,
+            kv_indices,
+            req_to_token.size(1),
+        )
+        return kv_indices, cum_kv_seq_len, qo_indptr, None

python/sglang/srt/speculative/eagle_utils.py

@@ -0,0 +1,892 @@
+from __future__ import annotations
+
+import logging
+import math
+from collections import defaultdict
+from enum import IntEnum
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.kernels.ops.speculative.spec_tree import (
+    sgl_build_tree_kernel_efficient_triton,
+    verify_tree_greedy_kernel_triton,
+)
+from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import (
+    alloc_paged_token_slots_extend_npu,
+)
+from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
+    maybe_build_dsv4_verify_bundle,
+)
+from sglang.srt.mem_cache.common import (
+    alloc_paged_token_slots_extend,
+    alloc_token_slots,
+    get_alloc_reserve_per_decode,
+    get_last_loc,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+)
+from sglang.srt.utils.async_probe import maybe_detect_oob
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_info import EagleVerifyInput
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+_is_musa = is_musa()
+_is_xpu = is_xpu()
+_is_cpu = is_cpu()
+
+logger = logging.getLogger(__name__)
+
+if _is_cuda or _is_hip or _is_musa:
+    from sgl_kernel import (
+        build_tree_kernel_efficient as sgl_build_tree_kernel_efficient,
+    )
+elif _is_cpu:
+    from sgl_kernel import (
+        build_tree_kernel_efficient_cpu as sgl_build_tree_kernel_efficient_cpu,
+    )
+    from sgl_kernel import verify_tree_greedy_cpu as sgl_verify_tree_greedy_cpu
+
+
+ALLOC_EXTEND_FUNCS = defaultdict(
+    lambda: alloc_paged_token_slots_extend,
+    {
+        "npu": alloc_paged_token_slots_extend_npu,
+    },
+)
+
+
+def per_step_draft_out_cache_loc(
+    out_cache_loc: torch.Tensor,
+    batch_size: int,
+    topk: int,
+    num_steps: int,
+) -> torch.Tensor:
+    """Per-step slice of the multi-step EAGLE draft out_cache_loc buffer.
+
+    Single source of truth for the layout shared by EagleWorkerV2.draft_forward
+    (per-step write target) and DeepseekV4AttnBackend (per-step compression
+    write target baked into metadata).
+    """
+    expected = batch_size * topk * num_steps
+    assert out_cache_loc.shape[0] == expected, (
+        f"out_cache_loc.shape[0]={out_cache_loc.shape[0]} != "
+        f"batch_size * topk * num_steps = {batch_size}*{topk}*{num_steps}={expected}"
+    )
+    return (
+        out_cache_loc.view(batch_size, topk, num_steps)
+        .permute(2, 0, 1)
+        .reshape(num_steps, -1)
+    )
+
+
+def _eagle_prefill_tail_tokens(
+    batch: ScheduleBatch, next_token_ids: torch.Tensor
+) -> torch.Tensor:
+    """Per-seq tail token for EAGLE prefill rotation; uses next prompt token for
+    non-final chunks (chunked-prefill chain consistency, see PR #26329)."""
+    tail_tokens = next_token_ids.to(batch.input_ids.dtype)
+    next_prompt_token = batch.chunked_req_next_prompt_token
+    if next_prompt_token is not None:
+        for i, r in enumerate(batch.reqs):
+            if r is batch.chunked_req:
+                tail_tokens = tail_tokens.clone()
+                tail_tokens[i] = next_prompt_token
+                break
+    return tail_tokens
+
+
+def organize_draft_results(
+    score_list: List[torch.Tensor],
+    token_list: List[torch.Tensor],
+    parents_list: List[torch.Tensor],
+    num_draft_token: int,
+):
+    score_list = torch.cat(score_list, dim=1).flatten(1)
+    ss_token_list = torch.cat(token_list, dim=1)
+    top_scores = torch.topk(score_list, num_draft_token - 1, dim=-1)
+    top_scores_index = top_scores.indices
+    top_scores_index = torch.sort(top_scores_index).values
+    maybe_detect_oob(
+        top_scores_index,
+        0,
+        ss_token_list.shape[1],
+        "organize_draft_results: top_scores_index OOB for gather on ss_token_list",
+    )
+    draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
+
+    if len(parents_list) > 1:
+        parent_list = torch.cat(parents_list[:-1], dim=1)
+    else:
+        batch_size = parents_list[0].shape[0]
+        parent_list = torch.empty(
+            batch_size, 0, dtype=torch.long, device=parents_list[0].device
+        )
+
+    return parent_list, top_scores_index, draft_tokens
+
+
+class TreeMaskMode(IntEnum):
+    FULL_MASK = 0
+    QLEN_ONLY = 1
+    QLEN_ONLY_BITPACKING = 2
+
+
+def default_tree_mask_mode() -> TreeMaskMode:
+    # The CPU verify attention kernel (intel_amx) consumes the qlen x qlen
+    # QLEN_ONLY tree mask directly; FULL_MASK is for the GPU kernels.
+    return TreeMaskMode.QLEN_ONLY if _is_cpu else TreeMaskMode.FULL_MASK
+
+
+def build_tree_kernel_efficient(
+    bonus_tokens: torch.Tensor,
+    parent_list: List[torch.Tensor],
+    top_scores_index: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_lens_sum: int,
+    topk: int,
+    spec_steps: int,
+    num_verify_tokens: int,
+    tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK,
+    tree_mask_buf: Optional[torch.Tensor] = None,
+    position_buf: Optional[torch.Tensor] = None,
+):
+    draft_tokens = torch.cat((bonus_tokens.unsqueeze(1), draft_tokens), dim=1).flatten()
+
+    # seq_lens_sum == sum(seq_lens); seq_lens: sequence length without draft tokens
+    bs = seq_lens.numel()
+    device = seq_lens.device
+    # e.g. for bs=1, tree_mask: num_draft_token, seq_lens_sum + num_draft_token (flattened)
+    # where each row indicates the attending pattern of each draft token
+    # if use_partial_packed_tree_mask is True, tree_mask: num_draft_token (flattened, packed)
+    if tree_mask_buf is not None:
+        tree_mask = tree_mask_buf
+        if tree_mask_mode == TreeMaskMode.QLEN_ONLY:
+            tree_mask.fill_(True)
+        elif tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING:
+            tree_mask.fill_(0)
+        elif tree_mask_mode == TreeMaskMode.FULL_MASK:
+            tree_mask.fill_(True)
+        else:
+            raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}")
+    elif tree_mask_mode == TreeMaskMode.QLEN_ONLY:
+        tree_mask = torch.full(
+            (num_verify_tokens * bs * num_verify_tokens,),
+            True,
+            dtype=torch.bool,
+            device=device,
+        )
+    elif tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING:
+        packed_dtypes = [torch.uint8, torch.uint16, torch.uint32]
+        packed_dtype_idx = int(math.ceil(math.log2((num_verify_tokens + 7) // 8)))
+        tree_mask = torch.zeros(
+            (num_verify_tokens * bs,),
+            dtype=packed_dtypes[packed_dtype_idx],
+            device=device,
+        )
+    elif tree_mask_mode == TreeMaskMode.FULL_MASK:
+        tree_mask = torch.full(
+            (
+                seq_lens_sum * num_verify_tokens
+                + num_verify_tokens * num_verify_tokens * bs,
+            ),
+            True,
+            device=device,
+        )
+    else:
+        raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}")
+
+    # TODO: make them torch.empty and fuse them into `sgl_build_tree_kernel`
+    retrieve_buf = torch.full(
+        (3, bs, num_verify_tokens), -1, device=device, dtype=torch.long
+    )
+    retrieve_index, retrieve_next_token, retrieve_next_sibling = retrieve_buf
+    # position: where each token belongs to
+    # e.g. if depth of each draft token is [0, 1, 1, 2] and the prompt length is 7
+    # then, positions = [7, 8, 8, 9]
+    if position_buf is not None:
+        positions = position_buf
+    else:
+        positions = torch.empty(
+            (bs * num_verify_tokens,), device=device, dtype=torch.long
+        )
+
+    if _is_npu:
+        torch.ops.npu.build_tree_kernel_efficient(
+            parent_list.to(dtype=torch.int64),
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
+    elif _is_xpu:
+        sgl_build_tree_kernel_triton(
+            parent_list,
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
+    elif _is_cpu:
+        sgl_build_tree_kernel_efficient_cpu(
+            parent_list,
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
+    else:
+        sgl_build_tree_kernel_efficient(
+            parent_list,
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
+    return (
+        tree_mask,
+        positions,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        draft_tokens,
+    )
+
+
+def sgl_build_tree_kernel_triton(
+    parent_list: torch.Tensor,
+    selected_index: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    tree_mask: torch.Tensor,
+    positions: torch.Tensor,
+    retrieve_index: torch.Tensor,
+    retrieve_next_token: torch.Tensor,
+    retrieve_next_sibling: torch.Tensor,
+    topk: int,
+    depth: int,
+    draft_token_num: int,
+    tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK,
+):
+    """Triton-based implementation."""
+    # TODO: Add support for QLEN_ONLY_BITPACKING mode
+    if tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING:
+        raise NotImplementedError(
+            "QLEN_ONLY_BITPACKING is not supported in Triton implementation"
+        )
+
+    batch_size = verified_seq_len.shape[0]
+    seq_len_prefix_sum = torch.cumsum(verified_seq_len, dim=0) - verified_seq_len
+
+    # Launch kernel with one program per batch item
+    grid = (batch_size,)
+
+    sgl_build_tree_kernel_efficient_triton[grid](
+        parent_list,
+        selected_index,
+        verified_seq_len,
+        seq_len_prefix_sum,
+        tree_mask,
+        positions,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        topk=topk,
+        depth=depth,
+        draft_token_num=draft_token_num,
+        tree_mask_mode=int(tree_mask_mode),
+        batch_size=batch_size,
+        parent_list_stride=(
+            parent_list.stride(0) if parent_list.dim() > 1 else parent_list.shape[0]
+        ),
+        selected_index_stride=selected_index.stride(0),
+    )
+
+
+def verify_tree_greedy_triton(
+    predicts: torch.Tensor,
+    accept_index: torch.Tensor,
+    accept_token_num: torch.Tensor,
+    candidates: torch.Tensor,
+    retrieve_index: torch.Tensor,
+    retrieve_next_token: torch.Tensor,
+    retrieve_next_sibling: torch.Tensor,
+    target_predict: torch.Tensor,
+):
+    """Triton-based implementation."""
+    batch_size = candidates.shape[0]
+    num_speculative_tokens = accept_index.shape[1]
+    num_draft_tokens = candidates.shape[1]
+
+    # Launch kernel with one program per batch item
+    grid = (batch_size,)
+
+    verify_tree_greedy_kernel_triton[grid](
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        target_predict,
+        batch_size=batch_size,
+        num_speculative_tokens=num_speculative_tokens,
+        num_draft_tokens=num_draft_tokens,
+    )
+
+
+def verify_tree_greedy_func(
+    predicts: torch.Tensor,
+    accept_index: torch.Tensor,
+    accept_token_num: torch.Tensor,
+    candidates: torch.Tensor,
+    retrieve_index: torch.Tensor,
+    retrieve_next_token: torch.Tensor,
+    retrieve_next_sibling: torch.Tensor,
+    target_predict: torch.Tensor,
+    topk: int = -1,
+):
+    if _is_cuda or _is_hip or _is_musa:
+        from sgl_kernel import verify_tree_greedy
+
+        verify_tree_greedy(
+            predicts=predicts,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=accept_token_num,  # mutable
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
+            retrive_index=retrieve_index,
+            retrive_next_token=retrieve_next_token,
+            retrive_next_sibling=retrieve_next_sibling,
+            target_predict=target_predict,
+        )
+
+    elif _is_cpu:
+        sgl_verify_tree_greedy_cpu(
+            predicts=predicts,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=accept_token_num,  # mutable
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match the CUDA op schema, so
+            # the CPU/CUDA call sites stay grep-symmetric.
+            retrive_index=retrieve_index,
+            retrive_next_token=retrieve_next_token,
+            retrive_next_sibling=retrieve_next_sibling,
+            target_predict=target_predict,
+        )
+
+    elif _is_npu:
+        from sgl_kernel_npu.sample.verify_tree_greedy import verify_tree_greedy
+
+        verify_tree_greedy(
+            predicts=predicts,
+            accept_index=accept_index,
+            accept_token_num=accept_token_num,
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
+            retrive_index=retrieve_index,
+            retrive_next_token=retrieve_next_token,
+            retrive_next_sibling=retrieve_next_sibling,
+            target_predict=target_predict,
+        )
+    elif _is_xpu:
+        verify_tree_greedy_triton(
+            predicts=predicts,
+            accept_index=accept_index,
+            accept_token_num=accept_token_num,
+            candidates=candidates,
+            retrieve_index=retrieve_index,
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            target_predict=target_predict,
+        )
+    return predicts, accept_index, accept_token_num
+
+
+def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int:
+    """Width of the target hidden states fed into the draft model.
+
+    This is the single source of truth and is derived entirely from config: for
+    EAGLE3 aux mode the draft consumes `num_aux` concatenated target layers
+    (each `target_hidden_size` wide); every other arch consumes the per-layer
+    `spec_hidden_size`.
+
+    Do NOT read this off a draft projection's `in_features` (e.g. an `fc`
+    layer): that width is arch-specific.
+
+    Note: read entirely from the *draft* `model_runner`'s config. The non-aux
+    branch assumes the draft's `spec_hidden_size` equals the target hidden width
+    fed to the draft (true for standard EAGLE, where the draft mirrors the
+    target hidden size); aux mode reads the explicit `target_hidden_size`.
+    """
+    model_config = model_runner.model_config
+    hf_config = model_config.hf_config
+    eagle_config = getattr(hf_config, "eagle_config", None) or {}
+    get_eagle_config = (
+        eagle_config.get
+        if isinstance(eagle_config, dict)
+        else lambda key, default=None: getattr(eagle_config, key, default)
+    )
+    use_aux = get_eagle_config("use_aux_hidden_state", True)
+    spec_algorithm = model_runner.spec_algorithm
+
+    if not (spec_algorithm is not None and spec_algorithm.is_eagle3() and use_aux):
+        return model_config.spec_hidden_size
+
+    target_hidden = getattr(hf_config, "target_hidden_size", None)
+    if target_hidden is None:
+        target_hidden = model_config.hidden_size
+    num_aux = getattr(hf_config, "num_aux_hidden_states", None)
+    if num_aux is None:
+        layer_ids = get_eagle_config("eagle_aux_hidden_state_layer_ids", None)
+        if layer_ids is None:
+            layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None)
+        num_aux = len(layer_ids) if layer_ids else 3
+    return target_hidden * num_aux
+
+
+def get_draft_recurrent_hidden_state_spec(
+    model_runner: ModelRunner,
+) -> tuple[Optional[int], Optional[torch.dtype]]:
+    """Return hidden_states width/dtype carried between draft decode steps."""
+    if model_runner.spec_algorithm.is_standalone():
+        return None, None
+    return model_runner.model_config.spec_hidden_size, model_runner.model_config.dtype
+
+
+def eagle_prepare_for_verify(
+    verify_input: EagleVerifyInput,
+    req_to_token_pool: ReqToTokenPool,
+    batch: ScheduleBatch,
+    target_worker: TpModelWorker,
+):
+    from sglang.kernels.ops.speculative.cache_locs import (
+        assign_extend_cache_locs_func,
+    )
+    from sglang.srt.model_executor.forward_batch_info import (
+        CaptureHiddenMode,
+        ForwardBatch,
+        ForwardMode,
+    )
+    from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify
+
+    if not batch.forward_mode.is_idle():
+        # Assign cache locations
+        bs = len(batch.req_pool_indices)
+        batch.input_ids = verify_input.draft_token
+        maybe_detect_oob(
+            batch.input_ids,
+            0,
+            batch.model_config.vocab_size,
+            "v2 prepare_for_verify input_ids",
+        )
+        device = batch.device
+        batch.out_cache_loc = assign_extend_cache_locs_func(
+            req_pool_indices=batch.req_pool_indices,
+            req_to_token=req_to_token_pool.req_to_token,
+            start_offset=batch.seq_lens,
+            end_offset=batch.seq_lens + verify_input.draft_token_num,
+            batch_size=bs,
+            draft_token_num=verify_input.draft_token_num,
+            device=device,
+        )
+
+        batch.out_cache_loc_dsv4 = maybe_build_dsv4_verify_bundle(
+            batch, verify_input.draft_token_num
+        )
+
+        prepare_mamba_track_for_verify(batch)
+
+        # TBO's split_spec_info reads these; no-verify-sync leaves both None.
+        verify_input.seq_lens_cpu = batch.seq_lens_cpu
+        verify_input.seq_lens_sum = (
+            int(batch.seq_lens_cpu.sum()) if batch.seq_lens_cpu is not None else None
+        )
+
+    # Get a forward batch
+    batch.forward_mode = (
+        ForwardMode.IDLE if batch.forward_mode.is_idle() else ForwardMode.TARGET_VERIFY
+    )
+    capture_mode = (
+        CaptureHiddenMode.NULL
+        if target_worker.model_runner.spec_algorithm.is_standalone()
+        else CaptureHiddenMode.FULL
+    )
+    batch.capture_hidden_mode = capture_mode
+    verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
+
+    # Run attention backend plan and cuda graph preparation
+    can_run_cuda_graph = bool(
+        target_worker.model_runner.decode_cuda_graph_runner
+        and target_worker.model_runner.decode_cuda_graph_runner.can_run_graph(
+            verify_forward_batch
+        )
+    )
+    if can_run_cuda_graph:
+        target_worker.model_runner.decode_cuda_graph_runner.load_batch(
+            verify_forward_batch
+        )
+        verify_forward_batch.mark_forward_metadata_ready()
+    # Non-cuda-graph: defer init to forward_extend, which runs after
+    # `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
+    # here would use pre-pad shapes and trip DSv4 indexer shape match.
+
+    return verify_forward_batch, can_run_cuda_graph
+
+
+def eagle_sample(
+    verify_input: EagleVerifyInput,
+    batch: ScheduleBatch,
+    logits_output: LogitsProcessorOutput,
+    vocab_mask: torch.Tensor = None,
+):
+    """
+    Verify and find accepted tokens based on logits output and batch
+    (which contains spec decoding information).
+    """
+    import torch.nn.functional as F
+
+    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.layers.dp_attention import (
+        is_dp_attention_enabled,
+    )
+    from sglang.srt.runtime_context import get_server_args
+    from sglang.srt.sampling.penaltylib.repetition_penalty import (
+        apply_scaling_penalties,
+    )
+    from sglang.srt.speculative.spec_utils import (
+        SIMULATE_ACC_LEN,
+        SIMULATE_ACC_TOKEN_MODE,
+        generate_simulated_accept_index,
+    )
+    from sglang.srt.utils.async_probe import maybe_detect_nan, sanitize_nan_logits
+
+    device = batch.device
+    if batch.forward_mode.is_idle():
+        predict = torch.empty(0, dtype=torch.int32, device=device)
+        num_correct_drafts = torch.empty(0, dtype=torch.int32, device=device)
+        accept_index = torch.empty(0, dtype=torch.int32, device=device)
+        return predict, num_correct_drafts, accept_index
+
+    bs = len(batch.seq_lens)
+    sampling_info = batch.sampling_info
+    next_token_logits = logits_output.next_token_logits
+
+    sanitize_nan_logits(next_token_logits, "verify: target model logits")
+
+    # Apply penalty
+    # This is a relaxed version of penalties for speculative decoding.
+    if sampling_info.acc_additive_penalties is not None:
+        next_token_logits.add_(
+            torch.repeat_interleave(
+                sampling_info.acc_additive_penalties,
+                verify_input.draft_token_num,
+                dim=0,
+            )
+        )
+    if sampling_info.acc_scaling_penalties is not None:
+        apply_scaling_penalties(
+            next_token_logits,
+            torch.repeat_interleave(
+                sampling_info.acc_scaling_penalties, verify_input.draft_token_num, dim=0
+            ),
+        )
+    if sampling_info.logit_bias is not None:
+        next_token_logits.add_(
+            torch.repeat_interleave(
+                sampling_info.logit_bias, verify_input.draft_token_num, dim=0
+            )
+        )
+
+    # Apply grammar mask if provided
+    if vocab_mask is not None:
+        assert verify_input.grammar is not None
+        verify_input.grammar.apply_vocab_mask(
+            logits=next_token_logits, vocab_mask=vocab_mask
+        )
+
+    candidates = verify_input.draft_token.reshape(bs, verify_input.draft_token_num)
+    predict_shape = list(next_token_logits.shape)[:-1]
+    predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
+    accept_index = torch.full(
+        (bs, verify_input.max_tree_depth), -1, dtype=torch.int32, device=device
+    )
+    num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
+
+    # Sample tokens
+    target_predict = None
+    if sampling_info.is_all_greedy or _is_cpu or _is_npu or _is_hip or _is_xpu:
+        target_predict = torch.argmax(next_token_logits, dim=-1)
+        target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
+        predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
+            predicts=predict,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=num_correct_drafts,  # mutable
+            candidates=candidates,
+            retrieve_index=verify_input.retrieve_index,
+            retrieve_next_token=verify_input.retrieve_next_token,
+            retrieve_next_sibling=verify_input.retrieve_next_sibling,
+            target_predict=target_predict,
+            topk=verify_input.tree_topk,
+        )
+    else:
+        from sgl_kernel import (
+            top_k_renorm_prob,
+            top_p_renorm_prob,
+            tree_speculative_sampling_target_only,
+        )
+
+        from sglang.srt.speculative.reject_sampling import (
+            chain_speculative_sampling_triton,
+        )
+
+        use_rejection_sampling = get_server_args().speculative_use_rejection_sampling
+
+        # Apply temperature and get target probs
+        expanded_temperature = torch.repeat_interleave(
+            sampling_info.temperatures, verify_input.draft_token_num, dim=0
+        )  # (bs * num_draft_tokens, 1)
+
+        target_probs = F.softmax(
+            next_token_logits / expanded_temperature, dim=-1
+        )  # (bs * num_draft_tokens, vocab_size)
+        maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
+        target_probs = top_k_renorm_prob(
+            target_probs,
+            torch.repeat_interleave(
+                sampling_info.top_ks, verify_input.draft_token_num, dim=0
+            ),
+        )  # (bs * num_draft_tokens, vocab_size)
+        maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
+        target_probs = top_p_renorm_prob(
+            target_probs,
+            torch.repeat_interleave(
+                sampling_info.top_ps, verify_input.draft_token_num, dim=0
+            ),
+        )
+        maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
+        target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1)
+        draft_probs = (
+            verify_input.draft_probs
+            if use_rejection_sampling
+            else torch.zeros_like(target_probs)
+        )
+        # Defense-in-depth behind the spec_hook startup allowlist: validate the
+        # actual kernel inputs (catches draft_probs plumbing regressions or a
+        # startup guard bypassed by a worker subclass) before the Triton kernel.
+        if use_rejection_sampling and (
+            draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1]
+        ):
+            raise ValueError(
+                "Rejection sampling requires a target-vocab draft proposal "
+                "distribution; the current speculative algorithm/draft worker "
+                "does not produce one (draft_probs missing or vocab-mismatched)."
+            )
+
+        # coins for rejection sampling
+        coins = torch.rand_like(candidates, dtype=torch.float32, device=device)
+        # coins for final sampling
+        coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device)
+
+        sampling_fn = (
+            chain_speculative_sampling_triton
+            if use_rejection_sampling
+            else tree_speculative_sampling_target_only
+        )
+        sampling_fn(
+            predicts=predict,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=num_correct_drafts,  # mutable
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
+            retrive_index=verify_input.retrieve_index,
+            retrive_next_token=verify_input.retrieve_next_token,
+            retrive_next_sibling=verify_input.retrieve_next_sibling,
+            uniform_samples=coins,
+            uniform_samples_for_final_sampling=coins_for_final_sampling,
+            target_probs=target_probs,
+            draft_probs=draft_probs,
+            threshold_single=get_server_args().speculative_accept_threshold_single,
+            threshold_acc=get_server_args().speculative_accept_threshold_acc,
+            deterministic=True,
+        )
+
+        # Sync sampling results across TP ranks: different GPUs may
+        # produce slightly different target_probs due to floating-point
+        # non-determinism in softmax/top_k/top_p, causing different
+        # sampled tokens. Broadcast from rank 0 to ensure consistency.
+        tp_group = (
+            get_parallel().attn_tp_group
+            if is_dp_attention_enabled()
+            else get_tp_group()
+        )
+        if tp_group.world_size > 1:
+            tp_group.broadcast(predict, src=0)
+            tp_group.broadcast(accept_index, src=0)
+            tp_group.broadcast(num_correct_drafts, src=0)
+
+    if SIMULATE_ACC_LEN > 0:
+        # Do simulation. The helper builds (and returns) a replacement
+        # accept_index of width spec_steps + 1, so pass max_tree_depth - 1
+        # to keep the simulated width identical to the real one.
+        if SIMULATE_ACC_TOKEN_MODE not in ("fixed", "real-draft-token"):
+            raise ValueError(
+                "Invalid SGLANG_SIMULATE_ACC_TOKEN_MODE "
+                f"{SIMULATE_ACC_TOKEN_MODE!r}; expected 'fixed' or "
+                "'real-draft-token'."
+            )
+
+        if SIMULATE_ACC_TOKEN_MODE == "real-draft-token":
+            if verify_input.tree_topk != 1:
+                raise ValueError(
+                    "SGLANG_SIMULATE_ACC_LEN with real draft tokens currently "
+                    "requires speculative_eagle_topk=1."
+                )
+
+            # Use target argmax as the synthetic bonus for non-greedy requests.
+            if target_predict is None:
+                target_predict = torch.argmax(next_token_logits, dim=-1).reshape(
+                    bs, verify_input.draft_token_num
+                )
+        accept_index = generate_simulated_accept_index(
+            accept_index=accept_index,
+            predict=predict,  # mutable
+            num_correct_drafts=num_correct_drafts,  # mutable
+            candidates=candidates,
+            target_predict=target_predict,
+            simulate_acc_len=SIMULATE_ACC_LEN,
+            simulate_acc_token_mode=SIMULATE_ACC_TOKEN_MODE,
+            bs=bs,
+            spec_steps=verify_input.max_tree_depth - 1,
+        )
+
+    # `num_correct_drafts` stays drafts-only inside this function; the returned
+    # tensor includes the trailing/bonus token via out-of-place +1 so the
+    # name no longer flips semantics mid-function (naming doc C2).
+    return predict, num_correct_drafts + 1, accept_index
+
+
+def eagle_prepare_for_decode(batch: ScheduleBatch):
+    batch.maybe_evict_swa()
+
+    from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
+
+    bs = batch.batch_size()
+
+    # Accumulate penalty
+    # This is a relaxed version of penalties for speculative decoding.
+    if batch.sampling_info.penalizer_orchestrator.is_required:
+        batch.cumulate_penalty_output_tokens()
+
+    page_size = batch.token_to_kv_pool_allocator.page_size
+    double_alloc = get_alloc_reserve_per_decode()
+
+    cur_kv_lens = [0] * bs
+    nxt_kv_lens = [0] * bs
+    num_needed_tokens = 0
+    for i, r in enumerate(batch.reqs):
+        cur = r.kv_allocated_len
+        # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur.
+        # kv_committed_len is honest (bonus committed in resolve, not here),
+        # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs.
+        nxt = max(cur, r.kv_committed_len + double_alloc)
+        cur_kv_lens[i] = cur
+        nxt_kv_lens[i] = nxt
+        num_needed_tokens += nxt - cur
+        r.kv_allocated_len = nxt
+        r.decode_batch_idx += 1
+
+    cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")
+    nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu")
+
+    # Fail fast if the page>1 + topk>1 draft over-allocation
+    # (get_alloc_reserve_per_decode) outgrows the req_to_token row: the write below
+    # would OOB and free would leak KV. The row is widened to hold it in _init_pools
+    # (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert.
+    from sglang.srt.runtime_context import get_server_args
+
+    if page_size > 1 and (get_server_args().speculative_eagle_topk or 1) > 1:
+        max_alloc_len = int(nxt_kv_lens_cpu.max())
+        row_width = batch.req_to_token_pool.req_to_token.shape[1]
+        assert max_alloc_len <= row_width, (
+            f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds "
+            f"req_to_token row width ({row_width}); page_size={page_size}. Widen the "
+            f"row to hold committed + get_alloc_reserve_per_decode (PR #26972)."
+        )
+
+    # non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR
+    # barrier has chained to the prev forward -> host stalls a full forward.
+    cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True)
+    nxt_kv_lens_device = nxt_kv_lens_cpu.to(device=batch.device, non_blocking=True)
+    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_device,
+        )
+        device_type = getattr(batch.device, "type", str(batch.device).split(":", 1)[0])
+        out_cache_loc = ALLOC_EXTEND_FUNCS[device_type](
+            batch.tree_cache,
+            cur_kv_lens_device,
+            cur_kv_lens_cpu,
+            nxt_kv_lens_device,
+            nxt_kv_lens_cpu,
+            last_loc,
+            num_needed_tokens,
+            req_pool_indices=batch.req_pool_indices,
+            batch=batch,
+        )
+    assign_req_to_token_pool_func(
+        batch.req_pool_indices,
+        batch.req_to_token_pool.req_to_token,
+        cur_kv_lens_device,
+        nxt_kv_lens_device,
+        out_cache_loc,
+        bs,
+    )

python/sglang/srt/speculative/external_corpus_manager.py

@@ -0,0 +1,110 @@
+"""Manages external SAM corpora for ngram speculative decoding.
+
+Handles add/remove/list operations and async background loading.
+Used by the Scheduler — not a mixin, a standalone manager object.
+"""
+
+import logging
+import threading
+from typing import Callable, Optional, Tuple
+
+from sglang.srt.managers.io_struct import (
+    AddExternalCorpusReqInput,
+    AddExternalCorpusReqOutput,
+    ListExternalCorporaReqInput,
+    ListExternalCorporaReqOutput,
+    RemoveExternalCorpusReqInput,
+    RemoveExternalCorpusReqOutput,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ExternalCorpusManager:
+    """Manages external SAM corpus lifecycle for a single scheduler.
+
+    Args:
+        draft_worker: the NGRAMWorker instance (must have add_external_corpus,
+            remove_external_corpus, list_external_corpora methods).
+        send_response: callable(output, recv_req) to send deferred responses
+            back to the tokenizer manager.
+    """
+
+    def __init__(self, draft_worker, send_response: Callable):
+        self._worker = draft_worker
+        self._send_response = send_response
+        self._pending_load: Optional[
+            Tuple[AddExternalCorpusReqInput, threading.Thread]
+        ] = None
+        self._load_result: Optional[AddExternalCorpusReqOutput] = None
+
+    def check_pending_load(self):
+        """Poll from the scheduler event loop. Sends response when done."""
+        if self._pending_load is None:
+            return
+        recv_req, thread = self._pending_load
+        if thread.is_alive():
+            return
+        self._pending_load = None
+        thread.join()  # formal happens-before for _load_result visibility
+        result = self._load_result
+        self._load_result = None
+        if result.success:
+            self._worker.commit_corpus_load(result.corpus_id, result.loaded_token_count)
+        self._send_response(result, recv_req)
+
+    def add(
+        self, recv_req: AddExternalCorpusReqInput
+    ) -> Optional[AddExternalCorpusReqOutput]:
+        if self._pending_load is not None:
+            return AddExternalCorpusReqOutput(
+                success=False,
+                message="Another corpus load is already in progress.",
+            )
+
+        def _build():
+            try:
+                loaded = self._worker.add_external_corpus(
+                    recv_req.corpus_id, recv_req.token_chunks
+                )
+                self._load_result = AddExternalCorpusReqOutput(
+                    success=True,
+                    corpus_id=recv_req.corpus_id,
+                    message=f"Loaded corpus '{recv_req.corpus_id}' with {loaded} tokens.",
+                    loaded_token_count=loaded,
+                )
+            except Exception as e:
+                self._load_result = AddExternalCorpusReqOutput(
+                    success=False, message=str(e)
+                )
+
+        thread = threading.Thread(target=_build, daemon=True)
+        self._pending_load = (recv_req, thread)
+        thread.start()
+        return None  # response sent later by check_pending_load
+
+    # FIXME(kpham-sgl): remove a corpus during a pending load is an undefined behaviour
+    # and should be explicitly prevented.
+    def remove(
+        self, recv_req: RemoveExternalCorpusReqInput
+    ) -> RemoveExternalCorpusReqOutput:
+        try:
+            self._worker.remove_external_corpus(recv_req.corpus_id)
+            return RemoveExternalCorpusReqOutput(
+                success=True,
+                message=f"Removed corpus '{recv_req.corpus_id}'.",
+            )
+        except Exception as e:
+            return RemoveExternalCorpusReqOutput(success=False, message=str(e))
+
+    def list(
+        self, recv_req: ListExternalCorporaReqInput
+    ) -> ListExternalCorporaReqOutput:
+        try:
+            token_counts = self._worker.list_external_corpora()
+            return ListExternalCorporaReqOutput(
+                success=True,
+                corpus_token_counts=token_counts,
+            )
+        except Exception as e:
+            return ListExternalCorporaReqOutput(success=False, message=str(e))

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -0,0 +1,449 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+from sglang.srt.model_executor.runner import (
+    DecodeCudaGraphRunner,
+    DeepEPCudaGraphRunnerAdapter,
+    ShapeKey,
+    get_batch_sizes_to_capture,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
+from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+from sglang.srt.runtime_context import get_flags
+from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPDraftInput
+from sglang.srt.utils import (
+    require_attn_tp_gather,
+    require_gathered_buffer,
+    require_mlp_sync,
+    require_mlp_tp_gather,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenKVMTPDraftWorker
+
+
+@dataclass
+class FrozenKVMTPInputBuffers(ForwardInputBuffers):
+    req_pool_indices: torch.Tensor
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    topk_p: torch.Tensor
+    topk_index: torch.Tensor
+    hidden_states: torch.Tensor
+    # Consumed by the captured seed iter; see `FrozenKVMTPDraftWorker.draft_forward`.
+    bonus_tokens: torch.Tensor
+    global_num_tokens_gpu: Optional[torch.Tensor]
+    global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+
+
+class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
+    """CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step.
+
+    Subclasses DecodeCudaGraphRunner to inherit the outer capture loop
+    (capture() / _capture_one_stream()), the bucket-padding helper
+    (_pad_to_bucket), and the backend-driven capture/replay scaffolding.
+    Frozen-KV-MTP-specific bits — the buffer dataclass, the dummy
+    ForwardBatch + FrozenKVMTPDraftInput built in capture_one_shape, the
+    target-KV-pool swap during capture, the worker's frozen-KV metadata
+    helpers, the topk*topk bucket math, the expanded-bs bookkeeping, and
+    the 3-tuple replay output — are overridden.
+
+    Like the EAGLE draft runner, it does NOT call
+    DecodeCudaGraphRunner.__init__ (that init sets up decode-only state);
+    it sets up its own fields directly while satisfying the parent's
+    capture() / backend contract.
+    """
+
+    def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPDraftWorker):
+        self.frozen_kv_mtp_worker = frozen_kv_mtp_worker
+        self.model_runner = model_runner = frozen_kv_mtp_worker.draft_model_runner
+
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.enable_torch_compile = get_flags().capture.enable_torch_compile
+        self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
+        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
+        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
+        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
+        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+        self.tp_size = self.model_runner.tp_size
+        self.dp_size = self.model_runner.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.speculative_num_steps = model_runner.server_args.speculative_num_steps
+        self.topk = model_runner.server_args.speculative_eagle_topk
+        self.draft_attn_backend = frozen_kv_mtp_worker.draft_attn_backend
+        self.enable_profile_cuda_graph = (
+            model_runner.server_args.enable_profile_cuda_graph
+        )
+
+        self.attn_backend = self.draft_attn_backend
+
+        self.compile_bs = []
+        self.enable_pdmux = False
+        self.record_nolora_graph = False
+        self.is_dllm = False
+
+        self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
+
+        self.capture_forward_mode = ForwardMode.DECODE
+        self.capture_hidden_mode = CaptureHiddenMode.LAST
+
+        self.num_tokens_per_bs = self.topk
+        self.capture_bs, _ = get_batch_sizes_to_capture(
+            model_runner, self.num_tokens_per_bs
+        )
+        self.max_bs = max(self.capture_bs)
+        self.max_num_token = self.max_bs * self.num_tokens_per_bs
+
+        self.draft_attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
+        self.seq_len_fill_value = (
+            self.draft_attn_backend.get_cuda_graph_seq_len_fill_value()
+        )
+        seq_lens_cpu = torch.full(
+            (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int64
+        )
+
+        if self.enable_torch_compile:
+            set_torch_compile_config()
+
+        with torch.device(model_runner.device):
+            req_pool_indices = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
+            seq_lens = torch.full(
+                (self.max_num_token,), self.seq_len_fill_value, dtype=torch.int64
+            )
+            topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
+            topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)
+            hidden_states = torch.zeros(
+                (self.max_bs, frozen_kv_mtp_worker._recurrent_hidden_size),
+                dtype=self.model_runner.dtype,
+            )
+            bonus_tokens = torch.zeros((self.max_bs,), dtype=torch.int64)
+
+            if self.require_gathered_buffer:
+                if self.require_mlp_tp_gather:
+                    global_num_tokens_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (self.dp_size,), dtype=torch.int32
+                    )
+                else:
+                    assert self.require_attn_tp_gather
+                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (1,), dtype=torch.int32
+                    )
+            else:
+                global_num_tokens_gpu = None
+                global_num_tokens_for_logprob_gpu = None
+
+        self.buffers = FrozenKVMTPInputBuffers(
+            req_pool_indices=req_pool_indices,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            topk_p=topk_p,
+            topk_index=topk_index,
+            hidden_states=hidden_states,
+            bonus_tokens=bonus_tokens,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+        )
+        self.buffers.share_buffers()
+
+        self.backend = resolve_decode_backend(self)
+
+        try:
+            with model_capture_mode():
+                self.capture()
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture frozen-KV MTP cuda graph failed: {e}\n"
+                f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+            )
+
+    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
+        return ShapeKey(size=bs)
+
+    def _replay_graph(self, shape_key, forward_batch):
+        return self.backend.replay(shape_key, forward_batch)
+
+    def can_run_graph(self, forward_batch: ForwardBatch):
+        if self.require_mlp_tp_gather:
+            cuda_graph_bs = max(forward_batch.global_num_tokens_cpu) // (
+                self.topk * self.topk
+            )
+        else:
+            cuda_graph_bs = (
+                forward_batch.batch_size // self.topk
+                if self.topk > 1
+                else forward_batch.batch_size
+            )
+
+        is_bs_supported = (
+            self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs))
+            if self.disable_padding
+            else cuda_graph_bs <= self.max_bs
+        )
+        if self.require_mlp_sync:
+            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+        return is_bs_supported
+
+    def capture_one_shape(
+        self,
+        size: int,
+        forward: Callable,
+        stream_idx: Optional[int] = None,
+        variant_label: Optional[str] = None,
+    ):
+        del forward, stream_idx, variant_label
+        buffers = self.buffers
+        request_bs = size
+        expanded_bs = request_bs * self.num_tokens_per_bs
+
+        req_pool_indices = buffers.req_pool_indices[:expanded_bs]
+        positions = buffers.positions[:expanded_bs]
+        mrope_positions = buffers.mrope_positions[:, :expanded_bs]
+        seq_lens = buffers.seq_lens[:expanded_bs]
+        seq_lens_cpu = buffers.seq_lens_cpu[:expanded_bs]
+        topk_p = buffers.topk_p[:request_bs]
+        topk_index = buffers.topk_index[:request_bs]
+        hidden_states = buffers.hidden_states[:request_bs]
+        bonus_tokens = buffers.bonus_tokens[:request_bs]
+
+        if self.require_mlp_tp_gather:
+            global_num_tokens_cpu = [expanded_bs] * self.dp_size
+        elif self.require_attn_tp_gather:
+            global_num_tokens_cpu = [expanded_bs]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            num_tokens_tensor = torch.tensor(
+                global_num_tokens_cpu,
+                dtype=torch.int32,
+                device=buffers.positions.device,
+            )
+            buffers.global_num_tokens_gpu.copy_(num_tokens_tensor)
+            buffers.global_num_tokens_for_logprob_gpu.copy_(num_tokens_tensor)
+            global_num_tokens = buffers.global_num_tokens_gpu
+            global_num_tokens_for_logprob = buffers.global_num_tokens_for_logprob_gpu
+        else:
+            global_dp_buffer_len = None
+            global_num_tokens = None
+            global_num_tokens_for_logprob = None
+
+        spec_info = FrozenKVMTPDraftInput(
+            topk_p=topk_p,
+            topk_index=topk_index,
+            hidden_states=hidden_states,
+            bonus_tokens=bonus_tokens,
+            capture_hidden_mode=CaptureHiddenMode.LAST,
+        )
+        spec_info.num_tokens_per_req = self.topk
+        spec_info.num_tokens_for_logprob_per_req = self.topk
+        spec_info.positions = positions
+
+        forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.DECODE,
+            batch_size=expanded_bs,
+            input_ids=None,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            out_cache_loc=None,
+            seq_lens_sum=seq_lens.sum().item(),
+            return_logprob=False,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            global_num_tokens_gpu=global_num_tokens,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
+            spec_algorithm=self.model_runner.spec_algorithm,
+            spec_info=spec_info,
+            capture_hidden_mode=CaptureHiddenMode.LAST,
+        )
+
+        def run_once():
+            # Record the metadata rebuild against the committed target-prefix
+            # geometry (spec_info nulled → plain target-length decode), matching
+            # every other frozen-KV metadata init. Without the view, backends
+            # that key seqlen offsets off spec_info (trtllm_mha's draft-decode
+            # branch adds speculative_step_id + 1) bake a +1 offset into the
+            # captured graph and replay reads one extra, never-written KV slot.
+            with self.frozen_kv_mtp_worker._frozen_kv_target_view(forward_batch):
+                self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
+
+            forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(
+                global_dp_buffer_len,
+                expanded_bs,
+                forward_batch.dp_padding_mode.is_max_len(),
+                global_num_tokens_cpu,
+            )
+            set_is_extend_in_batch(False)
+
+            hidden_states_backup = forward_batch.spec_info.hidden_states
+            # The capture batch is marked by the capture metadata helper
+            # below, so draft_forward skips its eager plan.
+            ret = self.frozen_kv_mtp_worker.draft_forward(forward_batch)
+            forward_batch.spec_info.hidden_states = hidden_states_backup
+            return ret
+
+        # Swap the draft backend's token_to_kv_pool to the frozen target pool
+        # for the capture; the single backend-attr swap is seen by both
+        # get_token_to_kv_pool() (via get_attn_backend()) and the
+        # backend's own reads.
+        target_pool = self.frozen_kv_mtp_worker.kv_context.target_token_to_kv_pool
+        saved_backend_pool = self.draft_attn_backend.token_to_kv_pool
+        self.draft_attn_backend.token_to_kv_pool = target_pool
+        try:
+            with forward_context(ForwardContext(attn_backend=self.draft_attn_backend)):
+                self.frozen_kv_mtp_worker._init_frozen_kv_metadata_capture_cuda_graph(
+                    forward_batch
+                )
+                self.deepep_adapter.capture(is_extend_in_batch=False)
+                shape_key = self._make_graph_key(request_bs)
+                post_warmup_hook = getattr(
+                    self.draft_attn_backend, "on_after_cuda_graph_warmup", None
+                )
+                maybe_flashinfer_autotune_speculative_draft(
+                    self,
+                    run_once,
+                    post_warmup_hook=post_warmup_hook,
+                    skip_logits=False,
+                )
+                self.backend.capture_one(
+                    shape_key,
+                    run_once,
+                    dummies=None,
+                    post_warmup_hook=post_warmup_hook,
+                )
+        finally:
+            self.draft_attn_backend.token_to_kv_pool = saved_backend_pool
+
+    def _postprocess_output_to_raw_bs(self, out, raw_bs):
+        parent_list, top_scores_index, draft_tokens = (t[:raw_bs] for t in out)
+        return parent_list, top_scores_index, draft_tokens
+
+    def execute(self, forward_batch: ForwardBatch):
+        self.deepep_adapter.replay()
+        buffers = self.buffers
+
+        raw_expanded_bs = forward_batch.batch_size
+        raw_bs = (
+            raw_expanded_bs // self.num_tokens_per_bs
+            if self.topk > 1
+            else raw_expanded_bs
+        )
+        raw_num_token = raw_expanded_bs
+
+        if self.require_mlp_tp_gather:
+            max_num_tokens = max(forward_batch.global_num_tokens_cpu)
+            max_batch_size = max_num_tokens // (
+                self.num_tokens_per_bs * self.num_tokens_per_bs
+            )
+            bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
+        else:
+            bs = self._pad_to_bucket(raw_bs, self.capture_bs)
+
+        expanded_bs = bs * self.num_tokens_per_bs
+        if bs != raw_bs:
+            buffers.seq_lens.fill_(self.seq_len_fill_value)
+            buffers.positions.zero_()
+            # Pair with seq_lens fill: padded rows must point at reserved
+            # req_pool slot 0 (req_to_token[0, :] is all zeros from init).
+            buffers.req_pool_indices.zero_()
+
+        num_tokens = expanded_bs
+        buffers.seq_lens[:raw_expanded_bs].copy_(forward_batch.seq_lens)
+        buffers.positions[:raw_num_token].copy_(forward_batch.positions)
+        if forward_batch.mrope_positions is not None:
+            buffers.mrope_positions[:, :raw_num_token].copy_(
+                forward_batch.mrope_positions
+            )
+        # `topk_p`/`topk_index` are produced by the captured seed iter.
+        buffers.bonus_tokens[:raw_bs].copy_(forward_batch.spec_info.bonus_tokens)
+        buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
+        buffers.req_pool_indices[:raw_expanded_bs].copy_(forward_batch.req_pool_indices)
+
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(expanded_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(expanded_bs)
+
+        if bs != raw_bs:
+            forward_batch.batch_size = expanded_bs
+            forward_batch.seq_lens = buffers.seq_lens[:expanded_bs]
+            forward_batch.req_pool_indices = buffers.req_pool_indices[:expanded_bs]
+            forward_batch.positions = buffers.positions[:num_tokens]
+            if forward_batch.mrope_positions is not None:
+                forward_batch.mrope_positions = buffers.mrope_positions[:, :num_tokens]
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+            buffers.seq_lens_cpu[:raw_expanded_bs].copy_(forward_batch.seq_lens_cpu)
+            forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:expanded_bs]
+
+        self.frozen_kv_mtp_worker._init_frozen_kv_metadata_replay_cuda_graph(
+            forward_batch,
+            expanded_bs,
+            forward_batch.seq_lens_sum
+            + (expanded_bs - raw_expanded_bs) * self.seq_len_fill_value,
+        )
+
+        self.raw_bs = raw_bs
+        self.bs = bs
+        shape_key = self._make_graph_key(bs)
+        # NVTX span: the graph bypasses `model_runner.forward`'s record_function.
+        span_name = f"step[DRAFT_LOOP raw_bs={raw_bs} bs={bs} topk={self.topk}]"
+        if torch.autograd._profiler_enabled():
+            with torch.profiler.record_function(span_name):
+                out = self._replay_graph(shape_key, forward_batch)
+        else:
+            out = self._replay_graph(shape_key, forward_batch)
+
+        if bs != raw_bs:
+            out = self._postprocess_output_to_raw_bs(out, raw_bs)
+            forward_batch.batch_size = raw_expanded_bs
+            forward_batch.positions = buffers.positions[:raw_num_token]
+            forward_batch.seq_lens = buffers.seq_lens[:raw_expanded_bs]
+            forward_batch.req_pool_indices = buffers.req_pool_indices[:raw_expanded_bs]
+            if forward_batch.mrope_positions is not None:
+                forward_batch.mrope_positions = buffers.mrope_positions[
+                    :, :raw_num_token
+                ]
+            if forward_batch.seq_lens_cpu is not None:
+                forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:raw_expanded_bs]
+
+        return out

python/sglang/srt/speculative/frozen_kv_mtp_info.py

@@ -0,0 +1,60 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+from sglang.srt.mem_cache.memory_pool import KVCache
+from sglang.srt.speculative.eagle_info import (
+    EagleDraftInput,
+    EagleVerifyInput,
+)
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
+
+
+@dataclass(frozen=True)
+class FrozenKVMTPContext:
+    """Target KV pool + assistant-logical -> target-physical layer map."""
+
+    target_token_to_kv_pool: KVCache
+    physical_layer_ids: Dict[int, int]
+
+    def get_physical_layer_id(self, idx: int) -> int:
+        if idx not in self.physical_layer_ids:
+            raise KeyError(
+                f"FrozenKVMTPContext has no physical layer id for assistant "
+                f"logical index {idx}; available: {sorted(self.physical_layer_ids)}"
+            )
+        return self.physical_layer_ids[idx]
+
+
+@dataclass
+class FrozenKVMTPDraftInput(EagleDraftInput):
+    """Draft input for Frozen-KV MTP.
+
+    Frozen-KV MTP currently reuses the EAGLE scheduler/attention contract, but
+    has a dedicated type so algorithm-specific behavior can move here over time.
+    """
+
+    def __post_init__(self):
+        SpecInput.__init__(self, SpecInputType.FROZEN_KV_MTP_DRAFT)
+
+
+@dataclass
+class FrozenKVMTPVerifyInput(EagleVerifyInput):
+    """Verify input for Frozen-KV MTP."""
+
+    def __post_init__(self):
+        SpecInput.__init__(self, SpecInputType.FROZEN_KV_MTP_VERIFY)

python/sglang/srt/speculative/frozen_kv_mtp_utils.py

@@ -0,0 +1,155 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+
+
+@contextmanager
+def frozen_kv_target_view(
+    forward_batch: ForwardBatch,
+    kv_context: FrozenKVMTPContext,
+    draft_attn_backend: AttentionBackend,
+):
+    """Build attention metadata against committed target-prefix geometry.
+
+    Swaps ``draft_attn_backend.token_to_kv_pool`` to the frozen target pool
+    so any helper that reads ``get_token_to_kv_pool()`` during metadata init
+    sees the frozen target pool. Pool refs are derived from
+    ``get_attn_backend().token_to_kv_pool`` — the single backend-attribute
+    swap is seen by both readers (``get_token_to_kv_pool()`` and the
+    backend's own ``self.token_to_kv_pool``).
+    """
+    if kv_context is None:
+        raise RuntimeError(
+            "Frozen-KV MTP target view called before the model was bound; "
+            "bind the frozen KV context first."
+        )
+    saved_spec_info = forward_batch.spec_info
+    forward_batch.spec_info = None
+    saved_backend_pool = draft_attn_backend.token_to_kv_pool
+    draft_attn_backend.token_to_kv_pool = kv_context.target_token_to_kv_pool
+    try:
+        yield
+    finally:
+        forward_batch.spec_info = saved_spec_info
+        draft_attn_backend.token_to_kv_pool = saved_backend_pool
+
+
+@contextmanager
+def target_kv_pool_view(
+    forward_batch: ForwardBatch,
+    kv_context: FrozenKVMTPContext,
+    draft_attn_backend: AttentionBackend,
+):
+    """Run the draft model's forward with the target's frozen KV pool.
+
+    Swaps ``draft_attn_backend.token_to_kv_pool`` to the frozen target pool.
+    The single backend-attribute swap is seen by both readers —
+    ``get_token_to_kv_pool()`` (because it resolves through
+    ``get_attn_backend()``) and the backend's own ``self.token_to_kv_pool``
+    reads (because ``self is draft_attn_backend``).
+    """
+    if kv_context is None:
+        raise RuntimeError(
+            "Frozen-KV MTP target KV pool view called before the model was bound; "
+            "bind the frozen KV context first."
+        )
+    saved_backend_pool = draft_attn_backend.token_to_kv_pool
+    draft_attn_backend.token_to_kv_pool = kv_context.target_token_to_kv_pool
+    try:
+        yield
+    finally:
+        draft_attn_backend.token_to_kv_pool = saved_backend_pool
+
+
+def set_frozen_kv_positions(forward_batch: ForwardBatch, topk: int) -> None:
+    """Rope phase = last written target slot, not advanced per draft step."""
+    seq_lens = forward_batch.seq_lens
+    positions = torch.clamp(seq_lens - 1, min=0).to(torch.int64)
+    if (
+        topk > 1
+        and forward_batch.positions is not None
+        and forward_batch.positions.numel() == positions.numel() * topk
+    ):
+        positions = positions.repeat_interleave(topk, dim=0)
+    if forward_batch.positions is None:
+        forward_batch.positions = positions
+    else:
+        if forward_batch.positions.shape == positions.shape:
+            forward_batch.positions.copy_(positions)
+        else:
+            forward_batch.positions = positions
+
+
+def expand_for_topk_draft(forward_batch: ForwardBatch, topk: int) -> None:
+    """Repeat committed-prefix metadata for the active ``B * topk`` frontier."""
+    if topk == 1 or forward_batch.batch_size == 0:
+        return
+
+    if forward_batch.batch_size != forward_batch.seq_lens.shape[0]:
+        raise RuntimeError(
+            "Frozen-KV MTP topk expansion expects an unexpanded forward "
+            "batch where batch_size == len(seq_lens)."
+        )
+
+    forward_batch.batch_size *= topk
+    forward_batch.req_pool_indices = forward_batch.req_pool_indices.repeat_interleave(
+        topk, dim=0
+    )
+    forward_batch.seq_lens = forward_batch.seq_lens.repeat_interleave(topk, dim=0)
+    if forward_batch.seq_lens_cpu is not None:
+        forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu.repeat_interleave(
+            topk, dim=0
+        )
+        forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()
+    else:
+        forward_batch.seq_lens_sum = torch.sum(forward_batch.seq_lens).item()
+
+    positions = torch.clamp(forward_batch.seq_lens - 1, min=0).to(torch.int64)
+    forward_batch.positions = positions
+    forward_batch.num_token_non_padded_cpu = positions.numel()
+    if forward_batch.num_token_non_padded is not None:
+        forward_batch.num_token_non_padded.fill_(positions.numel())
+    if (
+        forward_batch.mrope_positions is not None
+        and forward_batch.mrope_positions.shape[-1] * topk == positions.numel()
+    ):
+        forward_batch.mrope_positions = forward_batch.mrope_positions.repeat_interleave(
+            topk, dim=-1
+        )
+
+
+def position_for_batch(batch: ScheduleBatch) -> torch.Tensor:
+    return torch.clamp(batch.seq_lens - 1, min=0).to(torch.int64)
+
+
+def select_last_extend_hidden(
+    batch: ScheduleBatch, hidden_states: torch.Tensor
+) -> torch.Tensor:
+    if hidden_states.shape[0] == batch.batch_size():
+        return hidden_states
+    lens = torch.tensor(batch.extend_lens, device=hidden_states.device)
+    last_indices = torch.cumsum(lens, dim=0) - 1
+    return hidden_states[last_indices.to(torch.long)]

python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py

@@ -0,0 +1,781 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Spec-v2 worker for Frozen-KV MTP (two layers, like ``eagle_worker_v2``).
+
+The frozen draft reads the target KV cache read-only and owns no KV pool, so
+its "draft extend" is not a model forward: it selects the last accepted token +
+target hidden state as the next-iter seed, and the seed forward runs at the
+start of the next draft.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.moe.utils import (
+    speculative_moe_a2a_backend_context,
+    speculative_moe_backend_context,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.base_spec_worker import EagleDraftWorkerBase
+from sglang.srt.speculative.eagle_utils import (
+    build_tree_kernel_efficient,
+    organize_draft_results,
+)
+from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2, _get_plan_stream
+from sglang.srt.speculative.frozen_kv_mtp_info import (
+    FrozenKVMTPContext,
+    FrozenKVMTPDraftInput,
+    FrozenKVMTPVerifyInput,
+)
+from sglang.srt.speculative.frozen_kv_mtp_utils import (
+    expand_for_topk_draft,
+    frozen_kv_target_view,
+    position_for_batch,
+    select_last_extend_hidden,
+    set_frozen_kv_positions,
+    target_kv_pool_view,
+)
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import (
+    draft_tp_context,
+    fast_topk,
+    select_top_k_tokens,
+    spec_stage_span,
+)
+from sglang.srt.utils import empty_context
+from sglang.srt.utils.async_probe import (
+    maybe_detect_inf,
+    maybe_detect_nan,
+    maybe_detect_oob,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
+    """Frozen-KV MTP draft worker.
+
+    The assistant reads target KV only. It reuses EAGLE's verify input/output
+    contract, but owns the seed and recurrent draft loop because there is no
+    assistant-side KV extension.
+    """
+
+    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,
+    ):
+        self.server_args = server_args
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self.target_worker = target_worker
+        self.page_size = server_args.page_size
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+        assert self.speculative_algorithm.is_frozen_kv_mtp(), (
+            "FrozenKVMTPDraftWorker should only be instantiated for "
+            "SpeculativeAlgorithm.FROZEN_KV_MTP, got "
+            f"{self.speculative_algorithm.name}."
+        )
+
+        # Target pools (read-only) are bound in alloc_memory_pool(), not here, so
+        # the worker can be built before the target pool exists (see #29021).
+        self.req_to_token_pool = None
+        self.token_to_kv_pool_allocator = None
+        self.draft_pool_config: Optional[MemoryPoolConfig] = None
+
+        self.hot_token_id = None
+
+        with (
+            empty_context()
+        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            # NOTE: call TpModelWorker.__init__ explicitly -- EagleDraftWorkerBase is
+            # an ABC with no __init__, so cooperative super() would be ambiguous.
+            TpModelWorker.__init__(
+                self,
+                server_args=server_args,
+                gpu_id=gpu_id,
+                tp_rank=tp_rank,
+                pp_rank=0,
+                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,
+                is_draft_worker=True,
+            )
+
+        embed, head = self.target_worker.model_runner.model.get_embed_and_head()
+        if hasattr(self.draft_model_runner.model, "set_embed_and_head"):
+            self.draft_model_runner.model.set_embed_and_head(embed, head)
+        else:
+            logger.debug(
+                "Draft model %s does not implement set_embed_and_head; "
+                "skipping target-embedding bind in Frozen-KV MTP skeleton.",
+                type(self.draft_model_runner.model).__name__,
+            )
+
+        self.kv_context: Optional[FrozenKVMTPContext] = None
+
+        self.draft_tp_context = (
+            draft_tp_context if server_args.enable_dp_attention else empty_context
+        )
+
+        self.draft_attn_backend = None
+        self.cuda_graph_runner = None
+        # Frozen draft has no draft-extend forward (seed-select only); keep these
+        # None so inherited probes (spec_v2_attn_backends, adaptive) stay typed.
+        self.draft_extend_attn_backend = None
+        self.cuda_graph_runner_for_draft_extend = None
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+
+        self.draft_pool_config = MemoryPoolConfig(
+            max_total_num_tokens=64,  # Dummy value
+            max_running_requests=memory_pool_config.max_running_requests,
+        )
+
+        # NOTE: call TpModelWorker explicitly -- EagleDraftWorkerBase precedes it in
+        # the MRO and its alloc_memory_pool is a no-op stub.
+        TpModelWorker.alloc_memory_pool(
+            self,
+            memory_pool_config=self.draft_pool_config,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+
+        if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
+            self._bind_kv_context()
+
+    def init_attention_backends(self):
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
+            TpModelWorker.init_attention_backends(self)
+            self.draft_attn_backend = self._init_draft_attn_backend()
+            self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
+
+    def init_cuda_graphs(self):
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
+            TpModelWorker.init_cuda_graphs(self, capture_decode_cuda_graph=False)
+            self._capture_cuda_graphs()
+
+    @property
+    def draft_model_runner(self):
+        return self.model_runner
+
+    @property
+    def draft_runner(self):
+        # Alias for the inherited EAGLEWorkerV2 forward/verify skeleton, which
+        # reads `draft_worker.draft_runner`.
+        return self.model_runner
+
+    def get_attn_backend(self):  # pragma: no cover - exposed for adaptive
+        return self.draft_attn_backend
+
+    def clear_cache_pool(self):
+        pass
+
+    def _resolve_draft_backend_type(self) -> str:
+        return (
+            self.server_args.speculative_draft_attention_backend
+            or self.server_args.decode_attention_backend
+            or self.server_args.attention_backend
+        )
+
+    def _init_draft_attn_backend(self):
+        if self.topk == 1:
+            return self.draft_model_runner.attn_backend
+
+        backend_type = self._resolve_draft_backend_type()
+        if backend_type != "triton":
+            raise ValueError(
+                "Frozen-KV MTP topk > 1 currently supports only the triton "
+                f"attention backend, got {backend_type}."
+            )
+        return self._init_triton_draft_attn_backend()
+
+    def _init_triton_draft_attn_backend(self):
+        from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
+
+        max_bs = self.req_to_token_pool.size * self.topk
+        kv_indptr_buf = torch.zeros(
+            (max_bs + 1,), dtype=torch.int32, device=self.draft_model_runner.device
+        )
+        return TritonAttnBackend(
+            self.draft_model_runner,
+            skip_prefill=True,
+            kv_indptr_buf=kv_indptr_buf,
+        )
+
+    def _bind_kv_context(self) -> None:
+        draft_model = self.draft_model_runner.model
+        if not hasattr(draft_model, "build_frozen_kv_mtp_context") or not hasattr(
+            draft_model, "bind_frozen_kv_context"
+        ):
+            logger.debug(
+                "Draft model %s does not implement Frozen-KV MTP context hooks; "
+                "skipping frozen-kv bind.",
+                type(draft_model).__name__,
+            )
+            return
+
+        ctx = draft_model.build_frozen_kv_mtp_context(
+            target_model=self.target_worker.model_runner.model,
+            target_token_to_kv_pool=self.target_worker.model_runner.token_to_kv_pool,
+        )
+        draft_model.bind_frozen_kv_context(ctx)
+        self.kv_context = ctx
+
+    def _frozen_kv_target_view(self, forward_batch: ForwardBatch):
+        return frozen_kv_target_view(
+            forward_batch, self.kv_context, self.draft_attn_backend
+        )
+
+    def _target_kv_pool_view(self, forward_batch: ForwardBatch):
+        return target_kv_pool_view(
+            forward_batch, self.kv_context, self.draft_attn_backend
+        )
+
+    def _set_positions(self, forward_batch: ForwardBatch) -> None:
+        set_frozen_kv_positions(forward_batch, self.topk)
+
+    def _expand_for_topk_draft(self, forward_batch: ForwardBatch) -> None:
+        expand_for_topk_draft(forward_batch, self.topk)
+
+    def _position_for_batch(self, batch: ScheduleBatch) -> torch.Tensor:
+        return position_for_batch(batch)
+
+    @property
+    def _recurrent_hidden_size(self) -> int:
+        return int(self.draft_model_runner.model.backbone_hidden_size)
+
+    def _init_frozen_kv_metadata(self, forward_batch: ForwardBatch) -> None:
+        if forward_batch.forward_mode.is_idle():
+            return
+        if forward_batch.seq_lens_cpu is not None:
+            forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()
+        else:
+            forward_batch.seq_lens_sum = torch.sum(forward_batch.seq_lens).item()
+        with self._frozen_kv_target_view(forward_batch):
+            self.draft_attn_backend.init_forward_metadata(forward_batch)
+        forward_batch.mark_forward_metadata_ready()
+
+    def _init_frozen_kv_metadata_capture_cuda_graph(
+        self, forward_batch: ForwardBatch
+    ) -> None:
+        with self._frozen_kv_target_view(forward_batch):
+            self.draft_attn_backend.init_forward_metadata_out_graph(
+                forward_batch, in_capture=True
+            )
+        forward_batch.mark_forward_metadata_ready()
+
+    def _init_frozen_kv_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int, seq_lens_sum: int
+    ) -> None:
+        from types import SimpleNamespace
+
+        fb_view = SimpleNamespace(
+            batch_size=bs,
+            forward_mode=ForwardMode.DECODE,
+            input_ids=getattr(forward_batch, "input_ids", None),
+            req_pool_indices=forward_batch.req_pool_indices[:bs],
+            seq_lens=forward_batch.seq_lens[:bs],
+            seq_lens_sum=seq_lens_sum,
+            seq_lens_cpu=(
+                forward_batch.seq_lens_cpu[:bs]
+                if forward_batch.seq_lens_cpu is not None
+                else None
+            ),
+            encoder_lens=None,
+            out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
+            spec_info=None,
+        )
+        with self._frozen_kv_target_view(forward_batch):
+            self.draft_attn_backend.init_forward_metadata_out_graph(fb_view)
+
+    def _capture_cuda_graphs(self) -> None:
+        if cuda_graph_fully_disabled() or self.speculative_num_steps <= 1:
+            return
+        if self.target_worker.device != "cuda":
+            logger.info(
+                "Frozen-KV MTP draft CUDA graph is only supported on CUDA; "
+                "running the draft loop eagerly on %s.",
+                self.target_worker.device,
+            )
+            return
+
+        from sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner import (
+            FrozenKVMTPCudaGraphRunner,
+        )
+
+        logger.info("Capture Frozen-KV MTP draft cuda graph begin.")
+        self.cuda_graph_runner = FrozenKVMTPCudaGraphRunner(self)
+        logger.info("Capture Frozen-KV MTP draft cuda graph end.")
+
+    def _select_last_extend_hidden(
+        self, batch: ScheduleBatch, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        return select_last_extend_hidden(batch, hidden_states)
+
+    def _idle_seed(self) -> FrozenKVMTPDraftInput:
+        return FrozenKVMTPDraftInput.create_idle_input(
+            device=self.device,
+            hidden_size=self._recurrent_hidden_size,
+            dtype=self.model_config.dtype,
+            topk=self.topk,
+            capture_hidden_mode=CaptureHiddenMode.LAST,
+        )
+
+    def _build_seed_draft_input(
+        self,
+        last_token_ids: torch.Tensor,
+        last_hidden_states: torch.Tensor,
+    ) -> FrozenKVMTPDraftInput:
+        """Build the next-iter seed ``FrozenKVMTPDraftInput`` from (bonus token,
+        target hidden). No forward here -- the seed forward runs inside the
+        captured draft graph (see ``draft_forward``'s seed iter)."""
+        if last_token_ids.numel() == 0:
+            return self._idle_seed()
+
+        stashed = FrozenKVMTPDraftInput()
+        stashed.bonus_tokens = last_token_ids.to(torch.int64)
+        stashed.hidden_states = last_hidden_states
+        # Real-shaped zeros so inherited `filter_batch`/`merge_batch` can slice
+        # them between iters; overwritten by the captured seed iter.
+        bs = last_token_ids.shape[0]
+        device = last_token_ids.device
+        stashed.topk_p = torch.zeros(
+            (bs, self.topk), device=device, dtype=torch.float32
+        )
+        stashed.topk_index = torch.zeros(
+            (bs, self.topk), device=device, dtype=torch.int64
+        )
+        stashed.capture_hidden_mode = CaptureHiddenMode.LAST
+        stashed.num_tokens_per_req = 1
+        stashed.num_tokens_for_logprob_per_req = 1
+        return stashed
+
+    def draft(self, batch: ScheduleBatch):
+        if batch.forward_mode.is_idle():
+            return FrozenKVMTPVerifyInput.create_idle_input(
+                self.topk,
+                self.speculative_num_steps,
+                self.speculative_num_draft_tokens,
+            )
+
+        spec_info = batch.spec_info
+        assert isinstance(spec_info, FrozenKVMTPDraftInput)
+
+        # NOTE: per-iter bookkeeping (penalty cumulation, maybe_evict_swa,
+        # decode_batch_idx tick) is done by the scheduler-driven
+        # eagle_utils.eagle_prepare_for_decode (see
+        # ScheduleBatch.prepare_for_decode), not here -- matching EAGLE v2.
+        # Repeating evict/tick here would double-run them: the idx clock
+        # gates SWA eviction timing and the SWA prefix-lock release.
+
+        spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
+        spec_info.num_tokens_per_req = self.topk
+        spec_info.num_tokens_for_logprob_per_req = self.topk
+        spec_info.positions = self._position_for_batch(batch)
+        batch.seq_lens_sum = torch.sum(batch.seq_lens).item()
+        batch.return_hidden_states = False
+
+        forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
+        assert forward_batch.capture_hidden_mode == CaptureHiddenMode.LAST
+        self._set_positions(forward_batch)
+        self._expand_for_topk_draft(forward_batch)
+
+        # Frozen draft never writes KV; None signals fill_from to skip the slot.
+        forward_batch.out_cache_loc = None
+
+        can_run_cuda_graph = (
+            self.cuda_graph_runner
+            and self.cuda_graph_runner.can_run_graph(forward_batch)
+        )
+        if can_run_cuda_graph:
+            parent_list, top_scores_index, draft_tokens = (
+                self.cuda_graph_runner.execute(forward_batch)
+            )
+        else:
+            forward_batch.can_run_dp_cuda_graph = False
+            parent_list, top_scores_index, draft_tokens = self.draft_forward(
+                forward_batch
+            )
+
+        (
+            tree_mask,
+            position,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            draft_tokens,
+        ) = build_tree_kernel_efficient(
+            spec_info.bonus_tokens,
+            parent_list,
+            top_scores_index,
+            draft_tokens,
+            batch.seq_lens,
+            batch.seq_lens_sum,
+            self.topk,
+            self.speculative_num_steps,
+            self.speculative_num_draft_tokens,
+        )
+
+        return FrozenKVMTPVerifyInput(
+            draft_token=draft_tokens,
+            custom_mask=tree_mask,
+            positions=position,
+            retrieve_index=retrieve_index,
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            retrieve_cum_len=None,
+            spec_steps=self.speculative_num_steps,
+            topk=self.topk,
+            draft_token_num=self.speculative_num_draft_tokens,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+            seq_lens_sum=batch.seq_lens_sum,
+            seq_lens_cpu=batch.seq_lens_cpu,
+        )
+
+    def draft_forward(self, forward_batch: ForwardBatch):
+        spec_info = forward_batch.spec_info
+        assert isinstance(spec_info, FrozenKVMTPDraftInput)
+
+        score_list: list[torch.Tensor] = []
+        token_list: list[torch.Tensor] = []
+        parents_list: list[torch.Tensor] = []
+
+        # Seed + recurrent iters share the same `seq_lens - 1` rope position,
+        # so one init covers the loop. Must run even at num_steps == 1.
+        if forward_batch.needs_forward_metadata_init():
+            self._init_frozen_kv_metadata(forward_batch)
+
+        # Seed iter: assistant forward on (bonus_token, target_h) to produce
+        # iter-0 `(topk_p, topk_index, hidden_states)`. For topk>1, replicate
+        # to `bs*topk` to match kernel shapes, then slice back per-req.
+        bonus_tokens = spec_info.bonus_tokens
+        target_hidden = spec_info.hidden_states
+        if self.topk > 1:
+            seed_input_ids = bonus_tokens.repeat_interleave(self.topk, dim=0)
+            seed_prev_hidden = target_hidden.repeat_interleave(self.topk, dim=0)
+        else:
+            seed_input_ids = bonus_tokens
+            seed_prev_hidden = target_hidden
+
+        forward_batch.input_ids = seed_input_ids
+        forward_batch.spec_info.hidden_states = seed_prev_hidden
+        self._set_positions(forward_batch)
+
+        with (
+            self._target_kv_pool_view(forward_batch),
+            forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
+        ):
+            seed_output = self.draft_model_runner.forward(forward_batch).logits_output
+
+        maybe_detect_nan(
+            seed_output.next_token_logits, "frozen_kv_mtp_draft: seed iter"
+        )
+
+        if self.topk > 1:
+            seed_next_logits = seed_output.next_token_logits[:: self.topk]
+            seed_hidden_per_req = seed_output.hidden_states[:: self.topk]
+        else:
+            seed_next_logits = seed_output.next_token_logits
+            seed_hidden_per_req = seed_output.hidden_states
+
+        probs = torch.softmax(seed_next_logits, dim=-1)
+        topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
+        maybe_detect_oob(
+            topk_index,
+            0,
+            seed_next_logits.shape[-1],
+            "frozen_kv_mtp_draft: seed topk_index OOB",
+        )
+        hidden_states = seed_hidden_per_req
+
+        scores = None
+        for i in range(self.speculative_num_steps):
+            input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
+                i, topk_p, topk_index, hidden_states, scores, self.topk
+            )
+            score_list.append(tree_info[0])
+            token_list.append(tree_info[1])
+            parents_list.append(tree_info[2])
+
+            if i == self.speculative_num_steps - 1:
+                break
+
+            forward_batch.input_ids = input_ids
+            forward_batch.spec_info.hidden_states = hidden_states
+            self._set_positions(forward_batch)
+
+            with (
+                self._target_kv_pool_view(forward_batch),
+                forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
+            ):
+                logits_output = self.draft_model_runner.forward(
+                    forward_batch
+                ).logits_output
+
+            maybe_detect_nan(
+                logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
+            )
+            maybe_detect_inf(
+                logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
+            )
+            probs = torch.softmax(logits_output.next_token_logits, dim=-1)
+            topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
+            maybe_detect_oob(
+                topk_index,
+                0,
+                logits_output.next_token_logits.shape[-1],
+                "frozen_kv_mtp_draft: topk_index OOB",
+            )
+            hidden_states = logits_output.hidden_states
+
+        return organize_draft_results(
+            score_list, token_list, parents_list, self.speculative_num_draft_tokens
+        )
+
+    def draft_extend(self):
+        # EagleDraftWorkerBase contract. Frozen has no draft-KV extend forward; the
+        # orchestrator calls `_draft_extend_for_{prefill,decode}` directly.
+        pass
+
+    def _draft_extend_for_prefill(
+        self,
+        batch: ScheduleBatch,
+        target_hidden_states: torch.Tensor,
+        next_token_ids: torch.Tensor,
+        mm_input_embeds: Optional[torch.Tensor] = None,
+    ) -> FrozenKVMTPDraftInput:
+        """Seed for the first decode iter after prefill. Frozen draft writes no
+        KV (reads target KV), so unlike EAGLE there is no draft-extend forward:
+        just select the last prompt hidden + bonus token and stash the seed."""
+        del mm_input_embeds  # frozen seed needs no input embeds
+        if batch.forward_mode.is_idle():
+            return self._idle_seed()
+        last_hidden = self._select_last_extend_hidden(batch, target_hidden_states)
+        return self._build_seed_draft_input(next_token_ids, last_hidden)
+
+    def _draft_extend_for_decode(self, batch: ScheduleBatch, batch_result) -> None:
+        """Frozen 'draft extend': no forward. Pull the last accepted token's
+        target hidden from the verify output and stash it as the next-iter seed.
+
+        Replaces verify's `EagleDraftInput` with a `FrozenKVMTPDraftInput` so the
+        next draft passes the FROZEN_KV_MTP attn-backend assertions.
+        """
+        if batch.forward_mode.is_idle():
+            batch_result.next_draft_input = self._idle_seed()
+            return
+
+        bs = len(batch.seq_lens)
+        # Same per-req select_index EAGLE uses on its draft-extend output: the
+        # last accepted node (accept_lens - 1) in each per-req block of width
+        # num_draft_tokens. Verify already compacted the accepted path to the
+        # front (topk > 1) / it is the front chain (topk == 1).
+        select_index = (
+            torch.arange(bs, device=self.device) * self.speculative_num_draft_tokens
+            + batch_result.accept_lens
+            - 1
+        )
+        last_hidden = batch_result.logits_output.hidden_states[select_index]
+        bonus_tokens = batch_result.next_draft_input.bonus_tokens
+        batch_result.next_draft_input = self._build_seed_draft_input(
+            bonus_tokens, last_hidden
+        )
+
+
+class FrozenKVMTPWorkerV2(EAGLEWorkerV2):
+    """Spec-v2 (overlap) orchestrator for Frozen-KV MTP.
+
+    Reuses ``EAGLEWorkerV2``'s verify / ``move_accept_tokens`` / forward
+    skeleton verbatim; only the draft worker and the seed-based draft-extend
+    are frozen-specific.
+    """
+
+    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,
+    ):
+        # NOTE: intentionally does NOT call EAGLEWorkerV2.__init__ -- that builds
+        # an EagleDraftWorker (with its own draft KV pool). The frozen draft owns
+        # no KV, so we mirror the relevant setup and build a FrozenKVMTPDraftWorker.
+        self.server_args = server_args
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.tp_rank = tp_rank
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self._target_worker = target_worker
+        self.page_size = server_args.page_size
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+            target_worker.get_memory_pool()
+        )
+        # Match the draft context length to the target (assistant reads target KV).
+        server_args.override(
+            "spec_worker.match_target_context_length",
+            context_length=target_worker.model_runner.model_config.context_len,
+        )
+
+        self._draft_worker = FrozenKVMTPDraftWorker(
+            server_args,
+            gpu_id,
+            tp_rank,
+            dp_rank,
+            moe_ep_rank,
+            attn_cp_rank,
+            moe_dp_rank,
+            nccl_port,
+            target_worker,
+        )
+
+        # Frozen MTP does not wire the adaptive controller yet.
+        assert (
+            not server_args.speculative_adaptive
+        ), "Frozen-KV MTP does not support adaptive speculative decoding yet."
+        self.adaptive_controller = None
+
+        # Some dummy tensors (parity with EAGLEWorkerV2 init).
+        self.num_new_pages_per_topk = torch.empty(
+            (), dtype=torch.int64, device=self.device
+        )
+        self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
+
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+
+    @property
+    def spec_v2_attn_backends(self) -> tuple:
+        # Frozen draft touches no draft-extend backend; only target + draft.
+        return (
+            self._target_worker.model_runner.attn_backend,
+            self._draft_worker.draft_attn_backend,
+        )
+
+    def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None):
+        # Mirrors EAGLEWorkerV2.forward_batch_generation; the only frozen-specific
+        # change is the idle draft-input (FrozenKVMTPDraftInput + recurrent hidden
+        # size). The draft / seed-based draft-extend hooks are FrozenKVMTPDraftWorker's.
+        if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
+            # Target prefill (frozen is never standalone -> capture FULL hidden).
+            batch.capture_hidden_mode = CaptureHiddenMode.FULL
+            batch_output = self.target_worker.forward_batch_generation(batch)
+
+            # Spec_v2 convention: batch.seq_lens = length BEFORE this iter's tokens.
+            batch_output.new_seq_lens = batch.seq_lens
+            # Publish before draft-extend so the fence is at target-end.
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+
+            # Draft prefill seed (no forward).
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft_extend"),
+            ):
+                batch_output.next_draft_input = (
+                    self.draft_worker._draft_extend_for_prefill(
+                        batch,
+                        batch_output.logits_output.hidden_states,
+                        batch_output.next_token_ids,
+                        batch_output.logits_output.mm_input_embeds,
+                    )
+                )
+                return batch_output
+        else:
+            self.activate_step_by_batch(batch.seq_lens.shape[0])
+
+            if batch.spec_info is None:
+                batch.spec_info = self.draft_worker._idle_seed()
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft"),
+            ):
+                verify_input = self.draft_worker.draft(batch)
+            assert verify_input.is_verify_input()
+            batch.spec_info = verify_input
+            batch_output = self.verify(batch)
+            # Publish before draft-extend so the fence is at verify-end.
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft_extend"),
+            ):
+                self.draft_worker._draft_extend_for_decode(batch, batch_output)
+
+            return batch_output

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -0,0 +1,718 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import time
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+import torch
+
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    forward_context,
+)
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+from sglang.srt.model_executor.runner import (
+    DecodeCudaGraphRunner,
+    DeepEPCudaGraphRunnerAdapter,
+    ShapeKey,
+    get_batch_sizes_to_capture,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
+from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+from sglang.srt.runtime_context import get_flags
+from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
+from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim
+from sglang.srt.speculative.spec_utils import fast_topk
+from sglang.srt.utils import (
+    get_available_gpu_memory,
+    require_attn_tp_gather,
+    require_gathered_buffer,
+    require_mlp_sync,
+    require_mlp_tp_gather,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.multi_layer_eagle_worker_v2 import (
+        MultiLayerEagleDraftWorker,
+    )
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers):
+    """A single persistent buffer set shared by every MTP draft step."""
+
+    input_ids: torch.Tensor
+    out_cache_loc: torch.Tensor
+    positions: torch.Tensor
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    req_pool_indices: torch.Tensor
+    num_correct_drafts: torch.Tensor
+    num_accept_tokens: torch.Tensor
+    extend_seq_lens: torch.Tensor
+    extend_start_loc: torch.Tensor
+    # Flat index (into the token dimension) of each request's last accepted
+    # token. Used both by the in-graph top-k gather and by the worker's
+    # per-step input_ids rotation.
+    select_index: torch.Tensor
+    mrope_positions: torch.Tensor
+    hidden_states: torch.Tensor
+    next_token_logits_buffer: torch.Tensor
+    global_num_tokens_gpu: Optional[torch.Tensor]
+    global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+
+
+class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
+    """Per-step multi-layer EAGLE draft-extend runner.
+
+    Subclasses DecodeCudaGraphRunner. All steps share a single buffer set
+    owned by the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
+    so initialization is split: __init__ does basic field setup, and
+    init_buffers_and_capture (called by the composite once the shared buffers
+    exist) attaches them and runs capture.
+    """
+
+    def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int):
+        # Parse args
+        self.step = step
+        self.eagle_worker = eagle_worker
+        self.model_runner = model_runner = eagle_worker.mtp_model_runner(self.step)
+        self.forward_mode = ForwardMode.DRAFT_EXTEND_V2
+
+        # Fields the parent's capture() reads:
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.tp_size = model_runner.tp_size
+        self.dp_size = model_runner.server_args.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.enable_torch_compile = get_flags().capture.enable_torch_compile
+        self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
+        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
+        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
+        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
+        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+        self.enable_pdmux = model_runner.server_args.enable_pdmux
+        self.speculative_num_steps = model_runner.server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = (
+            model_runner.server_args.speculative_num_draft_tokens
+        )
+        self.topk = model_runner.server_args.speculative_eagle_topk
+        self.enable_profile_cuda_graph = (
+            model_runner.server_args.enable_profile_cuda_graph
+        )
+        self.attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
+
+        # Disable parent paths that don't apply.
+        self.compile_bs = []
+        self.record_nolora_graph = False
+        self.is_dllm = False
+
+        self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
+
+        self.capture_forward_mode = self.forward_mode
+        self.capture_hidden_mode = CaptureHiddenMode.FULL
+
+        self.capture_bs, _ = get_batch_sizes_to_capture(model_runner)
+        self.padded_static_len = -1
+
+        # Fixed window: every step extends each request by the same number of
+        # tokens, which lets all steps share one buffer set.
+        self.num_tokens_per_bs = self.speculative_num_draft_tokens
+        self.max_bs = max(self.capture_bs)
+        self.max_num_token = self.max_bs * self.num_tokens_per_bs
+        self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs
+
+        self.eagle_worker.draft_extend_attn_backend_list[
+            self.step
+        ].init_cuda_graph_state(self.max_bs, self.max_num_token)
+        self.seq_len_fill_value = self.eagle_worker.draft_extend_attn_backend_list[
+            self.step
+        ].get_cuda_graph_seq_len_fill_value()
+
+    def init_buffers_and_capture(self, buffers: MultiLayerEagleDraftExtendInputBuffers):
+        """Attach the shared buffer set and capture this step's graphs."""
+        self.buffers = buffers
+
+        if self.enable_torch_compile:
+            set_torch_compile_config()
+
+        self.backend = resolve_decode_backend(self)
+
+        try:
+            with model_capture_mode():
+                self.capture()
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+            )
+
+    def _replay_graph(self, shape_key, forward_batch):
+        return self.backend.replay(shape_key, forward_batch)
+
+    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
+        return ShapeKey(size=bs)
+
+    def can_run_graph(self, forward_batch: ForwardBatch):
+        if self.require_mlp_tp_gather:
+            cuda_graph_bs = (
+                max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
+                if self.model_runner.spec_algorithm.is_eagle()
+                else max(forward_batch.global_num_tokens_cpu)
+            )
+        else:
+            cuda_graph_bs = forward_batch.seq_lens.numel()
+
+        is_bs_supported = (
+            self.backend.can_run(forward_batch, cuda_graph_bs)
+            if self.disable_padding
+            else cuda_graph_bs <= self.max_bs
+        )
+
+        if self.require_mlp_sync:
+            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+
+        return is_bs_supported
+
+    def get_forward_batch(self, bs: int) -> ForwardBatch:
+        buffers = self.buffers
+        num_tokens = bs * self.num_tokens_per_bs
+
+        input_ids = buffers.input_ids[:num_tokens]
+        req_pool_indices = buffers.req_pool_indices[:bs]
+        seq_lens = buffers.seq_lens[:bs]
+        seq_lens_cpu = buffers.seq_lens_cpu[:bs]
+        extend_seq_lens = buffers.extend_seq_lens[:bs]
+        extend_seq_lens_cpu = self.extend_seq_lens_cpu[:bs]
+        extend_start_loc = buffers.extend_start_loc[:bs]
+        num_correct_drafts = buffers.num_correct_drafts[:bs]
+        num_accept_tokens = buffers.num_accept_tokens[:bs]
+        out_cache_loc = buffers.out_cache_loc[:num_tokens]
+        positions = buffers.positions[:num_tokens]
+        mrope_positions = buffers.mrope_positions[:, :num_tokens]
+        hidden_states = buffers.hidden_states[:num_tokens]
+        next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens]
+
+        if self.require_mlp_tp_gather:
+            global_num_tokens_cpu = [num_tokens] * self.dp_size
+            global_num_tokens_for_logprob_cpu = [num_tokens] * self.dp_size
+        elif self.require_attn_tp_gather:
+            global_num_tokens_cpu = [num_tokens]
+            # DRAFT_EXTEND_V2 produces logits for all tokens, not bs (see mlp branch above)
+            global_num_tokens_for_logprob_cpu = [num_tokens]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            buffers.global_num_tokens_gpu.copy_(
+                torch.tensor(
+                    global_num_tokens_cpu,
+                    dtype=torch.int32,
+                    device=buffers.input_ids.device,
+                )
+            )
+            buffers.global_num_tokens_for_logprob_gpu.copy_(
+                torch.tensor(
+                    global_num_tokens_for_logprob_cpu,
+                    dtype=torch.int32,
+                    device=buffers.input_ids.device,
+                )
+            )
+        else:
+            global_dp_buffer_len = None
+
+        spec_info = EagleDraftExtendInput(
+            hidden_states=hidden_states,
+            num_correct_drafts=num_correct_drafts,
+            num_accept_tokens=num_accept_tokens,
+        )
+        spec_info.positions = None
+
+        capture_mode = (
+            CaptureHiddenMode.NULL
+            if self.model_runner.spec_algorithm.is_standalone()
+            else CaptureHiddenMode.FULL
+        )
+
+        # Forward batch
+        forward_batch = ForwardBatch(
+            forward_mode=self.forward_mode,
+            batch_size=bs,
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            next_token_logits_buffer=next_token_logits_buffer,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=seq_lens.sum().item(),
+            return_logprob=False,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            global_num_tokens_gpu=buffers.global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
+            global_num_tokens_cpu=global_num_tokens_cpu,
+            spec_algorithm=self.model_runner.spec_algorithm,
+            spec_info=spec_info,
+            capture_hidden_mode=capture_mode,
+            extend_seq_lens=extend_seq_lens,
+            extend_seq_lens_cpu=extend_seq_lens_cpu,
+            padded_static_len=self.padded_static_len,
+            extend_start_loc=extend_start_loc,
+            extend_num_tokens=self.num_tokens_per_bs * bs,
+            num_token_non_padded_cpu=self.num_tokens_per_bs * bs,
+            return_hidden_states_before_norm=True,
+        )
+        return forward_batch
+
+    def _postprocess_forward_batch(self, forward_batch: ForwardBatch, bs: int):
+        """Hook for subclasses to mutate the captured forward batch."""
+        return forward_batch
+
+    def _compute_topk(self, ret, bs: int):
+        """Compute top-k on the last accepted token's logits and attach it to
+        ``ret``. The gather index lives in a persistent buffer, so the captured
+        graph reads the right rows on each replay. Overridable so distributed
+        (vocab-sharded) builds can plug in an all-reduce-aware sampler."""
+        buffers = self.buffers
+        probs = torch.softmax(ret.next_token_logits[buffers.select_index[:bs]], dim=-1)
+        ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
+
+    def capture_one_shape(
+        self,
+        size: int,
+        forward: Callable,
+        stream_idx: Optional[int] = None,
+        variant_label: Optional[str] = None,
+    ):
+        bs = size
+        buffers = self.buffers
+
+        num_tokens = bs * self.num_tokens_per_bs
+        forward_batch = self.get_forward_batch(bs)
+        forward_batch = self._postprocess_forward_batch(forward_batch, bs)
+        attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
+
+        def run_once():
+            attn_backend.init_forward_metadata_in_graph(forward_batch)
+
+            # Clean intermediate result cache for DP attention
+            forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(
+                forward_batch.global_dp_buffer_len,
+                num_tokens,
+                forward_batch.dp_padding_mode.is_max_len(),
+                forward_batch.global_num_tokens_cpu,
+            )
+            set_is_extend_in_batch(False)
+
+            output_cache_loc_backup = forward_batch.out_cache_loc
+            hidden_states_backup = forward_batch.spec_info.hidden_states
+
+            ret = self.model_runner.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+            )
+
+            if (
+                self.eagle_worker.chain_mtp_hidden_states
+                and ret.hidden_states is not None
+            ):
+                buffers.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens])
+
+            self._compute_topk(ret, bs)
+
+            forward_batch.out_cache_loc = output_cache_loc_backup
+            forward_batch.spec_info.hidden_states = hidden_states_backup
+            return ret
+
+        with forward_context(ForwardContext(attn_backend=attn_backend)):
+            attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
+            self.deepep_adapter.capture(is_extend_in_batch=True)
+            shape_key = self._make_graph_key(bs)
+            post_warmup_hook = getattr(
+                self.attn_backend, "on_after_cuda_graph_warmup", None
+            )
+            maybe_flashinfer_autotune_speculative_draft(
+                self,
+                run_once,
+                post_warmup_hook=post_warmup_hook,
+                skip_logits=False,
+            )
+            self.backend.capture_one(
+                shape_key,
+                run_once,
+                dummies=None,
+                post_warmup_hook=post_warmup_hook,
+            )
+
+    def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput):
+        """Init this step's attention metadata for the prepared bucket and
+        replay its graph. Buffers must already be populated by the composite
+        runner's ``prepare`` (step 0) or by the previous step's in-graph chain
+        write + worker-side rotation (steps > 0)."""
+        self.deepep_adapter.replay()
+        buffers = self.buffers
+        num_tokens = bs * self.num_tokens_per_bs
+
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(num_tokens)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(num_tokens)
+
+        fb_view = SimpleNamespace(
+            batch_size=bs,
+            forward_mode=self.forward_mode,
+            input_ids=buffers.input_ids[:num_tokens],
+            req_pool_indices=buffers.req_pool_indices,
+            seq_lens=buffers.seq_lens,
+            seq_lens_sum=seq_lens_sum,
+            seq_lens_cpu=buffers.seq_lens_cpu,
+            encoder_lens=None,
+            # per-step write target; out_cache_loc is frozen at prepare() time.
+            out_cache_loc=buffers.out_cache_loc[:num_tokens],
+            spec_info=spec_info,
+        )
+        self.eagle_worker.draft_extend_attn_backend_list[
+            self.step
+        ].init_forward_metadata_out_graph(fb_view)
+
+        self.bs = bs
+        shape_key = self._make_graph_key(bs)
+        return self._replay_graph(shape_key, fb_view)
+
+
+class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
+    """Owns one shared buffer set and the per-step runners.
+
+    Usage from the worker::
+
+        runner.prepare(forward_batch)
+        for step in range(num_steps):
+            _, topk_p, topk_index = runner.replay(step)
+            if step < num_steps - 1:
+                rotate_input_ids(...)  # advance the draft chain
+
+    Not itself a DecodeCudaGraphRunner -- it only routes work to the per-step
+    runners.
+    """
+
+    def __init__(self, eagle_worker: MultiLayerEagleDraftWorker):
+        self.eagle_worker = eagle_worker
+        self.device = eagle_worker.device
+        self.gpu_id = eagle_worker.gpu_id
+        self.speculative_num_steps = eagle_worker.speculative_num_steps
+        self.draft_extend_attn_backend_list = (
+            eagle_worker.draft_extend_attn_backend_list
+        )
+
+        self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = []
+        self.seq_len_fill_value = 1
+        self.max_bs = 1
+        self.num_tokens_per_bs = 1
+
+        self._init_and_capture()
+
+    def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendCudaGraphRunner:
+        return MultiLayerEagleDraftExtendCudaGraphRunner(self.eagle_worker, step)
+
+    def _capture_context(self, step: int):
+        """Context manager active while capturing ``step``'s graphs. Subclasses
+        can use it e.g. to temporarily expose a sharded local vocab size."""
+        return contextlib.nullcontext()
+
+    def _on_runners_created(self):
+        """Hook called after all per-step runners exist but before buffers are
+        allocated/captured (e.g. to allocate shared sconv buffers)."""
+
+    def _cuda_graph_disabled(self) -> bool:
+        return check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED)
+
+    def _init_and_capture(self):
+        if self._cuda_graph_disabled():
+            self.runners = [None] * self.speculative_num_steps
+            return
+
+        self.runners = []
+
+        # 1. Construct per-step runners (each initializes its own attn cuda
+        #    graph state). They share the same fixed window size.
+        for step in range(self.speculative_num_steps):
+            if self.draft_extend_attn_backend_list[step]:
+                runner = self._create_runner(step)
+                self.runners.append(runner)
+                self.seq_len_fill_value = runner.seq_len_fill_value
+                self.max_bs = runner.max_bs
+                self.num_tokens_per_bs = runner.num_tokens_per_bs
+                self.capture_bs = runner.capture_bs
+                self.require_gathered_buffer = runner.require_gathered_buffer
+                self.require_mlp_tp_gather = runner.require_mlp_tp_gather
+                self.require_mlp_sync = runner.require_mlp_sync
+                self.disable_padding = runner.disable_padding
+            else:
+                self.runners.append(None)
+
+        self._on_runners_created()
+
+        # 2. Allocate the single shared buffer set and capture each step in
+        #    reverse order.
+        self.buffers = self._allocate_buffers()
+        for step in range(self.speculative_num_steps - 1, -1, -1):
+            if self.runners[step] is not None:
+                tic = time.perf_counter()
+                before_mem = get_available_gpu_memory(self.device, self.gpu_id)
+                logger.info(
+                    f"Capture draft extend CUDA graph begin. step={step}, "
+                    f"avail mem={before_mem:.2f} GB"
+                )
+
+                with self._capture_context(step):
+                    self.runners[step].init_buffers_and_capture(self.buffers)
+
+                after_mem = get_available_gpu_memory(self.device, self.gpu_id)
+                logger.info(
+                    "Capture draft extend CUDA graph end. "
+                    f"step={step}, elapsed={time.perf_counter() - tic:.2f} s, "
+                    f"mem usage={(before_mem - after_mem):.2f} GB, "
+                    f"avail mem={after_mem:.2f} GB."
+                )
+
+    def _vocab_size(self) -> int:
+        model_runner = self.eagle_worker.mtp_model_runner(0)
+        if hasattr(model_runner.model_config.hf_config, "draft_vocab_size"):
+            return model_runner.model_config.hf_config.draft_vocab_size
+        if hasattr(model_runner.model_config.hf_config, "hot_vocab_size"):
+            return model_runner.model_config.hf_config.hot_vocab_size
+        return model_runner.model_config.vocab_size
+
+    def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers:
+        runner = next(r for r in self.runners if r is not None)
+        model_runner = runner.model_runner
+        max_bs = self.max_bs
+        num_tokens_per_bs = self.num_tokens_per_bs
+        max_num_token = max_bs * num_tokens_per_bs
+        hidden_size = get_draft_input_from_target_hidden_dim(model_runner)
+        dtype = model_runner.model_config.dtype
+        vocab_size = self._vocab_size()
+
+        seq_lens_cpu = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32)
+
+        with torch.device(self.device):
+            input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
+            out_cache_loc = torch.ones((max_num_token,), dtype=torch.int64)
+            positions = torch.zeros((max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
+            hidden_states = torch.zeros((max_num_token, hidden_size), dtype=dtype)
+
+            seq_lens = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32)
+            req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
+            num_correct_drafts = torch.full((max_bs,), 1, dtype=torch.int32)
+            num_accept_tokens = torch.full((max_bs,), 1, dtype=torch.int32)
+
+            # Fixed window: every request extends by exactly num_tokens_per_bs
+            # tokens, and start locs are a constant arange.
+            extend_seq_lens = torch.full(
+                (max_bs,), num_tokens_per_bs, dtype=torch.int32
+            )
+            extend_start_loc = torch.arange(
+                0, max_num_token, step=num_tokens_per_bs, dtype=torch.int32
+            )
+            select_index = torch.zeros((max_bs,), dtype=torch.int64)
+
+            next_token_logits_buffer = torch.zeros(
+                (max_num_token, vocab_size), dtype=torch.float
+            )
+
+            if self.require_gathered_buffer:
+                if self.require_mlp_tp_gather:
+                    dp_size = runner.dp_size
+                    global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (dp_size,), dtype=torch.int32
+                    )
+                else:
+                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (1,), dtype=torch.int32
+                    )
+            else:
+                global_num_tokens_gpu = None
+                global_num_tokens_for_logprob_gpu = None
+
+        return MultiLayerEagleDraftExtendInputBuffers(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            positions=positions,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            req_pool_indices=req_pool_indices,
+            num_correct_drafts=num_correct_drafts,
+            num_accept_tokens=num_accept_tokens,
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+            select_index=select_index,
+            mrope_positions=mrope_positions,
+            hidden_states=hidden_states,
+            next_token_logits_buffer=next_token_logits_buffer,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+        )
+
+    def _prepare_extra(self, forward_batch: ForwardBatch) -> None:
+        """Hook for subclasses to populate extra per-call buffers (e.g. sconv)."""
+
+    def prepare(self, forward_batch: ForwardBatch):
+        """Populate the shared buffers once from ``forward_batch`` and bucketize
+        the batch size. Subsequent ``replay(step)`` calls reuse this state."""
+        buffers = self.buffers
+        raw_bs = forward_batch.batch_size
+        num_tokens = raw_bs * self.num_tokens_per_bs
+
+        # Bucketize to a captured batch size (padding the tail).
+        if self.require_mlp_tp_gather:
+            max_batch_size = max(forward_batch.original_global_num_tokens_cpu)
+            bs = self.get_runner(0)._pad_to_bucket(int(max_batch_size), self.capture_bs)
+        else:
+            bs = self.get_runner(0)._pad_to_bucket(raw_bs, self.capture_bs)
+
+        # Reset padded slots, then copy the real values in.
+        buffers.input_ids.zero_()
+        buffers.out_cache_loc.zero_()
+        buffers.positions.zero_()
+        buffers.seq_lens.fill_(self.seq_len_fill_value)
+
+        buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids)
+        buffers.positions[:num_tokens].copy_(forward_batch.positions)
+        buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
+        buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
+        buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
+
+        if (
+            forward_batch.spec_info.hidden_states.shape[1]
+            == buffers.hidden_states.shape[1]
+        ):
+            buffers.hidden_states[:num_tokens].copy_(
+                forward_batch.spec_info.hidden_states
+            )
+
+        buffers.num_correct_drafts[:raw_bs].copy_(
+            forward_batch.spec_info.num_correct_drafts
+        )
+        buffers.num_accept_tokens[:raw_bs].copy_(
+            forward_batch.spec_info.num_accept_tokens
+        )
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+            buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+        # select_index[i] = i * window + num_correct_drafts[i]: the flat index
+        # of request i's last accepted token. Used by the in-graph top-k gather
+        # and by the worker's rotation.
+        arange = torch.arange(bs, device=self.device, dtype=torch.int64)
+        buffers.select_index[:bs].copy_(
+            arange * self.num_tokens_per_bs + buffers.num_correct_drafts[:bs]
+        )
+
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
+
+        # Reusable spec_info for per-step attention metadata.
+        padded_num_tokens = bs * self.num_tokens_per_bs
+        spec_info = EagleDraftExtendInput(
+            hidden_states=buffers.hidden_states[:padded_num_tokens],
+            num_correct_drafts=buffers.num_correct_drafts[:bs],
+            num_accept_tokens=buffers.num_accept_tokens[:bs],
+        )
+        spec_info.num_tokens_per_req = self.num_tokens_per_bs
+        spec_info.num_tokens_for_logprob_per_req = 1
+        spec_info.positions = buffers.positions[:padded_num_tokens]
+        spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs]
+        self._replay_spec_info = spec_info
+
+        self.raw_bs = raw_bs
+        self.bs = bs
+        self.raw_num_tokens = num_tokens
+        self.seq_lens_sum = (
+            forward_batch.seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
+        )
+
+        self._prepare_extra(forward_batch)
+
+    def replay(self, step: int):
+        """Replay ``step``'s graph at the prepared bucket. Returns
+        ``(LogitsProcessorOutput, topk_p, topk_index)`` sliced to the real
+        batch size."""
+        runner = self.runners[step]
+        runner.raw_bs = self.raw_bs
+        out = runner.replay(self.bs, self.seq_lens_sum, self._replay_spec_info)
+        raw_bs = self.raw_bs
+        raw_num_tokens = self.raw_num_tokens
+        logits_output = LogitsProcessorOutput(
+            next_token_logits=out.next_token_logits[:raw_num_tokens],
+            hidden_states=(
+                out.hidden_states[:raw_num_tokens]
+                if out.hidden_states is not None
+                else None
+            ),
+        )
+        return (
+            logits_output,
+            out.topk_p[:raw_bs],
+            out.topk_index[:raw_bs],
+        )
+
+    def get_runner(self, step):
+        return self.runners[step]
+
+    def get_last_runner(self):
+        return self.runners[-1] if self.runners else None
+
+    def can_run_graph(self, forward_batch):
+        return self.runners[0].can_run_graph(forward_batch)

python/sglang/srt/speculative/multi_layer_eagle_utils.py

@@ -0,0 +1,23 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from sglang.kernels.ops.speculative.multi_layer_eagle import (
+    rotate_input_ids,
+    rotate_input_ids_kernel,
+)
+
+__all__ = [
+    "rotate_input_ids",
+    "rotate_input_ids_kernel",
+]

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -0,0 +1,930 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import contextlib
+import logging
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import torch
+
+from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import (
+    MultiLayerEagleMultiStepDraftExtendNpuGraphRunner,
+)
+from sglang.srt.layers.moe.utils import speculative_moe_backend_context
+from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
+from sglang.srt.managers.io_struct import (
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightsFromIPCReqInput,
+)
+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.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+)
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
+from sglang.srt.speculative.draft_utils import DraftBackendFactory
+from sglang.srt.speculative.eagle_info import (
+    EagleDraftExtendInput,
+    EagleDraftInput,
+    EagleVerifyInput,
+)
+from sglang.srt.speculative.eagle_utils import (
+    build_tree_kernel_efficient,
+    default_tree_mask_mode,
+    eagle_prepare_for_verify,
+    eagle_sample,
+    get_draft_recurrent_hidden_state_spec,
+)
+from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
+    MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
+)
+from sglang.srt.speculative.multi_layer_eagle_utils import rotate_input_ids
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import (
+    draft_tp_context,
+    record_stream_each,
+    record_stream_for_v2_verify,
+    sample_draft_proposal,
+    select_top_k_tokens,
+)
+from sglang.srt.utils import is_cpu, is_npu
+from sglang.srt.utils.async_probe import (
+    maybe_detect_inf,
+    maybe_detect_nan,
+    maybe_detect_oob,
+)
+from sglang.srt.utils.common import empty_context, fast_topk
+
+_is_npu = is_npu()
+_is_cpu = is_cpu()
+
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput
+
+
+logger = logging.getLogger(__name__)
+
+
+def _get_plan_stream(
+    device: str,
+) -> Tuple[any, contextlib.AbstractContextManager]:
+    if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
+        plan_stream = torch.get_device_module(device).Stream()
+        plan_stream_ctx = torch.get_device_module(device).stream(plan_stream)
+        return plan_stream, plan_stream_ctx
+    else:
+        return None, contextlib.nullcontext()
+
+
+class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: int,
+        moe_ep_rank: int,
+        attn_cp_rank: int,
+        moe_dp_rank: int,
+        nccl_port: int,
+        target_worker: TpModelWorker,
+    ):
+        # copy args
+        self.server_args = server_args
+        self.gpu_id = gpu_id
+        self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.moe_ep_rank = moe_ep_rank
+        self.nccl_port = nccl_port
+        self.target_worker = target_worker
+        self.draft_extend_attn_backend_list = []
+        self.model_config = target_worker.model_config
+
+        # Args for easy access
+        self.device = server_args.device
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.use_rejection_sampling = server_args.speculative_use_rejection_sampling
+        assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, (
+            "multi-layer EAGLE requires speculative_num_draft_tokens == "
+            "speculative_num_steps + 1, "
+            f"got {self.speculative_num_draft_tokens} and {self.speculative_num_steps}"
+        )
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        # Set constant
+        EagleDraftInput.ALLOC_LEN_PER_DECODE = max(
+            self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
+        )
+
+        # Load draft model weights only.
+        with empty_context(), speculative_moe_backend_context():
+            self.draft_worker = TpModelWorker(
+                server_args=server_args,
+                gpu_id=gpu_id,
+                tp_rank=tp_rank,
+                pp_rank=0,  # spec workers don't support pipeline parallelism
+                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,
+                is_draft_worker=True,
+                is_multi_layer_eagle=True,
+            )
+
+        # Alias for better readability
+        self.draft_runner_list: List[ModelRunner] = self.draft_worker.model_runner_list
+        # Match `EagleDraftWorker.draft_runner` for generic draft-runner access.
+        self.draft_runner: ModelRunner = self.draft_runner_list[0]
+
+        # Chain-style MTP: each step propagates its own output hidden states to the
+        # next step.  Non-chain: each step uses the target model's hidden states.
+        draft_arch = self.draft_worker.model_config.hf_config.architectures[0]
+        self.chain_mtp_hidden_states = draft_arch in ["Step3p5MTP"]
+        self.draft_tp_context = (
+            draft_tp_context if server_args.enable_dp_attention else empty_context
+        )
+        self.tree_mask_mode = default_tree_mask_mode()
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        """Allocate draft KV cache pools (called by scheduler)."""
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.draft_worker.alloc_memory_pool(
+            memory_pool_config=memory_pool_config,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+        self.init_lm_head()
+
+    def init_attention_backends(self):
+        with (
+            self.draft_tp_context(self.draft_runner_list[0].tp_group),
+            speculative_moe_backend_context(),
+        ):
+            super().init_attention_backends()
+
+    def init_cuda_graphs(self):
+        with (
+            self.draft_tp_context(self.draft_runner_list[0].tp_group),
+            speculative_moe_backend_context(),
+        ):
+            super().init_cuda_graphs()
+
+    def mtp_model_runner(self, step: int):
+        return self.draft_runner_list[step]
+
+    def init_lm_head(self):
+        embed, head = self.target_worker.model_runner.model.get_embed_and_head()
+        # Share the embedding and lm_head
+        for i in range(self.speculative_num_steps):
+            self.draft_runner_list[i].model.set_embed_and_head(embed, head)
+
+    def init_attention_backend(self):
+        # Create attn backends
+        self.draft_extend_attn_backend_list = []
+        for step in range(self.speculative_num_steps):
+            draft_backend_factory = DraftBackendFactory(
+                self.server_args,
+                self.draft_runner_list[step],
+                self.topk,
+                self.speculative_num_steps,
+            )
+            self.draft_extend_attn_backend_list.append(
+                draft_backend_factory.create_draft_extend_backend()
+            )
+            if self.draft_extend_attn_backend_list[-1] is not None:
+                self.draft_runner_list[step].attn_backend = (
+                    self.draft_extend_attn_backend_list[-1]
+                )
+
+    def _capture_cuda_graphs(self):
+        self.cuda_graph_runner = None
+        self.cuda_graph_runner_for_draft_extend = None
+
+        if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
+            return
+
+        if envs.SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH.get():
+            return
+
+        if not _is_npu:
+            self.cuda_graph_runner_for_draft_extend = (
+                MultiLayerEagleMultiStepDraftExtendCudaGraphRunner(self)
+            )
+        else:
+            self.cuda_graph_runner_for_draft_extend = (
+                MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(self)
+            )
+
+    def draft(self, batch: ScheduleBatch):
+        draft_input: EagleDraftInput = batch.spec_info
+        forward_batch, can_cuda_graph = self.prepare_for_draft(
+            draft_input,
+            self.req_to_token_pool,
+            batch,
+            self.cuda_graph_runner,
+            self.draft_runner_list[0],
+            self.topk,
+            self.speculative_num_steps,
+        )
+
+        # Run draft
+        parent_list, top_scores_index, draft_tokens = self.draft_forward(forward_batch)
+
+        if batch.forward_mode.is_idle():
+            return EagleVerifyInput.create_idle_input(
+                self.topk,
+                self.speculative_num_steps,
+                self.speculative_num_draft_tokens,
+                self.device,
+            )
+
+        # Build tree mask
+        # Directly write to cuda graph buffers for verify attn
+        tree_mask_buf, position_buf = (
+            self.target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft()
+        )
+        (
+            tree_mask,
+            position,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            draft_tokens,
+        ) = build_tree_kernel_efficient(
+            draft_input.bonus_tokens,
+            parent_list,
+            top_scores_index,
+            draft_tokens,
+            batch.seq_lens,
+            batch.seq_lens_sum,
+            self.topk,
+            self.speculative_num_steps,
+            self.speculative_num_draft_tokens,
+            self.tree_mask_mode,
+            tree_mask_buf,
+            position_buf,
+        )
+
+        return EagleVerifyInput(
+            draft_token=draft_tokens,
+            custom_mask=tree_mask,
+            positions=position,
+            retrieve_index=retrieve_index,
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            retrieve_cum_len=None,
+            spec_steps=self.speculative_num_steps,
+            topk=self.topk,
+            draft_token_num=self.speculative_num_draft_tokens,
+            capture_hidden_mode=None,
+            seq_lens_sum=None,
+            seq_lens_cpu=None,
+            draft_probs=draft_input.draft_probs,
+        )
+
+    def draft_forward(self, forward_batch: ForwardBatch):
+        # Parse args
+        spec_info: EagleDraftInput = forward_batch.spec_info
+        topk_p, topk_index, hidden_states = (
+            spec_info.topk_p,
+            spec_info.topk_index,
+            spec_info.hidden_states,
+        )
+
+        maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info")
+
+        # Return values
+        score_list: List[torch.Tensor] = []
+        token_list: List[torch.Tensor] = []
+        parents_list: List[torch.Tensor] = []
+
+        # Forward multiple steps
+        scores = None
+        _, hidden_states, scores, tree_info = select_top_k_tokens(
+            0, topk_p, topk_index, hidden_states, scores, self.topk
+        )
+        if self.speculative_num_steps == 1:
+            score_list.append(tree_info[0])
+            token_list.append(tree_info[1])
+            parents_list.append(tree_info[2])
+        else:
+            for i in range(self.speculative_num_steps):
+                score_list.append(tree_info[0][:, :, i].unsqueeze(-1))
+                token_index = tree_info[1][:, i].unsqueeze(-1)
+                token_list.append(token_index)
+                if i == 0:
+                    parents_list.append(tree_info[2])
+                else:
+                    parents_list.append(
+                        torch.full(
+                            (tree_info[2].size(0), 1),
+                            i,
+                            dtype=torch.long,
+                            device=tree_info[2].device,
+                        )
+                    )
+
+        # Organize the results
+        score_list = torch.cat(score_list, dim=1).flatten(
+            1
+        )  # b, n, topk; n= 1 + (num_steps-1) * self.topk
+        ss_token_list = torch.cat(
+            token_list, dim=1
+        )  # b, (self.topk + (num_steps-1) * self.topk)
+        top_scores = torch.topk(
+            score_list, self.speculative_num_draft_tokens - 1, dim=-1
+        )
+        top_scores_index = top_scores.indices
+        top_scores_index = torch.sort(top_scores_index).values
+        maybe_detect_oob(
+            top_scores_index,
+            0,
+            ss_token_list.shape[1],
+            "draft_forward: top_scores_index OOB for gather on ss_token_list",
+        )
+        draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
+
+        if len(parents_list) > 1:
+            parent_list = torch.cat(parents_list[:-1], dim=1)
+        else:
+            batch_size = parents_list[0].shape[0]
+            parent_list = torch.empty(batch_size, 0, device=parents_list[0].device)
+
+        return parent_list, top_scores_index, draft_tokens
+
+    def draft_extend(self):
+        pass
+
+    def _draft_extend_for_prefill(
+        self,
+        batch: ScheduleBatch,
+        target_hidden_states: torch.Tensor,
+        next_token_ids: torch.Tensor,
+    ):
+        """
+        Run draft model extend to correctly fill the KV cache.
+
+        Args:
+            batch: The batch to run.
+            target_hidden_states: Hidden states from the target model forward
+            next_token_ids: Next token ids generated from the target forward.
+        """
+        # The draft embed clamps unconditionally (to tolerate multimodal pad
+        # sentinels), so probe next_token_ids here first -- otherwise a corrupted id
+        # would be clamped away instead of surfacing.
+        maybe_detect_oob(
+            next_token_ids,
+            0,
+            self.model_config.vocab_size,
+            "draft_extend_for_prefill: next_token_ids before draft embed",
+        )
+
+        # Draft-extend spec_info for the extend forward; carries only
+        # hidden_states + shape info.
+        extend_input = EagleDraftExtendInput(
+            hidden_states=target_hidden_states,
+            # draft mode is same with decode mode, only 1 token per req
+            num_tokens_per_req=1,
+            num_tokens_for_logprob_per_req=1,
+        )
+        batch.spec_info = extend_input
+
+        # Chain-style MTP needs FULL to get all-token hidden states;
+        # non-chain only needs LAST (the target model's hidden states).
+        # STANDALONE skips hidden states end-to-end.
+        if self.speculative_algorithm.is_standalone():
+            draft_capture_hidden_mode = CaptureHiddenMode.NULL
+        elif self.chain_mtp_hidden_states:
+            draft_capture_hidden_mode = CaptureHiddenMode.FULL
+        else:
+            draft_capture_hidden_mode = CaptureHiddenMode.LAST
+
+        # Run forward
+        batch.capture_hidden_mode = draft_capture_hidden_mode
+        batch.return_hidden_states_before_norm = True
+        forward_batch = ForwardBatch.init_new(batch, self.draft_runner_list[0])
+
+        # Construct input_ids
+        # TODO: same chunked-prefill chain divergence as PR #26329.
+        if not batch.forward_mode.is_idle():
+            rotate_input_ids(
+                forward_batch.input_ids,
+                forward_batch.extend_start_loc,
+                forward_batch.extend_seq_lens,
+                next_token_ids,
+            )
+
+        topk_p_list = []
+        topk_index_list = []
+        draft_probs_list = []
+        for step in range(self.speculative_num_steps):
+            output: ModelRunnerOutput = self.draft_runner_list[step].forward(
+                forward_batch
+            )
+            maybe_detect_nan(
+                output.logits_output.next_token_logits,
+                f"draft_extend_for_prefill step {step}",
+            )
+            maybe_detect_inf(
+                output.logits_output.next_token_logits,
+                f"draft_extend_for_prefill step {step}",
+            )
+            if self.use_rejection_sampling and self.topk == 1:
+                # Sample X ~ q and stash q for the first verify's Leviathan step.
+                probs, topk_p, topk_index = sample_draft_proposal(
+                    output.logits_output.next_token_logits,
+                    forward_batch.sampling_info.temperatures,
+                )
+                draft_probs_list.append(probs)
+            else:
+                probs = torch.softmax(output.logits_output.next_token_logits, dim=-1)
+                topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
+            topk_p_list.append(topk_p)
+            topk_index_list.append(topk_index)
+            # Chain-style: use this step's output hidden_states as next step's input
+            if (
+                self.chain_mtp_hidden_states
+                and step < self.speculative_num_steps - 1
+                and output.logits_output.hidden_states is not None
+            ):
+                forward_batch.spec_info.hidden_states = (
+                    output.logits_output.hidden_states
+                )
+            if forward_batch.extend_seq_lens is not None:
+                rotate_input_ids(
+                    forward_batch.input_ids,
+                    forward_batch.extend_start_loc,
+                    forward_batch.extend_seq_lens,
+                    topk_index,
+                )
+
+        next_draft_input = EagleDraftInput(
+            topk_p=torch.cat(topk_p_list, dim=1),
+            topk_index=torch.cat(topk_index_list, dim=1),
+            # Chain-style left the last step's hidden_states on the extend
+            # input; non-chain keeps the target hidden states.
+            hidden_states=extend_input.hidden_states,
+            bonus_tokens=next_token_ids,
+            num_tokens_per_req=1,
+            num_tokens_for_logprob_per_req=1,
+        )
+        # q [bs, num_steps, vocab] for the first verify's Leviathan step (RS only).
+        next_draft_input.draft_probs = (
+            torch.stack(draft_probs_list, dim=1)
+            if self.use_rejection_sampling and draft_probs_list
+            else None
+        )
+
+        return next_draft_input
+
+    def _draft_extend_for_decode(
+        self, batch: ScheduleBatch, batch_result: GenerationBatchResult
+    ):
+        # Batch 2: Draft extend
+        draft_extend_input = EagleDraftExtendInput(
+            hidden_states=batch_result.logits_output.hidden_states,
+            num_tokens_per_req=self.speculative_num_steps + 1,
+            num_tokens_for_logprob_per_req=1,
+        )
+
+        # Prepare for draft extend in a separate stream
+        # Notice that here we use batch_result.next_token_ids as the input ids
+        with self.plan_stream_ctx:
+            forward_batch = self.prepare_for_draft_extend(
+                draft_extend_input,
+                batch,
+                batch_result.next_token_ids,
+                self.speculative_num_draft_tokens,
+                self.draft_runner_list[0],
+                self.cuda_graph_runner_for_draft_extend,
+            )
+            forward_batch.return_hidden_states_before_norm = True
+
+        if self.plan_stream:
+            torch.get_device_module(self.device).current_stream().wait_stream(
+                self.plan_stream
+            )
+        # `batch_result.accept_lens` includes the bonus token, so drafts-only
+        # is accept_lens - 1. Stash on spec_info for the cuda-graph prepare().
+        forward_batch.spec_info.num_correct_drafts = batch_result.accept_lens - 1
+        forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens
+
+        # Run draft extend batch in the main compute stream
+        can_cuda_graph = (
+            self.cuda_graph_runner_for_draft_extend
+            and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch)
+        )
+        ret_topk_p_list = []
+        ret_topk_index_list = []
+        ret_draft_probs_list = []
+        next_token_ids_backup = batch_result.next_token_ids.clone()
+
+        if can_cuda_graph:
+            cgr = self.cuda_graph_runner_for_draft_extend
+            # Populate the single shared buffer set once; each step replays
+            # against it and the chain is advanced in place between steps.
+            cgr.prepare(forward_batch)
+            for step in range(self.speculative_num_steps):
+                _out, ret_topk_p, ret_topk_index = cgr.replay(step)
+                if self.use_rejection_sampling and self.topk == 1:
+                    # Re-pick X ~ q worker-side so the chain rotation carries it
+                    # to step N+1 (per-step graph does not sample in-graph).
+                    sel = cgr.buffers.select_index[: cgr.raw_bs]
+                    probs, ret_topk_p, ret_topk_index = sample_draft_proposal(
+                        _out.next_token_logits[sel],
+                        forward_batch.sampling_info.temperatures,
+                    )
+                    ret_draft_probs_list.append(probs)
+                ret_topk_p_list.append(ret_topk_p.clone())
+                ret_topk_index_list.append(ret_topk_index.clone())
+                # Advance the draft chain by rotating the shared input_ids window
+                # in place; step N+1's graph then reads the rotated values.
+                if step < self.speculative_num_steps - 1:
+                    rotate_input_ids(
+                        cgr.buffers.input_ids[: cgr.raw_num_tokens],
+                        cgr.buffers.extend_start_loc[: cgr.raw_bs],
+                        cgr.buffers.extend_seq_lens[: cgr.raw_bs],
+                        ret_topk_index,
+                        cgr.buffers.select_index[: cgr.raw_bs],
+                    )
+        else:
+            logger.warning_once(
+                "can't use cuda graph for draft extend! may have correctness issue!"
+            )
+            select_index = (
+                torch.arange(len(batch.seq_lens), device=self.device)
+                * self.speculative_num_draft_tokens
+                + batch_result.accept_lens
+                - 1
+            )
+            # NOTE: this non-graph path runs the per-step forwards without any
+            # pre-plan (see warning above). Mark the batch so the forward path
+            # keeps skipping metadata init — preserves the pre-existing
+            # behavior; the latent issue is tracked by the warning.
+            # On NPU with --disable-cuda-graph, leave each draft runner to init
+            # its own metadata in forward_extend (post-pad), otherwise
+            # per-runner attn_backend.forward_metadata is never initialized for
+            # draft_runner_list[1+].
+            if not _is_npu:
+                forward_batch.mark_forward_metadata_ready()
+
+            for step in range(self.speculative_num_steps):
+                draft_logits_output = self.draft_runner_list[step].forward(
+                    forward_batch
+                )
+                logits_sel = draft_logits_output.logits_output.next_token_logits[
+                    select_index
+                ]
+                if self.use_rejection_sampling and self.topk == 1:
+                    probs, ret_topk_p, ret_topk_index = sample_draft_proposal(
+                        logits_sel, forward_batch.sampling_info.temperatures
+                    )
+                    ret_draft_probs_list.append(probs)
+                else:
+                    probs = torch.softmax(logits_sel, dim=-1)
+                    ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1)
+                # Chain-style: use this step's output hidden_states as next step's input
+                if (
+                    self.chain_mtp_hidden_states
+                    and step < self.speculative_num_steps - 1
+                    and draft_logits_output.logits_output.hidden_states is not None
+                ):
+                    forward_batch.spec_info.hidden_states = (
+                        draft_logits_output.logits_output.hidden_states
+                    )
+                if forward_batch.extend_seq_lens is not None:
+                    rotate_input_ids(
+                        forward_batch.input_ids,
+                        forward_batch.extend_start_loc,
+                        forward_batch.extend_seq_lens,
+                        ret_topk_index,
+                        select_index,
+                    )
+                ret_topk_p_list.append(ret_topk_p)
+                ret_topk_index_list.append(ret_topk_index)
+
+        batch_result.next_token_ids = next_token_ids_backup
+        # Construct the return values
+        next_draft_input = batch_result.next_draft_input
+        (
+            next_draft_input.topk_p,
+            next_draft_input.topk_index,
+            next_draft_input.hidden_states,
+        ) = (
+            torch.cat(ret_topk_p_list, dim=1).clone(),
+            torch.cat(ret_topk_index_list, dim=1).clone(),
+            None,
+        )
+        # q [bs, num_steps, vocab] carries the per-chain-step draft distributions
+        # to the next verify's Leviathan step (accept iff coin*q < p). None
+        # otherwise (default target-only tree sampling).
+        next_draft_input.draft_probs = (
+            torch.stack(ret_draft_probs_list, dim=1) if ret_draft_probs_list else None
+        )
+
+
+class MultiLayerEagleWorkerV2(BaseSpecWorker):
+    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,
+    ):
+        # Parse arguments
+        self.server_args = server_args
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self._target_worker = target_worker
+        self.page_size = server_args.page_size
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        # Override the context length of the draft model to be the same as the target model.
+        server_args.override(
+            "spec_worker.match_target_context_length",
+            context_length=target_worker.model_runner.model_config.context_len,
+        )
+
+        self._draft_worker = MultiLayerEagleDraftWorker(
+            server_args,
+            gpu_id,
+            tp_rank,
+            dp_rank,
+            moe_ep_rank,
+            attn_cp_rank,
+            moe_dp_rank,
+            nccl_port,
+            target_worker,
+        )
+
+        # Some dummy tensors
+        self.num_new_pages_per_topk = torch.empty(
+            (), dtype=torch.int64, device=self.device
+        )
+        self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
+
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        self._draft_worker.alloc_memory_pool(
+            memory_pool_config, req_to_token_pool, token_to_kv_pool_allocator
+        )
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+
+    def init_attention_backends(self):
+        self._draft_worker.init_attention_backends()
+
+    def init_cuda_graphs(self):
+        self._draft_worker.init_cuda_graphs()
+
+    @property
+    def target_worker(self):
+        return self._target_worker
+
+    @property
+    def draft_worker(self):
+        return self._draft_worker
+
+    @property
+    def spec_v2_attn_backends(self) -> tuple:
+        return (
+            self._target_worker.model_runner.attn_backend,
+            *(
+                backend or runner.attn_backend
+                for backend, runner in zip(
+                    self._draft_worker.draft_extend_attn_backend_list,
+                    self._draft_worker.draft_runner_list,
+                )
+            ),
+        )
+
+    def clear_cache_pool(self):
+        # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
+        pass
+
+    def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None):
+        if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
+            # Target prefill
+            target_capture_mode = (
+                CaptureHiddenMode.NULL
+                if self.speculative_algorithm.is_standalone()
+                else CaptureHiddenMode.FULL
+            )
+            batch.capture_hidden_mode = target_capture_mode
+            batch_output = self.target_worker.forward_batch_generation(batch)
+
+            # Spec_v2 convention: batch.seq_lens = length BEFORE this iter's tokens.
+            # Extend processed L prompt tokens; next verify iter expects same L.
+            batch_output.new_seq_lens = batch.seq_lens
+            # Publish before draft_extend so the fence is at target-end.
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+
+            # Chain-style MTP needs FULL to get all-token hidden states;
+            # non-chain only needs LAST (the target model's hidden states).
+            batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill(
+                batch,
+                batch_output.logits_output.hidden_states,
+                batch_output.next_token_ids,
+            )
+            return batch_output
+        else:
+            if batch.spec_info is None:
+                capture_mode = (
+                    CaptureHiddenMode.NULL
+                    if self.speculative_algorithm.is_standalone()
+                    else CaptureHiddenMode.LAST
+                )
+                hidden_size, hidden_dtype = get_draft_recurrent_hidden_state_spec(
+                    self.draft_worker.draft_runner
+                )
+                batch.spec_info = EagleDraftInput.create_idle_input(
+                    device=self.device,
+                    hidden_size=hidden_size,
+                    dtype=hidden_dtype,
+                    topk=self.topk * self.speculative_num_steps,
+                    capture_hidden_mode=capture_mode,
+                )
+            verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
+            assert verify_input.is_verify_input()
+            batch.spec_info = verify_input
+            batch_output = self.verify(batch)
+            # Publish before draft_extend so the fence is at verify-end.
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+            self.draft_worker._draft_extend_for_decode(batch, batch_output)
+            return batch_output
+
+    def verify(
+        self,
+        batch: ScheduleBatch,
+    ):
+        fwd_stream = torch.get_device_module(self.device).current_stream()
+        verify_input: EagleVerifyInput = batch.spec_info
+        record_stream_for_v2_verify(batch, verify_input, fwd_stream)
+
+        bs = len(batch.seq_lens)
+
+        # Batch 1: Target verify
+        # Prepare for target verify in a separate stream
+        with self.plan_stream_ctx:
+            verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify(
+                verify_input,
+                self.req_to_token_pool,
+                batch,
+                self.target_worker,
+            )
+
+        # Cover post-prepare rebinds: draft_token, plan_stream-allocated out_cache_loc.
+        record_stream_each((batch.input_ids, batch.out_cache_loc), fwd_stream)
+
+        # Correct some buffers due to the overlap plan
+        if self.plan_stream:
+            torch.get_device_module(self.device).current_stream().wait_stream(
+                self.plan_stream
+            )
+
+            # Some values such as custom_mask and position depend on the output of draft,
+            # so the previous plan step used the wrong values. Here, we need to run the related
+            # computation again to update them to the correct values.
+            self.target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft(
+                verify_input,
+                (
+                    self.target_worker.model_runner.decode_cuda_graph_runner.bs
+                    if can_run_cuda_graph
+                    else None
+                ),
+            )
+        # NOTE: metadata init is skipped here unconditionally, although
+        # eagle_prepare_for_verify only plans when cuda-graph load_batch ran.
+        # eagle_worker_v2 re-inits the non-graph path instead (post-pad); this
+        # worker has not adopted that fix, so preserve its behavior verbatim.
+        # On NPU with --disable-cuda-graph, non-graph verify needs metadata init
+        # in forward_extend (post-pad); only mark ready for the cuda-graph path.
+        if not _is_npu or can_run_cuda_graph:
+            verify_forward_batch.mark_forward_metadata_ready()
+        # Run target verify batch in the main compute stream
+        forward_batch_output = self.target_worker.forward_batch_generation(
+            batch=None,
+            forward_batch=verify_forward_batch,
+            is_verify=True,
+        )
+        logits_output = forward_batch_output.logits_output
+
+        # Sample
+        maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits")
+        maybe_detect_inf(logits_output.next_token_logits, "verify: target model logits")
+        (
+            predict,
+            accept_lens,
+            accept_index,
+        ) = eagle_sample(verify_input, batch, logits_output)
+        new_seq_lens = batch.seq_lens + accept_lens
+
+        if not batch.forward_mode.is_idle():
+            accept_tokens = predict[accept_index]
+            bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
+            # stride = accept_tokens per-req width = accept_index.shape[1].
+            fill_bonus_tokens_func(
+                accept_tokens,
+                accept_lens,
+                bonus_tokens,
+                accept_index.shape[1],
+                bs,
+            )
+        else:
+            bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)
+
+        if batch.return_logprob and not batch.forward_mode.is_idle():
+            compute_spec_v2_logprobs(
+                batch, logits_output, predict, accept_index, self.speculative_num_steps
+            )
+
+        next_draft_input = EagleDraftInput(bonus_tokens=bonus_tokens)
+        # verify_forward_batch transitively holds verify-time GPU tensors that
+        # must outlive the imminent batch.input_ids rebind; scheduler pins it
+        # in batch_record_buf via extra_keep_alive_refs. See EAGLEWorkerV2.verify.
+        return GenerationBatchResult(
+            logits_output=logits_output,
+            next_token_ids=predict,
+            can_run_cuda_graph=can_run_cuda_graph,
+            speculative_num_draft_tokens=self.speculative_num_draft_tokens,
+            next_draft_input=next_draft_input,
+            accept_lens=accept_lens,
+            new_seq_lens=new_seq_lens,
+            routed_experts_output=forward_batch_output.routed_experts_output,
+            indexer_topk_output=forward_batch_output.indexer_topk_output,
+            extra_keep_alive_refs=[verify_forward_batch],
+        )
+
+    def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
+        for i in range(self.speculative_num_steps):
+            success, message = self._draft_worker.draft_runner_list[
+                i
+            ].update_weights_from_disk(
+                recv_req.model_path,
+                recv_req.load_format,
+                recapture_cuda_graph=recv_req.recapture_cuda_graph,
+            )
+            if not success:
+                return success, message
+        return True, "Succeeded to update model weights."
+
+    def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
+        for i in range(self.speculative_num_steps):
+            success, message = self._draft_worker.draft_runner_list[
+                i
+            ].update_weights_from_ipc(recv_req)
+            if not success:
+                return success, message
+        return True, "Succeeded to update model weights."

python/sglang/srt/speculative/ngram_info.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
+
+
+class NgramVerifyInput(SpecInput):
+    def __init__(
+        self,
+        draft_token: torch.Tensor = None,
+        custom_mask: torch.Tensor = None,
+        positions: torch.Tensor = None,
+        retrieve_index: torch.Tensor = None,
+        retrieve_next_token: torch.Tensor = None,
+        retrieve_next_sibling: torch.Tensor = None,
+        draft_token_num: int = None,
+        grammar: BaseGrammarObject = None,
+        future_indices: Optional[torch.Tensor] = None,
+        new_seq_lens: Optional[torch.Tensor] = None,
+        accept_tokens: Optional[torch.Tensor] = None,
+        accept_lens: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(SpecInputType.NGRAM_VERIFY)
+        self.draft_token = draft_token
+        self.custom_mask = custom_mask
+        self.positions = positions
+        self.retrieve_index = retrieve_index
+        self.retrieve_next_token = retrieve_next_token
+        self.retrieve_next_sibling = retrieve_next_sibling
+        self.draft_token_num = draft_token_num
+        self.grammar = grammar
+
+        # Inputs for V2 overlap worker
+        self.future_indices = future_indices
+        self.new_seq_lens = new_seq_lens
+        self.accept_tokens = accept_tokens
+        self.accept_lens = accept_lens
+
+        self.device = (
+            custom_mask.device if custom_mask is not None else new_seq_lens.device
+        )
+
+    @property
+    def max_tree_depth(self) -> int:
+        # NGRAM trees are node-budgeted with no depth cap: the corpus BFS only
+        # stops on the node budget, so a single long match can chain all
+        # draft_token_num nodes (spec_steps is meaningless for this tree).
+        return self.draft_token_num
+
+    @property
+    def tree_topk(self) -> int:
+        # Irregular tree: per-level branching follows the corpus matches.
+        return -1
+
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        return self.draft_token_num, self.draft_token_num
+
+    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,
+    ):
+        bs = len(req_pool_indices)
+
+        cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
+
+        paged_kernel_lens = paged_kernel_lens + self.draft_token_num
+        cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+
+        self.qo_indptr = (
+            torch.arange(0, bs + 1, dtype=torch.int32, device=self.device)
+            * self.draft_token_num
+        )
+
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum + self.draft_token_num * bs,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        create_flashinfer_kv_indices_triton[(bs,)](
+            req_to_token,
+            req_pool_indices,
+            paged_kernel_lens,
+            cum_kv_seq_len,
+            None,
+            kv_indices,
+            req_to_token.size(1),
+        )
+
+        # Pad custom_mask when CUDA graph pads batch size beyond the actual number of requests.
+        mask_numel = (
+            paged_kernel_lens_sum * self.draft_token_num
+            + (self.draft_token_num**2) * bs
+        )
+        custom_mask = self.custom_mask
+        if custom_mask.numel() < mask_numel:
+            custom_mask = torch.cat(
+                [
+                    custom_mask,
+                    torch.full(
+                        (mask_numel - custom_mask.numel(),),
+                        True,
+                        dtype=torch.bool,
+                        device=self.device,
+                    ),
+                ],
+                dim=0,
+            )
+
+        return kv_indices, cum_kv_seq_len, self.qo_indptr, custom_mask
+
+    def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
+        if self.future_indices is not None:
+            self.future_indices = self.future_indices[new_indices]
+        if self.new_seq_lens is not None:
+            self.new_seq_lens = self.new_seq_lens[new_indices]
+        self.accept_tokens = self.accept_tokens.reshape(-1, self.draft_token_num)[
+            new_indices, :
+        ]
+        self.accept_tokens = self.accept_tokens.flatten()
+        self.accept_lens = self.accept_lens[new_indices]
+
+    def merge_batch(self, spec_info: NgramVerifyInput):
+        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), dim=0
+            )
+        if self.new_seq_lens is not None:
+            assert spec_info.new_seq_lens is not None
+            self.new_seq_lens = torch.cat(
+                (self.new_seq_lens, spec_info.new_seq_lens), dim=0
+            )
+        self.accept_tokens = torch.cat(
+            (self.accept_tokens, spec_info.accept_tokens), dim=0
+        )
+        self.accept_lens = torch.cat((self.accept_lens, spec_info.accept_lens), dim=0)

python/sglang/srt/speculative/ngram_worker.py

@@ -0,0 +1,524 @@
+import logging
+from typing import List, Optional
+
+import numpy as np
+import torch
+from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
+
+from sglang.kernels.ops.speculative.cache_locs import (
+    assign_extend_cache_locs_func as assign_extend_cache_locs_func,
+)
+from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
+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 ForwardMode
+from sglang.srt.observability.req_time_stats import set_time_batch
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
+from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
+from sglang.srt.speculative.eagle_utils import eagle_sample
+from sglang.srt.speculative.ngram_info import NgramVerifyInput
+from sglang.srt.speculative.spec_utils import (
+    commit_mamba_states_after_verify,
+    generate_token_bitmask,
+    move_accept_tokens_to_target_kvcache,
+    prepare_mamba_track_for_verify,
+    record_stream_for_v2_verify,
+)
+from sglang.srt.utils import is_cpu
+from sglang.srt.utils.async_probe import maybe_detect_inf, maybe_detect_nan
+
+_is_cpu = is_cpu()
+
+logger = logging.getLogger(__name__)
+
+
+USE_FULL_MASK = True
+
+
+class NGRAMWorker(BaseSpecWorker):
+    def alloc_memory_pool(self, **kwargs):
+        # The target memory pool does not exist yet when __init__ runs.
+        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+            self._target_worker.get_memory_pool()
+        )
+        self.max_batch_size = self.model_runner.max_running_requests
+        self._init_preallocated_tensors()
+
+    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,
+    ):
+        self.server_args = server_args
+        self.enable_overlap = not server_args.disable_overlap_schedule
+        self._target_worker = target_worker
+        self.model_runner = target_worker.model_runner
+        self.tp_rank = tp_rank
+        self.page_size = server_args.page_size
+        self.draft_token_num: int = server_args.speculative_num_draft_tokens
+        self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        # req_to_token_pool / token_to_kv_pool_allocator are set in
+        # alloc_memory_pool(), after the target pools are allocated.
+        self.device = server_args.device
+
+        self.adaptive_controller = None
+        # rids of the last decode batch; used to erase corpus match state for
+        # requests that left the batch (see forward_batch_generation).
+        self._prev_decode_rids: set = set()
+
+        self.ngram_corpus = NgramCorpus(
+            min_bfs_breadth=server_args.speculative_ngram_min_bfs_breadth,
+            max_bfs_breadth=server_args.speculative_ngram_max_bfs_breadth,
+            match_type=server_args.speculative_ngram_match_type,
+            capacity=server_args.speculative_ngram_capacity,
+            max_trie_depth=server_args.speculative_ngram_max_trie_depth,
+            draft_token_num=server_args.speculative_num_draft_tokens,
+            external_sam_budget=server_args.speculative_ngram_external_sam_budget,
+            external_corpus_max_tokens=server_args.speculative_ngram_external_corpus_max_tokens,
+        )
+        if server_args.speculative_ngram_external_corpus_path is not None:
+            from sglang.srt.speculative.cpp_ngram.external_corpus import (
+                iter_external_corpus_chunks,
+            )
+
+            corpus_path = server_args.speculative_ngram_external_corpus_path
+            chunks = list(
+                iter_external_corpus_chunks(
+                    corpus_path,
+                    target_worker.tokenizer,
+                    server_args.speculative_ngram_external_corpus_max_tokens,
+                )
+            )
+            loaded = self.add_external_corpus(corpus_path, chunks)
+            self.commit_corpus_load(corpus_path, loaded)
+            logger.info(
+                "Loaded external ngram corpus '%s' (%d tokens).",
+                corpus_path,
+                loaded,
+            )
+
+    @property
+    def target_worker(self) -> TpModelWorker:
+        return self._target_worker
+
+    @property
+    def draft_worker(self) -> Optional[EagleDraftWorkerBase]:
+        # NGRAM has no draft model; drafts come from the CPU-side corpus.
+        return None
+
+    def clear_cache_pool(self):
+        self.ngram_corpus.reset()
+        self._prev_decode_rids = set()
+
+    def update_weights_from_tensor(self, recv_req):
+        # NGRAM has no draft weights of its own — the n-gram corpus is a CPU
+        # lookup structure built from request token streams — and its
+        # `model_runner` is shared with the target worker. The scheduler
+        # mixin dispatches via `self.draft_worker or self.tp_worker`, so
+        # without this method any caller of `update_weights_from_tensor`
+        # under `--speculative-algorithm NGRAM` raises AttributeError.
+        return self.target_worker.update_weights_from_tensor(recv_req)
+
+    def add_external_corpus(self, corpus_id: str, token_chunks: list[list[int]]) -> int:
+        return self.ngram_corpus.load_external_corpus_named(corpus_id, token_chunks)
+
+    def commit_corpus_load(self, corpus_id: str, loaded_token_count: int) -> None:
+        self.ngram_corpus.commit_external_corpus_load(corpus_id, loaded_token_count)
+
+    def remove_external_corpus(self, corpus_id: str) -> None:
+        self.ngram_corpus.remove_external_corpus(corpus_id)
+
+    def list_external_corpora(self) -> dict[str, int]:
+        return self.ngram_corpus.list_external_corpora()
+
+    def _efficient_concat_last_n(self, seq1: List[int], seq2: List[int], n: int):
+        seq2_len = len(seq2)
+        if seq2_len >= n:
+            return seq2[-n:]
+
+        need_from_seq1 = n - seq2_len
+        return seq1[-need_from_seq1:] + seq2
+
+    def _init_preallocated_tensors(self):
+        max_total_drafts = self.max_batch_size * self.draft_token_num
+        max_total_mask_size = (
+            self.max_batch_size * self.draft_token_num * self.draft_token_num
+        )
+
+        self.draft_tokens = torch.empty(
+            (max_total_drafts,), dtype=torch.int64, device=self.device
+        )
+        self.retrieve_indexes = torch.empty(
+            (self.max_batch_size, self.draft_token_num),
+            dtype=torch.int64,
+            device=self.device,
+        )
+        self.retrieve_next_token = torch.empty(
+            (self.max_batch_size, self.draft_token_num),
+            dtype=torch.int64,
+            device=self.device,
+        )
+        self.retrieve_next_sibling = torch.empty(
+            (self.max_batch_size, self.draft_token_num),
+            dtype=torch.int64,
+            device=self.device,
+        )
+        self.positions = torch.empty(
+            (max_total_drafts,), dtype=torch.int64, device=self.device
+        )
+        self.tree_mask = torch.empty(
+            (max_total_mask_size,), dtype=torch.bool, device=self.device
+        )
+
+        self.draft_tokens_batch = []
+        self.tree_mask_batch = []
+        self.retrieve_indexes_batch = []
+        self.retrieve_next_token_batch = []
+        self.retrieve_next_sibling_batch = []
+        self.positions_batch = []
+
+        for bs in range(0, self.max_batch_size + 1):
+            self.retrieve_indexes_batch.append(self.retrieve_indexes[:bs, :])
+            self.retrieve_next_token_batch.append(self.retrieve_next_token[:bs, :])
+            self.retrieve_next_sibling_batch.append(self.retrieve_next_sibling[:bs, :])
+            self.positions_batch.append(self.positions[: bs * self.draft_token_num])
+            self.draft_tokens_batch.append(
+                self.draft_tokens[: bs * self.draft_token_num]
+            )
+            self.tree_mask_batch.append(
+                self.tree_mask[: bs * self.draft_token_num * self.draft_token_num]
+            )
+
+    def on_verify_complete_cpu(
+        self, num_correct_drafts_per_req: list[int], batch_size: int = 0
+    ) -> None:
+        # Signature must match BaseSpecWorker.on_verify_complete_cpu; the
+        # result processor calls it with batch_size as a keyword argument.
+        if self.adaptive_controller is not None:
+            self.adaptive_controller.on_verify_complete(num_correct_drafts_per_req)
+
+    def _prepare_draft_tokens(
+        self, batch: ScheduleBatch
+    ) -> tuple[np.ndarray, np.ndarray]:
+        bs = len(batch.reqs)
+        stride = self.draft_token_num
+
+        prev_token_ids, prev_accept_lens = (
+            batch.spec_info.accept_tokens,
+            batch.spec_info.accept_lens,
+        )
+        if not prev_token_ids.is_cpu:
+            prev_token_ids = prev_token_ids.cpu()
+            prev_accept_lens = prev_accept_lens.cpu()
+        # Worker-level staging: written here at draft prep, consumed by
+        # _update_ngram_corpus after verify within the same forward call.
+        self.prev_token_ids = prev_token_ids.tolist()
+        self.prev_accept_lens = prev_accept_lens.tolist()
+
+        self.ngram_corpus.synchronize()
+        req_ids = []
+        batch_tokens = []
+        total_lens = []
+        assert len(batch.reqs) == len(self.prev_accept_lens)
+        # Overlap mode processes results one iteration behind, so the last
+        # round's accepted tokens are not yet in req.output_ids and must be
+        # spliced in from spec_info. Sync mode and grammar batches process
+        # results before the next draft prep, so output_ids is already
+        # complete and splicing would duplicate the tail.
+        use_prev_tokens = self.enable_overlap and not batch.has_grammar
+        i = 0
+        for req in batch.reqs:
+            prev_tokens = (
+                self.prev_token_ids[i * stride : i * stride + self.prev_accept_lens[i]]
+                if use_prev_tokens
+                else []
+            )
+            check_token = self._efficient_concat_last_n(
+                list(req.origin_input_ids),
+                list(req.output_ids[-self.max_trie_depth :]) + prev_tokens,
+                self.max_trie_depth,
+            )
+            req_ids.append(req.rid)
+            batch_tokens.append(check_token)
+            i += 1
+            total_lens.append(
+                len(req.origin_input_ids) + len(req.output_ids) + len(prev_tokens)
+            )
+        req_drafts, mask = self.ngram_corpus.batch_get(
+            req_ids, batch_tokens, total_lens
+        )
+        total_draft_token_num = len(req_drafts)
+
+        # Check if speculative decoding is needed; here we always enforce it
+        assert (
+            total_draft_token_num == bs * self.draft_token_num
+        ), f"{total_draft_token_num=}, {bs=}, {self.draft_token_num=}"
+        return req_drafts, mask
+
+    def _prepare_for_speculative_decoding(self, batch: ScheduleBatch):
+        # Decode-only: extend goes through the plain target forward, and an
+        # IDLE batch must keep its forward_mode instead of being rewritten to
+        # TARGET_VERIFY below (relevant once DP attention support lands).
+        if not batch.forward_mode.is_decode():
+            return
+
+        bs = len(batch.reqs)
+
+        retrieve_index = self.retrieve_indexes_batch[bs]
+        retrieve_next_token = self.retrieve_next_token_batch[bs]
+        retrieve_next_sibling = self.retrieve_next_sibling_batch[bs]
+        positions = self.positions_batch[bs]
+        tree_mask = self.tree_mask_batch[bs]
+        draft_tokens = self.draft_tokens_batch[bs]
+
+        req_drafts, mask = self._prepare_draft_tokens(batch)
+        tree_mask.copy_(torch.from_numpy(mask), non_blocking=True)
+        draft_tokens.copy_(torch.from_numpy(req_drafts), non_blocking=True)
+
+        # generate positions and some indices using tree_mask
+        reconstruct_indices_from_tree_mask(
+            tree_mask,
+            batch.seq_lens,
+            positions,  # mutable
+            retrieve_index,  # mutable
+            retrieve_next_token,  # mutable
+            retrieve_next_sibling,  # mutable
+            bs,
+            self.draft_token_num,
+        )
+
+        # NOTE: QLEN_MASK is faster than FULL_MASK, but requires corresponding changes in flashinfer.
+        # Testing shows about 8% performance improvement (the effect is roughly proportional to batch size).
+        if USE_FULL_MASK and not _is_cpu:
+            tree_mask = []
+            mask = mask.reshape(bs, self.draft_token_num, self.draft_token_num)
+            # TODO(siyuan): the for loop here leads to significant overhead in large batch size. Can be written into a kernel.
+            for i in range(bs):
+                seq_len = batch.seq_lens_cpu[i]
+                req_mask = torch.ones(
+                    (self.draft_token_num, seq_len), device=self.device
+                )
+                req_mask = torch.cat(
+                    (
+                        req_mask,
+                        torch.from_numpy(mask[i]).to(
+                            device=self.device, non_blocking=True
+                        ),
+                    ),
+                    dim=1,
+                ).to(torch.bool)
+                tree_mask.append(req_mask.flatten())
+            tree_mask = torch.cat(tree_mask, dim=0)
+
+        batch.forward_mode = ForwardMode.TARGET_VERIFY
+        batch.input_ids = draft_tokens
+        batch.out_cache_loc = assign_extend_cache_locs_func(
+            req_pool_indices=batch.req_pool_indices,
+            req_to_token=batch.req_to_token_pool.req_to_token,
+            start_offset=batch.seq_lens,
+            end_offset=batch.seq_lens + self.draft_token_num,
+            batch_size=bs,
+            draft_token_num=self.draft_token_num,
+            device=self.device,
+        )
+
+        prepare_mamba_track_for_verify(batch)
+
+        batch.spec_info = NgramVerifyInput(
+            draft_token=draft_tokens,
+            custom_mask=tree_mask,
+            positions=positions,
+            retrieve_index=retrieve_index,
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            draft_token_num=self.draft_token_num,
+        )
+
+    def _update_ngram_corpus(self, batch: ScheduleBatch):
+        batch_tokens = []
+        i, stride = 0, self.draft_token_num
+        # Same splice condition as _prepare_draft_tokens: only overlap mode
+        # has accepted tokens missing from req.output_ids.
+        use_prev_tokens = self.enable_overlap and not batch.has_grammar
+        for req in batch.reqs:
+            # FIXME: Whether to insert 'extend' into the cache or not, after testing,
+            # there is not much difference, so we will not insert it for now.
+            # if batch.forward_mode.is_extend():
+            #     put_ids = req.origin_input_ids + req.output_ids
+            # else:
+            prev_tokens = (
+                self.prev_token_ids[i * stride : i * stride + self.prev_accept_lens[i]]
+                if use_prev_tokens
+                else []
+            )
+            put_ids = self._efficient_concat_last_n(
+                list(req.origin_input_ids),
+                list(req.output_ids[-self.max_trie_depth :]) + prev_tokens,
+                self.max_trie_depth,
+            )
+            batch_tokens.append(put_ids)
+            i += 1
+        self.ngram_corpus.batch_put(batch_tokens)
+
+    def forward_batch_generation(
+        self, batch: ScheduleBatch, on_publish=None
+    ) -> GenerationBatchResult:
+        fwd_stream = torch.get_device_module(self.device).current_stream()
+        record_stream_for_v2_verify(batch, None, fwd_stream)
+        bs = len(batch.reqs)
+
+        set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
+        self._prepare_for_speculative_decoding(batch)
+        set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
+
+        verify_input: NgramVerifyInput = batch.spec_info
+        accept_lens = torch.ones(bs, dtype=torch.int32, device=self.device)
+
+        if batch.forward_mode.is_target_verify():
+            # Prepare grammar data on CPU if needed
+            if batch.has_grammar:
+                retrieve_next_token_cpu = verify_input.retrieve_next_token.cpu()
+                retrieve_next_sibling_cpu = verify_input.retrieve_next_sibling.cpu()
+                draft_tokens_cpu = verify_input.draft_token.view(
+                    verify_input.retrieve_next_token.shape
+                ).cpu()
+
+            batch_result = self.target_worker.forward_batch_generation(
+                batch, is_verify=True
+            )
+
+            logits_output, can_run_cuda_graph = (
+                batch_result.logits_output,
+                batch_result.can_run_cuda_graph,
+            )
+
+            verify_input: NgramVerifyInput = batch.spec_info
+            vocab_mask = None
+            if batch.has_grammar:
+                # Generate the logit mask for structured output.
+                # Overlap the CPU operations for bitmask generation with the forward pass.
+                vocab_mask = generate_token_bitmask(
+                    batch.reqs,
+                    verify_input,
+                    retrieve_next_token_cpu,
+                    retrieve_next_sibling_cpu,
+                    draft_tokens_cpu,
+                    batch.sampling_info.vocab_size,
+                )
+
+                if vocab_mask is not None:
+                    assert verify_input.grammar is not None
+                    vocab_mask = vocab_mask.to(verify_input.retrieve_next_token.device)
+                    # NOTE (sk): otherwise, this vocab mask will be the one from the previous extend stage
+                    # and will be applied to produce wrong results
+                    batch.sampling_info.vocab_mask = None
+
+            # Sample
+            maybe_detect_nan(
+                logits_output.next_token_logits, "verify: target model logits"
+            )
+            maybe_detect_inf(
+                logits_output.next_token_logits, "verify: target model logits"
+            )
+            (
+                predict,
+                accept_lens,
+                accept_index,
+            ) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
+            new_seq_lens = batch.seq_lens + accept_lens
+            commit_mamba_states_after_verify(
+                self.target_worker,
+                batch,
+                accept_lens,
+                accept_index,
+                self.draft_token_num,
+            )
+            accept_tokens = predict[accept_index].flatten()
+            next_token_ids = accept_tokens
+
+            # The KV mover expects drafts-only counts. NGRAM's
+            # accept_lens includes the bonus token, matching scheduler output.
+            num_correct_drafts_per_req = accept_lens - 1
+            move_accept_tokens_to_target_kvcache(
+                batch,
+                accept_index,
+                num_correct_drafts_per_req,
+                self.token_to_kv_pool_allocator,
+            )
+            if batch.return_logprob:
+                # The last arg is the accept_index row width minus 1. NGRAM's
+                # accept_index is (bs, draft_token_num) -- the tree depth is not
+                # bounded by spec_steps like EAGLE's (bs, spec_steps + 1).
+                compute_spec_v2_logprobs(
+                    batch,
+                    logits_output,
+                    predict,
+                    accept_index,
+                    self.draft_token_num - 1,
+                )
+
+            if on_publish is not None:
+                on_publish(new_seq_lens)
+
+            self._update_ngram_corpus(batch)
+            # Erase match state of requests that left the decode batch.
+            # req.finished() is unusable here: under overlap it flips at result
+            # processing, one iteration after the request left the batch.
+            # The last batch's entries persist while idle (bounded, small).
+            cur_rids = {req.rid for req in batch.reqs}
+            departed_rids = self._prev_decode_rids - cur_rids
+            if departed_rids:
+                self.ngram_corpus.erase_match_state(list(departed_rids))
+            self._prev_decode_rids = cur_rids
+            batch.forward_mode = ForwardMode.DECODE
+
+        else:
+            batch_result = self.target_worker.forward_batch_generation(batch)
+            logits_output, predict, can_run_cuda_graph = (
+                batch_result.logits_output,
+                batch_result.next_token_ids,
+                batch_result.can_run_cuda_graph,
+            )
+            new_seq_lens = batch.seq_lens.clone()
+
+            accept_tokens = torch.zeros(
+                bs, self.draft_token_num, dtype=torch.int32, device=self.device
+            )
+            accept_tokens[:, 0] = predict
+            accept_tokens = accept_tokens.flatten()
+            next_token_ids = predict
+
+            if on_publish is not None:
+                on_publish(new_seq_lens)
+
+        # Construct the next draft input
+        next_draft_input = NgramVerifyInput(
+            draft_token_num=self.draft_token_num,
+            new_seq_lens=new_seq_lens,
+            accept_tokens=accept_tokens,
+            accept_lens=accept_lens,
+        )
+        return GenerationBatchResult(
+            logits_output=logits_output,
+            next_token_ids=next_token_ids,
+            can_run_cuda_graph=can_run_cuda_graph,
+            accept_lens=accept_lens,
+            # Consumed by the non-overlap V2 scheduler branch to advance
+            # batch.seq_lens after the isolation restore; overlap mode relays
+            # it via on_publish instead.
+            new_seq_lens=new_seq_lens,
+            next_draft_input=next_draft_input,
+            speculative_num_draft_tokens=self.speculative_num_draft_tokens,
+        )

python/sglang/srt/speculative/ragged_verify.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import bisect
+from enum import Enum
+from typing import List, Optional, Sequence, Tuple
+
+import msgspec
+import torch
+
+from sglang.srt.environ import envs
+
+
+class RaggedVerifyMode(str, Enum):
+    STATIC = "static"
+    CAP_ACCEPT = "cap-accept"
+    COMPACT = "compact"
+
+
+def read_ragged_verify_mode() -> RaggedVerifyMode:
+    value = envs.SGLANG_RAGGED_VERIFY_MODE.get()
+    for mode in RaggedVerifyMode:
+        if value == mode.value:
+            return mode
+    raise ValueError(
+        f"invalid SGLANG_RAGGED_VERIFY_MODE={value!r}; expected one of "
+        f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}"
+    )
+
+
+def ragged_verify_compact_enabled() -> bool:
+    return read_ragged_verify_mode() == RaggedVerifyMode.COMPACT
+
+
+def round_up_grid(total: int, grid: Sequence[int]) -> int:
+    if not grid:
+        raise ValueError("round_up_grid requires a non-empty grid")
+    if total > grid[-1]:
+        raise ValueError(
+            f"total {total} exceeds max grid tier {grid[-1]}; "
+            "the caller must reject this batch before selecting a graph tier"
+        )
+    index = bisect.bisect_left(grid, total)
+    return grid[index]
+
+
+class RaggedVerifyLayout(msgspec.Struct, frozen=True):
+    verify_lens: torch.Tensor
+    graph_num_tokens: int
+    extend_start_loc: torch.Tensor
+    qo_indptr_device: torch.Tensor
+    verify_lens_cpu: Optional[list[int]] = None
+    total_verify_tokens: Optional[int] = None
+    qo_indptr_host: Optional[torch.Tensor] = None
+    kv_indptr_host: Optional[torch.Tensor] = None
+    kv_lens_host: Optional[torch.Tensor] = None
+    max_q_len: Optional[int] = None
+    max_kv_len: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        if self.verify_lens_cpu is None:
+            return
+        if not self.verify_lens_cpu:
+            raise ValueError("RaggedVerifyLayout requires at least one request")
+        if min(self.verify_lens_cpu) < 1:
+            raise ValueError(
+                f"every request must verify the anchor (verify_len >= 1), got "
+                f"{self.verify_lens_cpu}"
+            )
+        if self.total_verify_tokens != sum(self.verify_lens_cpu):
+            raise ValueError(
+                f"total_verify_tokens {self.total_verify_tokens} != "
+                f"sum(verify_lens_cpu) {sum(self.verify_lens_cpu)}"
+            )
+        if not (self.total_verify_tokens <= self.graph_num_tokens):
+            raise ValueError(
+                f"total_verify_tokens {self.total_verify_tokens} exceeds "
+                f"graph_num_tokens {self.graph_num_tokens}"
+            )
+
+    @property
+    def bs(self) -> int:
+        return int(self.verify_lens.shape[0])
+
+    @classmethod
+    def _assemble_device(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+        verify_lens_cpu: Optional[list[int]] = None,
+        total_verify_tokens: Optional[int] = None,
+    ) -> RaggedVerifyLayout:
+        from sglang.srt.speculative.ragged_verify_kernels import (
+            BuildQoIndptr,
+        )
+
+        verify_lens = verify_lens.to(torch.int32)
+        indptr = BuildQoIndptr.execute(verify_lens=verify_lens)
+        return cls(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            extend_start_loc=indptr.extend_start_loc,
+            qo_indptr_device=indptr.qo_indptr,
+            verify_lens_cpu=verify_lens_cpu,
+            total_verify_tokens=total_verify_tokens,
+        )
+
+    @classmethod
+    def _assemble(
+        cls,
+        *,
+        verify_lens_cpu: list[int],
+        total_verify_tokens: int,
+        graph_num_tokens: int,
+        device: torch.device,
+    ) -> RaggedVerifyLayout:
+        verify_lens = torch.tensor(verify_lens_cpu, dtype=torch.int32, device=device)
+        return cls._assemble_device(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            verify_lens_cpu=verify_lens_cpu,
+            total_verify_tokens=total_verify_tokens,
+        )
+
+    @classmethod
+    def from_verify_lens_device(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+    ) -> RaggedVerifyLayout:
+        return cls._assemble_device(
+            verify_lens=verify_lens, graph_num_tokens=graph_num_tokens
+        )
+
+    @classmethod
+    def from_verify_lens(
+        cls,
+        *,
+        verify_lens_cpu: Sequence[int],
+        device: torch.device,
+        grid: Sequence[int],
+        graph_num_tokens_floor: int = 0,
+    ) -> RaggedVerifyLayout:
+        verify_lens_list = [int(v) for v in verify_lens_cpu]
+        total_verify_tokens = sum(verify_lens_list)
+        bucket_input = max(total_verify_tokens, graph_num_tokens_floor)
+        graph_num_tokens = round_up_grid(total=bucket_input, grid=grid)
+
+        return cls._assemble(
+            verify_lens_cpu=verify_lens_list,
+            total_verify_tokens=total_verify_tokens,
+            graph_num_tokens=graph_num_tokens,
+            device=device,
+        )
+
+    def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout:
+        from sglang.srt.speculative.ragged_verify_kernels import (
+            PaddedToBucket,
+        )
+
+        padded = PaddedToBucket.execute(
+            verify_lens=self.verify_lens,
+            graph_num_tokens=self.graph_num_tokens,
+            bs=self.bs,
+            padded_bs=padded_bs,
+        )
+
+        return RaggedVerifyLayout._assemble_device(
+            verify_lens=padded,
+            graph_num_tokens=self.graph_num_tokens,
+            total_verify_tokens=self.graph_num_tokens,
+        )
+
+
+def build_capture_verify_lens(
+    *,
+    num_tokens: int,
+    num_slots: int,
+    num_draft_tokens: int,
+) -> list[int]:
+    if num_slots < 1 or num_tokens < num_slots:
+        raise ValueError(
+            f"capture layout needs 1 <= num_slots <= num_tokens, got "
+            f"num_slots={num_slots}, num_tokens={num_tokens}"
+        )
+    if num_tokens > num_slots * num_draft_tokens:
+        raise ValueError(
+            f"capture layout cannot pack num_tokens={num_tokens} into "
+            f"{num_slots} rows of at most {num_draft_tokens} tokens"
+        )
+    base = num_tokens // num_slots
+    rem = num_tokens - base * num_slots
+    return [base + 1] * rem + [base] * (num_slots - rem)
+
+
+def resolve_ragged_verify_layout(forward_batch) -> Optional[RaggedVerifyLayout]:
+    """Layout riding the batch's spec input, or None. Tolerates the runner's
+    ad-hoc replay batch views, which may not carry spec_info at all."""
+    spec_info = getattr(forward_batch, "spec_info", None)
+    if spec_info is None:
+        return None
+    return spec_info.ragged_verify_layout
+
+
+class RaggedTargetVerifyGeometry(msgspec.Struct):
+    cache_seqlens_int32: torch.Tensor
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+    max_seq_len_q: Optional[int]
+
+
+def build_ragged_target_verify_geometry(
+    *,
+    seq_lens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+) -> RaggedTargetVerifyGeometry:
+    cache_seqlens_int32 = (seq_lens + layout.verify_lens).to(torch.int32)
+    cu_seqlens_q = layout.qo_indptr_device.to(torch.int32)
+    cu_seqlens_k = torch.nn.functional.pad(
+        torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32), (1, 0)
+    )
+    max_seq_len_q = (
+        max(layout.verify_lens_cpu) if layout.verify_lens_cpu is not None else None
+    )
+    return RaggedTargetVerifyGeometry(
+        cache_seqlens_int32=cache_seqlens_int32,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+    )
+
+
+def compute_target_verify_graph_key(
+    *,
+    bs: int,
+    num_draft_tokens: int,
+    ragged_layout: Optional[RaggedVerifyLayout],
+) -> Tuple[int, int]:
+    num_tokens_full_block = num_draft_tokens * bs
+    if ragged_layout is None:
+        return bs, num_tokens_full_block
+    graph_num_tokens = ragged_layout.graph_num_tokens
+    assert graph_num_tokens <= num_tokens_full_block, (
+        f"ragged verify graph_num_tokens={graph_num_tokens} exceeds full block "
+        f"num_draft*bs={num_tokens_full_block}"
+    )
+    total_verify_tokens = ragged_layout.total_verify_tokens
+    if total_verify_tokens is not None:
+        assert total_verify_tokens <= graph_num_tokens, (
+            f"ragged verify total_verify_tokens={total_verify_tokens} exceeds the "
+            f"round-up bucket graph_num_tokens={graph_num_tokens}"
+        )
+    return graph_num_tokens, graph_num_tokens
+
+
+class VerifyExtendLengths(msgspec.Struct, frozen=True):
+    seq_lens_extended: torch.Tensor
+    seq_lens_cpu_extended: List[int]
+    extend_seq_lens_cpu: List[int]
+    num_tokens: int
+    extend_start_loc: Optional[torch.Tensor]
+
+
+def compute_uniform_extend_lengths(
+    *,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: List[int],
+    extend_len: int,
+) -> VerifyExtendLengths:
+    batch_size = len(seq_lens_cpu)
+    seq_lens_extended = seq_lens + extend_len
+    seq_lens_cpu_extended = [x + extend_len for x in seq_lens_cpu]
+    extend_seq_lens_cpu = [extend_len] * batch_size
+    num_tokens = extend_len * batch_size
+    return VerifyExtendLengths(
+        seq_lens_extended=seq_lens_extended,
+        seq_lens_cpu_extended=seq_lens_cpu_extended,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        num_tokens=num_tokens,
+        extend_start_loc=None,
+    )
+
+
+def compute_ragged_extend_lengths(
+    *,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: List[int],
+    ragged_layout: RaggedVerifyLayout,
+) -> VerifyExtendLengths:
+    extend_seq_lens_cpu = list(ragged_layout.verify_lens_cpu)
+    seq_lens_extended = seq_lens + ragged_layout.verify_lens
+    seq_lens_cpu_extended = [
+        raw + length for raw, length in zip(seq_lens_cpu, extend_seq_lens_cpu)
+    ]
+    num_tokens = ragged_layout.total_verify_tokens
+    extend_start_loc = ragged_layout.extend_start_loc
+    return VerifyExtendLengths(
+        seq_lens_extended=seq_lens_extended,
+        seq_lens_cpu_extended=seq_lens_cpu_extended,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        num_tokens=num_tokens,
+        extend_start_loc=extend_start_loc,
+    )

python/sglang/srt/speculative/ragged_verify_kernels.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+
+class PaddedToBucket:
+    @classmethod
+    def execute(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+        bs: int,
+        padded_bs: int,
+    ) -> torch.Tensor:
+        impl = cls.triton if verify_lens.is_cuda else cls.torch
+        return impl(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            bs=bs,
+            padded_bs=padded_bs,
+        )
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+        bs: int,
+        padded_bs: int,
+    ) -> torch.Tensor:
+        return pad_verify_lens_to_bucket(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            bs=bs,
+            padded_bs=padded_bs,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+        bs: int,
+        padded_bs: int,
+    ) -> torch.Tensor:
+        return pad_verify_lens_to_bucket_triton(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            bs=bs,
+            padded_bs=padded_bs,
+        )
+
+
+def pad_verify_lens_to_bucket(
+    *,
+    verify_lens: torch.Tensor,
+    graph_num_tokens: int,
+    bs: int,
+    padded_bs: int,
+) -> torch.Tensor:
+    assert padded_bs >= bs, (
+        f"padded_bs {padded_bs} < bs {bs}: the captured tier cannot hold this "
+        "batch's requests"
+    )
+    device = verify_lens.device
+    num_pad_reqs = padded_bs - bs
+    padded = verify_lens.to(torch.int32)
+    leftover = graph_num_tokens - padded.to(torch.int64).sum()
+    if num_pad_reqs > 0:
+        base = leftover // num_pad_reqs
+        rem = leftover - base * num_pad_reqs
+        pad_block = base + (
+            torch.arange(num_pad_reqs, device=device, dtype=torch.int64) < rem
+        )
+        padded = torch.cat([padded, pad_block.to(torch.int32)])
+    else:
+        padded = padded.clone()
+        padded[-1] = (padded[-1].to(torch.int64) + leftover).to(torch.int32)
+    return padded
+
+
+@triton.jit
+def _padded_to_bucket_kernel(
+    verify_lens_ptr,
+    out_ptr,
+    bs,
+    padded_bs,
+    graph_num_tokens,
+    BLOCK: tl.constexpr,
+):
+    idx = tl.arange(0, BLOCK)
+    valid = idx < padded_bs
+    is_real = idx < bs
+    vl = tl.load(verify_lens_ptr + idx, mask=is_real, other=0).to(tl.int64)
+    leftover = graph_num_tokens - tl.sum(vl)
+    num_pad = padded_bs - bs
+    num_pad_safe = tl.maximum(num_pad, 1)
+    base = leftover // num_pad_safe
+    rem = leftover - base * num_pad_safe
+    pad_len = base + tl.where((idx - bs) < rem, 1, 0)
+    final = tl.where(is_real, vl, pad_len)
+    final = final + tl.where((num_pad == 0) & (idx == bs - 1), leftover, 0)
+    tl.store(out_ptr + idx, final.to(tl.int32), mask=valid)
+
+
+def pad_verify_lens_to_bucket_triton(
+    *,
+    verify_lens: torch.Tensor,
+    graph_num_tokens: int,
+    bs: int,
+    padded_bs: int,
+) -> torch.Tensor:
+    assert padded_bs >= bs, (
+        f"padded_bs {padded_bs} < bs {bs}: the captured tier cannot hold this "
+        "batch's requests"
+    )
+    device = verify_lens.device
+    verify_lens = verify_lens.to(torch.int32).contiguous()
+    out = torch.empty(padded_bs, dtype=torch.int32, device=device)
+    BLOCK = triton.next_power_of_2(max(padded_bs, 1))
+    _padded_to_bucket_kernel[(1,)](
+        verify_lens,
+        out,
+        bs,
+        padded_bs,
+        graph_num_tokens,
+        BLOCK=BLOCK,
+    )
+    return out
+
+
+class QoIndptrResult(msgspec.Struct):
+    qo_indptr: torch.Tensor
+    extend_start_loc: torch.Tensor
+
+
+class BuildQoIndptr:
+    @classmethod
+    def execute(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult:
+        impl = cls.triton if verify_lens.is_cuda else cls.torch
+        return impl(verify_lens=verify_lens)
+
+    @classmethod
+    def torch(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult:
+        return build_qo_indptr(verify_lens=verify_lens)
+
+    @classmethod
+    def triton(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult:
+        return build_qo_indptr_triton(verify_lens=verify_lens)
+
+
+def build_qo_indptr(*, verify_lens: torch.Tensor) -> QoIndptrResult:
+    verify_lens = verify_lens.to(torch.int32)
+    cumsum = torch.cumsum(verify_lens, dim=0).to(torch.int32)
+    zero = torch.zeros(1, dtype=torch.int32, device=verify_lens.device)
+    qo_indptr = torch.cat([zero, cumsum])
+    extend_start_loc = qo_indptr[:-1].clone()
+    return QoIndptrResult(qo_indptr=qo_indptr, extend_start_loc=extend_start_loc)
+
+
+@triton.jit
+def _qo_indptr_kernel(
+    verify_lens_ptr,
+    qo_indptr_ptr,
+    extend_start_loc_ptr,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    idx = tl.arange(0, BLOCK)
+    valid = idx < bs
+    vl = tl.load(verify_lens_ptr + idx, mask=valid, other=0).to(tl.int32)
+    incl = tl.cumsum(vl, axis=0)
+    excl = incl - vl
+    tl.store(qo_indptr_ptr, 0)
+    tl.store(qo_indptr_ptr + 1 + idx, incl, mask=valid)
+    tl.store(extend_start_loc_ptr + idx, excl, mask=valid)
+
+
+def build_qo_indptr_triton(*, verify_lens: torch.Tensor) -> QoIndptrResult:
+    bs = verify_lens.shape[0]
+    device = verify_lens.device
+    verify_lens = verify_lens.contiguous()
+    qo_indptr = torch.empty(bs + 1, dtype=torch.int32, device=device)
+    extend_start_loc = torch.empty(bs, dtype=torch.int32, device=device)
+    BLOCK = triton.next_power_of_2(max(bs, 1))
+    _qo_indptr_kernel[(1,)](
+        verify_lens,
+        qo_indptr,
+        extend_start_loc,
+        bs,
+        BLOCK=BLOCK,
+    )
+    return QoIndptrResult(qo_indptr=qo_indptr, extend_start_loc=extend_start_loc)

python/sglang/srt/speculative/reject_sampling.py

@@ -0,0 +1,204 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def speculative_sampling_classic_kernel(
+    # Pointers
+    Predicts,
+    AcceptIndex,
+    AcceptTokenNum,
+    Candidates,
+    RetriveIndex,
+    UniformSamples,
+    UniformSamplesFinal,
+    TargetProbs,
+    DraftProbs,
+    # Strides
+    stride_cand_b,
+    stride_cand_s,
+    stride_idx_b,
+    stride_idx_s,
+    stride_uni_b,
+    stride_uni_s,
+    stride_tp_b,
+    stride_tp_s,
+    stride_tp_v,
+    stride_dp_b,
+    stride_dp_s,
+    stride_dp_v,
+    # Constants
+    NUM_SLOTS: tl.constexpr,
+    VOCAB_SIZE: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    cur_prob_row = 0
+
+    cand_ptr_base = Candidates + pid * stride_cand_b
+    idx_ptr_base = RetriveIndex + pid * stride_idx_b
+    uni_ptr_base = UniformSamples + pid * stride_uni_b
+
+    root_global_idx = tl.load(idx_ptr_base + 0 * stride_idx_s)
+    tl.store(AcceptIndex + pid * stride_idx_b + 0 * stride_idx_s, root_global_idx)
+    last_accepted_global_idx = root_global_idx
+
+    num_accept = 0
+
+    # Verification Loop
+    step = 1
+    continue_verifying = 1
+
+    while (step < NUM_SLOTS) and (continue_verifying == 1):
+        draft_token = tl.load(cand_ptr_base + step * stride_cand_s)
+
+        offset_prob = (
+            (pid * stride_tp_b)
+            + (cur_prob_row * stride_tp_s)
+            + (draft_token * stride_tp_v)
+        )
+        offset_draft = (
+            (pid * stride_dp_b)
+            + (cur_prob_row * stride_dp_s)
+            + (draft_token * stride_dp_v)
+        )
+
+        p = tl.load(TargetProbs + offset_prob)
+        q = tl.load(DraftProbs + offset_draft)
+
+        coin = tl.load(uni_ptr_base + (step - 1) * stride_uni_s)
+
+        if coin * q < p:
+            num_accept += 1
+            cur_prob_row = step
+            tl.store(Predicts + last_accepted_global_idx, draft_token)
+
+            curr_global_idx = tl.load(idx_ptr_base + step * stride_idx_s)
+            tl.store(
+                AcceptIndex + pid * stride_idx_b + num_accept * stride_idx_s,
+                curr_global_idx,
+            )
+            last_accepted_global_idx = curr_global_idx
+
+            step += 1
+        else:
+            continue_verifying = 0
+
+    tl.store(AcceptTokenNum + pid, num_accept)
+
+    # Final Sampling
+    all_drafts_accepted = continue_verifying
+    coin_final = tl.load(UniformSamplesFinal + pid)
+    norm_sum = 0.0
+
+    tp_base_ptr = TargetProbs + (pid * stride_tp_b) + (cur_prob_row * stride_tp_s)
+    # DraftProbs has only num_steps rows (TargetProbs has num_steps + 1). When
+    # all drafts are accepted cur_prob_row == num_steps is out of bounds for
+    # DraftProbs, but the all-accepted branch samples pure target p and never
+    # dereferences this pointer; on rejection cur_prob_row <= num_steps - 1.
+    dp_base_ptr_safe = DraftProbs + (pid * stride_dp_b) + (cur_prob_row * stride_dp_s)
+
+    # Pass 1: Sum
+    for v_start in range(0, VOCAB_SIZE, BLOCK_V):
+        v_offsets = v_start + tl.arange(0, BLOCK_V)
+        mask = v_offsets < VOCAB_SIZE
+
+        p_ptr = tp_base_ptr + v_offsets * stride_tp_v
+        p_val = tl.load(p_ptr, mask=mask, other=0.0)
+
+        if all_drafts_accepted:
+            val = p_val
+        else:
+            q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v
+            q_val = tl.load(q_ptr, mask=mask, other=0.0)
+            diff = p_val - q_val
+            val = tl.where(diff > 0.0, diff, 0.0)
+
+        norm_sum += tl.sum(val)
+
+    # Pass 2: CDF. Degenerate residual (norm_sum == 0, i.e. p == q everywhere on
+    # rejection) leaves the cumsum at 0 <= target_u, so final_token falls back to
+    # VOCAB_SIZE - 1; acceptable since this case is numerically near-impossible.
+    target_u = coin_final * norm_sum
+    cum_sum = 0.0
+    final_token = VOCAB_SIZE - 1
+    found = 0
+
+    for v_start in range(0, VOCAB_SIZE, BLOCK_V):
+        if found == 0:
+            v_offsets = v_start + tl.arange(0, BLOCK_V)
+            mask = v_offsets < VOCAB_SIZE
+
+            p_ptr = tp_base_ptr + v_offsets * stride_tp_v
+            p_val = tl.load(p_ptr, mask=mask, other=0.0)
+
+            if all_drafts_accepted:
+                val = p_val
+            else:
+                q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v
+                q_val = tl.load(q_ptr, mask=mask, other=0.0)
+                diff = p_val - q_val
+                val = tl.where(diff > 0.0, diff, 0.0)
+
+            block_cumsum = tl.cumsum(val, axis=0)
+            total_cumsum = cum_sum + block_cumsum
+
+            candidates_mask = total_cumsum > target_u
+            has_match = tl.max(candidates_mask, axis=0)
+
+            if has_match:
+                match_idx = tl.argmax(candidates_mask.to(tl.int32), axis=0)
+                final_token = v_start + match_idx
+                found = 1
+
+            cum_sum += tl.sum(val)
+
+    tl.store(Predicts + last_accepted_global_idx, final_token)
+
+
+def chain_speculative_sampling_triton(
+    predicts,
+    accept_index,
+    accept_token_num,
+    candidates,
+    retrive_index,
+    retrive_next_token,
+    retrive_next_sibling,  # not used in chain verification
+    uniform_samples,
+    uniform_samples_for_final_sampling,
+    target_probs,
+    draft_probs,
+    threshold_single,
+    threshold_acc,
+    deterministic,  # not used
+):
+    batch_size, num_slots = candidates.shape
+    vocab_size = target_probs.shape[-1]
+
+    grid = (batch_size,)
+    speculative_sampling_classic_kernel[grid](
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrive_index,
+        uniform_samples,
+        uniform_samples_for_final_sampling,
+        target_probs,
+        draft_probs,
+        candidates.stride(0),
+        candidates.stride(1),
+        retrive_index.stride(0),
+        retrive_index.stride(1),
+        uniform_samples.stride(0),
+        uniform_samples.stride(1),
+        target_probs.stride(0),
+        target_probs.stride(1),
+        target_probs.stride(2),
+        draft_probs.stride(0),
+        draft_probs.stride(1),
+        draft_probs.stride(2),
+        NUM_SLOTS=num_slots,
+        VOCAB_SIZE=vocab_size,
+        BLOCK_V=4096,
+    )

python/sglang/srt/speculative/spec_info.py

@@ -0,0 +1,402 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum, IntEnum, auto
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
+
+import torch
+
+from sglang.srt.speculative.spec_registry import (
+    CustomSpecAlgo,
+    ServerArgsValidator,
+    WorkerFactory,
+)
+from sglang.srt.speculative.spec_registry import get_spec as _get_registered_spec
+from sglang.srt.speculative.spec_registry import (
+    register_algorithm as _register_algorithm,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.overlap_utils import FutureMap
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+    from sglang.srt.speculative.ngram_worker import NGRAMWorker
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+class SpeculativeAlgorithm(Enum):
+    """Builtin speculative decoding algorithms. Plugin-registered ones are
+    ``CustomSpecAlgo`` instances; ``from_string`` returns either type, and
+    both expose the same ``is_*()`` / ``create_worker`` interface so callers
+    dispatch uniformly without isinstance checks.
+    """
+
+    DFLASH = auto()
+    DSPARK = auto()
+    EAGLE = auto()
+    EAGLE3 = auto()
+    FROZEN_KV_MTP = auto()
+    STANDALONE = auto()
+    NGRAM = auto()
+    NONE = auto()
+
+    @classmethod
+    def from_string(
+        cls, name: Optional[str]
+    ) -> Union[SpeculativeAlgorithm, CustomSpecAlgo]:
+        if name is None:
+            return cls.NONE
+        upper = name.upper()
+        try:
+            return cls[upper]
+        except KeyError:
+            pass
+        spec = _get_registered_spec(upper)
+        if spec is not None:
+            return spec
+        raise ValueError(f"Unknown speculative algorithm name: {name}")
+
+    @classmethod
+    def register(
+        cls,
+        name: str,
+        *,
+        supports_overlap: bool = False,
+        validate_server_args: Optional[ServerArgsValidator] = None,
+        spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo,
+    ) -> Callable[[WorkerFactory], WorkerFactory]:
+        """Decorator to register a plugin speculative algorithm. The factory
+        takes ``server_args`` and returns the worker class. Pass a
+        ``CustomSpecAlgo`` subclass via ``spec_class`` to override any
+        ``is_*()`` / ``create_worker`` method.
+
+        Example:
+            @SpeculativeAlgorithm.register("MY_SPEC", supports_overlap=False)
+            def _factory(server_args):
+                return MySpecWorker
+        """
+        return _register_algorithm(
+            name,
+            supports_overlap=supports_overlap,
+            validate_server_args=validate_server_args,
+            spec_class=spec_class,
+        )
+
+    def is_some(self) -> bool:
+        return self != SpeculativeAlgorithm.NONE
+
+    def is_none(self) -> bool:
+        return self == SpeculativeAlgorithm.NONE
+
+    def is_speculative(self) -> bool:
+        return self != SpeculativeAlgorithm.NONE
+
+    def is_eagle(self) -> bool:
+        # FIXME(kpham_sgl): Remove FROZEN_KV_MTP here once we
+        # have established support for it in the scheduler.
+        return self in (
+            SpeculativeAlgorithm.EAGLE,
+            SpeculativeAlgorithm.EAGLE3,
+            SpeculativeAlgorithm.FROZEN_KV_MTP,
+        )
+
+    def is_eagle3(self) -> bool:
+        return self == SpeculativeAlgorithm.EAGLE3
+
+    def is_frozen_kv_mtp(self) -> bool:
+        return self == SpeculativeAlgorithm.FROZEN_KV_MTP
+
+    def is_dflash(self) -> bool:
+        return self == SpeculativeAlgorithm.DFLASH
+
+    def is_dspark(self) -> bool:
+        return self == SpeculativeAlgorithm.DSPARK
+
+    def is_dflash_family(self) -> bool:
+        return self.is_dflash() or self.is_dspark()
+
+    def is_standalone(self) -> bool:
+        return self == SpeculativeAlgorithm.STANDALONE
+
+    def is_ngram(self) -> bool:
+        return self == SpeculativeAlgorithm.NGRAM
+
+    def supports_target_verify_for_draft(self) -> bool:
+        return self.is_dflash_family()
+
+    def supports_ragged_verify(self) -> bool:
+        """Whether this algorithm's verify step may carry a RaggedVerifyLayout
+        (per-request verify lengths); gates the token-bucket-keyed verify
+        graphs in the decode cuda graph runner."""
+        return self.is_dspark()
+
+    def has_draft_kv(self) -> bool:
+        """Whether the draft phase writes KV chains. NGRAM does not (its tree
+        lives only in the verify mask), so per-decode KV sizing needs no
+        per-topk page rounding; see get_alloc_len_per_decode."""
+        return not self.is_ngram()
+
+    def carries_draft_hidden_states(self) -> bool:
+        """Whether the disagg prefill->decode transfer carries draft hidden
+        states (EAGLE-family only; STANDALONE's vanilla draft ignores them)."""
+        return self.is_eagle()
+
+    def create_future_map(
+        self,
+        device: torch.device,
+        req_to_token_pool,
+        needs_cpu_seq_lens: bool = True,
+        needs_confidence_relay: bool = False,
+    ) -> FutureMap:
+        from sglang.srt.managers.overlap_utils import FutureMap
+
+        return FutureMap(
+            device,
+            self,
+            req_to_token_pool,
+            needs_cpu_seq_lens,
+            needs_confidence_relay,
+        )
+
+    def build_disagg_draft_input(
+        self,
+        batch: ScheduleBatch,
+        server_args: ServerArgs,
+        last_tokens_tensor: torch.Tensor,
+        future_map: FutureMap,
+    ) -> Optional[SpecInput]:
+        if self.is_eagle():
+            from sglang.srt.speculative.eagle_disaggregation import (
+                build_eagle_disagg_draft_input,
+            )
+
+            return build_eagle_disagg_draft_input(
+                batch, server_args, last_tokens_tensor, future_map
+            )
+        return None
+
+    def need_topk(self) -> bool:
+        return self.is_eagle() or self.is_standalone()
+
+    def handle_server_args(self, server_args: ServerArgs) -> None:
+        """Hook for per-algorithm server args mutation.
+
+        In-place updated.
+        """
+        from sglang.srt.arg_groups.speculative_hook import (
+            _handle_dflash,
+            _handle_dspark,
+            _handle_eagle_family,
+            _handle_frozen_kv_mtp,
+            _handle_ngram,
+        )
+
+        # Validate for every algorithm at startup: the metrics paths read the
+        # ragged-verify mode env and must not be where a typo'd value raises.
+        from sglang.srt.speculative.ragged_verify import read_ragged_verify_mode
+
+        read_ragged_verify_mode()
+
+        if self.is_dflash():
+            _handle_dflash(server_args)
+        elif self.is_dspark():
+            _handle_dspark(server_args)
+        elif self.is_frozen_kv_mtp():
+            _handle_frozen_kv_mtp(server_args)
+        elif self.is_eagle() or self.is_standalone():
+            _handle_eagle_family(server_args)
+        elif self.is_ngram():
+            _handle_ngram(server_args)
+
+    def get_num_tokens_per_bs_for_target_verify(
+        self, num_draft_tokens: int, is_draft_worker: bool
+    ) -> int:
+        # FIXME: Remove this after the forward mode refactor. Target verify is
+        # essentially a fixed sequence length prefill/extend with full cuda
+        # graph support. We can use it for target verify, or we can use it for
+        # other cases which is not target verify but fixed length prefill.
+        # Here, we expose this interface to allow the other use cases.
+        if self.is_dspark() and is_draft_worker:
+            return num_draft_tokens - 1
+        return num_draft_tokens
+
+    def create_worker(
+        self, server_args: ServerArgs
+    ) -> Optional[Union[Type[BaseSpecWorker], Type[TpModelWorker], Type[NGRAMWorker]]]:
+        assert (
+            not self.is_none()
+        ), "Cannot create worker for NONE speculative algorithm."
+
+        if self.is_dflash():
+            # V2 worker drives both overlap and non-overlap (scheduler runs it
+            # synchronously when overlap is disabled), same as EAGLE.
+            from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
+
+            return DFlashWorkerV2
+
+        if self.is_dspark():
+            from sglang.srt.speculative.dspark_components.dspark_worker_v2 import (
+                DSparkWorkerV2,
+            )
+
+            return DSparkWorkerV2
+
+        if self.is_frozen_kv_mtp():
+            # V2 worker drives both overlap and non-overlap (scheduler runs it
+            # synchronously when overlap is disabled), same as EAGLE.
+            from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import (
+                FrozenKVMTPWorkerV2,
+            )
+
+            return FrozenKVMTPWorkerV2
+
+        # EAGLE / EAGLE3 / STANDALONE / MULTI_LAYER always use the V2 worker,
+        # even with overlap disabled (scheduler drives it synchronously).
+        if self.is_eagle() and server_args.enable_multi_layer_eagle:
+            from sglang.srt.speculative.multi_layer_eagle_worker_v2 import (
+                MultiLayerEagleWorkerV2,
+            )
+
+            return MultiLayerEagleWorkerV2
+
+        elif self.is_eagle():
+            from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
+
+            return EAGLEWorkerV2
+        elif self.is_standalone():
+            from sglang.srt.speculative.standalone_worker_v2 import (
+                StandaloneWorkerV2,
+            )
+
+            return StandaloneWorkerV2
+        elif self.is_ngram():
+            from sglang.srt.speculative.ngram_worker import NGRAMWorker
+
+            return NGRAMWorker
+
+        raise ValueError("Unreachable code path in create_worker.")
+
+
+class SpecInputType(IntEnum):
+    # NOTE: introduce this to distinguish the SpecInput types of multiple algorithms when asserting in attention backends.
+    # If all algorithms can share the same datastrucutre of draft_input and verify_input, consider simplify it
+    EAGLE_DRAFT = auto()
+    EAGLE_DRAFT_EXTEND = auto()
+    EAGLE_VERIFY = auto()
+    FROZEN_KV_MTP_DRAFT = auto()
+    FROZEN_KV_MTP_VERIFY = auto()
+    DFLASH_DRAFT = auto()
+    DFLASH_VERIFY = auto()
+    NGRAM_VERIFY = auto()
+
+
+class SpecInput(ABC):
+    # Per-request verify lengths for the ragged-verify graphs (see
+    # sglang.srt.speculative.ragged_verify); verify inputs of algorithms with
+    # supports_ragged_verify() override it per step. Must stay a class-level
+    # default, not an __init__ assignment: dataclass subclasses declare it as
+    # a field and run __post_init__ -> super().__init__ *after* field
+    # assignment, so an init-time default would clobber the passed layout.
+    ragged_verify_layout: Optional[RaggedVerifyLayout] = None
+
+    def __init__(self, spec_input_type: SpecInputType):
+        self.spec_input_type = spec_input_type
+
+    # Cross-algorithm phase guards. Used by attention backends and
+    # ForwardBatch padding logic to dispatch on phase without hardcoding the
+    # specific algo class (EAGLE / FROZEN_KV_MTP / DFLASH / NGRAM each have
+    # their own draft / verify SpecInput subclasses).
+    def is_draft_input(self) -> bool:
+        return self.spec_input_type in {
+            SpecInputType.EAGLE_DRAFT,
+            SpecInputType.EAGLE_DRAFT_EXTEND,
+            SpecInputType.FROZEN_KV_MTP_DRAFT,
+            SpecInputType.DFLASH_DRAFT,
+        }
+
+    def is_verify_input(self) -> bool:
+        return self.spec_input_type in {
+            SpecInputType.EAGLE_VERIFY,
+            SpecInputType.FROZEN_KV_MTP_VERIFY,
+            SpecInputType.DFLASH_VERIFY,
+            SpecInputType.NGRAM_VERIFY,
+        }
+
+    @abstractmethod
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        pass
+
+    def get_spec_adjusted_global_num_tokens(
+        self, batch: ScheduleBatch
+    ) -> Tuple[List[int], List[int]]:
+        c1, c2 = self.get_spec_adjust_token_coefficient()
+        global_num_tokens = [x * c1 for x in batch.global_num_tokens]
+        global_num_tokens_for_logprob = [
+            x * c2 for x in batch.global_num_tokens_for_logprob
+        ]
+        return global_num_tokens, global_num_tokens_for_logprob
+
+
+def create_dummy_verify_input(
+    spec_algorithm: SpeculativeAlgorithm,
+    server_args: ServerArgs,
+    custom_mask: torch.Tensor,
+    num_tokens_per_bs: int,
+    is_draft_worker: bool,
+) -> Optional[SpecInput]:
+    """Dummy verify ``SpecInput`` for CUDA-graph capture (per-algorithm dispatch)."""
+    from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+
+    spec_info = None
+    if spec_algorithm.is_eagle() or spec_algorithm.is_standalone():
+        from sglang.srt.speculative.eagle_info import EagleVerifyInput
+
+        if is_draft_worker:
+            raise RuntimeError("This should not happen.")
+        else:
+            spec_info = EagleVerifyInput(
+                draft_token=None,
+                custom_mask=custom_mask,
+                positions=None,
+                retrieve_index=None,
+                retrieve_next_token=None,
+                retrieve_next_sibling=None,
+                retrieve_cum_len=None,
+                spec_steps=server_args.speculative_num_steps,
+                topk=server_args.speculative_eagle_topk,
+                draft_token_num=server_args.speculative_num_draft_tokens,
+                capture_hidden_mode=CaptureHiddenMode.FULL,
+                seq_lens_sum=None,
+                seq_lens_cpu=None,
+            )
+    elif spec_algorithm.is_dflash_family():
+        from sglang.srt.speculative.dflash_info import DFlashVerifyInput
+
+        # Dummy warmup only needs shape metadata; avoid forcing custom-mask mode.
+        spec_info = DFlashVerifyInput(
+            draft_token=None,
+            positions=None,
+            draft_token_num=server_args.speculative_num_draft_tokens,
+            custom_mask=None,
+            capture_hidden_mode=(
+                CaptureHiddenMode.NULL if is_draft_worker else CaptureHiddenMode.FULL
+            ),
+        )
+
+    elif spec_algorithm.is_ngram():
+        from sglang.srt.speculative.ngram_info import NgramVerifyInput
+
+        spec_info = NgramVerifyInput(
+            draft_token=None,
+            custom_mask=custom_mask,
+            positions=None,
+            retrieve_index=None,
+            retrieve_next_token=None,
+            retrieve_next_sibling=None,
+            draft_token_num=num_tokens_per_bs,
+        )
+        spec_info.capture_hidden_mode = CaptureHiddenMode.NULL
+
+    return spec_info

python/sglang/srt/speculative/spec_registry.py

@@ -0,0 +1,236 @@
+"""Internal storage backing ``SpeculativeAlgorithm.register``. Plugins
+should use that classmethod API; do not import from this module directly.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Type
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.overlap_utils import FutureMap
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.spec_info import SpecInput
+
+WorkerFactory = Callable[["ServerArgs"], Type]
+ServerArgsValidator = Callable[["ServerArgs"], None]
+
+logger = logging.getLogger(__name__)
+
+
+class CustomSpecAlgo:
+    """A plugin-registered speculative algorithm. Duck-types
+    ``SpeculativeAlgorithm`` enum values (same ``is_*()`` / ``create_worker``
+    interface).
+
+    Plugins may subclass this to override any ``is_*()`` / ``supports_*()`` /
+    ``create_worker`` method (e.g. to integrate with builtin-specific
+    branches like ``if spec_algorithm.is_eagle():`` in scheduler /
+    model_runner). Pass the subclass via ``spec_class=...`` at registration.
+
+    Defaults: all ``is_*()`` return ``False`` except ``is_speculative``.
+
+    ``supports_overlap=False`` is deprecated: the spec V1 worker path has been
+    removed, so such algorithms run on the V2 scheduler schema with overlap
+    disabled (synchronous). Migrate plugin workers to the V2 schema and
+    overlap scheduling.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        factory: WorkerFactory,
+        *,
+        supports_overlap: bool = False,
+        validate_server_args: Optional[ServerArgsValidator] = None,
+    ):
+        self.name = name
+        self.factory = factory
+        self.supports_overlap = supports_overlap
+        self.validate_server_args = validate_server_args
+
+    def __repr__(self) -> str:
+        return f"CustomSpecAlgo({self.name!r})"
+
+    def is_some(self) -> bool:
+        return True
+
+    def is_none(self) -> bool:
+        return False
+
+    def is_speculative(self) -> bool:
+        return True
+
+    def is_eagle(self) -> bool:
+        return False
+
+    def is_eagle3(self) -> bool:
+        return False
+
+    def is_frozen_kv_mtp(self) -> bool:
+        return False
+
+    def is_dflash(self) -> bool:
+        return False
+
+    def is_dspark(self) -> bool:
+        return False
+
+    def is_dflash_family(self) -> bool:
+        return False
+
+    def is_standalone(self) -> bool:
+        return False
+
+    def is_ngram(self) -> bool:
+        return False
+
+    def supports_target_verify_for_draft(self) -> bool:
+        return False
+
+    def supports_ragged_verify(self) -> bool:
+        return False
+
+    def has_draft_kv(self) -> bool:
+        # Conservative default: the larger KV reserve.
+        return True
+
+    def handle_server_args(self, server_args: ServerArgs) -> None:
+        pass
+
+    def create_worker(self, server_args: ServerArgs) -> Type:
+        if not server_args.disable_overlap_schedule and not self.supports_overlap:
+            raise ValueError(
+                f"Speculative algorithm {self.name} does not support overlap scheduling."
+            )
+        if not self.supports_overlap:
+            # Reached only when overlap is disabled, so the algorithm really
+            # does run synchronously on the V2 schema below.
+            logger.warning(
+                "Speculative algorithm %s is registered with "
+                "supports_overlap=False, which is deprecated: the spec V1 "
+                "worker path has been removed, and the algorithm now runs on "
+                "the V2 scheduler schema with overlap disabled (synchronous). "
+                "Migrate the plugin worker to support overlap scheduling.",
+                self.name,
+            )
+        return self.factory(server_args)
+
+    def get_num_tokens_per_bs_for_target_verify(
+        self, num_draft_tokens: int, is_draft_worker: bool
+    ) -> int:
+        # FIXME: Remove this after the forward mode refactor. Target verify is
+        # essentially a fixed sequence length prefill/extend with full cuda
+        # graph support. We can use it for target verify, or we can use it for
+        # other cases which is not target verify but fixed length prefill.
+        # Here, we expose this interface to allow the other use cases.
+        return num_draft_tokens
+
+    def build_disagg_draft_input(
+        self,
+        batch: ScheduleBatch,
+        server_args: ServerArgs,
+        last_tokens_tensor: torch.Tensor,
+        future_map: FutureMap,
+    ) -> Optional[SpecInput]:
+        return None
+
+
+_REGISTRY: Dict[str, CustomSpecAlgo] = {}
+
+# CLI spellings that are not ``SpeculativeAlgorithm`` members but still resolve
+# to a builtin (e.g. NEXTN -> EAGLE). Reserved alongside the enum members so
+# plugins cannot shadow them.
+_RESERVED_ALIASES = frozenset({"NEXTN"})
+
+
+def _reserved_names() -> frozenset:
+    """Names plugins cannot register under: every ``SpeculativeAlgorithm``
+    member plus ``_RESERVED_ALIASES``.
+
+    Derived from the enum (lazily, to avoid a circular import — ``spec_info``
+    imports this module) so any new builtin is reserved automatically without
+    editing a second list.
+    """
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+    return frozenset(algo.name for algo in SpeculativeAlgorithm) | _RESERVED_ALIASES
+
+
+def _assert_custom_spec_algo_conforms(spec_class: Type[CustomSpecAlgo]) -> None:
+    """Fail fast if ``spec_class`` drifts from the ``SpeculativeAlgorithm``
+    duck-typing contract.
+
+    ``from_string`` returns either type and callers dispatch on the shared
+    ``is_*()`` / ``supports_*()`` interface without isinstance checks, so every
+    such method on the enum must also exist on the registered spec class —
+    otherwise a plugin-registered algo hits ``AttributeError`` at a call site
+    (this is how ``is_some`` / ``is_frozen_kv_mtp`` silently went missing). New
+    predicates are covered automatically; no second list to maintain.
+
+    Called from ``register_algorithm`` rather than at import time because
+    ``spec_info`` imports this module, so ``SpeculativeAlgorithm`` does not yet
+    exist while this module is loading; at registration time it is fully
+    defined.
+    """
+    # NOTE: use ``vars()`` not ``dir()`` for the enum — ``EnumMeta.__dir__``
+    # hides instance methods, so ``dir(SpeculativeAlgorithm)`` would yield an
+    # empty interface and turn this guard into a silent no-op.
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+    interface = {
+        name
+        for name in vars(SpeculativeAlgorithm)
+        if name.startswith(("is_", "supports_"))
+    }
+    missing = sorted(interface - set(dir(spec_class)))
+    if missing:
+        raise TypeError(
+            f"{spec_class.__name__} is missing duck-typed methods from "
+            f"SpeculativeAlgorithm: {missing}. Add them to {spec_class.__name__} "
+            "so plugin-registered algorithms stay dispatchable."
+        )
+
+
+def register_algorithm(
+    name: str,
+    *,
+    supports_overlap: bool = False,
+    validate_server_args: Optional[ServerArgsValidator] = None,
+    spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo,
+) -> Callable[[WorkerFactory], WorkerFactory]:
+    """Return a decorator that registers a plugin algorithm under ``name``.
+
+    Pass a ``spec_class`` subclass of ``CustomSpecAlgo`` to override any
+    ``is_*()`` / ``supports_*()`` / ``create_worker`` method.
+    """
+    upper = name.upper()
+    if upper in _reserved_names():
+        raise ValueError(
+            f"'{upper}' is a reserved speculative algorithm name; cannot be re-registered."
+        )
+    if upper in _REGISTRY:
+        raise ValueError(f"Speculative algorithm '{upper}' already registered.")
+    _assert_custom_spec_algo_conforms(spec_class)
+
+    def decorator(factory: WorkerFactory) -> WorkerFactory:
+        _REGISTRY[upper] = spec_class(
+            name=upper,
+            factory=factory,
+            supports_overlap=supports_overlap,
+            validate_server_args=validate_server_args,
+        )
+        return factory
+
+    return decorator
+
+
+def get_spec(name: Optional[str]) -> Optional[CustomSpecAlgo]:
+    """Return the registered spec for ``name``, or ``None`` for builtin /
+    unknown names."""
+    if name is None:
+        return None
+    return _REGISTRY.get(name.upper())

python/sglang/srt/speculative/spec_utils.py

@@ -0,0 +1,719 @@
+from __future__ import annotations
+
+import logging
+import os
+import time
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+from huggingface_hub import snapshot_download
+
+from sglang.kernels.ops.speculative.cache_locs import (
+    align_evict_mask_to_page_size as align_evict_mask_to_page_size,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    assign_extend_cache_locs as assign_extend_cache_locs,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    assign_req_to_token_pool as assign_req_to_token_pool,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    assign_req_to_token_pool_func as assign_req_to_token_pool_func,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    generate_draft_decode_kv_indices as generate_draft_decode_kv_indices,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    get_src_tgt_cache_loc as get_src_tgt_cache_loc,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    get_target_cache_loc as get_target_cache_loc,
+)
+from sglang.kernels.ops.speculative.eagle import (
+    fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func,
+)
+from sglang.srt.distributed.parallel_state import (
+    GroupCoordinator,
+    patch_tensor_parallel_group,
+)
+from sglang.srt.environ import envs
+from sglang.srt.managers.schedule_batch import set_mamba_track_indices_from_reqs
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+    next_power_of_2,
+)
+from sglang.srt.utils.async_probe import maybe_detect_oob
+from sglang.srt.utils.nvtx_utils import profile_range
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+_is_musa = is_musa()
+_is_xpu = is_xpu()
+_is_cpu = is_cpu()
+
+if TYPE_CHECKING:
+    from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+    from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.eagle_info import EagleVerifyInput
+
+
+if _is_cuda:
+    from sgl_kernel import fast_topk
+elif _is_hip:
+    from sgl_kernel import fast_topk
+else:
+    from sglang.srt.utils.common import fast_topk
+
+if _is_cpu:
+    from sgl_kernel import assign_extend_cache_locs_cpu
+
+
+logger = logging.getLogger(__name__)
+
+
+def fast_sample(probs: torch.Tensor, num_samples: int = 1):
+    sample_index = torch.multinomial(probs, num_samples=num_samples)
+    sample_p = probs.gather(1, sample_index)
+    return sample_p, sample_index
+
+
+def renorm_draft_probs(
+    next_token_logits: torch.Tensor,
+    sampling_info,
+    use_rejection_sampling: bool,
+) -> torch.Tensor:
+    """Draft-side next-token distribution.
+
+    Plain softmax, except under rejection sampling where logits are
+    temperature-scaled so the draft proposal q tracks the target sampling
+    temperature (higher acceptance; correctness holds for any q).
+    """
+    if not use_rejection_sampling or not next_token_logits.size(0):
+        return torch.softmax(next_token_logits, dim=-1)
+    return torch.softmax(next_token_logits / sampling_info.temperatures, dim=-1)
+
+
+def sample_draft_proposal(next_token_logits: torch.Tensor, temperatures: torch.Tensor):
+    """Leviathan draft proposal: q = softmax(logits / T), X ~ q.
+
+    Returns (q, q(X), X). The verify's accept test coin*q(X) < p(X) is unbiased
+    only if q is exactly the distribution X was drawn from, so callers must hand
+    the returned q (not a recomputed one) to the verify.
+    """
+    probs = torch.softmax(next_token_logits / temperatures, dim=-1)
+    topk_p, topk_index = fast_sample(probs, num_samples=1)
+    return probs, topk_p, topk_index
+
+
+# Simulate acceptance length for benchmarking purposes
+SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get()  # turn off if < 0
+SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get()
+SIMULATE_ACC_TOKEN_MODE = envs.SGLANG_SIMULATE_ACC_TOKEN_MODE.get()
+
+TREE_TRAVERSE_TIME_THRESHOLD = 1  # TODO: set this properly
+TREE_SPEC_KERNEL_AVAILABLE = (
+    _is_cuda or _is_musa
+)  # This kernel is only available for CUDA and MUSA now
+
+
+def draft_kv_indices_buffer_width(
+    num_seqs: int, topk: int, max_context_len: int
+) -> int:
+    """Per-step row width of the EAGLE draft-decode kv_indices buffer.
+
+    num_seqs * topk branches each attend up to max_context_len KV slots; the topk
+    factor is mandatory -- dropping it under-allocates and overflows the row (#27338, #27460).
+    """
+    assert (
+        num_seqs * topk * max_context_len < 2**31
+    ), "kv_indices flat offset would overflow int32; reduce batch/topk/context"
+    return num_seqs * topk * max_context_len
+
+
+def draft_kv_indices_used_len(
+    seq_lens_sum: int, topk: int, bs: int, num_steps: int
+) -> int:
+    """kv_indices length used through num_steps draft-decode steps.
+
+    bs = topk * num_seqs branches, one index appended per branch per step. Called with
+    num_steps = i + 1 (per-step slice) and speculative_num_steps (capacity assert).
+    """
+    return seq_lens_sum * topk + bs * num_steps
+
+
+def record_stream_each(tensors, stream):
+    """Call record_stream(stream) on each cuda tensor in `tensors`, skipping
+    non-tensor / non-cuda entries. Tells the caching allocator that the
+    tensors are also used on `stream`, so memory is not recycled while
+    queued work is still in flight after Python refs drop.
+    """
+    for t in tensors:
+        if isinstance(t, torch.Tensor) and t.is_cuda:
+            t.record_stream(stream)
+
+
+def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
+    """Mark pre-prepare SB / verify_input GPU tensors as used on `fwd_stream`.
+
+    Spec V2 mutates SB mid-forward (`prepare_for_verify` rebinds
+    `batch.input_ids` / `out_cache_loc`; `_draft_extend_for_decode` later
+    replaces `batch.input_ids` again). Each rebind drops the only SB Python
+    ref to the old tensor while the verify forward kernel may still be
+    reading its memory on `fwd_stream`; `record_stream` tells the caching
+    allocator to wait for `fwd_stream` before recycling the block.
+
+    Covers pre-prepare tensors only; caller must also `record_stream_each`
+    the post-prepare rebinds (new `batch.input_ids` / `out_cache_loc`).
+    """
+    candidates = [
+        batch.seq_lens,
+        batch.req_pool_indices,
+        batch.input_ids,
+        batch.out_cache_loc,
+    ]
+    if verify_input is not None:
+        candidates.extend(
+            [
+                getattr(verify_input, attr, None)
+                for attr in (
+                    "draft_token",
+                    "custom_mask",
+                    "positions",
+                    "retrieve_index",
+                    "retrieve_next_token",
+                    "retrieve_next_sibling",
+                )
+            ]
+        )
+    record_stream_each(candidates, fwd_stream)
+
+
+def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
+    if server_args is None:
+        server_args = get_server_args()
+
+    # STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
+    # multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
+    # TODO(lsyin): also skip when step == 1.
+    if server_args.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
+        return False
+    return not server_args.enable_multi_layer_eagle
+
+
+@torch.compile(dynamic=True, disable=_is_npu or _is_xpu)
+def create_num_accept_tokens_filter(
+    num_correct_drafts: torch.Tensor,
+    unfinished_index_device: torch.Tensor,
+    seq_lens: torch.Tensor,
+):
+    num_accept_tokens_filter = torch.zeros_like(num_correct_drafts)
+    num_accept_tokens_filter[unfinished_index_device] = (
+        num_correct_drafts[unfinished_index_device] + 1
+    )
+    seq_lens.add_(num_correct_drafts + 1)
+    return num_accept_tokens_filter
+
+
+def _select_top_k_tokens_first(
+    topk_p: torch.Tensor,
+    topk_index: torch.Tensor,
+    hidden_states: Optional[torch.Tensor],
+    topk: int,
+):
+    input_ids = topk_index.flatten()
+    if hidden_states is not None:
+        hidden_states = hidden_states.repeat_interleave(topk, dim=0)
+
+    tree_info = (
+        topk_p.unsqueeze(1),  # (b, 1, topk)
+        topk_index,  # (b, topk)
+        torch.arange(-1, topk, dtype=torch.long, device=input_ids.device).expand(
+            topk_p.shape[0], -1
+        ),  # (b, topk + 1) — expand avoids the allocation of repeat
+    )
+    return input_ids, hidden_states, topk_p, tree_info
+
+
+@torch.compile(dynamic=True, disable=_is_npu or _is_xpu)
+def _select_top_k_tokens_later(
+    i: int,
+    topk_p: torch.Tensor,
+    topk_index: torch.Tensor,
+    hidden_states: torch.Tensor,
+    scores: torch.Tensor,
+    topk: int,
+):
+    topk_sq = topk * topk
+
+    expand_scores = scores.unsqueeze(2) * topk_p.view(-1, topk, topk)
+    # (b, topk, 1) * (b, topk, topk) -> (b, topk, topk)
+
+    topk_cs_p, topk_cs_index = fast_topk(
+        expand_scores.flatten(start_dim=1), topk, dim=-1
+    )  # (b, topk)
+
+    topk_index = topk_index.view(-1, topk_sq)
+    input_ids = torch.gather(topk_index, 1, topk_cs_index).flatten()
+
+    if hidden_states is not None and hidden_states.shape[0] > 0:
+        flat_cs = topk_cs_index.flatten()
+        batch_offsets = torch.arange(
+            0, hidden_states.shape[0], step=topk, device=flat_cs.device
+        )
+        selected_input_index = flat_cs // topk + batch_offsets.repeat_interleave(topk)
+        hidden_states = hidden_states[selected_input_index]
+
+    tree_info = (
+        expand_scores,  # (b, topk, topk)
+        topk_index,  # (b, topk * topk)
+        topk_cs_index + (topk_sq * (i - 1) + topk),  # (b, topk)
+    )
+    return input_ids, hidden_states, topk_cs_p, tree_info
+
+
+def select_top_k_tokens(
+    i: int,
+    topk_p: torch.Tensor,
+    topk_index: torch.Tensor,
+    hidden_states: torch.Tensor,
+    scores: torch.Tensor,
+    topk: int,
+):
+    if i == 0:
+        return _select_top_k_tokens_first(topk_p, topk_index, hidden_states, topk)
+    return _select_top_k_tokens_later(
+        i, topk_p, topk_index, hidden_states, scores, topk
+    )
+
+
+def _sample_simulated_acc_len(
+    simulate_acc_len: float,
+    simulate_acc_method: str,
+    max_len: int,
+) -> int:
+    """Sample a simulated acceptance length in [1, max_len]."""
+    if simulate_acc_method == "multinomial":
+        simulated_values = torch.normal(
+            mean=simulate_acc_len,
+            std=1.0,
+            size=(1,),
+            device="cpu",
+        )
+        # clamp simulated values to be between 1 and max_len
+        simulated_values = torch.clamp(simulated_values, min=1.0, max=max_len)
+        simulate_acc_len = int(simulated_values.round().item())
+    elif simulate_acc_method == "match-expected":
+        # multinomial sampling does not match the expected length
+        # we keep it for the sake of compatibility of existing tests
+        # but it's better to use "match-expected" for the cases that need to
+        # match the expected length, One caveat is that this will only sample
+        # either round down or round up of the expected length
+        simulate_acc_len = max(1.0, min(max_len, simulate_acc_len))
+        lower = int(simulate_acc_len // 1)
+        upper = lower + 1 if lower < max_len else lower
+        if lower == upper:
+            simulate_acc_len = lower
+        else:
+            weight_upper = simulate_acc_len - lower
+            weight_lower = 1.0 - weight_upper
+            probs = torch.tensor([weight_lower, weight_upper], device="cpu")
+            sampled_index = torch.multinomial(probs, num_samples=1)
+            simulate_acc_len = lower if sampled_index == 0 else upper
+    else:
+        raise ValueError(f"Invalid simulate_acc_method: {simulate_acc_method}")
+    return int(simulate_acc_len)
+
+
+def generate_simulated_accept_index(
+    accept_index,
+    predict,
+    num_correct_drafts,
+    candidates,
+    target_predict,
+    bs,
+    spec_steps,
+    simulate_acc_len: float = SIMULATE_ACC_LEN,
+    simulate_acc_method: str = SIMULATE_ACC_METHOD,
+    simulate_acc_token_mode: str = SIMULATE_ACC_TOKEN_MODE,
+):
+    use_real_draft_tokens = simulate_acc_token_mode == "real-draft-token"
+
+    assert simulate_acc_len > 0.0
+    simulate_acc_len = _sample_simulated_acc_len(
+        simulate_acc_len, simulate_acc_method, spec_steps + 1
+    )
+
+    accept_indx_first_col = accept_index[:, 0].view(-1, 1)
+    sim_accept_index = torch.full(
+        (bs, spec_steps + 1), -1, dtype=torch.int32, device=accept_index.device
+    )
+    sim_accept_index[:, :simulate_acc_len] = accept_indx_first_col + torch.arange(
+        simulate_acc_len, device=accept_index.device
+    )
+    num_correct_drafts.fill_(simulate_acc_len - 1)
+
+    if not use_real_draft_tokens:
+        predict.fill_(100)  # some legit token id
+        return sim_accept_index
+
+    # Use the topk=1 draft chain for forced acceptance, then a target-derived bonus.
+    if simulate_acc_len > 1:
+        draft_node_indices = sim_accept_index[:, : simulate_acc_len - 1].long()
+        predict[draft_node_indices] = candidates[:, 1:simulate_acc_len].to(
+            dtype=predict.dtype
+        )
+    bonus_node_indices = sim_accept_index[:, simulate_acc_len - 1].long()
+    predict[bonus_node_indices] = target_predict[:, simulate_acc_len - 1].to(
+        dtype=predict.dtype
+    )
+    return sim_accept_index
+
+
+def traverse_tree(
+    retrieve_next_token: torch.Tensor,
+    retrieve_next_sibling: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    grammar: BaseGrammarObject,
+    allocate_token_bitmask: torch.Tensor,
+    vocab_size: Optional[int] = None,
+):
+    """
+    Traverse the tree constructed by the draft model to generate the logits mask.
+    """
+    assert (
+        retrieve_next_token.shape == retrieve_next_sibling.shape == draft_tokens.shape
+    )
+
+    def dfs(
+        curr: int,
+        retrieve_next_token: torch.Tensor,
+        retrieve_next_sibling: torch.Tensor,
+        parent_pos: int,
+    ):
+        if curr == 0:
+            # the first token generated by the target model, and thus it is always
+            # accepted from the previous iteration
+            is_accepted = True
+        else:
+            parent_bitmask = allocate_token_bitmask[parent_pos]
+            current_token = draft_tokens[curr]
+            if vocab_size and current_token >= vocab_size:
+                is_accepted = False
+            else:
+                # 32 boolean bitmask values are packed into 32-bit integers
+                is_accepted = (
+                    parent_bitmask[current_token // 32] & (1 << (current_token % 32))
+                ) != 0
+
+        if is_accepted:
+            if curr != 0:
+                # Accept the current token
+                grammar.accept_token(int(draft_tokens[curr]))
+            if not grammar.is_terminated():
+                # Generate the bitmask for the current token
+                grammar.fill_vocab_mask(allocate_token_bitmask, curr)
+                if retrieve_next_token[curr] != -1:
+                    # Visit the child node
+                    dfs(
+                        int(retrieve_next_token[curr]),
+                        retrieve_next_token,
+                        retrieve_next_sibling,
+                        curr,
+                    )
+
+            if curr != 0:
+                # Rollback the current token
+                grammar.rollback(1)
+
+        if retrieve_next_sibling[curr] != -1:
+            # Visit the sibling node
+            dfs(
+                int(retrieve_next_sibling[curr]),
+                retrieve_next_token,
+                retrieve_next_sibling,
+                parent_pos,
+            )
+
+    dfs(0, retrieve_next_token, retrieve_next_sibling, -1)
+
+
+def generate_token_bitmask(
+    reqs: List[Req],
+    verify_input: EagleVerifyInput,
+    retrieve_next_token_cpu: torch.Tensor,
+    retrieve_next_sibling_cpu: torch.Tensor,
+    draft_tokens_cpu: torch.Tensor,
+    vocab_size: int,
+):
+    """
+    Generate the logit mask for structured output.
+    Draft model's token can be either valid or invalid with respect to the grammar.
+    We need to perform DFS to
+    1. figure out which tokens are accepted by the grammar.
+    2. if so, what is the corresponding logit mask.
+    """
+
+    num_draft_tokens = draft_tokens_cpu.shape[-1]
+
+    allocate_token_bitmask = None
+    assert len(reqs) == retrieve_next_token_cpu.shape[0]
+    grammar = None
+    for i, req in enumerate(reqs):
+        if req.grammar is not None:
+            if allocate_token_bitmask is None:
+                allocate_token_bitmask = req.grammar.allocate_vocab_mask(
+                    vocab_size=vocab_size,
+                    batch_size=draft_tokens_cpu.numel(),
+                    device="cpu",
+                )
+            grammar = req.grammar
+            s = time.perf_counter()
+            traverse_tree(
+                retrieve_next_token_cpu[i],
+                retrieve_next_sibling_cpu[i],
+                draft_tokens_cpu[i],
+                req.grammar,
+                allocate_token_bitmask[
+                    i * num_draft_tokens : (i + 1) * num_draft_tokens
+                ],
+                vocab_size=vocab_size,
+            )
+            tree_traverse_time = time.perf_counter() - s
+            if tree_traverse_time > TREE_TRAVERSE_TIME_THRESHOLD:
+                logger.warning(
+                    f"Bit mask generation took {tree_traverse_time} seconds with "
+                    f"grammar: {req.grammar}"
+                )
+
+    verify_input.grammar = grammar
+    return allocate_token_bitmask
+
+
+def load_token_map(token_map_path: str) -> List[int]:
+    if not os.path.exists(token_map_path):
+        repo_id = os.path.dirname(token_map_path)
+        file_name = os.path.basename(token_map_path)
+
+        cache_dir = None
+        if envs.SGLANG_USE_MODELSCOPE.get():
+            from modelscope.utils.file_utils import get_model_cache_root
+
+            cached_repo_path = os.path.join(get_model_cache_root(), repo_id)
+            if os.path.exists(cached_repo_path):
+                cache_dir = cached_repo_path
+
+        if cache_dir is None:
+            if envs.SGLANG_USE_MODELSCOPE.get():
+                from modelscope.hub.snapshot_download import (
+                    snapshot_download as download_func,
+                )
+            else:
+                download_func = snapshot_download
+            cache_dir = download_func(
+                repo_id,
+                ignore_patterns=["*.bin", "*.safetensors"],
+            )
+
+        token_map_path = os.path.join(cache_dir, file_name)
+    hot_token_id = torch.load(token_map_path, weights_only=True)
+    return torch.tensor(hot_token_id, dtype=torch.int64)
+
+
+@contextmanager
+def draft_tp_context(tp_group: GroupCoordinator):
+    # Draft model doesn't use dp and has its own tp group.
+    # We disable mscclpp now because it doesn't support 2 comm groups.
+    with patch_tensor_parallel_group(tp_group):
+        yield
+
+
+def spec_stage_span(name: str):
+    """Profiler span for a coarse speculative-decoding stage (``draft`` /
+    ``draft_extend`` / ``verify``).
+    """
+    return profile_range(name)
+
+
+def move_accept_tokens_to_target_kvcache(
+    batch: ScheduleBatch,
+    accept_index: torch.Tensor,
+    num_correct_drafts: torch.Tensor,
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+):
+    """
+    Move accepted tokens (drafts + bonus) to the target KV cache.
+
+    Args:
+        batch: The batch to run.
+        accept_index: The index of the accepted tokens (incl. bonus).
+        num_correct_drafts: Per-req count of correct drafts (excludes bonus);
+            seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot.
+    """
+    bs = len(batch.seq_lens)
+    device = batch.seq_lens.device
+    # accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the
+    # tree exceeds the accepted chain, over-reading accept_index (illegal memory).
+    size = bs * accept_index.shape[1]
+
+    # fill_accept_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
+    maybe_detect_oob(
+        accept_index,
+        -1,
+        batch.out_cache_loc.size(0),
+        "spec v2 move_accept_tokens accept_index",
+    )
+
+    tgt_cache_loc = torch.zeros(
+        size,
+        dtype=torch.int64,
+        device=device,
+    )
+    accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device)
+    if _is_cpu:
+        assign_extend_cache_locs_cpu(
+            batch.req_pool_indices,
+            batch.req_to_token_pool.req_to_token,
+            batch.seq_lens,
+            batch.seq_lens + num_correct_drafts + 1,
+            tgt_cache_loc,
+            batch.req_to_token_pool.req_to_token.shape[1],
+        )
+    else:
+        assign_extend_cache_locs[(bs,)](
+            batch.req_pool_indices,
+            batch.req_to_token_pool.req_to_token,
+            batch.seq_lens,
+            batch.seq_lens + num_correct_drafts + 1,
+            tgt_cache_loc,
+            batch.req_to_token_pool.req_to_token.shape[1],
+            next_power_of_2(bs),
+        )
+    fill_accept_out_cache_loc_func(
+        accept_index,
+        batch.out_cache_loc,
+        accept_out_cache_loc,
+        size,
+    )
+    token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
+        tgt_cache_loc, accept_out_cache_loc
+    )
+
+
+def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
+    """Rebuild mamba track indices from reqs before a TARGET_VERIFY forward.
+
+    Spec batches skip the refresh in prepare_for_decode, and filter/merge
+    null these fields, so they must be rebuilt right before verify. Clearing
+    the mask also keeps a stale extend-time mask from triggering in-forward
+    tracking during TARGET_VERIFY; tracking is done in
+    commit_mamba_states_after_verify instead.
+    """
+    if not get_server_args().enable_mamba_extra_buffer():
+        return
+    set_mamba_track_indices_from_reqs(batch)
+    batch.mamba_track_mask = None
+    batch.mamba_track_seqlens = None
+
+
+def commit_mamba_states_after_verify(
+    target_worker: TpModelWorker,
+    batch: ScheduleBatch,
+    accept_lens: torch.Tensor,
+    accept_index: torch.Tensor,
+    draft_token_num: int,
+) -> None:
+    """Commit accepted per-step mamba states into the persistent caches.
+
+    During TARGET_VERIFY, hybrid linear attention backends keep per-step
+    states in intermediate caches instead of advancing the persistent
+    conv/ssm caches. After acceptance, the state of each request's last
+    accepted step is committed back, plus the interval-crossing state used
+    for prefix-cache tracking (mamba extra_buffer mode).
+
+    No-op for models without mamba-style state or backends without the
+    commit hook.
+    """
+    model_runner = target_worker.model_runner
+    if model_runner.mambaish_config is None:
+        return
+    attn_backend = model_runner.attn_backend
+    if not hasattr(attn_backend, "update_mamba_state_after_mtp_verify"):
+        return
+
+    bs = accept_lens.shape[0]
+    # `accept_lens` already includes the bonus token (drafts + 1 per req).
+    if not batch.forward_mode.is_idle() and accept_index.numel() > 0:
+        accept_indices_offset = torch.arange(
+            0,
+            bs * draft_token_num,
+            step=draft_token_num,
+            dtype=accept_lens.dtype,
+            device=accept_lens.device,
+        )
+        req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
+        # Per-req tree step of the last accepted node, i.e. the step whose
+        # mamba state to commit; reduces to accept_lens - 1 for topk == 1.
+        last_correct_step_indices = (
+            accept_index[req_idx, (accept_lens - 1).to(torch.int64)]
+            - accept_indices_offset
+        )
+
+        if batch.mamba_track_indices is not None:
+            # If after verify, the request's seq_lens has crossed a mamba track interval,
+            # we need to update the mamba state for the request at the crossing point.
+            seq_lens_pre_verify = batch.seq_lens
+            seq_lens_post_verify = batch.seq_lens + accept_lens
+            mamba_track_interval = get_server_args().mamba_track_interval
+            to_track_mask = (
+                seq_lens_pre_verify // mamba_track_interval
+                != seq_lens_post_verify // mamba_track_interval
+            )
+            tracking_point = (
+                seq_lens_post_verify // mamba_track_interval * mamba_track_interval
+            )
+            to_track_ith = torch.clamp(
+                tracking_point - seq_lens_pre_verify - 1, min=0
+            ).to(torch.int64)
+            candidate_track_steps = (
+                accept_index[req_idx, to_track_ith] - accept_indices_offset
+            )
+            mamba_steps_to_track = torch.where(
+                to_track_mask,
+                candidate_track_steps,
+                torch.full_like(candidate_track_steps, -1),
+            )
+        else:
+            mamba_steps_to_track = None
+
+        attn_backend.update_mamba_state_after_mtp_verify(
+            last_correct_step_indices=last_correct_step_indices,
+            mamba_track_indices=batch.mamba_track_indices,
+            mamba_steps_to_track=mamba_steps_to_track,
+            model=model_runner.model,
+        )
+
+
+def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
+    """eagle/ngram share a stateless free function; dflash keeps stateful
+    prep on its draft input -- the dispatcher routes.
+    """
+    if batch.spec_algorithm.is_dflash_family():
+        batch.spec_info.prepare_for_decode(batch)
+    else:
+        from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode
+
+        eagle_prepare_for_decode(batch)

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -0,0 +1,248 @@
+import contextlib
+import logging
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.moe.utils import speculative_moe_backend_context
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.adaptive_runtime_state import (
+    AdaptiveController,
+)
+from sglang.srt.speculative.eagle_utils import default_tree_mask_mode
+from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import draft_tp_context
+from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
+
+if is_cuda():
+    from sgl_kernel import segment_packbits  # noqa: F401
+
+logger = logging.getLogger(__name__)
+SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
+
+
+def _get_plan_stream(
+    device: str,
+) -> Tuple[any, contextlib.AbstractContextManager]:
+    if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
+        plan_stream = torch.get_device_module(device).Stream()
+        plan_stream_ctx = torch.get_device_module(device).stream(plan_stream)
+        return plan_stream, plan_stream_ctx
+    else:
+        return None, contextlib.nullcontext()
+
+
+class StandaloneDraftWorker(EagleDraftWorker):
+    """Custom EagleDraftWorker that doesn't share embeddings/lm_head with target model."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: int,
+        moe_ep_rank: int,
+        attn_cp_rank: int,
+        moe_dp_rank: int,
+        nccl_port: int,
+        target_worker: TpModelWorker,
+    ):
+        # copy args
+        self.server_args = server_args
+        self.gpu_id = gpu_id
+        self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.moe_ep_rank = moe_ep_rank
+        self.nccl_port = nccl_port
+        self.target_worker = target_worker
+        self.attn_cp_rank = attn_cp_rank
+        self.moe_dp_rank = moe_dp_rank
+
+        # Args for easy access
+        self.device = server_args.device
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        # Pre-allocated constants for the topk=1 chain fast path in draft_forward.
+        self._topk1_parents_prealloc = None
+        self._topk1_score_indices_prealloc = None
+        self._rebuild_topk1_chain_buffers()
+
+        # Set constant
+        from sglang.srt.speculative.eagle_info import EagleDraftInput
+
+        EagleDraftInput.ALLOC_LEN_PER_DECODE = max(
+            self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
+        )
+
+        # Load draft model weights only.
+        with empty_context():
+            self.draft_worker = TpModelWorker(
+                server_args=server_args,
+                gpu_id=gpu_id,
+                tp_rank=tp_rank,
+                pp_rank=0,  # spec workers don't support pipeline parallelism
+                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,
+                is_draft_worker=True,
+            )
+
+        # Alias for better readability
+        self.draft_runner = self.draft_worker.model_runner
+        self.draft_tp_context = (
+            draft_tp_context if server_args.enable_dp_attention else empty_context
+        )
+        self.tree_mask_mode = default_tree_mask_mode()
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+        # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
+        self.index_share_for_mtp_iteration = (
+            getattr(
+                self.draft_runner.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
+        self.dsa_index_topk = None
+        self.seed_dsa_topk_from_draft_extend = False
+        self.dsa_extend_topk_buf = None
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        """Standalone: allocate pools without sharing embeddings."""
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.draft_worker.alloc_memory_pool(
+            memory_pool_config=memory_pool_config,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+        self.init_token_map()
+        self.init_lm_head()
+
+    def init_attention_backends(self):
+        with self.draft_tp_context(
+            self.draft_runner.tp_group
+        ), speculative_moe_backend_context():
+            super().init_attention_backends()
+
+    def init_cuda_graphs(self):
+        with self.draft_tp_context(
+            self.draft_runner.tp_group
+        ), speculative_moe_backend_context():
+            super().init_cuda_graphs()
+
+    def init_lm_head(self):
+        """Override to prevent sharing embeddings and lm_head with target model."""
+        # For standalone worker, we don't share embeddings and lm_head
+        # The draft model uses its own embeddings and lm_head
+        pass
+
+
+class StandaloneWorkerV2(EAGLEWorkerV2):
+
+    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,
+    ):
+        # Parse arguments
+        self.server_args = server_args
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self._target_worker = target_worker
+        self.page_size = server_args.page_size
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        # Override the context length of the draft model to be the same as the target model.
+        server_args.override(
+            "spec_worker.match_target_context_length",
+            context_length=target_worker.model_runner.model_config.context_len,
+        )
+
+        # Create our custom draft worker that doesn't share embeddings/lm_head
+        self._draft_worker = StandaloneDraftWorker(
+            server_args,
+            gpu_id,
+            tp_rank,
+            dp_rank,
+            moe_ep_rank,
+            attn_cp_rank,
+            moe_dp_rank,
+            nccl_port,
+            target_worker,
+        )
+
+        self._validate_vocab_compatibility(
+            target_vocab_size=target_worker.model_runner.model_config.vocab_size,
+            target_tokenizer=target_worker.tokenizer,
+        )
+
+        # Some dummy tensors
+        self.num_new_pages_per_topk = torch.empty(
+            (), dtype=torch.int64, device=self.device
+        )
+        self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
+
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+
+        # TODO: Adaptive speculative
+        self.adaptive_controller: Optional[AdaptiveController] = None
+
+    def _validate_vocab_compatibility(
+        self,
+        target_vocab_size: int,
+        target_tokenizer,
+    ) -> None:
+        """Raise ValueError if the draft and target vocabularies are incompatible."""
+        draft_vocab_size = self._draft_worker.draft_runner.model_config.vocab_size
+        draft_tokenizer = self._draft_worker.draft_worker.tokenizer
+        if target_vocab_size != draft_vocab_size:
+            raise ValueError(
+                f"STANDALONE speculative decoding requires the draft model to share the "
+                f"same vocabulary as the target model, but got "
+                f"target vocab_size={target_vocab_size} and "
+                f"draft vocab_size={draft_vocab_size}. "
+                f"Use a draft model with a matching vocabulary, or a speculative "
+                f"algorithm that supports heterogeneous vocabularies."
+            )
+        if (
+            target_tokenizer is not None
+            and draft_tokenizer is not None
+            and hasattr(target_tokenizer, "get_vocab")
+            and hasattr(draft_tokenizer, "get_vocab")
+            and target_tokenizer.get_vocab() != draft_tokenizer.get_vocab()
+        ):
+            raise ValueError(
+                "STANDALONE speculative decoding requires the draft model to share the "
+                "same vocabulary as the target model, but the two tokenizers have "
+                "different token-to-id mappings even though their vocab sizes match. "
+                "Use a draft model with a matching vocabulary, or a speculative "
+                "algorithm that supports heterogeneous vocabularies."
+            )