Sync upstream through PR #441: fix Mega MoE scheduler ordering

[murphymatt]

The Mega MoE scheduler could release a task-info slot before its metadata reads finished, allowing the scheduler to overwrite metadata still in use. Sync with upstream main at 78b69000794d0937b47ae3387eff7663410264d1 (deepseek-ai/DeepGEMM#441), which adds a CTA acquire/release fence before signaling slot reuse. Both BF16 and FP8/FP4 Mega MoE use this scheduler.

The runtime diff is exactly upstream’s seven-line fix; all existing fork changes are preserved. Update the sync documentation to record PR deepseek-ai#441.

History

The previous sync (#15) was squash-merged. Its tree is identical to the original sync merge 38eb89f, which already contains upstream through deepseek-ai#432. Resolve the resulting prior-release conflicts by retaining that integrated fork tree and applying only the upstream delta since deepseek-ai#432. Both the current fork main (8d785f7) and upstream main are ancestors of this merge.

Please land using a merge commit to retain upstream ancestry for future syncs.

Validation

• All 26 offline SM100 CUDA variants compiled successfully with CUDA 13.1: BF16, FP8/FP4, FP8/FP8, shared experts, training saves, backward, weight gradients, and strided PSUM packing.
• Verified the runtime patch matches upstream Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441 exactly, and both changed headers match upstream byte for byte.
• Ancestry checks and git diff --check passed.
• GPU numerical, synchronization, and performance tests were not run: all eight local B200 GPUs were occupied by active training workloads. Offline compilation does not validate the runtime race fix.

---

Note

Medium Risk
Targets a real concurrency bug in Mega MoE scheduling; the change is small but affects synchronization on all BF16/FP8/FP4 Mega MoE kernels.

Overview
Syncs upstream through PR deepseek-ai#441 to fix a Mega MoE scheduler race: consumers could signal a task-info pipeline slot as free before finishing metadata reads, so the scheduler could overwrite task_infos still in use.

Adds ptx::fence_acq_rel_cta() (fence.acq_rel.cta) in ld_st.cuh and invokes it in MegaMoEScheduler::release_task_info() immediately before the empty-barrier arrive, so metadata reads complete before slot reuse. BF16 and FP8/FP4 paths both call release_task_info() and pick up the fix.

docs/fork-upstream-sync.md is updated to record upstream commit 78b6900 and describe the ordering change.

^{Reviewed by Cursor Bugbot for commit c2e49f5. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Risk: medium. Left a non-blocking comment; no approval and no reviewers assigned. This CTA acquire/release fence in Mega MoE release_task_info is a GPU scheduler ordering change that exceeds the low-risk auto-approve threshold, so human review is needed. Bugbot completed with no findings; Cursor Security Agent was not present.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[yanzunh]

Do we usually add unit tests for these

[murphymatt]

this is just a copy of the upstream. should be fine.