[fp4_gemm_4wave] MXFP4 GEMM perf optimize - #990
Open
benenzhu wants to merge 7 commits into
Open
Conversation
Brings the gfx950 4-wave MXFP4 GEMM to within ~1-3% of aiter's hand-written assembly. Measured against aiter's f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256 on MI355X, same process, alternating, 5 rotating input sets (so nothing is served out of MALL), warmup 500 + 3x500 timed iterations: M=N K fly aiter delta 8192 8192 5034 5092 -1.13% 8192 16384 5456 5514 -1.05% 16384 8192 5002 5146 -2.81% 16384 16384 5499 5592 -1.66% Main changes, roughly in descending order of impact: * MFMA accumulators pinned in AGPR via inline asm, freeing arch VGPR for the deeper software pipeline. * Serpentine 2x2 MFMA emission order so consecutive MFMAs reuse XDL operands. * Depth-3 scale prefetch carried in VGPR; the per-K-step scale gather is a single dwordx4-to-LDS per wave, co-issued inside the MFMA execute shadow. * g2s uses the m0 set-once + s_add idiom, and the wave-uniform LDS base is readfirstlane'd once instead of per load. * b1 operand carried across K-steps, which erases the hot loop's independent lgkmcnt waits (68 -> 8 s_waitcnt). * Epilogue: dwordx4 stores via MFMA operand swap + permlane16, with the stores interleaved into the tail MFMAs. Back-to-back buffer_stores queue on L1, so spacing them out behind MFMAs is worth more than issuing them earlier. * Prologue: pinning M/N and the wave count at compile time removes both the software integer divides and the second s_waitcnt lgkmcnt(0), cutting the instructions ahead of the first buffer_load from 272 to 88. Note one behavior change: BLOCK_M/BLOCK_N/mn_aligned are replaced by a single optional MN=(M, N). The kernel is now specialized to a 256x256 block with an all-in-bounds epilogue, and the test asserts M % 256 == 0 and N % 256 == 0. The in-repo test was the only caller of the removed parameters. fp8_gemm_utils.py is deliberately untouched. The one helper that needed a change (compute_global_swizzle, to take a compile-time wave count) is 17 lines of pure arithmetic, so it is inlined here as _global_swizzle rather than altering a function the fp8 kernels share. Verified: tests/kernels/test_fp4_gemm_4wave.py passes at 8192^3 and 16384^3, cos=0.999999 against the torch mxfp4-dequant reference. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Compress the long rationale comments down to one-liners and drop the FP4_MAIN_VMCNT / FP4_SEG2_VMCNT / FP4_NO_ALIAS debug env knobs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
|
Awesome! CI failed. @benenzhu |
coderfeli
reviewed
Aug 10, 2026
coderfeli
reviewed
Aug 10, 2026
coderfeli
reviewed
Aug 10, 2026
coderfeli
reviewed
Aug 10, 2026
coderfeli
reviewed
Aug 10, 2026
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.
Brings the gfx950 4-wave MXFP4 GEMM to within ~1-3% of aiter's hand-written assembly.
Performance
Measured against aiter's
f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256on MI355X, and against the kernel as it stood before this branch (main@ 421935c).(TFLOPS, higher is better.) Still slows than aiter asm in prologue and epilogues
Compare scripts:
python3 tests/kernels/test_fp4_gemm_4wave.py --vs-aiterAll three run in the same process, alternating blocks over the same 5 rotating input sets (to clear the L2 & LLC cache), warmup 500 + 3x500 timed iterations under a single event pair.
Main changes
Roughly in descending order of impact:
buffer_load_b128...ldsnow.s_addidiom, and the wave-uniform LDS base isreadfirstlane'd once instead of per load.Cstores usedwordx4stores via MFMA operand swap +permlane16,buffer_stores queue on L1, so spacing them out behind MFMAs is worth more than issuing them earlier.s_waitcnt lgkmcnt(0), cutting the instructions ahead of the firstbuffer_loadfrom 272 to 88.Note:
FP4_DMA_INTRINSICto turn off the inline asm ofbuffer_load_lds. But it can't useinc_m0learned from https://github.com/carlushuang/gcnasm/blob/master/async_copy/main.hip.cc#L50-L53, and will be slightly slower.Behavior change
BLOCK_M/BLOCK_Nmust be256now.Also we can passin optional
MN=(M, N)to make M/N compile time constant.