Repository navigation
Triton mxfp4 (and a8w4) grouped GEMM kernel - #723
Merged
Merged
Conversation
matthiasdiener
changed the title
alextmagro
reviewed
Sep 8, 2026
| @@ -0,0 +1,601 @@ | |||
| # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. | |||
Contributor
There was a problem hiding this comment.
Primus kernels carry a 2025 watermark, not 100% sure of this but shouldn't we inherit that if still AMD?
matthiasdiener
marked this pull request as ready for review
September 10, 2026 22:50
matthiasdiener
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
September 10, 2026 22:50
matthiasdiener
requested review from
alextmagro
and
a balanced review from Copilot
September 10, 2026 22:50
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Missing shape and split validation can cause invalid memory access, while valid unaligned inputs currently fail quantization.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds gfx950 Triton kernels and functional APIs for grouped MXFP4 forward, input-gradient, and weight-gradient GEMMs.
Changes:
- Adds persistent grouped MXFP4 Triton kernels.
- Adds quantization, dispatch, and autograd integration.
- Adds numerical correctness and accumulation tests.
File summaries
| File | Description |
|---|---|
grouped_gemm_mxfp4.py |
Implements Triton kernels and launch wrappers. |
grouped_gemm_mxfp4_impl.py |
Adds quantization bridges and grouped operations. |
test_grouped_gemm_mxfp4.py |
Tests forward, backward, padding, and accumulation. |
Review details
Suppressed comments (2)
transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py:228
- This entry point also needs to reject negative splits and require
len(m_splits) == len(weights)andsum(m_splits) == grad_out.shape[0]. Passing unchecked offsets to the raw Triton kernel can cause out-of-bounds offset reads/stores or leave part of the returned gradient uninitialized.
group_offs = _prefix_offsets(m_splits, grad_out.device)
transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py:274
- Before slicing by
m_splits, require nonnegative splits and verify that their sum equals botha.shape[0]andgrad_out.shape[0]. As written, undersized splits silently drop trailing rows from wgrad, while oversized splits produce truncated groups and invalid quantizer shapes.
lhs_data, lhs_scale, go_pad = _col_operand_grouped_padded(grad_out, m_splits) # (N, Mpad/2)
rhs_data, rhs_scale, _ = _col_operand_grouped_padded(a, m_splits) # (K, Mpad/2)
- Files reviewed: 3/3 changed files
- Comments generated: 7
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
matthiasdiener
changed the title
ipanfilo
reviewed
Sep 18, 2026
wangye805
requested changes
Sep 18, 2026
matthiasdiener
requested review from
ipanfilo and
wangye805
and
a balanced review from Copilot
September 18, 2026 20:29
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
Architecture-specific autotuned GPU kernels and custom autograd require human validation alongside the unresolved correctness guards.
Review details
Suppressed comments (3)
transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py:1
- Use a single 2026 start year for this newly added file's AMD copyright, matching the new-file headers in
grouped_gemm_mxfp4_impl.py:1andtest_grouped_gemm_mxfp4.py:1.
# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py:1
- This is a newly added AMD-only file, but its header uses a 2025-2026 range. New AMD-only files in this repository use the current year alone; please change the header to 2026.
# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py:1
- This newly added file uses a
2025-2026copyright range, while new AMD files in this PR use the current-year-only header. Please use the required 2026 header for a new file.
# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
- Files reviewed: 4/4 changed files
- Comments generated: 4
- Review effort level: Balanced
| # Per-group wgrad M padding multiple; must be >= the largest wgrad BLOCK_K the | ||
| # kernel may pick. 256 lets the variable-K autotune try BLOCK_K=256 (at the cost | ||
| # of more padding for small groups). Default 128 = unchanged behavior. | ||
| WGRAD_PAD_MULTIPLE = int(os.environ.get("NVTE_MXFP4_WGRAD_PAD", str(BLOCK_SIZE_K))) |
Comment on lines
+81
to
+86
| def _require_gfx950() -> None: | ||
| # The kernels use the CDNA4 scaled-FP4 MFMA (tl.dot_scaled e2m1); no other arch has it. | ||
| if not _is_gfx950(): | ||
| raise RuntimeError( | ||
| "grouped MXFP4 GEMM requires gfx950 (CDNA4); the current device lacks the" | ||
| " scaled-FP4 MFMA used by tl.dot_scaled(..., \"e2m1\", ...)." |
Comment on lines
+248
to
+257
| if weight_row is not None: | ||
| b_data, b_scale = weight_row | ||
| N = b_data.shape[1] | ||
| _check_splits(m_splits, a.shape[0], b_data.shape[0]) | ||
| else: | ||
| N = weights[0].shape[0] | ||
| _check_splits(m_splits, a.shape[0], len(weights)) | ||
| b_datas, b_scales = [], [] | ||
| for w in weights: | ||
| d, s = _row_operand(w) # (N, K/2), (N, K/32) |
Comment on lines
+81
to
+83
| def _require_gfx950() -> None: | ||
| # The kernels use the CDNA4 scaled-FP4 MFMA (tl.dot_scaled e2m1); no other arch has it. | ||
| if not _is_gfx950(): |
ipanfilo
approved these changes
Sep 23, 2026
1 of 13 tasks
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
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
Based on the Primus-Turbo kernels at https://github.com/AMD-AGI/Primus-Turbo/tree/main/primus_turbo/triton/grouped_gemm.
Adds the kernel itself+functional API. Passes the added tests and the Grouped GEMM microbenchmarks.
Followup PRs:
Fixes https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/371.
Performance results: https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/371#issuecomment-5720191794
Type of change
Changes
Checklist: