[PR #350] Add shape-aware recommended alignment for SM90 small-M grouped GEMM #383

New issue

Open

opened 2026-10-26 21:17:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:06 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/350
Author: @qescccczmr
Created: 6/2/2026
Status: 🔄 Open

Base: main ← Head: sm90-small-m-alignment

📝 Commits (2)

146cd35 Add recommended alignment for SM90 small-M grouped GEMM
2ed7b01 Add SM90 small-K grouped GEMM alignment benchmark

📊 Changes

6 files changed (+288 additions, -0 deletions)

View changed files

📝 README.md (+1 -0)
📝 csrc/apis/layout.hpp (+10 -0)
📝 csrc/jit_kernels/heuristics/runtime.hpp (+12 -0)
📝 deep_gemm/utils/layout.py (+1 -0)
➕ tests/test_sm90_small_m_alignment.py (+80 -0)
➕ tools/bench_sm90_small_m_m_grouped.py (+184 -0)

📄 Description

Summary

This PR adds a shape-aware recommended M/K alignment helper for contiguous grouped layouts.

The existing get_theoretical_mk_alignment_for_contiguous_layout() keeps SM90 on the legacy 128 alignment. For SM90 m-grouped contiguous GEMM, this can introduce extra
padded rows for some small-M non-psum shapes.

This PR adds an opt-in recommendation helper that returns 64 only for the empirically validated shape family:

SM90
non-psum m-grouped contiguous layout
num_groups == 4
expected_m == 128
expected_k <= 256

All other shapes fall back to the existing theoretical alignment.

This does not change the default alignment.

Motivation

Grouped contiguous layout requires each expert segment to be aligned by get_mk_alignment_for_contiguous_layout(). In the target small-M case, using 128 alignment pads
the generated test shape from M=640 to M=768.

Reducing the recommended alignment to 64 for the validated shape reduces padded rows while preserving correctness.

Broader scans showed regressions for other shapes such as expected_m=64, expected_m=192, num_groups=8/16, and k>=512, so this helper intentionally keeps the 64
recommendation restricted.

Changes

Add get_recommended_mk_alignment_for_contiguous_layout(use_psum_layout, expected_m, expected_k, expected_num_groups).
Expose the helper through the Python layout API.
Document the helper in README.
Rename the benchmark/test wording from small-K to small-M.
Add SM90 correctness and guardrail coverage for:
- target shape returns 64
- psum layout falls back
- expected_m=64 falls back
- k=512 falls back
- num_groups=8 falls back

Benchmark

Hardware:

NVIDIA H200
SM90
CUDA 12.8
NVCC 12.8

Command:

python tools/bench_sm90_small_m_m_grouped.py \
  --compare-recommended \
  --num-groups-list 4,8,16 \
  --expected-m-list 128,192 \
  --n-list 3072,4096,6144,7168 \
  --k-list 128,256,512,1024 \
  --psum-list false,true \
  --repeat 20

Target shape results:

num_groups=4, expected_m=128, psum=false

n,k,baseline_alignment,recommended_alignment,baseline_m,recommended_m,speedup
3072,128,128,64,768,640,1.0361x
3072,256,128,64,768,640,1.0020x
4096,128,128,64,768,640,1.0944x
4096,256,128,64,768,640,1.0492x
6144,128,128,64,768,640,1.0975x
6144,256,128,64,768,640,1.0610x
7168,128,128,64,768,640,1.1163x
7168,256,128,64,768,640,1.0603x

Non-target shapes keep recommended_alignment=128, including:

- psum=true
- expected_m=192
- k=512/1024

This avoids the regressions observed when applying 64 alignment more broadly.

## Testing

python -m py_compile \
  tools/bench_sm90_small_m_m_grouped.py \
  tests/test_sm90_small_m_alignment.py

On H200:

python tests/test_sm90_small_m_alignment.py


---

<sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/350 **Author:** [@qescccczmr](https://github.com/qescccczmr) **Created:** 6/2/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm90-small-m-alignment` --- ### 📝 Commits (2) - [`146cd35`](https://github.com/deepseek-ai/DeepGEMM/commit/146cd35cd725e629608e40bdd8b3aec8b5a7ca0d) Add recommended alignment for SM90 small-M grouped GEMM - [`2ed7b01`](https://github.com/deepseek-ai/DeepGEMM/commit/2ed7b0197ef0cd749ecbec577796894f9afe171a) Add SM90 small-K grouped GEMM alignment benchmark ### 📊 Changes **6 files changed** (+288 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -0) 📝 `csrc/apis/layout.hpp` (+10 -0) 📝 `csrc/jit_kernels/heuristics/runtime.hpp` (+12 -0) 📝 `deep_gemm/utils/layout.py` (+1 -0) ➕ `tests/test_sm90_small_m_alignment.py` (+80 -0) ➕ `tools/bench_sm90_small_m_m_grouped.py` (+184 -0) </details> ### 📄 Description ## Summary This PR adds a shape-aware recommended M/K alignment helper for contiguous grouped layouts. The existing `get_theoretical_mk_alignment_for_contiguous_layout()` keeps SM90 on the legacy 128 alignment. For SM90 m-grouped contiguous GEMM, this can introduce extra padded rows for some small-M non-psum shapes. This PR adds an opt-in recommendation helper that returns 64 only for the empirically validated shape family: - SM90 - non-psum m-grouped contiguous layout - `num_groups == 4` - `expected_m == 128` - `expected_k <= 256` All other shapes fall back to the existing theoretical alignment. This does not change the default alignment. ## Motivation Grouped contiguous layout requires each expert segment to be aligned by `get_mk_alignment_for_contiguous_layout()`. In the target small-M case, using 128 alignment pads the generated test shape from `M=640` to `M=768`. Reducing the recommended alignment to 64 for the validated shape reduces padded rows while preserving correctness. Broader scans showed regressions for other shapes such as `expected_m=64`, `expected_m=192`, `num_groups=8/16`, and `k>=512`, so this helper intentionally keeps the 64 recommendation restricted. ## Changes - Add `get_recommended_mk_alignment_for_contiguous_layout(use_psum_layout, expected_m, expected_k, expected_num_groups)`. - Expose the helper through the Python layout API. - Document the helper in README. - Rename the benchmark/test wording from small-K to small-M. - Add SM90 correctness and guardrail coverage for: - target shape returns 64 - psum layout falls back - `expected_m=64` falls back - `k=512` falls back - `num_groups=8` falls back ## Benchmark Hardware: - NVIDIA H200 - SM90 - CUDA 12.8 - NVCC 12.8 Command: ```bash python tools/bench_sm90_small_m_m_grouped.py \ --compare-recommended \ --num-groups-list 4,8,16 \ --expected-m-list 128,192 \ --n-list 3072,4096,6144,7168 \ --k-list 128,256,512,1024 \ --psum-list false,true \ --repeat 20 Target shape results: num_groups=4, expected_m=128, psum=false n,k,baseline_alignment,recommended_alignment,baseline_m,recommended_m,speedup 3072,128,128,64,768,640,1.0361x 3072,256,128,64,768,640,1.0020x 4096,128,128,64,768,640,1.0944x 4096,256,128,64,768,640,1.0492x 6144,128,128,64,768,640,1.0975x 6144,256,128,64,768,640,1.0610x 7168,128,128,64,768,640,1.1163x 7168,256,128,64,768,640,1.0603x Non-target shapes keep recommended_alignment=128, including: - psum=true - expected_m=192 - k=512/1024 This avoids the regressions observed when applying 64 alignment more broadly. ## Testing python -m py_compile \ tools/bench_sm90_small_m_m_grouped.py \ tests/test_sm90_small_m_alignment.py On H200: python tests/test_sm90_small_m_alignment.py --- <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:17:06 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:20 +00:00

[PR #383] [MERGED] [SM90] Add FP8 MegaMoE support #403

bsakaguchi referenced this issue

2026-10-26 21:17:39 +00:00

[PR #411] feat: add SM90 FP8 x MXFP4 MegaMoE support #427

bsakaguchi referenced this issue

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