Add Graph-safe NVFP4 CUTLASS Group GEMM to unified entry points#3175
Open
cael-ling wants to merge 8 commits into
Open
Add Graph-safe NVFP4 CUTLASS Group GEMM to unified entry points#3175cael-ling wants to merge 8 commits into
cael-ling wants to merge 8 commits into
Conversation
Replace the per-expert multi-stream cuBLASLt loop in the per-tensor NVFP4
grouped path with one CUTLASS ptr-array grouped launch on SM100 (Blackwell).
Common:
- Add nvfp4_cutlass_grouped_gemm.{cuh,cu}: a single-launch per-tensor NVFP4
grouped kernel. Covers BF16 output (fprop/dgrad, overwrite), FP32 output
(wgrad, with optional in-place accumulate for Megatron wgrad fusion), and
optional fused per-group bias (fprop). Per-tensor scaling collapses the
second-level scale to one fp32 alpha per group, applied via the epilogue's
per-group alpha_ptr_array, Arch-gated on CUTLASS_ARCH_MMA_SM100_SUPPORTED.
- Wire it into nvte_multi_tensor_gemm behind the opt-in env
NVTE_NVFP4_CUTLASS_GROUPED_GEMM. M/N/K must be %128; empty (0-token)
experts schedule 0 tiles and no longer force a multi-stream fallback.
Anything unsupported falls through to the existing cuBLAS path, so default
behavior is unchanged.
- Add bench-only entry points nvte_nvfp4_grouped_per_tensor_compute_alpha /
nvte_nvfp4_grouped_per_tensor_gemm so a benchmark can precompute alpha
outside the timed region and time only the grouped GEMM launch.
PyTorch:
- Expose the two bench-only bindings (tex.nvfp4_grouped_per_tensor_compute_alpha
and tex.nvfp4_grouped_per_tensor_gemm). Not used by the production dispatch.
- Extend test_grouped_linear.py with NVFP4 cutlass-vs-multistream parity tests:
GEMM-level (uniform + uneven 128-aligned splits), empty groups, and
end-to-end GroupedLinear fwd+bwd (bias, fuse_wgrad_accumulation).
- Add a GEMM-level cutlass-vs-multistream comparison to
benchmark_grouped_linear.py (--compare-nvfp4-grouped-gemm): a DISPATCH row
(both backends via the shared dispatch) and a fair PURE row (operands
pre-swizzled, alpha precomputed; times only the GEMM).
Signed-off-by: Cael Ling <caell@nvidia.com>
…uped-tensor path Follow-up to NVIDIA#3134, which added the single-launch CUTLASS ptr-array grouped kernel for per-tensor NVFP4 but wired it into the per-expert multi-stream nvte_multi_tensor_gemm path. Per review feedback about grouped-GEMM entry-point fragmentation, move the CUTLASS kernel behind general_grouped_gemm_for_grouped_tensor as an backend and benchmark it head-to-head against the cuBLAS single-launch path (not multi-stream). Consuming the device-side grouped-tensor setup metadata also makes the CUTLASS path CUDA-graph safe. Common: - Route the CUTLASS kernel through the grouped-tensor GEMM (general_grouped_gemm_for_grouped_tensor / execute_grouped_gemm). It consumes the same device-side GroupedGemmSetupWorkspace arrays as the cuBLAS single-launch path, inheriting its graph-safety; selected by NVTE_NVFP4_CUTLASS_GROUPED_GEMM, falling back to cuBLAS when unsupported. - Replace the host-side vector launcher (run_grouped_per_tensor_gemm) and its persistent host buffer with the device-metadata entry point run_grouped_per_tensor_gemm_grouped_tensor. - Remove the old multi-stream CUTLASS dispatch branch and the two bench-only entry points (nvte_nvfp4_grouped_per_tensor_compute_alpha / nvte_nvfp4_grouped_per_tensor_gemm) and their declarations. PyTorch: - Drop the two bench-only bindings. - test_grouped_linear.py: replace the now-dead multi-stream CUTLASS parity tests with grouped-tensor backend tests -- test_nvfp4_grouped_tensor_cutlass_matches_cublas (fprop/dgrad/wgrad parity across balanced/imbalanced splits and fuse_wgrad_accumulation) and test_nvfp4_grouped_tensor_cutlass_cuda_graph_safe. - benchmark_grouped_linear.py: replace --compare-nvfp4-grouped-gemm with --compare-backends (eager per-GEMM) and --compare-backends --graph (CUDA-graphed whole-step) CUTLASS-vs-cuBLAS single-launch comparison. Signed-off-by: Cael Ling <caell@nvidia.com>
github-actions
Bot
added
the
community-contribution
Contributor
Greptile SummaryThis PR adds a CUDA-graph-safe CUTLASS grouped GEMM backend for NVFP4 on Blackwell (SM100), replacing the cuBLAS single-launch kernel when
Confidence Score: 5/5Safe to merge; no P0 or P1 issues found The implementation is well-structured with proper CUDA-graph safety (device-side metadata kernel, no host sync), the device/SM-count fix from the prior review thread is applied correctly, the cuBLAS→CUTLASS A/B swap is consistently applied across the codebase, and both a parity test and a graph-capture test are included. The one P2 comment is a documentation gap, not a correctness bug. grouped_gemm.cu dgrad call site (implicit beta=0 assumption) Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[nvte_grouped_gemm / nvte_grouped_gemm_with_discrete_inputA / nvte_grouped_gemm_with_discrete_out] --> B{NVFP4 dtype?}
B -- No --> C[execute_grouped_gemm cuBLAS backend]
B -- Yes --> D{NVTE_NVFP4_CUTLASS_GROUPED_GEMM=1 AND SM100?}
D -- No --> C
D -- Yes --> E[maybe_run_nvfp4_cutlass_grouped]
E --> F[build_grouped_metadata_kernel device kernel]
F --> G[Writes per-group GemmCoord, strides, SF layouts from d_rows/d_cols/a_rows/a_cols]
G --> H{fp32_output?}
H -- Yes --> I[run_impl_device float epilogue]
H -- No --> J[run_impl_device bf16 epilogue]
I --> K[CUTLASS Gemm::initialize + run]
J --> K
K --> L[D written, graph-safe no host sync]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[nvte_grouped_gemm / nvte_grouped_gemm_with_discrete_inputA / nvte_grouped_gemm_with_discrete_out] --> B{NVFP4 dtype?}
B -- No --> C[execute_grouped_gemm cuBLAS backend]
B -- Yes --> D{NVTE_NVFP4_CUTLASS_GROUPED_GEMM=1 AND SM100?}
D -- No --> C
D -- Yes --> E[maybe_run_nvfp4_cutlass_grouped]
E --> F[build_grouped_metadata_kernel device kernel]
F --> G[Writes per-group GemmCoord, strides, SF layouts from d_rows/d_cols/a_rows/a_cols]
G --> H{fp32_output?}
H -- Yes --> I[run_impl_device float epilogue]
H -- No --> J[run_impl_device bf16 epilogue]
I --> K[CUTLASS Gemm::initialize + run]
J --> K
K --> L[D written, graph-safe no host sync]
Reviews (5): Last reviewed commit: "Address review: split grouped GEMM commo..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…-grouped-tensor Resolve grouped_linear.py: upstream independently routes NVFP4 (with RHT) through the grouped-tensor path, which supersedes this branch's Python routing change. Take upstream's _is_grouped_tensor_path_supported; the CUTLASS backend hooks in at the C++ dispatch and is unaffected. Signed-off-by: Cael Ling <caell@nvidia.com>
…lerances Signed-off-by: Cael Ling <caell@nvidia.com>
for more information, see https://pre-commit.ci
vthumbe1503
reviewed
Jul 6, 2026
vthumbe1503
left a comment
vthumbe1503
left a comment
Collaborator
There was a problem hiding this comment.
Thanks @cael-ling for the updated PR. The PR mostly LGTM. I only have concern with cudaMallocAsync being done within the TE common. We could do cutlass workspace allocation similar to cublas workspace allocation
Comment on lines
+1097
to
+1102
| if (!transformer_engine::getenv<bool>("NVTE_NVFP4_CUTLASS_GROUPED_GEMM", false)) { | ||
| return false; | ||
| } | ||
| if (!transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { | ||
| return false; | ||
| } |
Collaborator
There was a problem hiding this comment.
Nit: Lets keep the order of getting env variable and checking nvfp4 scaling flipped so that the cost of getting env variable is not suffered by non-nvfp4 recipes
Suggested change
| if (!transformer_engine::getenv<bool>("NVTE_NVFP4_CUTLASS_GROUPED_GEMM", false)) { | |
| return false; | |
| } | |
| if (!transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { | |
| return false; | |
| } | |
| if (!transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { | |
| return false; | |
| } | |
| if (!transformer_engine::getenv<bool>("NVTE_NVFP4_CUTLASS_GROUPED_GEMM", false)) { | |
| return false; | |
| } |
Contributor
Author
There was a problem hiding this comment.
Done — the is_nvfp_scaling check now runs before the env-var lookup, so non-NVFP4 recipes never pay the getenv cost.
| // so this stays graph-safe. Requires fp32 output. | ||
| // Fused bias is not supported here (the grouped-tensor path applies bias via a | ||
| // separate nvte_grouped_bias_add); callers must exclude the fused-bias case. | ||
| void run_grouped_per_tensor_gemm_grouped_tensor(void **A_ptrs, void **B_ptrs, |
Collaborator
There was a problem hiding this comment.
Nit: I would rather describe it as being graph safe..Since the API doesnt really consume grouped tensor
Suggested change
| void run_grouped_per_tensor_gemm_grouped_tensor(void **A_ptrs, void **B_ptrs, | |
| void run_nvfp4_graph_safe_grouped_gemm(void **A_ptrs, void **B_ptrs, |
Contributor
Author
There was a problem hiding this comment.
Renamed to run_nvfp4_graph_safe_grouped_gemm across the header, implementation, and caller.
| # dispatches to (cuBLAS baseline vs CUTLASS). NVTE_GROUPED_LINEAR_USE_FUSED_ | ||
| # GROUPED_GEMM routes GroupedLinear onto that path. SM100-only. | ||
| # ============================================================================= | ||
| def _nvfp4_grouped_recipe(): |
Collaborator
There was a problem hiding this comment.
Can we get rid of this function? By the name of this function it would indicate something special about grouped thing.. warranting a different kind of recipe.
But that is not the case. So we can directly instantantiate nvfp4 block scaling recipe within the function.
Contributor
Author
There was a problem hiding this comment.
Removed the helper; both tests now instantiate recipe.NVFP4BlockScaling() directly (kept the RHT note as an inline comment).
| NVTE_CHECK_CUDA(cudaFreeAsync(bufs[device][which], stream)); | ||
| } | ||
| const size_t newcap = bytes + bytes / 2; // slack to avoid frequent regrows | ||
| NVTE_CHECK_CUDA(cudaMallocAsync(&bufs[device][which], newcap, stream)); |
Collaborator
There was a problem hiding this comment.
We dont generally allocate tensors inside the TE common framework. That should be resposibility of upstream framework like pytorch to do the allocation and pass in the NVTETensors to the framework, since otherwise it would be a surprise for the upstream framework.
Could you refer to how cublas uses the workspace? We can use something similar for cutlass as well
Contributor
Author
There was a problem hiding this comment.
Good catch. Removed persistent_buffer and all cudaMallocAsync/cudaFreeAsync from TE common. The CUTLASS path now carves both its on-device per-group metadata scratch and the CUTLASS GEMM workspace out of the caller-provided cuBLAS grouped-GEMM workspace, mirroring how the cuBLAS path consumes it.
Incorporate review feedback on the graph-safe NVFP4 CUTLASS grouped-GEMM
backend so it no longer allocates device memory inside TE common and
follows the cuBLAS workspace-ownership convention.
Changes:
- common/gemm/nvfp4_cutlass_grouped_gemm.cu:
* Remove persistent_buffer() and kMaxCudaDevices, dropping all
cudaMallocAsync/cudaFreeAsync from TE common.
* run_impl_device() now carves both the on-device per-group metadata
scratch and the CUTLASS GEMM workspace out of a caller-provided
buffer (ext_workspace/ext_workspace_bytes), with NVTE_CHECKs that it
is large enough. Upstream (PyTorch) owns the allocation, so the
launch stays CUDA-graph safe by construction. This also removes the
hardcoded-device-0 and device-unaware-buffer concerns.
* Rename run_grouped_per_tensor_gemm_grouped_tensor ->
run_nvfp4_graph_safe_grouped_gemm.
- common/gemm/nvfp4_cutlass_grouped_gemm.cuh:
* Match the rename, add the workspace parameters, and document that the
launcher performs no allocation and the buffer is caller-owned.
- common/gemm/cublaslt_grouped_gemm.cu:
* In maybe_run_nvfp4_cutlass_grouped(), check is_nvfp_scaling() before
the env-var lookup so non-NVFP4 recipes do not pay the getenv cost.
* Thread the cuBLAS grouped-GEMM workspace pointer (and its fixed size)
into the CUTLASS launcher, reusing it as scratch since the CUTLASS
branch returns before execute_grouped_gemm.
* Pass workspace.cublas_workspace_ptr at all three grouped-GEMM entry
points.
- tests/pytorch/test_grouped_linear.py:
* Drop the _nvfp4_grouped_recipe() helper; instantiate
recipe.NVFP4BlockScaling() directly in the two tests.
Default behavior is unchanged (backend remains opt-in via
NVTE_NVFP4_CUTLASS_GROUPED_GEMM).
Signed-off-by: Cael Ling <caell@nvidia.com>
Contributor
Author
|
Thanks for the review @vthumbe1503 , Addressed all the feedback in the latest push. |
Collaborator
|
/te-ci |
Signed-off-by: Cael Ling <caell@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follow-up to #3134.
Review feedback (from @vthumbe1503) was that TE already has grouped-GEMM entry points (cuBLAS single-launch, multi-stream cuBLAS, cuDSL single-launch, and now CUTLASS), and that the CUTLASS kernel should instead live behind the single grouped-tensor API. This PR migrates the kernel from #3134 accordingly:
general_grouped_gemm_for_grouped_tensor(the graph-safe grouped-tensor path), still gated byNVTE_NVFP4_CUTLASS_GROUPED_GEMM.NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1routesGroupedLinearonto that path.Default behavior is unchanged (backend is opt-in). #3134 is temporarily retained for comparison purposes.
Performance
SM100 (Blackwell), NVFP4 default recipe. cuBLAS single-launch vs CUTLASS single-launch, both through
general_grouped_gemm_for_grouped_tensor. Times in ms/iter (lower is better),iters=30. Reproduce with:python benchmarks/linear/benchmark_grouped_linear.py --compare-backends # eager, per-GEMM python benchmarks/linear/benchmark_grouped_linear.py --compare-backends --graph # CUDA-graphed whole step
Eager, isolated per-GEMM (cuBLAS / CUTLASS / speedup)
CUDA-graphed whole-step fwd+bwd (cuBLAS / CUTLASS / speedup)
Type of change
Changes
Common (CUDA/C++)
general_grouped_gemm_for_grouped_tensor/execute_grouped_gemm, consumingthe device-side
GroupedGemmSetupWorkspacearrays (graph-safe), selected byNVTE_NVFP4_CUTLASS_GROUPED_GEMM; fall back to cuBLAS when unsupported.run_grouped_per_tensor_gemm(and itspersistent host buffer) with
run_grouped_per_tensor_gemm_grouped_tensor.entry points.
PyTorch
tests/pytorch/test_grouped_linear.py: grouped-tensor backend parity(
test_nvfp4_grouped_tensor_cutlass_matches_cublas) + CUDA-graph safety(
test_nvfp4_grouped_tensor_cutlass_cuda_graph_safe); remove the deadmulti-stream CUTLASS tests.
benchmarks/linear/benchmark_grouped_linear.py:--compare-backends(eager per-GEMM) and
--compare-backends --graph(CUDA-graphed whole-step).Checklist:
tests/pytorch/test_grouped_linear.pypasses on SM100)