| 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.model_executor.layers.quantization.utils import replace_parameter |
| 9 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 10 | _upcast_e8m0_to_fp32, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 13 | convert_to_channelwise, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.utils import check_cpu_sgl_kernel |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.platforms.interface import CpuArchEnum |
| 18 | |
| 19 | from .BlockScaledMMLinearKernel import ( |
| 20 | Fp8BlockScaledMMLinearKernel, |
| 21 | FP8ScaledMMLinearLayerConfig, |
| 22 | ) |
| 23 | from .ScaledMMLinearKernel import ( |
| 24 | FP8ScaledMMLinearKernel, |
| 25 | Int8ScaledMMLinearKernel, |
| 26 | Int8ScaledMMLinearLayerConfig, |
| 27 | ) |
| 28 | |
| 29 | # BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size. |
| 30 | _BLOCK_SIZE_K = 128 |
| 31 | |
| 32 | |
| 33 | class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): |
| 34 | @classmethod |
| 35 | def is_supported( |
| 36 | cls, compute_capability: int | None = None |
| 37 | ) -> tuple[bool, str | None]: |
| 38 | if not current_platform.is_cpu(): |
| 39 | return False, "requires CPU." |
| 40 | return True, None |
| 41 | |
| 42 | @classmethod |
| 43 | def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 44 | return True, None |
| 45 | |
| 46 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 47 | w_q_name, _, _, _, _ = self.layer_param_names |
| 48 | weight = getattr(layer, w_q_name) |
| 49 | dtype = weight.dtype |
| 50 | N, K = weight.size() |
| 51 | if ( |
| 52 | current_platform.get_cpu_architecture() == CpuArchEnum.X86 |
| 53 | and self.config.input_symmetric |
| 54 | and check_cpu_sgl_kernel(N, K, dtype) |
| 55 | ): |
| 56 | self.linear_method = self._apply_weights_sgl |
| 57 | self.process_weights_for_sgl(layer) |
| 58 | else: |
| 59 | self.linear_method = self._apply_weights_onednn |
| 60 | self.process_weights_for_onednn(layer) |
| 61 | |
| 62 | def process_weights_for_onednn(self, layer: torch.nn.Module) -> None: |
| 63 | # WEIGHT |
| 64 | # Transpose to [K, N] for convenience |
| 65 | w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names |
| 66 | weight = getattr(layer, w_q_name) |
| 67 | replace_parameter( |
| 68 | layer, |
| 69 | w_q_name, |
| 70 | torch.nn.Parameter(weight.t().data, requires_grad=False), |
| 71 | ) |
| 72 | |
| 73 | # WEIGHT SCALE |
| 74 | # oneDNN kernels support only per-tensor and per-channel. |
| 75 | # If we have a fused module (QKV, MLP) with per tensor scales (thus N |
| 76 | # scales being passed to the kernel), convert to the per-channel case. |
| 77 | is_fused_module = len(layer.logical_widths) > 1 |
| 78 | weight_scale = getattr(layer, w_s_name) |
| 79 | if is_fused_module and not self.config.is_channelwise: |
| 80 | weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths) |
| 81 | replace_parameter( |
| 82 | layer, |
| 83 | w_s_name, |
| 84 | torch.nn.Parameter(weight_scale.data, requires_grad=False), |
| 85 | ) |
| 86 | |
| 87 | # INPUT SCALE |
| 88 | if self.config.is_static_input_scheme: |
| 89 | input_scale = getattr(layer, i_s_name) |
| 90 | |
| 91 | if self.config.input_symmetric: |
| 92 | replace_parameter( |
| 93 | layer, |
| 94 | i_s_name, |
| 95 | torch.nn.Parameter(input_scale.max(), requires_grad=False), |
| 96 | ) |
| 97 | else: |
| 98 | input_zero_point = getattr(layer, i_zp_name) |
| 99 | |
| 100 | # reconstruct the ranges |
| 101 | int8_traits = torch.iinfo(torch.int8) |
| 102 | azps = input_zero_point.to(dtype=torch.int32) |
| 103 | range_max = (input_scale * (int8_traits.max - azps)).max() |
| 104 | range_min = (input_scale * (int8_traits.min - azps)).min() |
| 105 | |
| 106 | scale = (range_max - range_min) / (int8_traits.max - int8_traits.min) |
| 107 | replace_parameter( |
| 108 | layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False) |
| 109 | ) |
| 110 | |
| 111 | azp = ( |
| 112 | (int8_traits.min - range_min / scale).round().to(dtype=torch.int32) |
| 113 | ) |
| 114 | replace_parameter( |
| 115 | layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False) |
| 116 | ) |
| 117 | |
| 118 | # Different from cutlass, oneDNN kernels only need the AZP adjustment |
| 119 | # term for dynamic quantization. And s_b should be folded into the |
| 120 | # term. Such as: |
| 121 | # s_a * s_b * [(A - zp_a)B] + bias = |
| 122 | # s_a * (s_b * AB) - s_a * s_b * zp_a * B + bias = |
| 123 | # s_a * GEMM_output - s_a * zp_a * adj + bias |
| 124 | if not (self.config.input_symmetric and self.config.is_static_input_scheme): |
| 125 | weight = getattr(layer, w_q_name) |
| 126 | weight_scale = getattr(layer, w_s_name) |
| 127 | azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.float32) |
| 128 | azp_adj = azp_adj * weight_scale.squeeze() |
| 129 | setattr( |
| 130 | layer, |
| 131 | azp_adj_name, |
| 132 | torch.nn.Parameter(azp_adj, requires_grad=False), |
| 133 | ) |
| 134 | |
| 135 | weight = getattr(layer, w_q_name) |
| 136 | self.dnnl_handler = ops.create_onednn_scaled_mm( |
| 137 | weight, |
| 138 | getattr(layer, w_s_name), |
| 139 | torch.get_default_dtype(), |
| 140 | getattr(layer, i_s_name) is None, |
| 141 | not self.config.input_symmetric, |
| 142 | 32, |
| 143 | ) |
| 144 | # weight is prepacked and maintained by the dnnl_handler, |
| 145 | # release the original weight |
| 146 | setattr(layer, w_q_name, None) |
| 147 | del weight |
| 148 | |
| 149 | def process_weights_for_sgl(self, layer: torch.nn.Module) -> None: |
| 150 | w_q_name, w_s_name, _, _, _ = self.layer_param_names |
| 151 | # WEIGHT |
| 152 | weight = getattr(layer, w_q_name) |
| 153 | packed_weight = torch.ops._C.convert_weight_packed(weight) |
| 154 | replace_parameter( |
| 155 | layer, w_q_name, torch.nn.Parameter(packed_weight, requires_grad=False) |
| 156 | ) |
| 157 | |
| 158 | if layer.bias is not None: |
| 159 | bias = layer.bias |
| 160 | layer.register_parameter( |
| 161 | "bias_fp32", torch.nn.Parameter(bias.float().data, requires_grad=False) |
| 162 | ) |
| 163 | |
| 164 | # WEIGHT SCALE |
| 165 | # CPU SGL kernels only support per-channel. |
| 166 | # For per-tensor quant, convert to the per-channel case. |
| 167 | weight_scale = getattr(layer, w_s_name) |
| 168 | if not self.config.is_channelwise: |
| 169 | weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths) |
| 170 | replace_parameter( |
| 171 | layer, |
| 172 | w_s_name, |
| 173 | torch.nn.Parameter(weight_scale.data, requires_grad=False), |
| 174 | ) |
| 175 | |
| 176 | def apply_weights( |
| 177 | self, |
| 178 | layer: torch.nn.Module, |
| 179 | x: torch.Tensor, |
| 180 | bias: torch.Tensor | None = None, |
| 181 | ) -> torch.Tensor: |
| 182 | return self.linear_method( |
| 183 | layer, |
| 184 | x, |
| 185 | bias, |
| 186 | ) |
| 187 | |
| 188 | def _apply_weights_onednn( |
| 189 | self, |
| 190 | layer: torch.nn.Module, |
| 191 | x: torch.Tensor, |
| 192 | bias: torch.Tensor | None = None, |
| 193 | ) -> torch.Tensor: |
| 194 | x_shape = x.shape |
| 195 | x = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x |
| 196 | w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) |
| 197 | |
| 198 | # ops.scaled_int8_quant supports both dynamic and static quant: |
| 199 | # * dynamic, i_s is None and x_s computed from x. |
| 200 | # * static, i_s is scalar and x_s is i_s. |
| 201 | x_q, x_s, x_zp = ops.onednn_scaled_int8_quant( |
| 202 | x, i_s, i_zp, self.config.input_symmetric |
| 203 | ) |
| 204 | |
| 205 | m = x.size(0) |
| 206 | n = self.dnnl_handler.n |
| 207 | out = torch.empty((m, n), dtype=x.dtype) |
| 208 | ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj, bias) |
| 209 | out = out.reshape(x_shape[:-1] + (n,)) if len(x_shape) > 2 else out |
| 210 | return out |
| 211 | |
| 212 | def _apply_weights_sgl( |
| 213 | self, |
| 214 | layer: torch.nn.Module, |
| 215 | x: torch.Tensor, |
| 216 | bias: torch.Tensor | None = None, |
| 217 | ) -> torch.Tensor: |
| 218 | w_q, w_s, _, _, _ = self._get_layer_params(layer) |
| 219 | return torch.ops._C.int8_scaled_mm_with_quant( |
| 220 | x, |
| 221 | w_q, |
| 222 | w_s, |
| 223 | layer.bias_fp32 if bias is not None else None, |
| 224 | x.dtype, |
| 225 | True, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 230 | """FP8 W8A16 block-quantized GEMM via AMX BRGEMM on CPU.""" |
| 231 | |
| 232 | # Input stays BF16 — no FP8 activation quantization. |
| 233 | apply_input_quant = False |
| 234 | |
| 235 | @classmethod |
| 236 | def is_supported( |
| 237 | cls, compute_capability: int | None = None |
| 238 | ) -> tuple[bool, str | None]: |
| 239 | if not current_platform.is_cpu(): |
| 240 | return False, "requires CPU platform." |
| 241 | if not torch.cpu._is_amx_tile_supported(): |
| 242 | return False, "requires AMX tile support (Sapphire Rapids or newer)." |
| 243 | if not ops._supports_cpu_fp8_w8a16: |
| 244 | return False, "fp8_scaled_mm_cpu op not available." |
| 245 | return True, None |
| 246 | |
| 247 | @classmethod |
| 248 | def can_implement( |
| 249 | cls, config: FP8ScaledMMLinearLayerConfig |
| 250 | ) -> tuple[bool, str | None]: |
| 251 | # Validate weight block shape |
| 252 | weight_gs = config.weight_quant_key.scale.group_shape |
| 253 | if weight_gs.col <= 0 or weight_gs.col != 128: |
| 254 | return False, ( |
| 255 | "CPU FP8 kernel requires K-dimension block size of 128, " |
| 256 | f"got {weight_gs.col}." |
| 257 | ) |
| 258 | if weight_gs.row <= 0 or weight_gs.row % 32 != 0: |
| 259 | return False, ( |
| 260 | "CPU FP8 kernel requires N-dimension block size to be " |
| 261 | f"a positive multiple of 32, got {weight_gs.row}." |
| 262 | ) |
| 263 | if config.out_dtype not in (torch.bfloat16, torch.float32): |
| 264 | return False, "Only bfloat16/float32 output dtype supported." |
| 265 | return True, None |
| 266 | |
| 267 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 268 | if getattr(layer, "_cpu_skip_gemm_dispatch", False): |
| 269 | return |
| 270 | # Skip the base class process (FP8 padding / fnuz normalization) |
| 271 | # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. |
| 272 | params = self._get_layer_params(layer) |
| 273 | |
| 274 | # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this |
| 275 | # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and |
| 276 | # let the kernel pack on the fly instead. |
| 277 | if not getattr(layer, "_cpu_skip_gemm_dispatch", False): |
| 278 | packed_weight = torch.ops._C.convert_weight_packed(params.weight) |
| 279 | replace_parameter( |
| 280 | layer, |
| 281 | params.WEIGHT, |
| 282 | torch.nn.Parameter(packed_weight, requires_grad=False), |
| 283 | ) |
| 284 | |
| 285 | # Re-wrap scale as a plain Parameter so the kernel can read it |
| 286 | # without weight-loader metadata interfering. |
| 287 | scale_attr = ( |
| 288 | params.WEIGHT_SCALE_INV |
| 289 | if params.weight_scale_inv is not None |
| 290 | else params.WEIGHT_SCALE |
| 291 | ) |
| 292 | weight_scale = ( |
| 293 | params.weight_scale_inv |
| 294 | if params.weight_scale_inv is not None |
| 295 | else params.weight_scale |
| 296 | ) |
| 297 | assert weight_scale is not None |
| 298 | if weight_scale.dtype in (torch.float8_e8m0fnu, torch.uint8): |
| 299 | weight_scale = _upcast_e8m0_to_fp32(weight_scale.data) |
| 300 | replace_parameter( |
| 301 | layer, |
| 302 | scale_attr, |
| 303 | torch.nn.Parameter(weight_scale.data.contiguous(), requires_grad=False), |
| 304 | ) |
| 305 | |
| 306 | def apply_weights( |
| 307 | self, |
| 308 | layer: torch.nn.Module, |
| 309 | x: torch.Tensor, |
| 310 | bias: torch.Tensor | None = None, |
| 311 | **kwargs, |
| 312 | ) -> torch.Tensor: |
| 313 | params = self._get_layer_params(layer) |
| 314 | weight_scale = ( |
| 315 | params.weight_scale_inv |
| 316 | if params.weight_scale_inv is not None |
| 317 | else params.weight_scale |
| 318 | ) |
| 319 | |
| 320 | x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x |
| 321 | out = torch.ops._C.fp8_scaled_mm_cpu( |
| 322 | x_2d, |
| 323 | params.weight, |
| 324 | weight_scale, |
| 325 | list(self.weight_group_shape), |
| 326 | bias, |
| 327 | x.dtype, |
| 328 | not getattr(layer, "_cpu_skip_gemm_dispatch", False), # is_vnni |
| 329 | ) |
| 330 | return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out |
| 331 | |
| 332 | def apply_block_scaled_mm( |
| 333 | self, |
| 334 | A: torch.Tensor, |
| 335 | B: torch.Tensor, |
| 336 | As: torch.Tensor, |
| 337 | Bs: torch.Tensor, |
| 338 | ) -> torch.Tensor: |
| 339 | raise NotImplementedError( |
| 340 | "CPUFp8BlockScaledMMKernel overrides apply_weights directly." |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 345 | """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU. |
| 346 | |
| 347 | Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single |
| 348 | synthetic block spanning the whole weight, so activations stay BF16/FP32 |
| 349 | — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel. |
| 350 | """ |
| 351 | |
| 352 | @classmethod |
| 353 | def is_supported( |
| 354 | cls, compute_capability: int | None = None |
| 355 | ) -> tuple[bool, str | None]: |
| 356 | if not current_platform.is_cpu(): |
| 357 | return False, "requires CPU platform." |
| 358 | if not torch.cpu._is_amx_tile_supported(): |
| 359 | return False, "requires AMX tile support (Sapphire Rapids or newer)." |
| 360 | if not ops._supports_cpu_fp8_w8a16: |
| 361 | return False, "fp8_scaled_mm_cpu op not available." |
| 362 | return True, None |
| 363 | |
| 364 | @classmethod |
| 365 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 366 | if not c.weight_quant_key.scale.group_shape.is_per_tensor(): |
| 367 | return False, "requires a per-tensor weight scale." |
| 368 | if not c.activation_quant_key.scale.group_shape.is_per_tensor(): |
| 369 | return False, "requires a per-tensor activation scale." |
| 370 | if c.out_dtype not in (torch.bfloat16, torch.float32): |
| 371 | return False, "Only bfloat16/float32 output dtype supported." |
| 372 | n = c.weight_shape[0] |
| 373 | if n % 32 != 0: |
| 374 | # The AMX tinygemm kernel tiles N in chunks of 32 and cannot |
| 375 | # handle a remainder tile smaller than that. |
| 376 | return False, f"requires weight output dim (N={n}) to be a multiple of 32." |
| 377 | return True, None |
| 378 | |
| 379 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 380 | w_name, w_s_name, _, _ = self.layer_param_names |
| 381 | # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style |
| 382 | # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as |
| 383 | # produced by CUDA nn.Linear-style checkpoints. |
| 384 | weight = getattr(layer, w_name).t().contiguous() |
| 385 | n, k = weight.shape |
| 386 | |
| 387 | packed_weight = torch.ops._C.convert_weight_packed(weight) |
| 388 | replace_parameter( |
| 389 | layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False) |
| 390 | ) |
| 391 | |
| 392 | # Synthesize a single-block "block scale" tensor from the scalar |
| 393 | # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is. |
| 394 | weight_scale = getattr(layer, w_s_name) |
| 395 | self._block_size_n = -(-n // 32) * 32 # round up to a multiple of 32 |
| 396 | num_k_blocks = -(-k // _BLOCK_SIZE_K) |
| 397 | block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous() |
| 398 | replace_parameter( |
| 399 | layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False) |
| 400 | ) |
| 401 | |
| 402 | def apply_weights( |
| 403 | self, |
| 404 | layer: torch.nn.Module, |
| 405 | x: torch.Tensor, |
| 406 | bias: torch.Tensor | None = None, |
| 407 | ) -> torch.Tensor: |
| 408 | w_name, w_s_name, _, _ = self.layer_param_names |
| 409 | weight = getattr(layer, w_name) |
| 410 | weight_scale = getattr(layer, w_s_name) |
| 411 | |
| 412 | x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x |
| 413 | out = torch.ops._C.fp8_scaled_mm_cpu( |
| 414 | x_2d, |
| 415 | weight, |
| 416 | weight_scale, |
| 417 | [self._block_size_n, _BLOCK_SIZE_K], |
| 418 | bias, |
| 419 | x.dtype, |
| 420 | True, # is_vnni (weight already prepacked) |
| 421 | ) |
| 422 | return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out |
| 423 | |
| 424 | def apply_scaled_mm( |
| 425 | self, |
| 426 | *, |
| 427 | A: torch.Tensor, |
| 428 | B: torch.Tensor, |
| 429 | out_dtype: torch.dtype, |
| 430 | As: torch.Tensor, |
| 431 | Bs: torch.Tensor, |
| 432 | bias: torch.Tensor | None, |
| 433 | output_shape: list, |
| 434 | ) -> torch.Tensor: |
| 435 | raise NotImplementedError( |
| 436 | "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly." |
| 437 | ) |
| 438 | |
| 439 | |
| 440 | class CPUFp8W8A8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 441 | """FP8 W8A8 block-quantized GEMM with dynamic per-token activation quant on CPU.""" |
| 442 | |
| 443 | # Activation quant is done inside fp8_scaled_mm_with_quant; |
| 444 | # the base class must not apply a second quant step. |
| 445 | apply_input_quant = False |
| 446 | |
| 447 | @classmethod |
| 448 | def is_supported( |
| 449 | cls, compute_capability: int | None = None |
| 450 | ) -> tuple[bool, str | None]: |
| 451 | if not current_platform.is_cpu(): |
| 452 | return False, "requires CPU platform." |
| 453 | if not torch.cpu._is_amx_tile_supported(): |
| 454 | return False, "requires AMX tile support (Sapphire Rapids or newer)." |
| 455 | if not ops.cpu_has_amx_fp8(): |
| 456 | return False, ( |
| 457 | "requires native AMX-FP8 hardware; falls back " |
| 458 | "to CPUFp8BlockScaledMMKernel (W8A16) otherwise." |
| 459 | ) |
| 460 | return True, None |
| 461 | |
| 462 | @classmethod |
| 463 | def can_implement( |
| 464 | cls, config: FP8ScaledMMLinearLayerConfig |
| 465 | ) -> tuple[bool, str | None]: |
| 466 | # Same block-shape requirements as the W8A16 sibling kernel. |
| 467 | weight_gs = config.weight_quant_key.scale.group_shape |
| 468 | if weight_gs.col <= 0 or weight_gs.col != 128: |
| 469 | return False, ( |
| 470 | "CPU FP8 W8A8 block kernel requires K-dimension block size of 128, " |
| 471 | f"got {weight_gs.col}." |
| 472 | ) |
| 473 | if weight_gs.row <= 0 or weight_gs.row % 32 != 0: |
| 474 | return False, ( |
| 475 | "CPU FP8 W8A8 block kernel requires N-dimension block size to be " |
| 476 | f"a positive multiple of 32, got {weight_gs.row}." |
| 477 | ) |
| 478 | if config.out_dtype not in (torch.bfloat16, torch.float32): |
| 479 | return False, "Only bfloat16/float32 output dtype supported." |
| 480 | n = config.weight_shape[0] |
| 481 | if n % 32 != 0: |
| 482 | # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot |
| 483 | # handle a remainder tile smaller than that. |
| 484 | return False, f"requires weight output dim (N={n}) to be a multiple of 32." |
| 485 | return True, None |
| 486 | |
| 487 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 488 | """Prepack block-quantized FP8 weight for W8A8 AMX BRGEMM. |
| 489 | |
| 490 | For block-quant FP8 models weight is stored as [N, K] (NOT transposed, |
| 491 | unlike non-block FP8 where vLLM stores [K, N]). |
| 492 | |
| 493 | float8_linear_prepack_cpu expects: |
| 494 | weight: [N, K] |
| 495 | scales: [N, G] where G = K / block_k |
| 496 | |
| 497 | weight_scale_inv from the checkpoint has shape [ceil(N/block_n), G], |
| 498 | so we expand rows from ceil(N/block_n) to N via repeat_interleave. |
| 499 | """ |
| 500 | params = self._get_layer_params(layer) |
| 501 | weight = params.weight # [N, K] |
| 502 | |
| 503 | # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this |
| 504 | # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and |
| 505 | # its block scales; apply_weights falls back to the W8A16 op. |
| 506 | if getattr(layer, "_cpu_skip_gemm_dispatch", False): |
| 507 | return |
| 508 | |
| 509 | weight_scale = ( |
| 510 | params.weight_scale_inv |
| 511 | if params.weight_scale_inv is not None |
| 512 | else params.weight_scale |
| 513 | ) # [ceil(N/block_n), ceil(K/block_k)] |
| 514 | assert weight_scale is not None, "FP8 W8A8 block kernel requires weight scale." |
| 515 | |
| 516 | N = weight.shape[0] |
| 517 | block_n = self.weight_group_shape.row # e.g. 128 |
| 518 | |
| 519 | # Expand scale rows: [ceil(N/block_n), G] → [N, G] |
| 520 | if weight_scale.size(0) != N: |
| 521 | weight_scale_expanded = weight_scale.repeat_interleave(block_n, dim=0)[ |
| 522 | :N, : |
| 523 | ].contiguous() |
| 524 | else: |
| 525 | weight_scale_expanded = weight_scale.contiguous() |
| 526 | |
| 527 | packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu( |
| 528 | weight, weight_scale_expanded.to(torch.float32) |
| 529 | ) |
| 530 | |
| 531 | scale_attr = ( |
| 532 | params.WEIGHT_SCALE_INV |
| 533 | if params.weight_scale_inv is not None |
| 534 | else params.WEIGHT_SCALE |
| 535 | ) |
| 536 | replace_parameter( |
| 537 | layer, |
| 538 | params.WEIGHT, |
| 539 | torch.nn.Parameter(packed_weight, requires_grad=False), |
| 540 | ) |
| 541 | replace_parameter( |
| 542 | layer, |
| 543 | scale_attr, |
| 544 | torch.nn.Parameter(packed_scales, requires_grad=False), |
| 545 | ) |
| 546 | |
| 547 | def apply_weights( |
| 548 | self, |
| 549 | layer: torch.nn.Module, |
| 550 | x: torch.Tensor, |
| 551 | bias: torch.Tensor | None = None, |
| 552 | **kwargs, |
| 553 | ) -> torch.Tensor: |
| 554 | params = self._get_layer_params(layer) |
| 555 | weight_scale = ( |
| 556 | params.weight_scale_inv |
| 557 | if params.weight_scale_inv is not None |
| 558 | else params.weight_scale |
| 559 | ) |
| 560 | |
| 561 | x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x |
| 562 | if getattr(layer, "_cpu_skip_gemm_dispatch", False): |
| 563 | # Unpacked weight: fp8_scaled_mm_with_quant has no is_vnni option, |
| 564 | # so fall back to the W8A16 op (BF16 activation, block-scaled FP8 |
| 565 | # weight) for this layer. |
| 566 | out = torch.ops._C.fp8_scaled_mm_cpu( |
| 567 | x_2d, |
| 568 | params.weight, |
| 569 | weight_scale, |
| 570 | list(self.weight_group_shape), |
| 571 | bias, |
| 572 | x.dtype, |
| 573 | False, # is_vnni |
| 574 | ) |
| 575 | return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out |
| 576 | |
| 577 | # act_scales=None + channelwise=True → dynamic per-token activation quant |
| 578 | out = ops.fp8_scaled_mm_with_quant( |
| 579 | x_2d, |
| 580 | None, |
| 581 | True, |
| 582 | params.weight, |
| 583 | weight_scale, |
| 584 | bias, |
| 585 | x.dtype, |
| 586 | ) |
| 587 | return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out |
| 588 | |
| 589 | def apply_block_scaled_mm( |
| 590 | self, |
| 591 | A: torch.Tensor, |
| 592 | B: torch.Tensor, |
| 593 | As: torch.Tensor, |
| 594 | Bs: torch.Tensor, |
| 595 | ) -> torch.Tensor: |
| 596 | raise NotImplementedError( |
| 597 | "CPUFp8W8A8BlockScaledMMKernel overrides apply_weights directly." |
| 598 | ) |
| 599 | |
| 600 | |
| 601 | class CPUFP8W8A8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 602 | """FP8 W8A8 GEMM with dynamic per-token activation quantization on CPU.""" |
| 603 | |
| 604 | @classmethod |
| 605 | def is_supported( |
| 606 | cls, compute_capability: int | None = None |
| 607 | ) -> tuple[bool, str | None]: |
| 608 | if not current_platform.is_cpu(): |
| 609 | return False, "requires CPU platform." |
| 610 | if not torch.cpu._is_amx_tile_supported(): |
| 611 | return False, "requires AMX tile support (Sapphire Rapids or newer)." |
| 612 | if not ops.cpu_has_amx_fp8(): |
| 613 | return False, ( |
| 614 | "requires native AMX-FP8 hardware (Diamond Rapids); falls back " |
| 615 | "to the Torch FP8 kernels otherwise." |
| 616 | ) |
| 617 | return True, None |
| 618 | |
| 619 | @classmethod |
| 620 | def can_implement( |
| 621 | cls, config: FP8ScaledMMLinearLayerConfig |
| 622 | ) -> tuple[bool, str | None]: |
| 623 | if config.out_dtype not in (torch.bfloat16, torch.float16, torch.float32): |
| 624 | return False, "Only bfloat16/float16/float32 output dtype supported." |
| 625 | # Reject block-quantized weights — those go to CPUFp8BlockScaledMMKernel |
| 626 | weight_gs = config.weight_quant_key.scale.group_shape |
| 627 | if weight_gs.col > 0 and weight_gs.row > 0: |
| 628 | return False, "Block-quantized weights not handled by FP8 W8A8 kernel." |
| 629 | n = config.weight_shape[0] |
| 630 | if n % 32 != 0: |
| 631 | # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot |
| 632 | # handle a remainder tile smaller than that. |
| 633 | return False, f"requires weight output dim (N={n}) to be a multiple of 32." |
| 634 | return True, None |
| 635 | |
| 636 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 637 | """Prepack weights with float8_linear_prepack_cpu (VNNI + block layout).""" |
| 638 | w_name, ws_name, xs_name, _ = self.layer_param_names |
| 639 | weight = getattr(layer, w_name) |
| 640 | weight_scale = getattr(layer, ws_name) |
| 641 | |
| 642 | # vLLM stores linear weight as [K, N], while the CPU prepack op expects [N, K]. |
| 643 | N = weight.shape[-1] |
| 644 | |
| 645 | # Ensure scale is 2D [N, G] |
| 646 | ws = weight_scale |
| 647 | if ws is None: |
| 648 | raise RuntimeError("FP8 W8A8 kernel requires weight scale.") |
| 649 | ws = ws.to(torch.float32) |
| 650 | |
| 651 | # Fused modules may carry per-shard scales (e.g. QKV/MLP shards). |
| 652 | # Expand those to per-output-channel scales before prepack. |
| 653 | if ( |
| 654 | len(getattr(layer, "logical_widths", [])) > 1 |
| 655 | and ws.dim() <= 1 |
| 656 | and ws.numel() == len(layer.logical_widths) |
| 657 | ): |
| 658 | ws = convert_to_channelwise(ws.reshape(-1), layer.logical_widths) |
| 659 | |
| 660 | if ws.dim() == 0: |
| 661 | ws = ws.view(1, 1).expand(N, 1) |
| 662 | elif ws.dim() == 1: |
| 663 | if ws.numel() == 1: |
| 664 | ws = ws.view(1, 1).expand(N, 1) |
| 665 | elif ws.numel() == N: |
| 666 | ws = ws.unsqueeze(1) # [N] -> [N, 1] |
| 667 | else: |
| 668 | raise RuntimeError( |
| 669 | "FP8 W8A8 kernel got unsupported 1D weight scale shape: " |
| 670 | f"{tuple(ws.shape)} for N={N}." |
| 671 | ) |
| 672 | |
| 673 | # Align scale rows to output channels N. |
| 674 | if ws.size(0) != N: |
| 675 | if ws.size(0) == 1: |
| 676 | ws = ws.expand(N, ws.size(1)).contiguous() |
| 677 | else: |
| 678 | raise RuntimeError( |
| 679 | "FP8 W8A8 kernel got mismatched weight scale rows: " |
| 680 | f"ws_rows={ws.size(0)}, expected N={N}." |
| 681 | ) |
| 682 | |
| 683 | w_2d = weight.t().contiguous() # [N, K] |
| 684 | packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(w_2d, ws) |
| 685 | |
| 686 | replace_parameter( |
| 687 | layer, |
| 688 | w_name, |
| 689 | torch.nn.Parameter(packed_weight, requires_grad=False), |
| 690 | ) |
| 691 | replace_parameter( |
| 692 | layer, |
| 693 | ws_name, |
| 694 | torch.nn.Parameter(packed_scales, requires_grad=False), |
| 695 | ) |
| 696 | |
| 697 | def apply_weights( |
| 698 | self, |
| 699 | layer: torch.nn.Module, |
| 700 | x: torch.Tensor, |
| 701 | bias: torch.Tensor | None = None, |
| 702 | **kwargs, |
| 703 | ) -> torch.Tensor: |
| 704 | w_name, ws_name, xs_name, _ = self.layer_param_names |
| 705 | weight = getattr(layer, w_name) |
| 706 | weight_scale = getattr(layer, ws_name) |
| 707 | # For static quantization, use pre-computed activation scale |
| 708 | act_scales = getattr(layer, xs_name, None) if xs_name else None |
| 709 | |
| 710 | x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x |
| 711 | # channelwise=True: per-token (per-row) quantization when act_scales is None |
| 712 | # channelwise=False: per-tensor quantization when act_scales is provided |
| 713 | channelwise = act_scales is None |
| 714 | out = ops.fp8_scaled_mm_with_quant( |
| 715 | x_2d, |
| 716 | act_scales, |
| 717 | channelwise, |
| 718 | weight, |
| 719 | weight_scale, |
| 720 | bias, |
| 721 | x.dtype, |
| 722 | ) |
| 723 | return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out |
| 724 | |
| 725 | def apply_scaled_mm(self, A, B, out_dtype, As, Bs, bias=None, **kwargs): |
| 726 | raise NotImplementedError( |
| 727 | "CPUFP8W8A8ScaledMMLinearKernel uses apply_weights directly." |
| 728 | ) |