A Mojo fp16 GEMM that beats hipBLASLt on a consumer RDNA3 card (RX 7900 XTX), with the receipts

Community Showcase
1

I spent the last few weeks writing inference kernels in Mojo for the one AMD GPU
most people actually own: the RX 7900 XTX, gfx1100. Repo is here, Apache-2:

The short version: on square fp16 GEMM the Mojo WMMA kernel is ahead of
hipBLASLt at all ten sizes I tested, 1.01x at 256 up to 1.35x at 1536, 1.10x at
4096. Same shim, same warm-up, same box.

size 256 512 1024 1536 2048 3072 4096
Mojo 6372 30642 74824 93632 91300 99307 90705
hipBLASLt 6288 26324 63623 69332 80203 97224 82437

(GFLOP/s. Full ten-column table and the exact command in the README.)

I want to be upfront that I did not believe this at first, and you shouldn’t
either until you run it. My first “2x faster than hipBLASLt” number was a bug in
my own shim: per-call workspace allocation, trusting the heuristic’s ordering,
never setting splitK. Fixing those took the vendor from 2497 to 5201 GFLOP/s and
wiped the lead out completely. That’s the moment I started preregistering every
bench: question, instrument, predicted range and the falsifier get committed
before the run, and the misses stay in the file. There are more misses than
hits in bench/. Three of the four rounds on the decode-shape kernel
falsified their own prediction.

Why I think it holds now: hipBLASLt’s tuning effort is on CDNA. RDNA3 has WMMA
but nobody at AMD is grinding tile shapes for a gaming card. A 4x2 warp layout
over a 128x128 block, two LDS buffers, one barrier per K-step, two-deep global
prefetch, XOR-swizzled A, transposed B. 188 VGPR, zero spills. Nothing exotic.
The thing that actually moved the needle was boring: dispatching on block
count so small sizes get a 64x64 tile instead of starving the GPU, and warming
the clocks for 10 s instead of 1 s. That warm-up alone was the difference
between 66k and 91k at 4096³ on the identical binary. If your AMD numbers look
noisy, check that before touching the kernel.

There’s also a working end-to-end decode engine on top of this, bf16 and q8
Qwen GGUF, 64 tokens bit-identical to llama.cpp on the same file. Trunk decode
is 0.94x llama.cpp without speculation. With MTP speculative decode it’s 1.33x
on the race prompt and 0.78x on real text, so llama.cpp still wins where it
counts there, and I say so in the README rather than picking the flattering
number. Persistent megakernel (one launch per decode token, 32 layers plus
head) is the current line of work and is what’s in main today.

Things I’d genuinely like from this forum:

  • Someone with a 7900 XTX or XT running bench/fp16-templates.sh and posting
    what they get. One box is one box.
  • Anyone who knows whether rocdl.flat_work_group_size via StaticTuple is the
    intended way to lift the 256-thread cap from Mojo, or whether I’m holding it
    wrong. It works, but it feels like I found a side door.
  • fp32 WMMA does not exist on gfx1100. I verified it against llvm-mc and the
    builtin table, but if Modular has a cleaner story for fp32 on RDNA3 I’d
    rather use it than my workaround.

Happy to answer anything about the protocol. The whole point of the repo is that
the numbers are checkable, so please check them.

1 Like