chore: import upstream snapshot with attribution

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