[TRTLLM-14054][perf] Load Qwen3-Next GDN in_proj in dense layout and add a multi-row gated RMSNorm

[kaiyux]

Description

HF gated-delta-net (GDN) checkpoints pack in_proj_qkvz rows per key-head group as [q_g|k_g|v_g|z_g] (and in_proj_ba as [b_g|a_g]), so every step ran a split/reshape kernel (fused_qkvzba_split_reshape_cat, one launch per GDN layer) to rearrange the projection output before the GDN kernels could consume it. This PR:

• Permutes the in_proj rows once at weight-load time into dense per-TP-rank [Q|K|V|Z] / [b|a] layouts (Qwen3NextHfWeightMapper). The projection output components become plain column slices consumed in place, and the per-step split/reshape kernel is deleted along with fix_query_key_value_ordering. The permutation handles row tensors of any dtype (packed FP4 included, via a uint8 view), FP8 2D-block weight_scale_inv tensors at scale-block granularity, per-row scales/bias, and TP sharding (rows are permuted per TP-rank chunk so the column-parallel contiguous row split still hands each rank its own heads).
• Teaches the downstream Triton helpers to read the resulting row-strided views in place — the gating kernels, the conv1d extract/transpose helper, and the QKV split take explicit row strides, so no packing copies reappear.
• Adds a multi-row gated RMSNorm kernel. At GDN decode shapes (thousands of short 128-wide rows) the single-row kernel is launch-limited; processing 4 rows per CTA with the same per-row reduction order is bitwise identical to the single-row kernel and roughly halves the kernel time (5.43 us → 3.11 us at batch 256). The gate z is read token-major directly from the projection slice.

The transformation is bitwise-lossless by construction: the same values flow through the same kernels, only their layout in the projection output changes. Net effect at concurrency 256: ~185 us and 30 kernel launches saved per decode step.

Test Coverage

• tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py extended: the row/scale-block permutation against the grouped reference (multiple head configs and TP sizes, quantized dtypes, FP8 scale-block consistency), strided column-slice inputs for the gating / transpose / split kernels against packed inputs (bitwise), and the multi-row gated RMSNorm against the reference and (bitwise) against the single-row path.
• Perf: trtllm-serve + benchmark_serving on Qwen3.6-35B-A3B-NVFP4 (hybrid GDN + MoE), 1x B200, ISL 1024 / OSL 6144, random dataset with --ignore-eos, one point per concurrency (paired num_prompts 8...1024), measured at base 8ef7190ebf on top of [TRTLLM-14054][perf] Reduce per-step host preparation overhead in the PyTorch executor decode path #16313:

concurrency	before (tok/s)	this PR (tok/s)	delta
1	345.15	352.14	+2.03%
2	638.76	653.96	+2.38%
4	1166.02	1197.43	+2.69%
8	2080.29	2152.15	+3.45%
16	3711.45	3774.30	+1.69%
32	6256.17	6334.38	+1.25%
64	10079.27	10273.00	+1.92%
128	15270.86	15695.72	+2.78%
256	20685.34	21454.05	+3.72%

Output throughput improves +2.44% on average across the 9-point curve (worst point +1.25%, no regressions). Independent of #16313 (no shared files); together the two compose to about +5.6% over the common baseline on this harness.

PR Checklist

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

🤖 Generated with Claude Code

Summary by CodeRabbit

• Performance

  • Improved Qwen3Next checkpoint loading and tensor layout handling, including quantized formats.
  • Optimized gated normalization and GDN processing for eligible workloads.
  • Improved kernel performance for strided tensor views.

• Bug Fixes

  • Fixed tensor ordering and scaling during distributed Qwen3Next execution.
  • Improved handling of non-contiguous inputs while preserving expected results.

• Tests

  • Added coverage for quantized layouts, distributed tensor permutations, optimized normalization, and strided inputs.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Qwen3Next GDN optimization

Layer / File(s)	Summary
Dense in-projection mapping tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py, tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py	Adds TP-aware grouped-to-dense row permutations for QKVZ and BA tensors, including packed dtypes and FP8 scale blocks.
Strided transpose and QKV kernels tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py, tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py	Updates Triton indexing and launch arguments to support arbitrary row strides while requiring contiguous final dimensions.
GDN token preparation and gating tensorrt_llm/_torch/modules/mamba/gdn_mixer.py, tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py	Slices dense per-rank projections directly, removes the prior ordering helper, and adds row-stride-aware fused gating kernels.
Multi-row gated RMSNorm dispatch tensorrt_llm/_torch/modules/mamba/layernorm_gated.py, tensorrt_llm/_torch/modules/mamba/gdn_mixer.py, tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py	Adds an eligible multi-row gated RMSNorm kernel path and routes GDN postprocessing through token-major normalization.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant HFCheckpoint
  participant WeightMapper
  participant GDNMixer
  participant TritonKernels
  participant GatedRMSNorm
  HFCheckpoint->>WeightMapper: load grouped in_proj tensors
  WeightMapper->>GDNMixer: provide dense per-rank QKVZ and BA tensors
  GDNMixer->>TritonKernels: prepare strided token inputs and gating values
  TritonKernels->>GatedRMSNorm: provide GDN output and gate tensor
  GatedRMSNorm->>GDNMixer: return normalized token-major output

Loading

Suggested reviewers: dc3671, pengbowang-nv, liji-nv, longlee0622, mingyanghao

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific and accurately summarizes the main performance and layout changes in the PR.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections with clear, relevant details.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/modules/mamba/layernorm_gated.py (1)

167-188: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Eligibility logic is duplicated and slightly divergent across the two entry points.

rms_norm_gated_token_major inlines its own eligibility check (x.stride(-1)==1 and z.stride(2)==1 and z.stride(1)==N and N<=_MULTIROW_MAX_N and pow2), while _layer_norm_fwd uses _multirow_gated_rmsnorm_eligible(...) (which omits the stride checks, relying on upstream asserts). Consider routing both through a single predicate so the constraints can't drift apart as the kernel evolves.

🤖 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/layernorm_gated.py` around lines 167 - 188,
Unify multi-row gated RMSNorm eligibility by extending or reusing
_multirow_gated_rmsnorm_eligible for the stride constraints currently inlined by
rms_norm_gated_token_major. Update both rms_norm_gated_token_major and
_layer_norm_fwd to use this single predicate, preserving their existing
eligibility inputs and ensuring all required shape, stride, normalization,
gating, and power-of-two constraints remain enforced.

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

Nitpick comments:
In `@tensorrt_llm/_torch/modules/mamba/layernorm_gated.py`:
- Around line 167-188: Unify multi-row gated RMSNorm eligibility by extending or
reusing _multirow_gated_rmsnorm_eligible for the stride constraints currently
inlined by rms_norm_gated_token_major. Update both rms_norm_gated_token_major
and _layer_norm_fwd to use this single predicate, preserving their existing
eligibility inputs and ensuring all required shape, stride, normalization,
gating, and power-of-two constraints remain enforced.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3c527273-05d0-4f35-99fd-5ff8101aeca9

📥 Commits

Reviewing files that changed from the base of the PR and between 53d2bde and 9d0ff11.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py
• tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py
• tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
• tensorrt_llm/_torch/modules/mamba/layernorm_gated.py
• tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py

[kaiyux]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58951 [ run ] triggered by Bot. Commit: 9d0ff11 Link to invocation

[tensorrt-cicd]

PR_Github #58951 [ run ] completed with state SUCCESS. Commit: 9d0ff11
/LLM/main/L0_MergeRequest_PR pipeline #47483 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

[kaiyux]

/bot run --disable-fail-fast

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
pytorch	CVE-2025-3000	A vulnerability classified as critical has been found in PyTorch 2.6.0. This affects the function torch.jit.script. The manipulation leads to memory corruption. It is possible to launch the attack on the local host. The exploit has been disclosed to the public and may be used.	MEDIUM

[kaiyux]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59065 [ run ] triggered by Bot. Commit: 9d0ff11 Link to invocation

[tensorrt-cicd]

PR_Github #59065 [ run ] completed with state SUCCESS. Commit: 9d0ff11
/LLM/main/L0_MergeRequest_PR pipeline #47588 completed with status: 'SUCCESS'

CI Report

Link to invocation