deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[PR #478] [MERGED] Fix mega_mhc split barriers under CUDA graph replay #478

New issue
Closed
opened 2026-10-26 21:18:21 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-26 21:18:21 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/478
Author: @lightseek-bot
Created: 10/19/2026
Status: ✅ Merged
Merged: 10/20/2026
Merged by: @zheanxu

Base: main ← Head: bot/mega-mhc-split-barrier-epoch

📝 Commits (1)

  • dd9cd5f Fix mega_mhc split barriers under CUDA graph replay

📊 Changes

3 files changed (+48 additions, -8 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh (+3 -0)
📝 deep_gemm/include/deep_gemm/layout/mega_mhc.cuh (+6 -0)
📝 deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh (+39 -8)

📄 Description

A CUDA graph kernel node keeps the same %gridid on every replay. The mega_mhc split barriers tag executions with %gridid, so after a replay a barrier line still holds encode_state(grid_idx, kNumSplits) — exactly what the next replay waits for. A Mix or Norm waiter whose m-block init CTA has not started yet proceeds at once and reads the previous replay's split partials (the graph-pool scratch is the same memory). wait_init only protects producers.

This tags each execution with an epoch instead of %gridid. The epoch and an exit count live in the unused tail of the first barrier line, so the host-side barrier buffer and the compiled extension are unchanged (JIT headers only). Every CTA reads the epoch after cudaGridDependencySynchronize(); the last CTA to retire resets the exit count and advances the epoch.

The MegaMoE combine readiness has the same stale-tag problem and is addressed separately in #451.

Validation on B200: one graph node is replayed while alternating two input sets, and every replay is compared bit for bit with eager results (fresh JIT cache per run). Measured on 66081d4, whose versions of these headers differ from main only by an added #include.

case before after
33 tokens, SMs contended by another stream 185/200 wrong 0/400
512 tokens, contended 27/200 0/400
8192 tokens, contended 3/200 0/400
512 tokens, idle GPU 1/200 0/400

Eager calls, and graph replays preceded by an eager launch that rewrites the tag, were 0/200 before the fix as well.

🔄 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/478 **Author:** [@lightseek-bot](https://github.com/lightseek-bot) **Created:** 10/19/2026 **Status:** ✅ Merged **Merged:** 10/20/2026 **Merged by:** [@zheanxu](https://github.com/zheanxu) **Base:** `main` ← **Head:** `bot/mega-mhc-split-barrier-epoch` --- ### 📝 Commits (1) - [`dd9cd5f`](https://github.com/deepseek-ai/DeepGEMM/commit/dd9cd5f7064e72c15ff043b4f7ed7b8cc97448cb) Fix mega_mhc split barriers under CUDA graph replay ### 📊 Changes **3 files changed** (+48 additions, -8 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh` (+3 -0) 📝 `deep_gemm/include/deep_gemm/layout/mega_mhc.cuh` (+6 -0) 📝 `deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh` (+39 -8) </details> ### 📄 Description A CUDA graph kernel node keeps the same `%gridid` on every replay. The `mega_mhc` split barriers tag executions with `%gridid`, so after a replay a barrier line still holds `encode_state(grid_idx, kNumSplits)` — exactly what the next replay waits for. A Mix or Norm waiter whose m-block init CTA has not started yet proceeds at once and reads the previous replay's split partials (the graph-pool scratch is the same memory). `wait_init` only protects producers. This tags each execution with an epoch instead of `%gridid`. The epoch and an exit count live in the unused tail of the first barrier line, so the host-side barrier buffer and the compiled extension are unchanged (JIT headers only). Every CTA reads the epoch after `cudaGridDependencySynchronize()`; the last CTA to retire resets the exit count and advances the epoch. The MegaMoE combine readiness has the same stale-tag problem and is addressed separately in #451. Validation on B200: one graph node is replayed while alternating two input sets, and every replay is compared bit for bit with eager results (fresh JIT cache per run). Measured on 66081d4, whose versions of these headers differ from `main` only by an added `#include`. | case | before | after | |---|---|---| | 33 tokens, SMs contended by another stream | 185/200 wrong | 0/400 | | 512 tokens, contended | 27/200 | 0/400 | | 8192 tokens, contended | 3/200 | 0/400 | | 512 tokens, idle GPU | 1/200 | 0/400 | Eager calls, and graph replays preceded by an eager launch that rewrites the tag, were 0/200 before the fix as well. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:18:21 +00:00
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
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#478
No description provided.