[TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5… - #15194
Conversation
6d86597 to
594584f
Compare
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR adds Gemma-style RMSNorm scaling support to fused AllReduce+RMSNorm kernels. It introduces a ChangesGemma AllReduce+RMSNorm Fusion Support
Sequence DiagramsequenceDiagram
participant PyTorchModel as PyTorch Decoder
participant EnvVarControl as Environment Variable
participant CppSetup as C++ All-Reduce Setup
participant CudaKernel as CUDA Kernel
PyTorchModel->>PyTorchModel: _fused_norm_weight() adjusts weight by (1 + w)
PyTorchModel->>CppSetup: call fused all-reduce with adjusted weight
EnvVarControl->>CppSetup: TLLM_QWEN3_AR_FUSION_GEMMA_NORM controls use_gemma flag
CppSetup->>CudaKernel: launch with use_gemma parameter
CudaKernel->>CudaKernel: if use_gemma: gamma = gamma + 1.0
CudaKernel->>CudaKernel: output = residual * scale * adjusted_gamma
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@cpp/tensorrt_llm/thop/allreduceOp.cpp`:
- Around line 688-691: The CUDA path is double-applying Gemma scaling because
_fused_norm_weight() in tensorrt_llm/_torch/models/modeling_qwen3_next.py
already folds Gemma as (1+weight) while allreduce_fusion_params.use_gemma is
being set from TLLM_QWEN3_AR_FUSION_GEMMA_NORM; fix by disabling the CUDA-side
Gemma folding here (set allreduce_fusion_params.use_gemma = false) so weights
are treated as pre-baked, and add a clear comment referencing _fused_norm_weight
and the env var to document the single-contract decision; alternatively, if you
prefer CUDA to apply Gemma, change the Python _fused_norm_weight to return raw
weights instead—pick one approach and make both sides consistent.
In `@tensorrt_llm/_torch/models/modeling_qwen3_next.py`:
- Around line 63-76: Replace the raw layernorm weight used in the linear PRE_MOE
fusion with the helper so Gemma scaling is consistent: where the fusion
currently passes norm_weight=self.post_attention_layernorm.weight, call
_fused_norm_weight(self.post_attention_layernorm) instead; ensure the PRE_MOE
fusion site uses the same helper as other fused attention paths (preserving
dtype handling and use_gemma checks performed inside _fused_norm_weight).
🪄 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: 06539bb4-81ce-4fd1-93bc-bf7fd36de495
📒 Files selected for processing (4)
cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cucpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.hcpp/tensorrt_llm/thop/allreduceOp.cpptensorrt_llm/_torch/models/modeling_qwen3_next.py
|
PR_Github #53207 [ run ] triggered by Bot. Commit: |
|
PR_Github #53207 [ run ] completed with state
|
594584f to
6ddda0f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53317 [ run ] triggered by Bot. Commit: |
|
PR_Github #53317 [ run ] completed with state
|
6ddda0f to
eb3f9c5
Compare
65e3a37 to
f7108ea
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59666 [ run ] triggered by Bot. Commit: |
f7108ea to
0046ca9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59726 [ run ] triggered by Bot. Commit: |
|
PR_Github #59666 [ run ] completed with state |
|
PR_Github #59726 [ run ] completed with state
|
|
PR_Github #61000 [ run ] triggered by Bot. Commit: |
|
PR_Github #61000 [ run ] completed with state |
2ez4bz
left a comment
There was a problem hiding this comment.
Approving modeling side changes.
c5f404b to
f2c0d15
Compare
…/Qwen3.5 TEP Enable eager AllReduce + RMSNorm fusion on the full-attention and GDN decoder layers in the tensor-parallel, non-attention-DP path. Defer each block reduction to a single fused AllReduce owner and keep MTP on its shared-head norm path. Precompute the Gemma weight offset in cache_derived_state so custom, NCCL, and symmetric-memory backends use one consistent norm contract across regular, GMS, staged, and reload loading paths. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
f2c0d15 to
68f89bf
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61214 [ run ] triggered by Bot. Commit: |
|
PR_Github #61214 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61300 [ run ] triggered by Bot. Commit: |
|
PR_Github #61300 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #61492 [ run ] triggered by Bot. Commit: |
|
PR_Github #61492 [ run ] completed with state |
…/Qwen3.5… (NVIDIA#15194) Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
…er GDN decode _flashinfer_gdn_decode passed initial_state_indices to the FlashInfer CuTe-DSL kernel through .int(), which is a no-op on an already-int32 tensor and keeps the pointer of a sliced view (e.g. the decode half state_indices[num_prefills:] of a mixed batch, offset 4*num_prefills bytes). The kernel asserts 32-byte data alignment on every tensor argument and rejects such views at runtime with 'Misaligned Tensor data on argument ... expected data alignment=32 bytes'; the dispatch gate does not check alignment, so there is no Triton fallback. Clone the index tensor into fresh, allocator-aligned storage when (and only when) it is misaligned, exactly like _flashinfer_gdn_verify (NVIDIA#15975) and the a/b activation guards in this same function (NVIDIA#15194). Adds a CUDA regression test mirroring test_fi_mtp_verify_misaligned_index_slice for the decode entry point. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
…er GDN decode _flashinfer_gdn_decode passed initial_state_indices to the FlashInfer CuTe-DSL kernel through .int(), which is a no-op on an already-int32 tensor and keeps the pointer of a sliced view (e.g. the decode half state_indices[num_prefills:] of a mixed batch, offset 4*num_prefills bytes). The kernel asserts 32-byte data alignment on every tensor argument and rejects such views at runtime with 'Misaligned Tensor data on argument ... expected data alignment=32 bytes'; the dispatch gate does not check alignment, so there is no Triton fallback. Clone the index tensor into fresh, allocator-aligned storage when (and only when) it is misaligned, exactly like _flashinfer_gdn_verify (NVIDIA#15975) and the a/b activation guards in this same function (NVIDIA#15194). Adds a CUDA regression test mirroring test_fi_mtp_verify_misaligned_index_slice for the decode entry point. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
Qwen3NextMTP.forward accepted all_rank_num_tokens and dropped it with `del`, so the MoE inside a draft iteration kept reading the target forward's per-rank token counts off attn_metadata. The draft loop rewrites the sequence layout in place after the first step (eagle3.py:1110 fills seq_lens with 1), so from step 1 the tensor carries one row per sequence while the counts still describe the target's max_draft_len + 1 rows per sequence. Step 0 is unaffected. spec_metadata.num_tokens is deliberately not reduced to the subseq shape for MTP Eagle (eagle3.py:519, "keep the 1st-iter shape (matches input_ids)"), so the value already on attn_metadata describes that step. With max_draft_len == 1 there is no divergence at all. The consumer decides what the staleness costs. AllGatherReduceScatter passes the list straight to allgather as exact per-rank sizes, and _allgather asserts input.shape[dim] == sizes[rank] (ops.py:247) before any NCCL call, so every rank raises AssertionError during dispatch: attention DP plus MTP Eagle with max_draft_len > 1 cannot run on that strategy at all. That combination is not exotic, because TRTLLM_FORCE_COMM_METHOD=ALLGATHER is the standard workaround for DeepEP and NVLink one-sided issues, so the fallback is broken exactly when it is needed. The alltoall strategies instead treat the value as an upper bound: reducescatter is skipped entirely (interface.py:1175) and NVLinkOneSided takes only max() for workspace capacity, so those paths already computed the correct result. What they pay is wasted work. calculate_num_chunks derives the chunk count from len(list) * max(list), so a stale value inflates it by up to max_draft_len + 1, and each surplus chunk is substituted with chunk 0's tokens (moe_scheduler.py:659) and run through a full dispatch, expert GEMM and combine before its result is discarded. The saving only materialises once the padded row count crosses moe_max_num_tokens, so it is configuration dependent and is not measured here. On SM120, where use_dp_padding is enabled, the same stale value truncates the output (moe_scheduler.py:216) and returns the target row count instead of the draft's. No numerical result changes on any validated deployment. Install the draft distribution for the duration of the call and restore the target value in `finally`, so an exception cannot leak the draft counts into later forwards. Passing None leaves the metadata untouched, which keeps every non-ADP and non-draft caller unchanged. This closes a leg left open by NVIDIA#12353 ("Merge Eagle3 and MTP-eagle one-model workers"), which moved save/restore from the caller to the callee: the Eagle3 leg gained the try/finally in modeling_speculative.py, but the MTP Eagle leg lands in modeling_qwen3_next.py, which that PR did not touch. The block added here has the same shape as modeling_speculative.py:426 and :681. test_qwen3_next_eager_fusion.py arrived with NVIDIA#15194 and has never appeared in a test list, so the cases added here would not have run in CI. Register the file in l0_a10.yml next to test_qwen3_next_moe_quant.py, which the comment above that entry already marks as the parking spot for CPU-only unit tests with no dedicated job. The CPU stage is not an option: it runs `-m cpu_only`, and conftest.py:228 skips any test file carrying no pytest.mark.cpu_only marker, which fails the entry with exit code 5 for collecting nothing. Marking the file is not the alternative it looks like. L0_Test.groovy:1578 hands every stage whose name does not start with CPU- the opposite expression, `-m 'not cpu_only'`, so a marked file would deselect all six cases on l0_a10 and fail there in exactly the same way. The two options are mutually exclusive: either the file carries the marker and is listed on the CPU stage, or it stays unmarked and is listed on a GPU stage. This takes the second. Both directions were checked against the file: `-m 'not cpu_only'` collects and passes all six, `-m cpu_only` deselects all six and exits 5. All six, including the four that predate this change, also pass with CUDA_VISIBLE_DEVICES empty. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Qwen3NextMTP.forward accepted all_rank_num_tokens and dropped it with `del`, so the MoE inside a draft iteration kept reading the target forward's per-rank token counts off attn_metadata. The draft loop rewrites the sequence layout in place after the first step (eagle3.py:1110 fills seq_lens with 1), so from step 1 the tensor carries one row per sequence while the counts still describe the target's max_draft_len + 1 rows per sequence. Step 0 is unaffected. spec_metadata.num_tokens is deliberately not reduced to the subseq shape for MTP Eagle (eagle3.py:519, "keep the 1st-iter shape (matches input_ids)"), so the value already on attn_metadata describes that step. With max_draft_len == 1 there is no divergence at all. The consumer decides what the staleness costs. AllGatherReduceScatter passes the list straight to allgather as exact per-rank sizes, and _allgather asserts input.shape[dim] == sizes[rank] (ops.py:247) before any NCCL call, so every rank raises AssertionError during dispatch: attention DP plus MTP Eagle with max_draft_len > 1 cannot run on that strategy at all. That combination is not exotic, because TRTLLM_FORCE_COMM_METHOD=ALLGATHER is the standard workaround for DeepEP and NVLink one-sided issues, so the fallback is broken exactly when it is needed. The alltoall strategies instead treat the value as an upper bound: reducescatter is skipped entirely (interface.py:1175) and NVLinkOneSided takes only max() for workspace capacity, so those paths already computed the correct result. What they pay is wasted work. calculate_num_chunks derives the chunk count from len(list) * max(list), so a stale value inflates it by up to max_draft_len + 1, and each surplus chunk is substituted with chunk 0's tokens (moe_scheduler.py:659) and run through a full dispatch, expert GEMM and combine before its result is discarded. The saving only materialises once the padded row count crosses moe_max_num_tokens, so it is configuration dependent and is not measured here. On SM120, where use_dp_padding is enabled, the same stale value truncates the output (moe_scheduler.py:216) and returns the target row count instead of the draft's. No numerical result changes on any validated deployment. Install the draft distribution for the duration of the call and restore the target value in `finally`, so an exception cannot leak the draft counts into later forwards. Passing None leaves the metadata untouched, which keeps every non-ADP and non-draft caller unchanged. This closes a leg left open by NVIDIA#12353 ("Merge Eagle3 and MTP-eagle one-model workers"), which moved save/restore from the caller to the callee: the Eagle3 leg gained the try/finally in modeling_speculative.py, but the MTP Eagle leg lands in modeling_qwen3_next.py, which that PR did not touch. The block added here has the same shape as modeling_speculative.py:426 and :681. test_qwen3_next_eager_fusion.py arrived with NVIDIA#15194 and has never appeared in a test list, so the cases added here would not have run in CI. Register the file in l0_a10.yml next to test_qwen3_next_moe_quant.py, which the comment above that entry already marks as the parking spot for CPU-only unit tests with no dedicated job. The CPU stage is not an option: it runs `-m cpu_only`, and conftest.py:228 skips any test file carrying no pytest.mark.cpu_only marker, which fails the entry with exit code 5 for collecting nothing. Marking the file is not the alternative it looks like. L0_Test.groovy:1578 hands every stage whose name does not start with CPU- the opposite expression, `-m 'not cpu_only'`, so a marked file would deselect all six cases on l0_a10 and fail there in exactly the same way. The two options are mutually exclusive: either the file carries the marker and is listed on the CPU stage, or it stays unmarked and is listed on a GPU stage. This takes the second. Both directions were checked against the file: `-m 'not cpu_only'` collects and passes all six, `-m cpu_only` deselects all six and exits 5. All six, including the four that predate this change, also pass with CUDA_VISIBLE_DEVICES empty. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…er GDN decode _flashinfer_gdn_decode passed initial_state_indices to the FlashInfer CuTe-DSL kernel through .int(), which is a no-op on an already-int32 tensor and keeps the pointer of a sliced view (e.g. the decode half state_indices[num_prefills:] of a mixed batch, offset 4*num_prefills bytes). The kernel asserts 32-byte data alignment on every tensor argument and rejects such views at runtime with 'Misaligned Tensor data on argument ... expected data alignment=32 bytes'; the dispatch gate does not check alignment, so there is no Triton fallback. Clone the index tensor into fresh, allocator-aligned storage when (and only when) it is misaligned, exactly like _flashinfer_gdn_verify (NVIDIA#15975) and the a/b activation guards in this same function (NVIDIA#15194). Adds a CUDA regression test mirroring test_fi_mtp_verify_misaligned_index_slice for the decode entry point. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
…er GDN decode _flashinfer_gdn_decode passed initial_state_indices to the FlashInfer CuTe-DSL kernel through .int(), which is a no-op on an already-int32 tensor and keeps the pointer of a sliced view (e.g. the decode half state_indices[num_prefills:] of a mixed batch, offset 4*num_prefills bytes). The kernel asserts 32-byte data alignment on every tensor argument and rejects such views at runtime with 'Misaligned Tensor data on argument ... expected data alignment=32 bytes'; the dispatch gate does not check alignment, so there is no Triton fallback. Clone the index tensor into fresh, allocator-aligned storage when (and only when) it is misaligned, exactly like _flashinfer_gdn_verify (NVIDIA#15975) and the a/b activation guards in this same function (NVIDIA#15194). Adds a CUDA regression test mirroring test_fi_mtp_verify_misaligned_index_slice for the decode entry point. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
… TEP
Enable the eager AllReduce + RMSNorm fusion on the Qwen3-Next/Qwen3.5 full-attention decoder layers in the tensor-parallel (non attention_dp) path, using a gemma-aware fused kernel.
Summary by CodeRabbit
New Features
Refactor
Description
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)
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.