[SM90] Add FP8 MegaMoE support - #383
Conversation
| // * Hidden dims must be multiples of 128 (per-128 SF + scheduler integer-tiling). | ||
| // * `l2_arrival_mask` is uint64, with one bit per L1-output N-block of size 64 in the | ||
| // intermediate dim, so `kNumL1BlockNs = intermediate_hidden / 64` must be ≤ 64. | ||
| DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); |
There was a problem hiding this comment.
🟡 warning: Validate the combine kernel's actual hidden alignment: For hidden values such as 128 or 384, this check passes, but the generated L2 kernel fails kCombineChunkElems % 256 == 0 during JIT compilation. These values satisfy the documented multiple-of-128 contract, so either handle partial combine vectors or reject them before invoking the compiler.
🤖 v6
| // * `l2_arrival_mask` is uint64, with one bit per L1-output N-block of size 64 in the | ||
| // intermediate dim, so `kNumL1BlockNs = intermediate_hidden / 64` must be ≤ 64. | ||
| DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); | ||
| DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); |
There was a problem hiding this comment.
🟡 warning: Remove the stale intermediate-size cap: Any intermediate_hidden > 4096 is rejected based on the stated l2_arrival_mask limit, but the new split SM90 kernel and scheduler never access that mask. Consequently otherwise valid multiples of 128, such as 4224, are rejected despite the README's contract.
🤖 v6
| check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, | ||
| num_experts_per_rank, false, true, torch::kFloat); | ||
| check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, | ||
| num_experts_per_rank, false, true, torch::kFloat); |
There was a problem hiding this comment.
🟡 warning: Reject transpose-contiguous weight scales: When a weight-SF tensor is transpose-contiguous—especially for square SF shapes such as L1 with hidden == 2 * intermediate_hidden—sm90_sfb_check=true accepts it, but this kernel indexes the raw pointer as (E, N, K) with K contiguous. The call then silently applies scales to transposed blocks; require stride(-1) == 1 here or account for the supplied layout.
🤖 v6
|
|
||
|
|
||
| def fp8_fp4_mega_moe(y: torch.Tensor, | ||
| l1_weights: Tuple[torch.Tensor, torch.Tensor], |
There was a problem hiding this comment.
🟡 warning: transform_weights_for_mega_moe_sm90 returns (_interleave_weights(l1_fp8), l1_sf), l2_weights. This yields the correct outer 2-tuple ONLY because l2_weights is already the (weight, sf) pair and l1 happens to need no SF/UTCCP transform. It does no dtype/shape/divisibility validation and is fragile/inconsistent with the sibling SM100 transform, which builds both pairs explicitly. Mirror that and validate each pair; otherwise a future caller passing a single tensor gets silently wrong shapes (a lambda/comment-only weight transform is the only thing standing between a working WGMMA run and a garbage/NaN result).
🤖 v4
| l2_sf_buffer.get_end_ptr()); | ||
|
|
||
| // Check SF buffer requirements | ||
| DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); |
There was a problem hiding this comment.
🟡 warning: Host entry aborts via DG_HOST_ASSERT on arch!=9, recipe!=(128,128,128), activation!='swiglu', use_fp8_dispatch==false, hidden/intermediate%128!=0 and intermediate_hidden/64>64, with no graceful degradation/fallback to the existing DeepEP/BF16 path. It is fully usable only on the high-SM H200/H20 tuning SKU; a smaller/consumer Hopper card or any unsupported-but-legal shape aborts rather than falling back. The benchmark was also run on the tuned SKU, so portability risk is invisible there.
🤖 v4
| @@ -0,0 +1,170 @@ | |||
| #pragma once | |||
There was a problem hiding this comment.
🟡 warning: The two-phase per-CTA traversal reconstructs (expert, m_block, n_block) from scratch per pass on the SAME grid/blockIdx.x with phase-specific kNumBlockNs/N-M-major and a RESET per-CTA state (the same CTA therefore does unrelated L1/L2 blocks). This is correct only if L1 and L2 are fully independent GLOBAL-memory producer/consumer passes over a byte-for-byte identical set of M-blocks. The per-expert block count/volume differs between phases (different N tiling), so the two passes must cover the same (expert, m_block) set exactly or a block is dropped/duplicated and the combine reader gets wrong/NaN data. Confirm this coverage equivalence before relying on it (the kernel itself could only be read in part here).
🤖 v4
| return requested_num_experts_per_wave; | ||
|
|
||
| // The tuned schedules use exact divisors. For a generic, previously unseen | ||
| // expert count, round upward to the first divisor so the fallback remains |
There was a problem hiding this comment.
🔵 suggestion: normalize_num_experts_per_wave_for_mega_moe_sm90 rounds a non-divisor requested experts-per-wave up to the next divisor, but the search loop runs candidate < num_experts_per_rank (strict), so num_experts_per_rank itself is never considered inside the loop; it is only reached via the final fallback return num_experts_per_rank. That happens to be correct here, but the loop bound looks like an off-by-one relative to intent. Consider candidate <= num_experts_per_rank (or documenting that the trailing return covers the top divisor) to make the invariant explicit and avoid a subtle regression if the fallback is ever changed.
🤖 v3
|
|
||
| static std::tuple<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(const torch::Tensor&)>> | ||
| get_symm_buffer_size_for_sm90_mega_moe( | ||
| const int& num_ranks, const int& num_experts, |
There was a problem hiding this comment.
🔵 suggestion: get_symm_buffer_size_for_sm90_mega_moe / fp8_mega_moe hard-assert use_fp8_dispatch and activation == "swiglu" via DG_HOST_ASSERT. Since these are the only supported modes for the SM90 path today this is acceptable, but the public Python signatures still accept use_fp8_dispatch and activation arguments, which invites callers to pass unsupported values and hit a raw host assert. A clearer error message (e.g. explicitly stating SM90 currently only supports fp8 dispatch + swiglu) would improve the failure mode.
🤖 v3
|
|
||
| # Skip on non-SM90 | ||
| if get_arch_major() != 9: | ||
| dist_print(f'[SKIP] test_mega_moe_sm90 requires SM90; got SM{get_arch_major()}0', |
There was a problem hiding this comment.
🔵 suggestion: The test suite correctly skips on non-SM90 devices and uses a PyTorch fp32/bf16 reference with calc_diff < 0.01. Because the reference is chunked (_CHUNK = 32) and not bitwise-identical to the kernel, correctness confidence depends on tolerance tuning. The layered coverage (smoke/heuristic/shape/edges/stress) is good; consider adding an assertion or note that Layer 3 production H200 shapes (Pro m128/m256) require a full node to exercise the BN512/BK256 paths, so a 2-rank default run may silently not cover those heuristic branches.
🤖 v3
|
|
||
| // Tensormap construction | ||
| // Acts/weights: standard 2D TMA descriptors (FP8 K-major). | ||
| // Activation SF: per-128 channel float for L1, per-64 for L2 (MN-major, no swizzle). |
There was a problem hiding this comment.
🔵 suggestion: The L1-output tensor map comment explains that a BK256 stage is two adjacent BK128 TMA tiles and the box N is capped at 256 (l1_tma_block_n / l2_tma_block_n via std::min(..., 256)). This is a subtle invariant tied to the kernel's stage advancement by 256 while issuing two 128 copies. It is documented here, but given how easy it is to break, a static/host assertion tying block_n<=512 and the 256 box cap together (rather than relying on the selector's legality check) would make the contract self-enforcing at the launch site.
🤖 v3
🤖 ds-review-bot Code Reviewv6The new SM90 path rejects documented-valid shapes, can fail only during JIT for accepted dimensions, and accepts a scale layout that produces silently incorrect results. v4The PR adds an SM90/Hopper FP8 MegaMoE kernel that runs the expert MLP as L1 (gate/up) + SwiGLU followed by a separate L2, connected by a global-memory FP8 activation staging buffer, using the same dispatch/combine contract as the existing SM100 path but with FP8 + block-(128,128) float weight SF + split SwiGLU/quant epilogue + two independent full-grid single-CTA phase passes. Most of the runtime tuning is in the host heuristics/launcher that the benchmark numbers were produced with. The core kernel (~2,400 lines) could only be partially read via git, so the kernel-internal/staging-equivalence risks below should be confirmed by a full read/run before merge. v3This PR adds an SM90 (Hopper) FP8xFP8 MegaMoE kernel to DeepGEMM, executing the expert MLP as L1 -> SwiGLU -> L2 while preserving the same overlapped EP dispatch/combine contract used by the existing SM100 FP8xFP4 path. The change is well-structured and cleanly layered end-to-end:
SM90-specific design is sound: FP8 e4m3 weights/activations with block-(128,128) float weight SF matching the DeepEP/DeepSeekV4FlashFp8 convention (so SF tensors can be physically shared), per-token/per-128-K activation SF for L1 and per-64-K SF for L2 activations, no TMEM/FP4/cluster-MMA (WGMMA accumulators register-resident, one CTA per work item). The kernel's SwiGLU clamp semantics (silu(min(gate,c)) * clamp(up,-c,c)), weight-application ordering, and per-64-column SF grouping match the test reference. The reported benchmarks show consistent +8% to +107% speedups over DeepEP HT/LL across H200 and H20 for Flash and Pro shapes. Overall the integration is coherent and correct in its wiring and public surface. Comments below are minor/robustness observations rather than blockers; note that this reviewer could not compile or execute the SM90 kernel (requires Hopper hardware), so correctness relies on the included test suite and the reported benchmarks. Files reviewed: 11 📍 未定位到 diff 的评论🔵 suggestion |
bbe9686 to
bc4f33a
Compare
| // Keep the SM90 MegaMoE barrier on the legacy trap-only timeout path. The | ||
| // shared SM100 barrier intentionally retains nv_dev's diagnostic printf, but | ||
| // compiling that printf into this register-heavy Hopper kernel creates a | ||
| // per-thread stack frame and slows both L1 and L2. This SM90-local adapter | ||
| // preserves the same synchronization and 300-second timeout semantics without | ||
| // changing the shared SM100 implementation. | ||
| template <uint32_t kNumRanks, uint32_t kNumSMs, uint32_t kNumThreads, | ||
| uint32_t kGridSyncIndex, uint32_t kTag, typename sync_scope_t> | ||
| CUTLASS_DEVICE void sm90_nvlink_barrier( | ||
| const layout::Workspace& workspace, | ||
| const layout::SymBuffer<kNumRanks>& sym_buffer, | ||
| const uint32_t& sm_idx, const uint32_t& thread_idx, | ||
| const sync_scope_t& sync_scope, | ||
| const bool& sync_prologue = true, | ||
| const bool& sync_epilogue = true) { | ||
| DG_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); | ||
|
|
||
| if (sync_prologue) | ||
| comm::grid_sync<kNumSMs, kGridSyncIndex>( | ||
| workspace, sm_idx, thread_idx, sync_scope); | ||
|
|
||
| if (sm_idx == 0) { | ||
| auto* counter_ptr = workspace.get_nvl_barrier_counter_ptr(); | ||
| const auto status = (*counter_ptr) & 3; | ||
| const auto signal_phase = status & 1, signal_sign = status >> 1; | ||
| auto* signal_ptr = workspace.get_nvl_barrier_signal_ptr(signal_phase); | ||
|
|
||
| if (thread_idx < kNumRanks) | ||
| ptx::red_add_rel_sys( | ||
| sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1); | ||
| sync_scope(); | ||
|
|
||
| constexpr int64_t kNumTimeoutCycles = 300ll * 2000000000ll; | ||
| if (thread_idx == 0) { | ||
| ptx::red_add(counter_ptr, 1); | ||
| const int target = signal_sign ? 0 : static_cast<int>(kNumRanks); | ||
| const auto start_clock = clock64(); | ||
| while (ptx::ld_acq_sys(signal_ptr) != target) { | ||
| if (clock64() - start_clock >= kNumTimeoutCycles) | ||
| DG_TRAP_ONLY_DEVICE_ASSERT(false and "NVLink barrier timeout"); | ||
| } |
There was a problem hiding this comment.
Could we reuse nvlink_barrier with a compile-time timeout policy? The synchronization logic here appears to be duplicated.
|
|
||
|
|
||
| float weight_r0 = 0.0f, weight_r1 = 0.0f; | ||
| if constexpr (kNumMaxTokensPerRank <= 1024) { |
There was a problem hiding this comment.
Why does kNumMaxTokensPerRank > 1024 switch from a single-lane load plus __shfl_sync to redundant per-lane loads? could we add a short comment？
There was a problem hiding this comment.
Added a short comment. For larger token capacities, each lane loads the same cached value directly, avoiding the wait for a single-lane load and shuffle broadcast.
|
Hi, @liz-badada DeepGEMM's sm100 MegaMoE path already has optimizations such as ring-buffered routed-expert buffers, L1/L2 wave interleaving, and fused shared-expert execution. Is there any plan to bring those optimizations into this PR's code path as well? |
Thanks for raising this. We considered the SM100 design, but it does not map directly to Hopper: without TMEM and the SM100 cluster/UMMA pipeline, combining everything puts too much pressure on registers and shared memory. That is why this PR uses two phase-specialized fused kernels. The listed optimizations would require a dedicated SM90 redesign and can be evaluated as follow-up work. |
|
@liz-badada @huangzhilin-hzl can you test against with the inital MegaMoE for SM 90 in sglang : sgl-project#36 |
Hi, We didn't compare that. But feel free to test these two. |
Hi, @yiakwy-xpu-ml-framework-team Flash
Pro
|
5cc19be to
4a825a2
Compare
Reuse shared barriers with compile-time timeout handling Harden buffer setup and add Hopper correctness/performance coverage
512ef39 to
788d5d9
Compare
6 participants
co-author： @jychen21
This PR adds an SM90 FP8 MegaMoE kernel for Hopper GPUs, executing the expert MLP as L1 + SwiGLU followed by L2.
Benchmark Results
Speedup is calculated as
DeepEP / MegaMoE - 1.H200 High Throughput
H200 Low Latency
H20 High Throughput
H20 Low Latency