[PR #314] [MERGED] fix: SM90 paged MQA logits IMA on next_n=4 (nv_dev rebase) #360

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/314
Author: @Barry-Delaney
Created: 4/22/2026
Status: ✅ Merged
Merged: 4/22/2026
Merged by: @RayWang96

Base: nv_dev ← Head: user/jinshik/nv_dev_rebase

📝 Commits (6)

d30fc36 Fix sync issue of TMEM alloc/dealloc (#292)
7f2a703 [Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)
f6d98f2 Enable bf16 bias output for FP8 GEMM on SM100
12cb7c0 SM90 paged MQA logits: support kv_block=32 and next_n=4
a97b74d Fix SM90 paged MQA logits for post-#304 metadata API mismatch
6033584 Sync with upstream main and integrate NV SM90/bf16-bias improvements

📊 Changes

110 files changed (+12337 additions, -3431 deletions)

View changed files

📝 CMakeLists.txt (+4 -5)
📝 README.md (+76 -44)
📝 csrc/apis/attention.hpp (+297 -138)
📝 csrc/apis/einsum.hpp (+35 -34)
📝 csrc/apis/gemm.hpp (+71 -67)
📝 csrc/apis/hyperconnection.hpp (+7 -7)
📝 csrc/apis/layout.hpp (+46 -23)
➕ csrc/apis/mega.hpp (+216 -0)
📝 csrc/apis/runtime.hpp (+20 -0)
📝 csrc/indexing/main.cu (+6 -1)
📝 csrc/jit/cache.hpp (+1 -1)
📝 csrc/jit/compiler.hpp (+70 -61)
📝 csrc/jit/device_runtime.hpp (+46 -7)
📝 csrc/jit/handle.hpp (+72 -18)
➕ csrc/jit/include_parser.hpp (+80 -0)
📝 csrc/jit/kernel_runtime.hpp (+77 -30)
📝 csrc/jit_kernels/heuristics/common.hpp (+34 -319)
➕ csrc/jit_kernels/heuristics/config.hpp (+171 -0)
➕ csrc/jit_kernels/heuristics/mega_moe.hpp (+211 -0)
➕ csrc/jit_kernels/heuristics/runtime.hpp (+62 -0)

...and 80 more files

📄 Description

Summary

Rebases nv_dev onto upstream main (brings in #292, #304), re-applies the two NV-side feature commits, and fixes a SM90 IMA introduced by the #304 metadata API change.

Commits

SHA	Kind	Title
`f6d98f2`	feat	bf16 bias output for FP8 GEMM on SM100
`12cb7c0`	feat	SM90 paged MQA logits: `block_kv=32` / `next_n=4`
`a97b74d`	fix	SM90 paged MQA logits IMA on `next_n=4` post-#304

`f6d98f2` — bf16 bias output for FP8 GEMM on SM100

sm100_fp8_gemm_1d1d / sm100_fp8_fp4_gemm_1d1d with accumulation can now produce FP32 or BF16 matching the bias/C dtype. Adds bf16 bias coverage to enumerate_normal in tests/generators.py.

`12cb7c0` — SM90 paged MQA logits: `block_kv=32` / `next_n=4`

MMA tile stays at kComputeBlockKV = 64; physical block_kv ∈ {32, 64} → kNumBlocksPerMMA = 64 / block_kv TMA copies per WGMMA.
kNumKVMulticast template for next_n = 4: 2-CTA cluster splits Q across CTAs (kNextNPerCTA = kNextN / kNumKVMulticast) to fit SM90 smem/register budgets. WGMMA N=256 otherwise exceeds kNumMathRegisters=104.
Extends the test_attention.py matrix.

`a97b74d` — SM90 paged MQA logits IMA on `next_n=4` post-#304

Symptom — IMA in sm90_fp8_paged_mqa_logits when running next_n=4 on the rebased tree.

Root cause — #304's metadata kernel hard-codes:

num_next_n_atoms = next_n / (next_n % 2 == 0 ? 2 : 1);

arch	strategy	kernel template `kNumNextNAtoms`	hard-coded host value
SM90 `next_n=4`	cluster multicast (no time-atom)	1	2 ❌
SM100 `next_n=4`	time-atomization	2	2 ✓

Metadata schedules 2× the tasks under the 2-atom scheme; SM90 scheduler decodes each entry as 1-atom → q_idx = q_atom_idx / 1 overruns batch_size → reads past context_lens.

Fix

Pipe num_next_n_atoms as a runtime parameter; host computes arch-aware (SM90 multicast → 1, SM100 → 2 for next_n=4).
Relax schedule_meta size to num_sms / num_kv_multicast + 1 for the SM90 cluster grid.
Relax get_paged_mqa_logits_metadata's block_kv assertion to {32, 64} (SM90 block_kv=32 is supported after 12cb7c0).

SM100 behavior unchanged — its num_next_n_atoms value equals the previously hard-coded formula.

Upstream (already reviewed upstream)

d30fc36 Fix sync issue of TMEM alloc/dealloc (#292)
7f2a703 [Public release 26/04] Mega MoE, FP4 Indexer and other features/fixes (#304)

Testing

Test	H200 (SM90)	GB200 (SM100)
`test_attention.py`	✅	✅ (80/80 paged MQA combinations)
`test_bf16.py`	✅	✅
`test_einsum.py`	✅	✅
`test_fp8_fp4.py`	✅	✅
`test_hyperconnection.py`	✅	✅
`test_layout.py`	✅	✅
`test_lazy_init.py`	✅	✅
`test_legacy.py`	✅	✅
`test_mega_moe.py`	⚠️ env	⚠️ env
`compute-sanitizer` (SM90 paged MQA path)	0 errors (24 sub-runs, 1h budget)	—

test_attention.py paged MQA logits covers block_kv ∈ {32,64} × next_n ∈ {1,2,4} × L ∈ {8k,32k} × dtype ∈ {float,bf16} on H200, and is_fp4 ∈ {T,F} × bf16 ∈ {T,F} × block_kv ∈ {32,64} × next_n ∈ {1..6} on GB200.
test_mega_moe.py fails with cuda:N out of range on both arch allocations (pre-existing multi-GPU test, not touched by this PR).

_{🔄 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/314 **Author:** [@Barry-Delaney](https://github.com/Barry-Delaney) **Created:** 4/22/2026 **Status:** ✅ Merged **Merged:** 4/22/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `user/jinshik/nv_dev_rebase` --- ### 📝 Commits (6) - [`d30fc36`](https://github.com/deepseek-ai/DeepGEMM/commit/d30fc36c8f229f4f873b90a492f6e19e6e610923) Fix sync issue of TMEM alloc/dealloc (#292) - [`7f2a703`](https://github.com/deepseek-ai/DeepGEMM/commit/7f2a703ed51ac1f7af07f5e1453b2d3267d37d50) [Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304) - [`f6d98f2`](https://github.com/deepseek-ai/DeepGEMM/commit/f6d98f22a25da9e6b6d1c06b59f3f05e5cc6e726) Enable bf16 bias output for FP8 GEMM on SM100 - [`12cb7c0`](https://github.com/deepseek-ai/DeepGEMM/commit/12cb7c09266890460b61ff10ae60e248eaac158f) SM90 paged MQA logits: support kv_block=32 and next_n=4 - [`a97b74d`](https://github.com/deepseek-ai/DeepGEMM/commit/a97b74d76482866be7e389c1d6f700d032d8e76e) Fix SM90 paged MQA logits for post-#304 metadata API mismatch - [`6033584`](https://github.com/deepseek-ai/DeepGEMM/commit/6033584997f7366dd2edc8f261fabbd7ccf38a72) Sync with upstream main and integrate NV SM90/bf16-bias improvements ### 📊 Changes **110 files changed** (+12337 additions, -3431 deletions) <details> <summary>View changed files</summary> 📝 `CMakeLists.txt` (+4 -5) 📝 `README.md` (+76 -44) 📝 `csrc/apis/attention.hpp` (+297 -138) 📝 `csrc/apis/einsum.hpp` (+35 -34) 📝 `csrc/apis/gemm.hpp` (+71 -67) 📝 `csrc/apis/hyperconnection.hpp` (+7 -7) 📝 `csrc/apis/layout.hpp` (+46 -23) ➕ `csrc/apis/mega.hpp` (+216 -0) 📝 `csrc/apis/runtime.hpp` (+20 -0) 📝 `csrc/indexing/main.cu` (+6 -1) 📝 `csrc/jit/cache.hpp` (+1 -1) 📝 `csrc/jit/compiler.hpp` (+70 -61) 📝 `csrc/jit/device_runtime.hpp` (+46 -7) 📝 `csrc/jit/handle.hpp` (+72 -18) ➕ `csrc/jit/include_parser.hpp` (+80 -0) 📝 `csrc/jit/kernel_runtime.hpp` (+77 -30) 📝 `csrc/jit_kernels/heuristics/common.hpp` (+34 -319) ➕ `csrc/jit_kernels/heuristics/config.hpp` (+171 -0) ➕ `csrc/jit_kernels/heuristics/mega_moe.hpp` (+211 -0) ➕ `csrc/jit_kernels/heuristics/runtime.hpp` (+62 -0) _...and 80 more files_ </details> ### 📄 Description ## Summary Rebases `nv_dev` onto upstream `main` (brings in `#292`, `#304`), re-applies the two NV-side feature commits, and fixes a SM90 IMA introduced by the `#304` metadata API change. ## Commits | SHA | Kind | Title | |---|---|---| | `f6d98f2` | feat | bf16 bias output for FP8 GEMM on SM100 | | `12cb7c0` | feat | SM90 paged MQA logits: `block_kv=32` / `next_n=4` | | `a97b74d` | fix | SM90 paged MQA logits IMA on `next_n=4` post-#304 | ### `f6d98f2` — bf16 bias output for FP8 GEMM on SM100 `sm100_fp8_gemm_1d1d` / `sm100_fp8_fp4_gemm_1d1d` with accumulation can now produce FP32 **or** BF16 matching the bias/C dtype. Adds bf16 bias coverage to `enumerate_normal` in `tests/generators.py`. ### `12cb7c0` — SM90 paged MQA logits: `block_kv=32` / `next_n=4` - MMA tile stays at `kComputeBlockKV = 64`; physical `block_kv ∈ {32, 64}` → `kNumBlocksPerMMA = 64 / block_kv` TMA copies per WGMMA. - `kNumKVMulticast` template for `next_n = 4`: 2-CTA cluster splits Q across CTAs (`kNextNPerCTA = kNextN / kNumKVMulticast`) to fit SM90 smem/register budgets. WGMMA `N=256` otherwise exceeds `kNumMathRegisters=104`. - Extends the `test_attention.py` matrix. ### `a97b74d` — SM90 paged MQA logits IMA on `next_n=4` post-#304 **Symptom** — IMA in `sm90_fp8_paged_mqa_logits` when running `next_n=4` on the rebased tree. **Root cause** — `#304`'s metadata kernel hard-codes: ```cpp num_next_n_atoms = next_n / (next_n % 2 == 0 ? 2 : 1); ``` | arch | strategy | kernel template `kNumNextNAtoms` | hard-coded host value | |---|---|---|---| | SM90 `next_n=4` | cluster multicast (no time-atom) | **1** | 2 ❌ | | SM100 `next_n=4` | time-atomization | 2 | 2 ✓ | Metadata schedules 2× the tasks under the 2-atom scheme; SM90 scheduler decodes each entry as 1-atom → `q_idx = q_atom_idx / 1` overruns `batch_size` → reads past `context_lens`. **Fix** - Pipe `num_next_n_atoms` as a runtime parameter; host computes arch-aware (SM90 multicast → 1, SM100 → 2 for `next_n=4`). - Relax `schedule_meta` size to `num_sms / num_kv_multicast + 1` for the SM90 cluster grid. - Relax `get_paged_mqa_logits_metadata`'s `block_kv` assertion to `{32, 64}` (SM90 `block_kv=32` is supported after `12cb7c0`). SM100 behavior unchanged — its `num_next_n_atoms` value equals the previously hard-coded formula. ## Upstream (already reviewed upstream) - `d30fc36` Fix sync issue of TMEM alloc/dealloc (#292) - `7f2a703` [Public release 26/04] Mega MoE, FP4 Indexer and other features/fixes (#304) ## Testing | Test | H200 (SM90) | GB200 (SM100) | |---|---|---| | `test_attention.py` | ✅ | ✅ (80/80 paged MQA combinations) | | `test_bf16.py` | ✅ | ✅ | | `test_einsum.py` | ✅ | ✅ | | `test_fp8_fp4.py` | ✅ | ✅ | | `test_hyperconnection.py` | ✅ | ✅ | | `test_layout.py` | ✅ | ✅ | | `test_lazy_init.py` | ✅ | ✅ | | `test_legacy.py` | ✅ | ✅ | | `test_mega_moe.py` | ⚠️ env | ⚠️ env | | `compute-sanitizer` (SM90 paged MQA path) | 0 errors (24 sub-runs, 1h budget) | — | - `test_attention.py` paged MQA logits covers `block_kv ∈ {32,64} × next_n ∈ {1,2,4} × L ∈ {8k,32k} × dtype ∈ {float,bf16}` on H200, and `is_fp4 ∈ {T,F} × bf16 ∈ {T,F} × block_kv ∈ {32,64} × next_n ∈ {1..6}` on GB200. - `test_mega_moe.py` fails with `cuda:N out of range` on both arch allocations (pre-existing multi-GPU test, not touched by this PR). --- <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:46 +00:00