spec_decode: consolidate DSpark correctness and load-aware serving

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,
+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()
+    if args.iters <= 0:
+        parser.error("--iters must be greater than zero")
+    if args.warmup < 0:
+        parser.error("--warmup must be non-negative")
+    if args.overhead_ms < 0:
+        parser.error("--overhead-ms must be non-negative")
+
+    # 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/test_config.py

@@ -1630,6 +1630,75 @@ 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, float("nan"), float("inf")):
+        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, float("nan"), float("inf")):
+        with pytest.raises(ValueError, match="dspark_budget_frac"):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                dspark_budget_frac=budget_frac,
+            )
+
+    for config_field, value in (
+        ("dspark_confidence_temperature", float("nan")),
+        ("dspark_confidence_temperature", float("inf")),
+        ("dspark_sps_overhead_ms", float("nan")),
+        ("dspark_sps_overhead_ms", float("inf")),
+    ):
+        with pytest.raises(ValueError, match=config_field):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                **{config_field: value},
+            )
+
+    with pytest.raises(ValueError, match="dspark_sps_curve"):
+        SpeculativeConfig(
+            method="ngram",
+            num_speculative_tokens=1,
+            dspark_sps_curve=[(1, float("nan"))],
+        )
+
+
 def test_ir_op_priority_default():
     """Test that IR op priority defaults are set correctly."""
     from vllm.config.kernel import IrOpPriorityConfig

tests/v1/attention/test_deepseek_v4_dspark_metadata.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm.v1.attention.backends.mla.sparse_swa import (
+    DeepseekSparseSWAMetadataBuilder,
+)
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+
+def test_dspark_swa_decode_threshold_matches_target_verification() -> None:
+    """DSpark verifies 1 + K target tokens, not the generic 1 + 2K."""
+    speculative_config = SimpleNamespace(
+        num_speculative_tokens=5,
+        parallel_drafting=True,
+        use_dspark=lambda: True,
+    )
+    hf_config = SimpleNamespace(sliding_window=128, compress_ratios=[1, 4, 128])
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=4096, hf_config=hf_config),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=16),
+        speculative_config=speculative_config,
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+    )
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.bfloat16,
+    )
+
+    builder = DeepseekSparseSWAMetadataBuilder(
+        kv_cache_spec,
+        ["placeholder"],
+        vllm_config,
+        torch.device("cpu"),
+    )
+
+    assert builder.decode_threshold == 6

tests/v1/spec_decode/test_acceptance_length_controller.py

@@ -268,6 +268,9 @@ def test_synthetic_scheduler_output_uses_default_speculative_depth():
     output.num_spec_tokens_to_schedule = 2
     assert output.resolve_num_spec_tokens_to_schedule(default=5) == 2
 
+    output.num_spec_tokens_to_schedule = 0
+    assert output.resolve_num_spec_tokens_to_schedule(default=5) == 0
+
 
 def test_runner_v2_autoregressive_drafter_stops_at_adaptive_depth(monkeypatch):
     monkeypatch.setattr(AutoRegressiveSpeculator, "__abstractmethods__", frozenset())