Add W8A8 (FP8xFP8) Mega MoE support

[xuzhiyuan1]

Motivation

Add W8A8 (FP8 activations x FP8 weights) support to Mega MoE, and relax the requirement that hidden / intermediate hidden sizes must be multiples of 512.

Changes

1. Add the SM100 W8A8 kernel (sm100_fp8_fp8_mega_moe.cuh) along with its JIT wrapper and host/Python API (fp8_fp8_mega_moe); extend tests/test_mega_moe.py with --mma-type fp8xfp8;
2. All supporting code is derived from the W4A8 (FP8xFP4) path: weights switch from e2m1 to e4m3, while the scale-factor handling (4-packed UE8M0) and the symmetric buffer layout stay identical, so buffer-size accounting reuses "fp8xfp4" (happy to introduce a dedicated MmaKind::MXFP8FP8 if preferred);
3. The second commit relaxes the 512-alignment assertion on the SF buffers.

Why relaxing the 512 alignment is safe

1. Where the requirement comes from: layout::Data asserts its byte count is 16B-aligned (the minimum TMA transfer alignment). The activation SF takes hidden / 32 (L1 input) or intermediate_hidden / 32 (L2 input) bytes per token, so both sizes end up being required to be multiples of 32 x 16 = 512;
2. What actually happens: the SF tensors are MN-major and TMA moves them in blocks along the M (token) dimension, so the alignment constraint applies to the M dimension — not to the per-token K-dimension byte count that this assertion checks;
3. Conclusion: the assertion is over-protective for SF buffer-size accounting (for both hidden and intermediate hidden). The change simply passes require_tma_alignment=false at the 3 accounting sites, with no behavioral change for shapes that are already 512-aligned.

Testing

# W8A8, default shape
python tests/test_mega_moe.py --mma-type fp8xfp8 --num-processes 2

# W8A8, hidden and intermediate hidden both not 512-aligned
python tests/test_mega_moe.py --mma-type fp8xfp8 --num-processes 2 \
    --num-experts 256 --num-topk 8 --num-tokens 2048 --num-max-tokens-per-rank 2048 \
    --hidden 768 --intermediate-hidden 1152

Both configurations produce bitwise-identical outputs against the non-overlapped baseline.

[ds-review-bot]

🟡 warning: The weight-tile barrier expectation is doubled relative to the FP4 path: arrive_and_expect_tx(sizeof(smem_b[0]) * 2 + sizeof(smem_sfb[0]) * 2). The comment "FP8 B tiles transfer twice the bytes of the FP4 path" is logically correct (FP8 = 1 B / logical element; FP4 is packed as 16U4 → 0.5 B / logical element, which the FP4 descriptor encodes via CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B). However, smem_b[0] is declared with the sub-byte type, and C++ sizeof(float_e2m1_unpacksmem_t) == 1 (its storage is a uint8_t in cutlass::float_exmy_base), so sizeof(smem_b[0]) for the FP4 path already represents the full (unpacked) smem footprint. A literal sizeof-based reading would therefore expect the FP8 path to use 1 * sizeof(smem_b[0]), and doubling it could over-count the barrier and risk a hang. The authors' empirical claim (bitwise-identical, no hang) strongly favors correctness, so I'm not asserting a bug — but SM100 sub-byte barrier-TX accounting is subtle enough that a focused runtime confirmation is the cleanest way to settle whether * 2 is exactly right vs. over/under by 2x.

_{🤖 v4}

[xuzhiyuan1]

Thanks, the factor is (bytes TMA delivers per CTA) x (2 CTAs): FP4 delivers packed data, sizeof(smem_b[0]) / 2 per CTA, so 0.5 x 2 = 1; FP8 is unpacked, so 1 x 2 = 2. All tested configs run to completion and are bitwise-identical to the baseline.

Also note your own v3 summary already confirmed this via the BF16 kernel cross-check ("non-packed B dtypes use * 2 for the 2-CTA broadcast... e4m3 is unpacked like BF16, so * 2 is right").

[ds-review-bot]

🔵 suggestion: The 512-alignment relaxation is sound: switching input_sf_layout/intermediate_sf_layout to layout::Data(<bytes>, false) bypasses the num_bytes % 16 == 0 assertion in Data::Data. This is well-justified — these are MN-major SF buffers where TMA strides along the M (token) dimension, so the alignment constraint applies to the block/M-dimension stride, not the per-token K-dim byte count. I verified the assertion logic at deep_gemm/include/deep_gemm/layout/mega_moe.cuh:208 and that this is the only source of the 512 constraint for the SF layouts (no remaining % 512 checks). No behavioral change for 512-aligned shapes. The test's unaligned config (768/1152) still satisfies the hidden % 128 == 0 / intermediate % 128 == 0 API contract.

_{🤖 v4}

[xuzhiyuan1]

Thanks! One addition: in practice hidden is constrained to multiples of 256 — the L2 GEMM has N = hidden and kNumL2BlockNs % 2 == 0 (scheduler/mega_moe.cuh:41, 2-CTA cluster) requires hidden % (2 * BLOCK_N) == 0.

[ds-review-bot]

🔵 suggestion: The alignment relaxation (require_tma_alignment=false) is safe. These Data objects feed only into buffer-size accounting (layout::Buffer), not TMA box dims. The one place SF is accessed with pointer arithmetic derived from num_bytes is the dispatch uint32_t read (sm100_fp8_fp8_mega_moe.cuh:569 remote_sf_ptr[j], offset src_token_idx * num_bytes), which requires num_bytes % 4 == 0. That is guaranteed by the pre-existing assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 (line 116), since num_bytes = hidden/32 = (hidden/128)*4. Good.

_{🤖 v3}

[ds-review-bot]

🤖 ds-review-bot Code Review

v4

This change adds W8A8 (FP8×FP8) Mega MoE support to DeepGEMM, derived from the existing W4A8 (FP8×FP4) path, and relaxes an over-protective 512-alignment assertion on the scale-factor buffer-size accounting. The work is split across two commits that exactly match the issue's scope: (1) the new SM100 kernel sm100_fp8_fp8_mega_moe (CUDA + JIT wrapper) with host/Python wiring and extended tests, and (2) relaxing the 512-alignment by passing require_tma_alignment=false at the three SF accounting sites. I performed a thorough static review in a sandbox where bash is restricted to git, so the kernel itself could not be executed; the analysis below is based on reading the diffs and the surrounding cutlass/deep_gemm sources.

v3

This PR adds W8A8 (FP8xFP8) Mega MoE support by deriving the SM100 kernel, JIT wrapper, and host/Python API from the existing W4A8 (FP8xFP4) path, and relaxes the 512-byte alignment assertion on the SF buffer-size accounting. I reviewed all 7 changed files against the FP8xFP4 and BF16 reference implementations. The implementation is correct and safe. Key verifications: (1) b_dtype_t is correctly switched from float_e2m1_unpacksmem_t to float_e4m3_t, and since the MMA/UMMA descriptors are templated on b_dtype_t (make_instr_desc_block_scaled, advance_umma_desc_lo) the FP8 B encoding is handled automatically; the SM100_MMA_MXF8F6F4 family supports e4m3. (2) The arrive_and_expect_tx(sizeof(smem_b[0]) * 2 + ...) change is CORRECT — I initially suspected a bug (sizeof(smem_b[0]) is identical for FP4/FP8 subbyte types), but cross-checking sm100_bf16_mega_moe.cuh confirmed non-packed B dtypes use '* 2' for the 2-CTA broadcast, whereas FP4 uses '* 1' only because packed FP4 in global memory transfers half the bytes via 16U4_ALIGN16B TMA. e4m3 is unpacked like BF16, so '* 2' is right. (3) The 512-alignment relaxation is safe: the SF Data objects are used only for buffer-size accounting; SF is written via plain byte stores and read via separately-constructed TMA descriptors (make_tma_sf_desc with its own get_tma_aligned_size). The dispatch path reads SF as uint32_t (offset token_idx * num_bytes), requiring num_bytes % 4 == 0, which is guaranteed because hidden % 128 == 0 and intermediate_hidden % 128 == 0 are enforced at mega.hpp, making num_bytes = hidden/32 = (hidden/128)*4 always a multiple of 4 (test shapes 768->24, 1152->36 confirm). (4) The host path correctly adds explicit e4m3 weight dtype guards (needed since FP4 weights also pass check_grouped_ab_fp8_fp4), and the test's reference GEMM (m_grouped_fp8_gemm_nt_contiguous, an alias dispatching on weight dtype) validates against genuine e4m3 weights. Overall: approve. Non-blocking suggestion: consider a dedicated MmaKind::MXFP8FP8 for clarity instead of overloading MXFP8FP4.

Files reviewed: 7
Issues found: 🟡 1 warning | 🔵 10 suggestion
Inline comments posted: 9
General comments (无法定位到 diff): 2

---

📍 未定位到 diff 的评论

🔵 suggestion csrc/jit_kernels/heuristics/mega_moe.hpp:L261: The FP8×FP8 path reuses MmaKind::MXFP8FP4 for the heuristics (get_mega_moe_config calls get_sf_uttcp_aligned_block_sizes(..., MmaKind::MXFP8FP4), and parse_mma_kind maps everything non-BF16 to MXFP8FP4). This is a correctness-safe choice — FP8 (e4m3×e4m3) and FP4 (e4m3×e2m1) have the same per-element byte count and smem footprint, so the per-stage smem accounting and num_stages count remain accurate (I verified: no smem-overflow risk, and the * 2 larger B tile is confined to the kernel). The trade-off is tunability: FP8 compute throughput on SM100 differs from FP4, so the block-M/N selection derived for FP4 is likely not FP8-optimal. This maps to the issue's explicit offer to introduce a dedicated MmaKind::MXFP8FP8. _{🤖 v4}

🔵 suggestion csrc/jit_kernels/heuristics/mega_moe.hpp:L68: Non-blocking design note: W8A8 reuses MmaKind::MXFP8FP4. This works because get_num_mma_elem_bytes (returns 1 for both FP8 activations and FP4-unpacked-to-8-bit weights) and is_mma_with_sf both yield correct values, and smem_b_size_per_stage = block_n * block_k * 1 matches the e4m3 smem_b size. However, parse_mma_kind still only accepts 'fp8xfp4'/'bf16xbf16', and the Python SymmBuffer is invoked with mma_type='fp8xfp4' for FP8xFP8, hiding the real type. A dedicated MmaKind::MXFP8FP8 (as the author offered) would be cleaner and safer against future FP4/FP8-specific divergence. _{🤖 v3}