[None][fix] Qwen3.5 DFlash#13782
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR updates TensorRT-LLM's speculative decoding (DFlash) to support Mamba models with MRoPE (Multi-headed Rotary Position Embeddings). Changes include adjusting draft token budgeting, adding 3D position_id tensor handling for MRoPE, integrating Mamba cache state synchronization into DFlash, and adding corresponding test coverage. ChangesDFlash with Mamba and MRoPE Support
Sequence DiagramsequenceDiagram
participant User as User/Test
participant LLM as LLM Engine
participant DFlash as DFlash Decoder
participant ModelEng as Model Engine
participant MambaCache as Mamba Cache Manager
participant GDNMixer as GDN Mixer
User->>LLM: Initialize with DFlashDecodingConfig
LLM->>MambaCache: Create with tokens_per_gen_step budget
LLM->>DFlash: Setup speculative decoding
User->>LLM: Generate tokens (with position_ids 3D)
loop Speculative Decoding
DFlash->>ModelEng: _preprocess_inputs (adjust 3D position_ids)
ModelEng->>GDNMixer: forward_decode/forward_extend
GDNMixer->>GDNMixer: Use max_total_draft_tokens for sizing
GDNMixer-->>DFlash: Draft tokens
DFlash->>DFlash: Verify tokens
end
DFlash->>MambaCache: update_mamba_states (after acceptance)
MambaCache->>MambaCache: Sync state with accepted tokens
DFlash->>ModelEng: _postprocess_inputs (adjust 3D position_ids)
DFlash-->>LLM: Accepted tokens + metadata
LLM-->>User: Generated sequence
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py (1)
831-833:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
pool_sizecan be under-allocated for parallel-draft modes.Line 850 now correctly uses
tokens_per_gen_step - 1for speculative state depth, but Line 833 still sizes the Mamba pool frommax_draft_len. When those diverge (parallel-draft modes), dummy-slot reservation can exceed the pool and hitrun out of mamba cache blocks.💡 Suggested fix
- max_draft_len = (spec_config.max_draft_len - if spec_config is not None else 0) - pool_size = max_batch_size + max_draft_len + 1 + if spec_config is not None: + max_runtime_draft_len = max( + spec_config.max_draft_len, + spec_config.tokens_per_gen_step - 1, + ) + else: + max_runtime_draft_len = 0 + pool_size = max_batch_size + max_runtime_draft_len + 1 @@ - speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 - if spec_config is not None else None), + speculative_num_draft_tokens=(max_runtime_draft_len + if spec_config is not None else None),Also applies to: 850-851
🤖 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/pyexecutor/mamba_cache_manager.py` around lines 831 - 833, The pool_size is being computed from max_draft_len but the code elsewhere uses speculative depth = tokens_per_gen_step - 1, so in Mamba cache sizing replace the use of max_draft_len with a computed speculative_depth (e.g., speculative_depth = spec_config.tokens_per_gen_step - 1 if spec_config is not None else 0) and set pool_size = max_batch_size + speculative_depth + 1; also update the other sizing/ reservation sites that currently use max_draft_len (the later usage around the speculative state depth calculation) to use this same speculative_depth variable so dummy-slot reservation cannot exceed the pool.
🧹 Nitpick comments (1)
tests/integration/test_lists/qa/llm_function_core.txt (1)
734-734: Please confirm perf guardrail coverage for this speculative-decoding change.Given the PR modifies performance-sensitive speculative decoding/attention execution paths, confirm there is a corresponding perf scenario in
tests/integration/test_lists/test-db/l0_perf.yml(and QAllm_perf_*lists if needed), or track it as explicit follow-up.As per coding guidelines: "If the PR touches performance-sensitive paths ... check whether a perf test entry is present or updated in (a) tests/integration/test_lists/test-db/l0_perf.yml ... and (b) tests/integration/test_lists/qa/llm_perf_*.yml ..."
🤖 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/integration/test_lists/qa/llm_function_core.txt` at line 734, Confirm that the speculative-decoding/attention execution changes touched by the test accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash are covered by a perf guardrail: check the L0 perf DB entry (l0_perf.yml) and the QA llm_perf_* lists for an existing scenario exercising speculative decoding/attention; if none exists, add a new perf scenario entry that measures latency/throughput for speculative decoding (referencing test_dflash or an equivalent workload) in l0_perf.yml and add/track it in the llm_perf_* QA lists, or create a follow-up tracking ticket if you cannot add it immediately.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`:
- Around line 831-833: The pool_size is being computed from max_draft_len but
the code elsewhere uses speculative depth = tokens_per_gen_step - 1, so in Mamba
cache sizing replace the use of max_draft_len with a computed speculative_depth
(e.g., speculative_depth = spec_config.tokens_per_gen_step - 1 if spec_config is
not None else 0) and set pool_size = max_batch_size + speculative_depth + 1;
also update the other sizing/ reservation sites that currently use max_draft_len
(the later usage around the speculative state depth calculation) to use this
same speculative_depth variable so dummy-slot reservation cannot exceed the
pool.
---
Nitpick comments:
In `@tests/integration/test_lists/qa/llm_function_core.txt`:
- Line 734: Confirm that the speculative-decoding/attention execution changes
touched by the test
accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash are covered by a
perf guardrail: check the L0 perf DB entry (l0_perf.yml) and the QA llm_perf_*
lists for an existing scenario exercising speculative decoding/attention; if
none exists, add a new perf scenario entry that measures latency/throughput for
speculative decoding (referencing test_dflash or an equivalent workload) in
l0_perf.yml and add/track it in the llm_perf_* QA lists, or create a follow-up
tracking ticket if you cannot add it immediately.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 850a5460-d612-4112-800c-f7996b0ef776
📒 Files selected for processing (8)
tensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/speculative/dflash.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/test-db/l0_h100.yml
zheyuf
approved these changes
May 8, 2026
amukkara
force-pushed
the
qwen3.5-dflash
branch
2 times, most recently
from
May 9, 2026 05:23
3ab415f to
35da95c
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #47477 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47477 [ run ] completed with state
|
ziyixiong-nv
added a commit
to ziyixiong-nv/TensorRT-LLM
that referenced
this pull request
May 9, 2026
Cherry-pick of NVIDIA#13782: - Fix mrope position_ids handling in DFlash - Fix GDN spec verify buffer sizing for DFlash - Fix GDN/Mamba2 state updates in DFlash - Add tests Original-author: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
ziyixiong-nv
added a commit
to ziyixiong-nv/TensorRT-LLM
that referenced
this pull request
May 11, 2026
Cherry-pick of NVIDIA#13782: - Fix mrope position_ids in DFlash - Fix GDN verify buffer size for DFlash - Fix GDN/Mamba2 state updates in DFlash - Add tests Original-author: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #47775 [ run ] triggered by Bot. Commit: |
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
Collaborator
Author
|
/bot kill |
Collaborator
|
PR_Github #47810 [ kill ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47810 [ kill ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #47816 [ run ] triggered by Bot. Commit: |
1 task
xinhe-nv
approved these changes
May 12, 2026
amukkara
enabled auto-merge (squash)
May 12, 2026 04:16
ziyixiong-nv
added a commit
to ziyixiong-nv/TensorRT-LLM
that referenced
this pull request
May 12, 2026
Cherry-pick of NVIDIA#13782: - Fix mrope position_ids in DFlash - Fix GDN verify buffer size for DFlash - Fix GDN/Mamba2 state updates in DFlash - Add tests Original-author: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Collaborator
|
PR_Github #47816 [ run ] completed with state
|
Wanli-Jiang
reviewed
May 12, 2026
Wanli-Jiang
reviewed
May 12, 2026
Wanli-Jiang
approved these changes
May 12, 2026
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #48005 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48005 [ run ] completed with state |
yufeiwu-nv
pushed a commit
to yufeiwu-nv/TensorRT-LLM
that referenced
this pull request
May 19, 2026
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
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
8 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
Improvements
Tests
Description
Minor fixes to combine DFlash with m-rope and linear attention in Qwen3.5:
Test Coverage
Added new test
tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflashand extended
tests/unittest/_torch/speculative/test_dflash.pyto test Qwen3.5-4B along with Qwen3-8B.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.