[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) 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: optional `dspark_sps_curve` config
  (profiled steps-per-second vs verification batch tokens) drives the
  theta = tau * SPS(B) argmax stopping rule; `dspark_budget_frac` remains
  as an admission upper bound. Includes
  benchmarks/profile_dspark_sps_curve.py to measure the curve from the
  captured cudagraph staircase.
- 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.

Co-authored-by: OpenAI Codex <codex@openai.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DHmbeNn4JRhg5zaZqPiGHm
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: 3 people

benchmarks/profile_dspark_sps_curve.py

@@ -0,0 +1,165 @@
+# 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
+encoding the resulting staircase (step time is constant within a cudagraph
+padding bucket and jumps 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 staircase_breakpoints(
+    ms_per_step: dict[int, float], overhead_ms: float
+) -> list[list[float]]:
+    """Convert per-capture-size step times into ``dspark_sps_curve``
+    breakpoints. Batches pad UP to the next capture size, so every B in
+    ``(prev_size, size]`` runs at ``size``'s rate; encode each bucket as a
+    flat segment with one-token jumps at the boundaries."""
+    curve: list[list[float]] = []
+    prev_size = 0
+    for size in sorted(ms_per_step):
+        rate = round(1000.0 / (ms_per_step[size] + overhead_ms), 3)
+        seg_start = prev_size + 1
+        if seg_start < size:
+            curve.append([seg_start, rate])
+        curve.append([size, rate])
+        prev_size = size
+    return curve
+
+
+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 = staircase_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/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/spec_decode/test_rejection_sampler_utils.py

@@ -150,7 +150,7 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl
 
     torch.manual_seed(42)
     device = "cuda"
-    num_trials = 10 * VOCAB_SIZE
+    num_trials = 64
 
     target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
     draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
@@ -288,7 +288,7 @@ def test_placeholder_draft_token_rejected():
     """
     torch.manual_seed(0)
     device = "cuda"
-    num_trials = 64
+    num_trials = 10 * VOCAB_SIZE
     K = 1
     temperature = 0.6
 

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]
+    ]