[zyongye]

Cherry-pick of deepseek-ai/DeepGEMM@5081f85 ("Fix mega_mhc split barriers under CUDA graph replay", deepseek-ai#478) onto dev. It applies cleanly, and the diff is identical to upstream.

Problem

mega_mhc tags each execution's split barriers with %gridid. A CUDA graph kernel node keeps the same %gridid on every replay, so after one replay a barrier line already holds the "all splits arrived" state that the next replay waits for. A Mix or Norm waiter whose m-block init CTA has not started yet then proceeds at once and reads the previous replay's split partials.

Fix (upstream)

Executions are tagged with an epoch instead. The epoch and an exit count live in the unused tail of the first barrier line, so the host-side barrier buffer is unchanged. Every CTA reads the epoch after cudaGridDependencySynchronize, and the last CTA to retire resets the exit count and advances the epoch. Only three JIT headers change: impls/sm100_mega_mhc.cuh, layout/mega_mhc.cuh and scheduler/mega_mhc.cuh.

Test

Run on GB200 (SM100) on exclusive Slurm nodes, comparing dev with and without this commit. The commit changes only JIT headers, so both trees share one extension build.

CUDA graph replay stress (not added to the repo). One mega_mhc call is captured in a graph and replayed 500 times per shape, with different inputs copied in before each replay. Each replay's outputs are compared bit for bit with an eager run on the same inputs. Shapes: T in {1, 64, 200, 1025, 4096}, H in {4096, 7168}, normal and shifted, BF16 output.

	dev	This PR
Plain replay	36 / 10000 mismatched	0 / 10000
A large GEMM running on a side stream during each replay	3 / 10000 mismatched	0 / 10000

On dev, the mismatches hit T=64/200/1025/4096 at H=4096 and H=7168, in both normal and shifted modes (up to 27/500 at normal T=64, H=7168).

tests/test_mega_mhc.py fails identically on dev and on this PR, at the capture-before-warm-up check. Because of that, the graph-replay, multi-stream and correctness sections never run. The stable-ABI migration (#21) changed the capture guard in csrc/apis/mega_mhc.hpp to not torch_compat::is_capturing(...). The test still expects CaptureStatus::None in the error message. The guard itself still rejects capture before warm-up. With that one string check changed to is_capturing in a scratch copy, the full test passes on both dev and this PR (API contract, single-replay graph check, multi-stream check, and all 56 correctness cases with 30 bitwise-repeat runs each). The stale check is not fixed here.

Performance (bench_kineto, eager, same node and job). Barrier calls now load the epoch from global memory instead of reading %gridid, and every CTA does a release atomic on exit. Together these add about 1–2 us per call:

Tokens	Change vs dev
T <= 200	+6% to +20% (about +0.8 to +2.3 us)
T = 1025	+3% to +8%
T = 4096	+1% to +3%
T = 32768	-0.5% to +1.7%

The mean over all 56 cases is +7%. Decode-sized calls pay the most.

AI assistance: Claude Code did the cherry-pick and the validation.

🤖 Generated with Claude Code