diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 15b9835d6c..da507e5caf 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -21,6 +21,7 @@ from vllm.model_executor.kernels.linear.nvfp4.base import ( NvFp4LinearLayerConfig, ) from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) @@ -59,6 +60,7 @@ from vllm.platforms import current_platform SUPPORTING = { CutlassFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel, diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index 54dcaa0b97..47c25d046d 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -13,11 +13,15 @@ from nvfp4_utils import ( from vllm import _custom_ops as ops from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCuteDslNvFp4W4A16LinearKernel, + FlashInferCutlassNvFp4LinearKernel, ) +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_to_dtype, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.flashinfer import ( flashinfer_scaled_fp4_mm, @@ -47,6 +51,54 @@ SEEDS = [42] CUDA_DEVICES = ["cuda:0"] +@pytest.mark.parametrize("token_shape", [(17,), (1, 17), (2, 17)]) +@pytest.mark.parametrize( + "kernel_cls", + [FlashInferCutlassNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel], +) +@torch.inference_mode() +def test_flashinfer_nvfp4_preserves_quantized_input_shape( + token_shape: tuple[int, ...], + kernel_cls: type[FlashInferCutlassNvFp4LinearKernel] + | type[FlashInferCuteDslNvFp4LinearKernel], +) -> None: + """Pre-quantized inputs must match the tensor path and retain token dimensions.""" + supported, reason = kernel_cls.is_supported() + if not supported: + pytest.skip(reason) + + set_random_seed(42) + hidden_size, output_size = 128, 128 + x = torch.randn(*token_shape, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(output_size, hidden_size, dtype=x.dtype, device=x.device) + input_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / x.abs().max()).float() + weight_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).float() + x_fp4, x_blockscale = ops.scaled_fp4_quant(x.reshape(-1, hidden_size), input_scale) + weight_fp4, weight_blockscale = ops.scaled_fp4_quant(weight, weight_scale) + quantized_input = QuantizedActivation( + data=x_fp4.reshape(*token_shape, hidden_size // 2), + scale=x_blockscale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kNvfp4Dynamic, + ) + layer = torch.nn.Module() + layer.input_size_per_partition = hidden_size + layer.output_size_per_partition = output_size + layer.weight = weight_fp4 + layer.weight_scale = weight_blockscale + layer.input_global_scale_inv = input_scale + layer.alpha = (input_scale * weight_scale).reciprocal() + kernel = kernel_cls(NvFp4LinearLayerConfig()) + + expected = kernel.apply_weights(layer, x.reshape(-1, hidden_size)) + actual = kernel.apply_weights(layer, quantized_input) + + assert actual.shape == (*token_shape, output_size) + assert actual.dtype == x.dtype + torch.testing.assert_close(actual.reshape(expected.shape), expected) + + def get_ref_results( a_fp4, b_fp4, diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 321cb6978e..90309dbd1d 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -115,13 +115,13 @@ class MockLinearRequiresUnquantized(MockLinearFp8Static): requires_unquantized_input = True -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_static( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -133,26 +133,29 @@ def test_maybe_fused_act_quant_fp8_static( scale = torch.tensor([0.5], device=device, dtype=torch.float32) linear = MockLinearFp8Static(scale) - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8StaticTensorSym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape - ref_out = ref_impl(act_fn, x, scale) + ref_out = ref_impl(act_fn, x.reshape(-1, hidden_size * 2), scale).reshape( + result.orig_shape + ) torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32)) -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_dynamic_block( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -166,31 +169,34 @@ def test_maybe_fused_act_quant_fp8_dynamic_block( linear = MockLinearFp8Dynamic128() scale = 1 / hidden_size - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) * scale result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8Dynamic128Sym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape num_groups = hidden_size // group_size - assert result.scale.shape == (num_tokens, num_groups) + assert result.scale.shape == (*token_shape, num_groups) gate, up = x.split(hidden_size, dim=-1) silu_out = F.silu(gate) * up ref_out, ref_scales = per_token_group_quant_fp8( - silu_out, group_size=group_size, use_ue8m0=False + silu_out.reshape(-1, hidden_size), group_size=group_size, use_ue8m0=False ) + ref_out = ref_out.reshape(result.orig_shape) + ref_scales = ref_scales.reshape(*token_shape, num_groups) torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5) ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index a555ca17bf..2a69a3be21 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -112,6 +112,9 @@ class FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearKernel): class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via FlashInfer's cutedsl backend.""" + def input_quant_key(self) -> QuantKey | None: + return kNvfp4Dynamic + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -137,19 +140,26 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: output_size = layer.output_size_per_partition - output_dtype = x.dtype - output_shape = [*x.shape[:-1], output_size] - - x_fp4, x_blockscale = scaled_fp4_quant( - x, - layer.input_global_scale_inv, - is_sf_swizzled_layout=True, - backend="flashinfer-cutedsl", - ) + qa = as_quantized_activation(x, self.input_quant_key()) + if qa is not None: + x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) + output_dtype = qa.orig_dtype + output_shape = [*qa.orig_shape[:-1], output_size] + else: + assert isinstance(x, torch.Tensor) + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-cutedsl", + ) x_fp4 = pad_nvfp4_activation_for_cutlass( x_fp4, nvfp4_weight_padding_bytes(layer) @@ -219,6 +229,7 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes) output_dtype = qa.orig_dtype output_shape = [*qa.orig_shape[:-1], output_size] diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index a9b14c6170..6261a1bdd0 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -63,10 +63,10 @@ def _silu_and_mul_fp8_dynamic_block( x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey ) -> QuantizedActivation: """SiluAndMul + FP8 dynamic per-block quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] num_groups = d // group_size @@ -102,10 +102,10 @@ def _silu_and_mul_nvfp4_dynamic( x: torch.Tensor, linear: LinearBase ) -> QuantizedActivation: """SiluAndMul + NVFP4 dynamic quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] # NVFP4 packs 2 values into 1 byte diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index fb10026ad1..4f96c0acfc 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -32,6 +32,7 @@ from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.conv import Conv2dLayer +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -758,7 +759,7 @@ class FeedForward(nn.Module): def forward(self, x): x, _ = self.gate_up_proj(x) - x = self.act_fn(x) + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) x, _ = self.down_proj(x) return x @@ -1259,7 +1260,7 @@ class PixtralHFMLP(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) - x = self.act_and_mul(gate_up) + x = maybe_fused_act_quant(self.act_and_mul, gate_up, self.down_proj) x, _ = self.down_proj(x) return x