diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index c3e73931e4..45bd116958 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -9,6 +9,9 @@ from .cuda_helpers import find_cuda_home, get_cuda_arch +_SM90_TINY_WEIGHT_SCALE_THRESHOLD = 1.0e-12 +_SM90_PRUNED_WEIGHT_SCALE = 1.0e-5 + if TYPE_CHECKING: from tvm_ffi.module import Module @@ -374,11 +377,65 @@ def get_symm_buffer_for_mega_moe(group, ) +def _sanitize_sm90_fp8_weight_blocks( + weight: torch.Tensor, + scale: torch.Tensor, + *, + threshold: float = _SM90_TINY_WEIGHT_SCALE_THRESHOLD, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Make effectively pruned FP8 blocks safe for SM90 MegaMoE.""" + if weight.ndim != 3 or scale.ndim != 3: + raise ValueError( + f"SM90 MegaMoE expects rank-3 weights/scales, got " + f"weight={tuple(weight.shape)}, scale={tuple(scale.shape)}" + ) + num_experts, n, k = weight.shape + expected_scale_shape = (num_experts, n // 128, k // 128) + if n % 128 != 0 or k % 128 != 0: + raise ValueError(f"SM90 FP8 weight shape must be 128-aligned, got {(n, k)}") + if tuple(scale.shape) != expected_scale_shape: + raise ValueError( + f"SM90 FP8 scale shape mismatch: got {tuple(scale.shape)}, " + f"expected {expected_scale_shape}" + ) + if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz): + raise ValueError(f"SM90 MegaMoE expects FP8 weights, got {weight.dtype}") + + tiny = scale.abs() < threshold + if not bool(tiny.any()): + return weight, scale + + sanitized_weight = weight.clone() + sanitized_scale = scale.clone() + weight_blocks = sanitized_weight.view( + num_experts, + n // 128, + 128, + k // 128, + 128, + ) + weight_blocks.view(torch.uint8).masked_fill_( + tiny[:, :, None, :, None], + 0, + ) + for expert in tiny.any(dim=(1, 2)).nonzero(as_tuple=False).flatten().tolist(): + normal_scales = sanitized_scale[expert][~tiny[expert]] + replacement = ( + normal_scales.median() + if normal_scales.numel() + else sanitized_scale.new_tensor(_SM90_PRUNED_WEIGHT_SCALE) + ) + sanitized_scale[expert].masked_fill_(tiny[expert], replacement) + return sanitized_weight, sanitized_scale + + def transform_weights_for_mega_moe_sm90( l1_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor] ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: l1_fp8, l1_sf = l1_weights + l1_fp8, l1_sf = _sanitize_sm90_fp8_weight_blocks(l1_fp8, l1_sf) + l2_fp8, l2_sf = _sanitize_sm90_fp8_weight_blocks(*l2_weights) def _interleave_one(t, gran: int = 8) -> torch.Tensor: g, n, *rest = t.shape @@ -387,7 +444,7 @@ def _interleave_one(t, gran: int = 8) -> torch.Tensor: up = t[:, half:].reshape(g, half // gran, gran, *rest) return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) - return (_interleave_one(l1_fp8), l1_sf), l2_weights + return (_interleave_one(l1_fp8), l1_sf), (l2_fp8, l2_sf) def fp8_mega_moe(y: torch.Tensor, diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index 86dba4b4c9..dcc3536b42 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -159,6 +159,7 @@ MEGA_MOE_BLACKWELL=( ) MEGA_MOE_HOPPER=( test_mega_moe_hopper.py + test_mega_moe_sm90_weight_transform.py test_mega_moe_pre_dispatch_sm90.py ) MEGA_MOE_ALL=( diff --git a/sgl_deep_gemm/tests/test_mega_moe_sm90_weight_transform.py b/sgl_deep_gemm/tests/test_mega_moe_sm90_weight_transform.py new file mode 100644 index 0000000000..a9bf36f049 --- /dev/null +++ b/sgl_deep_gemm/tests/test_mega_moe_sm90_weight_transform.py @@ -0,0 +1,51 @@ +"""Regression test for SM90 MegaMoE FP8 weight-block sanitization.""" + +import argparse + +import torch + +import deep_gemm + + +def main() -> None: + parser = argparse.ArgumentParser() + # Keep the common test-runner interface; this test is intentionally single-rank. + parser.add_argument("--num-processes", type=int, default=1) + parser.parse_args() + + if not torch.cuda.is_available(): + print("SKIP: CUDA is not available") + return + if torch.cuda.get_device_capability()[0] != 9: + print("SKIP: requires an SM90 CUDA device") + return + + l1_weight = torch.full( + (1, 256, 128), 7, dtype=torch.float8_e4m3fn, device="cuda" + ) + l2_weight = torch.full( + (1, 128, 128), 7, dtype=torch.float8_e4m3fn, device="cuda" + ) + l1_scale = torch.full((1, 2, 1), 1.0e-20, dtype=torch.float32, device="cuda") + l2_scale = torch.full((1, 1, 1), 1.0e-20, dtype=torch.float32, device="cuda") + + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( + (l1_weight, l1_scale), (l2_weight, l2_scale) + ) + transformed_l1_weight, transformed_l1_scale = transformed_l1 + transformed_l2_weight, transformed_l2_scale = transformed_l2 + + assert torch.count_nonzero(transformed_l1_weight).item() == 0 + assert torch.count_nonzero(transformed_l2_weight).item() == 0 + assert torch.isfinite(transformed_l1_scale).all() + assert torch.isfinite(transformed_l2_scale).all() + assert torch.all(transformed_l1_scale >= 1.0e-12) + assert torch.all(transformed_l2_scale >= 1.0e-12) + # The input tensors are caller-owned and must remain unchanged. + assert torch.count_nonzero(l1_weight).item() == l1_weight.numel() + assert torch.count_nonzero(l2_weight).item() == l2_weight.numel() + print("PASS: tiny SM90 MegaMoE weight blocks were sanitized") + + +if __name__ == "__main__": + main()