[TRTLLM-11585][feat] Add CUTEDSL moe backend for nemotron-h#12884
Conversation
015cca8 to
a573d2d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43962 [ run ] triggered by Bot. Commit: |
ee5b337 to
da25caa
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43973 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR generalizes a SwiGLU-specific gather-grouped GEMM kernel to support both gated (SwiGLU) and non-gated (ReLU²) activation modes. Custom ops are renamed from Changes
Sequence Diagram(s)sequenceDiagram
participant MoE as MoE Backend<br/>(FusedMoE)
participant QT as Quantization<br/>Transform
participant COP as Custom Op<br/>(act_fusion)
participant KER as Kernel<br/>(Epilogue)
MoE->>QT: Forward activation_type
alt is_gated_activation = True
QT->>QT: Interleave FC1/gate weights & scales
QT->>COP: Pass weights/scales
else is_gated_activation = False
QT->>COP: Pass weights/scales (unmodified)
end
MoE->>COP: Call with is_gated=activation_type
COP->>KER: Invoke kernel with is_gated
alt is_gated = True
KER->>KER: Load paired subtiles (up, gate)
KER->>KER: Compute: up * silu(gate)
else is_gated = False
KER->>KER: Load single subtile
KER->>KER: Compute: relu(alpha * acc)²
end
KER-->>COP: Output tensor
COP-->>MoE: Result
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_gather_grouped_gemm_act_fusion.py (1)
32-49:⚠️ Potential issue | 🟠 MajorThe renamed runner still only validates the gated/SwiGLU path and provides no way to test the newly added non-gated (Relu2) path.
The kernel class was generalized to support both gated (SwiGLU) and non-gated (Relu2) activations via the
is_gatedparameter, but this runner script still hardcodes SwiGLU semantics throughout: the output dimension is hardcoded asn // 2, the reference verification implements only the interleaved up/gate extraction and SwiGLU computation, and there is nois_gatedparameter or activation selector exposed to exercise the non-gated path. This leaves the main new functionality introduced by this PR unvalidated.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_gather_grouped_gemm_act_fusion.py` around lines 32 - 49, The runner currently assumes gated (SwiGLU) semantics (e.g., using n // 2 for output dim and the interleaved up/gate reference computation) and thus never exercises the new non-gated (Relu2) path; add a command-line flag (e.g., --is_gated or --activation with values "swi_glu"|"relu2") to the script, change the output-dimension logic to branch on is_gated instead of hardcoding n // 2, and update the reference verification routine to compute either the interleaved up/gate SwiGLU path (when is_gated true) or the single-path Relu2 activation (when is_gated false) so both kernel paths (the class parameter is_gated) are testable from the runner.tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py (1)
753-773:⚠️ Potential issue | 🟠 MajorPass
is_gatedthrough the DWDP FC1 path as well.Line 753 switches the DWDP path to the generalized act-fusion op, but this call still relies on the custom op’s default
is_gated=True. That meansactivation_type=Relu2will still execute the gated kernel on DWDP, producing the wrong epilogue and FC1 output width.Suggested fix
x, x_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell_multi_b( input=x.view(torch.float4_e2m1fn_x2), weight=[ w.view(torch.float4_e2m1fn_x2) for w in weight_view.w3_w1_weight ], input_scale=x_sf.view(torch.uint8), weight_scale=[ ws.view(torch.uint8) for ws in weight_view.fc1_weight_scale ], alpha=weight_view.fc1_global_scale, tile_idx_to_group_idx=tile_idx_to_expert_idx, tile_idx_to_mn_limit=tile_idx_to_mn_limit, permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, num_non_exiting_tiles=num_non_exiting_tiles, global_sf=self.fc2_input_scale, num_experts=self.num_slots, top_k=effective_top_k, num_local_experts=esp, local_expert_offset=slot_start, tile_size=tile_size, + is_gated=self.is_gated_activation, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py` around lines 753 - 773, The DWDP FC1 call to torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell_multi_b in fused_moe_cute_dsl.py currently omits the is_gated argument and therefore uses the op's default (True); update that call to pass the correct gating flag (is_gated) used elsewhere in this module/path so the DWDP FC1 path uses the same gating decision as the non-DWDP path (add is_gated=is_gated or the appropriate flag variable from the surrounding scope to the call to cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell_multi_b).
🧹 Nitpick comments (1)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py (1)
1-1: Update copyright year to 2026.Per coding guidelines, modified files should have the copyright year updated. Since this file is being modified and the current year is 2026, the copyright header should reflect this.
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines: "Add NVIDIA copyright header on ALL new files and update year on modified files"
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py` at line 1, Update the file header copyright year from 2025 to 2026: locate the top-of-file copyright comment (the existing "Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.") and change the year to 2026 so the header reads "Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.".
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Around line 3800-3808: The wrapper's local parameter is_gated can diverge from
the kernel instance flag self.is_gated causing interm_size and the output tensor
shape (c) to mismatch the kernel tiling set in __init__/_setup_attributes;
remove the is_gated parameter from wrapper and always use self.is_gated (or
alternatively add an explicit runtime assertion that is_gated == self.is_gated
at the start of wrapper) so interm_size computation and any c shape creation use
the instance's self.is_gated consistently; update references in wrapper that
compute interm_size and c shape to reference self.is_gated and adjust callers to
stop passing is_gated if you choose the remove option.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 6362-6363: test_nvfp4_8gpus was extended to accept "CUTEDSL" via
the moe_backend param but lacks the SM gating applied to other NVFP4 backends;
update the test to skip or restrict "CUTEDSL" to SMs 100 and 103 like the other
NVFP4 tests. Locate the test_nvfp4_8gpus function and add the same SM
check/pytest.skip logic used elsewhere in this file (or wrap the test with the
same requires-SM decorator) to ensure when moe_backend == "CUTEDSL" the test
only runs on SM 100 or 103.
---
Outside diff comments:
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py`:
- Around line 753-773: The DWDP FC1 call to
torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell_multi_b
in fused_moe_cute_dsl.py currently omits the is_gated argument and therefore
uses the op's default (True); update that call to pass the correct gating flag
(is_gated) used elsewhere in this module/path so the DWDP FC1 path uses the same
gating decision as the non-DWDP path (add is_gated=is_gated or the appropriate
flag variable from the surrounding scope to the call to
cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell_multi_b).
In
`@tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Around line 32-49: The runner currently assumes gated (SwiGLU) semantics
(e.g., using n // 2 for output dim and the interleaved up/gate reference
computation) and thus never exercises the new non-gated (Relu2) path; add a
command-line flag (e.g., --is_gated or --activation with values
"swi_glu"|"relu2") to the script, change the output-dimension logic to branch on
is_gated instead of hardcoding n // 2, and update the reference verification
routine to compute either the interleaved up/gate SwiGLU path (when is_gated
true) or the single-path Relu2 activation (when is_gated false) so both kernel
paths (the class parameter is_gated) are testable from the runner.
---
Nitpick comments:
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Line 1: Update the file header copyright year from 2025 to 2026: locate the
top-of-file copyright comment (the existing "Copyright (c) 2025 NVIDIA
CORPORATION & AFFILIATES. All rights reserved.") and change the year to 2026 so
the header reads "Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights
reserved.".
🪄 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: Pro Plus
Run ID: 09626fd7-be01-4541-ace7-2c1eb4adb660
📒 Files selected for processing (10)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pytensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/scripts/cute_dsl_kernels/moe_as_dense_gemm/run_moe_as_dense_gemm_fc1.pytests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_gather_grouped_gemm_act_fusion.pytests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
|
/bot run --disable-fail-fast |
|
PR_Github #44091 [ run ] triggered by Bot. Commit: |
|
PR_Github #44091 [ run ] completed with state
|
e320dea to
90f2504
Compare
90f2504 to
203a7e7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #47331 [ run ] completed with state
|
8e49a58 to
1db1271
Compare
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
The gather grouped-GEMM kernel now supports both SwiGLU (is_gated=True)
and Relu2 (is_gated=False), so the "swiglu" naming is inaccurate. Rename
the kernel file, Runner class, and torch ops to use "act_fusion", and
generalize the top-level/class/forward docstrings to describe both
activations. SwiGLU-specific code comments inside the is_gated=True
branch are kept verbatim.
No behavioral change: the non-gather and dense-FC1 SwiGLU kernels (which
do not have an is_gated path) are untouched.
Renames:
- blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py
-> blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
- Sm100BlockScaledContiguousGatherGroupedGemmSwigluFusionRunner
-> Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner
- trtllm::cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell[_multi_b]
-> trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell[_multi_b]
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…moe. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
1db1271 to
ac6d56a
Compare
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --stage-list "GB200-4_GPUs-PyTorch-1" --disable-fail-fast |
|
PR_Github #47484 [ run ] triggered by Bot. Commit: |
|
PR_Github #47484 [ run ] completed with state
|
|
/bot run --stage-list "GB200-4_GPUs-PyTorch-1" --disable-fail-fast |
|
PR_Github #47540 [ run ] triggered by Bot. Commit: |
|
PR_Github #47540 [ run ] completed with state |
|
/bot skip --comment "CI tests were passed with two runs" |
|
PR_Github #47561 [ skip ] triggered by Bot. Commit: |
|
PR_Github #47561 [ skip ] completed with state |
…2884) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Phase 1 of removing the DWDP IPC-era multi-B kernel infrastructure that the VA-based DWDP path no longer needs. Three follow-up phases will clean the runner-class plumbing and the kernel-side multi-B branches. Multi-B (multiple B tensors passed as a list to a single fused MoE NVFP4 gather kernel) was introduced in PR NVIDIA#12136 to let the original CUDA-IPC DWDP path GEMM N peer expert shards in one kernel call. In the VA-based pipeline (commit "Switch DWDP integration layer from IPC to VA") the composite VA region exposes peer experts as one contiguous [num_experts, ...] tensor via cuMemMap, so single-B handles every case. After that commit the only callers of the multi-B custom op were: * The single-B custom op itself, which had been rewired to wrap single tensors into single-element lists and forward to multi-B (see PR NVIDIA#12136 commits "make multi_b the inner op, original op as wrapper" and "unify single-B and multi-B kernel wrapper into single entry point"). * test_cute_dsl_moe.py, which exercised both paths. This commit: * Removes the trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion _blackwell_multi_b custom op and its register_fake registration. * Rewires the single-B trtllm::cute_dsl_nvfp4_gather_grouped_gemm _act_fusion_blackwell to call Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner directly, restoring it as the primary entry point. The runner / helper still expect list-form for weight, weight_scale, alpha; a one-line wrap-in-single-element-list keeps Phase 1 minimal. Phase 2 will convert the runner back to bare-tensor inputs. * Updates test_cute_dsl_moe.py to call the single-B op directly and drops the multi-B (num_local_experts > 1) test branch. Verified that nemotron-h, the only model that consumes the new ``activation_type=Relu2`` path added by PR NVIDIA#12884, calls the single-B op (fused_moe_cute_dsl.py:631), not multi-B — so multi-B has no remaining production callers. Both renames + ``activation_type`` parameter introduced by PR NVIDIA#12884 (rename swiglu_fusion → act_fusion, add Relu2 support) are preserved. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…r classes Phase 2 of removing the DWDP IPC-era multi-B kernel infrastructure that the VA-based DWDP path no longer needs. Phase 1 dropped the public multi-B custom op; this phase reverts the wrapper-level plumbing back to bare-tensor inputs (matching the pre-PR-NVIDIA#12136 shape) while preserving every unrelated upstream evolution — notably the swiglu_fusion -> act_fusion rename and ``activation_type`` parameter introduced by PR NVIDIA#12884 (CUTEDSL MoE backend for nemotron-h, the only production consumer of ``activation_type=Relu2``). Changes in tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: * GatherGroupedGemmInputsHelper.inputs_pre_hook: unpack inputs as bare tensors (``a, b, a_sf, b_sf, alpha, ...``) instead of list form (``a, b_list, a_sf, b_sf_list, alpha_list, ...``); update class docstring accordingly. * Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner: drop the ``b_tensor_l_sizes`` parameter from __init__, ``unique_id`` and the kernel cache key; rewrite ``get_valid_tactics`` / ``forward`` to take bare ``b`` / ``b_sf`` / ``alpha`` tensors and pass single pointers (not tuples) to the kernel. * Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner: same revert; in addition, remove the ``MAX_B_TENSORS = 4`` class constant. ``activation_type`` and the gated/non-gated branches added by PR NVIDIA#12884 are preserved. * cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell and cute_dsl_nvfp4_grouped_gemm_finalize_blackwell custom ops: drop the ``weight: List[torch.Tensor]`` / ``weight_scale: List[...]`` / ``alpha: List[...]`` type annotations in favour of bare ``torch.Tensor``; drop the runtime ``b_tensor_l_sizes`` derivation in finalize_inplace; both fake registrations follow. * cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell: drop the Phase 1 single-element-list wrapping that was carrying the runner through; the call site now passes bare tensors. Changes in tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py: * NvFp4WeightView: six ``List[torch.Tensor]`` fields become ``torch.Tensor`` since the VA composite-VA tensor swap leaves a single full ``[num_experts, ...]`` tensor in every field. ``_build_local_weight_view`` constructs accordingly; six call sites drop their ``[0]`` indexing. * cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell call site drops single-element list wrapping for weight / weight_scale / alpha. Changes in tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py: * Update finalize-blackwell test (around line 563) to call the bare- tensor signature. * Drop the ``num_local_experts > 1`` multi-B branch of the finalize test (the API is gone). Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…nels
Phase 3 of removing the DWDP IPC-era multi-B kernel infrastructure
that the VA-based DWDP path no longer needs. Phase 1 dropped the
public multi-B custom op; Phase 2 reverted the wrapper-level plumbing
to bare-tensor inputs. This phase reverts the cute.jit ``wrapper`` /
``__init__`` of the two affected kernels to a single-B-only shape and
drops the runner-side single-element-tuple bridging that Phase 2 had
left in place.
Affected kernels:
* ``blockscaled_contiguous_gather_grouped_gemm_act_fusion`` (the
nemotron-h activation-fusion gather GEMM)
* ``blockscaled_contiguous_grouped_gemm_finalize_fusion``
Changes per kernel (mirrors the diff in PR NVIDIA#12136-era commit
``b3cfd89065``, adapted onto the post-PR-NVIDIA#12884 ``act_fusion`` rename
and ``activation_type`` parameter):
* ``__init__``: drop the ``b_tensor_l_sizes`` parameter, the
``MAX_B_TENSORS = 4`` class constant, and the if/else block that
populated ``self.b_tensor_l_sizes`` / ``self.b_tensor_l_offsets``.
Set ``self.num_b_tensors = 1`` unconditionally so existing
downstream ``cutlass.const_expr(self.num_b_tensors >= 2/3/4)``
branches in ``__call__`` statically fold away at cute.jit compile
time.
* ``wrapper``: switch from ``b_ptr_tuple`` / ``b_sf_ptr_tuple`` /
``alpha_ptr_tuple`` to bare ``b_ptr`` / ``b_sf_ptr`` / ``alpha_ptr``
pointers and accept ``l: cutlass.Int64`` directly (instead of
reading ``self.b_tensor_l_sizes[0]``). Drop the up-to-four-element
unroll that built ``b_tuple`` / ``b_sf_tuple`` / ``alpha_tuple``;
the wrapper now constructs single ``b`` / ``b_sf`` / ``alpha``
tensors and forwards them to ``__call__``. ``__call__`` retains
its ``isinstance(b, tuple)`` defensive wrap that promotes a bare
tensor to a single-element tuple, so its multi-B-shaped internals
are untouched.
Corresponding changes in
``tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py``:
* Both ``forward`` methods (Sm100BlockScaledContiguousGroupedGemm
FinalizeFusionRunner, Sm100BlockScaledContiguousGatherGroupedGemm
ActFusionRunner) drop the Phase 2 single-element-tuple wrapping for
b_ptr / b_sf_ptr / alpha_ptr.
* Both kernel constructors no longer receive ``b_tensor_l_sizes=(l,)``.
* Both ``compile_args`` / ``exec_args`` lists pass ``l`` to the
wrapper as an explicit argument (since the kernel no longer caches
it inside ``self``).
Multi-B kernel code (Sm100BlockScaledContiguous*GroupedGemm runners
and the cute.jit kernels) was authored entirely by Zongfei Jing
inside PR NVIDIA#12136; with the IPC path gone there is no remaining
caller, so the cleanup is safe. Verified the only consumer of the
``activation_type`` extension introduced by PR NVIDIA#12884 (CUTEDSL MoE for
nemotron-h) takes the single-B code path
(fused_moe_cute_dsl.py:631), not the multi-B one.
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
7 participants
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
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)
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.