[ROCm] Expose kFp8DynamicTokenSym on AITER PTPC linears (#54248)

Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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,
 }

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

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 = (

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()

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)