[Inductor][NVGEMM] Automatically route small-M decode shapes to swap_ab kernels (matching vLLM/SGLang)
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
Currently, config.nvgemm_swap_ab is a manual, off-by-default flag. As a result, NVGEMM loses to cuBLAS on small-M decode shapes (e.g., M=8, 16, 32) because the standard CUTLASS tiles suffer from poor M-axis utilization.
When NVGEMM's swap_ab candidates are enabled, they easily beat cuBLAS at these decode shapes. In the original PR that introduced this (#189771), the flag was left off-by-default solely to save compile time ("it doubles the number of candidates").
But IMO Inductor should adopt the same logic as vLLM, SGLang, FlashInfer, DeepGEMM, where there's not really an advantage to not have swap_ab on by deafult. Under the proposed change, we can automatically generate swap candidates for M <= 64, where prefill compile is unaffected, while unlocking SOTA decode performance.
Here's the short survey of the other approaches:
- vLLM (CUTLASS SM100 FP8): M <= 64 uses swap configs; M > 64 uses non-swap.
- vLLM (CUTLASS SM90 FP8): M <= 64 uses N-specialized swap configs
- vLLM (Triton Fused MoE): Swaps when BLOCK_M < 64 and BLOCK_N >= 64
- vLLM / FlashInfer (DeepGEMM): M < 32 routes to swapAB kernel; M >= 32 routes to standard DeepGEMM
- SGLang (SM120 Blockwise): M <= 128/256 uses swapAB with narrower tiles (tuned via cold-L2 CUPTI + CUDA graphs). Disabled automatically under batch-invariant mode.
## NVGEMM `swap_ab` vs. cuBLAS on B300 (sm_103)
*Measured using CUPTI cold-L2 median of 50. BF16 `[-1, 4096, 4096]`.*
| M | cuBLAS (ATen) | NVGEMM (Default) | NVGEMM (w/ `swap_ab`) | Win vs cuBLAS? | Compile Time (Dflt ➔ Swap) |
| :--- | :--- | :--- | :--- | :--- | :--- |
| **8** | 11.7 µs | 15.2 µs | **10.8 µs** | **Yes (1.08x)** | 4.8s ➔ 18.1s |
| **16** | 13.9 µs | 11.0 µs | **10.8 µs** | **Yes (1.29x)** | 4.8s ➔ 18.1s |
| **32** | 13.3 µs | 11.1 µs | **10.8 µs** | **Yes (1.23x)** | 4.8s ➔ 18.8s |
| **64** | **10.8 µs** | 11.1 µs | 11.3 µs (loses) | No (0.96x) | 4.8s ➔ 7.3s |
| **128** | **10.9 µs** | 13.2 µs | 12.7 µs (loses) | No (0.86x) | 5.6s ➔ 13.1s |
At `M <= 32`, it's strictly better. The only cost is a 2x-3x compile-time penalty, which we can eliminate for large shapes by gating on `M <= 64`.
Our suggestion:
1. Make `nvgemm_swap_ab` a tri-state config: `None` (auto), `True` (always), `False` (never).
2. Introduce `nvgemm_swap_ab_max_m = 64` to establish the boundary.
3. Add a `_use_swap_ab(M)` validation helper in `nv_universal_gemm.py` that checks the size hint.
4. Disable swap automatically under `config.batch_invariant` or `config.deterministic` for this use case.
We are happy to open PR to help move this issue forward.
### Alternatives
_No response_
### Additional context
_No response_
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo
Contributor guide
Assessment
This issue has not been assessed yet.