[Spec] Add DSpark: confidence-scheduled speculative decoding (#30261)

Co-authored-by: sglang-bot <232288953+sglang-bot@users.noreply.github.com>
Co-authored-by: Claude Code <noreply@anthropic.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

python/sglang/benchmark/dspark_sts_fit.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import argparse
+import glob
+import logging
+import math
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_sts import (
+    DSparkStsCalibration,
+)
+
+logger = logging.getLogger(__name__)
+
+_EPS_PROB = 1e-8
+
+
+def default_temperature_grid() -> torch.Tensor:
+    return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41)
+
+
+def expected_calibration_error(
+    *,
+    probs: torch.Tensor,
+    targets: torch.Tensor,
+    num_bins: int,
+) -> float:
+    probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB)
+    targets = targets.reshape(-1).to(torch.float64)
+    total = probs.numel()
+    if total == 0:
+        return float("nan")
+    bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1)
+    count = torch.zeros(num_bins, dtype=torch.float64)
+    pred_sum = torch.zeros(num_bins, dtype=torch.float64)
+    target_sum = torch.zeros(num_bins, dtype=torch.float64)
+    count.scatter_add_(0, bin_index, torch.ones_like(probs))
+    pred_sum.scatter_add_(0, bin_index, probs)
+    target_sum.scatter_add_(0, bin_index, targets)
+    denom = count.clamp_min(1.0)
+    bin_error = (pred_sum / denom - target_sum / denom).abs()
+    return float((bin_error * count).sum().item() / total)
+
+
+def fit_sts_temperatures(
+    *,
+    logits: torch.Tensor,
+    prefix_mask: torch.Tensor,
+    grid: torch.Tensor,
+    num_bins: int = 15,
+) -> dict[str, list[float]]:
+    logits = logits.to(torch.float64)
+    prefix_mask = prefix_mask.to(torch.float64)
+    num_samples, gamma = logits.shape
+    if num_samples == 0:
+        raise ValueError("fit_sts_temperatures requires at least one sample.")
+    grid_values = grid.to(torch.float64).tolist()
+
+    temperatures: list[float] = []
+    ece_before: list[float] = []
+    ece_after: list[float] = []
+
+    survival_at_one = torch.ones(num_samples, dtype=torch.float64)
+    survival_fitted = torch.ones(num_samples, dtype=torch.float64)
+    for position in range(gamma):
+        position_logits = logits[:, position]
+        position_target = prefix_mask[:, position]
+
+        survival_at_one = survival_at_one * torch.sigmoid(position_logits)
+        ece_before.append(
+            expected_calibration_error(
+                probs=survival_at_one,
+                targets=position_target,
+                num_bins=num_bins,
+            )
+        )
+
+        best_temperature = grid_values[0]
+        best_survival = survival_fitted * torch.sigmoid(
+            position_logits / best_temperature
+        )
+        best_ece = expected_calibration_error(
+            probs=best_survival, targets=position_target, num_bins=num_bins
+        )
+        for temperature in grid_values[1:]:
+            candidate_survival = survival_fitted * torch.sigmoid(
+                position_logits / temperature
+            )
+            candidate_ece = expected_calibration_error(
+                probs=candidate_survival,
+                targets=position_target,
+                num_bins=num_bins,
+            )
+            if candidate_ece < best_ece:
+                best_ece = candidate_ece
+                best_temperature = temperature
+                best_survival = candidate_survival
+
+        temperatures.append(float(best_temperature))
+        ece_after.append(float(best_ece))
+        survival_fitted = best_survival
+
+    return {
+        "temperatures": temperatures,
+        "ece_before": ece_before,
+        "ece_after": ece_after,
+    }
+
+
+def load_collected_shards(*, data_glob: str) -> tuple[torch.Tensor, torch.Tensor]:
+    shard_paths = sorted(glob.glob(data_glob))
+    if not shard_paths:
+        raise ValueError(f"No STS data shards matched {data_glob!r}.")
+
+    logits_shards: list[torch.Tensor] = []
+    prefix_mask_shards: list[torch.Tensor] = []
+    shard_gamma: Optional[int] = None
+    for shard_path in shard_paths:
+        shard = torch.load(shard_path, map_location="cpu")
+        shard_logits = shard["logits"]
+        shard_prefix_mask = shard["prefix_mask"]
+        if shard_logits.shape != shard_prefix_mask.shape:
+            raise ValueError(
+                f"Shard {shard_path!r} logits / prefix_mask shape mismatch: "
+                f"{tuple(shard_logits.shape)} vs {tuple(shard_prefix_mask.shape)}."
+            )
+        if shard_gamma is None:
+            shard_gamma = int(shard_logits.shape[1])
+        elif int(shard_logits.shape[1]) != shard_gamma:
+            raise ValueError(
+                f"Shard {shard_path!r} gamma {int(shard_logits.shape[1])} disagrees "
+                f"with earlier shards' gamma {shard_gamma}."
+            )
+        logits_shards.append(shard_logits)
+        prefix_mask_shards.append(shard_prefix_mask)
+
+    return torch.cat(logits_shards, dim=0), torch.cat(prefix_mask_shards, dim=0)
+
+
+def fit(
+    *,
+    data_glob: str,
+    out: Path,
+    num_bins: int = 15,
+    gamma: Optional[int] = None,
+) -> None:
+    logits, prefix_mask = load_collected_shards(data_glob=data_glob)
+    resolved_gamma = int(logits.shape[1])
+    if gamma is not None and gamma != resolved_gamma:
+        raise ValueError(
+            f"Collected shards have gamma={resolved_gamma} but --gamma={gamma}."
+        )
+    num_samples = int(logits.shape[0])
+
+    result = fit_sts_temperatures(
+        logits=logits,
+        prefix_mask=prefix_mask,
+        grid=default_temperature_grid(),
+        num_bins=num_bins,
+    )
+    calibration = DSparkStsCalibration(
+        temperatures=result["temperatures"],
+        dataset=data_glob,
+        num_samples=num_samples,
+        ece_before=result["ece_before"],
+        ece_after=result["ece_after"],
+    )
+    out.write_text(calibration.to_json(), encoding="utf-8")
+
+    print(
+        f"Fit STS temperatures over {num_samples} samples (gamma={resolved_gamma}) "
+        f"-> {out}"
+    )
+    print("pos  temperature  ece_before  ece_after")
+    for position in range(resolved_gamma):
+        print(
+            f"{position:>3}  {result['temperatures'][position]:>11.4f}  "
+            f"{result['ece_before'][position]:>10.4f}  "
+            f"{result['ece_after'][position]:>9.4f}"
+        )
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser(
+        description="Fit DSpark Sequential Temperature Scaling (STS) calibration "
+        "temperatures from collected confidence shards."
+    )
+    parser.add_argument(
+        "--data-glob",
+        required=True,
+        help="Glob of collected .pt shards, each a dict with [n, gamma] "
+        "'logits' and 'prefix_mask' tensors.",
+    )
+    parser.add_argument(
+        "--out",
+        required=True,
+        type=Path,
+        help="Output STS calibration JSON path.",
+    )
+    parser.add_argument(
+        "--num-bins",
+        type=int,
+        default=15,
+        help="Number of equal-width ECE bins.",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=int,
+        default=None,
+        help="Optional gamma override to validate the shards against.",
+    )
+    args = parser.parse_args()
+
+    fit(
+        data_glob=args.data_glob,
+        out=args.out,
+        num_bins=args.num_bins,
+        gamma=args.gamma,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/benchmark/serving.py

@@ -109,6 +109,9 @@ class RequestFuncOutput:
     cached_tokens: int = 0
     cached_tokens_details: Optional[Dict[str, Any]] = None
     spec_accept_length: float = 0.0
+    spec_cap_length: float = 0.0
+    spec_block_accept_length: float = 0.0
+    spec_cap_lens_histogram: List[int] = field(default_factory=list)
 
     @staticmethod
     def init_new(request_func_input: RequestFuncInput):
@@ -483,6 +486,15 @@ async def async_request_openai_chat_completions(
                         output.spec_accept_length = (
                             _meta_info.get("spec_accept_length", 0.0) or 0.0
                         )
+                        output.spec_cap_length = (
+                            _meta_info.get("spec_cap_length", 0.0) or 0.0
+                        )
+                        output.spec_block_accept_length = (
+                            _meta_info.get("spec_block_accept_length", 0.0) or 0.0
+                        )
+                        output.spec_cap_lens_histogram = (
+                            _meta_info.get("spec_cap_lens_histogram", []) or []
+                        )
                         if getattr(args, "cache_report", False):
                             _extract_cache_from_sglext(response_json, output)
                     else:

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -53,9 +53,11 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None
         )
 
     if server_args.speculative_algorithm is not None:
-        assert (
-            server_args.speculative_algorithm == "EAGLE"
-        ), f"Only EAGLE speculative algorithm is supported for {model_arch}"
+        assert server_args.speculative_algorithm in (
+            "EAGLE",
+            "DSPARK",
+        ), f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}"
+        if server_args.speculative_algorithm == "EAGLE":
             assert (
                 server_args.speculative_eagle_topk == 1
             ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}"

python/sglang/srt/arg_groups/speculative_hook.py

@@ -265,6 +265,171 @@ def _handle_dflash(server_args: ServerArgs) -> None:
         )
 
 
+def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
+    from sglang.srt.speculative.dspark_components.dspark_config import (
+        checkpoint_bundles_dspark_draft,
+    )
+
+    return checkpoint_bundles_dspark_draft(server_args.get_model_config().hf_config)
+
+
+def _handle_dspark(server_args: ServerArgs) -> None:
+    if not server_args.device.startswith("cuda"):
+        raise ValueError("DSpark speculative decoding only supports CUDA device.")
+
+    if server_args.enable_dp_attention:
+        if not server_args.enable_dp_lm_head:
+            raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.")
+        if server_args.moe_a2a_backend != "none":
+            raise ValueError(
+                "DSpark with dp attention only supports the built-in TP MoE "
+                f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}."
+            )
+        if server_args.attn_cp_size > 1:
+            raise ValueError(
+                "DSpark with dp attention does not support context parallel "
+                f"(attn_cp_size={server_args.attn_cp_size})."
+            )
+        if (
+            server_args.speculative_moe_a2a_backend is not None
+            and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend
+        ):
+            raise ValueError(
+                "DSpark ignores --speculative-moe-a2a-backend; with dp attention it "
+                f"must match the target moe_a2a_backend={server_args.moe_a2a_backend!r} "
+                f"(got {server_args.speculative_moe_a2a_backend!r})."
+            )
+
+    if server_args.pp_size != 1:
+        raise ValueError(
+            "Currently DSpark speculative decoding only supports pp_size == 1."
+        )
+
+    if server_args.speculative_draft_model_path is None:
+        if _target_checkpoint_bundles_dspark_draft(server_args):
+            server_args.speculative_draft_model_path = server_args.model_path
+            server_args.speculative_draft_model_revision = server_args.revision
+            logger.info(
+                "DSpark draft weights are bundled in the target checkpoint; "
+                "defaulting --speculative-draft-model-path to --model-path (%s).",
+                server_args.model_path,
+            )
+        else:
+            raise ValueError(
+                "DSpark dense speculative decoding requires setting "
+                "--speculative-draft-model-path."
+            )
+
+    if server_args.speculative_num_steps is None:
+        server_args.speculative_num_steps = 1
+    elif int(server_args.speculative_num_steps) != 1:
+        logger.warning(
+            "DSpark only supports speculative_num_steps == 1; overriding speculative_num_steps=%s to 1.",
+            server_args.speculative_num_steps,
+        )
+        server_args.speculative_num_steps = 1
+
+    if server_args.speculative_eagle_topk is None:
+        server_args.speculative_eagle_topk = 1
+    elif int(server_args.speculative_eagle_topk) != 1:
+        logger.warning(
+            "DSpark only supports speculative_eagle_topk == 1; overriding speculative_eagle_topk=%s to 1.",
+            server_args.speculative_eagle_topk,
+        )
+        server_args.speculative_eagle_topk = 1
+
+    gamma: Optional[int] = None
+    if server_args.speculative_dspark_block_size is not None:
+        if int(server_args.speculative_dspark_block_size) <= 0:
+            raise ValueError(
+                "DSpark requires --speculative-dspark-block-size to be positive, "
+                f"got {server_args.speculative_dspark_block_size}."
+            )
+        gamma = int(server_args.speculative_dspark_block_size)
+    else:
+        from sglang.srt.speculative.dspark_components.dspark_config import (
+            DEFAULT_DSPARK_GAMMA,
+            read_draft_checkpoint_gamma,
+        )
+
+        try:
+            gamma = read_draft_checkpoint_gamma(server_args=server_args)
+        except Exception as e:
+            logger.warning(
+                "Failed to read DSpark gamma from draft model config; "
+                "cannot cross-check --speculative-num-draft-tokens. Error: %s",
+                e,
+            )
+        if gamma is None and server_args.speculative_num_draft_tokens is None:
+            gamma = DEFAULT_DSPARK_GAMMA
+            logger.warning(
+                "DSpark gamma is not set; defaulting to %d.",
+                gamma,
+            )
+
+    if gamma is not None:
+        verify_window = int(gamma) + 1
+        if (
+            server_args.speculative_num_draft_tokens is not None
+            and int(server_args.speculative_num_draft_tokens) != verify_window
+        ):
+            raise ValueError(
+                "DSpark speculative_num_draft_tokens must equal gamma + 1 "
+                f"(= {verify_window} for gamma={gamma}), but got "
+                f"speculative_num_draft_tokens={server_args.speculative_num_draft_tokens}."
+            )
+        server_args.speculative_num_draft_tokens = verify_window
+
+    if server_args.speculative_num_draft_tokens is None:
+        raise ValueError(
+            "DSpark could not resolve speculative_num_draft_tokens; set "
+            "--speculative-dspark-block-size (= gamma)."
+        )
+    if int(server_args.speculative_num_draft_tokens) < 2:
+        raise ValueError(
+            "DSpark speculative_num_draft_tokens must be >= 2 (= gamma + 1), "
+            f"got {server_args.speculative_num_draft_tokens}."
+        )
+
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    if server_args.enable_mixed_chunk:
+        server_args.enable_mixed_chunk = False
+        logger.warning(
+            "Mixed chunked prefill is disabled because of using dspark speculative decoding."
+        )
+
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+
+    ragged_mode = read_ragged_verify_mode()
+    if (
+        server_args.speculative_dspark_align_verify_tokens_to_graph_tier
+        and ragged_mode is not RaggedVerifyMode.COMPACT
+    ):
+        logger.warning(
+            "--speculative-dspark-align-verify-tokens-to-graph-tier only takes "
+            "effect with SGLANG_RAGGED_VERIFY_MODE=compact (got %r); it will be "
+            "a no-op.",
+            ragged_mode.value,
+        )
+    if (
+        server_args.speculative_dspark_sps_table_path
+        and ragged_mode is RaggedVerifyMode.STATIC
+    ):
+        logger.warning(
+            "--speculative-dspark-sps-table-path feeds the ragged-verify budget "
+            "scheduler, which is off under SGLANG_RAGGED_VERIFY_MODE=static; it "
+            "will be a no-op."
+        )
+
+
 def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
     """Resolve `speculative_draft_attention_backend` to a final, supported value.
 

python/sglang/srt/configs/model_config.py

@@ -122,6 +122,7 @@ def is_deepseek_v4(config) -> bool:
     return _hf_arch(config) in (
         "DeepseekV4ForCausalLM",
         "DeepseekV4ForCausalLMNextN",
+        "DeepseekV4ForCausalLMDSpark",
     )
 
 
@@ -247,12 +248,14 @@ class ModelConfig:
         language_only: bool = False,
         disable_hybrid_swa_memory: bool = False,
         model_config_parser: str = "auto",
+        speculative_algorithm: Optional[str] = None,
     ) -> None:
         # Parse args
         self.model_path = model_path
         self.revision = revision
         self.quantization = quantization
         self.is_draft_model = is_draft_model
+        self.speculative_algorithm = speculative_algorithm
         self.model_impl = model_impl
         self.sampling_defaults = sampling_defaults
         self.quantize_and_serve = quantize_and_serve
@@ -528,6 +531,7 @@ class ModelConfig:
             is_draft_model=is_draft_model,
             disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory,
             model_config_parser=server_args.model_config_parser,
+            speculative_algorithm=server_args.speculative_algorithm,
             **kwargs,
         )
 
@@ -547,6 +551,21 @@ class ModelConfig:
             is_draft_model
             and self.hf_config.architectures[0] == "DeepseekV4ForCausalLM"
         ):
+            from sglang.srt.speculative.dspark_components.dspark_config import (
+                checkpoint_bundles_dspark_draft,
+            )
+
+            # A dspark-bundled checkpoint may also carry MTP layers; the
+            # selected algorithm decides which draft arch to load.
+            if checkpoint_bundles_dspark_draft(self.hf_config) and (
+                self.speculative_algorithm in (None, "DSPARK")
+            ):
+                self.hf_config.architectures[0] = "DeepseekV4ForCausalLMDSpark"
+                logger.info(
+                    "Draft checkpoint bundles a DSpark head; loading draft arch "
+                    "DeepseekV4ForCausalLMDSpark."
+                )
+            else:
                 self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN"
                 self.hf_config.num_nextn_predict_layers = 1
 
@@ -636,7 +655,12 @@ class ModelConfig:
             logger.info(f"Hybrid swa model: {self.hf_config.architectures=}")
 
             self.is_deepseek_v4_arch = any(
-                arch in ["DeepseekV4ForCausalLM", "DeepseekV4ForCausalLMNextN"]
+                arch
+                in [
+                    "DeepseekV4ForCausalLM",
+                    "DeepseekV4ForCausalLMNextN",
+                    "DeepseekV4ForCausalLMDSpark",
+                ]
                 for arch in self.hf_config.architectures
             )
 
@@ -786,6 +810,7 @@ class ModelConfig:
         elif (
             "DeepseekV4ForCausalLM" in self.hf_config.architectures
             or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures
+            or "DeepseekV4ForCausalLMDSpark" in self.hf_config.architectures
         ):
             self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
             self.qk_nope_head_dim = self.hf_config.head_dim - self.qk_rope_head_dim
@@ -1715,6 +1740,7 @@ multimodal_model_archs = [
 piecewise_cuda_graph_disabled_model_archs = [
     "DeepseekV4ForCausalLM",
     "DeepseekV4ForCausalLMNextN",
+    "DeepseekV4ForCausalLMDSpark",
     "Qwen3NextForCausalLM",
     "BailingMoeV2_5ForCausalLM",
     "LLaDAModelLM",
@@ -1844,6 +1870,7 @@ def is_hybrid_swa_model(
         "Llama4ForConditionalGeneration",
         "DeepseekV4ForCausalLM",
         "DeepseekV4ForCausalLMNextN",
+        "DeepseekV4ForCausalLMDSpark",
         "GptOssForCausalLM",
         *MIMO_V2_MODEL_ARCHS,
         "MiMoV2MTP",

python/sglang/srt/disaggregation/decode.py

@@ -143,6 +143,10 @@ class DecodeReqToTokenPool:
             )
 
         self.free_slots = list(range(1, self._alloc_size))
+        # Slot-reuse generation counter; mirrors ReqToTokenPool. Required even
+        # here: HybridMambaDecodeReqToTokenPool borrows this __init__ while
+        # inheriting ReqToTokenPool.alloc, which bumps it.
+        self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64)
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -171,6 +175,7 @@ class DecodeReqToTokenPool:
         for r in reqs:
             if r.req_pool_idx is None:
                 r.req_pool_idx = select_index[offset]
+                self.req_generation[r.req_pool_idx] += 1
                 offset += 1
         return [r.req_pool_idx for r in reqs]
 
@@ -181,6 +186,7 @@ class DecodeReqToTokenPool:
 
     def clear(self):
         self.free_slots = list(range(1, self._alloc_size))
+        self.req_generation.zero_()
 
 
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):

python/sglang/srt/environ.py

@@ -215,6 +215,7 @@ class Envs:
     # Logging Options
     SGLANG_LOG_GC = EnvBool(False)
     SGLANG_LOG_FORWARD_ITERS = EnvBool(False)
+    SGLANG_LOG_DECODE_GRAPH_KEY = EnvBool(False)
     SGLANG_LOG_MS = EnvBool(False)
     SGLANG_LOG_REQUEST_EXCEEDED_MS = EnvInt(-1)
     SGLANG_LOG_REQUEST_HEADERS = EnvTuple(tuple())
@@ -262,6 +263,20 @@ class Envs:
     SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE = EnvBool(False)
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
+    SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER = EnvBool(False)
+    SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS = EnvBool(False)
+    SGLANG_DSPARK_DEBUG_DUMP = EnvTuple(tuple())
+    SGLANG_DSPARK_LOG_SPS_PRED_INTERVAL = EnvInt(0)
+    SGLANG_DSPARK_STS_COLLECT_PATH = EnvStr("")
+    SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH = EnvStr("")
+    SGLANG_DSPARK_BLOCK_ACCEPT_ONLINE_INTERVAL = EnvInt(0)
+    SGLANG_DSPARK_ENABLE_SPS_RECORD = EnvBool(False)
+    SGLANG_DSPARK_FAST_KERNEL = EnvBool(True)
+    SGLANG_DSPARK_FP32_LM_HEAD = EnvBool(False)
+    SGLANG_DSPARK_FAST_SAMPLING = EnvBool(True)
+    SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True)
+    SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True)
+    SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True)
     SGLANG_DEBUG_REVERT_PR = EnvInt(0)
     SGLANG_PHASE_CHECKER_DEBUG = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
@@ -705,6 +720,9 @@ class Envs:
 
     # Spec Config
     SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True)
+    SGLANG_RAGGED_VERIFY_MODE = EnvStr("static")
+    SGLANG_DSPARK_CONFIDENCE_RELAY_LAG_STEPS = EnvInt(2)
+    SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE = EnvBool(False)
     # Skip draft_extend while adaptive spec is at steps=0 (drafting disabled).
     # Saves the per-step draft forward, but the draft KV goes stale: an upshift
     # back to steps>0 starts from a cold draft state (low accept until it recovers).

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -42,6 +42,8 @@ class AttentionBackend(ABC):
     prefill_attention_backend_str: Optional[str] = None
     decode_attention_backend_str: Optional[str] = None
 
+    supports_ragged_verify_graph: bool = False
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Eager entry point. Default = ``_out_graph(fb) + _in_graph(fb)``.
 

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -21,6 +21,11 @@ import torch.nn.functional as F
 from sglang.jit_kernel.dsv4.online_c128_mtp import OnlineC128MTPController
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import (
+    BuildCausalSwaPageIndices,
+    BuildPageTablePositions,
+    ExpandPrefillCausally,
+)
 from sglang.srt.layers.attention.dsv4.compressor_v2 import (
     CompressorBackendMixin,
     FusedCompressMetadata,
@@ -49,7 +54,20 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
+from sglang.srt.speculative.dspark_components.kernels.dspark_attn_metadata import (
+    BuildBlockSeqLensCausal,
+    BuildDsparkSwaPageIndices,
+    ComputeDsparkWindowGather,
+)
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
+from sglang.srt.speculative.ragged_verify import (
+    RaggedVerifyMode,
+    compute_ragged_extend_lengths,
+    compute_target_verify_graph_key,
+    compute_uniform_extend_lengths,
+    read_ragged_verify_mode,
+    resolve_ragged_verify_layout,
+)
 from sglang.srt.utils import ceil_align, is_xpu
 from sglang.srt.utils.common import is_sm120_supported
 
@@ -58,6 +76,7 @@ if TYPE_CHECKING:
 
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
 
 _is_sm120 = is_sm120_supported()
 _is_xpu = is_xpu()
@@ -412,6 +431,10 @@ class DSV4RawVerifyMetadata:
     seq_lens_cpu: Optional[List[int]] = None
     c128_compress_metadata: Optional[FusedCompressMetadata] = None
 
+    extend_start_loc: Optional[torch.Tensor] = None
+    verify_lens: Optional[torch.Tensor] = None
+    total_verify_tokens: int = 0
+
     def copy_(self, other: DSV4RawVerifyMetadata):
         self.req_pool_indices.copy_(other.req_pool_indices)
         self.seq_lens.copy_(other.seq_lens)
@@ -423,6 +446,10 @@ class DSV4RawVerifyMetadata:
             self.c128_compress_metadata, other.c128_compress_metadata
         )
 
+        self.extend_start_loc = other.extend_start_loc
+        self.verify_lens = other.verify_lens
+        self.total_verify_tokens = other.total_verify_tokens
+
 
 @dataclass
 class DSV4RawDecodeMetadata:
@@ -456,7 +483,7 @@ class DeepseekV4AttnBackend(
     AttentionBackend, C4IndexerBackendMixin, CompressorBackendMixin
 ):
     use_captured_forward_metadata_for_breakable_cuda_graph: bool = True
-
+    supports_ragged_verify_graph: bool = True
     needs_cpu_seq_lens: bool = False
 
     def __init__(
@@ -504,6 +531,21 @@ class DeepseekV4AttnBackend(
         self.speculative_num_draft_tokens: int = (
             model_runner.server_args.speculative_num_draft_tokens
         )
+        if self.speculative_num_draft_tokens is not None:
+            # Persistent target-verify metadata buffers. Allocated here (not
+            # lazily) so they are ordinary tensors: the first touch of a lazy
+            # buffer would inherit the caller's context, and a creation inside
+            # an inference_mode forward would forbid the in-place updates the
+            # graph-capture path performs outside inference mode.
+            num_reqs = self.req_to_token.shape[0]
+            self.extend_seq_lens_buffer = torch.full(
+                (num_reqs,),
+                self.speculative_num_draft_tokens,
+                **self.cuda_int32_kwargs,
+            )
+            self.extend_start_loc_buffer = torch.zeros(
+                num_reqs, **self.cuda_int32_kwargs
+            )
         self.speculative_step_id = speculative_step_id
         self.forward_metadata: Union[
             DSV4Metadata,
@@ -514,14 +556,56 @@ class DeepseekV4AttnBackend(
         # Draft-extend and online-c128 verify metadata are host-planned, so
         # spec runs keep the relay publish (the mirror only exists under
         # spec-v2; without spec the flag has no consumer either way).
-        if model_runner.server_args.speculative_algorithm is not None:
+        # DSPARK is the exception: its draft path carries its own host lens
+        # (reserved_seq_lens_cpu) and its verify prep is device-side.
+        spec_alg = model_runner.spec_algorithm
+        if not spec_alg.is_none() and not spec_alg.is_dspark():
             self.needs_cpu_seq_lens = True
         self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
 
+        self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark()
+
     def _move_to_device(self, x: List[int]) -> torch.Tensor:
         pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True)
         return pin_tensor.to(self.device, non_blocking=True)
 
+    def _resolve_verify_layout(
+        self,
+        forward_batch: ForwardBatch,
+        bs: int,
+    ) -> Optional[RaggedVerifyLayout]:
+        layout = resolve_ragged_verify_layout(forward_batch)
+        if layout is None:
+            return None
+        if read_ragged_verify_mode() is not RaggedVerifyMode.COMPACT:
+            return None
+        if get_parallel().attn_cp_size > 1:
+            raise NotImplementedError(
+                "DSV4 ragged verify does not support context parallel (CP); "
+                "set SGLANG_RAGGED_VERIFY_MODE off for CP runs."
+            )
+        if self.online_c128_mtp.enabled():
+            raise NotImplementedError(
+                "DSV4 ragged verify does not support online c128 MTP; "
+                "set SGLANG_RAGGED_VERIFY_MODE off or disable online compress."
+            )
+        # Layout invariants (verify_lens >= 1, total == sum) are enforced in
+        # RaggedVerifyLayout.__post_init__; don't re-check the device tensor
+        # here -- that would D2H-sync the host-free verify prep path.
+        layout = layout.padded_to_bucket(padded_bs=bs)
+        return layout
+
+    def _target_verify_graph_key(
+        self,
+        bs: int,
+        ragged_layout: Optional[RaggedVerifyLayout],
+    ) -> Tuple[int, int]:
+        return compute_target_verify_graph_key(
+            bs=bs,
+            num_draft_tokens=self.speculative_num_draft_tokens,
+            ragged_layout=ragged_layout,
+        )
+
     def _make_target_verify_c128_metadata(
         self,
         req_pool_indices: torch.Tensor,
@@ -623,6 +707,7 @@ class DeepseekV4AttnBackend(
         need_compress: bool = True,
         use_prefill_cuda_graph: bool = False,
         online_c128_state_slot_offset: int = 0,
+        dspark_block_size: Optional[int] = None,
     ) -> DSV4Metadata:
         seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually(
             num_tokens=num_tokens,
@@ -642,6 +727,7 @@ class DeepseekV4AttnBackend(
             out_loc=out_cache_loc,
             need_compress=need_compress,
             is_prefill=True,
+            dspark_block_size=dspark_block_size,
         )
         indexer_metadata = (
             self.init_forward_metadata_indexer(
@@ -709,17 +795,29 @@ class DeepseekV4AttnBackend(
         out_cache_loc: Optional[torch.Tensor] = None,
         use_prefill_cuda_graph: bool = False,
         online_c128_state_slot_offset: int = 0,
+        ragged_layout: Optional[RaggedVerifyLayout] = None,
     ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]:
         if envs.SGLANG_PREP_IN_CUDA_GRAPH.get():
             assert out_cache_loc is not None
+            bs = len(seq_lens)
             seq_lens_cpu_list = (
                 seq_lens_cpu.tolist() if seq_lens_cpu is not None else None
             )
-            if not hasattr(self, "extend_seq_lens_buffer"):
-                self.extend_seq_lens_buffer = torch.tensor(
-                    [self.speculative_num_draft_tokens] * 1025, device=self.device
+            if ragged_layout is None:
+                self.extend_seq_lens_buffer[:bs].fill_(
+                    self.speculative_num_draft_tokens
                 )
-            extend_seq_lens = self.extend_seq_lens_buffer[: len(seq_lens)]
+                extend_seq_lens = self.extend_seq_lens_buffer[:bs]
+                extend_start_loc = None
+                verify_lens = None
+                total_verify_tokens = self.speculative_num_draft_tokens * bs
+            else:
+                self.extend_seq_lens_buffer[:bs].copy_(ragged_layout.verify_lens)
+                self.extend_start_loc_buffer[:bs].copy_(ragged_layout.extend_start_loc)
+                extend_seq_lens = self.extend_seq_lens_buffer[:bs]
+                extend_start_loc = self.extend_start_loc_buffer[:bs]
+                verify_lens = self.extend_seq_lens_buffer[:bs]
+                total_verify_tokens = ragged_layout.graph_num_tokens
 
             return DSV4RawVerifyMetadata(
                 req_pool_indices=req_pool_indices,
@@ -735,6 +833,9 @@ class DeepseekV4AttnBackend(
                     use_prefill_cuda_graph,
                     online_c128_state_slot_offset,
                 ),
+                extend_start_loc=extend_start_loc,
+                verify_lens=verify_lens,
+                total_verify_tokens=total_verify_tokens,
             )
         else:
             seq_lens_cpu = seq_lens.tolist()
@@ -746,6 +847,7 @@ class DeepseekV4AttnBackend(
                 out_cache_loc=out_cache_loc,
                 use_prefill_cuda_graph=use_prefill_cuda_graph,
                 online_c128_state_slot_offset=online_c128_state_slot_offset,
+                ragged_layout=ragged_layout,
             )
 
     def init_forward_metadata_target_verify_old(
@@ -757,13 +859,27 @@ class DeepseekV4AttnBackend(
         out_cache_loc: Optional[torch.Tensor] = None,
         use_prefill_cuda_graph: bool = False,
         online_c128_state_slot_offset: int = 0,
+        ragged_layout: Optional[RaggedVerifyLayout] = None,
     ) -> DSV4Metadata:
-        batch_size = len(seq_lens)
-        seq_lens = seq_lens + self.speculative_num_draft_tokens
-        seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu]
-        extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size
-        extend_seq_lens = self._move_to_device(extend_seq_lens_cpu)
-        num_tokens = self.speculative_num_draft_tokens * batch_size
+        if ragged_layout is None:
+            lengths = compute_uniform_extend_lengths(
+                seq_lens=seq_lens,
+                seq_lens_cpu=seq_lens_cpu,
+                extend_len=self.speculative_num_draft_tokens,
+            )
+            extend_seq_lens = self._move_to_device(lengths.extend_seq_lens_cpu)
+        else:
+            lengths = compute_ragged_extend_lengths(
+                seq_lens=seq_lens,
+                seq_lens_cpu=seq_lens_cpu,
+                ragged_layout=ragged_layout,
+            )
+            extend_seq_lens = ragged_layout.verify_lens
+        seq_lens = lengths.seq_lens_extended
+        seq_lens_cpu = lengths.seq_lens_cpu_extended
+        extend_seq_lens_cpu = lengths.extend_seq_lens_cpu
+        num_tokens = lengths.num_tokens
+        extend_start_loc = lengths.extend_start_loc
         if out_cache_loc is None:
             out_cache_loc = seq_lens.new_zeros(num_tokens)
         return self.init_forward_metadata_prefill(
@@ -775,12 +891,46 @@ class DeepseekV4AttnBackend(
             num_tokens=num_tokens,
             extend_seq_lens=extend_seq_lens,
             extend_seq_lens_cpu=extend_seq_lens_cpu,
-            extend_start_loc=None,
+            extend_start_loc=extend_start_loc,
             need_compress=True,
             use_prefill_cuda_graph=use_prefill_cuda_graph,
             online_c128_state_slot_offset=online_c128_state_slot_offset,
         )
 
+    def init_forward_metadata_dspark_draft_block(
+        self,
+        max_seq_len: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: Optional[torch.Tensor],
+        out_cache_loc: torch.Tensor,
+        block_size: int,
+    ) -> DSV4Metadata:
+        if seq_lens_cpu is None:
+            seq_lens_cpu_list = seq_lens.tolist()
+        else:
+            seq_lens_cpu_list = [int(x) for x in seq_lens_cpu.tolist()]
+        lengths = compute_uniform_extend_lengths(
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu_list,
+            extend_len=block_size,
+        )
+        extend_seq_lens = self._move_to_device(lengths.extend_seq_lens_cpu)
+        return self.init_forward_metadata_prefill(
+            max_seq_len=max_seq_len,
+            req_pool_indices=req_pool_indices,
+            seq_lens=lengths.seq_lens_extended,
+            seq_lens_cpu=lengths.seq_lens_cpu_extended,
+            out_cache_loc=out_cache_loc,
+            num_tokens=lengths.num_tokens,
+            extend_seq_lens=extend_seq_lens,
+            extend_seq_lens_cpu=lengths.extend_seq_lens_cpu,
+            extend_start_loc=lengths.extend_start_loc,
+            need_compress=False,
+            use_prefill_cuda_graph=False,
+            dspark_block_size=block_size,
+        )
+
     def make_forward_metadata_from_raw_verify(
         self,
         raw_metadata: DSV4RawVerifyMetadata,
@@ -791,13 +941,33 @@ class DeepseekV4AttnBackend(
         out_cache_loc = raw_metadata.out_cache_loc
 
         bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens
-        seq_lens = seq_lens + self.speculative_num_draft_tokens
         extend_seq_lens = raw_metadata.extend_seq_lens
         assert extend_seq_lens is not None
 
+        is_ragged = raw_metadata.verify_lens is not None
+        if is_ragged:
+            seq_lens = seq_lens + extend_seq_lens
+            num_q_tokens = raw_metadata.total_verify_tokens
+            assert num_q_tokens > 0, "ragged verify raw metadata is stale/empty"
+            seq_lens_casual, req_pool_indices_repeated = (
+                self._expand_prefill_casually_vectorized(
+                    num_tokens=num_q_tokens,
+                    seq_lens=seq_lens,
+                    extend_seq_lens=extend_seq_lens,
+                    extend_start_loc=raw_metadata.extend_start_loc,
+                    req_pool_indices=req_pool_indices,
+                    padded_num_tokens=out_cache_loc.shape[0],
+                )
+            )
+        else:
+            seq_lens = seq_lens + self.speculative_num_draft_tokens
+            num_q_tokens = num_draft_tokens * bs
             seq_lens_casual, req_pool_indices_repeated = (
                 self.expand_extend_with_same_length(
-                bs, num_draft_tokens, seq_lens, req_pool_indices
+                    bs=bs,
+                    qo_len=num_draft_tokens,
+                    seq_lens=seq_lens,
+                    req_pool_indices=req_pool_indices,
                 )
             )
         core_attn_metadata = self.make_core_attn_metadata(
@@ -820,7 +990,7 @@ class DeepseekV4AttnBackend(
             seq_lens_cpu=None,
             extend_lens_cpu=None,
             use_prefill_cuda_graph=True,
-            num_q_tokens=num_draft_tokens * bs,
+            num_q_tokens=num_q_tokens,
             online_state_slot_offset=online_c128_state_slot_offset,
         )
         c128_compress_metadata = raw_metadata.c128_compress_metadata
@@ -941,6 +1111,35 @@ class DeepseekV4AttnBackend(
                 )
             )
 
+            if self.is_dspark_draft and forward_batch.forward_mode.is_target_verify():
+                block_size = int(forward_batch.spec_info.draft_token_num)
+                seq_lens_casual = self._dspark_seq_lens_casual(
+                    seq_lens=forward_batch.seq_lens, block_size=block_size
+                )
+                req_pool_indices_repeated = (
+                    forward_batch.req_pool_indices.repeat_interleave(block_size)
+                )
+                (
+                    swa_page_indices,
+                    swa_topk_lengths,
+                ) = self.get_dspark_swa_page_indices(
+                    seq_lens_casual=seq_lens_casual,
+                    req_pool_indices_repeated=req_pool_indices_repeated,
+                    out_loc=out_cache_loc,
+                    block_size=block_size,
+                )
+                metadata.core_attn_metadata.swa_page_indices = swa_page_indices
+                metadata.core_attn_metadata.swa_topk_lengths = swa_topk_lengths
+
+    def _dspark_seq_lens_casual(
+        self, *, seq_lens: torch.Tensor, block_size: int
+    ) -> torch.Tensor:
+        return BuildBlockSeqLensCausal.execute(
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=self.cuda_int32_kwargs["device"],
+        )
+
     def init_forward_metadata_out_graph(
         self,
         forward_batch: ForwardBatch,
@@ -997,6 +1196,7 @@ class DeepseekV4AttnBackend(
             actual_max_seq_len = seq_lens_cpu.max().item()
             assert actual_max_seq_len <= chosen_max_seq_len
 
+        graph_key = bs
         if bucket == _GraphBucket.DECODE_OR_IDLE:
             assert out_cache_loc is not None
             assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}"
@@ -1017,16 +1217,48 @@ class DeepseekV4AttnBackend(
                 seq_lens=seq_lens,
                 out_cache_loc=out_cache_loc_padded,
             )
+        elif bucket == _GraphBucket.TARGET_VERIFY and self.is_dspark_draft:
+            block_size = self.speculative_num_draft_tokens - 1
+            num_tokens_block = block_size * bs
+            assert out_cache_loc is not None
+            out_cache_loc_padded = torch.nn.functional.pad(
+                out_cache_loc,
+                pad=(0, num_tokens_block - len(out_cache_loc)),
+                mode="constant",
+                value=0,
+            )
+            self.online_c128_mtp.prepare_forward(
+                actual_forward_mode,
+                req_pool_indices,
+                seq_lens,
+            )
+            temp_metadata = self.init_forward_metadata_dspark_draft_block(
+                max_seq_len=chosen_max_seq_len,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+                seq_lens_cpu=seq_lens_cpu,
+                out_cache_loc=out_cache_loc_padded,
+                block_size=block_size,
+            )
         elif bucket == _GraphBucket.TARGET_VERIFY:
             verify_bs = _get_target_verify_bs(forward_batch)
+            ragged_layout = self._resolve_verify_layout(forward_batch, bs=bs)
+            graph_key, num_tokens_v = self._target_verify_graph_key(
+                bs=bs, ragged_layout=ragged_layout
+            )
             if self.online_c128_mtp.enabled() and verify_bs == 0:
                 self.online_c128_mtp.clear()
                 self.forward_metadata = self.cuda_graph_metadata_of_bucket_and_bs[
                     bucket
-                ][bs]
+                ][graph_key]
                 return
             assert out_cache_loc is not None
-            num_tokens_v = self.speculative_num_draft_tokens * bs
+            assert num_tokens_v >= len(out_cache_loc), (
+                f"ragged verify token-keyed graph requires the decode cuda-graph "
+                f"runner to supply out_cache_loc sized to graph_num_tokens "
+                f"({num_tokens_v}), got {len(out_cache_loc)}; the decode graph "
+                "runner does not yet route token-keyed ragged captures."
+            )
             out_cache_loc_padded = torch.nn.functional.pad(
                 out_cache_loc,
                 pad=(0, num_tokens_v - len(out_cache_loc)),
@@ -1047,6 +1279,7 @@ class DeepseekV4AttnBackend(
                 out_cache_loc=out_cache_loc_padded,
                 use_prefill_cuda_graph=True,
                 online_c128_state_slot_offset=online_c128_state_slot_offset,
+                ragged_layout=ragged_layout,
             )
         elif bucket == _GraphBucket.DRAFT_EXTEND:
             self.online_c128_mtp.prepare_forward(
@@ -1081,7 +1314,7 @@ class DeepseekV4AttnBackend(
             raise NotImplementedError
 
         self.replay_cuda_graph_metadata_from(
-            bs=bs, temp_metadata=temp_metadata, bucket=bucket
+            bs=graph_key, temp_metadata=temp_metadata, bucket=bucket
         )
 
         if in_capture:
@@ -1152,7 +1385,18 @@ class DeepseekV4AttnBackend(
                 seq_lens=seq_lens,
                 out_cache_loc=out_cache_loc,
             )
+        elif self.is_dspark_draft and logical_forward_mode.is_target_verify():
+            block_size = int(forward_batch.spec_info.draft_token_num)
+            metadata = self.init_forward_metadata_dspark_draft_block(
+                max_seq_len=max_seq_len,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+                seq_lens_cpu=seq_lens_cpu,
+                out_cache_loc=forward_batch.out_cache_loc,
+                block_size=block_size,
+            )
         elif logical_forward_mode.is_target_verify():
+            ragged_layout = self._resolve_verify_layout(forward_batch, bs=len(seq_lens))
             metadata = self.init_forward_metadata_target_verify(
                 max_seq_len=max_seq_len,
                 req_pool_indices=req_pool_indices,
@@ -1160,6 +1404,7 @@ class DeepseekV4AttnBackend(
                 seq_lens_cpu=seq_lens_cpu,
                 out_cache_loc=forward_batch.out_cache_loc,
                 online_c128_state_slot_offset=online_c128_state_slot_offset,
+                ragged_layout=ragged_layout,
             )
         elif logical_forward_mode.is_prefill(include_draft_extend_v2=True):
             extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
@@ -1588,46 +1833,18 @@ class DeepseekV4AttnBackend(
         extend_seq_lens_tensor: Optional[torch.Tensor] = None,
         extend_start_loc: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if (
-            seq_lens_tensor is not None
-            and extend_seq_lens_tensor is not None
-            and extend_start_loc is not None
-        ):
-            return self._expand_prefill_casually_vectorized(
-                num_tokens=num_tokens,
+        assert seq_lens_tensor is not None and extend_seq_lens_tensor is not None
+        result = ExpandPrefillCausally.execute(
+            req_pool_indices=req_pool_indices,
             seq_lens=seq_lens_tensor,
             extend_seq_lens=extend_seq_lens_tensor,
             extend_start_loc=extend_start_loc,
-                req_pool_indices=req_pool_indices,
+            seq_lens_cpu=seq_lens,
+            extend_seq_lens_cpu=extend_seq_lens,
+            num_tokens=num_tokens,
             padded_num_tokens=padded_num_tokens,
         )
-
-        seq_lens_casual = torch.empty(num_tokens, **self.cuda_int32_kwargs)
-        idx_to_req_repeated = torch.empty(num_tokens, **self.cuda_int32_kwargs)
-        offset = 0
-        for i, (kv_len, qo_len) in enumerate(zip(seq_lens, extend_seq_lens)):
-            out = seq_lens_casual[offset : offset + qo_len]
-            offset += qo_len
-            torch.arange(kv_len - qo_len + 1, kv_len + 1, out=out)
-            idx_to_req_repeated[offset - qo_len : offset].fill_(i)
-
-        assert offset == num_tokens
-        req_pool_indices_repeated = req_pool_indices[idx_to_req_repeated]
-
-        if padded_num_tokens is not None and padded_num_tokens > num_tokens:
-            pad_size = padded_num_tokens - num_tokens
-            seq_lens_casual = torch.nn.functional.pad(
-                seq_lens_casual,
-                (0, pad_size),
-                value=1,
-            )
-            req_pool_indices_repeated = torch.nn.functional.pad(
-                req_pool_indices_repeated,
-                (0, pad_size),
-                value=req_pool_indices_repeated[-1].item(),
-            )
-
-        return seq_lens_casual, req_pool_indices_repeated
+        return result.seq_lens_casual, result.req_pool_indices_repeated
 
     def _expand_prefill_casually_vectorized(
         self,
@@ -1638,41 +1855,21 @@ class DeepseekV4AttnBackend(
         req_pool_indices: torch.Tensor,
         padded_num_tokens: Optional[int],
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        repeats = extend_seq_lens.to(torch.int64)
-        req_pool_indices_repeated = torch.repeat_interleave(
-            req_pool_indices, repeats, output_size=num_tokens
+        result = ExpandPrefillCausally.execute(
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+            seq_lens_cpu=None,
+            extend_seq_lens_cpu=None,
+            num_tokens=num_tokens,
+            padded_num_tokens=padded_num_tokens,
         )
-
-        start_positions = seq_lens.to(torch.int32) - extend_seq_lens.to(torch.int32) + 1
-        start_positions_repeated = torch.repeat_interleave(
-            start_positions, repeats, output_size=num_tokens
-        )
-        start_locs_repeated = torch.repeat_interleave(
-            extend_start_loc.to(torch.int32), repeats, output_size=num_tokens
-        )
-        token_offsets = (
-            torch.arange(num_tokens, **self.cuda_int32_kwargs) - start_locs_repeated
-        )
-        seq_lens_casual = start_positions_repeated + token_offsets
-
-        if padded_num_tokens is not None and padded_num_tokens > num_tokens:
-            pad_size = padded_num_tokens - num_tokens
-            seq_lens_casual = torch.nn.functional.pad(
-                seq_lens_casual,
-                (0, pad_size),
-                value=1,
-            )
-            req_pool_indices_repeated = torch.cat(
-                (
-                    req_pool_indices_repeated,
-                    req_pool_indices_repeated[-1:].expand(pad_size),
-                )
-            )
-
-        return seq_lens_casual, req_pool_indices_repeated
+        return result.seq_lens_casual, result.req_pool_indices_repeated
 
     def expand_extend_with_same_length(
         self,
+        *,
         bs: int,
         qo_len: int,
         seq_lens: torch.Tensor,
@@ -1697,27 +1894,49 @@ class DeepseekV4AttnBackend(
         out_loc: torch.Tensor,
         need_compress: bool = True,
         is_prefill: bool = False,
+        dspark_block_size: Optional[int] = None,
     ) -> DSV4AttnMetadata:
         assert self.swa_page_size == SWA_WINDOW
 
-        seq_lens_casual = seq_lens_casual.to(torch.int32)
+        prep = BuildPageTablePositions.execute(
+            req_to_token=req_to_token,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            seq_lens_casual=seq_lens_casual,
+            max_seq_len=max_seq_len,
+            page_size=self.page_size,
+            swa_window=SWA_WINDOW,
+        )
+        seq_lens_casual = prep.seq_lens_casual
 
-        swa_page_indices = self.get_swa_page_indices(
+        raw_positions = prep.positions_casual
+        if dspark_block_size is not None:
+            assert (
+                self.is_dspark_draft
+                and dspark_block_size == self.speculative_num_draft_tokens - 1
+            ), (
+                f"dspark_block_size={dspark_block_size} must equal gamma = "
+                f"speculative_num_draft_tokens-1={self.speculative_num_draft_tokens - 1} "
+                f"and is only valid on the DSpark draft backend "
+                f"(is_dspark_draft={self.is_dspark_draft})."
+            )
+            swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices(
                 seq_lens_casual=seq_lens_casual,
                 req_pool_indices_repeated=req_pool_indices_repeated,
+                out_loc=out_loc,
+                block_size=dspark_block_size,
             )
-
-        swa_page_indices = _pad_last_dim(
-            swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE
+        else:
+            swa_page_indices = BuildCausalSwaPageIndices.execute(
+                req_to_token=self.req_to_token,
+                full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping,
+                req_pool_indices_repeated=req_pool_indices_repeated,
+                seq_lens_casual=seq_lens_casual,
+                swa_window=SWA_WINDOW,
+                page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE,
             )
+            swa_topk_lengths = prep.swa_topk_lengths
 
-        raw_positions = seq_lens_casual - 1
-        swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW)
-
-        page_table = req_to_token[
-            req_pool_indices_repeated, : max_seq_len : self.page_size
-        ]
-        page_table = (page_table // self.page_size).to(torch.int32)
+        page_table = prep.page_table
 
         core_attn_metadata = DSV4AttnMetadata(
             page_size=self.page_size,
@@ -1743,24 +1962,34 @@ class DeepseekV4AttnBackend(
             core_attn_metadata.c128_flashmla_metadata = None
         return core_attn_metadata
 
-    def get_swa_page_indices(
+    def get_dspark_swa_page_indices(
         self,
+        *,
         seq_lens_casual: torch.Tensor,
         req_pool_indices_repeated: torch.Tensor,
-    ) -> torch.Tensor:
-        pos_causal = seq_lens_casual - 1
-        num_qo_tokens = seq_lens_casual.size(0)
-        offsets = pos_causal.unsqueeze(1) - torch.arange(
-            SWA_WINDOW, **self.cuda_int32_kwargs
-        ).unsqueeze(0)
-        invalid_offset_mask = offsets < 0
-        offsets.masked_fill_(invalid_offset_mask, 0)
-        raw_indices = self.req_to_token[req_pool_indices_repeated[:, None], offsets]
-        assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW)
-        raw_indices.masked_fill_(invalid_offset_mask, -1)
-        swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices)
-        # flash_mla attention requires int32 page indices.
-        return swa_indices.to(torch.int32)
+        out_loc: torch.Tensor,
+        block_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        gather = ComputeDsparkWindowGather.execute(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            block_size=block_size,
+            swa_window=SWA_WINDOW,
+        )
+
+        swa_page_indices, swa_topk_lengths = BuildDsparkSwaPageIndices.execute(
+            req_to_token=self.req_to_token,
+            full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping,
+            req_pool_indices_per_request=gather.req_pool_indices_per_request,
+            offsets=gather.offsets,
+            invalid=gather.invalid,
+            out_loc=out_loc[: gather.num_q],
+            context_lens=gather.context_lens,
+            block_size=block_size,
+            swa_window=SWA_WINDOW,
+            page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE,
+        )
+        return swa_page_indices, swa_topk_lengths
 
 
 class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend):

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -41,6 +41,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
+from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout
 from sglang.srt.utils import ceil_align
 
 if TYPE_CHECKING:
@@ -863,6 +864,12 @@ class DeepseekV4HipRadixBackend(
                 out_cache_loc=out_cache_loc_padded,
             )
         elif bucket == _GraphBucket.TARGET_VERIFY:
+            if resolve_ragged_verify_layout(forward_batch) is not None:
+                raise NotImplementedError(
+                    "DSV4 ragged verify is not supported on the HIP backend "
+                    "(DeepseekV4HipRadixBackend) cuda-graph path; disable "
+                    "SGLANG_RAGGED_VERIFY_MODE or use a CUDA device."
+                )
             assert out_cache_loc is not None
             num_tokens_v = self.speculative_num_draft_tokens * bs
             out_cache_loc_padded = torch.nn.functional.pad(
@@ -950,6 +957,12 @@ class DeepseekV4HipRadixBackend(
                 out_cache_loc=out_cache_loc,
             )
         elif forward_batch.forward_mode.is_target_verify():
+            if resolve_ragged_verify_layout(forward_batch) is not None:
+                raise NotImplementedError(
+                    "DSV4 ragged verify is not supported on the HIP backend "
+                    "(DeepseekV4HipRadixBackend); disable SGLANG_RAGGED_VERIFY_MODE "
+                    "or use a CUDA device."
+                )
             metadata = self.init_forward_metadata_target_verify(
                 max_seq_len=max_seq_len,
                 req_pool_indices=req_pool_indices,

python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py

@@ -0,0 +1,526 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+
+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")
+
+
+class ExpandPrefillCausallyResult(msgspec.Struct):
+    seq_lens_casual: torch.Tensor
+    req_pool_indices_repeated: torch.Tensor
+
+
+class ExpandPrefillCausally:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> ExpandPrefillCausallyResult:
+        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,
+        seq_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        extend_start_loc: Optional[torch.Tensor],
+        seq_lens_cpu: Optional[list[int]],
+        extend_seq_lens_cpu: Optional[list[int]],
+        num_tokens: int,
+        padded_num_tokens: Optional[int],
+    ) -> ExpandPrefillCausallyResult:
+        return expand_prefill_causally(
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+            seq_lens_cpu=seq_lens_cpu,
+            extend_seq_lens_cpu=extend_seq_lens_cpu,
+            num_tokens=num_tokens,
+            padded_num_tokens=padded_num_tokens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        extend_start_loc: Optional[torch.Tensor],
+        seq_lens_cpu: Optional[list[int]],
+        extend_seq_lens_cpu: Optional[list[int]],
+        num_tokens: int,
+        padded_num_tokens: Optional[int],
+    ) -> ExpandPrefillCausallyResult:
+        return expand_prefill_causally_triton(
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            extend_seq_lens=extend_seq_lens,
+            num_tokens=num_tokens,
+            padded_num_tokens=padded_num_tokens,
+        )
+
+
+def expand_prefill_causally(
+    *,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    extend_start_loc: Optional[torch.Tensor],
+    seq_lens_cpu: Optional[list[int]],
+    extend_seq_lens_cpu: Optional[list[int]],
+    num_tokens: int,
+    padded_num_tokens: Optional[int],
+) -> ExpandPrefillCausallyResult:
+    device = req_pool_indices.device
+    cuda_int32_kwargs = {"dtype": torch.int32, "device": device}
+
+    if extend_start_loc is not None:
+        repeats = extend_seq_lens.to(torch.int64)
+        req_pool_indices_repeated = torch.repeat_interleave(
+            req_pool_indices, repeats, output_size=num_tokens
+        )
+        start_positions = seq_lens.to(torch.int32) - extend_seq_lens.to(torch.int32) + 1
+        start_positions_repeated = torch.repeat_interleave(
+            start_positions, repeats, output_size=num_tokens
+        )
+        start_locs_repeated = torch.repeat_interleave(
+            extend_start_loc.to(torch.int32), repeats, output_size=num_tokens
+        )
+        token_offsets = (
+            torch.arange(num_tokens, **cuda_int32_kwargs) - start_locs_repeated
+        )
+        seq_lens_casual = start_positions_repeated + token_offsets
+
+        if padded_num_tokens is not None and padded_num_tokens > num_tokens:
+            pad_size = padded_num_tokens - num_tokens
+            seq_lens_casual = torch.nn.functional.pad(
+                seq_lens_casual, (0, pad_size), value=1
+            )
+            req_pool_indices_repeated = torch.cat(
+                (
+                    req_pool_indices_repeated,
+                    req_pool_indices_repeated[-1:].expand(pad_size),
+                )
+            )
+        return ExpandPrefillCausallyResult(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+        )
+
+    assert seq_lens_cpu is not None and extend_seq_lens_cpu is not None
+    seq_lens_casual = torch.empty(num_tokens, **cuda_int32_kwargs)
+    idx_to_req_repeated = torch.empty(num_tokens, **cuda_int32_kwargs)
+    offset = 0
+    for i, (kv_len, qo_len) in enumerate(zip(seq_lens_cpu, extend_seq_lens_cpu)):
+        out = seq_lens_casual[offset : offset + qo_len]
+        offset += qo_len
+        torch.arange(kv_len - qo_len + 1, kv_len + 1, out=out)
+        idx_to_req_repeated[offset - qo_len : offset].fill_(i)
+
+    assert offset == num_tokens
+    req_pool_indices_repeated = req_pool_indices[idx_to_req_repeated]
+
+    if padded_num_tokens is not None and padded_num_tokens > num_tokens:
+        pad_size = padded_num_tokens - num_tokens
+        seq_lens_casual = torch.nn.functional.pad(
+            seq_lens_casual, (0, pad_size), value=1
+        )
+        req_pool_indices_repeated = torch.nn.functional.pad(
+            req_pool_indices_repeated,
+            (0, pad_size),
+            value=req_pool_indices_repeated[-1].item(),
+        )
+    return ExpandPrefillCausallyResult(
+        seq_lens_casual=seq_lens_casual,
+        req_pool_indices_repeated=req_pool_indices_repeated,
+    )
+
+
+@triton.jit
+def _expand_prefill_causally_kernel(
+    req_pool_ptr,
+    seq_lens_ptr,
+    extend_seq_lens_ptr,
+    seq_lens_casual_ptr,
+    req_pool_repeated_ptr,
+    bs,
+    num_tokens,
+    total_tokens,
+    BLOCK: tl.constexpr,
+    BS_P2: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < total_tokens
+
+    b = tl.arange(0, BS_P2)
+    bmask = b < bs
+    extend = tl.load(extend_seq_lens_ptr + b, mask=bmask, other=0).to(tl.int32)
+    start_locs = tl.cumsum(extend, axis=0) - extend
+
+    is_real = offs < num_tokens
+    t = tl.where(is_real, offs, 0).to(tl.int32)
+    started = (start_locs[None, :] <= t[:, None]) & bmask[None, :]
+    r = tl.sum(started.to(tl.int32), axis=1) - 1
+    r = tl.where(is_real, r, bs - 1).to(tl.int64)
+
+    seq_len = tl.load(seq_lens_ptr + r, mask=mask, other=0).to(tl.int32)
+    ext = tl.load(extend_seq_lens_ptr + r, mask=mask, other=0).to(tl.int32)
+    start_loc = tl.sum(tl.where(started, extend[None, :], 0).to(tl.int32), axis=1) - ext
+    causal = (seq_len - ext + 1) + (t - start_loc)
+    causal = tl.where(is_real, causal, 1)
+
+    rp = tl.load(req_pool_ptr + r, mask=mask, other=0)
+    tl.store(seq_lens_casual_ptr + offs, causal, mask=mask)
+    tl.store(req_pool_repeated_ptr + offs, rp, mask=mask)
+
+
+def expand_prefill_causally_triton(
+    *,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    num_tokens: int,
+    padded_num_tokens: Optional[int],
+) -> ExpandPrefillCausallyResult:
+    bs = req_pool_indices.shape[0]
+    device = req_pool_indices.device
+    total_tokens = (
+        padded_num_tokens
+        if padded_num_tokens is not None and padded_num_tokens > num_tokens
+        else num_tokens
+    )
+
+    seq_lens_casual = torch.empty(total_tokens, dtype=torch.int32, device=device)
+    req_pool_indices_repeated = torch.empty(
+        total_tokens, dtype=req_pool_indices.dtype, device=device
+    )
+    BLOCK = 256
+    _expand_prefill_causally_kernel[(triton.cdiv(total_tokens, BLOCK),)](
+        req_pool_indices,
+        seq_lens,
+        extend_seq_lens,
+        seq_lens_casual,
+        req_pool_indices_repeated,
+        bs,
+        num_tokens,
+        total_tokens,
+        BLOCK=BLOCK,
+        BS_P2=triton.next_power_of_2(max(bs, 1)),
+    )
+    return ExpandPrefillCausallyResult(
+        seq_lens_casual=seq_lens_casual,
+        req_pool_indices_repeated=req_pool_indices_repeated,
+    )
+
+
+class PageTablePositionsResult(msgspec.Struct):
+    seq_lens_casual: torch.Tensor
+    positions_casual: torch.Tensor
+    page_table: torch.Tensor
+    swa_topk_lengths: torch.Tensor
+
+
+class BuildPageTablePositions:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> PageTablePositionsResult:
+        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,
+        req_pool_indices_repeated: torch.Tensor,
+        seq_lens_casual: torch.Tensor,
+        max_seq_len: int,
+        page_size: int,
+        swa_window: int,
+    ) -> PageTablePositionsResult:
+        return build_page_table_positions(
+            req_to_token=req_to_token,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            seq_lens_casual=seq_lens_casual,
+            max_seq_len=max_seq_len,
+            page_size=page_size,
+            swa_window=swa_window,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_to_token: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        seq_lens_casual: torch.Tensor,
+        max_seq_len: int,
+        page_size: int,
+        swa_window: int,
+    ) -> PageTablePositionsResult:
+        return build_page_table_positions_triton(
+            req_to_token=req_to_token,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            seq_lens_casual=seq_lens_casual,
+            max_seq_len=max_seq_len,
+            page_size=page_size,
+            swa_window=swa_window,
+        )
+
+
+def build_page_table_positions(
+    *,
+    req_to_token: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    seq_lens_casual: torch.Tensor,
+    max_seq_len: int,
+    page_size: int,
+    swa_window: int,
+) -> PageTablePositionsResult:
+    seq_lens_casual = seq_lens_casual.to(torch.int32)
+    positions_casual = seq_lens_casual - 1
+    page_table = req_to_token[
+        req_pool_indices_repeated.to(torch.int64), :max_seq_len:page_size
+    ]
+    page_table = (page_table // page_size).to(torch.int32)
+    swa_topk_lengths = torch.clamp(seq_lens_casual, max=swa_window)
+    return PageTablePositionsResult(
+        seq_lens_casual=seq_lens_casual,
+        positions_casual=positions_casual,
+        page_table=page_table,
+        swa_topk_lengths=swa_topk_lengths,
+    )
+
+
+@triton.jit
+def _page_table_positions_kernel(
+    req_to_token_ptr,
+    req_pool_ptr,
+    seq_lens_ptr,
+    seq_lens_out_ptr,
+    positions_out_ptr,
+    page_table_ptr,
+    topk_out_ptr,
+    rt_stride,
+    num_pages,
+    page_size,
+    swa_window,
+    BLOCK_P: tl.constexpr,
+):
+    row = tl.program_id(0)
+    seq_len = tl.load(seq_lens_ptr + row).to(tl.int32)
+    tl.store(seq_lens_out_ptr + row, seq_len)
+    tl.store(positions_out_ptr + row, seq_len - 1)
+    tl.store(topk_out_ptr + row, tl.minimum(seq_len, swa_window))
+
+    rp = tl.load(req_pool_ptr + row).to(tl.int64)
+    base = req_to_token_ptr + rp * rt_stride
+    out_base = page_table_ptr + row.to(tl.int64) * num_pages
+    for p0 in range(0, num_pages, BLOCK_P):
+        p = p0 + tl.arange(0, BLOCK_P)
+        pmask = p < num_pages
+        tok = tl.load(base + p.to(tl.int64) * page_size, mask=pmask, other=0).to(
+            tl.int32
+        )
+        tl.store(out_base + p, tok // page_size, mask=pmask)
+
+
+def build_page_table_positions_triton(
+    *,
+    req_to_token: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    seq_lens_casual: torch.Tensor,
+    max_seq_len: int,
+    page_size: int,
+    swa_window: int,
+) -> PageTablePositionsResult:
+    num_q = seq_lens_casual.shape[0]
+    num_pages = (max_seq_len + page_size - 1) // page_size
+    device = seq_lens_casual.device
+
+    seq_lens_out = torch.empty(num_q, dtype=torch.int32, device=device)
+    positions_out = torch.empty(num_q, dtype=torch.int32, device=device)
+    page_table = torch.empty((num_q, num_pages), dtype=torch.int32, device=device)
+    topk_out = torch.empty(num_q, dtype=torch.int32, device=device)
+    BLOCK_P = 256
+    _page_table_positions_kernel[(num_q,)](
+        req_to_token,
+        req_pool_indices_repeated,
+        seq_lens_casual,
+        seq_lens_out,
+        positions_out,
+        page_table,
+        topk_out,
+        req_to_token.stride(0),
+        num_pages,
+        page_size,
+        swa_window,
+        BLOCK_P=BLOCK_P,
+    )
+    return PageTablePositionsResult(
+        seq_lens_casual=seq_lens_out,
+        positions_casual=positions_out,
+        page_table=page_table,
+        swa_topk_lengths=topk_out,
+    )
+
+
+class BuildCausalSwaPageIndices:
+    @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,
+        *,
+        req_to_token: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        seq_lens_casual: torch.Tensor,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> torch.Tensor:
+        return build_causal_swa_page_indices(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            seq_lens_casual=seq_lens_casual,
+            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_repeated: torch.Tensor,
+        seq_lens_casual: torch.Tensor,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> torch.Tensor:
+        return build_causal_swa_page_indices_triton(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            seq_lens_casual=seq_lens_casual,
+            swa_window=swa_window,
+            page_index_aligned_size=page_index_aligned_size,
+        )
+
+
+def build_causal_swa_page_indices(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    seq_lens_casual: torch.Tensor,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> torch.Tensor:
+    device = seq_lens_casual.device
+    pos_causal = seq_lens_casual - 1
+    num_qo_tokens = seq_lens_casual.size(0)
+    offsets = pos_causal.unsqueeze(1) - torch.arange(
+        swa_window, dtype=torch.int32, device=device
+    ).unsqueeze(0)
+    invalid_offset_mask = offsets < 0
+    offsets.masked_fill_(invalid_offset_mask, 0)
+    raw_indices = req_to_token[req_pool_indices_repeated[:, None], offsets]
+    assert raw_indices.shape == (num_qo_tokens, swa_window)
+    raw_indices.masked_fill_(invalid_offset_mask, -1)
+    swa_indices = full_to_swa_mapping[raw_indices]
+    swa_indices = swa_indices.to(torch.int32)
+
+    padded_width = (
+        (swa_window + page_index_aligned_size - 1) // page_index_aligned_size
+    ) * page_index_aligned_size
+    if padded_width == swa_window:
+        return swa_indices
+    return torch.nn.functional.pad(
+        swa_indices, (0, padded_width - swa_window), value=-1
+    )
+
+
+@triton.jit
+def _causal_swa_page_indices_kernel(
+    req_to_token_ptr,
+    full_to_swa_ptr,
+    req_pool_ptr,
+    seq_lens_ptr,
+    out_ptr,
+    rt_stride,
+    swa_window,
+    padded_width,
+    BLOCK_K: tl.constexpr,
+):
+    row = tl.program_id(0)
+    pos = tl.load(seq_lens_ptr + row).to(tl.int64) - 1
+    rp = tl.load(req_pool_ptr + row).to(tl.int64)
+    base = req_to_token_ptr + rp * rt_stride
+    out_base = out_ptr + row.to(tl.int64) * padded_width
+
+    for k0 in range(0, padded_width, BLOCK_K):
+        k = k0 + tl.arange(0, BLOCK_K)
+        kmask = k < padded_width
+        off = pos - k.to(tl.int64)
+        valid = (k < swa_window) & (off >= 0) & kmask
+        full_loc = tl.load(base + tl.where(valid, off, 0), mask=valid, other=-1).to(
+            tl.int64
+        )
+        swa = tl.load(full_to_swa_ptr + full_loc, mask=valid, other=-1).to(tl.int32)
+        tl.store(out_base + k, tl.where(valid, swa, -1), mask=kmask)
+
+
+def build_causal_swa_page_indices_triton(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    seq_lens_casual: torch.Tensor,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> torch.Tensor:
+    num_qo_tokens = seq_lens_casual.size(0)
+    padded_width = (
+        (swa_window + page_index_aligned_size - 1) // page_index_aligned_size
+    ) * page_index_aligned_size
+    out = torch.empty(
+        (num_qo_tokens, padded_width),
+        dtype=torch.int32,
+        device=seq_lens_casual.device,
+    )
+    BLOCK_K = 256
+    _causal_swa_page_indices_kernel[(num_qo_tokens,)](
+        req_to_token,
+        full_to_swa_mapping,
+        req_pool_indices_repeated,
+        seq_lens_casual,
+        out,
+        req_to_token.stride(0),
+        swa_window,
+        padded_width,
+        BLOCK_K=BLOCK_K,
+    )
+    return out

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -29,6 +29,7 @@ from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_server_args
+from sglang.srt.speculative.ragged_verify import build_ragged_target_verify_geometry
 from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
 from sglang.srt.utils import get_compiler_backend
 
@@ -216,6 +217,8 @@ class FlashAttentionBackend(AttentionBackend):
     - For each forward batch, init_replay_cuda_graph will be called first and then replay the graph.
     """
 
+    supports_ragged_verify_graph: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -253,11 +256,12 @@ class FlashAttentionBackend(AttentionBackend):
         self.max_num_pages = (
             self.max_context_len + self.page_size - 1
         ) // self.page_size
-        # Opt out of the seq_lens_cpu D2H only for dflash (the worker adapted to
-        # the GPU-only relay); EAGLE/MTP/standalone/non-spec keep the CPU mirror.
+        # Opt out of the seq_lens_cpu D2H only for dflash/dspark (their workers
+        # adapted to the GPU-only relay); EAGLE/MTP/standalone/non-spec keep the
+        # CPU mirror.
         self.needs_cpu_seq_lens = not SpeculativeAlgorithm.from_string(
             model_runner.server_args.speculative_algorithm
-        ).is_dflash()
+        ).is_dflash_family()
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
         self.skip_prefill = skip_prefill
         self.attn_cp_size = model_runner.attn_cp_size
@@ -272,6 +276,15 @@ class FlashAttentionBackend(AttentionBackend):
         self.speculative_num_draft_tokens = (
             model_runner.server_args.speculative_num_draft_tokens
         )
+        if (
+            self.speculative_num_draft_tokens is not None
+            and model_runner.is_draft_worker
+        ):
+            self.speculative_num_draft_tokens = SpeculativeAlgorithm.from_string(
+                model_runner.server_args.speculative_algorithm
+            ).get_num_tokens_per_bs_for_target_verify(
+                int(self.speculative_num_draft_tokens), is_draft_worker=True
+            )
         self.speculative_step_id = speculative_step_id
 
         # Local attention settings
@@ -714,6 +727,27 @@ class FlashAttentionBackend(AttentionBackend):
             self._maybe_init_local_attn_metadata(forward_batch, metadata, device)
         elif forward_batch.forward_mode.is_target_verify():
             if self.topk <= 1:
+                ragged_layout = getattr(
+                    forward_batch.spec_info, "ragged_verify_layout", None
+                )
+                if ragged_layout is not None:
+                    geometry = build_ragged_target_verify_geometry(
+                        seq_lens=forward_batch.seq_lens, layout=ragged_layout
+                    )
+                    metadata.cache_seqlens_int32 = geometry.cache_seqlens_int32
+                    # Device-only layouts carry no host lens; the verify
+                    # window is a valid varlen upper bound.
+                    metadata.max_seq_len_q = (
+                        geometry.max_seq_len_q
+                        if geometry.max_seq_len_q is not None
+                        else self.speculative_num_draft_tokens
+                    )
+                    metadata.max_seq_len_k = int(
+                        metadata.cache_seqlens_int32.max().item()
+                    )
+                    metadata.cu_seqlens_q = geometry.cu_seqlens_q
+                    metadata.cu_seqlens_k = geometry.cu_seqlens_k
+                else:
                     metadata.cache_seqlens_int32 = (
                         forward_batch.seq_lens + self.speculative_num_draft_tokens
                     ).to(torch.int32)
@@ -2563,6 +2597,15 @@ class FlashAttentionBackend(AttentionBackend):
         elif forward_mode.is_target_verify():
             if self.topk <= 1:
                 metadata = self.target_verify_metadata[bs]
+                ragged_layout = getattr(spec_info, "ragged_verify_layout", None)
+                if ragged_layout is not None:
+                    padded = ragged_layout.padded_to_bucket(padded_bs=bs)
+                    geometry = build_ragged_target_verify_geometry(
+                        seq_lens=seq_lens, layout=padded
+                    )
+                    metadata.cache_seqlens_int32.copy_(geometry.cache_seqlens_int32)
+                    metadata.cu_seqlens_q.copy_(geometry.cu_seqlens_q)
+                else:
                     metadata.cache_seqlens_int32.copy_(
                         (seq_lens + self.speculative_num_draft_tokens)
                     )

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -657,6 +657,16 @@ class FlashInferAttnBackend(AttentionBackend):
         forward_mode = forward_batch.forward_mode
         spec_info = forward_batch.spec_info
 
+        if (
+            spec_info is not None
+            and spec_info.ragged_verify_layout is not None
+            and forward_mode.is_target_verify()
+        ):
+            raise NotImplementedError(
+                "FlashInfer does not support ragged verify in cuda graph; "
+                "disable SGLANG_RAGGED_VERIFY_MODE for this configuration."
+            )
+
         if in_capture:
             num_tokens = forward_batch.positions.numel()
             self._prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info)

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -32,6 +32,10 @@ from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_buffer
+from sglang.srt.speculative.ragged_verify import (
+    build_ragged_target_verify_geometry,
+    resolve_ragged_verify_layout,
+)
 from sglang.srt.utils import is_flashinfer_available
 from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported
 
@@ -43,6 +47,7 @@ if is_flashinfer_available():
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
     from sglang.srt.speculative.spec_info import SpecInput
 
 # Constants
@@ -69,6 +74,7 @@ class TRTLLMMHAMetadata:
     swa_page_table: torch.Tensor = None
     # full->SWA translated out_cache_loc (SWA KV-store write target)
     swa_out_cache_loc: torch.Tensor = None
+    is_ragged_verify: bool = False
 
 
 class TRTLLMHAAttnBackend(FlashInferAttnBackend):
@@ -79,6 +85,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
     # seq_lens_cpu D2H sync; opt out of it, matching trtllm_mla / triton.
     needs_cpu_seq_lens: bool = False
 
+    supports_ragged_verify_graph: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -353,6 +361,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 "cache_seqlens": torch.zeros(
                     max_bs, dtype=torch.int32, device=self.device
                 ),
+                # Static uniform preset (Q_MODE_NONE: the fused kernel never
+                # rewrites it). Ragged verify overwrites the [:bs+1] slice
+                # eagerly on every capture/replay-prep, and the ragged-verify
+                # mode is fixed for the whole server run, so the two never mix.
                 "cu_seqlens_q": torch.arange(
                     0,
                     max_bs * self.speculative_num_draft_tokens + 1,
@@ -449,7 +461,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 self.decode_cuda_graph_metadata[bs] = metadata
         elif forward_mode.is_target_verify():
             # Target Verify (topk = 1)
-            tokens_per_req = num_tokens // bs
             metadata.cache_seqlens_int32 = self.target_verify_metadata["cache_seqlens"][
                 :bs
             ]
@@ -459,7 +470,14 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][
                 : bs + 1
             ]
-            metadata.max_seq_len_q = tokens_per_req
+            metadata.is_ragged_verify = (
+                spec_info is not None and spec_info.ragged_verify_layout is not None
+            )
+            metadata.max_seq_len_q = (
+                self.speculative_num_draft_tokens
+                if metadata.is_ragged_verify
+                else num_tokens // bs
+            )
             metadata.page_table = self.target_verify_metadata["page_table"][:bs, :]
             self._bind_swa_page_table(
                 metadata,
@@ -536,6 +554,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         elif forward_mode.is_target_verify():
             # Here we only support topk = 1 for now.
             metadata = self.target_verify_metadata[bs]
+            if spec_info is not None and spec_info.ragged_verify_layout is not None:
+                # Ragged verify: the per-request k-extension is not a
+                # uniform scalar seqlen_offset, so the fused kernel cannot
+                # rebuild this metadata. It is written eagerly on every
+                # capture/replay-prep in init_forward_metadata_out_graph;
+                # record nothing here.
+                return
             seqlen_offset = metadata.max_seq_len_q
         elif forward_mode.is_draft_extend_v2():
             metadata = self.draft_extend_metadata[bs]
@@ -628,6 +653,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         forward_mode = forward_batch.forward_mode
         spec_info = forward_batch.spec_info
 
+        if (
+            forward_mode.is_target_verify()
+            and resolve_ragged_verify_layout(forward_batch) is not None
+        ):
+            self._assert_ragged_verify_supported()
+
         if in_capture:
             num_tokens = forward_batch.positions.numel()
             self._build_cuda_graph_metadata(
@@ -638,6 +669,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             self.forward_metadata = self.decode_cuda_graph_metadata[bs]
         elif forward_mode.is_target_verify():
             self.forward_metadata = self.target_verify_metadata[bs]
+            ragged_layout = resolve_ragged_verify_layout(forward_batch)
+            if ragged_layout is not None:
+                self._write_ragged_verify_graph_metadata(
+                    self.forward_metadata, forward_batch, ragged_layout, bs
+                )
         elif forward_mode.is_draft_extend_v2():
             self.forward_metadata = self.draft_extend_metadata[bs]
         else:
@@ -645,6 +681,53 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
             )
 
+    def _assert_ragged_verify_supported(self) -> None:
+        if self.is_xqa_impl:
+            raise NotImplementedError(
+                "Compact ragged verify (variable-length cum_seq_lens_q) "
+                "requires the trtllm-gen decode kernel; the xqa impl (sm90 / sm120) "
+                "rejects it. Disable SGLANG_RAGGED_VERIFY_MODE for this configuration."
+            )
+
+    def _write_ragged_verify_graph_metadata(
+        self,
+        metadata: TRTLLMMHAMetadata,
+        forward_batch: ForwardBatch,
+        ragged_layout: RaggedVerifyLayout,
+        bs: int,
+    ) -> None:
+        """Eagerly rebuild the target-verify graph metadata for ragged verify.
+
+        The per-request verify lengths make the k-extension non-uniform, which
+        the fused in-graph kernel cannot express (scalar ``seqlen_offset``
+        only), so this runs out-of-graph on every capture/replay-prep and
+        ``_apply_cuda_graph_metadata`` records nothing for ragged batches.
+        """
+        seq_lens = forward_batch.seq_lens[:bs]
+        req_pool_indices = forward_batch.req_pool_indices[:bs]
+        padded_layout = ragged_layout.padded_to_bucket(padded_bs=bs)
+        geometry = build_ragged_target_verify_geometry(
+            seq_lens=seq_lens, layout=padded_layout
+        )
+        metadata.cache_seqlens_int32.copy_(geometry.cache_seqlens_int32)
+        metadata.cu_seqlens_q.copy_(geometry.cu_seqlens_q)
+        metadata.cu_seqlens_k[1:].copy_(
+            torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
+        )
+        self._fill_page_table_device(
+            metadata, req_pool_indices, metadata.cache_seqlens_int32
+        )
+        # The fused in-graph kernel also skips ragged batches, so refill the
+        # SWA write-target buffer here (out_cache_loc -> SWA locs).
+        if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
+            n = forward_batch.out_cache_loc.shape[0]
+            self.cuda_graph_swa_out_cache_loc[n:].zero_()
+            self.cuda_graph_swa_out_cache_loc[:n].copy_(
+                self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                    forward_batch.out_cache_loc
+                )
+            )
+
     def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
         self._apply_cuda_graph_metadata(
             bs=forward_batch.batch_size,
@@ -689,11 +772,28 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
                 )
         elif forward_batch.forward_mode.is_target_verify():
-            # Only support topk = 1 for now.
-            tokens_per_req = forward_batch.input_ids.shape[0] // batch_size
-            metadata.cache_seqlens_int32 = (forward_batch.seq_lens + tokens_per_req).to(
-                torch.int32
+            ragged_layout = resolve_ragged_verify_layout(forward_batch)
+            if ragged_layout is not None:
+                self._assert_ragged_verify_supported()
+                geometry = build_ragged_target_verify_geometry(
+                    seq_lens=seqlens_in_batch, layout=ragged_layout
                 )
+                metadata.cache_seqlens_int32 = geometry.cache_seqlens_int32
+                # Device-only layouts carry no host lens; the verify window
+                # is a valid varlen upper bound.
+                metadata.max_seq_len_q = (
+                    geometry.max_seq_len_q
+                    if geometry.max_seq_len_q is not None
+                    else self.speculative_num_draft_tokens
+                )
+                metadata.cu_seqlens_q = geometry.cu_seqlens_q
+                metadata.cu_seqlens_k = geometry.cu_seqlens_k
+                metadata.is_ragged_verify = True
+            else:
+                tokens_per_req = forward_batch.input_ids.shape[0] // batch_size
+                metadata.cache_seqlens_int32 = (
+                    forward_batch.seq_lens + tokens_per_req
+                ).to(torch.int32)
                 metadata.max_seq_len_q = tokens_per_req
                 metadata.cu_seqlens_q = torch.arange(
                     0,
@@ -703,7 +803,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     device=device,
                 )
                 metadata.cu_seqlens_k = torch.nn.functional.pad(
-                torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32),
+                    torch.cumsum(
+                        metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                    ),
                     (1, 0),
                 )
 
@@ -906,6 +1008,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             forward_batch.forward_mode.is_target_verify()
             or forward_batch.forward_mode.is_draft_extend_v2()
         ):
+            if self.forward_metadata.is_ragged_verify:
                 o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
                     query=q,
                     kv_cache=kv_cache,
@@ -918,7 +1021,25 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     window_left=layer.sliding_window_size,
                     sinks=attention_sink,
                     skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
-                out_dtype=self.q_data_type,  # model_runner.dtype
+                    out_dtype=self.q_data_type,
+                    q_len_per_req=None,
+                    max_q_len=self.forward_metadata.max_seq_len_q,
+                    cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
+                )
+            else:
+                o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
+                    query=q,
+                    kv_cache=kv_cache,
+                    workspace_buffer=self.workspace_buffer,
+                    block_tables=page_table,
+                    seq_lens=self.forward_metadata.cache_seqlens_int32,
+                    max_seq_len=self.max_context_len,
+                    bmm1_scale=bmm1_scale,
+                    bmm2_scale=bmm2_scale,
+                    window_left=layer.sliding_window_size,
+                    sinks=attention_sink,
+                    skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
+                    out_dtype=self.q_data_type,
                     q_len_per_req=self.forward_metadata.max_seq_len_q,
                 )
         else:

python/sglang/srt/managers/detokenizer_manager.py

@@ -428,7 +428,10 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
             video_tokens=recv_obj.video_tokens,
             spec_verify_ct=recv_obj.spec_verify_ct,
             spec_num_correct_drafts=recv_obj.spec_num_correct_drafts,
+            spec_num_block_accept_tokens=recv_obj.spec_num_block_accept_tokens,
+            spec_num_cap_tokens=recv_obj.spec_num_cap_tokens,
             spec_correct_drafts_histogram=recv_obj.spec_correct_drafts_histogram,
+            spec_cap_lens_histogram=recv_obj.spec_cap_lens_histogram,
             input_token_logprobs_val=recv_obj.input_token_logprobs_val,
             input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
             output_token_logprobs_val=recv_obj.output_token_logprobs_val,

python/sglang/srt/managers/io_struct.py

@@ -1274,8 +1274,11 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
     spec_verify_ct: Optional[List[int]] = None
     # Accepted drafts
     spec_num_correct_drafts: Optional[List[int]] = None
+    spec_num_block_accept_tokens: Optional[List[int]] = None
+    spec_num_cap_tokens: Optional[List[int]] = None
     # Acceptance histogram
     spec_correct_drafts_histogram: Optional[List[List[int]]] = None
+    spec_cap_lens_histogram: Optional[List[List[int]]] = None
 
 
 class BatchStrOutput(BaseBatchReq, kw_only=True):
@@ -1349,8 +1352,11 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     spec_verify_ct: Optional[List[int]] = None
     # Accepted drafts
     spec_num_correct_drafts: Optional[List[int]] = None
+    spec_num_block_accept_tokens: Optional[List[int]] = None
+    spec_num_cap_tokens: Optional[List[int]] = None
     # Acceptance histogram
     spec_correct_drafts_histogram: Optional[List[List[int]]] = None
+    spec_cap_lens_histogram: Optional[List[List[int]]] = None
 
 
 class BatchEmbeddingOutput(BaseBatchReq, kw_only=True):

python/sglang/srt/managers/min_free_slots_delayer.py

@@ -4,7 +4,7 @@ from typing import Optional
 def resolve_min_free_slots(
     user_value: Optional[int],
     max_running_requests: int,
-    is_dflash: bool = False,
+    is_dflash_family: bool = False,
 ) -> Optional[int]:
     """Resolve the min-free-slots threshold (None = disabled).
 
@@ -16,7 +16,7 @@ def resolve_min_free_slots(
     max_running_requests = max(0, int(max_running_requests))
     formula = min(4, max(2, (max_running_requests + 5) // 6))
     if user_value is None:
-        user_value = formula if is_dflash else None
+        user_value = formula if is_dflash_family else None
 
     if user_value is None or user_value <= 1:
         return None

python/sglang/srt/managers/multi_tokenizer_mixin.py

@@ -161,6 +161,15 @@ def _handle_output_by_index(output, i):
             spec_correct_drafts_histogram=_extract_field_by_index(
                 output, "spec_correct_drafts_histogram", i
             ),
+            spec_num_block_accept_tokens=_extract_field_by_index(
+                output, "spec_num_block_accept_tokens", i
+            ),
+            spec_num_cap_tokens=_extract_field_by_index(
+                output, "spec_num_cap_tokens", i
+            ),
+            spec_cap_lens_histogram=_extract_field_by_index(
+                output, "spec_cap_lens_histogram", i
+            ),
             time_stats=_extract_field_by_index(output, "time_stats", i),
             finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
             decoded_texts=_extract_field_by_index(output, "decoded_texts", i),
@@ -263,6 +272,15 @@ def _handle_output_by_index(output, i):
             spec_correct_drafts_histogram=_extract_field_by_index(
                 output, "spec_correct_drafts_histogram", i
             ),
+            spec_num_block_accept_tokens=_extract_field_by_index(
+                output, "spec_num_block_accept_tokens", i
+            ),
+            spec_num_cap_tokens=_extract_field_by_index(
+                output, "spec_num_cap_tokens", i
+            ),
+            spec_cap_lens_histogram=_extract_field_by_index(
+                output, "spec_cap_lens_histogram", i
+            ),
             time_stats=_extract_field_by_index(output, "time_stats", i),
             finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
             output_strs=_extract_field_by_index(output, "output_strs", i),

python/sglang/srt/managers/overlap_utils.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Optional, Sequence
 
+import msgspec
 import torch
 
 from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras
@@ -35,7 +36,8 @@ def decide_needs_cpu_seq_lens(
     if server_args.enable_two_batch_overlap:
         # FIXME: support TBO without seq lens cpu value
         return True
-    if SpeculativeAlgorithm.from_string(server_args.speculative_algorithm).is_ngram():
+    algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+    if algo.is_ngram():
         # ngram's USE_FULL_MASK verify path reads seq_lens_cpu per req to size
         # the tree mask, regardless of the attn backend (e.g. Triton opts out).
         return True
@@ -46,6 +48,19 @@ def decide_needs_cpu_seq_lens(
     )
 
 
+def decide_needs_confidence_relay(server_args: ServerArgs) -> bool:
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+    algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+    if not algo.is_dspark():
+        return False
+    return read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
+
+
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
@@ -100,6 +115,16 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None:
         future_map._resolve_spec_extras(batch)
 
 
+CONFIDENCE_RELAY_RING_LAG: int = 2
+CONFIDENCE_RELAY_RING_DEPTH: int = CONFIDENCE_RELAY_RING_LAG + 1
+
+
+class ResolvedConfidence(msgspec.Struct):
+
+    confidence: torch.Tensor
+    generation: torch.Tensor
+
+
 @dataclass
 class RelayPayload:
     """Per-iteration stash payload for the FutureMap bufs. Non-spec fills only
@@ -125,6 +150,85 @@ class RelayPayload:
         )
 
 
+class ConfidenceRelay(msgspec.Struct):
+
+    device: torch.device
+    req_pool_size: int
+    pool: Any
+    confidence_buf: Optional[torch.Tensor] = None
+    conf_ring: Optional[torch.Tensor] = None
+    gen_ring: Optional[torch.Tensor] = None
+    copy_done: Optional[list] = None
+    ring_pos: int = 0
+    initialized: bool = False
+
+    def _lazy_init(self, confidence: torch.Tensor) -> None:
+        self.initialized = True
+        gamma = confidence.shape[-1]
+        self.confidence_buf = torch.empty(
+            (self.req_pool_size, gamma), dtype=torch.float32, device=self.device
+        )
+        if _is_cuda:
+            depth = CONFIDENCE_RELAY_RING_DEPTH
+            self.conf_ring = torch.empty(
+                (depth, self.req_pool_size, gamma),
+                dtype=torch.float32,
+                pin_memory=True,
+            )
+            self.gen_ring = torch.zeros((depth, self.req_pool_size), dtype=torch.int64)
+            self.copy_done = [
+                torch.get_device_module(self.device).Event() for _ in range(depth)
+            ]
+
+    def scatter(self, indices: torch.Tensor, confidence: torch.Tensor) -> None:
+        if not self.initialized:
+            self._lazy_init(confidence)
+        self.confidence_buf[indices] = confidence.to(self.confidence_buf.dtype)
+
+    def issue_ring_copy(self, *, stream, publish_ready) -> None:
+        if not self.initialized or stream is None or publish_ready is None:
+            return
+        slot = self.ring_pos % CONFIDENCE_RELAY_RING_DEPTH
+        stream.wait_event(publish_ready)
+        with torch.get_device_module(self.device).stream(stream):
+            self.conf_ring[slot].copy_(self.confidence_buf, non_blocking=True)
+            self.copy_done[slot].record()
+        self.gen_ring[slot].copy_(self.pool.req_generation)
+        self.ring_pos += 1
+
+    def resolve(
+        self, batch: ScheduleBatch, *, stream, publish_ready
+    ) -> Optional[ResolvedConfidence]:
+        if not self.initialized:
+            return None
+        draft_input = batch.spec_info
+        if draft_input is None:
+            return None
+        fi = draft_input.future_indices
+        if fi is None or fi.shape[0] == 0:
+            return None
+
+        if stream is None or publish_ready is None:
+            idx = batch.req_pool_indices
+            idx_cpu = batch.req_pool_indices_cpu
+            return ResolvedConfidence(
+                confidence=self.confidence_buf[idx].cpu(),
+                generation=self.pool.req_generation[idx_cpu].clone(),
+            )
+
+        if self.ring_pos < CONFIDENCE_RELAY_RING_LAG:
+            return None
+        slot = (self.ring_pos - CONFIDENCE_RELAY_RING_LAG) % CONFIDENCE_RELAY_RING_DEPTH
+        if not self.copy_done[slot].query():
+            return None
+
+        idx_cpu = batch.req_pool_indices_cpu
+        return ResolvedConfidence(
+            confidence=self.conf_ring[slot][idx_cpu],
+            generation=self.gen_ring[slot][idx_cpu],
+        )
+
+
 class FutureMap:
     """Always-on pool-indexed relay for cross-iter values. Forward writes via
     publish/stash; next iter reads via resolve_forward_inputs / resolve_seq_lens_cpu.
@@ -136,6 +240,7 @@ class FutureMap:
         spec_algo: SpeculativeAlgorithm,
         req_to_token_pool: ReqToTokenPool,
         needs_cpu_seq_lens: bool = True,
+        needs_confidence_relay: bool = False,
     ):
         # Bufs indexed by req_pool_idx; slot 0 mirrors KV padding row so
         # CUDA-graph padded batches (req_pool_idx == 0) are harmless.
@@ -144,6 +249,7 @@ class FutureMap:
         # Computed by decide_needs_cpu_seq_lens(); see that helper for the
         # full decision (per-backend flag + TBO / piecewise CG overrides).
         self.needs_cpu_seq_lens = needs_cpu_seq_lens
+        self.needs_confidence_relay = needs_confidence_relay
         self.req_pool_size = req_to_token_pool.req_to_token.shape[0]
 
         if _DEBUG_ASSERT:
@@ -181,6 +287,12 @@ class FutureMap:
         # resolve; arm/consume strictly alternate across all batch interleavings.
         self._publish_fresh = False
 
+        self.confidence_relay = ConfidenceRelay(
+            device=self.device,
+            req_pool_size=self.req_pool_size,
+            pool=req_to_token_pool,
+        )
+
     def _lazy_init_forward_buf(self, payload: RelayPayload):
         # Local import (see decide_needs_cpu_seq_lens): keep module-level deps leaf.
         from sglang.srt.speculative.spec_utils import spec_need_hidden_states
@@ -235,6 +347,17 @@ class FutureMap:
                 device=self.device,
             )
 
+    def resolve_confidence_cpu(
+        self, batch: ScheduleBatch
+    ) -> Optional[ResolvedConfidence]:
+        if not self.needs_confidence_relay:
+            return None
+        return self.confidence_relay.resolve(
+            batch,
+            stream=self.fwd_prepare_d2h_stream,
+            publish_ready=self.publish_ready,
+        )
+
     def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
         if self.spec_algo.is_ngram():
             # FIXME: remove once precomputed draft is supported.
@@ -339,11 +462,19 @@ class FutureMap:
             # mirror is not poisoned.
             _assert_nonneg_and_invalidate(batch.seq_lens, self.new_seq_lens_buf, fi)
 
-    def publish(self, future_indices: torch.Tensor, new_seq_lens: torch.Tensor) -> None:
+    def publish(
+        self,
+        future_indices: torch.Tensor,
+        new_seq_lens: torch.Tensor,
+        confidence: Optional[torch.Tensor] = None,
+    ) -> None:
         indices = future_indices
         if indices.shape[0] == 0:
             return  # DP idle
         self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype)
+        publish_confidence = self.needs_confidence_relay and confidence is not None
+        if publish_confidence:
+            self.confidence_relay.scatter(indices, confidence)
         # Only spec_v2 needs the event; it gates the seq_lens D2H on the private stream.
         if self.spec_algo.is_some():
             device_module = torch.get_device_module(self.device)
@@ -356,6 +487,11 @@ class FutureMap:
                 device_module.current_stream().wait_event(self.publish_ready)
             self.publish_ready.record()
             self._publish_fresh = True
+        if publish_confidence:
+            self.confidence_relay.issue_ring_copy(
+                stream=self.fwd_prepare_d2h_stream,
+                publish_ready=self.publish_ready,
+            )
 
     def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None:
         if self.spec_algo.is_ngram():

python/sglang/srt/managers/schedule_batch.py

@@ -965,11 +965,17 @@ class Req(ReqDllmMixin):
         # Per-request count of accepted draft tokens (excludes the bonus token).
         self.spec_num_correct_drafts = 0
 
+        self.spec_num_block_accept_tokens = 0
+
+        self.spec_num_cap_tokens = 0
+
         # Acceptance histogram for speculative decoding.
         # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens.
         # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens.
         self.spec_correct_drafts_histogram: List[int] = []
 
+        self.spec_cap_lens_histogram: List[int] = []
+
         # The number of times this request has been retracted / preempted.
         self.retraction_count = 0
         self.retraction_mb_id = None
@@ -1094,6 +1100,14 @@ class Req(ReqDllmMixin):
             )
         self.spec_correct_drafts_histogram[num_correct_drafts] += 1
 
+    def update_spec_cap_lens_histogram(self, cap_len: int):
+        cap_len = int(cap_len)
+        if len(self.spec_cap_lens_histogram) <= cap_len:
+            self.spec_cap_lens_histogram.extend(
+                [0] * (cap_len - len(self.spec_cap_lens_histogram) + 1)
+            )
+        self.spec_cap_lens_histogram[cap_len] += 1
+
     def extend_image_inputs(self, image_inputs):
         if self.multimodal_inputs is None:
             self.multimodal_inputs = image_inputs
@@ -1865,6 +1879,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     can_run_dp_cuda_graph: bool = False
     can_run_dp_breakable_cuda_graph: bool = False
     tbo_split_seq_index: Optional[int] = None
+    spec_verify_tier_num_tokens: int = -1
 
     # For processing logprobs
     return_logprob: bool = False
@@ -1910,6 +1925,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     # For DP attention
     global_num_tokens: Optional[List[int]] = None
     global_num_tokens_for_logprob: Optional[List[int]] = None
+    global_spec_verify_tier_num_tokens: Optional[List[int]] = None
 
     # === Compound crossing to ForwardBatch (carry their own device tensors) ===
     # Sampling info

python/sglang/srt/managers/scheduler.py

@@ -154,6 +154,7 @@ from sglang.srt.managers.min_free_slots_delayer import (
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
 from sglang.srt.managers.overlap_utils import (
     RelayPayload,
+    decide_needs_confidence_relay,
     decide_needs_cpu_seq_lens,
     resolve_forward_inputs,
 )
@@ -884,7 +885,7 @@ class Scheduler(
         min_free_slots = resolve_min_free_slots(
             self.server_args.min_free_slots_delay,
             self.max_running_requests,
-            is_dflash=self.spec_algorithm.is_dflash(),
+            is_dflash_family=self.spec_algorithm.is_dflash_family(),
         )
         if min_free_slots is not None:
             self.min_free_slots_delayer = MinFreeSlotsDelayer(
@@ -1252,10 +1253,22 @@ class Scheduler(
         else:
             attn_backends = (self.tp_worker.model_runner.attn_backend,)
         needs_cpu_seq_lens = decide_needs_cpu_seq_lens(self.server_args, attn_backends)
+        needs_confidence_relay = decide_needs_confidence_relay(self.server_args)
         self.future_map = self.spec_algorithm.create_future_map(
             self.device,
             self.req_to_token_pool,
             needs_cpu_seq_lens=needs_cpu_seq_lens,
+            needs_confidence_relay=needs_confidence_relay,
+        )
+
+        self._confidence_budget_prepare = None
+        if (
+            needs_confidence_relay
+            and self.enable_overlap
+            and self.draft_worker is not None
+        ):
+            self._confidence_budget_prepare = (
+                self.draft_worker.get_confidence_budget_prepare()
             )
 
         if use_mlx():
@@ -2173,7 +2186,7 @@ class Scheduler(
             self._add_request_to_queue(req)
             return
 
-        if self.spec_algorithm.is_dflash():
+        if self.spec_algorithm.is_dflash_family():
             error_msg = validate_dflash_request(req, self.enable_overlap)
             if error_msg is not None:
                 req.set_finish_with_abort(error_msg)
@@ -3242,6 +3255,8 @@ class Scheduler(
                 # Self-gates on batch.spec_info.future_indices; non-spec_v2
                 # no-ops (ForwardBatch.init_new lazily computes the sum).
                 self.future_map.resolve_seq_lens_cpu(batch)
+                if self._confidence_budget_prepare is not None:
+                    self._confidence_budget_prepare(batch, self.future_map)
 
                 with self.forward_stream_ctx:
                     self.forward_stream.wait_stream(self.schedule_stream)
@@ -3805,6 +3820,11 @@ class Scheduler(
         if RECORD_STEP_TIME:
             ret["step_time_dict"] = self.metrics_reporter.step_time_dict
 
+        if self.spec_algorithm.is_dspark() and self.draft_worker is not None:
+            info_record = self.draft_worker.dump_info_records()
+            if info_record is not None:
+                ret["dspark_info_record"] = info_record
+
         # This field is not serializable.
         ret.pop("model_config", None)
 
@@ -3817,6 +3837,8 @@ class Scheduler(
                 "pp_max_micro_batch_size",
                 "speculative_accept_threshold_single",
                 "speculative_accept_threshold_acc",
+                "dspark_force_budget_frac",
+                "dspark_clear_info_records",
             ]
         )
 
@@ -3834,6 +3856,30 @@ class Scheduler(
                 )
                 if_success = False
                 break
+            elif k == "dspark_force_budget_frac":
+                if not self.spec_algorithm.is_dspark() or not hasattr(
+                    self.draft_worker, "set_dspark_forced_budget_frac"
+                ):
+                    logging.warning(
+                        "dspark_force_budget_frac requires a DSpark draft worker."
+                    )
+                    if_success = False
+                    break
+                if v is not None and not (0.0 < float(v) <= 1.0):
+                    logging.warning(
+                        f"dspark_force_budget_frac must be in (0, 1] or null, got {v}."
+                    )
+                    if_success = False
+                    break
+            elif k == "dspark_clear_info_records":
+                if not self.spec_algorithm.is_dspark() or not hasattr(
+                    self.draft_worker, "clear_info_records"
+                ):
+                    logging.warning(
+                        "dspark_clear_info_records requires a DSpark draft worker."
+                    )
+                    if_success = False
+                    break
 
         if if_success:
             if (
@@ -3848,7 +3894,18 @@ class Scheduler(
             self.metrics_reporter.spec_total_num_accept_tokens = (
                 self.metrics_reporter.spec_total_num_forward_ct
             ) = 0
-            get_server_args().override(source="update_server_args", **server_args_dict)
+            # DSpark control keys are worker commands, not server args; route
+            # them to the draft worker and keep them out of the override.
+            remaining = dict(server_args_dict)
+            frac = remaining.pop("dspark_force_budget_frac", None)
+            if "dspark_force_budget_frac" in server_args_dict:
+                self.draft_worker.set_dspark_forced_budget_frac(
+                    None if frac is None else float(frac)
+                )
+            if remaining.pop("dspark_clear_info_records", None):
+                self.draft_worker.clear_info_records()
+            if remaining:
+                get_server_args().override(source="update_server_args", **remaining)
             logger.info(f"Global server args updated! {get_server_args()=}")
 
         server_args = dict(vars(get_server_args()))

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

@@ -18,6 +18,7 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
+    FINISH_MATCHED_TOKEN,
     Req,
     ScheduleBatch,
 )
@@ -26,6 +27,7 @@ from sglang.srt.mem_cache.common import (
     release_kv_cache,
 )
 from sglang.srt.runtime_context import get_server_args
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
 from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 
@@ -542,6 +544,17 @@ class SchedulerBatchResultProcessor:
         result.num_correct_drafts = sum(accept_lens) - len(batch.reqs)
         result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens]
 
+        block_accept_lens = (
+            result.block_accept_lens.tolist()
+            if result.block_accept_lens is not None
+            else None
+        )
+        result.num_block_accept_tokens = (
+            sum(block_accept_lens) if block_accept_lens else 0
+        )
+        cap_lens = result.cap_lens.tolist() if result.cap_lens is not None else None
+        result.num_cap_tokens = sum(cap_lens) if cap_lens else 0
+
         # Feed the adaptive controller now that accept_lens is on CPU,
         # instead of doing a synchronous GPU→CPU copy in the worker hot path.
         # BaseSpecWorker provides a no-op default for non-adaptive workers.
@@ -579,6 +592,12 @@ class SchedulerBatchResultProcessor:
                 req.spec_num_correct_drafts += num_correct_drafts
                 req.update_spec_correct_drafts_histogram(num_correct_drafts)
 
+                if block_accept_lens is not None:
+                    req.spec_num_block_accept_tokens += block_accept_lens[i]
+                if cap_lens is not None:
+                    req.spec_num_cap_tokens += cap_lens[i]
+                    req.update_spec_cap_lens_histogram(cap_lens[i])
+
             predict_tokens.append(accept_tokens)
 
         return predict_tokens
@@ -656,7 +675,10 @@ class SchedulerBatchResultProcessor:
         self.metrics_reporter.num_generated_tokens += len(batch.reqs)
         if not batch.spec_algorithm.is_none():
             self.metrics_reporter.update_spec_metrics(
-                batch.batch_size(), result.num_correct_drafts
+                batch.batch_size(),
+                result.num_correct_drafts,
+                num_block_accept_tokens=result.num_block_accept_tokens,
+                num_cap_tokens=result.num_cap_tokens,
             )
         if self.server_args.enable_metrics:
             self.metrics_collector.increment_decode_cuda_graph_pass(
@@ -828,6 +850,14 @@ class SchedulerBatchResultProcessor:
             self.decode_offload_manager.offload_kv_cache(req)
 
         if req.finished():
+            # isinstance narrowing: create_worker may also return plain
+            # TpModelWorker-based drafts, which carry no spec-worker hooks.
+            if isinstance(self.draft_worker, BaseSpecWorker):
+                self.draft_worker.note_request_finished(
+                    rid=req.rid,
+                    natural_stop=isinstance(req.finished_reason, FINISH_MATCHED_TOKEN),
+                )
+
             # delete feature to save memory
             if req.multimodal_inputs is not None and req.session is None:
                 req.multimodal_inputs.release_features()

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

@@ -139,6 +139,8 @@ class SchedulerMetricsReporter:
         self.spec_num_forward_ct = 0
         self.spec_total_num_accept_tokens = 0  # lifetime
         self.spec_total_num_forward_ct = 0
+        self.spec_num_block_accept_tokens = 0
+        self.spec_num_cap_tokens = 0
 
         # For PD disaggregation
         self.kv_transfer_speed_gb_s: float = 0.0
@@ -348,9 +350,17 @@ class SchedulerMetricsReporter:
             "num_draft_tokens": num_draft_tokens or 0,
         }
 
-    def update_spec_metrics(self, bs: int, num_correct_drafts: int):
+    def update_spec_metrics(
+        self,
+        bs: int,
+        num_correct_drafts: int,
+        num_block_accept_tokens: int = 0,
+        num_cap_tokens: int = 0,
+    ):
         self.spec_num_accept_tokens += num_correct_drafts + bs
         self.spec_num_forward_ct += bs
+        self.spec_num_block_accept_tokens += num_block_accept_tokens
+        self.spec_num_cap_tokens += num_cap_tokens
 
         # Bonus tokens updated elsewhere
         self.num_generated_tokens += num_correct_drafts
@@ -510,6 +520,8 @@ class SchedulerMetricsReporter:
         self.spec_num_forward_ct = 0
         self.spec_total_num_accept_tokens = 0
         self.spec_total_num_forward_ct = 0
+        self.spec_num_block_accept_tokens = 0
+        self.spec_num_cap_tokens = 0
 
     def report_prefill_stats(
         self,
@@ -733,6 +745,8 @@ class SchedulerMetricsReporter:
         if self.scheduler.spec_algorithm.is_none():
             spec_accept_length = 0
             spec_accept_rate = 0
+            spec_cap_length = 0
+            spec_block_accept_length = 0
         else:
             spec_accept_length = self.spec_num_accept_tokens / self.spec_num_forward_ct
             num_correct_drafts = self.spec_num_accept_tokens - self.spec_num_forward_ct
@@ -746,10 +760,38 @@ class SchedulerMetricsReporter:
             spec_accept_rate = (
                 num_correct_drafts / total_draft_tokens if total_draft_tokens > 0 else 0
             )
+            spec_cap_length = (
+                self.spec_num_cap_tokens / self.spec_num_forward_ct
+                if self.spec_num_forward_ct > 0
+                else 0
+            )
+            from sglang.srt.speculative.ragged_verify import (
+                RaggedVerifyMode,
+                read_ragged_verify_mode,
+            )
+
+            spec_block_accept_length = (
+                self.spec_num_block_accept_tokens / self.spec_num_forward_ct
+                if self.spec_num_forward_ct > 0
+                and read_ragged_verify_mode() is RaggedVerifyMode.CAP_ACCEPT
+                else 0
+            )
             self.spec_total_num_accept_tokens += self.spec_num_accept_tokens
             self.spec_total_num_forward_ct += self.spec_num_forward_ct
             self.spec_num_accept_tokens = self.spec_num_forward_ct = 0
+            self.spec_num_block_accept_tokens = 0
+            self.spec_num_cap_tokens = 0
             msg += f"accept len: {spec_accept_length:.2f}, accept rate: {spec_accept_rate:.2f}, "
+            if spec_cap_length > 0:
+                msg += f"cap len: {spec_cap_length:.2f}, "
+            if spec_block_accept_length > 0:
+                msg += f"block accept len: {spec_block_accept_length:.2f}, "
+            if self.scheduler.spec_algorithm.is_dspark():
+                draft_worker = self.scheduler.draft_worker
+                if draft_worker is not None:
+                    estimate_suffix = draft_worker.block_accept_estimate_log_suffix()
+                    if estimate_suffix:
+                        msg += f"{estimate_suffix}, "
 
             if self.current_scheduler_metrics_enabled:
                 spec_snapshot = self._active_spec_config_snapshot()
@@ -825,6 +867,8 @@ class SchedulerMetricsReporter:
             # Speculative decoding
             self.stats.spec_accept_length = spec_accept_length
             self.stats.spec_accept_rate = spec_accept_rate
+            self.stats.spec_cap_length = spec_cap_length
+            self.stats.spec_block_accept_length = spec_block_accept_length
             self.stats.spec_num_steps = spec_num_steps
             self.stats.spec_num_draft_tokens = spec_num_draft_tokens
 

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

@@ -278,7 +278,10 @@ class _GenerationStreamAccumulator:
     video_tokens: list = field(default_factory=list)
     spec_verify_ct: list = field(default_factory=list)
     spec_num_correct_drafts: list = field(default_factory=list)
+    spec_num_block_accept_tokens: list = field(default_factory=list)
+    spec_num_cap_tokens: list = field(default_factory=list)
     spec_correct_drafts_histogram: list = field(default_factory=list)
+    spec_cap_lens_histogram: list = field(default_factory=list)
     retraction_counts: list = field(default_factory=list)
     output_hidden_states: Optional[list] = None
     routed_experts: Optional[list] = None
@@ -406,7 +409,10 @@ class _GenerationStreamAccumulator:
         if not self.spec_algorithm.is_none():
             self.spec_verify_ct.append(req.spec_verify_ct)
             self.spec_num_correct_drafts.append(req.spec_num_correct_drafts)
+            self.spec_num_block_accept_tokens.append(req.spec_num_block_accept_tokens)
+            self.spec_num_cap_tokens.append(req.spec_num_cap_tokens)
             self.spec_correct_drafts_histogram.append(req.spec_correct_drafts_histogram)
+            self.spec_cap_lens_histogram.append(req.spec_cap_lens_histogram)
 
         if self.return_logprob:
             if (
@@ -528,7 +534,10 @@ class _GenerationStreamAccumulator:
             http_worker_ipcs=self.http_worker_ipcs,
             spec_verify_ct=self.spec_verify_ct,
             spec_num_correct_drafts=self.spec_num_correct_drafts,
+            spec_num_block_accept_tokens=self.spec_num_block_accept_tokens,
+            spec_num_cap_tokens=self.spec_num_cap_tokens,
             spec_correct_drafts_histogram=self.spec_correct_drafts_histogram,
+            spec_cap_lens_histogram=self.spec_cap_lens_histogram,
             time_stats=wrap_as_pickle(self.time_stats),
             finished_reasons=self.finished_reasons,
             decoded_texts=self.decoded_texts,

python/sglang/srt/managers/tokenizer_manager.py

@@ -32,6 +32,7 @@ from collections import deque
 from contextlib import nullcontext
 from datetime import datetime
 from enum import Enum
+from functools import lru_cache
 from http import HTTPStatus
 from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union
 
@@ -142,6 +143,19 @@ _REQUEST_STATE_WAIT_TIMEOUT = envs.SGLANG_REQUEST_STATE_WAIT_TIMEOUT.get()
 
 logger = logging.getLogger(__name__)
 
+
+@lru_cache(maxsize=1)
+def _ragged_verify_cap_accept() -> bool:
+    # The mode env is fixed at server launch; cache to keep it off the
+    # per-request metrics path.
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+
+    return read_ragged_verify_mode() is RaggedVerifyMode.CAP_ACCEPT
+
+
 _INCREMENTAL_STREAMING_META_INFO_KEYS = (
     "output_token_logprobs",
     "output_top_logprobs",
@@ -2366,6 +2380,25 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
                 meta_info["spec_num_proposed_drafts"] = num_proposed_drafts
                 meta_info["spec_verify_ct"] = recv_obj.spec_verify_ct[i]
 
+                if (
+                    getattr(recv_obj, "spec_num_cap_tokens", None) is not None
+                    and len(recv_obj.spec_num_cap_tokens) > i
+                    and recv_obj.spec_num_cap_tokens[i] > 0
+                ):
+                    meta_info["spec_cap_length"] = (
+                        recv_obj.spec_num_cap_tokens[i] / recv_obj.spec_verify_ct[i]
+                    )
+                if (
+                    _ragged_verify_cap_accept()
+                    and getattr(recv_obj, "spec_num_block_accept_tokens", None)
+                    is not None
+                    and len(recv_obj.spec_num_block_accept_tokens) > i
+                ):
+                    meta_info["spec_block_accept_length"] = (
+                        recv_obj.spec_num_block_accept_tokens[i]
+                        / recv_obj.spec_verify_ct[i]
+                    )
+
                 # FIXME: backward-compat aliases, remove in next release.
                 meta_info["spec_accepted_drafts"] = num_correct_drafts
                 meta_info["spec_proposed_drafts"] = num_proposed_drafts
@@ -2383,6 +2416,14 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
                 meta_info["spec_accept_histogram"] = (
                     recv_obj.spec_correct_drafts_histogram[i]
                 )
+            if (
+                getattr(recv_obj, "spec_cap_lens_histogram", None)
+                and len(recv_obj.spec_cap_lens_histogram) > i
+                and recv_obj.spec_cap_lens_histogram[i]
+            ):
+                meta_info["spec_cap_lens_histogram"] = recv_obj.spec_cap_lens_histogram[
+                    i
+                ]
 
     def _request_has_grammar(self, obj: GenerateReqInput) -> bool:
         return (

python/sglang/srt/managers/utils.py

@@ -44,6 +44,8 @@ class GenerationBatchResult:
     ] = None
     num_correct_drafts: int = 0  # no bonus included
     num_correct_drafts_per_req_cpu: Optional[List[int]] = None
+    num_block_accept_tokens: int = 0
+    num_cap_tokens: int = 0
     # FDFO dLLM batching: per-request accepted block length and carried algo state.
     accept_length_per_req_cpu: Optional[List[int]] = None
     dllm_algo_state: Optional[List[Any]] = None
@@ -68,6 +70,10 @@ class GenerationBatchResult:
     # sync path: forward stream -> output processor
     accept_lens: Optional[torch.Tensor] = None
 
+    block_accept_lens: Optional[torch.Tensor] = None
+
+    cap_lens: Optional[torch.Tensor] = None
+
     # Next-iter seq_lens; published via on_publish.
     new_seq_lens: Optional[torch.Tensor] = None
 
@@ -135,6 +141,12 @@ class GenerationBatchResult:
         if self.accept_lens is not None:
             self.accept_lens = _async_d2h(self.accept_lens)
 
+        if self.block_accept_lens is not None:
+            self.block_accept_lens = _async_d2h(self.block_accept_lens)
+
+        if self.cap_lens is not None:
+            self.cap_lens = _async_d2h(self.cap_lens)
+
         # Sub-objects only declare their device fields; the single copy+safety
         # primitive (_async_d2h: pinned D2H + record_stream) is injected here so
         # all device->host copying and lifetime safety lives in one place.

python/sglang/srt/mem_cache/memory_pool.py

@@ -264,6 +264,7 @@ class ReqToTokenPool:
                 (self._alloc_size, max_context_len), dtype=torch.int32, device=device
             )
         self.free_slots = list(range(1, self._alloc_size))
+        self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64)
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -295,6 +296,7 @@ class ReqToTokenPool:
         for r in reqs:
             if r.req_pool_idx is None:
                 r.req_pool_idx = select_index[offset]
+                self.req_generation[r.req_pool_idx] += 1
                 offset += 1
         return [r.req_pool_idx for r in reqs]
 
@@ -305,6 +307,7 @@ class ReqToTokenPool:
 
     def clear(self):
         self.free_slots = list(range(1, self._alloc_size))
+        self.req_generation.zero_()
 
 
 class MambaPool:

python/sglang/srt/model_executor/model_runner.py

@@ -423,9 +423,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # auxiliary hidden capture mode. TODO: expose this to server args?
         self.eagle_use_aux_hidden_state = False
         self.eagle_draft_num_layers = None
-        self.dflash_use_aux_hidden_state = False
-        self.dflash_target_layer_ids = None
-        self.dflash_draft_num_layers = None
+        self.dflash_family_use_aux_hidden_state = False
+        self.dflash_family_target_layer_ids = None
+        self.dflash_family_draft_num_layers = None
         if (
             (self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone())
             and not self.is_draft_worker
@@ -465,10 +465,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     # if there is no aux layer, set to None
                     self.eagle_aux_hidden_state_layer_ids = None
 
-        if self.spec_algorithm.is_dflash() and not self.is_draft_worker:
+        if self.spec_algorithm.is_dflash_family() and not self.is_draft_worker:
             from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
 
-            # Select target layers to capture for building DFlash context features.
+            # Select target layers to capture for building draft context features.
             draft_model_config = self._build_model_config(
                 server_args,
                 model_path=(server_args.speculative_draft_model_path),
@@ -486,8 +486,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             )
             if target_num_layers is None:
                 raise ValueError(
-                    "DFLASH requires target num_hidden_layers in config. "
-                    f"Got target={target_num_layers}."
+                    "Block-draft-with-target-kv spec requires target num_hidden_layers "
+                    f"in config. Got target={target_num_layers}."
                 )
             target_num_layers = int(target_num_layers)
 
@@ -496,19 +496,37 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 and trained_target_layers != target_num_layers
             ):
                 logger.warning(
-                    "DFLASH draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
+                    "Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
                     "selecting capture layers based on the runtime target model.",
                     trained_target_layers,
                     target_num_layers,
                 )
 
-            self.dflash_use_aux_hidden_state = True
-            self.dflash_draft_num_layers = int(draft_num_layers)
-            self.dflash_target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
+            target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
                 target_num_layers=int(target_num_layers),
                 draft_num_layers=int(draft_num_layers),
             )
 
+            if self.spec_algorithm.is_dspark():
+                from sglang.srt.speculative.dspark_components.dspark_config import (
+                    parse_dspark_draft_config,
+                )
+
+                dspark_draft_config = parse_dspark_draft_config(
+                    draft_hf_config=draft_model_config.hf_config
+                )
+                if not dspark_draft_config.require_markov():
+                    raise ValueError(
+                        "DSPARK requires markov_rank > 0 in the draft config, "
+                        f"got markov_rank={dspark_draft_config.markov_rank}."
+                    )
+                if dspark_draft_config.target_layer_ids is not None:
+                    target_layer_ids = list(dspark_draft_config.target_layer_ids)
+
+            self.dflash_family_use_aux_hidden_state = True
+            self.dflash_family_draft_num_layers = int(draft_num_layers)
+            self.dflash_family_target_layer_ids = target_layer_ids
+
         # Apply the rank zero filter to logger
         if server_args.show_time_cost:
             enable_show_time_cost()
@@ -1057,13 +1075,23 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             self.model.set_eagle3_layers_to_capture(
                 self.eagle_aux_hidden_state_layer_ids
             )
-        if self.dflash_use_aux_hidden_state:
-            if not hasattr(self.model, "set_dflash_layers_to_capture"):
+        if self.dflash_family_use_aux_hidden_state:
+            if self.spec_algorithm.is_dspark() and hasattr(
+                self.model, "set_dspark_layers_to_capture"
+            ):
+                self.model.set_dspark_layers_to_capture(
+                    self.dflash_family_target_layer_ids
+                )
+            elif hasattr(self.model, "set_dflash_layers_to_capture"):
+                self.model.set_dflash_layers_to_capture(
+                    self.dflash_family_target_layer_ids
+                )
+            else:
                 raise ValueError(
-                    f"Model {self.model.__class__.__name__} does not implement "
-                    "set_dflash_layers_to_capture, which is required for DFLASH."
+                    f"Model {self.model.__class__.__name__} implements neither "
+                    "set_dspark_layers_to_capture nor set_dflash_layers_to_capture, "
+                    "one of which is required for DFLASH/DSPARK."
                 )
-            self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids)
 
     def remote_instance_init_transfer_engine(self):
         try:

python/sglang/srt/model_executor/pool_configurator.py

@@ -154,13 +154,13 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                     * (1 + int(eagle_draft_num_layers) / int(num_layers))
                 )
 
-        # DFLASH: scale cell_size to account for draft model KV cache
-        if mr.spec_algorithm.is_dflash() and not mr.is_draft_worker:
+        # DFLASH/DSPARK: scale cell_size to account for draft model KV cache
+        if mr.spec_algorithm.is_dflash_family() and not mr.is_draft_worker:
             from sglang.srt.speculative.dflash_utils import (
                 scale_kv_cell_size_per_token_for_dflash,
             )
 
-            draft_num_layers = mr.dflash_draft_num_layers
+            draft_num_layers = mr.dflash_family_draft_num_layers
             if (
                 draft_num_layers is not None
                 and int(draft_num_layers) > 0

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -42,6 +42,7 @@ from sglang.srt.distributed.parallel_state import (
     set_pdmux_status,
 )
 from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
@@ -91,12 +92,11 @@ from sglang.srt.model_executor.runner_utils.deepep_adapter import (
 )
 from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
 from sglang.srt.runtime_context import get_flags, get_parallel
+from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout
 from sglang.srt.utils import (
     empty_context,
     get_available_gpu_memory,
     require_attn_tp_gather,
-    require_gathered_buffer,
-    require_mlp_sync,
     require_mlp_tp_gather,
 )
 from sglang.srt.utils.profile_utils import export_cuda_graph_capture_trace
@@ -112,6 +112,15 @@ logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+
+def ragged_verify_compact_graphs_enabled(spec_algorithm: SpeculativeAlgorithm) -> bool:
+    if not spec_algorithm.supports_ragged_verify():
+        return False
+    from sglang.srt.speculative.ragged_verify import ragged_verify_compact_enabled
+
+    return ragged_verify_compact_enabled()
 
 
 def build_replay_fb_view(
@@ -188,10 +197,19 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         self.enable_torch_compile = get_flags().capture.enable_torch_compile
         self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
         self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
-        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
-        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
-        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
+        self.require_mlp_tp_gather = require_mlp_tp_gather(
+            model_runner.server_args
+        ) and not self._forward_is_dp_local(model_runner)
         self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+        # Composite predicates derive from the instance values so the dp-local
+        # draft exemption above stays consistent (require_gathered_buffer ==
+        # mlp_tp_gather or attn_tp_gather; require_mlp_sync adds dp attention).
+        self.require_gathered_buffer = (
+            self.require_mlp_tp_gather or self.require_attn_tp_gather
+        )
+        self.require_mlp_sync = (
+            model_runner.server_args.enable_dp_attention or self.require_gathered_buffer
+        )
         self.enable_two_batch_overlap = (
             model_runner.server_args.enable_two_batch_overlap
         )
@@ -257,6 +275,29 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         if KTRANSFORMERS_AVAILABLE:
             KTMoEWrapper.set_capture_batch_sizes(self.capture_bs)
 
+        self.ragged_verify_mode = (
+            ragged_verify_compact_graphs_enabled(self.model_runner.spec_algorithm)
+            and (self.capture_forward_mode == ForwardMode.TARGET_VERIFY)
+            and not self.model_runner.is_draft_worker
+        )
+        self.capture_num_tokens: Optional[list[int]] = (
+            self._build_ragged_verify_token_buckets()
+            if self.ragged_verify_mode
+            else None
+        )
+        self._ragged_graph_size = 0
+        if self.ragged_verify_mode and (
+            self.enable_two_batch_overlap
+            or model_runner.server_args.enable_lora
+            or self.disable_padding
+        ):
+            raise ValueError(
+                "Compact ragged verify does not support two-batch-overlap, "
+                "LoRA, or disable-cuda-graph-padding (bs pads to the captured "
+                "tier); disable SGLANG_RAGGED_VERIFY_MODE or the conflicting "
+                "feature."
+            )
+
         # If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup
         if self.enable_return_hidden_states:
             self.capture_hidden_mode = CaptureHiddenMode.FULL
@@ -362,6 +403,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _build_ragged_verify_token_buckets(self) -> list[int]:
+        buckets = sorted({bs * self.num_tokens_per_bs for bs in self.capture_bs})
+        assert buckets and buckets[0] > 0, f"{buckets=}"
+        return buckets
+
     def _autotune_buffers(self):
         """Reuse these static decode buffers (sized to max_bs) for the warmup
         flashinfer-autotune dummy forward instead of allocating a throwaway set
@@ -385,13 +431,16 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
     def _cache_loc_dtype(self):
         return torch.int64
 
-    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
+    def _make_graph_key(self, size, stream_idx=None, variant_label=None):
         return ShapeKey(
-            size=bs,
+            size=size,
             stream_idx=stream_idx,
             variant_label=variant_label,
         )
 
+    def _capture_graph_size(self, *, bs: int, num_tokens: int) -> int:
+        return num_tokens if self.ragged_verify_mode else bs
+
     def _resolve_lora_variant(self, forward_batch: ForwardBatch):
         if not getattr(self, "record_nolora_graph", False):
             return None
@@ -401,16 +450,69 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             return "lora"
         return "nolora"
 
+    @staticmethod
+    def _forward_is_dp_local(model_runner) -> bool:
+        """The DSpark dense draft runs attn-TP-local (draft_tp_context): each
+        DP rank drafts independently with no cross-DP collective, so its
+        hand-built batches carry no dp-global metadata and must key graphs by
+        local batch size. Everything else keeps the dp-global padding path."""
+        if not model_runner.is_draft_worker:
+            return False
+        if not model_runner.spec_algorithm.is_dspark():
+            return False
+        from sglang.srt.speculative.dspark_components.dspark_config import (
+            draft_is_deepseek_v4,
+        )
+
+        return not draft_is_deepseek_v4(server_args=model_runner.server_args)
+
+    def _ragged_capture_slots(self, num_tokens: int) -> int:
+        if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get():
+            return num_tokens // self.num_tokens_per_bs
+        return min(num_tokens, self.max_bs)
+
+    def _capture_ragged_verify_layout(self, num_tokens: int):
+        if not self.ragged_verify_mode:
+            return None
+        if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get():
+            return None
+        from sglang.srt.speculative.ragged_verify import (
+            RaggedVerifyLayout,
+            build_capture_verify_lens,
+        )
+
+        verify_lens_cpu = build_capture_verify_lens(
+            num_tokens=num_tokens,
+            num_slots=self._ragged_capture_slots(num_tokens),
+            num_draft_tokens=self.num_tokens_per_bs,
+        )
+        return RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=verify_lens_cpu,
+            device=self.device,
+            grid=self.capture_num_tokens,
+        )
+
     def can_run_graph(self, forward_batch: ForwardBatch):
         # Disable for token embedding overrides (dynamic per-request)
         if forward_batch.replace_embeds is not None:
             return False
+
+        ragged_layout = (
+            resolve_ragged_verify_layout(forward_batch)
+            if self.ragged_verify_mode
+            else None
+        )
+        if ragged_layout is not None:
+            return self._can_run_ragged_verify_graph(forward_batch, ragged_layout)
+        if self.ragged_verify_mode and forward_batch.forward_mode.is_target_verify():
+            return False
+
         if self.require_mlp_tp_gather:
             cuda_graph_bs = (
                 max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
                 if self.model_runner.spec_algorithm.is_eagle()
                 or self.model_runner.spec_algorithm.is_standalone()
-                or self.model_runner.spec_algorithm.is_dflash()
+                or self.model_runner.spec_algorithm.is_dflash_family()
                 else max(forward_batch.global_num_tokens_cpu)
             )
         else:
@@ -472,6 +574,46 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             and is_ngram_supported
         )
 
+    def _can_run_ragged_verify_graph(self, forward_batch: ForwardBatch, ragged_layout):
+        if not self.attn_backend.supports_ragged_verify_graph:
+            return False
+
+        admission_tokens = ragged_layout.graph_num_tokens
+        is_tokens_supported = admission_tokens <= self.capture_num_tokens[
+            -1
+        ] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens)
+
+        is_dp_supported = (
+            forward_batch.can_run_dp_cuda_graph if self.require_mlp_sync else True
+        )
+
+        is_encoder_lens_supported = (
+            torch.all(forward_batch.encoder_lens > 0)
+            if self.is_encoder_decoder
+            else True
+        )
+
+        requested_capture_hidden_mode = max(
+            forward_batch.capture_hidden_mode,
+            (
+                forward_batch.spec_info.capture_hidden_mode
+                if getattr(forward_batch.spec_info, "capture_hidden_mode", None)
+                is not None
+                else CaptureHiddenMode.NULL
+            ),
+        )
+        capture_hidden_mode_matches = (
+            requested_capture_hidden_mode == CaptureHiddenMode.NULL
+            or requested_capture_hidden_mode == self.capture_hidden_mode
+        )
+
+        return (
+            is_tokens_supported
+            and is_dp_supported
+            and is_encoder_lens_supported
+            and capture_hidden_mode_matches
+        )
+
     def _init_profile_context_and_memory_record(self):
         profile_context = profile(
             activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
@@ -510,16 +652,21 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         self,
         size: int,
         stream_idx: Optional[int] = None,
+        num_tokens: Optional[int] = None,
     ):
         """Build the dummy decode ForwardBatch for capture at size (=bs),
         populate static input buffers, choose the active attn backend, and
         optionally build pp_proxy_tensors.
 
+        num_tokens defaults to the uniform bs * num_tokens_per_bs; ragged
+        verify capture passes the decoupled (slots, tier tokens) pair.
+
         Returns (forward_batch, attn_backend, pp_proxy_tensors);
         pp_proxy_tensors is None unless pp_size > 1.
         """
         bs = size
         buffers: DecodeInputBuffers = self.buffers
+        if num_tokens is None:
             num_tokens = bs * self.num_tokens_per_bs
 
         # Registry-owned FB-shared slots come through the registry (which
@@ -754,8 +901,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
     ):
-        bs = size
-        num_tokens = bs * self.num_tokens_per_bs
+        num_tokens = size * self.num_tokens_per_bs
+        bs = self._ragged_capture_slots(num_tokens) if self.ragged_verify_mode else size
 
         # Sanity-check: --debug-cuda-graph requires breakable backend.
         if self.model_runner.server_args.debug_cuda_graph:
@@ -764,7 +911,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             ), "Breakable CUDA graph is required for --debug-cuda-graph"
 
         forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare(
-            size, stream_idx=stream_idx
+            bs, stream_idx=stream_idx, num_tokens=num_tokens
         )
 
         # All setup hooks below read get_attn_backend() (TboForwardBatchPreparer,
@@ -808,9 +955,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                         {k: v.clone() for k, v in pp_proxy_tensors.tensors.items()}
                     )
                 if (
-                    self.model_runner.spec_algorithm.is_dflash()
+                    self.model_runner.spec_algorithm.is_dflash_family()
                     and self.model_runner.is_draft_worker
                     and "input_embeds" in inspect.signature(forward).parameters
+                    and not hasattr(self.model_runner.model, "forward_embed")
                 ):
                     kwargs["input_embeds"] = self.buffers.input_embeds[:num_tokens]
 
@@ -835,7 +983,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             # wires no buffer here. (SWA write loc rides the `swa_out_cache_loc` rail.)
 
             with canary_ctx:
-                shape_key = self._make_graph_key(bs, stream_idx, variant_label)
+                shape_key = self._make_graph_key(
+                    self._capture_graph_size(bs=bs, num_tokens=num_tokens),
+                    stream_idx,
+                    variant_label,
+                )
                 post_warmup_hook = getattr(
                     self.model_runner.attn_backend,
                     "on_after_cuda_graph_warmup",
@@ -892,15 +1044,35 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         forward_batch: ForwardBatch,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ):
+        ragged_layout = (
+            resolve_ragged_verify_layout(forward_batch)
+            if self.ragged_verify_mode
+            else None
+        )
+        is_ragged = ragged_layout is not None
+
         self.deepep_adapter.replay()
 
         if not forward_batch.needs_forward_metadata_init():
             # Pre-planned (plan-stream load_batch already ran).
             # In speculative decoding, these two fields are still needed.
+            graph_size_key = (
+                self._ragged_graph_size
+                if is_ragged
+                else self._capture_graph_size(
+                    bs=self.bs, num_tokens=self.bs * self.num_tokens_per_bs
+                )
+            )
+            if is_ragged:
+                assert self.raw_num_token == ragged_layout.graph_num_tokens, (
+                    f"stale ragged raw_num_token {self.raw_num_token} != "
+                    f"{ragged_layout.graph_num_tokens}"
+                )
             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
             if (
-                self.model_runner.spec_algorithm.is_dflash()
+                not is_ragged
+                and self.model_runner.spec_algorithm.is_dflash_family()
                 and self.model_runner.is_draft_worker
                 and forward_batch.input_embeds is not None
             ):
@@ -910,7 +1082,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             variant_label = self._resolve_lora_variant(forward_batch)
             stream_idx = get_current_stream_idx() if self.enable_pdmux else None
             self._replay_graph_key = self._make_graph_key(
-                self.bs, stream_idx, variant_label
+                graph_size_key, stream_idx, variant_label
             )
             return
 
@@ -918,37 +1090,57 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         self.recapture_if_needed(forward_batch)
 
         raw_bs = forward_batch.batch_size
-        raw_num_token = raw_bs * self.num_tokens_per_bs
 
+        if is_ragged:
+            raw_num_token = ragged_layout.graph_num_tokens
+            graph_size_key = self._ragged_graph_num_tokens(raw_num_token)
+            assert graph_size_key == ragged_layout.graph_num_tokens, (
+                f"ragged verify tier mismatch: runner tier {graph_size_key} != "
+                f"layout graph_num_tokens {ragged_layout.graph_num_tokens}"
+            )
+            bs = self._ragged_capture_slots(graph_size_key)
+            assert bs >= raw_bs, (
+                f"ragged capture slots {bs} (tier {graph_size_key}) < raw_bs "
+                f"{raw_bs}; the planner must reject this batch before replay"
+            )
+            padded_num_tokens = graph_size_key
+        else:
+            raw_num_token = raw_bs * self.num_tokens_per_bs
             if self.require_mlp_tp_gather:
                 max_num_tokens = max(forward_batch.global_num_tokens_cpu)
                 max_batch_size = (
                     max_num_tokens / self.num_tokens_per_bs
                     if self.model_runner.spec_algorithm.is_eagle()
                     or self.model_runner.spec_algorithm.is_standalone()
-                or self.model_runner.spec_algorithm.is_dflash()
+                    or self.model_runner.spec_algorithm.is_dflash_family()
                     else max_num_tokens
                 )
                 bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
             else:
                 bs = self._pad_to_bucket(raw_bs, self.capture_bs)
+            padded_num_tokens = bs * self.num_tokens_per_bs
+            graph_size_key = self._capture_graph_size(
+                bs=bs, num_tokens=padded_num_tokens
+            )
 
         self.buffer_registry.fill_from(
             forward_batch,
             raw_bs=raw_bs,
             padded_bs=bs,
             raw_num_tokens=raw_num_token,
-            padded_num_tokens=bs * self.num_tokens_per_bs,
+            padded_num_tokens=padded_num_tokens,
             pp_proxy_tensors=pp_proxy_tensors,
         )
 
         if (
-            self.model_runner.spec_algorithm.is_dflash()
+            not is_ragged
+            and self.model_runner.spec_algorithm.is_dflash_family()
             and self.model_runner.is_draft_worker
             and forward_batch.input_embeds is not None
         ):
             buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds)
-        # Padded tokens aren't read, so skip zeroing them.
+        # Padded tokens aren't read, so skip zeroing. Ragged input_ids arrive
+        # from the planner already padded to the tier, invalid slots zeroed.
         if self.enable_two_batch_overlap:
             self.tbo_plugin.replay_prepare(
                 forward_mode=self.capture_forward_mode,
@@ -956,7 +1148,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 num_token_non_padded=len(forward_batch.input_ids),
                 spec_info=forward_batch.spec_info,
             )
-        if forward_batch.forward_mode.is_idle() and forward_batch.spec_info is not None:
+        if (
+            not is_ragged
+            and forward_batch.forward_mode.is_idle()
+            and forward_batch.spec_info is not None
+        ):
             forward_batch.spec_info.custom_mask = buffers.custom_mask
         if self.enable_pdmux:
             stream_idx = get_current_stream_idx()
@@ -968,17 +1164,18 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             buffers=buffers,
             bs=bs,
             raw_bs=raw_bs,
-            num_tokens=bs * self.num_tokens_per_bs,
+            num_tokens=padded_num_tokens,
             seq_len_fill_value=self.seq_len_fill_value,
             capture_forward_mode=self.capture_forward_mode,
             is_encoder_decoder=self.is_encoder_decoder,
         )
         attn_backend.init_forward_metadata_out_graph(fb_view)
 
-        # Store fields
         self.raw_bs = raw_bs
         self.raw_num_token = raw_num_token
         self.bs = bs
+        if is_ragged:
+            self._ragged_graph_size = graph_size_key
 
         if self.model_runner.hisparse_coordinator is not None:
             self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs)
@@ -986,9 +1183,14 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         variant_label = self._resolve_lora_variant(forward_batch)
         stream_idx = get_current_stream_idx() if self.enable_pdmux else None
         self._replay_graph_key = self._make_graph_key(
-            self.bs, stream_idx, variant_label
+            graph_size_key, stream_idx, variant_label
         )
 
+    def _ragged_graph_num_tokens(self, total_verify_tokens: int) -> int:
+        from sglang.srt.speculative.ragged_verify import round_up_grid
+
+        return round_up_grid(total_verify_tokens, self.capture_num_tokens)
+
     def execute(
         self,
         forward_batch: ForwardBatch,
@@ -1001,19 +1203,46 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             if self.model_runner.device_timer
             else contextlib.nullcontext()
         )
-        with timer_ctx, self.backend.replay_session():
-            self.load_batch(forward_batch, pp_proxy_tensors)
         # Publish a read-done event for the WAR barrier: a cuda-graph forward
         # finishes its shared req_to_token / SWA reads at this pre-replay
-            # snapshot, so plain DECODE and DFLASH TARGET_VERIFY both qualify.
-            if forward_batch.forward_mode.is_decode() or (
+        # snapshot, so plain DECODE and block-draft TARGET_VERIFY qualify.
+        publish_read_done = forward_batch.forward_mode.is_decode() or (
             forward_batch.forward_mode.is_target_verify()
-                and self.model_runner.spec_algorithm.is_dflash()
-            ):
+            and self.model_runner.spec_algorithm.is_dflash_family()
+        )
+        # Exception: breakable-graph verify replays (captured forward metadata)
+        # re-read req_to_token *during* replay, so the pre-replay snapshot is
+        # too early -- record the event after replay instead.
+        read_done_post_replay = (
+            publish_read_done
+            and forward_batch.forward_mode.is_target_verify()
+            and self.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
+        )
+        with timer_ctx, self.backend.replay_session():
+            self.load_batch(forward_batch, pp_proxy_tensors)
+            if envs.SGLANG_LOG_DECODE_GRAPH_KEY.get():
+                logger.info(
+                    "Decode graph replay: worker=%s key_size=%s (%s) mode=%s raw_bs=%d%s",
+                    "draft" if self.model_runner.is_draft_worker else "target",
+                    self._replay_graph_key.size,
+                    "num_tokens" if self.ragged_verify_mode else "bs",
+                    forward_batch.forward_mode.name,
+                    forward_batch.batch_size,
+                    (
+                        f" slots={self._ragged_capture_slots(self._replay_graph_key.size)}"
+                        if self.ragged_verify_mode
+                        else ""
+                    ),
+                )
+            if publish_read_done and not read_done_post_replay:
                 read_done = self.device_module.Event()
                 read_done.record()
                 self.model_runner.war_fastpath_read_done_event = read_done
             output = self.backend.replay(self._replay_graph_key, forward_batch)
+            if read_done_post_replay:
+                read_done = self.device_module.Event()
+                read_done.record()
+                self.model_runner.war_fastpath_read_done_event = read_done
 
         if isinstance(output, LogitsProcessorOutput):
             if self.is_dllm:
@@ -1083,7 +1312,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     dtype=self.model_runner.dtype,
                     device=self.model_runner.device,
                 )
-        elif self.model_runner.spec_algorithm.is_dflash():
+        elif self.model_runner.spec_algorithm.is_dflash_family():
             from sglang.srt.speculative.dflash_info import DFlashVerifyInput
             from sglang.srt.speculative.dflash_utils import (
                 resolve_dflash_verify_mask_policy,
@@ -1097,7 +1326,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             spec_info = DFlashVerifyInput(
                 draft_token=None,
                 positions=None,
-                draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
+                draft_token_num=self.num_tokens_per_bs,
                 custom_mask=(
                     None
                     if (self.model_runner.is_draft_worker or not build_custom_mask)
@@ -1108,6 +1337,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     if self.model_runner.is_draft_worker
                     else CaptureHiddenMode.FULL
                 ),
+                ragged_verify_layout=self._capture_ragged_verify_layout(num_tokens),
             )
 
         elif self.model_runner.spec_algorithm.is_ngram():

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -182,7 +182,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # Ported from main #27468.
         if (
             model_runner.server_args.enable_return_hidden_states
-            or model_runner.spec_algorithm.is_dflash()
+            or model_runner.spec_algorithm.is_dflash_family()
         ):
             self.capture_hidden_mode = CaptureHiddenMode.FULL
         # EAGLE captures FULL hidden states for the target and LAST for the

python/sglang/srt/models/deepseek_v4.py

@@ -2035,6 +2035,8 @@ class DeepseekV4Model(nn.Module):
         if self.dsa_enable_prefill_cp:
             self.cp_size = get_parallel().attn_cp_size
 
+        self.dspark_layers_to_capture: Optional[List[int]] = None
+
     def get_input_embeddings(self) -> nn.Module:
         return self.embed_tokens
 
@@ -2220,7 +2222,18 @@ class DeepseekV4Model(nn.Module):
             if hasattr(forward_batch, _attr):
                 delattr(forward_batch, _attr)
 
-        if self._can_run_tbo(forward_batch):
+        capture_dspark = self.dspark_layers_to_capture is not None
+        if capture_dspark and dsa_use_prefill_cp(forward_batch):
+            raise NotImplementedError(
+                "DSpark aux hidden-state capture is not supported together with "
+                "DeepSeek-V4 prefill context parallelism (attn_cp_size > 1). Disable one "
+                "of them: DSpark static-verify is CP-off for v1."
+            )
+        dspark_aux_hidden_states: List[torch.Tensor] = []
+        # DSpark aux capture needs the per-layer eager loop (TBO's overlapped
+        # execution cannot expose per-layer completed hidden states), so skip
+        # TBO when capturing -- a perf-only downgrade, not a correctness one.
+        if self._can_run_tbo(forward_batch) and not capture_dspark:
             # Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is
             # disabled here (each layer self-contained), so no trailing hc_post.
             hidden_states = self._forward_layers_tbo(
@@ -2251,6 +2264,14 @@ class DeepseekV4Model(nn.Module):
                         prev_post=prev_post,
                         prev_comb=prev_comb,
                     )
+                if capture_dspark and i in self.dspark_layers_to_capture:
+                    if use_fused:
+                        completed = layer.hc_post(
+                            hidden_states, prev_residual, prev_post, prev_comb
+                        )
+                    else:
+                        completed = hidden_states
+                    dspark_aux_hidden_states.append(completed.mean(dim=1))
             if use_fused and last_layer is not None:
                 hidden_states = last_layer.hc_post(
                     hidden_states, prev_residual, prev_post, prev_comb
@@ -2276,6 +2297,9 @@ class DeepseekV4Model(nn.Module):
         )
         hidden_states = self.norm(hidden_states)
 
+        if capture_dspark:
+            return (hidden_states, pre_hc_head), dspark_aux_hidden_states
+
         return hidden_states, pre_hc_head
 
 
@@ -2352,6 +2376,16 @@ class DeepseekV4ForCausalLM(nn.Module):
     def get_input_embeddings(self) -> nn.Module:
         return self.model.get_input_embeddings()
 
+    def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None:
+        if not self.pp_group.is_last_rank:
+            return
+        if layer_ids is None:
+            raise ValueError(
+                "DSPARK requires explicit layer_ids for aux hidden capture."
+            )
+        self.capture_aux_hidden_states = True
+        self.model.dspark_layers_to_capture = list(layer_ids)
+
     def determine_num_fused_shared_experts(self):
         self.num_fused_shared_experts = 0
         if get_server_args().disable_shared_experts_fusion:
@@ -2429,7 +2463,9 @@ class DeepseekV4ForCausalLM(nn.Module):
             self.lm_head,
             forward_batch,
             aux_hidden_states,
-            hidden_states_before_norm=pre_hc_head,
+            hidden_states_before_norm=(
+                None if aux_hidden_states is not None else pre_hc_head
+            ),
         )
 
     def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None:

python/sglang/srt/models/deepseek_v4_dspark.py

@@ -0,0 +1,892 @@
+from __future__ import annotations
+
+import logging
+from typing import Iterable, List, Optional, Tuple
+
+import msgspec
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace
+from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
+from sglang.srt.environ import envs
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.dbrx import ReplicatedLinear
+from sglang.srt.models.deepseek_v4 import (
+    DEEPSEEK_V4_STACKED_PARAMS_MAPPING,
+    DeepseekV4DecoderLayer,
+    MqaAttentionBase,
+    _dequant_fp8_wo_a,
+    hc_head_torch,
+    make_hc_head_params,
+)
+from sglang.srt.models.dspark import (
+    DSparkConfidenceHead,
+    StepSampler,
+    gather_and_crop_vocab,
+    run_markov_block,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.speculative.dspark_components.dspark_config import (
+    parse_dspark_draft_config,
+)
+from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import (
+    BuildStepLocal,
+    CommitKvProj,
+)
+from sglang.srt.speculative.ragged_verify import (
+    RaggedVerifyMode,
+    read_ragged_verify_mode,
+)
+from sglang.srt.utils import add_prefix, is_blackwell_supported
+from sglang.srt.utils.async_probe import maybe_detect_in_closed_range
+
+logger = logging.getLogger(__name__)
+
+_PAD_NUM_HEADS = 64
+
+
+def apply_rotary_emb(
+    x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
+) -> torch.Tensor:
+    y = x
+    x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs_cis = freqs_cis.conj()
+    if x.ndim == 3:
+        freqs_cis = freqs_cis.view(x.size(0), 1, x.size(-1))
+    else:
+        freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+    x = torch.view_as_real(x * freqs_cis).flatten(-2)
+    y.copy_(x)
+    return y
+
+
+class DSparkAttention(MqaAttentionBase):
+
+    def __init__(
+        self,
+        config: DeepSeekV4Config,
+        layer_id: int,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        alt_streams: Optional[List[torch.cuda.Stream]] = None,
+    ) -> None:
+        super().__init__(
+            config,
+            layer_id,
+            quant_config,
+            prefix,
+            attn_tp_rank=get_parallel().attn_tp_rank,
+            attn_tp_size=get_parallel().attn_tp_size,
+            compress_ratio=0,
+            fuse_wqa_wkv=False,
+            wo_a_fp8=False,
+            wo_a_keeps_quant_config=False,
+            wo_b_reduce_results=True,
+            rope_original_seq_len=0,
+        )
+        assert (
+            self.compress_ratio == 0
+        ), "DSpark draft attention requires compress_ratio == 0."
+        self.window_size = int(
+            getattr(config, "sliding_window", None) or config.window_size
+        )
+
+        self.attn = RadixAttention(
+            self.n_local_heads,
+            self.head_dim,
+            self.softmax_scale,
+            num_kv_heads=1,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=add_prefix("attn", prefix),
+        )
+
+        self._use_fast_kernel = envs.SGLANG_DSPARK_FAST_KERNEL.get()
+        self.alt_streams = alt_streams
+        self._multi_stream_bs_limit = 128 if is_blackwell_supported() else 64
+
+    def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor:
+        kv, _ = self.wkv(x)
+        return kv
+
+    def _local_attn_sink(self) -> torch.Tensor:
+        if self.attn_tp_size == 1:
+            return self.attn_sink
+        if self._attn_sink_local is None:
+            rank = self.attn_tp_rank
+            num_heads = self.n_local_heads
+            sink = self.attn_sink.new_zeros(max(num_heads, _PAD_NUM_HEADS))
+            sink[:num_heads] = self.attn_sink[rank * num_heads : (rank + 1) * num_heads]
+            self._attn_sink_local = sink
+        return self._attn_sink_local
+
+    def _store_block_kv(
+        self,
+        *,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend,
+        pool: DeepSeekV4TokenToKVPool,
+    ) -> None:
+        pool.set_swa_key_buffer_radix_fused_norm_rope(
+            layer_id=self.layer_id,
+            swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
+            kv=kv,
+            kv_weight=self.kv_norm.weight.data,
+            eps=self.eps,
+            freqs_cis=self.freqs_cis,
+            positions=positions,
+        )
+
+    def _compute_q(
+        self,
+        x: torch.Tensor,
+        positions: torch.Tensor,
+        q_out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        q, _ = self.wq_a(x)
+        q = self.q_norm(q)
+        q, _ = self.wq_b(q)
+        q = q.view(-1, self.n_local_heads, self.head_dim)
+        if self._use_fast_kernel:
+            if q_out is None:
+                q_out = torch.empty_like(q)
+            fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions)
+            return q_out
+        else:
+            q = q * torch.rsqrt(
+                q.float().square().mean(-1, keepdim=True) + self.eps
+            ).to(q.dtype)
+            apply_rotary_emb(q[..., -self.rope_head_dim :], self.freqs_cis[positions])
+            if q_out is not None:
+                q_out.copy_(q)
+                return q_out
+            return q
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        from sglang.srt.model_executor.forward_context import get_attn_backend
+
+        pool = _resolve_dspark_pool()
+        attn_backend = get_attn_backend()
+        rd = self.rope_head_dim
+
+        enable_multi_stream = (
+            self.alt_streams is not None
+            and get_is_capture_mode()
+            and hidden_states.shape[0] <= self._multi_stream_bs_limit
+        )
+
+        q_padded: Optional[torch.Tensor] = None
+        q_out: Optional[torch.Tensor] = None
+        if self.n_local_heads < _PAD_NUM_HEADS:
+            q_padded = hidden_states.new_empty(
+                hidden_states.shape[0], _PAD_NUM_HEADS, self.head_dim
+            )
+            q_out = q_padded[:, : self.n_local_heads, :]
+
+        if enable_multi_stream:
+            current_stream = torch.cuda.current_stream()
+            stream_kv = self.alt_streams[0]
+            stream_kv.wait_stream(current_stream)
+            with torch.cuda.stream(stream_kv):
+                kv = self.kv_proj_only(hidden_states)
+                self._store_block_kv(
+                    kv=kv,
+                    positions=positions,
+                    forward_batch=forward_batch,
+                    attn_backend=attn_backend,
+                    pool=pool,
+                )
+            q = self._compute_q(hidden_states, positions, q_out=q_out)
+            current_stream.wait_stream(stream_kv)
+        else:
+            kv = self.kv_proj_only(hidden_states)
+            self._store_block_kv(
+                kv=kv,
+                positions=positions,
+                forward_batch=forward_batch,
+                attn_backend=attn_backend,
+                pool=pool,
+            )
+            q = self._compute_q(hidden_states, positions, q_out=q_out)
+
+        if q_padded is not None:
+            q = q_padded
+        attn_sink = self._local_attn_sink()
+
+        o = attn_backend.forward(
+            q=q,
+            k=kv,
+            v=kv,
+            layer=self.attn,
+            forward_batch=forward_batch,
+            compress_ratio=0,
+            attn_sink=attn_sink,
+            save_kv_cache=False,
+        )
+        if o.shape[1] != self.n_local_heads:
+            o = o[:, : self.n_local_heads, :]
+
+        if self._use_fast_kernel:
+            fused_rope_inplace(
+                o[..., -rd:], None, self.freqs_cis, positions=positions, inverse=True
+            )
+        else:
+            apply_rotary_emb(o[..., -rd:], self.freqs_cis[positions], inverse=True)
+
+        o = o.view(
+            o.shape[0],
+            self.n_local_groups,
+            o.shape[1] * o.shape[2] // self.n_local_groups,
+        )
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        if self._use_fast_kernel:
+            o = torch.einsum("bgd,grd->bgr", o, wo_a)
+        else:
+            o = torch.einsum("bgd,grd->bgr", o.float(), wo_a.float()).to(q.dtype)
+        out, _ = self.wo_b(o.reshape(o.shape[0], o.shape[1] * o.shape[2]))
+        return out
+
+
+def _resolve_dspark_pool() -> DeepSeekV4TokenToKVPool:
+    pool = get_token_to_kv_pool()
+    assert isinstance(pool, DeepSeekV4TokenToKVPool), (
+        "DSpark draft attention requires a DeepSeekV4TokenToKVPool, "
+        f"got {type(pool).__name__}."
+    )
+    return pool
+
+
+class MarkovW2ShardGeometry(msgspec.Struct, frozen=True):
+
+    tp_size: int
+    org_vocab_start: int
+    org_vocab_end: int
+    num_embeddings_per_partition: int
+    num_embeddings_padded: int
+
+
+class DSparkV4MarkovHead(nn.Module):
+
+    markov_head_type = "vanilla"
+
+    def __init__(self, *, vocab_size: int, markov_rank: int) -> None:
+        super().__init__()
+        self.vocab_size = int(vocab_size)
+        self.markov_rank = int(markov_rank)
+        if self.markov_rank <= 0:
+            raise ValueError(
+                f"DSparkV4MarkovHead requires markov_rank > 0, got {self.markov_rank}."
+            )
+        self.markov_w1 = VocabParallelEmbedding(
+            self.vocab_size, self.markov_rank, enable_tp=False
+        )
+        self._opt_markov_w2_bf16 = envs.SGLANG_DSPARK_OPT_MARKOV_W2_BF16.get()
+        self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get()
+        markov_w2_dtype = torch.bfloat16 if self._opt_markov_w2_bf16 else torch.float32
+        self.markov_w2 = nn.Linear(
+            self.markov_rank, self.vocab_size, bias=False, dtype=markov_w2_dtype
+        )
+        self._tp_shard: Optional[MarkovW2ShardGeometry] = None
+
+    def configure_tp_shard(self, *, lm_head: nn.Module) -> None:
+        if not self._opt_markov_w2_tp_shard:
+            return
+        if int(lm_head.org_vocab_size) != self.vocab_size:
+            raise ValueError(
+                "DSpark markov_w2 TP-shard requires lm_head.org_vocab_size == "
+                f"markov vocab_size, got {int(lm_head.org_vocab_size)} vs "
+                f"{self.vocab_size}."
+            )
+        tp_size = int(lm_head.tp_size)
+        per_partition = int(lm_head.num_embeddings_per_partition)
+        num_padded = int(lm_head.num_embeddings_padded)
+        if per_partition * tp_size != num_padded:
+            raise ValueError(
+                "DSpark markov_w2 TP-shard could not align to the lm_head partition: "
+                f"num_embeddings_per_partition({per_partition}) * tp_size({tp_size}) != "
+                f"num_embeddings_padded({num_padded})."
+            )
+        attn_tp_size = get_parallel().attn_tp_group.world_size
+        if attn_tp_size != tp_size:
+            raise ValueError(
+                "DSpark markov_w2 TP-shard needs the attn-TP group (used for the per-step "
+                f"all-gather) to equal the lm_head shard group, got attn_tp_size="
+                f"{attn_tp_size} vs lm_head tp_size={tp_size}. This config (e.g. DP "
+                "attention without --enable-dp-lm-head, where lm_head shards over the "
+                "global TP group) is unsupported; disable "
+                "SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD."
+            )
+        self._tp_shard = MarkovW2ShardGeometry(
+            tp_size=tp_size,
+            org_vocab_start=int(lm_head.shard_indices.org_vocab_start_index),
+            org_vocab_end=int(lm_head.shard_indices.org_vocab_end_index),
+            num_embeddings_per_partition=per_partition,
+            num_embeddings_padded=num_padded,
+        )
+
+    def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.markov_w1(token_ids.long())
+
+    def project_bias(
+        self, latent_states: torch.Tensor, *, weight: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        weight = self.markov_w2.weight if weight is None else weight
+        if self._opt_markov_w2_bf16:
+            return F.linear(latent_states.to(weight.dtype), weight).float()
+        return F.linear(latent_states.float(), weight)
+
+    def compute_step_bias(
+        self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor]
+    ) -> torch.Tensor:
+        del hidden_states
+        return self.project_bias(self.get_prev_embeddings(token_ids))
+
+    def apply_step_logits(
+        self,
+        logits: torch.Tensor,
+        *,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if self._tp_shard is not None:
+            return self._apply_step_logits_sharded(
+                base_local=logits, token_ids=token_ids
+            )
+        return logits + self.compute_step_bias(token_ids, hidden_states)
+
+    def _apply_step_logits_sharded(
+        self, *, base_local: torch.Tensor, token_ids: torch.Tensor
+    ) -> torch.Tensor:
+        shard = self._tp_shard
+        latent = self.get_prev_embeddings(token_ids)
+        weight_local = self.markov_w2.weight[
+            shard.org_vocab_start : shard.org_vocab_end
+        ]
+        if self._opt_markov_w2_bf16:
+            bias = F.linear(latent.to(weight_local.dtype), weight_local)
+        else:
+            bias = F.linear(latent.float(), weight_local)
+        step_local = BuildStepLocal.execute(bias=bias, base_local=base_local)
+        if shard.tp_size > 1:
+            full = get_parallel().attn_tp_group.all_gather(step_local, dim=-1)
+        else:
+            full = step_local
+        return full[..., : self.vocab_size]
+
+    def forward(self, token_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        embed = self.get_prev_embeddings(token_ids)
+        logits = self.project_bias(embed)
+        return logits, embed
+
+    def sample_block(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        first_prev_tokens: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+        sampler: StepSampler,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return run_markov_block(
+            self,
+            base_logits,
+            first_prev_tokens=first_prev_tokens,
+            hidden_states=hidden_states,
+            sampler=sampler,
+        )
+
+
+def build_dspark_v4_confidence_head(
+    *, config: DeepSeekV4Config, markov_rank: int
+) -> Optional[DSparkConfidenceHead]:
+    if read_ragged_verify_mode() is RaggedVerifyMode.STATIC:
+        return None
+    if not hasattr(config, "enable_confidence_head"):
+        logger.warning(
+            "DSpark draft config has no enable_confidence_head field; treating the "
+            "confidence head as enabled."
+        )
+    with_markov_cfg = getattr(config, "confidence_head_with_markov", None)
+    with_markov = (
+        (markov_rank > 0) if with_markov_cfg is None else bool(with_markov_cfg)
+    )
+    if with_markov and markov_rank <= 0:
+        raise ValueError(
+            "DSpark V4 confidence_head_with_markov requires markov_rank > 0, "
+            f"got markov_rank={markov_rank}."
+        )
+    return DSparkConfidenceHead(
+        hidden_size=int(config.hidden_size),
+        markov_rank=int(markov_rank),
+        with_markov=with_markov,
+        bias=False,
+    )
+
+
+class DSparkV4Stage(DeepseekV4DecoderLayer):
+
+    def __init__(
+        self,
+        config: DeepSeekV4Config,
+        layer_id: int,
+        stage_id: int,
+        num_stages: int,
+        num_target_layers: int,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        alt_streams: Optional[List[torch.cuda.Stream]] = None,
+    ) -> None:
+        super().__init__(
+            config=config,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=prefix,
+            is_nextn=True,
+            alt_streams=alt_streams,
+        )
+        self.stage_id = stage_id
+        self.dim = config.hidden_size
+
+        if stage_id == 0:
+            if num_target_layers <= 0:
+                raise ValueError(
+                    "DSpark needs target layers for the target-hidden projection."
+                )
+            self.main_proj = ReplicatedLinear(
+                config.hidden_size * num_target_layers,
+                config.hidden_size,
+                bias=False,
+                quant_config=quant_config,
+                prefix=add_prefix("main_proj", prefix),
+            )
+            self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        if stage_id == num_stages - 1:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            (
+                self.hc_head_fn,
+                self.hc_head_base,
+                self.hc_head_scale,
+            ) = make_hc_head_params(config.hc_mult, config.hidden_size)
+
+    def _build_self_attn(
+        self,
+        *,
+        config: DeepSeekV4Config,
+        layer_id: int,
+        quant_config: Optional[QuantizationConfig],
+        prefix: str,
+        alt_streams: Optional[List[torch.cuda.Stream]],
+        compress_ratio_override: Optional[int],
+    ) -> nn.Module:
+        del compress_ratio_override
+        return DSparkAttention(
+            config=config,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=prefix,
+            alt_streams=alt_streams,
+        )
+
+    def _hc_pre_block(
+        self,
+        x: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        y, post, comb, _ = self.hc_pre(x, hc_fn, hc_scale, hc_base)
+        return y, post, comb
+
+    def _hc_post_block(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.hc_post(x, residual, post, comb)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        x, post, comb = self._hc_pre_block(
+            hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
+        )
+        x = self.input_layernorm(x)
+        x = self.self_attn(positions, x, forward_batch)
+        x = self._hc_post_block(x, residual, post, comb)
+
+        residual = x
+        x, post, comb = self._hc_pre_block(
+            x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
+        )
+        x = self.post_attention_layernorm(x)
+        x = self._run_ffn(x, forward_batch)
+        x = self._hc_post_block(x, residual, post, comb)
+        return x
+
+    def _run_ffn(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor:
+        shape = x.shape
+        x = x.reshape(-1, self.dim)
+        y = self._run_moe_ffn_dp_sync(
+            x, forward_batch, input_ids=None, input_ids_global=None
+        )
+        return y.view(shape)
+
+
+class DeepseekV4ForCausalLMDSpark(nn.Module):
+
+    def __init__(
+        self,
+        config: DeepSeekV4Config,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+
+        dspark_config = parse_dspark_draft_config(draft_hf_config=config)
+        if not dspark_config.require_markov():
+            raise ValueError(
+                "DSpark V4 draft requires markov_rank > 0, "
+                f"got markov_rank={dspark_config.markov_rank}."
+            )
+        self.gamma = int(
+            dspark_config.resolve_gamma(default=int(config.num_hidden_layers))
+        )
+        self.block_size = self.gamma
+        if dspark_config.target_layer_ids is not None:
+            self.num_stages = len(dspark_config.target_layer_ids)
+        else:
+            self.num_stages = int(getattr(config, "num_nextn_predict_layers", 1) or 1)
+
+        target_num_layers = (
+            int(dspark_config.num_target_layers)
+            if dspark_config.num_target_layers is not None
+            else int(getattr(config, "num_hidden_layers", 1))
+        )
+        if dspark_config.target_layer_ids is not None:
+            self.num_target_features = len(dspark_config.target_layer_ids)
+        else:
+            self.num_target_features = target_num_layers
+
+        self.start_layer = 0
+        self.end_layer = self.num_stages
+        use_multi_stream = (
+            envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
+            and envs.SGLANG_DSPARK_ENABLE_MULTI_STREAM.get()
+            and torch.cuda.is_available()
+        )
+        self.alt_streams: Optional[List[torch.cuda.Stream]] = (
+            [torch.cuda.Stream()] if use_multi_stream else None
+        )
+        self.stages = nn.ModuleList(
+            [
+                DSparkV4Stage(
+                    config=config,
+                    layer_id=stage_id,
+                    stage_id=stage_id,
+                    num_stages=self.num_stages,
+                    num_target_layers=self.num_target_features,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"stages.{stage_id}", prefix),
+                    alt_streams=self.alt_streams,
+                )
+                for stage_id in range(self.num_stages)
+            ]
+        )
+        self.markov_head = DSparkV4MarkovHead(
+            vocab_size=int(config.vocab_size),
+            markov_rank=int(dspark_config.markov_rank),
+        )
+        self.confidence_head = build_dspark_v4_confidence_head(
+            config=config, markov_rank=int(dspark_config.markov_rank)
+        )
+        self.hc_mult = int(config.hc_mult)
+        self.norm_eps = float(config.rms_norm_eps)
+        self.hc_eps = float(config.hc_eps)
+
+        self.embed_tokens: Optional[nn.Module] = None
+        self.lm_head: Optional[nn.Module] = None
+        self._use_fp32_lm_head = envs.SGLANG_DSPARK_FP32_LM_HEAD.get()
+        self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get()
+
+    @property
+    def enable_confidence_head(self) -> bool:
+        return self.confidence_head is not None
+
+    def attach_shared_modules(
+        self, *, embed_tokens: nn.Module, lm_head: nn.Module
+    ) -> None:
+        self.embed_tokens = embed_tokens
+        self.lm_head = lm_head
+        self.markov_head.configure_tp_shard(lm_head=lm_head)
+
+    def project_target_hidden(self, main_hidden: torch.Tensor) -> torch.Tensor:
+        stage0 = self.stages[0]
+        projected, _ = stage0.main_proj(main_hidden)
+        return stage0.main_norm(projected)
+
+    def write_target_hidden_kv(
+        self,
+        *,
+        main_hidden: torch.Tensor,
+        swa_loc: torch.Tensor,
+        positions: torch.Tensor,
+        pool: DeepSeekV4TokenToKVPool,
+    ) -> None:
+        main_x = self.project_target_hidden(main_hidden)
+        swa_loc = swa_loc.to(torch.int32)
+        kvs = CommitKvProj.execute(
+            main_x=main_x,
+            wkv_linears=[stage.self_attn.wkv for stage in self.stages],
+        )
+        for stage, kv in zip(self.stages, kvs):
+            attn = stage.self_attn
+            pool.set_swa_key_buffer_radix_fused_norm_rope(
+                layer_id=attn.layer_id,
+                swa_loc=swa_loc,
+                kv=kv,
+                kv_weight=attn.kv_norm.weight.data,
+                eps=attn.eps,
+                freqs_cis=attn.freqs_cis,
+                positions=positions,
+            )
+
+    def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if self.embed_tokens is None:
+            raise ValueError(
+                "DeepseekV4ForCausalLMDSpark requires the target embed_tokens "
+                "(call attach_shared_modules first)."
+            )
+        x = self.embed_tokens(input_ids)
+        x = x.unsqueeze(1).repeat(1, self.hc_mult, 1)
+        return x
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+        get_embedding: bool = False,
+        pp_proxy_tensors=None,
+    ) -> LogitsProcessorOutput:
+        del get_embedding, pp_proxy_tensors
+        if input_embeds is None:
+            input_embeds = self.forward_embed(input_ids)
+        x = input_embeds
+        for stage in self.stages:
+            x = stage(positions, x, forward_batch)
+
+        return LogitsProcessorOutput(next_token_logits=None, hidden_states=x)
+
+    def collapse_hc_head(self, x: torch.Tensor) -> torch.Tensor:
+        last = self.stages[-1]
+        return hc_head_torch(
+            x,
+            last.hc_head_fn,
+            last.hc_head_scale,
+            last.hc_head_base,
+            norm_eps=self.norm_eps,
+            hc_eps=self.hc_eps,
+        )
+
+    def compute_base_logits(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+
+        x_post_hc = self.collapse_hc_head(x)
+        return self._logits_from_x_post_hc(x_post_hc), x_post_hc
+
+    def _logits_from_x_post_hc(self, x_post_hc: torch.Tensor) -> torch.Tensor:
+        if self.lm_head is None:
+            raise ValueError(
+                "DeepseekV4ForCausalLMDSpark requires the target lm_head "
+                "(call attach_shared_modules first)."
+            )
+        last = self.stages[-1]
+        x = last.norm(x_post_hc)
+        weight = self.lm_head.weight
+        if self._use_fp32_lm_head:
+            local_logits = F.linear(x.float(), weight.float())
+        else:
+            local_logits = torch.matmul(x.to(weight.dtype), weight.T)
+        if self._opt_markov_w2_tp_shard:
+            return local_logits
+        return gather_and_crop_vocab(local_logits, self.lm_head)
+
+    def compute_confidence(
+        self,
+        *,
+        anchor_tokens: torch.Tensor,
+        sampled_tokens: torch.Tensor,
+        x_post_hc: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        confidence_head = self.confidence_head
+        if confidence_head is None:
+            return None
+        bs = int(anchor_tokens.shape[0])
+        x_post_hc = x_post_hc.view(bs, self.gamma, -1)
+        if confidence_head.with_markov:
+            prev_seq = torch.cat(
+                [anchor_tokens.view(-1, 1), sampled_tokens[:, : self.gamma - 1]], dim=1
+            )
+            markov_embed_stack = self.markov_head.get_prev_embeddings(prev_seq)
+        else:
+            markov_embed_stack = None
+        confidence_raw = confidence_head(x_post_hc, markov_embed_stack)
+        confidence = confidence_head.apply_sts(confidence_raw)
+        maybe_detect_in_closed_range(
+            confidence, 0.0, 1.0, "DSpark confidence must lie in [0, 1]."
+        )
+        return confidence
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
+        params_dict = dict(self.named_parameters())
+        loaded_params = set()
+
+        weights = list(weights)
+        if any(name.endswith(".wo_a.scale") for name, _ in weights):
+            weights = list(_dequant_fp8_wo_a(weights))
+
+        stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+
+        expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.n_routed_experts,
+        )
+
+        for name, loaded_weight in weights:
+            mapped = self._remap_dspark_weight_name(name)
+            if mapped is None:
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in mapped:
+                    continue
+                candidate = mapped.replace(weight_name, param_name)
+                if candidate not in params_dict:
+                    continue
+                param = params_dict[candidate]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(candidate)
+                break
+            else:
+                for (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    shard_id,
+                ) in expert_params_mapping:
+                    if weight_name not in mapped:
+                        continue
+                    candidate = mapped.replace(weight_name, param_name)
+                    if candidate not in params_dict:
+                        continue
+                    param = params_dict[candidate]
+                    weight_loader = param.weight_loader
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        candidate,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                    )
+                    loaded_params.add(candidate)
+                    break
+                else:
+                    if mapped not in params_dict:
+                        logger.warning(
+                            "DSpark V4 draft: unexpected weight %r -> %r", name, mapped
+                        )
+                        continue
+                    param = params_dict[mapped]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(mapped)
+
+        self._assert_confidence_head_loaded(
+            params_dict=params_dict, loaded_params=loaded_params
+        )
+
+    def _assert_confidence_head_loaded(
+        self, *, params_dict: dict, loaded_params: set
+    ) -> None:
+        if self.confidence_head is None:
+            return
+        confidence_param_names = {
+            name for name in params_dict if name.startswith("confidence_head.")
+        }
+        missing = confidence_param_names - loaded_params
+        if missing:
+            raise ValueError(
+                f"DSpark V4 confidence head is enabled but the checkpoint is missing "
+                f"{sorted(missing)}. Provide a checkpoint with trained confidence weights, "
+                f"or disable the confidence head (enable_confidence_head=False)."
+            )
+
+    def _remap_dspark_weight_name(self, name: str) -> Optional[str]:
+        if name.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")):
+            return None
+        if "rotary_emb.inv_freq" in name:
+            return None
+
+        if not name.startswith("mtp."):
+            return None
+        parts = name.split(".", 2)
+        if len(parts) < 3:
+            return None
+        stage_id, rest = parts[1], parts[2]
+
+        if rest.startswith("markov_head."):
+            return f"markov_head.{rest[len('markov_head.'):]}"
+
+        if rest.startswith("confidence_head."):
+            if self.confidence_head is None:
+                return None
+            return f"confidence_head.{rest[len('confidence_head.'):]}"
+
+        mapped_rest = rest
+        mapped_rest = mapped_rest.replace("attn.", "self_attn.", 1)
+        mapped_rest = mapped_rest.replace("ffn.", "mlp.", 1)
+        mapped_rest = mapped_rest.replace("attn_norm.", "input_layernorm.", 1)
+        mapped_rest = mapped_rest.replace("ffn_norm.", "post_attention_layernorm.", 1)
+        mapped_rest = mapped_rest.replace(".w1.", ".gate_proj.")
+        mapped_rest = mapped_rest.replace(".w2.", ".down_proj.")
+        mapped_rest = mapped_rest.replace(".w3.", ".up_proj.")
+        mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid")
+        mapped_rest = mapped_rest.replace(".gate.bias", ".gate.e_score_correction_bias")
+        mapped_rest = mapped_rest.replace(".scale", ".weight_scale_inv")
+        return f"stages.{stage_id}.{mapped_rest}"
+
+
+EntryClass = [DeepseekV4ForCausalLMDSpark]

python/sglang/srt/models/dspark.py

@@ -0,0 +1,509 @@
+from __future__ import annotations
+
+import logging
+from typing import Callable, Iterable, Optional, Tuple
+
+import torch
+from torch import nn
+
+from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.dflash import DFlashDraftModel
+from sglang.srt.speculative.dspark_components.dspark_config import (
+    parse_dspark_draft_config,
+)
+from sglang.srt.speculative.ragged_verify import (
+    RaggedVerifyMode,
+    read_ragged_verify_mode,
+)
+
+logger = logging.getLogger(__name__)
+
+StepSampler = Callable[[torch.Tensor, int], torch.Tensor]
+
+
+def gather_and_crop_vocab(
+    local_logits: torch.Tensor, lm_head: nn.Module
+) -> torch.Tensor:
+    full_logits = tensor_model_parallel_all_gather(local_logits, dim=-1)
+    return full_logits[..., : int(lm_head.org_vocab_size)]
+
+
+def run_markov_block(
+    head: nn.Module,
+    base_logits: torch.Tensor,
+    *,
+    first_prev_tokens: torch.Tensor,
+    hidden_states: Optional[torch.Tensor],
+    sampler: StepSampler,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    batch_size, proposal_len = base_logits.shape[:2]
+    if proposal_len == 0:
+        empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device)
+        return empty, base_logits
+
+    sampled_tokens = []
+    corrected_logits = []
+    prev_tokens = first_prev_tokens.long()
+    for step_idx in range(proposal_len):
+        step_hidden = None if hidden_states is None else hidden_states[:, step_idx, ...]
+        step_logits = head.apply_step_logits(
+            base_logits[:, step_idx, :],
+            token_ids=prev_tokens,
+            hidden_states=step_hidden,
+        )
+        next_tokens = sampler(step_logits, step_idx)
+        sampled_tokens.append(next_tokens)
+        corrected_logits.append(step_logits.unsqueeze(1))
+        prev_tokens = next_tokens
+    return (
+        torch.stack(sampled_tokens, dim=1),
+        torch.cat(corrected_logits, dim=1),
+    )
+
+
+class VanillaMarkov(nn.Module):
+
+    markov_head_type = "vanilla"
+
+    def __init__(self, *, vocab_size: int, markov_rank: int) -> None:
+        super().__init__()
+        self.vocab_size = int(vocab_size)
+        self.markov_rank = int(markov_rank)
+        if self.markov_rank <= 0:
+            raise ValueError(
+                f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}."
+            )
+        self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank)
+        self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False)
+
+    def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.markov_w1(token_ids.long())
+
+    def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor:
+        return self.markov_w2(latent_states)
+
+    def compute_step_bias(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        del hidden_states
+        return self.project_bias(self.get_prev_embeddings(token_ids))
+
+    def apply_step_logits(
+        self,
+        logits: torch.Tensor,
+        *,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        return logits + self.compute_step_bias(token_ids, hidden_states)
+
+    def apply_block_logits(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if base_logits.size(-2) == 0:
+            return base_logits
+        return base_logits + self.compute_step_bias(token_ids, hidden_states)
+
+    def sample_block(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        first_prev_tokens: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+        sampler: StepSampler,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return run_markov_block(
+            self,
+            base_logits,
+            first_prev_tokens=first_prev_tokens,
+            hidden_states=hidden_states,
+            sampler=sampler,
+        )
+
+
+class GatedMarkovHead(VanillaMarkov):
+
+    markov_head_type = "gated"
+
+    def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None:
+        super().__init__(vocab_size=vocab_size, markov_rank=markov_rank)
+        self.gate_proj = nn.Linear(int(hidden_size) + markov_rank, markov_rank)
+
+    def compute_gate(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if hidden_states is None:
+            raise ValueError("GatedMarkovHead requires hidden_states.")
+        prev_embeddings = self.get_prev_embeddings(token_ids)
+        gate_inputs = torch.cat([hidden_states, prev_embeddings], dim=-1)
+        return torch.sigmoid(self.gate_proj(gate_inputs))
+
+    def compute_step_bias(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        prev_embeddings = self.get_prev_embeddings(token_ids)
+        gate = self.compute_gate(token_ids, hidden_states).to(
+            dtype=prev_embeddings.dtype
+        )
+        return self.project_bias(gate * prev_embeddings)
+
+
+class RNNHead(VanillaMarkov):
+
+    markov_head_type = "rnn"
+
+    def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None:
+        super().__init__(vocab_size=vocab_size, markov_rank=markov_rank)
+        self.hidden_size = int(hidden_size)
+        self.state_size = markov_rank
+        self.joint_proj = nn.Linear(2 * markov_rank + self.hidden_size, 3 * markov_rank)
+
+    def _rnn_step(
+        self,
+        state: torch.Tensor,
+        prev_embeddings: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        z = torch.cat([state, prev_embeddings, hidden_states], dim=-1)
+        gate_raw, candidate_raw, output_raw = self.joint_proj(z).chunk(3, dim=-1)
+        gate = torch.sigmoid(gate_raw)
+        candidate = torch.tanh(candidate_raw)
+        new_state = gate * state + (1.0 - gate) * candidate
+        bias = self.project_bias(torch.tanh(output_raw))
+        return new_state, bias
+
+    def compute_step_bias(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if hidden_states is None:
+            raise ValueError("RNNHead requires hidden_states.")
+        prev_embeddings = self.get_prev_embeddings(token_ids)
+        state = torch.zeros_like(prev_embeddings)
+        _, bias = self._rnn_step(state, prev_embeddings, hidden_states)
+        return bias
+
+    def apply_block_logits(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if hidden_states is None:
+            raise ValueError("RNNHead requires hidden_states.")
+        block_size = base_logits.size(-2)
+        if block_size == 0:
+            return base_logits
+        leading_shape = base_logits.shape[:-2]
+        state = torch.zeros(
+            *leading_shape,
+            self.markov_rank,
+            device=base_logits.device,
+            dtype=hidden_states.dtype,
+        )
+        output_logits = []
+        for k in range(block_size):
+            prev_emb = self.get_prev_embeddings(token_ids[..., k])
+            state, bias = self._rnn_step(state, prev_emb, hidden_states[..., k, :])
+            output_logits.append(base_logits[..., k, :] + bias)
+        return torch.stack(output_logits, dim=-2)
+
+    def sample_block(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        first_prev_tokens: torch.Tensor,
+        hidden_states: Optional[torch.Tensor],
+        sampler: StepSampler,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if hidden_states is None:
+            raise ValueError("RNNHead requires hidden_states.")
+        batch_size, proposal_len = base_logits.shape[:2]
+        if proposal_len == 0:
+            empty = torch.empty(
+                batch_size, 0, dtype=torch.long, device=base_logits.device
+            )
+            return empty, base_logits
+
+        state = torch.zeros(
+            batch_size,
+            self.markov_rank,
+            device=base_logits.device,
+            dtype=hidden_states.dtype,
+        )
+        sampled_tokens = []
+        corrected_logits = []
+        prev_tokens = first_prev_tokens.long()
+        for step_idx in range(proposal_len):
+            prev_emb = self.get_prev_embeddings(prev_tokens)
+            state, bias = self._rnn_step(state, prev_emb, hidden_states[:, step_idx, :])
+            step_logits = base_logits[:, step_idx, :] + bias
+            next_tokens = sampler(step_logits, step_idx)
+            sampled_tokens.append(next_tokens)
+            corrected_logits.append(step_logits.unsqueeze(1))
+            prev_tokens = next_tokens
+        return (
+            torch.stack(sampled_tokens, dim=1),
+            torch.cat(corrected_logits, dim=1),
+        )
+
+
+def build_markov_head(config) -> Optional[nn.Module]:
+    markov_rank = int(getattr(config, "markov_rank", 0))
+    if markov_rank <= 0:
+        raise ValueError(
+            "DSpark requires markov_rank > 0 (the Markov head is the core of the "
+            f"semi-AR draft); got markov_rank={markov_rank}."
+        )
+    markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower()
+    vocab_size = int(config.vocab_size)
+    hidden_size = int(config.hidden_size)
+    if markov_head_type == "vanilla":
+        return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank)
+    if markov_head_type == "gated":
+        return GatedMarkovHead(
+            vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size
+        )
+    if markov_head_type == "rnn":
+        return RNNHead(
+            vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size
+        )
+    raise ValueError(f"Unsupported DSpark markov_head_type={markov_head_type!r}.")
+
+
+class DSparkConfidenceHead(nn.Module):
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        markov_rank: int,
+        with_markov: bool = True,
+        bias: bool = True,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__()
+        self.with_markov = bool(with_markov)
+        input_dim = int(hidden_size) + (int(markov_rank) if self.with_markov else 0)
+        self.proj = nn.Linear(input_dim, 1, bias=bias, dtype=dtype)
+        self.register_buffer(
+            "sts_temperatures", torch.ones((), dtype=torch.float32), persistent=False
+        )
+        self._last_confidence_raw: Optional[torch.Tensor] = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        markov_embed_stack: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.with_markov:
+            if markov_embed_stack is None:
+                raise ValueError(
+                    "DSparkConfidenceHead(with_markov=True) requires markov_embed_stack."
+                )
+            features = torch.cat(
+                [hidden_states, markov_embed_stack.to(dtype=hidden_states.dtype)],
+                dim=-1,
+            )
+        else:
+            features = hidden_states
+        features = features.to(dtype=self.proj.weight.dtype)
+        return self.proj(features).squeeze(-1)
+
+    def apply_sts(self, confidence_raw: torch.Tensor) -> torch.Tensor:
+        self._last_confidence_raw = confidence_raw
+        return torch.sigmoid(confidence_raw.float() / self.sts_temperatures)
+
+
+def build_confidence_head(config) -> Optional[nn.Module]:
+    if read_ragged_verify_mode() is RaggedVerifyMode.STATIC:
+        return None
+    if not hasattr(config, "enable_confidence_head"):
+        logger.warning(
+            "DSpark draft config has no enable_confidence_head field; treating the "
+            "confidence head as enabled."
+        )
+    hidden_size = int(config.hidden_size)
+    markov_rank = int(getattr(config, "markov_rank", 0))
+    with_markov = bool(getattr(config, "confidence_head_with_markov", markov_rank > 0))
+    if with_markov and markov_rank <= 0:
+        raise ValueError(
+            "DSpark confidence_head_with_markov requires markov_rank > 0, "
+            f"got markov_rank={markov_rank}."
+        )
+    return DSparkConfidenceHead(
+        hidden_size=hidden_size,
+        markov_rank=markov_rank,
+        with_markov=with_markov,
+    )
+
+
+_DSPARK_SKIPPED_WEIGHT_PREFIXES = (
+    "embed_tokens.",
+    "lm_head.",
+    "rotary_emb.",
+)
+
+
+class DSparkDraftMixin:
+
+    def __init__(self, config, quant_config=None, prefix: str = "") -> None:
+        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
+        dspark_config = parse_dspark_draft_config(draft_hf_config=config)
+        if not dspark_config.require_markov():
+            raise ValueError(
+                "DSpark draft requires markov_rank > 0, "
+                f"got markov_rank={dspark_config.markov_rank}."
+            )
+        self.gamma = int(dspark_config.resolve_gamma(default=self.block_size))
+        self.markov_head = build_markov_head(config)
+        self.confidence_head = build_confidence_head(config)
+        self.lm_head: Optional[nn.Module] = None
+
+    def attach_shared_modules(
+        self, *, embed_tokens: nn.Module, lm_head: nn.Module
+    ) -> None:
+        del embed_tokens
+        self.lm_head = lm_head
+
+    def compute_base_logits(
+        self, hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if self.lm_head is None:
+            raise ValueError(
+                "DSpark dense draft requires the target lm_head "
+                "(call attach_shared_modules first)."
+            )
+        weight = self.lm_head.weight
+        if hidden.dtype != weight.dtype:
+            hidden = hidden.to(weight.dtype)
+        local_logits = torch.matmul(hidden, weight.T)
+        base_logits = gather_and_crop_vocab(local_logits, self.lm_head)
+        return base_logits, None
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        markov_weights = []
+        confidence_weights = []
+        backbone_weights = []
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in weights:
+            if any(name.startswith(p) for p in _DSPARK_SKIPPED_WEIGHT_PREFIXES):
+                continue
+            if name.startswith("confidence_head."):
+                if self.confidence_head is None:
+                    continue
+                confidence_weights.append((name, loaded_weight))
+            elif name.startswith("markov_head."):
+                markov_weights.append((name, loaded_weight))
+            else:
+                backbone_weights.append((name, loaded_weight))
+
+        super().load_weights(backbone_weights)
+
+        for name, loaded_weight in markov_weights:
+            if name not in params_dict:
+                raise ValueError(
+                    f"DSpark unexpected markov weight {name!r} not found in model "
+                    f"parameters (known markov params require a {type(self.markov_head).__name__} head)."
+                )
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+
+        self._load_confidence_weights(
+            confidence_weights=confidence_weights, params_dict=params_dict
+        )
+
+    def _load_confidence_weights(
+        self,
+        *,
+        confidence_weights: list,
+        params_dict: dict,
+    ) -> None:
+        if self.confidence_head is None:
+            return
+        loaded_names = set()
+        for name, loaded_weight in confidence_weights:
+            if name not in params_dict:
+                raise ValueError(
+                    f"DSpark unexpected confidence weight {name!r} not found in "
+                    "model parameters."
+                )
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_names.add(name)
+
+        confidence_param_names = {
+            name for name in params_dict if name.startswith("confidence_head.")
+        }
+        missing = confidence_param_names - loaded_names
+        if missing:
+            raise ValueError(
+                f"DSpark confidence head is enabled but the checkpoint is missing "
+                f"{sorted(missing)}. Provide a checkpoint with trained confidence weights, "
+                f"or disable the confidence head (enable_confidence_head=False)."
+            )
+
+    def write_target_hidden_kv(
+        self,
+        *,
+        target_hidden: torch.Tensor,
+        pool,
+        positions: torch.Tensor,
+        cache_loc: torch.Tensor,
+        cache_loc_2d: Optional[torch.Tensor] = None,
+        commit_lens: Optional[torch.Tensor] = None,
+    ) -> None:
+        ctx_hidden = self.project_target_hidden(target_hidden)
+        for layer in self.layers:
+            attn = layer.self_attn
+            k, v = attn.kv_proj_only(ctx_hidden)
+            k = attn.apply_k_norm(k)
+            k = attn.apply_k_rope(positions, k)
+            k = k.view(-1, attn.num_kv_heads, attn.head_dim)
+            v = v.view(-1, attn.num_kv_heads, attn.head_dim)
+            if cache_loc_2d is not None and commit_lens is not None:
+                pool.set_kv_buffer_prefix_valid(
+                    attn.attn,
+                    cache_loc_2d,
+                    commit_lens,
+                    k,
+                    v,
+                    attn.attn.k_scale,
+                    attn.attn.v_scale,
+                )
+            else:
+                pool.set_kv_buffer(
+                    attn.attn,
+                    cache_loc,
+                    k,
+                    v,
+                    attn.attn.k_scale,
+                    attn.attn.v_scale,
+                )
+
+
+class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel):
+
+    pass
+
+
+class Qwen3DSparkModel(DSparkDraftModel):
+    pass
+
+
+EntryClass = [Qwen3DSparkModel]

python/sglang/srt/observability/metrics_collector.py

@@ -109,6 +109,8 @@ class SchedulerStats:
     # Speculative decoding
     spec_accept_length: float = 0.0
     spec_accept_rate: float = 0.0
+    spec_cap_length: float = 0.0
+    spec_block_accept_length: float = 0.0
     # Adaptive speculative decoding (currently active tier).
     spec_num_steps: int = 0
     spec_num_draft_tokens: int = 0
@@ -423,6 +425,18 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
             labelnames=labels.keys(),
             multiprocess_mode="mostrecent",
         )
+        self.spec_cap_length = Gauge(
+            name="sglang:spec_cap_length",
+            documentation="Mean DSpark confidence-scheduled verify window per verify step, incl the bonus slot (0 when no cap is scheduled).",
+            labelnames=labels.keys(),
+            multiprocess_mode="mostrecent",
+        )
+        self.spec_block_accept_length = Gauge(
+            name="sglang:spec_block_accept_length",
+            documentation="Mean uncapped full-block accept length per verify step (accept + cap-trimmed drafts; exact only in DSpark cap-accept mode).",
+            labelnames=labels.keys(),
+            multiprocess_mode="mostrecent",
+        )
         self.spec_num_steps = Gauge(
             name="sglang:spec_num_steps",
             documentation="Currently active speculative_num_steps.",
@@ -1287,6 +1301,8 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
         # Speculative decoding
         self._log_gauge(self.spec_accept_length, stats.spec_accept_length)
         self._log_gauge(self.spec_accept_rate, stats.spec_accept_rate)
+        self._log_gauge(self.spec_cap_length, stats.spec_cap_length)
+        self._log_gauge(self.spec_block_accept_length, stats.spec_block_accept_length)
         self._log_gauge(self.spec_num_steps, stats.spec_num_steps)
         self._log_gauge(self.spec_num_draft_tokens, stats.spec_num_draft_tokens)
 

python/sglang/srt/sampling/sampling_batch_info.py

@@ -31,6 +31,8 @@ class SamplingBatchInfo:
     # Whether all requests use greedy sampling
     is_all_greedy: bool
 
+    is_any_greedy: bool
+
     # Whether any requests use top_p sampling
     need_top_p_sampling: bool
 
@@ -188,6 +190,7 @@ class SamplingBatchInfo:
             min_ps=min_ps,
             sampling_seed=sampling_seed,
             is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
+            is_any_greedy=any(r.sampling_params.top_k <= 1 for r in reqs),
             need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
             need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
             need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
@@ -409,6 +412,7 @@ class SamplingBatchInfo:
                 setattr(self, item, torch.cat([self_val, other_val]))
 
         self.is_all_greedy &= other.is_all_greedy
+        self.is_any_greedy |= other.is_any_greedy
         self.need_top_p_sampling |= other.need_top_p_sampling
         self.need_top_k_sampling |= other.need_top_k_sampling
         self.need_min_p_sampling |= other.need_min_p_sampling

python/sglang/srt/server_args.py

@@ -1655,7 +1655,7 @@ class ServerArgs:
     # -------------------------------------------------------------------------
     speculative_algorithm: A[
         Optional[str],
-        "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH. Or any name registered via `SpeculativeAlgorithm.register`.",
+        "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK. Or any name registered via `SpeculativeAlgorithm.register`.",
     ] = None
     speculative_draft_model_path: A[
         Optional[str],
@@ -1691,6 +1691,37 @@ class ServerArgs:
         Optional[int],
         "DFLASH only. Block size (verify window length). Alias of --speculative-num-draft-tokens for DFLASH.",
     ] = None
+    speculative_dspark_block_size: A[
+        Optional[int],
+        "DSPARK only. Draft block size gamma (number of proposed draft tokens). The verify window is gamma + 1, so this sets --speculative-num-draft-tokens = gamma + 1. Omit to auto-infer gamma from the draft checkpoint block_size.",
+    ] = None
+    speculative_dspark_sps_table_path: A[
+        Optional[str],
+        "DSPARK only. Path to a pre-profiled SPS cost table (JSON) built offline with "
+        "sglang.benchmark.dspark_sps_profiler, consumed by the ragged-verify "
+        "scheduler (cap-accept / compact). Omit for an uninitialized flat "
+        "constant-SPS table: the budget degenerates to verify-all (zero throughput "
+        "gain by itself).",
+    ] = None
+    speculative_dspark_confidence_sts_path: A[
+        Optional[str],
+        "DSPARK only. Optional path to a per-position STS (sequential temperature "
+        "scaling) calibration JSON, fit offline with sglang.benchmark.dspark_sts_fit. "
+        "Calibrates the confidence-head survival probabilities the ragged-verify "
+        "scheduler consumes. Omit to use identity (no calibration); losslessness is "
+        "unaffected either way.",
+    ] = None
+    speculative_dspark_align_verify_tokens_to_graph_tier: A[
+        bool,
+        "DSPARK compact ragged-verify only. Fill the per-request verify lengths so "
+        "the total verify-token count reaches the cuda-graph tier the forward is "
+        "already padded to: round the dp-max scheduled total up to the captured "
+        "token bucket and let the top-k allocator admit that many real draft tokens "
+        "(confidence-ordered). This recovers the padding the forward pays for anyway "
+        "-- both the cuda-graph bucket round-up and the dp cross-rank max -- turning "
+        "it into extra real verification at the same step time. Off by default; when "
+        "off the schedule is byte-for-byte unchanged.",
+    ] = False
     speculative_accept_threshold_single: A[
         float,
         "Accept a draft token if its probability in the target model is greater than this threshold.",

python/sglang/srt/speculative/base_spec_worker.py

@@ -338,6 +338,14 @@ class BaseSpecWorker(ABC):
         """
         pass
 
+    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        """Hook called by the batch-result processor when a request finishes.
+
+        Default no-op. DSpark overrides this to settle / censor its
+        block-accept estimator state for the finished request.
+        """
+        pass
+
     def activate_step_by_batch(self, batch_size: int) -> None:
         """Activate the optimal adaptive step for the current batch size.
 

python/sglang/srt/speculative/dflash_info.py

@@ -16,6 +16,7 @@ from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
 
 if TYPE_CHECKING:
     from sglang.srt.managers.tp_worker import TpModelWorker
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
 
 
 @dataclass
@@ -41,6 +42,8 @@ class DFlashVerifyInput(SpecInput):
     # Shape info for padding (e.g., DP attention / CUDA graph).
     num_tokens_per_req: int = -1
 
+    ragged_verify_layout: Optional[RaggedVerifyLayout] = None
+
     def __post_init__(self):
         super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY)
         if self.num_tokens_per_req == -1:
@@ -99,6 +102,9 @@ class DFlashVerifyInput(SpecInput):
         device = req_pool_indices.device
         bs = len(req_pool_indices)
 
+        layout = self.ragged_verify_layout
+
+        if layout is None:
             qo_indptr = torch.arange(
                 0,
                 (bs + 1) * self.draft_token_num,
@@ -106,13 +112,19 @@ class DFlashVerifyInput(SpecInput):
                 dtype=torch.int32,
                 device=device,
             )
+            verify_lens = self.draft_token_num
+            kv_indices_extra = self.draft_token_num * bs
+        else:
+            qo_indptr = layout.qo_indptr_device
+            verify_lens = layout.verify_lens
+            kv_indices_extra = layout.total_verify_tokens
 
         cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=device)
-        paged_kernel_lens = paged_kernel_lens + self.draft_token_num
+        paged_kernel_lens = paged_kernel_lens + verify_lens
         cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
 
         kv_indices = torch.empty(
-            paged_kernel_lens_sum + self.draft_token_num * bs,
+            paged_kernel_lens_sum + kv_indices_extra,
             dtype=torch.int32,
             device=device,
         )

python/sglang/srt/speculative/dflash_info_v2.py

@@ -59,6 +59,8 @@ class DFlashDraftInputV2(SpecInput):
     # Filled by scheduler after dispatch.
     future_indices: Optional[torch.Tensor] = None
 
+    verify_token_budget: Optional[int] = None
+
     def __post_init__(self):
         super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT)
 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -393,6 +393,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             top_ks=torch.full((num_seqs,), -1, dtype=torch.int32),
             min_ps=torch.zeros((num_seqs,), dtype=torch.float),
             is_all_greedy=False,
+            is_any_greedy=False,
             need_top_p_sampling=False,
             need_top_k_sampling=False,
             need_min_p_sampling=False,

python/sglang/srt/speculative/ragged_verify.py

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

python/sglang/srt/speculative/ragged_verify_kernels.py

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

python/sglang/srt/speculative/spec_info.py

@@ -23,6 +23,7 @@ if TYPE_CHECKING:
     from sglang.srt.server_args import ServerArgs
     from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
     from sglang.srt.speculative.ngram_worker import NGRAMWorker
+    from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
 
 
 class SpeculativeAlgorithm(Enum):
@@ -33,6 +34,7 @@ class SpeculativeAlgorithm(Enum):
     """
 
     DFLASH = auto()
+    DSPARK = auto()
     EAGLE = auto()
     EAGLE3 = auto()
     FROZEN_KV_MTP = auto()
@@ -109,6 +111,12 @@ class SpeculativeAlgorithm(Enum):
     def is_dflash(self) -> bool:
         return self == SpeculativeAlgorithm.DFLASH
 
+    def is_dspark(self) -> bool:
+        return self == SpeculativeAlgorithm.DSPARK
+
+    def is_dflash_family(self) -> bool:
+        return self.is_dflash() or self.is_dspark()
+
     def is_standalone(self) -> bool:
         return self == SpeculativeAlgorithm.STANDALONE
 
@@ -116,7 +124,13 @@ class SpeculativeAlgorithm(Enum):
         return self == SpeculativeAlgorithm.NGRAM
 
     def supports_target_verify_for_draft(self) -> bool:
-        return self.is_dflash()
+        return self.is_dflash_family()
+
+    def supports_ragged_verify(self) -> bool:
+        """Whether this algorithm's verify step may carry a RaggedVerifyLayout
+        (per-request verify lengths); gates the token-bucket-keyed verify
+        graphs in the decode cuda graph runner."""
+        return self.is_dspark()
 
     def has_draft_kv(self) -> bool:
         """Whether the draft phase writes KV chains. NGRAM does not (its tree
@@ -134,10 +148,17 @@ class SpeculativeAlgorithm(Enum):
         device: torch.device,
         req_to_token_pool,
         needs_cpu_seq_lens: bool = True,
+        needs_confidence_relay: bool = False,
     ) -> FutureMap:
         from sglang.srt.managers.overlap_utils import FutureMap
 
-        return FutureMap(device, self, req_to_token_pool, needs_cpu_seq_lens)
+        return FutureMap(
+            device,
+            self,
+            req_to_token_pool,
+            needs_cpu_seq_lens,
+            needs_confidence_relay,
+        )
 
     def build_disagg_draft_input(
         self,
@@ -166,13 +187,22 @@ class SpeculativeAlgorithm(Enum):
         """
         from sglang.srt.arg_groups.speculative_hook import (
             _handle_dflash,
+            _handle_dspark,
             _handle_eagle_family,
             _handle_frozen_kv_mtp,
             _handle_ngram,
         )
 
+        # Validate for every algorithm at startup: the metrics paths read the
+        # ragged-verify mode env and must not be where a typo'd value raises.
+        from sglang.srt.speculative.ragged_verify import read_ragged_verify_mode
+
+        read_ragged_verify_mode()
+
         if self.is_dflash():
             _handle_dflash(server_args)
+        elif self.is_dspark():
+            _handle_dspark(server_args)
         elif self.is_frozen_kv_mtp():
             _handle_frozen_kv_mtp(server_args)
         elif self.is_eagle() or self.is_standalone():
@@ -188,6 +218,8 @@ class SpeculativeAlgorithm(Enum):
         # graph support. We can use it for target verify, or we can use it for
         # other cases which is not target verify but fixed length prefill.
         # Here, we expose this interface to allow the other use cases.
+        if self.is_dspark() and is_draft_worker:
+            return num_draft_tokens - 1
         return num_draft_tokens
 
     def create_worker(
@@ -204,6 +236,13 @@ class SpeculativeAlgorithm(Enum):
 
             return DFlashWorkerV2
 
+        if self.is_dspark():
+            from sglang.srt.speculative.dspark_components.dspark_worker_v2 import (
+                DSparkWorkerV2,
+            )
+
+            return DSparkWorkerV2
+
         if self.is_frozen_kv_mtp():
             # V2 worker drives both overlap and non-overlap (scheduler runs it
             # synchronously when overlap is disabled), same as EAGLE.
@@ -254,6 +293,14 @@ class SpecInputType(IntEnum):
 
 
 class SpecInput(ABC):
+    # Per-request verify lengths for the ragged-verify graphs (see
+    # sglang.srt.speculative.ragged_verify); verify inputs of algorithms with
+    # supports_ragged_verify() override it per step. Must stay a class-level
+    # default, not an __init__ assignment: dataclass subclasses declare it as
+    # a field and run __post_init__ -> super().__init__ *after* field
+    # assignment, so an init-time default would clobber the passed layout.
+    ragged_verify_layout: Optional[RaggedVerifyLayout] = None
+
     def __init__(self, spec_input_type: SpecInputType):
         self.spec_input_type = spec_input_type
 
@@ -324,7 +371,7 @@ def create_dummy_verify_input(
                 seq_lens_sum=None,
                 seq_lens_cpu=None,
             )
-    elif spec_algorithm.is_dflash():
+    elif spec_algorithm.is_dflash_family():
         from sglang.srt.speculative.dflash_info import DFlashVerifyInput
 
         # Dummy warmup only needs shape metadata; avoid forcing custom-mask mode.

python/sglang/srt/speculative/spec_registry.py

@@ -76,6 +76,12 @@ class CustomSpecAlgo:
     def is_dflash(self) -> bool:
         return False
 
+    def is_dspark(self) -> bool:
+        return False
+
+    def is_dflash_family(self) -> bool:
+        return False
+
     def is_standalone(self) -> bool:
         return False
 
@@ -85,6 +91,9 @@ class CustomSpecAlgo:
     def supports_target_verify_for_draft(self) -> bool:
         return False
 
+    def supports_ragged_verify(self) -> bool:
+        return False
+
     def has_draft_kv(self) -> bool:
         # Conservative default: the larger KV reserve.
         return True

python/sglang/srt/speculative/spec_utils.py

@@ -207,9 +207,9 @@ def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
         server_args = get_server_args()
 
     # STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
-    # multi_layer_eagle and DFLASH don't relay hidden_states through FutureMap.
+    # multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
     # TODO(lsyin): also skip when step == 1.
-    if server_args.speculative_algorithm in ("STANDALONE", "DFLASH"):
+    if server_args.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
         return False
     return not server_args.enable_multi_layer_eagle
 
@@ -711,7 +711,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
     """eagle/ngram share a stateless free function; dflash keeps stateful
     prep on its draft input -- the dispatcher routes.
     """
-    if batch.spec_algorithm.is_dflash():
+    if batch.spec_algorithm.is_dflash_family():
         batch.spec_info.prepare_for_decode(batch)
     else:
         from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode

python/sglang/srt/utils/async_probe.py

@@ -76,6 +76,12 @@ def sanitize_nan_logits(logits: torch.Tensor, msg: str = ""):
     torch.nan_to_num_(logits, nan=-1e30, posinf=1e30, neginf=-1e30)
 
 
+def maybe_assert_async(cond: torch.Tensor, msg: str = ""):
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    torch._assert_async(cond, msg)
+
+
 def maybe_detect_nan(tensor: Optional[torch.Tensor], msg: str = ""):
     """Async NaN check — no GPU-CPU sync, error surfaces at next sync point."""
     if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
@@ -96,6 +102,19 @@ def maybe_detect_inf(tensor: Optional[torch.Tensor], msg: str = ""):
     torch._assert_async(~torch.any(torch.isinf(tensor)), f"Inf detected! {msg}")
 
 
+def maybe_detect_in_closed_range(
+    tensor: Optional[torch.Tensor], low: float, high: float, msg: str = ""
+):
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    if tensor is None or tensor.numel() == 0:
+        return
+    torch._assert_async(
+        ((tensor >= low) & (tensor <= high)).all(),
+        f"value outside [{low}, {high}]: {msg}",
+    )
+
+
 def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg: str):
     """Async OOB check — no GPU-CPU sync, error surfaces at next sync point.
 

test/registered/core/test_basic_sanity_dspark.py

@@ -0,0 +1,94 @@
+import unittest
+
+from sglang.srt.utils import is_sm100_supported, kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.basic_api_contract_kit import BasicAPIContractMixin
+from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
+from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixin
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.fwd_occupancy_kit import FwdOccupancyMixin
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
+
+TARGET_MODEL = "Qwen/Qwen3-14B"
+DRAFT_MODEL = "deepseek-ai/dspark_qwen3_14b_block7"
+
+# trtllm_mha prefill requires SM100 (Blackwell); use the Hopper-native pair elsewhere.
+if is_sm100_supported():
+    ATTENTION_BACKEND = "trtllm_mha"
+    DRAFT_ATTENTION_BACKEND = "fa4"
+else:
+    ATTENTION_BACKEND = "fa3"
+    DRAFT_ATTENTION_BACKEND = "fa3"
+
+
+class TestBasicSanityDSpark(
+    BasicAPIContractMixin,
+    BasicDecodeCorrectnessMixin,
+    BasicSchedulerStressMixin,
+    FwdOccupancyMixin,
+    GSM8KMixin,
+    CustomTestCase,
+):
+    served_model_name = TARGET_MODEL
+    model = TARGET_MODEL
+
+    fwd_occupancy_threshold = 60
+    fwd_occupancy_max_new_tokens = 4096
+    fwd_occupancy_acc_length_threshold: float = 2.0
+
+    gsm8k_num_questions = 200
+    gsm8k_accuracy_thres = 0.80
+    gsm8k_accept_length_thres = 2.0
+
+    attention_backend = ATTENTION_BACKEND
+    draft_attention_backend = DRAFT_ATTENTION_BACKEND
+
+    process = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            TARGET_MODEL,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--attention-backend",
+                cls.attention_backend,
+                "--speculative-draft-attention-backend",
+                cls.draft_attention_backend,
+                "--speculative-algorithm",
+                "DSPARK",
+                "--speculative-draft-model-path",
+                DRAFT_MODEL,
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--mem-fraction-static",
+                "0.7",
+                "--page-size",
+                "1",
+                "--enable-metrics",
+                "--disable-piecewise-cuda-graph",
+            ],
+            env={
+                "SGLANG_ENABLE_METRICS_DEVICE_TIMER": "1",
+                "SGLANG_RAGGED_VERIFY_MODE": "compact",
+            },
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.process is not None:
+            kill_process_tree(cls.process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/ops/test_aiter_greedy_sample_amd.py

@@ -52,6 +52,7 @@ def _make_sampling_info(batch_size, vocab_size, device="cuda"):
         top_ks=torch.zeros(batch_size, device=device, dtype=torch.int32),
         min_ps=torch.zeros(batch_size, device=device),
         is_all_greedy=True,
+        is_any_greedy=True,
         need_top_p_sampling=False,
         need_top_k_sampling=False,
         need_min_p_sampling=False,

test/registered/scheduler/test_min_free_slots_delayer.py

@@ -14,17 +14,17 @@ class TestResolveMinFreeSlots(unittest.TestCase):
 
     def test_unset_non_dflash_disables(self):
         # Unset + not DFlash -> trigger stays disabled.
-        self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash=False))
+        self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash_family=False))
 
     def test_unset_dflash_auto_enables(self):
         # Unset + DFlash -> falls back to the legacy formula (full mapping).
-        self.assertEqual(resolve_min_free_slots(None, 512, is_dflash=True), 4)
-        self.assertEqual(resolve_min_free_slots(None, 8, is_dflash=True), 2)
+        self.assertEqual(resolve_min_free_slots(None, 512, is_dflash_family=True), 4)
+        self.assertEqual(resolve_min_free_slots(None, 8, is_dflash_family=True), 2)
 
     def test_unset_dflash_small_cluster_disables(self):
         # DFlash auto-default still respects the < 8 guard.
-        self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash=True))
-        self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash=True))
+        self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash_family=True))
+        self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash_family=True))
 
     def test_le_one_disables(self):
         # <= 1 can never batch, so it is a no-op.
@@ -47,7 +47,7 @@ class TestResolveMinFreeSlots(unittest.TestCase):
 
     def test_user_value_overrides_dflash_default(self):
         # An explicit user value wins over the DFlash auto-default.
-        self.assertEqual(resolve_min_free_slots(3, 512, is_dflash=True), 3)
+        self.assertEqual(resolve_min_free_slots(3, 512, is_dflash_family=True), 3)
 
 
 class TestMinFreeSlotsDelayer(unittest.TestCase):

test/registered/spec/dspark/test_dspark_block_accept_estimator.py

@@ -0,0 +1,579 @@
+import json
+import math
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import (
+    BlockAcceptEstimateRecorder,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+_GAMMA = 3
+_VOCAB = 8
+
+
+class _FakeLayout:
+    def __init__(self, verify_lens: torch.Tensor):
+        self.verify_lens = verify_lens
+
+
+def _reference_logprob(
+    logits_row: torch.Tensor, token: int, temperature: float
+) -> float:
+    scaled = logits_row.to(torch.float32) / temperature
+    return float(scaled[token] - torch.logsumexp(scaled, dim=-1))
+
+
+def _make_recorder(tmp_dir: str) -> tuple[BlockAcceptEstimateRecorder, Path]:
+    path = Path(tmp_dir) / "estimate.jsonl"
+    recorder = BlockAcceptEstimateRecorder(path=str(path), gamma=_GAMMA, device="cpu")
+    return recorder, path
+
+
+class _FakeDelayed:
+
+    def __init__(self):
+        self._pending = None
+
+    def step(self, *, compute_on_device, postprocess_on_host):
+        if self._pending is not None:
+            result, post = self._pending
+            self._pending = None
+            if result is not None:
+                post(result)
+        result = compute_on_device()
+        self._pending = (result, postprocess_on_host) if result is not None else None
+
+
+def _observe(
+    recorder: BlockAcceptEstimateRecorder,
+    *,
+    forward_ct: int,
+    rid: str,
+    drafts: list[int],
+    corrected_logits: torch.Tensor,
+    target_logits: torch.Tensor,
+    verify_len: int,
+    correct_len: int,
+    bonus: int,
+    seq_len: int,
+    cap_trim: int = 0,
+    temperature: float = 1.0,
+) -> None:
+    recorder.observe_verify_step(
+        forward_ct=forward_ct,
+        rids=[rid],
+        draft_tokens=torch.tensor([drafts], dtype=torch.int64),
+        corrected_logits=corrected_logits.unsqueeze(0),
+        draft_temperatures=torch.tensor([temperature], dtype=torch.float32),
+        greedy_mask=torch.tensor([False]),
+        target_logits=target_logits,
+        target_temperatures=torch.tensor([[temperature]], dtype=torch.float32),
+        truncated_sampling_mask=None,
+        logits_adjustments_are_noop=True,
+        correct_len=torch.tensor([correct_len], dtype=torch.int32),
+        cap_trim_lens=torch.tensor([cap_trim], dtype=torch.int32),
+        bonus=torch.tensor([bonus], dtype=torch.int64),
+        prefix_lens=torch.tensor([seq_len], dtype=torch.int64),
+        layout=_FakeLayout(torch.tensor([verify_len], dtype=torch.int32)),
+    )
+
+
+def _read_records(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines()]
+
+
+class TestBlockAcceptEstimateRecorder(CustomTestCase):
+    def test_exact_block_when_rejected_inside_window(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected = torch.randn(_GAMMA, _VOCAB)
+            target = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=corrected,
+                target_logits=target,
+                verify_len=3,
+                correct_len=1,
+                bonus=5,
+                seq_len=10,
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            self.assertEqual(len(records), 1)
+            self.assertEqual(records[0]["w"], 2)
+            self.assertEqual(records[0]["cl"], 1)
+            self.assertNotIn("q_lp", records[0])
+            self.assertNotIn("pg", records[0])
+
+    def test_censored_block_gathers_q_and_same_step_bonus_row_p(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected = torch.randn(_GAMMA, _VOCAB)
+            target = torch.randn((_GAMMA + 1), _VOCAB)
+            drafts = [1, 2, 3]
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=drafts,
+                corrected_logits=corrected,
+                target_logits=target,
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            self.assertEqual(len(records), 1)
+            record = records[0]
+            self.assertEqual(record["w"], 1)
+            self.assertEqual(record["cl"], 1)
+            self.assertEqual(record["trimmed_tokens"], [2, 3])
+            self.assertEqual(len(record["q_lp"]), 2)
+            self.assertAlmostEqual(
+                record["q_lp"][0],
+                _reference_logprob(corrected[1], 2, 1.0),
+                places=4,
+            )
+            self.assertAlmostEqual(
+                record["q_lp"][1],
+                _reference_logprob(corrected[2], 3, 1.0),
+                places=4,
+            )
+            self.assertEqual(len(record["pg"]), 1)
+            src_fct, offset, p_lp, draft_token, realized_token = record["pg"][0]
+            self.assertEqual((src_fct, offset), (1, 2))
+            self.assertEqual(draft_token, 2)
+            self.assertEqual(realized_token, 2)
+            self.assertAlmostEqual(
+                p_lp, _reference_logprob(target[1], 2, 1.0), places=4
+            )
+
+    def test_pending_block_resolves_in_later_step(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected1 = torch.randn(_GAMMA, _VOCAB)
+            target1 = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=corrected1,
+                target_logits=target1,
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+            )
+
+            corrected2 = torch.randn(_GAMMA, _VOCAB)
+            target2 = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=2,
+                rid="r0",
+                drafts=[3, 6, 7],
+                corrected_logits=corrected2,
+                target_logits=target2,
+                verify_len=4,
+                correct_len=3,
+                bonus=4,
+                seq_len=12,
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            self.assertEqual(len(records), 2)
+            step2 = records[1]
+            self.assertEqual(len(step2["pg"]), 1)
+            src_fct, offset, p_lp, draft_token, realized_token = step2["pg"][0]
+            self.assertEqual((src_fct, offset), (1, 3))
+            self.assertEqual(draft_token, 3)
+            self.assertEqual(realized_token, 3)
+            self.assertAlmostEqual(
+                p_lp, _reference_logprob(target2[0], 3, 1.0), places=4
+            )
+            self.assertEqual(recorder._states["r0"].pending, [])
+
+    def test_divergence_drops_block_after_final_gather(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected1 = torch.randn(_GAMMA, _VOCAB)
+            target1 = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=corrected1,
+                target_logits=target1,
+                verify_len=2,
+                correct_len=1,
+                bonus=6,
+                seq_len=10,
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            src_fct, offset, p_lp, draft_token, realized_token = records[0]["pg"][0]
+            self.assertEqual(draft_token, 2)
+            self.assertEqual(realized_token, 6)
+            self.assertEqual(recorder._states["r0"].pending, [])
+
+    def test_temperature_scales_logprobs(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected = torch.randn(_GAMMA, _VOCAB)
+            target = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=corrected,
+                target_logits=target,
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+                temperature=0.7,
+            )
+            recorder._file.flush()
+
+            record = _read_records(path)[0]
+            self.assertAlmostEqual(
+                record["q_lp"][0],
+                _reference_logprob(corrected[1], 2, 0.7),
+                places=4,
+            )
+            self.assertAlmostEqual(
+                record["pg"][0][2],
+                _reference_logprob(target[1], 2, 0.7),
+                places=4,
+            )
+
+    def test_greedy_row_is_skipped_but_seq_len_bookkeeping_advances(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected = torch.randn(_GAMMA, _VOCAB)
+            target = torch.randn((_GAMMA + 1), _VOCAB)
+            recorder.observe_verify_step(
+                forward_ct=1,
+                rids=["r0"],
+                draft_tokens=torch.tensor([[1, 2, 3]], dtype=torch.int64),
+                corrected_logits=corrected.unsqueeze(0),
+                draft_temperatures=torch.tensor([1.0]),
+                greedy_mask=torch.tensor([True]),
+                target_logits=target,
+                target_temperatures=torch.tensor([[1.0]]),
+                truncated_sampling_mask=None,
+                logits_adjustments_are_noop=True,
+                correct_len=torch.tensor([2], dtype=torch.int32),
+                cap_trim_lens=torch.tensor([0], dtype=torch.int32),
+                bonus=torch.tensor([5], dtype=torch.int64),
+                prefix_lens=torch.tensor([10], dtype=torch.int64),
+                layout=_FakeLayout(torch.tensor([4], dtype=torch.int32)),
+            )
+            recorder._file.flush()
+
+            self.assertEqual(path.read_text(), "")
+            self.assertEqual(recorder._states["r0"].expected_seq_len, 13)
+
+    def test_seq_len_discontinuity_drops_pending_blocks(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected1 = torch.randn(_GAMMA, _VOCAB)
+            target1 = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=corrected1,
+                target_logits=target1,
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+            )
+            self.assertEqual(len(recorder._states["r0"].pending), 1)
+
+            corrected2 = torch.randn(_GAMMA, _VOCAB)
+            target2 = torch.randn((_GAMMA + 1), _VOCAB)
+            _observe(
+                recorder,
+                forward_ct=2,
+                rid="r0",
+                drafts=[3, 6, 7],
+                corrected_logits=corrected2,
+                target_logits=target2,
+                verify_len=4,
+                correct_len=3,
+                bonus=4,
+                seq_len=11,
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            self.assertNotIn("pg", records[1])
+            self.assertEqual(recorder._states["r0"].pending, [])
+            self.assertEqual(recorder._discontinuity_drop_ct, 1)
+
+    def test_truncated_sampling_row_is_excluded_while_clean_row_records(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, path = _make_recorder(tmp)
+            corrected = torch.randn(2, _GAMMA, _VOCAB)
+            target = torch.randn(2 * (_GAMMA + 1), _VOCAB)
+            recorder.observe_verify_step(
+                forward_ct=1,
+                rids=["r_clean", "r_top_p"],
+                draft_tokens=torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64),
+                corrected_logits=corrected,
+                draft_temperatures=torch.tensor([1.0, 1.0]),
+                greedy_mask=torch.tensor([False, False]),
+                target_logits=target,
+                target_temperatures=torch.tensor([[1.0], [1.0]]),
+                truncated_sampling_mask=torch.tensor([False, True]),
+                logits_adjustments_are_noop=True,
+                correct_len=torch.tensor([1, 1], dtype=torch.int32),
+                cap_trim_lens=torch.tensor([0, 0], dtype=torch.int32),
+                bonus=torch.tensor([2, 7], dtype=torch.int64),
+                prefix_lens=torch.tensor([10, 20], dtype=torch.int64),
+                layout=_FakeLayout(torch.tensor([2, 2], dtype=torch.int32)),
+            )
+            recorder._file.flush()
+
+            records = _read_records(path)
+            self.assertEqual([r["rid"] for r in records], ["r_clean"])
+            self.assertEqual(recorder._skipped_step_ct, 0)
+            self.assertEqual(recorder._states["r_top_p"].pending, [])
+            self.assertEqual(recorder._states["r_top_p"].expected_seq_len, 22)
+
+
+def _offline_estimate(path: Path, gamma: int) -> tuple[float, float, int]:
+    from collections import defaultdict
+
+    blocks: list[dict] = []
+    gathers: dict[tuple, list] = defaultdict(list)
+    for line in path.read_text().splitlines():
+        rec = json.loads(line)
+        blocks.append(rec)
+        for src_fct, offset, p_lp, draft_token, realized_token in rec.get("pg", []):
+            gathers[(rec["rid"], src_fct)].append(
+                [offset, p_lp, draft_token, realized_token]
+            )
+    los: list[float] = []
+    his: list[float] = []
+    for rec in blocks:
+        cl, w = rec["cl"], rec["w"]
+        if "q_lp" not in rec:
+            los.append(cl + 1.0)
+            his.append(cl + 1.0)
+            continue
+        q_lps = rec["q_lp"]
+        entries = {e[0]: e for e in gathers.get((rec["rid"], rec["fct"]), [])}
+        base, prod, lo_extra, tail = w + 1.0, 1.0, 0.0, 0.0
+        for offset in range(w + 1, gamma + 1):
+            entry = entries.get(offset)
+            if entry is None:
+                tail = prod * (gamma - offset + 1)
+                break
+            _, p_lp, draft_token, realized_token = entry
+            a = min(1.0, math.exp(p_lp - q_lps[offset - w - 1]))
+            prod *= a
+            lo_extra += prod
+            if draft_token != realized_token:
+                if offset < gamma:
+                    tail = prod * (gamma - offset)
+                break
+        los.append(base + lo_extra)
+        his.append(base + lo_extra + tail)
+    n = len(los)
+    return sum(los) / n, sum(his) / n, n
+
+
+class TestOnlineCeilingEstimate(CustomTestCase):
+    def test_online_estimate_matches_offline_aggregation(self):
+        bs, steps = 4, 14
+        gen = torch.Generator().manual_seed(11)
+        seq = [50 + 3 * b for b in range(bs)]
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "est.jsonl"
+            recorder = BlockAcceptEstimateRecorder(
+                path=str(path),
+                gamma=_GAMMA,
+                device="cpu",
+                online_log_interval=3,
+                online_window_steps=2,
+            )
+            for t in range(steps):
+                verify_lens, correct_lens, drafts, bonus, prefix = [], [], [], [], []
+                for b in range(bs):
+                    vl = int(torch.randint(1, _GAMMA + 2, (1,), generator=gen))
+                    window = vl - 1
+                    cl = int(torch.randint(0, window + 1, (1,), generator=gen))
+                    row = torch.randint(0, _VOCAB, (_GAMMA,), generator=gen).tolist()
+                    if cl < _GAMMA and int(torch.randint(0, 2, (1,), generator=gen)):
+                        bt = row[cl]
+                    else:
+                        bt = int(torch.randint(0, _VOCAB, (1,), generator=gen))
+                    verify_lens.append(vl)
+                    correct_lens.append(cl)
+                    drafts.append(row)
+                    bonus.append(bt)
+                    prefix.append(seq[b])
+                    seq[b] += cl + 1
+                recorder.observe_verify_step(
+                    forward_ct=t + 1,
+                    rids=[f"r{b}" for b in range(bs)],
+                    draft_tokens=torch.tensor(drafts, dtype=torch.int64),
+                    corrected_logits=torch.randn(bs, _GAMMA, _VOCAB, generator=gen),
+                    draft_temperatures=torch.ones(bs),
+                    greedy_mask=torch.zeros(bs, dtype=torch.bool),
+                    target_logits=torch.randn(bs * (_GAMMA + 1), _VOCAB, generator=gen),
+                    target_temperatures=torch.ones(bs),
+                    truncated_sampling_mask=None,
+                    logits_adjustments_are_noop=True,
+                    correct_len=torch.tensor(correct_lens, dtype=torch.int32),
+                    cap_trim_lens=torch.tensor(
+                        [_GAMMA - (v - 1) for v in verify_lens], dtype=torch.int32
+                    ),
+                    bonus=torch.tensor(bonus, dtype=torch.int64),
+                    prefix_lens=torch.tensor(prefix, dtype=torch.int64),
+                    layout=_FakeLayout(torch.tensor(verify_lens, dtype=torch.int32)),
+                )
+            recorder.drain_pending_online()
+            recorder._file.flush()
+
+            off_lo, off_hi, off_n = _offline_estimate(path, _GAMMA)
+            snap = recorder.online_estimate()
+            self.assertIsNotNone(snap)
+            self.assertEqual(snap.cumulative_blocks, off_n)
+            self.assertAlmostEqual(snap.cumulative_lo, off_lo, places=6)
+            self.assertAlmostEqual(snap.cumulative_hi, off_hi, places=6)
+            self.assertGreater(off_n, bs)
+            self.assertLessEqual(snap.window_blocks, snap.cumulative_blocks)
+
+    def test_online_window_evicts_forward_passes_outside_horizon(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "est.jsonl"
+            recorder = BlockAcceptEstimateRecorder(
+                path=str(path),
+                gamma=_GAMMA,
+                device="cpu",
+                online_log_interval=1,
+                online_window_steps=3,
+            )
+            target = torch.randn((_GAMMA + 1), _VOCAB)
+            seq = 10
+            for t in range(10):
+                cl = t % 3
+                _observe(
+                    recorder,
+                    forward_ct=t + 1,
+                    rid="r0",
+                    drafts=[1, 2, 3],
+                    corrected_logits=torch.randn(_GAMMA, _VOCAB),
+                    target_logits=target,
+                    verify_len=_GAMMA + 1,
+                    correct_len=cl,
+                    bonus=5,
+                    seq_len=seq,
+                )
+                seq += cl + 1
+            snap = recorder.online_estimate()
+            self.assertIsNotNone(snap)
+            self.assertLessEqual(snap.window_horizon, 3)
+            self.assertLessEqual(snap.window_blocks, 3)
+            self.assertEqual(snap.cumulative_blocks, 10)
+
+
+class TestNaturalStopEosTail(CustomTestCase):
+    def _finalize_kept_block(self, *, natural_stop: bool):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, _ = _make_recorder(tmp)
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=torch.randn(_GAMMA, _VOCAB),
+                target_logits=torch.randn((_GAMMA + 1), _VOCAB),
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+            )
+            self.assertEqual(len(recorder._states["r0"].pending), 1)
+            recorder.note_request_finished(rid="r0", natural_stop=natural_stop)
+            self.assertNotIn("r0", recorder._states)
+            return recorder.online_estimate()
+
+    def test_natural_eos_caps_tail_to_zero(self):
+        snap = self._finalize_kept_block(natural_stop=True)
+        self.assertEqual(snap.cumulative_blocks, 1)
+        self.assertAlmostEqual(snap.cumulative_lo, snap.cumulative_hi, places=6)
+
+    def test_external_finish_keeps_optimistic_tail(self):
+        snap = self._finalize_kept_block(natural_stop=False)
+        self.assertEqual(snap.cumulative_blocks, 1)
+        self.assertGreater(snap.cumulative_hi, snap.cumulative_lo)
+
+
+class TestAsyncFinishIntent(CustomTestCase):
+    def test_intent_buffered_then_applied_at_next_drain(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            recorder, _ = _make_recorder(tmp)
+            recorder._delayed = _FakeDelayed()
+
+            _observe(
+                recorder,
+                forward_ct=1,
+                rid="r0",
+                drafts=[1, 2, 3],
+                corrected_logits=torch.randn(_GAMMA, _VOCAB),
+                target_logits=torch.randn((_GAMMA + 1), _VOCAB),
+                verify_len=2,
+                correct_len=1,
+                bonus=2,
+                seq_len=10,
+            )
+            recorder.note_request_finished(rid="r0", natural_stop=True)
+            self.assertIn("r0", recorder._finish_intents)
+            self.assertIsNone(recorder.online_estimate())
+
+            _observe(
+                recorder,
+                forward_ct=2,
+                rid="r1",
+                drafts=[4, 5, 6],
+                corrected_logits=torch.randn(_GAMMA, _VOCAB),
+                target_logits=torch.randn((_GAMMA + 1), _VOCAB),
+                verify_len=_GAMMA + 1,
+                correct_len=1,
+                bonus=7,
+                seq_len=20,
+            )
+            self.assertNotIn("r0", recorder._finish_intents)
+            self.assertNotIn("r0", recorder._states)
+            snap = recorder.online_estimate()
+            self.assertEqual(snap.cumulative_blocks, 1)
+            self.assertAlmostEqual(snap.cumulative_lo, snap.cumulative_hi, places=6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_confidence_metrics.py

@@ -0,0 +1,159 @@
+import unittest
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.speculative.dspark_components.dspark_observability import (
+    ConfidenceMetricsProbe,
+    PerPositionConfidenceMetrics,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+
+def _cpu_metrics(gamma: int) -> PerPositionConfidenceMetrics:
+    return PerPositionConfidenceMetrics(gamma=gamma, device=torch.device("cpu"))
+
+
+class TestPerPositionConfidenceMetrics(CustomTestCase):
+    def test_perfectly_calibrated_has_low_ece(self):
+        torch.manual_seed(0)
+        n = 40000
+        survival = torch.full((n, 1), 0.3, dtype=torch.float64)
+        prefix_mask = (torch.rand(n, 1) < 0.3).to(torch.float64)
+        metrics = _cpu_metrics(gamma=1)
+        metrics.update(survival=survival, prefix_mask=prefix_mask)
+        row = metrics.compute()[0]
+        self.assertLess(row["ece"], 0.03)
+        self.assertAlmostEqual(row["pred_mean"], 0.3, places=4)
+
+    def test_overconfident_has_high_ece_and_pred_above_target(self):
+        torch.manual_seed(0)
+        n = 40000
+        survival = torch.full((n, 1), 0.9, dtype=torch.float64)
+        prefix_mask = (torch.rand(n, 1) < 0.3).to(torch.float64)
+        metrics = _cpu_metrics(gamma=1)
+        metrics.update(survival=survival, prefix_mask=prefix_mask)
+        row = metrics.compute()[0]
+        self.assertGreater(row["ece"], 0.4)
+        self.assertGreater(row["pred_mean"], row["target_mean"])
+
+    def test_separable_scores_give_auc_near_one(self):
+        torch.manual_seed(0)
+        n = 20000
+        pos = torch.rand(n, 1) * 0.3 + 0.7
+        neg = torch.rand(n, 1) * 0.3
+        survival = torch.cat([pos, neg], dim=0)
+        prefix_mask = torch.cat([torch.ones(n, 1), torch.zeros(n, 1)], dim=0)
+        metrics = _cpu_metrics(gamma=1)
+        metrics.update(survival=survival, prefix_mask=prefix_mask)
+        self.assertGreater(metrics.compute()[0]["auc"], 0.99)
+
+    def test_random_scores_give_auc_near_half(self):
+        torch.manual_seed(0)
+        n = 40000
+        survival = torch.rand(n, 1)
+        prefix_mask = (torch.rand(n, 1) < 0.5).to(torch.float64)
+        metrics = _cpu_metrics(gamma=1)
+        metrics.update(survival=survival, prefix_mask=prefix_mask)
+        auc = metrics.compute()[0]["auc"]
+        self.assertGreater(auc, 0.45)
+        self.assertLess(auc, 0.55)
+
+    def test_batched_update_matches_per_sample_update(self):
+        torch.manual_seed(0)
+        bs, gamma = 32, 5
+        survival = torch.rand(bs, gamma)
+        prefix_mask = (torch.rand(bs, gamma) < 0.5).to(torch.float64)
+
+        batched = _cpu_metrics(gamma=gamma)
+        batched.update(survival=survival, prefix_mask=prefix_mask)
+
+        per_sample = _cpu_metrics(gamma=gamma)
+        for row_idx in range(bs):
+            per_sample.update(
+                survival=survival[row_idx : row_idx + 1],
+                prefix_mask=prefix_mask[row_idx : row_idx + 1],
+            )
+
+        for name in (
+            "coarse_count",
+            "coarse_pred",
+            "coarse_target",
+            "fine_pos",
+            "fine_neg",
+            "brier_num",
+        ):
+            self.assertTrue(
+                torch.allclose(
+                    getattr(batched, name), getattr(per_sample, name), atol=1e-9
+                ),
+                msg=name,
+            )
+
+
+def _probe_inputs(bs: int, gamma: int):
+    torch.manual_seed(0)
+    verify_num_draft_tokens = gamma + 1
+    vocab = 16
+    verify_ids_2d = torch.randint(0, vocab, (bs, verify_num_draft_tokens))
+    target_logits = torch.randn(bs * verify_num_draft_tokens, vocab)
+    confidence_raw = torch.randn(bs, gamma)
+    return verify_ids_2d, target_logits, confidence_raw
+
+
+class TestConfidenceMetricsProbe(CustomTestCase):
+    def _observe(self, probe, *, carries_confidence=True, is_compact_mode=False):
+        verify_ids_2d, target_logits, confidence_raw = _probe_inputs(bs=3, gamma=4)
+        probe.maybe_observe(
+            carries_confidence=carries_confidence,
+            is_compact_mode=is_compact_mode,
+            confidence_raw=confidence_raw,
+            verify_ids_2d=verify_ids_2d,
+            target_logits=target_logits,
+            bs=3,
+        )
+
+    def test_disabled_env_is_noop(self):
+        probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0)
+        self._observe(probe)
+        self.assertIsNone(probe._metrics)
+        self.assertEqual(probe._step_ct, 0)
+
+    def test_non_rank0_is_noop(self):
+        probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=1)
+        with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True):
+            self._observe(probe)
+        self.assertIsNone(probe._metrics)
+
+    def test_missing_confidence_head_is_noop(self):
+        probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0)
+        with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True):
+            self._observe(probe, carries_confidence=False)
+        self.assertIsNone(probe._metrics)
+
+    def test_compact_mode_warns_once_and_skips(self):
+        probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0)
+        with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True):
+            self._observe(probe, is_compact_mode=True)
+            self.assertTrue(probe._compact_warned)
+            self._observe(probe, is_compact_mode=True)
+        self.assertIsNone(probe._metrics)
+        self.assertEqual(probe._step_ct, 0)
+
+    def test_enabled_path_accumulates(self):
+        probe = ConfidenceMetricsProbe(
+            gamma=4, verify_num_draft_tokens=5, tp_rank=0, print_every=2
+        )
+        with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True):
+            self._observe(probe)
+            self.assertIsInstance(probe._metrics, PerPositionConfidenceMetrics)
+            self.assertEqual(probe._step_ct, 1)
+            self._observe(probe)
+        self.assertEqual(probe._step_ct, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_dp_tier.py

@@ -0,0 +1,99 @@
+import random
+import unittest
+
+from sglang.srt.speculative.dspark_components.dspark_planner import (
+    dp_global_verify_tier_num_tokens,
+    local_verify_tier_num_tokens,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestLocalVerifyTierNumTokens(CustomTestCase):
+    def test_no_budget_returns_sentinel(self):
+        self.assertEqual(
+            local_verify_tier_num_tokens(
+                bs=8,
+                verify_token_budget=None,
+                verify_num_draft_tokens=6,
+                min_verify_len=1,
+            ),
+            -1,
+        )
+
+    def test_budget_adds_to_anchor_floor(self):
+        self.assertEqual(
+            local_verify_tier_num_tokens(
+                bs=8,
+                verify_token_budget=10,
+                verify_num_draft_tokens=6,
+                min_verify_len=1,
+            ),
+            18,
+        )
+
+    # Clamp/floor variants (verify-all clamp, min_verify_len floor, min=0) are
+    # covered by the TestBusyIdleGraphKeyIdentity sweep bounds.
+
+
+class TestDpGlobalVerifyTierNumTokens(CustomTestCase):
+    def test_any_sentinel_pins_everyone(self):
+        # The sweep never emits a -1 contribution, so this is the only guard
+        # on "any rank without a budget pins everyone"; losing it forks graph
+        # keys across DP ranks.
+        self.assertIsNone(
+            dp_global_verify_tier_num_tokens(global_tier_num_tokens=[100, -1, 50, 0])
+        )
+
+
+class TestBusyIdleGraphKeyIdentity(CustomTestCase):
+    def test_busy_and_idle_floors_agree_on_random_topologies(self):
+        rng = random.Random(20260703)
+        for _ in range(2000):
+            verify_num_draft_tokens = rng.randint(2, 8)
+            min_verify_len = rng.randint(0, verify_num_draft_tokens - 1)
+            effective_min = max(min_verify_len, 1)
+            num_ranks = rng.randint(1, 8)
+            contributions = []
+            num_reqs_per_rank = []
+            for _ in range(num_ranks):
+                if rng.random() < 0.3:
+                    num_reqs_per_rank.append(0)
+                    contributions.append(0)
+                    continue
+                bs = rng.randint(1, 512)
+                budget = rng.randint(0, bs * verify_num_draft_tokens)
+                num_reqs_per_rank.append(bs)
+                contributions.append(
+                    local_verify_tier_num_tokens(
+                        bs=bs,
+                        verify_token_budget=budget,
+                        verify_num_draft_tokens=verify_num_draft_tokens,
+                        min_verify_len=min_verify_len,
+                    )
+                )
+            tier_num_tokens = dp_global_verify_tier_num_tokens(
+                global_tier_num_tokens=contributions
+            )
+            global_num_reqs = max(num_reqs_per_rank)
+            if tier_num_tokens is None:
+                self.assertEqual(global_num_reqs, 0)
+                continue
+
+            self.assertGreaterEqual(tier_num_tokens, global_num_reqs * effective_min)
+            self.assertLessEqual(
+                tier_num_tokens, global_num_reqs * verify_num_draft_tokens
+            )
+
+            busy_floor = min(tier_num_tokens, global_num_reqs * verify_num_draft_tokens)
+            self.assertEqual(busy_floor, tier_num_tokens)
+
+            idle_lens_total = global_num_reqs
+            idle_bucket_input = max(idle_lens_total, tier_num_tokens)
+            self.assertEqual(idle_bucket_input, tier_num_tokens)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_draft_path_default.py

@@ -0,0 +1,88 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.arg_groups.speculative_hook import (
+    _handle_dspark,
+    _target_checkpoint_bundles_dspark_draft,
+)
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+_BUNDLED_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Flash-DSpark"
+_PLAIN_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Flash"
+
+
+def _bundled_hf_config() -> SimpleNamespace:
+    return SimpleNamespace(
+        architectures=["DeepseekV4ForCausalLM"],
+        dspark_block_size=5,
+        dspark_markov_rank=256,
+        dspark_target_layer_ids=[40, 41, 42],
+        dspark_noise_token_id=128799,
+    )
+
+
+def _plain_hf_config() -> SimpleNamespace:
+    return SimpleNamespace(architectures=["DeepseekV4ForCausalLM"])
+
+
+def _make_dspark_server_args(
+    *, model_path: str, hf_config: SimpleNamespace
+) -> ServerArgs:
+    server_args = ServerArgs(model_path="dummy")
+    server_args.model_path = model_path
+    server_args.device = "cuda"
+    server_args.speculative_algorithm = "DSPARK"
+    server_args.speculative_draft_model_path = None
+    server_args.speculative_dspark_block_size = 5
+    server_args.model_config = SimpleNamespace(hf_config=hf_config)
+    return server_args
+
+
+class TestTargetCheckpointBundlesDsparkDraft(CustomTestCase):
+    def test_bundled_dsv4_config_is_detected(self):
+        server_args = _make_dspark_server_args(
+            model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
+        )
+        self.assertTrue(_target_checkpoint_bundles_dspark_draft(server_args))
+
+    def test_plain_target_config_is_not_detected(self):
+        server_args = _make_dspark_server_args(
+            model_path=_PLAIN_MODEL_PATH, hf_config=_plain_hf_config()
+        )
+        self.assertFalse(_target_checkpoint_bundles_dspark_draft(server_args))
+
+
+class TestDsparkDraftPathDefaulting(CustomTestCase):
+    def test_bundled_checkpoint_defaults_draft_path_to_model_path(self):
+        server_args = _make_dspark_server_args(
+            model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
+        )
+        _handle_dspark(server_args)
+        self.assertEqual(server_args.speculative_draft_model_path, _BUNDLED_MODEL_PATH)
+        self.assertEqual(server_args.speculative_num_draft_tokens, 6)
+
+    def test_plain_target_without_draft_path_raises(self):
+        server_args = _make_dspark_server_args(
+            model_path=_PLAIN_MODEL_PATH, hf_config=_plain_hf_config()
+        )
+        with self.assertRaises(ValueError):
+            _handle_dspark(server_args)
+
+    def test_explicit_draft_path_is_not_overwritten(self):
+        server_args = _make_dspark_server_args(
+            model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
+        )
+        server_args.speculative_draft_model_path = "deepseek-ai/some-other-dspark-draft"
+        _handle_dspark(server_args)
+        self.assertEqual(
+            server_args.speculative_draft_model_path,
+            "deepseek-ai/some-other-dspark-draft",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_info_dumper.py

@@ -0,0 +1,398 @@
+import unittest
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.speculative.dspark_components.dspark_observability import (
+    DecodeStepObservation,
+    DsparkInfoDumper,
+    InfoComponent,
+    _PendingStep,
+    logger,
+    resolve_components,
+    resolve_enabled_components,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class FakeClock:
+    def __init__(self) -> None:
+        self.now = 100.0
+
+    def __call__(self) -> float:
+        return self.now
+
+    def advance(self, seconds: float) -> None:
+        self.now += seconds
+
+
+def make_dumper(components, **kwargs):
+    clock = FakeClock()
+    dumper = DsparkInfoDumper(
+        components=set(components),
+        gamma=5,
+        verify_num_draft_tokens=6,
+        attn_tp_rank=0,
+        device=torch.device("cpu"),
+        mode_value="static",
+        clock=clock,
+        **kwargs,
+    )
+    return dumper, clock
+
+
+def make_obs(
+    *,
+    forward_ct,
+    bs=4,
+    num_verify_tokens=24,
+    predicted_step_ms=None,
+    predicted_theta=None,
+):
+    return DecodeStepObservation(
+        forward_ct=forward_ct,
+        bs=bs,
+        mode="static",
+        budget=100,
+        lag_steps=0,
+        num_verify_tokens=num_verify_tokens,
+        verify_tokens_local=num_verify_tokens,
+        verify_tokens_dp_synced=num_verify_tokens,
+        verify_tokens_graph_key=num_verify_tokens,
+        predicted_step_ms=predicted_step_ms,
+        predicted_theta=predicted_theta,
+        verify_lens=torch.full((bs,), 6, dtype=torch.int32),
+        confidence=torch.full((bs, 5), 0.9),
+        req_pool_indices=torch.arange(bs, dtype=torch.int64),
+        prefix_lens=torch.full((bs,), 128, dtype=torch.int64),
+        draft_tokens=torch.zeros((bs, 5), dtype=torch.int64),
+        bonus_tokens=torch.zeros((bs,), dtype=torch.int64),
+        correct_len=torch.full((bs,), 3, dtype=torch.int32),
+        cap_trim_lens=torch.zeros((bs,), dtype=torch.int32),
+        commit_lens=torch.full((bs,), 4, dtype=torch.int32),
+        rids=[f"r{i}" for i in range(bs)],
+    )
+
+
+class TestResolveComponents(CustomTestCase):
+    def test_empty_disables(self):
+        self.assertEqual(resolve_components(()), set())
+
+    def test_all_expands_to_every_component(self):
+        self.assertEqual(resolve_components(("all",)), set(InfoComponent))
+
+    def test_subset_and_whitespace_are_kept(self):
+        self.assertEqual(
+            resolve_components((" core ", "reqs")),
+            {InfoComponent.CORE, InfoComponent.REQS},
+        )
+
+    def test_unknown_component_raises(self):
+        with self.assertRaises(ValueError):
+            resolve_components(("core", "bogus"))
+
+    def test_sps_record_env_enables_core_and_cpu_timing(self):
+        """SGLANG_DSPARK_ENABLE_SPS_RECORD=1 is the published SPS-profiling
+        switch; it must keep enabling the components the table fit reads."""
+        with envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.override(True):
+            self.assertEqual(
+                resolve_enabled_components(),
+                {InfoComponent.CORE, InfoComponent.STEP_CPU_TIME},
+            )
+
+    def test_sps_record_env_unions_with_debug_dump(self):
+        with envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.override(True):
+            with envs.SGLANG_DSPARK_DEBUG_DUMP.override("reqs"):
+                self.assertEqual(
+                    resolve_enabled_components(),
+                    {
+                        InfoComponent.CORE,
+                        InfoComponent.STEP_CPU_TIME,
+                        InfoComponent.REQS,
+                    },
+                )
+
+
+class TestCoreAndCpuTiming(CustomTestCase):
+    def test_disabled_dumper_records_nothing(self):
+        dumper, clock = make_dumper(set())
+        dumper.begin_step()
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        self.assertIsNone(dumper.dump())
+
+    def test_non_root_rank_is_disabled(self):
+        clock = FakeClock()
+        dumper = DsparkInfoDumper(
+            components={"core"},
+            gamma=5,
+            verify_num_draft_tokens=6,
+            attn_tp_rank=1,
+            device=torch.device("cpu"),
+            mode_value="static",
+            clock=clock,
+        )
+        self.assertFalse(dumper.enabled)
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        self.assertIsNone(dumper.dump())
+
+    def test_one_record_per_step_including_the_last(self):
+        dumper, clock = make_dumper({"core", "step_cpu_time"})
+        for forward_ct in range(1, 4):
+            dumper.observe_decode_step(make_obs(forward_ct=forward_ct))
+            clock.advance(0.01)
+        records = dumper.dump()["records"]
+        self.assertEqual([r["forward_ct"] for r in records], [1, 2, 3])
+
+    def test_step_cpu_ms_is_attributed_to_the_step_it_measures(self):
+        dumper, clock = make_dumper({"core", "step_cpu_time"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.02)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        records = dumper.dump()["records"]
+        first = next(r for r in records if r["forward_ct"] == 1)
+        second = next(r for r in records if r["forward_ct"] == 2)
+        self.assertNotIn("step_cpu_ms", first)
+        self.assertAlmostEqual(second["step_cpu_ms"], 20.0, places=3)
+
+    def test_core_fields_present(self):
+        dumper, _ = make_dumper({"core"})
+        dumper.observe_decode_step(make_obs(forward_ct=7, bs=3, num_verify_tokens=18))
+        record = dumper.dump()["records"][0]
+        self.assertEqual(record["bs"], 3)
+        self.assertEqual(record["num_running_reqs"], 3)
+        self.assertEqual(record["num_verify_tokens"], 18)
+        self.assertEqual(record["mode"], "static")
+
+    def test_core_only_omits_timing_fields(self):
+        dumper, clock = make_dumper({"core"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        for record in dumper.dump()["records"]:
+            self.assertNotIn("step_cpu_ms", record)
+
+    def test_non_decode_step_resets_cpu_pairing(self):
+        dumper, clock = make_dumper({"core", "step_cpu_time"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.02)
+        dumper.note_non_decode_step()
+        clock.advance(0.02)
+        dumper.observe_decode_step(make_obs(forward_ct=3))
+        records = dumper.dump()["records"]
+        self.assertEqual([r["forward_ct"] for r in records], [1, 3])
+        for record in records:
+            self.assertNotIn("step_cpu_ms", record)
+
+    def test_oversized_gap_nulls_cpu_ms_but_keeps_record(self):
+        dumper, clock = make_dumper({"core", "step_cpu_time"}, max_step_cpu_seconds=0.5)
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.6)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        records = dumper.dump()["records"]
+        self.assertEqual([r["forward_ct"] for r in records], [1, 2])
+        second = next(r for r in records if r["forward_ct"] == 2)
+        self.assertNotIn("step_cpu_ms", second)
+
+    def test_ring_buffer_evicts_oldest(self):
+        dumper, clock = make_dumper({"core"}, max_records=3)
+        for forward_ct in range(1, 8):
+            dumper.observe_decode_step(make_obs(forward_ct=forward_ct))
+            clock.advance(0.01)
+        records = dumper.dump()["records"]
+        self.assertEqual([r["forward_ct"] for r in records], [5, 6, 7])
+
+    def test_dump_is_repeatable(self):
+        dumper, clock = make_dumper({"core"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        self.assertEqual(dumper.dump(), dumper.dump())
+
+    def test_clear_drops_all_records_and_pending(self):
+        dumper, clock = make_dumper({"core"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        dumper.clear()
+        self.assertEqual(dumper.dump()["records"], [])
+        dumper.observe_decode_step(make_obs(forward_ct=9))
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=10))
+        self.assertEqual([r["forward_ct"] for r in dumper.dump()["records"]], [9, 10])
+
+
+class TestPredictedStepFields(CustomTestCase):
+    def test_predicted_fields_recorded_under_core(self):
+        dumper, clock = make_dumper({"core"})
+        dumper.observe_decode_step(
+            make_obs(forward_ct=1, predicted_step_ms=1.5, predicted_theta=200.0)
+        )
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1)
+        self.assertAlmostEqual(record["predicted_step_ms"], 1.5)
+        self.assertAlmostEqual(record["predicted_theta"], 200.0)
+
+    def test_predicted_fields_omitted_when_none(self):
+        dumper, clock = make_dumper({"core"})
+        dumper.observe_decode_step(make_obs(forward_ct=1))
+        clock.advance(0.01)
+        dumper.observe_decode_step(make_obs(forward_ct=2))
+        record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1)
+        self.assertNotIn("predicted_step_ms", record)
+        self.assertNotIn("predicted_theta", record)
+
+
+def _pending(*, bs, budget, num_verify_tokens, predicted_step_ms):
+    return _PendingStep(
+        forward_ct=1,
+        bs=bs,
+        mode="compact",
+        budget=budget,
+        lag_steps=1,
+        num_verify_tokens=num_verify_tokens,
+        verify_tokens_local=num_verify_tokens,
+        verify_tokens_dp_synced=num_verify_tokens,
+        verify_tokens_graph_key=num_verify_tokens,
+        predicted_step_ms=predicted_step_ms,
+        predicted_theta=1.0,
+        step_cpu_ms=None,
+        rids=None,
+        future=None,
+        segment_events={},
+    )
+
+
+class TestOnlineSpsReporter(CustomTestCase):
+    def test_report_interval_enables_dumper_and_gpu_timing(self):
+        dumper, _ = make_dumper(set(), sps_report_interval=2)
+        self.assertTrue(dumper.enabled)
+        self.assertIn(InfoComponent.STEP_GPU_TIME, dumper._components)
+
+    def test_report_interval_zero_leaves_dumper_disabled(self):
+        dumper, _ = make_dumper(set(), sps_report_interval=0)
+        self.assertFalse(dumper.enabled)
+
+    def test_reporter_logs_summary_every_interval_matched_steps(self):
+        dumper, _ = make_dumper(set(), sps_report_interval=2)
+        matched = dict(bs=4, budget=20, num_verify_tokens=24)
+        with self.assertLogs(logger, level="INFO") as cm:
+            dumper._report_sps_prediction(
+                pending=_pending(**matched, predicted_step_ms=10.0), step_gpu_ms=12.0
+            )
+            dumper._report_sps_prediction(
+                pending=_pending(**matched, predicted_step_ms=8.0), step_gpu_ms=9.0
+            )
+        self.assertEqual(sum("SPS prediction" in m for m in cm.output), 1)
+        self.assertEqual(dumper._sps_window, [])
+
+    def test_reporter_counts_mismatch_and_excludes_it_from_means(self):
+        dumper, _ = make_dumper(set(), sps_report_interval=1)
+        with self.assertLogs(logger, level="INFO") as cm:
+            dumper._report_sps_prediction(
+                pending=_pending(
+                    bs=4, budget=99, num_verify_tokens=24, predicted_step_ms=10.0
+                ),
+                step_gpu_ms=12.0,
+            )
+            dumper._report_sps_prediction(
+                pending=_pending(
+                    bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=10.0
+                ),
+                step_gpu_ms=12.0,
+            )
+        self.assertTrue(any("M_mismatch_rate=50.0%" in m for m in cm.output))
+
+    def test_reporter_skips_steps_missing_prediction_or_actual(self):
+        dumper, _ = make_dumper(set(), sps_report_interval=2)
+        dumper._report_sps_prediction(
+            pending=_pending(
+                bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=None
+            ),
+            step_gpu_ms=12.0,
+        )
+        dumper._report_sps_prediction(
+            pending=_pending(
+                bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=10.0
+            ),
+            step_gpu_ms=None,
+        )
+        self.assertEqual(dumper._sps_window, [])
+        self.assertEqual(dumper._sps_mismatched, 0)
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA for d2h staging")
+class TestReqsAndGpuTiming(CustomTestCase):
+    def _cuda_obs(self, *, forward_ct, bs=4):
+        obs = make_obs(forward_ct=forward_ct, bs=bs)
+        return DecodeStepObservation(
+            forward_ct=obs.forward_ct,
+            bs=obs.bs,
+            mode=obs.mode,
+            budget=obs.budget,
+            lag_steps=obs.lag_steps,
+            num_verify_tokens=obs.num_verify_tokens,
+            verify_tokens_local=obs.verify_tokens_local,
+            verify_tokens_dp_synced=obs.verify_tokens_dp_synced,
+            verify_tokens_graph_key=obs.verify_tokens_graph_key,
+            predicted_step_ms=obs.predicted_step_ms,
+            predicted_theta=obs.predicted_theta,
+            verify_lens=obs.verify_lens.cuda(),
+            confidence=obs.confidence.cuda(),
+            req_pool_indices=obs.req_pool_indices.cuda(),
+            prefix_lens=obs.prefix_lens.cuda(),
+            draft_tokens=obs.draft_tokens.cuda(),
+            bonus_tokens=obs.bonus_tokens.cuda(),
+            correct_len=obs.correct_len.cuda(),
+            cap_trim_lens=obs.cap_trim_lens.cuda(),
+            commit_lens=obs.commit_lens.cuda(),
+            rids=obs.rids,
+        )
+
+    def _make(self, components):
+        return DsparkInfoDumper(
+            components=set(components),
+            gamma=5,
+            verify_num_draft_tokens=6,
+            attn_tp_rank=0,
+            device=torch.device("cuda"),
+            mode_value="static",
+        )
+
+    def test_reqs_component_stages_per_request_detail(self):
+        dumper = self._make({"core", "reqs"})
+        dumper.observe_decode_step(self._cuda_obs(forward_ct=1, bs=3))
+        dumper.observe_decode_step(self._cuda_obs(forward_ct=2, bs=3))
+        record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1)
+        self.assertEqual(len(record["reqs"]), 3)
+        req = record["reqs"][0]
+        self.assertEqual(req["rid"], "r0")
+        self.assertEqual(req["verify_len"], 6)
+        self.assertEqual(req["acc_len"], 4)
+        self.assertEqual(req["correct_drafts"], 3)
+        self.assertEqual(len(req["survival"]), 5)
+
+    def test_gpu_timing_populates_segment_fields(self):
+        dumper = self._make(
+            {"step_gpu_time", "draft_gpu_time", "target_verify_gpu_time"}
+        )
+        for forward_ct in (1, 2):
+            dumper.begin_step()
+            with dumper.segment("draft"):
+                torch.zeros(1024, device="cuda").sum()
+            with dumper.segment("target_verify"):
+                torch.zeros(1024, device="cuda").sum()
+            dumper.observe_decode_step(self._cuda_obs(forward_ct=forward_ct))
+        record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1)
+        # The segments launch real kernels, so resolved event pairs must
+        # measure strictly positive time; 0.0 would mean the events never ran.
+        self.assertGreater(record["step_gpu_ms"], 0.0)
+        self.assertGreater(record["draft_gpu_ms"], 0.0)
+        self.assertGreater(record["target_verify_gpu_ms"], 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_kernel_parity.py

@@ -0,0 +1,562 @@
+"""Seeded triton-vs-torch parity sweep for the DSpark kernels.
+
+Guards against toolchain drift (triton/torch upgrades) silently diverging
+the triton implementations from their torch references. Every method calls
+the production kernel pair directly via `Cls.torch(...)` / `Cls.triton(...)`
+(no env-var dispatch) on a small set of adversarial inputs and compares
+exactly, or with the tolerance the kernel is specified to meet.
+"""
+
+import types
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.dsv4 import attn_metadata_kernels
+from sglang.srt.speculative import ragged_verify_kernels
+from sglang.srt.speculative.dspark_components.dspark_planner import (
+    DSparkScheduleConfig,
+)
+from sglang.srt.speculative.dspark_components.kernels import (
+    dspark_accept,
+    dspark_attn_metadata,
+    dspark_draft_model,
+    dspark_schedule,
+    dspark_verify_window,
+)
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
+
+DEVICE = torch.device("cuda")
+VOCAB = 129280
+
+
+def _ri(lo, hi, shape, dtype=torch.int64, g=None):
+    return torch.randint(lo, hi, shape, device=DEVICE, dtype=dtype, generator=g)
+
+
+def _layout(verify_lens, graph_num_tokens):
+    return RaggedVerifyLayout.from_verify_lens_device(
+        verify_lens=verify_lens, graph_num_tokens=graph_num_tokens
+    )
+
+
+class _Bf16Linear(torch.nn.Module):
+    quant_method = None
+
+    def __init__(self, weight):
+        super().__init__()
+        self.weight = weight
+
+    def forward(self, x):
+        return torch.nn.functional.linear(x, self.weight), None
+
+
+def _case_accept_greedy(tc):
+    torch.manual_seed(0)
+    bs, t = 8, 6
+    candidates = _ri(0, 200, (bs, t))
+    target_logits = torch.randn(bs * t, 200, device=DEVICE)
+    for cutoff in (None, _ri(1, t + 1, (bs,), torch.int32)):
+        tc._parity(
+            dspark_accept.AcceptGreedy,
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=t,
+            cutoff_verify_lens=cutoff,
+        )
+    # gather_row_bonus: bonus token at a per-row column index.
+    table, idx = _ri(0, VOCAB, (64, t)), _ri(0, t, (64,), torch.int32)
+    ref = table[torch.arange(64, device=DEVICE), idx.long()]
+    tc._eq(dspark_accept.gather_row_bonus_triton(table=table, idx=idx), ref)
+
+
+def _case_accept_sampling(tc):
+    torch.manual_seed(1)
+    bs, t = 64, 6
+    accept_index = _ri(0, bs * t, (bs, t))
+    predicts = _ri(0, VOCAB, (bs * t,))
+    correct_len = _ri(0, t, (bs,), torch.int32)
+    rows = torch.arange(bs, device=DEVICE)
+    ref = predicts[accept_index[rows, correct_len.long()].long()]
+    got = dspark_accept.gather_two_level_bonus_triton(
+        accept_index=accept_index, predicts=predicts, correct_len=correct_len
+    )
+    tc._eq(got, ref)
+
+
+def _case_build_block_seq_lens_causal(tc):
+    torch.manual_seed(2)
+    seq_lens = _ri(1, 100000, (128,))
+    for block_size in (1, 5, 7):
+        tc._parity(
+            dspark_attn_metadata.BuildBlockSeqLensCausal,
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=DEVICE,
+        )
+
+
+def _case_build_out_tokens(tc):
+    torch.manual_seed(3)
+    bs, gamma = 64, 5
+    for cl_dtype in (torch.int32, torch.int64):
+        # Bonus insertion swept through every position 0..gamma.
+        cl = (torch.arange(bs, device=DEVICE) % (gamma + 1)).to(cl_dtype)
+        tc._parity(
+            dspark_verify_window.BuildOutTokens,
+            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
+            correct_len=cl,
+            bonus=_ri(0, VOCAB, (bs,)),
+            verify_num_draft_tokens=gamma + 1,
+            gamma=gamma,
+        )
+
+
+def _case_build_ragged_verify_window(tc):
+    torch.manual_seed(4)
+    gamma, t, bs = 5, 6, 8
+    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
+    batch = types.SimpleNamespace(
+        seq_lens=_ri(1, 20, (bs,)),
+        req_pool_indices=torch.randperm(bs + 3, device=DEVICE)[:bs],
+    )
+    model_runner = types.SimpleNamespace(
+        req_to_token_pool=types.SimpleNamespace(
+            req_to_token=_ri(0, 1_000_000, (bs + 3, 64), torch.int32)
+        )
+    )
+    for graph_num_tokens in (bs * t, (bs + 3) * t):  # tight and bucket padding
+        tc._parity(
+            dspark_verify_window.BuildRaggedVerifyWindow,
+            batch=batch,
+            layout=_layout(verify_lens, graph_num_tokens),
+            draft_block_ids=_ri(0, VOCAB, (bs, gamma)),
+            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
+            bs=bs,
+            device=DEVICE,
+            verify_num_draft_tokens=t,
+            model_runner=model_runner,
+        )
+
+
+def _case_build_step_local(tc):
+    torch.manual_seed(5)
+    for org_width, per_partition, bias_dtype in (
+        (32320, 32384, torch.bfloat16),
+        (5000, 8192, torch.float32),
+    ):
+        bias = (torch.randn(3, org_width, device=DEVICE) * 3.0).to(bias_dtype)
+        base = torch.randn(3, per_partition, device=DEVICE)
+        got, _ = tc._parity(
+            dspark_draft_model.BuildStepLocal, bias=bias, base_local=base
+        )
+        # Padding columns beyond org_width must stay pure base.
+        tc.assertTrue(torch.equal(got[:, org_width:], base[:, org_width:]))
+
+
+def _case_cap_correct_len(tc):
+    torch.manual_seed(6)
+    bs, nd = 64, 6
+    verify_lens = _ri(1, nd + 1, (bs,), torch.int32)
+    for cl_dtype in (torch.int32, torch.int64):
+        cl = (torch.arange(bs, device=DEVICE) % (nd + 1)).to(cl_dtype)
+        tc._parity(dspark_accept.CapCorrectLen, correct_len=cl, verify_lens=verify_lens)
+
+
+def _case_causal_swa_page_indices(tc):
+    swa, num_pool, pool_len, num_q = 128, 64, 600, 40
+    g = torch.Generator(device=DEVICE).manual_seed(7)
+    kw = dict(
+        req_to_token=_ri(0, 40000, (num_pool, pool_len), torch.int32, g),
+        full_to_swa_mapping=_ri(0, 1 << 20, (40000,), torch.int64, g),
+        req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g),
+        swa_window=swa,
+        page_index_aligned_size=96,
+    )
+    # Lens short of / straddling / beyond the SWA window boundary.
+    for lo, hi in ((1, swa), (swa - 4, swa + 4), (swa + 1, pool_len)):
+        lens = _ri(lo, hi, (num_q,), torch.int32, g)
+        cls = attn_metadata_kernels.BuildCausalSwaPageIndices
+        ref = cls.torch(seq_lens_casual=lens, **kw)
+        got = cls.triton(seq_lens_casual=lens, **kw)
+        tc.assertEqual(got.shape, ref.shape)
+        tc.assertEqual(got.dtype, ref.dtype)
+        # Parity holds on the attended region; padding slots must be -1.
+        col = torch.arange(ref.shape[1], device=DEVICE).view(1, -1)
+        attended = col < torch.clamp(lens, max=swa).view(-1, 1)
+        tc.assertTrue(torch.equal(got[attended], ref[attended]))
+        tc.assertTrue(bool((got[~attended] == -1).all()))
+
+
+def _case_commit_inject_layout(tc):
+    stride, num_pool, pool_len, n_full, bs = 7, 300, 400, 50000, 64
+    g = torch.Generator(device=DEVICE).manual_seed(8)
+    pool_perm = torch.randperm(num_pool, device=DEVICE, generator=g)
+    kw = dict(
+        req_pool_indices=pool_perm[:bs],
+        req_to_token=_ri(0, n_full, (num_pool, pool_len), torch.int64, g),
+        prefix_lens=_ri(1, pool_len - stride, (bs,), torch.int64, g),
+        block_pos_offsets=torch.arange(stride, device=DEVICE),
+        full_to_swa_mapping=_ri(0, 1 << 20, (n_full,), torch.int64, g),
+        commit_lens=_ri(0, stride + 1, (bs,), torch.int32, g),
+        stride=stride,
+    )
+    tc._parity(dspark_verify_window.BuildCommitInjectLayout, **kw)
+    # commit_len edges: 0 masks the whole row to -1, stride keeps it all.
+    kw.update(
+        req_pool_indices=kw["req_pool_indices"][:2],
+        prefix_lens=kw["prefix_lens"][:2],
+        commit_lens=torch.tensor([0, stride], device=DEVICE, dtype=torch.int32),
+    )
+    edge = dspark_verify_window.BuildCommitInjectLayout.triton(**kw)
+    swa_2d = edge.swa_loc.view(2, stride)
+    tc.assertTrue(bool((swa_2d[0] == -1).all()))
+    tc.assertTrue(bool((swa_2d[1] >= 0).all()))
+
+
+def _case_commit_kv_proj(tc):
+    hidden, head_dim, num_stages = 1024, 576, 3
+    g = torch.Generator(device=DEVICE).manual_seed(9)
+    linears = [
+        _Bf16Linear(
+            (torch.randn(head_dim, hidden, device=DEVICE, generator=g) * 0.02).to(
+                torch.bfloat16
+            )
+        )
+        for _ in range(num_stages)
+    ]
+    main_x = (torch.randn(56, hidden, device=DEVICE, generator=g) * 0.5).to(
+        torch.bfloat16
+    )
+    cls = dspark_draft_model.CommitKvProj
+    ref = cls.torch(main_x=main_x, wkv_linears=linears)
+    got = cls.triton(main_x=main_x, wkv_linears=linears)
+    tc.assertEqual(len(got), num_stages)
+    for kv_got, kv_ref in zip(got, ref):
+        tc.assertEqual(kv_got.shape, kv_ref.shape)
+        tc.assertTrue(kv_got.is_contiguous())
+        torch.testing.assert_close(kv_got.float(), kv_ref.float(), rtol=2e-2, atol=2e-3)
+    # fp8 blockwise weight dequant path (2x3 grid of 128x128 blocks).
+    out_dim, in_dim, block = 192, 384, 128
+    w8 = torch.randn(out_dim, in_dim, device=DEVICE, generator=g).to(
+        torch.float8_e4m3fn
+    )
+    scale = torch.rand(2, 3, device=DEVICE, generator=g) + 0.5
+    sf = scale.repeat_interleave(block, 0)[:out_dim]
+    sf = sf.repeat_interleave(block, 1)[:, :in_dim]
+    expected = (w8.to(torch.float32) * sf).to(torch.bfloat16)
+    stub = types.SimpleNamespace(weight=w8, weight_scale_inv=scale)
+    tc._eq(dspark_draft_model._dequant_linear_weight(stub), expected)
+
+
+def _case_compact_layout(tc):
+    torch.manual_seed(10)
+    gamma, t, bs = 5, 6, 64
+    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
+    total = int(verify_lens.sum().item())
+    for padded_total in (total, bs * t):  # exact and bucket padding
+        tc._parity(
+            dspark_verify_window.CompactRowIndex,
+            verify_lens=verify_lens,
+            padded_total=padded_total,
+            device=DEVICE,
+        )
+        tc._parity(
+            dspark_verify_window.CompactVerifyIds,
+            draft_block_ids=_ri(0, VOCAB, (bs, gamma)),
+            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
+            layout=_layout(verify_lens, padded_total),
+            device=DEVICE,
+        )
+
+
+def _case_swa_page_indices(tc):
+    torch.manual_seed(11)
+    block_size, num_q, max_reqs, n_full = 5, 320, 300, 50000
+    _, gather = tc._parity(
+        dspark_attn_metadata.ComputeDsparkWindowGather,
+        seq_lens_casual=_ri(1, 300, (num_q,), torch.int32),
+        req_pool_indices_repeated=_ri(0, max_reqs, (num_q,)),
+        block_size=block_size,
+        swa_window=128,
+    )
+    tc._parity(
+        dspark_attn_metadata.BuildDsparkSwaPageIndices,
+        req_to_token=_ri(0, n_full, (max_reqs, 400), torch.int32),
+        full_to_swa_mapping=_ri(0, 20000, (n_full,), torch.int32),
+        req_pool_indices_per_request=gather.req_pool_indices_per_request,
+        offsets=gather.offsets,
+        invalid=gather.invalid,
+        out_loc=_ri(0, n_full, (num_q,)),
+        context_lens=gather.context_lens,
+        block_size=block_size,
+        swa_window=128,
+        page_index_aligned_size=64,
+    )
+
+
+def _case_expand_prefill_causally(tc):
+    torch.manual_seed(12)
+    # Vectorized branch: ragged extends with padded token count.
+    bs = 64
+    extend = _ri(1, 8, (bs,))
+    num_tokens = int(extend.sum())
+    req_pool_indices = torch.randperm(512, device=DEVICE)[:bs]
+    seq_lens = _ri(8, 500, (bs,))
+    tc._parity(
+        attn_metadata_kernels.ExpandPrefillCausally,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        extend_seq_lens=extend,
+        extend_start_loc=torch.cumsum(extend, dim=0) - extend,
+        seq_lens_cpu=None,
+        extend_seq_lens_cpu=None,
+        num_tokens=num_tokens,
+        padded_num_tokens=num_tokens + 5,
+    )
+    # Loop branch: uniform extend with CPU lens and no padding.
+    bs2, block = 8, 6
+    tc._parity(
+        attn_metadata_kernels.ExpandPrefillCausally,
+        req_pool_indices=req_pool_indices[:bs2],
+        seq_lens=seq_lens[:bs2],
+        extend_seq_lens=torch.full((bs2,), block, device=DEVICE),
+        extend_start_loc=None,
+        seq_lens_cpu=[int(x) for x in seq_lens[:bs2].tolist()],
+        extend_seq_lens_cpu=[block] * bs2,
+        num_tokens=bs2 * block,
+        padded_num_tokens=None,
+    )
+
+
+def _case_finalize_accept_lens(tc):
+    torch.manual_seed(13)
+    bs = 64
+    for prefix_dtype in (torch.int32, torch.int64):
+        tc._parity(
+            dspark_accept.FinalizeAcceptLens,
+            correct_len=_ri(0, 7, (bs,), torch.int32),
+            cap_trim_lens=_ri(0, 4, (bs,)),
+            prefix_lens=_ri(1, 4000, (bs,), prefix_dtype),
+        )
+
+
+def _case_mixed_accept_select(tc):
+    torch.manual_seed(14)
+    bs = 64
+    # Mixed dtypes between the greedy and sampling lanes.
+    tc._parity(
+        dspark_accept.SelectMixedAccept,
+        greedy_mask=torch.rand(bs, device=DEVICE) < 0.5,
+        greedy_len=_ri(0, 7, (bs,)),
+        greedy_bonus=_ri(0, 100000, (bs,)),
+        greedy_trim=_ri(0, 4, (bs,)),
+        sampling_len=_ri(0, 7, (bs,), torch.int32),
+        sampling_bonus=_ri(0, 100000, (bs,)),
+        sampling_trim=_ri(0, 4, (bs,), torch.int32),
+    )
+
+
+def _case_padded_to_bucket(tc):
+    torch.manual_seed(15)
+    for bs, padded_bs, graph_num_tokens in ((3, 6, 16), (2, 8, 16), (8, 128, 768)):
+        verify_lens = _ri(1, 7, (bs,), torch.int32)
+        if int(verify_lens.sum()) > graph_num_tokens:
+            verify_lens = torch.ones(bs, dtype=torch.int32, device=DEVICE)
+        got, _ = tc._parity(
+            ragged_verify_kernels.PaddedToBucket,
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            bs=bs,
+            padded_bs=padded_bs,
+        )
+        # Padding rows must absorb exactly the leftover budget.
+        tc.assertEqual(int(got.to(torch.int64).sum()), graph_num_tokens)
+        if padded_bs > bs:
+            tc.assertTrue(torch.equal(got[:bs], verify_lens))
+
+
+def _case_page_table_positions(tc):
+    num_pool, pool_len = 128, 4096
+    g = torch.Generator(device=DEVICE).manual_seed(16)
+    req_to_token = _ri(0, 1 << 20, (num_pool, pool_len), torch.int32, g)
+    # Large page + non-pool-aligned max_seq_len, then page_size 1.
+    for num_q, page_size, max_seq_len in ((300, 64, 4000), (56, 1, 4096)):
+        tc._parity(
+            attn_metadata_kernels.BuildPageTablePositions,
+            req_to_token=req_to_token,
+            req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g),
+            seq_lens_casual=_ri(1, pool_len, (num_q,), torch.int64, g),
+            max_seq_len=max_seq_len,
+            page_size=page_size,
+            swa_window=128,
+        )
+
+
+def _case_qo_indptr(tc):
+    torch.manual_seed(17)
+    cls = ragged_verify_kernels.BuildQoIndptr
+    for dtype in (torch.int32, torch.int64):
+        verify_lens = _ri(1, 8, (129,), dtype)  # straddles the 128 block
+        ref = cls.torch(verify_lens=verify_lens)
+        got = cls.triton(verify_lens=verify_lens.to(torch.int32))
+        tc._eq(got, ref)
+    # Aliasing regression: the two outputs must not share storage.
+    vl = torch.tensor([3, 1, 5], device=DEVICE, dtype=torch.int32)
+    got = cls.triton(verify_lens=vl)
+    got.extend_start_loc.fill_(-7)
+    tc.assertEqual(got.qo_indptr[:2].tolist(), [0, 3])
+
+
+def _case_sample_step_tokens(tc):
+    torch.manual_seed(18)
+    cls = dspark_draft_model.SampleStepTokens
+    # Injected noise makes stochastic sampling exactly comparable.
+    for vocab, dtype in ((130000, torch.bfloat16), (5003, torch.float32)):
+        bs = 3
+        tc._parity(
+            cls,
+            step_logits=(torch.randn(bs, vocab, device=DEVICE) * 4.0).to(dtype),
+            temperatures=torch.rand(bs, device=DEVICE) + 0.5,
+            greedy_mask=(torch.arange(bs, device=DEVICE) % 2) == 0,
+            exp_noise=torch.empty(bs, vocab, device=DEVICE).exponential_(1),
+        )
+    # Greedy tie straddling a triton block boundary picks the smaller index.
+    logits = torch.zeros(1, 2050, device=DEVICE)
+    logits[0, 1000] = logits[0, 1100] = 5.0
+    tokens = cls.triton(
+        step_logits=logits,
+        temperatures=torch.tensor([1.0], device=DEVICE),
+        greedy_mask=torch.tensor([True], device=DEVICE),
+        exp_noise=torch.ones(1, 2050, device=DEVICE),
+    )
+    tc.assertEqual(tokens.item(), 1000)
+    # Non-contiguous strided cropped view must match its contiguous copy.
+    view = (torch.randn(2, 129536, device=DEVICE) * 4.0)[:, :VOCAB]
+    tc.assertFalse(view.is_contiguous())
+    kw = dict(
+        temperatures=torch.rand(2, device=DEVICE) + 0.5,
+        greedy_mask=torch.tensor([True, False], device=DEVICE),
+        exp_noise=torch.empty(2, VOCAB, device=DEVICE).exponential_(1),
+    )
+    tc._eq(
+        cls.triton(step_logits=view, **kw),
+        cls.triton(step_logits=view.contiguous(), **kw),
+    )
+
+
+def _case_scatter_compact_to_strided(tc):
+    torch.manual_seed(19)
+    t, bs, dim = 6, 8, 4096
+    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
+    total = int(verify_lens.sum().item())
+    for graph_num_tokens in (total, bs * t):  # exact and bucket padding
+        compact = torch.randn(
+            graph_num_tokens, dim, dtype=torch.bfloat16, device=DEVICE
+        )
+        tc._parity(
+            dspark_verify_window.ScatterCompactToStrided,
+            compact=compact,
+            layout=_layout(verify_lens, graph_num_tokens),
+            fill_value=0.0,
+            verify_num_draft_tokens=t,
+        )
+
+
+def _case_schedule_verify_lens_topk(tc):
+    torch.manual_seed(20)
+    gamma, bs = 5, 64
+    cfg = DSparkScheduleConfig(gamma=gamma)
+    cls = dspark_schedule.ScheduleVerifyLensTopk
+    base = torch.rand(bs, gamma, device=DEVICE)
+    confidences = (
+        torch.full((bs, gamma), 0.5, device=DEVICE),  # all-ties
+        (base * 4).floor() / 4,  # coarse quantization
+        torch.where(base < 0.3, torch.zeros_like(base), base),  # invalid zeros
+    )
+    for confidence in confidences:
+        for budget in (0, 1, 3, 7, 1000):
+            tc._parity(cls, confidence=confidence, budget=budget, cfg=cfg)
+
+
+def _case_softmax_temp(tc):
+    g = torch.Generator(device=DEVICE).manual_seed(21)
+    cls = dspark_accept.SoftmaxTemp
+    # bf16 logits, non-power-of-two rows_per_request, full vocab.
+    logits = (torch.randn(56, VOCAB, device=DEVICE, generator=g) * 8.0).to(
+        torch.bfloat16
+    )
+    temps = (torch.rand(8, device=DEVICE, generator=g) * 1.5 + 0.05).float()
+    ref = cls.torch(logits=logits, temperatures=temps, rows_per_request=7)
+    got = cls.triton(logits=logits, temperatures=temps, rows_per_request=7)
+    tc.assertEqual(got.dtype, torch.float32)
+    torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-6)
+    torch.testing.assert_close(
+        got.sum(dim=-1), torch.ones_like(got.sum(dim=-1)), rtol=1e-5, atol=1e-5
+    )
+    # Column-shaped (bs, 1) temperatures.
+    logits2 = torch.randn(6, 512, device=DEVICE, generator=g).to(torch.bfloat16)
+    temps2 = (torch.rand(2, 1, device=DEVICE, generator=g) + 0.3).float()
+    ref2 = cls.torch(logits=logits2, temperatures=temps2, rows_per_request=3)
+    got2 = cls.triton(logits=logits2, temperatures=temps2, rows_per_request=3)
+    torch.testing.assert_close(got2, ref2, rtol=1e-5, atol=1e-7)
+
+
+_CASES = [
+    ("accept_greedy", _case_accept_greedy),
+    ("accept_sampling", _case_accept_sampling),
+    ("build_block_seq_lens_causal", _case_build_block_seq_lens_causal),
+    ("build_out_tokens", _case_build_out_tokens),
+    ("build_ragged_verify_window", _case_build_ragged_verify_window),
+    ("build_step_local", _case_build_step_local),
+    ("cap_correct_len", _case_cap_correct_len),
+    ("causal_swa_page_indices", _case_causal_swa_page_indices),
+    ("commit_inject_layout", _case_commit_inject_layout),
+    ("commit_kv_proj", _case_commit_kv_proj),
+    ("compact_layout", _case_compact_layout),
+    ("swa_page_indices", _case_swa_page_indices),
+    ("expand_prefill_causally", _case_expand_prefill_causally),
+    ("finalize_accept_lens", _case_finalize_accept_lens),
+    ("mixed_accept_select", _case_mixed_accept_select),
+    ("padded_to_bucket", _case_padded_to_bucket),
+    ("page_table_positions", _case_page_table_positions),
+    ("qo_indptr", _case_qo_indptr),
+    ("sample_step_tokens", _case_sample_step_tokens),
+    ("scatter_compact_to_strided", _case_scatter_compact_to_strided),
+    ("schedule_verify_lens_topk", _case_schedule_verify_lens_topk),
+    ("softmax_temp", _case_softmax_temp),
+]
+
+
+class TestDsparkKernelParity(CustomTestCase):
+    def _eq(self, got, ref):
+        """Exact comparison of tensors, tuples, and msgspec result structs."""
+        if isinstance(ref, tuple):
+            for g, r in zip(got, ref):
+                self._eq(g, r)
+        elif hasattr(ref, "__struct_fields__"):
+            for name in ref.__struct_fields__:
+                self._eq(getattr(got, name), getattr(ref, name))
+        elif isinstance(ref, torch.Tensor):
+            self.assertEqual(got.dtype, ref.dtype)
+            self.assertTrue(torch.equal(got, ref))
+        else:
+            self.assertEqual(got, ref)
+
+    def _parity(self, cls, **kw):
+        got, ref = cls.triton(**kw), cls.torch(**kw)
+        self._eq(got, ref)
+        return got, ref
+
+    def test_all_kernels_triton_matches_torch(self):
+        for name, case in _CASES:
+            with self.subTest(kernel=name):
+                case(self)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_scheduler.py

@@ -0,0 +1,558 @@
+import functools
+import types
+import unittest
+
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_planner import (
+    DSparkScheduleConfig,
+    HostConfidenceBudgetPlanner,
+    VerifyBudgetDecision,
+    compute_verify_token_budget,
+    graph_tier_fill_budget,
+)
+from sglang.srt.speculative.dspark_components.dspark_sps import (
+    SpsAdditiveCostTable,
+    SpsCostTable,
+)
+from sglang.srt.speculative.dspark_components.kernels.dspark_schedule import (
+    schedule_verify_lens_topk_from_survival,
+)
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+
+def _flat_table(
+    steps_per_sec: float = 1.0, max_batch_tokens: int = 4096
+) -> SpsCostTable:
+    return SpsCostTable(
+        sample_batch_tokens=[1],
+        sample_steps_per_sec=[steps_per_sec],
+        max_batch_tokens=max_batch_tokens,
+    )
+
+
+def _cliff_table() -> SpsCostTable:
+    return SpsCostTable(
+        sample_batch_tokens=[1, 2, 3, 4, 5, 6, 7, 8],
+        sample_steps_per_sec=[1.0, 1.0, 1.0, 0.5, 0.45, 0.44, 0.43, 0.42],
+        max_batch_tokens=64,
+    )
+
+
+def _additive_table() -> SpsAdditiveCostTable:
+    return SpsAdditiveCostTable(
+        bias_seconds=0.01,
+        bs_probes=[1, 100],
+        alpha_seconds=[0.0, 0.05],
+        m_probes=[1, 200],
+        theta_seconds=[0.0, 0.02],
+    )
+
+
+def _survival_from_confidence(confidence: torch.Tensor) -> torch.Tensor:
+    return torch.cumprod(confidence, dim=1)
+
+
+def _bruteforce_budget(
+    *,
+    history_survival_probs: torch.Tensor,
+    sps_table: SpsCostTable,
+    cfg: DSparkScheduleConfig,
+) -> int:
+    num_requests = history_survival_probs.shape[0]
+    max_len = cfg.resolved_max_verify_len()
+    candidates = history_survival_probs[:, :max_len].flatten()
+    candidates = [float(x) for x in candidates.tolist() if float(x) >= cfg.survival_eps]
+    candidates.sort(reverse=True)
+
+    best_extra, best_theta = 0, float("-inf")
+    for extra in range(len(candidates) + 1):
+        tau_star = num_requests + sum(candidates[:extra])
+        theta = tau_star * sps_table.lookup(num_requests + extra)
+        if theta > best_theta:
+            best_theta, best_extra = theta, extra
+    return best_extra
+
+
+def schedule_verify_lens_topk_vanilla(
+    *,
+    survival_probs: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    cfg.validate()
+    num_requests, _gamma = survival_probs.shape
+    max_len = cfg.resolved_max_verify_len()
+    device = survival_probs.device
+
+    valid_rows = (survival_probs >= cfg.survival_eps).tolist()
+    survival_rows = survival_probs.to(torch.float64).tolist()
+
+    candidates: list[tuple[float, int, int]] = []
+    for request in range(num_requests):
+        for position in range(min(max_len, len(valid_rows[request]))):
+            if valid_rows[request][position]:
+                candidates.append((survival_rows[request][position], position, request))
+
+    candidates.sort(key=lambda candidate: (-candidate[0], candidate[1], candidate[2]))
+
+    selected_extra = [0] * num_requests
+    for _survival, _position, request in candidates[: max(int(budget), 0)]:
+        selected_extra[request] += 1
+
+    lower_bound = max(cfg.min_verify_len, 1)
+    verify_lens = [
+        min(max(cfg.min_verify_len + extra, lower_bound), max_len)
+        for extra in selected_extra
+    ]
+    return torch.tensor(verify_lens, dtype=torch.int32, device=device)
+
+
+_TOPK_IMPLS = (
+    schedule_verify_lens_topk_from_survival,
+    schedule_verify_lens_topk_vanilla,
+)
+
+
+def _for_each_impl(test_method):
+    @functools.wraps(test_method)
+    def wrapper(self):
+        for impl in _TOPK_IMPLS:
+            with self.subTest(impl=impl.__name__):
+                test_method(self, impl)
+
+    return wrapper
+
+
+class TestComputeVerifyTokenBudget(CustomTestCase):
+    def test_budget_argmax_matches_bruteforce_scan_across_sps_cliffs(self):
+        torch.manual_seed(1)
+        cfg = DSparkScheduleConfig(gamma=7)
+        table = _cliff_table()
+        for _ in range(20):
+            confidence = torch.rand(3, 7, dtype=torch.float32) * 0.5 + 0.45
+            survival = _survival_from_confidence(confidence)
+            expected = _bruteforce_budget(
+                history_survival_probs=survival, sps_table=table, cfg=cfg
+            )
+            actual = compute_verify_token_budget(
+                history_survival_probs=survival, sps_table=table, cfg=cfg
+            ).budget
+            self.assertEqual(actual, expected)
+
+    def test_budget_pool_is_independent_of_min_verify_len(self):
+        survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32)
+        cfg_no_min = DSparkScheduleConfig(gamma=4, min_verify_len=0)
+        cfg_min2 = DSparkScheduleConfig(gamma=4, min_verify_len=2)
+        table = _flat_table()
+        budget_no_min = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg_no_min
+        ).budget
+        budget_min2 = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg_min2
+        ).budget
+        # The budget candidate pool spans all positions; the per-request floor
+        # is enforced later at verify-lens scheduling, not at budget time.
+        self.assertEqual(budget_no_min, 4)
+        self.assertEqual(budget_min2, 4)
+
+    def test_budget_drops_candidates_below_survival_eps(self):
+        survival = torch.tensor([[0.9, 1e-9, 1e-12]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=3, survival_eps=1e-6)
+        table = _flat_table()
+        budget = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg
+        ).budget
+        self.assertLessEqual(budget, 1)
+
+    def test_decision_predicted_step_matches_additive_table_at_budget(self):
+        survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=4)
+        table = _additive_table()
+        decision = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg
+        )
+        # Reference via the independent scalar interpolation path (step_time),
+        # not the tensor helper the implementation itself uses.
+        self.assertAlmostEqual(
+            decision.predicted_step_seconds,
+            table.step_time(num_reqs=1, budget=int(decision.budget)),
+            places=5,
+        )
+        self.assertGreater(decision.predicted_theta, 0.0)
+
+    def test_decision_predicted_step_is_inverse_sps_for_diagonal_table(self):
+        survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=4)
+        table = _cliff_table()
+        decision = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg
+        )
+        expected_sps = table.lookup(1 + decision.budget)
+        self.assertIsNotNone(decision.predicted_step_seconds)
+        self.assertAlmostEqual(
+            decision.predicted_step_seconds, 1.0 / expected_sps, places=9
+        )
+        self.assertGreater(decision.predicted_theta, 0.0)
+
+
+def _make_budget_planner() -> HostConfidenceBudgetPlanner:
+    return HostConfidenceBudgetPlanner(
+        sps_table=_flat_table(),
+        cfg=DSparkScheduleConfig(gamma=4),
+        model_runner=None,
+    )
+
+
+class TestBudgetDecisionLifecycle(CustomTestCase):
+    def test_take_last_decision_is_consume_once(self):
+        planner = _make_budget_planner()
+        planner.last_decision = VerifyBudgetDecision(
+            budget=3, predicted_step_seconds=0.01, predicted_theta=100.0
+        )
+        first = planner.take_last_decision()
+        self.assertEqual(first.budget, 3)
+        self.assertIsNone(planner.take_last_decision())
+
+    def test_note_non_decode_step_clears_decision(self):
+        planner = _make_budget_planner()
+        planner.last_decision = VerifyBudgetDecision(budget=1)
+        planner.note_non_decode_step()
+        self.assertIsNone(planner.take_last_decision())
+
+
+class TestScheduleVerifyLensTopk(CustomTestCase):
+    @_for_each_impl
+    def test_topk_does_not_exceed_budget(self, impl):
+        torch.manual_seed(2)
+        survival = _survival_from_confidence(torch.rand(5, 7) * 0.4 + 0.55)
+        cfg = DSparkScheduleConfig(gamma=7)
+        floor = max(cfg.min_verify_len, 1)
+        for budget in (0, 1, 5, 12, 100):
+            verify_lens = impl(survival_probs=survival, budget=budget, cfg=cfg)
+            total_extra = int((verify_lens.to(torch.int64) - floor).sum().item())
+            self.assertLessEqual(total_extra, budget)
+            self.assertGreaterEqual(int(verify_lens.min().item()), 1)
+
+    @_for_each_impl
+    def test_total_equals_anchors_plus_lens(self, impl):
+        survival = torch.tensor(
+            [[0.90, 0.80, 0.30, 0.20], [0.85, 0.70, 0.25, 0.15]],
+            dtype=torch.float32,
+        )
+        num_requests, max_len, budget = 2, 4, 2
+        cfg = DSparkScheduleConfig(gamma=max_len, min_verify_len=1)
+        verify_lens = impl(survival_probs=survival, budget=budget, cfg=cfg)
+        actual_total = num_requests + int(verify_lens.to(torch.int64).sum().item())
+        admitted = budget
+        expected_total = num_requests + num_requests * cfg.min_verify_len + admitted
+        self.assertEqual(actual_total, expected_total)
+
+    @_for_each_impl
+    def test_higher_confidence_admitted_first(self, impl):
+        survival = torch.tensor(
+            [[0.99, 0.98, 0.97, 0.96], [0.40, 0.30, 0.20, 0.10]],
+            dtype=torch.float32,
+        )
+        cfg = DSparkScheduleConfig(gamma=4)
+        verify_lens = impl(survival_probs=survival, budget=2, cfg=cfg)
+        extra = verify_lens.to(torch.int64) - cfg.min_verify_len
+        self.assertEqual(int(extra[0].item()), 2)
+        self.assertEqual(int(extra[1].item()), 0)
+
+    @_for_each_impl
+    def test_min_and_max_enter_the_budget(self, impl):
+        survival = torch.tensor([[0.99, 0.99, 0.99, 0.99, 0.99]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=5, min_verify_len=1, max_verify_len=3)
+        verify_lens = impl(survival_probs=survival, budget=100, cfg=cfg)
+        self.assertGreaterEqual(int(verify_lens.min().item()), 1)
+        self.assertLessEqual(int(verify_lens.max().item()), 3)
+
+    @_for_each_impl
+    def test_large_budget_selects_all_candidates(self, impl):
+        survival = torch.tensor([[0.9, 0.8, 0.7]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=3)
+        verify_lens = impl(survival_probs=survival, budget=1000, cfg=cfg)
+        # anchor (min_verify_len=1) + all 3 admitted drafts
+        self.assertEqual(int(verify_lens[0].item()), 4)
+
+    @_for_each_impl
+    def test_tie_break_is_value_independent(self, impl):
+        survival = torch.tensor([[0.8, 0.8, 0.8], [0.8, 0.8, 0.8]], dtype=torch.float32)
+        cfg = DSparkScheduleConfig(gamma=3)
+        floor = max(cfg.min_verify_len, 1)
+        verify_lens = impl(survival_probs=survival, budget=3, cfg=cfg)
+        total_extra = int((verify_lens.to(torch.int64) - floor).sum().item())
+        self.assertEqual(total_extra, 3)
+
+
+class TestVerifyLenAnchorContract(CustomTestCase):
+
+    @_for_each_impl
+    def test_explicit_zero_min_still_clamped_to_anchor(self, impl):
+        survival = _survival_from_confidence(
+            torch.tensor([[0.9, 0.8, 0.7], [0.6, 0.5, 0.4]], dtype=torch.float32)
+        )
+        cfg = DSparkScheduleConfig(gamma=3, min_verify_len=0)
+        verify_lens = impl(survival_probs=survival, budget=0, cfg=cfg)
+        self.assertGreaterEqual(int(verify_lens.min().item()), 1)
+        self.assertTrue(
+            torch.equal(verify_lens, torch.tensor([1, 1], dtype=torch.int32))
+        )
+
+    def test_non_flat_table_small_budget_feeds_ragged_layout(self):
+        table = SpsCostTable(
+            sample_batch_tokens=[2, 3],
+            sample_steps_per_sec=[1.0, 0.1],
+            max_batch_tokens=64,
+        )
+        cfg = DSparkScheduleConfig(gamma=3)
+        survival = _survival_from_confidence(
+            torch.tensor([[0.90, 0.80, 0.70], [0.85, 0.60, 0.40]], dtype=torch.float32)
+        )
+        budget = compute_verify_token_budget(
+            history_survival_probs=survival, sps_table=table, cfg=cfg
+        ).budget
+        self.assertEqual(budget, 0)
+        verify_lens = schedule_verify_lens_topk_from_survival(
+            survival_probs=survival, budget=budget, cfg=cfg
+        )
+        self.assertGreaterEqual(int(verify_lens.min().item()), 1)
+        verify_lens_cpu = verify_lens.to(torch.int64).tolist()
+        layout = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=verify_lens_cpu,
+            device=torch.device("cpu"),
+            grid=[sum(verify_lens_cpu)],
+        )
+        self.assertEqual(layout.verify_lens_cpu, verify_lens_cpu)
+
+
+class TestNonAnticipating(CustomTestCase):
+    @_for_each_impl
+    def test_lens_topk_non_anticipating_under_future_perturbation(self, impl):
+        base = torch.tensor(
+            [
+                [0.95, 0.90, 0.80, 0.40],
+                [0.92, 0.70, 0.30, 0.10],
+                [0.99, 0.98, 0.50, 0.05],
+            ],
+            dtype=torch.float32,
+        )
+        cfg = DSparkScheduleConfig(gamma=4)
+        budget = 5
+        baseline = impl(survival_probs=base, budget=budget, cfg=cfg)
+
+        request, cut = 1, 2
+        for delta in (-0.05, -0.2, 0.05, 0.0):
+            perturbed = base.clone()
+            future = perturbed[request, cut:]
+            perturbed[request, cut:] = torch.clamp(
+                torch.minimum(future + delta, base[request, cut - 1]), min=0.0
+            )
+            verify_lens = impl(survival_probs=perturbed, budget=budget, cfg=cfg)
+            admitted_prefix_unchanged = min(int(baseline[request].item()), cut) == min(
+                int(verify_lens[request].item()), cut
+            )
+            self.assertTrue(
+                admitted_prefix_unchanged,
+                msg=f"prefix admission changed under future perturbation delta={delta}",
+            )
+
+    @_for_each_impl
+    def test_other_requests_unaffected_by_one_request_future(self, impl):
+        base = torch.tensor(
+            [[0.95, 0.90, 0.20], [0.93, 0.88, 0.15]], dtype=torch.float32
+        )
+        cfg = DSparkScheduleConfig(gamma=3)
+        budget = 2
+        baseline = impl(survival_probs=base, budget=budget, cfg=cfg)
+        perturbed = base.clone()
+        perturbed[0, 2] = 0.01
+        verify_lens = impl(survival_probs=perturbed, budget=budget, cfg=cfg)
+        self.assertEqual(int(baseline[1].item()), int(verify_lens[1].item()))
+
+
+class TestVanillaMatchesReference(CustomTestCase):
+    def test_random_inputs_match_reference(self):
+        torch.manual_seed(20260630)
+        num_trials = 4000
+        for trial in range(num_trials):
+            num_requests = int(torch.randint(1, 6, ()).item())
+            gamma = int(torch.randint(1, 9, ()).item())
+
+            dtype = torch.float32 if trial % 2 == 0 else torch.float64
+            confidence = torch.rand(num_requests, gamma, dtype=dtype)
+            if trial % 3 == 0:
+                confidence = (confidence * 4).round() / 4
+            if trial % 7 == 0:
+                confidence = torch.ones(num_requests, gamma, dtype=dtype)
+            survival = torch.cumprod(confidence, dim=1)
+
+            min_verify_len = int(torch.randint(0, gamma + 1, ()).item())
+            if torch.rand(()).item() < 0.5:
+                max_verify_len = 0
+            else:
+                max_verify_len = int(
+                    torch.randint(min_verify_len, gamma + 1, ()).item()
+                )
+            survival_eps = float(
+                [1e-6, 1e-3, 0.1, 0.5][int(torch.randint(0, 4, ()).item())]
+            )
+            budget = int(torch.randint(0, num_requests * gamma + 3, ()).item())
+
+            cfg = DSparkScheduleConfig(
+                gamma=gamma,
+                min_verify_len=min_verify_len,
+                max_verify_len=max_verify_len,
+                survival_eps=survival_eps,
+            )
+            reference = schedule_verify_lens_topk_from_survival(
+                survival_probs=survival, budget=budget, cfg=cfg
+            )
+            vanilla = schedule_verify_lens_topk_vanilla(
+                survival_probs=survival, budget=budget, cfg=cfg
+            )
+            self.assertTrue(
+                torch.equal(reference, vanilla),
+                msg=(
+                    f"mismatch on trial {trial}: budget={budget} "
+                    f"min={min_verify_len} max={max_verify_len} eps={survival_eps} "
+                    f"survival={survival.tolist()} "
+                    f"reference={reference.tolist()} vanilla={vanilla.tolist()}"
+                ),
+            )
+
+
+class TestDSparkScheduleConfig(CustomTestCase):
+    def test_validate_rejects_min_greater_than_max(self):
+        with self.assertRaises(ValueError):
+            DSparkScheduleConfig(gamma=4, min_verify_len=3, max_verify_len=2).validate()
+
+    def test_validate_rejects_max_greater_than_gamma_plus_one(self):
+        with self.assertRaises(ValueError):
+            DSparkScheduleConfig(gamma=4, max_verify_len=6).validate()
+
+    def test_zero_max_resolves_to_gamma_plus_one(self):
+        cfg = DSparkScheduleConfig(gamma=7)
+        self.assertEqual(cfg.resolved_max_verify_len(), 8)
+
+
+class TestGraphTierFillBudget(CustomTestCase):
+    def test_floor_scales_with_min_verify_len(self):
+        """The subtracted floor is bs * max(min_verify_len, 1)."""
+        self.assertEqual(
+            graph_tier_fill_budget(
+                graph_num_tokens=60, bs=10, verify_num_draft_tokens=6, min_verify_len=2
+            ),
+            60 - 20,
+        )
+
+    def test_feeding_budget_fills_topk_total_to_tier(self):
+        """Feeding the fill budget to the top-k lifts the total to min(tier, bs*K)."""
+        bs = 6
+        cfg = DSparkScheduleConfig(gamma=6, min_verify_len=1)
+        cap = cfg.resolved_max_verify_len()
+        survival = torch.full((bs, cap), 0.99, dtype=torch.float32)
+        for graph_num_tokens in (bs, bs * cap // 2, bs * cap, bs * cap + 12):
+            budget = graph_tier_fill_budget(
+                graph_num_tokens=graph_num_tokens,
+                bs=bs,
+                verify_num_draft_tokens=cap,
+                min_verify_len=cfg.min_verify_len,
+            )
+            verify_lens = schedule_verify_lens_topk_from_survival(
+                survival_probs=survival, budget=budget, cfg=cfg
+            )
+            total = int(verify_lens.to(torch.int64).sum().item())
+            self.assertEqual(total, min(graph_num_tokens, bs * cap))
+
+
+class _FakeRaggedRunner(types.SimpleNamespace):
+    pass
+
+
+def _fake_model_runner(capture_num_tokens, max_bs):
+    runner = _FakeRaggedRunner(
+        ragged_verify_mode=True,
+        capture_num_tokens=capture_num_tokens,
+        max_bs=max_bs,
+    )
+    return types.SimpleNamespace(decode_cuda_graph_runner=runner)
+
+
+class TestBudgetTierSelection(CustomTestCase):
+    def test_floor_uses_tier_hint_capped_at_uniform_window(self):
+        from sglang.srt.speculative.dspark_components.dspark_planner import (
+            verify_layout_graph_num_tokens_floor,
+        )
+        from sglang.srt.speculative.ragged_verify import RaggedVerifyMode
+
+        model_runner = _fake_model_runner([8, 16, 1024], max_bs=128)
+        floor = verify_layout_graph_num_tokens_floor(
+            num_reqs=100,
+            ragged_verify_mode=RaggedVerifyMode.COMPACT,
+            verify_num_draft_tokens=8,
+            model_runner=model_runner,
+            tier_num_tokens=150,
+        )
+        self.assertEqual(floor, 150)
+        capped = verify_layout_graph_num_tokens_floor(
+            num_reqs=10,
+            ragged_verify_mode=RaggedVerifyMode.COMPACT,
+            verify_num_draft_tokens=8,
+            model_runner=model_runner,
+            tier_num_tokens=150,
+        )
+        self.assertEqual(capped, 80)
+        pinned = verify_layout_graph_num_tokens_floor(
+            num_reqs=100,
+            ragged_verify_mode=RaggedVerifyMode.COMPACT,
+            verify_num_draft_tokens=8,
+            model_runner=model_runner,
+        )
+        self.assertEqual(pinned, 800)
+
+    def test_exceeds_gate_checks_slots_and_tier(self):
+        from sglang.srt.speculative.dspark_components.dspark_planner import (
+            ragged_layout_exceeds_captured_grid,
+        )
+
+        model_runner = _fake_model_runner([8, 16, 1024], max_bs=128)
+        self.assertTrue(
+            ragged_layout_exceeds_captured_grid(
+                num_reqs=129,
+                verify_num_draft_tokens=8,
+                model_runner=model_runner,
+                tier_tokens_hint=200,
+            )
+        )
+        self.assertFalse(
+            ragged_layout_exceeds_captured_grid(
+                num_reqs=128,
+                verify_num_draft_tokens=8,
+                model_runner=model_runner,
+                tier_tokens_hint=512,
+            )
+        )
+        self.assertFalse(
+            ragged_layout_exceeds_captured_grid(
+                num_reqs=128,
+                verify_num_draft_tokens=8,
+                model_runner=model_runner,
+            )
+        )
+        self.assertTrue(
+            ragged_layout_exceeds_captured_grid(
+                num_reqs=128,
+                verify_num_draft_tokens=9,
+                model_runner=model_runner,
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_sps_profiler.py

@@ -0,0 +1,278 @@
+import unittest
+
+from sglang.benchmark.dspark_sps_profiler import (
+    LoadInfo,
+    ServerContext,
+    SpsRow,
+    build_request_count_sweep,
+    build_table_from_summaries,
+    count_aligned_steps,
+    postprocess_round,
+    resolve_cuda_graph_max_bs,
+    round_summary_dict,
+    validate_sweep_against_server,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+
+def make_load_info() -> LoadInfo:
+    return LoadInfo(
+        num_requests=4, max_new_tokens=1200, wall_seconds=1.0, reached_target=True
+    )
+
+
+def make_rows(
+    *,
+    num_rows: int = 30,
+    num_running_reqs: int = 4,
+    num_verify_tokens: int = 32,
+    step_time: float = 0.01,
+    first_forward_ct: int = 0,
+) -> list[SpsRow]:
+    return [
+        SpsRow(
+            forward_ct=first_forward_ct + index,
+            num_running_reqs=num_running_reqs,
+            num_verify_tokens=num_verify_tokens,
+            step_time=step_time,
+        )
+        for index in range(num_rows)
+    ]
+
+
+def make_context(**overrides) -> ServerContext:
+    values = dict(
+        base_url="http://localhost:30000",
+        tokenizer_path="dummy",
+        tp_size=4,
+        dp_size=1,
+        verify_num_draft_tokens=8,
+        simulate_acc_len=1.0,
+        cuda_graph_max_bs=128,
+        skip_max_running_requests_threshold=float("inf"),
+        skip_token_capacity_threshold=float("inf"),
+    )
+    values.update(overrides)
+    return ServerContext(**values)
+
+
+class TestPostprocessRound(CustomTestCase):
+    def test_single_rank_round_builds_probe_from_median_step_time(self):
+        outcome = postprocess_round(
+            rank_rows=[make_rows(step_time=0.01)],
+            batch_size_per_rank=4,
+            dp_size=1,
+            verify_num_draft_tokens=8,
+            min_steady_steps=16,
+            load_info=make_load_info(),
+        )
+        self.assertEqual(outcome.batch_tokens, 32)
+        self.assertAlmostEqual(outcome.steps_per_sec, 100.0)
+        self.assertEqual(outcome.match_fraction, 1.0)
+
+    def test_round_warmup_steps_are_dropped_from_timing(self):
+        slow_head = make_rows(num_rows=8, step_time=0.5, first_forward_ct=0)
+        steady_tail = make_rows(num_rows=20, step_time=0.01, first_forward_ct=8)
+        outcome = postprocess_round(
+            rank_rows=[slow_head + steady_tail],
+            batch_size_per_rank=4,
+            dp_size=1,
+            verify_num_draft_tokens=8,
+            min_steady_steps=16,
+            load_info=make_load_info(),
+        )
+        self.assertAlmostEqual(outcome.steps_per_sec, 100.0)
+
+    def test_off_target_batch_rows_are_filtered_out(self):
+        ramp = make_rows(num_rows=10, num_running_reqs=2, num_verify_tokens=16)
+        steady = make_rows(num_rows=30, first_forward_ct=10, step_time=0.02)
+        outcome = postprocess_round(
+            rank_rows=[ramp + steady],
+            batch_size_per_rank=4,
+            dp_size=1,
+            verify_num_draft_tokens=8,
+            min_steady_steps=16,
+            load_info=make_load_info(),
+        )
+        self.assertAlmostEqual(outcome.steps_per_sec, 50.0)
+        self.assertAlmostEqual(outcome.match_fraction, 1.0)
+
+    def test_mid_round_instability_raises(self):
+        head = make_rows(num_rows=15)
+        gap = make_rows(
+            num_rows=40, num_running_reqs=3, num_verify_tokens=24, first_forward_ct=15
+        )
+        tail = make_rows(num_rows=15, first_forward_ct=55)
+        with self.assertRaisesRegex(RuntimeError, "unstable mid-round"):
+            postprocess_round(
+                rank_rows=[head + gap + tail],
+                batch_size_per_rank=4,
+                dp_size=1,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+    def test_round_that_never_stabilizes_raises(self):
+        rows = make_rows(num_rows=50, num_running_reqs=3, num_verify_tokens=24)
+        rows += make_rows(num_rows=2, first_forward_ct=50)
+        with self.assertRaisesRegex(RuntimeError, "never stabilized"):
+            postprocess_round(
+                rank_rows=[rows],
+                batch_size_per_rank=4,
+                dp_size=1,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+
+class TestPostprocessRoundCrossRank(CustomTestCase):
+    def test_two_uniform_ranks_average_their_step_times(self):
+        outcome = postprocess_round(
+            rank_rows=[make_rows(step_time=0.01), make_rows(step_time=0.03)],
+            batch_size_per_rank=4,
+            dp_size=2,
+            verify_num_draft_tokens=8,
+            min_steady_steps=16,
+            load_info=make_load_info(),
+        )
+        self.assertEqual(outcome.batch_size_per_rank, 4)
+        self.assertEqual(outcome.batch_tokens, 32)
+        self.assertAlmostEqual(outcome.steps_per_sec, 50.0)
+        self.assertEqual(len(outcome.per_rank_median_step_time), 2)
+        self.assertAlmostEqual(outcome.per_rank_median_step_time[0], 0.01)
+        self.assertAlmostEqual(outcome.per_rank_median_step_time[1], 0.03)
+
+    def test_rank_with_no_new_records_raises(self):
+        with self.assertRaisesRegex(RuntimeError, "no new decode-step records"):
+            postprocess_round(
+                rank_rows=[make_rows(), []],
+                batch_size_per_rank=4,
+                dp_size=2,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+    def test_disjoint_forward_ct_ranges_raise(self):
+        with self.assertRaisesRegex(RuntimeError, "no common forward_ct"):
+            postprocess_round(
+                rank_rows=[
+                    make_rows(first_forward_ct=0),
+                    make_rows(first_forward_ct=1000),
+                ],
+                batch_size_per_rank=4,
+                dp_size=2,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+    def test_rank_below_expected_verify_tokens_raises(self):
+        # A rank reporting fewer verify tokens than bs_per_rank * K is not
+        # running the uniform static verify; above-expected counts are
+        # tolerated (the recorded count is the replayed graph tier).
+        with self.assertRaisesRegex(RuntimeError, "num_verify_tokens"):
+            postprocess_round(
+                rank_rows=[make_rows(), make_rows(num_verify_tokens=24)],
+                batch_size_per_rank=4,
+                dp_size=2,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+    def test_rank_count_mismatch_raises(self):
+        with self.assertRaisesRegex(RuntimeError, "DP ranks"):
+            postprocess_round(
+                rank_rows=[make_rows()],
+                batch_size_per_rank=4,
+                dp_size=2,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+
+
+class TestTableAssembly(CustomTestCase):
+    def test_repeats_take_the_median_per_batch_tokens(self):
+        rounds = [
+            postprocess_round(
+                rank_rows=[make_rows(step_time=step_time)],
+                batch_size_per_rank=4,
+                dp_size=1,
+                verify_num_draft_tokens=8,
+                min_steady_steps=16,
+                load_info=make_load_info(),
+            )
+            for step_time in (0.01, 0.02, 0.04)
+        ]
+        table = build_table_from_summaries(
+            summaries=[
+                round_summary_dict(outcome=outcome, repeat=repeat)
+                for repeat, outcome in enumerate(rounds)
+            ],
+            max_batch_tokens=None,
+            offdiag=False,
+        )
+        self.assertEqual(table.sample_batch_tokens, [32])
+        self.assertAlmostEqual(table.sample_steps_per_sec[0], 50.0)
+
+
+class TestSweepHelpers(CustomTestCase):
+    def test_request_count_sweep_tapers_and_hits_the_max(self):
+        sweep = build_request_count_sweep(100)
+        self.assertEqual(sweep[:4], [1, 2, 4, 8])
+        self.assertEqual(sweep[-1], 100)
+        self.assertIn(64, sweep)
+
+    def test_sweep_beyond_captured_cuda_graphs_raises(self):
+        with self.assertRaisesRegex(ValueError, "cuda graphs"):
+            validate_sweep_against_server(
+                context=make_context(cuda_graph_max_bs=64),
+                batch_sizes=[8, 128],
+            )
+
+    def test_sweep_within_captured_cuda_graphs_passes(self):
+        validate_sweep_against_server(
+            context=make_context(cuda_graph_max_bs=64, dp_size=2),
+            batch_sizes=[8, 64],
+        )
+
+    def test_resolve_cuda_graph_max_bs_prefers_captured_list(self):
+        internal_state = {
+            "cuda_graph_config": {"decode": {"bs": [1, 2, 160], "max_bs": 128}}
+        }
+        self.assertEqual(resolve_cuda_graph_max_bs(internal_state=internal_state), 160)
+
+    def test_resolve_cuda_graph_max_bs_handles_missing_config(self):
+        self.assertIsNone(resolve_cuda_graph_max_bs(internal_state={}))
+
+
+class TestCountAlignedSteps(CustomTestCase):
+    def test_off_target_steps_are_not_counted(self):
+        rows = make_rows(num_rows=10, num_running_reqs=3)
+        self.assertEqual(
+            count_aligned_steps(rank_rows=[rows], batch_size_per_rank=4), 0
+        )
+
+
+class TestMinSteadySteps(CustomTestCase):
+    def test_min_steady_steps_rejects_thin_probes(self):
+        with self.assertRaisesRegex(RuntimeError, "never stabilized"):
+            postprocess_round(
+                rank_rows=[make_rows(num_rows=20)],
+                batch_size_per_rank=4,
+                dp_size=1,
+                verify_num_draft_tokens=8,
+                min_steady_steps=32,
+                load_info=make_load_info(),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_sps_table.py

@@ -0,0 +1,200 @@
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+from sglang.srt.speculative.dspark_components.dspark_sps import (
+    SpsAdditiveCostTable,
+    SpsCostTable,
+    build_uninitialized_sps_table,
+    is_uninitialized_sps_table,
+    load_sps_table_from_path,
+    profile_sps_table,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+def _make_table() -> SpsCostTable:
+    return SpsCostTable(
+        sample_batch_tokens=[8, 16, 32, 64],
+        sample_steps_per_sec=[1000.0, 950.0, 500.0, 480.0],
+        max_batch_tokens=128,
+    )
+
+
+class TestSpsCostTableInvariants(CustomTestCase):
+    def test_rejects_non_increasing_batch_tokens(self):
+        with self.assertRaises(ValueError):
+            SpsCostTable(
+                sample_batch_tokens=[8, 8, 16],
+                sample_steps_per_sec=[1.0, 2.0, 3.0],
+                max_batch_tokens=16,
+            )
+
+    def test_rejects_unsorted_batch_tokens(self):
+        with self.assertRaises(ValueError):
+            SpsCostTable(
+                sample_batch_tokens=[16, 8],
+                sample_steps_per_sec=[1.0, 2.0],
+                max_batch_tokens=16,
+            )
+
+    def test_rejects_length_mismatch(self):
+        with self.assertRaises(ValueError):
+            SpsCostTable(
+                sample_batch_tokens=[8, 16],
+                sample_steps_per_sec=[1.0],
+                max_batch_tokens=16,
+            )
+
+    def test_rejects_empty_table(self):
+        with self.assertRaises(ValueError):
+            SpsCostTable(
+                sample_batch_tokens=[],
+                sample_steps_per_sec=[],
+                max_batch_tokens=0,
+            )
+
+    def test_rejects_max_below_largest_probe(self):
+        with self.assertRaises(ValueError):
+            SpsCostTable(
+                sample_batch_tokens=[8, 16],
+                sample_steps_per_sec=[1.0, 2.0],
+                max_batch_tokens=15,
+            )
+
+
+class TestSpsCostTableLookup(CustomTestCase):
+    def test_lookup_exact_probe_returns_that_sps(self):
+        table = _make_table()
+        self.assertEqual(table.lookup(8), 1000.0)
+        self.assertEqual(table.lookup(16), 950.0)
+        self.assertEqual(table.lookup(32), 500.0)
+        self.assertEqual(table.lookup(64), 480.0)
+
+    def test_lookup_floors_to_lower_captured_probe(self):
+        table = _make_table()
+        self.assertEqual(table.lookup(31), 950.0)
+        self.assertEqual(table.lookup(63), 500.0)
+
+    def test_lookup_below_first_probe_clamps_to_first(self):
+        table = _make_table()
+        self.assertEqual(table.lookup(1), 1000.0)
+        self.assertEqual(table.lookup(7), 1000.0)
+
+    def test_lookup_above_last_probe_clamps_to_last(self):
+        table = _make_table()
+        self.assertEqual(table.lookup(65), 480.0)
+        self.assertEqual(table.lookup(10_000), 480.0)
+
+
+class TestLoadSpsTableFromPath(CustomTestCase):
+    def test_load_from_path_round_trips_table_and_lookup(self):
+        table = _make_table()
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "sps.json"
+            path.write_text(table.to_json(), encoding="utf-8")
+            loaded = load_sps_table_from_path(str(path))
+        self.assertEqual(loaded.sample_batch_tokens, table.sample_batch_tokens)
+        self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec)
+        self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens)
+        for batch_tokens in (1, 8, 31, 64, 200):
+            self.assertEqual(loaded.lookup(batch_tokens), table.lookup(batch_tokens))
+
+
+class TestFlatTableLookupIsConstant(CustomTestCase):
+    def test_flat_table_lookup_is_one_for_any_batch(self):
+        flat = SpsCostTable(
+            sample_batch_tokens=[1],
+            sample_steps_per_sec=[1.0],
+            max_batch_tokens=4096,
+        )
+        for batch_tokens in (0, 1, 2, 17, 256, 100_000):
+            self.assertEqual(flat.lookup(batch_tokens), 1.0)
+
+
+class TestProfileSpsTable(CustomTestCase):
+    def test_profile_sorts_out_of_order_probes(self):
+        table = profile_sps_table(
+            probes=[(32, 500.0), (8, 1000.0), (16, 950.0)],
+        )
+        self.assertEqual(table.sample_batch_tokens, [8, 16, 32])
+        self.assertEqual(table.sample_steps_per_sec, [1000.0, 950.0, 500.0])
+
+    def test_profile_rejects_duplicate_batch_tokens(self):
+        with self.assertRaises(ValueError):
+            profile_sps_table(probes=[(8, 1000.0), (8, 900.0)])
+
+    def test_profile_rejects_empty_probes(self):
+        with self.assertRaises(ValueError):
+            profile_sps_table(probes=[])
+
+    def test_profile_max_batch_tokens_defaults_to_largest_probe(self):
+        table = profile_sps_table(probes=[(8, 1000.0), (64, 480.0), (16, 950.0)])
+        self.assertEqual(table.max_batch_tokens, 64)
+
+    def test_profile_honors_explicit_max_batch_tokens(self):
+        table = profile_sps_table(
+            probes=[(8, 1000.0), (16, 950.0)], max_batch_tokens=256
+        )
+        self.assertEqual(table.max_batch_tokens, 256)
+
+
+def _build_sps_cost_table_for(*, sps_table_path):
+    from sglang.srt.speculative.dspark_components.dspark_planner import (
+        build_sps_cost_table,
+    )
+
+    server_args = SimpleNamespace(
+        speculative_dspark_sps_table_path=sps_table_path,
+        max_running_requests=4,
+    )
+    return build_sps_cost_table(server_args=server_args, verify_num_draft_tokens=5)
+
+
+class TestBuildSpsCostTableContract(CustomTestCase):
+    def test_unset_table_path_returns_flat_table(self):
+        for sps_table_path in (None, ""):
+            table = _build_sps_cost_table_for(sps_table_path=sps_table_path)
+            self.assertEqual(table.sample_batch_tokens, [1])
+            self.assertEqual(table.sample_steps_per_sec, [1.0])
+            self.assertEqual(table.max_batch_tokens, 20)
+
+    def test_real_path_loads_table(self):
+        table = _make_table()
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "sps.json"
+            path.write_text(table.to_json(), encoding="utf-8")
+            loaded = _build_sps_cost_table_for(sps_table_path=str(path))
+        self.assertEqual(loaded.sample_batch_tokens, table.sample_batch_tokens)
+        self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec)
+        self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens)
+
+
+class TestIsUninitializedSpsTable(CustomTestCase):
+    def test_additive_table_is_never_uninitialized(self):
+        table = SpsAdditiveCostTable(
+            bias_seconds=0.1,
+            bs_probes=[128, 192, 256],
+            alpha_seconds=[0.0, 0.008, 0.016],
+            m_probes=[384, 512, 1024],
+            theta_seconds=[0.0, 0.02, 0.1],
+        )
+        self.assertFalse(is_uninitialized_sps_table(table))
+
+    def test_placeholder_diagonal_table_is_uninitialized(self):
+        self.assertTrue(
+            is_uninitialized_sps_table(
+                build_uninitialized_sps_table(max_batch_tokens=128)
+            )
+        )
+
+    def test_real_diagonal_table_is_initialized(self):
+        self.assertFalse(is_uninitialized_sps_table(_make_table()))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_sts.py

@@ -0,0 +1,142 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+
+from sglang.benchmark.dspark_sts_fit import (
+    default_temperature_grid,
+    expected_calibration_error,
+    fit_sts_temperatures,
+)
+from sglang.srt.models.dspark import DSparkConfidenceHead
+from sglang.srt.speculative.dspark_components.dspark_sts import (
+    DSparkStsCalibration,
+    StsDataRecorder,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestApplySts(CustomTestCase):
+    def test_default_buffer_is_identity_sigmoid(self):
+        head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False)
+        confidence_raw = torch.randn(3, 5) * 9.0
+        out = head.apply_sts(confidence_raw)
+        self.assertTrue(torch.equal(out, torch.sigmoid(confidence_raw.float())))
+
+    def test_per_position_temperature_scales_each_column(self):
+        head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False)
+        head.sts_temperatures = torch.tensor([0.5, 1.0, 2.0])
+        confidence_raw = torch.full((2, 3), 2.0)
+        out = head.apply_sts(confidence_raw)
+        # Hand-computed sigmoid(2.0 / T) per column: identical raw logits with
+        # distinct per-column T catch wrong-axis broadcasts, and dividing (not
+        # multiplying) by T is what separates 0.982 from 0.731 in column 0.
+        expected_row = [0.98201379, 0.88079708, 0.73105858]
+        for row in out.tolist():
+            for got, want in zip(row, expected_row):
+                self.assertAlmostEqual(got, want, places=6)
+
+    def test_apply_sts_stashes_raw_logit(self):
+        head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False)
+        confidence_raw = torch.randn(2, 5)
+        head.apply_sts(confidence_raw)
+        self.assertIs(head._last_confidence_raw, confidence_raw)
+
+
+class TestDSparkStsCalibration(CustomTestCase):
+    def test_json_round_trip_preserves_fields(self):
+        calibration = DSparkStsCalibration(
+            temperatures=[1.5, 2.0, 0.5],
+            dataset="shards.*.pt",
+            num_samples=1234,
+            ece_before=[0.3, 0.2, 0.1],
+            ece_after=[0.02, 0.01, 0.03],
+        )
+        restored = DSparkStsCalibration.from_json(calibration.to_json())
+        self.assertEqual(restored.temperatures, calibration.temperatures)
+        self.assertEqual(restored.dataset, calibration.dataset)
+        self.assertEqual(restored.num_samples, calibration.num_samples)
+        self.assertEqual(restored.ece_before, calibration.ece_before)
+        self.assertEqual(restored.ece_after, calibration.ece_after)
+
+    def test_rejects_empty_temperatures(self):
+        with self.assertRaises(ValueError):
+            DSparkStsCalibration(temperatures=[])
+
+    def test_rejects_non_positive_temperature(self):
+        with self.assertRaises(ValueError):
+            DSparkStsCalibration(temperatures=[1.0, 0.0, 2.0])
+        with self.assertRaises(ValueError):
+            DSparkStsCalibration(temperatures=[1.0, -0.5])
+
+
+class TestExpectedCalibrationError(CustomTestCase):
+    def test_perfectly_calibrated_probs_have_low_ece(self):
+        torch.manual_seed(0)
+        probs = torch.full((20000,), 0.3, dtype=torch.float64)
+        targets = (torch.rand(20000) < 0.3).to(torch.float64)
+        ece = expected_calibration_error(probs=probs, targets=targets, num_bins=15)
+        self.assertLess(ece, 0.02)
+
+    def test_overconfident_probs_have_high_ece(self):
+        probs = torch.full((20000,), 0.95, dtype=torch.float64)
+        targets = torch.full((20000,), 0.3, dtype=torch.float64)
+        ece = expected_calibration_error(probs=probs, targets=targets, num_bins=15)
+        self.assertGreater(ece, 0.5)
+
+
+class TestFitStsTemperatures(CustomTestCase):
+    def test_recovers_scale_and_reduces_ece(self):
+        torch.manual_seed(0)
+        num_samples, gamma, scale = 60000, 4, 2.5
+        base_logit = torch.tensor([2.0, 1.2, 0.8, 0.4])
+        true_logit = base_logit[None, :] + torch.randn(num_samples, gamma) * 0.5
+        true_prob = torch.sigmoid(true_logit)
+        accept = (torch.rand(num_samples, gamma) < true_prob).to(torch.float64)
+        prefix_mask = torch.cumprod(accept, dim=1)
+        overconfident_logits = true_logit * scale
+
+        result = fit_sts_temperatures(
+            logits=overconfident_logits,
+            prefix_mask=prefix_mask,
+            grid=default_temperature_grid(),
+            num_bins=15,
+        )
+
+        self.assertEqual(len(result["temperatures"]), gamma)
+        for temperature in result["temperatures"]:
+            self.assertGreater(temperature, scale / 1.5)
+            self.assertLess(temperature, scale * 1.5)
+        mean_before = sum(result["ece_before"]) / gamma
+        mean_after = sum(result["ece_after"]) / gamma
+        self.assertLess(mean_after, 0.25 * mean_before)
+
+
+class TestStsDataRecorder(CustomTestCase):
+    def test_builds_prefix_mask_and_writes_shard(self):
+        gamma = 4
+        confidence_raw = torch.randn(4, gamma)
+        num_correct_drafts = torch.tensor([0, 2, 4, 1], dtype=torch.int32)
+        expected_prefix_mask = torch.tensor(
+            [[0, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 1], [1, 0, 0, 0]],
+            dtype=torch.float32,
+        )
+        with tempfile.TemporaryDirectory() as tmp:
+            stem = str(Path(tmp) / "shard")
+            recorder = StsDataRecorder(path_stem=stem, gamma=gamma, flush_every=10)
+            recorder.record(
+                confidence_raw=confidence_raw,
+                num_correct_drafts=num_correct_drafts,
+            )
+            recorder.flush()
+            shard = torch.load(f"{stem}.0.pt")
+        self.assertTrue(torch.equal(shard["prefix_mask"], expected_prefix_mask))
+        self.assertTrue(torch.equal(shard["logits"], confidence_raw.to(torch.float32)))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_ragged_verify.py

@@ -0,0 +1,132 @@
+import unittest
+
+import torch
+
+from sglang.srt.speculative.ragged_verify import (
+    RaggedVerifyLayout,
+    build_ragged_target_verify_geometry,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+_DEVICE = torch.device("cpu")
+_GRID = [8, 16, 24, 32, 64]
+
+# The backend capability checks (supports_ragged_verify_graph) live in
+# test_ragged_verify_backend_capability.py: importing the backend modules
+# pulls GPU-only wheels, which fail to import on the CPU runners.
+
+
+class TestRaggedTargetVerifyGeometry(CustomTestCase):
+    def test_mixed_verify_lens_geometry(self):
+        layout = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 1, 3], device=_DEVICE, grid=_GRID
+        )
+        seq_lens = torch.tensor([10, 20, 30], dtype=torch.int32)
+        geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=layout)
+        self.assertEqual(geometry.cache_seqlens_int32.tolist(), [18, 21, 33])
+        self.assertEqual(geometry.cu_seqlens_q.tolist(), [0, 8, 9, 12])
+        self.assertEqual(geometry.cu_seqlens_k.tolist(), [0, 18, 39, 72])
+        self.assertEqual(geometry.max_seq_len_q, 8)
+
+    def test_geometry_dtypes_are_int32(self):
+        layout = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 1, 3], device=_DEVICE, grid=_GRID
+        )
+        seq_lens = torch.tensor([10, 20, 30], dtype=torch.int64)
+        geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=layout)
+        self.assertEqual(geometry.cache_seqlens_int32.dtype, torch.int32)
+        self.assertEqual(geometry.cu_seqlens_q.dtype, torch.int32)
+        self.assertEqual(geometry.cu_seqlens_k.dtype, torch.int32)
+
+
+class TestPaddedRaggedVerifyGeometry(CustomTestCase):
+    def test_padded_layout_grows_bs_and_fills_bucket(self):
+        raw = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 1, 3],
+            device=_DEVICE,
+            grid=[8, 16, 32, 64],
+            graph_num_tokens_floor=24,
+        )
+        self.assertEqual(raw.graph_num_tokens, 32)
+        padded = raw.padded_to_bucket(padded_bs=4)
+        self.assertEqual(padded.bs, 4)
+        self.assertEqual(padded.verify_lens.tolist(), [8, 1, 3, 20])
+        self.assertEqual(padded.qo_indptr_device.tolist(), [0, 8, 9, 12, 32])
+        seq_lens = torch.tensor([10, 20, 30, 1], dtype=torch.int32)
+        geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=padded)
+        self.assertEqual(geometry.cu_seqlens_q.tolist(), [0, 8, 9, 12, 32])
+        self.assertEqual(geometry.cache_seqlens_int32.tolist(), [18, 21, 33, 21])
+        self.assertEqual(int(geometry.cu_seqlens_k[-1]), 18 + 21 + 33 + 21)
+
+    def test_padded_layout_decoupled_slots_spread_slack(self):
+        raw = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 1, 3],
+            device=_DEVICE,
+            grid=[8, 16, 32, 64],
+            graph_num_tokens_floor=24,
+        )
+        padded = raw.padded_to_bucket(padded_bs=6)
+        self.assertEqual(padded.bs, 6)
+        self.assertEqual(padded.verify_lens.tolist(), [8, 1, 3, 7, 7, 6])
+        self.assertEqual(int(padded.qo_indptr_device[-1]), 32)
+
+    def test_padded_layout_budget_tier_below_uniform(self):
+        raw = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 1, 3],
+            device=_DEVICE,
+            grid=[8, 16, 32, 64],
+        )
+        self.assertEqual(raw.graph_num_tokens, 16)
+        padded = raw.padded_to_bucket(padded_bs=3)
+        self.assertEqual(padded.verify_lens.tolist(), [8, 1, 7])
+        self.assertEqual(int(padded.qo_indptr_device[-1]), 16)
+
+    def test_padded_layout_zero_len_pad_rows(self):
+        raw = RaggedVerifyLayout.from_verify_lens(
+            verify_lens_cpu=[8, 8],
+            device=_DEVICE,
+            grid=[8, 16, 32, 64],
+        )
+        self.assertEqual(raw.graph_num_tokens, 16)
+        padded = raw.padded_to_bucket(padded_bs=8)
+        self.assertEqual(padded.verify_lens.tolist(), [8, 8, 0, 0, 0, 0, 0, 0])
+        self.assertEqual(int(padded.qo_indptr_device[-1]), 16)
+
+
+class TestCaptureVerifyLens(CustomTestCase):
+    def test_small_tier_one_token_rows(self):
+        from sglang.srt.speculative.ragged_verify import build_capture_verify_lens
+
+        lens = build_capture_verify_lens(num_tokens=8, num_slots=8, num_draft_tokens=8)
+        self.assertEqual(lens, [1] * 8)
+
+    def test_large_tier_spreads_within_window(self):
+        from sglang.srt.speculative.ragged_verify import build_capture_verify_lens
+
+        lens = build_capture_verify_lens(
+            num_tokens=1024, num_slots=128, num_draft_tokens=8
+        )
+        self.assertEqual(sum(lens), 1024)
+        self.assertEqual(lens, [8] * 128)
+
+    def test_uneven_tier_rows_stay_legal(self):
+        from sglang.srt.speculative.ragged_verify import build_capture_verify_lens
+
+        lens = build_capture_verify_lens(num_tokens=24, num_slots=5, num_draft_tokens=8)
+        self.assertEqual(sum(lens), 24)
+        self.assertTrue(all(1 <= v <= 8 for v in lens))
+
+    def test_rejects_overpacked_tier(self):
+        from sglang.srt.speculative.ragged_verify import build_capture_verify_lens
+
+        with self.assertRaises(ValueError):
+            build_capture_verify_lens(num_tokens=64, num_slots=4, num_draft_tokens=8)
+        with self.assertRaises(ValueError):
+            build_capture_verify_lens(num_tokens=4, num_slots=8, num_draft_tokens=8)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_ragged_verify_backend_capability.py

@@ -0,0 +1,43 @@
+"""Backend opt-in flags for the ragged-verify graphs.
+
+Runs in the GPU suite because importing the backend modules pulls GPU-only
+wheels (sgl_kernel) at module scope, which fail to import on CPU runners.
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+
+
+class TestRaggedVerifyGraphCapability(CustomTestCase):
+    def test_base_backend_defaults_false(self):
+        from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+
+        self.assertFalse(AttentionBackend.supports_ragged_verify_graph)
+
+    def test_ragged_implementing_backends_declare_the_flag(self):
+        """Every backend with a ragged-verify metadata path must opt in; a
+        dropped flag silently disables ragged graphs for that backend (the
+        runner falls back to eager with no other test going red)."""
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            DeepseekV4AttnBackend,
+        )
+        from sglang.srt.layers.attention.flashattention_backend import (
+            FlashAttentionBackend,
+        )
+        from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
+
+        for backend in (
+            TRTLLMHAAttnBackend,
+            DeepseekV4AttnBackend,
+            FlashAttentionBackend,
+        ):
+            with self.subTest(backend=backend.__name__):
+                self.assertTrue(backend.supports_ragged_verify_graph)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/managers/test_batch_result_processor_spec_grammar.py

@@ -93,6 +93,8 @@ def _make_result(num_draft_tokens, accept_lens, flat_tokens):
         speculative_num_draft_tokens=num_draft_tokens,
         num_correct_drafts=None,
         num_correct_drafts_per_req_cpu=None,
+        block_accept_lens=None,
+        cap_lens=None,
     )
 
 

test/registered/unit/model_executor/test_pool_configurator.py

@@ -121,6 +121,7 @@ def _make_model_runner(
     spec.is_eagle.return_value = False
     spec.is_standalone.return_value = False
     spec.is_dflash.return_value = False
+    spec.is_dflash_family.return_value = False
     spec.is_none.return_value = True
     mr.spec_algorithm = spec
 

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -30,6 +30,7 @@ def _make_info(batch_size=2, **overrides):
         top_ks=torch.full((batch_size,), TOP_K_ALL, dtype=torch.int32),
         min_ps=torch.zeros(batch_size),
         is_all_greedy=False,
+        is_any_greedy=False,
         need_top_p_sampling=False,
         need_top_k_sampling=False,
         need_min_p_sampling=False,