diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index c239e93785b3..77c3ba5112d9 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/tests/compile/fusions_e2e/common.py @@ -43,6 +43,45 @@ class AttentionBackendCase(NamedTuple): """Are we running on Blackwell, a lot of tests depend on it""" +def nvfp4_kernel_exposes_input_quant_key() -> bool: + """Check if the NVFP4 kernel selected on this platform exposes input_quant_key. + + FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to + layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant) + output format and the cutedsl backend's expected input format. + + FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports + manual fusion. + """ + if not current_platform.is_cuda(): + return False + + try: + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, + ) + + is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported() + if is_supported: + return False + except ImportError: + pass + + try: + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCutlassNvFp4LinearKernel, + ) + + is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported() + if is_supported: + return True + except ImportError: + pass + + # Fallback: assume no manual fusion support + return False + + def custom_ops_combos(*custom_ops: str) -> Iterable[str]: """Generate all combinations of custom ops for parametrization.""" custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops] diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py index 3fab133f963e..fa1af8e1c6f7 100644 --- a/tests/compile/fusions_e2e/test_tp1_quant.py +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -14,6 +14,7 @@ Matches, custom_ops_combos, is_blackwell, + nvfp4_kernel_exposes_input_quant_key, ) from .models import ( FLASHINFER_ATTN, @@ -77,7 +78,6 @@ def test_tp1_fp8_fusions( inductor_graph_partition: bool, use_deepgemm: bool, run_e2e_fusion_test, - monkeypatch, ): if use_deepgemm and not current_platform.is_cuda(): pytest.skip("DeepGemm only supported on CUDA") @@ -109,7 +109,7 @@ def test_tp1_fp8_fusions( custom_ops=custom_ops.split(","), pass_config=PassConfig( fuse_norm_quant=True, - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_qk_norm_rope_fusion=True, ), @@ -119,7 +119,6 @@ def test_tp1_fp8_fusions( matches_check = [ "rms_quant_fusion", - "act_quant_fusion", "norm_rope_fusion", "attn_quant_fusion", ] @@ -168,6 +167,12 @@ def test_tp1_fp4_fusions( inductor_graph_partition: bool, run_e2e_fusion_test, ): + if nvfp4_kernel_exposes_input_quant_key(): + pytest.skip( + "NVFP4 kernel exposes input_quant_key; manual fusion fires " + "instead of compiler pass-based fusion" + ) + matches = matches_fn(n_layers) # Reduce size of model and skip weight loading time @@ -181,13 +186,13 @@ def test_tp1_fp4_fusions( custom_ops=custom_ops.split(","), pass_config=PassConfig( fuse_norm_quant=True, - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_qk_norm_rope_fusion=True, ), ) - matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"] + matches_check = ["attn_quant_fusion", "norm_rope_fusion"] run_e2e_fusion_test( model_name, diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index c88d47cd314f..4b3db06bc446 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -14,6 +14,7 @@ Matches, custom_ops_combos, is_blackwell, + nvfp4_kernel_exposes_input_quant_key, ) from .models import ( FLASHINFER_ATTN, @@ -68,7 +69,6 @@ def test_tp2_ar_rms_fp8_fusions( custom_ops: str, inductor_graph_partition: bool, run_e2e_fusion_test, - monkeypatch, ): matches = matches_fn(n_layers) @@ -89,7 +89,7 @@ def test_tp2_ar_rms_fp8_fusions( custom_ops=custom_ops.split(","), pass_config=PassConfig( fuse_norm_quant=True, - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_qk_norm_rope_fusion=True, fuse_allreduce_rms=True, @@ -98,7 +98,6 @@ def test_tp2_ar_rms_fp8_fusions( matches_check = [ "rms_quant_fusion", - "act_quant_fusion", "norm_rope_fusion", "attn_quant_fusion", "ar_rms_fusion", @@ -139,8 +138,13 @@ def test_tp2_ar_rms_fp4_fusions( custom_ops: str, inductor_graph_partition: bool, run_e2e_fusion_test, - monkeypatch, ): + if nvfp4_kernel_exposes_input_quant_key(): + pytest.skip( + "NVFP4 kernel exposes input_quant_key; manual fusion fires " + "instead of compiler pass-based fusion" + ) + matches = matches_fn(n_layers) # Reduce size of model and skip weight loading time @@ -154,14 +158,13 @@ def test_tp2_ar_rms_fp4_fusions( use_inductor_graph_partition=inductor_graph_partition, custom_ops=custom_ops.split(","), pass_config=PassConfig( - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, fuse_allreduce_rms=True, ), ) matches_check = [ - "act_quant_fusion", "attn_quant_fusion", "ar_rms_fusion", ] diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py index a22c68f4bf92..b0d32dad1c17 100644 --- a/tests/compile/fusions_e2e/test_tp2_async_tp.py +++ b/tests/compile/fusions_e2e/test_tp2_async_tp.py @@ -14,6 +14,7 @@ Matches, custom_ops_combos, is_blackwell, + nvfp4_kernel_exposes_input_quant_key, ) from .models import ( FLASHINFER_ATTN, @@ -61,7 +62,7 @@ def test_tp2_async_tp_fp8_fusions( custom_ops=custom_ops.split(","), pass_config=PassConfig( fuse_norm_quant=True, - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_qk_norm_rope_fusion=True, enable_sp=True, @@ -74,7 +75,6 @@ def test_tp2_async_tp_fp8_fusions( matches_check = [ "rms_quant_fusion", - "act_quant_fusion", "norm_rope_fusion", "attn_quant_fusion", "sequence_parallel", @@ -114,6 +114,12 @@ def test_tp2_async_tp_nvfp4_fusions( inductor_graph_partition: bool, run_e2e_fusion_test, ): + if nvfp4_kernel_exposes_input_quant_key(): + pytest.skip( + "NVFP4 kernel exposes input_quant_key; manual fusion fires " + "instead of compiler pass-based fusion" + ) + # NVFP4 currently wires the all-gather + GEMM path only. matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2) @@ -127,7 +133,7 @@ def test_tp2_async_tp_nvfp4_fusions( use_inductor_graph_partition=inductor_graph_partition, custom_ops=custom_ops.split(","), pass_config=PassConfig( - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_sp=True, fuse_gemm_comms=True, @@ -138,7 +144,6 @@ def test_tp2_async_tp_nvfp4_fusions( ) matches_check = [ - "act_quant_fusion", "attn_quant_fusion", "sequence_parallel", "async_tp", @@ -245,7 +250,7 @@ def test_tp2_sp_ar_rms_fp8_fusions( custom_ops=custom_ops.split(","), pass_config=PassConfig( fuse_norm_quant=True, - fuse_act_quant=True, + fuse_act_quant=False, fuse_attn_quant=True, enable_qk_norm_rope_fusion=True, enable_sp=True, @@ -258,7 +263,6 @@ def test_tp2_sp_ar_rms_fp8_fusions( matches_check = [ "rms_quant_fusion", - "act_quant_fusion", "norm_rope_fusion", "attn_quant_fusion", "ar_rms_fusion", diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py new file mode 100644 index 000000000000..e41b0877f68c --- /dev/null +++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for manual fusion via maybe_fused_act_quant. + +Tests all fusion paths in _FUSED_ACT_QUANT: +- kFp8StaticTensorSym: all platforms +- kFp8Dynamic128Sym: CUDA only +- kNvfp4Dynamic: CUDA SM100+ only +""" + +import pytest +import torch + +import vllm.envs as envs +from tests.utils import TestFP8Layer +from vllm.config import ( + CompilationConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.model_executor.kernels.linear import ( + CutlassFP8ScaledMMLinearKernel, + FlashInferFP8ScaledMMLinearKernel, + FP8ScaledMMLinearKernel, + PerTensorTorchFP8ScaledMMLinearKernel, + ROCmFP8ScaledMMLinearKernel, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fusion.fused_act_quant import ( + _FUSED_ACT_QUANT, + maybe_fused_act_quant, +) +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8StaticTensorSym, + kNvfp4Dynamic, +) +from vllm.platforms import current_platform + + +# Mock linear layer for testing fusion paths that don't have real kernel support +class MockLinearForFusion(torch.nn.Module): + """Mock linear layer that exposes input_quant_key for fusion testing.""" + + def __init__(self, quant_key, input_scale=None, input_global_scale=None): + super().__init__() + self.input_quant_key = quant_key + if input_scale is not None: + self.input_scale = input_scale + if input_global_scale is not None: + self.input_global_scale = input_global_scale + + +ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel] +CUDA_KERNELS = [ + FlashInferFP8ScaledMMLinearKernel, + CutlassFP8ScaledMMLinearKernel, + PerTensorTorchFP8ScaledMMLinearKernel, +] +TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS + + +@pytest.mark.parametrize("num_tokens", [32, 64]) +@pytest.mark.parametrize("hidden_size", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("force_kernel", TEST_KERNELS) +@pytest.mark.skipif( + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" +) +def test_manual_fusion_fp8_static_with_linear( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + force_kernel: FP8ScaledMMLinearKernel, +): + """Test manual fusion with real FP8 linear layer (kFp8StaticTensorSym). + + This is an end-to-end test that verifies the full flow: + unfused (silu_and_mul -> in-kernel quant) vs fused (silu_and_mul_quant). + """ + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + x = torch.rand(num_tokens, hidden_size * 2) + + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + + with set_current_vllm_config(config): + silu_and_mul = SiluAndMul() + fp8_linear = TestFP8Layer( + weight_shape=(hidden_size, hidden_size), + activation_quant_key=kFp8StaticTensorSym, + weight_quant_key=kFp8StaticTensorSym, + force_kernel=force_kernel, + input_dtype=dtype, + ) + + # Run without fusion: silu_and_mul returns plain tensor + y_unfused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear) + assert isinstance(y_unfused, torch.Tensor) + result_unfused = fp8_linear(y_unfused) + + # Enable fusion + expose_input_quant_key(fp8_linear, fp8_linear.kernel) + + if not hasattr(fp8_linear, "input_quant_key"): + pytest.skip( + f"Kernel {force_kernel.__name__} doesn't support input_quant_key" + ) + + # Run with fusion: silu_and_mul returns QuantizedActivation + y_fused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear) + assert isinstance(y_fused, QuantizedActivation) + assert y_fused.quant_key == kFp8StaticTensorSym + assert y_fused.data.dtype == current_platform.fp8_dtype() + assert y_fused.data.shape == (num_tokens, hidden_size) + assert y_fused.orig_dtype == dtype + assert y_fused.orig_shape == (num_tokens, hidden_size) + + result_fused = fp8_linear(y_fused) + + torch.testing.assert_close( + result_fused.to(dtype=dtype), + result_unfused.to(dtype=dtype), + atol=5e-2, + rtol=5e-2, + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Dynamic block quant CUDA only" +) +@pytest.mark.skipif( + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" +) +def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype): + """Test kFp8Dynamic128Sym fusion path (group_size=128). + + Compares fused (silu_and_mul_per_block_quant) vs unfused (silu_and_mul) + by dequantizing the fused result and comparing with unfused. + """ + if (SiluAndMul, kFp8Dynamic128Sym) not in _FUSED_ACT_QUANT: + pytest.skip("kFp8Dynamic128Sym fusion not available") + + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + # hidden_size must be divisible by group_size (128) + num_tokens, hidden_size = 32, 256 + group_size = 128 + x = torch.rand(num_tokens, hidden_size * 2) + + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + + with set_current_vllm_config(config): + silu_and_mul = SiluAndMul() + + # Unfused path: just apply silu_and_mul + mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size) + result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key) + assert isinstance(result_unfused, torch.Tensor) + + # Fused path: apply silu_and_mul + per-block quantization + mock_linear_with_key = MockLinearForFusion(kFp8Dynamic128Sym) + result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key) + + # Verify fused result structure + assert isinstance(result_fused, QuantizedActivation) + assert result_fused.quant_key == kFp8Dynamic128Sym + assert result_fused.data.dtype == current_platform.fp8_dtype() + assert result_fused.data.shape == (num_tokens, hidden_size) + assert result_fused.orig_dtype == dtype + + # Check scale shape + expected_num_groups = hidden_size // group_size + assert result_fused.scale.shape == (num_tokens, expected_num_groups) + + # Dequantize fused result and compare with unfused + # Per-block dequant: data * scale (broadcast scale across group) + dequant_data = result_fused.data.to(dtype).view( + num_tokens, expected_num_groups, group_size + ) + scales_expanded = result_fused.scale.unsqueeze( + -1 + ) # (num_tokens, num_groups, 1) + dequant_result = ( + (dequant_data * scales_expanded).view(num_tokens, hidden_size).to(dtype) + ) + + torch.testing.assert_close( + dequant_result, + result_unfused, + atol=5e-2, + rtol=5e-2, + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="NVFP4 CUDA only") +@pytest.mark.skipif( + not current_platform.has_device_capability(100), reason="NVFP4 requires SM100+" +) +@pytest.mark.skipif( + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" +) +def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype): + """Test kNvfp4Dynamic fusion path. + + Compares fused (silu_and_mul_nvfp4_quant) vs unfused (silu_and_mul) + by dequantizing the fused result and comparing with unfused. + """ + if (SiluAndMul, kNvfp4Dynamic) not in _FUSED_ACT_QUANT: + pytest.skip("kNvfp4Dynamic fusion not available") + + from tests.kernels.quantization.nvfp4_utils import dequantize_nvfp4_to_dtype + + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + # NVFP4 requires hidden_size divisible by 16 (block size) and by 2 (packing). + # M (num_tokens) must be >= 128 for the 128x4 swizzled scale layout. + num_tokens, hidden_size = 128, 256 + x = torch.rand(num_tokens, hidden_size * 2) + input_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") + + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + + with set_current_vllm_config(config): + silu_and_mul = SiluAndMul() + + # Unfused path: just apply silu_and_mul + mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size) + result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key) + assert isinstance(result_unfused, torch.Tensor) + + # Fused path: apply silu_and_mul + NVFP4 quantization + mock_linear_with_key = MockLinearForFusion( + kNvfp4Dynamic, input_global_scale=input_global_scale + ) + result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key) + + # Verify fused result structure + assert isinstance(result_fused, QuantizedActivation) + assert result_fused.quant_key == kNvfp4Dynamic + # NVFP4 packs 2 values into 1 byte + assert result_fused.data.dtype == torch.uint8 + assert result_fused.data.shape == (num_tokens, hidden_size // 2) + assert result_fused.orig_dtype == dtype + assert result_fused.orig_shape == (num_tokens, hidden_size) + + # Dequantize fused result and compare with unfused + dequant_result = dequantize_nvfp4_to_dtype( + tensor_fp4=result_fused.data, + tensor_sf=result_fused.scale, + global_scale=input_global_scale, + dtype=dtype, + device="cuda", + block_size=16, + is_sf_128x4_layout=True, + ) + + torch.testing.assert_close( + dequant_result, + result_unfused, + atol=3e-1, + rtol=3e-1, + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif( + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" +) +def test_manual_fusion_fallback_no_key(dtype: torch.dtype): + """Test that maybe_fused_act_quant falls back when no input_quant_key.""" + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + x = torch.rand(32, 256) + + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + + with set_current_vllm_config(config): + silu_and_mul = SiluAndMul() + # Linear without input_quant_key attribute + mock_linear = torch.nn.Linear(128, 128) + + result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) + + # Should fall back to plain silu_and_mul + assert isinstance(result, torch.Tensor) + assert result.shape == (32, 128) + assert result.dtype == dtype + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif( + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" +) +def test_manual_fusion_fallback_unsupported_key(dtype: torch.dtype): + """Test that maybe_fused_act_quant falls back for unsupported quant keys.""" + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + QuantKey, + ScaleDesc, + ) + + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + x = torch.rand(32, 256) + + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + + with set_current_vllm_config(config): + # Create an unsupported quant key + unsupported_key = QuantKey( + dtype=torch.int8, + scale=ScaleDesc( + dtype=torch.float32, static=False, group_shape=GroupShape(1, 1) + ), + ) + + silu_and_mul = SiluAndMul() + mock_linear = MockLinearForFusion(unsupported_key) + + result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) + + # Should fall back to plain silu_and_mul since key not in _FUSED_ACT_QUANT + assert isinstance(result, torch.Tensor) + assert result.shape == (32, 128) + assert result.dtype == dtype diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 48d492b8d2ee..a794ff1e3a67 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -25,6 +25,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( FlashInferFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.pytorch import ( + PerTensorTorchFP8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( FP8ScaledMMLinearLayerConfig, Int8ScaledMMLinearKernel, @@ -46,6 +49,7 @@ CutlassFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, + PerTensorTorchFP8ScaledMMLinearKernel, } diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 0696ebb8d556..267484e43361 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -2,10 +2,20 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch +import torch.nn.functional as F import vllm._custom_ops as ops from tests.kernels.utils import opcheck from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8StaticTensorSym, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -69,3 +79,140 @@ def test_silu_and_mul( ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32) ) opcheck(torch.ops._C.silu_and_mul_quant, (ops_out, x, scale)) + + +# --------------------------------------------------------------------------- +# Tests for maybe_fused_act_quant interface +# --------------------------------------------------------------------------- + + +class MockLinearFp8Static(torch.nn.Module): + """Mock linear layer advertising kFp8StaticTensorSym.""" + + def __init__(self, input_scale: torch.Tensor): + super().__init__() + self.input_quant_key = kFp8StaticTensorSym + self.input_scale = input_scale + + +class MockLinearFp8Dynamic128(torch.nn.Module): + """Mock linear layer advertising kFp8Dynamic128Sym.""" + + def __init__(self): + super().__init__() + self.input_quant_key = kFp8Dynamic128Sym + + +class MockLinearNoQuant(torch.nn.Module): + """Mock linear layer with no input_quant_key (no fusion).""" + + pass + + +@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@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, + hidden_size: int, + dtype: torch.dtype, +) -> None: + """Test maybe_fused_act_quant with FP8 static per-tensor quantization.""" + device = "cuda:0" + torch.set_default_device(device) + + act_fn = SiluAndMul() + 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) + 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) + + ref_out = ref_impl(act_fn, x, scale) + 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("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, + hidden_size: int, + dtype: torch.dtype, +) -> None: + """Test maybe_fused_act_quant with FP8 dynamic per-block quantization.""" + group_size = 128 # We only support 128 for now + + device = "cuda:0" + torch.set_default_device(device) + + act_fn = SiluAndMul() + linear = MockLinearFp8Dynamic128() + + scale = 1 / hidden_size + x = torch.randn(num_tokens, 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) + + num_groups = hidden_size // group_size + assert result.scale.shape == (num_tokens, 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 + ) + + 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 + ) + result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave( + group_size, dim=1 + ) + torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2) + + +@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("hidden_size", [128, 512]) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_maybe_fused_act_quant_fallback( + default_vllm_config, + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, +) -> None: + """Test maybe_fused_act_quant falls back when no input_quant_key.""" + device = "cuda:0" + torch.set_default_device(device) + + act_fn = SiluAndMul() + linear = MockLinearNoQuant() + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) + + result = maybe_fused_act_quant(act_fn, x, linear) + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedActivation) + assert result.dtype == dtype + assert result.shape == (num_tokens, hidden_size) + + ref_out = act_fn(x) + torch.testing.assert_close(result, ref_out) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 7cc9d8c5144a..4fc7ace5a1fe 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -7,7 +7,11 @@ from vllm.config import CompilationMode, get_current_vllm_config from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + QuantKey, + kFp8StaticTensorSym, +) from vllm.platforms import current_platform from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel @@ -88,6 +92,11 @@ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | Non return False, "requires per tensor activation and weight scales." return True, None + def input_quant_key(self) -> QuantKey | None: + if self.config.activation_quant_key == kFp8StaticTensorSym: + return kFp8StaticTensorSym + return None + def apply_scaled_mm( self, *, diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py new file mode 100644 index 000000000000..534f94406a46 --- /dev/null +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Producer side of the QuantizedActivation contract for activation layers. + +Given an activation module and the downstream linear it feeds, fuse the +activation with that linear's input quantization into a single kernel when the +linear advertises a consumable input_quant_key (see quant_activation.py). +Falls back to the plain activation when nothing matches, so a model forward can +always call maybe_fused_act_quant unconditionally. + +This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion +fires here the silu_and_mul pattern is already consumed, so the compiler pass +finds nothing to rewrite and the two never double-fuse. +""" + +from collections.abc import Callable + +import torch + +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8StaticTensorSym, + kNvfp4Dynamic, +) +from vllm.platforms import current_platform + +FP8_DTYPE = current_platform.fp8_dtype() +FP4_DTYPE = torch.uint8 + + +def _silu_and_mul_fp8_static( + x: torch.Tensor, linear: LinearBase +) -> QuantizedActivation: + """SiluAndMul + FP8 static per-tensor quantization.""" + d = x.shape[-1] // 2 + out_shape = x.shape[:-1] + (d,) + result = torch.empty(out_shape, dtype=FP8_DTYPE, device=x.device) + # TODO(mgoin): read the consumer scale via the contract instead of reaching + # into the kernel-specific input_scale attribute. + scale = linear.input_scale + torch.ops._C.silu_and_mul_quant(result, x, scale) + return QuantizedActivation( + data=result, + scale=scale, + orig_dtype=x.dtype, + orig_shape=out_shape, + quant_key=kFp8StaticTensorSym, + ) + + +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}" + + d = x.shape[-1] // 2 + out_shape = x.shape[:-1] + (d,) + num_tokens = x.shape[0] + num_groups = d // group_size + + result = torch.empty((num_tokens, d), dtype=FP8_DTYPE, device=x.device) + scales = torch.empty((num_tokens, num_groups), dtype=torch.float32, device=x.device) + + torch.ops._C.silu_and_mul_per_block_quant( + out=result, + input=x, + scales=scales, + group_size=group_size, + scale_ub=None, + is_scale_transposed=False, + ) + + return QuantizedActivation( + data=result.view(out_shape), + scale=scales.view(out_shape[:-1] + (num_groups,)), + orig_dtype=x.dtype, + orig_shape=out_shape, + quant_key=quant_key, + ) + + +def _silu_and_mul_fp8_dynamic_128( + x: torch.Tensor, linear: LinearBase +) -> QuantizedActivation: + """SiluAndMul + FP8 dynamic per-group (group=128) quantization.""" + return _silu_and_mul_fp8_dynamic_block(x, linear, 128, kFp8Dynamic128Sym) + + +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}" + + d = x.shape[-1] // 2 + out_shape = x.shape[:-1] + (d,) + num_tokens = x.shape[0] + + # NVFP4 packs 2 values into 1 byte + result = torch.empty((num_tokens, d // 2), dtype=FP4_DTYPE, device=x.device) + + # Block scale output shape: swizzled layout for tensor cores + # Each group of 16 elements shares one FP8 scale + num_k_tiles = (d + 63) // 64 + block_scale = torch.empty( + (num_tokens, num_k_tiles * 4), dtype=FP8_DTYPE, device=x.device + ) + + input_global_scale = getattr(linear, "input_global_scale", None) + assert input_global_scale is not None, ( + "input_global_scale is required for NVFP4 quantization" + ) + + torch.ops._C.silu_and_mul_nvfp4_quant(result, block_scale, x, input_global_scale) + + return QuantizedActivation( + data=result.view(out_shape[:-1] + (d // 2,)), + scale=block_scale, + orig_dtype=x.dtype, + orig_shape=out_shape, + quant_key=kNvfp4Dynamic, + ) + + +# (activation module type, consumer input_quant_key) -> fused producer. +# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme. +_FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = { + (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static, +} + +# Add CUDA-specific entries for dynamic block quantization +if current_platform.is_cuda_alike(): + _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128 + +# Add NVFP4 if supported (requires SM100+) +if current_platform.is_cuda() and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant"): + _FUSED_ACT_QUANT[(SiluAndMul, kNvfp4Dynamic)] = _silu_and_mul_nvfp4_dynamic + + +def maybe_fused_act_quant( + act_fn: torch.nn.Module, + x: torch.Tensor, + linear: LinearBase, +) -> "torch.Tensor | QuantizedActivation": + """Apply act_fn, fusing the downstream linear's input quant when possible. + + Returns a QuantizedActivation when a fused kernel matches + (act_fn, linear.input_quant_key), else the plain activated tensor. + """ + key = getattr(linear, "input_quant_key", None) + if key is not None: + producer = _FUSED_ACT_QUANT.get((type(act_fn), key)) + if producer is not None: + return producer(x, linear) + return act_fn(x) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index ba2e3fef09fe..a595ee713e67 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -39,6 +39,7 @@ Attention, EncoderOnlyAttention, ) +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, @@ -114,7 +115,7 @@ def __init__( 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