Skip to content

Writing high-performance matrix multiplication kernels for Blackwell

docs.jax.dev
63 pointslairv6 comments
On HN

Comments

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

I wonder if the same person wrote it.

The interesting part is this is done in Pallas!

Seems like the Pallas of old has completely been upgraded

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.

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

Tilus/CUTLASS I assume

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

AboutSource Built by g1lg1l

Hackerly is an independent reader for Hacker News, built on the public HN API. Not affiliated with Y Combinator.