diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 45e52e004d..900b281fbf 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -808,16 +808,20 @@ def _make_mxfp8_moe_weights( @pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"]) +@pytest.mark.parametrize("transposed", [True, False], ids=["transposed", "row-major"]) @pytest.mark.parametrize("m", [1, 16]) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Requires TRTLLM-Gen FP8 MoE (SM100)", ) def test_trtllm_fp8_block_moe_deferred_finalize( - m: int, block_shape: list[int], workspace_init + m: int, block_shape: list[int], transposed: bool, workspace_init ): """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k - finalize to the caller.""" + finalize to the caller. For fp8-block, ``transposed`` toggles the oracle's + is_scale_transposed column-major activation-scale layout.""" + if block_shape != [128, 128] and transposed: + pytest.skip("column-major scales apply to fp8 [128, 128] block only") e, topk, n, k = 32, 4, 1024, 1024 set_random_seed(7) with set_current_vllm_config(vllm_config): @@ -851,6 +855,9 @@ def test_trtllm_fp8_block_moe_deferred_finalize( a2_scale=None, block_shape=block_shape, ) + is_block = block_shape == [128, 128] + if is_block and not transposed: + quant_config = replace(quant_config, is_scale_transposed=False) # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( num_experts=e, @@ -868,19 +875,27 @@ def test_trtllm_fp8_block_moe_deferred_finalize( routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), ) - kernel = mk.FusedMoEKernel( - maybe_make_prepare_finalize( - moe=moe_config, quant_config=quant_config, allow_new_interface=True - ), - TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config), - ) + + def make_kernel(quant_config: FusedMoEQuantConfig) -> mk.FusedMoEKernel: + return mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + ), + TrtLlmFp8ExpertsModular( + moe_config=moe_config, quant_config=quant_config + ), + ) + + kernel = make_kernel(quant_config) a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) - check_deferred_moe_finalize( - moe_config, - lambda: kernel.apply( + + def apply(k_: mk.FusedMoEKernel) -> torch.Tensor: + return k_.apply( hidden_states=a, w1=w1, w2=w2, @@ -890,7 +905,34 @@ def test_trtllm_fp8_block_moe_deferred_finalize( global_num_experts=e, expert_map=None, apply_router_weight_on_input=False, - ), + ) + + if is_block and transposed: + # The oracle gates column-major scales on TRTLLM fp8-block. + assert quant_config.is_scale_transposed + assert not make_fp8_moe_quant_config( + Fp8MoeBackend.DEEPGEMM, + w1_scale, + w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=block_shape, + ).is_scale_transposed + # no_dp_ep emits column-major [M, G] scales (physical [G, M]), so + # the experts' a1q_scale.t().contiguous() is a no-op. + _, a1q_scale, _, _, _ = kernel.prepare_finalize.prepare( + a, topk_weights, topk_ids, e, None, False, quant_config + ) + assert a1q_scale.t().is_contiguous() + # Bitwise-equal to running with row-major scales. + row_major = replace(quant_config, is_scale_transposed=False) + torch.testing.assert_close( + apply(kernel), apply(make_kernel(row_major)), atol=0, rtol=0 + ) + + check_deferred_moe_finalize( + moe_config, + lambda: apply(kernel), router_weights=topk_weights, ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e79700ae44..e8bd822bec 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -249,6 +249,8 @@ class FusedMoEQuantConfig: _w1: FusedMoEQuantDesc _w2: FusedMoEQuantDesc is_scale_swizzled: bool = True + # Emit block-quant activation scales column-major (physical [G, M] layout). + is_scale_transposed: bool = False # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping. # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit @@ -503,6 +505,7 @@ class FusedMoEQuantConfig: w2_zp: torch.Tensor | None = None, weight_dtype: torch.dtype | str | None = None, is_scale_swizzled: bool = True, + is_scale_transposed: bool = False, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, @@ -537,6 +540,9 @@ class FusedMoEQuantConfig: - is_scale_swizzled: Whether the activation scale-factor layout is swizzled. Pass through to the underlying quantization kernel for dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True. + - is_scale_transposed: Whether block-quantized activation scales are + emitted column-major (physical [num_groups, num_tokens] layout) by + the quantize step. Defaults to False. - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter. - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter. - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit. @@ -573,6 +579,7 @@ class FusedMoEQuantConfig: weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias ), is_scale_swizzled=is_scale_swizzled, + is_scale_transposed=is_scale_transposed, gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, @@ -597,6 +604,7 @@ def fp8_w8a8_moe_quant_config( a2_gscale: torch.Tensor | None = None, g1_alphas: torch.Tensor | None = None, g2_alphas: torch.Tensor | None = None, + is_scale_transposed: bool = False, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, @@ -617,6 +625,7 @@ def fp8_w8a8_moe_quant_config( per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, block_shape=block_shape, + is_scale_transposed=is_scale_transposed, gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index aef5c5b5b5..07e8815fc7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -841,6 +841,11 @@ def make_fp8_moe_quant_config( block_shape=block_shape, per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, + # TRTLLM fp8-block consumes activation scales as [K/128, M]; emit them + # transposed from the quant kernel instead of copying afterwards. + is_scale_transposed=( + fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM and block_shape == [128, 128] + ), gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py index 660807fea3..fe5d518cf3 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -31,6 +31,7 @@ def _quantize_input( per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, is_scale_swizzled=quant_config.is_scale_swizzled, + is_scale_transposed=quant_config.is_scale_transposed, mx_alignment=quant_config.mx_alignment, ) diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b1a6dd5632..776c2064fa 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -325,6 +325,7 @@ def _fp8_quantize( A_scale: torch.Tensor | None, per_act_token: bool, block_shape: list[int] | None = None, + is_scale_transposed: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Perform fp8 quantization on the inputs. If a block_shape is provided, the output will be blocked. @@ -339,7 +340,9 @@ def _fp8_quantize( assert not per_act_token assert len(block_shape) == 2 _, block_k = block_shape[0], block_shape[1] - A, A_scale = per_token_group_quant_fp8(A, block_k) + A, A_scale = per_token_group_quant_fp8( + A, block_k, column_major_scales=is_scale_transposed + ) assert cdiv(A.size(-1), block_k) == A_scale.size(-1) return A, A_scale @@ -462,6 +465,7 @@ def moe_kernel_quantize_input( per_act_token_quant: bool, block_shape: list[int] | None = None, is_scale_swizzled: bool = True, + is_scale_transposed: bool = False, quantization_emulation: bool = False, mx_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -469,7 +473,9 @@ def moe_kernel_quantize_input( if quantization_emulation: return _fp8_quantize_dequantize(A, A_scale) else: - return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape) + return _fp8_quantize( + A, A_scale, per_act_token_quant, block_shape, is_scale_transposed + ) elif quant_dtype == torch.int8: if quantization_emulation: raise NotImplementedError(