From 9db4650e5e4c726eb5ae29330cd55e796567469c Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Tue, 21 Apr 2026 18:12:36 -0400 Subject: [PATCH] [MoE Refactor] Add more MoE layer tests (#39349) Signed-off-by: Bill Nell Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../modular_kernel_tools/parallel_utils.py | 2 +- tests/kernels/moe/test_moe_layer.py | 137 ++++++++++++------ .../model_executor/layers/fused_moe/config.py | 3 +- 3 files changed, 98 insertions(+), 44 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index 95004fa0ab..07f244451b 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -77,8 +77,8 @@ def _worker_parallel_launch( *args: Any, ) -> None: rank = node_rank * world_local_size + local_rank - torch.accelerator.set_device_index(local_rank) device = torch.device("cuda", local_rank) + torch.accelerator.set_device_index(device) torch.distributed.init_process_group( backend="cpu:gloo,cuda:nccl", init_method=init_method, diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index fca4096b08..838674db58 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -65,8 +65,8 @@ fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype SHAPE_COMBOS = [ (1, 128, 256), - (32, 1024, 512), - (222, 2048, 2048), + (32, 512, 512), + (222, 1024, 2048), ] MAX_M = max([x[0] for x in SHAPE_COMBOS]) @@ -95,7 +95,7 @@ if has_flashinfer_nvlink_one_sided(): BACKENDS += ["flashinfer_nvlink_one_sided"] if has_deep_ep(): - BACKENDS += ["deepep_low_latency", "deepep_high_throughput"] + BACKENDS += ["deepep_high_throughput", "deepep_low_latency"] if has_nixl_ep(): BACKENDS += ["nixl_ep"] @@ -103,6 +103,7 @@ if has_nixl_ep(): QUANT_METHODS = [ None, "fp8", + "fp8_blocked", "modelopt_fp8", "modelopt_fp4", ] @@ -114,10 +115,21 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { "mori": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"}, "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"}, - "deepep_low_latency": {None, "modelopt_fp8", "modelopt_fp4"}, - "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, + "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, + "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 "nixl_ep": {None, "fp8", "modelopt_fp8"}, } + +# Map from backend -> (DP/EP support, DP support, TP support) +BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool]] = { + "allgather_reducescatter": (True, True, True), + "mori": (True, False, False), + "flashinfer_nvlink_two_sided": (False, True, False), + "flashinfer_nvlink_one_sided": (False, True, False), + "deepep_low_latency": (True, False, False), + "deepep_high_throughput": (True, False, False), + "nixl_ep": (True, False, False), +} # fmt: on # Which quantization methods support EPLB. @@ -424,27 +436,35 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: f"Skipping unsupported K {config.k} in {config.backend} w/o EP.", ) - if config.enable_eplb and config.ep_size == 1: - return False, "EPLB requires EP." + if config.backend is not None: + supports_ep_dp, supports_dp, supports_tp = BACKEND_EP_DP_TP_SUPPORT[ + config.backend + ] - if config.enable_eplb and config.quantization not in EPLB_SUPPORTED_QUANTS: - return False, f"EPLB not supported with {config.quantization} quantization." + if config.tp_size > 1 and not supports_tp: + return False, f"{config.backend} does not support TP." - if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS: - return False, f"EPLB not supported with {config.backend}." + if config.dp_size > 1 and config.ep_size == 1 and not supports_dp: + return False, f"{config.backend} does not support DP." - if ( - config.backend is not None - and config.backend.startswith("flashinfer_nvlink") - and config.ep_size > 1 - ): - return False, "flashinfer_nvlink EP not yet supported." + if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp: + return False, f"{config.backend} does not support EP/DP." + else: + if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1: + return False, "An all2all backend is required for parallelism." - if config.enable_eplb and config.num_experts % config.dp_size != 0: - return False, "EPLB requires num_experts divisible by ep_size" + if config.enable_eplb: + if config.ep_size == 1: + return False, "EPLB requires EP." - if config.enable_eplb and config.ep_size == 1: - return False, "EPLB only works with EP+DP" + if config.quantization not in EPLB_SUPPORTED_QUANTS: + return False, f"EPLB not supported with {config.quantization} quantization." + + if config.backend not in EPLB_SUPPORTED_BACKENDS: + return False, f"EPLB not supported with {config.backend}." + + if config.num_experts % config.dp_size != 0: + return False, "EPLB requires num_experts divisible by ep_size" # Disable fp4 tests until flashinfer is updated or the Dockerfile is # modified to install cublasLt.h. See #39525. @@ -507,27 +527,48 @@ class QuantizedWeights: def _quantize_fp8_halves( w1: torch.Tensor, w2: torch.Tensor, + block_shape: list[int] | None = None, ) -> QuantizedWeights: """Quantize w13 gate/up halves separately to FP8, producing per-shard scales.""" half = w1.shape[1] // 2 w1q_a, w1s_a, _ = moe_quantize_weights( - w1[:, :half, :], None, fp8_dtype, False, None + w1[:, :half, :], + None, + fp8_dtype, + False, + block_shape, ) w1q_b, w1s_b, _ = moe_quantize_weights( - w1[:, half:, :], None, fp8_dtype, False, None + w1[:, half:, :], + None, + fp8_dtype, + False, + block_shape, ) assert w1s_a is not None and w1s_b is not None - w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, None) + w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, block_shape) assert w2s is not None + if block_shape is not None: + # Blocked quantization: scales have shape (E, n_tiles, k_tiles) + # Concatenate gate and up scales along the n_tiles dimension (dim=1) + # to match the concatenation of gate and up weights + w13_weight_scale = torch.cat([w1s_a, w1s_b], dim=1) + # w2 scales keep their blocked shape (E, k_tiles, n_tiles) + w2_weight_scale = w2s + else: + # Non-blocked quantization: scales have shape (E, 1, 1) + # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2) + w13_weight_scale = torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1) + # w2s is (E, 1, 1) -> reshape to (E,) + w2_weight_scale = w2s.view(-1) + return QuantizedWeights( w13_weight=torch.cat([w1q_a, w1q_b], dim=1), w2_weight=w2q, - # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2) - w13_weight_scale=torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1), - # w2s is (E, 1, 1) -> reshape to (E,) - w2_weight_scale=w2s.view(-1), + w13_weight_scale=w13_weight_scale, + w2_weight_scale=w2_weight_scale, ) @@ -536,7 +577,7 @@ def quantization_to_quant_dtype( ) -> torch.dtype | str | None: if quantization is None: return None - elif quantization in ["fp8", "modelopt_fp8"]: + elif quantization in ["fp8", "fp8_blocked", "modelopt_fp8"]: return fp8_dtype elif quantization in ["modelopt_fp4"]: return "nvfp4" @@ -558,6 +599,12 @@ def make_quant_config( if quantization == "fp8": return Fp8Config(True), _quantize_fp8_halves(w1, w2) + if quantization == "fp8_blocked": + block_shape = [128, 128] + return Fp8Config(True, weight_block_size=block_shape), _quantize_fp8_halves( + w1, w2, block_shape + ) + if quantization == "modelopt_fp8": qw = _quantize_fp8_halves(w1, w2) # why? @@ -896,11 +943,13 @@ def make_fused_moe_layer( **kwargs, ) + weight_scale_name = getattr(layer.quant_method, "weight_scale_name", "weight_scale") + for name, value in [ ("w13_weight", qw.w13_weight), ("w2_weight", qw.w2_weight), - ("w13_weight_scale", qw.w13_weight_scale), - ("w2_weight_scale", qw.w2_weight_scale), + (f"w13_{weight_scale_name}", qw.w13_weight_scale), + (f"w2_{weight_scale_name}", qw.w2_weight_scale), ("w13_weight_scale_2", qw.w13_weight_scale_2), ("w2_weight_scale_2", qw.w2_weight_scale_2), ("w13_input_scale", qw.w13_input_scale), @@ -922,7 +971,7 @@ def make_fake_moe_layer( top_k: int, global_num_experts: int, in_dtype: torch.dtype, - quant_dtype: torch.dtype | None, + quantization: str | None, renormalize: bool = False, shared_experts_config: SharedExpertsConfig | None = None, use_grouped_topk: bool = False, @@ -948,6 +997,7 @@ def make_fake_moe_layer( dp_size: int = 1, ep_size: int = 1, ) -> Callable: + quant_dtype = None activation = MoEActivation.from_str(activation) router = create_fused_moe_router( @@ -1139,7 +1189,6 @@ def _test_body_eplb( routed_output_transform=routed_output_transform, ) - # Necessary? if eplb_moe_layer._expert_map is not None: eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device) @@ -1267,6 +1316,7 @@ def _run_one_config( gate = test_data.gate routed_input_transform = test_data.routed_input_transform routed_output_transform = test_data.routed_output_transform + activation = "silu" baseline_layer = make_fake_moe_layer( w1=w1, @@ -1274,7 +1324,7 @@ def _run_one_config( top_k=top_k, global_num_experts=num_experts, in_dtype=in_dtype, - quant_dtype=None, # quantization_to_quant_dtype(quantization), + quantization=quantization, renormalize=False, shared_experts_config=shared_experts_config, gate=gate, @@ -1284,6 +1334,7 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, + activation=activation, ) baseline_output = baseline_layer(hidden_states, router_logits) @@ -1328,9 +1379,9 @@ def _run_one_config( gate=gate, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, + activation=activation, ) - # Necessary? if moe_layer._expert_map is not None: moe_layer._expert_map = moe_layer._expert_map.to(device) @@ -1377,13 +1428,17 @@ def _run_one_config( atol, rtol = 7.6e-2, 7.6e-2 else: atol, rtol = 3.5e-2, 3.5e-2 - elif quantization in ("fp8", "modelopt_fp8"): - if k >= 2048: - atol, rtol = 7.6e-2, 7.6e-2 - else: - atol, rtol = 6e-2, 6e-2 + elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"): + atol, rtol = 6e-2, 6e-2 elif quantization == "modelopt_fp4": - atol = rtol = 1e-1 + k * 5e-4 + if k >= 2048: + atol = rtol = 1e-1 + (k * 1e-4) + else: + atol = rtol = 1e-1 + + if backend == "allgather_reducescatter" and tp_size > 1: + atol += 2e-1 + rtol += 2e-1 else: atol, rtol = 6e-2, 6e-2 diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 231c5652e4..00d7d8e789 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -990,8 +990,7 @@ class FusedMoEParallelConfig: @property def use_batched_activation_format(self): - # TODO(bnell): nixl also uses batched format - return self.use_deepep_ll_kernels + return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels @property def use_ag_rs_all2all_kernels(self):