Writing high-performance matrix multiplication kernels for Blackwell

[lairv]

https://docs.jax.dev/en/latest/pallas/gpu/blackwell_matmul.html

[m00x]

This is essentially the same as https://siboehm.com/articles/22/CUDA-MMM

I wonder if the same person wrote it.

[arjvik]

The interesting part is this is done in Pallas!

Seems like the Pallas of old has completely been upgraded

[reasonableklout → arjvik]

Pallas has a couple backends, this is the new-ish Mosaic GPU one. AAUI it provides a bunch of low-level APIs for interacting directly with NVIDIA-specific and new Blackwell features like SMEM, TMEM, collective MMA, etc.

What's interesting is that the MGPU team has achieved SOTA Blackwell GEMM performance before Triton (which IIUC is trying to bring up Gluon to reach the same level). All the big players are coming up with their own block-based low-level-ish DSLs for CUDA: OpenAI, NVIDIA, and now Google.

[flakiness → reasonableklout]

So OpenAI has Triton and Google has Pallas. What's the NVIDIA counterpart?

[saagarjha → flakiness]

Tilus/CUTLASS I assume

[flakiness → saagarjha]

Interesting: https://github.com/NVIDIA/tilus Thanks for the pointer!