[PR #380] SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32 #401

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/380
Author: @leavelet
Created: 7/15/2026
Status: 🔄 Open

Base: nv_dev ← Head: sm120-opt

📝 Commits (4)

f22c312 SM120: skip M-padding stores + BLOCK_M=32 for m-grouped decode
c52a9cb SM120: BLOCK_N=64 for the BLOCK_M=32 decode layout
5034208 SM120: skip M-padding A reads via warp-cooperative cp.async
91c8c78 SM120: harden padding-skip paths for edge inputs

📊 Changes

6 files changed (+168 additions, -36 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/runtime.hpp (+5 -2)
📝 csrc/jit_kernels/heuristics/sm120.hpp (+24 -5)
📝 csrc/jit_kernels/impls/sm120_bf16_gemm.hpp (+3 -3)
📝 csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp (+18 -3)
📝 deep_gemm/include/deep_gemm/common/sm120_utils.cuh (+26 -0)
📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh (+92 -23)

📄 Description

Motivation

MoE grouped GEMM at decode batch sizes is weight-bandwidth-bound on SM120: almost all DRAM time goes to streaming expert weights, while the M-padding of the contiguous/masked layouts adds pure-waste traffic (padded A-row reads, padded D-row writes) and the BLOCK_M=64 floor inflates padded compute.

This PR removes the padding waste and lifts the BLOCK_M floor, after it, DeepGEMM leads CUTLASS
on all tested MoE shapes (see numbers below) and streams weights at 98%~99% of
the device's measured pure-read DRAM peak.

What's in here (3 commits)

Skip M-padding row stores (e18f68b).
BLOCK_M=32 via a 2x4 warp layout (e18f68b, 5c018f3).
Skip M-padding A reads via warp-cooperative cp.async (7fdc998).

Performance

RTX PRO 6000 Blackwell (188 SMs), M=32 decode tokens, multinomial top-k
routing.

shape (per-expert weights)	DeepGEMM (Before)	DeepGEMM (Now)	CUTLASS	margin
128E topk2, gate_up N=5120 K=2048 (FP8)	378.7	350.2	357.7	+2.1%
128E topk2, down N=2048 K=2560 (FP8)	197.5	182.5	194.4	+6.1%
DSv4-like 128E topk6, gate_up N=4096 K=4096 (FP8xMXFP4)	702.2	645	678	+4.9%
DSv4-like, down N=4096 K=2048 (FP8xMXFP4)	371.1	336	366	+8.2%
GLM5-TP8-like 256E topk8, gate_up N=512 K=6144 (FP8)	397.1	330.8	346.7	+4.6%
GLM5-TP8-like, down N=6144 K=256 (FP8)	213.0	171.1	183.1	+6.6%
GLM5-EP8-like 32E topk8, gate_up N=4096 K=6144 (FP8)	593.3	558.6	571.7	+2.3%
GLM5-EP8-like, down N=6144 K=2048 (FP8)	301.5	283.3	290.1	+2.4%

_{🔄 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/380 **Author:** [@leavelet](https://github.com/leavelet) **Created:** 7/15/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `sm120-opt` --- ### 📝 Commits (4) - [`f22c312`](https://github.com/deepseek-ai/DeepGEMM/commit/f22c31210599a4d6aa4d8ac901de35505e9fa79c) SM120: skip M-padding stores + BLOCK_M=32 for m-grouped decode - [`c52a9cb`](https://github.com/deepseek-ai/DeepGEMM/commit/c52a9cb4907855801184bba456e19e8337c3a0fe) SM120: BLOCK_N=64 for the BLOCK_M=32 decode layout - [`5034208`](https://github.com/deepseek-ai/DeepGEMM/commit/5034208c5b2d275ba4ac619d4a9ce4b301e0397b) SM120: skip M-padding A reads via warp-cooperative cp.async - [`91c8c78`](https://github.com/deepseek-ai/DeepGEMM/commit/91c8c78758b9ca1c37cf404c533f84f7863484d7) SM120: harden padding-skip paths for edge inputs ### 📊 Changes **6 files changed** (+168 additions, -36 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/runtime.hpp` (+5 -2) 📝 `csrc/jit_kernels/heuristics/sm120.hpp` (+24 -5) 📝 `csrc/jit_kernels/impls/sm120_bf16_gemm.hpp` (+3 -3) 📝 `csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp` (+18 -3) 📝 `deep_gemm/include/deep_gemm/common/sm120_utils.cuh` (+26 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh` (+92 -23) </details> ### 📄 Description ## Motivation MoE grouped GEMM at decode batch sizes is weight-bandwidth-bound on SM120: almost all DRAM time goes to streaming expert weights, while the M-padding of the contiguous/masked layouts adds pure-waste traffic (padded A-row reads, padded D-row writes) and the BLOCK_M=64 floor inflates padded compute. This PR removes the padding waste and lifts the BLOCK_M floor, after it, DeepGEMM leads CUTLASS on all tested MoE shapes (see numbers below) and streams weights at 98%~99% of the device's measured pure-read DRAM peak. ## What's in here (3 commits) 1. **Skip M-padding row stores** (`e18f68b`). 2. **BLOCK_M=32 via a 2x4 warp layout** (`e18f68b`, `5c018f3`). 3. **Skip M-padding A reads via warp-cooperative `cp.async`** (`7fdc998`). ## Performance RTX PRO 6000 Blackwell (188 SMs), M=32 decode tokens, multinomial top-k routing. | shape (per-expert weights)| DeepGEMM (Before) | DeepGEMM (Now) | CUTLASS | margin | |---|---|---|---|---| | 128E topk2, gate_up N=5120 K=2048 (FP8)| 378.7 | **350.2** | 357.7 | +2.1% | | 128E topk2, down N=2048 K=2560 (FP8)| 197.5 | **182.5** | 194.4 | +6.1% | | DSv4-like 128E topk6, gate_up N=4096 K=4096 (FP8xMXFP4) | 702.2 | **645** | 678 | +4.9% | | DSv4-like, down N=4096 K=2048 (FP8xMXFP4) | 371.1 | **336** | 366 | +8.2% | | GLM5-TP8-like 256E topk8, gate_up N=512 K=6144 (FP8)| 397.1 | **330.8** | 346.7 | +4.6% | | GLM5-TP8-like, down N=6144 K=256 (FP8) | 213.0| **171.1** | 183.1 | +6.6% | | GLM5-EP8-like 32E topk8, gate_up N=4096 K=6144 (FP8) | 593.3 | **558.6** | 571.7 | +2.3% | | GLM5-EP8-like, down N=6144 K=2048 (FP8)| 301.5 | **283.3** | 290.1 | +2.4% | --- <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:16 +00:00

bsakaguchi referenced this issue

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

[PR #401] feat(native): add standalone CUDA core #419