[SM90] Add FP8 MegaMoE support - #383
Conversation
Optimize SM90 MegaMoE split path
Optimize MegaMoE split path
Keep the verified swapAB and profile-aware policies, add stable Flash/Pro benchmarks and SM90 coverage, and remove experimental tuning paths and process artifacts.
- Add parameterized H200 schedules and numerical guards. - Add DeepEP low-latency and high-throughput baselines. - Record correctness, tuning, and benchmark validation.
| // * 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
3 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