[PR #316] [MERGED] Add various optimizations and Mega MoE benchmarks #361

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/316
Author: @zheanxu
Created: 4/24/2026
Status: ✅ Merged
Merged: 4/24/2026
Merged by: @LyricZhao

Base: main ← Head: mega-update


📝 Commits (4)

📊 Changes

21 files changed (+1272 additions, -368 deletions)

View changed files

📝 README.md (+29 -29)
📝 csrc/apis/attention.hpp (+37 -14)
📝 csrc/apis/mega.hpp (+37 -18)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+52 -23)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+11 -1)
📝 csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp (+33 -15)
📝 deep_gemm/__init__.py (+1 -1)
📝 deep_gemm/include/deep_gemm/comm/barrier.cuh (+9 -0)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh (+2 -1)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh (+72 -58)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+45 -29)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh (+69 -44)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh (+7 -4)
📝 deep_gemm/include/deep_gemm/layout/mega_moe.cuh (+15 -10)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+6 -2)
📝 deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh (+153 -28)
📝 deep_gemm/mega/__init__.py (+4 -4)
➕ scripts/quick_plot_pm.py (+448 -0)
➕ scripts/run_ncu_mega_moe.sh (+89 -0)
📝 tests/test_attention.py (+56 -21)

...and 1 more files

📄 Description

We benchmarked Mega MoE on DeepSeek-V4-Flash and DeepSeek-V4-Pro under 8-way expert parallelism (EP8), testing at various batch sizes (i.e., the number of tokens per rank) to cover different serving scenarios. All values are averaged across 8 ranks.

DeepSeek-V4-Flash

DeepSeek-V4-Flash has 256 experts with top-k=6 (each token is routed to 6 experts), a hidden dimension of 4096, and an intermediate hidden dimension of 2048.

Batch Size Time (us) Compute (TFLOPS) Global Memory (GB/s) Interconnect (GB/s) Speedup (vs legacy)
1 56.5 5 1311 1 1.96x
512 146.5 1056 3192 266 1.73x
8192 1283.1 1928 998 499 1.56x
32768 4855.5 2038 794 529 1.62x

DeepSeek-V4-Pro

DeepSeek-V4-Pro has 384 experts with top-k=6, a hidden dimension of 7168, and an intermediate hidden dimension of 3072.

Batch Size Time (us) Compute (TFLOPS) Global Memory (GB/s) Interconnect (GB/s) Speedup (vs legacy)
1 108.1 7 1758 1 1.61x
512 369.6 1098 4619 182 1.54x
8192 2818.5 2304 1094 393 1.50x
32768 10655.2 2438 692 417 1.54x

🔄 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/316 **Author:** [@zheanxu](https://github.com/zheanxu) **Created:** 4/24/2026 **Status:** ✅ Merged **Merged:** 4/24/2026 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `mega-update` --- ### 📝 Commits (4) - [`a266725`](https://github.com/deepseek-ai/DeepGEMM/commit/a266725c2dbada1547845ef16afe0f6fa003ed80) Merge with private repo - [`e1b8e01`](https://github.com/deepseek-ai/DeepGEMM/commit/e1b8e0128448a8e8fb677b2a597343debe945cab) Add Mega MoE Benchmark - [`c1115cc`](https://github.com/deepseek-ai/DeepGEMM/commit/c1115ccb16c88e8b0e24f420e3ddf4ac1e6d2848) Minor fix - [`9fd9460`](https://github.com/deepseek-ai/DeepGEMM/commit/9fd94601a07a2a91f73c6434705db29e1aa382ee) Update ### 📊 Changes **21 files changed** (+1272 additions, -368 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+29 -29) 📝 `csrc/apis/attention.hpp` (+37 -14) 📝 `csrc/apis/mega.hpp` (+37 -18) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+52 -23) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+11 -1) 📝 `csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp` (+33 -15) 📝 `deep_gemm/__init__.py` (+1 -1) 📝 `deep_gemm/include/deep_gemm/comm/barrier.cuh` (+9 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh` (+2 -1) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh` (+72 -58) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+45 -29) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh` (+69 -44) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh` (+7 -4) 📝 `deep_gemm/include/deep_gemm/layout/mega_moe.cuh` (+15 -10) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+6 -2) 📝 `deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh` (+153 -28) 📝 `deep_gemm/mega/__init__.py` (+4 -4) ➕ `scripts/quick_plot_pm.py` (+448 -0) ➕ `scripts/run_ncu_mega_moe.sh` (+89 -0) 📝 `tests/test_attention.py` (+56 -21) _...and 1 more files_ </details> ### 📄 Description We benchmarked Mega MoE on DeepSeek-V4-Flash and DeepSeek-V4-Pro under 8-way expert parallelism (EP8), testing at various batch sizes (i.e., the number of tokens per rank) to cover different serving scenarios. All values are averaged across 8 ranks. ### DeepSeek-V4-Flash DeepSeek-V4-Flash has 256 experts with top-k=6 (each token is routed to 6 experts), a hidden dimension of 4096, and an intermediate hidden dimension of 2048. | Batch Size | Time (us) | Compute (TFLOPS) | Global Memory (GB/s) | Interconnect (GB/s) | Speedup (vs legacy) | |:-----:|:------:|:----:|:----:|:---:|:-----:| | 1 | 56.5 | 5 | 1311 | 1 | 1.96x | | 512 | 146.5 | 1056 | 3192 | 266 | 1.73x | | 8192 | 1283.1 | 1928 | 998 | 499 | 1.56x | | 32768 | 4855.5 | 2038 | 794 | 529 | 1.62x | ### DeepSeek-V4-Pro DeepSeek-V4-Pro has 384 experts with top-k=6, a hidden dimension of 7168, and an intermediate hidden dimension of 3072. | Batch Size | Time (us) | Compute (TFLOPS) | Global Memory (GB/s) | Interconnect (GB/s) | Speedup (vs legacy) | |:-----:|:-------:|:----:|:----:|:---:|:-----:| | 1 | 108.1 | 7 | 1758 | 1 | 1.61x | | 512 | 369.6 | 1098 | 4619 | 182 | 1.54x | | 8192 | 2818.5 | 2304 | 1094 | 393 | 1.50x | | 32768 | 10655.2 | 2438 | 692 | 417 | 1.54x | --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:16:47 +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#361
No description provided.