From 007c645f87e0493c6d0b18dfc630a700cd1b02c3 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 18:48:04 +0800 Subject: [PATCH] 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. --- deep_gemm/mega/__init__.py | 46 ++++++++++++++----- deep_gemm/utils/math.py | 9 ++-- .../glm52_nvfp4_group16_notes.md | 18 +++++++- tests/test_layout.py | 11 +++++ tests/test_mega_moe.py | 29 +++++++----- 5 files changed, 86 insertions(+), 27 deletions(-) diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index fdbb4c0..1affb3c 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, requant_fp4_to_gran_k +from ..utils.math import align, requant_fp4_to_gran_k, unpack_ue8m0_from_int # noinspection PyBroadException try: @@ -176,21 +176,45 @@ 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]]: +def _prepare_raw_fp4_weight_for_mega_moe( + weights: Tuple[torch.Tensor, torch.Tensor], + source_weight_gran_k: int, + runtime_weight_gran_k: int, + source_scale_packed_ue8m0: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + weight, weight_sf = weights 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) + weight, weight_sf = requant_fp4_to_gran_k( + weight, weight_sf, + source_weight_gran_k, runtime_weight_gran_k, + src_scale_packed_ue8m0=source_scale_packed_ue8m0) + source_scale_packed_ue8m0 = False + + if source_scale_packed_ue8m0: + weight_sf = unpack_ue8m0_from_int(weight_sf) + num_groups, mn, packed_k = weight.shape + weight_sf = _C.transform_sf_into_required_layout( + weight_sf, mn, packed_k * 2, (1, runtime_weight_gran_k), num_groups) + return weight, weight_sf + + +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, + source_scale_packed_ue8m0: bool = False, +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + l1_weights = _prepare_raw_fp4_weight_for_mega_moe( + l1_weights, source_weight_gran_k, runtime_weight_gran_k, + source_scale_packed_ue8m0) + l2_weights = _prepare_raw_fp4_weight_for_mega_moe( + l2_weights, source_weight_gran_k, runtime_weight_gran_k, + source_scale_packed_ue8m0) return transform_weights_for_mega_moe( l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k) diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index 1e745e0..9289d82 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -144,7 +144,8 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128 def requant_fp4_to_gran_k(packed: torch.Tensor, sf: torch.Tensor, src_gran_k: int, - dst_gran_k: int, use_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + dst_gran_k: int, src_scale_packed_ue8m0: bool = False, + dst_scale_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: assert packed.dtype == torch.int8 assert packed.dim() >= 2 assert sf.shape[:-1] == packed.shape[:-1] @@ -153,10 +154,10 @@ def requant_fp4_to_gran_k(packed: torch.Tensor, sf: torch.Tensor, src_gran_k: in packed_2d = packed.reshape(-1, packed.size(-1)) sf_2d = sf.reshape(-1, sf.size(-1)) - restored = cast_back_from_fp4(packed_2d, sf_2d, src_gran_k, use_packed_ue8m0) + restored = cast_back_from_fp4(packed_2d, sf_2d, src_gran_k, src_scale_packed_ue8m0) repacked, rescaled = per_token_cast_to_fp4( - restored, use_ue8m0=use_packed_ue8m0, gran_k=dst_gran_k, - use_packed_ue8m0=use_packed_ue8m0) + restored, use_ue8m0=True, gran_k=dst_gran_k, + use_packed_ue8m0=dst_scale_packed_ue8m0) return ( repacked.reshape(original_packed_shape), diff --git a/megamoe-research-reports/glm52_nvfp4_group16_notes.md b/megamoe-research-reports/glm52_nvfp4_group16_notes.md index d0c03d4..bda0e63 100644 --- a/megamoe-research-reports/glm52_nvfp4_group16_notes.md +++ b/megamoe-research-reports/glm52_nvfp4_group16_notes.md @@ -20,7 +20,8 @@ carried an implicit group32 assumption. 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. + 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`. @@ -39,6 +40,21 @@ group16 correctly requires auditing at least: 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: + +1. raw checkpoint scale: per-weight-group scale from model loading or synthetic + quantization, either float UE8M0 values or packed UE8M0 integers; +2. DeepGEMM runtime scale layout: output of `transform_sf_into_required_layout`, + packed, MN-major, TMA-aligned; +3. 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 diff --git a/tests/test_layout.py b/tests/test_layout.py index a96851d..80b5ea8 100644 --- a/tests/test_layout.py +++ b/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() diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index b0bc575..51214ef 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 +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