From 89ba17ad056d36aa6e2e4f66176e873fcb12734a Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 11:42:51 +0800 Subject: [PATCH] feat(moe): add fp8 megamoe fused path and nvfp4 layout guards Wire MegaMoE FP8 through DeepGEMM's fused fp8_mega_moe path, preserve fallback runner layouts, and add explicit NVFP4 group-size guardrails for unsupported DeepGEMM scale transforms. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py python/sglang/srt/layers/quantization/fp8.py python/sglang/srt/layers/quantization/modelopt_quant.py test/registered/unit/moe/test_glm_megamoe.py Tested: git diff --check Not-tested: python3 -m unittest test/registered/unit/moe/test_glm_megamoe.py (local Python environments do not have torch installed). --- python/sglang/srt/layers/moe/mega_moe.py | 335 +++++++++++++++++- .../moe/moe_runner/flashinfer_trtllm.py | 13 + python/sglang/srt/layers/quantization/fp8.py | 43 ++- .../srt/layers/quantization/modelopt_quant.py | 11 + test/registered/unit/moe/test_glm_megamoe.py | 78 +++- 5 files changed, 465 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index c6730a161..db7ce6871 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -17,12 +17,14 @@ from __future__ import annotations import os from contextlib import nullcontext +from math import ceil from typing import TYPE_CHECKING, Optional import torch from sglang.srt.environ import envs from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.dp_attention import get_dp_global_num_tokens from sglang.srt.layers.moe.topk import select_experts from sglang.srt.layers.moe.utils import get_moe_a2a_backend @@ -57,6 +59,7 @@ def _get_mega_moe_symm_buffer( num_topk: int, hidden: int, intermediate_hidden: int, + mma_type: str = "fp8xfp4", ) -> SymmBuffer: import deep_gemm @@ -69,6 +72,7 @@ def _get_mega_moe_symm_buffer( num_topk, hidden, intermediate_hidden, + mma_type, ) buf = _MEGA_MOE_SYMM_BUFFER.get(key) if buf is None: @@ -80,6 +84,7 @@ def _get_mega_moe_symm_buffer( hidden, intermediate_hidden, use_fp8_dispatch=True, + mma_type=mma_type, activation="swiglu", ) _MEGA_MOE_SYMM_BUFFER[key] = buf @@ -91,8 +96,11 @@ def should_use_mega_moe(moe, hidden_states: torch.Tensor) -> bool: return False if not getattr(moe.experts, "_mega_moe_weights_built", False): return False - # This fork does not carry the DSV4 JIT FP8 pre-dispatch path yet. Keep the - # fast path gated to DeepGEMM's FP4 activation pre-dispatch API. + weight_format = getattr(moe.experts, "_mega_moe_weight_format", None) + if weight_format == "fp8": + return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + if weight_format != "fp4": + return False if not envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get(): return False if get_is_capture_mode(): @@ -209,6 +217,18 @@ def _run_mega_routed( input_ids_global: Optional[torch.Tensor], num_tokens: int, ) -> torch.Tensor: + weight_format = getattr(moe.experts, "_mega_moe_weight_format", None) + if weight_format == "fp8": + return _run_mega_fp8_routed( + moe, hidden_states, forward_batch, input_ids_global, num_tokens + ) + if weight_format != "fp4": + raise RuntimeError( + f"Unsupported MegaMoE weight format: {weight_format!r}. " + "Build weights with build_mega_moe_experts_weights() or " + "build_mega_moe_fp8_experts_weights()." + ) + import deep_gemm from sglang.srt.distributed.parallel_state import get_moe_ep_group @@ -258,6 +278,7 @@ def _run_mega_routed( num_topk=top_k, hidden=hidden_size, intermediate_hidden=intermediate_size, + mma_type="fp8xfp4", ) if num_tokens > 0: @@ -304,6 +325,101 @@ def _run_mega_routed( return y +def _run_mega_fp8_routed( + moe, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch], + input_ids_global: Optional[torch.Tensor], + num_tokens: int, +) -> torch.Tensor: + import deep_gemm + + from sglang.srt.distributed.parallel_state import get_moe_ep_group + + if num_tokens > 0: + router_logits = _call_gate(moe, hidden_states, forward_batch) + topk_output = _select_mega_topk( + moe, hidden_states, router_logits, forward_batch, input_ids_global + ) + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + else: + topk_ids = None + topk_weights = None + + experts = moe.experts + hidden_size = getattr(experts, "_mega_moe_hidden_size", hidden_states.shape[-1]) + ep_group = get_moe_ep_group().device_group + num_experts = experts.num_experts + top_k = _get_mega_top_k(moe) + intermediate_size = getattr( + experts, + "_mega_moe_intermediate_size", + getattr(moe.config, "moe_intermediate_size"), + ) + num_max_tokens_per_rank = ( + envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() + ) + assert num_tokens <= num_max_tokens_per_rank, ( + f"mega MoE: num_tokens={num_tokens} exceeds cap " + f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" + f"{num_max_tokens_per_rank}; raise the env var or shrink " + f"cuda_graph_max_bs / chunked_prefill_size accordingly" + ) + + buf = _get_mega_moe_symm_buffer( + ep_group, + num_experts=num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_topk=top_k, + hidden=hidden_size, + intermediate_hidden=intermediate_size, + mma_type="fp8xfp8", + ) + + if num_tokens > 0: + topk_ids_in = topk_ids.to(torch.int32) + topk_weights_in = topk_weights.to(torch.float32) + else: + topk_ids_in = hidden_states.new_empty((0, top_k), dtype=torch.int32) + topk_weights_in = hidden_states.new_empty((0, top_k), dtype=torch.float32) + + deep_gemm.mega_moe_pre_dispatch( + hidden_states, + topk_ids_in, + topk_weights_in, + buf.x, + buf.x_sf, + buf.topk_idx, + buf.topk_weights, + num_tokens=num_tokens, + group_size=128, + use_fp4_acts=False, + ) + + y = torch.empty( + (max(num_tokens, 1), hidden_size), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + swiglu_limit = getattr(moe.config, "swiglu_limit", None) + deep_gemm.fp8_mega_moe( + y, + experts.mega_l1_weights, + experts.mega_l2_weights, + buf, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=swiglu_limit, + fast_math=True, + ) + y = y[:num_tokens] + + if not getattr(experts, "should_fuse_routed_scaling_factor_in_topk", False): + y.mul_(getattr(moe, "routed_scaling_factor", 1.0)) + return y + + def _as_deep_gemm_packed_fp4_weights( weights: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: @@ -371,6 +487,37 @@ def _infer_mega_moe_scale_group_size( return k // num_scale_cols +def _transform_mega_moe_sf( + transform_sf_into_required_layout, + sf: torch.Tensor, + *, + mn: int, + k: int, + group_size: int, + num_groups: int, + name: str, +) -> torch.Tensor: + try: + return transform_sf_into_required_layout( + sf.to(torch.float32), + mn=mn, + k=k, + recipe=(1, group_size), + num_groups=num_groups, + disable_ue8m0_cast=False, + ) + except RuntimeError as exc: + if "Unknown SF transformation" not in str(exc): + raise + raise RuntimeError( + f"DeepGEMM does not support the MegaMoE {name} scale layout " + f"recipe=(1,{group_size}) for mn={mn}, k={k}, " + f"num_groups={num_groups}. GLM NVFP4 checkpoints use group16 " + "scales, so this requires an sgl-deep-gemm/DeepGEMM build with " + "group16 MegaMoE SF transformation support." + ) from exc + + def build_mega_moe_experts_weights( experts, *, @@ -397,21 +544,23 @@ def build_mega_moe_experts_weights( f"{group_k1=} and {group_k2=}." ) - w13_sf = transform_sf_into_required_layout( - w13_sf_src.to(torch.float32), + w13_sf = _transform_mega_moe_sf( + transform_sf_into_required_layout, + w13_sf_src, mn=n1, k=k1, - recipe=(1, group_k1), + group_size=group_k1, num_groups=num_groups, - disable_ue8m0_cast=False, + name="w13", ) - w2_sf = transform_sf_into_required_layout( - w2_sf_src.to(torch.float32), + w2_sf = _transform_mega_moe_sf( + transform_sf_into_required_layout, + w2_sf_src, mn=n2, k=k2, - recipe=(1, group_k2), + group_size=group_k2, num_groups=num_groups, - disable_ue8m0_cast=False, + name="w2", ) if preserve_runner_layout: @@ -455,4 +604,170 @@ def build_mega_moe_experts_weights( experts._mega_moe_hidden_size = k1 experts._mega_moe_intermediate_size = k2 experts._mega_moe_weight_group_size = group_k1 + experts._mega_moe_weight_format = "fp4" + experts._mega_moe_weights_built = True + + +def _swap_gate_up_halves(weight: torch.Tensor) -> torch.Tensor: + half = weight.shape[1] // 2 + return torch.cat((weight[:, half:], weight[:, :half]), dim=1).contiguous() + + +def _normalize_mega_moe_fp8_block_shape(block_shape) -> tuple[int, int]: + if block_shape is None: + return 128, 128 + block_mn, block_k = int(block_shape[0]), int(block_shape[1]) + if (block_mn, block_k) != (128, 128): + raise ValueError( + "MegaMoE FP8 fused path currently requires 128x128 weight " + f"block scales, got {block_shape}." + ) + return block_mn, block_k + + +def _expand_mega_moe_fp8_scale( + scale: torch.Tensor, + *, + mn: int, + k: int, + block_shape: tuple[int, int], + num_groups: int, + name: str, +) -> torch.Tensor: + block_mn, block_k = block_shape + expected_shape = (num_groups, ceil(mn / block_mn), ceil(k / block_k)) + if scale.dim() == 1: + if scale.shape[0] != num_groups: + raise ValueError( + f"MegaMoE FP8 {name} per-expert scale must have " + f"{num_groups} entries, got shape={tuple(scale.shape)}." + ) + return scale.view(num_groups, 1, 1).expand(expected_shape).contiguous() + if tuple(scale.shape) != expected_shape: + raise ValueError( + f"MegaMoE FP8 {name} scale shape must be {expected_shape}, " + f"got {tuple(scale.shape)}." + ) + return scale.contiguous() + + +def _transform_mega_moe_fp8_sf( + transform_sf_into_required_layout, + sf: torch.Tensor, + *, + mn: int, + k: int, + block_shape: tuple[int, int], + num_groups: int, + name: str, +) -> torch.Tensor: + block_mn, block_k = block_shape + sf = _expand_mega_moe_fp8_scale( + sf, + mn=mn, + k=k, + block_shape=block_shape, + num_groups=num_groups, + name=name, + ) + try: + return transform_sf_into_required_layout( + sf.to(torch.float32), + mn=mn, + k=k, + recipe=(block_mn, block_k), + num_groups=num_groups, + disable_ue8m0_cast=False, + ) + except RuntimeError as exc: + if "Unknown SF transformation" not in str(exc): + raise + raise RuntimeError( + f"DeepGEMM does not support the MegaMoE FP8 {name} scale layout " + f"recipe=({block_mn},{block_k}) for mn={mn}, k={k}, " + f"num_groups={num_groups}." + ) from exc + + +def build_mega_moe_fp8_experts_weights( + experts, + *, + preserve_runner_layout: bool = True, + swap_w13_halves: bool = False, +) -> None: + from deep_gemm import ( + transform_sf_into_required_layout, + transform_weights_for_mega_moe, + ) + + if getattr(experts, "_mega_moe_weights_built", False): + return + + w13_weight = experts.w13_weight.data + w2_weight = experts.w2_weight.data + if swap_w13_halves: + w13_weight = _swap_gate_up_halves(w13_weight) + elif preserve_runner_layout: + w13_weight = w13_weight.clone() + if preserve_runner_layout: + w2_weight = w2_weight.clone() + + block_quant = hasattr(experts, "w13_weight_scale_inv") and hasattr( + experts, "w2_weight_scale_inv" + ) + if block_quant: + w13_scale = experts.w13_weight_scale_inv.data + w2_scale = experts.w2_weight_scale_inv.data + if preserve_runner_layout: + w13_scale = w13_scale.clone() + w2_scale = w2_scale.clone() + quant_config = getattr( + getattr(experts, "quant_method", None), "quant_config", None + ) + block_shape = getattr(quant_config, "weight_block_size", None) + else: + w13_scale = experts.w13_weight_scale.data + w2_scale = experts.w2_weight_scale.data + if preserve_runner_layout: + w13_scale = w13_scale.clone() + w2_scale = w2_scale.clone() + block_shape = [128, 128] + + block_shape = _normalize_mega_moe_fp8_block_shape(block_shape) + num_groups, n1, k1 = w13_weight.shape + _, n2, k2 = w2_weight.shape + w13_sf = _transform_mega_moe_fp8_sf( + transform_sf_into_required_layout, + w13_scale, + mn=n1, + k=k1, + block_shape=block_shape, + num_groups=num_groups, + name="w13", + ) + w2_sf = _transform_mega_moe_fp8_sf( + transform_sf_into_required_layout, + w2_scale, + mn=n2, + k=k2, + block_shape=block_shape, + num_groups=num_groups, + name="w2", + ) + l1_pair, l2_pair = transform_weights_for_mega_moe( + (w13_weight, w13_sf), (w2_weight, w2_sf) + ) + + experts.mega_fp8_w13_weight = w13_weight + experts.mega_fp8_w2_weight = w2_weight + experts.mega_fp8_w13_scale = w13_scale + experts.mega_fp8_w2_scale = w2_scale + experts.mega_l1_weights = l1_pair + experts.mega_l2_weights = l2_pair + experts._mega_moe_fp8_block_quant = block_quant + experts._mega_moe_fp8_block_shape = block_shape + experts._mega_moe_preserve_runner_layout = preserve_runner_layout + experts._mega_moe_hidden_size = w13_weight.shape[2] + experts._mega_moe_intermediate_size = w2_weight.shape[2] + experts._mega_moe_weight_format = "fp8" experts._mega_moe_weights_built = True diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index c8a88f8fb..7d5b076e1 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -1258,6 +1258,19 @@ def fused_experts_none_to_flashinfer_trtllm( ) +@register_fused_func("megamoe", "flashinfer_trtllm") +def fused_experts_megamoe_to_flashinfer_trtllm( + dispatch_output: StandardDispatchOutput, + quant_info: MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + # MegaMoE uses the standard dispatcher in this fork. Keep FlashInfer TRTLLM + # available as the fallback path when the MegaMoE fast path is not selected. + return fused_experts_none_to_flashinfer_trtllm( + dispatch_output, quant_info, runner_config + ) + + @register_fused_func("none", "flashinfer_trtllm_routed") def fused_experts_none_to_flashinfer_trtllm_routed( dispatch_output: StandardDispatchOutput, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 987162922..e0e158c38 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -27,7 +27,11 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmFp8MoeQuantInfo, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + get_moe_a2a_backend, + get_moe_runner_backend, +) from sglang.srt.layers.parameter import ( BlockQuantScaleParameter, ModelWeightParameter, @@ -757,6 +761,8 @@ class Fp8MoEMethod(FusedMoEMethodBase): moe_runner_backend = get_moe_runner_backend() if moe_runner_backend.is_deep_gemm(): return True + if get_moe_a2a_backend().is_megamoe(): + return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM if moe_runner_backend.is_auto(): return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and ( get_moe_a2a_backend().is_deepep() @@ -1038,9 +1044,13 @@ class Fp8MoEMethod(FusedMoEMethodBase): and will_use_deepgemm and not layer.w13_weight_scale_inv.format_ue8m0 ): - assert isinstance( - layer, DeepEPMoE - ), "DeepGemm MoE is only supported with DeepEPMoE" + assert ( + isinstance(layer, DeepEPMoE) + or get_moe_a2a_backend().is_megamoe() + ), ( + "DeepGemm MoE is only supported with DeepEPMoE or the " + "MegaMoE fast path" + ) weight_block_size = self.quant_config.weight_block_size requant_weight_ue8m0_inplace( layer.w13_weight, layer.w13_weight_scale_inv, weight_block_size @@ -1237,6 +1247,14 @@ class Fp8MoEMethod(FusedMoEMethodBase): # Block quant doesn't need to process weights after loading if self.block_quant: self.process_weights_after_loading_block_quant(layer) + if get_moe_a2a_backend().is_megamoe(): + from sglang.srt.layers.moe.mega_moe import ( + build_mega_moe_fp8_experts_weights, + ) + + build_mega_moe_fp8_experts_weights( + layer, preserve_runner_layout=True + ) return # If checkpoint is fp16 or bfloat16, quantize in place. @@ -1267,6 +1285,14 @@ class Fp8MoEMethod(FusedMoEMethodBase): if _is_hip: self.process_weights_hip_scale_padding(layer) + if get_moe_a2a_backend().is_megamoe(): + from sglang.srt.layers.moe.mega_moe import ( + build_mega_moe_fp8_experts_weights, + ) + + build_mega_moe_fp8_experts_weights( + layer, preserve_runner_layout=True + ) return # If checkpoint is fp8, we need to handle that the @@ -1351,6 +1377,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): if _is_hip: self.process_weights_hip_scale_padding(layer) + if get_moe_a2a_backend().is_megamoe(): + from sglang.srt.layers.moe.mega_moe import ( + build_mega_moe_fp8_experts_weights, + ) + + build_mega_moe_fp8_experts_weights( + layer, preserve_runner_layout=True + ) + # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled if get_moe_runner_backend().is_flashinfer_trtllm(): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index bf0636c03..216a847c8 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -958,6 +958,17 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): layer.w2_input_scale.max(), requires_grad=False ) + if get_moe_a2a_backend().is_megamoe(): + from sglang.srt.layers.moe.mega_moe import ( + build_mega_moe_fp8_experts_weights, + ) + + # ModelOpt FP8 stores W13 in [Up, Gate] order. Preserve a DeepGEMM + # compatible [Gate, Up] copy before FlashInfer mutates runner layout. + build_mega_moe_fp8_experts_weights( + layer, preserve_runner_layout=True, swap_w13_halves=True + ) + # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled if get_moe_runner_backend().is_flashinfer_trtllm(): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( diff --git a/test/registered/unit/moe/test_glm_megamoe.py b/test/registered/unit/moe/test_glm_megamoe.py index 5df6c0609..73a7a738d 100644 --- a/test/registered/unit/moe/test_glm_megamoe.py +++ b/test/registered/unit/moe/test_glm_megamoe.py @@ -33,7 +33,10 @@ class TestGLMMegaMoE(unittest.TestCase): def test_should_use_mega_moe_respects_env_and_token_cap(self): hidden_states = torch.empty((2, 32)) moe = SimpleNamespace( - experts=SimpleNamespace(_mega_moe_weights_built=True), + experts=SimpleNamespace( + _mega_moe_weights_built=True, + _mega_moe_weight_format="fp4", + ), ) with patch.object( @@ -67,6 +70,22 @@ class TestGLMMegaMoE(unittest.TestCase): ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(2): self.assertTrue(mega_moe.should_use_mega_moe(moe, hidden_states)) + def test_should_use_mega_moe_fp8_does_not_require_fp4_act_env(self): + hidden_states = torch.empty((2, 32)) + moe = SimpleNamespace( + experts=SimpleNamespace( + _mega_moe_weights_built=True, + _mega_moe_weight_format="fp8", + ), + ) + + with patch.object( + mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend() + ), patch.object( + mega_moe.deep_gemm_wrapper, "ENABLE_JIT_DEEPGEMM", True + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(False): + self.assertTrue(mega_moe.should_use_mega_moe(moe, hidden_states)) + def test_glm_forward_uses_megamoe_fast_path(self): block = object.__new__(glm4_moe.Glm4MoeSparseMoeBlock) hidden_states = torch.empty((1, 32)) @@ -125,6 +144,7 @@ class TestGLMMegaMoE(unittest.TestCase): self.assertTrue(experts._mega_moe_weights_built) self.assertTrue(experts._mega_moe_preserve_runner_layout) self.assertEqual(experts._mega_moe_weight_group_size, 16) + self.assertEqual(experts._mega_moe_weight_format, "fp4") self.assertTrue(torch.equal(experts.w13_weight.data, original_w13)) self.assertNotEqual( experts.mega_l1_weights[0].data_ptr(), @@ -155,6 +175,62 @@ class TestGLMMegaMoE(unittest.TestCase): self.assertIs(same_data, int8_data) self.assertIs(same_scale, scale) + def test_build_fp8_mega_moe_weights_preserves_runner_layout(self): + experts = SimpleNamespace( + w13_weight=_Param(torch.arange(2 * 8 * 4).reshape(2, 8, 4)), + w2_weight=_Param(torch.arange(2 * 4 * 4).reshape(2, 4, 4)), + w13_weight_scale=_Param(torch.arange(2, dtype=torch.float32) + 1), + w2_weight_scale=_Param(torch.arange(2, dtype=torch.float32) + 3), + quant_method=SimpleNamespace( + quant_config=SimpleNamespace(weight_block_size=None) + ), + ) + original_w13 = experts.w13_weight.data.clone() + fake_deep_gemm = types.SimpleNamespace( + transform_sf_into_required_layout=( + lambda sf, *, mn, k, recipe, num_groups, disable_ue8m0_cast: torch.zeros( + (num_groups, mn, max(k // recipe[1], 1)), dtype=torch.int32 + ) + ), + transform_weights_for_mega_moe=lambda l1, l2: (l1, l2), + ) + + with patch.dict(sys.modules, {"deep_gemm": fake_deep_gemm}): + mega_moe.build_mega_moe_fp8_experts_weights( + experts, preserve_runner_layout=True, swap_w13_halves=True + ) + + self.assertTrue(experts._mega_moe_weights_built) + self.assertEqual(experts._mega_moe_weight_format, "fp8") + self.assertFalse(experts._mega_moe_fp8_block_quant) + self.assertTrue(torch.equal(experts.w13_weight.data, original_w13)) + self.assertTrue( + torch.equal( + experts.mega_fp8_w13_weight, + torch.cat((original_w13[:, 4:], original_w13[:, :4]), dim=1), + ) + ) + self.assertNotEqual( + experts.mega_fp8_w13_weight.data_ptr(), + experts.w13_weight.data.data_ptr(), + ) + + def test_transform_sf_reports_group16_dependency_gap(self): + def fake_transform_sf(*args, **kwargs): + del args, kwargs + raise RuntimeError("Unknown SF transformation") + + with self.assertRaisesRegex(RuntimeError, "group16"): + mega_moe._transform_mega_moe_sf( + fake_transform_sf, + torch.ones((2, 256, 2), dtype=torch.float32), + mn=256, + k=32, + group_size=16, + num_groups=2, + name="w13", + ) + if __name__ == "__main__": unittest.main()