Grouped GEMM for Imbalanced Experts on Blackwell: A WIP Worklog – Gaurav's corner
This post is a worklog on implementing and optimizing a ragged-$M$ grouped GEMM on an NVIDIA B300. It is a work in progress, and I hope to add continuations ...
03 Aug, 2026 This post is a worklog on implementing and optimizing a ragged-M grouped GEMM on an NVIDIA B300. It is a work in progress, and I hope to add continuations as I explore the Blackwell architecture more deeply. Why grouped GEMM? GEMMs found in Mixture-of-experts are heavily dependent on the number of tokens each expert caters to. This turns an otherwise regular matrix multiplication into an irregular scheduling problem as the distribution of tokens can be highly skewed. For a given expert-FFN projection, expert e computes: Ae[Me,K]×We[K,N]=Ce[Me,N] Here, K and N remain fixed,…
saved by
related reading
- How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklogsiboehm.com
- Worklog: Optimising GEMM on NVIDIA H100 for cuBLAS-like Performance (WIP) – Hamza's Bloghamzaelshafie.bearblog.dev
- Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordićaleksagordic.com
- Implementing a fast Tensor Core matmul on the Ada Architecture | spatters.caspatters.ca
- 1.5x faster MoE training with custom MXFP8 kernels · Cursorcursor.com
- Outperforming cuBLAS on H100: a Worklogcudaforfun.substack.com
- GitHub - wangzyon/NVIDIA_SGEMM_PRACTICE: Step-by-step optimization of CUDA SGEMMgithub.com
- AMD GPUs go brrr · Hazy Researchhazyresearch.stanford.edu
- My 2 cents on Fusing GEMM + Top-K + Softmax on SM100 – Subho's research at your service 🫡ighoshsubho.bearblog.dev
- We reverse-engineered Flash Attention 4modal.com
- ⭐️ Fast LLM Inference From Scratchandrewkchan.dev
- Strangely, Matrix Multiplications on GPUs Run Faster When Given "Predictable" Data! [short]thonking.ai