[PR #371] Add W8A8 (FP8xFP8) Mega MoE support #395

New issue

Open

opened 2026-10-26 21:17:13 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:13 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/371
Author: @xuzhiyuan1
Created: 7/3/2026
Status: 🔄 Open

Base: main ← Head: w8a8-megamoe

📝 Commits (2)

f53e0b2 Add W8A8 (FP8xFP8) Mega MoE support
7bc58b9 Support unaligned intermediate hidden sizes

📊 Changes

7 files changed (+1785 additions, -11 deletions)

View changed files

📝 csrc/apis/mega.hpp (+104 -2)
➕ csrc/jit_kernels/impls/sm100_fp8_fp8_mega_moe.hpp (+234 -0)
📝 deep_gemm/__init__.py (+1 -0)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+2 -2)
➕ deep_gemm/include/deep_gemm/impls/sm100_fp8_fp8_mega_moe.cuh (+1397 -0)
📝 deep_gemm/mega/__init__.py (+24 -0)
📝 tests/test_mega_moe.py (+23 -7)

📄 Description

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

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;
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);
The second commit relaxes the 512-alignment assertion on the SF buffers.

Why relaxing the 512 alignment is safe

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;
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;
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.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/371 **Author:** [@xuzhiyuan1](https://github.com/xuzhiyuan1) **Created:** 7/3/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `w8a8-megamoe` --- ### 📝 Commits (2) - [`f53e0b2`](https://github.com/deepseek-ai/DeepGEMM/commit/f53e0b20d16742af6c03a6c78870ced62e670734) Add W8A8 (FP8xFP8) Mega MoE support - [`7bc58b9`](https://github.com/deepseek-ai/DeepGEMM/commit/7bc58b926df91fee6ed0be6a1cf177d4ee7bad8c) Support unaligned intermediate hidden sizes ### 📊 Changes **7 files changed** (+1785 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/mega.hpp` (+104 -2) ➕ `csrc/jit_kernels/impls/sm100_fp8_fp8_mega_moe.hpp` (+234 -0) 📝 `deep_gemm/__init__.py` (+1 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+2 -2) ➕ `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp8_mega_moe.cuh` (+1397 -0) 📝 `deep_gemm/mega/__init__.py` (+24 -0) 📝 `tests/test_mega_moe.py` (+23 -7) </details> ### 📄 Description ## 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 ```bash # 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. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:17:13 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:31 +00:00

[PR #395] Add situ activation #413