Files

LuminolT 8ad348fb11 feat(megamoe): add fp4 group16 to group32 requant path

Add a utility and synthetic benchmark path for converting FP4 group16 tensors into group32 tensors that the existing SM100 block32 MegaMoE kernels can consume.

Document that this is a requantization path rather than a lossless metadata rewrite, so GLM-5.2 accuracy validation is still required before production use.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/utils/math.py tests/test_layout.py tests/test_mega_moe.py

Tested: git diff --check

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

2026-07-08 18:34:33 +08:00

2.6 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.
requant_fp4_to_gran_k(...) provides an explicit experiment path for converting a group16 FP4 checkpoint tensor into a group32 tensor that the existing block32 kernels can consume.
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.

Requantization path

If the SM100 MXF4 MMA path cannot consume group16 scales directly, the fallback engineering path is checkpoint conversion:

dequantize group16 FP4 values with their original scales;
requantize the restored values with group32 scales;
run the existing recipe=(1, 1, 32) MegaMoE kernel.

This is not a lossless metadata rewrite. It changes the quantized checkpoint and must go through GLM-5.2 accuracy validation before it can be used as a production answer.

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;
tests/test_mega_moe.py --weight-gran-k 16 --requant-group16-to-group32 for the checkpoint conversion path against existing block32 kernels;
SGLang GLM-5.2 NVFP4 real-weight layout build and 8-card e2e smoke.

2.6 KiB Raw Blame History