| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | import vllm.config |
| 9 | import vllm.ir.ops |
| 10 | import vllm.plugins |
| 11 | from tests.compile.backend import TestBackend |
| 12 | from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor |
| 13 | from tests.utils import TestFP8Layer |
| 14 | from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops |
| 15 | from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant |
| 16 | from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS |
| 17 | from vllm.compilation.passes.fusion.rms_quant_fusion import ( |
| 18 | FUSED_OPS, |
| 19 | FusedRMSQuantKey, |
| 20 | RMSNormQuantFusionPass, |
| 21 | ) |
| 22 | from vllm.compilation.passes.fx_utils import find_auto_fn_maybe |
| 23 | from vllm.compilation.passes.utility.fix_functionalization import ( |
| 24 | FixFunctionalizationPass, |
| 25 | ) |
| 26 | from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass |
| 27 | from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass |
| 28 | from vllm.config import ( |
| 29 | CompilationConfig, |
| 30 | CompilationMode, |
| 31 | ModelConfig, |
| 32 | PassConfig, |
| 33 | VllmConfig, |
| 34 | ) |
| 35 | from vllm.model_executor.kernels.linear import ( |
| 36 | AiterFp8BlockScaledMMKernel, |
| 37 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 38 | CutlassFp8BlockScaledMMKernel, |
| 39 | CutlassFP8ScaledMMLinearKernel, |
| 40 | DeepGemmFp8BlockScaledMMKernel, |
| 41 | FlashInferFp8DeepGEMMDynamicBlockScaledKernel, |
| 42 | FlashInferFP8ScaledMMLinearKernel, |
| 43 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 44 | ROCmFP8ScaledMMLinearKernel, |
| 45 | RowWiseTorchFP8ScaledMMLinearKernel, |
| 46 | TritonFp8BlockScaledMMKernel, |
| 47 | _KernelT, |
| 48 | ) |
| 49 | from vllm.model_executor.layers.layernorm import RMSNorm, RMSNormGated |
| 50 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 51 | GroupShape, |
| 52 | create_fp8_quant_key, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 55 | cutlass_block_fp8_supported, |
| 56 | ) |
| 57 | from vllm.platforms import current_platform |
| 58 | from vllm.utils.deep_gemm import ( |
| 59 | is_deep_gemm_e8m0_used, |
| 60 | is_deep_gemm_supported, |
| 61 | ) |
| 62 | |
| 63 | # Kernel and group_shape combinations: (kernel, group_shape) |
| 64 | # CUDA kernels |
| 65 | CUDA_KERNEL_GROUPSHAPE_COMBINATIONS = [ |
| 66 | # FlashInferFP8ScaledMMLinearKernel supports both per-tensor only |
| 67 | (FlashInferFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR), |
| 68 | # CutlassFP8ScaledMMLinearKernel supports both per-tensor and per-token |
| 69 | (CutlassFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN), |
| 70 | (CutlassFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR), |
| 71 | # PerTensorTorchFP8ScaledMMLinearKernel only supports per-tensor |
| 72 | (PerTensorTorchFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR), |
| 73 | # ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token |
| 74 | (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN), |
| 75 | # Blockwise group shapes |
| 76 | (FlashInferFp8DeepGEMMDynamicBlockScaledKernel, GroupShape(1, 128)), |
| 77 | (CutlassFp8BlockScaledMMKernel, GroupShape(1, 128)), |
| 78 | (DeepGemmFp8BlockScaledMMKernel, GroupShape(1, 128)), |
| 79 | (TritonFp8BlockScaledMMKernel, GroupShape(1, 128)), |
| 80 | (TritonFp8BlockScaledMMKernel, GroupShape(1, 64)), |
| 81 | ] |
| 82 | |
| 83 | # ROCm kernels |
| 84 | ROCM_KERNEL_GROUPSHAPE_COMBINATIONS = [ |
| 85 | # ROCmFP8ScaledMMLinearKernel supports per-tensor only |
| 86 | (ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR), |
| 87 | # RowWiseTorchFP8ScaledMMLinearKernel only supports per-token |
| 88 | (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN), |
| 89 | # ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token |
| 90 | (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN), |
| 91 | # Blockwise group shapes (no kernel abstraction) |
| 92 | (TritonFp8BlockScaledMMKernel, GroupShape(1, 128)), |
| 93 | (TritonFp8BlockScaledMMKernel, GroupShape(1, 64)), |
| 94 | ] |
| 95 | |
| 96 | KERNEL_GROUPSHAPE_COMBINATIONS = ( |
| 97 | CUDA_KERNEL_GROUPSHAPE_COMBINATIONS |
| 98 | if current_platform.is_cuda() |
| 99 | else ROCM_KERNEL_GROUPSHAPE_COMBINATIONS |
| 100 | ) |
| 101 | |
| 102 | # For Aiter tests we toggle use_aiter_quant_op |
| 103 | AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [ |
| 104 | # Per-token with RowWiseTorchFP8ScaledMMLinearKernel |
| 105 | (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True), |
| 106 | (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False), |
| 107 | # Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel |
| 108 | (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True), |
| 109 | (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False), |
| 110 | # Blockwise |
| 111 | (AiterFp8BlockScaledMMKernel, GroupShape(1, 128), True), |
| 112 | ] |
| 113 | |
| 114 | |
| 115 | class TestModel(torch.nn.Module): |
| 116 | def __init__( |
| 117 | self, |
| 118 | hidden_size: int, |
| 119 | eps: float, |
| 120 | force_kernel: type[_KernelT] | None, |
| 121 | group_shape: GroupShape, |
| 122 | dtype: torch.dtype, |
| 123 | use_aiter_fusion: bool = False, |
| 124 | use_aiter_quant: bool = False, |
| 125 | *args, |
| 126 | **kwargs, |
| 127 | ): |
| 128 | super().__init__(*args, **kwargs) |
| 129 | self.fp8_linear_layers: list[torch.nn.Module] |
| 130 | self.group_shape = group_shape |
| 131 | self.use_aiter_quant_op = use_aiter_quant |
| 132 | self.use_aiter_fusion = use_aiter_fusion |
| 133 | self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)] |
| 134 | self.enable_rms_norm_custom_op = self.norm[0].enabled() |
| 135 | |
| 136 | # Determine if blockwise based on group_shape |
| 137 | is_blockwise = group_shape.is_per_group() |
| 138 | |
| 139 | if is_blockwise: |
| 140 | block_size = group_shape.col |
| 141 | self.activation_quant_key = create_fp8_quant_key( |
| 142 | static=False, group_shape=group_shape |
| 143 | ) |
| 144 | self.weight_quant_key = create_fp8_quant_key( |
| 145 | static=True, group_shape=GroupShape(block_size, block_size) |
| 146 | ) |
| 147 | |
| 148 | else: |
| 149 | is_static = group_shape == GroupShape.PER_TENSOR |
| 150 | self.activation_quant_key = create_fp8_quant_key( |
| 151 | is_static, group_shape=group_shape |
| 152 | ) |
| 153 | self.weight_quant_key = create_fp8_quant_key( |
| 154 | static=True, group_shape=group_shape |
| 155 | ) |
| 156 | |
| 157 | self.fp8_linear_layers = [ |
| 158 | TestFP8Layer( |
| 159 | weight_shape=(hidden_size, hidden_size), |
| 160 | activation_quant_key=self.activation_quant_key, |
| 161 | weight_quant_key=self.weight_quant_key, |
| 162 | force_kernel=force_kernel, |
| 163 | transpose_weights=use_aiter_fusion, |
| 164 | input_dtype=dtype, |
| 165 | ) |
| 166 | for _ in range(3) |
| 167 | ] |
| 168 | |
| 169 | # Enable aiter quantization if requested |
| 170 | for layer in self.fp8_linear_layers: |
| 171 | layer.kernel.quant_fp8.use_aiter = use_aiter_quant |
| 172 | |
| 173 | self.enable_quant_fp8_custom_op = self.fp8_linear_layers[ |
| 174 | 0 |
| 175 | ].is_quant_fp8_enabled() |
| 176 | |
| 177 | def forward(self, x): |
| 178 | # avoid having graph input be an arg to a pattern directly |
| 179 | x = resid = torch.relu(x) |
| 180 | y = self.norm[0](x) |
| 181 | |
| 182 | x2 = self.fp8_linear_layers[0](y) |
| 183 | # make sure resid is used for replacement to work |
| 184 | y2, resid = self.norm[1](x2, resid) |
| 185 | |
| 186 | x3 = self.fp8_linear_layers[1](y2) |
| 187 | |
| 188 | y3, resid = self.norm[2](x3, resid) # use resid here |
| 189 | |
| 190 | x4 = self.fp8_linear_layers[2](y3) |
| 191 | |
| 192 | y4, resid = self.norm[3](x4, resid) # use resid here |
| 193 | return y4 |
| 194 | |
| 195 | def ops_in_model_before(self): |
| 196 | if self.group_shape.is_per_group(): |
| 197 | # Blockwise path |
| 198 | if self.use_aiter_fusion and self.use_aiter_quant_op: |
| 199 | return [rocm_aiter_ops.get_group_quant_op()] |
| 200 | else: |
| 201 | if self.use_aiter_quant_op: |
| 202 | return [rocm_aiter_ops.get_per_token_quant_op()] |
| 203 | |
| 204 | # Common path |
| 205 | return ( |
| 206 | [QUANT_OPS[self.activation_quant_key]] |
| 207 | if self.enable_quant_fp8_custom_op |
| 208 | else [torch.ops.aten.reciprocal] |
| 209 | ) |
| 210 | |
| 211 | def ops_in_model_after(self): |
| 212 | if self.use_aiter_fusion: |
| 213 | if self.group_shape.is_per_group(): |
| 214 | # Blockwise aiter fusion |
| 215 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 216 | AiterFusedAddRMSFp8GroupQuantPattern, |
| 217 | AiterRMSFp8GroupQuantPattern, |
| 218 | ) |
| 219 | |
| 220 | return [ |
| 221 | AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP, |
| 222 | AiterRMSFp8GroupQuantPattern.FUSED_OP, |
| 223 | ] |
| 224 | else: |
| 225 | # Per-token aiter fusion |
| 226 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 227 | AiterFusedAddRMSNormDynamicQuantPattern, |
| 228 | AiterRMSNormDynamicQuantPattern, |
| 229 | ) |
| 230 | |
| 231 | return [ |
| 232 | AiterFusedAddRMSNormDynamicQuantPattern.FUSED_OP, |
| 233 | AiterRMSNormDynamicQuantPattern.FUSED_OP, |
| 234 | ] |
| 235 | |
| 236 | # Regular fusion |
| 237 | return [ |
| 238 | FUSED_OPS[FusedRMSQuantKey(self.activation_quant_key, True)], |
| 239 | FUSED_OPS[FusedRMSQuantKey(self.activation_quant_key, False)], |
| 240 | ] |
| 241 | |
| 242 | def ops_in_model_before_partial(self): |
| 243 | return [ |
| 244 | torch.ops.vllm_ir.rms_norm, |
| 245 | torch.ops.vllm_ir.fused_add_rms_norm.default, |
| 246 | ] |
| 247 | |
| 248 | |
| 249 | class AddRMSNormNvfp4Model(torch.nn.Module): |
| 250 | def __init__(self, hidden_size: int, eps: float): |
| 251 | super().__init__() |
| 252 | self.norm = RMSNorm(hidden_size, eps) |
| 253 | self.activation_scale = torch.rand((1, 1), dtype=torch.float32) |
| 254 | |
| 255 | weight = torch.rand((hidden_size, hidden_size)) |
| 256 | self.weight, self.weight_scale, weight_global_scale = quant_nvfp4_tensor(weight) |
| 257 | self.alpha = (1 / (weight_global_scale * self.activation_scale)).reshape(1) |
| 258 | |
| 259 | def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 260 | x = residual = torch.relu(x) |
| 261 | x, residual = self.norm(x, residual) |
| 262 | x, block_scale = scaled_fp4_quant(x, self.activation_scale) |
| 263 | x = cutlass_scaled_fp4_mm( |
| 264 | x, |
| 265 | self.weight, |
| 266 | block_scale, |
| 267 | self.weight_scale, |
| 268 | self.alpha, |
| 269 | out_dtype=torch.get_default_dtype(), |
| 270 | ) |
| 271 | return x, residual |
| 272 | |
| 273 | |
| 274 | def _run_fusion_test( |
| 275 | model, |
| 276 | fusion_pass, |
| 277 | vllm_config, |
| 278 | dtype, |
| 279 | hidden_size, |
| 280 | num_tokens, |
| 281 | ): |
| 282 | """Helper function for common fusion test logic. |
| 283 | |
| 284 | Must be called within vllm_config context. |
| 285 | """ |
| 286 | noop_pass = NoOpEliminationPass(vllm_config) |
| 287 | cleanup_pass = PostCleanupPass(vllm_config) |
| 288 | |
| 289 | backend = TestBackend(noop_pass, fusion_pass, cleanup_pass) |
| 290 | backend2 = TestBackend(noop_pass, cleanup_pass) |
| 291 | |
| 292 | x = torch.rand(num_tokens, hidden_size) |
| 293 | torch._dynamo.mark_dynamic(x, 0) |
| 294 | |
| 295 | model_fused = torch.compile(model, backend=backend) |
| 296 | result_fused = model_fused(x) |
| 297 | |
| 298 | model_unfused = torch.compile(model, backend=backend2) |
| 299 | result_unfused = model_unfused(x) |
| 300 | |
| 301 | if dtype == torch.float16: |
| 302 | ATOL, RTOL = (2e-3, 2e-3) |
| 303 | else: |
| 304 | ATOL, RTOL = (1e-2, 1e-2) |
| 305 | |
| 306 | torch.testing.assert_close(result_fused, result_unfused, atol=ATOL, rtol=RTOL) |
| 307 | |
| 308 | assert fusion_pass.matched_count == 3 |
| 309 | backend.check_before_ops(model.ops_in_model_before()) |
| 310 | backend.check_after_ops(model.ops_in_model_after()) |
| 311 | |
| 312 | return backend, backend2 |
| 313 | |
| 314 | |
| 315 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 316 | @pytest.mark.parametrize("hidden_size", [256]) |
| 317 | @pytest.mark.parametrize("num_tokens", [257]) |
| 318 | @pytest.mark.parametrize("eps", [1e-5, 1e-6]) |
| 319 | @pytest.mark.parametrize("kernel_groupshape", KERNEL_GROUPSHAPE_COMBINATIONS) |
| 320 | @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False]) |
| 321 | @pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False]) |
| 322 | @pytest.mark.skipif( |
| 323 | not current_platform.is_cuda_alike(), reason="Only test on CUDA and ROCm" |
| 324 | ) |
| 325 | def test_fusion_rmsnorm_quant( |
| 326 | dtype, |
| 327 | hidden_size, |
| 328 | num_tokens, |
| 329 | eps, |
| 330 | kernel_groupshape, |
| 331 | enable_rms_norm_custom_op, |
| 332 | enable_quant_fp8_custom_op, |
| 333 | ): |
| 334 | force_kernel, group_shape = kernel_groupshape |
| 335 | |
| 336 | if not enable_quant_fp8_custom_op and group_shape.is_per_group(): |
| 337 | pytest.skip("Unsupported unwrapped quant fp8 op for blockwise quantization") |
| 338 | |
| 339 | if group_shape == GroupShape(1, 64) and ( |
| 340 | cutlass_block_fp8_supported() or is_deep_gemm_supported() |
| 341 | ): |
| 342 | pytest.skip("Unsupported group shape 64 for CUTLASS/DeepGemm") |
| 343 | |
| 344 | # TODO(quant-rms-fusion): DeepGEMM UE8M0 activation quant on B200 lowers |
| 345 | # to a packed int32-scale op (per_token_group_quant_fp8_packed_for_deepgemm), |
| 346 | # but the rms+quant fusion pattern only matches the fp32-scale variant, so |
| 347 | # the fused output gets a mismatched scale layout and produces NaN. Only |
| 348 | # reproduces on bf16 (DeepGEMM UE8M0 on B200 is bf16-only). |
| 349 | # To re-enable: make rms_norm_per_block_quant emit packed UE8M0 scales |
| 350 | # and extend the fusion pattern to rewrite the packed activation quant. |
| 351 | deepgemm_kernels = ( |
| 352 | DeepGemmFp8BlockScaledMMKernel, |
| 353 | FlashInferFp8DeepGEMMDynamicBlockScaledKernel, |
| 354 | ) |
| 355 | if ( |
| 356 | dtype == torch.bfloat16 |
| 357 | and force_kernel in deepgemm_kernels |
| 358 | and is_deep_gemm_e8m0_used() |
| 359 | ): |
| 360 | pytest.skip( |
| 361 | "rms+quant fusion does not yet match the packed UE8M0 DeepGEMM path" |
| 362 | ) |
| 363 | |
| 364 | custom_ops = [] |
| 365 | if enable_rms_norm_custom_op: |
| 366 | custom_ops.append("+rms_norm") |
| 367 | if enable_quant_fp8_custom_op: |
| 368 | custom_ops.append("+quant_fp8") |
| 369 | |
| 370 | vllm_config = VllmConfig( |
| 371 | model_config=ModelConfig(dtype=dtype), |
| 372 | compilation_config=CompilationConfig( |
| 373 | mode=CompilationMode.VLLM_COMPILE, |
| 374 | custom_ops=custom_ops, |
| 375 | pass_config=PassConfig( |
| 376 | fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True |
| 377 | ), |
| 378 | ), |
| 379 | ) |
| 380 | |
| 381 | with ( |
| 382 | vllm.config.set_current_vllm_config(vllm_config), |
| 383 | vllm_config.kernel_config.ir_op_priority.set_priority(), |
| 384 | ): |
| 385 | # Setup device before model creation |
| 386 | torch.set_default_device("cuda") |
| 387 | torch.set_default_dtype(dtype) |
| 388 | torch.manual_seed(1) |
| 389 | |
| 390 | fusion_pass = RMSNormQuantFusionPass(vllm_config) |
| 391 | |
| 392 | model = TestModel( |
| 393 | hidden_size=hidden_size, |
| 394 | eps=eps, |
| 395 | force_kernel=force_kernel, |
| 396 | group_shape=group_shape, |
| 397 | dtype=dtype, |
| 398 | use_aiter_fusion=False, |
| 399 | use_aiter_quant=False, |
| 400 | ) |
| 401 | if any( |
| 402 | type(layer.kernel) is not force_kernel for layer in model.fp8_linear_layers |
| 403 | ): |
| 404 | pytest.skip(f"{force_kernel.__name__} is not supported on this platform") |
| 405 | |
| 406 | backend, _ = _run_fusion_test( |
| 407 | model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens |
| 408 | ) |
| 409 | backend.check_before_ops( |
| 410 | model.ops_in_model_before_partial(), fully_replaced=False |
| 411 | ) |
| 412 | |
| 413 | |
| 414 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 415 | @pytest.mark.parametrize("eps", [1e-5, 1e-6]) |
| 416 | @pytest.mark.parametrize( |
| 417 | ("num_tokens", "hidden_size"), |
| 418 | [(1, 256), (32, 1024), (257, 4096)], |
| 419 | ) |
| 420 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") |
| 421 | def test_fusion_add_rmsnorm_nvfp4_quant( |
| 422 | dtype: torch.dtype, eps: float, num_tokens: int, hidden_size: int |
| 423 | ): |
| 424 | fused_op = getattr( |
| 425 | torch.ops.vllm, |
| 426 | "flashinfer_fused_add_rms_norm_nvfp4_quant", |
| 427 | None, |
| 428 | ) |
| 429 | if not current_platform.has_device_capability(100) or fused_op is None: |
| 430 | pytest.skip("FlashInfer add-RMSNorm NVFP4 fusion is not available") |
| 431 | assert fused_op is not None |
| 432 | fused_op_default = fused_op.default |
| 433 | |
| 434 | vllm_config = VllmConfig( |
| 435 | model_config=ModelConfig(dtype=dtype), |
| 436 | compilation_config=CompilationConfig( |
| 437 | mode=CompilationMode.VLLM_COMPILE, |
| 438 | pass_config=PassConfig( |
| 439 | fuse_norm_quant=True, |
| 440 | eliminate_noops=True, |
| 441 | ), |
| 442 | ), |
| 443 | ) |
| 444 | |
| 445 | with ( |
| 446 | vllm.config.set_current_vllm_config(vllm_config), |
| 447 | vllm_config.kernel_config.ir_op_priority.set_priority(), |
| 448 | ): |
| 449 | torch.set_default_device("cuda") |
| 450 | torch.set_default_dtype(dtype) |
| 451 | torch.manual_seed(1) |
| 452 | |
| 453 | fusion_pass = RMSNormQuantFusionPass(vllm_config) |
| 454 | fused_backend = TestBackend( |
| 455 | NoOpEliminationPass(vllm_config), |
| 456 | fusion_pass, |
| 457 | PostCleanupPass(vllm_config), |
| 458 | FixFunctionalizationPass(vllm_config), |
| 459 | ) |
| 460 | unfused_backend = TestBackend( |
| 461 | NoOpEliminationPass(vllm_config), |
| 462 | PostCleanupPass(vllm_config), |
| 463 | FixFunctionalizationPass(vllm_config), |
| 464 | ) |
| 465 | |
| 466 | model = AddRMSNormNvfp4Model(hidden_size=hidden_size, eps=eps) |
| 467 | x = torch.rand(num_tokens, hidden_size) |
| 468 | torch._dynamo.mark_dynamic(x, 0) |
| 469 | |
| 470 | result_fused = torch.compile(model, backend=fused_backend)(x) |
| 471 | result_unfused = torch.compile(model, backend=unfused_backend)(x) |
| 472 | |
| 473 | torch.testing.assert_close(result_fused, result_unfused, atol=2e-1, rtol=2e-1) |
| 474 | assert fusion_pass.matched_count == 1 |
| 475 | fused_backend.check_before_ops([torch.ops._C.scaled_fp4_quant.out]) |
| 476 | fused_backend.check_after_ops([fused_op_default]) |
| 477 | assert ( |
| 478 | find_auto_fn_maybe( |
| 479 | fused_backend.graph_post_pass.nodes, |
| 480 | fused_op_default, |
| 481 | ) |
| 482 | is None |
| 483 | ) |
| 484 | |
| 485 | |
| 486 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 487 | @pytest.mark.parametrize("hidden_size", [256]) |
| 488 | @pytest.mark.parametrize("num_tokens", [257]) |
| 489 | @pytest.mark.parametrize("eps", [1e-5, 1e-6]) |
| 490 | @pytest.mark.parametrize( |
| 491 | "kernel_groupshape_quant", AITER_KERNEL_GROUPSHAPE_COMBINATIONS |
| 492 | ) |
| 493 | @pytest.mark.skipif( |
| 494 | (not current_platform.is_rocm() or not IS_AITER_FOUND), |
| 495 | reason="Only test on ROCm with aiter package installed", |
| 496 | ) |
| 497 | def test_aiter_fusion_rmsnorm_quant( |
| 498 | dtype: torch.dtype, |
| 499 | hidden_size: int, |
| 500 | num_tokens: int, |
| 501 | eps: float, |
| 502 | kernel_groupshape_quant: tuple, |
| 503 | monkeypatch: pytest.MonkeyPatch, |
| 504 | ): |
| 505 | force_kernel, group_shape, use_aiter_quant_op = kernel_groupshape_quant |
| 506 | vllm_config = VllmConfig( |
| 507 | model_config=ModelConfig(dtype=dtype), |
| 508 | compilation_config=CompilationConfig( |
| 509 | mode=CompilationMode.VLLM_COMPILE, |
| 510 | custom_ops=["+rms_norm", "+quant_fp8"], |
| 511 | pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True), |
| 512 | ), |
| 513 | ) |
| 514 | |
| 515 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 516 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 517 | RocmAiterRMSNormQuantFusionPass, |
| 518 | ) |
| 519 | |
| 520 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 521 | |
| 522 | rocm_aiter_ops.refresh_env_variables() |
| 523 | |
| 524 | torch.set_default_device("cuda") |
| 525 | torch.set_default_dtype(dtype) |
| 526 | torch.manual_seed(1) |
| 527 | |
| 528 | fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config) |
| 529 | |
| 530 | model = TestModel( |
| 531 | hidden_size=hidden_size, |
| 532 | eps=eps, |
| 533 | force_kernel=force_kernel, |
| 534 | group_shape=group_shape, |
| 535 | dtype=dtype, |
| 536 | use_aiter_fusion=True, # Always use aiter fusion ops in aiter test |
| 537 | use_aiter_quant=use_aiter_quant_op, # Toggle aiter quantization |
| 538 | ) |
| 539 | |
| 540 | _run_fusion_test( |
| 541 | model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens |
| 542 | ) |
| 543 | |
| 544 | |
| 545 | class TestGatedModel(torch.nn.Module): |
| 546 | """Model that uses RMSNormGated + reshape + group FP8 quant + linear. |
| 547 | |
| 548 | Mimics GatedDeltaNetAttention's output projection path where: |
| 549 | - RMSNormGated operates on per-head tensors, either (N, H, D) as |
| 550 | QwenGatedDeltaNetAttention passes them or flattened to (N*H, D) as |
| 551 | OlmoHybridGatedDeltaNetAttention does |
| 552 | - Output is reshaped to (N, H*D) before group quantization + linear |
| 553 | """ |
| 554 | |
| 555 | def __init__( |
| 556 | self, |
| 557 | num_heads: int, |
| 558 | head_dim: int, |
| 559 | eps: float, |
| 560 | force_kernel: type[_KernelT], |
| 561 | group_shape: GroupShape, |
| 562 | dtype: torch.dtype, |
| 563 | use_aiter_quant: bool = True, |
| 564 | flatten_heads: bool = True, |
| 565 | ): |
| 566 | super().__init__() |
| 567 | self.num_heads = num_heads |
| 568 | self.head_dim = head_dim |
| 569 | self.flatten_heads = flatten_heads |
| 570 | hidden_dim = num_heads * head_dim |
| 571 | |
| 572 | self.norm = RMSNormGated( |
| 573 | head_dim, |
| 574 | eps=eps, |
| 575 | group_size=None, |
| 576 | norm_before_gate=True, |
| 577 | ) |
| 578 | |
| 579 | self.activation_quant_key = create_fp8_quant_key( |
| 580 | static=False, group_shape=group_shape |
| 581 | ) |
| 582 | self.weight_quant_key = create_fp8_quant_key( |
| 583 | static=True, group_shape=GroupShape(group_shape.col, group_shape.col) |
| 584 | ) |
| 585 | |
| 586 | self.fp8_linear = TestFP8Layer( |
| 587 | weight_shape=(hidden_dim, hidden_dim), |
| 588 | activation_quant_key=self.activation_quant_key, |
| 589 | weight_quant_key=self.weight_quant_key, |
| 590 | force_kernel=force_kernel, |
| 591 | transpose_weights=True, |
| 592 | input_dtype=dtype, |
| 593 | ) |
| 594 | self.fp8_linear.kernel.quant_fp8.use_aiter = use_aiter_quant |
| 595 | |
| 596 | def forward(self, x, z): |
| 597 | num_heads = self.num_heads |
| 598 | head_dim = self.head_dim |
| 599 | hidden_dim = num_heads * head_dim |
| 600 | x = torch.relu(x) |
| 601 | z = torch.relu(z) |
| 602 | x_heads = x.reshape(-1, num_heads, head_dim) |
| 603 | z_heads = z.reshape(-1, num_heads, head_dim) |
| 604 | if self.flatten_heads: |
| 605 | x_heads = x_heads.reshape(-1, head_dim) |
| 606 | z_heads = z_heads.reshape(-1, head_dim) |
| 607 | normed = self.norm(x_heads, z_heads) |
| 608 | merged = normed.reshape(-1, hidden_dim) |
| 609 | out = self.fp8_linear(merged) |
| 610 | return out |
| 611 | |
| 612 | def ops_in_model_after(self): |
| 613 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 614 | AiterRMSNormGatedFp8GroupQuantPattern, |
| 615 | ) |
| 616 | |
| 617 | return [AiterRMSNormGatedFp8GroupQuantPattern.FUSED_OP] |
| 618 | |
| 619 | |
| 620 | class _MockGDNLayer: |
| 621 | """Minimal mock to populate static_forward_context for pass discovery. |
| 622 | |
| 623 | Uses __class__ assignment to pass isinstance checks against |
| 624 | GatedDeltaNetAttention without requiring a full config-based init. |
| 625 | """ |
| 626 | |
| 627 | def __init__(self, num_v_heads: int, head_v_dim: int, tp_size: int = 1): |
| 628 | self.num_v_heads = num_v_heads |
| 629 | self.head_v_dim = head_v_dim |
| 630 | self.tp_size = tp_size |
| 631 | |
| 632 | from vllm.model_executor.layers.mamba.gdn.base import ( |
| 633 | GatedDeltaNetAttention, |
| 634 | ) |
| 635 | |
| 636 | self.__class__ = GatedDeltaNetAttention |
| 637 | |
| 638 | |
| 639 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 640 | @pytest.mark.parametrize("num_heads", [2]) |
| 641 | @pytest.mark.parametrize("head_dim", [128]) |
| 642 | @pytest.mark.parametrize("num_tokens", [8]) |
| 643 | @pytest.mark.parametrize("eps", [1e-5, 1e-6]) |
| 644 | @pytest.mark.parametrize("flatten_heads", [True, False]) |
| 645 | @pytest.mark.parametrize("quant_fp8_op", ["-quant_fp8", "+quant_fp8"]) |
| 646 | @pytest.mark.skipif( |
| 647 | (not current_platform.is_rocm() or not IS_AITER_FOUND), |
| 648 | reason="Only test on ROCm with aiter package installed", |
| 649 | ) |
| 650 | def test_aiter_fusion_rmsnorm_gated_quant( |
| 651 | dtype: torch.dtype, |
| 652 | num_heads: int, |
| 653 | head_dim: int, |
| 654 | num_tokens: int, |
| 655 | eps: float, |
| 656 | flatten_heads: bool, |
| 657 | quant_fp8_op: str, |
| 658 | monkeypatch: pytest.MonkeyPatch, |
| 659 | ): |
| 660 | group_shape = GroupShape(1, 128) |
| 661 | vllm_config = VllmConfig( |
| 662 | model_config=ModelConfig(dtype=dtype), |
| 663 | compilation_config=CompilationConfig( |
| 664 | mode=CompilationMode.VLLM_COMPILE, |
| 665 | custom_ops=["-rms_norm", "-silu_and_mul", quant_fp8_op], |
| 666 | pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True), |
| 667 | ), |
| 668 | ) |
| 669 | |
| 670 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 671 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 672 | RocmAiterRMSNormQuantFusionPass, |
| 673 | ) |
| 674 | |
| 675 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 676 | rocm_aiter_ops.refresh_env_variables() |
| 677 | |
| 678 | # Register a mock GDN layer so the pass discovers num_heads/head_dim |
| 679 | mock_gdn = _MockGDNLayer(num_v_heads=num_heads, head_v_dim=head_dim, tp_size=1) |
| 680 | vllm_config.compilation_config.static_forward_context["mock_gdn_layer"] = ( |
| 681 | mock_gdn |
| 682 | ) |
| 683 | |
| 684 | torch.set_default_device("cuda") |
| 685 | torch.set_default_dtype(dtype) |
| 686 | torch.manual_seed(1) |
| 687 | |
| 688 | fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config) |
| 689 | |
| 690 | model = TestGatedModel( |
| 691 | num_heads=num_heads, |
| 692 | head_dim=head_dim, |
| 693 | eps=eps, |
| 694 | force_kernel=AiterFp8BlockScaledMMKernel, |
| 695 | group_shape=group_shape, |
| 696 | dtype=dtype, |
| 697 | use_aiter_quant=True, |
| 698 | flatten_heads=flatten_heads, |
| 699 | ) |
| 700 | |
| 701 | noop_pass = NoOpEliminationPass(vllm_config) |
| 702 | cleanup_pass = PostCleanupPass(vllm_config) |
| 703 | |
| 704 | backend = TestBackend(noop_pass, fusion_pass, cleanup_pass) |
| 705 | backend2 = TestBackend(noop_pass, cleanup_pass) |
| 706 | |
| 707 | hidden_dim = num_heads * head_dim |
| 708 | x = torch.rand(num_tokens, hidden_dim) |
| 709 | z = torch.rand(num_tokens, hidden_dim) |
| 710 | torch._dynamo.mark_dynamic(x, 0) |
| 711 | torch._dynamo.mark_dynamic(z, 0) |
| 712 | |
| 713 | model_fused = torch.compile(model, backend=backend) |
| 714 | result_fused = model_fused(x, z) |
| 715 | |
| 716 | model_unfused = torch.compile(model, backend=backend2) |
| 717 | result_unfused = model_unfused(x, z) |
| 718 | |
| 719 | # The opaque AITER quant op preserves a BF16 intermediate in the |
| 720 | # unfused path, while the fused kernel normalizes and quantizes |
| 721 | # directly. Allow for that extra rounding before the FP8 GEMM. |
| 722 | tol = 2e-2 if quant_fp8_op == "+quant_fp8" else 1e-2 |
| 723 | torch.testing.assert_close(result_fused, result_unfused, atol=tol, rtol=tol) |
| 724 | |
| 725 | assert fusion_pass.matched_count == 1 |
| 726 | backend.check_after_ops(model.ops_in_model_after()) |
| 727 | |
| 728 | |
| 729 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 730 | @pytest.mark.parametrize("num_heads", [2]) |
| 731 | @pytest.mark.parametrize("head_dim", [128]) |
| 732 | @pytest.mark.parametrize("num_tokens", [8]) |
| 733 | @pytest.mark.parametrize("eps", [1e-6]) |
| 734 | @pytest.mark.skipif( |
| 735 | (not current_platform.is_rocm() or not IS_AITER_FOUND), |
| 736 | reason="Only test on ROCm with aiter package installed", |
| 737 | ) |
| 738 | def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers( |
| 739 | dtype: torch.dtype, |
| 740 | num_heads: int, |
| 741 | head_dim: int, |
| 742 | num_tokens: int, |
| 743 | eps: float, |
| 744 | monkeypatch: pytest.MonkeyPatch, |
| 745 | ): |
| 746 | """Verify that without GDN layers in static_forward_context, |
| 747 | the gated pattern is not registered and no matches occur.""" |
| 748 | group_shape = GroupShape(1, 128) |
| 749 | vllm_config = VllmConfig( |
| 750 | model_config=ModelConfig(dtype=dtype), |
| 751 | compilation_config=CompilationConfig( |
| 752 | mode=CompilationMode.VLLM_COMPILE, |
| 753 | custom_ops=["-rms_norm", "-silu_and_mul", "-quant_fp8"], |
| 754 | pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True), |
| 755 | ), |
| 756 | ) |
| 757 | |
| 758 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 759 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 760 | RocmAiterRMSNormQuantFusionPass, |
| 761 | ) |
| 762 | |
| 763 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 764 | rocm_aiter_ops.refresh_env_variables() |
| 765 | |
| 766 | torch.set_default_device("cuda") |
| 767 | torch.set_default_dtype(dtype) |
| 768 | torch.manual_seed(1) |
| 769 | |
| 770 | # No mock GDN layer registered -- pass should not register gated pattern |
| 771 | fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config) |
| 772 | |
| 773 | model = TestGatedModel( |
| 774 | num_heads=num_heads, |
| 775 | head_dim=head_dim, |
| 776 | eps=eps, |
| 777 | force_kernel=AiterFp8BlockScaledMMKernel, |
| 778 | group_shape=group_shape, |
| 779 | dtype=dtype, |
| 780 | use_aiter_quant=True, |
| 781 | ) |
| 782 | |
| 783 | noop_pass = NoOpEliminationPass(vllm_config) |
| 784 | cleanup_pass = PostCleanupPass(vllm_config) |
| 785 | |
| 786 | backend = TestBackend(noop_pass, fusion_pass, cleanup_pass) |
| 787 | |
| 788 | hidden_dim = num_heads * head_dim |
| 789 | x = torch.rand(num_tokens, hidden_dim) |
| 790 | z = torch.rand(num_tokens, hidden_dim) |
| 791 | torch._dynamo.mark_dynamic(x, 0) |
| 792 | torch._dynamo.mark_dynamic(z, 0) |
| 793 | |
| 794 | model_fused = torch.compile(model, backend=backend) |
| 795 | model_fused(x, z) |
| 796 | |
| 797 | assert fusion_pass.matched_count == 0 |