[Kernel][Perf] Add RDNA3 INT8 WMMA GEMM, 1.23-2.80x over torch

[vlluvia]

Summary

Adds an INT8 WMMA GEMM for RDNA3 (gfx11), built on v_wmma_i32_16x16x16_iu8. Supports i32, f32, bf16 and f16 output, plus an optional per-row/per-column dequantisation fused into the epilogue.

On gfx1100 it sustains 90–93% of the hardware WMMA issue ceiling on compute-bound shapes. It delivers 1.23–2.00× over torch._int_mm on 15 of 17 tested shapes, with up to 4.33× on skinny-M workloads, and a median 1.36× over Triton 3.8.

Motivation

torch._int_mm always routes to hipBLASLt regardless of torch.backends.cuda.preferred_blas_library, which costs 25–75% against rocBLAS depending on the shape. rocBLAS is not a drop-in either: its INT8 throughput swings up to 3.0× with the operand layout. This gives FlyDSL a native INT8 GEMM that consumes the K-major layout quantised inference already produces, with no repacking and no tuning run.

Changes

• Double-buffered LDS and 128-bit vectorized copies for the WMMA main loop.
• Shape-aware tile selection that avoids padded work and underfilled grids.
• Atomic split-K accumulation for small-M, large-K workloads.
• i32, f32, bf16 and f16 output, with optional fused row/column dequantisation.
• Optional measured autotuning, benchmark coverage against torch._int_mm and best rocBLAS, and 44 correctness tests.

All files are new, nothing existing is touched, and no new third-party dependency is added.

Two things are worth calling out. Split-K lets several workgroups share one output tile and accumulate into it atomically, so a small M with a long K can still fill the device; integer adds stay exact under any order, so the result is bit-identical to the unsplit kernel, and a scaled epilogue is excluded because it would need the full sum first. pick_tile rejects tiles taller than M — a 64-row tile on a 32-row problem makes the WMMA multiply 32 rows of padding — then takes the widest remaining tile whose grid still puts two workgroups on every processor, counting split-K because it multiplies the same grid.

Performance

gfx1100, 96 CU, i32 output, median of five CUDA Graph rounds after 60 warmups, every result checked bit-exact against torch._int_mm before timing. The table reports only the fastest rocBLAS result for each shape. Triton uses native tl.dot(int8, int8) -> int32, the same A[M,K] and B_T[N,K] inputs without repacking, and the fastest of 10 autotuned tile configurations.

The ceiling is 102.5 TOP/s: RDNA3 retires one 16×16×16 INT8 WMMA per 32 wave32 cycles, and under a saturating GEMM the card holds 2085 MHz at 241 W without throttling.

Shape	FlyDSL TOP/s	% peak	rocBLAS best TOP/s	vs rocBLAS	torch TOP/s	Triton TOP/s
128³	0.91	0.9%	0.98	0.93×	0.71	1.03
256³	5.61	5.5%	6.42	0.87×	4.41	6.74
384³	15.86	15.5%	18.88	0.84×	11.43	16.69
512³	27.27	26.6%	32.70	0.83×	20.45	26.68
1024³	65.86	64.3%	66.62	0.99×	46.58	47.49
2048³	84.37	82.3%	73.35	1.15×	64.41	62.61
4096³	92.53	90.3%	92.72	1.00×	74.98	77.48
6144³	95.10	92.8%	97.38	0.98×	66.67	79.60
32×4096×4096	47.68	46.5%	31.47	1.52×	11.01	18.57
64×4096×4096	62.02	60.5%	52.45	1.18×	21.60	36.41
128×4096×4096	70.50	68.8%	67.80	1.04×	35.28	49.54
256×4096×4096	75.86	74.0%	68.15	1.11×	55.21	55.30
64×8192×8192	73.38	71.6%	71.14	1.03×	37.11	52.14
128×8192×8192	78.66	76.8%	68.31	1.15×	42.68	55.43
4096×512×4096	78.46	76.6%	81.72	0.96×	55.21	37.83
8192×8192×1024	92.48	90.2%	90.90	1.02×	70.63	77.48
1024×8192×8192	90.61	88.4%	94.81	0.96×	73.48	75.55

Median 1.39× over torch._int_mm, 1.36× over Triton, and 1.00× against best rocBLAS. Triton leads at 128³ and 256³ and is within measurement noise at 384³; FlyDSL wins the other 14 shapes, by up to 2.59×. Against rocBLAS, FlyDSL wins the skinny-M shapes by up to 1.52× because split-K and a matched tile fill a grid rocBLAS leaves short; rocBLAS wins the small squares and stays 2–4% ahead on 6144³ and 1024×8192×8192.

The percentages below 90% are not tuning left on the table. Efficiency tracks how much work each workgroup has to amortise its prologue and epilogue over, and converges from below: 81.6% at 2048³, 90.4% at 4096³, 91.8% at 6144³, 92.7% at 8192³.

Nor is the config space: 234 configurations were built and timed at 4096³ and the default pick_tile already returns is the fastest of them. Register block, K depth and wave grid span 2.3% over 60 configs; swizzle group, LDS layout and stagger span 0.3% over 160, which is the informative one, since a memory-bound kernel would respond to them and this one does not. What remains is inside the main loop — one unrolled iteration is 64 v_wmma_i32_16x16x16_iu8 against 48 ds_load_b128 and 45 s_waitcnt, at 146 VGPRs with no spills — and closing it means a deeper software pipeline in codegen rather than a tuning knob.

Testing

• Unit tests added — 44 tests in tests/kernels/test_rdna3_int8_gemm.py, all passing on gfx1100
• Performance benchmarks run — scripts/bench_rdna3_int8_gemm.py --rocblas --torch-int-mm --check

Breaking Changes

None.

[coderfeli]

why inline asm here? we have swait hint. has_side_effects=True, could cause data hazard bug

[vlluvia]

My mistake. Fixed by using rocdl.s_waitcnt(lgkmcnt=0) and gpu.barrier(). Thanks for catching it.

[vlluvia]

/rerun-ci

[coderfeli]

why these scripts?

[vlluvia]

I’ll remove these unnecessary scripts.