| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING, Any, Literal |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.envs as envs |
| 9 | from vllm.config import get_current_vllm_config |
| 10 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 11 | from vllm.model_executor.kernels.linear import ( |
| 12 | init_fp8_linear_kernel, |
| 13 | ) |
| 14 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 15 | CutlassFP8ScaledMMLinearKernel, |
| 16 | MarlinFP8ScaledMMLinearKernel, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.attention import Attention |
| 19 | from vllm.model_executor.layers.fused_moe import ( |
| 20 | FusedMoEMethodBase, |
| 21 | FusedMoeWeightScaleSupported, |
| 22 | RoutedExperts, |
| 23 | SharedExperts, |
| 24 | UnquantizedFusedMoEMethod, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 27 | from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput |
| 28 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 29 | convert_to_fp8_moe_kernel_format, |
| 30 | make_fp8_moe_kernel, |
| 31 | make_fp8_moe_quant_config, |
| 32 | resolve_fp8_moe_weight_block_shape, |
| 33 | select_fp8_moe_backend, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.linear import ( |
| 36 | LinearBase, |
| 37 | LinearMethodBase, |
| 38 | UnquantizedLinearMethod, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 41 | from vllm.model_executor.layers.quantization.base_config import ( |
| 42 | QuantizationConfig, |
| 43 | QuantizeMethodBase, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod |
| 46 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 47 | create_fp8_input_scale, |
| 48 | create_fp8_scale_parameter, |
| 49 | create_fp8_weight_parameter, |
| 50 | process_fp8_input_tensor_strategy_moe, |
| 51 | process_fp8_weight_tensor_strategy, |
| 52 | process_fp8_weight_tensor_strategy_moe, |
| 53 | validate_fp8_block_shape, |
| 54 | validate_fp8_block_shape_moe, |
| 55 | ) |
| 56 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 57 | get_marlin_input_dtype, |
| 58 | ) |
| 59 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 60 | GroupShape, |
| 61 | create_fp8_quant_key, |
| 62 | is_layer_skipped, |
| 63 | kFp8Dynamic128Sym, |
| 64 | kFp8DynamicTensorSym, |
| 65 | kFp8DynamicTokenSym, |
| 66 | kFp8StaticTensorSym, |
| 67 | ) |
| 68 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 69 | cutlass_fp8_supported, |
| 70 | normalize_e4m3fn_to_e4m3fnuz, |
| 71 | ) |
| 72 | from vllm.model_executor.parameter import ( |
| 73 | BlockQuantScaleParameter, |
| 74 | PerTensorScaleParameter, |
| 75 | ) |
| 76 | from vllm.model_executor.utils import ( |
| 77 | is_weights_pre_processed, |
| 78 | replace_parameter, |
| 79 | set_weight_attrs, |
| 80 | ) |
| 81 | from vllm.platforms import current_platform |
| 82 | from vllm.utils.deep_gemm import ( |
| 83 | is_deep_gemm_supported, |
| 84 | ) |
| 85 | |
| 86 | if TYPE_CHECKING: |
| 87 | from vllm.model_executor.models.utils import WeightsMapper |
| 88 | |
| 89 | ACTIVATION_SCHEMES = ["static", "dynamic"] |
| 90 | |
| 91 | |
| 92 | class Fp8Config(QuantizationConfig): |
| 93 | """Config class for FP8.""" |
| 94 | |
| 95 | def __init__( |
| 96 | self, |
| 97 | is_checkpoint_fp8_serialized: bool = True, |
| 98 | activation_scheme: str = "dynamic", |
| 99 | ignored_layers: list[str] | None = None, |
| 100 | weight_block_size: list[int] | None = None, |
| 101 | store_dtype: str | None = None, |
| 102 | ) -> None: |
| 103 | super().__init__() |
| 104 | |
| 105 | self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized |
| 106 | if not is_checkpoint_fp8_serialized: |
| 107 | raise ValueError( |
| 108 | "The `fp8` quantization method no longer supports online " |
| 109 | "quantization. Please use `--quantization fp8_per_tensor` " |
| 110 | "instead. See " |
| 111 | "https://docs.vllm.ai/en/stable/features/quantization/online/" |
| 112 | ) |
| 113 | |
| 114 | if activation_scheme not in ACTIVATION_SCHEMES: |
| 115 | raise ValueError(f"Unsupported activation scheme {activation_scheme}") |
| 116 | self.activation_scheme = activation_scheme |
| 117 | self.ignored_layers = ignored_layers or [] |
| 118 | self.ignored_layers_match_mode: Literal["exact", "substring", "suffix"] = ( |
| 119 | "exact" |
| 120 | ) |
| 121 | self.store_dtype = store_dtype |
| 122 | if weight_block_size is not None: |
| 123 | if len(weight_block_size) != 2: |
| 124 | raise ValueError( |
| 125 | "The quantization block size of weight must have 2 " |
| 126 | f"dimensions, but got {len(weight_block_size)} dimensions" |
| 127 | ) |
| 128 | if activation_scheme != "dynamic": |
| 129 | raise ValueError( |
| 130 | "The block-wise quantization only supports " |
| 131 | "dynamic activation scheme for now, but got " |
| 132 | f"{activation_scheme} activation scheme." |
| 133 | ) |
| 134 | self.weight_block_size = weight_block_size |
| 135 | self.use_deep_gemm: bool | None = None |
| 136 | |
| 137 | @classmethod |
| 138 | def get_name(cls) -> QuantizationMethods: |
| 139 | return "fp8" |
| 140 | |
| 141 | @classmethod |
| 142 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 143 | return [torch.bfloat16, torch.half] |
| 144 | |
| 145 | @classmethod |
| 146 | def get_min_capability(cls) -> int: |
| 147 | return 75 |
| 148 | |
| 149 | @classmethod |
| 150 | def get_config_filenames(cls) -> list[str]: |
| 151 | return [] |
| 152 | |
| 153 | def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): |
| 154 | if self.ignored_layers is not None: |
| 155 | self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) |
| 156 | |
| 157 | @classmethod |
| 158 | def from_config(cls, config: dict[str, Any]) -> "Fp8Config": |
| 159 | quant_method = cls.get_from_keys(config, ["quant_method"]) |
| 160 | is_checkpoint_fp8_serialized = "fp8" in quant_method |
| 161 | activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) |
| 162 | ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) |
| 163 | weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) |
| 164 | store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None) |
| 165 | if not ignored_layers: |
| 166 | ignored_layers = cls.get_from_keys_or( |
| 167 | config, ["modules_to_not_convert"], None |
| 168 | ) |
| 169 | return cls( |
| 170 | is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, |
| 171 | activation_scheme=activation_scheme, |
| 172 | ignored_layers=ignored_layers, |
| 173 | weight_block_size=weight_block_size, |
| 174 | store_dtype=store_dtype, |
| 175 | ) |
| 176 | |
| 177 | def get_quant_method( |
| 178 | self, layer: torch.nn.Module, prefix: str |
| 179 | ) -> "QuantizeMethodBase | None": |
| 180 | if isinstance(layer, LinearBase): |
| 181 | if is_layer_skipped( |
| 182 | prefix=prefix, |
| 183 | ignored_layers=self.ignored_layers, |
| 184 | fused_mapping=self.packed_modules_mapping, |
| 185 | match_mode=self.ignored_layers_match_mode, |
| 186 | ): |
| 187 | return UnquantizedLinearMethod() |
| 188 | offline_method = Fp8LinearMethod(self) |
| 189 | offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix) |
| 190 | return offline_method |
| 191 | elif isinstance(layer, RoutedExperts): |
| 192 | if is_layer_skipped( |
| 193 | prefix=prefix, |
| 194 | ignored_layers=self.ignored_layers, |
| 195 | fused_mapping=self.packed_modules_mapping, |
| 196 | match_mode=self.ignored_layers_match_mode, |
| 197 | ): |
| 198 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 199 | if self.store_dtype == "mxfp4": |
| 200 | from vllm.model_executor.layers.quantization.mxfp4 import ( |
| 201 | Mxfp4MoEMethod, |
| 202 | ) |
| 203 | |
| 204 | return Mxfp4MoEMethod(layer.moe_config) |
| 205 | return Fp8MoEMethod(self, layer) |
| 206 | elif isinstance(layer, Attention): |
| 207 | return Fp8KVCacheMethod(self) |
| 208 | return None |
| 209 | |
| 210 | @staticmethod |
| 211 | def get_cache_scale_mapper() -> "WeightsMapper": |
| 212 | """Map compressed-tensors KV-cache scale names to vLLM names.""" |
| 213 | from vllm.model_executor.models.utils import WeightsMapper |
| 214 | |
| 215 | orig_to_new_suffix = { |
| 216 | ".k_proj.output_scale": ".attn.k_scale", |
| 217 | ".v_proj.output_scale": ".attn.v_scale", |
| 218 | ".q_proj.output_scale": ".attn.q_scale", |
| 219 | ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", |
| 220 | } |
| 221 | cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix) |
| 222 | return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper() |
| 223 | |
| 224 | |
| 225 | class Fp8LinearMethod(LinearMethodBase): |
| 226 | """Linear method for FP8. |
| 227 | Supports loading FP8 checkpoints with static weight scale and |
| 228 | dynamic/static activation scale. |
| 229 | |
| 230 | Limitations: |
| 231 | 1. Only support float8_e4m3fn data type due to the limitation of |
| 232 | torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) |
| 233 | |
| 234 | Args: |
| 235 | quant_config: The quantization config. |
| 236 | |
| 237 | """ |
| 238 | |
| 239 | supports_pre_processed_weights = True |
| 240 | |
| 241 | def __init__(self, quant_config: Fp8Config): |
| 242 | self.quant_config = quant_config |
| 243 | self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False) |
| 244 | self.out_dtype = torch.get_default_dtype() |
| 245 | self.input_dtype = get_current_vllm_config().model_config.dtype |
| 246 | |
| 247 | # For GPUs that lack FP8 hardware support, we can leverage the Marlin |
| 248 | # kernel for fast weight-only FP8 quantization |
| 249 | self.marlin_input_dtype = None |
| 250 | self.use_marlin = False |
| 251 | |
| 252 | if self.quant_config.use_deep_gemm is not None: |
| 253 | self.use_deep_gemm = self.quant_config.use_deep_gemm |
| 254 | else: |
| 255 | self.use_deep_gemm = is_deep_gemm_supported() |
| 256 | |
| 257 | self.weight_block_size = self.quant_config.weight_block_size |
| 258 | self.block_quant = self.weight_block_size is not None |
| 259 | self.act_q_static = self.quant_config.activation_scheme == "static" |
| 260 | |
| 261 | if self.block_quant: |
| 262 | assert not self.act_q_static |
| 263 | assert self.weight_block_size is not None |
| 264 | |
| 265 | self.activation_quant_key = create_fp8_quant_key( |
| 266 | static=self.act_q_static, |
| 267 | group_shape=GroupShape(1, self.weight_block_size[0]), |
| 268 | ) |
| 269 | self.weight_quant_key = create_fp8_quant_key( |
| 270 | static=True, group_shape=GroupShape(*self.weight_block_size) |
| 271 | ) |
| 272 | else: |
| 273 | self.weight_quant_key = kFp8StaticTensorSym |
| 274 | # Use per-token quantization for better perf if dynamic and cutlass |
| 275 | if self.act_q_static: |
| 276 | self.activation_quant_key = kFp8StaticTensorSym |
| 277 | elif cutlass_fp8_supported(): |
| 278 | self.activation_quant_key = kFp8DynamicTokenSym |
| 279 | else: |
| 280 | self.activation_quant_key = kFp8DynamicTensorSym |
| 281 | |
| 282 | def create_weights( |
| 283 | self, |
| 284 | layer: torch.nn.Module, |
| 285 | input_size_per_partition: int, |
| 286 | output_partition_sizes: list[int], |
| 287 | input_size: int, |
| 288 | output_size: int, |
| 289 | params_dtype: torch.dtype, |
| 290 | **extra_weight_attrs, |
| 291 | ): |
| 292 | output_size_per_partition = sum(output_partition_sizes) |
| 293 | weight_loader = extra_weight_attrs.get("weight_loader") |
| 294 | layer.logical_widths = output_partition_sizes |
| 295 | layer.input_size_per_partition = input_size_per_partition |
| 296 | layer.output_size_per_partition = output_size_per_partition |
| 297 | layer.orig_dtype = params_dtype |
| 298 | layer.weight_block_size = None |
| 299 | |
| 300 | if self.block_quant: |
| 301 | assert self.weight_block_size is not None |
| 302 | layer.weight_block_size = self.weight_block_size |
| 303 | validate_fp8_block_shape( |
| 304 | layer, |
| 305 | input_size, |
| 306 | output_size, |
| 307 | input_size_per_partition, |
| 308 | output_partition_sizes, |
| 309 | self.weight_block_size, |
| 310 | ) |
| 311 | |
| 312 | weight = create_fp8_weight_parameter( |
| 313 | output_size_per_partition, input_size_per_partition, weight_loader |
| 314 | ) |
| 315 | layer.register_parameter("weight", weight) |
| 316 | |
| 317 | # WEIGHT SCALE |
| 318 | if not self.block_quant: |
| 319 | scale = create_fp8_scale_parameter( |
| 320 | PerTensorScaleParameter, |
| 321 | output_partition_sizes, |
| 322 | input_size_per_partition, |
| 323 | None, |
| 324 | weight_loader, |
| 325 | ) |
| 326 | layer.register_parameter("weight_scale", scale) |
| 327 | else: |
| 328 | assert not self.act_q_static |
| 329 | assert self.weight_block_size is not None |
| 330 | scale = create_fp8_scale_parameter( |
| 331 | BlockQuantScaleParameter, |
| 332 | output_partition_sizes, |
| 333 | input_size_per_partition, |
| 334 | self.weight_block_size, |
| 335 | weight_loader, |
| 336 | scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None), |
| 337 | ) |
| 338 | # The weight_scale_inv name is intentional for deepseekv3 |
| 339 | layer.register_parameter("weight_scale_inv", scale) |
| 340 | |
| 341 | # INPUT ACTIVATION SCALE |
| 342 | if self.act_q_static: |
| 343 | scale = create_fp8_input_scale(output_partition_sizes, weight_loader) |
| 344 | set_weight_attrs(scale, {"scale_type": "input_scale"}) |
| 345 | layer.register_parameter("input_scale", scale) |
| 346 | |
| 347 | self.fp8_linear = init_fp8_linear_kernel( |
| 348 | activation_quant_key=self.activation_quant_key, |
| 349 | weight_quant_key=self.weight_quant_key, |
| 350 | weight_shape=layer.weight.shape, |
| 351 | input_dtype=self.input_dtype, |
| 352 | out_dtype=self.out_dtype, |
| 353 | module_name=self.__class__.__name__, |
| 354 | ) |
| 355 | |
| 356 | self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) |
| 357 | |
| 358 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 359 | if is_weights_pre_processed(): |
| 360 | # Weights are already in runtime format; the only state tensor |
| 361 | # export cannot carry is the `input_scale = None` stamp that dynamic |
| 362 | # activation quantization relies on, since create_weights never |
| 363 | # registered that parameter. |
| 364 | if not self.act_q_static and not hasattr(layer, "input_scale"): |
| 365 | layer.input_scale = None |
| 366 | # Marlin reads marlin_input_dtype in apply_weights; it lives on the |
| 367 | # method (not exported with the weights), so restore it here too. |
| 368 | if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"): |
| 369 | self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype |
| 370 | return |
| 371 | |
| 372 | if self.use_marlin: |
| 373 | if not self.block_quant: |
| 374 | # Canonicalize to (K, N) for the kernel. |
| 375 | replace_parameter(layer, "weight", layer.weight.t()) |
| 376 | # Only Marlin kernels support `marlin_input_dtype`; guard to avoid |
| 377 | # AttributeError if backend selection changes. |
| 378 | if hasattr(self.fp8_linear, "marlin_input_dtype"): |
| 379 | self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype |
| 380 | self.fp8_linear.process_weights_after_loading(layer) |
| 381 | return |
| 382 | |
| 383 | input_scale = None |
| 384 | # TODO(rob): refactor block quant into separate class. |
| 385 | if self.block_quant: |
| 386 | assert not self.act_q_static |
| 387 | |
| 388 | # If checkpoint not serialized fp8, quantize the weights. |
| 389 | else: |
| 390 | # If checkpoint is fp8 per-tensor, handle that there are N scales for N |
| 391 | # shards in a fused module |
| 392 | weight = layer.weight |
| 393 | weight_scale = layer.weight_scale |
| 394 | |
| 395 | # If using w8a8, torch._scaled_mm needs per tensor, so |
| 396 | # requantize the logical shards as a single weight. |
| 397 | weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy( |
| 398 | weight, |
| 399 | weight_scale, |
| 400 | layer.logical_widths, |
| 401 | getattr(layer, "input_scale", None), |
| 402 | ) |
| 403 | if self.act_q_static: |
| 404 | assert input_scale is not None |
| 405 | input_scale = input_scale.max() |
| 406 | |
| 407 | # Update layer with new values. |
| 408 | replace_parameter(layer, "weight", weight.data) |
| 409 | replace_parameter(layer, "weight_scale", weight_scale.data) |
| 410 | |
| 411 | if input_scale is not None: |
| 412 | replace_parameter(layer, "input_scale", input_scale) |
| 413 | else: |
| 414 | layer.input_scale = None |
| 415 | |
| 416 | self.fp8_linear.process_weights_after_loading(layer) |
| 417 | |
| 418 | def apply( |
| 419 | self, |
| 420 | layer: torch.nn.Module, |
| 421 | x: torch.Tensor, |
| 422 | bias: torch.Tensor | None = None, |
| 423 | ) -> torch.Tensor: |
| 424 | # if batch invariant mode is enabled, prefer direct FP8 path |
| 425 | # we will use BF16 dequant when direct FP8 is not supported. |
| 426 | if envs.VLLM_BATCH_INVARIANT: |
| 427 | if self.block_quant: |
| 428 | assert self.weight_block_size is not None |
| 429 | return self.fp8_linear.apply_weights( |
| 430 | layer, |
| 431 | x, |
| 432 | bias, |
| 433 | ) |
| 434 | else: |
| 435 | if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): |
| 436 | return self.fp8_linear.apply_weights(layer, x, bias) |
| 437 | |
| 438 | # per-tensor/channel: dequant to BF16 and run GEMM |
| 439 | weight_fp8 = layer.weight.to(torch.bfloat16) |
| 440 | weight_scale = layer.weight_scale.to(torch.bfloat16) |
| 441 | if weight_scale.numel() == 1: |
| 442 | # Per-tensor: simple scalar multiplication |
| 443 | weight_bf16 = weight_fp8 * weight_scale |
| 444 | else: |
| 445 | # Multiple scales (fused modules like QKV) |
| 446 | # Try to infer correct broadcasting |
| 447 | # weight is [K, N], scale could be [num_logical_weights] |
| 448 | # Need to figure out how to broadcast - for now just try |
| 449 | # direct multiplication |
| 450 | if ( |
| 451 | weight_scale.dim() == 1 |
| 452 | and weight_scale.shape[0] == weight_fp8.shape[0] |
| 453 | ): |
| 454 | # Per-row scaling |
| 455 | weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1) |
| 456 | else: |
| 457 | # Fallback |
| 458 | weight_bf16 = weight_fp8 * weight_scale |
| 459 | return torch.nn.functional.linear(x, weight_bf16.t(), bias) |
| 460 | |
| 461 | return self.fp8_linear.apply_weights(layer, x, bias) |
| 462 | |
| 463 | |
| 464 | class Fp8MoEMethod(FusedMoEMethodBase): |
| 465 | """MoE method for FP8. |
| 466 | Supports loading FP8 checkpoints with static weight scale and |
| 467 | dynamic/static activation scale. |
| 468 | |
| 469 | Also supports loading quantized FP16/BF16 model checkpoints with dynamic |
| 470 | activation scaling. The weight scaling factor will be initialized after |
| 471 | the model weights are loaded. |
| 472 | |
| 473 | Args: |
| 474 | quant_config: The quantization config. |
| 475 | |
| 476 | """ |
| 477 | |
| 478 | supports_pre_processed_weights = True |
| 479 | |
| 480 | def __init__(self, quant_config: Fp8Config, layer: RoutedExperts): |
| 481 | super().__init__(layer.moe_config) |
| 482 | self.quant_config = quant_config |
| 483 | self.weight_block_size = self.quant_config.weight_block_size |
| 484 | self.block_quant: bool = self.weight_block_size is not None |
| 485 | self.weight_scale_name = ( |
| 486 | "weight_scale_inv" if self.block_quant else "weight_scale" |
| 487 | ) |
| 488 | |
| 489 | self.weight_scale_refine: tuple[int, int] | None = None |
| 490 | self.moe_block_shape = self.weight_block_size |
| 491 | |
| 492 | # Set weight key and activation key for kernel compatibility |
| 493 | if self.block_quant: |
| 494 | assert self.weight_block_size is not None |
| 495 | # Adapt the checkpoint block shape to TP sharding before |
| 496 | # building the weight key. |
| 497 | self.moe_block_shape, self.weight_scale_refine = ( |
| 498 | resolve_fp8_moe_weight_block_shape( |
| 499 | self.moe, |
| 500 | self.weight_block_size, |
| 501 | kFp8Dynamic128Sym, |
| 502 | self.quant_config.is_checkpoint_fp8_serialized, |
| 503 | ) |
| 504 | ) |
| 505 | weight_key = create_fp8_quant_key( |
| 506 | static=True, group_shape=GroupShape(*self.moe_block_shape) |
| 507 | ) |
| 508 | activation_key = kFp8Dynamic128Sym |
| 509 | else: |
| 510 | weight_key = kFp8StaticTensorSym |
| 511 | activation_key = ( |
| 512 | kFp8StaticTensorSym |
| 513 | if self.quant_config.activation_scheme == "static" |
| 514 | else kFp8DynamicTensorSym |
| 515 | ) |
| 516 | |
| 517 | # Select Fp8 MoE backend |
| 518 | self.fp8_backend, self.experts_cls = select_fp8_moe_backend( |
| 519 | config=self.moe, |
| 520 | weight_key=weight_key, |
| 521 | activation_key=activation_key, |
| 522 | allow_vllm_cutlass=False, |
| 523 | ) |
| 524 | |
| 525 | def create_weights( |
| 526 | self, |
| 527 | layer: RoutedExperts, |
| 528 | num_experts: int, |
| 529 | hidden_size: int, |
| 530 | intermediate_size_per_partition: int, |
| 531 | params_dtype: torch.dtype, |
| 532 | **extra_weight_attrs, |
| 533 | ): |
| 534 | layer.num_experts = num_experts |
| 535 | layer.orig_dtype = params_dtype |
| 536 | layer.weight_block_size = None |
| 537 | |
| 538 | params_dtype = torch.float8_e4m3fn |
| 539 | |
| 540 | if self.block_quant: |
| 541 | assert self.weight_block_size is not None |
| 542 | assert self.moe_block_shape is not None |
| 543 | moe_block_shape = self.moe_block_shape |
| 544 | layer.weight_block_size = self.weight_block_size |
| 545 | block_n, block_k = ( |
| 546 | self.weight_block_size[0], |
| 547 | self.weight_block_size[1], |
| 548 | ) |
| 549 | if self.weight_scale_refine is None: |
| 550 | validate_fp8_block_shape_moe( |
| 551 | intermediate_size_per_partition, |
| 552 | self.weight_block_size, |
| 553 | ) |
| 554 | else: |
| 555 | # Use the refined block grid for the scale parameters; the |
| 556 | # loader upsamples the checkpoint scales accordingly. |
| 557 | tp_size = get_tensor_model_parallel_world_size() |
| 558 | if intermediate_size_per_partition % block_n != 0: |
| 559 | block_n, block_k = moe_block_shape |
| 560 | if tp_size > 1 and intermediate_size_per_partition % block_k != 0: |
| 561 | block_n, block_k = moe_block_shape |
| 562 | |
| 563 | # WEIGHTS |
| 564 | w13_weight = torch.nn.Parameter( |
| 565 | torch.empty( |
| 566 | num_experts, |
| 567 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 568 | hidden_size, |
| 569 | dtype=params_dtype, |
| 570 | ), |
| 571 | requires_grad=False, |
| 572 | ) |
| 573 | layer.register_parameter("w13_weight", w13_weight) |
| 574 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 575 | |
| 576 | w2_weight = torch.nn.Parameter( |
| 577 | torch.empty( |
| 578 | num_experts, |
| 579 | hidden_size, |
| 580 | intermediate_size_per_partition, |
| 581 | dtype=params_dtype, |
| 582 | ), |
| 583 | requires_grad=False, |
| 584 | ) |
| 585 | layer.register_parameter("w2_weight", w2_weight) |
| 586 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 587 | |
| 588 | # BIASES (for models like GPT-OSS that have biased MoE) |
| 589 | if self.moe.has_bias: |
| 590 | w13_bias = torch.nn.Parameter( |
| 591 | torch.zeros( |
| 592 | num_experts, |
| 593 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 594 | dtype=layer.orig_dtype, |
| 595 | ), |
| 596 | requires_grad=False, |
| 597 | ) |
| 598 | layer.register_parameter("w13_bias", w13_bias) |
| 599 | set_weight_attrs(w13_bias, extra_weight_attrs) |
| 600 | w2_bias = torch.nn.Parameter( |
| 601 | torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype), |
| 602 | requires_grad=False, |
| 603 | ) |
| 604 | layer.register_parameter("w2_bias", w2_bias) |
| 605 | set_weight_attrs(w2_bias, extra_weight_attrs) |
| 606 | |
| 607 | # WEIGHT_SCALES |
| 608 | if not self.block_quant: |
| 609 | # For per-tensor quant, the scales are per expert and weight. |
| 610 | w13_scale_data = torch.ones( |
| 611 | num_experts, self.moe.w13_num_shards, dtype=torch.float32 |
| 612 | ) |
| 613 | w2_scale_data = torch.ones(num_experts, dtype=torch.float32) |
| 614 | else: |
| 615 | # For block quant, the scales are per block (typically 128x128). |
| 616 | w13_scale_data = torch.ones( |
| 617 | num_experts, |
| 618 | self.moe.w13_num_shards |
| 619 | * ((intermediate_size_per_partition + block_n - 1) // block_n), |
| 620 | (hidden_size + block_k - 1) // block_k, |
| 621 | dtype=torch.float32, |
| 622 | ) |
| 623 | w2_scale_data = torch.ones( |
| 624 | num_experts, |
| 625 | (hidden_size + block_n - 1) // block_n, |
| 626 | (intermediate_size_per_partition + block_k - 1) // block_k, |
| 627 | dtype=torch.float32, |
| 628 | ) |
| 629 | w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False) |
| 630 | w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False) |
| 631 | # Note: name is weight_scale for tensor, weight_scale_inv for block. |
| 632 | layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale) |
| 633 | layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale) |
| 634 | |
| 635 | # Add the quantization method used (per tensor/grouped/channel) |
| 636 | # to ensure the weight scales are loaded in properly |
| 637 | extra_weight_attrs.update( |
| 638 | {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} |
| 639 | if self.block_quant |
| 640 | else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} |
| 641 | ) |
| 642 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 643 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 644 | |
| 645 | # INPUT_SCALES |
| 646 | if self.quant_config.activation_scheme == "static": |
| 647 | assert not self.block_quant |
| 648 | w13_input_scale = torch.nn.Parameter( |
| 649 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 650 | ) |
| 651 | layer.register_parameter("w13_input_scale", w13_input_scale) |
| 652 | set_weight_attrs(w13_input_scale, extra_weight_attrs) |
| 653 | |
| 654 | w2_input_scale = torch.nn.Parameter( |
| 655 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 656 | ) |
| 657 | layer.register_parameter("w2_input_scale", w2_input_scale) |
| 658 | set_weight_attrs(w2_input_scale, extra_weight_attrs) |
| 659 | |
| 660 | else: |
| 661 | layer.w13_input_scale = None |
| 662 | layer.w2_input_scale = None |
| 663 | |
| 664 | def _setup_kernel( |
| 665 | self, |
| 666 | layer: RoutedExperts, |
| 667 | w13: torch.Tensor, |
| 668 | w2: torch.Tensor, |
| 669 | w13_scale: torch.Tensor, |
| 670 | w2_scale: torch.Tensor, |
| 671 | w13_input_scale: torch.Tensor | None, |
| 672 | w2_input_scale: torch.Tensor | None, |
| 673 | ) -> None: |
| 674 | # Shuffle weights to runtime format. |
| 675 | w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 676 | fp8_backend=self.fp8_backend, |
| 677 | layer=layer, |
| 678 | w13=w13, |
| 679 | w2=w2, |
| 680 | w13_scale=w13_scale, |
| 681 | w2_scale=w2_scale, |
| 682 | w13_input_scale=w13_input_scale, |
| 683 | w2_input_scale=w2_input_scale, |
| 684 | ) |
| 685 | |
| 686 | # Replace parameters with updated versions. Note that this helper |
| 687 | # function ensures the replacement is compatible with RL weight reloads. |
| 688 | replace_parameter(layer, "w13_weight", w13) |
| 689 | replace_parameter(layer, "w2_weight", w2) |
| 690 | replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) |
| 691 | replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) |
| 692 | |
| 693 | self._init_moe_kernel(layer) |
| 694 | |
| 695 | def _init_moe_kernel(self, layer: RoutedExperts) -> None: |
| 696 | """Build the MoE kernel from the layer's current (converted) weights.""" |
| 697 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 698 | assert self.moe_quant_config is not None |
| 699 | assert self.experts_cls is not None |
| 700 | self.moe_kernel = make_fp8_moe_kernel( |
| 701 | moe_quant_config=self.moe_quant_config, |
| 702 | moe_config=self.moe, |
| 703 | fp8_backend=self.fp8_backend, |
| 704 | experts_cls=self.experts_cls, |
| 705 | routing_tables=layer._expert_routing_tables(), |
| 706 | ) |
| 707 | |
| 708 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 709 | if is_weights_pre_processed(): |
| 710 | # Weights are already in kernel format; rebuild the kernel only. |
| 711 | self._init_moe_kernel(layer) |
| 712 | return |
| 713 | |
| 714 | # Allow for accessing weights and scales in standard way. |
| 715 | w13 = layer.w13_weight |
| 716 | w2 = layer.w2_weight |
| 717 | w13_scale = getattr(layer, f"w13_{self.weight_scale_name}") |
| 718 | w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") |
| 719 | w13_input_scale = layer.w13_input_scale |
| 720 | w2_input_scale = layer.w2_input_scale |
| 721 | |
| 722 | # MI300x and MI325x use FNUZ format for FP8. Convert if needed. |
| 723 | if current_platform.is_fp8_fnuz(): |
| 724 | w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz( |
| 725 | w13, |
| 726 | w13_scale, |
| 727 | w13_input_scale, |
| 728 | ) |
| 729 | w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( |
| 730 | w2, |
| 731 | w2_scale, |
| 732 | w2_input_scale, |
| 733 | ) |
| 734 | |
| 735 | # Per tensor kernels require single activation scale. Use the max. |
| 736 | if self.quant_config.activation_scheme == "static": |
| 737 | assert not self.block_quant |
| 738 | assert w13_input_scale is not None and w2_input_scale is not None |
| 739 | w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( |
| 740 | w13_input_scale, |
| 741 | w2_input_scale, |
| 742 | layer.moe_config.moe_parallel_config.enable_eplb, |
| 743 | ) |
| 744 | replace_parameter(layer, "w13_input_scale", w13_input_scale) |
| 745 | replace_parameter(layer, "w2_input_scale", w2_input_scale) |
| 746 | |
| 747 | # Per tensor kernels require single weight scale for w13 per expert, but |
| 748 | # on disk there is a scale for w1 and w3. Use the max to requantize. |
| 749 | if not self.block_quant: |
| 750 | shard_size = layer.intermediate_size_per_partition |
| 751 | w13, w13_scale = process_fp8_weight_tensor_strategy_moe( |
| 752 | w13, |
| 753 | w13_scale, |
| 754 | shard_size, |
| 755 | layer.local_num_experts, |
| 756 | is_act_and_mul=self.moe.is_act_and_mul, |
| 757 | ) |
| 758 | |
| 759 | # Shuffle weights to runtime format and setup kernel. |
| 760 | self._setup_kernel( |
| 761 | layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale |
| 762 | ) |
| 763 | |
| 764 | def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: |
| 765 | w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") |
| 766 | w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") |
| 767 | a1_scale = layer.w13_input_scale |
| 768 | a2_scale = layer.w2_input_scale |
| 769 | |
| 770 | quant_config = make_fp8_moe_quant_config( |
| 771 | fp8_backend=self.fp8_backend, |
| 772 | w1_scale=w1_scale, |
| 773 | w2_scale=w2_scale, |
| 774 | a1_scale=a1_scale, |
| 775 | a2_scale=a2_scale, |
| 776 | block_shape=self.moe_block_shape, |
| 777 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 778 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 779 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 780 | layer=layer, |
| 781 | ) |
| 782 | |
| 783 | # Inject biases into the quant config if the model has them |
| 784 | # (e.g. GPT-OSS biased MoE) |
| 785 | if quant_config is not None and self.moe.has_bias: |
| 786 | w13_bias = getattr(layer, "w13_bias", None) |
| 787 | w2_bias = getattr(layer, "w2_bias", None) |
| 788 | if w13_bias is not None: |
| 789 | quant_config._w1.bias = w13_bias |
| 790 | if w2_bias is not None: |
| 791 | quant_config._w2.bias = w2_bias |
| 792 | |
| 793 | return quant_config |
| 794 | |
| 795 | @property |
| 796 | def supports_eplb(self) -> bool: |
| 797 | return True |
| 798 | |
| 799 | def apply_monolithic( |
| 800 | self, |
| 801 | layer: RoutedExperts, |
| 802 | x: torch.Tensor, |
| 803 | router_logits: torch.Tensor, |
| 804 | input_ids: torch.Tensor | None = None, |
| 805 | ) -> torch.Tensor: |
| 806 | assert self.is_monolithic |
| 807 | assert self.moe_kernel is not None |
| 808 | return self.moe_kernel.apply_monolithic( |
| 809 | x, |
| 810 | layer.w13_weight, |
| 811 | layer.w2_weight, |
| 812 | router_logits, |
| 813 | activation=layer.activation, |
| 814 | global_num_experts=layer.global_num_experts, |
| 815 | expert_map=layer.expert_map, |
| 816 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 817 | num_expert_group=layer.num_expert_group, |
| 818 | topk_group=layer.topk_group, |
| 819 | e_score_correction_bias=layer.e_score_correction_bias, |
| 820 | routed_scaling_factor=layer.routed_scaling_factor, |
| 821 | ) |
| 822 | |
| 823 | def apply( |
| 824 | self, |
| 825 | layer: RoutedExperts, |
| 826 | x: torch.Tensor, |
| 827 | topk_weights: torch.Tensor, |
| 828 | topk_ids: torch.Tensor, |
| 829 | shared_experts: SharedExperts | None, |
| 830 | shared_experts_input: torch.Tensor | None, |
| 831 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 832 | assert not self.is_monolithic |
| 833 | assert self.moe_kernel is not None |
| 834 | return self.moe_kernel.apply( |
| 835 | x, |
| 836 | layer.w13_weight, |
| 837 | layer.w2_weight, |
| 838 | topk_weights, |
| 839 | topk_ids, |
| 840 | activation=layer.activation, |
| 841 | global_num_experts=layer.global_num_experts, |
| 842 | expert_map=layer.expert_map, |
| 843 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 844 | shared_experts=shared_experts, |
| 845 | shared_experts_input=shared_experts_input, |
| 846 | ) |
| 847 | |
| 848 | |
| 849 | class Fp8KVCacheMethod(BaseKVCacheMethod): |
| 850 | """Supports loading kv-cache scaling factors from FP8 checkpoints.""" |
| 851 | |
| 852 | def __init__(self, quant_config: Fp8Config): |
| 853 | super().__init__(quant_config) |