[Kernel][Perf] Speed up the RDNA3 GEMM up to 1.85x by choosing the tile from the shape - #1028
Open
vlluvia wants to merge 7 commits into
Open
[Kernel][Perf] Speed up the RDNA3 GEMM up to 1.85x by choosing the tile from the shape#1028vlluvia wants to merge 7 commits into
vlluvia wants to merge 7 commits into
Conversation
This was referenced Aug 19, 2026
Collaborator
Collaborator
|
Large trunks of comments. Clean slightly? |
coderfeli
reviewed
Aug 20, 2026
coderfeli
left a comment
coderfeli
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. @vivienfanghuagood also take a look?
vivienfanghuagood
previously approved these changes
Aug 20, 2026
vivienfanghuagood
left a comment
vivienfanghuagood
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, great work!
Collaborator
There was a problem hiding this comment.
should remove some internal infomations like comments
Contributor
Author
There was a problem hiding this comment.
I'll clean them up.
vlluvia
force-pushed
the
feat/rdna3-gemm-tile-selection
branch
from
August 20, 2026 07:01
1f811d4 to
88f327e
Compare
rdna3_f16_gemm builds whatever tile it is handed and defaults to 128x128x32. That tile is right once the problem fills the grid, but it cuts only 4 workgroups at 256x256 and 16 at 512x512, so on a 96-CU part most CUs idle no matter how good the inner loop is. Choosing the tile from the shape is worth up to 3.0x there. rdna3_f16_gemm_autotune owns that decision in two layers. pick_tile is a heuristic fitted to a sweep of every feasible tile on 27 shapes; it needs no GPU and no measurement, and it is what a call resolves to with nothing configured, so the wrapper benchmarks nothing by default. Above it sits the shared autotuner: FLYDSL_AUTOTUNE=1 sweeps feasible_tiles for real, and the result can be frozen into an offline artifact. The heuristic defaults to 64x64x64 rather than the widest tile that covers the machine. Measured on gfx1100 it is fastest on 16 of the 27 shapes and holds 50-59 TFLOP/s throughout, where 128x128x32 swings between 40 and 72. Taking the widest covering tile cost up to 37% and averaged 6.5%; against the per-shape fastest tile this averages 0.6%, worst case 8.1%. Two limits worth knowing. NUM_CU is hard-coded for gfx1100, so the thresholds do not transfer to a gfx11 part with a different CU count, and shapes outside the fitted set are extrapolation -- the search exists for both cases. And _graph_bench, which captures a CUDA graph to get under the ~90us launch overhead that would otherwise swamp these kernels, still reads the multi-wave tiles a few us high below about 50us, so a tuned result for a short kernel is a hypothesis to confirm rather than a fact. feasible_tiles doubles as the search space: anything it excludes does not divide the shape, cannot fill the prefetch pipeline, or does not fit in LDS, so benchmarking it would only measure a build failure. Points the gfx11 benchmark path at the wrapper so its numbers reflect the chosen tile rather than the default. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Extend the RDNA3 GEMM tile ladder with wide large-shape tiles that fit through an unpadded LDS layout. Signed-off-by: vlluvia <Haodian.Feng@amd.com>
Every workgroup walked the k-tiles in the same order, so the whole machine read the same k-slice of A and B at once. Where the row stride is a power of two those reads pile onto a narrow set of memory channels. Pinning M and N at 2048 so the grid and the ISA never change and scanning only K isolates it: K=2048 (4096-byte rows) runs 67.4 TFLOP/s while 1920, 1984, 2112, 2176 and 2304 all sit at 70.0-71.9. Starting each workgroup at a different k-tile and wrapping around decorrelates them. The wraparound is a mask, so it is one scalar add and one scalar and per trip, and it is only wired up when the k-tile count is a power of two -- which is exactly when the row stride is one too. rocBLAS carries StaggerU=32 on every solution it picks at this shape. The scheduling hints move at the same time, because which tile wants them has reversed. They were once required at 256x256x32 to stop LLVM assigning every A fragment the same register quad; against flydsl 0.3.1 that is fixed and the hint now costs 6% there (8192, 12288 and 16384 square, 0.1-0.4% spread). The two narrower tiles are the ones that want it now. Together on 16 shapes: worst 1.015x, best 1.159x, nothing slower. 2048 cubed goes 254.8 -> 228.9 us, which is 1.117x rocBLAS's best solution under the same tight-leading-dimension constraint. 16384x16384x4096 reaches 89.0 TFLOP/s, 85% of peak. 18-shape numerical regression passes. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com>
A tile reads BLOCK_K of a row at a time, which at K=2048 in bf16 is a 64-byte run every 4096 bytes -- exactly the spacing that camps on one set of L2. Taking the row strides from the caller's tensors instead of assuming they are tight lets a caller who has room to pad step out of that pattern; rocBLAS gains 1.20x on this shape from ld+32 alone (255.71 -> 212.86 us). A padded stride and the k-loop stagger break up the same camping, so a padded caller does not also pay for the stagger: at 2048 cubed, tight with stagger is 224.22 us, ld+64 without it 221.70, and the two together 227.70 -- worse than either alone. Persistent workgroups give each workgroup a band of whole tiles, which hides the grid tail when the tile count is not a multiple of the slot count. Only 0 or num_tiles is accepted, because the effect does not generalise across tiles and going wider is where it turns: 256x256x32 loses 14% at both 16384x16384x2048 and x4096, and the small tiles lose 1-2.5% where the whole kernel is tens of microseconds and one more loop around the body never amortizes. The strides are compile-time arguments of the built module, so they belong in the resolve cache key. Without them a padded slice and a tight operand of the same M, N, K share one entry, and whichever is served from the other's kernel is read at the wrong pitch -- silently wrong results, not a failure. Keyed through a single helper so the lookup and the insert cannot drift apart again. Also drops the unused 128x256 tile from the ladder and derives the LDS K-padding from lds_layout instead of carrying separate pad arguments. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com>
scripts/bench_rdna3_gemm_compare.py times the RDNA3 kernel against Triton max-autotune and torch.mm over one shape list, so the numbers quoted in a performance discussion can be reproduced without reassembling a harness. Triton is compiled per shape with static dims and returns its result rather than copying into an out tensor, so cudagraphs can arm. Sharing one dynamic-shape compile across sizes or timing an in-place copy understates it, which is worth avoiding when the comparison is the point. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: vlluvia <Haodian.Feng@amd.com>
Address review feedback on PR ROCm#1028: remove benchmark tables, sweep numbers, and other internal notes while keeping brief API and algorithm documentation. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
Author
|
Trimmed the internal tuning comments in |
Remove verbose notes from stagger, persistent workgroups, lds_layout, sched_hint, and leading-dimension additions only; leave pre-existing comments unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
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
3 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
rdna3_f16_gemmbuilt whatever tile it was handed and defaulted to 128x128x32. Thisadds
rdna3_f16_gemm_autotune, which picks the tile from the shape, plus threememory-side options the chosen tile drives: a per-workgroup stagger of the k-loop
start, caller-supplied row strides, and persistent workgroups.
On gfx1100 this is 1.2-1.85x below 1024 cubed, where the fixed tile left most of the
machine idle, and 1.1-1.2x from 2048 cubed up, where the win comes from breaking L2
set camping rather than from occupancy. That also brings the kernel level with
rocBLAS at the large squares and ahead of it through 2048-4096.
Motivation
128x128x32 is the right tile once the problem fills the grid, but it cuts only 4
workgroups at 256x256 and 16 at 512x512, so on a 96-CU part most CUs idle no matter
how good the inner loop is.
At the other end the limit is not occupancy but address distribution. A tile reads
BLOCK_K of a row at a time, which at K=2048 in bf16 is a 64-byte run every 4096
bytes -- exactly the spacing that camps on one set of L2. Pinning M and N at 2048 so
the grid and the ISA never change and scanning only K isolates it: K=2048 runs 67.4
TFLOP/s while 1920, 1984, 2112, 2176 and 2304 all sit at 70.0-71.9. Every workgroup
walking the k-tiles in the same order is what concentrates those reads, and rocBLAS
carries
StaggerU=32on every solution it picks at this shape.Changes
kernels/gemm/rdna3_f16_gemm_autotune.py(new).pick_tileis a heuristicfitted to a sweep of every feasible tile on 27 shapes; it needs no GPU and no
measurement, so the wrapper benchmarks nothing by default. Above it sits the shared
autotuner:
FLYDSL_AUTOTUNE=1sweepsfeasible_tilesfor real and the result canbe frozen into an offline artifact.
feasible_tilesdoubles as the search space --anything it excludes does not divide the shape, cannot fill the prefetch pipeline,
or does not fit in LDS.
kernels/gemm/rdna3_f16_gemm.py. Adds the wide tiles, the k-loopstagger,lda/ldb/ldc, andpersistent_wgs. The stagger wraparound is a mask, so itcosts one scalar add and one scalar and per trip, and it is only wired up when the
k-tile count is a power of two -- which is exactly when the row stride is one too.
persistent_wgsaccepts only 0 ornum_tiles.reversed. They were once required at 256x256x32 to stop LLVM assigning every A
fragment the same register quad; against flydsl 0.3.1 that is fixed and the hint
now costs 6% there. The two narrower tiles are the ones that want it now.
not also pay for the stagger: at 2048 cubed, tight with stagger is 224.22 us, ld+64
without it 221.70, and the two together 227.70 -- worse than either alone.
tile rather than the default.
scripts/bench_rdna3_gemm_compare.pyso the numbers below can be reproduced.Impact on existing code
create_wmma_gemm_modulekeeps its old defaults, so a caller that spells out a tileis unaffected. The new arguments are keyword-only with inert defaults
(
stagger=0,persistent_wgs=0,lda=ldb=ldc=Nonemeaning tight).Performance
AMD Radeon PRO W7900D (gfx1100, 96 CU), bf16, TN (
C = A @ B_T.T), one GPU.Baseline is
mainat its default 128x128x32; both sides are timed identically as acaptured-graph replay, because below about 1024 cubed a launch costs more host time
than the kernel takes on the device and a plain Python loop would report dispatch
instead of the kernel.
Nothing regresses; 1536 cubed is the one shape that is already well served by the
fixed tile.
Against the libraries on the same card
First, what the card can actually do, because it is power limited rather than clock
limited: under sustained bf16 WMMA it sits at 240-241 W of a 241 W cap and 1870-1899
MHz at 53 C, and because RDNA3 issues WMMA on the vector ALU rather than a separate
matrix pipe, that clock puts the arithmetic ceiling at roughly 92.6 TFLOP/s, not the
122.6 of the nominal boost clock. Every number below should be read against 92.6.
bf16, M=N=K, each library at whichever of NN/NT is better for it, all on the same
card in one session. rocBLAS is
rocblas-bench --function gemm_ex, Triton isInductor at
max-autotune,torch.mmis this build's default route.torch.mmWhere that leaves the kernel: ahead of rocBLAS through 2048-4096, level with it from
5120 up, behind it at 1536, and ahead of Triton by 11-25% throughout. Against the
~92.6 ceiling above, 87 TFLOP/s is 94%, so most of what is left at the large shapes
is the power limit rather than the schedule.
These numbers are each library's own harness wall clock, not the captured-graph
replay of the table above, so they are quoted only from 1536 up where the two agree
for this kernel within 1.3% (2048: 74.99 plain against 75.24 replay; 8192: 85.88
against 86.98). Below 1536 a launch costs more than the kernel and every row would
be measuring a different host path instead: at 512 cubed this kernel reports 3.30
through a Python loop and 25.77 as a replay, and rocBLAS's 21.05 comes from a tight
C++ loop. Small-shape comparisons need kernel-only timing for all four and are not
claimed here.
rocBLAS NN falls off at 7168 and 8192 -- both multiples of 512 elements, so
ldbisa multiple of 1 KB and the tile loads collide in one L2 set (NN 74.55 and 73.40
against NT 86.18 and 85.83). That is the same effect the stagger addresses here, and
it is why this kernel holds within 1% between TN and NN while rocBLAS spreads 12
points -- so rocBLAS's "best" column above depends on the caller being able to reach
NT or pad the stride.
Testing
tests/kernels/test_rdna_gemm.py: the tilethe heuristic picks for each measured shape, the two properties that make
selection safe to leave on by default (it never returns a tile the shape cannot
build, and it never moves a shape off 128x128x32 while that grid is deep enough),
and a regression test that a padded and a tight operand of one shape do not share
a built module.
scripts/bench_rdna3_gemm_compare.py)Two limits worth stating.
NUM_CUis hard-coded for gfx1100, so the thresholds donot transfer to a gfx11 part with a different CU count, and shapes outside the fitted
set are extrapolation -- the search exists for both cases.
Dependencies
Breaking Changes
None. All new arguments default to the previous behaviour.