| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for manual fusion via maybe_fused_act_quant. |
| 4 | |
| 5 | Tests all fusion paths in _FUSED_ACT_QUANT: |
| 6 | - kFp8StaticTensorSym: all platforms |
| 7 | - kFp8Dynamic128Sym: CUDA only |
| 8 | - kNvfp4Dynamic: CUDA SM100+ only |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | import vllm.envs as envs |
| 15 | from tests.utils import TestFP8Layer |
| 16 | from vllm.config import ( |
| 17 | CompilationConfig, |
| 18 | VllmConfig, |
| 19 | set_current_vllm_config, |
| 20 | ) |
| 21 | from vllm.model_executor.kernels.linear import ( |
| 22 | CutlassFP8ScaledMMLinearKernel, |
| 23 | FlashInferFP8ScaledMMLinearKernel, |
| 24 | FP8ScaledMMLinearKernel, |
| 25 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 26 | ROCmFP8ScaledMMLinearKernel, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.activation import SiluAndMul |
| 29 | from vllm.model_executor.layers.fusion.fused_act_quant import ( |
| 30 | _FUSED_ACT_QUANT, |
| 31 | maybe_fused_act_quant, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fusion.quant_activation import ( |
| 34 | QuantizedActivation, |
| 35 | expose_input_quant_key, |
| 36 | get_input_quant_key, |
| 37 | ) |
| 38 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 39 | kFp8Dynamic128Sym, |
| 40 | kFp8StaticTensorSym, |
| 41 | kNvfp4Dynamic, |
| 42 | ) |
| 43 | from vllm.platforms import current_platform |
| 44 | |
| 45 | |
| 46 | # Mock linear layer for testing fusion paths that don't have real kernel support |
| 47 | class MockLinearForFusion(torch.nn.Module): |
| 48 | """Mock linear layer that exposes an input key for fusion testing.""" |
| 49 | |
| 50 | def __init__( |
| 51 | self, |
| 52 | quant_key, |
| 53 | input_scale=None, |
| 54 | input_global_scale=None, |
| 55 | input_global_scale_inv=None, |
| 56 | ): |
| 57 | super().__init__() |
| 58 | self._input_quant_key = quant_key |
| 59 | if input_scale is not None: |
| 60 | self.input_scale = input_scale |
| 61 | if input_global_scale is not None: |
| 62 | self.input_global_scale = input_global_scale |
| 63 | if input_global_scale_inv is not None: |
| 64 | self.input_global_scale_inv = input_global_scale_inv |
| 65 | |
| 66 | |
| 67 | ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel] |
| 68 | CUDA_KERNELS = [ |
| 69 | FlashInferFP8ScaledMMLinearKernel, |
| 70 | CutlassFP8ScaledMMLinearKernel, |
| 71 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 72 | ] |
| 73 | TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS |
| 74 | |
| 75 | |
| 76 | @pytest.mark.parametrize("num_tokens", [32, 64]) |
| 77 | @pytest.mark.parametrize("hidden_size", [128, 256]) |
| 78 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 79 | @pytest.mark.parametrize("force_kernel", TEST_KERNELS) |
| 80 | @pytest.mark.skipif( |
| 81 | envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" |
| 82 | ) |
| 83 | def test_manual_fusion_fp8_static_with_linear( |
| 84 | num_tokens: int, |
| 85 | hidden_size: int, |
| 86 | dtype: torch.dtype, |
| 87 | force_kernel: type[FP8ScaledMMLinearKernel], |
| 88 | ): |
| 89 | """Test manual fusion with real FP8 linear layer (kFp8StaticTensorSym). |
| 90 | |
| 91 | This is an end-to-end test that verifies the full flow: |
| 92 | unfused (silu_and_mul -> in-kernel quant) vs fused (silu_and_mul_quant). |
| 93 | """ |
| 94 | torch.set_default_device("cuda") |
| 95 | torch.set_default_dtype(dtype) |
| 96 | |
| 97 | x = torch.rand(num_tokens, hidden_size * 2) |
| 98 | |
| 99 | config = VllmConfig( |
| 100 | compilation_config=CompilationConfig(custom_ops=["none"]), |
| 101 | ) |
| 102 | |
| 103 | with set_current_vllm_config(config): |
| 104 | silu_and_mul = SiluAndMul() |
| 105 | fp8_linear = TestFP8Layer( |
| 106 | weight_shape=(hidden_size, hidden_size), |
| 107 | activation_quant_key=kFp8StaticTensorSym, |
| 108 | weight_quant_key=kFp8StaticTensorSym, |
| 109 | force_kernel=force_kernel, |
| 110 | input_dtype=dtype, |
| 111 | ) |
| 112 | |
| 113 | # Run without fusion: silu_and_mul returns plain tensor |
| 114 | y_unfused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear) |
| 115 | assert isinstance(y_unfused, torch.Tensor) |
| 116 | result_unfused = fp8_linear(y_unfused) |
| 117 | |
| 118 | # Enable fusion |
| 119 | expose_input_quant_key(fp8_linear, fp8_linear.kernel) |
| 120 | |
| 121 | if get_input_quant_key(fp8_linear) is None: |
| 122 | pytest.skip( |
| 123 | f"Kernel {force_kernel.__name__} doesn't support input_quant_key" |
| 124 | ) |
| 125 | |
| 126 | # Run with fusion: silu_and_mul returns QuantizedActivation |
| 127 | y_fused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear) |
| 128 | assert isinstance(y_fused, QuantizedActivation) |
| 129 | assert y_fused.quant_key == kFp8StaticTensorSym |
| 130 | assert y_fused.data.dtype == current_platform.fp8_dtype() |
| 131 | assert y_fused.data.shape == (num_tokens, hidden_size) |
| 132 | assert y_fused.orig_dtype == dtype |
| 133 | assert y_fused.orig_shape == (num_tokens, hidden_size) |
| 134 | |
| 135 | result_fused = fp8_linear(y_fused) |
| 136 | |
| 137 | torch.testing.assert_close( |
| 138 | result_fused.to(dtype=dtype), |
| 139 | result_unfused.to(dtype=dtype), |
| 140 | atol=5e-2, |
| 141 | rtol=5e-2, |
| 142 | ) |
| 143 | |
| 144 | |
| 145 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 146 | @pytest.mark.skipif( |
| 147 | not current_platform.is_cuda_alike(), reason="Dynamic block quant CUDA only" |
| 148 | ) |
| 149 | @pytest.mark.skipif( |
| 150 | envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" |
| 151 | ) |
| 152 | def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype): |
| 153 | """Test kFp8Dynamic128Sym fusion path (group_size=128). |
| 154 | |
| 155 | Compares fused (silu_and_mul_per_block_quant) vs unfused (silu_and_mul) |
| 156 | by dequantizing the fused result and comparing with unfused. |
| 157 | """ |
| 158 | if (SiluAndMul, kFp8Dynamic128Sym) not in _FUSED_ACT_QUANT: |
| 159 | pytest.skip("kFp8Dynamic128Sym fusion not available") |
| 160 | |
| 161 | torch.set_default_device("cuda") |
| 162 | torch.set_default_dtype(dtype) |
| 163 | |
| 164 | # hidden_size must be divisible by group_size (128) |
| 165 | num_tokens, hidden_size = 32, 256 |
| 166 | group_size = 128 |
| 167 | x = torch.rand(num_tokens, hidden_size * 2) |
| 168 | |
| 169 | config = VllmConfig( |
| 170 | compilation_config=CompilationConfig(custom_ops=["none"]), |
| 171 | ) |
| 172 | |
| 173 | with set_current_vllm_config(config): |
| 174 | silu_and_mul = SiluAndMul() |
| 175 | |
| 176 | # Unfused path: just apply silu_and_mul |
| 177 | mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size) |
| 178 | result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key) |
| 179 | assert isinstance(result_unfused, torch.Tensor) |
| 180 | |
| 181 | # Fused path: apply silu_and_mul + per-block quantization |
| 182 | mock_linear_with_key = MockLinearForFusion(kFp8Dynamic128Sym) |
| 183 | result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key) |
| 184 | |
| 185 | # Verify fused result structure |
| 186 | assert isinstance(result_fused, QuantizedActivation) |
| 187 | assert result_fused.quant_key == kFp8Dynamic128Sym |
| 188 | assert result_fused.data.dtype == current_platform.fp8_dtype() |
| 189 | assert result_fused.data.shape == (num_tokens, hidden_size) |
| 190 | assert result_fused.orig_dtype == dtype |
| 191 | |
| 192 | # Check scale shape |
| 193 | expected_num_groups = hidden_size // group_size |
| 194 | assert result_fused.scale.shape == (num_tokens, expected_num_groups) |
| 195 | |
| 196 | # Dequantize fused result and compare with unfused |
| 197 | # Per-block dequant: data * scale (broadcast scale across group) |
| 198 | dequant_data = result_fused.data.to(dtype).view( |
| 199 | num_tokens, expected_num_groups, group_size |
| 200 | ) |
| 201 | scales_expanded = result_fused.scale.unsqueeze( |
| 202 | -1 |
| 203 | ) # (num_tokens, num_groups, 1) |
| 204 | dequant_result = ( |
| 205 | (dequant_data * scales_expanded).view(num_tokens, hidden_size).to(dtype) |
| 206 | ) |
| 207 | |
| 208 | torch.testing.assert_close( |
| 209 | dequant_result, |
| 210 | result_unfused, |
| 211 | atol=5e-2, |
| 212 | rtol=5e-2, |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 217 | # Cover M < 128 (padded swizzled SF tiles) and hidden_size that makes the |
| 218 | # kernel launch grid.y > 1 (num_packed_cols = hidden_size / 16 > 512). |
| 219 | @pytest.mark.parametrize("num_tokens", [1, 127, 128]) |
| 220 | @pytest.mark.parametrize("hidden_size", [256, 14336]) |
| 221 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="NVFP4 CUDA only") |
| 222 | @pytest.mark.skipif( |
| 223 | not current_platform.has_device_capability(100), reason="NVFP4 requires SM100+" |
| 224 | ) |
| 225 | @pytest.mark.skipif( |
| 226 | envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" |
| 227 | ) |
| 228 | def test_manual_fusion_nvfp4_dynamic( |
| 229 | dtype: torch.dtype, num_tokens: int, hidden_size: int |
| 230 | ): |
| 231 | """Test kNvfp4Dynamic fusion path. |
| 232 | |
| 233 | Compares fused (silu_and_mul_nvfp4_quant) vs unfused (silu_and_mul) |
| 234 | by dequantizing the fused result and comparing with unfused. |
| 235 | """ |
| 236 | if (SiluAndMul, kNvfp4Dynamic) not in _FUSED_ACT_QUANT: |
| 237 | pytest.skip("kNvfp4Dynamic fusion not available") |
| 238 | |
| 239 | from tests.kernels.quantization.nvfp4_utils import dequantize_nvfp4_to_dtype |
| 240 | |
| 241 | torch.set_default_device("cuda") |
| 242 | torch.set_default_dtype(dtype) |
| 243 | |
| 244 | # NVFP4 requires hidden_size divisible by 16 (block size) and by 2 (packing). |
| 245 | x = torch.rand(num_tokens, hidden_size * 2) |
| 246 | # Non-1.0 global scale so the test is sensitive to the scale direction: |
| 247 | # the fused producer must quantize with input_global_scale_inv (the GEMM's |
| 248 | # alpha divides by input_global_scale), matching the unfused path. |
| 249 | input_global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda") |
| 250 | input_global_scale_inv = 1.0 / input_global_scale |
| 251 | |
| 252 | config = VllmConfig( |
| 253 | compilation_config=CompilationConfig(custom_ops=["none"]), |
| 254 | ) |
| 255 | |
| 256 | with set_current_vllm_config(config): |
| 257 | silu_and_mul = SiluAndMul() |
| 258 | |
| 259 | # Unfused path: just apply silu_and_mul |
| 260 | mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size) |
| 261 | result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key) |
| 262 | assert isinstance(result_unfused, torch.Tensor) |
| 263 | |
| 264 | # Fused path: apply silu_and_mul + NVFP4 quantization |
| 265 | mock_linear_with_key = MockLinearForFusion( |
| 266 | kNvfp4Dynamic, |
| 267 | input_global_scale=input_global_scale, |
| 268 | input_global_scale_inv=input_global_scale_inv, |
| 269 | ) |
| 270 | result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key) |
| 271 | |
| 272 | # Verify fused result structure |
| 273 | assert isinstance(result_fused, QuantizedActivation) |
| 274 | assert result_fused.quant_key == kNvfp4Dynamic |
| 275 | # NVFP4 packs 2 values into 1 byte |
| 276 | assert result_fused.data.dtype == torch.uint8 |
| 277 | assert result_fused.data.shape == (num_tokens, hidden_size // 2) |
| 278 | assert result_fused.orig_dtype == dtype |
| 279 | assert result_fused.orig_shape == (num_tokens, hidden_size) |
| 280 | |
| 281 | # The scale tensor must use the padded 128x4-tile swizzled layout, |
| 282 | # same as scaled_fp4_quant's allocator; the kernel writes a full |
| 283 | # 128-row tile even when M is smaller. |
| 284 | from vllm.utils.math_utils import round_up |
| 285 | |
| 286 | assert result_fused.scale.dtype == current_platform.fp8_dtype() |
| 287 | assert result_fused.scale.shape == ( |
| 288 | round_up(num_tokens, 128), |
| 289 | round_up(hidden_size // 16, 4), |
| 290 | ) |
| 291 | |
| 292 | # Dequantize fused result and compare with unfused. The block scales |
| 293 | # must have been folded with input_global_scale_inv, so dequantize |
| 294 | # with the same inverse scale to recover the activation. |
| 295 | dequant_result = dequantize_nvfp4_to_dtype( |
| 296 | tensor_fp4=result_fused.data, |
| 297 | tensor_sf=result_fused.scale, |
| 298 | global_scale=input_global_scale_inv, |
| 299 | dtype=dtype, |
| 300 | device="cuda", |
| 301 | block_size=16, |
| 302 | is_sf_128x4_layout=True, |
| 303 | ) |
| 304 | |
| 305 | torch.testing.assert_close( |
| 306 | dequant_result, |
| 307 | result_unfused, |
| 308 | atol=3e-1, |
| 309 | rtol=3e-1, |
| 310 | ) |
| 311 | |
| 312 | |
| 313 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 314 | @pytest.mark.skipif( |
| 315 | envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" |
| 316 | ) |
| 317 | def test_manual_fusion_fallback_no_key(dtype: torch.dtype): |
| 318 | """Test that maybe_fused_act_quant falls back without an input key.""" |
| 319 | torch.set_default_device("cuda") |
| 320 | torch.set_default_dtype(dtype) |
| 321 | |
| 322 | x = torch.rand(32, 256) |
| 323 | |
| 324 | config = VllmConfig( |
| 325 | compilation_config=CompilationConfig(custom_ops=["none"]), |
| 326 | ) |
| 327 | |
| 328 | with set_current_vllm_config(config): |
| 329 | silu_and_mul = SiluAndMul() |
| 330 | # Linear without an input quantization key |
| 331 | mock_linear = torch.nn.Linear(128, 128) |
| 332 | |
| 333 | result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) |
| 334 | |
| 335 | # Should fall back to plain silu_and_mul |
| 336 | assert isinstance(result, torch.Tensor) |
| 337 | assert result.shape == (32, 128) |
| 338 | assert result.dtype == dtype |
| 339 | |
| 340 | |
| 341 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 342 | @pytest.mark.skipif( |
| 343 | envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" |
| 344 | ) |
| 345 | def test_manual_fusion_fallback_unsupported_key(dtype: torch.dtype): |
| 346 | """Test that maybe_fused_act_quant falls back for unsupported quant keys.""" |
| 347 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 348 | GroupShape, |
| 349 | QuantKey, |
| 350 | ScaleDesc, |
| 351 | ) |
| 352 | |
| 353 | torch.set_default_device("cuda") |
| 354 | torch.set_default_dtype(dtype) |
| 355 | |
| 356 | x = torch.rand(32, 256) |
| 357 | |
| 358 | config = VllmConfig( |
| 359 | compilation_config=CompilationConfig(custom_ops=["none"]), |
| 360 | ) |
| 361 | |
| 362 | with set_current_vllm_config(config): |
| 363 | # Create an unsupported quant key |
| 364 | unsupported_key = QuantKey( |
| 365 | dtype=torch.int8, |
| 366 | scale=ScaleDesc( |
| 367 | dtype=torch.float32, static=False, group_shape=GroupShape(1, 1) |
| 368 | ), |
| 369 | ) |
| 370 | |
| 371 | silu_and_mul = SiluAndMul() |
| 372 | mock_linear = MockLinearForFusion(unsupported_key) |
| 373 | |
| 374 | result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) |
| 375 | |
| 376 | # Should fall back to plain silu_and_mul since key not in _FUSED_ACT_QUANT |
| 377 | assert isinstance(result, torch.Tensor) |
| 378 | assert result.shape == (32, 128) |
| 379 | assert result.dtype == dtype |