[weireweire]

Problem

deepseek-ai#441 fixed the task-info slot release ordering by adding fence.acq_rel.cta to MegaMoEScheduler::release_task_info(). The epilogue calls this once per block.

In SASS the fence is MEMBAR.ALL.CTA. It waits for all of the thread's outstanding memory operations, not only the slot reads. The epilogue thread has just issued the previous block's remote combine stores and counter updates, so every release waits for them.

At a 32-rank decode shape this costs ~6 us per fp8xfp4 call, about 7% of the kernel. The release only has to come after the slot reads.

Change

• get_next_task() copies the 32-byte slot with two ld.acquire.cta.shared::cta.v4.u32 loads, right after the full-barrier wait.
• release_task_info() drops the fence.
• The acquire orders every later operation of the thread after the slot reads, including the release arrive. The scope is CTA, as with Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441.
• ptxas emits no MEMBAR or CCTL for it. The SASS keeps the same MEMBAR/CCTL/SYNCS.ARRIVE counts as the unfenced kernel; the slot copy becomes ordered scalar LDS.
• The asm is volatile with a "memory" clobber, so the copy also stays below the full-barrier wait. This covers the wait side, which Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441 does not touch.

Benchmark

Setup:

• fp8_fp4_mega_moe across 32 ranks on 8 GB300 nodes in one NVLink domain.
• 416 experts with 13 per rank, top-6, 1 shared expert, hidden 5120, intermediate 2304.
• 108 tokens per rank, BLOCK_M 64, 10 calls per CUDA graph, 50 replays.
• Metric: slowest rank's median per-call time. "Cross" uses routing across ranks; "local" keeps it on each rank's own experts.
• Variants run in rotating order within one allocation. Three separate allocations; one value per allocation, in us/call.

variant	cross	local	vs deepseek-ai#441
deepseek-ai#441 (fence.acq_rel.cta, current main)	89.62 / 89.63 / 89.76	82.73 / 82.08 / 82.69	-
no fence (pre-deepseek-ai#441)	83.52 / - / 83.74	76.72 / - / 76.90	-6.0
this PR (ld.acquire.cta slot loads)	83.81 / 83.46 / 83.78	77.18 / 76.53 / 77.06	-6.0
compiler barriers only (asm volatile("" ::: "memory"))	- / - / 83.74	- / - / 76.82	-6.0
ld.acquire.cluster slot loads	- / - / 84.56	- / - / 76.96	-5.2
mbarrier.arrive.release.cluster on the release, no fence	85.79 / - / -	80.29 / - / -	-3.8

Correctness:

• The output is bitwise identical across all variants.
• 200 graph replays x 10 calls on every rank compared bitwise against the first replay: 0 mismatches for every variant.

Scope note

• For the leader CTA, the scheduler and the consumer share a CTA. The acquire loads plus the arrive's default .release.cta order the reads before the scheduler's next write.
• The peer CTA arrives remotely on the leader's barrier. Strictly, that needs a cluster-scope release on the arrive and a cluster-scope acquire on the scheduler's empty wait. Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441's CTA-scope fence does not provide that either.
• Measured cost of the cluster-scope options:
  • release.cluster on the arrive: +2.3 us per call over no fence.
  • ld.acquire.cluster: +0.8 us, and it adds CCTL.
• This PR therefore keeps Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441's scope at no cost. If you want the formal cluster-scope version, adding .release.cluster to the remote arrive is a one-line follow-up, at that price.

Alternative considered

Compiler barriers alone (asm volatile("" ::: "memory")) after the full wait and before the arrive:

• The SASS is identical to the unfenced kernel, and so is the speed (row above).
• It relies on the compiler-chosen order and on the slot loads having completed by the time the arrive executes.
• The acquire loads make the hardware order explicit at the same measured cost, so this PR uses them.

Serving

DeepSeek-V4.1-Flash, disaggregated serving on GB300:

• 32-way expert-parallel decode with data-parallel attention, MTP, concurrency ~3000, 1-hour run.
• Two same-day control/patched pairs, averaged:
  • Decode step: 1.7-1.9% shorter.
  • Mean interactivity (tokens/s per user): +2.4%. p90: +2.2%.
  • p90 TPOT: -2.2%. Throughput: +0.4%.
  • Both patched runs beat both controls on each of these.
• GSM8K accuracy is unchanged (96.8 / 97.1%).

Test plan

• 32-rank fp8xfp4 benchmark and bitwise replay check (above).
• SASS diff against Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441, unfenced, compiler barriers, and ld.acquire.cluster.
• tests/test_mega_moe.py at 4 ranks for fp8xfp4 and bf16xbf16, with and without the change:
  • Both compile and run; the per-call times match within 1-2 us on one node.
  • The legacy baseline (DeepEP) was not installed, so the test's own correctness comparison was skipped.
  • bf16xbf16 shares the scheduler but is not benchmarked across ranks.

[zcnrex]

Can we reuse the existing CUTLASS_DEVICE uint32_t ld_acquire_cta_shared(const uint32_t* ptr)

[weireweire]

ld_acquire_cta_shared only read 4 bytes which need 8 reads and be slower, but we can override it, changed