deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[PR #95] [MERGED] Weight gradient kernels for dense and MoE models #245

New issue
Closed
opened 2026-10-26 21:15:16 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-26 21:15:16 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/95
Author: @zheanxu
Created: 5/6/2025
Status: ✅ Merged
Merged: 5/14/2025
Merged by: @LyricZhao

Base: main ← Head: wgrad-gemm

📝 Commits (7)

📊 Changes

12 files changed (+911 additions, -72 deletions)

View changed files

📝 README.md (+4 -3)
📝 deep_gemm/__init__.py (+2 -0)
📝 deep_gemm/include/deep_gemm/fp8_gemm.cuh (+0 -15)
➕ deep_gemm/include/deep_gemm/fp8_wgrad_gemm.cuh (+362 -0)
📝 deep_gemm/include/deep_gemm/mma_utils.cuh (+26 -0)
📝 deep_gemm/jit_kernels/__init__.py (+4 -0)
📝 deep_gemm/jit_kernels/gemm.py (+43 -21)
📝 deep_gemm/jit_kernels/m_grouped_gemm.py (+18 -14)
📝 deep_gemm/jit_kernels/runtime.py (+131 -10)
➕ deep_gemm/jit_kernels/wgrad_gemm.py (+179 -0)
📝 deep_gemm/utils.py (+12 -5)
📝 tests/test_core.py (+130 -4)

📄 Description

This Pull Request introduces deepgemm.wgrad_gemm_fp8_fp8_fp32_nt and k_grouped_wgrad_gemm_fp8_fp8_fp32_nt, optimized weight gradient kernels for dense and MoE models. These kernels achieve a ~20% speedup compared to the internal CUTLASS implementation.

For detailed usage, refer to the function documentation.

Weight gradient GEMMs for dense models

M N K Opti BMxBN Computation (TFLOPS) Memory Bandwidth (GB/s)
7168 2112 4096 128x152 920 507
1536 24576 4096 128x152 986 582
512 32768 4096 128x152 878 1086
16384 7168 4096 128x152 994 342
7168 4096 4096 128x152 942 411
2048 7168 4096 128x152 920 513
7168 2112 8192 128x152 1052 451
1536 24576 8192 128x152 1092 511
512 32768 8192 128x152 1014 1129
16384 7168 8192 128x152 1079 240
7168 4096 8192 128x152 1061 333
2048 7168 8192 128x152 1037 452

Grouped weight gradient GEMMs for MoE models

Groups M N K Opti BMxBN Computation (TFLOPS) Memory Bandwidth (GB/s)
4 7168 4096 4096 128x152 939 409
4 2048 7168 4096 128x152 900 502
4 7168 4096 8192 128x152 1044 328
4 2048 7168 8192 128x152 1033 450
8 7168 4096 4096 128x152 942 411
8 2048 7168 4096 128x152 902 503

🔄 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/95 **Author:** [@zheanxu](https://github.com/zheanxu) **Created:** 5/6/2025 **Status:** ✅ Merged **Merged:** 5/14/2025 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `wgrad-gemm` --- ### 📝 Commits (7) - [`d5470d3`](https://github.com/deepseek-ai/DeepGEMM/commit/d5470d3b4e7a6e3bbe426887e34dfc2f05aa49ee) Init weight gradient kernels. - [`adf5de0`](https://github.com/deepseek-ai/DeepGEMM/commit/adf5de02445c4b525deca757544411bd9caf1c2e) Merge branch 'main' into wgrad-gemm - [`919f55b`](https://github.com/deepseek-ai/DeepGEMM/commit/919f55be9c77afdd593a0207c4bb42d7d7b99023) Support unaligned n,k and gmem stride - [`6233709`](https://github.com/deepseek-ai/DeepGEMM/commit/6233709c67d50ec4583d33c54dd3b1a4bec9b5f2) Update docs - [`c4a7116`](https://github.com/deepseek-ai/DeepGEMM/commit/c4a7116e0aee19f881076bbb66355508f5dcfe44) Several cleanups - [`279eb03`](https://github.com/deepseek-ai/DeepGEMM/commit/279eb03190639fc267471a15e00da4c665f2202a) Remove restrictions on N - [`a6ced6f`](https://github.com/deepseek-ai/DeepGEMM/commit/a6ced6f2079003934c7993457ac1585859043de8) Add stride(0) assertions ### 📊 Changes **12 files changed** (+911 additions, -72 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+4 -3) 📝 `deep_gemm/__init__.py` (+2 -0) 📝 `deep_gemm/include/deep_gemm/fp8_gemm.cuh` (+0 -15) ➕ `deep_gemm/include/deep_gemm/fp8_wgrad_gemm.cuh` (+362 -0) 📝 `deep_gemm/include/deep_gemm/mma_utils.cuh` (+26 -0) 📝 `deep_gemm/jit_kernels/__init__.py` (+4 -0) 📝 `deep_gemm/jit_kernels/gemm.py` (+43 -21) 📝 `deep_gemm/jit_kernels/m_grouped_gemm.py` (+18 -14) 📝 `deep_gemm/jit_kernels/runtime.py` (+131 -10) ➕ `deep_gemm/jit_kernels/wgrad_gemm.py` (+179 -0) 📝 `deep_gemm/utils.py` (+12 -5) 📝 `tests/test_core.py` (+130 -4) </details> ### 📄 Description This Pull Request introduces `deepgemm.wgrad_gemm_fp8_fp8_fp32_nt` and `k_grouped_wgrad_gemm_fp8_fp8_fp32_nt`, optimized weight gradient kernels for dense and MoE models. These kernels achieve a ~20% speedup compared to the internal CUTLASS implementation. For detailed usage, refer to the function documentation. ### Weight gradient GEMMs for dense models | M | N | K | Opti BMxBN | Computation (TFLOPS) | Memory Bandwidth (GB/s) | |:-----:|:-----:|:-----:|:-----------:|:----------------------:|:-------------------------:| | 7168 | 2112 | 4096 | 128x152 | 920 | 507 | | 1536 | 24576 | 4096 | 128x152 | 986 | 582 | | 512 | 32768 | 4096 | 128x152 | 878 | 1086 | | 16384 | 7168 | 4096 | 128x152 | 994 | 342 | | 7168 | 4096 | 4096 | 128x152 | 942 | 411 | | 2048 | 7168 | 4096 | 128x152 | 920 | 513 | | 7168 | 2112 | 8192 | 128x152 | 1052 | 451 | | 1536 | 24576 | 8192 | 128x152 | 1092 | 511 | | 512 | 32768 | 8192 | 128x152 | 1014 | 1129 | | 16384 | 7168 | 8192 | 128x152 | 1079 | 240 | | 7168 | 4096 | 8192 | 128x152 | 1061 | 333 | | 2048 | 7168 | 8192 | 128x152 | 1037 | 452 | ### Grouped weight gradient GEMMs for MoE models | Groups | M | N | K | Opti BMxBN | Computation (TFLOPS) | Memory Bandwidth (GB/s) | |:------:|:-----:|:-----:|:-----:|:-----------:|:----------------------:|:-------------------------:| | 4 | 7168 | 4096 | 4096 | 128x152 | 939 | 409 | | 4 | 2048 | 7168 | 4096 | 128x152 | 900 | 502 | | 4 | 7168 | 4096 | 8192 | 128x152 | 1044 | 328 | | 4 | 2048 | 7168 | 8192 | 128x152 | 1033 | 450 | | 8 | 7168 | 4096 | 4096 | 128x152 | 942 | 411 | | 8 | 2048 | 7168 | 4096 | 128x152 | 902 | 503 | --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:15:16 +00:00
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
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#245
No description provided.