From 2c7543130beb410d6f49b40bdc415c7abe75671f Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 18:36:40 +0800 Subject: [PATCH] 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 | 1 + deep_gemm/mega/__init__.py | 21 ++++++++++++++++++- .../glm52_nvfp4_group16_notes.md | 3 +++ tests/test_mega_moe.py | 15 +++++++------ 4 files changed, 33 insertions(+), 7 deletions(-) diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index aad7df2..d6ad870 100644 --- a/deep_gemm/__init__.py +++ b/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, ) diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 9fb5c7f..fdbb4c0 100644 --- a/deep_gemm/mega/__init__.py +++ b/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], diff --git a/megamoe-research-reports/glm52_nvfp4_group16_notes.md b/megamoe-research-reports/glm52_nvfp4_group16_notes.md index 841be2e..d0c03d4 100644 --- a/megamoe-research-reports/glm52_nvfp4_group16_notes.md +++ b/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`. diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 63f74ea..b0bc575 100644 --- a/tests/test_mega_moe.py +++ b/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