| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import torch |
| 7 | from torch.nn import Module |
| 8 | |
| 9 | if TYPE_CHECKING: |
| 10 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 11 | from vllm.model_executor.layers.fused_moe.config import ( |
| 12 | FusedMoEQuantConfig, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 16 | |
| 17 | import vllm.envs as envs |
| 18 | from vllm import _custom_ops as ops |
| 19 | from vllm.config import get_current_vllm_config |
| 20 | from vllm.model_executor.kernels.linear import init_fp8_linear_kernel |
| 21 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 22 | CutlassFP8ScaledMMLinearKernel, |
| 23 | MarlinFP8ScaledMMLinearKernel, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 26 | 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 | ) |
| 37 | from vllm.model_executor.layers.quantization.online.moe_base import ( |
| 38 | OnlineMoEMethodBase, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 41 | GroupShape, |
| 42 | amax_for_moe_weight_quant, |
| 43 | amax_for_tp_weight_quant, |
| 44 | create_fp8_quant_key, |
| 45 | get_fp8_min_max, |
| 46 | kFp8Dynamic128Sym, |
| 47 | kFp8DynamicTensorSym, |
| 48 | kFp8DynamicTokenSym, |
| 49 | kFp8Static128BlockSym, |
| 50 | kFp8StaticChannelSym, |
| 51 | kFp8StaticTensorSym, |
| 52 | weight_amax, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 55 | cutlass_fp8_supported, |
| 56 | ) |
| 57 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 58 | initialize_online_processing, |
| 59 | ) |
| 60 | from vllm.model_executor.parameter import ModelWeightParameter |
| 61 | from vllm.model_executor.utils import replace_parameter |
| 62 | from vllm.platforms import current_platform |
| 63 | from vllm.utils.deep_gemm import per_block_cast_to_fp8 |
| 64 | from vllm.utils.math_utils import round_up |
| 65 | |
| 66 | # --------------------------------------------------------------------------- |
| 67 | # Online FP8 Linear Methods |
| 68 | # --------------------------------------------------------------------------- |
| 69 | |
| 70 | |
| 71 | def _fp8_max(device: torch.device) -> torch.Tensor: |
| 72 | # A 0-d tensor divisor keeps the true division the quant kernels do; |
| 73 | # dividing by a Python float lowers to multiply-by-reciprocal instead. |
| 74 | return torch.tensor(get_fp8_min_max()[1], device=device, dtype=torch.float32) |
| 75 | |
| 76 | |
| 77 | def _fp8_scale(amax: torch.Tensor) -> torch.Tensor: |
| 78 | amax = amax.to(torch.float32) |
| 79 | return amax / _fp8_max(amax.device) |
| 80 | |
| 81 | |
| 82 | def _fp8_channel_scale(amax: torch.Tensor) -> torch.Tensor: |
| 83 | fp8_max = _fp8_max(amax.device) |
| 84 | return _fp8_scale(amax).clamp_min(1.0 / (fp8_max * 512)) |
| 85 | |
| 86 | |
| 87 | # Chunk so the fp32 intermediate from the divide stays under ~64 MB, regardless |
| 88 | # of how wide the weight is. |
| 89 | _QUANT_CHUNK_ELEMS = 16 * 1024 * 1024 |
| 90 | |
| 91 | |
| 92 | def _fp8_quant_per_channel(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| 93 | # True division, matching the dynamic per-token kernel; the static kernel |
| 94 | # multiplies by the reciprocal instead and can round a value that sits on a |
| 95 | # representable midpoint the other way. |
| 96 | fp8_min, fp8_max = get_fp8_min_max() |
| 97 | out = torch.empty_like(weight, dtype=current_platform.fp8_dtype()) |
| 98 | chunk_rows = max(1, _QUANT_CHUNK_ELEMS // weight.shape[-1]) |
| 99 | for start in range(0, weight.shape[0], chunk_rows): |
| 100 | rows = slice(start, start + chunk_rows) |
| 101 | out[rows] = (weight[rows] / scale[rows]).clamp_(fp8_min, fp8_max).to(out.dtype) |
| 102 | return out |
| 103 | |
| 104 | |
| 105 | def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: |
| 106 | """Whether the weight is sharded along a dim its ``amax`` reduces over. |
| 107 | |
| 108 | Row parallel shards the input dim, column parallel the output dim. |
| 109 | Per-output-channel scales only reduce the input dim, so they are already |
| 110 | unsharded-equivalent on column-parallel layers. |
| 111 | """ |
| 112 | if layer.tp_size == 1: |
| 113 | return False |
| 114 | is_row_parallel = layer.input_size != layer.input_size_per_partition |
| 115 | is_column_parallel = layer.output_size != layer.output_size_per_partition |
| 116 | return is_row_parallel or (reduces_output_dim and is_column_parallel) |
| 117 | |
| 118 | |
| 119 | class OnlineLinearBase(LinearMethodBase): |
| 120 | """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint |
| 121 | weights onto meta device and materializes them just-in-time.""" |
| 122 | |
| 123 | uses_meta_device: bool = True |
| 124 | |
| 125 | def __init__(self): |
| 126 | self.out_dtype = torch.get_default_dtype() |
| 127 | self.input_dtype = get_current_vllm_config().model_config.dtype |
| 128 | |
| 129 | def create_weights( |
| 130 | self, |
| 131 | layer: torch.nn.Module, |
| 132 | input_size_per_partition: int, |
| 133 | output_partition_sizes: list[int], |
| 134 | input_size: int, |
| 135 | output_size: int, |
| 136 | params_dtype: torch.dtype, |
| 137 | **extra_weight_attrs, |
| 138 | ): |
| 139 | output_size_per_partition = sum(output_partition_sizes) |
| 140 | weight_loader = extra_weight_attrs.get("weight_loader") |
| 141 | layer.logical_widths = output_partition_sizes |
| 142 | layer.input_size_per_partition = input_size_per_partition |
| 143 | layer.output_size_per_partition = output_size_per_partition |
| 144 | layer.orig_dtype = params_dtype |
| 145 | layer.weight_block_size = None |
| 146 | |
| 147 | weight = ModelWeightParameter( |
| 148 | data=torch.empty( |
| 149 | output_size_per_partition, |
| 150 | input_size_per_partition, |
| 151 | device="meta", # materialized and processed during loading |
| 152 | dtype=params_dtype, |
| 153 | ), |
| 154 | input_dim=1, |
| 155 | output_dim=0, |
| 156 | weight_loader=weight_loader, |
| 157 | ) |
| 158 | layer.register_parameter("weight", weight) |
| 159 | |
| 160 | initialize_online_processing(layer) |
| 161 | |
| 162 | |
| 163 | class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): |
| 164 | """Online tensorwise FP8 linear quantization. |
| 165 | Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" |
| 166 | |
| 167 | def __init__(self): |
| 168 | super().__init__() |
| 169 | |
| 170 | self.block_quant = False |
| 171 | self.use_deep_gemm = False |
| 172 | self.use_marlin = False |
| 173 | self.marlin_input_dtype = None |
| 174 | self.weight_quant_key = kFp8StaticTensorSym |
| 175 | # Use per-token quantization for better perf if dynamic and cutlass |
| 176 | if cutlass_fp8_supported(): |
| 177 | self.activation_quant_key = kFp8DynamicTokenSym |
| 178 | else: |
| 179 | self.activation_quant_key = kFp8DynamicTensorSym |
| 180 | |
| 181 | def create_weights( |
| 182 | self, |
| 183 | layer: torch.nn.Module, |
| 184 | input_size_per_partition: int, |
| 185 | output_partition_sizes: list[int], |
| 186 | input_size: int, |
| 187 | output_size: int, |
| 188 | params_dtype: torch.dtype, |
| 189 | **extra_weight_attrs, |
| 190 | ): |
| 191 | super().create_weights( |
| 192 | layer, |
| 193 | input_size_per_partition, |
| 194 | output_partition_sizes, |
| 195 | input_size, |
| 196 | output_size, |
| 197 | params_dtype, |
| 198 | **extra_weight_attrs, |
| 199 | ) |
| 200 | |
| 201 | self.fp8_linear = init_fp8_linear_kernel( |
| 202 | activation_quant_key=self.activation_quant_key, |
| 203 | weight_quant_key=self.weight_quant_key, |
| 204 | weight_shape=layer.weight.shape, |
| 205 | input_dtype=self.input_dtype, |
| 206 | out_dtype=self.out_dtype, |
| 207 | module_name=self.__class__.__name__, |
| 208 | ) |
| 209 | self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) |
| 210 | |
| 211 | def process_weights_after_loading(self, layer: Module) -> None: |
| 212 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 213 | return |
| 214 | |
| 215 | layer.input_scale = None |
| 216 | amax = weight_amax(layer.weight).reshape(1) |
| 217 | amax = amax_for_tp_weight_quant(amax, _is_tp_sharded(layer)) |
| 218 | weight_scale = _fp8_scale(amax) |
| 219 | qweight, _ = ops.scaled_fp8_quant(layer.weight, scale=weight_scale) |
| 220 | |
| 221 | # Update layer with new values. |
| 222 | replace_parameter(layer, "weight", qweight.t().data) |
| 223 | replace_parameter(layer, "weight_scale", weight_scale.data) |
| 224 | |
| 225 | if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"): |
| 226 | self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype |
| 227 | self.fp8_linear.process_weights_after_loading(layer) |
| 228 | |
| 229 | # Prevent duplicate processing (e.g., during weight reload) |
| 230 | layer._already_called_process_weights_after_loading = True |
| 231 | |
| 232 | def apply( |
| 233 | self, |
| 234 | layer: torch.nn.Module, |
| 235 | x: torch.Tensor, |
| 236 | bias: torch.Tensor | None = None, |
| 237 | ) -> torch.Tensor: |
| 238 | # if batch invariant mode is enabled, use BF16 dequant |
| 239 | if envs.VLLM_BATCH_INVARIANT: |
| 240 | if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): |
| 241 | return self.fp8_linear.apply_weights(layer, x, bias) |
| 242 | |
| 243 | weight_fp8 = layer.weight.to(torch.bfloat16) |
| 244 | weight_scale = layer.weight_scale.to(torch.bfloat16) |
| 245 | if weight_scale.numel() == 1: |
| 246 | # Per-tensor: simple scalar multiplication |
| 247 | weight_bf16 = weight_fp8 * weight_scale |
| 248 | else: |
| 249 | # Multiple scales (fused modules like QKV) |
| 250 | if ( |
| 251 | weight_scale.dim() == 1 |
| 252 | and weight_scale.shape[0] == weight_fp8.shape[0] |
| 253 | ): |
| 254 | # Per-row scaling |
| 255 | weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1) |
| 256 | else: |
| 257 | # Fallback |
| 258 | weight_bf16 = weight_fp8 * weight_scale |
| 259 | return torch.nn.functional.linear(x, weight_bf16.t(), bias) |
| 260 | |
| 261 | return self.fp8_linear.apply_weights(layer, x, bias) |
| 262 | |
| 263 | |
| 264 | class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): |
| 265 | """Online blockwise FP8 linear quantization. |
| 266 | Loads fp16/bf16 weights and quantizes them per-block during loading.""" |
| 267 | |
| 268 | def __init__(self): |
| 269 | super().__init__() |
| 270 | self.weight_block_size = [128, 128] |
| 271 | self.activation_quant_key = create_fp8_quant_key( |
| 272 | static=False, |
| 273 | group_shape=GroupShape(1, self.weight_block_size[0]), |
| 274 | ) |
| 275 | self.weight_quant_key = create_fp8_quant_key( |
| 276 | static=True, group_shape=GroupShape(*self.weight_block_size) |
| 277 | ) |
| 278 | |
| 279 | def create_weights( |
| 280 | self, |
| 281 | layer: torch.nn.Module, |
| 282 | input_size_per_partition: int, |
| 283 | output_partition_sizes: list[int], |
| 284 | input_size: int, |
| 285 | output_size: int, |
| 286 | params_dtype: torch.dtype, |
| 287 | **extra_weight_attrs, |
| 288 | ): |
| 289 | super().create_weights( |
| 290 | layer, |
| 291 | input_size_per_partition, |
| 292 | output_partition_sizes, |
| 293 | input_size, |
| 294 | output_size, |
| 295 | params_dtype, |
| 296 | **extra_weight_attrs, |
| 297 | ) |
| 298 | layer.weight_block_size = self.weight_block_size |
| 299 | |
| 300 | self.fp8_linear = init_fp8_linear_kernel( |
| 301 | activation_quant_key=self.activation_quant_key, |
| 302 | weight_quant_key=self.weight_quant_key, |
| 303 | weight_shape=layer.weight.shape, |
| 304 | input_dtype=self.input_dtype, |
| 305 | out_dtype=self.out_dtype, |
| 306 | module_name=self.__class__.__name__, |
| 307 | ) |
| 308 | |
| 309 | def process_weights_after_loading(self, layer: Module) -> None: |
| 310 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 311 | return |
| 312 | |
| 313 | layer.input_scale = None |
| 314 | block_size = self.weight_block_size |
| 315 | |
| 316 | qweight, weight_scale_inv = per_block_cast_to_fp8( |
| 317 | layer.weight, block_size=block_size, use_ue8m0=False |
| 318 | ) |
| 319 | |
| 320 | replace_parameter(layer, "weight", qweight.data) |
| 321 | replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data) |
| 322 | |
| 323 | self.fp8_linear.process_weights_after_loading(layer) |
| 324 | |
| 325 | # Prevent duplicate processing (e.g., during weight reload) |
| 326 | layer._already_called_process_weights_after_loading = True |
| 327 | |
| 328 | def apply( |
| 329 | self, |
| 330 | layer: torch.nn.Module, |
| 331 | x: torch.Tensor, |
| 332 | bias: torch.Tensor | None = None, |
| 333 | ) -> torch.Tensor: |
| 334 | assert self.weight_block_size is not None |
| 335 | |
| 336 | # Note: batch invariance already handled in the function below |
| 337 | return self.fp8_linear.apply_weights( |
| 338 | layer, |
| 339 | x, |
| 340 | bias, |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | class Fp8PtpcOnlineLinearMethod(OnlineLinearBase): |
| 345 | """Online PTPC FP8 linear quantization. |
| 346 | |
| 347 | Per-output-channel weight scale + dynamic per-token activation scale. The |
| 348 | layout matches the llmcompressor's FP8_DYNAMIC recipe, so accuracy |
| 349 | is comparable but no pre-quantized checkpoint is required. |
| 350 | """ |
| 351 | |
| 352 | weight_quant_key = kFp8StaticChannelSym |
| 353 | activation_quant_key = kFp8DynamicTokenSym |
| 354 | |
| 355 | def create_weights( |
| 356 | self, |
| 357 | layer: torch.nn.Module, |
| 358 | input_size_per_partition: int, |
| 359 | output_partition_sizes: list[int], |
| 360 | input_size: int, |
| 361 | output_size: int, |
| 362 | params_dtype: torch.dtype, |
| 363 | **extra_weight_attrs, |
| 364 | ): |
| 365 | super().create_weights( |
| 366 | layer, |
| 367 | input_size_per_partition, |
| 368 | output_partition_sizes, |
| 369 | input_size, |
| 370 | output_size, |
| 371 | params_dtype, |
| 372 | **extra_weight_attrs, |
| 373 | ) |
| 374 | |
| 375 | self.fp8_linear = init_fp8_linear_kernel( |
| 376 | activation_quant_key=self.activation_quant_key, |
| 377 | weight_quant_key=self.weight_quant_key, |
| 378 | weight_shape=layer.weight.shape, |
| 379 | input_dtype=self.input_dtype, |
| 380 | out_dtype=self.out_dtype, |
| 381 | module_name=self.__class__.__name__, |
| 382 | ) |
| 383 | # PTPC requires per-token activation FP8; MarlinFP8 is W8A16 and |
| 384 | # would silently produce a weight-only fp8 model. |
| 385 | if isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel): |
| 386 | raise ValueError( |
| 387 | "FP8 PTPC online quant requires a kernel that honors " |
| 388 | "per-token activation quantization; MarlinFP8 is W8A16 " |
| 389 | "weight-only. Requires SM89+ for Cutlass FP8 or ROCm MI3xx " |
| 390 | "for rowwise scaled_mm." |
| 391 | ) |
| 392 | |
| 393 | def process_weights_after_loading(self, layer: Module) -> None: |
| 394 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 395 | return |
| 396 | |
| 397 | layer.input_scale = None |
| 398 | amax = weight_amax(layer.weight, dim=-1, keepdim=True) |
| 399 | amax = amax_for_tp_weight_quant( |
| 400 | amax, _is_tp_sharded(layer, reduces_output_dim=False) |
| 401 | ) |
| 402 | weight_scale = _fp8_channel_scale(amax) |
| 403 | qweight = _fp8_quant_per_channel(layer.weight, weight_scale) |
| 404 | |
| 405 | replace_parameter(layer, "weight", qweight.t()) |
| 406 | replace_parameter(layer, "weight_scale", weight_scale) |
| 407 | |
| 408 | self.fp8_linear.process_weights_after_loading(layer) |
| 409 | expose_input_quant_key(layer, self.fp8_linear) |
| 410 | |
| 411 | layer._already_called_process_weights_after_loading = True |
| 412 | |
| 413 | def apply( |
| 414 | self, |
| 415 | layer: torch.nn.Module, |
| 416 | x: torch.Tensor, |
| 417 | bias: torch.Tensor | None = None, |
| 418 | ) -> torch.Tensor: |
| 419 | # if batch invariant mode is enabled dequant |
| 420 | if ( |
| 421 | envs.VLLM_BATCH_INVARIANT |
| 422 | and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) |
| 423 | and not isinstance(x, QuantizedActivation) |
| 424 | ): |
| 425 | weight_dequant = ( |
| 426 | layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t() |
| 427 | ) |
| 428 | return torch.nn.functional.linear(x, weight_dequant.t(), bias) |
| 429 | |
| 430 | return self.fp8_linear.apply_weights(layer, x, bias) |
| 431 | |
| 432 | |
| 433 | # --------------------------------------------------------------------------- |
| 434 | # Online FP8 MoE Methods |
| 435 | # --------------------------------------------------------------------------- |
| 436 | |
| 437 | |
| 438 | class _Fp8OnlineMoEBase(OnlineMoEMethodBase): |
| 439 | """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint |
| 440 | weights onto meta device and materializes them just-in-time.""" |
| 441 | |
| 442 | # Declared here for mypy; actual values are set in __init__. |
| 443 | fp8_backend: "Fp8MoeBackend" |
| 444 | experts_cls: "type[mk.FusedMoEExperts] | None" |
| 445 | weight_scale_name: str |
| 446 | weight_block_size: list[int] | None |
| 447 | per_act_token_quant: bool = False |
| 448 | per_out_ch_quant: bool = False |
| 449 | |
| 450 | def __init__( |
| 451 | self, |
| 452 | *, |
| 453 | weight_block_size: list[int] | None, |
| 454 | moe: FusedMoEConfig, |
| 455 | weight_key: "QuantKey | None" = None, |
| 456 | activation_key: "QuantKey | None" = None, |
| 457 | allow_vllm_cutlass: bool = False, |
| 458 | ): |
| 459 | super().__init__(moe) |
| 460 | self.weight_block_size = weight_block_size |
| 461 | self.block_quant: bool = self.weight_block_size is not None |
| 462 | self.weight_scale_name = ( |
| 463 | "weight_scale_inv" if self.block_quant else "weight_scale" |
| 464 | ) |
| 465 | |
| 466 | # Subclasses may pass explicit kernel keys (PTPC needs channelwise + |
| 467 | # per-token). |
| 468 | if weight_key is None or activation_key is None: |
| 469 | if self.block_quant: |
| 470 | weight_key = kFp8Static128BlockSym |
| 471 | activation_key = kFp8Dynamic128Sym |
| 472 | else: |
| 473 | weight_key = kFp8StaticTensorSym |
| 474 | activation_key = kFp8DynamicTensorSym |
| 475 | |
| 476 | # Select Fp8 MoE backend |
| 477 | self.fp8_backend, self.experts_cls = select_fp8_moe_backend( |
| 478 | config=self.moe, |
| 479 | weight_key=weight_key, |
| 480 | activation_key=activation_key, |
| 481 | allow_vllm_cutlass=allow_vllm_cutlass, |
| 482 | ) |
| 483 | |
| 484 | def _setup_kernel( |
| 485 | self, |
| 486 | layer: RoutedExperts, |
| 487 | w13: torch.Tensor, |
| 488 | w2: torch.Tensor, |
| 489 | w13_scale: torch.Tensor, |
| 490 | w2_scale: torch.Tensor, |
| 491 | w13_input_scale: torch.Tensor | None, |
| 492 | w2_input_scale: torch.Tensor | None, |
| 493 | ) -> None: |
| 494 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 495 | convert_to_fp8_moe_kernel_format, |
| 496 | make_fp8_moe_kernel, |
| 497 | ) |
| 498 | |
| 499 | # Shuffle weights to runtime format. |
| 500 | w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 501 | fp8_backend=self.fp8_backend, |
| 502 | layer=layer, |
| 503 | w13=w13, |
| 504 | w2=w2, |
| 505 | w13_scale=w13_scale, |
| 506 | w2_scale=w2_scale, |
| 507 | w13_input_scale=w13_input_scale, |
| 508 | w2_input_scale=w2_input_scale, |
| 509 | ) |
| 510 | |
| 511 | # Replace parameters with updated versions. Note that this helper |
| 512 | # function ensures the replacement is compatible with RL weight reloads. |
| 513 | replace_parameter(layer, "w13_weight", w13) |
| 514 | replace_parameter(layer, "w2_weight", w2) |
| 515 | replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) |
| 516 | replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) |
| 517 | |
| 518 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 519 | if self.moe_quant_config: |
| 520 | assert self.experts_cls is not None |
| 521 | self.moe_kernel = make_fp8_moe_kernel( |
| 522 | moe_quant_config=self.moe_quant_config, |
| 523 | moe_config=self.moe, |
| 524 | fp8_backend=self.fp8_backend, |
| 525 | experts_cls=self.experts_cls, |
| 526 | routing_tables=layer._expert_routing_tables(), |
| 527 | ) |
| 528 | |
| 529 | def get_fused_moe_quant_config( |
| 530 | self, layer: torch.nn.Module |
| 531 | ) -> "FusedMoEQuantConfig": |
| 532 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 533 | make_fp8_moe_quant_config, |
| 534 | ) |
| 535 | |
| 536 | w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") |
| 537 | w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") |
| 538 | a1_scale = layer.w13_input_scale |
| 539 | a2_scale = layer.w2_input_scale |
| 540 | |
| 541 | return make_fp8_moe_quant_config( |
| 542 | fp8_backend=self.fp8_backend, |
| 543 | w1_scale=w1_scale, |
| 544 | w2_scale=w2_scale, |
| 545 | a1_scale=a1_scale, |
| 546 | a2_scale=a2_scale, |
| 547 | w1_bias=getattr(layer, "w13_bias", None), |
| 548 | w2_bias=getattr(layer, "w2_bias", None), |
| 549 | block_shape=self.weight_block_size, |
| 550 | per_act_token_quant=self.per_act_token_quant, |
| 551 | per_out_ch_quant=self.per_out_ch_quant, |
| 552 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 553 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 554 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 555 | layer=layer, |
| 556 | ) |
| 557 | |
| 558 | |
| 559 | class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): |
| 560 | """Online tensorwise FP8 MoE quantization. |
| 561 | Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" |
| 562 | |
| 563 | def __init__( |
| 564 | self, |
| 565 | *, |
| 566 | moe: FusedMoEConfig, |
| 567 | ): |
| 568 | super().__init__( |
| 569 | weight_block_size=None, |
| 570 | moe=moe, |
| 571 | ) |
| 572 | |
| 573 | def process_weights_after_loading(self, layer: Module) -> None: |
| 574 | # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones |
| 575 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 576 | return |
| 577 | |
| 578 | # If checkpoint is fp16, quantize in place. |
| 579 | fp8_dtype = current_platform.fp8_dtype() |
| 580 | w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype) |
| 581 | w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype) |
| 582 | layer.w13_input_scale = None |
| 583 | layer.w2_input_scale = None |
| 584 | |
| 585 | moe_tp_size = self.moe.tp_size |
| 586 | w13_amax = weight_amax(layer.w13_weight.flatten(1), dim=-1) |
| 587 | w13_amax = amax_for_moe_weight_quant(w13_amax, moe_tp_size) |
| 588 | w13_scale = _fp8_scale(w13_amax) |
| 589 | w2_amax = weight_amax(layer.w2_weight.flatten(1), dim=-1) |
| 590 | w2_amax = amax_for_moe_weight_quant(w2_amax, moe_tp_size) |
| 591 | w2_scale = _fp8_scale(w2_amax) |
| 592 | |
| 593 | for expert in range(layer.local_num_experts): |
| 594 | w13[expert, :, :], _ = ops.scaled_fp8_quant( |
| 595 | layer.w13_weight[expert, :, :], scale=w13_scale[expert] |
| 596 | ) |
| 597 | w2[expert, :, :], _ = ops.scaled_fp8_quant( |
| 598 | layer.w2_weight[expert, :, :], scale=w2_scale[expert] |
| 599 | ) |
| 600 | |
| 601 | # Shuffle weights to runtime format and setup kernel. |
| 602 | self._setup_kernel( |
| 603 | layer, |
| 604 | w13, |
| 605 | w2, |
| 606 | w13_scale, |
| 607 | w2_scale, |
| 608 | w13_input_scale=layer.w13_input_scale, |
| 609 | w2_input_scale=layer.w2_input_scale, |
| 610 | ) |
| 611 | |
| 612 | # Prevent duplicate processing (e.g., during weight reload) |
| 613 | layer._already_called_process_weights_after_loading = True |
| 614 | |
| 615 | |
| 616 | class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): |
| 617 | """Online blockwise FP8 MoE quantization. |
| 618 | Loads fp16/bf16 weights and quantizes them per-block during loading.""" |
| 619 | |
| 620 | def __init__( |
| 621 | self, |
| 622 | *, |
| 623 | moe: FusedMoEConfig, |
| 624 | ): |
| 625 | super().__init__( |
| 626 | weight_block_size=[128, 128], |
| 627 | moe=moe, |
| 628 | ) |
| 629 | |
| 630 | def maybe_roundup_sizes( |
| 631 | self, |
| 632 | hidden_size: int, |
| 633 | intermediate_size_per_partition: int, |
| 634 | act_dtype: torch.dtype, |
| 635 | moe_parallel_config, |
| 636 | ) -> tuple[int, int]: |
| 637 | hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes( |
| 638 | hidden_size=hidden_size, |
| 639 | intermediate_size_per_partition=intermediate_size_per_partition, |
| 640 | act_dtype=act_dtype, |
| 641 | moe_parallel_config=moe_parallel_config, |
| 642 | ) |
| 643 | assert self.weight_block_size is not None |
| 644 | block_size = self.weight_block_size[0] |
| 645 | return ( |
| 646 | round_up(hidden_size, block_size), |
| 647 | round_up(intermediate_size_per_partition, block_size), |
| 648 | ) |
| 649 | |
| 650 | def process_weights_after_loading(self, layer: Module) -> None: |
| 651 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 652 | return |
| 653 | |
| 654 | self._zero_padding(layer) |
| 655 | |
| 656 | fp8_dtype = current_platform.fp8_dtype() |
| 657 | w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype) |
| 658 | w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype) |
| 659 | |
| 660 | block_size = self.weight_block_size |
| 661 | assert block_size is not None |
| 662 | block_n, block_k = block_size |
| 663 | |
| 664 | # Create block-shaped scales (computed here rather than in |
| 665 | # create_weights because online quant doesn't need them until now). |
| 666 | num_experts = layer.local_num_experts |
| 667 | _, w13_out, w13_in = layer.w13_weight.shape |
| 668 | _, w2_out, w2_in = layer.w2_weight.shape |
| 669 | |
| 670 | w13_scale = torch.ones( |
| 671 | num_experts, |
| 672 | (w13_out + block_n - 1) // block_n, |
| 673 | (w13_in + block_k - 1) // block_k, |
| 674 | dtype=torch.float32, |
| 675 | device=w13.device, |
| 676 | ) |
| 677 | w2_scale = torch.ones( |
| 678 | num_experts, |
| 679 | (w2_out + block_n - 1) // block_n, |
| 680 | (w2_in + block_k - 1) // block_k, |
| 681 | dtype=torch.float32, |
| 682 | device=w2.device, |
| 683 | ) |
| 684 | |
| 685 | for expert in range(num_experts): |
| 686 | w13[expert], w13_scale[expert] = per_block_cast_to_fp8( |
| 687 | layer.w13_weight[expert], |
| 688 | block_size=block_size, |
| 689 | use_ue8m0=False, |
| 690 | ) |
| 691 | w2[expert], w2_scale[expert] = per_block_cast_to_fp8( |
| 692 | layer.w2_weight[expert], |
| 693 | block_size=block_size, |
| 694 | use_ue8m0=False, |
| 695 | ) |
| 696 | |
| 697 | layer.weight_block_size = block_size |
| 698 | |
| 699 | # Shuffle weights to runtime format and setup kernel. |
| 700 | self._setup_kernel( |
| 701 | layer, |
| 702 | w13, |
| 703 | w2, |
| 704 | w13_scale, |
| 705 | w2_scale, |
| 706 | layer.w13_input_scale, |
| 707 | layer.w2_input_scale, |
| 708 | ) |
| 709 | |
| 710 | # Prevent duplicate processing (e.g., during weight reload) |
| 711 | layer._already_called_process_weights_after_loading = True |
| 712 | |
| 713 | |
| 714 | class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): |
| 715 | """Online PTPC FP8 MoE quantization. |
| 716 | |
| 717 | Quantizes each expert's weights per output channel during loading. |
| 718 | Activations are quantized dynamically per token at runtime. |
| 719 | """ |
| 720 | |
| 721 | per_act_token_quant: bool = True |
| 722 | per_out_ch_quant: bool = True |
| 723 | |
| 724 | def __init__( |
| 725 | self, |
| 726 | *, |
| 727 | moe: FusedMoEConfig, |
| 728 | ): |
| 729 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend |
| 730 | |
| 731 | super().__init__( |
| 732 | weight_block_size=None, |
| 733 | moe=moe, |
| 734 | weight_key=kFp8StaticChannelSym, |
| 735 | activation_key=kFp8DynamicTokenSym, |
| 736 | allow_vllm_cutlass=True, |
| 737 | ) |
| 738 | # Reject backends whose make_fp8_moe_quant_config branch silently |
| 739 | # drops per_act_token_quant / per_out_ch_quant or collapses scales: |
| 740 | # MARLIN / CPU route through fp8_w8a16_moe_quant_config, CPU_W8A8 |
| 741 | # ignores both flags, and FLASHINFER_* fold scales into a per-tensor |
| 742 | # alpha (oracle/fp8.py). |
| 743 | if self.fp8_backend in ( |
| 744 | Fp8MoeBackend.MARLIN, |
| 745 | Fp8MoeBackend.CPU, |
| 746 | Fp8MoeBackend.CPU_W8A8, |
| 747 | Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 748 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 749 | ): |
| 750 | raise ValueError( |
| 751 | f"FP8 PTPC online MoE quant is not supported with the " |
| 752 | f"{self.fp8_backend.value} backend, which does not implement " |
| 753 | "per-output-channel weight scales." |
| 754 | ) |
| 755 | |
| 756 | def process_weights_after_loading(self, layer: Module) -> None: |
| 757 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 758 | return |
| 759 | |
| 760 | fp8_dtype = current_platform.fp8_dtype() |
| 761 | w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype) |
| 762 | w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype) |
| 763 | # Scale's leading dim is taken from the fp8 weight tensor by |
| 764 | # construction, so it cannot drift from the weight's expert count |
| 765 | # under EP / padded MoE. |
| 766 | n_w13 = layer.w13_weight.shape[1] |
| 767 | w13_scale = torch.ones( |
| 768 | w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32 |
| 769 | ) |
| 770 | layer.w13_input_scale = None |
| 771 | layer.w2_input_scale = None |
| 772 | |
| 773 | w2_amax = weight_amax(layer.w2_weight, dim=-1, keepdim=True) |
| 774 | w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size) |
| 775 | w2_scale = _fp8_channel_scale(w2_amax) |
| 776 | |
| 777 | for expert in range(layer.local_num_experts): |
| 778 | w13[expert], w13_scale[expert] = ops.scaled_fp8_quant( |
| 779 | layer.w13_weight[expert], |
| 780 | scale=None, |
| 781 | use_per_token_if_dynamic=True, |
| 782 | ) |
| 783 | w2[expert] = _fp8_quant_per_channel( |
| 784 | layer.w2_weight[expert], w2_scale[expert] |
| 785 | ) |
| 786 | |
| 787 | self._setup_kernel( |
| 788 | layer, |
| 789 | w13, |
| 790 | w2, |
| 791 | w13_scale, |
| 792 | w2_scale, |
| 793 | w13_input_scale=None, |
| 794 | w2_input_scale=None, |
| 795 | ) |
| 796 | |
| 797 | layer._already_called_process_weights_after_loading = True |