From 5fe97637df12aee6794415848787f77f697b8696 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Sun, 2 Aug 2026 21:48:13 -0700 Subject: [PATCH] Support DeepGEMM for standard MoE dispatch (#33128) Co-authored-by: Sam Li --- .../sglang/kernels/ops/moe/ep_moe_kernels.py | 19 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 185 ++++++++++++-- python/sglang/srt/layers/quantization/fp8.py | 20 +- .../ops/moe/test_minimax_quant_scatter.py | 241 +++++++++++++++++- .../test_deepgemm_ue8m0_requant.py | 57 ++++- 5 files changed, 485 insertions(+), 37 deletions(-) diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 252d2fad9..0cebb241b 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1019,6 +1019,7 @@ def post_reorder_triton_kernel( @triton.jit def _fwd_kernel_ep_scatter_1( num_recv_tokens_per_expert, + num_valid_tokens_per_expert, expert_start_loc, m_indices, num_experts: tl.constexpr, @@ -1037,15 +1038,17 @@ def _fwd_kernel_ep_scatter_1( tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts) cur_expert_start = tl.load(expert_start_loc + cur_expert) - cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + cur_expert_padded_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + cur_expert_valid_token_num = tl.load(num_valid_tokens_per_expert + cur_expert) m_indices_start_ptr = m_indices + cur_expert_start off_expert = tl.arange(0, BLOCK_E) - for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4): + for start_m in tl.range(0, cur_expert_padded_token_num, BLOCK_E, num_stages=4): + offsets = start_m + off_expert tl.store( - m_indices_start_ptr + start_m + off_expert, - cur_expert, + m_indices_start_ptr + offsets, + tl.where(offsets < cur_expert_valid_token_num, cur_expert, -1), ) @@ -1137,6 +1140,7 @@ def ep_scatter( recv_x_scale: torch.Tensor, recv_topk: torch.Tensor, num_recv_tokens_per_expert: torch.Tensor, + num_valid_tokens_per_expert: torch.Tensor, expert_start_loc: torch.Tensor, output_tensor: torch.Tensor, output_tensor_scale: torch.Tensor, @@ -1172,6 +1176,7 @@ def ep_scatter( _fwd_kernel_ep_scatter_1[(grid,)]( num_recv_tokens_per_expert, + num_valid_tokens_per_expert, expert_start_loc, m_indices, num_experts=num_experts, @@ -1575,7 +1580,11 @@ def moe_ep_deepgemm_preprocess( assert len(block_shape) == 2 block_n, block_k = block_shape[0], block_shape[1] is_fp8 = output_dtype == torch.float8_e4m3fn - if is_fp8 and use_mxfp8: + # Quantize FP8 values with the UE8M0 scale directly. Rounding only the + # scale afterward can change the represented activation by up to 2x. + from sglang.srt.layers import deep_gemm_wrapper + + if is_fp8 and (use_mxfp8 or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0): from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import ( per_token_quant_fp8_ue8m0_scatter, ) diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index e773b6b37..52658f656 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -98,6 +98,24 @@ def copy_list_to_gpu_no_ce(arr: List[int]): return tensor_gpu +def _should_use_masked_standard_layout(runner_config: MoeRunnerConfig) -> bool: + """Use masked GEMM when expert parallelism keeps its buffer small.""" + return ( + runner_config.num_experts > runner_config.num_local_experts + and runner_config.num_local_experts <= 32 + ) + + +def _get_compact_all_tokens( + num_assignments: int, num_experts: int, block_e: int = 128 +) -> int: + """Return the maximum padded rows over all routings of the assignments.""" + max_nonempty_experts = min(num_assignments, num_experts) + return block_e * ( + max_nonempty_experts + (num_assignments - max_nonempty_experts) // block_e + ) + + @dataclass class DeepGemmRunnerInput(RunnerInput): hidden_states: torch.Tensor @@ -670,7 +688,11 @@ def pre_permute_standard_to_deep_gemm( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepGemmRunnerInput: - from sglang.kernels.ops.moe.ep_moe_kernels import moe_ep_deepgemm_preprocess + from sglang.kernels.ops.moe.ep_moe_kernels import ( + ep_scatter, + fused_moe_dispatch_index, + moe_ep_deepgemm_preprocess, + ) hidden_states, topk_output = ( dispatch_output.hidden_states, @@ -685,25 +707,151 @@ def pre_permute_standard_to_deep_gemm( topk_weights, topk_ids = topk_weights, topk_ids - # PreReorder + if _should_use_masked_standard_layout(runner_config): + output_dtype = ( + torch.bfloat16 + if quant_info.w13_weight.dtype == torch.bfloat16 + else torch.float8_e4m3fn + ) + masked_m, _, src2dst, hidden_states, hidden_states_scale = ( + moe_ep_deepgemm_preprocess( + topk_ids, + runner_config.num_local_experts, + hidden_states, + runner_config.top_k, + quant_info.block_shape, + output_dtype=output_dtype, + use_mxfp8=quant_info.use_mxfp8, + ) + ) + # Use the global expert count because expected_m is a tuning hint, not + # the per-rank buffer capacity. + expected_m = max( + 1, + ceil_div( + hidden_states_shape[0] * runner_config.top_k, + runner_config.num_experts, + ), + ) + + if runner_config.inplace: + dispose_tensor(hidden_states_ref) + + running_state["topk_ids"] = topk_ids + running_state["topk_weights"] = topk_weights + running_state["hidden_states_shape"] = hidden_states_shape + running_state["hidden_states_dtype"] = hidden_states_dtype + running_state["hidden_states_device"] = hidden_states_device + running_state["src2dst"] = src2dst + running_state["mxfp8_act_gran_k"] = ( + quant_info.block_shape[1] if quant_info.block_shape else 128 + ) + + return DeepGemmRunnerInput( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + use_masked_gemm=True, + masked_m=masked_m, + expected_m=expected_m, + ) + + # The compact layout avoids scaling masked buffers with the expert count. + # Scatter and post-permute skip non-local experts mapped to -1. + block_e = 128 + num_experts = runner_config.num_local_experts + num_assignments = topk_ids.numel() + all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e) + + tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index( + topk_ids, num_experts, 1 + ) + dispose_tensor(unused_masked_dst) + valid_tokens_per_expert = tokens_per_expert + tokens_per_expert = (ceil_div(tokens_per_expert, block_e) * block_e).to(torch.int32) + # Keep graph-static shapes by appending padding to the final segment. + # Its m_indices stay -1, so DeepGEMM skips those rows. + tokens_per_expert[-1].add_(all_tokens - tokens_per_expert.sum()) + + k = hidden_states.size(1) output_dtype = ( torch.bfloat16 if quant_info.w13_weight.dtype == torch.bfloat16 else torch.float8_e4m3fn ) - masked_m, expected_m, src2dst, hidden_states, hidden_states_scale = ( - moe_ep_deepgemm_preprocess( - topk_ids, - runner_config.num_local_experts, - hidden_states, - runner_config.top_k, - quant_info.block_shape, - output_dtype=output_dtype, - use_mxfp8=quant_info.use_mxfp8, + if output_dtype == torch.bfloat16: + packed_input_source = hidden_states + packed_input_source_scale = None + packed_input = torch.empty( + (all_tokens, k), device=hidden_states_device, dtype=torch.bfloat16 + ) + # ep_scatter ignores scales for BF16, but a real tensor keeps its + # Triton signature uniform across the existing DeepEP caller. + packed_input_scale = torch.empty( + (all_tokens, 1), device=hidden_states_device, dtype=torch.float32 + ) + else: + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, ) - ) - dispose_tensor(hidden_states_ref) + block_k = quant_info.block_shape[1] if quant_info.block_shape else 128 + packed_input_source, packed_input_source_scale = ( + sglang_per_token_group_quant_fp8( + hidden_states, + block_k, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + ) + packed_input = torch.zeros( + (all_tokens, k), + device=hidden_states_device, + dtype=torch.float8_e4m3fn, + ) + scale_width = k // block_k + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + scale_width = ceil_div(scale_width, 4) + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + packed_input_scale = torch.zeros( + (scale_width, all_tokens), + device=hidden_states_device, + dtype=packed_input_source_scale.dtype, + ).transpose(0, 1) + else: + packed_input_scale = torch.zeros( + (all_tokens, scale_width), + device=hidden_states_device, + dtype=packed_input_source_scale.dtype, + ) + + expert_start_loc = torch.empty( + num_experts, device=hidden_states_device, dtype=torch.int32 + ) + m_indices = torch.empty(all_tokens, device=hidden_states_device, dtype=torch.int32) + src2dst = torch.empty_like(topk_ids, dtype=torch.int32) + ep_scatter( + packed_input_source, + packed_input_source_scale, + topk_ids, + tokens_per_expert, + valid_tokens_per_expert, + expert_start_loc, + packed_input, + packed_input_scale, + m_indices, + src2dst, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128), + ) + if packed_input_source is not hidden_states: + dispose_tensor(packed_input_source) + if packed_input_source_scale is not None: + dispose_tensor(packed_input_source_scale) + + # Preserve the input when a shared expert or its gate may still use it. + if runner_config.inplace: + dispose_tensor(hidden_states_ref) running_state["topk_ids"] = topk_ids running_state["topk_weights"] = topk_weights @@ -711,16 +859,16 @@ def pre_permute_standard_to_deep_gemm( running_state["hidden_states_dtype"] = hidden_states_dtype running_state["hidden_states_device"] = hidden_states_device running_state["src2dst"] = src2dst + running_state["all_tokens"] = all_tokens running_state["mxfp8_act_gran_k"] = ( quant_info.block_shape[1] if quant_info.block_shape else 128 ) return DeepGemmRunnerInput( - hidden_states=hidden_states, - hidden_states_scale=hidden_states_scale, - use_masked_gemm=True, - masked_m=masked_m, - expected_m=expected_m, + hidden_states=packed_input, + hidden_states_scale=packed_input_scale, + use_masked_gemm=False, + m_indices=m_indices, ) @@ -891,6 +1039,7 @@ def pre_permute_deepep_normal_to_deep_gemm( hidden_states_scale, topk_ids, num_recv_tokens_per_expert_gpu, + num_recv_tokens_per_expert_gpu, expert_start_loc, input_tensor, input_tensor_scale, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 7fa220b32..3e142d29d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1603,7 +1603,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): else: # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0 from sglang.srt.layers import deep_gemm_wrapper - from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE # Check if MoE will actually use DeepGEMM runner will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled() @@ -1677,20 +1676,17 @@ class Fp8MoEMethod(FusedMoEMethodBase): if not self.is_fp4_expert: weight_block_size = self.quant_config.weight_block_size - if requant_block_scale_ue8m0_for_deepgemm( - layer.w13_weight, - layer.w13_weight_scale_inv, - weight_block_size, - use_deepgemm_runner=will_use_deepgemm, + for weight, weight_scale in ( + (layer.w13_weight, layer.w13_weight_scale_inv), + (layer.w2_weight, layer.w2_weight_scale_inv), ): - assert isinstance( - layer, DeepEPMoE - ), "DeepGemm MoE is only supported with DeepEPMoE" requant_block_scale_ue8m0_for_deepgemm( - layer.w2_weight, - layer.w2_weight_scale_inv, + weight, + weight_scale, weight_block_size, - use_deepgemm_runner=True, + use_deepgemm_runner=will_use_deepgemm, + output_dtype=torch.bfloat16, + weight_shape=weight.shape[-2:], ) def _convert_mxfp8_moe_to_block_fp8(self, layer: Module) -> None: diff --git a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py index aaead0331..3b9821835 100644 --- a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py +++ b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py @@ -1,14 +1,31 @@ import random import sys +from contextlib import nullcontext import pytest import torch -from sglang.kernels.ops.moe.ep_moe_kernels import fill_gateup_input_triton_kernel +import sglang.srt.layers.moe.moe_runner.deep_gemm as deep_gemm_runner +from sglang.kernels.ops.moe.ep_moe_kernels import ( + fill_gateup_input_triton_kernel, + moe_ep_deepgemm_preprocess, +) from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import ( per_token_quant_fp8_ue8m0, per_token_quant_fp8_ue8m0_scatter, ) +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.deep_gemm import ( + DeepGemmMoeQuantInfo, + DeepGemmRunnerCore, + post_permute_deep_gemm_to_standard, + pre_permute_standard_to_deep_gemm, +) +from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput +from sglang.srt.layers.quantization.fp8_utils import ( + quant_weight_ue8m0, + transform_scale_ue8m0, +) from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") @@ -80,5 +97,227 @@ def test_quant_scatter_matches_quant_plus_fill(num_tokens, topk, hidden, group): ), f"scale mismatch token={t} slot={j} expert={e}" +def test_standard_deepgemm_preprocess_quantizes_with_ue8m0_scale(): + arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device()) + if arch_major <= 9: + pytest.skip("UE8M0 fusion is Blackwell-only") + + num_tokens, topk, hidden, group, num_experts = 7, 4, 2048, 128, 8 + torch.manual_seed(1234) + x = torch.randn(num_tokens, hidden, device=dev, dtype=torch.bfloat16) + topk_ids = torch.stack( + [ + (torch.arange(topk, device=dev, dtype=torch.int32) + token) % num_experts + for token in range(num_tokens) + ] + ) + + _, _, src2dst, grouped_x, grouped_scale = moe_ep_deepgemm_preprocess( + topk_ids=topk_ids, + num_local_experts=num_experts, + hidden_states=x, + top_k=topk, + block_shape=[group, group], + output_dtype=torch.float8_e4m3fn, + use_mxfp8=False, + ) + direct_x, direct_scale = per_token_quant_fp8_ue8m0(x, group) + + assert grouped_scale.dtype == torch.int32 + for token in range(num_tokens): + for slot in range(topk): + dst = int(src2dst[token * topk + slot]) + expert, row = divmod(dst, grouped_x.shape[1]) + assert torch.equal( + grouped_x[expert, row].view(torch.uint8), + direct_x[token].view(torch.uint8), + ) + assert torch.equal( + grouped_scale[expert, row], + direct_scale[token], + ) + + +@pytest.mark.parametrize( + "num_assignments,num_experts,expected", + [ + (14, 2, 256), + (10, 512, 1280), + (20, 512, 2560), + (320, 512, 40960), + (640, 512, 65664), + (1280, 512, 66304), + ], +) +def test_compact_all_tokens_uses_tight_routing_independent_bound( + num_assignments, num_experts, expected +): + assert ( + deep_gemm_runner._get_compact_all_tokens(num_assignments, num_experts) + == expected + ) + + +@pytest.mark.parametrize("weight_dtype", ["fp8", "bf16"]) +def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_dtype): + """Exercise both production grouped GEMMs through the standard path.""" + arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device()) + if arch_major <= 9: + pytest.skip("DeepGEMM UE8M0 is Blackwell-only") + + # This kernel test runs outside a model-parallel process. Bypass only the + # symmetric-allocation context; all pre-permute, DeepGEMM, activation, + # quantization, down-GEMM, and post-permute kernels remain real. + monkeypatch.setattr(deep_gemm_runner, "get_tp_group", lambda: None) + monkeypatch.setattr( + deep_gemm_runner, + "use_symmetric_memory", + lambda *args, **kwargs: nullcontext(), + ) + + # UE8M0 packs four 128-wide scale groups into each int32. Use the smallest + # legal K for both the gate/up and down GEMMs. + num_tokens, hidden, intermediate, topk, num_local_experts = 7, 512, 512, 2, 2 + torch.manual_seed(20260730) + hidden_states = torch.randn(num_tokens, hidden, device=dev, dtype=torch.bfloat16) + topk_ids = torch.tensor( + [ + [0, -1], + [1, -1], + [0, 1], + [-1, 1], + [0, -1], + [1, 0], + [-1, 1], + ], + device=dev, + dtype=torch.int32, + ) + topk_weights = torch.tensor( + [ + [0.8, 0.2], + [0.7, 0.3], + [0.6, 0.4], + [0.1, 0.9], + [0.75, 0.25], + [0.55, 0.45], + [0.35, 0.65], + ], + device=dev, + dtype=torch.float32, + ) + + weight_std = hidden**-0.5 + w13_bf16 = ( + torch.randn( + num_local_experts, + 2 * intermediate, + hidden, + device=dev, + dtype=torch.bfloat16, + ) + * weight_std + ) + w2_bf16 = ( + torch.randn( + num_local_experts, + hidden, + intermediate, + device=dev, + dtype=torch.bfloat16, + ) + * weight_std + ) + if weight_dtype == "fp8": + w13, w13_scale = quant_weight_ue8m0(w13_bf16, [128, 128]) + w2, w2_scale = quant_weight_ue8m0(w2_bf16, [128, 128]) + quant_info = DeepGemmMoeQuantInfo( + w13_weight=w13, + w2_weight=w2, + use_fp8=True, + w13_scale=transform_scale_ue8m0(w13_scale, mn=w13.shape[-2]), + w2_scale=transform_scale_ue8m0(w2_scale, mn=w2.shape[-2]), + block_shape=[128, 128], + ) + else: + quant_info = DeepGemmMoeQuantInfo( + w13_weight=w13_bf16, + w2_weight=w2_bf16, + use_fp8=False, + ) + + dispatch_output = StandardDispatchOutput( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_output=(topk_weights, topk_ids, None), + ) + + def run_with_num_experts(num_experts): + config = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden, + intermediate_size_per_partition=intermediate, + top_k=topk, + activation="silu", + is_gated=True, + inplace=False, + ) + running_state = {} + runner_input = pre_permute_standard_to_deep_gemm( + dispatch_output, + quant_info, + config, + running_state, + ) + runner_output = DeepGemmRunnerCore(config).run( + runner_input, + quant_info, + running_state, + ) + return ( + runner_input.use_masked_gemm, + running_state.get("all_tokens"), + runner_input.m_indices, + post_permute_deep_gemm_to_standard( + runner_output, + quant_info, + config, + running_state, + ).hidden_states, + ) + + compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = ( + run_with_num_experts(num_local_experts) + ) + masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = ( + run_with_num_experts(8) + ) + torch.cuda.synchronize() + + assert not compact_is_masked + assert masked_is_masked + assert compact_all_tokens == 256 + assert masked_all_tokens is None + assert masked_m_indices is None + valid_assignments = topk_ids[topk_ids >= 0] + assert torch.equal( + torch.bincount( + compact_m_indices[compact_m_indices >= 0], + minlength=num_local_experts, + ), + torch.bincount(valid_assignments, minlength=num_local_experts), + ) + assert (compact_m_indices == -1).sum() == compact_all_tokens - len( + valid_assignments + ) + torch.testing.assert_close( + masked_output, + compact_output, + rtol=5e-2, + atol=5e-2, + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-x"])) diff --git a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py index 412dbad5b..e96b68a84 100644 --- a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py +++ b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py @@ -5,13 +5,14 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") import unittest -from unittest.mock import patch +from unittest.mock import call, patch import torch from compressed_tensors.quantization import QuantizationStrategy import sglang.srt.layers.quantization.fp8_utils as fp8_utils from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.quantization import fp8 as fp8_quant from sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import ( CompressedTensorsW8A8Fp8, ) @@ -162,6 +163,60 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): self.assertTrue(layer.weight_scale.format_ue8m0) requant.assert_called_once() + def test_fp8_moe_requants_standard_layer_for_deepgemm(self): + method = fp8_quant.Fp8MoEMethod.__new__(fp8_quant.Fp8MoEMethod) + method.convert_mxfp8_to_block = False + method.use_mxfp8 = False + method.is_fp4_expert = False + method.dequant_fp4_to_fp8 = False + method.quant_config = unittest.mock.Mock(weight_block_size=BLOCK_SIZE) + + layer = torch.nn.Module() + layer.w13_weight, layer.w13_weight_scale_inv = _make_params() + layer.w2_weight, layer.w2_weight_scale_inv = _make_params() + + def _mark_ue8m0(weight, weight_scale, *args, **kwargs): + weight_scale.format_ue8m0 = True + return True + + with patch.multiple( + fp8_quant, + _is_cpu=False, + _is_fp8_fnuz=False, + _use_aiter=False, + ), patch.object( + method, "is_deepgemm_moe_runner_backend_enabled", return_value=True + ), patch.object( + fp8_quant, + "requant_block_scale_ue8m0_for_deepgemm", + side_effect=_mark_ue8m0, + ) as requant: + method.process_weights_after_loading_block_quant(layer) + + self.assertEqual( + requant.call_args_list, + [ + call( + layer.w13_weight, + layer.w13_weight_scale_inv, + BLOCK_SIZE, + use_deepgemm_runner=True, + output_dtype=torch.bfloat16, + weight_shape=layer.w13_weight.shape[-2:], + ), + call( + layer.w2_weight, + layer.w2_weight_scale_inv, + BLOCK_SIZE, + use_deepgemm_runner=True, + output_dtype=torch.bfloat16, + weight_shape=layer.w2_weight.shape[-2:], + ), + ], + ) + self.assertTrue(layer.w13_weight_scale_inv.format_ue8m0) + self.assertTrue(layer.w2_weight_scale_inv.format_ue8m0) + if __name__ == "__main__": unittest.main(verbosity=3)