| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Configure, prepare weights for, and assemble Humming MoE kernels.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 11 | import vllm.model_executor.layers.quantization.utils.humming.schema as humming_schema |
| 12 | from vllm import envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 15 | maybe_make_prepare_finalize, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.fused_moe.config import ( |
| 18 | FusedMoEConfig, |
| 19 | FusedMoEQuantConfig, |
| 20 | FusedMoEQuantDesc, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 23 | BatchedHummingGroupedExperts, |
| 24 | HummingGroupedExperts, |
| 25 | HummingIndexedExperts, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization.utils.humming.schema import ( |
| 28 | humming_is_layer_skipped, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 31 | FP8_DTYPE, |
| 32 | GroupShape, |
| 33 | QuantKey, |
| 34 | ) |
| 35 | from vllm.utils.import_utils import has_humming |
| 36 | |
| 37 | if TYPE_CHECKING: |
| 38 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 39 | from vllm.utils.humming import LayerConfig |
| 40 | |
| 41 | logger = init_logger(__name__) |
| 42 | |
| 43 | |
| 44 | @dataclass(kw_only=True) |
| 45 | class HummingMoEQuantConfig(FusedMoEQuantConfig): |
| 46 | w1_humming_config: "LayerConfig" |
| 47 | w2_humming_config: "LayerConfig" |
| 48 | w1_hadamard_block_size: int = 1 |
| 49 | w2_hadamard_block_size: int = 1 |
| 50 | w1_input_scale: torch.Tensor | None = None |
| 51 | w2_input_scale: torch.Tensor | None = None |
| 52 | w1_input_scale_2: torch.Tensor | None = None |
| 53 | w2_input_scale_2: torch.Tensor | None = None |
| 54 | |
| 55 | |
| 56 | def make_humming_moe_quant_config( |
| 57 | quant_dtype: torch.dtype | str | None, |
| 58 | weight_dtype: torch.dtype | str | None, |
| 59 | weight_group_shape: GroupShape | None = None, |
| 60 | activation_group_shape: GroupShape | None = None, |
| 61 | w1_scale: torch.Tensor | None = None, |
| 62 | w2_scale: torch.Tensor | None = None, |
| 63 | w1_zp: torch.Tensor | None = None, |
| 64 | w2_zp: torch.Tensor | None = None, |
| 65 | w1_bias: torch.Tensor | None = None, |
| 66 | w2_bias: torch.Tensor | None = None, |
| 67 | w1_gscale: torch.Tensor | None = None, |
| 68 | w2_gscale: torch.Tensor | None = None, |
| 69 | gemm1_alpha: float | None = None, |
| 70 | gemm1_beta: float | None = None, |
| 71 | gemm1_clamp_limit: float | None = None, |
| 72 | humming_configs: dict[str, "LayerConfig"] | None = None, |
| 73 | ) -> HummingMoEQuantConfig: |
| 74 | assert humming_configs is not None |
| 75 | if quant_dtype is None: |
| 76 | a_quant_desc = FusedMoEQuantDesc(dtype=None) |
| 77 | elif activation_group_shape is not None: |
| 78 | # Pre-dispatch quantization. |
| 79 | a_quant_desc = FusedMoEQuantDesc( |
| 80 | dtype=quant_dtype, shape=activation_group_shape |
| 81 | ) |
| 82 | else: |
| 83 | # Deferred path: Humming quantizes the activation internally, so the |
| 84 | # descriptor only needs a non-None dtype to mark it as quantized. |
| 85 | shape = GroupShape(row=1, col=-1) |
| 86 | a_quant_desc = FusedMoEQuantDesc(dtype=quant_dtype, shape=shape) |
| 87 | |
| 88 | w1_quant_desc = FusedMoEQuantDesc( |
| 89 | dtype=weight_dtype, |
| 90 | shape=weight_group_shape, |
| 91 | scale=w1_scale, |
| 92 | alpha_or_gscale=w1_gscale, |
| 93 | zp=w1_zp, |
| 94 | bias=w1_bias, |
| 95 | ) |
| 96 | |
| 97 | w2_quant_desc = FusedMoEQuantDesc( |
| 98 | dtype=weight_dtype, |
| 99 | shape=weight_group_shape, |
| 100 | scale=w2_scale, |
| 101 | alpha_or_gscale=w2_gscale, |
| 102 | zp=w2_zp, |
| 103 | bias=w2_bias, |
| 104 | ) |
| 105 | |
| 106 | return HummingMoEQuantConfig( |
| 107 | _a1=a_quant_desc, |
| 108 | _a2=a_quant_desc, |
| 109 | _w1=w1_quant_desc, |
| 110 | _w2=w2_quant_desc, |
| 111 | gemm1_alpha=gemm1_alpha, |
| 112 | gemm1_beta=gemm1_beta, |
| 113 | gemm1_clamp_limit=gemm1_clamp_limit, |
| 114 | w1_humming_config=humming_configs["w13"], |
| 115 | w2_humming_config=humming_configs["w2"], |
| 116 | ) |
| 117 | |
| 118 | |
| 119 | def get_humming_moe_quant_config( |
| 120 | layer: "RoutedExperts", |
| 121 | humming_configs: dict[str, "LayerConfig"] | None = None, |
| 122 | gemm1_alpha: float | None = None, |
| 123 | gemm1_beta: float | None = None, |
| 124 | gemm1_clamp_limit: float | None = None, |
| 125 | ): |
| 126 | if humming_configs is None: |
| 127 | humming_configs = layer.humming_configs |
| 128 | input_schema = layer.input_schemas["w13"] |
| 129 | weight_schema = layer.weight_schemas["w13"] |
| 130 | |
| 131 | if input_schema.a_dtype is None or input_schema.a_dtype.num_bits == 16: |
| 132 | q_dtype = None |
| 133 | else: |
| 134 | q_dtype = str(input_schema.a_dtype) |
| 135 | |
| 136 | # Block-FP8 (group-128) activations are quantized *before* the EP all-to-all |
| 137 | # dispatch (so FP8 rather than BF16 crosses the interconnect) and consumed by |
| 138 | # Humming as-is. |
| 139 | activation_group_shape: GroupShape | None = None |
| 140 | input_scale_group_size = getattr(input_schema, "input_scale_group_size", 0) or 0 |
| 141 | if ( |
| 142 | q_dtype is not None |
| 143 | and q_dtype.startswith("float8") |
| 144 | and input_scale_group_size == 128 |
| 145 | and weight_schema.hadamard_block_size <= 1 |
| 146 | and input_schema.input_quant_mode in (None, "dynamic_group") |
| 147 | ): |
| 148 | q_dtype = humming_schema._HUMMING_TO_QUANT_DTYPE.get( |
| 149 | input_schema.a_dtype, FP8_DTYPE |
| 150 | ) |
| 151 | activation_group_shape = GroupShape(row=1, col=input_scale_group_size) |
| 152 | |
| 153 | weight_scale_group_size = weight_schema.weight_scale_group_size |
| 154 | weight_scale_group_size_n = weight_schema.weight_scale_group_size_n |
| 155 | weight_group_shape: tuple[int, ...] = () |
| 156 | if weight_scale_group_size_n > 1: |
| 157 | weight_group_shape = GroupShape( |
| 158 | row=weight_scale_group_size, |
| 159 | col=weight_scale_group_size_n, |
| 160 | ) |
| 161 | elif weight_scale_group_size == 0: |
| 162 | weight_group_shape = GroupShape(row=-1, col=1) |
| 163 | else: |
| 164 | weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) |
| 165 | |
| 166 | config = make_humming_moe_quant_config( |
| 167 | quant_dtype=q_dtype, |
| 168 | weight_dtype=str(weight_schema.b_dtype), |
| 169 | weight_group_shape=weight_group_shape, |
| 170 | activation_group_shape=activation_group_shape, |
| 171 | w1_scale=getattr(layer, "w13_weight_scale", None), |
| 172 | w1_gscale=getattr(layer, "w13_weight_scale_2", None), |
| 173 | w1_zp=getattr(layer, "w13_zero_point", None), |
| 174 | w1_bias=getattr(layer, "w13_bias", None), |
| 175 | w2_scale=getattr(layer, "w2_weight_scale", None), |
| 176 | w2_gscale=getattr(layer, "w2_weight_scale_2", None), |
| 177 | w2_zp=getattr(layer, "w2_zero_point", None), |
| 178 | w2_bias=getattr(layer, "w2_bias", None), |
| 179 | gemm1_alpha=gemm1_alpha, |
| 180 | gemm1_beta=gemm1_beta, |
| 181 | gemm1_clamp_limit=gemm1_clamp_limit, |
| 182 | humming_configs=humming_configs, |
| 183 | ) |
| 184 | for prefix, sublayer in (("w1", "w13"), ("w2", "w2")): |
| 185 | setattr( |
| 186 | config, |
| 187 | f"{prefix}_hadamard_block_size", |
| 188 | layer.weight_schemas[sublayer].hadamard_block_size, |
| 189 | ) |
| 190 | for name in ("input_scale", "input_scale_2"): |
| 191 | key = f"{prefix}_{name}" |
| 192 | value = getattr(layer, f"{sublayer}_{name}", None) |
| 193 | setattr(config, key, value) |
| 194 | return config |
| 195 | |
| 196 | |
| 197 | def select_humming_moe_experts( |
| 198 | config: FusedMoEConfig, |
| 199 | weight_key: QuantKey | None, |
| 200 | activation_key: QuantKey | None, |
| 201 | ) -> type[mk.FusedMoEExperts] | None: |
| 202 | """Select the primary Humming MoE Experts class |
| 203 | Note: Shape-specific fallbacks may still occur at runtime. |
| 204 | """ |
| 205 | if not has_humming(): |
| 206 | return None |
| 207 | |
| 208 | # NOTE: the kernels are selected in the following order. |
| 209 | AVAILABLE_EXPERTS: list[type[mk.FusedMoEExperts]] = [ |
| 210 | BatchedHummingGroupedExperts, |
| 211 | HummingGroupedExperts, |
| 212 | HummingIndexedExperts, |
| 213 | ] |
| 214 | |
| 215 | # NOTE(rob): We need to peak into the P/F selection to determine |
| 216 | # if we are using the batched or standard expert format, which |
| 217 | # if not ideal. Once we unify TP + DP/EP, we can select P/F first. |
| 218 | activation_format = ( |
| 219 | mk.FusedMoEActivationFormat.BatchedExperts |
| 220 | if config.moe_parallel_config.use_batched_activation_format |
| 221 | else mk.FusedMoEActivationFormat.Standard |
| 222 | ) |
| 223 | |
| 224 | def _make_log_backend(experts_cls: type[mk.FusedMoEExperts]): |
| 225 | return f"Using {experts_cls.__name__} Humming MoE backend." |
| 226 | |
| 227 | def _make_log_unsupported( |
| 228 | experts_cls: type[mk.FusedMoEExperts], reason: str | None |
| 229 | ) -> str: |
| 230 | if reason: |
| 231 | return ( |
| 232 | f"Humming MoE experts {experts_cls.__name__} does not support the " |
| 233 | f"deployment configuration since {reason}." |
| 234 | ) |
| 235 | else: |
| 236 | return ( |
| 237 | f"Humming MoE experts '{experts_cls.__name__}' does not support the " |
| 238 | "deployment configuration." |
| 239 | ) |
| 240 | |
| 241 | for k_cls in AVAILABLE_EXPERTS: |
| 242 | supported, reason = k_cls.is_supported_config( |
| 243 | k_cls, |
| 244 | config, |
| 245 | weight_key, |
| 246 | activation_key, |
| 247 | activation_format, |
| 248 | ) |
| 249 | if supported: |
| 250 | logger.info_once(_make_log_backend(k_cls)) |
| 251 | return k_cls |
| 252 | else: |
| 253 | logger.debug_once(_make_log_unsupported(k_cls, reason)) |
| 254 | |
| 255 | return None |
| 256 | |
| 257 | |
| 258 | def make_humming_moe_kernel( |
| 259 | moe_quant_config: FusedMoEQuantConfig, |
| 260 | moe_config: FusedMoEConfig, |
| 261 | experts_cls: type[mk.FusedMoEExperts], |
| 262 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 263 | ) -> mk.FusedMoEKernel: |
| 264 | # Create Prepare/Finalize. |
| 265 | prepare_finalize = maybe_make_prepare_finalize( |
| 266 | moe=moe_config, |
| 267 | quant_config=moe_quant_config, |
| 268 | routing_tables=routing_tables, |
| 269 | allow_new_interface=True, |
| 270 | use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), |
| 271 | ) |
| 272 | assert prepare_finalize is not None |
| 273 | |
| 274 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 275 | |
| 276 | # Create Experts. |
| 277 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 278 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 279 | assert max_num_tokens is not None |
| 280 | experts = experts_cls( |
| 281 | moe_config=moe_config, |
| 282 | quant_config=moe_quant_config, |
| 283 | max_num_tokens=max_num_tokens, |
| 284 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 285 | ) |
| 286 | else: |
| 287 | experts = experts_cls( |
| 288 | moe_config=moe_config, |
| 289 | quant_config=moe_quant_config, |
| 290 | ) |
| 291 | |
| 292 | kernel = mk.FusedMoEKernel( |
| 293 | prepare_finalize, |
| 294 | experts, |
| 295 | ) |
| 296 | |
| 297 | return kernel |
| 298 | |
| 299 | |
| 300 | def _extract_sublayer_tensors( |
| 301 | layer: "RoutedExperts", |
| 302 | sublayer_name: str, |
| 303 | ) -> dict[str, torch.Tensor]: |
| 304 | """Extract tensors for a specific sublayer from the layer's state dict.""" |
| 305 | return dict( |
| 306 | (key.removeprefix(sublayer_name + "_"), value) |
| 307 | for key, value in layer.state_dict().items() |
| 308 | if key.startswith(sublayer_name + "_") |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | def _replace_layer_parameters( |
| 313 | layer: "RoutedExperts", |
| 314 | sublayer_name: str, |
| 315 | tensors: dict[str, torch.Tensor], |
| 316 | preserve_bias: bool = False, |
| 317 | ) -> None: |
| 318 | """Replace layer parameters for a sublayer with new tensors. |
| 319 | |
| 320 | Args: |
| 321 | layer: The RoutedExperts layer |
| 322 | sublayer_name: Name of the sublayer (e.g., "w13", "w2") |
| 323 | tensors: Dict of parameter name to tensor |
| 324 | preserve_bias: If True, don't delete bias parameters |
| 325 | |
| 326 | """ |
| 327 | # Delete old parameters |
| 328 | for name, _ in list(layer.named_parameters()): |
| 329 | if not name.startswith(sublayer_name + "_"): |
| 330 | continue |
| 331 | if preserve_bias and name == sublayer_name + "_bias": |
| 332 | continue |
| 333 | delattr(layer, name) |
| 334 | |
| 335 | # Set new parameters |
| 336 | for name, tensor in tensors.items(): |
| 337 | param_name = f"{sublayer_name}_{name}" |
| 338 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 339 | setattr(layer, param_name, param) |
| 340 | |
| 341 | |
| 342 | def _convert_sublayer_to_humming( |
| 343 | layer: "RoutedExperts", |
| 344 | sublayer_name: str, |
| 345 | shape_n: int, |
| 346 | shape_k: int, |
| 347 | weight_schema: Any, |
| 348 | input_schema: Any, |
| 349 | num_experts: int, |
| 350 | param_dtype: torch.dtype, |
| 351 | ) -> tuple[Any, Any]: |
| 352 | """Convert a sublayer's weights from checkpoint format to Humming format. |
| 353 | |
| 354 | Returns: |
| 355 | Tuple of (converted_weight_schema, converted_input_schema) |
| 356 | |
| 357 | """ |
| 358 | tensors = _extract_sublayer_tensors(layer, sublayer_name) |
| 359 | |
| 360 | shape_k_stacks = [shape_k] |
| 361 | shape_n_stacks = [shape_n] |
| 362 | if sublayer_name == "w13" and layer.moe_config.activation.is_gated: |
| 363 | shape_n_stacks = [shape_n // 2] * 2 |
| 364 | |
| 365 | converted_weight_schema, converted_tensors = weight_schema.convert_humming( |
| 366 | tensors=tensors, |
| 367 | shape_n_stacks=shape_n_stacks, |
| 368 | shape_k_stacks=shape_k_stacks, |
| 369 | param_dtype=param_dtype, |
| 370 | num_experts=num_experts, |
| 371 | ) |
| 372 | |
| 373 | converted_input_schema, input_tensors = input_schema.convert_humming( |
| 374 | tensors=tensors, |
| 375 | shape_n_stacks=shape_n_stacks, |
| 376 | shape_k_stacks=shape_k_stacks, |
| 377 | param_dtype=param_dtype, |
| 378 | num_experts=num_experts, |
| 379 | ) |
| 380 | converted_tensors.update(input_tensors) |
| 381 | |
| 382 | _replace_layer_parameters(layer, sublayer_name, converted_tensors) |
| 383 | |
| 384 | return converted_weight_schema, converted_input_schema |
| 385 | |
| 386 | |
| 387 | def _prepare_and_transform_sublayer( |
| 388 | layer: "RoutedExperts", |
| 389 | sublayer_name: str, |
| 390 | shape_n: int, |
| 391 | shape_k: int, |
| 392 | weight_schema: Any, |
| 393 | input_schema: Any, |
| 394 | has_bias: bool, |
| 395 | num_experts: int, |
| 396 | param_dtype: torch.dtype, |
| 397 | ) -> "LayerConfig": |
| 398 | """Prepare Humming configuration and transform one sublayer's tensors.""" |
| 399 | from vllm.utils.humming import ( |
| 400 | prepare_layer_config, |
| 401 | transform_humming_tensors, |
| 402 | ) |
| 403 | |
| 404 | config = prepare_layer_config( |
| 405 | shape_n=shape_n, |
| 406 | shape_k=shape_k, |
| 407 | pad_n_to_multiple=256, |
| 408 | pad_k_to_multiple=128, |
| 409 | input_schema=input_schema, |
| 410 | weight_schema=weight_schema, |
| 411 | has_bias=has_bias, |
| 412 | num_experts=num_experts, |
| 413 | torch_dtype=param_dtype, |
| 414 | device=getattr(layer, f"{sublayer_name}_weight").device, |
| 415 | ) |
| 416 | source_tensors = _extract_sublayer_tensors(layer, sublayer_name) |
| 417 | for name in ("input_scale", "input_scale_2"): |
| 418 | if name != input_schema.static_tensor_scale_name: |
| 419 | source_tensors.pop(name, None) |
| 420 | tensors = transform_humming_tensors(config, source_tensors) |
| 421 | for name in ("input_scale", "input_scale_2"): |
| 422 | if name in source_tensors: |
| 423 | tensors[name] = source_tensors[name] |
| 424 | _replace_layer_parameters(layer, sublayer_name, tensors) |
| 425 | return config |
| 426 | |
| 427 | |
| 428 | def _process_single_sublayer( |
| 429 | layer: "RoutedExperts", |
| 430 | sublayer_name: str, |
| 431 | shape_n: int, |
| 432 | shape_k: int, |
| 433 | weight_schema: Any, |
| 434 | input_schema: Any, |
| 435 | has_bias: bool, |
| 436 | num_experts: int, |
| 437 | param_dtype: torch.dtype, |
| 438 | force_weight_schema: Any | None = None, |
| 439 | allow_input_schema_fallback: bool = True, |
| 440 | ) -> tuple[Any, Any, "LayerConfig"]: |
| 441 | """Process a single sublayer: convert, optionally requant, prepare, and transform. |
| 442 | |
| 443 | This combines the common logic from convert_to_humming_moe_kernel_format |
| 444 | for processing a single sublayer. |
| 445 | |
| 446 | Args: |
| 447 | layer: The RoutedExperts layer |
| 448 | sublayer_name: Name of the sublayer (e.g., "w13", "w2") |
| 449 | shape_n: Output dimension size |
| 450 | shape_k: Input dimension size |
| 451 | weight_schema: Initial weight quantization schema |
| 452 | input_schema: Initial input quantization schema |
| 453 | has_bias: Whether the layer has bias terms |
| 454 | num_experts: Number of experts |
| 455 | param_dtype: Parameter data type |
| 456 | force_weight_schema: Optional schema to force requantization to |
| 457 | allow_input_schema_fallback: Whether incompatible input schemas may be replaced. |
| 458 | |
| 459 | Returns: |
| 460 | Tuple of the final weight schema, input schema, and Humming layer config. |
| 461 | |
| 462 | """ |
| 463 | from vllm.utils.humming import HummingWeightSchema |
| 464 | |
| 465 | # Step 1: Convert from checkpoint format to humming format if needed |
| 466 | current_weight_schema, current_input_schema = _convert_sublayer_to_humming( |
| 467 | layer=layer, |
| 468 | sublayer_name=sublayer_name, |
| 469 | shape_n=shape_n, |
| 470 | shape_k=shape_k, |
| 471 | weight_schema=weight_schema, |
| 472 | input_schema=input_schema, |
| 473 | num_experts=num_experts, |
| 474 | param_dtype=param_dtype, |
| 475 | ) |
| 476 | |
| 477 | # Step 2: Force requant if needed |
| 478 | assert isinstance(current_weight_schema, HummingWeightSchema) |
| 479 | if force_weight_schema is not None and current_weight_schema != force_weight_schema: |
| 480 | source_tensors = _extract_sublayer_tensors(layer, sublayer_name) |
| 481 | |
| 482 | tensors = current_weight_schema.requant_tensors( |
| 483 | tensors=source_tensors, |
| 484 | target_weight_schema=force_weight_schema, |
| 485 | param_dtype=param_dtype, |
| 486 | ) |
| 487 | name = current_input_schema.static_tensor_scale_name |
| 488 | if name is not None: |
| 489 | tensors[name] = source_tensors[name] |
| 490 | |
| 491 | current_weight_schema = force_weight_schema |
| 492 | _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True) |
| 493 | del tensors |
| 494 | |
| 495 | # Step 3: Prepare layer metadata and transform weights |
| 496 | current_input_schema = humming_schema.check_and_fallback_input_schema( |
| 497 | weight_schema=current_weight_schema, |
| 498 | input_schema=current_input_schema, |
| 499 | param_dtype=param_dtype, |
| 500 | allow_fallback=allow_input_schema_fallback, |
| 501 | ) |
| 502 | config = _prepare_and_transform_sublayer( |
| 503 | layer=layer, |
| 504 | sublayer_name=sublayer_name, |
| 505 | shape_n=shape_n, |
| 506 | shape_k=shape_k, |
| 507 | weight_schema=current_weight_schema, |
| 508 | input_schema=current_input_schema, |
| 509 | has_bias=has_bias, |
| 510 | num_experts=num_experts, |
| 511 | param_dtype=param_dtype, |
| 512 | ) |
| 513 | |
| 514 | return current_weight_schema, current_input_schema, config |
| 515 | |
| 516 | |
| 517 | def convert_to_humming_moe_kernel_format( |
| 518 | layer: "RoutedExperts", |
| 519 | quant_config: dict | None = None, |
| 520 | sublayer_configs: dict[str, Any] | None = None, |
| 521 | weight_schema: Any | None = None, |
| 522 | input_schema: Any | None = None, |
| 523 | force_weight_schema: Any | None = None, |
| 524 | allow_input_schema_fallback: bool = True, |
| 525 | ) -> dict[str, "LayerConfig"]: |
| 526 | """Convert MoE weights from checkpoint format to Humming kernel format. |
| 527 | |
| 528 | This function processes weights for each sublayer (w13, w2) by: |
| 529 | 1. Converting from checkpoint format to humming format if needed |
| 530 | 2. Force requanting if a different quantization schema is specified |
| 531 | 3. Preparing layer metadata for the Humming kernel |
| 532 | 4. Transforming weights for inference |
| 533 | |
| 534 | Args: |
| 535 | layer: The RoutedExperts layer containing weights to process |
| 536 | quant_config: Optional quantization config dict. Required if weight_schema |
| 537 | or input_schema are None. Used to build schemas via |
| 538 | BaseWeightSchema.from_config(). |
| 539 | sublayer_configs: Optional configuration dict for each sublayer (w13, w2). |
| 540 | Each config must have "shape_n" and "shape_k" keys. |
| 541 | If None, configs are built from layer.moe_config properties. |
| 542 | weight_schema: Optional initial weight quantization schema. |
| 543 | If None, built from quant_config. |
| 544 | input_schema: Optional initial input quantization schema. |
| 545 | If None, built from quant_config or env vars. |
| 546 | force_weight_schema: Optional schema to force requantization to |
| 547 | allow_input_schema_fallback: Whether incompatible input schemas may be replaced. |
| 548 | |
| 549 | Side effects: |
| 550 | - Modifies layer parameters in place |
| 551 | - Sets layer.weight_schemas and layer.input_schemas |
| 552 | - Sets layer.humming_configs for quant config construction |
| 553 | |
| 554 | """ |
| 555 | # Build schemas from quant_config if not provided |
| 556 | has_bias = layer.moe_config.has_bias |
| 557 | num_experts = layer.moe_config.num_local_experts |
| 558 | param_dtype = layer.params_dtype |
| 559 | |
| 560 | if weight_schema is None or input_schema is None: |
| 561 | if quant_config is None: |
| 562 | raise ValueError( |
| 563 | "Must provide either weight_schema/input_schema or quant_config" |
| 564 | ) |
| 565 | |
| 566 | from vllm.utils.humming import BaseWeightSchema, HummingInputSchema |
| 567 | |
| 568 | if weight_schema is None: |
| 569 | weight_schema = BaseWeightSchema.from_config(quant_config) |
| 570 | |
| 571 | if input_schema is None: |
| 572 | input_quant_config = (envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}).copy() |
| 573 | if humming_is_layer_skipped(input_quant_config, layer.layer_name): |
| 574 | input_schema = HummingInputSchema() |
| 575 | else: |
| 576 | # TODO: read input_quant_config from quant_config |
| 577 | input_quant_config = humming_schema.resolve_humming_layer_config( |
| 578 | input_quant_config, layer.layer_name |
| 579 | ) |
| 580 | allow_input_schema_fallback = input_quant_config.pop( |
| 581 | "allow_fallback", False |
| 582 | ) |
| 583 | input_schema = HummingInputSchema.from_config(input_quant_config) |
| 584 | |
| 585 | # Build sublayer configs from layer properties if not provided |
| 586 | if sublayer_configs is None: |
| 587 | is_gated = layer.moe_config.activation.is_gated |
| 588 | intermediate_size = layer.moe_config.intermediate_size_per_partition |
| 589 | sublayer_configs = { |
| 590 | "w13": { |
| 591 | "shape_n": intermediate_size * (2 if is_gated else 1), |
| 592 | "shape_k": layer.moe_config.hidden_dim, |
| 593 | }, |
| 594 | "w2": { |
| 595 | "shape_n": layer.moe_config.hidden_dim, |
| 596 | "shape_k": intermediate_size, |
| 597 | }, |
| 598 | } |
| 599 | |
| 600 | layer.weight_schemas = {} |
| 601 | layer.input_schemas = {} |
| 602 | humming_configs = {} |
| 603 | |
| 604 | for sublayer_name, configs in sublayer_configs.items(): |
| 605 | final_weight_schema, final_input_schema, humming_config = ( |
| 606 | _process_single_sublayer( |
| 607 | layer=layer, |
| 608 | sublayer_name=sublayer_name, |
| 609 | shape_n=configs["shape_n"], |
| 610 | shape_k=configs["shape_k"], |
| 611 | weight_schema=weight_schema, |
| 612 | input_schema=input_schema, |
| 613 | has_bias=has_bias, |
| 614 | num_experts=num_experts, |
| 615 | param_dtype=param_dtype, |
| 616 | force_weight_schema=force_weight_schema, |
| 617 | allow_input_schema_fallback=allow_input_schema_fallback, |
| 618 | ) |
| 619 | ) |
| 620 | |
| 621 | layer.weight_schemas[sublayer_name] = final_weight_schema |
| 622 | layer.input_schemas[sublayer_name] = final_input_schema |
| 623 | humming_configs[sublayer_name] = humming_config |
| 624 | |
| 625 | layer.humming_configs = humming_configs |
| 626 | return humming_configs |