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.

tests/test_layout.py

@@ -113,11 +113,22 @@ def test_fp4_requant_granularity() -> None:
         repacked, rescaled = requant_fp4_to_gran_k(fp4_g16[0], fp4_g16[1], 16, 32)
         restored = cast_back_from_fp4(fp4_g16[0], fp4_g16[1], gran_k=16)
         ref_repacked, ref_rescaled = per_token_cast_to_fp4(restored, use_ue8m0=True, gran_k=32)
+        fp4_g16_packed = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16, use_packed_ue8m0=True)
+        repacked_from_packed, rescaled_from_packed = requant_fp4_to_gran_k(
+            fp4_g16_packed[0], fp4_g16_packed[1], 16, 32, src_scale_packed_ue8m0=True)
+        ref_repacked_packed, ref_rescaled_packed = per_token_cast_to_fp4(
+            restored, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
+        repacked_to_packed, rescaled_to_packed = requant_fp4_to_gran_k(
+            fp4_g16[0], fp4_g16[1], 16, 32, dst_scale_packed_ue8m0=True)
 
         assert repacked.shape == fp4_g16[0].shape
         assert rescaled.shape[-1] == ceil_div(n, 32)
         assert torch.equal(repacked, ref_repacked)
         assert torch.equal(rescaled, ref_rescaled)
+        assert torch.equal(repacked_from_packed, ref_repacked)
+        assert torch.equal(rescaled_from_packed, ref_rescaled)
+        assert torch.equal(repacked_to_packed, ref_repacked_packed)
+        assert torch.equal(rescaled_to_packed, ref_rescaled_packed)
         print(f' > Requant ({m=}, {n=}): group16 -> group32')
     print()
 

tests/test_mega_moe.py

@@ -7,7 +7,7 @@ import torch.distributed as dist
 from typing import Tuple
 
 import deep_gemm
-from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8
+from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8, requant_fp4_to_gran_k
 from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather
 from deep_gemm.testing import bench_kineto
 
@@ -97,18 +97,25 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
             for i in range(num_groups):
                 w[i], w_sf[i] = per_token_cast_to_fp4(
                     bf16_weights[i], use_ue8m0=True, gran_k=args.weight_gran_k)
-            if not args.requant_group16_to_group32:
-                w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, runtime_weight_gran_k), num_groups)
             return w, w_sf
 
-        l1_weights = cast_grouped_weights_to_fp4(l1_weights)
-        l2_weights = cast_grouped_weights_to_fp4(l2_weights)
-        if args.requant_group16_to_group32:
-            transformed_l1_weights, transformed_l2_weights = deep_gemm.prepare_fp4_weights_for_mega_moe(
-                l1_weights, l2_weights, source_weight_gran_k=16, runtime_weight_gran_k=32)
-        else:
-            transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(
-                l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
+        def prepare_grouped_runtime_weights(weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+            w, w_sf = weights
+            if args.requant_group16_to_group32:
+                w, w_sf = requant_fp4_to_gran_k(w, w_sf, src_gran_k=16, dst_gran_k=32)
+            num_groups, n, packed_k = w.shape
+            w_sf = deep_gemm.transform_sf_into_required_layout(
+                w_sf, n, packed_k * 2, (1, runtime_weight_gran_k), num_groups)
+            return w, w_sf
+
+        source_l1_weights = cast_grouped_weights_to_fp4(l1_weights)
+        source_l2_weights = cast_grouped_weights_to_fp4(l2_weights)
+        l1_weights = prepare_grouped_runtime_weights(source_l1_weights)
+        l2_weights = prepare_grouped_runtime_weights(source_l2_weights)
+        transformed_l1_weights, transformed_l2_weights = deep_gemm.prepare_fp4_weights_for_mega_moe(
+            source_l1_weights, source_l2_weights,
+            source_weight_gran_k=args.weight_gran_k,
+            runtime_weight_gran_k=runtime_weight_gran_k)
 
     # Run fused mega MoE
     # NOTES: copy x into buffer before each call because debug mode zeros the entire buffer