From dc8729444730314bb344ba8cdedf18df1119f8ec Mon Sep 17 00:00:00 2001 From: Yinghan Ma Date: Wed, 9 Sep 2026 06:47:35 +0000 Subject: [PATCH] mega_moe: configurable SM headroom on cooperative mega grids (DG_MEGA_MOE_SM_HEADROOM) The SM100 FP8/FP4 and BF16 mega MoE kernels (forward + backward) launch cooperative grids whose grid_sync / nvlink_barrier (comm/barrier.cuh) require ALL grid CTAs simultaneously resident, with a hard 60s timeout that traps and takes down the process (SIGABRT) and the GPU (Xid 43 fallout). An ordinary launch gives no gang-scheduling guarantee: CUDA places CTAs greedily, so a concurrent kernel holding some (but not all) SMs leaves part of the grid queued while resident CTAs spin against the deadline and never yield their slots -- mutual deadlock. Reproduced locally (train-firetitan:0.638.337, 2x B300 -- same SKU as the production node): a single concurrent smem-heavy kernel on one rank's GPU delays that rank's grid; its own CTAs time out at barrier.cuh:39 ('Grid sync timeout') and peer ranks die at barrier.cuh:80 ('NVLink barrier timeout') -- the exact INC-1291 signature (figma RLOR xy98ftpyg7kovb4m, AP_MALAYSIA_2, 2026-09-08, forward dispatch barrier tag=1 per production Chronosphere logs). With the production BF16 kernel and real FSDP-shaped collectives, all three collective types (reduce-scatter / all-gather / HSDP all-reduce) abort the stock build; none abort with headroom + pinned NCCL channels. Reserve headroom via a shared get_mega_moe_num_sms() helper used by the fp8_fp4 forward, all four fp8_fp4 backward launches, and the bf16 default (the absolute DG_BF16_MEGA_MOE_NUM_SMS override is retained). The amount is controlled exclusively by DG_MEGA_MOE_SM_HEADROOM; it has NO built-in default -- unset means 0, the historical full-device grid -- because the right value is deployment-specific (serving SendRecv: 2; training FSDP2 collectives with 32 NCCL channels: 8 alongside NCCL_MAX_NCHANNELS=8). The value is rounded up to even, since these are 2-CTA cluster launches and several sites assert num_sms % 2 == 0. Measured sizing for reference (8x B300, production BF16 kernel): DG_MEGA_MOE_SM_HEADROOM=8 + NCCL_MAX_NCHANNELS=8 + NCCL_MIN_NCHANNELS=8 + NCCL_NVLS_ENABLE=0: zero aborts across 90 mixed + 200 e2e iterations, no measurable slowdown (collectives ~9% faster at 64 MiB payloads). Forward-path credit: Ying Zhang's yingz/mega-sm-headroom (95046f0, opened as #14) -- same mechanism, fixed value 2; this PR generalizes it to all cooperative sites and makes it a deployment knob. Known gap: reservation narrows the window but does not close it -- a preemptive long-lived kernel occupying SMs before the mega launch still starves the grid. A structural fix needs residency-agnostic synchronization or driver-validated cooperative launch. Co-Authored-By: Claude Code --- csrc/jit_kernels/impls/runtime_utils.hpp | 50 +++++++++++++++++++ .../jit_kernels/impls/sm100_bf16_mega_moe.hpp | 4 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe_backward.hpp | 8 +-- 4 files changed, 58 insertions(+), 6 deletions(-) diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 2e617d11a5..e8e81025d4 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -4,6 +4,7 @@ #include #include "../heuristics/sm90.hpp" +#include "../../jit/device_runtime.hpp" #include "../../jit/handle.hpp" #include "../../utils/math.hpp" #include "../../utils/system.hpp" @@ -11,6 +12,55 @@ namespace deep_gemm { +// Grid size for the cooperative mega-MoE kernels, with SM headroom reserved for +// co-resident kernels. +// +// These kernels launch a persistent grid of one CTA per SM joined by a software +// whole-grid barrier (`comm/barrier.cuh`) with a hard 60s timeout. The barrier is +// only satisfiable if every CTA is *simultaneously resident*, but an ordinary +// launch gives no gang-scheduling guarantee: CUDA places CTAs greedily, so a +// concurrent kernel holding some (but not all) SMs leaves part of the grid queued +// while the resident CTAs spin against the deadline and never yield their slots. +// The result is a mutual deadlock that traps at 60s -> 'DeepGEMM grid sync +// timeout' (barrier.cuh:39) / 'NVLink barrier timeout' (barrier.cuh:80), which in +// turn produces Xid 43 and kills the process. +// +// Reserving SMs gives a later-arriving co-resident kernel somewhere to land other +// than an SM the grid needs. How much headroom is enough depends entirely on the +// competing kernel's footprint, which is deployment-specific: +// - disaggregated serving: NCCL `ncclDevKernel_SendRecv` (KV transfer) — 2 was +// measured sufficient (yingz/mega-sm-headroom) +// - training: FSDP2 all-gather / reduce-scatter / HSDP all-reduce on dedicated +// comm streams; a 32-channel collective touches ~8 SMs (one channel = one +// 512-thread block, 4 blocks/SM), so 8 (with NCCL_MAX_NCHANNELS=8 pinned) +// was measured for the INC-1291 shape +// +// There is deliberately NO built-in default: the headroom is set exclusively via +// `DG_MEGA_MOE_SM_HEADROOM`, and an unset variable means 0 (no reservation), i.e. +// the historical full-device grid. Deployments that run comm kernels concurrently +// with the mega grid MUST set it; see the INC-1291 RCA (fw-ai/fireworks#47752) +// for how to size it. +// +// The value is clamped even, because these kernels launch 2-CTA clusters (2-SM +// MMA with a paired TMEM accumulator layout) and several call sites assert +// `num_sms % 2 == 0`. +static int get_mega_moe_num_sms() { + const int num_device_sms = device_runtime->get_num_sms(); + // No default: unset env == 0 == the historical full-device grid. The + // deployment owns the value (see comment above). + int headroom = get_env("DG_MEGA_MOE_SM_HEADROOM", 0); + + DG_HOST_ASSERT(headroom >= 0 and "DG_MEGA_MOE_SM_HEADROOM must be non-negative"); + // Round up to keep the grid even for the 2-CTA cluster launch. + headroom = align(headroom, 2); + DG_HOST_ASSERT(headroom < num_device_sms and + "DG_MEGA_MOE_SM_HEADROOM leaves no SMs for the mega-MoE grid"); + + const int num_sms = num_device_sms - headroom; + DG_HOST_ASSERT(num_sms % 2 == 0); + return num_sms; +} + static std::pair get_inner_outer_dims(const cute::UMMA::Major& major, const int& k, const int& mn) { return major == cute::UMMA::Major::K ? std::make_pair(k, mn) : std::make_pair(mn, k); } diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index a3b23f6c0e..705e71e94d 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -296,9 +296,11 @@ static void sm100_bf16_mega_moe( // Launch const auto physical_num_sms = device_runtime->get_num_sms(); + // Default to the shared headroom-reserving grid size (see + // `get_mega_moe_num_sms`); the explicit absolute override is retained. const auto num_sms = get_env( "DG_BF16_MEGA_MOE_NUM_SMS", - physical_num_sms); + get_mega_moe_num_sms()); DG_HOST_ASSERT(num_sms > 0 && num_sms <= physical_num_sms); const SM100BF16MegaMoERuntime::Args args = { .num_max_tokens_per_rank = num_max_tokens_per_rank, diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 5a7106da50..255285891d 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -326,7 +326,7 @@ static void sm100_fp8_fp4_mega_moe( cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); // Launch - const auto num_sms = device_runtime->get_num_sms(); + const int num_sms = get_mega_moe_num_sms(); const SM100FP8FP4MegaMoERuntime::Args args = { .num_max_tokens_per_rank = num_max_tokens_per_rank, .hidden = hidden, .intermediate_hidden = intermediate_hidden, diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_backward.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_backward.hpp index baa82a94d4..1989266da0 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_backward.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_backward.hpp @@ -743,7 +743,7 @@ static void sm100_fp8_fp4_mega_moe_backward_dgrad_swiglu( dgrad_block_k, load_block_m, static_cast(grad_gate_up_output.stride(-2)), 128); // Each launch gets a unique readiness epoch; no host memset is required. - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = get_mega_moe_num_sms(); DG_HOST_ASSERT(num_sms % 2 == 0); static std::atomic next_launch_epoch{1}; uint32_t launch_epoch = @@ -953,7 +953,7 @@ static void sm100_mega_moe_backward_combine_grad_x( DG_HOST_ASSERT(topk_ids->size(1) == num_topk); } - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = get_mega_moe_num_sms(); const SM100MegaMoEBackwardCombineRuntime::Args args = { .num_ranks = num_ranks, .num_local_experts = num_local_experts, @@ -1145,7 +1145,7 @@ static void sm100_bf16_mega_moe_backward_post_down_prelude( } } - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = get_mega_moe_num_sms(); const auto backward_sym_buffer = layout::SymBuffer<>( backward_sym_buffer_ptrs, backward_rank); const auto backward_workspace = layout::Workspace( @@ -1605,7 +1605,7 @@ static void sm100_bf16_mega_moe_backward_dgrad( static_cast( grad_gate_up_output.stride(-2)), 128); - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = get_mega_moe_num_sms(); DG_HOST_ASSERT(num_sms % 2 == 0); constexpr int num_trace_sites = 22; constexpr int num_trace_values = 5;