MegaMoE: add cluster-paired FP4 weight packets

[foraxe]

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.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

No actionable correctness issues were identified in the changed code.

v4

⚠️ 未完成评审（no_result_file：模型未产出结果文件）

v3

This MR adds an opt-in cluster-paired FP4 weight-packet layout for FP8xFP4 Mega MoE, implemented across 9 files in three commits (feat c99da9c, test 1097cba, test 779e283). The change is complete, internally consistent, and follows the described access order (expert -> CTA pair -> K tile -> CTA peer).

Key elements verified:

• Python _pack_fp4_weight_tiles packs each 128x256 logical FP4 tile immediately followed by its scale words in one shared allocation via torch.as_strided; transform_weights_for_mega_moe(..., tile_pack=True) exposes it opt-in and leaves BF16 / row-major paths untouched.
• Host validation in csrc/apis/mega.hpp detects the packet layout from shape+strides, asserts exact strides and weight/SF adjacency, derives intermediate_hidden from packed shapes, and cleanly falls back to the row-major K-major contiguous path. Shared-expert weights retain the original layout and checks.
• Launch code selects make_tma_3d_desc for packet weights/SF vs 2D descriptors for row-major, and sets evict_first_weights = tile_packed_weights && block_m == 16 (the smallest-token heuristic branch).
• Heuristics thread tile_packed_weights through and use one epilogue warpgroup (1 vs 2) only in the smallest-token packed case.
• Kernel adds kEvictFirstWeights/kTilePackedWeights template params (ordering matches the compile-test stubs), a 3D TMA copy path with cache-hint plumbing, and weight_tile_idx computation matching the Python packing order. Static assert kWeightPacketK % BLOCK_K == 0 supports both K256 and K128 compute tiles.
• Tests cover packet unit layout, K256/K128 compile stubs, a bitwise A/B correctness comparison (--compare-tile-packed uses torch.equal on output and stats, the correct gate for a layout-only transform, and correctly asserts num_shared_experts == 0 and not bf16xbf16), an alternating A/B benchmark, and a deterministic --seed-offset.

The implementation matches the PR description and I found no correctness or consistency defects. Only minor maintainability suggestions are noted below. Note: the review environment permits git-only command execution, so I could not re-run the build/pytest/four-GPU steps; I relied on static review plus the MR's recorded validation results.

Files reviewed: 9
Issues found: 🔵 3 suggestion
Inline comments posted: 3

⚠️ Parse warning: [v4] no_result_file：模型未产出结果文件

[foraxe]

🤖 ds-review-bot Code Review

v6

No actionable correctness issues were identified in the changed code.

v4

⚠️ 未完成评审（no_result_file：模型未产出结果文件）

v3

This MR adds an opt-in cluster-paired FP4 weight-packet layout for FP8xFP4 Mega MoE, implemented across 9 files in three commits (feat c99da9c, test 1097cba, test 779e283). The change is complete, internally consistent, and follows the described access order (expert -> CTA pair -> K tile -> CTA peer).

Key elements verified:

• Python _pack_fp4_weight_tiles packs each 128x256 logical FP4 tile immediately followed by its scale words in one shared allocation via torch.as_strided; transform_weights_for_mega_moe(..., tile_pack=True) exposes it opt-in and leaves BF16 / row-major paths untouched.
• Host validation in csrc/apis/mega.hpp detects the packet layout from shape+strides, asserts exact strides and weight/SF adjacency, derives intermediate_hidden from packed shapes, and cleanly falls back to the row-major K-major contiguous path. Shared-expert weights retain the original layout and checks.
• Launch code selects make_tma_3d_desc for packet weights/SF vs 2D descriptors for row-major, and sets evict_first_weights = tile_packed_weights && block_m == 16 (the smallest-token heuristic branch).
• Heuristics thread tile_packed_weights through and use one epilogue warpgroup (1 vs 2) only in the smallest-token packed case.
• Kernel adds kEvictFirstWeights/kTilePackedWeights template params (ordering matches the compile-test stubs), a 3D TMA copy path with cache-hint plumbing, and weight_tile_idx computation matching the Python packing order. Static assert kWeightPacketK % BLOCK_K == 0 supports both K256 and K128 compute tiles.
• Tests cover packet unit layout, K256/K128 compile stubs, a bitwise A/B correctness comparison (--compare-tile-packed uses torch.equal on output and stats, the correct gate for a layout-only transform, and correctly asserts num_shared_experts == 0 and not bf16xbf16), an alternating A/B benchmark, and a deterministic --seed-offset.

The implementation matches the PR description and I found no correctness or consistency defects. Only minor maintainability suggestions are noted below. Note: the review environment permits git-only command execution, so I could not re-run the build/pytest/four-GPU steps; I relied on static review plus the MR's recorded validation results.

Files reviewed: 9 Issues found: 🔵 3 suggestion Inline comments posted: 3

⚠️ Parse warning: [v4] no_result_file：模型未产出结果文件

Addressed all three review suggestions.

Changes:

• Documented the fixed packet width beside the heuristic.
• Added a clear error for incomplete CTA pairs.
• Replaced the indirect block_m == 16 check with a shared named constant.
• Added a unit test for the new error message.
• Updated the PR validation count.

Validation:

• Host extension build passed.
• Packet tests: 3/3 passed.
• K128 and K256 SM100 compilation passed.

Pushed commit: d1e7368