Files

LuminolT 79fcfd6abf feat(megamoe): add nvfp4 group16 capability gate

Allow SM100 FP4 scale layout transforms to accept group16 and thread weight granularity through the MegaMoE Python wrapper, API checks, and synthetic benchmark entrypoint.

Keep fused SM100 MegaMoE compute behind an explicit group16 capability gate until the SFB/TMEM/MMA scale path is updated and validated.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/mega/__init__.py tests/test_mega_moe.py tests/generators.py

Tested: git diff --check

Not-tested: CUDA build and SM100/B300 runtime validation are not available locally.

2026-07-08 18:29:09 +08:00

1.8 KiB

Raw Blame History

GLM-5.2 NVFP4 group16 MegaMoE notes

Context

GLM-5.2 NVFP4 checkpoints use FP4 weight scales with group_size=16. The existing SM100 FP8xFP4 MegaMoE path was developed and tested with recipe=(1, 1, 32), so both the weight scale layout and fused compute path carried an implicit group32 assumption.

Current patch scope

transform_sf_into_required_layout now accepts SM100 packed UE8M0 scale transforms with gran_k=16.
transform_weights_for_mega_moe(..., weight_gran_k=...) can apply the UTCCP scale transpose with a group16-aware 128-element tiling.
tests/test_mega_moe.py exposes --weight-gran-k 16|32 so synthetic runs can reproduce GLM-style group16 inputs without loading model weights.
The fused SM100 MegaMoE compute API now performs an explicit capability check for recipe=(1, 1, 16) instead of failing earlier with Unknown SF transformation.

Remaining kernel work

The fused compute kernel still uses the SM100 MXF4 block-scale path and its current shared-memory/TMEM scale layout is group32-equivalent. Supporting group16 correctly requires auditing at least:

weight scale TMA width per K block;
SFB shared-memory and tensor-memory column allocation;
scale id selection passed to the MMA instruction;
the UTCCP scale transpose layout consumed by SM100_UTCCP_4x32dp128bit_2cta.

Until that kernel work is complete and validated on B300/SM100, group16 should be treated as layout-supported but fused-compute unsupported.

Validation target

After kernel support is added, validate with:

existing group32 MegaMoE tests unchanged;
tests/test_layout.py on SM100 for gran_k=16;
tests/test_mega_moe.py --weight-gran-k 16 --ncu-profile-only for synthetic fused execution;
SGLang GLM-5.2 NVFP4 real-weight layout build and 8-card e2e smoke.

1.8 KiB Raw Blame History