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