4wave fp4 gemm opt#798
Open
benenzhu wants to merge 122 commits into
Open
Conversation
Extract fixed-shape mxfp4 MoE runners for M=16384, wire bench/profile entry points, and document correctness/performance status. Current FlyDSL scatter, GEMM1, and GEMM2 paths are runnable against the aiter mxfp4 pipeline.
Keep the fixed-shape FlyDSL GEMM1 single-buffered but wait only for the older A-to-LDS VMEM loads before the cross-wave barrier. Document grid differences and the validated v1 profiling result.
Current optimized FlyDSL GEMM1 path compiles and has been validated runnable; profiler/resource scratch outputs are intentionally left untracked.
experiments/, handoff*.md, resource_inspect/snapshots/, and run_bf16_gemm_att.sh are local scratch/WIP, not for tracking.
…-asm g2s Merge the 8 separate LDS buffers into one contiguous symbol so cross-buffer ds_read offsets fold into the 16-bit ds offset: immediate (hot-loop ds_reads now use 4 base VGPRs / 2 windows per operand instead of ~20 materialized addresses), dropping arch VGPR high-water from 252 to 220 (headroom 4->36, off the spill cliff). A single LDS symbol breaks the compiler's alias analysis between the g2s global->LDS DMA writes and the ds_reads, which would insert an s_waitcnt vmcnt(0) drain before every ds_read (-40% perf); emitting g2s as inline-asm buffer_load_dwordx4 ... lds (G2SLoaderAsm) hides the LDS write so no drain is added, while ds_read stays a plain llvm.load so the compiler manages it with fine-grained lgkmcnt (matching the baseline) rather than conservative VMEM vmcnt. K-step offset goes in the buffer instruction's scalar soffset (hardware-free) to keep voffset loop-invariant. ~4270 TFLOPS @ 8192^3, 4812 @ 16384^3. Also add the vgpr-pressure-analysis skill (ISA dump -> high-water + per-purpose arch-VGPR breakdown + peak region) used to drive this work. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…unting Emit the per-K-step E8M0 scale buffer_load as inline-asm (like g2s) instead of the plain buffer_load op. With BOTH g2s and scale opaque to the compiler, it no longer folds the scale VMEM into the same dependence graph as the g2s LDS DMA and stops emitting the coarse combined `vmcnt(N) lgkmcnt(M)` waits that were the dominant hot-loop stall (ATT VMEM-wait 3.7%->14.3% in the prior commit); the hot-loop waitcnts are now clean separate vmcnt(16)/lgkmcnt. Consequence: since the compiler now sees no buffer_load at all, it uses our literal wait_barrier vmcnt verbatim (no recompute). The main-loop count must therefore include the in-flight scale VMEM. The scale is a depth-1 prefetch (8 loads issued at the top of each step, consumed next iter); the nan boundary is vmcnt(11) (>=11 lets a not-yet-complete g2s LDS write be read -> nan at large K, K>=16384), so use 2*(n_ga)+2*(n_gb)=8 with a 2-slot safety margin. Verified correct across 8192/16384/24576 K and non-square shapes. ~4281 TFLOPS @ 8192^3, 4805 @ 16384^3. arch VGPR still 220 (headroom 36). A depth-2 scale prefetch (next) should let the count relax toward 16. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…->4385 Move the E8M0 scale to a depth-2 prefetch issued at the END of each _one_step (prefetch scale[kc+2] after all g2s) instead of depth-1 at the step top. With the scale buffer_loads pulled out of the FIFO window right before the binding wait_barrier, the vmcnt nan-boundary rises: depth-1 forced wait_barrier(8) (the 8 step-top scale loads occupied vmcnt slots ahead of the binding g2s, so >=11 read a not-yet-complete g2s LDS write -> nan at large K), whereas depth-2 end-of-step lets wait_barrier(16) keep more g2s in flight while staying correct. Carry two scale groups (sc_cur, sc_nxt) through the scf.for; the tail consumes both with no extra load. Result: 4281->4385 TFLOPS @ 8192^3, 4805->4881 @ 16384^3 (matches the original 8-symbol baseline 4385 but at much lower register pressure). arch VGPR 220->204 (headroom 36->52, spill 0) -- the clean single vmcnt(16) per barrier (no merged vmcnt/lgkmcnt) also improves scheduling. Verified correct across K=8192..24576 and non-square shapes. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
The two tail K-steps issued their b1/a1 (and next-step a0/b0) LDS reads as a single up-front batch before the MFMAs, exposing the ds_read latency as a serial stall (the main loop already hides these via interleave thunks). Restructure the tail to mirror the main loop: issue the b1/a1 ds_reads as _s2r_thunks in the shadow of the c00 MFMA cluster (c00 only depends on the already-loaded a0/b0), and the next-step a0'/b0' reads in the c10 shadow. Every non-MFMA op now overlaps an MFMA's ~12-cycle execute window instead of standing alone. Mostly helps large K where the tail's data movement is larger: 4881->4920 @ 16384^3; 8192^3 ~flat within noise (~4400). Correct across K=8192..24576. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…opy idiom) The g2s inline-asm recomputed v_readfirstlane + s_mov m0 for every step of every load (32 readfirstlane + 32 s_mov per main-loop iter), all materializing the same loop-invariant LDS base + a constant step stride. Following the gcnasm async_copy idiom (carlushuang/gcnasm), set m0 once for step 0 (s_mov) then advance it with `s_add_u32 m0, <stride>, m0` for the remaining steps. The N_TILES steps of one load are issued back-to-back inside one MFMA cluster, so the m0 add-chain stays coherent; the s_add is volatile inline-asm so the compiler can't reorder across it. Step stride uses the compile-time wave count (block=256 -> 4 waves) so it bakes into the asm string. Main-loop scalar ops: v_readfirstlane 32->8, s_mov 32->8 (replaced by 25 s_add). Frees scalar-issue slots (toward cyc/mfma 16) and lowers power; arch VGPR also drops 201->184. Perf flat within noise (~4400 @ 8192^3, ~4900 @ 16384^3) since the scalars already hid in the MFMA shadow -- this is headroom/efficiency, not a direct speedup. Correct across K=8192..24576. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…ring Scale now loads via ScaleLoaderLDS: one buffer_load_dwordx4...lds per operand per K-step gathers a wave's 4 non-contiguous scale blocks into a contiguous LDS region (lane-contiguous write, m0+L*16), then per-lane ds_read feeds the MFMA -- replacing the 8 per-step buffer_load_dword (each with a dur-6 voffset v_add). Fix: prologue scale gathers must be issued BEFORE the 32 g2s loads so they are the oldest outstanding VMEM and the main-loop wait_barrier(17) drains them; issued last, vmcnt(17) couldn't reach them past 32 newer g2s -> step-0 read saw un-landed LDS -> big-K (16384) cos 0.9995. Also needs 4 scale LDS slots (depth-2 x unroll-2 keeps 4 K-steps live; 3 -> nan). cos 0.999999 all shapes. 4556 TFLOPS @ 8192^3, 5164 @ 16384^3 (vs baseline 6788c9e: 4561 / 5093, +0% / +1.4%). Also add the mfma-coverage-analysis skill. Co-Authored-By: Claude <noreply@anthropic.com>
The end-of-step scale dwordx4-lds gather pair was issued back-to-back after all MFMAs (ATT: exposed, not hidden). Make them thunks appended to the last two mfma.call interleaves so they co-issue in the MFMA execute shadow. cos 0.999999. 4616 TFLOPS @ 8192^3 (was 4556), 5176 @ 16384^3 (was 5164). Co-Authored-By: Claude <noreply@anthropic.com>
Scale ds_read was 8 back-to-back ds_read2st64 before the first MFMA (ATT: exposed). Split read by half: half-0 (saR0/sbC0) feeds c00/c01 immediately; half-1 (saR1/sbC1) is only needed by c10/c11, so issue it as thunks in c00's MFMA execute shadow. cos 0.999999. 4648 TFLOPS @ 8192^3 (was 4616), 5166 @ 16384^3. Co-Authored-By: Claude <noreply@anthropic.com>
Scale read moves into VGPR loop-carry: step kc uses scale[kc] already in registers (read in the prior step's MFMA shadow), so the loop top has no scale ds_read / lgkmcnt wait. Requires depth-3 gather (scale[kc+3], clamped to K_ITERS-1 to avoid OOB fault) so scale[kc+1] is landed when read. +12 arch VGPR (180->192, no spill). Also add interleave_stride to mfma.call: issue a load/read thunk every N MFMAs so each buffer_load/ds_read gets >1 MFMA execute-shadow instead of bunching after the first few MFMAs. Main loop + tail use stride=2. cos 0.999999 all shapes. 4571 TFLOPS @ 8192^3, 4962 @ 16384^3 (vs step2 f89d596: 4531 / 4926, +0.9% / +0.7%). Co-Authored-By: Claude <noreply@anthropic.com>
…firstlane The scale gather m0 was readfirstlane'd every gather (m0 derived from wave_id, a VGPR, forced by the "s" constraint) -- 4 readfirstlane/iter, ATT-exposed ~19%. Precompute the wave-uniform LDS base (lds_base + wave_id*2048 + region) into an SGPR ONCE via set_wave_base(); gather() then builds m0 = wave_base_s + slot*8192 with scalar arithmetic (no readfirstlane). soffset=0 also readfirstlane'd once. Loop readfirstlane 12->8. cos 0.999999. 4573 @ 8192^3, 4977 @ 16384^3. Co-Authored-By: Claude <noreply@anthropic.com>
The old union-of-[issue,issue+EXEC) model capped overlapping windows and mislabeled shadow-hidden loads as exposed (blamed ds_read_b128 33% when dense ~8-cyc-apart fp4 MFMAs actually hid all of them). The next_free model tracks the matrix unit as a pipeline: an MFMA issuing after the unit frees exposes only the idle gap. On fp4_gemm_4wave this corrected EXPOSED 28%->14% and revealed the real blockers: s_add address arithmetic 34% + idle/s_waitcnt, ds_read ~0%. Co-Authored-By: Claude <noreply@anthropic.com>
…gap geometry ATT wave rows are [cycle, issue_dur, STALL, total_dur, code_id]; r[2] is the real per-instruction stall. Every gap-geometry heuristic (first-in-gap / longest / occupancy-overlap) gave contradictory answers on the SAME trace -- variously blaming ds_read_b128 33%, s_add 34%, buffer_load 30%, all wrong. The stall field shows ds_read_b128 and s_add have stall==0 (pure fill between MFMAs, never block); the true blockers are s_barrier 39% + s_waitcnt(lgkmcnt) 35% + buffer_load 11% + readfirstlane 7% + s_waitcnt(vmcnt) 5%. Keep next_free only for the EXPOSED %. Co-Authored-By: Claude <noreply@anthropic.com>
… + pure-idle Two bugs in the prior stall-only attribution: (1) it summed the whole r[2] stall even when the stall overlapped MFMA execute (hidden, pipe busy), and (2) the "MFMA busy" mask must be the next_free gap complement, not a geometric union of [issue,issue+16) -- consecutive ~8-cyc-apart MFMAs keep the unit busy past any single 16-cyc window. Fix: a stall counts as exposed only for its overlap with a next_free idle gap; the uncovered remainder is PURE IDLE (scheduling/dependency slack), reported as its own bucket. On fp4_gemm_4wave this corrected the picture: EXPOSED 6344 = pure-idle 72% + s_barrier 15% + lgkmcnt 7% (was over-reported 35%) + buffer_load 3% + vmcnt 0%. The real lever is tighter MFMA packing, not cutting a single op. Co-Authored-By: Claude <noreply@anthropic.com>
Drop the stall/blocking views entirely. For MFMA-bound speedup what matters is which non-MFMA ops steal issue bandwidth between MFMAs, so tile every EXPOSED (matrix-idle) cycle by the OCCUPYING instruction = its issue_dur + stall, intersected with the next_free idle gaps. Every idle cycle gets an owner (%exp sums ~100). Add a %all column (of the whole window incl. MFMA) = the real wall-clock lever. fp4_gemm_4wave (cyc/mfma 18.73, EXPOSED 14.4%): buffer_load 30%exp/4.3%all, readfirstlane 17%/2.5%, s_barrier 16%/2.3%, lgkmcnt 8%/1.2%, s_add 6%, ds_read 6%. Reads as: gather + readfirstlane + address arithmetic stealing issue slots. Co-Authored-By: Claude <noreply@anthropic.com>
Same fix as the scale gather: the g2s step-0 m0 base was ptrtoint(lds)+wave_id* 1024 with wave_id a VGPR, forcing a readfirstlane per load (8/iter, ATT occupancy 17%exp / 2.5%all). set_wave_base() readfirstlanes the wave-uniform base once; _lds_base_sgpr adds the per-buffer compile-time byte_off with scalar arithmetic. Loop readfirstlane 8->0. cyc/mfma 18.73->18.06, EXPOSED 14.4%->11.4%. cos 0.999999. 4710 @ 8192^3, 5264 @ 16384^3 (GPU6, idle). Co-Authored-By: Claude <noreply@anthropic.com>
# Conflicts: # .gitignore # kernels/fp4_gemm_4wave.py # kernels/mfma_preshuffle_pipeline.py
…tion suffixes) Session scratch had annotated identifiers with __<value> suffixes (BLOCK_M__256, N_TILES_A__4, ...). The BLOCK_M__256 / BLOCK_N__256 ones were keyword parameters of compile_fp4_gemm_4w, so the __N suffix broke the public API -- the upstream test tests/kernels/test_fp4_gemm_4wave.py calls compile_fp4_gemm_4w(BLOCK_M=..., BLOCK_N=...) and failed with "unexpected keyword argument". Strip all __N suffixes back to clean names and black-format. Upstream fp4 test now passes (2/2), cos 0.999999. Co-Authored-By: Claude <noreply@anthropic.com>
coderfeli
reviewed
Jul 3, 2026
| --- | ||
| name: 4wave-mfma-coverage-analysis | ||
| description: From an ATT trace, find which hot-loop instructions are NOT hidden behind MFMA execution -- i.e. what is keeping cyc/mfma above the MFMA execute floor (16 for fp4, 32 for fp8 16x16x128). Models the matrix unit as a pipeline (next_free) for the EXPOSED cycles, then tiles those idle cycles by OCCUPYING instruction (issue_dur + stall, intersected with the idle gaps) so the ranking says which non-MFMA ops steal issue bandwidth between MFMAs. Reports %exp and %all. Use when a GEMM/attention kernel is MFMA-bound and you want to push cyc/mfma toward the floor, or the user asks "what isn't hidden behind MFMA", "哪些指令没被 mfma 掩盖", "暴露的指令", or which non-MFMA op is the biggest exposed stall. | ||
| --- |
Contributor
Author
There was a problem hiding this comment.
Updated to remove the Chinese words, this skill can print how many mfma cycles cost in the hot loops. And how many instrcutions we can optimize to (co-issue/hide the latency between the mfma istructions).
For this kernels. It costs about 18.06 cycles/mfma, while the lower bound is 16 cycles. Pretty close.
However, one drawback is that it can't optimize frequency reduction/optimize the power.
coderfeli
reviewed
Jul 3, 2026
| @@ -0,0 +1,76 @@ | |||
| --- | |||
| name: vgpr-pressure-analysis | |||
| description: Analyze AMDGCN VGPR/AGPR register pressure of a FlyDSL kernel from its ISA dump — report the compiler high-water (vgpr/agpr/accum_offset/spill/headroom) AND a per-purpose breakdown of where the arch VGPRs go (A/B fragment ds_read, address ALU, scale loads, epilogue), plus the true peak-pressure region. Use when the user asks "how many VGPRs does X use", "where is the register pressure", "为啥 vgpr 这么高 / 占用统计 / 寄存器压力", is chasing spill, or wants to free registers (deeper prefetch, more occupancy) on CDNA GEMM/attention kernels. | |||
Contributor
Author
There was a problem hiding this comment.
Sorry forget to delete this one.
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.
Summary
Optimize the 4-wave MXFP4 GEMM (
kernels/fp4_gemm_4wave.py, gfx950/MI355X) scale path and scalar overhead.Scale via LDS gather: replace the 8 per-K-step
buffer_load_dword(into VGPR) with onebuffer_load_dwordx4 ... ldsthat gathers a wave's 4 non-contiguous scale blocks into contiguous LDS, then per-laneds_readfeeds the MFMA. Gather + half-1 read co-issue in the MFMA execute shadow; scale is carried in VGPR with depth-3 prefetch so the loop top has no scaleds_read/lgkmcntwait.Eliminate per-load
v_readfirstlane: precompute the wave-uniform LDS base (both g2s m0 and scale-gather m0) into an SGPR once viaset_wave_base(), cutting loop readfirstlane 12 to 0.Strided MFMA interleave spreads loads across the execute shadow.
Net: cyc/mfma 20.25 to 18.06, MFMA-exposed fraction 23% to 11.4%.
Perf (gfx950 / MI355X, idle GPU, bf16 out, per-1x32 e8m0, cos 0.999999)
TFLOPS. This PR beats the pre-opt kernel by +2.3% / +3.2%
Still have some rooms for the large shapes. Maybe need presistent kernels.
Test plan
tests/kernels/test_fp4_gemm_4wave.pypasses (2/2)Also adds two analysis skills used for this work:
4wave-mfma-coverage-analysis(next_free pipeline exposure + per-instruction occupancy attribution) andvgpr-pressure-analysis.