Skip to content

Fix mega_mhc split barriers under CUDA graph replay - #478

Merged
zheanxu merged 1 commit into
deepseek-ai:mainfrom
lightseek-bot:bot/mega-mhc-split-barrier-epoch
Oct 20, 2026
Merged

zheanxu merged 1 commit into
deepseek-ai:mainfrom
lightseek-bot:bot/mega-mhc-split-barrier-epoch

Conversation

@lightseek-bot

@lightseek-bot lightseek-bot commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

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.

A CUDA graph kernel node keeps the same %gridid on every replay. The
mega_mhc split barriers tag each execution with %gridid, so after a replay
a barrier line still holds encode_state(grid_idx, kNumSplits), which is
exactly the state the next replay waits for. A Mix or Norm waiter whose
m-block init CTA has not started yet proceeds immediately and reads the
previous replay's split partials.

Tag executions with an epoch instead. The epoch and an exit count live in
the unused tail of the first barrier line, so the host-side buffer is
unchanged. Every CTA reads the epoch after cudaGridDependencySynchronize,
and the last CTA to retire resets the exit count and advances the epoch.

Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
@lightseek-bot

Copy link
Copy Markdown
Contributor Author

@LyricZhao PTAL Thanks!

@zheanxu
zheanxu merged commit 5081f85 into deepseek-ai:main Oct 20, 2026
zyongye added a commit to vllm-project/DeepGEMM that referenced this pull request Oct 29, 2026
…#48)

A CUDA graph kernel node keeps the same %gridid on every replay. The
mega_mhc split barriers tag each execution with %gridid, so after a replay
a barrier line still holds encode_state(grid_idx, kNumSplits), which is
exactly the state the next replay waits for. A Mix or Norm waiter whose
m-block init CTA has not started yet proceeds immediately and reads the
previous replay's split partials.

Tag executions with an epoch instead. The epoch and an exit count live in
the unused tail of the first barrier line, so the host-side buffer is
unchanged. Every CTA reads the epoch after cudaGridDependencySynchronize,
and the last CTA to retire resets the exit count and advances the epoch.


(cherry picked from commit 5081f85)

Co-authored-by: LightSeek Foundation <243258330+lightseek-bot@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
zcnrex added a commit to sgl-project/DeepGEMM that referenced this pull request Nov 2, 2026
…#131)

A CUDA graph kernel node keeps the same %gridid on every replay. The
mega_mhc split barriers tag each execution with %gridid, so after a replay
a barrier line still holds encode_state(grid_idx, kNumSplits), which is
exactly the state the next replay waits for. A Mix or Norm waiter whose
m-block init CTA has not started yet proceeds immediately and reads the
previous replay's split partials.

Tag executions with an epoch instead. The epoch and an exit count live in
the unused tail of the first barrier line, so the host-side buffer is
unchanged. Every CTA reads the epoch after cudaGridDependencySynchronize,
and the last CTA to retire resets the exit count and advances the epoch.

Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
Co-authored-by: LightSeek Foundation <243258330+lightseek-bot@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
zcnrex added a commit to sgl-project/DeepGEMM that referenced this pull request Nov 4, 2026
* Fix mega_mhc split barriers under CUDA graph replay (deepseek-ai#478)

A CUDA graph kernel node keeps the same %gridid on every replay. The
mega_mhc split barriers tag each execution with %gridid, so after a replay
a barrier line still holds encode_state(grid_idx, kNumSplits), which is
exactly the state the next replay waits for. A Mix or Norm waiter whose
m-block init CTA has not started yet proceeds immediately and reads the
previous replay's split partials.

Tag executions with an epoch instead. The epoch and an exit count live in
the unused tail of the first barrier line, so the host-side buffer is
unchanged. Every CTA reads the epoch after cudaGridDependencySynchronize,
and the last CTA to retire resets the exit count and advances the epoch.

Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>

* Update instructions for deepgemm fork development

* Fix SM100 kernels to use the allocated tensor memory base

* Keep the pull request limited to the tensor memory base fix

* Make the HC-prenorm regression test require a nonzero TMEM base

* Skip the nonzero TMEM base HC-prenorm test on CC 10.7

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
Co-authored-by: LightSeek Foundation <243258330+lightseek-bot@users.noreply.github.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
kevmo314 added a commit to kevmo314/DeepGEMM that referenced this pull request Nov 11, 2026
SM90 grouped 1D2D heuristics, thread-local alignment, MoE precompile, and deepseek-ai#455/deepseek-ai#478
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants