@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9f771b3ab9
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 16:15 UTC

[Quantization] add humming quantization kernel (#34556)

Jinzhen Lin committed Apr 24, 2026 at 21:29 UTC 9f771b3ab92d26a7d91a8255572c5d8d2b3ad601
10 files changed +1948 -9
vllm/config/model.py
+4
@@ -953,8 +953,12 @@ class ModelConfig:
953 "mxfp4",
954 "gpt_oss_mxfp4",
955 "cpu_awq",
956 + "humming",
957 "gguf",
958 ]
959 + # if the user specifies humming, we should always use humming
960 + if self.quantization == "humming":
961 + overrides = ["humming"] + overrides
962 quantization_methods = [
963 q for q in supported_quantization if q not in overrides
964 ]
vllm/envs.py
+32
@@ -152,6 +152,10 @@ if TYPE_CHECKING:
152 VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: str = ""
153 VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
154 VLLM_MARLIN_INPUT_DTYPE: Literal["int8", "fp8"] | None = None
155 + VLLM_HUMMING_ONLINE_QUANT_CONFIG: dict[str, Any] | None = None
156 + VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None
157 + VLLM_HUMMING_USE_F16_ACCUM: bool = False
158 + VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None
159 VLLM_MXFP4_USE_MARLIN: bool | None = None
160 VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False
161 VLLM_V1_USE_OUTLINES_CACHE: bool = False
@@ -285,6 +289,15 @@ def maybe_convert_bool(value: str | None) -> bool | None:
289 return bool(int(value))
290
291
292 +def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None:
293 + if value is None:
294 + return None
295 + if os.path.exists(value):
296 + with open(value) as f:
297 + return json.load(f)
298 + return json.loads(value)
299 +
300 +
301 def disable_compile_cache() -> bool:
302 return bool(int(os.getenv("VLLM_DISABLE_COMPILE_CACHE", "0")))
303
@@ -1193,6 +1206,25 @@ environment_variables: dict[str, Callable[[], Any]] = {
1206 "VLLM_MARLIN_INPUT_DTYPE": env_with_choices(
1207 "VLLM_MARLIN_INPUT_DTYPE", None, ["int8", "fp8"]
1208 ),
1209 + # The online quantization dtype for humming kernel
1210 + "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
1211 + os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None)
1212 + ),
1213 + # The activation dtype config for humming kernel
1214 + "VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
1215 + os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None)
1216 + ),
1217 + # Whether to use fp16 accumulator mma
1218 + "VLLM_HUMMING_USE_F16_ACCUM": lambda: maybe_convert_bool(
1219 + os.environ.get("VLLM_HUMMING_USE_F16_ACCUM", "0")
1220 + ),
1221 + # Whether to use indexed gemm for humming moe
1222 + # if 1, force use indexed gemm
1223 + # if 0, force use grouped gemm
1224 + # if None, choose better gemm type automatically
1225 + "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: maybe_convert_bool(
1226 + os.environ.get("VLLM_HUMMING_MOE_GEMM_TYPE", None)
1227 + ),
1228 # Whether to use DeepEPLL kernels for NVFP4 quantization and dispatch method
1229 # only supported on Blackwell GPUs and with
1230 # https://github.com/deepseek-ai/DeepEP/pull/341
vllm/model_executor/layers/fused_moe/fused_humming_moe.py new
+690
@@ -0,0 +1,690 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused MoE utilities for Humming."""
4 +
5 +import json
6 +import math
7 +from typing import TYPE_CHECKING, Any
8 +
9 +import torch
10 +from humming import dtypes
11 +from humming.config import GemmType as HummingGemmType
12 +from humming.layer import HummingLayerMeta, HummingMethod
13 +
14 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15 +from vllm import envs
16 +from vllm.forward_context import get_forward_context
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 +from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
20 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
21 + moe_align_block_size,
22 +)
23 +from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum
24 +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
25 + moe_permute,
26 + moe_unpermute,
27 +)
28 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
29 + TopKWeightAndReduceDelegate,
30 + TopKWeightAndReduceNoOP,
31 +)
32 +from vllm.model_executor.layers.fused_moe.utils import _resize_cache
33 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
34 +from vllm.platforms import current_platform
35 +from vllm.v1.worker.workspace import current_workspace_manager
36 +
37 +if TYPE_CHECKING:
38 + from vllm.model_executor.layers.quantization.humming import HummingMoEMethod
39 +
40 +
41 +logger = init_logger(__name__)
42 +
43 +
44 +def get_humming_moe_gemm_type() -> str:
45 + env_gemm_type: str = envs.VLLM_HUMMING_MOE_GEMM_TYPE or ""
46 + env_gemm_type = env_gemm_type.lower()
47 + if env_gemm_type in ["indexed", "grouped"]:
48 + gemm_type = env_gemm_type
49 + elif current_platform.has_device_capability(90):
50 + # for device that supports TMA, use grouped gemm
51 + gemm_type = "grouped"
52 + else:
53 + gemm_type = "indexed"
54 +
55 + logger.info_once(f"Using {gemm_type} gemm for humming moe") # noqa
56 + return gemm_type
57 +
58 +
59 +class HummingExpertsBase(mk.FusedMoEExpertsModular):
60 + def __init__(
61 + self,
62 + layer: torch.nn.Module,
63 + quant_method: "HummingMoEMethod",
64 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular | None = None,
65 + ):
66 + self.layer = layer
67 + self.num_experts = self.layer.num_experts
68 + self.global_num_experts = self.layer.global_num_experts
69 + self.init_humming_moe()
70 +
71 + if prepare_finalize is not None:
72 + max_num_tokens: int | None = None
73 + num_dispatchers: int | None = None
74 + if self.is_batched:
75 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
76 + num_dispatchers = prepare_finalize.num_dispatchers()
77 +
78 + assert quant_method.moe_quant_config is not None
79 + super().__init__(
80 + moe_config=quant_method.moe,
81 + quant_config=quant_method.moe_quant_config,
82 + max_num_tokens=max_num_tokens,
83 + num_dispatchers=num_dispatchers,
84 + )
85 + else:
86 + assert not self.is_batched
87 +
88 + def init_humming_moe(self):
89 + self.compute_config = {
90 + "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
91 + "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
92 + "gemm_type": self.humming_gemm_type.value,
93 + }
94 + self.w13_tuning_config = HummingMethod.get_default_tuning_configs(
95 + layer=self.layer,
96 + use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,
97 + use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
98 + gemm_type=self.humming_gemm_type,
99 + sublayer_name="w13",
100 + )
101 + self.w2_tuning_config = HummingMethod.get_default_tuning_configs(
102 + layer=self.layer,
103 + use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,
104 + use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
105 + gemm_type=self.humming_gemm_type,
106 + sublayer_name="w2",
107 + )
108 + self.compute_config_str = json.dumps(self.compute_config)
109 + self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
110 + self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)
111 +
112 + def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
113 + num_tokens = topk_ids.size(0)
114 + ctx = get_forward_context()
115 + if ctx.dp_metadata is not None:
116 + num_tokens = ctx.dp_metadata.num_tokens_across_dp_cpu.sum().item()
117 +
118 + return num_tokens * topk_ids.size(1)
119 +
120 + def estimate_local_valid_shape_m(self, topk_ids: torch.Tensor):
121 + # estimate shape_m for kernel tuning
122 + global_valid_shape_m = self.get_global_valid_shape_m(topk_ids)
123 + num_experts = self.num_experts
124 + global_num_experts = self.global_num_experts
125 + return math.ceil(global_valid_shape_m * num_experts / global_num_experts)
126 +
127 + @property
128 + def humming_gemm_type(self) -> HummingGemmType:
129 + raise NotImplementedError
130 +
131 + @property
132 + def is_batched(self) -> bool:
133 + return self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts
134 +
135 + @staticmethod
136 + def _supports_quant_scheme(
137 + weight_key: QuantKey | None,
138 + activation_key: QuantKey | None,
139 + ) -> bool:
140 + return True
141 +
142 + def supports_expert_map(self) -> bool:
143 + return True
144 +
145 + @staticmethod
146 + def _supports_current_device() -> bool:
147 + platform = current_platform
148 + return platform.is_cuda() and platform.has_device_capability((7, 5))
149 +
150 + @staticmethod
151 + def _supports_no_act_and_mul() -> bool:
152 + return True
153 +
154 + @staticmethod
155 + def _supports_activation(activation: MoEActivation) -> bool:
156 + # Humming uses apply_moe_activation() callback for activation,
157 + # so any activation supported there can be used here.
158 + return activation in [
159 + MoEActivation.SILU,
160 + MoEActivation.GELU,
161 + MoEActivation.SWIGLUOAI,
162 + MoEActivation.SWIGLUSTEP,
163 + MoEActivation.SILU_NO_MUL,
164 + MoEActivation.GELU_NO_MUL,
165 + MoEActivation.RELU2_NO_MUL,
166 + ]
167 +
168 + @staticmethod
169 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
170 + return not (
171 + moe_parallel_config.use_fi_nvl_two_sided_kernels
172 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
173 + )
174 +
175 + def moe_problem_size(
176 + self,
177 + a1: torch.Tensor,
178 + w1: torch.Tensor,
179 + w2: torch.Tensor,
180 + topk_ids: torch.Tensor,
181 + ) -> tuple[int, int, int, int, int]:
182 + meta1: HummingLayerMeta = self.layer.humming_metas["w13"]
183 + meta2: HummingLayerMeta = self.layer.humming_metas["w2"]
184 +
185 + assert meta1.num_experts == meta2.num_experts
186 +
187 + num_experts = meta1.num_experts
188 + top_k = topk_ids.size(1)
189 + assert w1.size(0) == num_experts
190 + assert w2.size(0) == num_experts
191 +
192 + if not self.is_batched:
193 + num_tokens = a1.size(0)
194 + assert topk_ids.size(0) == num_tokens
195 + else:
196 + assert a1.dim() == 3
197 + assert a1.size(0) == num_experts
198 + num_tokens = a1.size(1)
199 +
200 + return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k
201 +
202 + def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
203 + num_experts = self.num_experts
204 + N = self.layer.intermediate_size
205 + K = self.layer.hidden_size
206 + assert isinstance(num_experts, int)
207 + assert isinstance(N, int)
208 + assert isinstance(K, int)
209 +
210 + # hidden_states
211 + # (-> quanted_gate_up_input) (if not BF16/FP16 activation)
212 + # -> gate_up_output
213 + # -> activation_output
214 + # (-> quanted_down_input) (if not BF16/FP16 activation)
215 + # -> down_output
216 + # (-> output) (if not is_batched)
217 + # Neighboring nodes are required to utilize distinct workspaces.
218 + # The output must be derived from workspace1.
219 +
220 + output_shape: tuple[int, ...]
221 + if self.is_batched:
222 + max_num_tokens = self.max_num_tokens
223 + num_dispatchers = self.num_dispatchers
224 + assert max_num_tokens is not None and num_dispatchers is not None
225 + input_shape_m = num_experts * max_num_tokens
226 + real_shape_m = num_experts * max_num_tokens * num_dispatchers
227 + output_shape = (num_experts, max_num_tokens * num_dispatchers, K)
228 + else:
229 + input_shape_m = M
230 + if self.humming_gemm_type != HummingGemmType.INDEXED:
231 + input_shape_m = M * topk
232 + real_shape_m = M * topk
233 + output_shape = (M, K)
234 +
235 + down_input_size = N if activation.is_gated else (N * 2)
236 + a_dtype = self.layer.humming_metas["w13"].a_dtype
237 + c_dtype = self.layer.humming_metas["w13"].c_dtype
238 + num_bits = a_dtype.num_bits
239 + torch_dtype_map = {
240 + dtypes.float16: torch.float16,
241 + dtypes.bfloat16: torch.bfloat16,
242 + dtypes.float8e4m3: torch.float8_e4m3fn,
243 + dtypes.int8: torch.int8,
244 + dtypes.int4: torch.uint8,
245 + }
246 +
247 + buffer_metas = {
248 + "quanted_gate_up_input": {
249 + "shape": (input_shape_m, K),
250 + "dtype": torch_dtype_map[a_dtype],
251 + },
252 + "gate_up_output": {
253 + "shape": (real_shape_m, N * 2),
254 + "dtype": torch_dtype_map[c_dtype],
255 + },
256 + "activation_output": {
257 + "shape": (real_shape_m, down_input_size),
258 + "dtype": torch_dtype_map[c_dtype],
259 + },
260 + "quanted_down_input": {
261 + "shape": (real_shape_m, down_input_size),
262 + "dtype": torch_dtype_map[a_dtype],
263 + },
264 + "down_output": {
265 + "shape": output_shape if self.is_batched else (real_shape_m, K),
266 + "dtype": torch_dtype_map[c_dtype],
267 + },
268 + "output": {
269 + "shape": output_shape,
270 + "dtype": torch_dtype_map[c_dtype],
271 + },
272 + }
273 +
274 + for key in buffer_metas:
275 + meta = buffer_metas[key]
276 + if "quanted" in key and a_dtype.num_bits == 4:
277 + meta["shape"] = meta["shape"][:-1] + (meta["shape"][-1] // 2,)
278 +
279 + if num_bits == 16:
280 + required_buffers = ["gate_up_output", "activation_output", "down_output"]
281 + else:
282 + required_buffers = [
283 + "quanted_gate_up_input",
284 + "gate_up_output",
285 + "activation_output",
286 + "quanted_down_input",
287 + "down_output",
288 + ]
289 +
290 + # batched moe use down_output as output
291 + if not self.is_batched:
292 + required_buffers.append("output")
293 +
294 + return buffer_metas, required_buffers
295 +
296 + def _workspace_shapes(self, M: int, topk: int, activation: MoEActivation):
297 + buffer_metas, required_buffers = self.get_buffer_metas(M, topk, activation)
298 +
299 + workspace1_nbytes = 0
300 + workspace2_nbytes = 0
301 +
302 + for index, name in enumerate(required_buffers[::-1]):
303 + buffer_meta = buffer_metas[name]
304 + nelement = math.prod(buffer_meta["shape"])
305 + nbytes = nelement * buffer_meta["dtype"].itemsize
306 + if index % 2 == 0:
307 + workspace1_nbytes = max(workspace1_nbytes, nbytes)
308 + else:
309 + workspace2_nbytes = max(workspace2_nbytes, nbytes)
310 +
311 + output_key = "down_output" if self.is_batched else "output"
312 + output_shape = buffer_metas[output_key]["shape"]
313 +
314 + return (workspace1_nbytes // 2,), (workspace2_nbytes // 2,), output_shape
315 +
316 + def workspace_shapes(
317 + self,
318 + M: int,
319 + N: int,
320 + K: int,
321 + topk: int,
322 + global_num_experts: int,
323 + local_num_experts: int,
324 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
325 + activation: MoEActivation,
326 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
327 + return self._workspace_shapes(M, topk, activation)
328 +
329 + def make_workspaces(self, M: int, topk: int, activation: MoEActivation):
330 + shapes = self._workspace_shapes(M, topk, activation)
331 + workspace1_shape, workspace2_shape, output_shape = shapes
332 + torch_dtype = self.layer.param_dtype
333 + workspace1, workspace2 = current_workspace_manager().get_simultaneous(
334 + (workspace1_shape, torch_dtype),
335 + (workspace2_shape, torch_dtype),
336 + )
337 + output = _resize_cache(workspace1, output_shape)
338 + return workspace1, workspace2, output
339 +
340 + def prepare_buffers(
341 + self,
342 + workspace1: torch.Tensor,
343 + workspace2: torch.Tensor,
344 + M: int,
345 + topk: int,
346 + activation: MoEActivation,
347 + ) -> dict[str, torch.Tensor]:
348 + buffer_metas, required_buffers = self.get_buffer_metas(M, topk, activation)
349 + buffers = {}
350 + for index, name in enumerate(required_buffers[::-1]):
351 + buffer_meta = buffer_metas[name]
352 + workspace = workspace1 if index % 2 == 0 else workspace2
353 + workspace = workspace.view(buffer_meta["dtype"])
354 + buffers[name] = _resize_cache(workspace, buffer_meta["shape"])
355 +
356 + return buffers
357 +
358 + def apply(
359 + self,
360 + output: torch.Tensor,
361 + hidden_states: torch.Tensor,
362 + w1: torch.Tensor,
363 + w2: torch.Tensor,
364 + topk_weights: torch.Tensor,
365 + topk_ids: torch.Tensor,
366 + activation: MoEActivation,
367 + global_num_experts: int,
368 + expert_map: torch.Tensor | None,
369 + a1q_scale: torch.Tensor | None,
370 + a2_scale: torch.Tensor | None,
371 + workspace13: torch.Tensor,
372 + workspace2: torch.Tensor,
373 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
374 + apply_router_weight_on_input: bool,
375 + ):
376 + assert not apply_router_weight_on_input
377 +
378 + self.main_apply(
379 + hidden_states=hidden_states,
380 + topk_weights=topk_weights,
381 + topk_ids=topk_ids,
382 + workspace1=workspace13,
383 + workspace2=workspace2,
384 + expert_tokens_meta=expert_tokens_meta,
385 + )
386 +
387 + def main_apply(
388 + self,
389 + hidden_states: torch.Tensor,
390 + topk_weights: torch.Tensor,
391 + topk_ids: torch.Tensor,
392 + workspace1: torch.Tensor,
393 + workspace2: torch.Tensor,
394 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
395 + ):
396 + raise NotImplementedError
397 +
398 +
399 +class HummingIndexedExperts(HummingExpertsBase):
400 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
401 + return TopKWeightAndReduceNoOP()
402 +
403 + @staticmethod
404 + def activation_format() -> mk.FusedMoEActivationFormat:
405 + return mk.FusedMoEActivationFormat.Standard
406 +
407 + @property
408 + def humming_gemm_type(self) -> HummingGemmType:
409 + return HummingGemmType.INDEXED
410 +
411 + def prepare_humming_moe_kwargs(
412 + self,
413 + topk_ids: torch.Tensor,
414 + expert_map: torch.Tensor | None,
415 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
416 + ) -> tuple[dict[str, Any], dict[str, Any]]:
417 + valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
418 +
419 + for min_shape_m, max_shape_m, config in self.w13_tuning_config:
420 + if valid_shape_m > min_shape_m and valid_shape_m <= max_shape_m:
421 + moe_block_size = config["block_shape"][0]
422 + break
423 + else:
424 + raise ValueError(f"cannot found moe_block_size for shape {valid_shape_m}")
425 +
426 + sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size(
427 + topk_ids=topk_ids,
428 + block_size=moe_block_size,
429 + num_experts=self.global_num_experts,
430 + expert_map=expert_map,
431 + ignore_invalid_experts=True,
432 + )
433 +
434 + moe_common_kwargs = {
435 + "sorted_ids": sorted_ids,
436 + "expert_ids": expert_ids,
437 + "num_tokens_padded": num_tokens_padded,
438 + "compute_config": self.compute_config_str,
439 + "valid_shape_m": valid_shape_m,
440 + }
441 +
442 + top_k = topk_ids.size(1)
443 + moe_kwargs1 = {"top_k": top_k, "tuning_config": self.w13_tuning_config_str}
444 + moe_kwargs2 = {"top_k": 1, "tuning_config": self.w2_tuning_config_str}
445 + moe_kwargs1.update(moe_common_kwargs)
446 + moe_kwargs2.update(moe_common_kwargs)
447 +
448 + return moe_kwargs1, moe_kwargs2
449 +
450 + def main_apply(
451 + self,
452 + hidden_states: torch.Tensor,
453 + topk_weights: torch.Tensor,
454 + topk_ids: torch.Tensor,
455 + workspace1: torch.Tensor,
456 + workspace2: torch.Tensor,
457 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
458 + ):
459 + hidden_states = hidden_states.view(-1, hidden_states.size(-1))
460 + buffers = self.prepare_buffers(
461 + workspace1,
462 + workspace2,
463 + topk_ids.size(0),
464 + topk_ids.size(1),
465 + self.layer.activation,
466 + )
467 +
468 + moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs(
469 + topk_ids=topk_ids,
470 + expert_map=self.layer.expert_map,
471 + expert_tokens_meta=expert_tokens_meta,
472 + )
473 +
474 + inputs, input_scale = HummingMethod.may_quant_input(
475 + layer=self.layer,
476 + inputs=hidden_states,
477 + quanted_input=buffers.get("quanted_gate_up_input", None),
478 + sublayer_name="w13",
479 + )
480 +
481 + HummingMethod.forward_layer(
482 + layer=self.layer,
483 + inputs=inputs,
484 + input_scale=input_scale,
485 + outputs=buffers["gate_up_output"],
486 + sublayer_name="w13",
487 + **moe_kwargs1,
488 + )
489 +
490 + self.activation(
491 + activation=self.layer.activation,
492 + input=buffers["gate_up_output"],
493 + output=buffers["activation_output"],
494 + )
495 +
496 + inputs, input_scale = HummingMethod.may_quant_input(
497 + layer=self.layer,
498 + inputs=buffers["activation_output"],
499 + quanted_input=buffers.get("quanted_down_input", None),
500 + sublayer_name="w2",
501 + )
502 +
503 + HummingMethod.forward_layer(
504 + layer=self.layer,
505 + inputs=inputs,
506 + input_scale=input_scale,
507 + outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
508 + sublayer_name="w2",
509 + **moe_kwargs2,
510 + )
511 +
512 + moe_fused_mul_sum(
513 + inputs=buffers["down_output"].view(*topk_ids.shape, -1),
514 + topk_weights=topk_weights,
515 + topk_ids=topk_ids,
516 + expert_map=self.layer.expert_map,
517 + outputs=buffers["output"],
518 + )
519 +
520 +
521 +class HummingGroupedExperts(HummingExpertsBase):
522 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
523 + return TopKWeightAndReduceNoOP()
524 +
525 + @staticmethod
526 + def activation_format() -> mk.FusedMoEActivationFormat:
527 + return mk.FusedMoEActivationFormat.Standard
528 +
529 + @property
530 + def humming_gemm_type(self) -> HummingGemmType:
531 + return HummingGemmType.GROUPED_CONTIGUOUS
532 +
533 + def main_apply(
534 + self,
535 + hidden_states: torch.Tensor,
536 + topk_weights: torch.Tensor,
537 + topk_ids: torch.Tensor,
538 + workspace1: torch.Tensor,
539 + workspace2: torch.Tensor,
540 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
541 + ):
542 + valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
543 +
544 + buffers = self.prepare_buffers(
545 + workspace1,
546 + workspace2,
547 + topk_ids.size(0),
548 + topk_ids.size(1),
549 + self.layer.activation,
550 + )
551 +
552 + hidden_states, _, expert_first_token_offset, inv_perm, _ = moe_permute(
553 + hidden_states=hidden_states,
554 + a1q_scale=None,
555 + topk_ids=topk_ids,
556 + n_expert=self.global_num_experts,
557 + n_local_expert=self.num_experts,
558 + expert_map=self.layer.expert_map,
559 + )
560 +
561 + inputs, input_scale = HummingMethod.may_quant_input(
562 + layer=self.layer,
563 + inputs=hidden_states,
564 + quanted_input=buffers.get("quanted_gate_up_input", None),
565 + sublayer_name="w13",
566 + )
567 +
568 + HummingMethod.forward_layer(
569 + layer=self.layer,
570 + inputs=inputs,
571 + input_scale=input_scale,
572 + outputs=buffers["gate_up_output"],
573 + valid_shape_m=valid_shape_m,
574 + expert_layout=expert_first_token_offset,
575 + compute_config=self.compute_config_str,
576 + tuning_config=self.w13_tuning_config_str,
577 + sublayer_name="w13",
578 + )
579 +
580 + self.activation(
581 + activation=self.layer.activation,
582 + input=buffers["gate_up_output"],
583 + output=buffers["activation_output"],
584 + )
585 +
586 + inputs, input_scale = HummingMethod.may_quant_input(
587 + layer=self.layer,
588 + inputs=buffers["activation_output"],
589 + quanted_input=buffers.get("quanted_down_input", None),
590 + sublayer_name="w2",
591 + )
592 +
593 + HummingMethod.forward_layer(
594 + layer=self.layer,
595 + inputs=inputs,
596 + input_scale=input_scale,
597 + outputs=buffers["down_output"],
598 + valid_shape_m=valid_shape_m,
599 + expert_layout=expert_first_token_offset,
600 + compute_config=self.compute_config_str,
601 + tuning_config=self.w2_tuning_config_str,
602 + sublayer_name="w2",
603 + )
604 +
605 + moe_unpermute(
606 + out=buffers["output"],
607 + permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1),
608 + topk_weights=topk_weights,
609 + inv_permuted_idx=inv_perm,
610 + expert_first_token_offset=expert_first_token_offset,
611 + )
612 +
613 +
614 +class BatchedHummingGroupedExperts(HummingExpertsBase):
615 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
616 + return TopKWeightAndReduceDelegate()
617 +
618 + @staticmethod
619 + def activation_format() -> mk.FusedMoEActivationFormat:
620 + return mk.FusedMoEActivationFormat.BatchedExperts
621 +
622 + @property
623 + def humming_gemm_type(self) -> HummingGemmType:
624 + return HummingGemmType.GROUPED_MASKED
625 +
626 + def main_apply(
627 + self,
628 + hidden_states: torch.Tensor,
629 + topk_weights: torch.Tensor,
630 + topk_ids: torch.Tensor,
631 + workspace1: torch.Tensor,
632 + workspace2: torch.Tensor,
633 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
634 + ):
635 + assert expert_tokens_meta is not None
636 + hidden_states = hidden_states.view(-1, hidden_states.size(-1))
637 + valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
638 + expert_num_tokens = expert_tokens_meta.expert_num_tokens
639 +
640 + buffers = self.prepare_buffers(
641 + workspace1,
642 + workspace2,
643 + topk_ids.size(0),
644 + topk_ids.size(1),
645 + self.layer.activation,
646 + )
647 +
648 + inputs, input_scale = HummingMethod.may_quant_input(
649 + layer=self.layer,
650 + inputs=hidden_states,
651 + quanted_input=buffers.get("quanted_gate_up_input", None),
652 + sublayer_name="w13",
653 + )
654 +
655 + HummingMethod.forward_layer(
656 + layer=self.layer,
657 + inputs=inputs,
658 + input_scale=input_scale,
659 + outputs=buffers["gate_up_output"],
660 + valid_shape_m=valid_shape_m,
661 + expert_layout=expert_num_tokens,
662 + compute_config=self.compute_config_str,
663 + tuning_config=self.w13_tuning_config_str,
664 + sublayer_name="w13",
665 + )
666 +
667 + self.activation(
668 + activation=self.layer.activation,
669 + input=buffers["gate_up_output"],
670 + output=buffers["activation_output"],
671 + )
672 +
673 + inputs, input_scale = HummingMethod.may_quant_input(
674 + layer=self.layer,
675 + inputs=buffers["activation_output"],
676 + quanted_input=buffers.get("quanted_down_input", None),
677 + sublayer_name="w2",
678 + )
679 +
680 + HummingMethod.forward_layer(
681 + layer=self.layer,
682 + inputs=inputs,
683 + input_scale=input_scale,
684 + outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
685 + valid_shape_m=valid_shape_m,
686 + expert_layout=expert_num_tokens,
687 + compute_config=self.compute_config_str,
688 + tuning_config=self.w2_tuning_config_str,
689 + sublayer_name="w2",
690 + )
vllm/model_executor/layers/fused_moe/layer.py
+5 -1
@@ -1097,7 +1097,11 @@ class FusedMoE(PluggableLayer):
1097 expert_id: int,
1098 return_success: bool = False,
1099 ) -> bool | None:
1100 - if self.quant_config and self.quant_config.get_name() == "gpt_oss_mxfp4":
1100 + quant_config_name = self.quant_config and self.quant_config.get_name()
1101 + if quant_config_name == "humming":
1102 + assert hasattr(self.quant_method, "weight_schema")
1103 + quant_config_name = self.quant_method.weight_schema.quant_method
1104 + if quant_config_name == "gpt_oss_mxfp4":
1105 # (FIXME) for gpt-oss all experts are combined
1106 if "bias" in weight_name:
1107 dim1 = loaded_weight.shape[1]
vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py new
+202
@@ -0,0 +1,202 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from torch._subclasses.fake_tensor import FakeTensor
5 +
6 +from vllm.platforms import current_platform
7 +from vllm.triton_utils import tl, triton
8 +
9 +
10 +@triton.jit
11 +def moe_fused_mul_sum_kernel(
12 + inputs_ptr,
13 + topk_weights_ptr,
14 + outputs_ptr,
15 + top_ids_ptr,
16 + expert_map_ptr,
17 + num_tokens,
18 + stride_m,
19 + has_expert_map: tl.constexpr,
20 + top_k: tl.constexpr,
21 + size: tl.constexpr,
22 + BLOCK_M: tl.constexpr,
23 + BLOCK_K: tl.constexpr,
24 +):
25 + pid_k = tl.program_id(0)
26 + pid_m = tl.program_id(1)
27 +
28 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
29 + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
30 +
31 + m_mask = offs_m < num_tokens
32 + k_mask = offs_k < size
33 + mask = m_mask[:, None] & k_mask[None, :]
34 +
35 + a_base = inputs_ptr + (offs_m * stride_m)[:, None] + offs_k[None, :]
36 + b_base = topk_weights_ptr + offs_m * top_k
37 +
38 + acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32)
39 +
40 + for n in tl.static_range(top_k):
41 + b_val = tl.load(b_base + n, mask=m_mask, other=0.0).to(tl.float32)
42 + if has_expert_map:
43 + id_val = tl.load(top_ids_ptr + offs_m * top_k + n, mask=m_mask, other=0)
44 + expert_mask = tl.load(expert_map_ptr + id_val) >= 0
45 + a_vec = tl.load(
46 + a_base + n * size,
47 + mask=mask & expert_mask[:, None],
48 + other=0.0,
49 + ).to(tl.float32)
50 + else:
51 + a_vec = tl.load(
52 + a_base + n * size,
53 + mask=mask,
54 + other=0.0,
55 + ).to(tl.float32)
56 + acc += a_vec * b_val[:, None]
57 +
58 + out_ptrs = outputs_ptr + (offs_m * size)[:, None] + offs_k[None, :]
59 + tl.store(
60 + out_ptrs,
61 + acc.to(outputs_ptr.dtype.element_ty),
62 + mask=mask,
63 + )
64 +
65 +
66 +def _heuristic_config(
67 + num_tokens: int,
68 + top_k: int,
69 + size: int,
70 + element_size: int,
71 +):
72 + is_fp32 = element_size > 2
73 + is_sm90_plus = current_platform.has_device_capability(90)
74 + is_sm80_before = not current_platform.has_device_capability(80)
75 +
76 + if current_platform.has_device_capability(90):
77 + # SM90/SM100+: prefer small tiles + many CTAs.
78 + if is_fp32:
79 + BLOCK_M = 1 if num_tokens <= 4 else 2
80 + else:
81 + if num_tokens <= 4:
82 + BLOCK_M = 1
83 + elif num_tokens <= 128:
84 + BLOCK_M = 2
85 + else:
86 + BLOCK_M = 4
87 + elif is_fp32:
88 + if num_tokens <= 4:
89 + BLOCK_M = 1
90 + elif num_tokens <= 32:
91 + BLOCK_M = 2
92 + elif num_tokens <= 128:
93 + BLOCK_M = 4
94 + else:
95 + BLOCK_M = 4
96 + else:
97 + if num_tokens <= 4:
98 + BLOCK_M = 1
99 + elif num_tokens <= 32:
100 + BLOCK_M = 2
101 + elif num_tokens <= 128:
102 + BLOCK_M = 4
103 + elif num_tokens <= 1024:
104 + BLOCK_M = 16
105 + else:
106 + BLOCK_M = 8
107 +
108 + if is_fp32:
109 + max_block_k = 256
110 + elif is_sm80_before or is_sm90_plus:
111 + max_block_k = 512
112 + else:
113 + max_block_k = 1024
114 + BLOCK_K = min(triton.next_power_of_2(size), max_block_k)
115 + BLOCK_K = max(BLOCK_K, 256)
116 +
117 + total = BLOCK_M * BLOCK_K
118 + if is_fp32:
119 + num_warps = max(8, min(16, total // 64))
120 + else:
121 + num_warps = max(4, min(16, total // 256))
122 +
123 + if is_sm80_before:
124 + num_warps = min(num_warps, 8)
125 + num_stages = 2
126 + elif is_sm90_plus:
127 + num_warps = min(num_warps, 8)
128 + num_stages = 4 if total <= 2048 else 2
129 + else:
130 + num_stages = 4 if total <= 2048 else 2
131 +
132 + return BLOCK_M, BLOCK_K, num_warps, num_stages
133 +
134 +
135 +def moe_fused_mul_sum(
136 + inputs: torch.Tensor,
137 + topk_weights: torch.Tensor,
138 + outputs: torch.Tensor | None = None,
139 + topk_ids: torch.Tensor | None = None,
140 + expert_map: torch.Tensor | None = None,
141 +) -> torch.Tensor:
142 + """
143 + Fused kernel for MoE (Mixture of Experts) to perform weighted summation
144 + of expert outputs.
145 +
146 + Args:
147 + inputs: The output from experts.
148 + Shape: (num_tokens, top_k, hidden_size).
149 + topk_weights: The weights assigned to each expert for each token.
150 + Shape: (num_tokens, top_k).
151 + outputs: Optional pre-allocated output tensor.
152 + Shape: (num_tokens, hidden_size).
153 + topk_ids: Optional indices of the top-k experts. Used when
154 + `expert_map` is provided. Shape: (num_tokens, top_k).
155 + expert_map: Optional mapping for Expert Parallelism. A value < 0
156 + indicates an invalid token/expert pair that will be skipped.
157 +
158 + Returns:
159 + The fused weighted sum of expert outputs.
160 + Shape: (num_tokens, hidden_size).
161 + """
162 + assert inputs.ndim == 3
163 + assert topk_weights.ndim == 2
164 + assert inputs.is_contiguous()
165 + assert topk_weights.is_contiguous()
166 + assert inputs.dtype in (torch.float32, torch.float16, torch.bfloat16)
167 + assert topk_weights.dtype in (torch.float32, torch.float16, torch.bfloat16)
168 +
169 + num_tokens, top_k, size = inputs.shape
170 + output_shape = (num_tokens, size)
171 + if outputs is None:
172 + outputs = torch.empty(output_shape, dtype=inputs.dtype, device=inputs.device)
173 +
174 + assert outputs.shape == output_shape
175 + assert topk_weights.shape == (num_tokens, top_k)
176 +
177 + if not isinstance(inputs, FakeTensor):
178 + BLOCK_M, BLOCK_K, num_warps, num_stages = _heuristic_config(
179 + num_tokens,
180 + top_k,
181 + size,
182 + inputs.element_size(),
183 + )
184 + grid = (triton.cdiv(size, BLOCK_K), triton.cdiv(num_tokens, BLOCK_M))
185 + moe_fused_mul_sum_kernel[grid](
186 + inputs,
187 + topk_weights,
188 + outputs,
189 + topk_ids,
190 + expert_map,
191 + num_tokens,
192 + top_k * size,
193 + expert_map is not None,
194 + top_k,
195 + size,
196 + BLOCK_M,
197 + BLOCK_K,
198 + num_warps=num_warps,
199 + num_stages=num_stages,
200 + )
201 +
202 + return outputs
vllm/model_executor/layers/linear.py
+13 -6
@@ -60,6 +60,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
60 "ModelOptFp8PbWoLinearMethod",
61 "QuarkLinearMethod",
62 "ModelOptNvFp4LinearMethod",
63 + "HummingLinearMethod",
64 ]
65
66
@@ -245,6 +246,7 @@ class LinearBase(PluggableLayer):
246 self,
247 input_size: int,
248 output_size: int,
249 + bias: bool = False,
250 skip_bias_add: bool = False,
251 params_dtype: torch.dtype | None = None,
252 quant_config: QuantizationConfig | None = None,
@@ -258,6 +260,7 @@ class LinearBase(PluggableLayer):
260 # Keep input parameters
261 self.input_size = input_size
262 self.output_size = output_size
263 + self.has_bias = bias
264 self.skip_bias_add = skip_bias_add
265 if params_dtype is None:
266 params_dtype = torch.get_default_dtype()
@@ -323,6 +326,7 @@ class ReplicatedLinear(LinearBase):
326 super().__init__(
327 input_size,
328 output_size,
329 + bias,
330 skip_bias_add,
331 params_dtype,
332 quant_config,
@@ -458,6 +462,7 @@ class ColumnParallelLinear(LinearBase):
462 super().__init__(
463 input_size,
464 output_size,
465 + bias,
466 skip_bias_add,
467 params_dtype,
468 quant_config,
@@ -483,6 +488,7 @@ class ColumnParallelLinear(LinearBase):
488 else self.weight_loader
489 ),
490 )
491 +
492 if bias:
493 self.bias = Parameter(
494 torch.empty(self.output_size_per_partition, dtype=params_dtype)
@@ -817,8 +823,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
823 # for the packing.
824 packed_dim = getattr(param, "packed_dim", None)
825 if packed_dim == output_dim:
820 - shard_size = shard_size // param.packed_factor
821 - shard_offset = shard_offset // param.packed_factor
826 + shard_size = round(shard_size // param.packed_factor)
827 + shard_offset = round(shard_offset // param.packed_factor)
828 # Special case for Marlin.
829 shard_size, shard_offset = adjust_marlin_shard(
830 param, shard_size, shard_offset
@@ -1252,8 +1258,8 @@ class QKVParallelLinear(ColumnParallelLinear):
1258 )
1259
1260 if packed_dim == output_dim:
1255 - shard_size = shard_size // param.packed_factor
1256 - shard_offset = shard_offset // param.packed_factor
1261 + shard_size = round(shard_size // param.packed_factor)
1262 + shard_offset = round(shard_offset // param.packed_factor)
1263
1264 # Special case for Marlin.
1265 shard_size, shard_offset = adjust_marlin_shard(
@@ -1315,8 +1321,8 @@ class QKVParallelLinear(ColumnParallelLinear):
1321 # for the packing.
1322 packed_dim = getattr(param, "packed_dim", None)
1323 if packed_dim == output_dim:
1318 - shard_size = shard_size // param.packed_factor
1319 - shard_offset = shard_offset // param.packed_factor
1324 + shard_size = round(shard_size // param.packed_factor)
1325 + shard_offset = round(shard_offset // param.packed_factor)
1326
1327 # Special case for Marlin.
1328 shard_size, shard_offset = adjust_marlin_shard(
@@ -1440,6 +1446,7 @@ class RowParallelLinear(LinearBase):
1446 super().__init__(
1447 input_size,
1448 output_size,
1449 + bias,
1450 skip_bias_add,
1451 params_dtype,
1452 quant_config,
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -22,6 +22,7 @@ QuantizationMethods = Literal[
22 "gptq_marlin",
23 "awq_marlin",
24 "gptq",
25 + "humming",
26 "compressed-tensors",
27 "bitsandbytes",
28 "experts_int8",
@@ -126,6 +127,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
127 from .gguf import GGUFConfig
128 from .gptq import GPTQConfig
129 from .gptq_marlin import GPTQMarlinConfig
130 + from .humming import HummingConfig
131 from .inc import INCConfig
132 from .modelopt import (
133 ModelOptFp8Config,
@@ -162,6 +164,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
164 "mxfp4": Mxfp4Config,
165 "gpt_oss_mxfp4": GptOssMxfp4Config,
166 "cpu_awq": CPUAWQConfig,
167 + "humming": HummingConfig,
168 "online": OnlineQuantizationConfig,
169 }
170
vllm/model_executor/layers/quantization/humming.py new
+962
@@ -0,0 +1,962 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import json
4 +import math
5 +from collections.abc import Callable
6 +from typing import TYPE_CHECKING, Any
7 +
8 +import regex as re
9 +import torch
10 +
11 +from vllm import envs
12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FusedMoEConfig,
15 + FusedMoEQuantConfig,
16 + FusedMoEQuantDesc,
17 +)
18 +from vllm.model_executor.layers.fused_moe.layer import (
19 + FusedMoE,
20 + FusedMoEMethodBase,
21 +)
22 +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
23 + UnquantizedFusedMoEMethod,
24 +)
25 +from vllm.model_executor.layers.linear import (
26 + LinearBase,
27 + LinearMethodBase,
28 + UnquantizedLinearMethod,
29 +)
30 +from vllm.model_executor.layers.quantization import QuantizationMethods
31 +from vllm.model_executor.layers.quantization.base_config import (
32 + QuantizationConfig,
33 + QuantizeMethodBase,
34 +)
35 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
36 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37 +from vllm.model_executor.parameter import (
38 + BasevLLMParameter,
39 + BlockQuantScaleParameter,
40 + ChannelQuantScaleParameter,
41 + GroupQuantScaleParameter,
42 + ModelWeightParameter,
43 + PackedvLLMParameter,
44 + PerTensorScaleParameter,
45 + RowvLLMParameter,
46 +)
47 +from vllm.model_executor.utils import set_weight_attrs
48 +
49 +if TYPE_CHECKING:
50 + from vllm.model_executor.models.utils import WeightsMapper
51 +
52 +
53 +try:
54 + from humming.dtypes import DataType
55 + from humming.layer import HummingMethod
56 + from humming.schema import (
57 + BaseInputSchema,
58 + BaseWeightSchema,
59 + HummingInputSchema,
60 + HummingWeightSchema,
61 + )
62 + from humming.utils.weight import quantize_weight
63 +
64 + from vllm.model_executor.layers.fused_moe.fused_humming_moe import (
65 + BatchedHummingGroupedExperts,
66 + HummingGroupedExperts,
67 + HummingIndexedExperts,
68 + get_humming_moe_gemm_type,
69 + )
70 +except ModuleNotFoundError:
71 + HummingMethod = None
72 +
73 +
74 +def assert_humming_available():
75 + assert HummingMethod is not None, (
76 + "humming is not available, please run "
77 + "'pip install git+https://github.com/inclusionAI/humming' to install it."
78 + )
79 +
80 +
81 +def prepare_padded_shape(shape, x):
82 + padded_shape = math.ceil(shape / x) * x
83 + return padded_shape, padded_shape - shape
84 +
85 +
86 +def prepare_param(tensor, name, extra_attrs):
87 + extra_attrs = extra_attrs.copy()
88 + scale_type = extra_attrs.pop("scale_type", None)
89 + param_cls_name_map = {
90 + "block": BlockQuantScaleParameter,
91 + "tensor": PerTensorScaleParameter,
92 + "group": GroupQuantScaleParameter,
93 + "channel": ChannelQuantScaleParameter,
94 + "input_scale": PerTensorScaleParameter,
95 + }
96 +
97 + param_cls: type[BasevLLMParameter]
98 + if "packed_dim" in extra_attrs:
99 + param_cls = PackedvLLMParameter
100 + elif scale_type in param_cls_name_map:
101 + param_cls = param_cls_name_map[scale_type]
102 + elif "output_dim" in extra_attrs and "input_dim" in extra_attrs:
103 + param_cls = ModelWeightParameter
104 + elif "input_dim" in extra_attrs:
105 + param_cls = RowvLLMParameter
106 + elif "output_dim" in extra_attrs:
107 + param_cls = ChannelQuantScaleParameter
108 + else:
109 + param_cls = BasevLLMParameter
110 +
111 + kwargs_keys = [
112 + "input_dim",
113 + "output_dim",
114 + "packed_dim",
115 + "packed_factor",
116 + "weight_loader",
117 + ]
118 + cls_kwargs = {}
119 + for key in extra_attrs.copy():
120 + if key in kwargs_keys:
121 + cls_kwargs[key] = extra_attrs.pop(key)
122 +
123 + param = param_cls(data=tensor, **cls_kwargs)
124 + set_weight_attrs(param, extra_attrs)
125 +
126 + param.param_name = name
127 + param.ignore_warning = True
128 + if scale_type in ["tensor", "input_scale"]:
129 + param.needs_scalar_to_array = True
130 +
131 + return param
132 +
133 +
134 +def prepare_moe_param(tensor, name, extra_attrs):
135 + param = torch.nn.Parameter(tensor, requires_grad=False)
136 + if "scale_type" in extra_attrs:
137 + extra_attrs["quant_method"] = extra_attrs["scale_type"]
138 +
139 + if "input_dim" in extra_attrs and "output_dim" in extra_attrs:
140 + input_dim = extra_attrs["input_dim"]
141 + output_dim = extra_attrs["output_dim"]
142 + extra_attrs["is_transposed"] = input_dim < output_dim
143 +
144 + set_weight_attrs(param, extra_attrs)
145 + param.param_name = name
146 + return param
147 +
148 +
149 +def may_pad_loaded_weight(param, loaded_weight):
150 + pad_shape = getattr(param, "pad_shape", None)
151 + if pad_shape is None:
152 + return loaded_weight
153 + value = 1 if loaded_weight.dtype == torch.float8_e8m0fnu else 0
154 + padding = []
155 + for x in pad_shape[::-1][: loaded_weight.ndim]:
156 + padding += [0, x]
157 + loaded_weight = torch.nn.functional.pad(
158 + input=loaded_weight,
159 + pad=padding,
160 + value=value,
161 + )
162 + return loaded_weight
163 +
164 +
165 +def compressed_tensors_get_config(config: dict[str, Any], key: str):
166 + assert key in ["weights", "input_activations"]
167 + target_group_config = None
168 + for group_config in config["config_groups"].values():
169 + if "Linear" in group_config["targets"]:
170 + if "weights" not in group_config:
171 + return None
172 + if key not in group_config or group_config[key] is None:
173 + return None
174 + target_group_config = group_config[key].copy()
175 + break
176 +
177 + if target_group_config is None:
178 + return None
179 + target_group_config["quant_method"] = config["quant_method"]
180 + if config["quant_method"] == "compressed-tensors":
181 + target_group_config["format"] = config["format"]
182 + elif config["quant_method"] == "modelopt":
183 + target_group_config["quant_algo"] = config["quant_algo"]
184 + return target_group_config
185 +
186 +
187 +class HummingConfig(QuantizationConfig):
188 + packed_modules_mapping = {}
189 +
190 + def __init__(self, full_config: dict[str, Any] | None = None):
191 + assert_humming_available()
192 + self.full_config: dict[str, Any] = full_config or {}
193 +
194 + @classmethod
195 + def get_name(cls) -> QuantizationMethods:
196 + return "humming"
197 +
198 + @classmethod
199 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:
200 + return [torch.bfloat16, torch.half]
201 +
202 + @classmethod
203 + def get_min_capability(cls) -> int:
204 + return 75
205 +
206 + @classmethod
207 + def get_config_filenames(cls) -> list[str]:
208 + return []
209 +
210 + @classmethod
211 + def from_config(cls, config: dict[str, Any]) -> "HummingConfig":
212 + return cls(full_config=config)
213 +
214 + @classmethod
215 + def override_quantization_method(
216 + cls, hf_quant_cfg, user_quant, hf_config=None
217 + ) -> QuantizationMethods | None:
218 + return "humming" if user_quant == "humming" else None
219 +
220 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
221 + self.hf_to_vllm_mapper = hf_to_vllm_mapper
222 +
223 + def is_layer_skipped(self, config: dict[str, Any], prefix: str):
224 + keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
225 + ignored_layers = self.get_from_keys_or(config, keys, []) or []
226 + if hasattr(self, "hf_to_vllm_mapper"):
227 + ignored_layers = self.hf_to_vllm_mapper.apply_list(ignored_layers)
228 +
229 + if any(module_name in prefix for module_name in ignored_layers):
230 + return True
231 + if "lm_head" in prefix:
232 + return True
233 +
234 + for regex in config.get("dynamic", {}):
235 + if regex[:1] != "-":
236 + continue
237 + if re.match(regex[2:], prefix):
238 + return True
239 +
240 + return False
241 +
242 + def get_layer_weight_schema(self, config: dict[str, Any], prefix: str):
243 + if self.is_layer_skipped(config, prefix):
244 + return None
245 +
246 + if config["quant_method"] in ["compressed-tensors", "modelopt"]:
247 + group_config = compressed_tensors_get_config(config, "weights")
248 + if group_config is None:
249 + return None
250 + config = group_config
251 +
252 + layer_config = config
253 + layer_dynamic = config.get("dynamic", {})
254 + if not isinstance(layer_dynamic, dict):
255 + layer_dynamic = {}
256 + for regex, override_config in layer_dynamic.items():
257 + if regex[:1] != "+":
258 + continue
259 + if re.match(regex[2:], prefix):
260 + layer_config = config.copy()
261 + layer_config.update(override_config)
262 + break
263 +
264 + if "quant_method" in layer_config:
265 + return BaseWeightSchema.from_config(layer_config)
266 + return None
267 +
268 + def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
269 + if self.is_layer_skipped(config, prefix):
270 + return None
271 + if config["quant_method"] in ["compressed-tensors", "modelopt"]:
272 + group_config = compressed_tensors_get_config(config, "input_activations")
273 + if group_config is None:
274 + return None
275 + config = group_config
276 +
277 + if config.get("quant_method", None) in BaseInputSchema.INPUT_SCHEMA_MAP:
278 + return BaseInputSchema.from_config(config)
279 + return None
280 +
281 + def get_quant_config_for_layer(
282 + self, prefix: str, layer_type: str
283 + ) -> "HummingLayerQuantizationConfig | None":
284 + weight_schema: BaseWeightSchema | None = None
285 + force_weight_schema: HummingWeightSchema | None = None
286 +
287 + if self.full_config:
288 + weight_schema = self.get_layer_weight_schema(self.full_config, prefix)
289 +
290 + is_online_quant = False
291 + online_quant_config = envs.VLLM_HUMMING_ONLINE_QUANT_CONFIG or {}
292 + if not self.full_config or online_quant_config.get("force_requant", False):
293 + online_quant_config["quant_method"] = "humming"
294 + schema = self.get_layer_weight_schema(online_quant_config, prefix)
295 + if not self.full_config:
296 + weight_schema = schema
297 + is_online_quant = True
298 + else:
299 + force_weight_schema = schema
300 +
301 + if weight_schema is not None:
302 + if weight_schema.quant_method == "gpt_oss_mxfp4" and layer_type != "moe":
303 + return None
304 + input_schema = None
305 + force_input_schema = None
306 +
307 + if self.full_config:
308 + input_schema = self.get_layer_input_schema(self.full_config, prefix)
309 +
310 + if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
311 + quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy()
312 + quant_config["quant_method"] = "humming"
313 + force_input_schema = self.get_layer_input_schema(quant_config, prefix)
314 + if input_schema is None:
315 + input_schema = force_input_schema
316 +
317 + if force_weight_schema is not None and force_input_schema is None:
318 + force_input_schema = HummingInputSchema()
319 +
320 + return HummingLayerQuantizationConfig(
321 + weight_schema=weight_schema,
322 + input_schema=input_schema,
323 + force_weight_schema=force_weight_schema,
324 + force_input_schema=force_input_schema,
325 + is_online_quant=is_online_quant,
326 + )
327 + return None
328 +
329 + def get_quant_method(
330 + self, layer: torch.nn.Module, prefix: str
331 + ) -> "QuantizeMethodBase | None":
332 + layer_type = "other"
333 + if isinstance(layer, FusedMoE):
334 + layer_type = "moe"
335 + elif isinstance(layer, LinearBase):
336 + layer_type = "linear"
337 +
338 + # TODO: remove this after humming moe backend is ready
339 + quant_method = self.full_config.get("quant_method", None)
340 + moe_activation = getattr(layer, "activation", None)
341 + if quant_method == "mxfp4" and moe_activation == MoEActivation.SWIGLUOAI:
342 + self.full_config["quan_method"] = "gpt_oss_mxfp4"
343 +
344 + quant_config = self.get_quant_config_for_layer(prefix, layer_type)
345 + if quant_config is None:
346 + if isinstance(layer, FusedMoE):
347 + return UnquantizedFusedMoEMethod(layer.moe_config)
348 + elif isinstance(layer, LinearBase):
349 + return UnquantizedLinearMethod()
350 + elif isinstance(layer, LinearBase):
351 + return HummingLinearMethod(quant_config)
352 + elif isinstance(layer, FusedMoE):
353 + return HummingMoEMethod(quant_config, layer.moe_config)
354 + return None
355 +
356 +
357 +class HummingLayerQuantizationConfig(HummingConfig):
358 + def __init__(
359 + self,
360 + weight_schema: "BaseWeightSchema",
361 + input_schema: "BaseInputSchema | None" = None,
362 + force_weight_schema: "HummingWeightSchema | None" = None,
363 + force_input_schema: "HummingInputSchema | None" = None,
364 + is_online_quant: bool = False,
365 + ):
366 + self.weight_schema = weight_schema
367 + if input_schema is None:
368 + input_schema = HummingInputSchema()
369 + self.input_schema = input_schema
370 + self.force_weight_schema = force_weight_schema
371 + self.force_input_schema = force_input_schema
372 + self.is_online_quant = is_online_quant
373 +
374 + @classmethod
375 + def from_config(cls, config):
376 + weight_schema = BaseWeightSchema.from_config(config)
377 + return cls(weight_schema)
378 +
379 + def get_quant_method(
380 + self, layer: torch.nn.Module, prefix: str
381 + ) -> QuantizeMethodBase | None:
382 + raise NotImplementedError
383 +
384 +
385 +class HummingLinearMethod(LinearMethodBase):
386 + def __init__(self, quant_config: HummingLayerQuantizationConfig):
387 + self.quant_config = quant_config
388 + self.weight_schema = quant_config.weight_schema
389 + self.input_schema = quant_config.input_schema
390 + self.force_weight_schema = quant_config.force_weight_schema
391 + self.force_input_schema = quant_config.force_input_schema
392 + self.is_online_quant = self.quant_config.is_online_quant
393 +
394 + def prepare_weight_loader(self, layer: torch.nn.Module, weight_loader: Callable):
395 + def new_weight_loader(
396 + param: torch.nn.Parameter,
397 + loaded_weight: torch.Tensor,
398 + shard_id: str | int | None = None,
399 + ):
400 + name = param.param_name
401 + float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
402 + is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
403 + if is_unquantized and self.is_online_quant:
404 + # online quant (fp16/bf16 -> quant_type)
405 + assert isinstance(self.weight_schema, HummingWeightSchema)
406 + f16_dtype = DataType.from_torch_dtype(layer.param_dtype)
407 + has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
408 + tensor_list = quantize_weight(
409 + weight=loaded_weight,
410 + dtype=self.weight_schema.b_dtype,
411 + scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
412 + group_size=self.weight_schema.weight_scale_group_size,
413 + has_zero_point=self.weight_schema.has_zero_point,
414 + has_global_scale=has_global_scale,
415 + is_fp_zero_point=self.weight_schema.is_fp_zero_point,
416 + pack=True,
417 + )
418 +
419 + key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
420 + for key, tensor in zip(key_list, tensor_list):
421 + if tensor is None or tensor.nelement() == 0:
422 + continue
423 + param = getattr(layer, key)
424 + param.weight_loader(param, tensor, shard_id)
425 +
426 + return None
427 + elif is_unquantized and not self.is_online_quant:
428 + # fallback to unquantized linear
429 + # some model skip some layer when quantizing model, but
430 + # don't mark the layer as unquantized.
431 + if not layer.is_fallback:
432 + layer.is_fallback = True
433 + for name, _ in list(layer.named_parameters()):
434 + if name != "bias":
435 + delattr(layer, name)
436 + delattr(layer, "locks")
437 + self.__class__ = UnquantizedLinearMethod # type: ignore
438 + tensor = torch.empty(
439 + (
440 + layer.output_partition_sizes_sum,
441 + layer.input_size_per_partition,
442 + ),
443 + dtype=layer.param_dtype,
444 + device=param.device,
445 + )
446 + extra_weight_attrs = layer.extra_weight_attrs.copy()
447 + orig_weight_loader = extra_weight_attrs.pop("weight_loader")
448 + layer.weight = ModelWeightParameter(
449 + data=tensor,
450 + input_dim=1,
451 + output_dim=0,
452 + weight_loader=orig_weight_loader,
453 + )
454 + layer.weight.tp_size = layer.tp_size
455 + layer.weight.tp_rank = layer.tp_rank
456 + set_weight_attrs(layer.weight, extra_weight_attrs)
457 +
458 + param = layer.weight
459 + if shard_id is not None:
460 + return layer.weight.weight_loader(param, loaded_weight, shard_id)
461 + return layer.weight.weight_loader(param, loaded_weight)
462 +
463 + # weight processing logic for specific quantization schema
464 + loaded_weight = self.weight_schema.process_loaded_weight(
465 + tensor=loaded_weight,
466 + name=name,
467 + )
468 + if shard_id is not None:
469 + return weight_loader(param, loaded_weight, shard_id)
470 + return weight_loader(param, loaded_weight)
471 +
472 + return new_weight_loader
473 +
474 + def create_weights(
475 + self,
476 + layer: torch.nn.Module,
477 + input_size_per_partition: int,
478 + output_partition_sizes: list[int],
479 + input_size: int,
480 + output_size: int,
481 + params_dtype: torch.dtype,
482 + **extra_weight_attrs,
483 + ):
484 + layer.is_fallback = False
485 + layer.param_dtype = params_dtype
486 + layer.input_size = input_size
487 + layer.output_size = output_size
488 + layer.input_size_per_partition = input_size_per_partition
489 + layer.output_partition_sizes_sum = sum(output_partition_sizes)
490 + layer.output_partition_sizes = output_partition_sizes
491 + layer.extra_weight_attrs = extra_weight_attrs.copy()
492 +
493 + weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
494 + new_weight_loader = self.prepare_weight_loader(layer, weight_loader)
495 + extra_weight_attrs["weight_loader"] = new_weight_loader
496 +
497 + for key in ["weight_block_size", "block_structure"]:
498 + block_size = getattr(self.weight_schema, key, None)
499 + if block_size is not None:
500 + layer.weight_block_size = block_size
501 +
502 + weight_tensor_attrs = self.weight_schema.get_tensors_attrs(
503 + shape_n=layer.output_partition_sizes_sum,
504 + shape_k=layer.input_size_per_partition,
505 + param_dtype=params_dtype,
506 + stack_size=len(layer.output_partition_sizes),
507 + )
508 +
509 + input_tensor_attrs = self.input_schema.get_tensors_attrs(
510 + shape_k=layer.input_size_per_partition,
511 + param_dtype=params_dtype,
512 + stack_size=len(layer.output_partition_sizes),
513 + )
514 +
515 + tensors_attrs = weight_tensor_attrs | input_tensor_attrs
516 +
517 + for name, attrs in tensors_attrs.items():
518 + tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
519 + extra_attrs = attrs.get("extra_attrs", {}).copy()
520 + extra_attrs.update(extra_weight_attrs)
521 + param = prepare_param(tensor, name, extra_attrs)
522 + setattr(layer, name, param)
523 +
524 + locks = torch.zeros(1024, dtype=torch.int32)
525 + layer.register_buffer("locks", locks)
526 +
527 + if self.force_input_schema is not None:
528 + self.input_schema = self.force_input_schema
529 +
530 + if not hasattr(layer, "weight"):
531 + param = prepare_param(torch.tensor(0), "weight", extra_weight_attrs)
532 + layer.weight = param
533 +
534 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
535 + if layer.is_fallback:
536 + return None
537 +
538 + # convert from checkpoint format to humming format
539 + if not isinstance(self.weight_schema, HummingWeightSchema):
540 + self.weight_schema, tensors = self.weight_schema.convert_humming(
541 + tensors=layer.state_dict(),
542 + shape_n_stacks=layer.output_partition_sizes,
543 + shape_k_stacks=[layer.input_size_per_partition],
544 + param_dtype=layer.param_dtype,
545 + )
546 +
547 + self.input_schema, _ = self.input_schema.convert_humming(
548 + tensors=layer.state_dict(),
549 + shape_n_stacks=layer.output_partition_sizes,
550 + shape_k_stacks=[layer.input_size_per_partition],
551 + param_dtype=layer.param_dtype,
552 + )
553 +
554 + for name, _ in list(layer.named_parameters()):
555 + delattr(layer, name)
556 +
557 + for name, tensor in tensors.items():
558 + param = torch.nn.Parameter(tensor, requires_grad=False)
559 + setattr(layer, name, param)
560 +
561 + del tensors
562 +
563 + # force requant (origin quant setting -> fp16/bf16 -> new_quant setting)
564 + assert isinstance(self.weight_schema, HummingWeightSchema)
565 + force_requant = self.force_weight_schema is not None
566 + if force_requant and self.weight_schema != self.force_weight_schema:
567 + tensors = self.weight_schema.requant_tensors(
568 + tensors=layer.state_dict(),
569 + target_weight_schema=self.force_weight_schema,
570 + param_dtype=layer.param_dtype,
571 + )
572 +
573 + self.weight_schema = self.force_weight_schema
574 +
575 + for name, _ in list(layer.named_parameters()):
576 + if name != "bias":
577 + delattr(layer, name)
578 +
579 + for name, tensor in tensors.items():
580 + param = torch.nn.Parameter(tensor, requires_grad=False)
581 + setattr(layer, name, param)
582 +
583 + del tensors
584 +
585 + # prepare layer config from humming kernel
586 + HummingMethod.prepare_layer_meta(
587 + layer=layer,
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 + )
597 +
598 + # preprocess weight for inference
599 + HummingMethod.transform_humming_layer(layer)
600 +
601 + # compute_config: kernel configs that do not directly affect weights
602 + # but significantly impact kernel behavior or computation precision.
603 + # see https://github.com/inclusionAI/humming/blob/main/docs/config.md
604 + compute_config = {
605 + "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
606 + "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
607 + "gemm_type": "dense",
608 + }
609 + self.compute_config = json.dumps(compute_config)
610 +
611 + def apply(
612 + self,
613 + layer: torch.nn.Module,
614 + x: torch.Tensor,
615 + bias: torch.Tensor | None = None,
616 + ) -> torch.Tensor:
617 + flatten_inputs = x.view(-1, x.size(-1))
618 + output = HummingMethod.forward_layer(
619 + layer=layer,
620 + inputs=flatten_inputs,
621 + compute_config=self.compute_config,
622 + )
623 + output = output.view(*x.shape[:-1], output.size(-1))
624 + return output
625 +
626 +
627 +class HummingMoEMethod(FusedMoEMethodBase):
628 + def __init__(
629 + self, quant_config: HummingLayerQuantizationConfig, moe: "FusedMoEConfig"
630 + ) -> None:
631 + super().__init__(moe)
632 + self.quant_config = quant_config
633 + self.moe = moe
634 + self.weight_schema = quant_config.weight_schema
635 + self.input_schema = quant_config.input_schema
636 + self.force_weight_schema = quant_config.force_weight_schema
637 + self.force_input_schema = quant_config.force_input_schema
638 +
639 + def prepare_weight_loader(self, layer, weight_loader):
640 + def new_weight_loader(
641 + param: torch.nn.Parameter,
642 + loaded_weight: torch.Tensor,
643 + weight_name: str,
644 + shard_id: str,
645 + expert_id: int | None = None,
646 + return_success: bool = False,
647 + ):
648 + name = param.param_name
649 + float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
650 + is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
651 + # online quant (fp16/bf16 -> quant_type)
652 + if is_unquantized:
653 + assert isinstance(self.weight_schema, HummingWeightSchema)
654 + f16_dtype = DataType.from_torch_dtype(layer.param_dtype)
655 + has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
656 + tensor_list = quantize_weight(
657 + weight=loaded_weight,
658 + dtype=self.weight_schema.b_dtype,
659 + scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
660 + group_size=self.weight_schema.weight_scale_group_size,
661 + has_zero_point=self.weight_schema.has_zero_point,
662 + has_global_scale=has_global_scale,
663 + is_fp_zero_point=self.weight_schema.is_fp_zero_point,
664 + pack=True,
665 + )
666 +
667 + key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
668 + success = True
669 + for key, tensor in zip(key_list, tensor_list):
670 + if tensor is None or tensor.nelement() == 0:
671 + continue
672 + sublayer_name = "w2" if shard_id == "w2" else "w13"
673 +
674 + param = getattr(layer, sublayer_name + "_" + key)
675 + part_subccess = param.weight_loader(
676 + param=param,
677 + loaded_weight=tensor.cpu(),
678 + weight_name=shard_id + "_" + key,
679 + shard_id=shard_id,
680 + expert_id=expert_id,
681 + return_success=return_success,
682 + )
683 + success = success and part_subccess
684 +
685 + return success if return_success else None
686 +
687 + # weight processing logic for specific quantization schema
688 + loaded_weight = self.weight_schema.process_loaded_weight(
689 + tensor=loaded_weight,
690 + name=name,
691 + )
692 + return weight_loader(
693 + param,
694 + loaded_weight,
695 + weight_name,
696 + shard_id=shard_id,
697 + expert_id=expert_id,
698 + return_success=return_success,
699 + )
700 +
701 + return new_weight_loader
702 +
703 + def create_weights(
704 + self,
705 + layer: torch.nn.Module,
706 + num_experts: int,
707 + hidden_size: int,
708 + intermediate_size_per_partition: int,
709 + params_dtype: torch.dtype,
710 + **extra_weight_attrs,
711 + ):
712 + layer.num_experts = num_experts
713 + layer.param_dtype = params_dtype
714 + layer.intermediate_size = intermediate_size_per_partition
715 + weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
716 + weight_loader = self.prepare_weight_loader(layer, weight_loader)
717 + extra_weight_attrs["weight_loader"] = weight_loader
718 +
719 + # sublayer: a layer contains multiple sets of weights for quantized GEMM
720 + # (e.g., weight, weight_scale, etc.).
721 + # The weight names of sublayer start with the prefix "{sublayer_name}_"
722 + layer.sublayer_configs = {
723 + "w13": {
724 + "shape_n": intermediate_size_per_partition * 2,
725 + "shape_k": hidden_size,
726 + "tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
727 + shape_n=intermediate_size_per_partition * 2,
728 + shape_k=hidden_size,
729 + num_experts=num_experts,
730 + param_dtype=params_dtype,
731 + has_bias=self.moe.has_bias,
732 + ),
733 + },
734 + "w2": {
735 + "shape_n": hidden_size,
736 + "shape_k": intermediate_size_per_partition,
737 + "tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
738 + shape_n=hidden_size,
739 + shape_k=intermediate_size_per_partition,
740 + num_experts=num_experts,
741 + param_dtype=params_dtype,
742 + has_bias=self.moe.has_bias,
743 + ),
744 + },
745 + }
746 +
747 + for sublayer_name, configs in layer.sublayer_configs.items():
748 + for name, attrs in configs["tensors_attrs"].items():
749 + tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
750 + param = torch.nn.Parameter(tensor, requires_grad=False)
751 + extra_attrs = attrs.get("extra_attrs", {}).copy()
752 + extra_attrs.update(extra_weight_attrs)
753 + param = prepare_moe_param(tensor, name, extra_attrs)
754 + setattr(layer, f"{sublayer_name}_{name}", param)
755 +
756 + if self.force_input_schema is not None:
757 + self.input_schema = self.force_input_schema
758 +
759 + locks = torch.zeros(1024, dtype=torch.int32)
760 + layer.register_buffer("locks", locks)
761 +
762 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
763 + self.process_weights_after_loading(layer)
764 +
765 + input_schema = self.input_schemas["w13"]
766 + weight_schema = self.weight_schemas["w13"]
767 +
768 + a_dtype = input_schema.a_dtype
769 + if a_dtype is None or a_dtype.num_bits == 16:
770 + a_quant_desc = FusedMoEQuantDesc(dtype=None)
771 + else:
772 + shape = GroupShape(row=1, col=-1)
773 + a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape)
774 +
775 + weight_scale_group_size = weight_schema.weight_scale_group_size
776 + weight_scale_group_size_n = weight_schema.weight_scale_group_size_n
777 + weight_group_shape: tuple[int, ...] = ()
778 + if weight_scale_group_size_n > 1:
779 + weight_group_shape = GroupShape(
780 + row=weight_scale_group_size,
781 + col=weight_scale_group_size_n,
782 + )
783 + elif weight_scale_group_size == 0:
784 + weight_group_shape = GroupShape(row=-1, col=1)
785 + else:
786 + weight_group_shape = GroupShape(row=weight_scale_group_size, col=1)
787 +
788 + w1_quant_desc = FusedMoEQuantDesc(
789 + dtype=str(weight_schema.b_dtype),
790 + shape=weight_group_shape,
791 + scale=getattr(layer, "w13_weight_scale", None),
792 + alpha_or_gscale=getattr(layer, "w13_global_scale", None),
793 + zp=getattr(layer, "w13_zero_point", None),
794 + bias=getattr(layer, "w13_bias", None),
795 + )
796 +
797 + w2_quant_desc = FusedMoEQuantDesc(
798 + dtype=str(weight_schema.b_dtype),
799 + shape=weight_group_shape,
800 + scale=getattr(layer, "w2_weight_scale", None),
801 + alpha_or_gscale=getattr(layer, "w2_global_scale", None),
802 + zp=getattr(layer, "w2_zero_point", None),
803 + bias=getattr(layer, "w2_bias", None),
804 + )
805 +
806 + return FusedMoEQuantConfig(
807 + _a1=a_quant_desc,
808 + _a2=a_quant_desc,
809 + _w1=w1_quant_desc,
810 + _w2=w2_quant_desc,
811 + )
812 +
813 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
814 + if getattr(self, "processed", False):
815 + return
816 + self.processed = True
817 + self.weight_schemas = {}
818 + self.input_schemas = {}
819 + for sublayer_name, configs in layer.sublayer_configs.items():
820 + input_schema = self.input_schema
821 + weight_schema = self.weight_schema
822 + # convert from checkpoint format to humming format
823 + if not isinstance(weight_schema, HummingWeightSchema):
824 + tensors: dict[str, torch.Tensor] = dict(
825 + (key.removeprefix(sublayer_name + "_"), value)
826 + for key, value in layer.state_dict().items()
827 + if key.startswith(sublayer_name + "_")
828 + )
829 +
830 + shape_k_stacks = [configs["shape_k"]]
831 + shape_n_stacks = [configs["shape_n"]]
832 + if sublayer_name == "w13":
833 + shape_n_stacks = [configs["shape_n"] // 2] * 2
834 +
835 + weight_schema, tensors = weight_schema.convert_humming(
836 + tensors=tensors,
837 + shape_n_stacks=shape_n_stacks,
838 + shape_k_stacks=shape_k_stacks,
839 + param_dtype=layer.param_dtype,
840 + num_experts=layer.num_experts,
841 + )
842 +
843 + input_schema, _ = input_schema.convert_humming(
844 + tensors=tensors,
845 + shape_n_stacks=shape_n_stacks,
846 + shape_k_stacks=shape_k_stacks,
847 + param_dtype=layer.param_dtype,
848 + num_experts=layer.num_experts,
849 + )
850 +
851 + for name, _ in list(layer.named_parameters()):
852 + if not name.startswith(sublayer_name + "_"):
853 + continue
854 + delattr(layer, name)
855 +
856 + for name, tensor in tensors.items():
857 + name = f"{sublayer_name}_{name}"
858 + param = torch.nn.Parameter(tensor, requires_grad=False)
859 + setattr(layer, name, param)
860 +
861 + self.weight_schemas[sublayer_name] = weight_schema
862 + self.input_schemas[sublayer_name] = input_schema
863 +
864 + # force requant (origin quant setting -> fp16/bf16 -> new_quant setting)
865 + assert isinstance(weight_schema, HummingWeightSchema)
866 + force_requant = self.force_weight_schema is not None
867 + if force_requant and weight_schema != self.force_weight_schema:
868 + tensors = dict(
869 + (key.removeprefix(sublayer_name + "_"), value)
870 + for key, value in layer.state_dict().items()
871 + if key.startswith(sublayer_name + "_")
872 + )
873 +
874 + tensors = weight_schema.requant_tensors(
875 + tensors=tensors,
876 + target_weight_schema=self.force_weight_schema,
877 + param_dtype=layer.param_dtype,
878 + )
879 +
880 + weight_schema = self.force_weight_schema
881 +
882 + for name, _ in list(layer.named_parameters()):
883 + if not name.startswith(sublayer_name + "_"):
884 + continue
885 + if name == sublayer_name + "_bias":
886 + continue
887 + delattr(layer, name)
888 +
889 + for name, tensor in tensors.items():
890 + name = f"{sublayer_name}_{name}"
891 + param = torch.nn.Parameter(tensor, requires_grad=False)
892 + setattr(layer, name, param)
893 +
894 + del tensors
895 +
896 + # prepare layer config from humming kernel
897 + HummingMethod.prepare_layer_meta(
898 + layer=layer,
899 + shape_n=configs["shape_n"],
900 + shape_k=configs["shape_k"],
901 + pad_n_to_multiple=256,
902 + pad_k_to_multiple=128,
903 + input_schema=input_schema,
904 + weight_schema=weight_schema,
905 + has_bias=self.moe.has_bias,
906 + num_experts=layer.num_experts,
907 + torch_dtype=layer.param_dtype,
908 + sublayer_name=sublayer_name,
909 + )
910 +
911 + # preprocess weight for inference
912 + HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name)
913 +
914 + # use moe modular
915 + experts: HummingIndexedExperts | HummingGroupedExperts
916 + if get_humming_moe_gemm_type() == "indexed":
917 + experts = HummingIndexedExperts(layer, self)
918 + else:
919 + experts = HummingGroupedExperts(layer, self)
920 + self.experts = experts
921 +
922 + def select_gemm_impl(
923 + self,
924 + prepare_finalize,
925 + layer: torch.nn.Module,
926 + ):
927 + from vllm.model_executor.layers.fused_moe import modular_kernel as mk
928 +
929 + activation_format = prepare_finalize.activation_format
930 + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
931 + return BatchedHummingGroupedExperts(layer, self, prepare_finalize)
932 + elif get_humming_moe_gemm_type() == "indexed":
933 + return HummingIndexedExperts(layer, self, prepare_finalize)
934 + else:
935 + return HummingGroupedExperts(layer, self, prepare_finalize)
936 +
937 + def apply(
938 + self,
939 + layer: FusedMoE,
940 + x: torch.Tensor,
941 + topk_weights: torch.Tensor,
942 + topk_ids: torch.Tensor,
943 + shared_experts_input: torch.Tensor | None,
944 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
945 + workspace1, workspace2, output = self.experts.make_workspaces(
946 + M=topk_ids.size(0),
947 + topk=topk_ids.size(1),
948 + activation=layer.activation,
949 + )
950 +
951 + assert workspace1.data_ptr() == output.data_ptr()
952 +
953 + self.experts.main_apply(
954 + hidden_states=x,
955 + topk_weights=topk_weights,
956 + topk_ids=topk_ids,
957 + workspace1=workspace1,
958 + workspace2=workspace2,
959 + expert_tokens_meta=None,
960 + )
961 +
962 + return output
vllm/model_executor/layers/quantization/utils/humming_moe_utils.py new
+35
@@ -0,0 +1,35 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
6 + moe_align_block_size,
7 +)
8 +
9 +
10 +def humming_moe_align(
11 + configs: list[int],
12 + topk_ids: torch.Tensor,
13 + num_experts: int,
14 + expert_map: torch.Tensor | None = None,
15 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
16 + assert len(configs) > 0 and len(configs) % 3 == 0
17 + # NOTE: we choose moe_block_size based on
18 + # num_tokens * top_k (= topk_ids.nelement())
19 + shape_m = topk_ids.nelement()
20 +
21 + for i in range(len(configs) // 3):
22 + if shape_m > configs[i * 3] and shape_m <= configs[i * 3 + 1]:
23 + block_size = configs[i * 3 + 2]
24 + break
25 + else:
26 + raise ValueError(f"Could not find a matching block_size for shape_m={shape_m}")
27 +
28 + return moe_align_block_size(
29 + topk_ids=topk_ids,
30 + block_size=block_size,
31 + num_experts=num_experts,
32 + expert_map=expert_map,
33 + pad_sorted_ids=False,
34 + ignore_invalid_experts=True,
35 + )
vllm/model_executor/parameter.py
+2 -2
@@ -605,8 +605,8 @@ def _adjust_shard_indexes_for_marlin(shard_size, shard_offset, marlin_tile_size)
605 def _adjust_shard_indexes_for_packing(
606 shard_size, shard_offset, packed_factor, marlin_tile_size
607 ):
608 - shard_size = shard_size // packed_factor
609 - shard_offset = shard_offset // packed_factor
608 + shard_size = round(shard_size // packed_factor)
609 + shard_offset = round(shard_offset // packed_factor)
610 if marlin_tile_size is not None:
611 return _adjust_shard_indexes_for_marlin(
612 shard_size=shard_size,