Files

LuminolT 007c645f87 fix(megamoe): normalize fp4 weight preparation contract

Separate source and destination FP4 scale packing in requant_fp4_to_gran_k so group16-to-group32 conversion always recomputes UE8M0 runtime scales by default.

Make prepare_fp4_weights_for_mega_moe accept raw grouped FP4 weights and scales, then perform optional requantization, DeepGEMM scale layout transform, and MegaMoE UTCCP weight transform internally.

Update the MegaMoE synthetic benchmark so baseline grouped GEMM uses runtime-layout weights while fused MegaMoE uses transformed weights from the same raw source tensors.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/__init__.py deep_gemm/mega/__init__.py 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:48:04 +08:00

3.5 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.
prepare_fp4_weights_for_mega_moe(...) wraps the source/runtime granularity decision for callers such as SGLang: source group16 can be requantized to runtime group32, converted into DeepGEMM scale layout, and then passed through the normal MegaMoE weight transform.
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.

Scale layout contract

There are three distinct scale states:

raw checkpoint scale: per-weight-group scale from model loading or synthetic quantization, either float UE8M0 values or packed UE8M0 integers;
DeepGEMM runtime scale layout: output of transform_sf_into_required_layout, packed, MN-major, TMA-aligned;
MegaMoE weight scale layout: runtime scale layout after the UTCCP transpose required by the fused MegaMoE kernel.

prepare_fp4_weights_for_mega_moe(...) accepts state 1 and returns state 3. transform_weights_for_mega_moe(...) remains a lower-level helper that accepts state 2 and returns state 3.

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.

3.5 KiB Raw Blame History