[PR #186] Add split-k optimization for sm90, reduce through DSMEM. #284

New issue

Open

opened 2026-10-26 21:15:46 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:15:46 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/186
Author: @Insideyyy
Created: 9/5/2025
Status: 🔄 Open

Base: main ← Head: Insideyyy/split-k-sm90

📝 Commits (2)

c6e6b82 Add split-k optimization for sm90, reducing through DSMEM.
aecf90d Resolve conflicts; support split-k for sm90_fp8_gemm_1d1d.

📊 Changes

14 files changed (+410 additions, -140 deletions)

View changed files

📝 csrc/jit/handle.hpp (+12 -16)
📝 csrc/jit/kernel_runtime.hpp (+11 -8)
📝 csrc/jit_kernels/heuristics/common.hpp (+101 -46)
📝 csrc/jit_kernels/heuristics/sm100.hpp (+9 -0)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+26 -0)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+8 -6)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp (+4 -2)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp (+8 -6)
📝 deep_gemm/include/deep_gemm/common/scheduler.cuh (+7 -5)
📝 deep_gemm/include/deep_gemm/common/sm90_utils.cuh (+4 -4)
➕ deep_gemm/include/deep_gemm/common/split_k_reduce.cuh (+65 -0)
📝 deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh (+47 -12)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh (+48 -15)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh (+60 -20)

📄 Description

This PR adds split-k optimization for sm90, reduce partitioned d through DSMEM.
Currently support fp8 & bf16 Normal, MGroupedContiguous, MGroupedMasked gemms on sm90.

fp8_gemm_1d2d on H20:

m x n x k	TFLOPS w/o split-k	TFLOPS w/ split-k (optional)
128 x 64 x 8192	12	21
128 x 128 x 8192	24	35
128 x 256 x 8192	47	64
128 x 1024 x 8192	137	137
128 x 1280 x 8192	137	151
256 x 64 x 8192	24	32
256 x 128 x 8192	47	64
256 x 256 x 8192	93	93
256 x 1024 x 8192	181	180
256 x 1280 x 8192	190	198

bf16_gemm on H20:

m x n x k	TFLOPS w/o split-k	TFLOPS w/ split-k (optional)
128 x 64 x 8192	7	15
128 x 128 x 8192	13	25
128 x 256 x 8192	26	42
128 x 1024 x 8192	76	76
128 x 1280 x 8192	76	90
256 x 64 x 8192	13	21
256 x 128 x 8192	26	41
256 x 256 x 8192	52	51
256 x 1024 x 8192	99	99
256 x 1280 x 8192	104	112

Notes:

Split-k is enabled automatically if possible to improve SM utilization.
The k_slices partitions of same (m_block_idx, n_block_idx) are assigned to k_slices SMs within a thread block cluster, so that the intermediate results could be reduced through DSMEM.

_{🔄 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/186 **Author:** [@Insideyyy](https://github.com/Insideyyy) **Created:** 9/5/2025 **Status:** 🔄 Open **Base:** `main` ← **Head:** `Insideyyy/split-k-sm90` --- ### 📝 Commits (2) - [`c6e6b82`](https://github.com/deepseek-ai/DeepGEMM/commit/c6e6b82d0d30cf4724009732f3f48feafcc8297a) Add split-k optimization for sm90, reducing through DSMEM. - [`aecf90d`](https://github.com/deepseek-ai/DeepGEMM/commit/aecf90d4ff9dde57d8aca54458b26b003369eb4a) Resolve conflicts; support split-k for sm90_fp8_gemm_1d1d. ### 📊 Changes **14 files changed** (+410 additions, -140 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit/handle.hpp` (+12 -16) 📝 `csrc/jit/kernel_runtime.hpp` (+11 -8) 📝 `csrc/jit_kernels/heuristics/common.hpp` (+101 -46) 📝 `csrc/jit_kernels/heuristics/sm100.hpp` (+9 -0) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+26 -0) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+8 -6) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp` (+4 -2) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp` (+8 -6) 📝 `deep_gemm/include/deep_gemm/common/scheduler.cuh` (+7 -5) 📝 `deep_gemm/include/deep_gemm/common/sm90_utils.cuh` (+4 -4) ➕ `deep_gemm/include/deep_gemm/common/split_k_reduce.cuh` (+65 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh` (+47 -12) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh` (+48 -15) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh` (+60 -20) </details> ### 📄 Description This PR adds split-k optimization for sm90, reduce partitioned d through DSMEM. Currently support fp8 & bf16 `Normal`, `MGroupedContiguous`, `MGroupedMasked` gemms on sm90. fp8_gemm_1d2d on H20: | m x n x k | TFLOPS w/o split-k | TFLOPS w/ split-k (optional)| | :--- | :--- | :--- | | 128 x 64 x 8192 | 12 | 21 | | 128 x 128 x 8192 |24|35 | | 128 x 256 x 8192 |47|64 | | 128 x 1024 x 8192 |137|137 | | 128 x 1280 x 8192 |137|151 | | 256 x 64 x 8192 |24|32 | | 256 x 128 x 8192 |47|64 | | 256 x 256 x 8192 |93|93 | | 256 x 1024 x 8192 |181|180 | | 256 x 1280 x 8192 |190|198 | bf16_gemm on H20: | m x n x k | TFLOPS w/o split-k | TFLOPS w/ split-k (optional)| | :--- | :--- | :--- | 128 x 64 x 8192 | 7 | 15 128 x 128 x 8192 | 13 | 25 128 x 256 x 8192 | 26 | 42 128 x 1024 x 8192 | 76 | 76 128 x 1280 x 8192 | 76 | 90 256 x 64 x 8192 | 13 | 21 256 x 128 x 8192 | 26 | 41 256 x 256 x 8192 | 52 | 51 256 x 1024 x 8192 | 99 | 99 256 x 1280 x 8192 | 104 | 112 Notes: - Split-k is enabled automatically if possible to improve SM utilization. - The `k_slices` partitions of same `(m_block_idx, n_block_idx)` are assigned to `k_slices` SMs within a thread block cluster, so that the intermediate results could be reduced through DSMEM. --- <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:15:46 +00:00

bsakaguchi referenced this issue

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

[PR #284] Windows support #342