feat(megamoe): expose fp4 weight preparation helper

Add a top-level MegaMoE helper that handles source/runtime FP4 weight granularity before applying the existing MegaMoE weight layout transform.

Use the helper from the synthetic MegaMoE benchmark so SGLang can later follow the same contract for GLM-5.2 NVFP4 group16 checkpoints.

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

deep_gemm/__init__.py

@@ -85,6 +85,7 @@ from .mega import (
     SymmBuffer,
     get_symm_buffer_for_mega_moe,
     transform_weights_for_mega_moe,
+    prepare_fp4_weights_for_mega_moe,
     fp8_mega_moe,
     fp8_fp4_mega_moe,
 )

deep_gemm/mega/__init__.py

@@ -1,6 +1,6 @@
 import torch
 from typing import Tuple, Optional
-from ..utils.math import align
+from ..utils.math import align, requant_fp4_to_gran_k
 
 # noinspection PyBroadException
 try:
@@ -176,6 +176,25 @@ def transform_weights_for_mega_moe(
     return l1_weights, l2_weights
 
 
+def prepare_fp4_weights_for_mega_moe(
+    l1_weights: Tuple[torch.Tensor, torch.Tensor],
+    l2_weights: Tuple[torch.Tensor, torch.Tensor],
+    source_weight_gran_k: int = 32,
+    runtime_weight_gran_k: int = 32,
+) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
+    if source_weight_gran_k != runtime_weight_gran_k:
+        if source_weight_gran_k != 16 or runtime_weight_gran_k != 32:
+            raise RuntimeError(
+                f'Unsupported MegaMoE FP4 weight granularity conversion: '
+                f'{source_weight_gran_k} -> {runtime_weight_gran_k}')
+        l1_weights = requant_fp4_to_gran_k(
+            l1_weights[0], l1_weights[1], source_weight_gran_k, runtime_weight_gran_k)
+        l2_weights = requant_fp4_to_gran_k(
+            l2_weights[0], l2_weights[1], source_weight_gran_k, runtime_weight_gran_k)
+    return transform_weights_for_mega_moe(
+        l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
+
+
 def fp8_mega_moe(y: torch.Tensor,
                  l1_weights: Tuple[torch.Tensor, torch.Tensor],
                  l2_weights: Tuple[torch.Tensor, torch.Tensor],

megamoe-research-reports/glm52_nvfp4_group16_notes.md

@@ -18,6 +18,9 @@ carried an implicit group32 assumption.
 - `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 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`.

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, requant_fp4_to_gran_k
+from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8
 from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather
 from deep_gemm.testing import bench_kineto
 
@@ -97,15 +97,18 @@ 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 args.requant_group16_to_group32:
-                w, w_sf = requant_fp4_to_gran_k(w, w_sf, src_gran_k=16, dst_gran_k=32)
-            w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, runtime_weight_gran_k), num_groups)
+            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)
-        transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(
-            l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
+        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)
 
     # Run fused mega MoE
     # NOTES: copy x into buffer before each call because debug mode zeros the entire buffer