mega_moe: reserve 2 SMs of headroom (grid = num_sms - 2) — original yingz/mega-sm-headroom branch

[mayinghan]

What this is

This opens Ying Zhang's existing yingz/mega-sm-headroom branch as a PR. The work and the code
are Ying's (commit 95046f0, 2026-06-27); I am only surfacing it for review because the same
defect just hit us again in training. @ipiszy — please take this over / correct anything below.

Net diff against main is one file, +11/-1, and merges cleanly (the intermediate
DEBUG reserve 16 SMs commit is reverted by the last commit, so the branch nets out to the 2-SM
reservation on the fp8_fp4 forward launch).

Why it matters now

The fix was written for disaggregated serving — a concurrent NCCL ncclDevKernel_SendRecv (KV
transfer) grabbing an SM the cooperative mega grid needs, leaving it 1-2 CTAs short so the
grid/NVLink barrier never completes (the "signal=4 target=8" crash on GB300 pf8).

It was never opened as a PR and never merged, so it is absent from affab4d — the DeepGEMM
commit shipped in the production trainer image. On 2026-09-08 the same defect class took down an
RLOR training job (INC-1291): DeepGEMM barrier.cuh:39 "Grid sync timeout" + barrier.cuh:80
"NVLink barrier timeout" mid-forward_backward, then SIGABRT and Xid 43 on all 8 GPUs.

I reproduced it end-to-end on healthy hardware (including with the production BF16 kernel and
real FSDP-shaped collectives — all three collective types are fatal; see the comment below for
the full analysis and the production-log decode). Full RCA in fw-ai/fireworks#47752.

Relationship to #13

#13 re-applies the same one-line change on the current base, makes the amount configurable
(DG_MEGA_MOE_SM_HEADROOM, default 2), and extends it to the launches this branch does not
cover (the 4 fp8_fp4 backward launches + the bf16 default).

Suggested order: merge this PR first so the original fix lands with correct provenance, then
rebase #13 down to only the extension. Happy to do that rebase, or to close #13 entirely if you'd
rather fold the extension in here.

[mayinghan]

INC-1291: the same defect, reproduced end-to-end in training

Posting the evidence that brought me to this branch, in case it's useful for deciding how far to
take the fix.

The production failure

Figma RLOR job xy98ftpyg7kovb4m (DSV4 Flash mxfp4, DP=2 CP=4, ~262k ctx, AP_MALAYSIA_2, node
e02-sg-e1n4r2c626n, image train-firetitan:0.638.337 → DeepGEMM affab4d), 2026-09-08T19:26:31Z,
mid-run forward_backward call_id=392 — after 391 successful calls:

DeepGEMM grid sync timeout: sm=..., thread=0, grid_sync_idx=0, old=..., current=..., expected_tag=...
Assertion failed: .../deep_gemm/comm/barrier.cuh:39, condition: false and "Grid sync timeout"
Assertion failed: .../deep_gemm/comm/barrier.cuh:80, condition: false and "NVLink barrier timeout"

then NCCL watchdogs failing with unspecified launch failure, torchrun SIGABRT, container exit 1,
and NVRM Xid 43 on all 8 GPUs. It was initially triaged as a possible bad node.

It is not hardware. DCGM showed Xid 0 on all 8 GPUs right up to the crash and flipped to 43 only
after the assert — Xid 43 is the fallout of the device-side trap, not a chip fault. No ECC, no
NVLink CRC, no Xid 48/74/79/94. The node was never tainted and a later trainer ran on it healthy
with 0 restarts.

Local reproduction

2× B300, the production image, a 2-rank EP group via torchrun, plus a synthetic SM-occupancy kernel
(dynamic smem ~227 KB/CTA, spinning ~200 s) on one rank's GPU before its MegaMoE call:

• baseline (no competing kernel): both ranks complete in ~3 s
• with the competing kernel: abort at ~60 s with the verbatim production signature — both
  timeout printfs, both asserts, Signal 6 (SIGABRT) from torchrun, and
  NVRM: Xid (PCI:...): 43, name=python in the host kernel log, on otherwise-idle healthy GPUs

One occupied SM on one rank is enough, and it kills both ranks. That is the cross-rank cascade:
rank 1 dies at barrier.cuh:39, rank 0 — whose own grid was fully resident and had already cleared
its grid sync — dies at barrier.cuh:80 waiting for rank 1's NVLink signals. Exactly the split seen
in production.

The part that surprised me: full occupancy is benign

Competing kernel occupies	Resident CTAs	Outcome
all SMs	0 of the grid	completes — 263 s / 336 s / 224 s across three hog variants
some SMs	partial	aborts at 60 s

When the whole grid is queued, nobody is spinning, so no clock64() deadline is running; the hog
drains, the grid enters together, and the barrier clears instantly. The kernel is only fatal in the
partial state — resident CTAs spin against a deadline while holding SM slots they never yield,
so the queued CTAs can never get in. Mutual deadlock.

Why I think reservation is mitigation rather than fix

The kernel semantically needs gang scheduling — all kNumSMs CTAs resident or none started —
because the whole-grid barrier is only satisfiable under full co-residency. But it uses an ordinary
launch, and CUDA schedules CTAs greedily and incrementally. The kernel depends on a guarantee it
never asked the driver for, and greedy scheduling is precisely what manufactures the fatal partial
state whenever a competitor holds some but not all SMs.

Reserving 2 SMs narrows that window — a later-arriving comm kernel lands in the headroom instead
of stealing an SM the grid needs — which is the interleaving that matters in practice and is why
this branch fixed the serving case. It does not close it: with the hog launched ahead of the mega
grid I still reproduce the abort on a patched build, with as little as one occupied SM, because
cluster_dim = 2 means an occupied SM also wastes its pairing partner (enough free SMs ≠ enough
pairable SM pairs).

And cross-rank, patience doesn't help at all: a rank whose grid is merely queued is locally benign,
but its peers are already resident and spinning in nvlink_barrier against their own independent
60 s deadlines.

Structural options, if you want to go further

1. cudaLaunchCooperativeKernel — the driver validates co-residency up front and returns an error
   instead of trapping 60 s later and killing a multi-hour job
2. residency-agnostic synchronization — a barrier that doesn't assume simultaneous residency
3. cheapest hardening regardless: make the timeout a recoverable CUDA error rather than a
   device-side trap, so the framework sees something retryable instead of Xid 43 + SIGABRT

Repro scripts and the full RCA are in fw-ai/fireworks#47752
(train-firetitan-py/tools/inc1291_repro/,
docs/incidents/2026-09-08-inc1291-deepgemm-grid-sync-timeout.md).

[cursor]

Risk: medium. Left a non-blocking comment; no reviewers were assigned. Human review is needed: this changes cooperative mega-kernel launch occupancy (get_num_sms() - 2) on a production-critical path, which is above the low-risk auto-approve threshold.

  

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

[mayinghan]

Follow-up evidence (final): the production Chronosphere logs for INC-1291 were retrieved
directly, and they pin the crash beyond the earlier analysis:

DeepGEMM NVLink barrier timeout: rank=0, counter=286245, signal=3, target=4, phase=0, sign=0, tag=1
DeepGEMM grid sync timeout: sm=96, thread=0, grid_sync_idx=0, old=2147483710, current=2147483795, expected_tag=62

• tag=1 = kBeforeDispatchPullBarrierTag (sm100_bf16_mega_moe.cuh:287) — the crash was in
  the forward kernel's dispatch barrier (the job ran megamoe_precision=bf16), not backward
• All 405 grid-sync prints at grid_sync_idx=0 — the first barrier, before any MMA work
• signal=3, target=4 — one EP rank's NVLink signal never arrived; that rank's grid could not
  become resident (its counters show 38 CTAs short)
• 907 crash lines across all 8 ranks, then Xid 43 + NCCL watchdog + SIGABRT

Also reproduced with the production BF16 kernel (not just fp8_fp4) and real funcol
collectives in the incident job shape (TP=1 PP=1 CP=4 EP=4, HSDP=2): all three FSDP-shaped
collective types (reduce-scatter, all-gather, HSDP all-reduce) abort the stock build with the
verbatim signature; zero aborts with headroom + pinned NCCL channels, at no measurable cost
(collectives ~9% faster at 64 MiB). The BF16 forward is more exposed than fp8_fp4 (6 cooperative
barriers vs 4).

One note for the training case: this branch's fixed - 2 is calibrated for serving (SendRecv,
few SMs). In the training shape a 32-channel collective touches ~8 SMs, so the measured setting
there is headroom 8 + NCCL_MAX_NCHANNELS=8 — which is why #13 makes the amount configurable
with this branch's 2 as the default. Full data in fw-ai/fireworks#47752.