[PR #328] [MERGED] Sync nv_dev with upstream #316 (Mega MoE optimizations & benchmarks) #371

New issue

Closed

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

bsakaguchi commented

2026-10-26 21:16:57 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/328
Author: @Barry-Delaney
Created: 5/7/2026
Status: ✅ Merged
Merged: 5/8/2026
Merged by: @RayWang96

Base: nv_dev ← Head: user/jinshik/nv_dev_316

📝 Commits (2)

891d57b Add various optimizations and Mega MoE benchmarks (#316)
5201957 Merge upstream commit 891d57b (#316: Add various optimizations and Mega MoE benchmarks) into nv_dev

📊 Changes

21 files changed (+1302 additions, -381 deletions)

View changed files

📝 README.md (+29 -29)
📝 csrc/apis/attention.hpp (+54 -21)
📝 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 (+41 -21)
📝 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 (+8 -5)
📝 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 (+155 -27)
📝 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 (+58 -21)

...and 1 more files

📄 Description

Summary

Cherry-picks upstream 891d57b (deepseek-ai PR #316: "Add various optimizations and Mega MoE benchmarks") onto nv_dev (currently at c491439, the merge of PR #314).

Conflicts resolved against the NV-side commits from PR #314 (f6d98f2, 12cb7c0, a97b74d):

deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh — keep NV's host-passed num_next_n_atoms runtime parameter, plus 891d57b's kIsVarlen template + binary-search SM work distribution + varlen path.
deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh — keep NV's kNumKVMulticast template + cluster multicast logic; add 891d57b's kIsVarlen template parameter with static_assert(!kIsVarlen) (SM90 doesn't support varlen).
csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp — unified Args struct (num_next_n_atoms + is_varlen + indices); SM90 instantiation passes is_varlen=false explicitly; SM100 TMA-box rule uses 891d57b's (is_varlen or next_n >= 2) ? 2 : 1.
csrc/apis/attention.hpp — host computes num_next_n_atoms to match the post-#316 SM100 kernel (kNextNAtom = (is_varlen or next_n >= 2) ? 2 : 1, kNumNextNAtoms = ceil_div(next_n, kNextNAtom)); SM90 multicast hard-codes num_next_n_atoms = 1. Without this fix next_n=5 over-schedules by 5/3 and triggers IMA.
tests/test_attention.py — keep NV's block_kv ∈ {32, 64} and next_n=4 for SM90 multicast; add 891d57b's outer is_varlen loop (varlen is SM100-only).

Test plan

Validated on B300 SXM6 (sm103, arch_major == 10, 8-GPU) and H200 NVL (sm90, arch_major == 9, 1-GPU).

_{🔄 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/328 **Author:** [@Barry-Delaney](https://github.com/Barry-Delaney) **Created:** 5/7/2026 **Status:** ✅ Merged **Merged:** 5/8/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `user/jinshik/nv_dev_316` --- ### 📝 Commits (2) - [`891d57b`](https://github.com/deepseek-ai/DeepGEMM/commit/891d57b4db1071624b5c8fa0d1e51cb317fa709f) Add various optimizations and Mega MoE benchmarks (#316) - [`5201957`](https://github.com/deepseek-ai/DeepGEMM/commit/520195729c219a58951accfa695953725698a632) Merge upstream commit 891d57b (#316: Add various optimizations and Mega MoE benchmarks) into nv_dev ### 📊 Changes **21 files changed** (+1302 additions, -381 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+29 -29) 📝 `csrc/apis/attention.hpp` (+54 -21) 📝 `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` (+41 -21) 📝 `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` (+8 -5) 📝 `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` (+155 -27) 📝 `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` (+58 -21) _...and 1 more files_ </details> ### 📄 Description ## Summary Cherry-picks upstream `891d57b` (deepseek-ai PR #316: "Add various optimizations and Mega MoE benchmarks") onto `nv_dev` (currently at `c491439`, the merge of PR #314). Conflicts resolved against the NV-side commits from PR #314 (`f6d98f2`, `12cb7c0`, `a97b74d`): - **`deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh`** — keep NV's host-passed `num_next_n_atoms` runtime parameter, plus 891d57b's `kIsVarlen` template + binary-search SM work distribution + varlen path. - **`deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh`** — keep NV's `kNumKVMulticast` template + cluster multicast logic; add 891d57b's `kIsVarlen` template parameter with `static_assert(!kIsVarlen)` (SM90 doesn't support varlen). - **`csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp`** — unified Args struct (`num_next_n_atoms` + `is_varlen` + `indices`); SM90 instantiation passes `is_varlen=false` explicitly; SM100 TMA-box rule uses 891d57b's `(is_varlen or next_n >= 2) ? 2 : 1`. - **`csrc/apis/attention.hpp`** — host computes `num_next_n_atoms` to **match the post-#316 SM100 kernel** (`kNextNAtom = (is_varlen or next_n >= 2) ? 2 : 1`, `kNumNextNAtoms = ceil_div(next_n, kNextNAtom)`); SM90 multicast hard-codes `num_next_n_atoms = 1`. Without this fix `next_n=5` over-schedules by 5/3 and triggers IMA. - **`tests/test_attention.py`** — keep NV's `block_kv ∈ {32, 64}` and `next_n=4` for SM90 multicast; add 891d57b's outer `is_varlen` loop (varlen is SM100-only). ## Test plan Validated on **B300 SXM6** (sm103, `arch_major == 10`, 8-GPU) and **H200 NVL** (sm90, `arch_major == 9`, 1-GPU). --- <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:57 +00:00