Fp8 mxscale bmm bpreshuffle opt - #4747
Draft
yzhou103 wants to merge 21 commits into
Draft
Conversation
- gfx950 kernel family with codegen'd launchers, split-K and split-K reduce - dispatch for tile-unaligned M: padded-M lookup, sub-tile gating, M-split fallback - tuned-row lookup in the a8w8 batched family layer, libtype-dispatched and lru_cached - kid -> M alignment derived from codegen, asserted at build time and in the op test - self-contained tuner plus the DeepSeek-V4 tuned CSV gsm8k 5-shot exact_match 0.953 on DeepSeek-V4-Pro (tp8, fp8 kv, mxscale on). Co-authored-by: Cursor <cursoragent@cursor.com>
Guard the tuned lookup behind a functional custom op while keeping focused dispatch and tileN regression coverage. Co-authored-by: Cursor <cursoragent@cursor.com>
The 16x16 shuffle_weight order already is the mfma_16x16x128 B fragment order, so a consumer wave's B operand is 2048 contiguous bytes with one naturally aligned dwordx4 per lane. Consumers therefore buffer_load B straight into their MFMA registers, which drops the B half of the producers' async copies, the B ds_reads and the B LDS buffers; the per-tile scale wait doubles as the retire point for B, so the vmcnt accounting stays uniform. Per-subtile e8m0 scales are picked with the hardware scale_op_sel immediate instead of a broadcast pack. Three kids ship, each mirroring the geometry of the row-major kid the tuner picks in that M band so the only variable is where B comes from: 179 (16x32x512) is 5-38% ahead of kid311 for M <= 64, 171 (64x32x256) is ~9% ahead of kid321/kid653 around M 256-512, and 178 (128x64x256) is ~8% ahead of kid653 near M 1024. From M ~2048 up, kid325 and kid158 stay 1.4-2.3x ahead of anything this family can do -- half its waves are producers, so a compute-bound tile only gets two MFMA waves. Negative results are recorded next to the tile table so they are not retried: spending the freed LDS on a deeper A pipeline or a higher WG_PER_CU both lose (occupancy is VGPR-bound at 2 waves/SIMD), scaling only M or only N does nothing, and the scale prefetch kid311 ships is worth nothing once the direct-B vmcnt wait retires the scale loads. Co-authored-by: Cursor <cursoragent@cursor.com>
…up size scale_shuffle=true used to mean one layout, the MFMA scale-load swizzle for V_MFMA_SCALE_F32_16x16x128_F8. That swizzle only earns its keep at quant_group_size 32, where the four bytes of a lane's dword are the op_sel iterations. A 1x128 scale instead spans the whole MFMA K step, so the consumer broadcasts the single byte to all four 32-blocks and all four bytes of the dword would be equal -- there is no tile to swizzle, only an axis order to choose. So at other group sizes emit the plain [G, Ks, S] transpose, which is what dynamic_per_group_scaled_quant already produces there. The win is on the consumer's side: its 16 M lanes read 16 adjacent bytes instead of bytes Ks apart. Both layouts occupy the same bytes per group, so the row addressing stays shared. Also fill the S_pad tail rows the kernel never writes with 0x7F (E8M0 1.0): a caller that padded S for its GEMM's M tile has its consumer read those rows unconditionally, and they have to dequantize harmlessly rather than hold whatever was in the allocation. Co-authored-by: Cursor <cursoragent@cursor.com>
Grows the preshuffled-B work from one family into four, adding a wave8 pipeline
(wave8n4, wavetm1) alongside families on the existing flatmm-splitk (blds) and
plain BMM pipelines, and takes the catalog from 45 to 93 instances. B arriving in
(16, 16) MFMA-fragment order is what buys this: the consumer waves buffer_load it
straight into the MFMA registers, so B never touches LDS, and every MFMA picks
its e8m0 byte with the hardware scale_op_sel immediate instead of a broadcast
pack. The shuffle_scale layout reads both scale panels from global, which
compiles the LDS panel and its K bound out, so those kids run any K.
PREFETCH_SCALE loads the next K-tile-pair's scale word one iteration ahead. It is
gated on COM_REP_K == 1 and either <=16 register tiles, or <=32 when B_M <= 128,
because register headroom -- not tile count -- is what decides whether it pays.
Measured: 128x256 gains 5-12% on a 2x4 wave grid and 1.3-5.7% on 1x4 (250 -> 254
VGPR, no spills), while 256x256 has the same 32 tiles yet already spills 10
before prefetching and loses 11-16% with it. The traits cannot express "has six
VGPR to spare", so B_M stands in for it.
Two smaller fixes ride along:
* The launcher now AITER_CHECKs the LDS scale-panel K bound. Past it the kernel
returned without writing Y, which a caller cannot tell apart from a GEMM that
legitimately produced zeros.
* alignas(16) __shared__ rather than __shared__ alignas(16) in three pipelines.
clang 20 rejects the latter ("'alignas' attribute cannot be applied to
types"); clang 22 emits identical ISA either way.
Co-authored-by: Cursor <cursoragent@cursor.com>
The tuner listed three preshuffled-B kids as candidates but handed every candidate the same row-major B, so those three read the right bytes in the wrong order, failed the correctness gate, and were dropped -- silently, and on every shape. That is why the shipped table names none of them. It now shuffles B per candidate, which makes all four families evaluable for the first time, and opens the full set of them at splitK=1 (none carries the flatmm-splitk launcher's fused reduce tail). Results go to a second table rather than the shipped one, since a row naming a preshuffled kid is only correct for a caller whose weights really are baked that way: --bpreshuffle retunes the shipped shapes into dsv4_batched_gemm_a8w8_blockscale_mxscale_bpreshuffle_tuned.csv, added here with 133 rows over g=2/4/8/16 x n1024 x k1024/4096. Its filename deliberately does not match the glob that merges model_configs tables, so it stays opt-in behind AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE. --pool rowb|preb splits the candidate pool by B layout, which is how to price preshuffling: tune the same shapes twice and compare the two tables cell by cell. Also adds kid159 and kid164 to the row-major pool -- narrow-N wo_a leaves the 256x256 tile only 4 N-tiles, so mid-M shapes idle half the CUs, and these fill them from the M and N sides (164 also covers n128, which the 256-wide tiles reject). Co-authored-by: Cursor <cursoragent@cursor.com>
A serving stack bakes wo_a into the (16, 16) MFMA-fragment layout at load time and has no way to say so, so b_preshuffled goes on batched_gemm_a8w8_mxscale itself. It is the one caller-facing exception to keeping kernel choice out of that entry, and not a tuning knob: it describes the data. It also cannot be detected -- a shuffled weight has the same shape, dtype and strides as a row-major one, so a kernel mismatched to it reads the right bytes in the wrong order and returns a plausible wrong answer instead of failing. The backend therefore never runs a kernel whose B layout disagrees with the declaration. It drops such an id the same way it already drops one that cannot run this M, which is right for row-major B (the heuristic then answers correctly, so pointing the config at the preshuffle table without passing True degrades to untuned rather than to wrong). Under True there is nothing to fall back to, since every heuristic kid reads B row-major, so that raises. Scales get the same treatment: 7 of the 93 instances read them through an M-packed panel or the shuffle_scale layout, and this entry passes the plain arrays through, so those ids are dropped too. No tuned row names one today, but the preshuffled wave8 kids are among them and are now tunable, so a re-tune could put one in the CSV -- where, unlike an M-alignment mismatch, it would not throw. Adopts the tuned-CSV lookup that lives one layer up in the family entry and deletes the backend's private copy. The private one rounded M to the nearest bucket in its own table while the shared one uses the C++ getPaddedM every other lookup uses, so the same shape could resolve to different kernels depending on which entry a caller reached. The public entry is the torch.compile-guarded custom op over that; a caller needing to write into its own batch-major buffer uses the backend directly, which keeps out=. Adds a dispatch check for all six routing outcomes. It spies on the raw binding over meta tensors, because none of them is visible in the output tensor. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
The sweep already handed every kid the B layout its family reads and the M-packed A scale panel, but not the shuffle_scale layout, so the six kids wanting it (210 and 213-217) computed against plain scales and missed the reference at every M -- reported as m_align disagreeing with the launcher, which said nothing about M and buried the 87 kids the guard does cover under 72 lines of false failure. Both callers now share one picker keyed on the kid, so a kid's scale layout is looked up in the catalog exactly where its weight layout already is. The preshuffle op_test had grown its own copy of this logic, including the strides that make the shuffle_scale slabs addressable; that copy is what the sweep was missing, and a second copy is how it would go missing again. The guard now passes for all 93 kids. Co-authored-by: Cursor <cursoragent@cursor.com>
The preshuffled-B tuned rows were kept in their own CSV, but reaching them meant pointing AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE -- the row-major table's own env var -- at that file. That replaced the row-major table for every caller in the process: a b_preshuffled=False call then read rows naming preshuffled kids, dropped each one for the heuristic, and silently lost the shipped table's pick. It also meant a process serving both layouts could hold only one of the two tables. Register the preshuffle table as its own config family instead, the way AITER_CONFIG_GEMM_A8W8[_BLOCKSCALE]_BPRESHUFFLE already splits the non-batched GEMM, and let the layout the caller declares pick the table rather than an env var. Both tables now load at once, and neither glob can see the other's files (_bpreshuffle_ sits between mxscale and tuned). The shipped preshuffle CSV becomes that entry's default, so b_preshuffled=True needs nothing set. Checked on gfx950: 77 rows / 8 kids row-major, 133 / 21 preshuffled, no kid on either side wanting the other layout -- the new tables: check in the op test reads both the way the entry does, so a rename that let one glob swallow the other would fail it. The config-collision test covers the new family too. Also fix two findings ruff's default rule set reports on earlier commits here: the test's import grouping (I001), and two implicit string concatenations inside tuple literals in gen_instances.py (ISC004), which read exactly like a missing comma. The reflowed template-parameter decls are byte-identical, and the rebuilt kernels still match the reference on all 47 kids. Co-authored-by: Cursor <cursoragent@cursor.com>
The table was tuned one timing pass per candidate, which crowns the wrong kid wherever several sit within a couple of percent of each other. Re-swept all 133 cells against every kid the entry can dispatch, in interleaved rounds scored on their mins, then re-ran each hit head to head in the tuner's own regime (10 warmup / 101 iters) before believing it: 26 cells looked at least 2% better in the sweep, 7 survived the confirm, 4 survived the re-measure that wrote them. g8/m64/k4096 kid172 -> kid228 +5.6% bdirect -> the blds twin g16/m32/k4096 kid172 -> kid226 +3.9% g16/m32/k1024 kid244 -> kid226 +4.8% g8/m4096/k1024 kid203 -> kid194 +3.6% wavetm1 -> wave8n4 Every column comes from the tuner's pipeline (gen_bmm_mxscale_data for the inputs and reference, checkAllclose at its tolerances for errRatio, its own tflops / bw formulas), so a re-pointed row is what the tuner would have written had its pass landed on the winner. The candidate set is narrower than the sweep behind the original table, which also swept the kids reading their scales through an M-packed panel or the shuffle_scale layout. The entry passes x_scale / w_scale straight through and drops such a kid, so those cells' apparent winners were never dispatchable -- kid215 "beating" a row by 1.4% could not have been taken. The op test's table check now asserts that invariant next to the B-layout one, on both tables. Five cells stay slower than the shipped table is with row-major B; the comment on BPRESHUFFLE_CSV records them, because two are twin-vs-twin (same tile, same sfpreload, only B's layout and its LDS hop differ) and so measure what preshuffling costs at the 128-wide tiles rather than anything a re-tune fixes. Co-authored-by: Cursor <cursoragent@cursor.com>
Four of the 133 cells trail the shipped row-major table, not five. The fifth, g2/m32768/k4096, was never a layout cost: kid196 -- kid158's own pipeline reading a preshuffled B -- and kid205 both land within 1% of row-major there, and the row named kid194, the slowest of the three. The candidates sit inside 2% of each other, which is closer than the single-pass sweep that wrote the row can rank. Re-point it to kid205, and g16/m4096/k4096, the same mis-rank in the same family, to kid196. The other 15 rows naming kid194 keep it; it is right on all of them. Getting there needed two measurement traps ruled out, so both are written down next to the numbers they would otherwise quietly falsify. run_perftest rotates deep copies of the arguments it is handed, and operands captured in a closure are not arguments -- that times one cache-resident weight, worth 14% to the 8-wave kids, and it is what made kid175 look like it beat row-major at g16/m128 when it and kid230 are a wash. And the K=4096 weight-placement effect opus_gemm_common.py records had to be measured rather than assumed away, since the two sides of this comparison necessarily hold different buffers: over 8 draws that move both, every remaining gap holds its sign and no kernel varies by more than 4%. Co-authored-by: Cursor <cursoragent@cursor.com>
main carries this same BMM work as the squashed ROCm#4320, so every conflict was our branch's newer form of a line against the version that landed upstream (hard-coded SFA_K_MAX=8192 before the per-traits bound, the silent kid fallback before the AITER_CHECK, and so on). No third party had touched the 14 conflicted files, so all of them resolve to ours; opus BMM paths come out byte-identical to the pre-merge HEAD. Co-authored-by: Cursor <cursoragent@cursor.com>
…tch dir Two comments named the scratch harness they were measured with, which is not in the tree and tells a reader nothing they can open. The numbers are the point; keep those and drop the file names. Co-authored-by: Cursor <cursoragent@cursor.com>
…n two more families Instances are deduplicated by name, so a flag with no suffix in _name collapses a tile onto its plain sibling and the kid is emitted as the sibling. The bdirect branch already guards against that; the flatmm split-K branch did not, and there the trap was live -- that kernel has implemented SHUFFLE_SCALE all along and the codegen spells it on every one of its kids, so a shuffle_scale instance would have measured as "the layout changes nothing" with nothing raised. wave8n4 has the same hole for XCD_WGM. Adds the two suffixes, the shuffle_scale kwarg to the flatmm split-K factory (asserting the kernel's own static_asserts on the Python side) and xcd_wgm to the wave8n4 factory. No instance sets either flag yet, so all 676 emitted kernel names are unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
… LDS budget
SF_PRELOAD_K_MAX was a flat 8192 copied into three traits. That figure is
kid158/196's arithmetic -- 151,680 of 163,840 bytes, with a 2*(B_M+B_N)*B_K
double buffer as the staging -- and it does not describe the wave8 family, whose
staging is prefetch_k_iter A slots. Reading .group_segment_fixed_size out of the
built code objects puts 19 of 25 panel kids at 3-884x spare: kid205 sits at
59,012 bytes and has room for ~111,000 of per-split K, so the flat bound was 13x
short and kept it out of the split_k=1 column, which on large-K machine-filling
shapes is where the fastest kernel runs.
So the wave8 traits now computes the bound from what its staging leaves over.
Two things the naive version of that got wrong, both kept as comments:
* budget against the LDS share that keeps the workgroups resident, not
max_lds_size_per_wg. WG_PER_CU is a declared attribute, not what the CU
schedules: kid203/kid205 are 256-thread workgroups that fit twice at 59,012
bytes and once at 83,972, and spending the spare half on reach cost them
1.19-1.20x at m>=1536. Preserving residency still yields 30,464 of per-split
K for them and 30,848 for kid194.
* cap the result and reserve 256 bytes, because the panel array is sized from
this constant rather than from the runtime K, and this arithmetic is not the
allocator's (58,960 modelled against 58,948 real on kid338).
A static_assert holds the derived value at or above the 8192 the flat constant
promised, so no kid loses reach. The other two traits keep the constant and now
say whose arithmetic it is.
Co-authored-by: Cursor <cursoragent@cursor.com>
… its flatmm kids
With the name suffix in place the layout is measurable on the flatmm split-K
family for the first time, and the answer over the pool is no: best
shuffle_scale kid against best plain-scale kid, both drawing the whole pool, is
-0.80% at K=1024, -4.25% at K=4096 and -5.74% at K=8192, 4 of 36 cells better by
more than 1% and none at K=8192. So kid328-333 come out again; the block that
replaces them keeps the numbers and the dict, commented, for whoever reopens it.
Two mis-framed comparisons inverted that verdict twice on the way, and both are
worth keeping:
* the first pass had PREFETCH_SCALE off, which is the axis the two mechanisms
compete on -- it is worth 1.139x to the layout on flatmm, enough to turn its
best tile from a loss into a tie. kid334/335 are the bdirect twins with it
on, and there it is worth only 1-2%.
* comparing against the table's pick prices the table's sub-optimality along
with the layout. Full pool on both sides is what the numbers above use.
Measured at its best, the verdict is per tile rather than global: the shuffled
read beats the LDS panel on kid334's 64x32x256 (1.108x/1.041x/1.038x) and loses
a fifth on kid335's 128x128x128 (0.984x/0.805x/0.793x). Those two differ in four
terms at once, so kids 336-345 walk the path one term at a time and find B_K
carrying the largest attributable share (0.907x at K=8192), B_M little (0.979x),
WG_PER_CU exactly nothing, and a 0.86x residual B_N cannot be separated from --
B_N=128 is only viable at B_K=128, so the orthogonal walk is not available. The
four kids that establish that last point run 105,000-332,000us against 19us for
the baseline geometry, a VGPR spill rather than a traits rejection; 340/341 are
therefore not wired at all and 344/345 are kept only as the record.
The tuner sweeps 336/338/342, the plain halves, since they are ordinary
preload_sf kids at tiles the pool has only at a different B_K or WG_PER_CU.
kid344 stays out: it cannot win a cell and sweeping it costs more than the rest
of the pool together.
Co-authored-by: Cursor <cursoragent@cursor.com>
… to wave8n4 Two instances this family was missing, for different reasons. kid348/349 are the 128x128 tile, which wave8n4 had at 64 columns and at 256 but never between. They win 4 cells of the m=128..512 band by 1.025-1.057x, and below g8 they are far off the pace (9.1us against 4.9 at g2/m128/k1024), which is what a 128-row eight-wave tile should do on a grid of 16 workgroups. They were built to test an explanation of the g16/m256/k4096 gap to Triton, and they refute it. Triton's swept winner there is 31.1us against 40.3us for the best of the preshuffle pool, and its config is a 128x128 tile 256 deep at eight warps -- the same tile at four warps measures 40.34us, kid229's number to three digits, which read as the wave count at a grid of exactly one workgroup per CU. It is not the wave count: kid349 is that tile at eight waves and lands at 39.4us, still 1.27x behind, and the K ordering inverts between the two (Triton wants 256 over 128, kid348 loses to kid349), so they are not responding to the same thing. What is left is the scale path, and it has a measured size at this cell -- the same tile with the panel preloaded is 40.3us and without it 51.7-52.3us -- so the gap is the price of microscaling here rather than a tile that was missing. The comment says so, to keep the next person out of the tile table. kid346 is kid194 plus the banded tile map that only wavetm1 has had. It is worth trying now because the swizzle is live only at split_k=1, and before the derived SF_PRELOAD_K_MAX kid194 could not reach K=16384 on that column at all. It works, and what gates it is the width of the tile grid: +0.4% at n=1024 (the shipped envelope, i.e. noise) rising to +2.2% at n=4096 and n=8192, because the whole effect is the aspect ratio of the tile run an XCD walks. Band 2 was measured alongside and is noise at every width, so kid347 stays commented out. Co-authored-by: Cursor <cursoragent@cursor.com>
…set collision Two of the four cells recorded as slower under preshuffled B are twin against twin -- g16/m128/k4096 and g16/m256/k4096, same tile and same sfpreload, only B's layout differing -- and the note called them the cost of the layout. They are not. The pair compiles to identical VGPR/AGPR/LDS with no spill and issues the same 210 ds_read / 86 buffer_load / 288 MFMA, and across a g x m grid at both K they are a wash on 39 of 40 cells. The exception is the cell where the grid is exactly one occupancy wave, 256 workgroups on 256 CUs. There every volume counter matches within 0.1% and both spread evenly over the 128 channels, so it is not camping; the preshuffled side even takes 3x fewer tag stalls. What the counters cannot see is the L2 set index, which advances per 2 KiB and wraps at 256 KiB: the 64 KiB panel stride puts the tile's 8 chunks on 4 sets, and the n-tile and batch strides are whole multiples of the wrap, so every workgroup picks the same 4. Padding stride_b to 72 KiB takes the cell from 0.87x to 0.98x and does nothing elsewhere, and across a wider sweep every stride landing on 8 sets runs 0.98-1.00x against 0.87-0.94x for <=4. It is the stride and not the shuffle: the shuffle only multiplies B's stride by 16, and forcing a 16 KiB row stride on the row-major baseline costs it 20%, more than preshuffle ever loses. Not shipped -- 12.5% of weight memory for one cell. Co-authored-by: Cursor <cursoragent@cursor.com>
… re-sweep win Four cells go to the tiles added here -- kid349 takes g8/m512 at both K and g16/m256/k4096, kid348 takes g16/m256/k1024, kid346 takes g4/m3584/k4096 and g16/m1024/k1024 -- by 1.025-1.118x over the incumbent. The other six are the same pool measured better. Stepping the candidates in kid order gives whoever runs first a several-percent head start, so these were re-drawn with the order rotated and ranked on per-draw values rather than on a median over a full-pool pass: kid194 takes two large-M k4096 cells off kid205 by 1.12-1.14x, and kid179/kid243 take one each at small m. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
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
1 participant
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist