[PR #455] [Fix] SM90 heuristic offers block M 16/32 to the FP8 1D1D kernel (JIT failure for M <= 32) #459

New issue

Open

opened 2026-10-26 21:18:05 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:05 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/455
Author: @LiRunGuo
Created: 9/22/2026
Status: 🔄 Open

Base: main ← Head: fix-sm90-1d1d-small-m

📝 Commits (2)

f84c92a Fix SM90 heuristic offering block M 16/32 to the FP8 1D1D kernel
95880a3 Address review: name small-M wgrad shape, add SM90 small-M k-grouped case

📊 Changes

2 files changed (+25 additions, -11 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/sm90.hpp (+6 -2)
📝 tests/generators.py (+19 -9)

📄 Description

Problem

On SM90, the layout heuristic (csrc/jit_kernels/heuristics/sm90.hpp) adds block M candidates of 16 and 32 when m <= 32 for Normal, Batched and KGroupedContiguous GEMMs. The SM90 FP8 1D2D and BF16 kernels support these small tiles, but the SM90 FP8 1D1D kernel (used for FP32-output / weight-gradient GEMMs and for the k-grouped GEMM) only supports BLOCK_M of 64 or 128:

deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh(61): error: static assertion failed with "Invalid block size"
      static_assert(BLOCK_M % WGMMA::M == 0, "Invalid block size");
deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh(264): error: static assertion failed with "Invalid block sizes"
      static_assert(BLOCK_M == WGMMA::M * (BLOCK_M <= 64 ? 1 : 2), "Invalid block sizes");

For m <= 32 the cost model always prefers the smaller tile (e.g. it picks block_m=32, block_n=56 for m=27, n=5744, k=7936), so JIT compilation fails and the call raises. This affects, on Hopper:

fp8_gemm_nt(a, b, d, c=d) with an FP32 d (the 1D1D wgrad path) whenever m <= 32
k_grouped_fp8_gemm_nt_contiguous whenever m <= 32

Neither shape was covered by the tests (the wgrad shapes only use m from nk_list, all >= 576, and the k-grouped shapes all have m >= 768). Found with a randomized shape fuzzer on an H200.

Fix

Only offer the 16/32 candidates to kernels that support them (kernel_type != Kernel1D1D). With the fix, the heuristic picks block_m=64 for these shapes. Shapes with m > 32 are unaffected, since the candidate list only changes when m <= 32.

Tests (tests/generators.py):

Small-M weight-gradient shapes (m in {16, 32}, n=4096, k=7168, FP32 accumulate) in enumerate_normal. On SM90 they use the FP8 1D1D kernel and the BF16 kernel. On SM100 they should go through swap-AB.
A small-M k-grouped shape (num_groups=8, m=16, n=2048) in enumerate_k_grouped_contiguous, enabled on SM90 only. 8 groups keep TMA multicast disabled, so this shape does not depend on the separate k-grouped multicast fix in #415 / #429.

Testing (H200, CUDA 13.0, PyTorch 2.11.0+cu130, at the PR head)

Suite	Result
`tests/test_fp8_fp4.py`	pass
`tests/test_bf16.py`	pass
`tests/test_layout.py`	pass
`tests/test_einsum.py`	pass
`tests/test_legacy.py`	pass
`tests/test_attention.py`	pass

New test lines:

# test_fp8_fp4.py (FP8 1D1D wgrad and k-grouped)
 > Perf (m=    16, n=  4096, k=  7168, 1D1D, layout=NT, FP32, acc=1):   19.2 us |   49 TFLOPS | 1610 GB/s | 0.81x cuBLAS speedup
 > Perf (m=    32, n=  4096, k=  7168, 1D1D, layout=NT, FP32, acc=1):   17.8 us |  106 TFLOPS | 1778 GB/s | 0.90x cuBLAS speedup
 > Perf (FP8, num_groups= 8, m=   16, n= 2048, k= 1927, gran_k=128, k_alignment=128, psum=0, acc=1, FP32):    7 us |   17 TFLOPS |  952 GB/s
# test_bf16.py
 > Perf (m=    16, n=  4096, k=  7168, layout=TN, FP32, acc=1):    32.7 us |   29 TFLOPS | 1817 GB/s | 0.76x cuBLAS
 > Perf (m=    32, n=  4096, k=  7168, layout=TN, FP32, acc=1):    26.7 us |   70 TFLOPS | 2254 GB/s | 0.93x cuBLAS
 > Perf (num_groups= 8, m=   16, n= 2048, k= 1927, align=128, psum=0, acc=1, FP32):    8 us |   16 TFLOPS | 1513 GB/s
# test_legacy.py (k-grouped)
 > Perf (num_groups=8, m=   16, n= 2048, k= 2304, layout=TN):    9 us |   16 TFLOPS | 1146 GB/s

Before the fix, the same shapes fail at JIT compile time with the static assertions quoted above.

I could only test on SM90. The k-grouped shape is gated to SM90 because I could not check whether the SM100 k-grouped kernel accepts the swap-AB layout its heuristic picks for small M. It could be enabled there if it does.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/455 **Author:** [@LiRunGuo](https://github.com/LiRunGuo) **Created:** 9/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-sm90-1d1d-small-m` --- ### 📝 Commits (2) - [`f84c92a`](https://github.com/deepseek-ai/DeepGEMM/commit/f84c92a5def36a13f442f659f15019a581e488b1) Fix SM90 heuristic offering block M 16/32 to the FP8 1D1D kernel - [`95880a3`](https://github.com/deepseek-ai/DeepGEMM/commit/95880a39794f6e2b1569dd8a2674ed453074ebd0) Address review: name small-M wgrad shape, add SM90 small-M k-grouped case ### 📊 Changes **2 files changed** (+25 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+6 -2) 📝 `tests/generators.py` (+19 -9) </details> ### 📄 Description ## Problem On SM90, the layout heuristic (`csrc/jit_kernels/heuristics/sm90.hpp`) adds block M candidates of 16 and 32 when `m <= 32` for `Normal`, `Batched` and `KGroupedContiguous` GEMMs. The SM90 FP8 1D2D and BF16 kernels support these small tiles, but the SM90 FP8 **1D1D** kernel (used for FP32-output / weight-gradient GEMMs and for the k-grouped GEMM) only supports `BLOCK_M` of 64 or 128: ``` deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh(61): error: static assertion failed with "Invalid block size" static_assert(BLOCK_M % WGMMA::M == 0, "Invalid block size"); deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh(264): error: static assertion failed with "Invalid block sizes" static_assert(BLOCK_M == WGMMA::M * (BLOCK_M <= 64 ? 1 : 2), "Invalid block sizes"); ``` For `m <= 32` the cost model always prefers the smaller tile (e.g. it picks `block_m=32, block_n=56` for `m=27, n=5744, k=7936`), so JIT compilation fails and the call raises. This affects, on Hopper: - `fp8_gemm_nt(a, b, d, c=d)` with an FP32 `d` (the 1D1D wgrad path) whenever `m <= 32` - `k_grouped_fp8_gemm_nt_contiguous` whenever `m <= 32` Neither shape was covered by the tests (the wgrad shapes only use `m` from `nk_list`, all `>= 576`, and the k-grouped shapes all have `m >= 768`). Found with a randomized shape fuzzer on an H200. ## Fix Only offer the 16/32 candidates to kernels that support them (`kernel_type != Kernel1D1D`). With the fix, the heuristic picks `block_m=64` for these shapes. Shapes with `m > 32` are unaffected, since the candidate list only changes when `m <= 32`. Tests (`tests/generators.py`): - Small-M weight-gradient shapes (`m in {16, 32}`, `n=4096, k=7168`, FP32 accumulate) in `enumerate_normal`. On SM90 they use the FP8 1D1D kernel and the BF16 kernel. On SM100 they should go through swap-AB. - A small-M k-grouped shape (`num_groups=8, m=16, n=2048`) in `enumerate_k_grouped_contiguous`, enabled on SM90 only. 8 groups keep TMA multicast disabled, so this shape does not depend on the separate k-grouped multicast fix in #415 / #429. ## Testing (H200, CUDA 13.0, PyTorch 2.11.0+cu130, at the PR head) | Suite | Result | |---|---| | `tests/test_fp8_fp4.py` | pass | | `tests/test_bf16.py` | pass | | `tests/test_layout.py` | pass | | `tests/test_einsum.py` | pass | | `tests/test_legacy.py` | pass | | `tests/test_attention.py` | pass | New test lines: ``` # test_fp8_fp4.py (FP8 1D1D wgrad and k-grouped) > Perf (m= 16, n= 4096, k= 7168, 1D1D, layout=NT, FP32, acc=1): 19.2 us | 49 TFLOPS | 1610 GB/s | 0.81x cuBLAS speedup > Perf (m= 32, n= 4096, k= 7168, 1D1D, layout=NT, FP32, acc=1): 17.8 us | 106 TFLOPS | 1778 GB/s | 0.90x cuBLAS speedup > Perf (FP8, num_groups= 8, m= 16, n= 2048, k= 1927, gran_k=128, k_alignment=128, psum=0, acc=1, FP32): 7 us | 17 TFLOPS | 952 GB/s # test_bf16.py > Perf (m= 16, n= 4096, k= 7168, layout=TN, FP32, acc=1): 32.7 us | 29 TFLOPS | 1817 GB/s | 0.76x cuBLAS > Perf (m= 32, n= 4096, k= 7168, layout=TN, FP32, acc=1): 26.7 us | 70 TFLOPS | 2254 GB/s | 0.93x cuBLAS > Perf (num_groups= 8, m= 16, n= 2048, k= 1927, align=128, psum=0, acc=1, FP32): 8 us | 16 TFLOPS | 1513 GB/s # test_legacy.py (k-grouped) > Perf (num_groups=8, m= 16, n= 2048, k= 2304, layout=TN): 9 us | 16 TFLOPS | 1146 GB/s ``` Before the fix, the same shapes fail at JIT compile time with the static assertions quoted above. I could only test on SM90. The k-grouped shape is gated to SM90 because I could not check whether the SM100 k-grouped kernel accepts the swap-AB layout its heuristic picks for small M. It could be enabled there if it does. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:05 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:08 +00:00

[PR #459] [MERGED] Sparse MQA logits follow-ups for nv_dev: 16-byte paged KV stride, workspace reuse under graph capture #463