[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints) - #16813
[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints)#16813chungen04 wants to merge 3 commits into
Conversation
|
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 (3)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughChangesQwen3 DSpark drafting is added with rolling K/V windows, batched execution, checkpoint loading, and shared target heads. DSpark mask-token resolution and speculative dispatch are updated. Golden tests cover protocol execution, prefix reuse, batching, and ring-buffer wraparound. Qwen3 DSpark drafter
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SpeculativeConfig
participant Qwen3DSparkForCausalLM
participant Qwen3DSparkDraftModel
participant ContextRing
participant ProposalHeads
SpeculativeConfig->>Qwen3DSparkForCausalLM: architecture, block_size, mask_token_id
Qwen3DSparkForCausalLM->>Qwen3DSparkDraftModel: construct and load checkpoint
Qwen3DSparkDraftModel->>ContextRing: seed projected context K/V rows
Qwen3DSparkDraftModel->>ContextRing: read valid rolling window
Qwen3DSparkDraftModel->>ProposalHeads: produce draft hidden states
ProposalHeads-->>Qwen3DSparkForCausalLM: proposed tokens and logits
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
| return self.proj(features.float()).squeeze(-1) | ||
|
|
||
|
|
||
| class Qwen3DSparkDraftModel(nn.Module): |
There was a problem hiding this comment.
Hey @chungen04, thanks for the contribution! Is it possible to merge any of this code with the existing dspark drafter? Would be nice to consolidate all logic related to dspark drafting.
There was a problem hiding this comment.
@mikeiovine Thank you for the review. I am refactoring the previous modeling_dspark_qwen3.py into existing frameworks and creating the base class, class DSparkForCausalLMBase. Let me know if you have any suggestions while I refine the PR.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/modeling_dspark.py (2)
1717-1724: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the other cross-module public names to
__all__.
modeling_speculative.get_draft_modelimportscount_dspark_stagesfrom this module, and it is absent from__all__. Direct imports still work, so nothing breaks now, but the list no longer describes the module's public interface.Based on the guideline "keep
__all__updated for public interfaces".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_dspark.py` around lines 1717 - 1724, Add count_dspark_stages to the __all__ declaration in modeling_dspark.py so the public export list includes the cross-module symbol imported by modeling_speculative.get_draft_model. Preserve all existing public names.Source: Coding guidelines
1462-1468: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid the explicit GQA materialization.
repeat_interleavematerializes full K/V tensors at every layer and decode step. The supported PyTorch range acceptsenable_gqa=Truewithattn_mask, but SDPA's fallback also performsrepeat_interleave; the flag alone does not remove the copy. Use a tested attention path or tensor layout that consumes grouped K/V without materialization.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_dspark.py` around lines 1462 - 1468, Update the attention flow around scaled_dot_product_attention to avoid materializing grouped K/V via repeat_interleave; replace the current path with the tested attention implementation or tensor layout that directly consumes grouped K/V while preserving attn_mask, softmax_scale, and output behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_dspark.py`:
- Around line 1189-1191: Replace the assertion validating
self.num_capture_layers in the Qwen3 DSpark drafter configuration with an
explicit ValueError, preserving the existing message and ensuring invalid
target_layer_ids fail before self.fc is constructed.
- Around line 1505-1519: In the wrapper containing the shown initialization
logic, add an explicit validation that every normalized start_pos value is
strictly greater than zero before calling forward_batched or writing context
state. Reject scalar and batched inputs consistently, preserving valid positive
positions and failing immediately for any start_pos <= 0.
- Around line 1329-1338: Constrain the interim back-fill around the
stage_windows write in the layer loop so max_draft_len never exceeds the
configured window_size, or explicitly handle wrapped positions without duplicate
indexed slots. Preserve the existing absolute-position conversion and ensure
each stage_windows slot receives the correct context when interim rows exceed
one window.
In `@tensorrt_llm/_torch/speculative/dspark.py`:
- Around line 250-271: Initialize `_batch_to_slot` in the DSpark setup alongside
`_dummy_slot` using the scratch-slot index rather than zero. Preserve updates
for real request slots during preparation, and add a regression test covering
the first forward call with mixed context and generation rows before
`_dspark_worker` is assigned.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_dspark.py`:
- Around line 1717-1724: Add count_dspark_stages to the __all__ declaration in
modeling_dspark.py so the public export list includes the cross-module symbol
imported by modeling_speculative.get_draft_model. Preserve all existing public
names.
- Around line 1462-1468: Update the attention flow around
scaled_dot_product_attention to avoid materializing grouped K/V via
repeat_interleave; replace the current path with the tested attention
implementation or tensor layout that directly consumes grouped K/V while
preserving attn_mask, softmax_scale, and output behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f346db2e-39b7-4de1-84d0-38ae3cdb06d6
📒 Files selected for processing (7)
tensorrt_llm/_torch/models/dspark/draft.pytensorrt_llm/_torch/models/dspark/heads.pytensorrt_llm/_torch/models/modeling_dspark.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/speculative/dspark.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/models/modeling_speculative.py
| assert self.num_capture_layers > 0, ( | ||
| "Qwen3 DSpark drafter config must provide target_layer_ids" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Raise ValueError instead of asserting on the config.
python -O removes the assertion. Without target_layer_ids, self.num_capture_layers becomes 0 and self.fc is built with in_features=0, so the failure surfaces later as a confusing shape error. Raise an explicit exception for the invalid configuration.
🛠️ Proposed fix
- assert self.num_capture_layers > 0, (
- "Qwen3 DSpark drafter config must provide target_layer_ids"
- )
+ if self.num_capture_layers == 0:
+ raise ValueError("Qwen3 DSpark drafter config must provide target_layer_ids")Based on the guideline "raise ValueError rather than assertions".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| assert self.num_capture_layers > 0, ( | |
| "Qwen3 DSpark drafter config must provide target_layer_ids" | |
| ) | |
| if self.num_capture_layers == 0: | |
| raise ValueError("Qwen3 DSpark drafter config must provide target_layer_ids") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/modeling_dspark.py` around lines 1189 - 1191,
Replace the assertion validating self.num_capture_layers in the Qwen3 DSpark
drafter configuration with an explicit ValueError, preserving the existing
message and ensuring invalid target_layer_ids fail before self.fc is
constructed.
Source: Coding guidelines
| # Reserve slot index max_batch as a scratch slot for cuda-graph padding | ||
| # rows and warmup dummies; real requests only draw slots 0..max_batch-1, | ||
| # so their windows can never be written by a padded row (which the | ||
| # per-step context write would otherwise do, unmasked, through the | ||
| # duplicated slot index). | ||
| self._dummy_slot = max_batch | ||
| num_slots = max_batch + 1 | ||
|
|
||
| self._kv_windows = torch.zeros( | ||
| (max_batch, num_stages, self._win, head_dim), | ||
| (num_slots, num_stages, self._win, head_dim), | ||
| dtype=torch.bfloat16, | ||
| device="cuda", | ||
| ) | ||
| self._ctx_len = torch.zeros(max_batch, dtype=torch.long, device="cuda") | ||
| self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") | ||
| self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") | ||
| self._free_slots = deque(range(max_batch)) | ||
| self._req_to_slot = {} | ||
| self._win_inited = True | ||
| logger.info( | ||
| f"DSpark: allocated rolling KV windows " | ||
| f"[{max_batch}, {num_stages}, {self._win}, {head_dim}]" | ||
| f"[{num_slots}, {num_stages}, {self._win}, {head_dim}] " | ||
| f"({max_batch} request slots + 1 scratch)" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 6 \
'_dspark_worker\s*=|\.prepare\(\)|_lazy_init\(' \
tensorrt_llm/_torch/speculative \
tests/unittest/_torch/speculativeRepository: NVIDIA/TensorRT-LLM
Length of output: 28812
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/_torch/speculative/dspark.py
rg -n -C 12 \
'class DSparkSpecMetadata|def prepare\(|def _forward_impl|def forward\(|_dspark_worker|prepare\(\)' \
tensorrt_llm/_torch/speculative/dspark.py \
tensorrt_llm/_torch/speculative/speculative_interface.py \
tensorrt_llm/_torch/speculative/speculative_decoding.py \
tensorrt_llm/_torch/speculative/drafting_loops.py \
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 35293
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,230p' tensorrt_llm/_torch/speculative/dspark.py
sed -n '230,520p' tensorrt_llm/_torch/speculative/dspark.py
sed -n '1,180p' tensorrt_llm/_torch/speculative/drafting_loops.py
sed -n '1,330p' tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 44693
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 10 \
'DSparkSpecMetadata|DSparkWorker|spec_metadata\.prepare|draft_model.*forward|worker.*forward|forward\(.*spec_metadata' \
tensorrt_llm tests \
-g '*.py' | head -n 1200Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 18 'spec_metadata\.prepare\(\)|\.prepare\(\).*spec_metadata|prepare_spec|prepare\(' \
tensorrt_llm/_torch/model_engine.py \
tensorrt_llm/_torch/models/modeling_speculative.py \
tensorrt_llm/_torch/speculative \
tensorrt_llm | rg -v '(^|/)(test|tests)/' | head -n 1200
sed -n '1960,2085p' tensorrt_llm/_torch/models/modeling_speculative.py
sed -n '1020,1110p' tensorrt_llm/_torch/speculative/interface.py
sed -n '430,550p' tensorrt_llm/_torch/speculative/dflash.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '480,545p' tensorrt_llm/_torch/speculative/model_drafter.py
rg -n -C 25 'spec_worker\.forward|self\.spec_worker\(|spec_worker' \
tensorrt_llm/_torch/models/modeling_speculative.py \
tensorrt_llm/_torch/model_engine.py \
tensorrt_llm/_torch | head -n 1400
rg -n -C 12 'def forward\(|_forward_impl\(' tensorrt_llm/_torch/speculative/interface.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '490,535p' tensorrt_llm/_torch/speculative/model_drafter.py
rg -n -C 20 'spec_worker\.forward|spec_worker\(' tensorrt_llm/_torch --glob '*.py'
rg -n -C 20 'def forward\(|_forward_impl\(' tensorrt_llm/_torch/speculative/interface.py | head -n 500Repository: NVIDIA/TensorRT-LLM
Length of output: 29021
Initialize _batch_to_slot with _dummy_slot
should_forward_draft_model() skips DSpark’s first context chunk, so the first worker call can contain both context and generation rows. prepare() cannot map these rows before _dspark_worker is assigned; zero-initialization can route a generation row to a real context slot. Add a first-forward mixed-batch regression test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/speculative/dspark.py` around lines 250 - 271, Initialize
`_batch_to_slot` in the DSpark setup alongside `_dummy_slot` using the
scratch-slot index rather than zero. Preserve updates for real request slots
during preparation, and add a regression test covering the first forward call
with mixed context and generation rows before `_dspark_worker` is assigned.
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>
…ramework Fold modeling_dspark_qwen3.py into the shared DSpark modeling path instead of keeping a parallel Qwen3-specific implementation, so drafter construction, heads and worker plumbing go through one code path. Signed-off-by: chungen04 <b09901027@ntu.edu.tw>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (4)
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py (4)
332-334: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop the unused
weightsunpack.
test_batched_matches_eager_singletonsnever usesweights. Ruff reports RUF059 on Line 334. If RUF rules gate CI, this fails the lint stage.♻️ Proposed fix
- weights, model, device = setup + _, model, device = setup🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py` around lines 332 - 334, Remove the unused weights unpacking from test_batched_matches_eager_singletons by binding only the model and device values returned by setup, while preserving the test’s existing behavior.Source: Linters/SAST tools
395-469: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
_Ref.draftinstead of duplicating it inline.Lines 416-469 re-implement
_Ref.draftverbatim._WrapRefadds nothing over_Ref._Ref.draftalready derivespos_cfromfirst_pos, sofirst_pos = start_pos - winproduces exactly the same context positions (total+1-win .. total) and satisfies the internal length assertion. The duplicated block will drift from_Refwhen the reference path changes.♻️ Proposed fix
- # Reference limited to the last `win` context positions. - ref = _Ref(weights, device) - ref.append_ctx(h) - ref.ctx_x = ref.ctx_x[-win:] - bonus = torch.tensor([42], device=device) start = torch.tensor([total + 1], device=device) main = _rand_hidden(gen, 1).to(device) - # keep ref in sync: main_hidden row is position `total` (= start-1) + + # Reference limited to the last `win` context positions; the main_hidden + # row is position `total` (= start-1). + ref = _Ref(weights, device) + ref.append_ctx(h) ref.append_ctx(main) ref.ctx_x = ref.ctx_x[-win:] + ref.first_pos = int(start[0]) - win got_toks, _, got_logits = model.forward_batched( main, bonus, start, kv_windows=kv_windows, slots=torch.tensor([0], device=device), return_logits=True, ) - - # Reference drafts with positions: ctx = last `win` absolute positions. - class _WrapRef(_Ref): - pass - - wref = _WrapRef(weights, device) - wref.ctx_x = ref.ctx_x - # override position bookkeeping: ctx positions are total+1-win .. total - w = wref.w - ids = torch.full((BLOCK,), MASK_ID, dtype=torch.long, device=device) - ids[0] = bonus[0] - hh = F.embedding(ids, w["embed_tokens.weight"]) - pos_q = start[0] + torch.arange(BLOCK, device=device) - pos_c = torch.arange(start[0] - win, start[0], device=device) - for i in range(N_LAYERS): - p = f"layers.{i}." - x = wref._norm(hh, p + "input_layernorm.weight") - q = F.linear(x, w[p + "self_attn.q_proj.weight"]).view(BLOCK, N_HEADS, HEAD_DIM) - q = wref._norm(q, p + "self_attn.q_norm.weight") - q = _apply_rope(q, wref.cos[pos_q].to(DTYPE), wref.sin[pos_q].to(DTYPE)) - src = torch.cat([wref.ctx_x, x], dim=0) - pos_k = torch.cat([pos_c, pos_q]) - k = F.linear(src, w[p + "self_attn.k_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) - k = wref._norm(k, p + "self_attn.k_norm.weight") - k = _apply_rope(k, wref.cos[pos_k].to(DTYPE), wref.sin[pos_k].to(DTYPE)) - v = F.linear(src, w[p + "self_attn.v_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) - rep = N_HEADS // N_KV_HEADS - kk = k.transpose(0, 1).repeat_interleave(rep, dim=0) - vv = v.transpose(0, 1).repeat_interleave(rep, dim=0) - o = F.scaled_dot_product_attention(q.transpose(0, 1), kk, vv, scale=HEAD_DIM**-0.5) - o = o.transpose(0, 1).reshape(BLOCK, N_HEADS * HEAD_DIM) - hh = hh + F.linear(o, w[p + "self_attn.o_proj.weight"]) - x = wref._norm(hh, p + "post_attention_layernorm.weight") - mlp = F.linear( - F.silu(F.linear(x, w[p + "mlp.gate_proj.weight"])) - * F.linear(x, w[p + "mlp.up_proj.weight"]), - w[p + "mlp.down_proj.weight"], - ) - hh = hh + mlp - hh = wref._norm(hh, "norm.weight") - base = F.linear(hh, w["lm_head.weight"]) - toks, logits = [], [] - prev = bonus.long() - for kstep in range(BLOCK): - bias = F.linear( - F.embedding(prev, w["markov_head.markov_w1.weight"]), w["markov_head.markov_w2.weight"] - ) - step = base[kstep : kstep + 1] + bias - logits.append(step) - prev = step.argmax(dim=-1) - toks.append(prev) - torch.testing.assert_close( - got_logits[0].float(), torch.cat(logits).float(), atol=0.05, rtol=0.05 - ) - assert torch.equal(got_toks[0].long(), torch.cat(toks).long()) + exp_toks, exp_logits = ref.draft(bonus[0], int(start[0])) + torch.testing.assert_close(got_logits[0].float(), exp_logits.float(), atol=0.05, rtol=0.05) + assert torch.equal(got_toks[0].long().cpu(), exp_toks.long().cpu())🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py` around lines 395 - 469, Replace the inline reference forward computation and no-op _WrapRef subclass with a call to _Ref.draft, passing the existing context, bonus tokens, and first_pos derived as start[0] - win. Preserve the current logits and token comparisons while relying on _Ref.draft to derive the context positions and validate the window length.
372-373: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse explicit tolerances for the batched-vs-singleton logit comparison.
This call relies on the default
assert_closetolerances, while every other logit assertion in this file usesatol=0.05, rtol=0.05. Batched and singleton calls use different tensor shapes, so GEMM reduction order can differ and produce small numeric drift. Set explicit tolerances to keep the test stable.♻️ Proposed fix
- torch.testing.assert_close(got_logits, torch.cat(exp_logits, dim=0)) + torch.testing.assert_close( + got_logits.float(), torch.cat(exp_logits, dim=0).float(), atol=0.05, rtol=0.05 + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py` around lines 372 - 373, Update the torch.testing.assert_close call comparing got_logits with torch.cat(exp_logits, dim=0) to pass explicit atol=0.05 and rtol=0.05, matching the other logit assertions in this test file while leaving the token comparison unchanged.
229-469: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winTest coverage summary.
test_worker_protocol_golden— added and registered.test_prefix_reuse_masks_unseeded_rows— added and registered.test_batched_matches_eager_singletons— added and registered.test_ring_window_wraparound— added and registered.The module is listed in
tests/integration/test_lists/test-db/l0_cpu.ymlandtests/integration/test_lists/test-db/l0_h100.yml. Coverage verdict: sufficient.Consider adding assertions for
return_logits=Falseand confidence-head outputs.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py` around lines 229 - 469, Extend the existing forward_batched coverage in test_worker_protocol_golden or test_batched_matches_eager_singletons to exercise return_logits=False and verify its returned confidence-head outputs against the corresponding return_logits=True behavior. Preserve the current token assertions and ensure both batched and singleton paths remain consistent.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py`:
- Around line 332-334: Remove the unused weights unpacking from
test_batched_matches_eager_singletons by binding only the model and device
values returned by setup, while preserving the test’s existing behavior.
- Around line 395-469: Replace the inline reference forward computation and
no-op _WrapRef subclass with a call to _Ref.draft, passing the existing context,
bonus tokens, and first_pos derived as start[0] - win. Preserve the current
logits and token comparisons while relying on _Ref.draft to derive the context
positions and validate the window length.
- Around line 372-373: Update the torch.testing.assert_close call comparing
got_logits with torch.cat(exp_logits, dim=0) to pass explicit atol=0.05 and
rtol=0.05, matching the other logit assertions in this test file while leaving
the token comparison unchanged.
- Around line 229-469: Extend the existing forward_batched coverage in
test_worker_protocol_golden or test_batched_matches_eager_singletons to exercise
return_logits=False and verify its returned confidence-head outputs against the
corresponding return_logits=True behavior. Preserve the current token assertions
and ensure both batched and singleton paths remain consistent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3fb8a693-9ae5-4dd7-8b2a-c86cd5eff1e8
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/dspark/draft.pytensorrt_llm/_torch/models/dspark/heads.pytensorrt_llm/_torch/models/modeling_dspark.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/models/dspark/draft.py
- tensorrt_llm/_torch/models/dspark/heads.py
- tensorrt_llm/_torch/models/modeling_speculative.py
Signed-off-by: chungen04 <b09901027@ntu.edu.tw>
|
There's no CI accuracy gate on the real checkpoint. Could we add a single-GPU test mirroring test_dflash (tests/integration/defs/accuracy/test_llm_api_pytorch.py:656), plus a test-list entry? The existing TestDeepSeekV4ProDSpark test needs 8 GPUs, but a Qwen3-8B drafter fits the standard single-GPU harness. |
| self._build_rope_tables() | ||
|
|
||
| # Ring-window length for the worker-owned context K/V buffer. | ||
| window = int(os.environ.get(_CTX_WINDOW_ENV, _DEFAULT_CTX_WINDOW)) |
There was a problem hiding this comment.
Is this required? Can we make it flow through the decoding config?
|
/bot run |
|
PR_Github #64645 [ run ] triggered by Bot. Commit: |
|
PR_Github #64645 [ run ] completed with state
|
|
Looking into the context length argument issue, I found another thing that worth a revisit. The code I made earlier is trying to indicate the context window for the DSpark speculator, which is assigned to 128 in DSv4 (see, for example, https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731/blob/main/config.json, Yet, in an earlier issue and PR also made by me (#16005, #16150 ), I pointed that since current drafter context is a contiguous buffer rather than paged, the buffer can be too large for long So the support of DSpark drafter without sliding window is related to #16005, #16150, while that PR requires a mid-size code change. I would like to know the feedbacks in the upstream team regarding this. Possible directions and design choices include:
cc @zhaoyangwang-nvidia as you reviewed #16150, would appreciate your suggestions. |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Solid addition — the golden tests against a torch-only DeepSpec reference port are exactly the right coverage for a pure-torch drafter, and I verified the worker-side contracts the new model relies on (_seed_context_windows slices each chunk to min(win, chunk_len), so write_context_windows never sees duplicate ring indices in one scatter; interim back-fill frames and the bonus row are disjoint, so the masked torch.where scatter in write_context_windows_batched is well-defined). Note the new test file needs no test-list change: unittest/_torch/speculative/hw_agnostic is already registered as a directory in l0_cpu.yml and l0_h100.yml, so the CodeRabbit "insufficient test-list coverage" note in the description is wrong.
Two whole-PR notes:
TRTLLM_DSPARK_QWEN3_CTX_WINDOWis user-facing but undocumented — it appears only in the PR description. DSpark has no feature docs yet, so there's no obvious home, but at minimum a mention wherever DSpark docs eventually land, including the memory cost: the worker allocates(max_batch+1) × num_layers × window × 2·kv_dimbf16, which for a multi-layer drafter at window 2048 is nontrivial per batch slot.- The inline comments are all robustness/diagnosability items (out-of-range noise-token fallback, silent RoPE clamp past
max_position_embeddings, unguarded env-var parse, config-key fallback scope) — none block merge given the E2E validation.
| if mask_token_id is None: | ||
| mask_token_id = getattr(config, ckpt_attr, None) | ||
| if mask_token_id is None: | ||
| mask_token_id = config.vocab_size |
There was a problem hiding this comment.
The config.vocab_size fallback produces an out-of-range token id for the Qwen3 drafter: embed_tokens is the target's nn.Embedding(vocab_size, hidden), so embedding id vocab_size triggers a device-side assert deep inside the first draft forward — a hard-to-diagnose failure mode. Since the DeepSpec checkpoints always carry mask_token_id, raising a clear ValueError when both the spec config and the checkpoint attribute are missing would fail fast with an actionable message. (The fallback predates this PR on the V4 path, but now that the helper is shared it's worth hardening.) Separately, the docstring's first sentence is truncated: "either indicated in DSparkDecodingConfig validation" doesn't parse.
| def _gather_cos_sin(self, positions: torch.Tensor, dtype: torch.dtype): | ||
| # Clamp for graph-safety: masked-out entries may carry arbitrary | ||
| # (already clamped by the worker) positions. | ||
| p = positions.long().clamp(min=0, max=self._freqs_cap - 1) |
There was a problem hiding this comment.
Positions past _freqs_cap - 1 (≈ max_position_embeddings, 40960 for these checkpoints) silently clamp to the last RoPE row, so once a sequence runs past the drafter's trained range every draft query/context position collapses to the same rotary phase. Output stays correct via target verification, so this shows up only as an unexplained acceptance-rate/perf cliff at long context. A warning inside this method isn't graph-safe, but a one-time construction-time warning when the serving max sequence length exceeds the drafter's max_position_embeddings would make the cliff attributable. Related: only rope_theta is read here — a rope_scaling/YaRN entry in a drafter config would be silently ignored; worth asserting it's absent.
| # DeepSpec-released dense drafter checkpoints | ||
| # (e.g. Qwen3DSparkModel) use unprefixed top-level | ||
| # keys in their own config.json. | ||
| value = draft_cfg.get(key) |
There was a problem hiding this comment.
This unprefixed fallback runs for every DSpark checkpoint, not just the DeepSpec dense drafters: a V4-style checkpoint that happens to carry an unrelated top-level block_size or mask_token_id (both plausible generic key names) would silently adopt it — best case a confusing block_size != max_draft_len validation error, worst case a wrong mask token that only degrades acceptance. Consider gating this branch on the checkpoint's architectures containing Qwen3DSpark, mirroring the dispatch in get_draft_model.
| self._build_rope_tables() | ||
|
|
||
| # Ring-window length for the worker-owned context K/V buffer. | ||
| window = int(os.environ.get(_CTX_WINDOW_ENV, _DEFAULT_CTX_WINDOW)) |
There was a problem hiding this comment.
int(os.environ.get(...)) raises a bare ValueError on a non-integer value, and the next line silently clamps out-of-range values into [block_size + 2, max_position_embeddings]. A logger.warning when the clamp changes the requested value (and a clearer error on unparseable input) would save a user who sets the env var and wonders why it had no effect. The env var itself is user-facing and currently documented only in the PR description.
5 participants
Dev Engineer Review
deepseek-ai/dspark_qwen3_{4b,8b,14b}_block7.get_draft_model.QA Engineer Review
Added test functions:
test_worker_protocol_goldentest_prefix_reuse_masks_unseeded_rowstest_batched_matches_eager_singletonstest_ring_window_wraparoundThese tests cover worker-protocol drafting, prefix reuse, batched and singleton parity, and ring-buffer wraparound. No corresponding
tests/integration/test_lists/coverage is listed. Verdict: insufficient.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):
Engine launch command:
with
cfg_dspark.yamlQwen3-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.