@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6ca2b23e22

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

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

Rebecca Lee committed Sep 16, 2026 at 09:48 UTC 6ca2b23e22aab8534e00a85e8ec5222508a6ecf3

5 files changed +104 -4

tests/fusion/test_quant_activation_contract.py

@@ -24,6 +24,11 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (

24 FlashInferCutlassNvFp4LinearKernel,

25 FlashInferTrtllmNvFp4LinearKernel,

26 )

27 +from vllm.model_executor.kernels.linear.scaled_mm.aiter import (

28 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,

29 + AiterPerTokenFp8ScaledMMLinearKernel,

30 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

31 +)

32 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (

33 CutlassFP8ScaledMMLinearKernel,

34 )

@@ -56,6 +61,9 @@ SUPPORTING = {

61 FlashInferFP8ScaledMMLinearKernel,

62 FlashInferCutlassNvFp4LinearKernel,

63 PerTensorTorchFP8ScaledMMLinearKernel,

64 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,

65 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

66 + AiterPerTokenFp8ScaledMMLinearKernel,

67 FlashInferCutedslMxfp8LinearKernel,

68 FlashInferCutlassMxfp8LinearKernel,

69 }

tests/quantization/test_quark.py

+64 -1

@@ -21,9 +21,18 @@ from packaging import version

22 from tests.quantization.utils import load_model_without_vllm_runner

23 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

24 -from vllm.config import set_current_vllm_config

24 +from vllm.config import VllmConfig, set_current_vllm_config

25 from vllm.config.cache import CacheConfig

26 from vllm.forward_context import set_forward_context

27 +from vllm.model_executor import parameter

28 +from vllm.model_executor.kernels.linear.scaled_mm.aiter import (

29 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,

30 + AiterPerTokenFp8ScaledMMLinearKernel,

31 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

32 +)

33 +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (

34 + FP8ScaledMMLinearLayerConfig,

35 +)

36 from vllm.model_executor.layers.attention import Attention

37 from vllm.model_executor.layers.fused_moe import (

38 FusedMoeWeightScaleSupported,

@@ -113,6 +122,12 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(

122

123 AITER_AVAILABLE = is_aiter_found_and_supported()

124

125 +AITER_PTPC_KERNELS = (

126 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,

127 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

128 + AiterPerTokenFp8ScaledMMLinearKernel,

129 +)

130 +

131 DEVICE_TYPE = current_platform.device_type

132

133

@@ -1095,6 +1110,54 @@ def test_quark_w8a8_fp8_moe_per_block_rejects_misaligned_partition():

1110 )

1111

1112

1113 +def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch):

1114 + """QuarkW8A8Fp8 must advertise the key its kernel consumes pre-quantized.

1115 +

1116 + Only kernel selection runs, no GEMM. The shape matters: AITER PTPC

1117 + kernels decline untuned shapes, leaving a torch kernel that has no key.

1118 + """

1119 + # Llama-3.1-70B qkv_proj at TP1 (aiter-tuned N, K on gfx950).

1120 + N, K = 10240, 8192

1121 + dtype = torch.bfloat16

1122 + kernel_config = FP8ScaledMMLinearLayerConfig(

1123 + weight_quant_key=kFp8StaticChannelSym,

1124 + activation_quant_key=kFp8DynamicTokenSym,

1125 + weight_shape=(N, K),

1126 + input_dtype=dtype,

1127 + out_dtype=dtype,

1128 + )

1129 + if not any(

1130 + cls.is_supported()[0] and cls.can_implement(kernel_config)[0]

1131 + for cls in AITER_PTPC_KERNELS

1132 + ):

1133 + pytest.skip("no AITER PTPC kernel is usable for this shape and environment")

1134 +

1135 + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)

1136 + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)

1137 +

1138 + scheme = QuarkW8A8Fp8.__new__(QuarkW8A8Fp8)

1139 + scheme.weight_qscheme = "per_channel"

1140 + scheme.is_static_input_scheme = False

1141 + scheme.input_qscheme = "per_channel"

1142 + scheme.activation_quant_key = kFp8DynamicTokenSym

1143 + scheme.weight_quant_key = kFp8StaticChannelSym

1144 + scheme.out_dtype = dtype

1145 + scheme.input_dtype = dtype

1146 +

1147 + layer = torch.nn.Module()

1148 + with set_current_vllm_config(VllmConfig()):

1149 + scheme.create_weights(

1150 + layer,

1151 + output_partition_sizes=[N],

1152 + input_size_per_partition=K,

1153 + params_dtype=dtype,

1154 + weight_loader=lambda *args, **kwargs: None,

1155 + )

1156 +

1157 + assert isinstance(scheme.fp8_linear, AITER_PTPC_KERNELS)

1158 + assert layer.input_quant_key == kFp8DynamicTokenSym

1159 +

1160 +

1161 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])

1162 def test_quark_fp8_w_per_tensor_a_per_tensor(

1163 kv_cache_dtype: str, monkeypatch, dist_init, workspace_init

vllm/model_executor/kernels/linear/scaled_mm/aiter.py

+16

@@ -16,6 +16,8 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

16 )

17 from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 GroupShape,

19 + QuantKey,

20 + kFp8DynamicTokenSym,

21 )

22 from vllm.model_executor.utils import replace_parameter

23 from vllm.platforms import current_platform

@@ -150,6 +152,12 @@ class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

152 return False, "requires aiter library to be installed."

153 return True, None

154

155 + def input_quant_key(self) -> QuantKey | None:

156 + # Does not call get_output_padding() - torch fallbacks

157 + # resolve padding from compilation_config, which is unset in

158 + # profile_run. This kernel does not pad activations.

159 + return kFp8DynamicTokenSym

160 +

161 @classmethod

162 def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

163 is_ptpc = (

@@ -242,6 +250,11 @@ class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

250

251 return True, None

252

253 + def input_quant_key(self) -> QuantKey | None:

254 + # Same per-token FP8 consume path as

255 + # AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.

256 + return kFp8DynamicTokenSym

257 +

258 @classmethod

259 def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

260 is_ptpc = (

@@ -319,6 +332,9 @@ class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

332 compute_capability

333 )

334

335 + def input_quant_key(self) -> QuantKey | None:

336 + return kFp8DynamicTokenSym

337 +

338 @classmethod

339 def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

340 is_ptpc = (

vllm/model_executor/layers/quantization/online/fp8.py

+9 -2

@@ -27,6 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig

27 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

28 select_fp8_moe_backend,

29 )

30 +from vllm.model_executor.layers.fusion.quant_activation import (

31 + QuantizedActivation,

32 + expose_input_quant_key,

33 +)

34 from vllm.model_executor.layers.linear import (

35 LinearMethodBase,

36 )

@@ -405,6 +409,7 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase):

409 replace_parameter(layer, "weight_scale", weight_scale)

410

411 self.fp8_linear.process_weights_after_loading(layer)

412 + expose_input_quant_key(layer, self.fp8_linear)

413

414 layer._already_called_process_weights_after_loading = True

415

@@ -415,8 +420,10 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase):

420 bias: torch.Tensor | None = None,

421 ) -> torch.Tensor:

422 # if batch invariant mode is enabled dequant

418 - if envs.VLLM_BATCH_INVARIANT and not isinstance(

419 - self.fp8_linear, CutlassFP8ScaledMMLinearKernel

423 + if (

424 + envs.VLLM_BATCH_INVARIANT

425 + and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel)

426 + and not isinstance(x, QuantizedActivation)

427 ):

428 weight_dequant = (

429 layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t()

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

+7 -1

@@ -9,6 +9,10 @@ from torch.nn import Parameter

9 from vllm.config import get_current_vllm_config

10 from vllm.logger import init_logger

11 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel

12 +from vllm.model_executor.layers.fusion.quant_activation import (

13 + QuantizedActivation,

14 + expose_input_quant_key,

15 +)

16 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme

17 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

18 create_fp8_scale_parameter,

@@ -195,10 +199,12 @@ class QuarkW8A8Fp8(QuarkScheme):

199 module_name=self.__class__.__name__,

200 )

201

202 + expose_input_quant_key(layer, self.fp8_linear)

203 +

204 def apply_weights(

205 self,

206 layer: torch.nn.Module,

201 - x: torch.Tensor,

207 + x: torch.Tensor | QuantizedActivation,

208 bias: torch.Tensor | None = None,

209 ) -> torch.Tensor:

210 return self.fp8_linear.apply_weights(layer, x, bias)