[PR #384] [MERGED] Sync nv_dev with upstream #377 #407

New issue

Closed

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

bsakaguchi commented

2026-10-26 21:17:21 +00:00

Owner

📋 Pull Request Information

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

Base: nv_dev ← Head: user/jinshik/nv_dev_377

📝 Commits (5)

54e2261 Multiple updates and refactorings (#364)
1f6f3f3 Fix a race condition where in-flight tensormaps are updated in-place (#343)
559d79f Public release 26/07 (#377)
17780e5 Merge upstream commit 559d79f (#377: Public release 26/07) into nv_dev
96fbced Fix SM100 FP16 MQA synchronization

📊 Changes

77 files changed (+7890 additions, -7136 deletions)

View changed files

📝 csrc/apis/attention.hpp (+176 -181)
📝 csrc/apis/einsum.hpp (+1 -1)
📝 csrc/apis/gemm.hpp (+113 -50)
📝 csrc/apis/layout.hpp (+41 -22)
📝 csrc/apis/mega.hpp (+263 -94)
📝 csrc/indexing/main.cu (+2 -4)
📝 csrc/jit_kernels/heuristics/config.hpp (+4 -0)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+60 -86)
📝 csrc/jit_kernels/heuristics/runtime.hpp (+2 -2)
📝 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 (+236 -0)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp (+30 -20)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+111 -18)
➖ csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp (+0 -416)
➕ csrc/jit_kernels/impls/sm100_mqa_logits.hpp (+490 -0)
➕ csrc/jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp (+179 -0)
➕ csrc/jit_kernels/impls/sm120_mqa_logits.hpp (+691 -0)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+5 -5)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp (+8 -8)

...and 57 more files

📄 Description

Summary

Merges upstream 559d79f (#377, Public release 26/07) into nv_dev at a6b593d (#369). This also brings in upstream #364 and the race fix from #343; the merge base is 88965b0.

The branch uses the same explicit non-fast-forward topology as the previous nv_dev sync PRs (#314, #328, and #349), keeping the upstream main commit as the merge commit's second parent.

B200 validation exposed two SM100 FP16-weight MQA synchronization gaps in the PTX-wrapper migration. Follow-up 96fbced elects one lane for raw-PTX 2x1SM UMMA issue, waits for each CTA's local KV-scale TMA before consumption, and fences the TMEM producer/consumer handoff.

Conflict resolution

Adopted #377's unified SM100 MQA and common MMA/PTX refactor while retaining nv_dev's SM90 next_n=4 multicast and block_kv=32, dedicated SM120 paths, FP16/BF16 weights, and 16-head indexer support.
Preserved metadata safety for empty/zero-context requests and made producer Q prefetch follow the scheduler's actual next task.
Ported the FP16-weights path to the new headers, with padded host offsets, correct per-warpgroup row indexing, and per-row compressed-window bounds.
Guarded the SM90 block_kv=32 odd physical-page tail from reading one page past the block-table row.
Combined upstream grouped scale-factor, PSUM, zero-padding, and Mega-MoE changes with nv_dev's SM120 dense/grouped GEMM dispatch.
Merged and extended attention/GEMM/layout coverage for SM90, SM100, and SM120.

Validation

Tests were run from isolated Docker worktrees on NVIDIA B200 GPUs (SM100):

The final FP16-weight MQA fix passed a clean CUDA extension build/JIT.
Core suites passed: test_bf16.py, test_fp8_fp4.py, test_layout.py, test_einsum.py, test_legacy.py, and test_hyperconnection.py.
SM100 attention prefill passed all 1,824 generated cases. Each completed case includes 20 repeated bitwise-consistency checks plus mask and numerical-reference validation.
Targeted SM100 FP16-weight MQA coverage passed head dimensions 32/64/128 with CP=0/1; the original H64/D32/CP0 failure passed 500 consecutive consistency runs.
SM100 paged MQA passed all 6,624 generated configurations. The exhaustive matrix was completed across isolated resume shards with two-case overlaps at every interrupted handoff; each completed case includes 20 repeated bitwise-consistency checks plus mask and numerical-reference validation.
Six-GPU Mega-MoE smoke runs passed for both fp8xfp4 and bf16xbf16. The fused kernels launched and benchmarked successfully; the deep_ep baseline comparison was skipped because that module was unavailable in the image.
All Git-tracked Python files pass py_compile; no unresolved conflicts, conflict markers, missing project-local includes, or references to removed headers were found.

GPU validation is SM100/B200 only. SM90 and SM120 GPU paths were not covered.

Three git diff --check warnings are inherited byte-for-byte from upstream #377 (Mega-MoE trailing whitespace); the conflict resolutions and follow-up fix introduce no new whitespace warnings.

_{🔄 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/384 **Author:** [@Barry-Delaney](https://github.com/Barry-Delaney) **Created:** 7/17/2026 **Status:** ✅ Merged **Merged:** 7/20/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `user/jinshik/nv_dev_377` --- ### 📝 Commits (5) - [`54e2261`](https://github.com/deepseek-ai/DeepGEMM/commit/54e22612409371d6364144b69086735beb54e98b) Multiple updates and refactorings (#364) - [`1f6f3f3`](https://github.com/deepseek-ai/DeepGEMM/commit/1f6f3f378920ccb5cc036ef43eb3f5972e921713) Fix a race condition where in-flight tensormaps are updated in-place (#343) - [`559d79f`](https://github.com/deepseek-ai/DeepGEMM/commit/559d79fb6994a58b8a15b4b93bf13ccc16edf247) Public release 26/07 (#377) - [`17780e5`](https://github.com/deepseek-ai/DeepGEMM/commit/17780e5d29ddf3178cd8faae1e6ff33cd57286fe) Merge upstream commit 559d79f (#377: Public release 26/07) into nv_dev - [`96fbced`](https://github.com/deepseek-ai/DeepGEMM/commit/96fbced2c06b01b353ba4f58e44c45b0e799b340) Fix SM100 FP16 MQA synchronization ### 📊 Changes **77 files changed** (+7890 additions, -7136 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+176 -181) 📝 `csrc/apis/einsum.hpp` (+1 -1) 📝 `csrc/apis/gemm.hpp` (+113 -50) 📝 `csrc/apis/layout.hpp` (+41 -22) 📝 `csrc/apis/mega.hpp` (+263 -94) 📝 `csrc/indexing/main.cu` (+2 -4) 📝 `csrc/jit_kernels/heuristics/config.hpp` (+4 -0) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+60 -86) 📝 `csrc/jit_kernels/heuristics/runtime.hpp` (+2 -2) 📝 `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` (+236 -0) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp` (+30 -20) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+111 -18) ➖ `csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp` (+0 -416) ➕ `csrc/jit_kernels/impls/sm100_mqa_logits.hpp` (+490 -0) ➕ `csrc/jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp` (+179 -0) ➕ `csrc/jit_kernels/impls/sm120_mqa_logits.hpp` (+691 -0) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+5 -5) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp` (+8 -8) _...and 57 more files_ </details> ### 📄 Description ## Summary Merges upstream `559d79f` (`#377`, Public release 26/07) into `nv_dev` at `a6b593d` (`#369`). This also brings in upstream `#364` and the race fix from `#343`; the merge base is `88965b0`. The branch uses the same explicit non-fast-forward topology as the previous `nv_dev` sync PRs (`#314`, `#328`, and `#349`), keeping the upstream main commit as the merge commit's second parent. B200 validation exposed two SM100 FP16-weight MQA synchronization gaps in the PTX-wrapper migration. Follow-up `96fbced` elects one lane for raw-PTX 2x1SM UMMA issue, waits for each CTA's local KV-scale TMA before consumption, and fences the TMEM producer/consumer handoff. ## Conflict resolution - Adopted `#377`'s unified SM100 MQA and common MMA/PTX refactor while retaining nv_dev's SM90 `next_n=4` multicast and `block_kv=32`, dedicated SM120 paths, FP16/BF16 weights, and 16-head indexer support. - Preserved metadata safety for empty/zero-context requests and made producer Q prefetch follow the scheduler's actual next task. - Ported the FP16-weights path to the new headers, with padded host offsets, correct per-warpgroup row indexing, and per-row compressed-window bounds. - Guarded the SM90 `block_kv=32` odd physical-page tail from reading one page past the block-table row. - Combined upstream grouped scale-factor, PSUM, zero-padding, and Mega-MoE changes with nv_dev's SM120 dense/grouped GEMM dispatch. - Merged and extended attention/GEMM/layout coverage for SM90, SM100, and SM120. ## Validation Tests were run from isolated Docker worktrees on NVIDIA B200 GPUs (SM100): - The final FP16-weight MQA fix passed a clean CUDA extension build/JIT. - Core suites passed: `test_bf16.py`, `test_fp8_fp4.py`, `test_layout.py`, `test_einsum.py`, `test_legacy.py`, and `test_hyperconnection.py`. - SM100 attention prefill passed all 1,824 generated cases. Each completed case includes 20 repeated bitwise-consistency checks plus mask and numerical-reference validation. - Targeted SM100 FP16-weight MQA coverage passed head dimensions 32/64/128 with `CP=0/1`; the original H64/D32/CP0 failure passed 500 consecutive consistency runs. - SM100 paged MQA passed all 6,624 generated configurations. The exhaustive matrix was completed across isolated resume shards with two-case overlaps at every interrupted handoff; each completed case includes 20 repeated bitwise-consistency checks plus mask and numerical-reference validation. - Six-GPU Mega-MoE smoke runs passed for both `fp8xfp4` and `bf16xbf16`. The fused kernels launched and benchmarked successfully; the `deep_ep` baseline comparison was skipped because that module was unavailable in the image. - All Git-tracked Python files pass `py_compile`; no unresolved conflicts, conflict markers, missing project-local includes, or references to removed headers were found. GPU validation is SM100/B200 only. SM90 and SM120 GPU paths were not covered. Three `git diff --check` warnings are inherited byte-for-byte from upstream `#377` (Mega-MoE trailing whitespace); the conflict resolutions and follow-up fix introduce no new whitespace warnings. --- <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:21 +00:00