Skip to content

[TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5… - #15194

Merged
nv-guomingz merged 1 commit into
NVIDIA:mainfrom
nv-guomingz:user/guomingz/qwen3next_post_moe_fusion_fix
Jul 24, 2026
Merged

[TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5…#15194
nv-guomingz merged 1 commit into
NVIDIA:mainfrom
nv-guomingz:user/guomingz/qwen3next_post_moe_fusion_fix

Conversation

@nv-guomingz

@nv-guomingz nv-guomingz commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

… 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.

image

Summary by CodeRabbit

  • New Features

    • Added support for Gemma-style RMSNorm scaling in fused all-reduce operations, controllable via environment variable configuration.
  • Refactor

    • Improved weight handling alignment between fused and non-fused computation paths.
    • Enhanced tensor-parallel fusion configuration logic for improved MoE support.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

@nv-guomingz
nv-guomingz requested a review from a team as a code owner June 10, 2026 02:54
@nv-guomingz
nv-guomingz requested a review from 2ez4bz June 10, 2026 02:54
@nv-guomingz nv-guomingz changed the title [None][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5… [TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5… Jun 10, 2026
@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from 6d86597 to 594584f Compare June 10, 2026 02:59
@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This PR adds Gemma-style RMSNorm scaling support to fused AllReduce+RMSNorm kernels. It introduces a use_gemma parameter to the kernel contract, implements conditional (1 + gamma) adjustment in the CUDA kernel, wires environment-variable control in C++, and updates PyTorch decoder layers to pass pre-adjusted weights to fusion paths.

Changes

Gemma AllReduce+RMSNorm Fusion Support

Layer / File(s) Summary
Kernel parameter contract and CUDA implementation
cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h, cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
AllReduceFusionParams adds use_gemma flag (defaults false). CUDA rms_norm method now conditionally applies 1.0f adjustment to gamma when flag is set before scaling the residual output.
C++ environment variable integration
cpp/tensorrt_llm/thop/allreduceOp.cpp
Adds environment utilities header and wires TLLM_QWEN3_AR_FUSION_GEMMA_NORM env-var to set use_gemma flag in the all-reduce fusion execution path.
PyTorch weight adjustment helper and decoder integration
tensorrt_llm/_torch/models/modeling_qwen3_next.py
New _fused_norm_weight() helper bakes Gemma's (1 + weight) into RMSNorm weights (fp32 computation, cast back to dtype). Linear-attention and full-attention decoder layers pass adjusted weights to fused all-reduce paths instead of raw weights.
Full-attention fusion configuration and forward finalization logic
tensorrt_llm/_torch/models/modeling_qwen3_next.py
Full-attention decoder reworks POST_MOE fusion config logic to depend on PRE_MOE_FUSION, pipeline parallelism, and tensor-parallel rank (not enable_attention_dp). Forward path forces do_finalize = True when PRE_MOE_FUSION is enabled to avoid unimplemented fully-fused MoE flow.

Sequence Diagram

sequenceDiagram
  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
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 10.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check ⚠️ Warning The template sections for Description and Test Coverage are mostly empty, so the PR rationale and validation details are incomplete. Add a short issue/solution summary under Description and list the relevant tests or validation steps under Test Coverage.
✅ Passed checks (3 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title matches the main change: fusing Gemma-aware RMSNorm into AllReduce for Qwen3-Next/Qwen3.5.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 8e40515 and 6d86597.

📒 Files selected for processing (4)
  • cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
  • cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h
  • cpp/tensorrt_llm/thop/allreduceOp.cpp
  • tensorrt_llm/_torch/models/modeling_qwen3_next.py

Comment thread cpp/tensorrt_llm/thop/allreduceOp.cpp Outdated
Comment thread tensorrt_llm/_torch/models/modeling_qwen3_next.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53207 [ run ] triggered by Bot. Commit: 594584f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53207 [ run ] completed with state SUCCESS. Commit: 594584f
/LLM/main/L0_MergeRequest_PR pipeline #42403 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from 594584f to 6ddda0f Compare June 10, 2026 13:06
@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53317 [ run ] triggered by Bot. Commit: 6ddda0f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53317 [ run ] completed with state FAILURE. Commit: 6ddda0f
/LLM/main/L0_MergeRequest_PR pipeline #42503 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from 6ddda0f to eb3f9c5 Compare July 16, 2026 05:39
@nv-guomingz
nv-guomingz requested review from a team as code owners July 16, 2026 05:39
@nv-guomingz
nv-guomingz requested a review from zongfeijing July 16, 2026 05:39
@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch 2 times, most recently from 65e3a37 to f7108ea Compare July 16, 2026 07:54
@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59666 [ run ] triggered by Bot. Commit: f7108ea Link to invocation

@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from f7108ea to 0046ca9 Compare July 16, 2026 14:36
@nv-guomingz
nv-guomingz requested a review from a team as a code owner July 16, 2026 14:36
@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59726 [ run ] triggered by Bot. Commit: 0046ca9 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59666 [ run ] completed with state ABORTED. Commit: f7108ea

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59726 [ run ] completed with state FAILURE. Commit: 0046ca9
/LLM/main/L0_MergeRequest_PR pipeline #48153 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61000 [ run ] triggered by Bot. Commit: c5f404b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61000 [ run ] completed with state SUCCESS. Commit: c5f404b
/LLM/main/L0_MergeRequest_PR pipeline #49258 completed with status: 'SUCCESS'

CI Report

Link to invocation

@2ez4bz 2ez4bz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving modeling side changes.

Comment thread tensorrt_llm/_torch/models/modeling_qwen3_next.py
Comment thread tensorrt_llm/_torch/models/modeling_qwen3_next.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_qwen3_next.py
Comment thread tensorrt_llm/_torch/models/modeling_qwen3_next.py
Comment thread tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py Outdated
Comment thread tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py Outdated
Comment thread tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py Outdated
Comment thread tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py
Comment thread tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py
@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from c5f404b to f2c0d15 Compare July 23, 2026 06:33
…/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>
@nv-guomingz
nv-guomingz force-pushed the user/guomingz/qwen3next_post_moe_fusion_fix branch from f2c0d15 to 68f89bf Compare July 23, 2026 06:34
@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61214 [ run ] triggered by Bot. Commit: 68f89bf Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61214 [ run ] completed with state FAILURE. Commit: 68f89bf
/LLM/main/L0_MergeRequest_PR pipeline #49455 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61300 [ run ] triggered by Bot. Commit: 68f89bf Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61300 [ run ] completed with state FAILURE. Commit: 68f89bf
/LLM/main/L0_MergeRequest_PR pipeline #49531 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@nv-guomingz

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61492 [ run ] triggered by Bot. Commit: 68f89bf Link to invocation

@nv-guomingz
nv-guomingz enabled auto-merge (squash) July 24, 2026 04:48
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61492 [ run ] completed with state SUCCESS. Commit: 68f89bf
/LLM/main/L0_MergeRequest_PR pipeline #49711 completed with status: 'SUCCESS'

CI Report

Link to invocation

@nv-guomingz
nv-guomingz merged commit 1fae43c into NVIDIA:main Jul 24, 2026
8 checks passed
yuanjingx87 pushed a commit to yuanjingx87/TensorRT-LLM that referenced this pull request Jul 26, 2026
…/Qwen3.5… (NVIDIA#15194)

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
Navjot10 added a commit to Navjot10/TensorRT-LLM that referenced this pull request Jul 29, 2026
…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>
Navjot10 added a commit to Navjot10/TensorRT-LLM that referenced this pull request Jul 30, 2026
…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>
Wanli-Jiang added a commit to Wanli-Jiang/TensorRT-LLM that referenced this pull request Aug 17, 2026
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>
Wanli-Jiang added a commit to Wanli-Jiang/TensorRT-LLM that referenced this pull request Aug 17, 2026
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>
@nv-guomingz
nv-guomingz deleted the user/guomingz/qwen3next_post_moe_fusion_fix branch August 28, 2026 02:43
Navjot10 added a commit to Navjot10/TensorRT-LLM that referenced this pull request Sep 4, 2026
…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>
Navjot10 added a commit to Navjot10/TensorRT-LLM that referenced this pull request Sep 4, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants