| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm import _custom_ops as ops |
| 8 | from vllm._aiter_ops import ( |
| 9 | rocm_aiter_ops, |
| 10 | ) |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 13 | _upcast_e8m0_to_fp32, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 16 | GroupShape, |
| 17 | QuantKey, |
| 18 | kFp8DynamicTokenSym, |
| 19 | ) |
| 20 | from vllm.model_executor.utils import replace_parameter |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | from .BlockScaledMMLinearKernel import ( |
| 24 | FP8BlockParams, |
| 25 | Fp8BlockScaledMMLinearKernel, |
| 26 | ) |
| 27 | from .cutlass import CutlassInt8ScaledMMLinearKernel |
| 28 | from .ScaledMMLinearKernel import ( |
| 29 | FP8ScaledMMLinearKernel, |
| 30 | FP8ScaledMMLinearLayerConfig, |
| 31 | Int8ScaledMMLinearLayerConfig, |
| 32 | ) |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): |
| 38 | @classmethod |
| 39 | def is_supported( |
| 40 | cls, compute_capability: int | None = None |
| 41 | ) -> tuple[bool, str | None]: |
| 42 | if not current_platform.is_rocm(): |
| 43 | return False, "Requires ROCm." |
| 44 | |
| 45 | if compute_capability is not None and compute_capability < 90: |
| 46 | return False, "requires compute capability 90 and above." |
| 47 | |
| 48 | try: |
| 49 | import aiter # noqa: F401 # deliberately attempt to import aiter |
| 50 | except Exception: |
| 51 | return False, "requires `aiter` to be installed." |
| 52 | |
| 53 | if not rocm_aiter_ops.is_linear_enabled(): |
| 54 | return ( |
| 55 | False, |
| 56 | "requires setting `VLLM_ROCM_USE_AITER=1` " |
| 57 | "and `VLLM_ROCM_USE_AITER_LINEAR=1`. " |
| 58 | "`VLLM_ROCM_USE_AITER_LINEAR` default is True.", |
| 59 | ) |
| 60 | return True, None |
| 61 | |
| 62 | @classmethod |
| 63 | def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 64 | if not c.input_symmetric: |
| 65 | return False, "supports symmetric quantization only." |
| 66 | return True, None |
| 67 | |
| 68 | def apply_weights( |
| 69 | self, |
| 70 | layer: torch.nn.Module, |
| 71 | x: torch.Tensor, |
| 72 | bias: torch.Tensor | None = None, |
| 73 | ) -> torch.Tensor: |
| 74 | """`AiterInt8ScaledMMLinearKernel` implements a fused version of |
| 75 | `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)` |
| 76 | where scale_a * a and scale_b * b are implemented using numpy-style |
| 77 | broadcasting. |
| 78 | Currently only support per-tensor-per-tensor GEMM |
| 79 | and per-token-per-channel GEMM through AITER |
| 80 | w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` also does not support |
| 81 | ATIER block scaled GEMM and mix-precision GEMM. |
| 82 | """ |
| 83 | w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) |
| 84 | |
| 85 | # ops.scaled_int8_quant supports both dynamic and static quant: |
| 86 | # * dynamic, i_s is None and x_s computed from x. |
| 87 | # * static, i_s is scalar and x_s is i_s. |
| 88 | symmetric = azp_adj is None |
| 89 | assert symmetric, ( |
| 90 | "AiterInt8ScaledMMLinearKernel only supports symmetric quantization." |
| 91 | ) |
| 92 | x_q, x_s, x_zp = ops.scaled_int8_quant(x, i_s, i_zp, symmetric=symmetric) |
| 93 | |
| 94 | assert x_zp is None, ( |
| 95 | "AiterInt8ScaledMMLinearKernel only supports symmetric quantization." |
| 96 | ) |
| 97 | out_dtype = x.dtype |
| 98 | |
| 99 | assert w_q.shape[0] % 16 == 0 and w_q.shape[1] % 16 == 0 |
| 100 | assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 |
| 101 | assert bias is None or bias.shape[0] == w_q.shape[1] and bias.dtype == out_dtype |
| 102 | |
| 103 | m = x_q.shape[0] # a |
| 104 | n = w_q.shape[1] # b |
| 105 | |
| 106 | per_tensor_scale_a = x_s.numel() == 1 |
| 107 | per_tensor_scale_b = w_s.numel() == 1 |
| 108 | per_token_scale_a = x_s.numel() == m |
| 109 | per_channel_scale_b = w_s.numel() == n |
| 110 | |
| 111 | # @TODO: |
| 112 | # Maybe broadcast the per-tensor-scale into per-channel-scale |
| 113 | # if one of the scale is a per-channel-scale. |
| 114 | # For now, it only supports: |
| 115 | # - per-tensor-per-tensor a8w8 scaled GEMM, and |
| 116 | # - per-token-per-channel a8w8 scaled GEMM |
| 117 | assert (per_tensor_scale_a and per_tensor_scale_b) or ( |
| 118 | per_token_scale_a and per_channel_scale_b |
| 119 | ), ( |
| 120 | "Currently only support per-tensor-per-tensor GEMM " |
| 121 | " and per-token-per-channel GEMM through AITER" |
| 122 | " w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` " |
| 123 | "does not support AITER block scaled GEMM." |
| 124 | ) |
| 125 | |
| 126 | # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects |
| 127 | # a to be [M, K] |
| 128 | # b to be [N, K] |
| 129 | # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format |
| 130 | return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype) |
| 131 | |
| 132 | |
| 133 | class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 134 | @classmethod |
| 135 | def is_supported( |
| 136 | cls, compute_capability: int | None = None |
| 137 | ) -> tuple[bool, str | None]: |
| 138 | if not current_platform.is_rocm(): |
| 139 | return False, "requires ROCm." |
| 140 | if not rocm_aiter_ops.is_linear_fp8_enabled(): |
| 141 | return ( |
| 142 | False, |
| 143 | "requires setting `VLLM_ROCM_USE_AITER=1` " |
| 144 | "and `VLLM_ROCM_USE_AITER_LINEAR=1`. " |
| 145 | "`VLLM_ROCM_USE_AITER_LINEAR` default is True.", |
| 146 | ) |
| 147 | try: |
| 148 | import aiter # noqa: F401 |
| 149 | except Exception: |
| 150 | return False, "requires aiter library to be installed." |
| 151 | return True, None |
| 152 | |
| 153 | def input_quant_key(self) -> QuantKey | None: |
| 154 | # Does not call get_output_padding() - torch fallbacks |
| 155 | # resolve padding from compilation_config, which is unset in |
| 156 | # profile_run. This kernel does not pad activations. |
| 157 | return kFp8DynamicTokenSym |
| 158 | |
| 159 | @classmethod |
| 160 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 161 | is_ptpc = ( |
| 162 | c.activation_quant_key.scale.group_shape.is_per_token() |
| 163 | and c.weight_quant_key.scale.group_shape.is_per_channel() |
| 164 | ) |
| 165 | if c.weight_shape is None: |
| 166 | return False, "weight_shape is required for Aiter kernels" |
| 167 | N, K = c.weight_shape |
| 168 | fp8_dtype = current_platform.fp8_dtype() |
| 169 | |
| 170 | if c.out_dtype is not torch.bfloat16: |
| 171 | return False, "requires bfloat16 output dtype." |
| 172 | |
| 173 | if not is_ptpc: |
| 174 | return ( |
| 175 | False, |
| 176 | "requires per token activation scales and per channel weight scales.", |
| 177 | ) |
| 178 | |
| 179 | if not (N % 16 == 0 and K % 16 == 0): |
| 180 | return ( |
| 181 | False, |
| 182 | f"requires N and K dimensions divisible by 16, received " |
| 183 | f"N={N} and K={K}.", |
| 184 | ) |
| 185 | |
| 186 | # Aiter's shuffled per-token Gemm performs better than torch only when its |
| 187 | # tuned. |
| 188 | if not rocm_aiter_ops.is_shuffled_per_token_w8a8_gemm_tuned(N, K, fp8_dtype): |
| 189 | return ( |
| 190 | False, |
| 191 | f"requires a tuned configuration for N: {N} and K: {K} " |
| 192 | f"and fp8 dtype {fp8_dtype}.", |
| 193 | ) |
| 194 | |
| 195 | return True, None |
| 196 | |
| 197 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 198 | w_name, *_ = self.layer_param_names |
| 199 | w, *_ = self._get_layer_params(layer) |
| 200 | |
| 201 | replace_parameter( |
| 202 | layer, |
| 203 | w_name, |
| 204 | torch.nn.Parameter( |
| 205 | rocm_aiter_ops.shuffle_weight(w.t().contiguous()).data, |
| 206 | requires_grad=False, |
| 207 | ), |
| 208 | ) |
| 209 | |
| 210 | def apply_scaled_mm( |
| 211 | self, |
| 212 | *, |
| 213 | A: torch.Tensor, |
| 214 | B: torch.Tensor, |
| 215 | out_dtype: torch.dtype, |
| 216 | As: torch.Tensor, |
| 217 | Bs: torch.Tensor, |
| 218 | bias: torch.Tensor | None, |
| 219 | output_shape: list, |
| 220 | ) -> torch.Tensor: |
| 221 | return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm( |
| 222 | A, B, As, Bs, bias, out_dtype |
| 223 | ) |
| 224 | |
| 225 | |
| 226 | class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 227 | @classmethod |
| 228 | def is_supported( |
| 229 | cls, compute_capability: int | None = None |
| 230 | ) -> tuple[bool, str | None]: |
| 231 | if not current_platform.is_rocm(): |
| 232 | return False, "requires ROCm." |
| 233 | |
| 234 | if not rocm_aiter_ops.is_linear_hipbmm_enabled(): |
| 235 | return ( |
| 236 | False, |
| 237 | "requires setting `VLLM_ROCM_USE_AITER=1`, " |
| 238 | "`VLLM_ROCM_USE_AITER_LINEAR=1`, " |
| 239 | "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`.", |
| 240 | ) |
| 241 | try: |
| 242 | import aiter # noqa: F401 |
| 243 | except Exception: |
| 244 | return False, "requires aiter library to be installed." |
| 245 | |
| 246 | if not hasattr(aiter, "hipb_mm"): |
| 247 | return False, "requires aiter hipb_mm support." |
| 248 | |
| 249 | return True, None |
| 250 | |
| 251 | def input_quant_key(self) -> QuantKey | None: |
| 252 | # Same per-token FP8 consume path as |
| 253 | # AiterPreshuffledPerTokenFp8ScaledMMLinearKernel. |
| 254 | return kFp8DynamicTokenSym |
| 255 | |
| 256 | @classmethod |
| 257 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 258 | is_ptpc = ( |
| 259 | c.activation_quant_key.scale.group_shape.is_per_token() |
| 260 | and c.weight_quant_key.scale.group_shape.is_per_channel() |
| 261 | ) |
| 262 | if c.weight_shape is None: |
| 263 | return False, "weight_shape is required for Aiter kernels" |
| 264 | N, K = c.weight_shape |
| 265 | |
| 266 | if c.out_dtype is not torch.bfloat16: |
| 267 | return False, "requires bfloat16 output dtype." |
| 268 | |
| 269 | if not is_ptpc: |
| 270 | return ( |
| 271 | False, |
| 272 | "requires per token activation scales and per channel weight scales.", |
| 273 | ) |
| 274 | |
| 275 | if not (N >= 16 and N % 16 == 0 and K % 16 == 0): |
| 276 | return ( |
| 277 | False, |
| 278 | "requires N >= 16 and both N and K divisible by 16, " |
| 279 | f"received N={N} and K={K}.", |
| 280 | ) |
| 281 | |
| 282 | return True, None |
| 283 | |
| 284 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 285 | w_name, w_s_name, *_ = self.layer_param_names |
| 286 | w, w_s, *_ = self._get_layer_params(layer) |
| 287 | |
| 288 | # Pre-apply the transposes that used to live in |
| 289 | # _rocm_aiter_hipb_mm_fp8_impl so the kernel can consume B/Bs directly. |
| 290 | # The `.t()` on the shuffled weight is kept as a non-contiguous view — |
| 291 | # materializing it with `.contiguous()` would re-arrange the bytes and |
| 292 | # break the `bpreshuffle` layout. |
| 293 | shuffled_w = rocm_aiter_ops.shuffle_weight(w.t().contiguous()) |
| 294 | replace_parameter( |
| 295 | layer, |
| 296 | w_name, |
| 297 | torch.nn.Parameter(shuffled_w.t(), requires_grad=False), |
| 298 | ) |
| 299 | |
| 300 | if w_s.ndim > 1: |
| 301 | replace_parameter( |
| 302 | layer, |
| 303 | w_s_name, |
| 304 | torch.nn.Parameter(w_s.t().contiguous(), requires_grad=False), |
| 305 | ) |
| 306 | |
| 307 | def apply_scaled_mm( |
| 308 | self, |
| 309 | *, |
| 310 | A: torch.Tensor, |
| 311 | B: torch.Tensor, |
| 312 | out_dtype: torch.dtype, |
| 313 | As: torch.Tensor, |
| 314 | Bs: torch.Tensor, |
| 315 | bias: torch.Tensor | None, |
| 316 | output_shape: list, |
| 317 | ) -> torch.Tensor: |
| 318 | output_shape[-1] = B.shape[1] |
| 319 | return rocm_aiter_ops.hipb_mm_fp8(A, B, As, Bs, bias, out_dtype).view( |
| 320 | *output_shape |
| 321 | ) |
| 322 | |
| 323 | |
| 324 | class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 325 | @classmethod |
| 326 | def is_supported( |
| 327 | cls, compute_capability: int | None = None |
| 328 | ) -> tuple[bool, str | None]: |
| 329 | return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported( |
| 330 | compute_capability |
| 331 | ) |
| 332 | |
| 333 | def input_quant_key(self) -> QuantKey | None: |
| 334 | return kFp8DynamicTokenSym |
| 335 | |
| 336 | @classmethod |
| 337 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 338 | is_ptpc = ( |
| 339 | c.activation_quant_key.scale.group_shape.is_per_token() |
| 340 | and c.weight_quant_key.scale.group_shape.is_per_channel() |
| 341 | ) |
| 342 | if c.weight_shape is None: |
| 343 | return False, "weight_shape is required for Aiter kernels" |
| 344 | N, K = c.weight_shape |
| 345 | fp8_dtype = current_platform.fp8_dtype() |
| 346 | |
| 347 | if not is_ptpc: |
| 348 | return ( |
| 349 | False, |
| 350 | "requires per token activation scales and per channel weight scales.", |
| 351 | ) |
| 352 | |
| 353 | # Aiter's per-token Gemm performs better than torch oonly when its |
| 354 | # tuned. |
| 355 | if not rocm_aiter_ops.is_per_token_w8a8_gemm_tuned(N, K, fp8_dtype): |
| 356 | return ( |
| 357 | False, |
| 358 | f"requires a tuned configuration for N: {N} and K: {K} " |
| 359 | f"and fp8 dtype {fp8_dtype}.", |
| 360 | ) |
| 361 | return True, None |
| 362 | |
| 363 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 364 | w_name, *_ = self.layer_param_names |
| 365 | w, *_ = self._get_layer_params(layer) |
| 366 | |
| 367 | replace_parameter( |
| 368 | layer, |
| 369 | w_name, |
| 370 | torch.nn.Parameter(w.t(), requires_grad=False), |
| 371 | ) |
| 372 | |
| 373 | def apply_scaled_mm( |
| 374 | self, |
| 375 | *, |
| 376 | A: torch.Tensor, |
| 377 | B: torch.Tensor, |
| 378 | out_dtype: torch.dtype, |
| 379 | As: torch.Tensor, |
| 380 | Bs: torch.Tensor, |
| 381 | bias: torch.Tensor | None, |
| 382 | output_shape: list, |
| 383 | ) -> torch.Tensor: |
| 384 | return rocm_aiter_ops.w8a8_gemm(A, B, As, Bs, bias, out_dtype) |
| 385 | |
| 386 | |
| 387 | class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 388 | def __init__(self, config: FP8ScaledMMLinearLayerConfig): |
| 389 | super().__init__(config) |
| 390 | n, k = config.weight_shape |
| 391 | |
| 392 | _on_gfx1250 = False |
| 393 | if current_platform.is_rocm(): |
| 394 | from vllm.platforms.rocm import on_gfx1250 |
| 395 | |
| 396 | _on_gfx1250 = on_gfx1250() |
| 397 | |
| 398 | self.use_triton = not current_platform.is_fp8_fnuz() and ( |
| 399 | rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250 |
| 400 | ) |
| 401 | |
| 402 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 403 | super().process_weights_after_loading(layer) |
| 404 | |
| 405 | params = FP8BlockParams.from_layer(layer) |
| 406 | if params.weight_scale_inv is not None: |
| 407 | ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV |
| 408 | else: |
| 409 | ws, attr = params.weight_scale, params.WEIGHT_SCALE |
| 410 | if ws is not None and ws.dtype == torch.float8_e8m0fnu: |
| 411 | replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) |
| 412 | |
| 413 | @classmethod |
| 414 | def is_supported(cls, compute_capability=None): |
| 415 | if ( |
| 416 | rocm_aiter_ops.is_linear_enabled() |
| 417 | or rocm_aiter_ops.is_rdna_linear_enabled() |
| 418 | ): |
| 419 | return True, None |
| 420 | return ( |
| 421 | False, |
| 422 | "Only supported on ROCm platform with aiter package installed.", |
| 423 | ) |
| 424 | |
| 425 | @classmethod |
| 426 | def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): |
| 427 | can_implement_base, reason = super().can_implement(config) |
| 428 | if not can_implement_base: |
| 429 | return can_implement_base, reason |
| 430 | |
| 431 | act_quant_desc = config.activation_quant_key.scale |
| 432 | if act_quant_desc.group_shape != GroupShape(1, 128): |
| 433 | return ( |
| 434 | False, |
| 435 | "Supports only dynamic per token group activation " |
| 436 | "quantization with group_shape=(1,128).", |
| 437 | ) |
| 438 | |
| 439 | # RDNA4 (gfx12) only has the aiter Triton blockscale backend, which |
| 440 | # needs a per-(N,K) tune. Reject untuned shapes so the dispatcher falls |
| 441 | # through to the generic backend. |
| 442 | if rocm_aiter_ops.is_rdna_linear_enabled(): |
| 443 | n, k = config.weight_shape |
| 444 | if not rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k): |
| 445 | return ( |
| 446 | False, |
| 447 | f"(N={n}, K={k}) is not in the aiter Triton blockscale tuned list.", |
| 448 | ) |
| 449 | return True, None |
| 450 | |
| 451 | def apply_block_scaled_mm( |
| 452 | self, |
| 453 | A: torch.Tensor, |
| 454 | B: torch.Tensor, |
| 455 | As: torch.Tensor, |
| 456 | Bs: torch.Tensor, |
| 457 | ) -> torch.Tensor: |
| 458 | if As.dtype != Bs.dtype: |
| 459 | if As.dtype == torch.float8_e8m0fnu: |
| 460 | As = _upcast_e8m0_to_fp32(As).contiguous() |
| 461 | else: |
| 462 | As = As.to(torch.float32) |
| 463 | |
| 464 | Bs = Bs.to(torch.float32) |
| 465 | |
| 466 | out_dtype = self.config.out_dtype |
| 467 | if self.use_triton: |
| 468 | gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale |
| 469 | else: |
| 470 | gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale |
| 471 | |
| 472 | return gemm_a8w8_blockscale_op( |
| 473 | A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype |
| 474 | ) |