[PR #332] feat: support sm90_fp8_fp4 kernel #370

Open
opened 2026-10-26 21:16:57 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/332
Author: @zhangxiaolei123456
Created: 5/11/2026
Status: 🔄 Open

Base: main ← Head: main_hopper_fp8_fp4


📝 Commits (10+)

  • 6e12aeb Create test_sm90_fp8_fp4.py
  • b118d88 Create sm90_fp8_fp4_gemm_1d1d.cuh
  • 561cf21 Create sm90_fp8_fp4_gemm_1d2d.cuh
  • 4122a7d Update gemm.hpp
  • b3d92e0 Update init.py
  • 825cbdb Update init.py
  • 9625f4f Update test_sm90_fp8_fp4.py
  • 6a3b1ee Merge pull request #1 from zhangxiaolei123456/zhangxiaolei123456-patch-1
  • 043f007 Create sm90_fp8_fp4_gemm_1d2d.hpp
  • eebc4f4 Update sm90_fp8_fp4_gemm_1d2d.hpp

📊 Changes

15 files changed (+8134 additions, -10 deletions)

View changed files

📝 csrc/apis/gemm.hpp (+53 -0)
📝 csrc/apis/layout.hpp (+17 -2)
➕ csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp (+1492 -0)
➕ csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp (+296 -0)
📝 csrc/jit_kernels/impls/smxx_layout.hpp (+19 -3)
📝 csrc/utils/layout.hpp (+14 -4)
📝 deep_gemm/__init__.py (+19 -0)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh (+396 -0)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh (+713 -0)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh (+3377 -0)
📝 deep_gemm/include/deep_gemm/mma/sm90.cuh (+70 -0)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+40 -0)
📝 deep_gemm/utils/math.py (+82 -1)
➕ tests/test_sm90_fp8_fp4.py (+1022 -0)
➕ tests/test_sm90_int4_a8.py (+524 -0)

📄 Description

For contiguous kernel
direct FP32 B scale case: b.second shape = [groups, N, K/128]

groups m/group n k W4 us W4 GB/s W4 diff FP8 us FP8 GB/s FP8 diff Speedup
8 256 4096 7168 529 296 0.0000 512 536 0.0335 0.97x
8 512 4096 7168 1036 182 0.0000 925 331 0.0338 0.89x
8 1024 4096 7168 1978 128 0.0000 1884 196 0.0338 0.95x
8 2048 4096 7168 3875 98 0.0000 3643 137 0.0336 0.94x
16 256 4096 7168 1041 301 0.0000 927 592 0.0339 0.89x
16 512 4096 7168 1988 190 0.0000 1890 324 0.0335 0.95x
16 1024 4096 7168 3872 130 0.0000 3657 202 0.0339 0.94x
16 2048 4096 7168 7689 99 0.0000 7174 139 0.0337 0.93x
24 256 4096 7168 1489 316 0.0000 1408 584 0.0336 0.95x
24 512 4096 7168 2947 192 0.0000 2781 330 0.0337 0.94x
24 1024 4096 7168 5762 131 0.0000 5426 205 0.0340 0.94x
24 2048 4096 7168 11379 100 0.0000 10608 141 0.0337 0.93x
32 256 4096 7168 1979 317 0.0000 1882 583 0.0339 0.95x
32 512 4096 7168 3852 196 0.0000 3637 337 0.0338 0.94x
32 1024 4096 7168 7630 132 0.0000 7126 208 0.0339 0.93x
32 2048 4096 7168 15137 100 0.0000 14094 141 0.0338 0.93x

For mask kernel
direct FP32 B scale case: b.second shape = [groups, N, K/128]

groups m/group n k W4 us W4 GB/s W4 diff FP8 us FP8 GB/s FP8 diff Speedup
8 1 4096 7168 97 1289 0.0000 134 1803 0.0344 1.39x
8 4 4096 7168 97 1287 0.0000 134 1817 0.0350 1.37x
8 8 4096 7168 97 1301 0.0000 133 1824 0.0344 1.38x
8 16 4096 7168 101 1251 0.0000 134 1829 0.0337 1.32x
8 32 4096 7168 113 1142 0.0000 133 1846 0.0344 1.18x
8 1 7168 2048 52 1195 0.0000 79 1540 0.0363 1.50x
8 4 7168 2048 53 1196 0.0000 79 1548 0.0345 1.49x
8 8 7168 2048 52 1219 0.0000 79 1555 0.0340 1.51x
8 16 7168 2048 56 1161 0.0000 79 1567 0.0348 1.42x
8 32 7168 2048 60 1110 0.0000 79 1596 0.0347 1.31x

direct E8M0 B scale case: b.second shape = [groups, N, K/32]

groups m/group n k W4 us W4 GB/s W4 diff FP8 us FP8 GB/s FP8 diff Speedup
8 1 4096 7168 130 1127 0.0000 132 1843 0.0344 1.01x
8 4 4096 7168 131 1122 0.0000 131 1849 0.0351 1.00x
8 8 4096 7168 131 1125 0.0000 132 1849 0.0344 1.00x
8 16 4096 7168 161 925 0.0000 131 1858 0.0337 0.82x
8 1 7168 2048 68 1086 0.0000 77 1566 0.0349 1.14x
8 4 7168 2048 68 1085 0.0000 77 1573 0.0356 1.13x
8 8 7168 2048 68 1097 0.0000 77 1584 0.0346 1.14x
8 16 7168 2048 80 947 0.0000 78 1589 0.0352 0.97x

🔄 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/332 **Author:** [@zhangxiaolei123456](https://github.com/zhangxiaolei123456) **Created:** 5/11/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `main_hopper_fp8_fp4` --- ### 📝 Commits (10+) - [`6e12aeb`](https://github.com/deepseek-ai/DeepGEMM/commit/6e12aebefb6b9c120945e7f3b3464a807c541856) Create test_sm90_fp8_fp4.py - [`b118d88`](https://github.com/deepseek-ai/DeepGEMM/commit/b118d880cc3bff1e2dc41965c86a36e51d872e61) Create sm90_fp8_fp4_gemm_1d1d.cuh - [`561cf21`](https://github.com/deepseek-ai/DeepGEMM/commit/561cf215ea3464bfec014b49a0ffbdd47eff0b3b) Create sm90_fp8_fp4_gemm_1d2d.cuh - [`4122a7d`](https://github.com/deepseek-ai/DeepGEMM/commit/4122a7d7bd6687be8538871e74c9d562de6ce2d7) Update gemm.hpp - [`b3d92e0`](https://github.com/deepseek-ai/DeepGEMM/commit/b3d92e05b599d12d3e564d3220873779e93c4c5b) Update __init__.py - [`825cbdb`](https://github.com/deepseek-ai/DeepGEMM/commit/825cbdbe2a954764d35eecd3bc0b358577949537) Update __init__.py - [`9625f4f`](https://github.com/deepseek-ai/DeepGEMM/commit/9625f4f8e2dc5470bdc6fe150ff8086418d2af98) Update test_sm90_fp8_fp4.py - [`6a3b1ee`](https://github.com/deepseek-ai/DeepGEMM/commit/6a3b1eea52826591dc05817a3f05efa42d446ff2) Merge pull request #1 from zhangxiaolei123456/zhangxiaolei123456-patch-1 - [`043f007`](https://github.com/deepseek-ai/DeepGEMM/commit/043f007110d13b2154bfbbca053d7388e8733e4b) Create sm90_fp8_fp4_gemm_1d2d.hpp - [`eebc4f4`](https://github.com/deepseek-ai/DeepGEMM/commit/eebc4f473a8b87728fb8f452477b5647efeb261f) Update sm90_fp8_fp4_gemm_1d2d.hpp ### 📊 Changes **15 files changed** (+8134 additions, -10 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/gemm.hpp` (+53 -0) 📝 `csrc/apis/layout.hpp` (+17 -2) ➕ `csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp` (+1492 -0) ➕ `csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp` (+296 -0) 📝 `csrc/jit_kernels/impls/smxx_layout.hpp` (+19 -3) 📝 `csrc/utils/layout.hpp` (+14 -4) 📝 `deep_gemm/__init__.py` (+19 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh` (+396 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh` (+713 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh` (+3377 -0) 📝 `deep_gemm/include/deep_gemm/mma/sm90.cuh` (+70 -0) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+40 -0) 📝 `deep_gemm/utils/math.py` (+82 -1) ➕ `tests/test_sm90_fp8_fp4.py` (+1022 -0) ➕ `tests/test_sm90_int4_a8.py` (+524 -0) </details> ### 📄 Description **For contiguous kernel** direct FP32 B scale case: b.second shape = [groups, N, K/128] groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- 8 | 256 | 4096 | 7168 | 529 | 296 | 0.0000 | 512 | 536 | 0.0335 | 0.97x 8 | 512 | 4096 | 7168 | 1036 | 182 | 0.0000 | 925 | 331 | 0.0338 | 0.89x 8 | 1024 | 4096 | 7168 | 1978 | 128 | 0.0000 | 1884 | 196 | 0.0338 | 0.95x 8 | 2048 | 4096 | 7168 | 3875 | 98 | 0.0000 | 3643 | 137 | 0.0336 | 0.94x 16 | 256 | 4096 | 7168 | 1041 | 301 | 0.0000 | 927 | 592 | 0.0339 | 0.89x 16 | 512 | 4096 | 7168 | 1988 | 190 | 0.0000 | 1890 | 324 | 0.0335 | 0.95x 16 | 1024 | 4096 | 7168 | 3872 | 130 | 0.0000 | 3657 | 202 | 0.0339 | 0.94x 16 | 2048 | 4096 | 7168 | 7689 | 99 | 0.0000 | 7174 | 139 | 0.0337 | 0.93x 24 | 256 | 4096 | 7168 | 1489 | 316 | 0.0000 | 1408 | 584 | 0.0336 | 0.95x 24 | 512 | 4096 | 7168 | 2947 | 192 | 0.0000 | 2781 | 330 | 0.0337 | 0.94x 24 | 1024 | 4096 | 7168 | 5762 | 131 | 0.0000 | 5426 | 205 | 0.0340 | 0.94x 24 | 2048 | 4096 | 7168 | 11379 | 100 | 0.0000 | 10608 | 141 | 0.0337 | 0.93x 32 | 256 | 4096 | 7168 | 1979 | 317 | 0.0000 | 1882 | 583 | 0.0339 | 0.95x 32 | 512 | 4096 | 7168 | 3852 | 196 | 0.0000 | 3637 | 337 | 0.0338 | 0.94x 32 | 1024 | 4096 | 7168 | 7630 | 132 | 0.0000 | 7126 | 208 | 0.0339 | 0.93x 32 | 2048 | 4096 | 7168 | 15137 | 100 | 0.0000 | 14094 | 141 | 0.0338 | 0.93x **For mask kernel** direct FP32 B scale case: b.second shape = [groups, N, K/128] groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- 8 | 1 | 4096 | 7168 | 97 | 1289 | 0.0000 | 134 | 1803 | 0.0344 | 1.39x 8 | 4 | 4096 | 7168 | 97 | 1287 | 0.0000 | 134 | 1817 | 0.0350 | 1.37x 8 | 8 | 4096 | 7168 | 97 | 1301 | 0.0000 | 133 | 1824 | 0.0344 | 1.38x 8 | 16 | 4096 | 7168 | 101 | 1251 | 0.0000 | 134 | 1829 | 0.0337 | 1.32x 8 | 32 | 4096 | 7168 | 113 | 1142 | 0.0000 | 133 | 1846 | 0.0344 | 1.18x 8 | 1 | 7168 | 2048 | 52 | 1195 | 0.0000 | 79 | 1540 | 0.0363 | 1.50x 8 | 4 | 7168 | 2048 | 53 | 1196 | 0.0000 | 79 | 1548 | 0.0345 | 1.49x 8 | 8 | 7168 | 2048 | 52 | 1219 | 0.0000 | 79 | 1555 | 0.0340 | 1.51x 8 | 16 | 7168 | 2048 | 56 | 1161 | 0.0000 | 79 | 1567 | 0.0348 | 1.42x 8 | 32 | 7168 | 2048 | 60 | 1110 | 0.0000 | 79 | 1596 | 0.0347 | 1.31x direct E8M0 B scale case: b.second shape = [groups, N, K/32] groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- 8 | 1 | 4096 | 7168 | 130 | 1127 | 0.0000 | 132 | 1843 | 0.0344 | 1.01x 8 | 4 | 4096 | 7168 | 131 | 1122 | 0.0000 | 131 | 1849 | 0.0351 | 1.00x 8 | 8 | 4096 | 7168 | 131 | 1125 | 0.0000 | 132 | 1849 | 0.0344 | 1.00x 8 | 16 | 4096 | 7168 | 161 | 925 | 0.0000 | 131 | 1858 | 0.0337 | 0.82x 8 | 1 | 7168 | 2048 | 68 | 1086 | 0.0000 | 77 | 1566 | 0.0349 | 1.14x 8 | 4 | 7168 | 2048 | 68 | 1085 | 0.0000 | 77 | 1573 | 0.0356 | 1.13x 8 | 8 | 7168 | 2048 | 68 | 1097 | 0.0000 | 77 | 1584 | 0.0346 | 1.14x 8 | 16 | 7168 | 2048 | 80 | 947 | 0.0000 | 78 | 1589 | 0.0352 | 0.97x --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
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#370
No description provided.