diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 629759de40..45e52e004d 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -52,6 +52,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( ) from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -65,7 +66,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.models.llama4 import Llama4MoE from vllm.platforms import current_platform -from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.math_utils import next_power_of_2, round_up from vllm.utils.torch_utils import set_random_seed try: @@ -788,7 +789,7 @@ def test_trtllm_fp8_swiglu_clamp_support( def _make_mxfp8_moe_weights( - e: int, n: int, k: int + e: int, n: int, k: int, is_gated: bool = True ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: q, s = zip( @@ -796,8 +797,9 @@ def _make_mxfp8_moe_weights( ) return torch.stack(q), torch.stack(s) + w1_rows = 2 * n if is_gated else n w1, w1_scale = quantize( - torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10 + torch.randn(e, w1_rows, k, device="cuda", dtype=torch.bfloat16) / 10 ) w2, w2_scale = quantize( torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10 @@ -986,3 +988,123 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in ) torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize( + ("activation", "n", "k"), + [ + (MoEActivation.SILU, 1024, 1024), + (MoEActivation.SILU, 928, 1024), + (MoEActivation.SILU, 1024, 2688), + (MoEActivation.RELU2_NO_MUL, 1024, 1024), + # Nemotron-3.5 Lightning's non-gated MoE: intermediate 1856, hidden 2688. + (MoEActivation.RELU2_NO_MUL, 1856, 2688), + ], +) +@pytest.mark.parametrize("m", [1, 32]) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Requires TRTLLM-Gen FP8 MoE (SM100)", +) +def test_trtllm_mxfp8_moe_pads_intermediate( + m: int, activation: MoEActivation, n: int, k: int, workspace_init +): + """TRTLLM-Gen MXFP8 experts need the intermediate size to be a multiple + of 128 and the hidden size a multiple of 256. Misaligned experts, gated or + not, are padded at load and, with the activations padded and the output + cropped as the MoE runner does, must still match the unquantized reference + on the dequantized weights.""" + e, topk = 8, 2 + is_gated = activation.is_gated + block_shape = [1, 32] + set_random_seed(7) + with set_current_vllm_config(vllm_config): + w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k, is_gated) + w1_ref = dequant_mxfp8_to_bf16(w1, w1_scale) + w2_ref = dequant_mxfp8_to_bf16(w2, w2_scale) + + layer = SimpleNamespace( + weight_block_size=block_shape, + moe_config=SimpleNamespace( + is_act_and_mul=is_gated, intermediate_size_per_partition=n + ), + activation=activation, + ) + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( + Fp8MoeBackend.FLASHINFER_TRTLLM, + layer, + w1, + w2, + w1_scale, + w2_scale, + w13_input_scale=None, + w2_input_scale=None, + ) + padded_n = layer.moe_config.intermediate_size_per_partition + padded_k = layer.moe_config.hidden_dim + assert padded_n == round_up(n, 128) + assert padded_k == round_up(k, 256) + + quant_config = make_fp8_moe_quant_config( + Fp8MoeBackend.FLASHINFER_TRTLLM, + w1_scale, + w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=block_shape, + ) + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=padded_k, + intermediate_size=padded_n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + routing_method=RoutingMethodType.Renormalize, + max_num_tokens=next_power_of_2(m), + ) + # The monolithic experts pass the activation type to the kernel; the + # modular ones always run it gated. + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=True, + ), + TrtLlmFp8ExpertsMonolithic( + moe_config=moe_config, quant_config=quant_config + ), + ) + + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 + logits = torch.randn((m, e), device="cuda", dtype=torch.float32) + output = kernel.apply_monolithic( + hidden_states=torch.nn.functional.pad(a, (0, padded_k - k)), + w1=w1, + w2=w2, + router_logits=logits, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + routed_scaling_factor=1.0, + routing_sink=None, + )[:, :k] + topk_weights, topk_ids, _ = fused_topk(a, logits, topk, renormalize=True) + reference = fused_experts( + a, + w1_ref, + w2_ref, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=e, + expert_map=None, + ) + + check_accuracy(ref_output=reference, actual_output=output) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 5a6ee76428..4e113d6912 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -379,6 +379,79 @@ def align_moe_weights_for_fi( return padded_w13, padded_w2, padded_intermediate +def align_mxfp8_moe_weights_for_fi_trtllm( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + is_act_and_mul: bool, + min_alignment: int, + hidden_alignment: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """align_moe_weights_for_fi for MXFP8 weights and their UE8M0 block scales + for the FlashInfer TRT-LLM MoE kernel, which also pads the hidden size to a + multiple of hidden_alignment. + + Padded weights are zero and padded scales are 127 (UE8M0 for 1.0), so the + padded rows and columns contribute nothing. Scales may be 3D + [E, rows, cols] or flattened [E, rows * cols]; their layout is kept. + Returns the padded intermediate and hidden sizes. + """ + num_experts, hidden_size, intermediate = w2.shape + padded_hidden = round_up(hidden_size, hidden_alignment) + w13, w2, padded_intermediate = align_moe_weights_for_fi( + w13, w2, is_act_and_mul, min_alignment + ) + if padded_intermediate == intermediate and padded_hidden == hidden_size: + return w13, w2, w13_scale, w2_scale, intermediate, hidden_size + + if padded_hidden != hidden_size: + logger.info_once( + "Padding hidden size from %d to %d for MXFP8 TRTLLM MoE kernel.", + hidden_size, + padded_hidden, + ) + padded_w13 = w13.new_zeros((num_experts, w13.shape[1], padded_hidden)) + padded_w13[:, :, :hidden_size] = w13 + padded_w2 = w2.new_zeros((num_experts, padded_hidden, w2.shape[2])) + padded_w2[:, :hidden_size] = w2 + w13, w2 = padded_w13, padded_w2 + + up_mult = 2 if is_act_and_mul else 1 + w13_scale_3d = w13_scale.reshape(num_experts, up_mult * intermediate, -1) + w2_scale_3d = w2_scale.reshape(num_experts, hidden_size, -1) + block = hidden_size // w13_scale_3d.shape[2] + assert padded_hidden % block == 0 and padded_intermediate % block == 0 + + # Pad each of the gate and up halves (or the single non-gated block) on its + # own, like the weights. + padded_w13_scale = w13_scale_3d.new_full( + (num_experts, up_mult * padded_intermediate, padded_hidden // block), 127 + ) + for i in range(up_mult): + dst = i * padded_intermediate + padded_w13_scale[:, dst : dst + intermediate, : w13_scale_3d.shape[2]] = ( + w13_scale_3d[:, i * intermediate : (i + 1) * intermediate] + ) + padded_w2_scale = w2_scale_3d.new_full( + (num_experts, padded_hidden, padded_intermediate // block), 127 + ) + padded_w2_scale[:, :hidden_size, : w2_scale_3d.shape[2]] = w2_scale_3d + + if w13_scale.ndim == 2: + padded_w13_scale = padded_w13_scale.reshape(num_experts, -1) + if w2_scale.ndim == 2: + padded_w2_scale = padded_w2_scale.reshape(num_experts, -1) + return ( + w13, + w2, + padded_w13_scale, + padded_w2_scale, + padded_intermediate, + padded_hidden, + ) + + def _shuffle_deepseek_fp8_moe_weights( w13: torch.Tensor, w2: torch.Tensor, @@ -486,8 +559,9 @@ def prepare_fp8_moe_layer_for_fi( Note that for trtllm we update the model state dict with the scale format needed for these kernels. - Note that for per-tensor, we update the layer's - intermediate size if the weights needed padding. + Note that for per-tensor and for MXFP8 TRT-LLM, we update the layer's + intermediate size (and for MXFP8 TRT-LLM its hidden size) if the weights + needed padding. """ assert hasattr(layer.moe_config, "is_act_and_mul") block_quant = ( @@ -499,6 +573,21 @@ def prepare_fp8_moe_layer_for_fi( # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle. if is_mxfp8 and is_trtllm: + # The kernel needs intermediate % 128 == 0 and hidden % 256 == 0. The + # MoE runner pads the activations to the padded hidden size. + w13, w2, w13_scale, w2_scale, new_intermediate, new_hidden = ( + align_mxfp8_moe_weights_for_fi_trtllm( + w13, + w2, + w13_scale, + w2_scale, + layer.moe_config.is_act_and_mul, + min_alignment=128, + hidden_alignment=256, + ) + ) + layer.moe_config.intermediate_size_per_partition = new_intermediate + layer.moe_config.hidden_dim = new_hidden # FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores # [gate; up]. Swap both weights and scales before interleaving. if layer.moe_config.is_act_and_mul: