[Spec Decode] DSpark capacity reallocation with varlen full-CG verification

Implements DSpark (arXiv 2607.05147) confidence-scheduled verification with
two capacity enforcement modes and full CUDA graph support:

- Capacity manager with `mask` (pad pruned verify rows; defaults
  VLLM_MOE_SKIP_PADDING=1 so MoE kernels skip the pruned rows) and
  `varlen` (compact the verifier batch) modes; varlen replays FULL CUDA graphs for
  both the target verify step and the DSpark draft query step.
- Paper-faithful Algorithm 1 allocator: capacities are the greedy admission
  counts (sum(capacities) == spent budget, hard cap; zero-survival tokens
  are never candidates), fixing a threshold-recount tie escape that
  disabled the budget under saturated confidence logits.
- Hardware-aware prefix scheduler on the HOST, as in the paper: the
  draft step D2H-copies its confidence logits asynchronously, and at
  dispatch the host consumes the copy staged one step earlier (fixed
  offset — decided by step counting so every TP rank makes the identical
  choice, and the copy has had a full step to land, so allocation never
  blocks on the GPU). Algorithm 1 (greedy global admission by
  prefix-survival, theta = tau * SPS(R, B) argmax stopping rule,
  dspark_budget_frac cap) runs in numpy; capacities scatter into the
  host array that already shapes batch trimming/dispatch. Per-request
  verify budgets lag drafts by one step; new requests default to full
  capacity. An admit-everything ablation measures the enforcement tax at
  +0.1-0.3 ms/step vs no-capacity (the previous in-graph GPU allocator +
  capacity D2H flush cost ~2.2 ms/step of lost async-scheduling overlap).
  The step-rate table is 2D over (running requests R, verify batch
  tokens B): draft and host costs scale with R and are paid regardless
  of pruning, so a 1D curve over B alone overestimates pruning payoff
  ~2x at high concurrency and over-prunes. `dspark_sps_curve="auto"`
  measures the surface at engine init: REAL decode steps are timed over
  an (R pow2, k in {0,2,4,7}) grid with capacities forced through the
  mode's real enforcement path (median over chunks of pipelined steps;
  plain SamplingParams; rank-0 broadcast; online-STS reset afterwards).
  Explicit 1D breakpoint lists are still accepted (replicated across R);
  `dspark_sps_overhead_ms` adds scheduler/IPC time above the worker.
  benchmarks/profile_dspark_sps_curve.py remains for offline measurement.
- Online Sequential Temperature Scaling (`dspark_online_sts`, on by
  default with capacity modes): per-position temperatures fitted online by
  an ECE grid search over binned rejection-sampler outcomes
  (order-preserving, per the paper Sec 3.2.1; identity until observations
  accumulate). Runs on the host: outcomes ride the same staged confidence
  copies, joined to the previous copy's logits by request id, so TP-rank
  agreement is plain deterministic numpy.
- Varlen full-CG correctness fixes: per-request token bound in cudagraph
  dispatch, capture/replay buffer-address consistency in the DSA indexer
  varlen decode path (forced flatten + persistent indices buffer),
  padded-row sizing in the indexer build, TP-deterministic capacity
  flushes, and correct handling of the scheduler's -1 draft placeholder
  ids in capacity accounting.
- TP-rank determinism fixes for padded draft FULL-graph replays: padding
  rows of sample_idx_mapping now carry a -1 inert-row sentinel so replays
  never scatter draft logits through stale req-state slot ids
  (duplicate-index scatters have undefined write order and silently
  diverge per-rank state; with varlen capacity this became a
  collective-size-mismatch deadlock), and the online-STS proposal staging
  moved out of the captured graph. Worker slot recycling now iterates
  finished_req_ids in sorted order and fallback per-request sampling
  seeds come from a dedicated RNG stream, removing two more per-rank
  divergence hazards. Debug guard: VLLM_DSPARK_TP_CHECK={1,2}
  (capacity.py::check_dspark_tp_consistency) cross-checks request-keyed
  capacity/STS state across TP ranks each step and fails fast with
  per-rank state dumps.

Co-authored-by: OpenAI Codex <codex@openai.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: 3 people

benchmarks/profile_dspark_sps_curve.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Profile the engine step-rate curve for the DSpark prefix scheduler.
+
+Times the captured FULL cudagraph replays of the target verification step at
+every captured batch token count and emits ``dspark_sps_curve`` breakpoints,
+one per capture size. The scheduler linearly interpolates between
+breakpoints, which amortizes cudagraph padding smoothly instead of
+concentrating it into thresholds at capture-size boundaries.
+
+Example:
+    python benchmarks/profile_dspark_sps_curve.py <target-model> \\
+        --speculative-config '{"method": "dspark", "model": "...", ...}' \\
+        --engine-args '{"tensor_parallel_size": 4, "max_num_seqs": 32}'
+
+Paste the printed ``dspark_sps_curve`` entry into --speculative-config.
+
+Caveats: replays run on whatever (dummy) buffer state capture left behind, so
+data-dependent kernels (e.g. MoE routing) may be timed on unrepresentative
+inputs, and per-step CPU/draft overhead is modeled only through the constant
+``--overhead-ms``. Only the curve's shape matters to the scheduler.
+"""
+
+import argparse
+import json
+
+
+def _time_fullgraph_replays(worker, iters: int, warmup: int) -> dict[int, float]:
+    """Worker-side: time FULL graph replay per batch token count (ms/step).
+
+    Runs on every TP rank via collective_rpc so the collectives captured in
+    the graphs stay matched; every rank replays the same descs in the same
+    sorted order. Before timing each descriptor the input buffers are
+    refreshed into the same coherent dummy state capture used, so replays
+    never read stale metadata.
+    """
+    import torch
+
+    from vllm.v1.worker.gpu.cudagraph_utils import prepare_inputs_to_capture
+
+    runner = worker.model_runner
+    mgr = runner.cudagraph_manager
+    assert mgr is not None and mgr.graphs, (
+        "No FULL cudagraphs captured; run with a cudagraph_mode that captures "
+        "FULL decode graphs."
+    )
+    # Prefer varlen spec-decode descs; fall back to all captured graphs.
+    descs = [d for d in mgr.graphs if d.max_req_tokens is not None]
+    if not descs:
+        descs = list(mgr.graphs.keys())
+    # One desc per token count: the largest request count is the most
+    # representative shape under load.
+    by_tokens: dict[int, object] = {}
+    for d in descs:
+        cur = by_tokens.get(d.num_tokens)
+        if cur is None or (d.num_reqs or 0) > (cur.num_reqs or 0):
+            by_tokens[d.num_tokens] = d
+
+    results: dict[int, float] = {}
+    for num_tokens in sorted(by_tokens):
+        desc = by_tokens[num_tokens]
+        num_reqs = desc.num_reqs or min(num_tokens, mgr.max_num_reqs)
+        prepare_inputs_to_capture(
+            num_reqs,
+            num_tokens,
+            runner.model_state,
+            runner.input_buffers,
+            runner.block_tables,
+            runner.attn_groups,
+            runner.kv_cache_config,
+            max_req_tokens=desc.max_req_tokens,
+        )
+        graph = mgr.graphs[desc]
+        for _ in range(warmup):
+            graph.replay()
+        torch.accelerator.synchronize()
+        start = torch.Event(enable_timing=True)
+        end = torch.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            graph.replay()
+        end.record()
+        torch.accelerator.synchronize()
+        results[num_tokens] = start.elapsed_time(end) / iters
+    return results
+
+
+def curve_breakpoints(
+    ms_per_step: dict[int, float], overhead_ms: float
+) -> list[list[float]]:
+    """Convert per-capture-size step times into ``dspark_sps_curve``
+    breakpoints, one per capture size. The scheduler's table linearly
+    interpolates between them (and clamps at the ends)."""
+    return [
+        [size, round(1000.0 / (ms_per_step[size] + overhead_ms), 3)]
+        for size in sorted(ms_per_step)
+    ]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model", help="Target model (path or HF id)")
+    parser.add_argument(
+        "--speculative-config",
+        required=True,
+        help="JSON speculative config (same value you pass to vllm serve)",
+    )
+    parser.add_argument(
+        "--engine-args",
+        default="{}",
+        help="JSON dict of extra vllm.LLM kwargs "
+        '(e.g. \'{"tensor_parallel_size": 4, "max_num_seqs": 32}\')',
+    )
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument(
+        "--overhead-ms",
+        type=float,
+        default=0.0,
+        help="Constant per-step overhead (draft forward, sampling, CPU gap) "
+        "added to every measured step time before converting to a rate.",
+    )
+    parser.add_argument("--output", help="Write the curve JSON to this file")
+    args = parser.parse_args()
+
+    # The timing callable is shipped to the workers via collective_rpc, which
+    # requires the pickle fallback. Local profiling tool, trusted input.
+    import os
+
+    os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    from vllm import LLM
+
+    llm = LLM(
+        model=args.model,
+        speculative_config=json.loads(args.speculative_config),
+        **json.loads(args.engine_args),
+    )
+    per_rank = llm.collective_rpc(
+        _time_fullgraph_replays, kwargs={"iters": args.iters, "warmup": args.warmup}
+    )
+    ms_per_step = per_rank[0]
+
+    print("\nMeasured FULL-graph step times (rank 0):")
+    for size in sorted(ms_per_step):
+        print(f"  B={size:5d} tokens: {ms_per_step[size]:8.3f} ms/step")
+
+    curve = curve_breakpoints(ms_per_step, args.overhead_ms)
+    entry = {"dspark_sps_curve": curve}
+    print("\nAdd to --speculative-config:")
+    print(json.dumps(entry))
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(entry, f, indent=2)
+        print(f"\nWritten to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

tests/engine/test_arg_utils.py

@@ -225,6 +225,43 @@ def test_jit_monitor_mode_arg(mode):
     assert engine_args.create_observability_config().jit_monitor_mode == mode
 
 
+def test_dspark_capacity_verification_mode_arg():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = parser.parse_args(
+        [
+            "--spec-method",
+            "ngram",
+            "--spec-tokens",
+            "1",
+            "--dspark-capacity-verification-mode",
+            "mask",
+        ]
+    )
+
+    engine_args = EngineArgs.from_cli_args(args)
+    assert engine_args.dspark_capacity_verification_mode == "mask"
+    speculative_config = engine_args.create_speculative_config(None, None)
+    assert speculative_config is not None
+    assert speculative_config.dspark_capacity_verification_mode == "mask"
+
+
+def test_dspark_capacity_verification_mode_conflicts_with_speculative_config():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = parser.parse_args(
+        [
+            "--speculative-config",
+            '{"method":"ngram","num_speculative_tokens":1,'
+            '"dspark_capacity_verification_mode":"mask"}',
+            "--dspark-capacity-verification-mode",
+            "varlen",
+        ]
+    )
+
+    engine_args = EngineArgs.from_cli_args(args)
+    with pytest.raises(ValueError, match="dspark_capacity_verification_mode"):
+        engine_args.create_speculative_config(None, None)
+
+
 def test_hf_token_get_kwargs():
     kwargs = get_kwargs(ModelConfig)["hf_token"]
 

tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml

@@ -0,0 +1,22 @@
+model_name: "deepseek-ai/DeepSeek-V4-Flash-DSpark"
+accuracy_threshold: 0.92
+num_questions: 1319
+num_fewshot: 5
+startup_max_wait_seconds: 1800
+server_args: >-
+  --tokenizer-mode deepseek_v4
+  --trust-remote-code
+  --dtype bfloat16
+  --max-model-len 8192
+  --tensor-parallel-size 4
+  --enable-expert-parallel
+  --block-size 256
+  --gpu-memory-utilization 0.5
+  --kv-cache-dtype fp8
+  --max-num-batched-tokens 16384
+  --max-num-seqs 32
+  --speculative-config '{"method":"dspark",
+  "model":"deepseek-ai/DeepSeek-V4-Flash-DSpark",
+  "attention_backend":"FLASH_ATTN","num_speculative_tokens":7,
+  "draft_sample_method":"probabilistic","dspark_confidence_threshold":0.0,
+  "dspark_budget_frac":0.5,"dspark_capacity_verification_mode":"varlen"}'

tests/evals/gsm8k/configs/dspark_list.txt

@@ -0,0 +1 @@
+DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml

tests/test_config.py

@@ -1572,6 +1572,55 @@ def test_draft_sample_method_gumbel_is_rejected():
         )
 
 
+def test_dspark_capacity_config_validation():
+    speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=1,
+        dspark_confidence_threshold=0.25,
+        dspark_budget_frac=0.5,
+        dspark_capacity_verification_mode="mask",
+    )
+    assert speculative_config.dspark_confidence_threshold == 0.25
+    assert speculative_config.dspark_budget_frac == 0.5
+    assert speculative_config.dspark_capacity_verification_mode == "mask"
+    assert (
+        SpeculativeConfig(
+            method="ngram", num_speculative_tokens=1
+        ).dspark_capacity_verification_mode
+        == "varlen"
+    )
+    assert (
+        SpeculativeConfig(
+            method="ngram",
+            num_speculative_tokens=1,
+            dspark_capacity_verification_mode="compact",
+        ).dspark_capacity_verification_mode
+        == "varlen"
+    )
+    assert (
+        SpeculativeConfig(
+            method="ngram", num_speculative_tokens=1
+        ).dspark_confidence_threshold
+        == 0.0
+    )
+
+    for threshold in (-0.1, 1.1):
+        with pytest.raises(ValueError, match="dspark_confidence_threshold"):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                dspark_confidence_threshold=threshold,
+            )
+
+    for budget_frac in (0.0, -0.1, 1.1):
+        with pytest.raises(ValueError, match="dspark_budget_frac"):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                dspark_budget_frac=budget_frac,
+            )
+
+
 def test_ir_op_priority_default():
     """Test that IR op priority defaults are set correctly."""
     from vllm.config.kernel import IrOpPriorityConfig

tests/v1/worker/test_gpu_block_table.py

@@ -5,6 +5,7 @@
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.block_table import BlockTables
 
 pytestmark = pytest.mark.skipif(
@@ -130,3 +131,49 @@ def test_block_tables_apply_staged_writes_single_group():
         block_tables.block_tables[0].gpu[0, :2],
         torch.tensor([1, 2], dtype=torch.int32, device=device),
     )
+
+
+def test_compute_slot_mappings_applies_padding_mask():
+    device = torch.device("cuda")
+    block_tables = BlockTables(
+        block_sizes=[16],
+        max_num_reqs=2,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[16],
+    )
+
+    block_tables.append_block_ids(
+        req_index=0,
+        new_block_ids=([2],),
+        overwrite=True,
+    )
+    block_tables.append_block_ids(
+        req_index=1,
+        new_block_ids=([3],),
+        overwrite=True,
+    )
+    block_tables.apply_staged_writes()
+
+    idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
+    positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device)
+    is_padding = torch.tensor(
+        [False, True, False, False, True, False, False, False],
+        dtype=torch.bool,
+        device=device,
+    )
+
+    slot_mappings = block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        positions,
+        num_tokens_padded=8,
+        is_padding=is_padding,
+    )
+    torch.accelerator.synchronize()
+
+    assert slot_mappings.cpu().tolist() == [
+        [32, PAD_SLOT_ID, 34, 48, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID]
+    ]