diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index ec9bb6110d..7e346cc767 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -24,6 +24,11 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( CutlassFP8ScaledMMLinearKernel, ) @@ -56,6 +61,9 @@ SUPPORTING = { FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, } diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 9f00adf783..ced034cb1a 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -21,9 +21,18 @@ from packaging import version from tests.quantization.utils import load_model_without_vllm_runner from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops -from vllm.config import set_current_vllm_config +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.cache import CacheConfig from vllm.forward_context import set_forward_context +from vllm.model_executor import parameter +from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoeWeightScaleSupported, @@ -113,6 +122,12 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse( AITER_AVAILABLE = is_aiter_found_and_supported() +AITER_PTPC_KERNELS = ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, +) + DEVICE_TYPE = current_platform.device_type @@ -1095,6 +1110,54 @@ def test_quark_w8a8_fp8_moe_per_block_rejects_misaligned_partition(): ) +def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch): + """QuarkW8A8Fp8 must advertise the key its kernel consumes pre-quantized. + + Only kernel selection runs, no GEMM. The shape matters: AITER PTPC + kernels decline untuned shapes, leaving a torch kernel that has no key. + """ + # Llama-3.1-70B qkv_proj at TP1 (aiter-tuned N, K on gfx950). + N, K = 10240, 8192 + dtype = torch.bfloat16 + kernel_config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticChannelSym, + activation_quant_key=kFp8DynamicTokenSym, + weight_shape=(N, K), + input_dtype=dtype, + out_dtype=dtype, + ) + if not any( + cls.is_supported()[0] and cls.can_implement(kernel_config)[0] + for cls in AITER_PTPC_KERNELS + ): + pytest.skip("no AITER PTPC kernel is usable for this shape and environment") + + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + + scheme = QuarkW8A8Fp8.__new__(QuarkW8A8Fp8) + scheme.weight_qscheme = "per_channel" + scheme.is_static_input_scheme = False + scheme.input_qscheme = "per_channel" + scheme.activation_quant_key = kFp8DynamicTokenSym + scheme.weight_quant_key = kFp8StaticChannelSym + scheme.out_dtype = dtype + scheme.input_dtype = dtype + + layer = torch.nn.Module() + with set_current_vllm_config(VllmConfig()): + scheme.create_weights( + layer, + output_partition_sizes=[N], + input_size_per_partition=K, + params_dtype=dtype, + weight_loader=lambda *args, **kwargs: None, + ) + + assert isinstance(scheme.fp8_linear, AITER_PTPC_KERNELS) + assert layer.input_quant_key == kFp8DynamicTokenSym + + @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) def test_quark_fp8_w_per_tensor_a_per_tensor( kv_cache_dtype: str, monkeypatch, dist_init, workspace_init diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 2408f4ad6c..a57e7510de 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -16,6 +16,8 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, + kFp8DynamicTokenSym, ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -150,6 +152,12 @@ class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): return False, "requires aiter library to be installed." return True, None + def input_quant_key(self) -> QuantKey | None: + # Does not call get_output_padding() - torch fallbacks + # resolve padding from compilation_config, which is unset in + # profile_run. This kernel does not pad activations. + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( @@ -242,6 +250,11 @@ class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): return True, None + def input_quant_key(self) -> QuantKey | None: + # Same per-token FP8 consume path as + # AiterPreshuffledPerTokenFp8ScaledMMLinearKernel. + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( @@ -319,6 +332,9 @@ class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): compute_capability ) + def input_quant_key(self) -> QuantKey | None: + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index c1524e78b6..0e12799304 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -27,6 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( select_fp8_moe_backend, ) +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) from vllm.model_executor.layers.linear import ( LinearMethodBase, ) @@ -405,6 +409,7 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase): replace_parameter(layer, "weight_scale", weight_scale) self.fp8_linear.process_weights_after_loading(layer) + expose_input_quant_key(layer, self.fp8_linear) layer._already_called_process_weights_after_loading = True @@ -415,8 +420,10 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase): bias: torch.Tensor | None = None, ) -> torch.Tensor: # if batch invariant mode is enabled dequant - if envs.VLLM_BATCH_INVARIANT and not isinstance( - self.fp8_linear, CutlassFP8ScaledMMLinearKernel + if ( + envs.VLLM_BATCH_INVARIANT + and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) + and not isinstance(x, QuantizedActivation) ): weight_dequant = ( layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t() diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py index 57b0836128..1ef7c1dcc9 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -9,6 +9,10 @@ from torch.nn import Parameter from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, @@ -195,10 +199,12 @@ class QuarkW8A8Fp8(QuarkScheme): module_name=self.__class__.__name__, ) + expose_input_quant_key(layer, self.fp8_linear) + def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: return self.fp8_linear.apply_weights(layer, x, bias)