| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING, Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 9 | from vllm import _custom_ops as ops |
| 10 | from vllm._aiter_ops import rocm_aiter_ops |
| 11 | from vllm.config import get_current_vllm_config |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.layers.fused_moe import ( |
| 14 | FusedMoEConfig, |
| 15 | FusedMoEMethodBase, |
| 16 | FusedMoeWeightScaleSupported, |
| 17 | RoutedExperts, |
| 18 | SharedExperts, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.config import ( |
| 21 | FusedMoEParallelConfig, |
| 22 | FusedMoEQuantConfig, |
| 23 | fp8_w8a8_moe_quant_config, |
| 24 | int8_w8a8_moe_quant_config, |
| 25 | mxfp4_w4a8_moe_quant_config, |
| 26 | mxfp4_w4a16_moe_quant_config, |
| 27 | ocp_mx_moe_quant_config, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput |
| 30 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 31 | convert_to_fp8_moe_kernel_format, |
| 32 | make_fp8_moe_kernel, |
| 33 | make_fp8_moe_quant_config, |
| 34 | select_fp8_moe_backend, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.fused_moe.oracle.int8 import ( |
| 37 | Int8MoeBackend, |
| 38 | convert_to_int8_moe_kernel_format, |
| 39 | make_int8_moe_kernel, |
| 40 | make_int8_moe_quant_config, |
| 41 | select_int8_moe_backend, |
| 42 | ) |
| 43 | from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( |
| 44 | convert_to_wna16_moe_kernel_format, |
| 45 | make_wna16_moe_kernel, |
| 46 | make_wna16_moe_quant_config, |
| 47 | select_wna16_moe_backend, |
| 48 | ) |
| 49 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 50 | TRITON_BACKENDS, |
| 51 | Mxfp4MoeBackend, |
| 52 | backend_to_kernel_cls, |
| 53 | convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, |
| 54 | make_mxfp4_moe_kernel, |
| 55 | make_mxfp4_moe_quant_config, |
| 56 | mxfp4_round_up_hidden_size_and_intermediate_size, |
| 57 | select_mxfp4_moe_backend, |
| 58 | ) |
| 59 | from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( |
| 60 | convert_to_nvfp4_moe_kernel_format, |
| 61 | make_nvfp4_moe_kernel, |
| 62 | make_nvfp4_moe_quant_config, |
| 63 | select_nvfp4_moe_backend, |
| 64 | ) |
| 65 | from vllm.model_executor.layers.quantization.quark.utils import ( |
| 66 | QuarkQTensorHint, |
| 67 | parse_w4a16_int4_weight_config, |
| 68 | ) |
| 69 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 70 | validate_fp8_block_shape_moe, |
| 71 | ) |
| 72 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( |
| 73 | _ACTIVATION_QUANT_KEY_MAP, |
| 74 | _WEIGHT_QUANT_KEY_MAP, |
| 75 | OCP_MX_BLOCK_SIZE, |
| 76 | OCP_MX_Scheme, |
| 77 | ) |
| 78 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 79 | GroupShape, |
| 80 | QuantKey, |
| 81 | kFp8Dynamic128Sym, |
| 82 | kFp8DynamicTensorSym, |
| 83 | kFp8DynamicTokenSym, |
| 84 | kFp8Static128BlockSym, |
| 85 | kFp8StaticChannelSym, |
| 86 | kFp8StaticTensorSym, |
| 87 | kInt4Static, |
| 88 | kInt4Static32, |
| 89 | kInt4Static32Asym, |
| 90 | kInt4StaticAsym, |
| 91 | kInt4W4A8StaticChannelSym, |
| 92 | kInt8DynamicTensorAsym, |
| 93 | kInt8DynamicTensorSym, |
| 94 | kInt8DynamicTokenAsym, |
| 95 | kInt8DynamicTokenSym, |
| 96 | kInt8StaticChannelSym, |
| 97 | kInt8StaticTensorAsym, |
| 98 | kInt8StaticTensorSym, |
| 99 | kMxfp4Dynamic, |
| 100 | kNvfp4Dynamic, |
| 101 | kNvfp4Static, |
| 102 | ) |
| 103 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 104 | all_close_1d, |
| 105 | normalize_e4m3fn_to_e4m3fnuz, |
| 106 | per_tensor_dequantize, |
| 107 | ) |
| 108 | from vllm.model_executor.utils import replace_parameter, set_weight_attrs |
| 109 | from vllm.platforms import current_platform |
| 110 | |
| 111 | if TYPE_CHECKING: |
| 112 | from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig |
| 113 | |
| 114 | logger = init_logger(__name__) |
| 115 | |
| 116 | __all__ = [ |
| 117 | "QuarkMoEMethod", |
| 118 | "QuarkW4A16Int4MoEMethod", |
| 119 | "QuarkW8A8Fp8MoEMethod", |
| 120 | "QuarkOCP_MX_MoEMethod", |
| 121 | "QuarkNvfp4MoEMethod", |
| 122 | ] |
| 123 | |
| 124 | |
| 125 | class QuarkMoEMethod(FusedMoEMethodBase): |
| 126 | supported_activation_quant_keys: list[QuantKey | None] = [] |
| 127 | supported_weight_quant_keys: list[QuantKey] = [] |
| 128 | |
| 129 | def __init__( |
| 130 | self, |
| 131 | moe: FusedMoEConfig, |
| 132 | weight_quant_key: QuantKey, |
| 133 | activation_quant_key: QuantKey | None, |
| 134 | ): |
| 135 | super().__init__(moe) |
| 136 | if activation_quant_key not in self.supported_activation_quant_keys: |
| 137 | raise ValueError( |
| 138 | f"Unsupported activation quant key: {activation_quant_key}" |
| 139 | ) |
| 140 | if weight_quant_key not in self.supported_weight_quant_keys: |
| 141 | raise ValueError(f"Unsupported weight quant key: {weight_quant_key}") |
| 142 | self.activation_quant_key = activation_quant_key |
| 143 | self.weight_quant_key = weight_quant_key |
| 144 | self.has_bias = self.moe.has_bias |
| 145 | |
| 146 | @staticmethod |
| 147 | def get_moe_method_target( |
| 148 | quant_config: "QuarkConfig", |
| 149 | layer_type: type[torch.nn.Module], |
| 150 | layer_name: str, |
| 151 | ) -> tuple[ |
| 152 | QuantKey | None, |
| 153 | QuantKey | None, |
| 154 | type["QuarkMoEMethod"] | None, |
| 155 | ]: |
| 156 | """Return the quantization target for a routed-experts layer.""" |
| 157 | layer_quant_config = quant_config._find_matched_config(layer_name, layer_type) |
| 158 | if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"): |
| 159 | raise NotImplementedError( |
| 160 | "Currently, Quark models with " |
| 161 | "output_tensors and bias " |
| 162 | "quantized are not supported" |
| 163 | ) |
| 164 | weight_config: QuarkQTensorHint = layer_quant_config.get("weight") |
| 165 | input_config: QuarkQTensorHint = layer_quant_config.get("input_tensors") |
| 166 | |
| 167 | if ( |
| 168 | isinstance(weight_config, list) |
| 169 | and len(weight_config) > 1 |
| 170 | or isinstance(input_config, list) |
| 171 | and len(input_config) > 1 |
| 172 | ): |
| 173 | if match := quant_config._is_fp8_w4a8(weight_config, input_config): |
| 174 | return ( |
| 175 | match.weight_quant_key, |
| 176 | match.activation_quant_key, |
| 177 | QuarkW4A8Fp8MoEMethod, |
| 178 | ) |
| 179 | if match := quant_config._is_nvfp4(weight_config, input_config): |
| 180 | return ( |
| 181 | match.weight_quant_key, |
| 182 | match.activation_quant_key, |
| 183 | QuarkNvfp4MoEMethod, |
| 184 | ) |
| 185 | raise RuntimeError("Unsupported FusedMoe scheme") |
| 186 | |
| 187 | weight_config = quant_config._unwrap_single_quant_config(weight_config) |
| 188 | input_config = quant_config._unwrap_single_quant_config(input_config) |
| 189 | |
| 190 | if match := quant_config._is_fp8_w8a8(weight_config, input_config): |
| 191 | return ( |
| 192 | match.weight_quant_key, |
| 193 | match.activation_quant_key, |
| 194 | QuarkW8A8Fp8MoEMethod, |
| 195 | ) |
| 196 | if match := quant_config._is_w4a16_int4(weight_config, input_config): |
| 197 | return ( |
| 198 | match.weight_quant_key, |
| 199 | match.activation_quant_key, |
| 200 | QuarkW4A16Int4MoEMethod, |
| 201 | ) |
| 202 | if match := quant_config._is_w_ocp_mx_a_x( |
| 203 | weight_config, input_config, allow_static_fp8=True |
| 204 | ): |
| 205 | return ( |
| 206 | match.weight_quant_key, |
| 207 | match.activation_quant_key, |
| 208 | QuarkOCP_MX_MoEMethod, |
| 209 | ) |
| 210 | if match := quant_config._is_w8a8_int8(weight_config, input_config): |
| 211 | return ( |
| 212 | match.weight_quant_key, |
| 213 | match.activation_quant_key, |
| 214 | QuarkW8A8Int8MoEMethod, |
| 215 | ) |
| 216 | raise RuntimeError("Unsupported FusedMoe scheme") |
| 217 | |
| 218 | @staticmethod |
| 219 | def get_moe_method( |
| 220 | quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821 |
| 221 | module: RoutedExperts, |
| 222 | method_cls: type["QuarkMoEMethod"], |
| 223 | weight_quant_key: QuantKey | None, |
| 224 | activation_quant_key: QuantKey | None, |
| 225 | layer_name: str | None = None, |
| 226 | ) -> "QuarkMoEMethod": |
| 227 | if weight_quant_key is None: |
| 228 | raise RuntimeError("Unsupported FusedMoe scheme") |
| 229 | |
| 230 | if method_cls is QuarkW4A8Fp8MoEMethod: |
| 231 | return QuarkW4A8Fp8MoEMethod(module.moe_config, activation_quant_key) |
| 232 | elif method_cls is QuarkNvfp4MoEMethod: |
| 233 | return QuarkNvfp4MoEMethod( |
| 234 | module.moe_config, quant_config, activation_quant_key |
| 235 | ) |
| 236 | elif method_cls is QuarkW8A8Fp8MoEMethod: |
| 237 | return QuarkW8A8Fp8MoEMethod( |
| 238 | module.moe_config, weight_quant_key, activation_quant_key |
| 239 | ) |
| 240 | elif method_cls is QuarkOCP_MX_MoEMethod: |
| 241 | # All OCP MX schemes (W4A16, W4A8, etc.) handled by QuarkOCP_MX_MoEMethod |
| 242 | # Backend selection happens inside via oracle |
| 243 | return QuarkOCP_MX_MoEMethod( |
| 244 | module.moe_config, weight_quant_key, activation_quant_key |
| 245 | ) |
| 246 | elif method_cls is QuarkW8A8Int8MoEMethod: |
| 247 | return QuarkW8A8Int8MoEMethod( |
| 248 | module.moe_config, weight_quant_key, activation_quant_key |
| 249 | ) |
| 250 | elif method_cls is QuarkW4A16Int4MoEMethod: |
| 251 | assert layer_name is not None |
| 252 | layer_quant_config = quant_config._find_matched_config( |
| 253 | layer_name, type(module) |
| 254 | ) |
| 255 | weight_config = quant_config._unwrap_single_quant_config( |
| 256 | layer_quant_config.get("weight") |
| 257 | ) |
| 258 | if weight_config is None: |
| 259 | raise RuntimeError( |
| 260 | f"Missing weight quantization config for layer {layer_name}" |
| 261 | ) |
| 262 | return QuarkW4A16Int4MoEMethod( |
| 263 | weight_quant_key, |
| 264 | activation_quant_key, |
| 265 | weight_config, |
| 266 | quant_config.pack_method, |
| 267 | module.moe_config, |
| 268 | quant_config=quant_config, |
| 269 | ) |
| 270 | else: |
| 271 | raise RuntimeError("Unsupported FusedMoe scheme") |
| 272 | |
| 273 | |
| 274 | class QuarkW4A16Int4MoEMethod(QuarkMoEMethod): |
| 275 | """Quark packed INT4 weight-only MoE method.""" |
| 276 | |
| 277 | supported_activation_quant_keys: list[QuantKey | None] = [None] |
| 278 | supported_weight_quant_keys: list[QuantKey] = [ |
| 279 | kInt4Static, |
| 280 | kInt4Static32, |
| 281 | kInt4StaticAsym, |
| 282 | kInt4Static32Asym, |
| 283 | ] |
| 284 | |
| 285 | def __init__( |
| 286 | self, |
| 287 | weight_quant_key: QuantKey, |
| 288 | activation_quant_key: QuantKey | None, |
| 289 | weight_config: dict[str, Any], |
| 290 | pack_method: str, |
| 291 | moe: FusedMoEConfig, |
| 292 | quant_config: "QuarkConfig | None" = None, # type: ignore # noqa F821 |
| 293 | ): |
| 294 | super().__init__(moe, weight_quant_key, activation_quant_key) |
| 295 | |
| 296 | self.quant_config = quant_config |
| 297 | # group_size is read from the Quark config rather than weight_quant_key: |
| 298 | # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4 |
| 299 | # checkpoints commonly use group_size 128, which the QuantKey cannot |
| 300 | # represent. Symmetry is unambiguous, so it comes from the key. |
| 301 | self.group_size = parse_w4a16_int4_weight_config(weight_config)[0] |
| 302 | self.is_symmetric = weight_quant_key.symmetric |
| 303 | self.bit8_pack_factor = 2 |
| 304 | self.pack_reorder = pack_method == "reorder" |
| 305 | |
| 306 | self.use_wna16_backend = True |
| 307 | self.wna16_backend, self.experts_cls = select_wna16_moe_backend( |
| 308 | config=self.moe, |
| 309 | weight_key=weight_quant_key, |
| 310 | quant_config=self.quant_config, |
| 311 | may_have_zp=not self.is_symmetric, |
| 312 | may_have_bias=False, |
| 313 | ) |
| 314 | |
| 315 | def create_weights( |
| 316 | self, |
| 317 | layer: RoutedExperts, |
| 318 | num_experts: int, |
| 319 | hidden_size: int, |
| 320 | intermediate_size_per_partition: int, |
| 321 | params_dtype: torch.dtype, |
| 322 | **extra_weight_attrs, |
| 323 | ): |
| 324 | group_size = self.group_size |
| 325 | group_size_div_factor = 1 |
| 326 | |
| 327 | while intermediate_size_per_partition % group_size or hidden_size % group_size: |
| 328 | group_size = group_size // 2 |
| 329 | group_size_div_factor *= 2 |
| 330 | assert group_size >= 32 |
| 331 | layer.group_size = group_size |
| 332 | layer.group_size_div_factor = group_size_div_factor |
| 333 | |
| 334 | strategy = FusedMoeWeightScaleSupported.GROUP.value |
| 335 | extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False}) |
| 336 | |
| 337 | assert "weight_loader" in extra_weight_attrs |
| 338 | weight_loader = extra_weight_attrs["weight_loader"] |
| 339 | extra_weight_attrs["weight_loader"] = self.get_weight_loader( |
| 340 | layer, weight_loader |
| 341 | ) |
| 342 | |
| 343 | w13_weight = torch.nn.Parameter( |
| 344 | torch.empty( |
| 345 | num_experts, |
| 346 | 2 * intermediate_size_per_partition, |
| 347 | hidden_size // self.bit8_pack_factor, |
| 348 | dtype=torch.uint8, |
| 349 | ), |
| 350 | requires_grad=False, |
| 351 | ) |
| 352 | layer.register_parameter("w13_weight", w13_weight) |
| 353 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 354 | |
| 355 | w2_weight = torch.nn.Parameter( |
| 356 | torch.empty( |
| 357 | num_experts, |
| 358 | hidden_size, |
| 359 | intermediate_size_per_partition // self.bit8_pack_factor, |
| 360 | dtype=torch.uint8, |
| 361 | ), |
| 362 | requires_grad=False, |
| 363 | ) |
| 364 | layer.register_parameter("w2_weight", w2_weight) |
| 365 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 366 | |
| 367 | w13_weight_scale = torch.nn.Parameter( |
| 368 | torch.zeros( |
| 369 | num_experts, |
| 370 | 2 * intermediate_size_per_partition, |
| 371 | hidden_size // group_size, |
| 372 | dtype=params_dtype, |
| 373 | ), |
| 374 | requires_grad=False, |
| 375 | ) |
| 376 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 377 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 378 | |
| 379 | w2_weight_scale = torch.nn.Parameter( |
| 380 | torch.zeros( |
| 381 | num_experts, |
| 382 | hidden_size, |
| 383 | intermediate_size_per_partition // group_size, |
| 384 | dtype=params_dtype, |
| 385 | ), |
| 386 | requires_grad=False, |
| 387 | ) |
| 388 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 389 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 390 | |
| 391 | w13_weight_zero_point = torch.nn.Parameter( |
| 392 | torch.zeros( |
| 393 | num_experts, |
| 394 | 2 * intermediate_size_per_partition // self.bit8_pack_factor, |
| 395 | hidden_size // group_size, |
| 396 | dtype=torch.uint8, |
| 397 | ), |
| 398 | requires_grad=False, |
| 399 | ) |
| 400 | layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point) |
| 401 | set_weight_attrs(w13_weight_zero_point, extra_weight_attrs) |
| 402 | |
| 403 | w2_weight_zero_point = torch.nn.Parameter( |
| 404 | torch.zeros( |
| 405 | num_experts, |
| 406 | hidden_size // self.bit8_pack_factor, |
| 407 | intermediate_size_per_partition // group_size, |
| 408 | dtype=torch.uint8, |
| 409 | ), |
| 410 | requires_grad=False, |
| 411 | ) |
| 412 | layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point) |
| 413 | set_weight_attrs(w2_weight_zero_point, extra_weight_attrs) |
| 414 | |
| 415 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 416 | w13_qzeros = layer.w13_weight_zero_point if not self.is_symmetric else None |
| 417 | w2_qzeros = layer.w2_weight_zero_point if not self.is_symmetric else None |
| 418 | |
| 419 | converted = convert_to_wna16_moe_kernel_format( |
| 420 | backend=self.wna16_backend, |
| 421 | layer=layer, |
| 422 | quant_config=self.quant_config, |
| 423 | input_dtype=None, |
| 424 | w13=layer.w13_weight, |
| 425 | w2=layer.w2_weight, |
| 426 | w13_scale=layer.w13_weight_scale, |
| 427 | w2_scale=layer.w2_weight_scale, |
| 428 | w13_qzeros=w13_qzeros, |
| 429 | w2_qzeros=w2_qzeros, |
| 430 | ) |
| 431 | |
| 432 | if converted is None: |
| 433 | # Backend rewrote the layer's params in place (e.g. Humming). |
| 434 | self._setup_kernel(layer) |
| 435 | return |
| 436 | |
| 437 | ( |
| 438 | w13_qweight, |
| 439 | w2_qweight, |
| 440 | w13_scales, |
| 441 | w2_scales, |
| 442 | _, # w13_qzeros |
| 443 | _, # w2_qzeros |
| 444 | _, # w13_input_global_scale |
| 445 | _, # w2_input_global_scale |
| 446 | _, # w13_bias |
| 447 | _, # w2_bias |
| 448 | ) = converted |
| 449 | |
| 450 | replace_parameter(layer, "w13_weight", w13_qweight) |
| 451 | replace_parameter(layer, "w2_weight", w2_qweight) |
| 452 | replace_parameter(layer, "w13_weight_scale", w13_scales) |
| 453 | replace_parameter(layer, "w2_weight_scale", w2_scales) |
| 454 | |
| 455 | self._setup_kernel(layer) |
| 456 | |
| 457 | def _setup_kernel(self, layer: RoutedExperts) -> None: |
| 458 | assert self.experts_cls is not None |
| 459 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 460 | assert self.moe_quant_config is not None |
| 461 | self.moe_kernel = make_wna16_moe_kernel( |
| 462 | moe_quant_config=self.moe_quant_config, |
| 463 | moe_config=self.moe, |
| 464 | experts_cls=self.experts_cls, |
| 465 | backend=self.wna16_backend, |
| 466 | routing_tables=layer._expert_routing_tables(), |
| 467 | ) |
| 468 | |
| 469 | def get_fused_moe_quant_config( |
| 470 | self, layer: RoutedExperts |
| 471 | ) -> FusedMoEQuantConfig | None: |
| 472 | w1_zp = layer.w13_weight_zero_point if not self.is_symmetric else None |
| 473 | w2_zp = layer.w2_weight_zero_point if not self.is_symmetric else None |
| 474 | return make_wna16_moe_quant_config( |
| 475 | w1_scale=layer.w13_weight_scale, |
| 476 | w2_scale=layer.w2_weight_scale, |
| 477 | w1_zp=w1_zp, |
| 478 | w2_zp=w2_zp, |
| 479 | group_size=layer.group_size, |
| 480 | num_bits=4, |
| 481 | ) |
| 482 | |
| 483 | def apply( |
| 484 | self, |
| 485 | layer: RoutedExperts, |
| 486 | x: torch.Tensor, |
| 487 | topk_weights: torch.Tensor, |
| 488 | topk_ids: torch.Tensor, |
| 489 | shared_experts: SharedExperts | None, |
| 490 | shared_experts_input: torch.Tensor | None, |
| 491 | ) -> torch.Tensor: |
| 492 | assert self.moe_kernel is not None |
| 493 | return self.moe_kernel.apply( |
| 494 | hidden_states=x, |
| 495 | w1=layer.w13_weight, |
| 496 | w2=layer.w2_weight, |
| 497 | topk_weights=topk_weights, |
| 498 | topk_ids=topk_ids, |
| 499 | activation=layer.activation, |
| 500 | global_num_experts=layer.global_num_experts, |
| 501 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 502 | expert_map=layer.expert_map, |
| 503 | shared_experts=shared_experts, |
| 504 | shared_experts_input=shared_experts_input, |
| 505 | ) |
| 506 | |
| 507 | def get_weight_loader(self, layer: RoutedExperts, weight_loader): |
| 508 | def convert_quark_tensor(tensor: torch.Tensor, tensor_type: str): |
| 509 | size0 = tensor.size(0) |
| 510 | tensor = tensor.view(torch.uint8) |
| 511 | |
| 512 | shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) |
| 513 | tensor = (tensor[:, :, None] >> shifter) & 0xF |
| 514 | |
| 515 | if self.pack_reorder: |
| 516 | reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7] |
| 517 | tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order] |
| 518 | else: |
| 519 | tensor = tensor.view(-1, 8) |
| 520 | if self.is_symmetric: |
| 521 | tensor = tensor ^ 0x8 |
| 522 | tensor = tensor.view(size0, -1) |
| 523 | |
| 524 | tensor = tensor.T.contiguous() |
| 525 | if tensor_type == "weight": |
| 526 | tensor = tensor[:, 1::2] * 16 + tensor[:, ::2] |
| 527 | elif tensor_type == "zero_point": |
| 528 | tensor = tensor[1::2, :] * 16 + tensor[::2, :] |
| 529 | return tensor |
| 530 | |
| 531 | def quark_int4_weight_loader( |
| 532 | param: torch.nn.Parameter, |
| 533 | loaded_weight: torch.Tensor, |
| 534 | weight_name: str, |
| 535 | shard_id: str, |
| 536 | expert_id: int, |
| 537 | return_success: bool = False, |
| 538 | ): |
| 539 | loaded_weight = loaded_weight.to(param.data.device) |
| 540 | shard_size = layer.intermediate_size_per_partition |
| 541 | |
| 542 | if weight_name.endswith("_weight"): |
| 543 | loaded_weight = convert_quark_tensor(loaded_weight, "weight") |
| 544 | elif weight_name.endswith("_weight_zero_point"): |
| 545 | loaded_weight = convert_quark_tensor(loaded_weight, "zero_point") |
| 546 | elif weight_name.endswith("_weight_scale"): |
| 547 | loaded_weight = loaded_weight.T |
| 548 | |
| 549 | if layer.group_size_div_factor > 1 and ( |
| 550 | "weight_zero_point" in weight_name or "weight_scale" in weight_name |
| 551 | ): |
| 552 | loaded_weight = loaded_weight.repeat_interleave( |
| 553 | layer.group_size_div_factor, 1 |
| 554 | ) |
| 555 | |
| 556 | if "w13_weight_zero_point" in weight_name: |
| 557 | tensor = loaded_weight.view( |
| 558 | layer.moe_config.tp_size, -1, loaded_weight.size(1) |
| 559 | )[layer.moe_config.tp_rank] |
| 560 | if shard_id == "w1": |
| 561 | param.data[expert_id, : shard_size // 2] = tensor |
| 562 | else: |
| 563 | param.data[expert_id, shard_size // 2 :] = tensor |
| 564 | return True if return_success else None |
| 565 | elif "w2_weight_zero_point" in weight_name: |
| 566 | param.data[expert_id] = loaded_weight.view( |
| 567 | loaded_weight.size(0), layer.moe_config.tp_size, -1 |
| 568 | )[:, layer.moe_config.tp_rank] |
| 569 | return True if return_success else None |
| 570 | |
| 571 | return weight_loader( |
| 572 | param, |
| 573 | loaded_weight, |
| 574 | weight_name, |
| 575 | shard_id, |
| 576 | expert_id, |
| 577 | return_success=return_success, |
| 578 | ) |
| 579 | |
| 580 | return quark_int4_weight_loader |
| 581 | |
| 582 | |
| 583 | class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): |
| 584 | supported_activation_quant_keys = [ |
| 585 | kFp8StaticTensorSym, |
| 586 | kFp8DynamicTensorSym, |
| 587 | kFp8DynamicTokenSym, |
| 588 | kFp8Dynamic128Sym, |
| 589 | ] |
| 590 | supported_weight_quant_keys = [ |
| 591 | kFp8StaticChannelSym, |
| 592 | kFp8StaticTensorSym, |
| 593 | kFp8Static128BlockSym, |
| 594 | ] |
| 595 | |
| 596 | def __init__( |
| 597 | self, |
| 598 | moe: FusedMoEConfig, |
| 599 | weight_quant_key: QuantKey, |
| 600 | activation_quant_key: QuantKey | None, |
| 601 | ): |
| 602 | super().__init__(moe, weight_quant_key, activation_quant_key) |
| 603 | self.weight_dtype = "fp8" |
| 604 | if weight_quant_key == kFp8StaticChannelSym: |
| 605 | self.weight_qscheme = "per_channel" |
| 606 | elif weight_quant_key == kFp8Static128BlockSym: |
| 607 | self.weight_qscheme = "per_block" |
| 608 | else: |
| 609 | self.weight_qscheme = "per_tensor" |
| 610 | if activation_quant_key == kFp8DynamicTokenSym: |
| 611 | self.input_qscheme = "per_channel" |
| 612 | elif activation_quant_key == kFp8Dynamic128Sym: |
| 613 | self.input_qscheme = "per_group" |
| 614 | else: |
| 615 | self.input_qscheme = "per_tensor" |
| 616 | self.static_input_scales = activation_quant_key == kFp8StaticTensorSym |
| 617 | per_tensor = ( |
| 618 | self.weight_qscheme == "per_tensor" and self.input_qscheme == "per_tensor" |
| 619 | ) |
| 620 | per_channel = ( |
| 621 | self.weight_qscheme == "per_channel" and self.input_qscheme == "per_channel" |
| 622 | ) |
| 623 | per_block = ( |
| 624 | self.weight_qscheme == "per_block" and self.input_qscheme == "per_group" |
| 625 | ) |
| 626 | if not (per_tensor or per_channel or per_block): |
| 627 | raise ValueError( |
| 628 | "For FP8 Fused MoE layers, only per-tensor, per-channel and " |
| 629 | "per-block scales for weights and activations are supported. " |
| 630 | f"Found {self.weight_qscheme}, {self.input_qscheme}" |
| 631 | ) # noqa E501 |
| 632 | |
| 633 | # One scale per (block_n, block_k) tile of the weight, so the block |
| 634 | # shape is carried by the weight key's group shape. |
| 635 | self.weight_block_size: list[int] | None = ( |
| 636 | list(weight_quant_key.scale.group_shape) if per_block else None |
| 637 | ) |
| 638 | self.block_quant = self.weight_block_size is not None |
| 639 | |
| 640 | if per_channel: |
| 641 | self.act_quant_group_shape = GroupShape.PER_TOKEN |
| 642 | elif per_block: |
| 643 | assert self.weight_block_size is not None |
| 644 | self.act_quant_group_shape = GroupShape(1, self.weight_block_size[1]) |
| 645 | else: |
| 646 | self.act_quant_group_shape = GroupShape.PER_TENSOR |
| 647 | |
| 648 | self.fp8_backend, self.experts_cls = select_fp8_moe_backend( |
| 649 | config=moe, |
| 650 | weight_key=self.weight_quant_key, |
| 651 | activation_key=self.activation_quant_key, |
| 652 | ) |
| 653 | |
| 654 | self.model_type = getattr( |
| 655 | get_current_vllm_config().model_config.hf_config, "model_type", None |
| 656 | ) |
| 657 | |
| 658 | def create_weights( |
| 659 | self, |
| 660 | layer: RoutedExperts, |
| 661 | num_experts: int, |
| 662 | hidden_size: int, |
| 663 | intermediate_size_per_partition: int, |
| 664 | params_dtype: torch.dtype, |
| 665 | **extra_weight_attrs, |
| 666 | ): |
| 667 | layer.num_experts = num_experts |
| 668 | layer.orig_dtype = params_dtype |
| 669 | layer.weight_block_size = None |
| 670 | params_dtype = torch.float8_e4m3fn |
| 671 | |
| 672 | if self.block_quant: |
| 673 | assert self.weight_block_size is not None |
| 674 | layer.weight_block_size = self.weight_block_size |
| 675 | block_n, block_k = self.weight_block_size[0], self.weight_block_size[1] |
| 676 | validate_fp8_block_shape_moe( |
| 677 | intermediate_size_per_partition, |
| 678 | self.weight_block_size, |
| 679 | ) |
| 680 | |
| 681 | # WEIGHTS |
| 682 | w13_weight = torch.nn.Parameter( |
| 683 | torch.zeros( |
| 684 | num_experts, |
| 685 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 686 | hidden_size, |
| 687 | dtype=params_dtype, |
| 688 | ), |
| 689 | requires_grad=False, |
| 690 | ) |
| 691 | layer.register_parameter("w13_weight", w13_weight) |
| 692 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 693 | |
| 694 | w2_weight = torch.nn.Parameter( |
| 695 | torch.zeros( |
| 696 | num_experts, |
| 697 | hidden_size, |
| 698 | intermediate_size_per_partition, |
| 699 | dtype=params_dtype, |
| 700 | ), |
| 701 | requires_grad=False, |
| 702 | ) |
| 703 | layer.register_parameter("w2_weight", w2_weight) |
| 704 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 705 | |
| 706 | # WEIGHT_SCALES |
| 707 | if self.weight_qscheme == "per_tensor": |
| 708 | # Allocate 2 scales for w1 and w3 respectively. |
| 709 | # They are combined to a single scale after weight loading. |
| 710 | if self.model_type != "gpt_oss": |
| 711 | w13_weight_scale = torch.nn.Parameter( |
| 712 | torch.ones( |
| 713 | num_experts, self.moe.w13_num_shards, dtype=torch.float32 |
| 714 | ), |
| 715 | requires_grad=False, |
| 716 | ) |
| 717 | else: |
| 718 | # For gpt_oss, the w1(gate) & w3(up) are fused as one. |
| 719 | # Therefore, only one weight scale for each expert. |
| 720 | w13_weight_scale = torch.nn.Parameter( |
| 721 | torch.ones(num_experts, 1, dtype=torch.float32), requires_grad=False |
| 722 | ) |
| 723 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 724 | w2_weight_scale = torch.nn.Parameter( |
| 725 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 726 | ) |
| 727 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 728 | # Add PER-TENSOR quantization for RoutedExperts.weight_loader. |
| 729 | extra_weight_attrs.update( |
| 730 | {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} |
| 731 | ) |
| 732 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 733 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 734 | elif self.weight_qscheme == "per_channel": |
| 735 | # quark's scale is 1 dim. |
| 736 | w13_weight_scale = torch.nn.Parameter( |
| 737 | torch.ones( |
| 738 | num_experts, |
| 739 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 740 | dtype=torch.float32, |
| 741 | ), |
| 742 | requires_grad=False, |
| 743 | ) |
| 744 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 745 | w2_weight_scale = torch.nn.Parameter( |
| 746 | torch.ones(num_experts, hidden_size, dtype=torch.float32), |
| 747 | requires_grad=False, |
| 748 | ) |
| 749 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 750 | # Add PER-CHANNEL quantization for RoutedExperts.weight_loader. |
| 751 | extra_weight_attrs.update( |
| 752 | {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} |
| 753 | ) |
| 754 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 755 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 756 | elif self.weight_qscheme == "per_block": |
| 757 | # One scale per (block_n, block_k) tile of each expert's weight. |
| 758 | w13_weight_scale = torch.nn.Parameter( |
| 759 | torch.ones( |
| 760 | num_experts, |
| 761 | self.moe.w13_num_shards |
| 762 | * ((intermediate_size_per_partition + block_n - 1) // block_n), |
| 763 | (hidden_size + block_k - 1) // block_k, |
| 764 | dtype=torch.float32, |
| 765 | ), |
| 766 | requires_grad=False, |
| 767 | ) |
| 768 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 769 | w2_weight_scale = torch.nn.Parameter( |
| 770 | torch.ones( |
| 771 | num_experts, |
| 772 | (hidden_size + block_n - 1) // block_n, |
| 773 | (intermediate_size_per_partition + block_k - 1) // block_k, |
| 774 | dtype=torch.float32, |
| 775 | ), |
| 776 | requires_grad=False, |
| 777 | ) |
| 778 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 779 | # Add PER-BLOCK quantization for RoutedExperts.weight_loader. |
| 780 | extra_weight_attrs.update( |
| 781 | {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} |
| 782 | ) |
| 783 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 784 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 785 | |
| 786 | # INPUT_SCALES |
| 787 | if self.static_input_scales: |
| 788 | w13_input_scale = torch.nn.Parameter( |
| 789 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 790 | ) |
| 791 | layer.register_parameter("w13_input_scale", w13_input_scale) |
| 792 | set_weight_attrs(w13_input_scale, extra_weight_attrs) |
| 793 | |
| 794 | w2_input_scale = torch.nn.Parameter( |
| 795 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 796 | ) |
| 797 | layer.register_parameter("w2_input_scale", w2_input_scale) |
| 798 | set_weight_attrs(w2_input_scale, extra_weight_attrs) |
| 799 | else: |
| 800 | layer.w13_input_scale = None |
| 801 | layer.w2_input_scale = None |
| 802 | |
| 803 | if self.has_bias: |
| 804 | w13_bias = torch.nn.Parameter( |
| 805 | torch.zeros( |
| 806 | num_experts, |
| 807 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 808 | dtype=torch.float32, |
| 809 | ), |
| 810 | requires_grad=False, |
| 811 | ) |
| 812 | layer.register_parameter("w13_bias", w13_bias) |
| 813 | set_weight_attrs(w13_bias, extra_weight_attrs) |
| 814 | |
| 815 | w2_bias = torch.nn.Parameter( |
| 816 | torch.zeros(num_experts, hidden_size, dtype=torch.float32), |
| 817 | requires_grad=False, |
| 818 | ) |
| 819 | layer.register_parameter("w2_bias", w2_bias) |
| 820 | set_weight_attrs(w2_bias, extra_weight_attrs) |
| 821 | else: |
| 822 | layer.w13_bias, layer.w2_bias = None, None |
| 823 | |
| 824 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 825 | # Fp8 moe kernels require a single activation scale. |
| 826 | # We take the max of all the scales in case they differ. |
| 827 | if self.static_input_scales: |
| 828 | if layer.w13_input_scale is None or layer.w2_input_scale is None: |
| 829 | raise ValueError( |
| 830 | "QuantConfig has static quantization, but found " |
| 831 | "activation scales are None." |
| 832 | ) |
| 833 | if not all_close_1d(layer.w13_input_scale) or not all_close_1d( |
| 834 | layer.w2_input_scale |
| 835 | ): |
| 836 | logger.warning_once( |
| 837 | "Found input_scales that are not equal for " |
| 838 | "fp8 MoE layer. Using the maximum across experts " |
| 839 | "for each layer. " |
| 840 | ) |
| 841 | layer.w13_input_scale = torch.nn.Parameter( |
| 842 | layer.w13_input_scale.max(), requires_grad=False |
| 843 | ) |
| 844 | layer.w2_input_scale = torch.nn.Parameter( |
| 845 | layer.w2_input_scale.max(), requires_grad=False |
| 846 | ) |
| 847 | |
| 848 | if current_platform.is_fp8_fnuz(): |
| 849 | # Normalize the weights and scales |
| 850 | w13_weight, w13_weight_scale, w13_input_scale = ( |
| 851 | normalize_e4m3fn_to_e4m3fnuz( |
| 852 | layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale |
| 853 | ) |
| 854 | ) |
| 855 | w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( |
| 856 | layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale |
| 857 | ) |
| 858 | # Reset the parameter |
| 859 | layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) |
| 860 | layer.w13_weight_scale = torch.nn.Parameter( |
| 861 | w13_weight_scale, requires_grad=False |
| 862 | ) |
| 863 | if w13_input_scale is not None: |
| 864 | layer.w13_input_scale = torch.nn.Parameter( |
| 865 | w13_input_scale, requires_grad=False |
| 866 | ) |
| 867 | layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) |
| 868 | layer.w2_weight_scale = torch.nn.Parameter( |
| 869 | w2_weight_scale, requires_grad=False |
| 870 | ) |
| 871 | if w2_input_scale is not None: |
| 872 | layer.w2_input_scale = torch.nn.Parameter( |
| 873 | w2_input_scale, requires_grad=False |
| 874 | ) |
| 875 | |
| 876 | # For per-tensor case, Fp8 moe kernel needs single weight scale |
| 877 | # for w13 per expert. Use max then dequant and requant each expert. |
| 878 | if self.weight_qscheme == "per_tensor": |
| 879 | assert layer.w13_weight_scale is not None |
| 880 | shard_size = layer.intermediate_size_per_partition |
| 881 | max_w13_scales = layer.w13_weight_scale.max(dim=1).values |
| 882 | |
| 883 | # For gpt_oss, w1 and w3 are fused into a single combined |
| 884 | # gate_up_proj tensor with size 2*intermediate_size_per_partition |
| 885 | # and only one scale per expert. |
| 886 | # Process the entire weight tensor as one shard. |
| 887 | if self.model_type == "gpt_oss": |
| 888 | for expert_id in range(layer.local_num_experts): |
| 889 | # Process all 2*intermediate_size_per_partition rows at once |
| 890 | dq_weight = per_tensor_dequantize( |
| 891 | layer.w13_weight[expert_id], |
| 892 | layer.w13_weight_scale[expert_id][0], |
| 893 | ) |
| 894 | layer.w13_weight[expert_id], _ = ops.scaled_fp8_quant( |
| 895 | dq_weight, max_w13_scales[expert_id] |
| 896 | ) |
| 897 | else: |
| 898 | # For non-gpt_oss, process w1 and w3 shards separately |
| 899 | for expert_id in range(layer.local_num_experts): |
| 900 | start = 0 |
| 901 | for shard_id in range(self.moe.w13_num_shards): |
| 902 | dq_weight = per_tensor_dequantize( |
| 903 | layer.w13_weight[expert_id][start : start + shard_size, :], |
| 904 | layer.w13_weight_scale[expert_id][shard_id], |
| 905 | ) |
| 906 | ( |
| 907 | layer.w13_weight[expert_id][start : start + shard_size, :], |
| 908 | _, |
| 909 | ) = ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) |
| 910 | start += shard_size |
| 911 | |
| 912 | layer.w13_weight_scale = torch.nn.Parameter( |
| 913 | max_w13_scales, requires_grad=False |
| 914 | ) |
| 915 | |
| 916 | # quark's scale is 1 dim. |
| 917 | elif self.weight_qscheme == "per_channel": |
| 918 | if self.act_quant_group_shape == GroupShape.PER_TOKEN: |
| 919 | w13_weight_scale = layer.w13_weight_scale.unsqueeze(-1) |
| 920 | layer.w13_weight_scale = torch.nn.Parameter( |
| 921 | w13_weight_scale, requires_grad=False |
| 922 | ) |
| 923 | w2_weight_scale = layer.w2_weight_scale.unsqueeze(-1) |
| 924 | layer.w2_weight_scale = torch.nn.Parameter( |
| 925 | w2_weight_scale, requires_grad=False |
| 926 | ) |
| 927 | self._setup_kernel(layer) |
| 928 | |
| 929 | def _setup_kernel(self, layer: RoutedExperts) -> None: |
| 930 | w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 931 | fp8_backend=self.fp8_backend, |
| 932 | layer=layer, |
| 933 | w13=layer.w13_weight, |
| 934 | w2=layer.w2_weight, |
| 935 | w13_scale=layer.w13_weight_scale, |
| 936 | w2_scale=layer.w2_weight_scale, |
| 937 | w13_input_scale=layer.w13_input_scale, |
| 938 | w2_input_scale=layer.w2_input_scale, |
| 939 | ) |
| 940 | replace_parameter(layer, "w13_weight", w13) |
| 941 | replace_parameter(layer, "w2_weight", w2) |
| 942 | replace_parameter(layer, "w13_weight_scale", w13_scale) |
| 943 | replace_parameter(layer, "w2_weight_scale", w2_scale) |
| 944 | |
| 945 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 946 | assert self.moe_quant_config is not None |
| 947 | assert self.experts_cls is not None |
| 948 | self.moe_kernel = make_fp8_moe_kernel( |
| 949 | moe_quant_config=self.moe_quant_config, |
| 950 | moe_config=self.moe, |
| 951 | fp8_backend=self.fp8_backend, |
| 952 | experts_cls=self.experts_cls, |
| 953 | routing_tables=layer._expert_routing_tables(), |
| 954 | ) |
| 955 | |
| 956 | def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: |
| 957 | return make_fp8_moe_quant_config( |
| 958 | fp8_backend=self.fp8_backend, |
| 959 | w1_scale=layer.w13_weight_scale, |
| 960 | w2_scale=layer.w2_weight_scale, |
| 961 | a1_scale=layer.w13_input_scale, |
| 962 | a2_scale=layer.w2_input_scale, |
| 963 | w1_bias=getattr(layer, "w13_bias", None), |
| 964 | w2_bias=getattr(layer, "w2_bias", None), |
| 965 | block_shape=self.weight_block_size, |
| 966 | per_act_token_quant=self.input_qscheme == "per_channel", |
| 967 | per_out_ch_quant=self.weight_qscheme == "per_channel", |
| 968 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 969 | ) |
| 970 | |
| 971 | def apply( |
| 972 | self, |
| 973 | layer: RoutedExperts, |
| 974 | x: torch.Tensor, |
| 975 | topk_weights: torch.Tensor, |
| 976 | topk_ids: torch.Tensor, |
| 977 | shared_experts: SharedExperts | None, |
| 978 | shared_experts_input: torch.Tensor | None, |
| 979 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 980 | assert self.moe_kernel is not None |
| 981 | return self.moe_kernel.apply( |
| 982 | hidden_states=x, |
| 983 | w1=layer.w13_weight, |
| 984 | w2=layer.w2_weight, |
| 985 | topk_weights=topk_weights, |
| 986 | topk_ids=topk_ids, |
| 987 | activation=layer.activation, |
| 988 | global_num_experts=layer.global_num_experts, |
| 989 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 990 | expert_map=layer.expert_map, |
| 991 | shared_experts=shared_experts, |
| 992 | shared_experts_input=shared_experts_input, |
| 993 | ) |
| 994 | |
| 995 | |
| 996 | class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): |
| 997 | """Quark W8A8 INT8 MoE method.""" |
| 998 | |
| 999 | supported_activation_quant_keys = [ |
| 1000 | kInt8StaticTensorSym, |
| 1001 | kInt8StaticTensorAsym, |
| 1002 | kInt8DynamicTensorSym, |
| 1003 | kInt8DynamicTensorAsym, |
| 1004 | kInt8DynamicTokenSym, |
| 1005 | kInt8DynamicTokenAsym, |
| 1006 | ] |
| 1007 | supported_weight_quant_keys = [ |
| 1008 | kInt8StaticChannelSym, |
| 1009 | kInt8StaticTensorSym, |
| 1010 | ] |
| 1011 | |
| 1012 | def __init__( |
| 1013 | self, |
| 1014 | moe: FusedMoEConfig, |
| 1015 | weight_quant_key: QuantKey, |
| 1016 | activation_quant_key: QuantKey | None, |
| 1017 | ): |
| 1018 | super().__init__(moe, weight_quant_key, activation_quant_key) |
| 1019 | |
| 1020 | # TODO: oracle needs to support asym int8. |
| 1021 | if self.activation_quant_key == kInt8StaticTensorAsym: |
| 1022 | self.activation_quant_key = kInt8StaticTensorSym |
| 1023 | if self.activation_quant_key == kInt8DynamicTensorAsym: |
| 1024 | self.activation_quant_key = kInt8DynamicTensorSym |
| 1025 | if self.activation_quant_key == kInt8DynamicTokenAsym: |
| 1026 | self.activation_quant_key = kInt8DynamicTokenSym |
| 1027 | |
| 1028 | self.weight_qscheme = ( |
| 1029 | "per_channel" if weight_quant_key == kInt8StaticChannelSym else "per_tensor" |
| 1030 | ) |
| 1031 | |
| 1032 | assert self.activation_quant_key is not None |
| 1033 | self.static_input_scales = self.activation_quant_key.scale.static |
| 1034 | |
| 1035 | self.moe_quant_config: FusedMoEQuantConfig | None = None |
| 1036 | self.moe_kernel: mk.FusedMoEKernel | None = None |
| 1037 | self.int8_backend: Int8MoeBackend | None = None |
| 1038 | self.experts_cls: type[mk.FusedMoEExperts] | None = None |
| 1039 | |
| 1040 | # Dynamic-activation INT8 MoE goes through the oracle + modular kernel. |
| 1041 | # The modular TritonExperts kernel consumes float activations and |
| 1042 | # quantizes them to int8 itself, so it cannot apply a loaded static |
| 1043 | # activation scale (this matches CompressedTensorsW8A8Int8MoEMethod). |
| 1044 | # TODO: Static-activation INT8 therefore stays on the legacy fused_experts |
| 1045 | # path (see apply()) for now, preserving pre-refactor behavior. |
| 1046 | # Needs to be migrated to expert backend. |
| 1047 | if not self.static_input_scales: |
| 1048 | # Map the Quark weight scheme to oracle quant keys. Per-channel |
| 1049 | # weights pair with dynamic per-token activations; per-tensor |
| 1050 | # weights with dynamic per-tensor activations. |
| 1051 | self.int8_backend, self.experts_cls = select_int8_moe_backend( |
| 1052 | config=moe, |
| 1053 | weight_key=self.weight_quant_key, |
| 1054 | activation_key=self.activation_quant_key, |
| 1055 | ) |
| 1056 | |
| 1057 | def create_weights( |
| 1058 | self, |
| 1059 | layer: torch.nn.Module, |
| 1060 | num_experts: int, |
| 1061 | hidden_size: int, |
| 1062 | intermediate_size_per_partition: int, |
| 1063 | params_dtype: torch.dtype, |
| 1064 | **extra_weight_attrs, |
| 1065 | ): |
| 1066 | layer.num_experts = num_experts |
| 1067 | layer.orig_dtype = params_dtype |
| 1068 | layer.weight_block_size = None |
| 1069 | params_dtype = torch.int8 |
| 1070 | |
| 1071 | # WEIGHTS |
| 1072 | w13_weight = torch.nn.Parameter( |
| 1073 | torch.empty( |
| 1074 | num_experts, |
| 1075 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1076 | hidden_size, |
| 1077 | dtype=params_dtype, |
| 1078 | ), |
| 1079 | requires_grad=False, |
| 1080 | ) |
| 1081 | layer.register_parameter("w13_weight", w13_weight) |
| 1082 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 1083 | |
| 1084 | w2_weight = torch.nn.Parameter( |
| 1085 | torch.empty( |
| 1086 | num_experts, |
| 1087 | hidden_size, |
| 1088 | intermediate_size_per_partition, |
| 1089 | dtype=params_dtype, |
| 1090 | ), |
| 1091 | requires_grad=False, |
| 1092 | ) |
| 1093 | layer.register_parameter("w2_weight", w2_weight) |
| 1094 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 1095 | |
| 1096 | # WEIGHT_SCALES |
| 1097 | if self.weight_qscheme == "per_channel": |
| 1098 | w13_weight_scale = torch.nn.Parameter( |
| 1099 | torch.ones( |
| 1100 | num_experts, |
| 1101 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1102 | dtype=torch.float32, |
| 1103 | ), |
| 1104 | requires_grad=False, |
| 1105 | ) |
| 1106 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 1107 | w2_weight_scale = torch.nn.Parameter( |
| 1108 | torch.ones(num_experts, hidden_size, dtype=torch.float32), |
| 1109 | requires_grad=False, |
| 1110 | ) |
| 1111 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 1112 | extra_weight_attrs.update( |
| 1113 | {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} |
| 1114 | ) |
| 1115 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 1116 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 1117 | else: |
| 1118 | # per-tensor: one scalar per expert (two for the fused w1/w3) |
| 1119 | w13_weight_scale = torch.nn.Parameter( |
| 1120 | torch.ones(num_experts, self.moe.w13_num_shards, dtype=torch.float32), |
| 1121 | requires_grad=False, |
| 1122 | ) |
| 1123 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 1124 | w2_weight_scale = torch.nn.Parameter( |
| 1125 | torch.ones(num_experts, dtype=torch.float32), |
| 1126 | requires_grad=False, |
| 1127 | ) |
| 1128 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 1129 | extra_weight_attrs.update( |
| 1130 | {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} |
| 1131 | ) |
| 1132 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 1133 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 1134 | |
| 1135 | # INPUT_SCALES |
| 1136 | if self.static_input_scales: |
| 1137 | # Static activations: the per-expert scales are loaded from the |
| 1138 | # checkpoint (used by the legacy fused_experts path). |
| 1139 | w13_input_scale = torch.nn.Parameter( |
| 1140 | torch.ones(num_experts, dtype=torch.float32), |
| 1141 | requires_grad=False, |
| 1142 | ) |
| 1143 | layer.register_parameter("w13_input_scale", w13_input_scale) |
| 1144 | set_weight_attrs(w13_input_scale, extra_weight_attrs) |
| 1145 | |
| 1146 | w2_input_scale = torch.nn.Parameter( |
| 1147 | torch.ones(num_experts, dtype=torch.float32), |
| 1148 | requires_grad=False, |
| 1149 | ) |
| 1150 | layer.register_parameter("w2_input_scale", w2_input_scale) |
| 1151 | set_weight_attrs(w2_input_scale, extra_weight_attrs) |
| 1152 | else: |
| 1153 | # Dynamic activations are quantized in-kernel (no stored scale). |
| 1154 | layer.w13_input_scale = None |
| 1155 | layer.w2_input_scale = None |
| 1156 | |
| 1157 | # ZERO POINTS (loaded but discarded after loading; kernel uses symmetric) |
| 1158 | w13_input_zero_point = torch.nn.Parameter( |
| 1159 | torch.zeros(num_experts, self.moe.w13_num_shards, dtype=torch.int8), |
| 1160 | requires_grad=False, |
| 1161 | ) |
| 1162 | layer.register_parameter("w13_input_zero_point", w13_input_zero_point) |
| 1163 | set_weight_attrs(w13_input_zero_point, extra_weight_attrs) |
| 1164 | |
| 1165 | w2_input_zero_point = torch.nn.Parameter( |
| 1166 | torch.zeros(num_experts, dtype=torch.int8), |
| 1167 | requires_grad=False, |
| 1168 | ) |
| 1169 | layer.register_parameter("w2_input_zero_point", w2_input_zero_point) |
| 1170 | set_weight_attrs(w2_input_zero_point, extra_weight_attrs) |
| 1171 | |
| 1172 | if self.weight_qscheme == "per_channel": |
| 1173 | w13_weight_zero_point = torch.nn.Parameter( |
| 1174 | torch.zeros( |
| 1175 | num_experts, |
| 1176 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1177 | dtype=torch.int8, |
| 1178 | ), |
| 1179 | requires_grad=False, |
| 1180 | ) |
| 1181 | w2_weight_zero_point = torch.nn.Parameter( |
| 1182 | torch.zeros(num_experts, hidden_size, dtype=torch.int8), |
| 1183 | requires_grad=False, |
| 1184 | ) |
| 1185 | else: |
| 1186 | w13_weight_zero_point = torch.nn.Parameter( |
| 1187 | torch.zeros(num_experts, self.moe.w13_num_shards, dtype=torch.int8), |
| 1188 | requires_grad=False, |
| 1189 | ) |
| 1190 | w2_weight_zero_point = torch.nn.Parameter( |
| 1191 | torch.zeros(num_experts, dtype=torch.int8), |
| 1192 | requires_grad=False, |
| 1193 | ) |
| 1194 | layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point) |
| 1195 | set_weight_attrs(w13_weight_zero_point, extra_weight_attrs) |
| 1196 | layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point) |
| 1197 | set_weight_attrs(w2_weight_zero_point, extra_weight_attrs) |
| 1198 | |
| 1199 | # BIAS |
| 1200 | if self.has_bias: |
| 1201 | w13_bias = torch.nn.Parameter( |
| 1202 | torch.zeros( |
| 1203 | num_experts, |
| 1204 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1205 | dtype=torch.float32, |
| 1206 | ), |
| 1207 | requires_grad=False, |
| 1208 | ) |
| 1209 | layer.register_parameter("w13_bias", w13_bias) |
| 1210 | set_weight_attrs(w13_bias, extra_weight_attrs) |
| 1211 | w2_bias = torch.nn.Parameter( |
| 1212 | torch.zeros(num_experts, hidden_size, dtype=torch.float32), |
| 1213 | requires_grad=False, |
| 1214 | ) |
| 1215 | layer.register_parameter("w2_bias", w2_bias) |
| 1216 | set_weight_attrs(w2_bias, extra_weight_attrs) |
| 1217 | else: |
| 1218 | layer.w13_bias, layer.w2_bias = None, None |
| 1219 | |
| 1220 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 1221 | # Discard zero points (INT8 fused MoE kernel uses symmetric quant) |
| 1222 | for attr in ( |
| 1223 | "w13_input_zero_point", |
| 1224 | "w2_input_zero_point", |
| 1225 | "w13_weight_zero_point", |
| 1226 | "w2_weight_zero_point", |
| 1227 | ): |
| 1228 | if hasattr(layer, attr): |
| 1229 | delattr(layer, attr) |
| 1230 | |
| 1231 | # For static input scales, collapse the per-expert scales to a single |
| 1232 | # value (the legacy fused_experts path expects one scale per layer). |
| 1233 | if self.static_input_scales: |
| 1234 | if layer.w13_input_scale is None or layer.w2_input_scale is None: |
| 1235 | raise ValueError( |
| 1236 | "QuantConfig has static quantization, but found " |
| 1237 | "activation scales are None." |
| 1238 | ) |
| 1239 | if not all_close_1d(layer.w13_input_scale) or not all_close_1d( |
| 1240 | layer.w2_input_scale |
| 1241 | ): |
| 1242 | logger.warning_once( |
| 1243 | "Found input_scales that are not equal for " |
| 1244 | "INT8 MoE layer. Using the maximum across experts " |
| 1245 | "for each layer." |
| 1246 | ) |
| 1247 | layer.w13_input_scale = torch.nn.Parameter( |
| 1248 | layer.w13_input_scale.max(), requires_grad=False |
| 1249 | ) |
| 1250 | layer.w2_input_scale = torch.nn.Parameter( |
| 1251 | layer.w2_input_scale.max(), requires_grad=False |
| 1252 | ) |
| 1253 | |
| 1254 | # Per-channel scales: 2D [E, N] -> 3D [E, N, 1] for the int8 MoE kernel. |
| 1255 | if self.weight_qscheme == "per_channel": |
| 1256 | for attr in ("w13_weight_scale", "w2_weight_scale"): |
| 1257 | param = getattr(layer, attr, None) |
| 1258 | if param is not None and param.dim() == 2: |
| 1259 | replace_parameter( |
| 1260 | layer, |
| 1261 | attr, |
| 1262 | torch.nn.Parameter( |
| 1263 | param.data.unsqueeze(-1).contiguous(), |
| 1264 | requires_grad=False, |
| 1265 | ), |
| 1266 | ) |
| 1267 | |
| 1268 | # For per-tensor weights, merge the w1/w3 scales into a single |
| 1269 | # per-expert scale (dequant -> requant at the max scale). |
| 1270 | if self.weight_qscheme == "per_tensor": |
| 1271 | assert layer.w13_weight_scale is not None |
| 1272 | shard_size = layer.intermediate_size_per_partition |
| 1273 | max_w13_scales = layer.w13_weight_scale.max(dim=1).values |
| 1274 | |
| 1275 | for expert_id in range(layer.local_num_experts): |
| 1276 | start = 0 |
| 1277 | for shard_id in range(self.moe.w13_num_shards): |
| 1278 | dq_weight = per_tensor_dequantize( |
| 1279 | layer.w13_weight[expert_id][start : start + shard_size, :], |
| 1280 | layer.w13_weight_scale[expert_id][shard_id], |
| 1281 | ) |
| 1282 | layer.w13_weight[expert_id][start : start + shard_size, :], _, _ = ( |
| 1283 | ops.scaled_int8_quant( |
| 1284 | dq_weight, |
| 1285 | scale=max_w13_scales[expert_id], |
| 1286 | ) |
| 1287 | ) |
| 1288 | start += shard_size |
| 1289 | |
| 1290 | layer.w13_weight_scale = torch.nn.Parameter( |
| 1291 | max_w13_scales, requires_grad=False |
| 1292 | ) |
| 1293 | |
| 1294 | # Dynamic activations run through the oracle's modular kernel; static |
| 1295 | # activations use the legacy fused_experts path in apply(). |
| 1296 | if not self.static_input_scales: |
| 1297 | assert self.int8_backend is not None |
| 1298 | assert self.experts_cls is not None |
| 1299 | w13, w2 = convert_to_int8_moe_kernel_format( |
| 1300 | int8_backend=self.int8_backend, |
| 1301 | w13=layer.w13_weight, |
| 1302 | w2=layer.w2_weight, |
| 1303 | layer=layer, |
| 1304 | w13_scale=layer.w13_weight_scale, |
| 1305 | ) |
| 1306 | replace_parameter(layer, "w13_weight", w13) |
| 1307 | replace_parameter(layer, "w2_weight", w2) |
| 1308 | |
| 1309 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 1310 | assert self.moe_quant_config is not None |
| 1311 | |
| 1312 | if not self.static_input_scales: |
| 1313 | assert self.int8_backend is not None |
| 1314 | assert self.experts_cls is not None |
| 1315 | self.moe_kernel = make_int8_moe_kernel( |
| 1316 | int8_backend=self.int8_backend, |
| 1317 | moe_quant_config=self.moe_quant_config, |
| 1318 | moe_config=self.moe, |
| 1319 | experts_cls=self.experts_cls, |
| 1320 | routing_tables=layer._expert_routing_tables(), |
| 1321 | ) |
| 1322 | |
| 1323 | def get_fused_moe_quant_config( |
| 1324 | self, layer: torch.nn.Module |
| 1325 | ) -> FusedMoEQuantConfig | None: |
| 1326 | # Static-activation INT8 has no oracle backend (it uses the legacy |
| 1327 | # fused_experts path); build its config directly. |
| 1328 | if self.int8_backend is None: |
| 1329 | return int8_w8a8_moe_quant_config( |
| 1330 | w1_scale=layer.w13_weight_scale, |
| 1331 | w2_scale=layer.w2_weight_scale, |
| 1332 | a1_scale=layer.w13_input_scale, |
| 1333 | a2_scale=layer.w2_input_scale, |
| 1334 | w1_bias=getattr(layer, "w13_bias", None), |
| 1335 | w2_bias=getattr(layer, "w2_bias", None), |
| 1336 | per_act_token_quant=False, |
| 1337 | ) |
| 1338 | return make_int8_moe_quant_config( |
| 1339 | int8_backend=self.int8_backend, |
| 1340 | w1_scale=layer.w13_weight_scale, |
| 1341 | w2_scale=layer.w2_weight_scale, |
| 1342 | a1_scale=layer.w13_input_scale, |
| 1343 | a2_scale=layer.w2_input_scale, |
| 1344 | w1_bias=getattr(layer, "w13_bias", None), |
| 1345 | w2_bias=getattr(layer, "w2_bias", None), |
| 1346 | per_act_token_quant=(self.weight_qscheme == "per_channel"), |
| 1347 | layer=layer, |
| 1348 | ) |
| 1349 | |
| 1350 | def apply( |
| 1351 | self, |
| 1352 | layer: RoutedExperts, |
| 1353 | x: torch.Tensor, |
| 1354 | topk_weights: torch.Tensor, |
| 1355 | topk_ids: torch.Tensor, |
| 1356 | shared_experts: SharedExperts | None, |
| 1357 | shared_experts_input: torch.Tensor | None, |
| 1358 | ) -> torch.Tensor: |
| 1359 | if self.moe_kernel is not None: |
| 1360 | return self.moe_kernel.apply( |
| 1361 | hidden_states=x, |
| 1362 | w1=layer.w13_weight, |
| 1363 | w2=layer.w2_weight, |
| 1364 | topk_weights=topk_weights, |
| 1365 | topk_ids=topk_ids, |
| 1366 | activation=layer.activation, |
| 1367 | global_num_experts=layer.global_num_experts, |
| 1368 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 1369 | expert_map=layer.expert_map, |
| 1370 | shared_experts_input=shared_experts_input, |
| 1371 | ) |
| 1372 | |
| 1373 | # Static-activation INT8 MoE: legacy monolithic path (the modular kernel |
| 1374 | # quantizes activations dynamically and cannot apply a loaded scale). |
| 1375 | from vllm.model_executor.layers.fused_moe import fused_experts |
| 1376 | |
| 1377 | return fused_experts( |
| 1378 | hidden_states=x, |
| 1379 | w1=layer.w13_weight, |
| 1380 | w2=layer.w2_weight, |
| 1381 | topk_weights=topk_weights, |
| 1382 | topk_ids=topk_ids, |
| 1383 | activation=layer.activation, |
| 1384 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 1385 | global_num_experts=layer.global_num_experts, |
| 1386 | expert_map=layer.expert_map, |
| 1387 | quant_config=self.moe_quant_config, |
| 1388 | ) |
| 1389 | |
| 1390 | |
| 1391 | class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod): |
| 1392 | supported_activation_quant_keys = [ |
| 1393 | kFp8DynamicTokenSym, |
| 1394 | kFp8StaticTensorSym, |
| 1395 | ] |
| 1396 | supported_weight_quant_keys = [kInt4W4A8StaticChannelSym] |
| 1397 | |
| 1398 | def __init__( |
| 1399 | self, |
| 1400 | moe: FusedMoEConfig, |
| 1401 | activation_quant_key: QuantKey | None, |
| 1402 | ): |
| 1403 | super().__init__(moe, kInt4W4A8StaticChannelSym, activation_quant_key) |
| 1404 | |
| 1405 | assert rocm_aiter_ops.is_fused_moe_enabled(), ( |
| 1406 | "W4A8 FP8 MoE requires ROCm AITER fused MoE support." |
| 1407 | ) |
| 1408 | |
| 1409 | def create_weights( |
| 1410 | self, |
| 1411 | layer: RoutedExperts, |
| 1412 | num_experts: int, |
| 1413 | hidden_size: int, |
| 1414 | intermediate_size_per_partition: int, |
| 1415 | params_dtype: torch.dtype, |
| 1416 | **extra_weight_attrs, |
| 1417 | ): |
| 1418 | params_dtype = torch.uint32 |
| 1419 | w13_weight = torch.nn.Parameter( |
| 1420 | torch.zeros( |
| 1421 | num_experts, |
| 1422 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1423 | hidden_size // 8, # INT32 packing for W4 |
| 1424 | dtype=params_dtype, |
| 1425 | ), |
| 1426 | requires_grad=False, |
| 1427 | ) |
| 1428 | w2_weight = torch.nn.Parameter( |
| 1429 | torch.zeros( |
| 1430 | num_experts, |
| 1431 | hidden_size, |
| 1432 | intermediate_size_per_partition // 8, # INT32 packing for W4 |
| 1433 | dtype=params_dtype, |
| 1434 | ), |
| 1435 | requires_grad=False, |
| 1436 | ) |
| 1437 | layer.register_parameter("w13_weight", w13_weight) |
| 1438 | layer.register_parameter("w2_weight", w2_weight) |
| 1439 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 1440 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 1441 | |
| 1442 | # Per-tensor fp8 weight scales |
| 1443 | w13_weight_scale = torch.nn.Parameter( |
| 1444 | torch.ones(num_experts, self.moe.w13_num_shards, dtype=torch.float32), |
| 1445 | requires_grad=False, |
| 1446 | ) |
| 1447 | w2_weight_scale = torch.nn.Parameter( |
| 1448 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 1449 | ) |
| 1450 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 1451 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 1452 | extra_weight_attrs.update( |
| 1453 | {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} |
| 1454 | ) |
| 1455 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 1456 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 1457 | |
| 1458 | # Per-channel int4 weight scales |
| 1459 | w13_weight_scale_2 = torch.nn.Parameter( |
| 1460 | torch.ones( |
| 1461 | num_experts, |
| 1462 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1463 | dtype=torch.float32, |
| 1464 | ), |
| 1465 | requires_grad=False, |
| 1466 | ) |
| 1467 | w2_weight_scale_2 = torch.nn.Parameter( |
| 1468 | torch.ones(num_experts, hidden_size, dtype=torch.float32), |
| 1469 | requires_grad=False, |
| 1470 | ) |
| 1471 | layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2) |
| 1472 | layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) |
| 1473 | extra_weight_attrs.update( |
| 1474 | {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} |
| 1475 | ) |
| 1476 | set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) |
| 1477 | set_weight_attrs(w2_weight_scale_2, extra_weight_attrs) |
| 1478 | |
| 1479 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 1480 | shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights( |
| 1481 | layer.w13_weight.data, layer.w2_weight.data |
| 1482 | ) |
| 1483 | layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False) |
| 1484 | layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False) |
| 1485 | |
| 1486 | # INT4-FP8 : offset INT4 w13_weight_scale1 to single w13_weight_scale |
| 1487 | # Fp8 moe kernel needs single fp8 w13_weight_scale for w13 per expert. |
| 1488 | # We won't do requant each expert's fp8 weight (not direct available), |
| 1489 | # instead we adjust half of INT4 w13_weight_scale1 numbers |
| 1490 | shard_size = layer.intermediate_size_per_partition |
| 1491 | max_w13_scales = layer.w13_weight_scale.max(dim=1).values |
| 1492 | assert torch.all(max_w13_scales != 0), "fp8 weight scale cannot be zero." |
| 1493 | for expert_id in range(layer.local_num_experts): |
| 1494 | start = 0 |
| 1495 | max_w13_scale_fp8 = max_w13_scales[expert_id] |
| 1496 | for shard_id in range(self.moe.w13_num_shards): |
| 1497 | if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8: |
| 1498 | int4_rescale = ( |
| 1499 | layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8 |
| 1500 | ) |
| 1501 | layer.w13_weight_scale_2[expert_id][start : start + shard_size] *= ( |
| 1502 | int4_rescale |
| 1503 | ) |
| 1504 | start += shard_size |
| 1505 | |
| 1506 | layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False) |
| 1507 | |
| 1508 | # special hack to asm_moe, which takes (weight_scale1 * weight_scale) as post |
| 1509 | # GEMM scaling optimal design - shall apply per-column weight_scale1 before |
| 1510 | # GEMM, and weight_scale post |
| 1511 | for expert_id in range(layer.local_num_experts): |
| 1512 | layer.w13_weight_scale_2[expert_id] *= max_w13_scales[expert_id] |
| 1513 | layer.w2_weight_scale_2[expert_id] *= layer.w2_weight_scale[expert_id] |
| 1514 | |
| 1515 | def get_fused_moe_quant_config(self, layer): |
| 1516 | return fp8_w8a8_moe_quant_config( |
| 1517 | w1_scale=layer.w13_weight_scale_2, |
| 1518 | w2_scale=layer.w2_weight_scale_2, |
| 1519 | per_out_ch_quant=True, |
| 1520 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 1521 | ) |
| 1522 | |
| 1523 | def apply( |
| 1524 | self, |
| 1525 | layer: RoutedExperts, |
| 1526 | x: torch.Tensor, |
| 1527 | topk_weights: torch.Tensor, |
| 1528 | topk_ids: torch.Tensor, |
| 1529 | shared_experts: SharedExperts | None, |
| 1530 | shared_experts_input: torch.Tensor | None, |
| 1531 | ) -> torch.Tensor: |
| 1532 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 1533 | rocm_aiter_fused_experts, |
| 1534 | ) |
| 1535 | |
| 1536 | return rocm_aiter_fused_experts( |
| 1537 | hidden_states=x, |
| 1538 | w1=layer.w13_weight, |
| 1539 | w2=layer.w2_weight, |
| 1540 | topk_weights=topk_weights, |
| 1541 | topk_ids=topk_ids, |
| 1542 | activation=layer.activation, |
| 1543 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 1544 | quant_config=self.moe_quant_config, |
| 1545 | moe_config=layer.moe_config, |
| 1546 | expert_mask=layer.expert_mask, |
| 1547 | ) |
| 1548 | |
| 1549 | |
| 1550 | class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): |
| 1551 | supported_activation_quant_keys = [ |
| 1552 | *_ACTIVATION_QUANT_KEY_MAP.values(), |
| 1553 | kFp8DynamicTensorSym, |
| 1554 | kFp8StaticTensorSym, |
| 1555 | None, |
| 1556 | ] |
| 1557 | supported_weight_quant_keys = [*_WEIGHT_QUANT_KEY_MAP.values()] |
| 1558 | |
| 1559 | def __init__( |
| 1560 | self, |
| 1561 | moe: FusedMoEConfig, |
| 1562 | weight_quant_key: QuantKey, |
| 1563 | activation_quant_key: QuantKey | None, |
| 1564 | ): |
| 1565 | super().__init__(moe, weight_quant_key, activation_quant_key) |
| 1566 | self.weight_dtype = next( |
| 1567 | dtype |
| 1568 | for dtype, quant_key in _WEIGHT_QUANT_KEY_MAP.items() |
| 1569 | if quant_key == weight_quant_key |
| 1570 | ) |
| 1571 | if activation_quant_key in {kFp8DynamicTensorSym, kFp8StaticTensorSym}: |
| 1572 | self.input_dtype: str | None = "fp8" |
| 1573 | elif activation_quant_key is None: |
| 1574 | self.input_dtype = None |
| 1575 | else: |
| 1576 | self.input_dtype = next( |
| 1577 | dtype |
| 1578 | for dtype, quant_key in _ACTIVATION_QUANT_KEY_MAP.items() |
| 1579 | if quant_key == activation_quant_key |
| 1580 | ) |
| 1581 | |
| 1582 | self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype( |
| 1583 | self.input_dtype, self.weight_dtype |
| 1584 | ) |
| 1585 | |
| 1586 | if self.ocp_mx_scheme is None: |
| 1587 | raise ValueError( |
| 1588 | f"Unsupported OCP MX dtype combination for MoE: " |
| 1589 | f"input_dtype={self.input_dtype}, weight_dtype={self.weight_dtype}. " |
| 1590 | f"MXFP8 experts are not supported through the MoE path; only " |
| 1591 | f"the linear path handles MXFP8. Supported MoE weight dtypes " |
| 1592 | f"are mxfp4, mxfp6_e3m2, and mxfp6_e2m3." |
| 1593 | ) |
| 1594 | |
| 1595 | # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes, |
| 1596 | # use kernel abstraction for all OCP MX MOE implementations. |
| 1597 | self.mxfp4_backend: Mxfp4MoeBackend = Mxfp4MoeBackend.NONE |
| 1598 | self.experts_cls: type[mk.FusedMoEExperts] | None = None |
| 1599 | self.moe_kernel: mk.FusedMoEKernel | None = None |
| 1600 | |
| 1601 | # Used for triton kernel precision configs (W4A8, TRITON backends) |
| 1602 | self.w13_precision_config = None |
| 1603 | self.w2_precision_config = None |
| 1604 | |
| 1605 | self.static_input_scales = activation_quant_key == kFp8StaticTensorSym |
| 1606 | |
| 1607 | # Select backend based on OCP MX scheme |
| 1608 | if self.ocp_mx_scheme == "w_mxfp4": |
| 1609 | # W4A16: weight-only MXFP4 |
| 1610 | self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) |
| 1611 | elif self.ocp_mx_scheme == "w_mxfp4_a_fp8" and self.static_input_scales: |
| 1612 | # W4A8: MXFP4 weights + static FP8 activations |
| 1613 | self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( |
| 1614 | moe, activation_key=kFp8StaticTensorSym |
| 1615 | ) |
| 1616 | elif self.ocp_mx_scheme == "w_mxfp4_a_mxfp4": |
| 1617 | # W4A4: MXFP4 weights + MXFP4 activations |
| 1618 | self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( |
| 1619 | moe, activation_key=kMxfp4Dynamic |
| 1620 | ) |
| 1621 | |
| 1622 | # Validation for unsupported schemes |
| 1623 | if any( |
| 1624 | self.ocp_mx_scheme.endswith(a_scheme) |
| 1625 | for a_scheme in ["a_mxfp4", "a_mxfp6_e3m2", "a_mxfp6_e2m3"] |
| 1626 | ): |
| 1627 | if self.static_input_scales: |
| 1628 | raise NotImplementedError( |
| 1629 | "QuarkOCP_MX_MoEMethod with static input scales is currently " |
| 1630 | f"not implemented for OCP MX scheme {self.ocp_mx_scheme}. " |
| 1631 | "Please open an issue." |
| 1632 | ) |
| 1633 | elif self.ocp_mx_scheme.endswith("a_fp8") and not self.static_input_scales: |
| 1634 | raise NotImplementedError( |
| 1635 | "QuarkOCP_MX_MoEMethod with dynamic input scales is currently " |
| 1636 | f"not implemented for OCP MX scheme {self.ocp_mx_scheme}. " |
| 1637 | "Please open an issue." |
| 1638 | ) |
| 1639 | |
| 1640 | self.model_type = getattr( |
| 1641 | get_current_vllm_config().model_config.hf_config, "model_type", None |
| 1642 | ) |
| 1643 | |
| 1644 | # If no native backend available, use emulation. |
| 1645 | if self.mxfp4_backend is Mxfp4MoeBackend.NONE: |
| 1646 | self.mxfp4_backend = Mxfp4MoeBackend.EMULATION |
| 1647 | |
| 1648 | self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] |
| 1649 | |
| 1650 | logger.info_once( |
| 1651 | f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}" |
| 1652 | ) |
| 1653 | |
| 1654 | def maybe_roundup_sizes( |
| 1655 | self, |
| 1656 | hidden_size: int, |
| 1657 | intermediate_size_per_partition: int, |
| 1658 | act_dtype: torch.dtype, |
| 1659 | moe_parallel_config: FusedMoEParallelConfig, |
| 1660 | ) -> tuple[int, int]: |
| 1661 | hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes( |
| 1662 | hidden_size=hidden_size, |
| 1663 | intermediate_size_per_partition=intermediate_size_per_partition, |
| 1664 | act_dtype=act_dtype, |
| 1665 | moe_parallel_config=moe_parallel_config, |
| 1666 | ) |
| 1667 | # Round per-partition sizes up to each backend's requirement. Emulation is |
| 1668 | # handled inside the helper too (OCP MX block alignment), so no special-case. |
| 1669 | if self.mxfp4_backend is not None: |
| 1670 | hidden_size, intermediate_size_per_partition = ( |
| 1671 | mxfp4_round_up_hidden_size_and_intermediate_size( |
| 1672 | self.mxfp4_backend, hidden_size, intermediate_size_per_partition |
| 1673 | ) |
| 1674 | ) |
| 1675 | return hidden_size, intermediate_size_per_partition |
| 1676 | |
| 1677 | def get_packed_dim(self, dim: int, quant_dtype: str): |
| 1678 | if quant_dtype == "mxfp4": |
| 1679 | assert dim % 2 == 0 |
| 1680 | return dim // 2 |
| 1681 | elif quant_dtype in {"mxfp6_e3m2", "mxfp6_e2m3"}: |
| 1682 | # FP6 packs 4 * 6 = 24 bits on 3 bytes. |
| 1683 | assert (dim * 3) % 4 == 0 |
| 1684 | return (dim * 3) // 4 |
| 1685 | else: |
| 1686 | raise NotImplementedError( |
| 1687 | f"Unsupported quant_dtype in QuarkOCP_MX_MoEMethod." |
| 1688 | f"get_packed_dim: {quant_dtype}. MXFP8 is not supported for " |
| 1689 | f"MoE experts." |
| 1690 | ) |
| 1691 | |
| 1692 | def create_weights( |
| 1693 | self, |
| 1694 | layer: RoutedExperts, |
| 1695 | num_experts: int, |
| 1696 | hidden_size: int, |
| 1697 | intermediate_size_per_partition: int, |
| 1698 | params_dtype: torch.dtype, |
| 1699 | **extra_weight_attrs, |
| 1700 | ): |
| 1701 | # Add the quantization method used (per tensor/grouped/channel) |
| 1702 | # to ensure the weight scales are loaded in properly |
| 1703 | extra_weight_attrs.update( |
| 1704 | {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} |
| 1705 | ) |
| 1706 | |
| 1707 | params_dtype = torch.uint8 |
| 1708 | |
| 1709 | # WEIGHTS |
| 1710 | w13_weight = torch.nn.Parameter( |
| 1711 | torch.zeros( |
| 1712 | num_experts, |
| 1713 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1714 | self.get_packed_dim(hidden_size, self.weight_dtype), |
| 1715 | dtype=params_dtype, |
| 1716 | ), |
| 1717 | requires_grad=False, |
| 1718 | ) |
| 1719 | layer.register_parameter("w13_weight", w13_weight) |
| 1720 | |
| 1721 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 1722 | |
| 1723 | w2_weight = torch.nn.Parameter( |
| 1724 | torch.zeros( |
| 1725 | num_experts, |
| 1726 | hidden_size, |
| 1727 | self.get_packed_dim(intermediate_size_per_partition, self.weight_dtype), |
| 1728 | dtype=params_dtype, |
| 1729 | ), |
| 1730 | requires_grad=False, |
| 1731 | ) |
| 1732 | layer.register_parameter("w2_weight", w2_weight) |
| 1733 | |
| 1734 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 1735 | |
| 1736 | # WEIGHT_SCALES |
| 1737 | w13_weight_scale = torch.nn.Parameter( |
| 1738 | torch.ones( |
| 1739 | num_experts, |
| 1740 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1741 | hidden_size // OCP_MX_BLOCK_SIZE, |
| 1742 | dtype=params_dtype, |
| 1743 | ), |
| 1744 | requires_grad=False, |
| 1745 | ) |
| 1746 | w2_weight_scale = torch.nn.Parameter( |
| 1747 | torch.ones( |
| 1748 | num_experts, |
| 1749 | hidden_size, |
| 1750 | intermediate_size_per_partition // OCP_MX_BLOCK_SIZE, |
| 1751 | dtype=params_dtype, |
| 1752 | ), |
| 1753 | requires_grad=False, |
| 1754 | ) |
| 1755 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 1756 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 1757 | |
| 1758 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 1759 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 1760 | |
| 1761 | if self.has_bias: |
| 1762 | w13_bias = torch.nn.Parameter( |
| 1763 | torch.zeros( |
| 1764 | num_experts, |
| 1765 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 1766 | dtype=torch.float32, |
| 1767 | ), |
| 1768 | requires_grad=False, |
| 1769 | ) |
| 1770 | layer.register_parameter("w13_bias", w13_bias) |
| 1771 | set_weight_attrs(w13_bias, extra_weight_attrs) |
| 1772 | |
| 1773 | w2_bias = torch.nn.Parameter( |
| 1774 | torch.zeros(num_experts, hidden_size, dtype=torch.float32), |
| 1775 | requires_grad=False, |
| 1776 | ) |
| 1777 | layer.register_parameter("w2_bias", w2_bias) |
| 1778 | set_weight_attrs(w2_bias, extra_weight_attrs) |
| 1779 | else: |
| 1780 | layer.w13_bias, layer.w2_bias = None, None |
| 1781 | |
| 1782 | # INPUT_SCALES |
| 1783 | if self.static_input_scales: |
| 1784 | w13_input_scale = torch.nn.Parameter( |
| 1785 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 1786 | ) |
| 1787 | layer.register_parameter("w13_input_scale", w13_input_scale) |
| 1788 | set_weight_attrs(w13_input_scale, extra_weight_attrs) |
| 1789 | |
| 1790 | w2_input_scale = torch.nn.Parameter( |
| 1791 | torch.ones(num_experts, dtype=torch.float32), requires_grad=False |
| 1792 | ) |
| 1793 | layer.register_parameter("w2_input_scale", w2_input_scale) |
| 1794 | set_weight_attrs(w2_input_scale, extra_weight_attrs) |
| 1795 | else: |
| 1796 | layer.w13_input_scale = None |
| 1797 | layer.w2_input_scale = None |
| 1798 | |
| 1799 | def process_weights_after_loading(self, layer): |
| 1800 | self._setup_kernel(layer) |
| 1801 | |
| 1802 | def _setup_kernel(self, layer: RoutedExperts): |
| 1803 | """Setup kernel using oracle functions for MXFP4 schemes (W4A16, W4A8).""" |
| 1804 | w13_bias = getattr(layer, "w13_bias", None) |
| 1805 | w2_bias = getattr(layer, "w2_bias", None) |
| 1806 | |
| 1807 | # Convert weights to kernel format (handles all backend-specific logic) |
| 1808 | w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( |
| 1809 | convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( |
| 1810 | mxfp4_backend=self.mxfp4_backend, |
| 1811 | layer=layer, |
| 1812 | w13_weight=layer.w13_weight, |
| 1813 | w2_weight=layer.w2_weight, |
| 1814 | w13_weight_scale=layer.w13_weight_scale, |
| 1815 | w2_weight_scale=layer.w2_weight_scale, |
| 1816 | w13_bias=w13_bias, |
| 1817 | w2_bias=w2_bias, |
| 1818 | w13_input_scale=layer.w13_input_scale, |
| 1819 | w2_input_scale=layer.w2_input_scale, |
| 1820 | ) |
| 1821 | ) |
| 1822 | |
| 1823 | # Handle weight/scale assignment based on backend type |
| 1824 | if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( |
| 1825 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 1826 | ): |
| 1827 | # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor |
| 1828 | # Store on layer for apply(), scales are PrecisionConfig |
| 1829 | layer.w13_weight = w13 |
| 1830 | layer.w2_weight = w2 |
| 1831 | self.w13_precision_config = w13_scale |
| 1832 | self.w2_precision_config = w2_scale |
| 1833 | else: |
| 1834 | # Standard backends: replace parameters |
| 1835 | replace_parameter(layer, "w13_weight", w13) |
| 1836 | replace_parameter(layer, "w2_weight", w2) |
| 1837 | replace_parameter(layer, "w13_weight_scale", w13_scale) |
| 1838 | replace_parameter(layer, "w2_weight_scale", w2_scale) |
| 1839 | |
| 1840 | if w13_bias is not None and w2_bias is not None: |
| 1841 | replace_parameter(layer, "w13_bias", w13_bias) |
| 1842 | replace_parameter(layer, "w2_bias", w2_bias) |
| 1843 | |
| 1844 | torch.accelerator.empty_cache() |
| 1845 | |
| 1846 | # Build quant config and kernel |
| 1847 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 1848 | if self.moe_quant_config is not None and self.experts_cls is not None: |
| 1849 | self.moe_kernel = make_mxfp4_moe_kernel( |
| 1850 | moe_quant_config=self.moe_quant_config, |
| 1851 | moe_config=self.moe, |
| 1852 | mxfp4_backend=self.mxfp4_backend, |
| 1853 | experts_cls=self.experts_cls, |
| 1854 | routing_tables=layer._expert_routing_tables(), |
| 1855 | ) |
| 1856 | self.moe_kernel.fused_experts.process_weights_after_loading(layer) |
| 1857 | |
| 1858 | def get_fused_moe_quant_config( |
| 1859 | self, layer: RoutedExperts |
| 1860 | ) -> FusedMoEQuantConfig | None: |
| 1861 | # For oracle-based backends (W4A16, W4A8), use make_mxfp4_moe_quant_config |
| 1862 | if self.mxfp4_backend not in (Mxfp4MoeBackend.NONE, Mxfp4MoeBackend.EMULATION): |
| 1863 | # Determine scale source based on backend type |
| 1864 | if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( |
| 1865 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 1866 | ): |
| 1867 | w1_scale = self.w13_precision_config |
| 1868 | w2_scale = self.w2_precision_config |
| 1869 | else: |
| 1870 | w1_scale = layer.w13_weight_scale |
| 1871 | w2_scale = layer.w2_weight_scale |
| 1872 | |
| 1873 | return make_mxfp4_moe_quant_config( |
| 1874 | mxfp4_backend=self.mxfp4_backend, |
| 1875 | w1_scale=w1_scale, |
| 1876 | w2_scale=w2_scale, |
| 1877 | w1_bias=getattr(layer, "w13_bias", None), |
| 1878 | w2_bias=getattr(layer, "w2_bias", None), |
| 1879 | a1_scale=getattr(layer, "w13_input_scale", None), |
| 1880 | a2_scale=getattr(layer, "w2_input_scale", None), |
| 1881 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 1882 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 1883 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 1884 | layer=layer, |
| 1885 | ) |
| 1886 | |
| 1887 | # Emulation and other schemes |
| 1888 | if self.ocp_mx_scheme == "w_mxfp4": |
| 1889 | return mxfp4_w4a16_moe_quant_config( |
| 1890 | w1_scale=layer.w13_weight_scale, |
| 1891 | w2_scale=layer.w2_weight_scale, |
| 1892 | w1_bias=layer.w13_bias, |
| 1893 | w2_bias=layer.w2_bias, |
| 1894 | ) |
| 1895 | elif self.ocp_mx_scheme == "w_mxfp4_a_fp8": |
| 1896 | return mxfp4_w4a8_moe_quant_config( |
| 1897 | w1_scale=layer.w13_weight_scale, |
| 1898 | w2_scale=layer.w2_weight_scale, |
| 1899 | a1_scale=layer.w13_input_scale, |
| 1900 | a2_scale=layer.w2_input_scale, |
| 1901 | w1_bias=layer.w13_bias, |
| 1902 | w2_bias=layer.w2_bias, |
| 1903 | block_shape=None, |
| 1904 | ) |
| 1905 | elif self.ocp_mx_scheme in ["w_mxfp6_e3m2_a_fp8", "w_mxfp6_e2m3_a_fp8"]: |
| 1906 | raise NotImplementedError( |
| 1907 | "Currently there is no corresponding fused moe quant config configured " |
| 1908 | f"in vLLM for OCP MX scheme {self.ocp_mx_scheme}. Please open an issue." |
| 1909 | ) |
| 1910 | else: |
| 1911 | assert self.input_dtype is not None |
| 1912 | return ocp_mx_moe_quant_config( |
| 1913 | quant_dtype=self.input_dtype, |
| 1914 | weight_dtype=self.weight_dtype, |
| 1915 | w1_scale=layer.w13_weight_scale, |
| 1916 | w2_scale=layer.w2_weight_scale, |
| 1917 | w1_bias=layer.w13_bias, |
| 1918 | w2_bias=layer.w2_bias, |
| 1919 | a1_scale=None, |
| 1920 | a2_scale=None, |
| 1921 | block_shape=None, |
| 1922 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 1923 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 1924 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 1925 | ) |
| 1926 | |
| 1927 | @property |
| 1928 | def supports_eplb(self) -> bool: |
| 1929 | # AITER shuffle keeps expert dim outermost, so EPLB row moves are layout-safe. |
| 1930 | return True |
| 1931 | |
| 1932 | @property |
| 1933 | def is_monolithic(self) -> bool: |
| 1934 | if self.moe_kernel is not None: |
| 1935 | return self.moe_kernel.is_monolithic |
| 1936 | return False |
| 1937 | |
| 1938 | def apply( |
| 1939 | self, |
| 1940 | layer: RoutedExperts, |
| 1941 | x: torch.Tensor, |
| 1942 | topk_weights: torch.Tensor, |
| 1943 | topk_ids: torch.Tensor, |
| 1944 | shared_experts: SharedExperts | None, |
| 1945 | shared_experts_input: torch.Tensor | None, |
| 1946 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1947 | assert self.moe_kernel is not None |
| 1948 | return self.moe_kernel.apply( |
| 1949 | hidden_states=x, |
| 1950 | w1=layer.w13_weight, |
| 1951 | w2=layer.w2_weight, |
| 1952 | topk_weights=topk_weights, |
| 1953 | topk_ids=topk_ids, |
| 1954 | activation=layer.activation, |
| 1955 | global_num_experts=layer.global_num_experts, |
| 1956 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 1957 | expert_map=layer.expert_map, |
| 1958 | shared_experts_input=shared_experts_input, |
| 1959 | ) |
| 1960 | |
| 1961 | def apply_monolithic( |
| 1962 | self, |
| 1963 | layer: RoutedExperts, |
| 1964 | x: torch.Tensor, |
| 1965 | router_logits: torch.Tensor, |
| 1966 | input_ids: torch.Tensor | None = None, |
| 1967 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1968 | assert self.is_monolithic |
| 1969 | assert self.moe_kernel is not None |
| 1970 | return self.moe_kernel.apply_monolithic( |
| 1971 | hidden_states=x, |
| 1972 | w1=layer.w13_weight, |
| 1973 | w2=layer.w2_weight, |
| 1974 | router_logits=router_logits, |
| 1975 | activation=layer.activation, |
| 1976 | global_num_experts=layer.global_num_experts, |
| 1977 | expert_map=layer.expert_map, |
| 1978 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 1979 | num_expert_group=layer.num_expert_group, |
| 1980 | topk_group=layer.topk_group, |
| 1981 | e_score_correction_bias=layer.e_score_correction_bias, |
| 1982 | routed_scaling_factor=layer.routed_scaling_factor, |
| 1983 | ) |
| 1984 | |
| 1985 | |
| 1986 | class QuarkNvfp4MoEMethod(QuarkMoEMethod): |
| 1987 | supported_activation_quant_keys = [kNvfp4Dynamic] |
| 1988 | supported_weight_quant_keys = [kNvfp4Static] |
| 1989 | |
| 1990 | def __init__( |
| 1991 | self, |
| 1992 | moe: FusedMoEConfig, |
| 1993 | quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821 |
| 1994 | activation_quant_key: QuantKey | None, |
| 1995 | ): |
| 1996 | super().__init__(moe, kNvfp4Static, activation_quant_key) |
| 1997 | self.quant_config = quant_config |
| 1998 | self.group_size = 16 |
| 1999 | |
| 2000 | # Select experts implementation. |
| 2001 | self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( |
| 2002 | config=self.moe, |
| 2003 | weight_key=self.weight_quant_key, |
| 2004 | activation_key=self.activation_quant_key, |
| 2005 | ) |
| 2006 | |
| 2007 | def create_weights( |
| 2008 | self, |
| 2009 | layer: torch.nn.Module, |
| 2010 | num_experts: int, |
| 2011 | hidden_size: int, |
| 2012 | intermediate_size_per_partition: int, |
| 2013 | params_dtype: torch.dtype, |
| 2014 | **extra_weight_attrs, |
| 2015 | ): |
| 2016 | layer.num_experts = num_experts |
| 2017 | layer.params_dtype = params_dtype |
| 2018 | layer.quant_config = self.quant_config |
| 2019 | weight_dtype = torch.uint8 |
| 2020 | weight_scale_dtype = torch.float8_e4m3fn |
| 2021 | |
| 2022 | # GEMM 1 - w13 weight |
| 2023 | w13_weight = torch.nn.Parameter( |
| 2024 | torch.empty( |
| 2025 | num_experts, |
| 2026 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 2027 | # 2 fp4 items are packed in the input dimension |
| 2028 | hidden_size // 2, |
| 2029 | dtype=weight_dtype, |
| 2030 | ), |
| 2031 | requires_grad=False, |
| 2032 | ) |
| 2033 | layer.register_parameter("w13_weight", w13_weight) |
| 2034 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 2035 | |
| 2036 | # GEMM 2 - w2 weight |
| 2037 | w2_weight = torch.nn.Parameter( |
| 2038 | torch.empty( |
| 2039 | num_experts, |
| 2040 | hidden_size, |
| 2041 | # 2 fp4 items are packed in the input dimension |
| 2042 | intermediate_size_per_partition // 2, |
| 2043 | dtype=weight_dtype, |
| 2044 | ), |
| 2045 | requires_grad=False, |
| 2046 | ) |
| 2047 | layer.register_parameter("w2_weight", w2_weight) |
| 2048 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 2049 | |
| 2050 | # Weight scales (per-group FP8 scales) |
| 2051 | w13_weight_scale = torch.nn.Parameter( |
| 2052 | torch.empty( |
| 2053 | num_experts, |
| 2054 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 2055 | hidden_size // self.group_size, |
| 2056 | dtype=weight_scale_dtype, |
| 2057 | ), |
| 2058 | requires_grad=False, |
| 2059 | ) |
| 2060 | layer.register_parameter("w13_weight_scale", w13_weight_scale) |
| 2061 | extra_weight_attrs.update( |
| 2062 | {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} |
| 2063 | ) |
| 2064 | set_weight_attrs(w13_weight_scale, extra_weight_attrs) |
| 2065 | |
| 2066 | w2_weight_scale = torch.nn.Parameter( |
| 2067 | torch.empty( |
| 2068 | num_experts, |
| 2069 | hidden_size, |
| 2070 | intermediate_size_per_partition // self.group_size, |
| 2071 | dtype=weight_scale_dtype, |
| 2072 | ), |
| 2073 | requires_grad=False, |
| 2074 | ) |
| 2075 | layer.register_parameter("w2_weight_scale", w2_weight_scale) |
| 2076 | set_weight_attrs(w2_weight_scale, extra_weight_attrs) |
| 2077 | |
| 2078 | # Global weight scales (per-tensor FP32 scales) |
| 2079 | extra_weight_attrs.update( |
| 2080 | {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} |
| 2081 | ) |
| 2082 | |
| 2083 | w13_weight_scale_2 = torch.nn.Parameter( |
| 2084 | torch.empty(num_experts, self.moe.w13_num_shards, dtype=torch.float32), |
| 2085 | requires_grad=False, |
| 2086 | ) |
| 2087 | layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2) |
| 2088 | set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) |
| 2089 | |
| 2090 | w2_weight_scale_2 = torch.nn.Parameter( |
| 2091 | torch.empty(num_experts, dtype=torch.float32), |
| 2092 | requires_grad=False, |
| 2093 | ) |
| 2094 | layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) |
| 2095 | set_weight_attrs(w2_weight_scale_2, extra_weight_attrs) |
| 2096 | |
| 2097 | # Input global scales (per-tensor FP32 scales) |
| 2098 | w13_input_scale_2 = torch.nn.Parameter( |
| 2099 | torch.empty(num_experts, self.moe.w13_num_shards, dtype=torch.float32), |
| 2100 | requires_grad=False, |
| 2101 | ) |
| 2102 | layer.register_parameter("w13_input_scale_2", w13_input_scale_2) |
| 2103 | set_weight_attrs(w13_input_scale_2, extra_weight_attrs) |
| 2104 | |
| 2105 | w2_input_scale_2 = torch.nn.Parameter( |
| 2106 | torch.empty(num_experts, dtype=torch.float32), |
| 2107 | requires_grad=False, |
| 2108 | ) |
| 2109 | layer.register_parameter("w2_input_scale_2", w2_input_scale_2) |
| 2110 | set_weight_attrs(w2_input_scale_2, extra_weight_attrs) |
| 2111 | |
| 2112 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 2113 | """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" |
| 2114 | # Match existing NVFP4 MoE paths: fused w13 uses the w1 global scale. |
| 2115 | if self.moe.is_act_and_mul and not torch.allclose( |
| 2116 | layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1] |
| 2117 | ): |
| 2118 | logger.warning_once( |
| 2119 | "w1_weight_scale_2 must match w3_weight_scale_2. " |
| 2120 | "Accuracy may be affected." |
| 2121 | ) |
| 2122 | |
| 2123 | w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0].contiguous() |
| 2124 | |
| 2125 | w2_weight_scale_2 = layer.w2_weight_scale_2 |
| 2126 | |
| 2127 | ( |
| 2128 | w13, |
| 2129 | w13_scale, |
| 2130 | w13_scale_2, |
| 2131 | a13_scale, |
| 2132 | w2, |
| 2133 | w2_scale, |
| 2134 | w2_scale_2, |
| 2135 | a2_scale, |
| 2136 | ) = convert_to_nvfp4_moe_kernel_format( |
| 2137 | nvfp4_backend=self.nvfp4_backend, |
| 2138 | layer=layer, |
| 2139 | w13=layer.w13_weight, |
| 2140 | w13_scale=layer.w13_weight_scale, |
| 2141 | w13_scale_2=w13_weight_scale_2, |
| 2142 | a13_scale=layer.w13_input_scale_2, |
| 2143 | w2=layer.w2_weight, |
| 2144 | w2_scale=layer.w2_weight_scale, |
| 2145 | w2_scale_2=w2_weight_scale_2, |
| 2146 | a2_scale=layer.w2_input_scale_2, |
| 2147 | is_act_and_mul=self.moe.is_act_and_mul, |
| 2148 | ) |
| 2149 | |
| 2150 | replace_parameter(layer, "w13_weight", w13) |
| 2151 | replace_parameter(layer, "w13_weight_scale", w13_scale) |
| 2152 | replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) |
| 2153 | replace_parameter(layer, "w13_input_scale_2", a13_scale) |
| 2154 | |
| 2155 | replace_parameter(layer, "w2_weight", w2) |
| 2156 | replace_parameter(layer, "w2_weight_scale", w2_scale) |
| 2157 | replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) |
| 2158 | replace_parameter(layer, "w2_input_scale_2", a2_scale) |
| 2159 | |
| 2160 | # Setup modular kernel. |
| 2161 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 2162 | if self.moe_quant_config: |
| 2163 | assert self.experts_cls is not None |
| 2164 | self.moe_kernel = make_nvfp4_moe_kernel( |
| 2165 | moe_quant_config=self.moe_quant_config, |
| 2166 | moe_config=self.moe, |
| 2167 | experts_cls=self.experts_cls, |
| 2168 | backend=self.nvfp4_backend, |
| 2169 | routing_tables=layer._expert_routing_tables(), |
| 2170 | ) |
| 2171 | |
| 2172 | def get_fused_moe_quant_config( |
| 2173 | self, layer: torch.nn.Module |
| 2174 | ) -> FusedMoEQuantConfig | None: |
| 2175 | return make_nvfp4_moe_quant_config( |
| 2176 | backend=self.nvfp4_backend, |
| 2177 | w13_scale=layer.w13_weight_scale, |
| 2178 | w2_scale=layer.w2_weight_scale, |
| 2179 | w13_scale_2=layer.w13_weight_scale_2, |
| 2180 | w2_scale_2=layer.w2_weight_scale_2, |
| 2181 | a13_scale=layer.w13_input_scale_2, |
| 2182 | a2_scale=layer.w2_input_scale_2, |
| 2183 | swiglu_limit=getattr(layer, "swiglu_limit", None), |
| 2184 | swiglu_alpha=getattr(layer, "swiglu_alpha", None), |
| 2185 | swiglu_beta=getattr(layer, "swiglu_beta", None), |
| 2186 | layer=layer, |
| 2187 | ) |
| 2188 | |
| 2189 | def apply( |
| 2190 | self, |
| 2191 | layer: RoutedExperts, |
| 2192 | x: torch.Tensor, |
| 2193 | topk_weights: torch.Tensor, |
| 2194 | topk_ids: torch.Tensor, |
| 2195 | shared_experts: SharedExperts | None, |
| 2196 | shared_experts_input: torch.Tensor | None, |
| 2197 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 2198 | assert self.moe_kernel is not None |
| 2199 | return self.moe_kernel.apply( |
| 2200 | x, |
| 2201 | layer.w13_weight, |
| 2202 | layer.w2_weight, |
| 2203 | topk_weights, |
| 2204 | topk_ids, |
| 2205 | activation=layer.activation, |
| 2206 | global_num_experts=layer.global_num_experts, |
| 2207 | expert_map=layer.expert_map, |
| 2208 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 2209 | shared_experts=shared_experts, |
| 2210 | shared_experts_input=shared_experts_input, |
| 2211 | ) |