| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import CompilationMode, get_current_vllm_config |
| 9 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 10 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 11 | GroupShape, |
| 12 | QuantKey, |
| 13 | kFp8StaticTensorSym, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel |
| 18 | from .ScaledMMLinearKernel import ( |
| 19 | FP8ScaledMMLinearKernel, |
| 20 | FP8ScaledMMLinearLayerConfig, |
| 21 | ) |
| 22 | |
| 23 | |
| 24 | def _get_num_tokens(output_shape: list) -> int: |
| 25 | # torch._scaled_mm works with 2D tensors, so input tensors are |
| 26 | # flattened if they are 3D. If output_shape is 3D, num_tokens is |
| 27 | # the product of all dims except the last (hidden dim). |
| 28 | return math.prod(output_shape[:-1]) |
| 29 | |
| 30 | |
| 31 | def _rocm_torch_fp8_scaled_mm_supported() -> bool: |
| 32 | from vllm.platforms.rocm import on_gfx12x, on_gfx942, on_gfx950, on_gfx1250 |
| 33 | |
| 34 | return on_gfx942() or on_gfx950() or on_gfx12x() or on_gfx1250() |
| 35 | |
| 36 | |
| 37 | def _supports_torch_fp8_scaled_mm() -> bool: |
| 38 | if current_platform.is_cpu(): |
| 39 | return True |
| 40 | if current_platform.is_xpu(): |
| 41 | return True |
| 42 | if not current_platform.is_cuda_alike(): |
| 43 | return False |
| 44 | |
| 45 | # TODO: Use torch.cuda.is_scaled_mm_supported once it is in the supported |
| 46 | # PyTorch baseline. |
| 47 | if current_platform.is_rocm(): |
| 48 | return _rocm_torch_fp8_scaled_mm_supported() |
| 49 | return current_platform.supports_fp8() |
| 50 | |
| 51 | |
| 52 | class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 53 | """Base class for FP8 linear kernels using Torch. |
| 54 | Each subclass represents a kernel variant for |
| 55 | specific device capabilities and torch versions. |
| 56 | """ |
| 57 | |
| 58 | @classmethod |
| 59 | def is_supported( |
| 60 | cls, compute_capability: int | None = None |
| 61 | ) -> tuple[bool, str | None]: |
| 62 | if _supports_torch_fp8_scaled_mm(): |
| 63 | return True, None |
| 64 | |
| 65 | return False, "requires a platform with torch FP8 scaled-MM support." |
| 66 | |
| 67 | def get_output_padding(self) -> int | None: |
| 68 | # Note: we pad the input because torch._scaled_mm is more performant |
| 69 | # for matrices with batch dimension > 16. |
| 70 | # This could change in the future. |
| 71 | # We also don't pad when using torch.compile, |
| 72 | # as it breaks with dynamic shapes. |
| 73 | # |
| 74 | # The perf gain is still relevant as of 16/1/2026 |
| 75 | # torch version == 2.9.0. More details in the link below: |
| 76 | # https://github.com/vllm-project/vllm/issues/32269 |
| 77 | vllm_config = get_current_vllm_config().compilation_config |
| 78 | pad_output = vllm_config.mode < CompilationMode.VLLM_COMPILE |
| 79 | return 17 if pad_output else None |
| 80 | |
| 81 | |
| 82 | class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): |
| 83 | @classmethod |
| 84 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 85 | per_tensor_activation_scales = ( |
| 86 | c.activation_quant_key.scale.group_shape.is_per_tensor() |
| 87 | ) |
| 88 | per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor() |
| 89 | |
| 90 | if not (per_tensor_activation_scales and per_tensor_weight_scales): |
| 91 | return False, "requires per tensor activation and weight scales." |
| 92 | return True, None |
| 93 | |
| 94 | def input_quant_key(self) -> QuantKey | None: |
| 95 | if self.config.activation_quant_key == kFp8StaticTensorSym: |
| 96 | return kFp8StaticTensorSym |
| 97 | return None |
| 98 | |
| 99 | def apply_scaled_mm( |
| 100 | self, |
| 101 | *, |
| 102 | A: torch.Tensor, |
| 103 | B: torch.Tensor, |
| 104 | out_dtype: torch.dtype, |
| 105 | As: torch.Tensor, |
| 106 | Bs: torch.Tensor, |
| 107 | bias: torch.Tensor | None, |
| 108 | output_shape: list, |
| 109 | ) -> torch.Tensor: |
| 110 | # torch._scaled_mm under torch.compile does not support 0-D scales |
| 111 | if As.dim() == 0: |
| 112 | As = As.view(1) |
| 113 | if Bs.dim() == 0: |
| 114 | Bs = Bs.view(1) |
| 115 | |
| 116 | output = torch._scaled_mm( |
| 117 | A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias |
| 118 | ) |
| 119 | # A fix for discrepancy in scaled_mm which returns tuple |
| 120 | # for torch < 2.5 and a single value in torch >= 2.5 |
| 121 | if type(output) is tuple and len(output) == 2: |
| 122 | output = output[0] |
| 123 | |
| 124 | num_tokens = _get_num_tokens(output_shape) |
| 125 | return torch.narrow(output, 0, 0, num_tokens).view(*output_shape) |
| 126 | |
| 127 | |
| 128 | class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): |
| 129 | @classmethod |
| 130 | def is_supported( |
| 131 | cls, compute_capability: int | None = None |
| 132 | ) -> tuple[bool, str | None]: |
| 133 | if not current_platform.is_rocm(): |
| 134 | return False, "requires ROCm." |
| 135 | |
| 136 | if not _supports_torch_fp8_scaled_mm(): |
| 137 | return False, "requires platform with torch FP8 scaled-MM support." |
| 138 | |
| 139 | return True, None |
| 140 | |
| 141 | @classmethod |
| 142 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 143 | per_tensor_activation_scales = ( |
| 144 | c.activation_quant_key.scale.group_shape.is_per_tensor() |
| 145 | ) |
| 146 | per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor() |
| 147 | |
| 148 | if c.out_dtype == torch.float16: |
| 149 | # hipblaslt rowwise _scaled_mm only supports BFloat16 |
| 150 | return False, "supports BFloat16 output data type only." |
| 151 | |
| 152 | if per_tensor_activation_scales or per_tensor_weight_scales: |
| 153 | return False, "cannot be used with per tensor activation and weight scales." |
| 154 | |
| 155 | return True, None |
| 156 | |
| 157 | def apply_scaled_mm( |
| 158 | self, |
| 159 | *, |
| 160 | A: torch.Tensor, |
| 161 | B: torch.Tensor, |
| 162 | out_dtype: torch.dtype, |
| 163 | As: torch.Tensor, |
| 164 | Bs: torch.Tensor, |
| 165 | bias: torch.Tensor | None, |
| 166 | output_shape: list, |
| 167 | ) -> torch.Tensor: |
| 168 | # Note: |
| 169 | # For now it has only been validated on ROCm platform. |
| 170 | # fp8 rowwise scaling in torch._scaled_mm is introduced in |
| 171 | # https://github.com/pytorch/pytorch/pull/144432 using |
| 172 | # hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above. |
| 173 | # |
| 174 | # For CUDA platform please validate if the torch._scaled_mm supports |
| 175 | # rowwise scaled GEMM before using it |
| 176 | |
| 177 | # torch._scaled_mm rowwise requires scale_a = (m, 1), scale_b = (1, n). |
| 178 | # CompressedTensors stores weight_scale as (n, 1), so `.t()` yields (1, n). |
| 179 | # ModelOpt FP8_PER_CHANNEL_PER_TOKEN stores it as 1-D (n,); reshape to |
| 180 | # (1, n) so both paths satisfy the rowwise contract. |
| 181 | scale_b = Bs.view(1, -1) if Bs.dim() == 1 else Bs.t() |
| 182 | if As.dim() == 1: |
| 183 | As = As.view(-1, 1) |
| 184 | |
| 185 | # Fused GEMM_DQ Rowwise GEMM |
| 186 | output = torch._scaled_mm( |
| 187 | A, |
| 188 | B, |
| 189 | out_dtype=out_dtype, |
| 190 | scale_a=As, |
| 191 | scale_b=scale_b, |
| 192 | bias=bias, |
| 193 | ) |
| 194 | |
| 195 | num_tokens = _get_num_tokens(output_shape) |
| 196 | return torch.narrow(output, 0, 0, num_tokens).view(*output_shape) |
| 197 | |
| 198 | |
| 199 | class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): |
| 200 | @classmethod |
| 201 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 202 | per_tensor_activation_scales = ( |
| 203 | c.activation_quant_key.scale.group_shape.is_per_tensor() |
| 204 | ) |
| 205 | per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor() |
| 206 | |
| 207 | if per_tensor_activation_scales and per_tensor_weight_scales: |
| 208 | return False, "cannot be used with per tensor activation and weight scales." |
| 209 | |
| 210 | return True, None |
| 211 | |
| 212 | def apply_scaled_mm( |
| 213 | self, |
| 214 | *, |
| 215 | A: torch.Tensor, |
| 216 | B: torch.Tensor, |
| 217 | out_dtype: torch.dtype, |
| 218 | As: torch.Tensor, |
| 219 | Bs: torch.Tensor, |
| 220 | bias: torch.Tensor | None, |
| 221 | output_shape: list, |
| 222 | ) -> torch.Tensor: |
| 223 | # Use unfused DQ due to limitations with scaled_mm |
| 224 | |
| 225 | # Symmetric quantized GEMM by definition computes the following: |
| 226 | # C = (s_x * X) (s_w * W) + bias |
| 227 | # This is equivalent to dequantizing the weights and activations |
| 228 | # before applying a GEMM. |
| 229 | # |
| 230 | # In order to compute quantized operands, a quantized kernel |
| 231 | # will rewrite the above like so: |
| 232 | # C = s_w * s_x * (X * W) + bias |
| 233 | # |
| 234 | # For the scaled_mm fallback case, we break this down, since it |
| 235 | # does not support s_w being a vector. |
| 236 | |
| 237 | # Input scaling factors are no longer optional in _scaled_mm starting |
| 238 | # from pytorch 2.5. Allocating a dummy tensor to pass as scales |
| 239 | dummy_tensor = torch.ones(1, dtype=torch.float32, device=A.device) |
| 240 | |
| 241 | # GEMM |
| 242 | # This computes C = (X * W). |
| 243 | # Output in fp32 to allow subsequent ops to happen in-place |
| 244 | output = torch._scaled_mm( |
| 245 | A, |
| 246 | B, |
| 247 | scale_a=dummy_tensor, |
| 248 | scale_b=dummy_tensor, |
| 249 | out_dtype=torch.float32, |
| 250 | ) |
| 251 | # A fix for discrepancy in scaled_mm which returns tuple |
| 252 | # for torch < 2.5 and a single value in torch >= 2.5 |
| 253 | if type(output) is tuple and len(output) == 2: |
| 254 | output = output[0] |
| 255 | |
| 256 | # Unpad (undo num_token_padding) |
| 257 | num_tokens = _get_num_tokens(output_shape) |
| 258 | output = torch.narrow(output, 0, 0, num_tokens) |
| 259 | x_scale = torch.narrow(As, 0, 0, num_tokens) |
| 260 | |
| 261 | # DQ |
| 262 | # C = sw * sx * (X * W) + bias |
| 263 | output = output * x_scale * Bs.t() |
| 264 | if bias is not None: |
| 265 | output = output + bias |
| 266 | return output.to(out_dtype).view(*output_shape) |
| 267 | |
| 268 | |
| 269 | class BlockWiseTorchFP8ScaledMMLinearKernel(Fp8BlockScaledMMLinearKernel): |
| 270 | """FP8 block-scaled linear kernel using ``torch._scaled_mm``. |
| 271 | |
| 272 | Implements the block-scaled path of ``torch._scaled_mm``, which |
| 273 | dispatches on the shapes of the scale tensors. For ``A = [M, K]`` and |
| 274 | ``B = [K, N]`` (both fp8) the op's block path requires, with float32 |
| 275 | scales: |
| 276 | * 1x128 activation: ``scale_a = [M, ceil(K / 128)]`` |
| 277 | * 128x128 weight: ``scale_b = [ceil(K / 128), ceil(N / 128)]`` |
| 278 | """ |
| 279 | |
| 280 | def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None: |
| 281 | super().__init__(config) |
| 282 | act_scale_descriptor = config.activation_quant_key.scale |
| 283 | self.quant_fp8 = QuantFP8( |
| 284 | static=act_scale_descriptor.static, |
| 285 | group_shape=act_scale_descriptor.group_shape, |
| 286 | num_token_padding=self.get_output_padding(), |
| 287 | use_ue8m0=False, |
| 288 | column_major_scales=current_platform.is_cuda_alike(), |
| 289 | ) |
| 290 | |
| 291 | @classmethod |
| 292 | def is_supported( |
| 293 | cls, compute_capability: int | None = None |
| 294 | ) -> tuple[bool, str | None]: |
| 295 | if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): |
| 296 | return False, "requires CUDA, ROCm or XPU." |
| 297 | # torch._scaled_mm DeepSeek-style (1x128, 128x128) block scaling is |
| 298 | # only implemented for SM90 (Hopper) on NVIDIA CUDA. See PyTorch's |
| 299 | # _check_deepseek_support(): |
| 300 | # https://github.com/pytorch/pytorch/blob/33812ece06e2b0d597f73fbe41de03a83f9109f9/aten/src/ATen/native/cuda/ScaledBlas.cpp#L804-L820 # noqa: E501 |
| 301 | if current_platform.is_cuda(): |
| 302 | is_sm90 = ( |
| 303 | compute_capability == 90 |
| 304 | if compute_capability is not None |
| 305 | else current_platform.is_device_capability(90) |
| 306 | ) |
| 307 | if not is_sm90: |
| 308 | return ( |
| 309 | False, |
| 310 | "DeepSeek-style (1x128, 128x128) block scaling on CUDA " |
| 311 | "requires compute capability 90 (Hopper).", |
| 312 | ) |
| 313 | return True, None |
| 314 | |
| 315 | @classmethod |
| 316 | def can_implement( |
| 317 | cls, config: FP8ScaledMMLinearLayerConfig |
| 318 | ) -> tuple[bool, str | None]: |
| 319 | can_implement_base, reason = super().can_implement(config) |
| 320 | if not can_implement_base: |
| 321 | return can_implement_base, reason |
| 322 | |
| 323 | act_group_shape = config.activation_quant_key.scale.group_shape |
| 324 | if act_group_shape != GroupShape(1, 128): |
| 325 | return ( |
| 326 | False, |
| 327 | "requires 1x128 activation quantization.", |
| 328 | ) |
| 329 | weight_group_shape = config.weight_quant_key.scale.group_shape |
| 330 | if weight_group_shape != GroupShape(128, 128): |
| 331 | return ( |
| 332 | False, |
| 333 | "requires 128x128 block weight quantization.", |
| 334 | ) |
| 335 | return True, None |
| 336 | |
| 337 | def apply_block_scaled_mm( |
| 338 | self, |
| 339 | A: torch.Tensor, |
| 340 | B: torch.Tensor, |
| 341 | As: torch.Tensor, |
| 342 | Bs: torch.Tensor, |
| 343 | ) -> torch.Tensor: |
| 344 | # torch._scaled_mm implemented by cuBLASLt requires the M dimension to be a |
| 345 | # multiple of 4; pad A (and its column-major scale) up and slice back. |
| 346 | # https://docs.nvidia.com/cuda/cublas/index.html?highlight=128%2520element%25201D%2520and%2520128x128%25202D#scaling-factors-layouts # noqa: E501 |
| 347 | M = A.shape[0] |
| 348 | pad = (-M) % 4 if current_platform.is_cuda() else 0 |
| 349 | if pad: |
| 350 | Ap = A.new_empty((M + pad, A.shape[1])) |
| 351 | Ap[:M] = A |
| 352 | A = Ap |
| 353 | # As is column-major [M, ceil(K/128)] (stride (1, M)); build a padded |
| 354 | # column-major buffer and copy only the valid rows. Padded rows are left |
| 355 | # uninitialized — they don't affect the sliced-back output. |
| 356 | Asp = As.new_empty((As.shape[1], M + pad)).t() |
| 357 | Asp[:M] = As |
| 358 | As = Asp |
| 359 | |
| 360 | output = torch._scaled_mm( |
| 361 | A, |
| 362 | B.t(), |
| 363 | scale_a=As, |
| 364 | scale_b=Bs.t(), |
| 365 | out_dtype=self.config.out_dtype, |
| 366 | ) |
| 367 | if type(output) is tuple and len(output) == 2: |
| 368 | output = output[0] |
| 369 | if pad: |
| 370 | output = output[:M, :] |
| 371 | return output |