From 8ad348fb1185ae873f2180263fa87b681d87eeba Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 18:34:33 +0800 Subject: [PATCH] feat(megamoe): add fp4 group16 to group32 requant path Add a utility and synthetic benchmark path for converting FP4 group16 tensors into group32 tensors that the existing SM100 block32 MegaMoE kernels can consume. Document that this is a requantization path rather than a lossless metadata rewrite, so GLM-5.2 accuracy validation is still required before production use. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile 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/utils/math.py | 23 ++++++++++++++++++- .../glm52_nvfp4_group16_notes.md | 18 +++++++++++++++ tests/test_layout.py | 19 +++++++++++++++ tests/test_mega_moe.py | 19 ++++++++++----- 4 files changed, 72 insertions(+), 7 deletions(-) diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index f1582ed..1e745e0 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -140,4 +140,25 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128 x_dequantized = _dequantize_from_fp4_e2m1(unpacked) group_idx = torch.arange(n, device=packed.device) // gran_k x_restored = x_dequantized * sf[:, group_idx] - return x_restored \ No newline at end of file + return x_restored + + +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]: + assert packed.dtype == torch.int8 + assert packed.dim() >= 2 + assert sf.shape[:-1] == packed.shape[:-1] + original_packed_shape = packed.shape + original_sf_shape = sf.shape + + 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) + repacked, rescaled = per_token_cast_to_fp4( + restored, use_ue8m0=use_packed_ue8m0, gran_k=dst_gran_k, + use_packed_ue8m0=use_packed_ue8m0) + + return ( + repacked.reshape(original_packed_shape), + rescaled.reshape(*original_sf_shape[:-1], rescaled.size(-1)), + ) diff --git a/megamoe-research-reports/glm52_nvfp4_group16_notes.md b/megamoe-research-reports/glm52_nvfp4_group16_notes.md index 3bfa775..841be2e 100644 --- a/megamoe-research-reports/glm52_nvfp4_group16_notes.md +++ b/megamoe-research-reports/glm52_nvfp4_group16_notes.md @@ -15,6 +15,9 @@ carried an implicit group32 assumption. 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. - 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`. @@ -33,6 +36,19 @@ 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. +## Requantization path + +If the SM100 MXF4 MMA path cannot consume group16 scales directly, the fallback +engineering path is checkpoint conversion: + +1. dequantize group16 FP4 values with their original scales; +2. requantize the restored values with group32 scales; +3. 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: @@ -41,4 +57,6 @@ After kernel support is added, validate with: - `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. diff --git a/tests/test_layout.py b/tests/test_layout.py index a0d4a02..a96851d 100644 --- a/tests/test_layout.py +++ b/tests/test_layout.py @@ -4,6 +4,7 @@ from deep_gemm.testing import bench_kineto, count_bytes, get_arch_major from deep_gemm.utils import ( align, ceil_div, per_token_cast_to_fp8, per_channel_cast_to_fp8, + per_token_cast_to_fp4, cast_back_from_fp4, requant_fp4_to_gran_k, get_tma_aligned_size, get_mn_major_tma_aligned_tensor, get_mn_major_tma_aligned_packed_ue8m0_tensor, @@ -104,9 +105,27 @@ def test_k_grouped_sf_layout_kernels() -> None: print() +def test_fp4_requant_granularity() -> None: + print('Testing FP4 requant granularity:') + for m, n in ((128, 7168), (257, 3072)): + x = torch.randn((m, n), dtype=torch.bfloat16, device='cuda') + fp4_g16 = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16) + 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) + + 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) + print(f' > Requant ({m=}, {n=}): group16 -> group32') + print() + + if __name__ == '__main__': torch.manual_seed(1) random.seed(1) test_sf_layout_kernels() test_k_grouped_sf_layout_kernels() + test_fp4_requant_granularity() diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index b25a236..63f74ea 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 @@ -47,7 +47,9 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): hidden, intermediate_hidden = args.hidden, args.intermediate_hidden num_experts, num_topk = args.num_experts, args.num_topk num_experts_per_rank = num_experts // num_ranks + runtime_weight_gran_k = 32 if args.requant_group16_to_group32 else args.weight_gran_k assert num_tokens <= num_max_tokens_per_rank + assert not args.requant_group16_to_group32 or args.weight_gran_k == 16 # Allocate symmetric memory buffer = deep_gemm.get_symm_buffer_for_mega_moe( @@ -95,13 +97,15 @@ 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) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, args.weight_gran_k), num_groups) + 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) 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=args.weight_gran_k) + 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 @@ -118,7 +122,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): transformed_l1_weights, transformed_l2_weights, buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, - recipe=(1, 1, args.weight_gran_k), + recipe=(1, 1, runtime_weight_gran_k), activation_clamp=args.activation_clamp, fast_math=bool(args.fast_math) ) @@ -129,6 +133,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): dist_print(f' > Hidden: {hidden}', once_in_node=True) dist_print(f' > Intermediate: {intermediate_hidden}', once_in_node=True) dist_print(f' > Experts: {num_topk}/{num_experts}', once_in_node=True) + dist_print(f' > Weight granularity: source={args.weight_gran_k}, runtime={runtime_weight_gran_k}', once_in_node=True) dist_print(f' > Buffer: {buffer.buffer.nbytes / 2 ** 30:.3f} GiB', once_in_node=True) dist_print(once_in_node=True) @@ -170,7 +175,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l1_y = torch.empty((n, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda') deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, args.weight_gran_k)) + use_psum_layout=True, recipe=(1, 1, runtime_weight_gran_k)) # noinspection PyCallingNonCallable l1_y = tilelang_ops.swiglu_apply_weight_to_fp8( x=l1_y, @@ -187,7 +192,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, args.weight_gran_k)) + use_psum_layout=True, recipe=(1, 1, runtime_weight_gran_k)) return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline # Check correctness (must be bitwise identical) @@ -281,6 +286,8 @@ if __name__ == '__main__': parser.add_argument('--fast-math', type=int, default=1, help='Enable fast math (0 or 1, default: 1)') parser.add_argument('--weight-gran-k', type=int, default=32, choices=(16, 32), help='FP4 weight scale granularity along K') + parser.add_argument('--requant-group16-to-group32', action='store_true', + help='Simulate loading a group16 FP4 checkpoint and requantizing it to DeepGEMM group32') # Test settings parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test')