@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 08834cc3ce

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, 08:45 UTC

[Quantization] add humming mxfp4 moe backend (#41083)

Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com>

Jinzhen Lin committed May 3, 2026 at 09:36 UTC 08834cc3ceb86d77231666a5e07847f741637355

9 files changed +384 -157

vllm/config/kernel.py

@@ -118,6 +118,7 @@ MoEBackend = Literal[

118 "flashinfer_cutlass",

119 "flashinfer_cutedsl",

120 "marlin",

121 + "humming",

122 "aiter",

123 "emulation",

124 ]

@@ -148,6 +149,7 @@ class KernelConfig:

149 - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels

150 - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)

151 - "marlin": Use Marlin kernels (weight-only quantization)

152 + - "humming": Use Humming Mixed Precision kernels

153 - "aiter": Use AMD AITer kernels (ROCm only)

154 - "emulation": use BF16/FP16 GEMM, dequantizing weights and

155 running QDQ on activations.

vllm/envs.py

+2 -2

@@ -1231,8 +1231,8 @@ environment_variables: dict[str, Callable[[], Any]] = {

1231 # if 1, force use indexed gemm

1232 # if 0, force use grouped gemm

1233 # if None, choose better gemm type automatically

1234 - "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: maybe_convert_bool(

1235 - os.environ.get("VLLM_HUMMING_MOE_GEMM_TYPE", None)

1234 + "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: os.environ.get(

1235 + "VLLM_HUMMING_MOE_GEMM_TYPE", None

1236 ),

1237 # Whether to use DeepEPLL kernels for NVFP4 quantization and dispatch method

1238 # only supported on Blackwell GPUs and with

vllm/model_executor/layers/fused_moe/fused_humming_moe.py

+70 -49

@@ -4,7 +4,7 @@

5 import json

6 import math

7 -from typing import TYPE_CHECKING, Any

7 +from typing import Any

9 import torch

10 from humming import dtypes

@@ -16,7 +16,11 @@ 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

19 +from vllm.model_executor.layers.fused_moe.config import (

20 + FusedMoEConfig,

21 + FusedMoEParallelConfig,

22 + FusedMoEQuantConfig,

23 +)

24 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

25 moe_align_block_size,

26 )

@@ -34,21 +38,16 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

38 from vllm.platforms import current_platform

39 from vllm.v1.worker.workspace import current_workspace_manager

37 -if TYPE_CHECKING:

38 - from vllm.model_executor.layers.quantization.humming import HummingMoEMethod

39 -

40 -

41 logger = init_logger(__name__)

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"]:

47 + if env_gemm_type == "indexed":

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"

49 + elif env_gemm_type in ["grouped_contiguous", "grouped"]:

50 + gemm_type = "grouped_contiguous"

51 else:

52 gemm_type = "indexed"

@@ -60,49 +59,44 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

59 def __init__(

60 self,

61 layer: torch.nn.Module,

63 - quant_method: "HummingMoEMethod",

64 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular | None = None,

62 + moe_config: FusedMoEConfig,

63 + quant_config: FusedMoEQuantConfig,

64 + max_num_tokens: int | None = None,

65 + num_dispatchers: int | None = None,

66 ):

67 self.layer = layer

68 self.num_experts = self.layer.num_experts

69 self.global_num_experts = self.layer.global_num_experts

70 self.init_humming_moe()

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

72 + if self.is_batched():

73 + assert max_num_tokens is not None and num_dispatchers is not None

74 +

75 + super().__init__(

76 + moe_config=moe_config,

77 + quant_config=quant_config,

78 + max_num_tokens=max_num_tokens,

79 + num_dispatchers=num_dispatchers,

80 + )

82 def init_humming_moe(self):

83 self.compute_config = {

84 "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,

85 "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,

92 - "gemm_type": self.humming_gemm_type.value,

86 + "gemm_type": self.humming_gemm_type().value,

87 }

88 self.w13_tuning_config = HummingMethod.get_default_tuning_configs(

89 layer=self.layer,

90 use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,

91 use_batch_invariant=envs.VLLM_BATCH_INVARIANT,

98 - gemm_type=self.humming_gemm_type,

92 + gemm_type=self.humming_gemm_type(),

93 sublayer_name="w13",

94 )

95 self.w2_tuning_config = HummingMethod.get_default_tuning_configs(

96 layer=self.layer,

97 use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,

98 use_batch_invariant=envs.VLLM_BATCH_INVARIANT,

105 - gemm_type=self.humming_gemm_type,

99 + gemm_type=self.humming_gemm_type(),

100 sublayer_name="w2",

101 )

102 self.compute_config_str = json.dumps(self.compute_config)

@@ -124,13 +118,13 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

118 global_num_experts = self.global_num_experts

119 return math.ceil(global_valid_shape_m * num_experts / global_num_experts)

120

127 - @property

128 - def humming_gemm_type(self) -> HummingGemmType:

121 + @staticmethod

122 + def humming_gemm_type() -> HummingGemmType:

123 raise NotImplementedError

124

131 - @property

132 - def is_batched(self) -> bool:

133 - return self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts

125 + @classmethod

126 + def is_batched(cls) -> bool:

127 + return cls.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts

128

129 @staticmethod

130 def _supports_quant_scheme(

@@ -189,7 +183,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

183 assert w1.size(0) == num_experts

184 assert w2.size(0) == num_experts

185

192 - if not self.is_batched:

186 + if not self.is_batched():

187 num_tokens = a1.size(0)

188 assert topk_ids.size(0) == num_tokens

189 else:

@@ -201,7 +195,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

195

196 def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):

197 num_experts = self.num_experts

204 - N = self.layer.intermediate_size

198 + N = self.layer.intermediate_size_per_partition

199 K = self.layer.hidden_size

200 assert isinstance(num_experts, int)

201 assert isinstance(N, int)

@@ -218,7 +212,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

212 # The output must be derived from workspace1.

213

214 output_shape: tuple[int, ...]

221 - if self.is_batched:

215 + if self.is_batched():

216 max_num_tokens = self.max_num_tokens

217 num_dispatchers = self.num_dispatchers

218 assert max_num_tokens is not None and num_dispatchers is not None

@@ -227,7 +221,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

221 output_shape = (num_experts, max_num_tokens * num_dispatchers, K)

222 else:

223 input_shape_m = M

230 - if self.humming_gemm_type != HummingGemmType.INDEXED:

224 + if self.humming_gemm_type() != HummingGemmType.INDEXED:

225 input_shape_m = M * topk

226 real_shape_m = M * topk

227 output_shape = (M, K)

@@ -262,7 +256,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

256 "dtype": torch_dtype_map[a_dtype],

257 },

258 "down_output": {

265 - "shape": output_shape if self.is_batched else (real_shape_m, K),

259 + "shape": output_shape if self.is_batched() else (real_shape_m, K),

260 "dtype": torch_dtype_map[c_dtype],

261 },

262 "output": {

@@ -288,7 +282,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

282 ]

283

284 # batched moe use down_output as output

291 - if not self.is_batched:

285 + if not self.is_batched():

286 required_buffers.append("output")

287

288 return buffer_metas, required_buffers

@@ -308,7 +302,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

302 else:

303 workspace2_nbytes = max(workspace2_nbytes, nbytes)

304

311 - output_key = "down_output" if self.is_batched else "output"

305 + output_key = "down_output" if self.is_batched() else "output"

306 output_shape = buffer_metas[output_key]["shape"]

307

308 return (workspace1_nbytes // 2,), (workspace2_nbytes // 2,), output_shape

@@ -395,6 +389,33 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

389 ):

390 raise NotImplementedError

391

392 + @staticmethod

393 + def is_supported_config(

394 + cls: type[mk.FusedMoEExperts],

395 + moe_config: FusedMoEConfig,

396 + weight_key: QuantKey | None,

397 + activation_key: QuantKey | None,

398 + activation_format: mk.FusedMoEActivationFormat,

399 + ) -> tuple[bool, str | None]:

400 + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

401 + supported = cls.activation_format() == activation_format

402 + reason = "activation_format mismatched"

403 + elif activation_format == mk.FusedMoEActivationFormat.Standard:

404 + if cls.activation_format() != mk.FusedMoEActivationFormat.Standard:

405 + supported = False

406 + reason = "activation_format mismatched"

407 + else:

408 + assert hasattr(cls, "humming_gemm_type")

409 + gemm_type = cls.humming_gemm_type().value.lower()

410 + preferred_gemm_type = get_humming_moe_gemm_type().lower()

411 + supported = preferred_gemm_type == gemm_type

412 + reason = "preferred gemm type mismatched"

413 + else:

414 + supported = False

415 + reason = "unsupported activation_format"

416 +

417 + return supported, None if supported else reason

418 +

419

420 class HummingIndexedExperts(HummingExpertsBase):

421 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

@@ -404,8 +425,8 @@ class HummingIndexedExperts(HummingExpertsBase):

425 def activation_format() -> mk.FusedMoEActivationFormat:

426 return mk.FusedMoEActivationFormat.Standard

427

407 - @property

408 - def humming_gemm_type(self) -> HummingGemmType:

428 + @staticmethod

429 + def humming_gemm_type() -> HummingGemmType:

430 return HummingGemmType.INDEXED

431

432 def prepare_humming_moe_kwargs(

@@ -526,8 +547,8 @@ class HummingGroupedExperts(HummingExpertsBase):

547 def activation_format() -> mk.FusedMoEActivationFormat:

548 return mk.FusedMoEActivationFormat.Standard

549

529 - @property

530 - def humming_gemm_type(self) -> HummingGemmType:

550 + @staticmethod

551 + def humming_gemm_type() -> HummingGemmType:

552 return HummingGemmType.GROUPED_CONTIGUOUS

553

554 def main_apply(

@@ -619,8 +640,8 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):

640 def activation_format() -> mk.FusedMoEActivationFormat:

641 return mk.FusedMoEActivationFormat.BatchedExperts

642

622 - @property

623 - def humming_gemm_type(self) -> HummingGemmType:

643 + @staticmethod

644 + def humming_gemm_type() -> HummingGemmType:

645 return HummingGemmType.GROUPED_MASKED

646

647 def main_apply(

vllm/model_executor/layers/fused_moe/layer.py

-3

@@ -1103,9 +1103,6 @@ class FusedMoE(PluggableLayer):

1103 return_success: bool = False,

1104 ) -> bool | None:

1105 quant_config_name = self.quant_config and self.quant_config.get_name()

1106 - if quant_config_name == "humming":

1107 - assert hasattr(self.quant_method, "weight_schema")

1108 - quant_config_name = self.quant_method.weight_schema.quant_method

1106 if quant_config_name == "gpt_oss_mxfp4":

1107 # (FIXME) for gpt-oss all experts are combined

1108 if "bias" in weight_name:

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+63 -1

@@ -68,6 +68,8 @@ class Mxfp4MoeBackend(Enum):

68 XPU = "XPU"

69 # Emulation

70 EMULATION = "EMULATION"

71 + # Humming

72 + HUMMING = "HUMMING"

75 # Backends that share the same TRTLLM weight format

@@ -130,6 +132,19 @@ def backend_to_kernel_cls(

132

133 return [UnfusedOAITritonExperts]

134

135 + elif backend == Mxfp4MoeBackend.HUMMING:

136 + from vllm.model_executor.layers.fused_moe.fused_humming_moe import (

137 + BatchedHummingGroupedExperts,

138 + HummingGroupedExperts,

139 + HummingIndexedExperts,

140 + )

141 +

142 + return [

143 + BatchedHummingGroupedExperts,

144 + HummingGroupedExperts,

145 + HummingIndexedExperts,

146 + ]

147 +

148 elif backend == Mxfp4MoeBackend.MARLIN:

149 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

150 MarlinExperts,

@@ -177,6 +192,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:

192 "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,

193 "triton": Mxfp4MoeBackend.TRITON,

194 "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED,

195 + "humming": Mxfp4MoeBackend.HUMMING,

196 "marlin": Mxfp4MoeBackend.MARLIN,

197 "aiter": Mxfp4MoeBackend.AITER,

198 "xpu": Mxfp4MoeBackend.XPU,

@@ -573,7 +589,21 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

589

590 sf_block_size = 32 # mxfp4 block size

591

576 - if mxfp4_backend in (

592 + if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

593 + from vllm.model_executor.layers.quantization.utils.humming_utils import (

594 + prepare_humming_moe_layer,

595 + )

596 +

597 + prepare_humming_moe_layer(layer, {"quant_method": "gpt_oss_mxfp4"})

598 + return (

599 + layer.w13_weight,

600 + layer.w2_weight,

601 + layer.w13_weight_scale,

602 + layer.w2_weight_scale,

603 + getattr(layer, "w13_bias", None),

604 + getattr(layer, "w2_bias", None),

605 + )

606 + elif mxfp4_backend in (

607 Mxfp4MoeBackend.MARLIN,

608 Mxfp4MoeBackend.BATCHED_MARLIN,

609 ):

@@ -970,6 +1000,21 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1000 w2_bias,

1001 )

1002

1003 + if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

1004 + from vllm.model_executor.layers.quantization.utils.humming_utils import (

1005 + prepare_humming_moe_layer,

1006 + )

1007 +

1008 + prepare_humming_moe_layer(layer, {"quant_method": "mxfp4"})

1009 + return (

1010 + layer.w13_weight,

1011 + layer.w2_weight,

1012 + layer.w13_weight_scale,

1013 + layer.w2_weight_scale,

1014 + getattr(layer, "w13_bias", None),

1015 + getattr(layer, "w2_bias", None),

1016 + )

1017 +

1018 if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):

1019 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

1020 prepare_moe_mxfp4_layer_for_marlin,

@@ -1175,6 +1220,7 @@ def make_mxfp4_moe_quant_config(

1220 swiglu_limit: float | None = None,

1221 w1_bias: torch.Tensor | None = None,

1222 w2_bias: torch.Tensor | None = None,

1223 + layer: torch.nn.Module | None = None,

1224 ) -> FusedMoEQuantConfig | None:

1225 """Create a FusedMoEQuantConfig for the given MXFP4 backend."""

1226 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:

@@ -1234,6 +1280,14 @@ def make_mxfp4_moe_quant_config(

1280 gemm1_beta=gemm1_beta,

1281 gemm1_clamp_limit=swiglu_limit,

1282 )

1283 + elif mxfp4_backend == Mxfp4MoeBackend.HUMMING:

1284 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE

1285 + from vllm.model_executor.layers.quantization.utils.humming_utils import (

1286 + get_humming_moe_quant_config,

1287 + )

1288 +

1289 + assert isinstance(layer, FusedMoE)

1290 + return get_humming_moe_quant_config(layer)

1291 else:

1292 return ocp_mx_moe_quant_config(

1293 quant_dtype="mxfp4",

@@ -1254,6 +1308,7 @@ def make_mxfp4_moe_kernel(

1308 mxfp4_backend: Mxfp4MoeBackend,

1309 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1310 shared_experts: torch.nn.Module | None = None,

1311 + layer: torch.nn.Module | None = None,

1312 ) -> mk.FusedMoEKernel:

1313 """Create a FusedMoEKernel for the given MXFP4 backend."""

1314 is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)

@@ -1269,6 +1324,11 @@ def make_mxfp4_moe_kernel(

1324

1325 logger.info_once("Using %s", prepare_finalize.__class__.__name__)

1326

1327 + extra_kwargs = {}

1328 + if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

1329 + assert layer is not None

1330 + extra_kwargs["layer"] = layer

1331 +

1332 # Create Experts.

1333 if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

1334 max_num_tokens = prepare_finalize.max_num_tokens_per_rank()

@@ -1278,11 +1338,13 @@ def make_mxfp4_moe_kernel(

1338 quant_config=moe_quant_config,

1339 max_num_tokens=max_num_tokens,

1340 num_dispatchers=prepare_finalize.num_dispatchers(),

1341 + **extra_kwargs,

1342 )

1343 else:

1344 experts = experts_cls(

1345 moe_config=moe_config,

1346 quant_config=moe_quant_config,

1347 + **extra_kwargs,

1348 )

1349

1350 kernel = mk.FusedMoEKernel(

vllm/model_executor/layers/quantization/humming.py

+29 -67

@@ -9,11 +9,9 @@ import regex as re

9 import torch

11 from vllm import envs

12 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation

12 from vllm.model_executor.layers.fused_moe.config import (

13 FusedMoEConfig,

14 FusedMoEQuantConfig,

16 - FusedMoEQuantDesc,

15 )

16 from vllm.model_executor.layers.fused_moe.layer import (

17 FusedMoE,

@@ -32,7 +30,6 @@ from vllm.model_executor.layers.quantization.base_config import (

30 QuantizationConfig,

31 QuantizeMethodBase,

32 )

35 -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

33 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

34 from vllm.model_executor.parameter import (

35 BasevLLMParameter,

@@ -215,6 +212,15 @@ class HummingConfig(QuantizationConfig):

212 def override_quantization_method(

213 cls, hf_quant_cfg, user_quant, hf_config=None

214 ) -> QuantizationMethods | None:

215 + if user_quant == "humming" and hf_config is not None:

216 + model_type = hf_config.model_type

217 + quant_method = hf_quant_cfg.get("quant_method", None)

218 + if model_type == "gpt_oss" and quant_method == "mxfp4":

219 + msg = (

220 + "For gpt-oss model, use '--moe-backend humming' "

221 + "instead of '--quantization humming'."

222 + )

223 + raise ValueError(msg)

224 return "humming" if user_quant == "humming" else None

225

226 def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

@@ -299,8 +305,6 @@ class HummingConfig(QuantizationConfig):

305 force_weight_schema = schema

306

307 if weight_schema is not None:

302 - if weight_schema.quant_method == "gpt_oss_mxfp4" and layer_type != "moe":

303 - return None

308 input_schema = None

309 force_input_schema = None

310

@@ -335,12 +339,6 @@ class HummingConfig(QuantizationConfig):

339 elif isinstance(layer, LinearBase):

340 layer_type = "linear"

341

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 -

342 quant_config = self.get_quant_config_for_layer(prefix, layer_type)

343 if quant_config is None:

344 if isinstance(layer, FusedMoE):

@@ -760,62 +758,18 @@ class HummingMoEMethod(FusedMoEMethodBase):

758 layer.register_buffer("locks", locks)

759

760 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),

761 + from vllm.model_executor.layers.quantization.utils.humming_utils import (

762 + get_humming_moe_quant_config,

763 )

764

806 - return FusedMoEQuantConfig(

807 - _a1=a_quant_desc,

808 - _a2=a_quant_desc,

809 - _w1=w1_quant_desc,

810 - _w2=w2_quant_desc,

811 - )

765 + return get_humming_moe_quant_config(layer)

766

767 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

768 if getattr(self, "processed", False):

769 return

770 self.processed = True

817 - self.weight_schemas = {}

818 - self.input_schemas = {}

771 + layer.weight_schemas = {}

772 + layer.input_schemas = {}

773 for sublayer_name, configs in layer.sublayer_configs.items():

774 input_schema = self.input_schema

775 weight_schema = self.weight_schema

@@ -858,8 +812,8 @@ class HummingMoEMethod(FusedMoEMethodBase):

812 param = torch.nn.Parameter(tensor, requires_grad=False)

813 setattr(layer, name, param)

814

861 - self.weight_schemas[sublayer_name] = weight_schema

862 - self.input_schemas[sublayer_name] = input_schema

815 + layer.weight_schemas[sublayer_name] = weight_schema

816 + layer.input_schemas[sublayer_name] = input_schema

817

818 # force requant (origin quant setting -> fp16/bf16 -> new_quant setting)

819 assert isinstance(weight_schema, HummingWeightSchema)

@@ -913,10 +867,11 @@ class HummingMoEMethod(FusedMoEMethodBase):

867

868 # use moe modular

869 experts: HummingIndexedExperts | HummingGroupedExperts

870 + assert self.moe_quant_config is not None

871 if get_humming_moe_gemm_type() == "indexed":

917 - experts = HummingIndexedExperts(layer, self)

872 + experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config)

873 else:

919 - experts = HummingGroupedExperts(layer, self)

874 + experts = HummingGroupedExperts(layer, self.moe, self.moe_quant_config)

875 self.experts = experts

876

877 def select_gemm_impl(

@@ -927,12 +882,19 @@ class HummingMoEMethod(FusedMoEMethodBase):

882 from vllm.model_executor.layers.fused_moe import modular_kernel as mk

883

884 activation_format = prepare_finalize.activation_format

885 + assert self.moe_quant_config is not None

886 if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

931 - return BatchedHummingGroupedExperts(layer, self, prepare_finalize)

887 + return BatchedHummingGroupedExperts(

888 + layer=layer,

889 + moe_config=self.moe,

890 + quant_config=self.moe_quant_config,

891 + max_num_tokens=prepare_finalize.max_num_tokens_per_rank(),

892 + num_dispatchers=prepare_finalize.num_dispatchers(),

893 + )

894 elif get_humming_moe_gemm_type() == "indexed":

933 - return HummingIndexedExperts(layer, self, prepare_finalize)

895 + return HummingIndexedExperts(layer, self.moe, self.moe_quant_config)

896 else:

935 - return HummingGroupedExperts(layer, self, prepare_finalize)

897 + return HummingGroupedExperts(layer, self.moe, self.moe_quant_config)

898

899 def apply(

900 self,

vllm/model_executor/layers/quantization/mxfp4.py

@@ -366,6 +366,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

366 experts_cls=self.experts_cls,

367 routing_tables=layer._maybe_init_expert_routing_tables(),

368 shared_experts=layer.shared_experts,

369 + layer=layer,

370 )

371

372 def process_weights_after_loading(self, layer):

@@ -404,6 +405,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

405 gemm1_alpha=1.702,

406 gemm1_beta=1.0,

407 swiglu_limit=7.0,

408 + layer=layer,

409 )

410

411 def select_gemm_impl(

@@ -692,6 +694,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

694 experts_cls=self.experts_cls,

695 routing_tables=layer._maybe_init_expert_routing_tables(),

696 shared_experts=layer.shared_experts,

697 + layer=layer,

698 )

699

700 def process_weights_after_loading(self, layer):

@@ -729,6 +732,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

732 w1_bias=w1_bias,

733 w2_bias=w2_bias,

734 swiglu_limit=swiglu_limit,

735 + layer=layer,

736 )

737

738 def select_gemm_impl(

vllm/model_executor/layers/quantization/utils/humming_moe_utils.py deleted

-35

@@ -1,35 +0,0 @@

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/layers/quantization/utils/humming_utils.py new

+214

@@ -0,0 +1,214 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any

4 +

5 +import regex as re

6 +import torch

7 +from humming.layer import HummingInputSchema, HummingMethod

8 +from humming.schema import BaseWeightSchema

9 +

10 +from vllm import envs

11 +from vllm.model_executor.layers.fused_moe.config import (

12 + FusedMoEQuantConfig,

13 + FusedMoEQuantDesc,

14 +)

15 +from vllm.model_executor.layers.fused_moe.layer import FusedMoE

16 +from vllm.model_executor.layers.linear import LinearBase

17 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

18 +

19 +

20 +def humming_is_layer_skipped(config: dict[str, Any], prefix: str):

21 + if not config:

22 + return True

23 +

24 + keys = ["ignored_layers", "ignore", "modules_to_not_convert"]

25 + ignored_layers: list[str] = []

26 + for key in keys:

27 + ignored_layers = config.get(key, []) or []

28 + if not ignored_layers:

29 + break

30 +

31 + if any(module_name in prefix for module_name in ignored_layers):

32 + return True

33 + if "lm_head" in prefix:

34 + return True

35 +

36 + for regex in config.get("dynamic", {}):

37 + if regex[:1] != "-":

38 + continue

39 + if re.match(regex[2:], prefix):

40 + return True

41 +

42 + return False

43 +

44 +

45 +def prepare_humming_layer(layer: LinearBase, quant_config: dict):

46 + weight_schema = BaseWeightSchema.from_config(quant_config)

47 + input_schema = HummingInputSchema()

48 +

49 + shape_k_stacks = [layer.input_size_per_partition]

50 + shape_n_stacks = layer.output_partition_sizes

51 +

52 + # Step 1: convert weight to humming standard format

53 + weight_schema, tensors = weight_schema.convert_humming(

54 + tensors=layer.named_parameters(),

55 + shape_n_stacks=shape_n_stacks,

56 + shape_k_stacks=shape_k_stacks,

57 + param_dtype=layer.params_dtype,

58 + )

59 +

60 + layer.weight_schema = weight_schema

61 +

62 + for name, _ in list(layer.named_parameters()):

63 + delattr(layer, name)

64 +

65 + for name, tensor in tensors.items():

66 + param = torch.nn.Parameter(tensor, requires_grad=False)

67 + setattr(layer, name, param)

68 +

69 + # Step 2: transform weight (humming standard format) for forwarding

70 + HummingMethod.prepare_layer_meta(

71 + layer=layer,

72 + shape_n=layer.output_partition_sizes_sum,

73 + shape_k=layer.input_size_per_partition,

74 + weight_schema=weight_schema,

75 + input_schema=input_schema,

76 + pad_n_to_multiple=256,

77 + pad_k_to_multiple=128,

78 + has_bias=layer.has_bias,

79 + torch_dtype=layer.param_dtype,

80 + )

81 +

82 + HummingMethod.transform_humming_layer(layer)

83 +

84 +

85 +def prepare_humming_moe_layer(layer: FusedMoE, quant_config: dict):

86 + weight_schema = BaseWeightSchema.from_config(quant_config)

87 + input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}

88 + if humming_is_layer_skipped(input_quant_config, layer.layer_name):

89 + input_schema = HummingInputSchema()

90 + else:

91 + # TODO: read input_quant_config from quant_config

92 + input_schema = HummingInputSchema.from_config(input_quant_config)

93 +

94 + is_gated = layer.activation.is_gated

95 + shape_config = {

96 + "w13": (

97 + layer.moe_config.intermediate_size_per_partition * 2,

98 + layer.moe_config.hidden_dim,

99 + ),

100 + "w2": (

101 + layer.moe_config.hidden_dim,

102 + layer.moe_config.intermediate_size_per_partition * (1 if is_gated else 2),

103 + ),

104 + }

105 +

106 + layer.weight_schemas = {}

107 + layer.input_schemas = {}

108 +

109 + for sublayer_name in shape_config:

110 + # Step 1: convert weight to humming standard format

111 + tensors: dict[str, torch.Tensor] = dict(

112 + (key.removeprefix(sublayer_name + "_"), value)

113 + for key, value in layer.state_dict().items()

114 + if key.startswith(sublayer_name + "_")

115 + )

116 +

117 + shape_n, shape_k = shape_config[sublayer_name]

118 + shape_n_stacks = [shape_n]

119 + shape_k_stacks = [shape_k]

120 + if sublayer_name == "w13":

121 + shape_n_stacks = [shape_n // 2] * 2

122 +

123 + weight_schema_new, tensors = weight_schema.convert_humming(

124 + tensors=tensors,

125 + shape_n_stacks=shape_n_stacks,

126 + shape_k_stacks=shape_k_stacks,

127 + num_experts=layer.local_num_experts,

128 + param_dtype=layer.params_dtype,

129 + )

130 +

131 + layer.weight_schemas[sublayer_name] = weight_schema_new

132 + layer.input_schemas[sublayer_name] = input_schema

133 +

134 + for name, _ in list(layer.named_parameters()):

135 + if not name.startswith(sublayer_name + "_"):

136 + continue

137 + delattr(layer, name)

138 +

139 + for name, tensor in tensors.items():

140 + name = f"{sublayer_name}_{name}"

141 + param = torch.nn.Parameter(tensor, requires_grad=False)

142 + setattr(layer, name, param)

143 +

144 + # Step 2: transform weight (humming standard format) for forwarding

145 + HummingMethod.prepare_layer_meta(

146 + layer=layer,

147 + shape_n=shape_n,

148 + shape_k=shape_k,

149 + pad_n_to_multiple=256,

150 + pad_k_to_multiple=128,

151 + input_schema=input_schema,

152 + weight_schema=weight_schema_new,

153 + has_bias=layer.moe_config.has_bias,

154 + num_experts=layer.num_experts,

155 + torch_dtype=layer.params_dtype,

156 + sublayer_name=sublayer_name,

157 + )

158 +

159 + HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name)

160 +

161 + if not hasattr(layer, "locks"):

162 + device = layer.w13_weight.device

163 + locks = torch.zeros(1024, dtype=torch.int32, device=device)

164 + layer.register_buffer("locks", locks)

165 +

166 +

167 +def get_humming_moe_quant_config(layer: FusedMoE):

168 + input_schema = layer.input_schemas["w13"]

169 + weight_schema = layer.weight_schemas["w13"]

170 +

171 + a_dtype = input_schema.a_dtype

172 + if a_dtype is None or a_dtype.num_bits == 16:

173 + a_quant_desc = FusedMoEQuantDesc(dtype=None)

174 + else:

175 + shape = GroupShape(row=1, col=-1)

176 + a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape)

177 +

178 + weight_scale_group_size = weight_schema.weight_scale_group_size

179 + weight_scale_group_size_n = weight_schema.weight_scale_group_size_n

180 + weight_group_shape: tuple[int, ...] = ()

181 + if weight_scale_group_size_n > 1:

182 + weight_group_shape = GroupShape(

183 + row=weight_scale_group_size,

184 + col=weight_scale_group_size_n,

185 + )

186 + elif weight_scale_group_size == 0:

187 + weight_group_shape = GroupShape(row=-1, col=1)

188 + else:

189 + weight_group_shape = GroupShape(row=weight_scale_group_size, col=1)

190 +

191 + w1_quant_desc = FusedMoEQuantDesc(

192 + dtype=str(weight_schema.b_dtype),

193 + shape=weight_group_shape,

194 + scale=getattr(layer, "w13_weight_scale", None),

195 + alpha_or_gscale=getattr(layer, "w13_global_scale", None),

196 + zp=getattr(layer, "w13_zero_point", None),

197 + bias=getattr(layer, "w13_bias", None),

198 + )

199 +

200 + w2_quant_desc = FusedMoEQuantDesc(

201 + dtype=str(weight_schema.b_dtype),

202 + shape=weight_group_shape,

203 + scale=getattr(layer, "w2_weight_scale", None),

204 + alpha_or_gscale=getattr(layer, "w2_global_scale", None),

205 + zp=getattr(layer, "w2_zero_point", None),

206 + bias=getattr(layer, "w2_bias", None),

207 + )

208 +

209 + return FusedMoEQuantConfig(

210 + _a1=a_quant_desc,

211 + _a2=a_quant_desc,

212 + _w1=w1_quant_desc,

213 + _w2=w2_quant_desc,

214 + )