[Kernel] Add compile_mxfp6_gemm: MXFP6×MXFP4 preshuffle GEMM (gfx950)#780
Merged
Conversation
amd-satre
force-pushed
the
feat/mxfp6-gemm
branch
3 times, most recently
from
June 30, 2026 21:06
e54c5a2 to
b88ebab
Compare
Collaborator
|
@amd-satre same perf as the original preshuffle version? |
coderfeli
reviewed
Jul 1, 2026
coderfeli
reviewed
Jul 1, 2026
coderfeli
reviewed
Jul 1, 2026
coderfeli
requested changes
Jul 1, 2026
amd-satre
force-pushed
the
feat/mxfp6-gemm
branch
7 times, most recently
from
July 1, 2026 16:54
5f225b0 to
85f9e54
Compare
Collaborator
|
@amd-satre CI failed. |
amd-satre
force-pushed
the
feat/mxfp6-gemm
branch
2 times, most recently
from
July 8, 2026 18:44
3668526 to
53d9f0a
Compare
compile_mxfp4_gemm and compile_mxfp6_gemm share a single private
implementation _compile_mxfp_blockscale_gemm(a_dtype='fp4'|'fp6').
use_async_copy (default True) selects buffer_load_lds vs coop gmem->VGPR->LDS.
A-operand (fp6): per_1x32_f6_quant / pack_fp6_e2m3 produce tight-packed
fp6 (24 B per K=32 chunk) + 8 B zero-pad. read_a takes the first 6 of 8
DWORDs from LDS → i32[6] for MFMA_Scale(Float6E2M3FN, Float4E2M1FN).
Tile selection: MXFP6_TUNED_CONFIGS + _pick_mxfp6_tiles heuristic.
compile_mxfp6_gemm(M_hint, N, K) auto-selects tiles when not provided.
Performance on gfx950, (M,N,K)=(x,8192,8192), TFLOPS:
fp4-sync fp4-dma fp6-sync fp6-dma
M=32: 185 220 90 204
M=64: 352 394 181 367
M=128: 632 699 335 659
M=256: 521 1333 671 856
M=512: 1459 1658 948 1267
M=1024: 2916 2737 1708 1643
M=2048: 3612 3591 2337 2661
Tests: test_mfma_a6w4_preshuffle (10 cases, gfx950 only).
Signed-off-by: Shreyas Atre <satre@amd.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
amd-satre
force-pushed
the
feat/mxfp6-gemm
branch
from
July 8, 2026 18:46
53d9f0a to
b42e725
Compare
Contributor
Author
Resolved, thank you! |
coderfeli
reviewed
Jul 9, 2026
Collapse _compile_mxfp_blockscale_gemm + the fp4/fp6 pass-through wrappers into a single public compile_mxfp4_gemm(a_dtype='fp4'|'fp6'); compile_mxfp6_gemm is now a thin **kw delegator. Drop redundant _raw() calls in the fp6 read_a repack (Vec.from_elements already coerces elements) and use a comprehension. Net -79 LOC. Verified no codegen change: final gfx950 ISA and LLVM IR are byte-for-byte identical for both fp4 and fp6 (tile 64x128x128, M/N/K=64/8192/8192). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
pushed a commit
that referenced
this pull request
Jul 10, 2026
- api/dsl.rst: document fx.gather/scatter (#805), fx.Basis/E (#735), fx.SyncScope (#742), fx.fastmath + FastMathFlags (#789), arith.maxnumf and chained comparisons (#799/#778), and the current buffer_load signature (cache_modifier/is_scalar, #778). - prebuilt_kernels_guide.md: document rmsnorm backward + store_rstd (#795, build_rmsnorm_bwd_module) and the unified gfx950 launch_gemm (a_dtype fp4/fp6/fp8; replaced standalone compile_mxfp6_gemm, #780). - extern_integration_guide.md: fix path lib/Runtime -> lib/Runtime/ROCm. All symbols verified exported/importable from source. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.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.
compile_mxfp4_gemm and compile_mxfp6_gemm share a single private implementation _compile_mxfp_blockscale_gemm(a_dtype='fp4'|'fp6'). use_async_copy (default True) selects buffer_load_lds vs coop gmem->VGPR->LDS.
A-operand (fp6): per_1x32_f6_quant / pack_fp6_e2m3 produce tight-packed fp6 (24 B per K=32 chunk) + 8 B zero-pad. read_a takes the first 6 of 8 DWORDs from LDS → i32[6] for MFMA_Scale(Float6E2M3FN, Float4E2M1FN). Use case is mentioned in https://rocm.blogs.amd.com/software-tools-optimization/mxfp4-mxfp6-quantization/README.html
Tile selection: MXFP6_TUNED_CONFIGS + _pick_mxfp6_tiles heuristic. compile_mxfp6_gemm(M_hint, N, K) auto-selects tiles when not provided.
Performance on gfx950, (M,N,K)=(x,8192,8192), TFLOPS:
fp4-sync fp4-dma fp6-sync fp6-dma
M=32: 185 220 90 204
M=64: 352 394 181 367
M=128: 632 699 335 659
M=256: 521 1333 671 856
M=512: 1459 1658 948 1267
M=1024: 2916 2737 1708 1643
M=2048: 3612 3591 2337 2661
Motivation
Previous changes were removed. Please see #767
Test Plan
Added back the tests. Tests: test_mfma_a6w4_preshuffle (10 cases, gfx950 only).
Test Result
Tests pass
Submission Checklist