freenode
AI & ML

PyTorch Inductor to auto-route small decode GEMMs via swap_ab

NVGEMM will match vLLM and SGLang on M<=64 shapes instead of losing to cuBLAS by default.

PyTorch Inductor is set to automatically enable NVGEMM swap_ab candidates for small-M matrix multiplies, closing a decode-path gap where the compiler's default CUTLASS tiles were slower than cuBLAS.

The manual nvgemm_swap_ab flag has been off by default since it was introduced, mainly because enabling it roughly doubles the kernel candidates and compile time. That choice left standard tiles with poor M-axis utilization on typical LLM decode shapes such as M=8, 16, or 32. With swap_ab enabled, the same kernels beat cuBLAS on those sizes; measurements on B300 (sm_103) BF16 GEMMs show wins of about 1.08x to 1.29x for M<=32, while larger M values still favor cuBLAS or non-swap paths.

mlazos argues Inductor should follow the same policy already used in vLLM, SGLang, FlashInfer, and DeepGEMM: generate swap candidates only when M is at most 64. Prefill and large-batch compiles stay unaffected, and the extra compile cost is confined to the shapes that actually benefit. The change is aimed at state-of-the-art decode latency without a permanent global compile-time penalty.