deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #324] [MERGED] feat: add sm120 support for DeepGEMM #366

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/324
Author: @leavelet
Created: 5/1/2026
Status: ✅ Merged
Merged: 6/24/2026
Merged by: @RayWang96

Base: nv_dev ← Head: sm120

📝 Commits (10+)

373f073 Add SM120a (RTX PRO 6000) dense FP8 GEMM support
6d8c50e Add SM120a dense FP4 GEMM support and consolidate kernel to warp-spec only
18553d2 Add SM120a K-grouped contiguous FP8 GEMM and use custom MMA wrappers
8726331 Add SM120a Einsum (batched GEMM), M-grouped BF16, and HC prenorm TF32
c8a545e Add SM120a FP8 MQA logits (dense + paged) for attention
c346b33 feat: bf16 mn gemm
d0a56e3 Add SM120a skip_head_mid, MN-major B for M-grouped, and FP8 stride_d fix
c5d92cd Optimize MN-major B with ldmatrix.trans.x2, replacing 4 scalar loads
05d2818 Enable BLOCK_N=128 for MN-major B via multi-atom ldmatrix.trans
9996ceb Add SM120a Einsum: BF16 bmk,bnk->mn kernel, FP8 bhd,hdr->bhr and bhd,bhr->hdr

📊 Changes

38 files changed (+6934 additions, -161 deletions)

View changed files

📝 csrc/apis/attention.hpp (+20 -7)
📝 csrc/apis/einsum.hpp (+65 -11)
📝 csrc/apis/gemm.hpp (+156 -23)
📝 csrc/apis/hyperconnection.hpp (+4 -1)
📝 csrc/apis/layout.hpp (+20 -12)
📝 csrc/jit/compiler.hpp (+11 -4)
📝 csrc/jit/device_runtime.hpp (+17 -2)
📝 csrc/jit_kernels/heuristics/config.hpp (+10 -1)
📝 csrc/jit_kernels/heuristics/runtime.hpp (+23 -6)
➕ csrc/jit_kernels/heuristics/sm120.hpp (+333 -0)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+6 -4)
➕ csrc/jit_kernels/impls/sm120_bf16_gemm.hpp (+494 -0)
➕ csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp (+130 -0)
➕ csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp (+682 -0)
➕ csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp (+126 -0)
📝 csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp (+194 -8)
📝 csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp (+213 -9)
📝 csrc/jit_kernels/impls/smxx_layout.hpp (+3 -1)
📝 csrc/utils/layout.hpp (+3 -3)
📝 csrc/utils/math.hpp (+28 -0)

...and 18 more files

📄 Description

This PR adds first-class sm_120 support to DeepGEMM, maintained on the nv_dev branch by the NVIDIA DevTech APAC team. It brings the full DeepGEMM kernel surface — dense, grouped/MoE, einsum, hyper-connection, and MQA-logits — to the sm120 and sm121 devices like RTX Pro 6000 and DGX Spark.

co-authored with @lucifer1004.

Feature surface

Dense GEMM (sm120_fp8_fp4_gemm_1d1d, sm120_bf16_gemm)

FP8 (e4m3, UE8M0 block-scaled, 1D1D), all layouts NT/NN/TN/TT
FP4 (e2m1, UE8M0 mxf4nvf4, gran_k=32)
BF16 and TF32
Mixed precision: FP8_A × FP4_B (mxf8f6f4 .e4m3.e2m1) and FP4_A × FP8_B
(.e2m1.e4m3) — both directions
AB-swap small-M path (M = 1..16) via TMA + per-element strided epilogue
Wave-packing latency heuristic (BLOCK_M 64/128) and split-K for small dims

Grouped GEMM / MoE

M-grouped contiguous, M-grouped masked, K-grouped contiguous
FP8 / FP4 / BF16

Batched GEMM / Einsum (sm120_bmk_bnk_mn)

K-major B (bhr,hdr->bhd) and MN-major B (bhd,hdr->bhr)
Split-S reduction (bmk,bnk->mn); FP8 / BF16; small-M AB-swap

HC Prenorm (sm120_tf32_hc_prenorm_gemm)

Fused TF32 GEMM + square-sum with split-K

**MQA Logits ** (sm120_fp8_mqa_logits, sm120_fp4_mqa_logits + paged)

FP8 / FP4; dense (ragged, warp-specialized, L2-cached KV, split-KV) and
paged (next_n / kPadOddN support)

Performance (RTX PRO 6000 Blackwell, single GPU, `tests/test_*.py`)

Roofline used: FP8 ≈ 814 TFLOPS, FP4 ≈ 1628 TFLOPS (2× FP8), BF16 ≈ 445 TFLOPS. GB/s is effective bandwidth (cache-inclusive for small/cached shapes).

Dense GEMM

dtype	peak TFLOPS	% roofline	peak GB/s
FP8	778	96%	1235
FP4	1561	96%	1337
BF16	444	~98%	1497

Grouped GEMM (MoE)

variant	FP8 TFLOPS	FP4 TFLOPS	peak GB/s
M-grouped contiguous (prefill)	814	1571	—
M-grouped masked	670	1363	836
K-grouped contiguous (EP)	543 (gk32)	666 (gk128)	1011

Einsum

Peak 742 TFLOPS (b=8192,h=8,r=4096,d=1024); up to 2893 GB/s (small-batch, cache-bound); Split-S path scales down for small problems.

HC Prenorm

m=8192, n=24, k=28672, splits=16: 29 TFLOPS / 1248 GB/s (bandwidth-bound).

MQA Logits

variant	peak TFLOPS
Ragged FP4	1049
Paged FP4	685
Ragged FP8	636
Paged FP8	547

_{🔄 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/324 **Author:** [@leavelet](https://github.com/leavelet) **Created:** 5/1/2026 **Status:** ✅ Merged **Merged:** 6/24/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `sm120` --- ### 📝 Commits (10+) - [`373f073`](https://github.com/deepseek-ai/DeepGEMM/commit/373f0738de7a320d71d1e7d1824f012a10416e8b) Add SM120a (RTX PRO 6000) dense FP8 GEMM support - [`6d8c50e`](https://github.com/deepseek-ai/DeepGEMM/commit/6d8c50e38cbfe3af3e2688bbc16a79221c7ff3b2) Add SM120a dense FP4 GEMM support and consolidate kernel to warp-spec only - [`18553d2`](https://github.com/deepseek-ai/DeepGEMM/commit/18553d233ecbc7c3a901ff4b01fc33dc4baabd78) Add SM120a K-grouped contiguous FP8 GEMM and use custom MMA wrappers - [`8726331`](https://github.com/deepseek-ai/DeepGEMM/commit/872633141c30a1ce1d098bdade9ff27682876871) Add SM120a Einsum (batched GEMM), M-grouped BF16, and HC prenorm TF32 - [`c8a545e`](https://github.com/deepseek-ai/DeepGEMM/commit/c8a545e590a58a7a8870ff828583b040b16505b2) Add SM120a FP8 MQA logits (dense + paged) for attention - [`c346b33`](https://github.com/deepseek-ai/DeepGEMM/commit/c346b33d15e04eb804364ebc63e6bd345cef10e4) feat: bf16 mn gemm - [`d0a56e3`](https://github.com/deepseek-ai/DeepGEMM/commit/d0a56e3cfb5f3221530f42e45d5f6231874b2ba8) Add SM120a skip_head_mid, MN-major B for M-grouped, and FP8 stride_d fix - [`c5d92cd`](https://github.com/deepseek-ai/DeepGEMM/commit/c5d92cd3ed52d35eeb2b4f6b56886a5ab64380c7) Optimize MN-major B with ldmatrix.trans.x2, replacing 4 scalar loads - [`05d2818`](https://github.com/deepseek-ai/DeepGEMM/commit/05d2818104b0232153ab6c4a8a1b2f21456eca2d) Enable BLOCK_N=128 for MN-major B via multi-atom ldmatrix.trans - [`9996ceb`](https://github.com/deepseek-ai/DeepGEMM/commit/9996cebcc9bb29a88f09416047dcb15c559230dc) Add SM120a Einsum: BF16 bmk,bnk->mn kernel, FP8 bhd,hdr->bhr and bhd,bhr->hdr ### 📊 Changes **38 files changed** (+6934 additions, -161 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+20 -7) 📝 `csrc/apis/einsum.hpp` (+65 -11) 📝 `csrc/apis/gemm.hpp` (+156 -23) 📝 `csrc/apis/hyperconnection.hpp` (+4 -1) 📝 `csrc/apis/layout.hpp` (+20 -12) 📝 `csrc/jit/compiler.hpp` (+11 -4) 📝 `csrc/jit/device_runtime.hpp` (+17 -2) 📝 `csrc/jit_kernels/heuristics/config.hpp` (+10 -1) 📝 `csrc/jit_kernels/heuristics/runtime.hpp` (+23 -6) ➕ `csrc/jit_kernels/heuristics/sm120.hpp` (+333 -0) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+6 -4) ➕ `csrc/jit_kernels/impls/sm120_bf16_gemm.hpp` (+494 -0) ➕ `csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp` (+130 -0) ➕ `csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp` (+682 -0) ➕ `csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp` (+126 -0) 📝 `csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp` (+194 -8) 📝 `csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp` (+213 -9) 📝 `csrc/jit_kernels/impls/smxx_layout.hpp` (+3 -1) 📝 `csrc/utils/layout.hpp` (+3 -3) 📝 `csrc/utils/math.hpp` (+28 -0) _...and 18 more files_ </details> ### 📄 Description This PR adds first-class `sm_120` support to DeepGEMM, maintained on the `nv_dev` branch by the NVIDIA DevTech APAC team. It brings the full DeepGEMM kernel surface — dense, grouped/MoE, einsum, hyper-connection, and MQA-logits — to the sm120 and sm121 devices like RTX Pro 6000 and DGX Spark. co-authored with @lucifer1004. ## Feature surface **Dense GEMM** (`sm120_fp8_fp4_gemm_1d1d`, `sm120_bf16_gemm`) - FP8 (e4m3, UE8M0 block-scaled, 1D1D), all layouts NT/NN/TN/TT - FP4 (e2m1, UE8M0 `mxf4nvf4`, gran_k=32) - BF16 and TF32 - Mixed precision: FP8_A × FP4_B (`mxf8f6f4 .e4m3.e2m1`) and FP4_A × FP8_B (`.e2m1.e4m3`) — both directions - AB-swap small-M path (M = 1..16) via TMA + per-element strided epilogue - Wave-packing latency heuristic (BLOCK_M 64/128) and split-K for small dims **Grouped GEMM / MoE** - M-grouped contiguous, M-grouped masked, K-grouped contiguous - FP8 / FP4 / BF16 **Batched GEMM / Einsum** (`sm120_bmk_bnk_mn`) - K-major B (`bhr,hdr->bhd`) and MN-major B (`bhd,hdr->bhr`) - Split-S reduction (`bmk,bnk->mn`); FP8 / BF16; small-M AB-swap **HC Prenorm** (`sm120_tf32_hc_prenorm_gemm`) - Fused TF32 GEMM + square-sum with split-K **MQA Logits ** (`sm120_fp8_mqa_logits`, `sm120_fp4_mqa_logits` + paged) - FP8 / FP4; dense (ragged, warp-specialized, L2-cached KV, split-KV) and paged (next_n / kPadOddN support) ## Performance (RTX PRO 6000 Blackwell, single GPU, `tests/test_*.py`) Roofline used: FP8 ≈ 814 TFLOPS, FP4 ≈ 1628 TFLOPS (2× FP8), BF16 ≈ 445 TFLOPS. GB/s is effective bandwidth (cache-inclusive for small/cached shapes). ### Dense GEMM | dtype | peak TFLOPS | % roofline | peak GB/s | |-------|------------:|-----------:|----------:| | FP8 | 778 | 96% | 1235 | | FP4 | 1561 | 96% | 1337 | | BF16 | 444 | ~98% | 1497 | ### Grouped GEMM (MoE) | variant | FP8 TFLOPS | FP4 TFLOPS | peak GB/s | |-------------------------------|-----------:|-----------:|----------:| | M-grouped contiguous (prefill)| 814 | 1571 | — | | M-grouped masked | 670 | 1363 | 836 | | K-grouped contiguous (EP) | 543 (gk32) | 666 (gk128)| 1011 | ### Einsum Peak 742 TFLOPS (`b=8192,h=8,r=4096,d=1024`); up to 2893 GB/s (small-batch, cache-bound); Split-S path scales down for small problems. ### HC Prenorm `m=8192, n=24, k=28672`, splits=16: 29 TFLOPS / 1248 GB/s (bandwidth-bound). ### MQA Logits | variant | peak TFLOPS | |-------------|------------:| | Ragged FP4 | 1049 | | Paged FP4 | 685 | | Ragged FP8 | 636 | | Paged FP8 | 547 | --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#366

No description provided.

Rows
Columns