[PR #364] [MERGED] Multiple updates and refactorings #393

Closed
opened 2026-10-26 21:17:13 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/364
Author: @zheanxu
Created: 6/23/2026
Status: ✅ Merged
Merged: 6/23/2026
Merged by: @LyricZhao

Base: main ← Head: public-release-2606


📝 Commits (2)

📊 Changes

64 files changed (+5087 additions, -5802 deletions)

View changed files

📝 csrc/apis/attention.hpp (+52 -35)
📝 csrc/apis/gemm.hpp (+101 -45)
📝 csrc/apis/layout.hpp (+25 -15)
📝 csrc/apis/mega.hpp (+154 -45)
📝 csrc/indexing/main.cu (+2 -4)
📝 csrc/jit_kernels/heuristics/config.hpp (+4 -0)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+96 -50)
📝 csrc/jit_kernels/heuristics/runtime.hpp (+1 -1)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+1 -0)
📝 csrc/jit_kernels/impls/sm100_bf16_gemm.hpp (+16 -15)
➕ csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp (+185 -0)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp (+30 -20)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+16 -12)
➖ csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp (+0 -416)
➕ csrc/jit_kernels/impls/sm100_mqa_logits.hpp (+517 -0)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+5 -5)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp (+8 -8)
➕ csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp (+399 -0)
➖ csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp (+0 -328)
➖ csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp (+0 -464)

...and 44 more files

📄 Description

In this PR, we merged the four SM100 fp8/fp4 mqa logits kernels (paged/non-paged) into a single kernel and added flexible support for various combinations of num_heads, head_dim, and reduce precision (float32/bfloat16). This yields ~2% speedup for fp4 non-paged cases, ~10% for fp4 paged varlen (due to scheduler and load balancing improvements), and ~20% when using bf16 precision for reduction.

For Mega MoE, we now support BF16 precision, which also brings significant acceleration to BF16 MoE. Additionally, we implemented a ring buffer that substantially reduces memory usage compared to the previous version.

There are also several other minor improvements.


🔄 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/364 **Author:** [@zheanxu](https://github.com/zheanxu) **Created:** 6/23/2026 **Status:** ✅ Merged **Merged:** 6/23/2026 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `public-release-2606` --- ### 📝 Commits (2) - [`0e4266e`](https://github.com/deepseek-ai/DeepGEMM/commit/0e4266ed652166aa19509f060f894340c6e8ffda) Public release 26/06 - [`feaf586`](https://github.com/deepseek-ai/DeepGEMM/commit/feaf5860cc784df3664b9c00e6032ca7365df43c) Minor fix ### 📊 Changes **64 files changed** (+5087 additions, -5802 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+52 -35) 📝 `csrc/apis/gemm.hpp` (+101 -45) 📝 `csrc/apis/layout.hpp` (+25 -15) 📝 `csrc/apis/mega.hpp` (+154 -45) 📝 `csrc/indexing/main.cu` (+2 -4) 📝 `csrc/jit_kernels/heuristics/config.hpp` (+4 -0) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+96 -50) 📝 `csrc/jit_kernels/heuristics/runtime.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+1 -0) 📝 `csrc/jit_kernels/impls/sm100_bf16_gemm.hpp` (+16 -15) ➕ `csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp` (+185 -0) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp` (+30 -20) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+16 -12) ➖ `csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp` (+0 -416) ➕ `csrc/jit_kernels/impls/sm100_mqa_logits.hpp` (+517 -0) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+5 -5) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp` (+8 -8) ➕ `csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp` (+399 -0) ➖ `csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp` (+0 -328) ➖ `csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp` (+0 -464) _...and 44 more files_ </details> ### 📄 Description In this PR, we merged the four SM100 fp8/fp4 mqa logits kernels (paged/non-paged) into a single kernel and added flexible support for various combinations of `num_heads`, `head_dim`, and reduce precision `(float32/bfloat16)`. This yields ~2% speedup for fp4 non-paged cases, ~10% for fp4 paged varlen (due to scheduler and load balancing improvements), and ~20% when using bf16 precision for reduction. For Mega MoE, we now support BF16 precision, which also brings significant acceleration to BF16 MoE. Additionally, we implemented a ring buffer that substantially reduces memory usage compared to the previous version. There are also several other minor improvements. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:17:13 +00:00
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#393
No description provided.