[TRTLLM-14345][feat] Support GDN MTP Replay#16464
Conversation
Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com>
JadoTu
changed the title
📝 WalkthroughWalkthroughAdds Triton cached-replay kernels for GDN verification and history commits, wires them into Qwen3-Next speculative decoding, adds runtime gating and cache-manager coordination, and validates numerical correctness, state updates, invalid slots, and CUDA-graph execution. ChangesGDN cached replay
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant ForwardCore
participant GDNMixer
participant ReplayUpdate
participant MambaCacheManager
participant LayerCommit
ForwardCore->>GDNMixer: pass replay metadata and cached state
GDNMixer->>ReplayUpdate: run cached recurrent verification
ReplayUpdate-->>GDNMixer: return verification output
MambaCacheManager->>LayerCommit: commit replay history layers
LayerCommit->>MambaCacheManager: update SSM state buffers
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tensorrt_llm/_torch/modules/fla/cached_replay.py (2)
745-745: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winWrite replay results directly into the provided output buffer.
The current hot path always allocates a full output tensor and then performs a second full-size device copy.
tensorrt_llm/_torch/modules/fla/cached_replay.py#L745-L745: accept an optional output tensor and launch directly into it.tensorrt_llm/_torch/modules/mamba/gdn_mixer.py#L518-L520: passoutput_dthrough and remove the post-kernel copy.🤖 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/modules/fla/cached_replay.py` at line 745, Update tensorrt_llm/_torch/modules/fla/cached_replay.py lines 745-745 to accept an optional output tensor and use it as the replay destination, allocating a new tensor only when none is provided. Update tensorrt_llm/_torch/modules/mamba/gdn_mixer.py lines 518-520 to pass output_d into the replay call and remove the subsequent full-size device copy.
494-508: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public replay interfaces and tensor contracts.
These exported functions omit argument shapes, dtypes, mutation behavior, and return semantics despite relying on strict CUDA layouts.
As per coding guidelines, “document public function arguments, and include tensor dimensions and constrained dtypes where applicable.”
Also applies to: 619-650
🤖 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/modules/fla/cached_replay.py` around lines 494 - 508, Document the public replay interfaces, including commit_gdn_cached_replay_history_layers and the functions around its related replay entry points, with each argument’s tensor shape, required dtype, mutation behavior, and return semantics. Specify constrained layouts and scalar/configuration types for parameters such as replay_work_items, n_writes, history_size, and commit tuning options, while preserving the existing API behavior.Source: Coding guidelines
tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py (1)
379-513: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftExtend coverage to every enabled dtype and the cache-manager dispatch.
Kernel-level bfloat16 coverage is sufficient, but end-to-end coverage is insufficient:
- Parametrize this file’s all-layer commit and CUDA-graph tests for
float16,bfloat16, andfloat32, matching_util.py.- Add
tests/unittest/_torch/pyexecutor/test_mamba_cache_manager.pycases around 15/16 decodes to verify commit ordering, PNAT/buffer flips, and disabled replay.As per path instructions, test feedback must identify concrete files and state whether coverage is sufficient.
🤖 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/modules/mamba/test_gdn_replay_recurrent.py` around lines 379 - 513, Extend the relevant all-layer commit and CUDA-graph tests, including test_gdn_cached_replay_dispatch_cuda_graph_matches_eager, with a dtype parameter covering float16, bfloat16, and float32 as defined in _util.py, and create cache-manager tests around the 15/16 decode boundary in test_mamba_cache_manager.py. Verify commit ordering, PNAT and buffer-index flips, and behavior when replay is disabled; report concrete test files and whether the resulting coverage is sufficient.Source: Path instructions
tensorrt_llm/_torch/modules/mamba/gdn_mixer.py (1)
457-475: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the replay verification contract.
Add precise tensor, metadata, optional-buffer, integer, boolean, and return types. This method crosses kernel, cache-manager, and mixer boundaries, so untyped parameters obscure its strict contract.
As per coding guidelines, “Annotate every function” and avoid imprecise types.
🤖 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/modules/mamba/gdn_mixer.py` around lines 457 - 475, Annotate the complete signature of _replay_verify_recurrent with precise types for all tensor inputs, replay metadata, cache/work-item objects, integer and boolean arguments, optional buffers, and the return value. Reuse established tensor, metadata, cache-manager, and mixer type aliases from the surrounding module or project, and express optional parameters explicitly rather than using imprecise types.Source: Coding guidelines
🤖 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/pyexecutor/_util.py`:
- Around line 1996-2000: Require TRTLLM_USE_GDN_REPLAY to equal exactly "1"
before enabling replay in the _util.py replay configuration, treating all other
values as disabled. Apply the same exact-value parsing and consistent disabled
logging in gdn_mixer.py so both sites share identical opt-in behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fla/cached_replay.py`:
- Line 745: Update tensorrt_llm/_torch/modules/fla/cached_replay.py lines
745-745 to accept an optional output tensor and use it as the replay
destination, allocating a new tensor only when none is provided. Update
tensorrt_llm/_torch/modules/mamba/gdn_mixer.py lines 518-520 to pass output_d
into the replay call and remove the subsequent full-size device copy.
- Around line 494-508: Document the public replay interfaces, including
commit_gdn_cached_replay_history_layers and the functions around its related
replay entry points, with each argument’s tensor shape, required dtype, mutation
behavior, and return semantics. Specify constrained layouts and
scalar/configuration types for parameters such as replay_work_items, n_writes,
history_size, and commit tuning options, while preserving the existing API
behavior.
In `@tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`:
- Around line 457-475: Annotate the complete signature of
_replay_verify_recurrent with precise types for all tensor inputs, replay
metadata, cache/work-item objects, integer and boolean arguments, optional
buffers, and the return value. Reuse established tensor, metadata,
cache-manager, and mixer type aliases from the surrounding module or project,
and express optional parameters explicitly rather than using imprecise types.
In `@tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py`:
- Around line 379-513: Extend the relevant all-layer commit and CUDA-graph
tests, including test_gdn_cached_replay_dispatch_cuda_graph_matches_eager, with
a dtype parameter covering float16, bfloat16, and float32 as defined in
_util.py, and create cache-manager tests around the 15/16 decode boundary in
test_mamba_cache_manager.py. Verify commit ordering, PNAT and buffer-index
flips, and behavior when replay is disabled; report concrete test files and
whether the resulting coverage is sufficient.
🪄 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: Enterprise
Run ID: 9296ea1d-a7b3-4ca5-a1ef-28911b44bcfa
📒 Files selected for processing (5)
tensorrt_llm/_torch/modules/fla/cached_replay.pytensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py
nv-guomingz
changed the title
Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com>
Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com>
Signed-off-by: Jian Tu <107457950+JadoTu@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59976 [ run ] triggered by Bot. Commit: |
|
PR_Github #59976 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60102 [ run ] triggered by Bot. Commit: |
|
PR_Github #60102 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60138 [ run ] triggered by Bot. Commit: |
|
PR_Github #60138 [ run ] completed with state |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Reviewed the GDN MTP cached replay change. The feature is opt-in via TRTLLM_USE_GDN_REPLAY=1 and remains disabled by default; enablement is additionally gated by the speculative/cache/hardware constraints, so the default execution path is not broadened. The earlier env-var ambiguity was fixed by centralizing is_gdn_replay_enabled() as an exact "1" check, with a parameterized test.\n\nThe patch includes meaningful coverage: replay-vs-reference/legacy numerical checks, invalid-slot behavior, all-layer commit reference checks, CUDA-graph-vs-eager checks, and a B200 e2e accuracy entry. Final Blossom CI on the head commit is green and there are already current member approvals. I also ran an independent Claude cross-check before approving; it agreed this is safe to approve.\n\nNon-blocking follow-ups: document the public replay tensor contracts more explicitly, and consider extending dtype/cache-manager boundary coverage as CodeRabbit suggested.
VALLIS-NERIA
left a comment
VALLIS-NERIA
left a comment
There was a problem hiding this comment.
LGTM for tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
8 participants
Summary by CodeRabbit
New Features
Performance
Bug Fixes
Tests
-->
Description
This change adds an opt-in cached replay path for Gated DeltaNet state updates during one-model MTP verification in the PyTorch backend.
Instead of materializing and retaining a full recurrent state for every draft position, cached replay stores a compact update history and reconstructs the accepted checkpoint state only when needed. This reduces MTP state-cache memory and avoids redundant replay computation while preserving the existing state precision.
The implementation includes:
Cached replay is disabled by default because its end-to-end benefit is workload-dependent. Set TRTLLM_USE_GDN_REPLAY=1 to enable it.
The replay function can dramatically saving the buffer of SSM used.
Estimated GDN MTP cache HBM saving
Qwen3.5-397B-A17B-NVFP4 | B200 | TEP4 | T=3 | HIST=16 | BF16 state
end-to-end perf comparisons.
kernel microbenchmarks
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.