Skip to content

Add DSpark block drafter support - #2462

Closed
Tianlei Wu (tianleiwu) wants to merge 1 commit into
tlwu/qwen_gdn_paged_dflash2from
tlwu/qwen_gdn_paged_dspark
Closed

Add DSpark block drafter support#2462
Tianlei Wu (tianleiwu) wants to merge 1 commit into
tlwu/qwen_gdn_paged_dflash2from
tlwu/qwen_gdn_paged_dspark

Conversation

@tianleiwu

@tianleiwu Tianlei Wu (tianleiwu) commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Description

Adds end-to-end DSpark speculative decoding support for Qwen3.8 models. The implementation
reuses the DFlash2 block-drafter lattice runtime introduced by #2460, while adding DSpark
graph export, Markov scoring, shared target weights, and cache accounting for both full and
sliding-window draft models.

This PR is intentionally stacked on #2460 and should be reviewed after that PR.

Summary of Changes

DSpark model export

File Change
src/python/py/models/builders/dspark.py Builds the BF16 DSpark graph with dual-source non-causal attention, YaRN RoPE, target embedding and LM-head reuse, top-k candidates, and FP32 Markov lattice scores.
src/python/py/models/builders/qwen.py Adds DSpark export options, validates target hidden-state taps, saves the auxiliary graph, shares external initializers, and emits model.dspark configuration.

Runtime integration

File Change
src/config.cpp Parses model.dspark through the common block-drafter configuration.
src/dflash2_drafter.h Exposes cache sizing helpers used by engine pool planning.
src/dflash2_drafter.cpp Accepts DSpark's all-predicting block geometry and supports fixed-ring or full-context cache allocation according to sliding_window.
src/engine/engine.cpp Budgets full-attention drafter KV bytes with the main paged cache and retains fixed per-request rings for windowed drafters.

Testing

  • cmake --build build/cuda130/Release --parallel
  • ./build/cuda130/Release/engine_unit_tests
  • python3 -m py_compile src/python/py/models/builders/dspark.py src/python/py/models/builders/qwen.py
  • Standalone graph validation: 7/7 top-1 candidates matched; 109/112 top-k overlap.
  • Greedy losslessness: 0 mismatches across 2,042 generated tokens.
  • Eight-prompt acceptance test: 3.586 tokens per verification step over 1,862 steps.
  • H200 benchmark matrix: 2.09x to 5.08x decode throughput over the same target graph without a drafter across context lengths 512-8192 and batch sizes 1-16.
  • MMLU-Pro 800 samples: 82.62% with DSpark versus 83.00% without drafting (p=0.508).
  • GPQA Diamond: 76.26% with DSpark versus 72.73% without drafting (p=0.143).

Motivation and Context

DSpark extends the DFlash-style parallel block drafter with a low-rank Markov head. Both
models expose the same candidate lattice to the runtime, so sharing the proposal and
verification path avoids a second speculative-decoding implementation.

The evaluated Qwen3.8 checkpoint uses full attention and therefore grows its draft KV cache
with resident context. DSpark is not inherently limited to full attention: positive
sliding_window configurations use the existing fixed-ring path.

Checklist

  • Tests added/updated through the shared block-drafter test coverage
  • No breaking changes
  • Export and runtime behavior validated end to end

Reuse the DFlash2 lattice runtime while adding DSpark graph export, shared initializer plumbing, and full-attention cache accounting.
@tianleiwu

Copy link
Copy Markdown
Contributor Author

Superseded by #2497, which reuses the compact #2496 block-drafter runtime for DSpark without duplicated builder changes, dense recurrent checkpoints, or an ONNX Runtime dependency bump.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant