deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #394] MegaMoE: add cluster-paired FP4 weight packets #414

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/394
Author: @foraxe
Created: 7/28/2026
Status: 🔄 Open

Base: main ← Head: pr/mega-moe-fp4-weight-packets

📝 Commits (4)

c99da9c feat(mega-moe): add cluster-paired FP4 weight packets
1097cba test(mega-moe): add packet-layout A/B coverage
779e283 test(mega-moe): expose deterministic A/B seed
d1e7368 refactor(mega-moe): clarify packet layout invariants

📊 Changes

9 files changed (+522 additions, -59 deletions)

View changed files

📝 csrc/apis/mega.hpp (+71 -13)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+16 -5)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+79 -23)
📝 deep_gemm/include/deep_gemm/common/tma_copy.cuh (+9 -7)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+54 -4)
📝 deep_gemm/mega/__init__.py (+73 -3)
➕ tests/compile_mega_moe_tile_packets.cu (+47 -0)
📝 tests/test_mega_moe.py (+74 -4)
➕ tests/test_mega_moe_tile_packets.py (+99 -0)

📄 Description

Summary

This change adds an optional weight layout for FP8xFP4 Mega MoE:

pack each routed expert's 128x256 logical FP4 weight tile together with the
scale words consumed by that tile;
place the two tiles read by a cooperating CTA pair next to each other;
load packet weights and scales through 3D TMA descriptors;
use the same packet allocation with both K256 and K128 compute tiles.

The packet and scale views share one allocation. Expert ownership remains
sharded by rank, so this does not replicate expert weights.

Motivation

Sparse Mega MoE batches touch many experts but read only a small part of each
expert at a time. In the row-major layout, cooperating CTAs and their scale
loads can be far apart in memory.

The packet layout follows the kernel's actual access order:

expert -> CTA pair -> K tile -> CTA peer

This keeps each CTA pair's weight bodies adjacent and places the matching
scale words directly after each body.

For sparse small-token work, the packet path also uses one output warpgroup
and marks one-use routed weight loads as EVICT_FIRST. Other workloads keep
the existing production heuristic.

API

The layout is opt-in during the existing model-load weight transform:

l1_weights, l2_weights = transform_weights_for_mega_moe(
    l1_weights,
    l2_weights,
    tile_pack=True,
)

The runtime detects the packet layout from its shape and strides. Existing
row-major weights continue to use the original path.

Correctness

On four GB200 GPUs, row-major and packet layouts produced bitwise-identical
output tensors and cumulative expert statistics for every reported workload.

Performance

Configuration: FP8xFP4, EP4, 384 experts, top-k 6, hidden size 7168,
intermediate size 3072, no shared experts, deterministic random routing.
The metric is the maximum rank median from alternating same-process A/B runs.

Tokens per rank	Row-major	Packet	Improvement
1	98.718 us	89.329 us	9.51%
8	281.477 us	269.167 us	4.37%
32	540.133 us	534.621 us	1.02%
128	591.242 us	582.190 us	1.53%
2,048	1421.0 us	1415.5 us	0.39%
8,192	4287.0 us	4276.0 us	0.26%

The small-token candidate includes the packet layout, one output warpgroup,
and EVICT_FIRST for routed packet loads. The large-token comparisons use
the same default compute setting on both sides and differ only in weight
layout.

Reproduce the correctness and alternating A/B measurements with:

while read -r tokens repeats calls; do
  CUDA_VISIBLE_DEVICES=0,1,2,3 \
  python tests/test_mega_moe.py \
    --num-processes 4 \
    --num-max-tokens-per-rank "${tokens}" \
    --num-tokens "${tokens}" \
    --hidden 7168 \
    --intermediate-hidden 3072 \
    --num-shared-experts 0 \
    --num-experts 384 \
    --num-topk 6 \
    --mma-type fp8xfp4 \
    --num-correctness-tests 0 \
    --seed-offset 1009 \
    --compare-tile-packed \
    --benchmark-tile-packed-ab \
    --ab-repeats "${repeats}" \
    --ab-num-tests "${calls}"
done <<'EOF'
1 5 30
8 5 20
32 5 20
128 4 20
2048 4 20
8192 3 10
EOF

Scope and fallback

FP8xFP4 routed expert weights can use packets.
Shared-expert weights retain the existing layout and load policy.
Existing row-major weights remain fully supported.
Unsupported packet shapes fail closed during host validation.

Validation

Host extension build: passed.
K256 and K128 packet kernel compilation for SM100 family: passed.
Packet layout unit tests: 3 passed.
Four-GPU rerun on the clean PR branch: pending an idle node.

_{🔄 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/394 **Author:** [@foraxe](https://github.com/foraxe) **Created:** 7/28/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `pr/mega-moe-fp4-weight-packets` --- ### 📝 Commits (4) - [`c99da9c`](https://github.com/deepseek-ai/DeepGEMM/commit/c99da9c2ca2728de20f741e33b5ae87f65fe470c) feat(mega-moe): add cluster-paired FP4 weight packets - [`1097cba`](https://github.com/deepseek-ai/DeepGEMM/commit/1097cbae5e048ff339c94455399594143e743cb0) test(mega-moe): add packet-layout A/B coverage - [`779e283`](https://github.com/deepseek-ai/DeepGEMM/commit/779e28303c5e3d44ea287e58d7e834b21af1e972) test(mega-moe): expose deterministic A/B seed - [`d1e7368`](https://github.com/deepseek-ai/DeepGEMM/commit/d1e73685349137b5019ec7b75fe640509e89fbd4) refactor(mega-moe): clarify packet layout invariants ### 📊 Changes **9 files changed** (+522 additions, -59 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/mega.hpp` (+71 -13) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+16 -5) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+79 -23) 📝 `deep_gemm/include/deep_gemm/common/tma_copy.cuh` (+9 -7) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+54 -4) 📝 `deep_gemm/mega/__init__.py` (+73 -3) ➕ `tests/compile_mega_moe_tile_packets.cu` (+47 -0) 📝 `tests/test_mega_moe.py` (+74 -4) ➕ `tests/test_mega_moe_tile_packets.py` (+99 -0) </details> ### 📄 Description ## Summary This change adds an optional weight layout for FP8xFP4 Mega MoE: - pack each routed expert's 128x256 logical FP4 weight tile together with the scale words consumed by that tile; - place the two tiles read by a cooperating CTA pair next to each other; - load packet weights and scales through 3D TMA descriptors; - use the same packet allocation with both K256 and K128 compute tiles. The packet and scale views share one allocation. Expert ownership remains sharded by rank, so this does not replicate expert weights. ## Motivation Sparse Mega MoE batches touch many experts but read only a small part of each expert at a time. In the row-major layout, cooperating CTAs and their scale loads can be far apart in memory. The packet layout follows the kernel's actual access order: `expert -> CTA pair -> K tile -> CTA peer` This keeps each CTA pair's weight bodies adjacent and places the matching scale words directly after each body. For sparse small-token work, the packet path also uses one output warpgroup and marks one-use routed weight loads as `EVICT_FIRST`. Other workloads keep the existing production heuristic. ## API The layout is opt-in during the existing model-load weight transform: ```python l1_weights, l2_weights = transform_weights_for_mega_moe( l1_weights, l2_weights, tile_pack=True, ) ``` The runtime detects the packet layout from its shape and strides. Existing row-major weights continue to use the original path. ## Correctness On four GB200 GPUs, row-major and packet layouts produced bitwise-identical output tensors and cumulative expert statistics for every reported workload. ## Performance Configuration: FP8xFP4, EP4, 384 experts, top-k 6, hidden size 7168, intermediate size 3072, no shared experts, deterministic random routing. The metric is the maximum rank median from alternating same-process A/B runs. | Tokens per rank | Row-major | Packet | Improvement | |---:|---:|---:|---:| | 1 | 98.718 us | 89.329 us | 9.51% | | 8 | 281.477 us | 269.167 us | 4.37% | | 32 | 540.133 us | 534.621 us | 1.02% | | 128 | 591.242 us | 582.190 us | 1.53% | | 2,048 | 1421.0 us | 1415.5 us | 0.39% | | 8,192 | 4287.0 us | 4276.0 us | 0.26% | The small-token candidate includes the packet layout, one output warpgroup, and `EVICT_FIRST` for routed packet loads. The large-token comparisons use the same default compute setting on both sides and differ only in weight layout. Reproduce the correctness and alternating A/B measurements with: ```bash while read -r tokens repeats calls; do CUDA_VISIBLE_DEVICES=0,1,2,3 \ python tests/test_mega_moe.py \ --num-processes 4 \ --num-max-tokens-per-rank "${tokens}" \ --num-tokens "${tokens}" \ --hidden 7168 \ --intermediate-hidden 3072 \ --num-shared-experts 0 \ --num-experts 384 \ --num-topk 6 \ --mma-type fp8xfp4 \ --num-correctness-tests 0 \ --seed-offset 1009 \ --compare-tile-packed \ --benchmark-tile-packed-ab \ --ab-repeats "${repeats}" \ --ab-num-tests "${calls}" done <<'EOF' 1 5 30 8 5 20 32 5 20 128 4 20 2048 4 20 8192 3 10 EOF ``` ## Scope and fallback - FP8xFP4 routed expert weights can use packets. - Shared-expert weights retain the existing layout and load policy. - Existing row-major weights remain fully supported. - Unsupported packet shapes fail closed during host validation. ## Validation - Host extension build: passed. - K256 and K128 packet kernel compilation for SM100 family: passed. - Packet layout unit tests: 3 passed. - Four-GPU rerun on the clean PR branch: pending an idle node. --- <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:31 +00:00

bsakaguchi referenced this issue

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

[PR #420] [CLOSED] fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) #432

bsakaguchi referenced this issue

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

[PR #421] fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) #433

bsakaguchi referenced this issue

2026-10-26 21:18:15 +00:00