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 #192] support swapAB for m_grouped_fp8_gemm_nt_masked #286

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/192
Author: @Wangzheee
Created: 9/11/2025
Status: 🔄 Open

Base: main ← Head: swap_ab_for_mask

📝 Commits (1)

  • 2991c77 support swapAB for m_grouped_fp8_gemm_nt_masked

📊 Changes

7 files changed (+456 additions, -6 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/common.hpp (+6 -0)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+17 -4)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp (+8 -1)
📝 deep_gemm/include/deep_gemm/common/sm90_utils.cuh (+11 -0)
📝 deep_gemm/include/deep_gemm/common/utils.cuh (+6 -0)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh (+407 -0)
📝 tests/generators.py (+1 -1)

📄 Description

SwapAB: Significantly improve the performance for M%64<32

Description

  • Significantly improve the performance for BLOCK_M = 32 or M%64<32
  • Swap A B: WGMMA::wgmma(desc_b, desc_a, accum, k)
  • Use multi-math_warp_groups and multi-wave MMAs support BLOCK_N(64, 128, 256): BLOCK_N=256 for H20

How to use

  • export ENABLE_SWAPAB=1

Improvements (H20)

Aligned M, desired state: masked_m[j] = int(expected_m_per_group * random.uniform(1, 1))

  • Original

Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 257 us | 117 TFLOPS | 1862 GB/s
Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 127 us | 118 TFLOPS | 1914 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 497 us | 181 TFLOPS | 993 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 240 us | 188 TFLOPS | 1085 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 921 us | 163 TFLOPS | 554 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 472 us | 159 TFLOPS | 587 GB/s

  • SwapAB

Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 154 us | 195 TFLOPS | 3100 GB/s
Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 84 us | 180 TFLOPS | 2906 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 340 us | 266 TFLOPS | 1454 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 191 us | 236 TFLOPS | 1364 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 572 us | 263 TFLOPS | 891 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 303 us | 248 TFLOPS | 915 GB/s

Other case (original test): masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3))

  • Original

Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 349 us | 214 TFLOPS | 141 GB/s
Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 159 us | 215 TFLOPS | 213 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 348 us | 178 TFLOPS | 216 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 159 us | 191 TFLOPS | 292 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 347 us | 174 TFLOPS | 384 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 217 us | 146 TFLOPS | 353 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 406 us | 153 TFLOPS | 1199 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 172 us | 164 TFLOPS | 1455 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 740 us | 165 TFLOPS | 678 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 354 us | 168 TFLOPS | 758 GB/s

  • SwapAB

Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 304 us | 246 TFLOPS | 162 GB/s
Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 145 us | 235 TFLOPS | 233 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 271 us | 229 TFLOPS | 278 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 136 us | 224 TFLOPS | 342 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 271 us | 223 TFLOPS | 493 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 146 us | 217 TFLOPS | 523 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 308 us | 201 TFLOPS | 1580 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 161 us | 175 TFLOPS | 1556 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 539 us | 227 TFLOPS | 931 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 272 us | 218 TFLOPS | 986 GB/s

TODO

🔄 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/192 **Author:** [@Wangzheee](https://github.com/Wangzheee) **Created:** 9/11/2025 **Status:** 🔄 Open **Base:** `main` ← **Head:** `swap_ab_for_mask` --- ### 📝 Commits (1) - [`2991c77`](https://github.com/deepseek-ai/DeepGEMM/commit/2991c774d0ae7ee066622bcbec27f7a71ff31ef3) support swapAB for m_grouped_fp8_gemm_nt_masked ### 📊 Changes **7 files changed** (+456 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/common.hpp` (+6 -0) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+17 -4) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp` (+8 -1) 📝 `deep_gemm/include/deep_gemm/common/sm90_utils.cuh` (+11 -0) 📝 `deep_gemm/include/deep_gemm/common/utils.cuh` (+6 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh` (+407 -0) 📝 `tests/generators.py` (+1 -1) </details> ### 📄 Description ## SwapAB: Significantly improve the performance for M%64<32 ### Description - Significantly improve the performance for BLOCK_M = 32 or M%64<32 - Swap A B: WGMMA::wgmma(desc_b, desc_a, accum, k) - Use multi-math_warp_groups and multi-wave MMAs support BLOCK_N(64, 128, 256): **BLOCK_N=256 for H20** ### How to use - export ENABLE_SWAPAB=1 ### Improvements (**H20**) #### Aligned M, desired state: masked_m[j] = int(expected_m_per_group * random.uniform(1, 1)) - Original > Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 257 us | 117 TFLOPS | 1862 GB/s > Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 127 us | 118 TFLOPS | 1914 GB/s > Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 497 us | 181 TFLOPS | 993 GB/s > Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 240 us | 188 TFLOPS | 1085 GB/s > Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 921 us | 163 TFLOPS | 554 GB/s > Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 472 us | 159 TFLOPS | 587 GB/s - SwapAB > Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 154 us | 195 TFLOPS | 3100 GB/s > Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 84 us | 180 TFLOPS | 2906 GB/s > Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 340 us | 266 TFLOPS | 1454 GB/s > Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 191 us | 236 TFLOPS | 1364 GB/s > Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 572 us | 263 TFLOPS | 891 GB/s > Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 303 us | 248 TFLOPS | 915 GB/s #### Other case (original test): masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3)) - Original > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 349 us | 214 TFLOPS | 141 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 159 us | 215 TFLOPS | 213 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 348 us | 178 TFLOPS | 216 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 159 us | 191 TFLOPS | 292 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 347 us | 174 TFLOPS | 384 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 217 us | 146 TFLOPS | 353 GB/s > Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 406 us | 153 TFLOPS | 1199 GB/s > Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 172 us | 164 TFLOPS | 1455 GB/s > Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 740 us | 165 TFLOPS | 678 GB/s > Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 354 us | 168 TFLOPS | 758 GB/s - SwapAB > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 304 us | 246 TFLOPS | 162 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 145 us | 235 TFLOPS | 233 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 271 us | 229 TFLOPS | 278 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 136 us | 224 TFLOPS | 342 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 271 us | 223 TFLOPS | 493 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 146 us | 217 TFLOPS | 523 GB/s > Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 308 us | 201 TFLOPS | 1580 GB/s > Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 161 us | 175 TFLOPS | 1556 GB/s > Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 539 us | 227 TFLOPS | 931 GB/s > Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 272 us | 218 TFLOPS | 986 GB/s ### TODO - In LowLatency, expected_m does not represent the maximum possible value of actual shape M in GroupsGEMM. How to choose BLOCK_M and BLOCK_N? - Manually set BLOCK_M and BLOCK_N: https://github.com/deepseek-ai/DeepGEMM/pull/192#discussion_r2341320465 - Set BLOCK_M and BLOCK_N to improve performance for **H100/H800** --- <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 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#286
No description provided.