[TRTLLM-11228][feat] Support DFlash in one-model spec dec#12794
Merged
Conversation
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
3 times, most recently
from
April 8, 2026 05:59
9e7c58a to
1cc5cef
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #42288 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #42288 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #42342 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #42342 [ run ] completed with state
|
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
from
April 9, 2026 02:53
f614377 to
f9badbf
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #42442 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #42442 [ run ] completed with state
|
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
from
April 9, 2026 10:40
f9badbf to
a4aef3d
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #42530 [ run ] triggered by Bot. Commit: |
Contributor
📝 WalkthroughWalkthroughThis PR introduces DFlash speculative decoding to TensorRT-LLM. It adds a new draft model implementation, speculative metadata and worker classes, configuration support, pipeline integration, and comprehensive tests. DFlash performs cross-attention between draft and target models during generation to efficiently produce token candidates for speculative decoding. Changes
Sequence Diagram(s)sequenceDiagram
participant Target as Target Model
participant SpecWorker as DFlashWorker
participant DraftModel as DFlashForCausalLM
participant Sampler as Token Sampler
rect rgba(0, 100, 200, 0.5)
note right of Target: Context/Prefill Phase
Target->>SpecWorker: forward(hidden_states, ...)
SpecWorker->>SpecWorker: capture hidden states from target_layer_ids
SpecWorker->>DraftModel: load_weights_from_target_model()
SpecWorker->>SpecWorker: project captured states via fc + hidden_norm
SpecWorker->>SpecWorker: append to context buffers
end
rect rgba(100, 200, 0, 0.5)
note right of SpecWorker: Generation Phase (per step)
SpecWorker->>SpecWorker: retrieve draft_tokens from spec_metadata
SpecWorker->>SpecWorker: construct noise embeddings (mask + bonus tokens)
SpecWorker->>DraftModel: dflash_forward(noise_embedding, target_hidden, ...)
DraftModel->>DraftModel: cross-attention: Q from noise, K/V from target context
DraftModel->>DraftModel: per-layer attention + mlp
DraftModel-->>SpecWorker: logits at mask positions
SpecWorker->>Sampler: sample draft tokens
SpecWorker->>Sampler: accept/verify tokens
SpecWorker->>SpecWorker: update context buffers with accepted tokens
SpecWorker-->>SpecWorker: return draft logits & verification results
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (4)
tensorrt_llm/_torch/models/modeling_speculative.py (2)
1034-1034: Rename unused loop variable to suppress linter warning.The static analysis tool flags
layer_idxas unused. Per PEP 8 convention, prefix unused variables with underscore.♻️ Proposed fix
- for layer_idx, layer in enumerate(self.model.layers): + for _layer_idx, layer in enumerate(self.model.layers):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_speculative.py` at line 1034, The for-loop in modeling_speculative.py uses an unused loop variable layer_idx (for layer_idx, layer in enumerate(self.model.layers)), which triggers a linter warning; rename layer_idx to _layer_idx (or simply use _ ) to indicate it's intentionally unused and suppress the warning, updating the loop header in the scope where self.model.layers is iterated (the enumerate loop in the relevant method in the SpeculativeModel class) so existing logic referencing only layer continues to work.
924-932: Consider using the customRMSNormmodule for consistency.The code uses PyTorch's
nn.RMSNormforhidden_norm, while the rest of the codebase uses the customRMSNormfrom..modules.rms_norm. While functionally similar, this could lead to subtle differences in behavior (e.g., quantization support, dtype handling) or maintenance overhead.♻️ Suggested change for consistency
+ from ..modules.rms_norm import RMSNorm + if 'hidden_norm.weight' in remapped: rms_norm_eps = getattr(self.config, 'rms_norm_eps', 1e-6) - self.hidden_norm = nn.RMSNorm( - remapped['hidden_norm.weight'].shape[0], + self.hidden_norm = RMSNorm( + hidden_size=remapped['hidden_norm.weight'].shape[0], eps=rms_norm_eps, - device='cuda', - elementwise_affine=True, dtype=remapped['hidden_norm.weight'].dtype) self.hidden_norm.weight.data.copy_(remapped['hidden_norm.weight']) del remapped['hidden_norm.weight']🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_speculative.py` around lines 924 - 932, The hidden_norm is instantiated with torch.nn.RMSNorm which is inconsistent with the codebase; replace nn.RMSNorm with the project’s custom RMSNorm (imported from ..modules.rms_norm) in the hidden_norm creation flow (the block that checks 'hidden_norm.weight' in remapped), passing the same eps (use self.config.rms_norm_eps or 1e-6), device, elementwise_affine and dtype semantics required by the custom RMSNorm API, then copy weights from remapped['hidden_norm.weight'] into self.hidden_norm.weight.data as before so behavior and dtype/quantization support remain consistent.tensorrt_llm/_torch/speculative/dflash.py (2)
572-576: Fragile deep attribute access - consider defensive accessor.The access
draft_model.draft_model_full.model.embed_tokensassumes a specific structure that may not hold for all draft model types. If the structure differs, this will fail with an unclearAttributeError.Consider adding a helper method or more defensive access:
♻️ Suggested defensive access
- # Get the embed_tokens layer from the draft model - embed_tokens = draft_model.draft_model_full.model.embed_tokens + # Get the embed_tokens layer from the draft model + if hasattr(draft_model, 'draft_model_full'): + embed_tokens = draft_model.draft_model_full.model.embed_tokens + elif hasattr(draft_model, 'model') and hasattr(draft_model.model, 'embed_tokens'): + embed_tokens = draft_model.model.embed_tokens + else: + raise AttributeError( + f"Cannot find embed_tokens in draft model of type {type(draft_model).__name__}" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/speculative/dflash.py` around lines 572 - 576, The direct deep attribute access draft_model.draft_model_full.model.embed_tokens is fragile; implement a defensive accessor (e.g., get_embed_tokens(draft_model)) that attempts the expected paths (draft_model.draft_model_full.model.embed_tokens, draft_model.draft_model_full.embed_tokens, draft_model.model.embed_tokens, draft_model.embed_tokens) using getattr/hasattr checks and returns the found object or raises a clear ValueError mentioning draft_model and available attributes; then replace the direct access in the code with a call to get_embed_tokens(draft_model) before computing hidden_dim so hidden_dim calculation (using spec_metadata.hidden_size or hidden_states.shape[-1]) uses a reliably-resolved embed_tokens.
338-339: Minor performance concern: CUDA synchronization in prefill path.
self._ctx_len[slot].item()causes a CUDA synchronization point. While this is in the prefill path (not the hot generation loop), it could add latency during context processing. Consider keepingcuras a tensor if possible, or batching the operations to minimize sync points.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 1006-1008: The expression computing num_heads_per_rank //
num_kv_heads_per_rank is dead code; assign it to a meaningful variable (e.g.,
gqa_factor = num_heads_per_rank // num_kv_heads_per_rank) and then use
gqa_factor where the generalised query attention factor is intended, or if it
was not needed remove the standalone expression; update references around attn0,
num_heads_per_rank and num_kv_heads_per_rank so the computed GQA factor is
either stored as gqa_factor or the line is deleted to avoid the discarded
calculation.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 3747-3767: The DFlash resolution currently only falls back for
target_layer_ids and assumes speculative_model is a local path; update the logic
in the block handling isinstance(self.speculative_config, DFlashDecodingConfig)
so that both target_layer_ids and mask_token_id are resolved independently from
the draft config (load config.json and set each field if it is None), support
speculative_model values that are HF hub IDs by using huggingface_hub (e.g.,
hf_hub_download or snapshot_dir) to fetch config.json when speculative_model is
not a local directory, and after attempting both fallbacks, raise a clear
exception (fail fast) if either self.speculative_config.target_layer_ids or
self.speculative_config.mask_token_id remain unset. Ensure you reference the
existing symbols speculative_config, DFlashDecodingConfig, target_layer_ids,
mask_token_id, and speculative_model so the change is localized to that branch.
- Around line 1584-1631: DFlashDecodingConfig must validate that
speculative_model is provided and max_draft_len is > 0 at construction; add
checks inside the class (e.g., extend the existing
`@model_validator`(mode="after") method set_max_total_draft_tokens or add a new
`@model_validator`) to raise a clear ValueError if self.speculative_model is None
or if self.max_draft_len is None or <= 0, then set self.max_total_draft_tokens =
self.max_draft_len as before; ensure the error messages reference
DFlashDecodingConfig/speculative_model and max_draft_len so users fail fast
instead of breaking later in TorchLlmArgs or worker setup.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 497-519: The test_dflash currently assumes the speculative_model
checkpoint (dflash_model_dir) exists and will error during LLM(...) setup if
missing; update the test to check for the presence of dflash_model_dir before
constructing LLM and call pytest.skip with a clear message when the
directory/checkpoint is not available so the test cleanly skips; reference the
test_dflash function, the dflash_model_dir variable, and the
LLM/DFlashDecodingConfig.speculative_model when adding the guard.
- Around line 4479-4501: The test_dflash function assumes the local DFlash
artifact exists; before creating dflash_model_dir and using it in
LLM/speculative_config, check for the artifact's presence (e.g., filesystem
existence of dflash_model_dir) and skip the test when it's missing (use
pytest.skip or the test harness's skip/skipTest) so environments without
Qwen3-8B-DFlash-b16 don't hard-fail; place this guard early in test_dflash,
before constructing spec_config or entering the with LLM(...) block.
- Around line 5206-5239: The test_dflash creates a local artifact dependency
dflash_model_dir ("gpt-oss-20b-DFlash") with no skip path; add the same
availability guard used elsewhere to skip the test when that DFlash checkpoint
is missing: before building dflash_model_dir/target_model_dir or instantiating
LLM, check for the presence of the DFlash model (using the existing helper or
simple filesystem check against llm_models_root()/gpt-oss-20b-DFlash) and call
the existing test-skip mechanism if absent, so test_dflash short-circuits rather
than failing when the checkpoint isn't present.
In `@tests/unittest/_torch/speculative/test_dflash.py`:
- Around line 15-22: The test never activates the enforce_single_worker fixture,
so TLLM_WORKER_USE_SINGLE_PROCESS isn't set during test_dflash; update the test
to request the fixture by either adding enforce_single_worker to the test_dflash
signature (def test_dflash(disable_overlap_scheduler: bool,
enforce_single_worker):) or apply
pytest.mark.usefixtures("enforce_single_worker") above test_dflash to ensure the
single-worker env var is set when running DFlash-related code.
- Around line 36-38: The test hardcodes dflash_model_dir to
"/code/llm-models/..." which breaks on different filesystems; instead compute
dflash_model_dir from llm_models_root() like models_path does so it follows the
same root resolution. Update the assignment of dflash_model_dir to use
llm_models_root() (via models_path or calling llm_models_root() directly) and
join the relative segments for "Qwen3-8B-DFlash-b16" (e.g., with os.path.join or
an f-string) so that dflash_model_dir and target_model_dir are derived
consistently from the same root.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Line 1034: The for-loop in modeling_speculative.py uses an unused loop
variable layer_idx (for layer_idx, layer in enumerate(self.model.layers)), which
triggers a linter warning; rename layer_idx to _layer_idx (or simply use _ ) to
indicate it's intentionally unused and suppress the warning, updating the loop
header in the scope where self.model.layers is iterated (the enumerate loop in
the relevant method in the SpeculativeModel class) so existing logic referencing
only layer continues to work.
- Around line 924-932: The hidden_norm is instantiated with torch.nn.RMSNorm
which is inconsistent with the codebase; replace nn.RMSNorm with the project’s
custom RMSNorm (imported from ..modules.rms_norm) in the hidden_norm creation
flow (the block that checks 'hidden_norm.weight' in remapped), passing the same
eps (use self.config.rms_norm_eps or 1e-6), device, elementwise_affine and dtype
semantics required by the custom RMSNorm API, then copy weights from
remapped['hidden_norm.weight'] into self.hidden_norm.weight.data as before so
behavior and dtype/quantization support remain consistent.
In `@tensorrt_llm/_torch/speculative/dflash.py`:
- Around line 572-576: The direct deep attribute access
draft_model.draft_model_full.model.embed_tokens is fragile; implement a
defensive accessor (e.g., get_embed_tokens(draft_model)) that attempts the
expected paths (draft_model.draft_model_full.model.embed_tokens,
draft_model.draft_model_full.embed_tokens, draft_model.model.embed_tokens,
draft_model.embed_tokens) using getattr/hasattr checks and returns the found
object or raises a clear ValueError mentioning draft_model and available
attributes; then replace the direct access in the code with a call to
get_embed_tokens(draft_model) before computing hidden_dim so hidden_dim
calculation (using spec_metadata.hidden_size or hidden_states.shape[-1]) uses a
reliably-resolved embed_tokens.
🪄 Autofix (Beta)
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: Pro
Run ID: b8f46231-d247-45d4-9175-3f4c137cb7d1
📒 Files selected for processing (18)
tensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/speculative/dflash.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_core_sanity.txttests/integration/test_lists/qa/llm_function_l20.txttests/integration/test_lists/qa/llm_function_rtx6k.txttests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/speculative/test_dflash.pytests/unittest/api_stability/references_committed/llm.yaml
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
2 times, most recently
from
April 24, 2026 05:30
069d12c to
6f8c917
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #45332 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #45332 [ run ] completed with state
|
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
from
April 24, 2026 07:10
6f8c917 to
71ed33f
Compare
Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
ziyixiong-nv
force-pushed
the
dev-fxiong-dflash-impl
branch
from
April 24, 2026 07:12
71ed33f to
0966340
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #45365 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #45365 [ run ] completed with state
|
Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #45414 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #45414 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #45465 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #45465 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #45566 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #45566 [ run ] completed with state |
nv-guomingz
reviewed
Apr 27, 2026
nv-guomingz
approved these changes
Apr 27, 2026
lucaslie
approved these changes
Apr 27, 2026
lucaslie
left a comment
lucaslie
left a comment
Contributor
There was a problem hiding this comment.
API changes look good to me
arysef
approved these changes
Apr 27, 2026
4 tasks
yufeiwu-nv
pushed a commit
to yufeiwu-nv/TensorRT-LLM
that referenced
this pull request
May 19, 2026
Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
3 tasks
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
7 participants
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.
Summary by CodeRabbit
New Features
Tests
Description
The PR is to support DFlash (https://arxiv.org/pdf/2602.06036 ) in one-model spec dec.
In the mtbench testiing on Qwen3-8B-FP8 on 8xB200, the perf of DFlash is not as good as PARD, and the AR is also lower than PARD.
Test Coverage
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)
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.