| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from math import gcd |
| 5 | from typing import TYPE_CHECKING, Any |
| 6 | |
| 7 | import regex as re |
| 8 | import torch |
| 9 | |
| 10 | import vllm.utils.humming as _hm |
| 11 | from vllm import envs |
| 12 | from vllm.model_executor.layers.fused_moe import ( |
| 13 | FusedMoEConfig, |
| 14 | FusedMoEMethodBase, |
| 15 | FusedMoEQuantConfig, |
| 16 | RoutedExperts, |
| 17 | SharedExperts, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 20 | UnquantizedFusedMoEMethod, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.linear import ( |
| 23 | LinearBase, |
| 24 | LinearMethodBase, |
| 25 | UnquantizedLinearMethod, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 28 | from vllm.model_executor.layers.quantization.base_config import ( |
| 29 | QuantizationConfig, |
| 30 | QuantizeMethodBase, |
| 31 | ) |
| 32 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 33 | check_and_fallback_input_schema, |
| 34 | convert_to_humming_moe_kernel_format, |
| 35 | get_humming_linear_compute_config, |
| 36 | get_humming_moe_quant_config, |
| 37 | humming_update_schema_hadamard_block_size, |
| 38 | input_schema_to_quant_key, |
| 39 | make_humming_moe_kernel, |
| 40 | resolve_humming_layer_config, |
| 41 | select_humming_moe_experts, |
| 42 | weight_schema_to_quant_key, |
| 43 | ) |
| 44 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 45 | from vllm.model_executor.parameter import ( |
| 46 | BasevLLMParameter, |
| 47 | BlockQuantScaleParameter, |
| 48 | ChannelQuantScaleParameter, |
| 49 | GroupQuantScaleParameter, |
| 50 | ModelWeightParameter, |
| 51 | PackedvLLMParameter, |
| 52 | PerTensorScaleParameter, |
| 53 | RowvLLMParameter, |
| 54 | ) |
| 55 | from vllm.model_executor.utils import set_weight_attrs |
| 56 | |
| 57 | if TYPE_CHECKING: |
| 58 | from vllm.model_executor.models.utils import WeightsMapper |
| 59 | from vllm.utils.humming import ( |
| 60 | BaseInputSchema, |
| 61 | BaseWeightSchema, |
| 62 | HummingInputSchema, |
| 63 | HummingWeightSchema, |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def prepare_param(tensor, name, extra_attrs): |
| 68 | extra_attrs = extra_attrs.copy() |
| 69 | scale_type = extra_attrs.pop("scale_type", None) |
| 70 | param_cls_name_map = { |
| 71 | "block": BlockQuantScaleParameter, |
| 72 | "tensor": PerTensorScaleParameter, |
| 73 | "group": GroupQuantScaleParameter, |
| 74 | "channel": ChannelQuantScaleParameter, |
| 75 | "input_scale": PerTensorScaleParameter, |
| 76 | } |
| 77 | |
| 78 | param_cls: type[BasevLLMParameter] |
| 79 | if "packed_dim" in extra_attrs: |
| 80 | param_cls = PackedvLLMParameter |
| 81 | elif scale_type in param_cls_name_map: |
| 82 | param_cls = param_cls_name_map[scale_type] |
| 83 | elif "output_dim" in extra_attrs and "input_dim" in extra_attrs: |
| 84 | param_cls = ModelWeightParameter |
| 85 | elif "input_dim" in extra_attrs: |
| 86 | param_cls = RowvLLMParameter |
| 87 | elif "output_dim" in extra_attrs: |
| 88 | param_cls = ChannelQuantScaleParameter |
| 89 | else: |
| 90 | param_cls = BasevLLMParameter |
| 91 | |
| 92 | kwargs_keys = [ |
| 93 | "input_dim", |
| 94 | "output_dim", |
| 95 | "packed_dim", |
| 96 | "packed_factor", |
| 97 | "weight_loader", |
| 98 | ] |
| 99 | cls_kwargs = {} |
| 100 | for key in extra_attrs.copy(): |
| 101 | if key in kwargs_keys: |
| 102 | cls_kwargs[key] = extra_attrs.pop(key) |
| 103 | |
| 104 | param = param_cls(data=tensor, **cls_kwargs) |
| 105 | set_weight_attrs(param, extra_attrs) |
| 106 | |
| 107 | param.param_name = name |
| 108 | param.ignore_warning = True |
| 109 | if scale_type in ["tensor", "input_scale"]: |
| 110 | param.needs_scalar_to_array = True |
| 111 | |
| 112 | return param |
| 113 | |
| 114 | |
| 115 | def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, Any]): |
| 116 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 117 | if "scale_type" in extra_attrs: |
| 118 | extra_attrs["quant_method"] = extra_attrs["scale_type"] |
| 119 | |
| 120 | if "input_dim" in extra_attrs and "output_dim" in extra_attrs: |
| 121 | input_dim = extra_attrs["input_dim"] |
| 122 | output_dim = extra_attrs["output_dim"] |
| 123 | extra_attrs["is_transposed"] = input_dim < output_dim |
| 124 | |
| 125 | set_weight_attrs(param, extra_attrs) |
| 126 | param.param_name = name |
| 127 | return param |
| 128 | |
| 129 | |
| 130 | def compressed_tensors_get_config(config: dict[str, Any], key: str): |
| 131 | assert key in ["weights", "input_activations"] |
| 132 | target_group_config = None |
| 133 | for group_config in config["config_groups"].values(): |
| 134 | if "Linear" in group_config["targets"]: |
| 135 | if "weights" not in group_config: |
| 136 | return None |
| 137 | if key not in group_config or group_config[key] is None: |
| 138 | return None |
| 139 | target_group_config = group_config[key].copy() |
| 140 | break |
| 141 | |
| 142 | if target_group_config is None: |
| 143 | return None |
| 144 | target_group_config["quant_method"] = config["quant_method"] |
| 145 | if config["quant_method"] == "compressed-tensors": |
| 146 | target_group_config["format"] = config["format"] |
| 147 | elif config["quant_method"] == "modelopt": |
| 148 | target_group_config["quant_algo"] = config["quant_algo"] |
| 149 | return target_group_config |
| 150 | |
| 151 | |
| 152 | class HummingConfig(QuantizationConfig): |
| 153 | packed_modules_mapping: dict[str, list[str]] = {} |
| 154 | |
| 155 | def __init__(self, full_config: dict[str, Any] | None = None): |
| 156 | self.full_config: dict[str, Any] = full_config or {} |
| 157 | |
| 158 | @classmethod |
| 159 | def get_name(cls) -> QuantizationMethods: |
| 160 | return "humming" |
| 161 | |
| 162 | @classmethod |
| 163 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 164 | return [torch.bfloat16, torch.half] |
| 165 | |
| 166 | @classmethod |
| 167 | def get_min_capability(cls) -> int: |
| 168 | return 75 |
| 169 | |
| 170 | @classmethod |
| 171 | def get_config_filenames(cls) -> list[str]: |
| 172 | return [] |
| 173 | |
| 174 | @classmethod |
| 175 | def from_config(cls, config: dict[str, Any]) -> "HummingConfig": |
| 176 | return cls(full_config=config) |
| 177 | |
| 178 | @classmethod |
| 179 | def override_quantization_method( |
| 180 | cls, hf_quant_cfg, user_quant, hf_config=None |
| 181 | ) -> QuantizationMethods | None: |
| 182 | if user_quant == "humming" and hf_config is not None: |
| 183 | model_type = hf_config.model_type |
| 184 | quant_method = hf_quant_cfg.get("quant_method", None) |
| 185 | if model_type == "gpt_oss" and quant_method == "mxfp4": |
| 186 | msg = ( |
| 187 | "For gpt-oss model, use '--moe-backend humming' " |
| 188 | "instead of '--quantization humming'." |
| 189 | ) |
| 190 | raise ValueError(msg) |
| 191 | return "humming" if user_quant == "humming" else None |
| 192 | |
| 193 | def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): |
| 194 | self.hf_to_vllm_mapper = hf_to_vllm_mapper |
| 195 | |
| 196 | def is_layer_skipped(self, config: dict[str, Any], prefix: str): |
| 197 | keys = ["ignored_layers", "ignore", "modules_to_not_convert"] |
| 198 | ignored_layers = self.get_from_keys_or(config, keys, []) or [] |
| 199 | if hasattr(self, "hf_to_vllm_mapper"): |
| 200 | ignored_layers = self.hf_to_vllm_mapper.apply_list(ignored_layers) |
| 201 | |
| 202 | for module_name in ignored_layers: |
| 203 | # compressed-tensors style entries may be regex patterns prefixed |
| 204 | # with "re:" (e.g. "re:vision_tower.*"). These must be regex-matched; |
| 205 | # plain substring matching never matches the literal "re:..." string, |
| 206 | # so ignored layers get silently quantized. Non-"re:" entries keep the |
| 207 | # existing substring behavior (e.g. bitsandbytes modules_to_not_convert). |
| 208 | if module_name.startswith("re:"): |
| 209 | if re.match(module_name[3:], prefix): |
| 210 | return True |
| 211 | elif module_name in prefix: |
| 212 | return True |
| 213 | if "lm_head" in prefix: |
| 214 | return True |
| 215 | |
| 216 | for regex in config.get("dynamic", {}): |
| 217 | if regex[:1] != "-": |
| 218 | continue |
| 219 | if re.match(regex[2:], prefix): |
| 220 | return True |
| 221 | |
| 222 | return False |
| 223 | |
| 224 | def get_layer_weight_schema(self, config: dict[str, Any], prefix: str): |
| 225 | if self.is_layer_skipped(config, prefix): |
| 226 | return None |
| 227 | |
| 228 | if config["quant_method"] in ["compressed-tensors", "modelopt"]: |
| 229 | group_config = compressed_tensors_get_config(config, "weights") |
| 230 | if group_config is None: |
| 231 | return None |
| 232 | config = group_config |
| 233 | |
| 234 | layer_config = resolve_humming_layer_config(config, prefix) |
| 235 | |
| 236 | if "quant_method" in layer_config: |
| 237 | return _hm.BaseWeightSchema.from_config(layer_config) |
| 238 | return None |
| 239 | |
| 240 | def get_layer_input_config(self, config: dict[str, Any], prefix: str): |
| 241 | if self.is_layer_skipped(config, prefix): |
| 242 | return None |
| 243 | if config["quant_method"] in ["compressed-tensors", "modelopt"]: |
| 244 | group_config = compressed_tensors_get_config(config, "input_activations") |
| 245 | if group_config is None: |
| 246 | return None |
| 247 | config = group_config |
| 248 | |
| 249 | config = resolve_humming_layer_config(config, prefix) |
| 250 | if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP: |
| 251 | return config |
| 252 | return None |
| 253 | |
| 254 | def get_layer_input_schema(self, config: dict[str, Any], prefix: str): |
| 255 | config = self.get_layer_input_config(config, prefix) |
| 256 | return None if config is None else _hm.BaseInputSchema.from_config(config) |
| 257 | |
| 258 | def get_quant_config_for_layer( |
| 259 | self, prefix: str, layer_type: str |
| 260 | ) -> "HummingLayerQuantizationConfig | None": |
| 261 | weight_schema: BaseWeightSchema | None = None |
| 262 | force_weight_schema: HummingWeightSchema | None = None |
| 263 | |
| 264 | if self.full_config: |
| 265 | weight_schema = self.get_layer_weight_schema(self.full_config, prefix) |
| 266 | |
| 267 | is_online_quant = False |
| 268 | online_quant_config = envs.VLLM_HUMMING_ONLINE_QUANT_CONFIG or {} |
| 269 | if not self.full_config or online_quant_config.get("force_requant", False): |
| 270 | online_quant_config["quant_method"] = "humming" |
| 271 | schema = self.get_layer_weight_schema(online_quant_config, prefix) |
| 272 | if not self.full_config: |
| 273 | weight_schema = schema |
| 274 | is_online_quant = True |
| 275 | else: |
| 276 | force_weight_schema = schema |
| 277 | |
| 278 | if weight_schema is not None: |
| 279 | input_schema = None |
| 280 | force_input_schema = None |
| 281 | allow_input_schema_fallback = True |
| 282 | |
| 283 | if self.full_config: |
| 284 | input_config = self.get_layer_input_config(self.full_config, prefix) |
| 285 | if input_config is not None: |
| 286 | allow_input_schema_fallback = input_config.pop( |
| 287 | "allow_fallback", True |
| 288 | ) |
| 289 | input_schema = _hm.BaseInputSchema.from_config(input_config) |
| 290 | |
| 291 | if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG: |
| 292 | quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy() |
| 293 | quant_config["quant_method"] = "humming" |
| 294 | input_config = self.get_layer_input_config(quant_config, prefix) |
| 295 | if input_config is not None: |
| 296 | allow_input_schema_fallback = input_config.pop( |
| 297 | "allow_fallback", False |
| 298 | ) |
| 299 | force_input_schema = _hm.BaseInputSchema.from_config(input_config) |
| 300 | if input_schema is None: |
| 301 | input_schema = force_input_schema |
| 302 | |
| 303 | if force_weight_schema is not None and force_input_schema is None: |
| 304 | force_input_schema = _hm.HummingInputSchema() |
| 305 | |
| 306 | return HummingLayerQuantizationConfig( |
| 307 | weight_schema=weight_schema, |
| 308 | input_schema=input_schema, |
| 309 | force_weight_schema=force_weight_schema, |
| 310 | force_input_schema=force_input_schema, |
| 311 | is_online_quant=is_online_quant, |
| 312 | allow_input_schema_fallback=allow_input_schema_fallback, |
| 313 | ) |
| 314 | return None |
| 315 | |
| 316 | def get_quant_method( |
| 317 | self, layer: torch.nn.Module, prefix: str |
| 318 | ) -> "QuantizeMethodBase | None": |
| 319 | layer_type = "other" |
| 320 | if isinstance(layer, RoutedExperts): |
| 321 | layer_type = "moe" |
| 322 | elif isinstance(layer, LinearBase): |
| 323 | layer_type = "linear" |
| 324 | |
| 325 | quant_config = self.get_quant_config_for_layer(prefix, layer_type) |
| 326 | if quant_config is None: |
| 327 | if isinstance(layer, RoutedExperts): |
| 328 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 329 | elif isinstance(layer, LinearBase): |
| 330 | return UnquantizedLinearMethod() |
| 331 | elif isinstance(layer, LinearBase): |
| 332 | input_size = layer.input_size |
| 333 | partition_size = getattr(layer, "input_size_per_partition", input_size) |
| 334 | is_online_quant = quant_config.is_online_quant |
| 335 | if is_online_quant and (input_size % 32 != 0 or partition_size % 32 != 0): |
| 336 | return UnquantizedLinearMethod() |
| 337 | return HummingLinearMethod(quant_config) |
| 338 | elif isinstance(layer, RoutedExperts): |
| 339 | return HummingMoEMethod(quant_config, layer.moe_config) |
| 340 | return None |
| 341 | |
| 342 | |
| 343 | class HummingLayerQuantizationConfig(HummingConfig): |
| 344 | def __init__( |
| 345 | self, |
| 346 | weight_schema: "BaseWeightSchema", |
| 347 | input_schema: "BaseInputSchema | None" = None, |
| 348 | force_weight_schema: "HummingWeightSchema | None" = None, |
| 349 | force_input_schema: "HummingInputSchema | None" = None, |
| 350 | is_online_quant: bool = False, |
| 351 | allow_input_schema_fallback: bool = True, |
| 352 | ): |
| 353 | self.weight_schema = weight_schema |
| 354 | if input_schema is None: |
| 355 | input_schema = _hm.HummingInputSchema() |
| 356 | self.input_schema = input_schema |
| 357 | self.force_weight_schema = force_weight_schema |
| 358 | self.force_input_schema = force_input_schema |
| 359 | self.is_online_quant = is_online_quant |
| 360 | self.allow_input_schema_fallback = allow_input_schema_fallback |
| 361 | |
| 362 | @classmethod |
| 363 | def from_config(cls, config): |
| 364 | weight_schema = _hm.BaseWeightSchema.from_config(config) |
| 365 | return cls(weight_schema) |
| 366 | |
| 367 | def get_quant_method( |
| 368 | self, layer: torch.nn.Module, prefix: str |
| 369 | ) -> QuantizeMethodBase | None: |
| 370 | raise NotImplementedError |
| 371 | |
| 372 | |
| 373 | class HummingLinearMethod(LinearMethodBase): |
| 374 | def __init__(self, quant_config: HummingLayerQuantizationConfig): |
| 375 | self.quant_config = quant_config |
| 376 | self.weight_schema = quant_config.weight_schema |
| 377 | self.input_schema = quant_config.input_schema |
| 378 | self.force_weight_schema = quant_config.force_weight_schema |
| 379 | self.force_input_schema = quant_config.force_input_schema |
| 380 | self.is_online_quant = self.quant_config.is_online_quant |
| 381 | |
| 382 | def prepare_weight_loader(self, layer: torch.nn.Module, weight_loader: Callable): |
| 383 | def new_weight_loader( |
| 384 | param: torch.nn.Parameter, |
| 385 | loaded_weight: torch.Tensor, |
| 386 | shard_id: str | int | None = None, |
| 387 | ): |
| 388 | name = param.param_name |
| 389 | float_dtypes = [torch.float16, torch.bfloat16, torch.float32] |
| 390 | is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes |
| 391 | if is_unquantized and self.is_online_quant: |
| 392 | # online quant (fp16/bf16 -> quant_type) |
| 393 | assert isinstance(self.weight_schema, _hm.HummingWeightSchema) |
| 394 | tensors = self.weight_schema.quant_tensor( |
| 395 | loaded_weight.to(param.device), |
| 396 | self.weight_schema, |
| 397 | layer.param_dtype, |
| 398 | ) |
| 399 | |
| 400 | for key, tensor in tensors.items(): |
| 401 | if tensor is None or tensor.nelement() == 0: |
| 402 | continue |
| 403 | param = getattr(layer, key) |
| 404 | param.weight_loader(param, tensor, shard_id) |
| 405 | |
| 406 | return None |
| 407 | elif is_unquantized and not self.is_online_quant: |
| 408 | # fallback to unquantized linear |
| 409 | # some model skip some layer when quantizing model, but |
| 410 | # don't mark the layer as unquantized. |
| 411 | if not layer.is_fallback: |
| 412 | layer.is_fallback = True |
| 413 | for name, _ in list(layer.named_parameters()): |
| 414 | if name != "bias": |
| 415 | delattr(layer, name) |
| 416 | self.__class__ = UnquantizedLinearMethod # type: ignore |
| 417 | tensor = torch.empty( |
| 418 | ( |
| 419 | layer.output_partition_sizes_sum, |
| 420 | layer.input_size_per_partition, |
| 421 | ), |
| 422 | dtype=layer.param_dtype, |
| 423 | device=param.device, |
| 424 | ) |
| 425 | extra_weight_attrs = layer.extra_weight_attrs.copy() |
| 426 | orig_weight_loader = extra_weight_attrs.pop("weight_loader") |
| 427 | layer.weight = ModelWeightParameter( |
| 428 | data=tensor, |
| 429 | input_dim=1, |
| 430 | output_dim=0, |
| 431 | weight_loader=orig_weight_loader, |
| 432 | ) |
| 433 | layer.weight.tp_size = layer.tp_size |
| 434 | layer.weight.tp_rank = layer.tp_rank |
| 435 | set_weight_attrs(layer.weight, extra_weight_attrs) |
| 436 | |
| 437 | param = layer.weight |
| 438 | if shard_id is not None: |
| 439 | return layer.weight.weight_loader(param, loaded_weight, shard_id) |
| 440 | return layer.weight.weight_loader(param, loaded_weight) |
| 441 | |
| 442 | # weight processing logic for specific quantization schema |
| 443 | loaded_weight = self.weight_schema.process_loaded_weight( |
| 444 | tensor=loaded_weight, |
| 445 | name=name, |
| 446 | ) |
| 447 | if shard_id is not None: |
| 448 | return weight_loader(param, loaded_weight, shard_id) |
| 449 | return weight_loader(param, loaded_weight) |
| 450 | |
| 451 | return new_weight_loader |
| 452 | |
| 453 | def create_weights( |
| 454 | self, |
| 455 | layer: torch.nn.Module, |
| 456 | input_size_per_partition: int, |
| 457 | output_partition_sizes: list[int], |
| 458 | input_size: int, |
| 459 | output_size: int, |
| 460 | params_dtype: torch.dtype, |
| 461 | **extra_weight_attrs, |
| 462 | ): |
| 463 | layer.is_fallback = False |
| 464 | layer.param_dtype = params_dtype |
| 465 | layer.input_size = input_size |
| 466 | layer.output_size = output_size |
| 467 | layer.input_size_per_partition = input_size_per_partition |
| 468 | layer.output_partition_sizes_sum = sum(output_partition_sizes) |
| 469 | layer.output_partition_sizes = output_partition_sizes |
| 470 | layer.extra_weight_attrs = extra_weight_attrs.copy() |
| 471 | |
| 472 | input_schema = self.force_input_schema or self.input_schema |
| 473 | input_schema = input_schema.to_humming_schema(params_dtype) |
| 474 | for name in ("weight_schema", "force_weight_schema"): |
| 475 | schema = getattr(self, name) |
| 476 | if getattr(schema, "hadamard_block_size", None) == -1: |
| 477 | schema = humming_update_schema_hadamard_block_size( |
| 478 | weight_schema=schema, |
| 479 | input_schema=input_schema, |
| 480 | shape_k=input_size_per_partition, |
| 481 | ) |
| 482 | setattr(self, name, schema) |
| 483 | |
| 484 | weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) |
| 485 | new_weight_loader = self.prepare_weight_loader(layer, weight_loader) |
| 486 | extra_weight_attrs["weight_loader"] = new_weight_loader |
| 487 | |
| 488 | for key in ["weight_block_size", "block_structure"]: |
| 489 | block_size = getattr(self.weight_schema, key, None) |
| 490 | if block_size is not None: |
| 491 | layer.weight_block_size = block_size |
| 492 | |
| 493 | weight_tensor_attrs = self.weight_schema.get_tensors_attrs( |
| 494 | shape_n=layer.output_partition_sizes_sum, |
| 495 | shape_k=layer.input_size_per_partition, |
| 496 | param_dtype=params_dtype, |
| 497 | stack_size=len(layer.output_partition_sizes), |
| 498 | ) |
| 499 | |
| 500 | input_tensor_attrs = self.input_schema.get_tensors_attrs( |
| 501 | shape_k=layer.input_size_per_partition, |
| 502 | param_dtype=params_dtype, |
| 503 | stack_size=len(layer.output_partition_sizes), |
| 504 | ) |
| 505 | |
| 506 | tensors_attrs = weight_tensor_attrs | input_tensor_attrs |
| 507 | |
| 508 | for name, attrs in tensors_attrs.items(): |
| 509 | tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"]) |
| 510 | extra_attrs = attrs.get("extra_attrs", {}).copy() |
| 511 | extra_attrs.update(extra_weight_attrs) |
| 512 | param = prepare_param(tensor, name, extra_attrs) |
| 513 | setattr(layer, name, param) |
| 514 | |
| 515 | if self.force_input_schema is not None: |
| 516 | self.input_schema = self.force_input_schema |
| 517 | |
| 518 | if not hasattr(layer, "weight"): |
| 519 | param = prepare_param(torch.tensor(0), "weight", extra_weight_attrs) |
| 520 | layer.weight = param |
| 521 | |
| 522 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 523 | if layer.is_fallback: |
| 524 | return None |
| 525 | |
| 526 | # convert from checkpoint format to humming format |
| 527 | is_humming_weight = isinstance(self.weight_schema, _hm.HummingWeightSchema) |
| 528 | is_humming_input = isinstance(self.input_schema, _hm.HummingInputSchema) |
| 529 | |
| 530 | if not is_humming_weight or not is_humming_input: |
| 531 | self.weight_schema, tensors = self.weight_schema.convert_humming( |
| 532 | tensors=layer.state_dict(), |
| 533 | shape_n_stacks=layer.output_partition_sizes, |
| 534 | shape_k_stacks=[layer.input_size_per_partition], |
| 535 | param_dtype=layer.param_dtype, |
| 536 | ) |
| 537 | |
| 538 | self.input_schema, input_tensors = self.input_schema.convert_humming( |
| 539 | tensors=layer.state_dict(), |
| 540 | shape_n_stacks=layer.output_partition_sizes, |
| 541 | shape_k_stacks=[layer.input_size_per_partition], |
| 542 | param_dtype=layer.param_dtype, |
| 543 | ) |
| 544 | tensors.update(input_tensors) |
| 545 | |
| 546 | for name, _ in list(layer.named_parameters()): |
| 547 | delattr(layer, name) |
| 548 | |
| 549 | for name, tensor in tensors.items(): |
| 550 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 551 | setattr(layer, name, param) |
| 552 | |
| 553 | del tensors |
| 554 | |
| 555 | # force requant (origin quant setting -> fp16/bf16 -> new_quant setting) |
| 556 | assert isinstance(self.weight_schema, _hm.HummingWeightSchema) |
| 557 | force_requant = self.force_weight_schema is not None |
| 558 | if force_requant and self.weight_schema != self.force_weight_schema: |
| 559 | source_tensors = layer.state_dict() |
| 560 | tensors = self.weight_schema.requant_tensors( |
| 561 | tensors=source_tensors, |
| 562 | target_weight_schema=self.force_weight_schema, |
| 563 | param_dtype=layer.param_dtype, |
| 564 | ) |
| 565 | name = self.input_schema.static_tensor_scale_name |
| 566 | if name is not None: |
| 567 | tensors[name] = source_tensors[name] |
| 568 | |
| 569 | self.weight_schema = self.force_weight_schema |
| 570 | |
| 571 | for name, _ in list(layer.named_parameters()): |
| 572 | if name != "bias": |
| 573 | delattr(layer, name) |
| 574 | |
| 575 | for name, tensor in tensors.items(): |
| 576 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 577 | setattr(layer, name, param) |
| 578 | |
| 579 | del tensors |
| 580 | |
| 581 | self.input_schema = check_and_fallback_input_schema( |
| 582 | weight_schema=self.weight_schema, |
| 583 | input_schema=self.input_schema, |
| 584 | param_dtype=layer.param_dtype, |
| 585 | allow_fallback=self.quant_config.allow_input_schema_fallback, |
| 586 | ) |
| 587 | self.layer_config = _hm.prepare_layer_config( |
| 588 | shape_n=layer.output_partition_sizes_sum, |
| 589 | shape_k=layer.input_size_per_partition, |
| 590 | weight_schema=self.weight_schema, |
| 591 | input_schema=self.input_schema, |
| 592 | pad_n_to_multiple=256, |
| 593 | pad_k_to_multiple=128, |
| 594 | has_bias=layer.has_bias, |
| 595 | torch_dtype=layer.param_dtype, |
| 596 | device=layer.weight.device, |
| 597 | ) |
| 598 | source_tensors = dict(layer.named_parameters()) |
| 599 | tensors = _hm.transform_humming_tensors(self.layer_config, source_tensors) |
| 600 | name = self.input_schema.static_tensor_scale_name |
| 601 | if name is not None: |
| 602 | tensors[name] = source_tensors[name] |
| 603 | if "bias" in tensors: |
| 604 | zero_bias = torch.zeros_like(tensors["bias"]) |
| 605 | layer.register_buffer("zero_bias", zero_bias) |
| 606 | for name, _ in list(layer.named_parameters()): |
| 607 | delattr(layer, name) |
| 608 | for name, tensor in tensors.items(): |
| 609 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 610 | setattr(layer, name, param) |
| 611 | |
| 612 | self.compute_config = get_humming_linear_compute_config() |
| 613 | self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) |
| 614 | |
| 615 | def apply( |
| 616 | self, |
| 617 | layer: torch.nn.Module, |
| 618 | x: torch.Tensor, |
| 619 | bias: torch.Tensor | None = None, |
| 620 | ) -> torch.Tensor: |
| 621 | flatten_inputs = x.reshape(-1, x.size(-1)) |
| 622 | output = _hm.humming_forward( |
| 623 | self.layer_config, |
| 624 | inputs=flatten_inputs, |
| 625 | weight=layer.weight, |
| 626 | weight_scale=getattr(layer, "weight_scale", None), |
| 627 | zero_point=getattr(layer, "zero_point", None), |
| 628 | bias=getattr(layer, "bias" if bias is not None else "zero_bias", None), |
| 629 | weight_scale_2=getattr(layer, "weight_scale_2", None), |
| 630 | input_scale=getattr(layer, "input_scale", None), |
| 631 | input_scale_2=getattr(layer, "input_scale_2", None), |
| 632 | hadamard_block_size=self.weight_schema.hadamard_block_size, |
| 633 | locks=self.locks, |
| 634 | compute_config=self.compute_config, |
| 635 | ) |
| 636 | output = output.view(*x.shape[:-1], output.size(-1)) |
| 637 | return output |
| 638 | |
| 639 | |
| 640 | class HummingMoEMethod(FusedMoEMethodBase): |
| 641 | def __init__( |
| 642 | self, quant_config: HummingLayerQuantizationConfig, moe: "FusedMoEConfig" |
| 643 | ) -> None: |
| 644 | super().__init__(moe) |
| 645 | self.quant_config = quant_config |
| 646 | self.weight_schema = quant_config.weight_schema |
| 647 | self.input_schema = quant_config.input_schema |
| 648 | self.force_weight_schema = quant_config.force_weight_schema |
| 649 | self.force_input_schema = quant_config.force_input_schema |
| 650 | |
| 651 | # Derive QuantKeys from humming schemas. |
| 652 | # Prefer force schemas (the final format after requant) over base. |
| 653 | weight_key = weight_schema_to_quant_key( |
| 654 | self.force_weight_schema or self.weight_schema, moe.in_dtype |
| 655 | ) |
| 656 | runtime_input_schema = check_and_fallback_input_schema( |
| 657 | weight_schema=self.force_weight_schema or self.weight_schema, |
| 658 | input_schema=self.force_input_schema or self.input_schema, |
| 659 | param_dtype=moe.in_dtype, |
| 660 | allow_fallback=quant_config.allow_input_schema_fallback, |
| 661 | ) |
| 662 | activation_key = input_schema_to_quant_key(runtime_input_schema, moe.in_dtype) |
| 663 | |
| 664 | # Select Humming MoE experts |
| 665 | self.experts_cls = select_humming_moe_experts( |
| 666 | config=self.moe, |
| 667 | weight_key=weight_key, |
| 668 | activation_key=activation_key, |
| 669 | ) |
| 670 | |
| 671 | def prepare_weight_loader(self, layer, weight_loader): |
| 672 | def new_weight_loader( |
| 673 | param: torch.nn.Parameter, |
| 674 | loaded_weight: torch.Tensor, |
| 675 | weight_name: str, |
| 676 | shard_id: str, |
| 677 | expert_id: int | None = None, |
| 678 | return_success: bool = False, |
| 679 | ): |
| 680 | name = param.param_name |
| 681 | float_dtypes = [torch.float16, torch.bfloat16, torch.float32] |
| 682 | is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes |
| 683 | # online quant (fp16/bf16 -> quant_type) |
| 684 | if is_unquantized: |
| 685 | assert isinstance(self.weight_schema, _hm.HummingWeightSchema) |
| 686 | tensors = self.weight_schema.quant_tensor( |
| 687 | loaded_weight.to(param.device), |
| 688 | self.weight_schema, |
| 689 | layer.param_dtype, |
| 690 | ) |
| 691 | |
| 692 | success = True |
| 693 | for key, tensor in tensors.items(): |
| 694 | if tensor is None or tensor.nelement() == 0: |
| 695 | continue |
| 696 | sublayer_name = "w2" if shard_id == "w2" else "w13" |
| 697 | |
| 698 | param = getattr(layer, sublayer_name + "_" + key) |
| 699 | part_success = param.weight_loader( |
| 700 | param=param, |
| 701 | loaded_weight=tensor.cpu(), |
| 702 | weight_name=shard_id + "_" + key, |
| 703 | shard_id=shard_id, |
| 704 | expert_id=expert_id, |
| 705 | return_success=return_success, |
| 706 | ) |
| 707 | success = success and part_success |
| 708 | |
| 709 | return success if return_success else None |
| 710 | |
| 711 | # weight processing logic for specific quantization schema |
| 712 | loaded_weight = self.weight_schema.process_loaded_weight( |
| 713 | tensor=loaded_weight, |
| 714 | name=name, |
| 715 | ) |
| 716 | return weight_loader( |
| 717 | param, |
| 718 | loaded_weight, |
| 719 | weight_name, |
| 720 | shard_id=shard_id, |
| 721 | expert_id=expert_id, |
| 722 | return_success=return_success, |
| 723 | ) |
| 724 | |
| 725 | return new_weight_loader |
| 726 | |
| 727 | def create_weights( |
| 728 | self, |
| 729 | layer: RoutedExperts, |
| 730 | num_experts: int, |
| 731 | hidden_size: int, |
| 732 | intermediate_size_per_partition: int, |
| 733 | params_dtype: torch.dtype, |
| 734 | **extra_weight_attrs, |
| 735 | ): |
| 736 | layer.num_experts = num_experts |
| 737 | layer.param_dtype = params_dtype |
| 738 | layer.intermediate_size = intermediate_size_per_partition |
| 739 | input_schema = self.force_input_schema or self.input_schema |
| 740 | input_schema = input_schema.to_humming_schema(params_dtype) |
| 741 | shape_k = gcd(hidden_size, intermediate_size_per_partition) |
| 742 | for name in ("weight_schema", "force_weight_schema"): |
| 743 | schema = getattr(self, name) |
| 744 | if getattr(schema, "hadamard_block_size", None) == -1: |
| 745 | schema = humming_update_schema_hadamard_block_size( |
| 746 | weight_schema=schema, |
| 747 | input_schema=input_schema, |
| 748 | shape_k=shape_k, |
| 749 | ) |
| 750 | setattr(self, name, schema) |
| 751 | weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) |
| 752 | weight_loader = self.prepare_weight_loader(layer, weight_loader) |
| 753 | extra_weight_attrs["weight_loader"] = weight_loader |
| 754 | |
| 755 | # sublayer: a layer contains multiple sets of weights for quantized GEMM |
| 756 | # (e.g., weight, weight_scale, etc.). |
| 757 | # The weight names of sublayer start with the prefix "{sublayer_name}_" |
| 758 | w13_output_size = intermediate_size_per_partition * self.moe.w13_num_shards |
| 759 | layer.sublayer_configs = { |
| 760 | "w13": { |
| 761 | "shape_n": w13_output_size, |
| 762 | "shape_k": hidden_size, |
| 763 | "tensors_attrs": self.weight_schema.get_padded_tensors_attrs( |
| 764 | shape_n=w13_output_size, |
| 765 | shape_k=hidden_size, |
| 766 | num_experts=num_experts, |
| 767 | param_dtype=params_dtype, |
| 768 | has_bias=self.moe.has_bias, |
| 769 | ), |
| 770 | }, |
| 771 | "w2": { |
| 772 | "shape_n": hidden_size, |
| 773 | "shape_k": intermediate_size_per_partition, |
| 774 | "tensors_attrs": self.weight_schema.get_padded_tensors_attrs( |
| 775 | shape_n=hidden_size, |
| 776 | shape_k=intermediate_size_per_partition, |
| 777 | num_experts=num_experts, |
| 778 | param_dtype=params_dtype, |
| 779 | has_bias=self.moe.has_bias, |
| 780 | ), |
| 781 | }, |
| 782 | } |
| 783 | |
| 784 | for sublayer_name, configs in layer.sublayer_configs.items(): |
| 785 | for name, attrs in configs["tensors_attrs"].items(): |
| 786 | tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"]) |
| 787 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 788 | extra_attrs = attrs.get("extra_attrs", {}).copy() |
| 789 | extra_attrs.update(extra_weight_attrs) |
| 790 | param = prepare_moe_param(tensor, name, extra_attrs) |
| 791 | setattr(layer, f"{sublayer_name}_{name}", param) |
| 792 | |
| 793 | if self.force_input_schema is not None: |
| 794 | self.input_schema = self.force_input_schema |
| 795 | |
| 796 | def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: |
| 797 | return get_humming_moe_quant_config( |
| 798 | layer, |
| 799 | self.humming_configs, |
| 800 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 801 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 802 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 803 | ) |
| 804 | |
| 805 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 806 | if getattr(self, "processed", False): |
| 807 | return |
| 808 | self.processed = True |
| 809 | |
| 810 | # Convert weights to Humming kernel format |
| 811 | self.humming_configs = convert_to_humming_moe_kernel_format( |
| 812 | layer=layer, |
| 813 | sublayer_configs=layer.sublayer_configs, |
| 814 | weight_schema=self.weight_schema, |
| 815 | input_schema=self.input_schema, |
| 816 | force_weight_schema=self.force_weight_schema, |
| 817 | allow_input_schema_fallback=self.quant_config.allow_input_schema_fallback, |
| 818 | ) |
| 819 | |
| 820 | # Build the MoE kernel |
| 821 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 822 | assert self.moe_quant_config is not None |
| 823 | assert self.experts_cls is not None |
| 824 | self.moe_kernel = make_humming_moe_kernel( |
| 825 | self.moe_quant_config, |
| 826 | self.moe, |
| 827 | self.experts_cls, |
| 828 | routing_tables=layer._expert_routing_tables(), |
| 829 | ) |
| 830 | |
| 831 | def apply( |
| 832 | self, |
| 833 | layer: RoutedExperts, |
| 834 | x: torch.Tensor, |
| 835 | topk_weights: torch.Tensor, |
| 836 | topk_ids: torch.Tensor, |
| 837 | shared_experts: SharedExperts | None, |
| 838 | shared_experts_input: torch.Tensor | None, |
| 839 | ) -> torch.Tensor: |
| 840 | """Apply Humming-quantized MoE computation using the standard kernel flow. |
| 841 | |
| 842 | This method uses FusedMoEKernel.apply() which orchestrates: |
| 843 | 1. Preparation (quantization if needed - skipped for Humming via |
| 844 | expects_unquantized_inputs=True to prevent double quantization) |
| 845 | 2. Expert computation (via experts.apply()) |
| 846 | 3. Finalization (weight application & reduction - no-op for Humming |
| 847 | since it's already done internally) |
| 848 | |
| 849 | Humming handles quantization, weight application, and reduction |
| 850 | internally in the experts implementation. |
| 851 | |
| 852 | Humming experts consume the weights and quantization tensors passed |
| 853 | through the standard modular kernel interface. |
| 854 | """ |
| 855 | assert self.moe_kernel is not None |
| 856 | return self.moe_kernel.apply( |
| 857 | hidden_states=x, |
| 858 | w1=layer.w13_weight, |
| 859 | w2=layer.w2_weight, |
| 860 | topk_ids=topk_ids, |
| 861 | topk_weights=topk_weights, |
| 862 | activation=layer.activation, |
| 863 | global_num_experts=layer.global_num_experts, |
| 864 | expert_map=layer.expert_map, |
| 865 | apply_router_weight_on_input=False, |
| 866 | shared_experts=shared_experts, |
| 867 | shared_experts_input=shared_experts_input, |
| 868 | ) |