flâneur

Grouped GEMM for Imbalanced Experts on Blackwell: A WIP Worklog – Gaurav's corner

gauravjain.bearblog.dev · 5,200 words · saved by 1 readers

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