[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints)#16813
Open
chungen04 wants to merge 1 commit into
Open
[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints)#16813chungen04 wants to merge 1 commit into
chungen04 wants to merge 1 commit into
Conversation
Support the DeepSpec-released dense DSpark drafters for Qwen3 targets (e.g. deepseek-ai/dspark_qwen3_8b_block7) alongside the existing DeepSeek-V4 mtp.*-namespace drafter: - modeling_dspark_qwen3.py: Qwen3DSparkDraftModel/Qwen3DSparkForCausalLM, a pure-torch dense GQA draft backbone (fc+hidden_norm context projection, Qwen3 layers with captured-context K/V and bidirectional block attention, Markov head refinement) implementing the same worker-facing protocol as DSparkDraftModel, so DSparkWorker / DSparkSpecMetadata / CUDA-graph plumbing are reused unchanged. The worker rolling buffer holds per-layer context K/V as a ring over the last TRTLLM_DSPARK_QWEN3_CTX_WINDOW (default 2048) committed positions. - get_draft_model: dispatch on the drafter checkpoint's Qwen3DSparkModel architecture. - llm_args DSpark validation: also resolve unprefixed top-level config keys (block_size / target_layer_ids / mask_token_id / markov_rank) used by the DeepSpec drafter checkpoints (no schema change; golden manifest unchanged). - tests: golden tests vs a torch-only port of the DeepSpec reference (worker frame conventions across multi-step decode, batched-vs-singleton parity, ring wraparound). Validation: golden unit tests vs a torch-only DeepSpec reference port; end-to-end + GSM8K/MATH-500/HumanEval benchmarks (aiperf, 1xB300, conc 1..64) were run on the v1.3.0rc22 backport of this change (branch dspark-qwen3-rc22, same diff) inside the 1.3.0rc22 release container: 3.3-3.7x per-user speedup over vanilla at concurrency 1 (~750 tok/s/user), 6.17 avg decoded tokens/iter on GSM8K (block=7). Signed-off-by: chungen04 <cho322@gatech.edu>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
WalkthroughAdds a complete Qwen3 DSpark draft model with projected context K/V ring storage, batched decoding, target-model weight sharing, speculative-decoding dispatch, configuration fallback, and golden tests for protocol, batching, and wraparound behavior. ChangesQwen3 DSpark drafter
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TargetModel
participant Qwen3DSparkForCausalLM
participant Qwen3DSparkDraftModel
participant RingBuffer
TargetModel->>Qwen3DSparkForCausalLM: load shared embeddings and LM head
TargetModel->>Qwen3DSparkForCausalLM: provide captured hidden states
Qwen3DSparkForCausalLM->>Qwen3DSparkDraftModel: write context windows
Qwen3DSparkDraftModel->>RingBuffer: store per-layer context K/V
Qwen3DSparkForCausalLM->>Qwen3DSparkDraftModel: forward_batched
Qwen3DSparkDraftModel->>RingBuffer: read committed context
Qwen3DSparkDraftModel-->>Qwen3DSparkForCausalLM: proposed tokens and logits
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dev Engineer Review
Qwen3DSparkForCausalLMinto DSpark model dispatch and added support for unprefixed DeepSpec configuration keys.QA Engineer Review
test_worker_protocol_goldentest_batched_matches_eager_singletonstest_ring_window_wraparoundsetupfixturetests/integration/test_lists/,test-db/, orqa/coverage entry was found.Description
DSpark support today (#15808) covers only the DeepSeek-V4 drafter, whose draft weights live in the target checkpoint and in V4 blocks. DeepSeek's release (https://github.com/deepseek-ai/DeepSpec) also ships standalone dense DSpark drafters for Qwen3 targets,
deepseek-ai/dspark_qwen3_{4b,8b,14b}_block7, which this PR enables. This PR also serves as a stepping stone to serve other Qwen3-based DSpark drafter, e.g. novita/kimi-k2.6-dsparkmodeling_dspark_qwen3.py(new):Qwen3DSparkDraftModel/Qwen3DSparkForCausalLM: a pure-torch dense GQA draft backbone (fc+hidden_normcaptured-context projection, Qwen3 decoder layers with per-head q/k RMSNorm and RoPE, bidirectional block attention over a per-layer context-K/V ring cache, Markov-head block refinement). It implements the same worker-facing protocol as the V4DSparkDraftModel, soDSparkWorker,DSparkSpecMetadata, and the CUDA-graph plumbing are reused unchanged. The worker-owned rolling buffer holds per-layer context K/V as a ring over the lastTRTLLM_DSPARK_QWEN3_CTX_WINDOW(default 2048) committed positions.get_draft_model: dispatch on the drafter checkpoint'sQwen3DSparkModelarchitecture (mirrors the DFlash/Laguna pattern).llm_argsDSpark validation: additionally resolve the unprefixed top-level config keys (block_size/target_layer_ids/mask_token_id/markov_rank) used by the DeepSpec drafter checkpoints.Usage:
Test Coverage
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py(new): golden tests against a torch-only port of the DeepSpec reference (deepspec/modeling/dspark/qwen3/modeling.py+eval/dspark/draft_ops.py) — token-exact agreement across multi-step decode driven through the exactDSparkWorkerconventions (prefill seeding, interim back-fill, frame offsets), batched-vs-singleton parity, and ring-window wraparound.test_llm_args.pyDSpark validation tests (9) pass unchanged.1.3.0rc22backport of this diff): results below.trtllm-serve(1xB300, bf16, greedy, chat template with thinking disabled, natural EOS, max_tokens 1024,max_batch_size 64, CUDA graphs on, overlap scheduler off for all configs), GSM8K / MATH-500 / HumanEval in the DeepSpec eval prompt format. Eagle3 baseline swept at draft length 1 and 3.Per-user decode rate, tok/s (speedup vs vanilla):
Qwen3-8B
Qwen3-4B
Qwen3-14B
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.