Repository navigation
Conversation
The narrowest tile the ladder could express was 32x32x64, so any M or N
that was a multiple of 16 but not 32, and any K no 64-deep k-tile
divides, had no feasible tile at all and failed the divisibility assert
at build time. Model shapes hit this constantly: 16 or 48 tokens against
a weight matrix, or a K of 48.
A 16x16 tile makes reg_m * reg_n == 1, and scf.for hands a single-value
loop carry back as the bare accumulator rather than a one-element
sequence, so the k-loop epilogue raised on the slice. That was the only
thing stopping the kernel from building a tile this narrow.
Adds the 32x32x32, 32x16x64, 16x32x64, 16x16x64, 16x16x32 and 16x16x16
rungs -- 16x16x16 is the floor, one wave holding a single WMMA atom --
and routes shapes that cannot keep 96 CUs busy on a 64-wide tile to a
fallback that weighs occupancy and last-round fill against the work a
tile does per byte of LDS traffic.
Measured over 39 shapes against the whole ladder, graph-timed:
* 34 land within 5% of the best tile in the ladder; the 5 that do not
are the pre-existing 64x64x64-vs-128x128x32 cases.
* Of the 17 shapes MEASURED_TILES pins, 2 move: 256x256x4096 and
256x256x1024, both onto a tile 1.10x and 1.31x faster.
* 256x256x256 1.48x, 32x4096x4096 1.23x, 16x4096x4096 1.25x.
* 4096^3 unchanged at 88.6 TF, 2048^3 unchanged.
Signed-off-by: vlluvia <Haodian.Feng@amd.com>
M still had to divide the block tile, so a decode batch of 1 or a prefill
of 3402 tokens had no feasible tile at all. Rounding M up to the tile is
free on this hardware if the descriptors are sized right: a raw buffer
access past num_records is suppressed by the hardware, so the A rows a
partial tile reaches past read as zero -- contributing nothing to the
product -- and the C stores for those rows are dropped. No predicate in
the loop, no extra instruction in the k-pipeline.
So create_wmma_gemm_module takes m_valid, the real row count behind a
tiled M, and bounds the A and C descriptors at it. _tile_workgroups no
longer requires M to divide; N and K still do, because a column past the
end of a row is a valid offset into the next one and would corrupt it.
Padded rows cost full WMMA issue and full A traffic, so feasible_tiles
drops a tile that would spend more than 25% of its rows on padding,
except when nothing else is available. That keeps M=48 on 16x16 and lets
M=3402 reach 128x128 for 1.6% waste.
Measured over 50 shapes against the whole ladder, graph-timed:
* 3402x4096x1024 83.3 TF, 3402x2048x2048 86.0 TF -- a ragged M now runs
at the same tile and nearly the same rate as an aligned one.
* M=1 and M=8 run at all; 1x11008x4096 is 3.54x torch.
* 44 of 50 shapes land within 5% of the best tile in the ladder.
* 528x528x1024 moves to 32x16x64, 1.05x the 16-row tile it used to get.
* 4096^3 unchanged at 88.1 TF, 2048^3 unchanged.
The 6 shapes still off the best are all the pre-existing
64x64x64-vs-128x128x32 case at K<=1024, now visible on 333x4096x1024 and
1536x1536x1024 as well as the 1152 square the table already noted.
Signed-off-by: vlluvia <Haodian.Feng@amd.com>
M already ran on any tile by padding to the block and bounding the A and C descriptors at the real row count. N could not borrow that trick: a column past the end of a row is a valid offset into the next one, so the padded columns would race with that row's own stores instead of falling off the descriptor. Give the C store a predicate built from the identity-coordinate tensor, and bound the B_T descriptor at the real row count so the B rows past N stay free the way the A rows do. The predicate is const_expr-gated on N actually being padded, so an aligned N compiles to the same store it did before. With both edges served, only K has to divide the tile. The autotuner's padding budget now covers the padded output area rather than the padded rows, which is the same rule for an N the tile divides. Measured on W7900 (gfx1100), graph-timed, 57 shapes against the whole ladder: 528x528x1024 23.0 -> 16.8 us padding N to 576 buys the 64-wide tile 1024x50257x512 729.9 us 72.2 TF, 1.25x torch (gpt2 lm_head) 2048x11007x1024 568.0 us 81.3 TF, 1.28x torch 4095x4095x1024 414.9 us 82.8 TF, 1.28x torch 3402x3402x1024 303.0 us 78.2 TF, 1.23x torch 4096x4096x4096 1582.4 us 86.9 TF unchanged The only pinned tile that moves is 528x528x1024, which now reaches 64x64x64. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com>
The padding and masking tests were checking with verify_output, which falls back to a cosine similarity over the whole tensor when the elementwise check fails. An unmasked edge column clobbers twenty elements out of a hundred thousand, which moves that metric by nothing: disabling the N mask left the suite green at a 0.43 max difference. Check those tests elementwise instead, since padding and masking are either exact or broken. With a check that can fail, cover the irregular cases that were missing. One parametrised test does it: back C with (M + guard, ldc) and slice it, so a tight ldc puts a missed column on the next row where the product check sees it, and a padded one puts it in the row's own pad where the guard does. The nineteen shapes run both extents ragged at once, so the corner tile is past the descriptor bound and under the column mask together; a ragged N written into a column slice; f32 output, which stores through a different retile; and shapes on both K ladders. A second test runs every tile of both ladders against a shape that divides none of them, built with the tile's own options the way the autotuner builds it -- the mask comes from the tile's thread partitioning, so a tile whose wave layout it got wrong would otherwise only surface if pick_tile chose it, which for a tile 2.5x wider than the shape it never does. Each tile appears once, at the first K whose ladder offers it: the two ladders share all but one tile each, and the mask sits in the C store rather than the k-loop, so a second K would rebuild the same store. Disabling the N mask fails 23 cases; disabling the C descriptor bound fails 11, so both mechanisms are covered. 121 pass, 17 skip. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The kernel carried 60 comment lines for 616 lines of code, and some described
code that no longer exists: the module docstring documented a store-back row
formula ("g_row = base + 2*si + klane") that moved into the tiled_mma, and two
more pointed at hand-rolled offset tables and a flat _v8_store that tiled
copies replaced.
Also drop what the code already says. BLOCK_M = WMMA_M * reg_m * waves_m was
annotated "# 128", which is only true of the default tile and misleading now
that the ladder runs from 16x16 to 256x256. The gfx11 v16 lane mapping was
spelled out in the docstring, again at lane16, and again above the LDS read
helpers; the G2S thread geometry twice; the M/N edge asymmetry three times.
Each fact stays at the one place it is load-bearing -- the edge asymmetry
where the descriptor bounds are computed, the measured numbers next to the
tiles they pin. Banner comments shrink to one line.
The autotuner had little to cut: most of its comments carry a measured number
or a cache-correctness constraint, and those stay. Gone is a docstring
paragraph re-listing the screens _tile_workgroups already documents, a line
narrating the replay loop under it, and three asides that restated the code.
Kernel 60 comment lines to 31, 616 lines to 586. Comments only: for both
files the AST with docstrings stripped is unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: vlluvia <Haodian.Feng@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This branch has not been deployed
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.
Summary
The RDNA3 WMMA GEMM could only run shapes whose M and N were multiples of
the block tile, and the narrowest tile the ladder could express was 32x32.
It now serves any M and any N on any tile, down to a 16x16 block. K is
the one extent that still has to divide the k-tile.
Motivation
Model shapes do not come pre-rounded. A 3402-token batch, a 50257-wide
lm_head, 16 or 48 tokens against a weight matrix -- none of these satisfied
the old granularity, and the kernel failed a divisibility assert at build
time rather than running slowly. Callers had to pad the operands themselves,
which costs a copy of the whole matrix to avoid a few wasted tiles.
Three separate things enforced the 32x32 floor, and all three are fixed here:
reg_m * reg_n == 1, andscf.forhands a single-value loop carry back as the bare accumulatorrather than a one-element sequence, so the k-loop epilogue raised on the
slice.
Changes
kernels/gemm/rdna3_f16_gemm.pyscf.forcarry so a 16x16 block builds; 16x16x16is the floor, one wave holding a single WMMA atom.
m_valid/n_validparameters: build over whole tiles and tell thekernel how much of the last one is real.
the bound are suppressed by the hardware -- zero on load, which a product
ignores, and dropped on store -- so the M edge costs no instructions.
offset into the next one. The C store takes a real predicate, built from
an identity coordinate tensor through the same thread partitioning as the
store itself. It sits outside the k-loop, so it is paid once per tile.
kernels/gemm/rdna3_f16_gemm_autotune.py16x16x64, 16x16x32, 16x16x16).
feasible_tilesno longer requires M and N to divide; it screens onpadded area instead, rejecting a tile that would pad past 1.25x the real
output. Shapes that cannot keep 96 CUs busy on a 64-wide tile route to a
fallback weighing occupancy and last-round fill against work per byte of
LDS traffic.
tests/kernels/test_rdna_gemm.py-- see Testing.Performance
Hardware: Radeon PRO W7900 (gfx1100, 96 CUs, wave32), ROCm, bf16 in / bf16
out, CUDA-graph timed on an otherwise idle GPU. Baseline is the fastest tile
the old ladder could legally pick for that shape, taken from a sweep of every
tile against every shape.
Shapes that already built. Nothing regresses, and the narrow rungs pick
up some small-output shapes that the 32-wide floor was leaving on the table:
Shapes that could not build at all. 528x528x1024 shows what each half of
the change is worth: no tile fit it before, M padding alone gets it onto
32x16x64 at 23.0 us, and with the N mask it reaches 64x64x64 at 16.7 us --
padding 528 to 576 in both dimensions wastes 19% of the output and is still
1.38x faster than the narrow tile that needed no N padding.
Against
torch.mm, a sample of the shapes that are new here:Across 34 shapes measured against
torch.mm, FlyDSL wins everywhere, from1.05x to 3.39x.
Testing
tests/kernels/test_rdna_gemm.py: 121 passed, 17 skipped.The existing padding tests were checking with
verify_output, which fallsback to a cosine similarity over the whole tensor. An unmasked edge column
clobbers twenty elements out of a hundred thousand, which does not move that
metric: with the N mask disabled the suite stayed green at a 0.43 max
difference. Those tests now check elementwise, since padding and masking are
either exact or broken.
Coverage added on top of that check:
ragged at once so the corner tile is past the descriptor bound and under
the column mask together. C is backed by a larger buffer and sliced, so a
tight
ldcputs a missed column on the next row where the product checksees it, and a padded one puts it in the row's own pad where a guard does.
Includes f32 output, which stores through a different retile.
with the tile's own options the way the autotuner builds it. The mask comes
from the tile's thread partitioning, so a tile whose wave layout it got
wrong would otherwise only surface if
pick_tilechose it -- which for atile 2.5x wider than the shape it never does.
Both mechanisms are covered by tests that fail when they are removed:
disabling the N mask fails 23 cases, disabling the C descriptor bound fails
11.
Dependencies
None.
Breaking Changes
None.
m_validandn_validdefault toNone, which is the old behaviour,and every shape that worked before builds the same tile and runs in the same
time.
Not covered
K still has to be a multiple of the k-tile (16, 32 or 64, whichever the tile
uses). K has no free edge on either operand -- a k past the end of a row is a
valid offset into the next one, on both A and B_T -- and the predicate would
sit inside the main loop, where masking costs the most.
Separately, and not introduced here:
pick_tileprefers 64x64x64 on somelarge shapes where 128x128x32 measures faster at K <= 1024 (1.26x on
1536x1536x1024, 1.19x on 333x4096x1024). The sweep makes it easy to see now
that it covers more shapes, but it is a large-shape tuning question rather
than a tile-granularity one, so it is left alone.