Repository navigation
perf(gemm-ar): mxfp8 with row-major scales staged through LDS - #806
Conversation
… the loop The fused GEMM's peeled tail pre-reads the last K step's B0 (end of step K_ITERS-2) and A0 (top of step K_ITERS-1), both before that step's wait_barrier(0). Nothing drains those tiles' global->LDS loads first, so the reads can take stale LDS: a few wrong rows, no NaN, intermittent. The main loop's end-of-step wait was lowered from 2A+B to A+B-1 to fix large-grid corruption, but that only made the tail's tiles land sooner. At K=256 the main loop runs zero times, so it cannot help: ptpc 4096x7168x256 is wrong on 12/12 single-GPU launches with it in place. Drain the producers in the tail itself, and restore the loop's original count. With the tail fixed, 2A+B is clean on [16384, 7168] K=512 -- the shape that motivated A+B-1, and which still corrupts 8/12 under 2A+B without the tail change. The last step's wait_barrier(0) stays; it still covers the final A1, issued after it. Also add read_done_barrier() (lgkmcnt(0) + s_barrier) at the loop's S0: the leading wave group refills B0 one stage after the barrier the lagging group's S0 pairs with, and LDS reads retire under LGKM, which the VM waits do not cover. This one is a statically found gap; the tail drain alone clears every failure reproduced here, and it costs nothing measurable. Same diagnosis as ROCm/aiter#6368, which this pipeline was copied from. Measured on 8x MI355X, gemm-only, idle-gated, arms alternating: ptpc 32768x7168x1792 437.1 -> 412.8 us -5.5% blockscale 16384x7168x2048 377.9 -> 362.3 us -4.1% mxfp8 32768x7168x1792 476.6 -> 462.2 us -3.0% mxfp8 16384x7168x2048 246.8 -> 246.2 us -0.2% The fused gemm_ar path is unchanged within noise (communication-bound). Co-Authored-By: Claude <noreply@anthropic.com>
Add quant="mxfp8_raw": the same 32-wide ue8m0 GEMM, taking A's scales as [M, K/32] and B's as [N/32, K/32] uint8 -- what the quantiser and the checkpoint already hold -- and staging them through LDS, after aiter's ScaleRawLds. "mxfp8" (packed, preshuffle_a_scale) is unchanged and stays the default, so existing callers keep working. Reading the packed layout from global once per K step costs twice: the MFMA's vmcnt wait on the scale also drains the tile prefetches issued before it (63 vmcnt waits at K=2048, 20 here), and the scale registers spill in a kernel already at 256 VGPRs (8 at K=512, none here). Each wave now copies its share of four K steps into LDS with buffer_load ... lds, double-buffered, and the pipeline's own wait_barrier covers the copies. GEMM only, MI355X, CUDA-graph timed, idle-gated, modes alternating, cold weights: -9% to -15% at M=64/1024 on the 256x128 tile, -1% to -12% on the 256x256 tile across 10 shapes, +0.9% at K=640 where five K steps do not amortise two chunks. Fused-sdma at TP4 wo_b is unchanged (+-0.3%): that layer is bound by the collective. No padding: any M goes in. Co-Authored-By: Claude <noreply@anthropic.com>
…-lds-scale # Conflicts: # python/mori/ops/gemm_ar/kernels_fused.py
jhchouuu
left a comment
jhchouuu
left a comment
There was a problem hiding this comment.
Tested on 8x MI355X with the CI image (ROCm 7.14) and CI build flags, under both FlyDSL 0.2.4 and 0.3.4.1:
test_gemm_ar.py+test_gemm_ar_op.py+test_flydsl_ar.py: 136 passed, 1 skipped (main: 114 / 1);test_gemm_ar_bench.py: 4/4.- I walked the double-buffer sync by hand and it holds: the refill at step
4ccomes afterread_done_barrier, and even with the one-barrier skew between the wave groups the lagging group has already drained its last read of chunkc-1(its ownlgkmcnt(0)at step4c-1); a chunk is issued at least four steps before its first read; and for everyK_ITERS % 4the last chunk is issued inside the main loop, so the peeled tail never reads an unstaged chunk. The tests' random per-block scales would catch a stale chunk. - GEMM-only (
Mxfp8GemmOp,timing.cold_hot_us, mean of 3 alternating passes, <=1.5% spread) reproduces your table: cold -3.5% to -14%. fused-sdma TP4 is flat (426.9 vs 426.5 us at M=4096, 790.7 vs 789.8 us at M=8192).
Two inline notes. Also: this conflicts with #748 in kernels_fused.py and op.py, so whichever lands second needs a rebase.
| #: ``"mxfp8"`` 32-wide ue8m0 on both, fed to the scaled MFMA as | ||
| #: instruction operands. DeepSeek-V4.1-Flash. | ||
| #: ``"mxfp8_raw"`` the same numbers, in the quantiser's own row-major bytes, | ||
| #: which the kernel stages through LDS. Faster than |
There was a problem hiding this comment.
"Faster than "mxfp8"" holds for most shapes but not all. Measured on MI355X (same method as above), the raw form is slower hot at small M:
| M x N x K | hot, FlyDSL 0.3.4.1 | hot, 0.2.4 |
|---|---|---|
| 64 x 8192 x 1280 | +4.5% | +5.6% |
| 1024 x 5120 x 2048 | +2.0% | +2.7% |
| 8192 x 7168 x 640 | +0.5% | +3.3% |
Maybe qualify it, e.g. "faster cold and at large M; can be slightly slower hot at small M or short K".
There was a problem hiding this comment.
Agreed, the claim was too broad. op.py, the README and the benchmark doc now say it is faster cold and at large M, and can be a few percent slower hot at small M or short K (c74ef85).
| self.lds_ptr_t = fx.PointerType.get(fp8_ir_t, 2, 512) | ||
| self.srcs = [] | ||
| for arg, first, rows_per in ((A_scale, row0, 1), (B_scale, col0, 32)): | ||
| # The descriptor takes the tensor's own extent, so rows past M or |
There was a problem hiding this comment.
Is the "last chunk's bytes past K read as zero" part right? The chunk offset goes in through soffset=chunk * 16 (L150), and as far as I know the raw-buffer range check on CDNA compares only the VGPR offset (+ inst offset) against num_records; soffset is not part of it. Then for the tensor's last row the bytes past K come from just beyond the tensor (up to 12 bytes at K=640) rather than zeros. They are never consumed, so results are unaffected, but the comment could say "unused" instead, or the chunk offset could move into the VGPR offset if a zero read is wanted.
There was a problem hiding this comment.
Right: soffset is outside the range check, so the last row's bytes past K were read from beyond the tensor. Never consumed, but the read could reach past the allocation, so I fixed it rather than reword it: the chunk offset now goes in the per-lane VGPR offset (c74ef85). Every scale copy in the ISA is now 0 offen lds and the zero-fill comment is true. K % 128 == 0 keeps each row whole dwords, so a dword is either in the tensor or zeroed.
I also tried the immediate offset to avoid the extra add, but on an LDS-direct load it moves the LDS destination too (chunk >= 1 landed 16 bytes off), and slot 0 starts at LDS 0 with all 160 KB in use, so M0 can't compensate. The docstring records that.
Cost, same node, old vs new alternating, 9 runs each: within ±1% on five shapes, +2-3% at 8192x7168x640 (two chunks, both in the prologue). 140 passed / 1 skipped on the merged tree.
The chunk offset went in through soffset, which the raw buffer range check does not cover, so the last row's bytes past K were read from beyond the scale tensor rather than zero-filled. They are never consumed, but the read could reach past the allocation. With the offset in the per-lane address the check covers it, and the comment that claimed zero-fill is now true. Also qualify "faster than mxfp8": it holds cold and at large M; hot at small M or short K it can be a few percent slower. Co-Authored-By: Claude <noreply@anthropic.com>
|
@jhchouuu Thanks for the review, and for walking the double-buffer sync by hand. Both inline points are addressed in c74ef85 (replies inline); the branch is also merged with current main. On 8x MI355X (ROCm 7.14, FlyDSL 0.3.4.1): On the conflict with #748: whichever lands second, I'll merge main into it and resolve there. |
jhchouuu
left a comment
jhchouuu
left a comment
There was a problem hiding this comment.
Re-tested c74ef85 merged with current main on 8x MI355X with the CI image. Current main includes #817, so this is on FlyDSL 0.4.0.
test_gemm_ar*.py+test_flydsl_ar.py: 136 passed, 1 skipped on FlyDSL 0.4.0, and the same on 0.3.0.test_gemm_ar_bench.py: 4/4.- The default paths compile to the same ISA as main on every rank.
- The raw-mode scale copies are all
0 offen ldson 0.4.0 too. - GEMM-only timing matches the description.
Nit: the 8192 x 7168 x 640 row in the description predates c74ef85. I now measure +2.9% cold and +3.5% hot, consistent with the +2-3% in your reply.
LGTM.
…-lds-scale # Conflicts: # python/mori/ops/gemm_ar/op.py
What
Adds
quant="mxfp8_raw": the same 32-wide ue8m0 GEMM, but it takes the scales as the quantiser and the checkpoint already store them — A as[M, K/32], B as[N/32, K/32], uint8, row major — and stages them through LDS. The reader is adapted from aiter'sScaleRawLds."mxfp8"(packed A scale throughpreshuffle_a_scale, int32) is unchanged and stays the default, so existing callers keep working.Why it is faster
"mxfp8"reads its scales from global once per K step. That costs in two ways:s_waitcnt vmcnt, which retires in issue order. The scale loads are issued between the tile prefetches, so that wait also drains the prefetches ahead of them. The K=2048 kernel has 63vmcntwaits; most arevmcnt(3..5)sitting in front of an MFMA."mxfp8_raw"has each wave copy its share of four K steps into LDS withbuffer_load ... lds, double-buffered, and read operands back withds_read_u8. The copies count undervmcntlike the tile loads, so the pipeline's ownwait_barriercovers them.mxfp8mxfp8_rawvmcntwaits, 256x256, K=2048vmcntwaits, 256x128, K=2048Moving the global scale loads one step earlier, without LDS, does not get the same result. It removes the coupling but adds live registers: K=2048 spills went from 3 to 11 and the GEMM got slower on 8 of 12 shapes (1–6.6%, on a run the idle check flagged as busy by the end, so treat those numbers as directional).
Changes
_mxfp8_scale_lds.py(new):Mxfp8ScaleLds, trimmed from aiter'sScaleRawLdsto the 1x32 row-major form, with aiter's notice kept as attribution, as for_gemm_a8w8_8wave.py. One change: the staged rows start at the block's first row rather thanblock * 256, so the 256x128 tile stages the right B rows.kernels_fused.py: themxfp8_rawmode. The scale slots go first inSharedStorage(16 KB). The prologue copies chunks 0 and 1 ahead of its tile loads. The main loop copies chunkk/4 + 1right after S0'sread_done_barrier: the slot it overwrites held chunkk/4 - 1, whose last reads both wave groups have drained by then. The kernel name gets anRtag.op.py,gemm.py: both ops acceptquant="mxfp8_raw". The scales are checked for dtype, shape and device, then flattened.Mxfp8GemmOptakes any M in this mode; there is no packed group to pad to.soffset: the raw buffer range check covers only the VGPR and immediate offsets, so viasoffsetthe last row's bytes past K were read from beyond the tensor. The immediate offset cannot carry it either: on an LDS-direct load it also moves the LDS destination.bench_gemm_ar.py:--quant mxfp8_raw.docs/MORI-GEMM-AR-BENCHMARK.md: the mode and its numbers; "faster" is qualified to cold and large M.Measured
MI355X (crsuse2 v3 node, ROCm 7.14, FlyDSL 0.4.0 as CI now pins), on this branch merged with main. GEMM only (
fuse=False), CUDA-graph timed withbenchmark/cco/flydsl/gemm_ar/timing.py. Each run starts only after 5 consecutive samples show every GPU under 5% and is discarded if any is busy at the end. Min of three passes with the modes alternating. Times in µs. Cold rotates the weight past the LLC, which is what a forward pass sees.mxfp8coldmxfp8_rawcoldWorst pass-to-pass spread within an arm was 2.1%, under 1% for most rows. The same run under FlyDSL 0.3.4.1 on the same node agrees within about a point on every row. It loses where the staging is not amortised: K=640 is five K steps split across two chunks (and pays most for the VGPR-offset fix below, +2-3% in a same-node A/B), and the hot M=64 case copies 256 rows of A scale for 64 real ones. 8192 x 7168 x 4096 is also slightly slower (+1.7%); I have not pinned down why.
Fused-sdma, TP4,
wo_b(N=5120, K=2048),max_rank_time_us, two passes alternating:mxfp8mxfp8_rawNo change. The fused layer's time is mostly the collective, so a few percent off the GEMM does not show there. The gain is for the standalone GEMM (
Mxfp8GemmOp). relL2 is identical between the two modes on every fused point (bf16 2.37e-3, fp8 2.31e-2).Test
No new test file; the cases join the existing ones:
test_mxfp8_lands_at_the_fp8_floorruns both modes at K = 256, 512, 640, 1280, 2048. These cover the chunk edges: a single chunk, only the prologue's two chunks, and chunks refilled in the loop.test_standalone_mxfp8_gemmruns both modes on both tiles (M = 64, 192, 1024 run the 256x128 tile; 4096 runs 256x256).test_standalone_mxfp8_raw_takes_any_mruns M = 1, 100 and 1000 unpadded.test_mxfp8_raw_quant_runs_end_to_endruns the 2-rank fused op with row-major scales.On 8x MI355X with FlyDSL 0.4.0, merged with main:
test_gemm_ar.py+test_gemm_ar_op.py+test_gemm_ar_schedule.py(including #748's five worker runs) +test_flydsl_ar.py, 162 passed, 1 skipped;test_gemm_ar_bench.py4/4. The benchmark also checks every result element-wise against an fp32 reference, with the output filled with NaN before each of 6 launches: 0 elements off on all 16 shapes.Not covered
Mxfp8GemvOp) is untouched; it already takes row-major scales.wo_bwrites the packed layout with its own quantiser and keeps"mxfp8". Moving it to"mxfp8_raw"is a separate change, and it would remove that quantiser's packing.🤖 Generated with Claude Code