| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Online MXFP4 (microscaling FP4, block-32) quantization methods.""" |
| 5 | |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | import torch |
| 9 | from torch.nn import Module |
| 10 | |
| 11 | if TYPE_CHECKING: |
| 12 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 13 | from vllm.model_executor.layers.fused_moe import ( |
| 14 | FusedMoEParallelConfig, |
| 15 | FusedMoEQuantConfig, |
| 16 | RoutedExperts, |
| 17 | ) |
| 18 | |
| 19 | from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel |
| 20 | from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig |
| 21 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 22 | TRITON_BACKENDS, |
| 23 | Mxfp4MoeBackend, |
| 24 | convert_weight_to_mxfp4_moe_kernel_format, |
| 25 | make_mxfp4_moe_kernel, |
| 26 | make_mxfp4_moe_quant_config, |
| 27 | mxfp4_round_up_hidden_size_and_intermediate_size, |
| 28 | select_mxfp4_moe_backend, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.online.fp8 import ( |
| 31 | OnlineLinearBase, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.online.moe_base import ( |
| 34 | OnlineMoEMethodBase, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 37 | mxfp4_quantize, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic |
| 40 | from vllm.model_executor.utils import replace_parameter |
| 41 | |
| 42 | # MXFP4 block size (elements per shared e8m0 scale), same as the checkpoint |
| 43 | # and MoE oracle paths. |
| 44 | MXFP4_BLOCK_SIZE = 32 |
| 45 | |
| 46 | |
| 47 | def _quantize_mxfp4_moe_weight( |
| 48 | weight: torch.Tensor, |
| 49 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 50 | """Batch quantization: bf16/fp16 stacked expert weights -> MXFP4. |
| 51 | |
| 52 | Returns packed FP4 weights (uint8, two values per byte) and per-block |
| 53 | (group-32) e8m0 scales (uint8), one pair per expert. |
| 54 | """ |
| 55 | num_experts = weight.size(0) |
| 56 | first_quant, first_scale = mxfp4_quantize(weight[0]) |
| 57 | # Pre-allocate the output tensors rather than stacking. |
| 58 | # This is important for consistent memory layout. |
| 59 | w_quant = torch.empty( |
| 60 | (num_experts, *first_quant.shape), |
| 61 | dtype=first_quant.dtype, |
| 62 | device=weight.device, |
| 63 | ) |
| 64 | w_scales = torch.empty( |
| 65 | (num_experts, *first_scale.shape), |
| 66 | dtype=first_scale.dtype, |
| 67 | device=weight.device, |
| 68 | ) |
| 69 | w_quant[0] = first_quant |
| 70 | w_scales[0] = first_scale |
| 71 | for i in range(1, num_experts): |
| 72 | w_quant[i], w_scales[i] = mxfp4_quantize(weight[i]) |
| 73 | |
| 74 | return w_quant, w_scales |
| 75 | |
| 76 | |
| 77 | class Mxfp4OnlineLinearMethod(OnlineLinearBase): |
| 78 | """Online MXFP4 linear method. |
| 79 | Loads bf16/fp16 checkpoints and quantizes weights to MXFP4 (microscaling |
| 80 | FP4 with block-32 scales) during weight loading. |
| 81 | """ |
| 82 | |
| 83 | activation_quant_key = kMxfp4Dynamic |
| 84 | |
| 85 | def __init__(self): |
| 86 | super().__init__() |
| 87 | self.kernel = init_mxfp4_linear_kernel( |
| 88 | activation_quant_key=self.activation_quant_key |
| 89 | ) |
| 90 | |
| 91 | def create_weights( |
| 92 | self, |
| 93 | layer: torch.nn.Module, |
| 94 | input_size_per_partition: int, |
| 95 | output_partition_sizes: list[int], |
| 96 | input_size: int, |
| 97 | output_size: int, |
| 98 | params_dtype: torch.dtype, |
| 99 | **extra_weight_attrs, |
| 100 | ): |
| 101 | if input_size_per_partition % MXFP4_BLOCK_SIZE != 0: |
| 102 | raise ValueError( |
| 103 | f"MXFP4 requires input_size_per_partition " |
| 104 | f"({input_size_per_partition}) to be divisible by " |
| 105 | f"{MXFP4_BLOCK_SIZE}." |
| 106 | ) |
| 107 | |
| 108 | super().create_weights( |
| 109 | layer, |
| 110 | input_size_per_partition, |
| 111 | output_partition_sizes, |
| 112 | input_size, |
| 113 | output_size, |
| 114 | params_dtype, |
| 115 | **extra_weight_attrs, |
| 116 | ) |
| 117 | |
| 118 | def process_weights_after_loading(self, layer: Module) -> None: |
| 119 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 120 | return |
| 121 | |
| 122 | weight_fp4, weight_scale = mxfp4_quantize(layer.weight.contiguous()) |
| 123 | |
| 124 | layer.input_scale = None |
| 125 | replace_parameter(layer, "weight", weight_fp4.data) |
| 126 | replace_parameter(layer, "weight_scale", weight_scale.data) |
| 127 | |
| 128 | self.kernel.process_weights_after_loading(layer) |
| 129 | |
| 130 | layer._already_called_process_weights_after_loading = True |
| 131 | |
| 132 | def apply( |
| 133 | self, |
| 134 | layer: torch.nn.Module, |
| 135 | x: torch.Tensor, |
| 136 | bias: torch.Tensor | None = None, |
| 137 | ) -> torch.Tensor: |
| 138 | return self.kernel.apply_weights(layer, x, bias) |
| 139 | |
| 140 | |
| 141 | class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): |
| 142 | """MoE method for online MXFP4 (block) quantization.""" |
| 143 | |
| 144 | mxfp4_backend: Mxfp4MoeBackend |
| 145 | experts_cls: "type[mk.FusedMoEExperts] | None" |
| 146 | activation_quant_key = kMxfp4Dynamic |
| 147 | |
| 148 | def __init__(self, *, moe: FusedMoEConfig): |
| 149 | super().__init__(moe) |
| 150 | self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE] |
| 151 | self.weight_scale_name = "weight_scale" |
| 152 | |
| 153 | self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( |
| 154 | config=self.moe, activation_key=self.activation_quant_key |
| 155 | ) |
| 156 | |
| 157 | def maybe_roundup_sizes( |
| 158 | self, |
| 159 | hidden_size: int, |
| 160 | intermediate_size_per_partition: int, |
| 161 | act_dtype: torch.dtype, |
| 162 | moe_parallel_config: "FusedMoEParallelConfig", |
| 163 | ) -> tuple[int, int]: |
| 164 | hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes( |
| 165 | hidden_size=hidden_size, |
| 166 | intermediate_size_per_partition=intermediate_size_per_partition, |
| 167 | act_dtype=act_dtype, |
| 168 | moe_parallel_config=moe_parallel_config, |
| 169 | ) |
| 170 | return mxfp4_round_up_hidden_size_and_intermediate_size( |
| 171 | self.mxfp4_backend, hidden_size, intermediate_size_per_partition |
| 172 | ) |
| 173 | |
| 174 | def create_weights( |
| 175 | self, |
| 176 | layer: Module, |
| 177 | num_experts: int, |
| 178 | hidden_size: int, |
| 179 | intermediate_size_per_partition: int, |
| 180 | params_dtype: torch.dtype, |
| 181 | **extra_weight_attrs, |
| 182 | ): |
| 183 | if ( |
| 184 | hidden_size % MXFP4_BLOCK_SIZE != 0 |
| 185 | or intermediate_size_per_partition % MXFP4_BLOCK_SIZE != 0 |
| 186 | ): |
| 187 | raise ValueError( |
| 188 | "Online MXFP4 MoE requires hidden/intermediate sizes divisible " |
| 189 | f"by {MXFP4_BLOCK_SIZE}." |
| 190 | ) |
| 191 | |
| 192 | super().create_weights( |
| 193 | layer=layer, |
| 194 | num_experts=num_experts, |
| 195 | hidden_size=hidden_size, |
| 196 | intermediate_size_per_partition=intermediate_size_per_partition, |
| 197 | params_dtype=params_dtype, |
| 198 | **extra_weight_attrs, |
| 199 | ) |
| 200 | |
| 201 | layer.weight_block_size = [1, MXFP4_BLOCK_SIZE] |
| 202 | |
| 203 | def _setup_kernel( |
| 204 | self, |
| 205 | layer: "RoutedExperts", |
| 206 | w13: torch.Tensor, |
| 207 | w2: torch.Tensor, |
| 208 | w13_scale: torch.Tensor, |
| 209 | w2_scale: torch.Tensor, |
| 210 | w13_bias: torch.Tensor | None, |
| 211 | w2_bias: torch.Tensor | None, |
| 212 | ) -> None: |
| 213 | w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( |
| 214 | convert_weight_to_mxfp4_moe_kernel_format( |
| 215 | mxfp4_backend=self.mxfp4_backend, |
| 216 | layer=layer, |
| 217 | w13_weight=w13, |
| 218 | w2_weight=w2, |
| 219 | w13_weight_scale=w13_scale, |
| 220 | w2_weight_scale=w2_scale, |
| 221 | w13_bias=w13_bias, |
| 222 | w2_bias=w2_bias, |
| 223 | use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4, |
| 224 | ) |
| 225 | ) |
| 226 | |
| 227 | # Handle weight/scale assignment based on backend type |
| 228 | # Copied from `quark_moe.py`. |
| 229 | # TODO: This should not be here, replace_parameter should handle |
| 230 | # triton_kernels.tensor.Tensor. |
| 231 | if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( |
| 232 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 233 | ): |
| 234 | # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor |
| 235 | # Store on layer for apply(), scales are PrecisionConfig |
| 236 | layer.w13_weight = w13 |
| 237 | layer.w2_weight = w2 |
| 238 | self.w13_precision_config = w13_scale |
| 239 | self.w2_precision_config = w2_scale |
| 240 | else: |
| 241 | # Standard backends: replace parameters |
| 242 | replace_parameter(layer, "w13_weight", w13) |
| 243 | replace_parameter(layer, "w2_weight", w2) |
| 244 | replace_parameter(layer, "w13_weight_scale", w13_scale) |
| 245 | replace_parameter(layer, "w2_weight_scale", w2_scale) |
| 246 | |
| 247 | if w13_bias is not None and w2_bias is not None: |
| 248 | replace_parameter(layer, "w13_bias", w13_bias) |
| 249 | replace_parameter(layer, "w2_bias", w2_bias) |
| 250 | |
| 251 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 252 | if self.moe_quant_config is not None and self.experts_cls is not None: |
| 253 | self.moe_kernel = make_mxfp4_moe_kernel( |
| 254 | moe_quant_config=self.moe_quant_config, |
| 255 | moe_config=self.moe, |
| 256 | mxfp4_backend=self.mxfp4_backend, |
| 257 | experts_cls=self.experts_cls, |
| 258 | routing_tables=layer._expert_routing_tables(), |
| 259 | ) |
| 260 | self.moe_kernel.fused_experts.process_weights_after_loading(layer) |
| 261 | |
| 262 | def get_fused_moe_quant_config( |
| 263 | self, layer: torch.nn.Module |
| 264 | ) -> "FusedMoEQuantConfig | None": |
| 265 | # NOTE: unlike `QuarkOCP_MX_MoEMethod.get_fused_moe_quant_config`, this |
| 266 | # does not branch on `self.mxfp4_backend in TRITON_BACKENDS or |
| 267 | # AITER_MXFP4_FP8` to read `self.w13_precision_config`/ |
| 268 | # `self.w2_precision_config` instead of `layer.w13_weight_scale`/ |
| 269 | # `layer.w2_weight_scale`. For those backends, |
| 270 | # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so |
| 271 | # `getattr` below would raise `AttributeError`. This is currently |
| 272 | # unreachable because `activation_quant_key` is hardcoded to |
| 273 | # `kMxfp4Dynamic` |
| 274 | # TODO: When supporting online activation quant key override, |
| 275 | # fix this convert_weight_to_mxfp4_moe_kernel_format issue. |
| 276 | w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") |
| 277 | w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") |
| 278 | |
| 279 | return make_mxfp4_moe_quant_config( |
| 280 | mxfp4_backend=self.mxfp4_backend, |
| 281 | w1_scale=w1_scale, |
| 282 | w2_scale=w2_scale, |
| 283 | w1_bias=getattr(layer, "w13_bias", None), |
| 284 | w2_bias=getattr(layer, "w2_bias", None), |
| 285 | a1_scale=layer.w13_input_scale, |
| 286 | a2_scale=layer.w2_input_scale, |
| 287 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 288 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 289 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 290 | layer=layer, |
| 291 | ) |
| 292 | |
| 293 | def process_weights_after_loading(self, layer: Module) -> None: |
| 294 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 295 | return |
| 296 | |
| 297 | self._zero_padding(layer) |
| 298 | |
| 299 | if self.mxfp4_backend == Mxfp4MoeBackend.NONE: |
| 300 | layer._already_called_process_weights_after_loading = True |
| 301 | return |
| 302 | |
| 303 | layer.w13_input_scale = None |
| 304 | layer.w2_input_scale = None |
| 305 | |
| 306 | w13, w13_scale = _quantize_mxfp4_moe_weight(layer.w13_weight) |
| 307 | w2, w2_scale = _quantize_mxfp4_moe_weight(layer.w2_weight) |
| 308 | |
| 309 | self._setup_kernel( |
| 310 | layer, |
| 311 | w13, |
| 312 | w2, |
| 313 | w13_scale, |
| 314 | w2_scale, |
| 315 | getattr(layer, "w13_bias", None), |
| 316 | getattr(layer, "w2_bias", None), |
| 317 | ) |
| 318 | |
| 319 | layer._already_called_process_weights_after_loading = True |