Add DeepSeek V4 DSpark proposer

[alexbi29]

Summary

This is a stacked DSpark prototype branch on top of Jasl DeepSeek V4 Flash SM120 enablement branch.

Stack base: vllm-project#41834 (jasl:codex/ds4-sm120-min-enable, currently a5ebb5f66b).

This PR adds:

• method="dspark" speculative config and proposer wiring.
• DeepSeek V4 DSpark draft model support for the hybrid DSpark checkpoint.
• DSpark Triton kernels for QKV postprocess, context KV store, materialized attention helpers, Markov sampling, and fused activation/quant paths.
• A fused probabilistic Markov sampler with TP-deterministic Philox seeding.
• Stable DSpark fast-path kernels default-on for method="dspark"; obsolete env gates and dead BF16/runtime-layer ablations were removed. Only the draft-forward CUDA graph remains env-gated for now.
• Shared target embed_tokens / lm_head storage for DSpark to recover KV-cache headroom.
• Unused speculative buffers are skipped for text-only draft models and disabled DSpark materialized attention, preserving more KV-cache headroom.
• Focused DSpark tests covering config predicates, sampler correctness/fallbacks, fused sampler probability matching, and TP seed determinism.

Validation

Unit/local:

• compileall on touched config/env/proposer/runner/sampler/test files
• git diff --check
• PYTHONPATH=/home/ubuntu/llm-src-dss-pr-minimal FLASHINFER_DISABLE_VERSION_CHECK=1 /home/ubuntu/vllm-env-ds4/bin/python -m pytest tests/v1/spec_decode/test_dspark.py -q -> 14 passed
• PYTHONPATH=/home/ubuntu/llm-src-dss-pr-minimal FLASHINFER_DISABLE_VERSION_CHECK=1 /home/ubuntu/vllm-env-ds4/bin/python -m pytest tests/v1/sample/test_rejection_sampler.py::test_perfect_match tests/v1/sample/test_rejection_sampler.py::test_early_mismatch tests/v1/sample/test_rejection_sampler.py::test_multiple_sequences -q -> 3 passed

Live on 2x RTX PRO 6000 Blackwell, TP=2, DSpark hybrid checkpoint:

• TPS, 6 runs, thinking off, temperature=1.0, server --max-num-seqs 10: avg 305.0 tok/s, median 308.2 tok/s, all runs full 2048 tokens

[jasl]

Could you share your vllm serve? I tested it on my RTX Pro 6000 PC, and it got stuck on initializing the drafter

[jasl]

Oh, it seems my local model DeepSeek-V4-Flash-DSpark is corrupted.

[jasl]

vllm main has merged DSpark support, I'm trying to rebase first, and see what we are lacking

[alexbi29]

vllm-start.sh:
export VLLM_SKIP_P2P_CHECK=1
export NCCL_P2P_LEVEL=SYS
export NCCL_IB_DISABLE=1
export NCCL_NET_GDR_LEVEL=SYS
export NCCL_MIN_NCHANNELS=8
export NCCL_ALLOC_P2P_NET_LL_BUFFERS=1
export TORCH_FLOAT32_MATMUL_PRECISION=high
export NVIDIA_TF32_OVERRIDE=1

exec "$_VLLM_BIN" serve "$@" >> "$LOG_FILE" 2>&1

---

  - PYTHONPATH=/home/ubuntu/llm-src-dss-pr-minimal
  - VLLM_DSPARK_FORWARD_CUDAGRAPH=1
  - VLLM_DSPARK_FORWARD_CUDAGRAPH_ALLOW_TP=1
  - VLLM_TRITON_MLA_SPARSE_HEAD_BLOCK_SIZE=4
  - VLLM_USE_SIMPLE_KV_OFFLOAD=1
  - FLASHINFER_DISABLE_VERSION_CHECK=1
  - VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE=1
  - VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL=1
cmd: >-
  vllm-start.sh /var/log/vllm/deepseek-v4-flash-dspark.log
  /mnt/data/vllm-cache/huggingface/hybrid/DeepSeek-V4-Flash-DSpark-hybrid
  --host 127.0.0.1
  --port ${PORT}
  --served-model-name deepseek-v4-flash
  --no-use-tqdm-on-load
  --trust-remote-code
  --tokenizer-mode deepseek_v4
  --reasoning-parser deepseek_v4
  --reasoning-config '{"reasoning_start_str":"<think>","reasoning_end_str":"</think>"}'
  --enable-auto-tool-choice
  --tool-call-parser deepseek_v4
  --tensor-parallel-size 2
  --pipeline-parallel-size 1
  --enable-expert-parallel
  --kv-cache-dtype fp8
  --block-size 256
  --max-model-len 524288
  --max-num-seqs 12
  --gpu-memory-utilization 0.975
  --no-disable-hybrid-kv-cache-manager
  --async-scheduling
  --stream-interval 20
  --enable-prefix-caching
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'
  --default-chat-template-kwargs '{"thinking":true,"reasoning_effort":"high"}'
  --speculative-config '{"method":"dspark","model":"/mnt/data/vllm-cache/huggingface/hybrid/DeepSeek-V4-Flash-DSpark-hybrid","num_speculative_tokens":5}'
  --kv-offloading-backend native
  --kv-offloading-size 32

[alexbi29]

I also have some local edits (see force pushed [3411a4d] (3411a4d) - mainly parser and kv cache offload fix (soon to be merged, but not yet)

[alexbi29]

Updated in 370b45ad57: replaced the old DSpark attention v1 Triton kernel with the v2 KV-shared tensor-core kernel and removed the v1 fallback.

The v1 kernel reread the same [window + block, head_dim] KV once per (draft token, head). v2 tiles the block_size * n_heads query rows per batch element, streams the shared KV once per tile, and folds the attention sink into the online softmax denominator.

Kernel microbench on RTX PRO 6000 Blackwell:

batch	v1	v2	speedup
1	24.7 us	14.8 us	1.7x
2	34.8 us	14.9 us	2.3x
4	69.2 us	14.9 us	4.6x
8	136.2 us	15.1 us	9.0x
12	201.8 us	15.2 us	13.3x

Upstream DSpark uses a different path: topk-index sparse-MLA reuse on SM90/SM100 and FlashInfer sparse-MLA on SM100/SM120. I did not benchmark upstream on SM100, so I am not claiming a direct SM100 win. The closest measured comparison for this V4-Flash shape was the general FlashInfer sparse-MLA path on SM120, which was flat at ~133 us/call; v2 is ~15 us/call on the same shape. Upstream still has the cleaner long-term framework/reuse story, but this kernel is the better fit for this PR's current SM120 V4-Flash path.

This is a kernel cleanup/latency win rather than a claimed end-to-end TPS win for V4-Flash; earlier live A/B was throughput-neutral because target forward dominates the step.

Local validation:

• ruff format / ruff check on touched files
• git diff --check
• py_compile on touched DSpark files
• tests/v1/spec_decode/test_dspark.py -q: 17 passed
• Added CUDA coverage comparing dspark_triton_attention against a PyTorch reference at batch sizes 1, 4, and 12

[jasl]

Updated in 370b45ad57: replaced the old DSpark attention v1 Triton kernel with the v2 KV-shared tensor-core kernel and removed the v1 fallback.

The v1 kernel reread the same [window + block, head_dim] KV once per (draft token, head). v2 tiles the block_size * n_heads query rows per batch element, streams the shared KV once per tile, and folds the attention sink into the online softmax denominator.

Kernel microbench on 2x RTX PRO 6000 Blackwell Max-Q:

batch v1 v2 speedup
1 24.7 us 14.8 us 1.7x
2 34.8 us 14.9 us 2.3x
4 69.2 us 14.9 us 4.6x
8 136.2 us 15.1 us 9.0x
12 201.8 us 15.2 us 13.3x
This is a kernel cleanup/latency win rather than a claimed end-to-end TPS win for V4-Flash; earlier live A/B was throughput-neutral because target forward dominates the step.

Local validation:

• ruff format / ruff check on touched files
• git diff --check
• py_compile on touched DSpark files
• tests/v1/spec_decode/test_dspark.py -q: 17 passed
• Added CUDA coverage comparing dspark_triton_attention against a PyTorch reference at batch sizes 1, 4, and 12

I'm rebasing upstream now.
When I validated and pushed, please help to rebase, ETA about 30 minutes

[jasl]

Rebased, please rebase

[alexbi29]

rebased and fixed

[jasl]

I need more time to verify.
This PR works in most cases, but it failed on a user-reported case.

And I'm also testing Model Runner 1 vs 2

[jasl]

Update: I took the branch further and I think the headline finding is good news for your design — your V1 DSparkProposer works, and on our SM12x stack it's the better path than the V2 speculator.

I checked it out, activated the V1 path, and validated on 2×RTX PRO 6000 (SM120, TP=2) with the DSpark checkpoint.

The V1 path is complete and just needed un-gating

The V1 DSparkProposer (vllm/v1/spec_decode/dspark.py) is already wired in the V1 runner (gpu_model_runner.py via use_dspark()); it was only dead because vllm/config/vllm.py:use_v2_model_runner force-routes method=="dspark" → V2. Since the env is read before that branch, VLLM_USE_V2_MODEL_RUNNER=0 already runs your V1 proposer with zero code changes — that's how I tested it.

Why V1 is worth defaulting to

The decisive reason is correctness: on our stack the V2 runner has a DeepSeek-V4 long-context recall bug (needle recall collapses under concurrency): plain V4-Flash and DSpark on V2 both score ~3/16 on an arthur-style needle test at ~28k context / concurrency 8, while V1 is 16/16. So DSpark-on-V2 inherits it; DSpark-on-V1 doesn't.

On perf (same build, same workload, fixed seeds, fresh serve per cell; single-sample spec-decode throughput is very noisy so treat tok/s ranges with care):

metric (2×RTX SM120, DSpark ckpt)	V1 runner	V2 runner
arthur recall @28k conc8	16/16	3/16
mean draft acceptance, temp 0	3.39	2.19
mean draft acceptance, temp 0.7	2.72	2.06
single-stream tok/s, temp 0 (median of 5)	199 (185–246)	224 (187–228)

Interesting tradeoff: V1's draft acceptance is markedly higher (your V1 _sample_markov accepts ~55% more tokens/step than the V2 speculator), but V2's per-step is faster (~1.7×), so V2 edges out V1 on raw tok/s despite lower acceptance — i.e. the tok/s difference is general V1-vs-V2 runner overhead, not the DSpark sampler. Given V2's recall bug, V1 is still the right default; the ~10% raw-throughput gap is the price of correctness (and it's within the measurement noise here).

(The DSpark draft model reports it doesn't support torch.compile, so it runs eager on both runners — but V1 still gets FULL_AND_PIECEWISE cudagraph for the target. VLLM_DSPARK_FORWARD_CUDAGRAPH_ALLOW_TP=1 for the draft made no measurable difference on TP=2.)

Three small edits to make V1 the clean default (branch feature/dspark-v1-optin on jasl/vllm, on top of this PR)

1. vllm/config/vllm.py — drop the method=="dspark" → V2 force so DSpark follows normal runner selection (V1 default; V2 still reachable via VLLM_USE_V2_MODEL_RUNNER=1 for A/B).
2. vllm/model_executor/warmup/kernel_warmup.py — _deepseek_v4_slot_mapping_warmup was rewritten to the V2 interface (block_tables/input_buffers) with an early return when block_tables is None, so on the V1 runner it silently no-ops and reintroduces first-request JIT for all DeepSeek-V4 serving (MTP/DFlash/plain). I split it into V1 (input_batch/query_start_loc) + V2 branches. This one is a production regression independent of DSpark, worth folding in regardless.
3. vllm/models/deepseek_v4/nvidia/model.py — dspark_fused_shared_experts_quant defaults True on every SpeculativeConfig, and the MTP block sits at layer_idx == num_hidden_layers, so _is_dspark_runtime_layer is true for it → the fused FP8 shared-experts kernel would engage on production MTP numerics. I gated _is_dspark_runtime_layer on method=="dspark".

I verified MTP2 (W4A16) is unregressed with these: warmup fires, coherent, 16/16 recall, ~same tok/s.

Happy to open these as a PR against your branch (or you can cherry-pick feature/dspark-v1-optin). Either way, nice work — the V1 proposer is the piece that makes this land cleanly on the SM12x DeepSeek-V4 stack.

[alexbi29]

V2 vs V1 - good catch, i did not run high concurrency tests! Debugged and fixed though, PR coming.