@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bc8481af09

[MoE Refactor] Standardize Humming MoE experts + utilities (#43373)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed Jun 29, 2026 at 09:19 UTC bc8481af09cd4c7f7272ba7bc1913f1051649813

8 files changed +1178 -392

vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

+221 -104

@@ -36,28 +36,42 @@ from vllm.model_executor.layers.fused_moe.utils import (

36 _resize_cache,

37 swiglu_limit_func,

38 )

39 -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

39 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

40 + QuantKey,

41 + kFp8DynamicTokenSym,

42 + kFp8Static128BlockSym,

43 + kFp8StaticChannelSym,

44 + kInt4Static,

45 + kInt8Static,

46 + kMxfp4Dynamic,

47 + kMxfp4Static,

48 + kMxfp8Dynamic,

49 + kMxfp8Static,

50 + kNvfp4Static,

51 +)

52 from vllm.platforms import current_platform

41 -from vllm.utils.humming import GemmType as HummingGemmType

42 -from vllm.utils.humming import HummingLayerMeta, HummingMethod, dtypes

53 +from vllm.utils.import_utils import has_humming

54 from vllm.v1.worker.workspace import current_workspace_manager

56 if TYPE_CHECKING:

57 from vllm.model_executor.layers.fused_moe import RoutedExperts

58 + from vllm.utils.humming import GemmType as HummingGemmType

61 logger = init_logger(__name__)

52 -def get_humming_moe_gemm_type() -> str:

53 - env_gemm_type: str = envs.VLLM_HUMMING_MOE_GEMM_TYPE or ""

54 - env_gemm_type = env_gemm_type.lower()

55 - if env_gemm_type == "indexed":

56 - gemm_type = env_gemm_type

57 - elif env_gemm_type in ["grouped_contiguous", "grouped"]:

58 - gemm_type = "grouped_contiguous"

59 - else:

60 - gemm_type = "indexed"

64 +def get_humming_moe_gemm_type() -> str | None:

65 + env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE

66 + gemm_type = None

67 + if env_gemm_type is not None:

68 + env_gemm_type = env_gemm_type.lower()

69 + if env_gemm_type == "indexed":

70 + gemm_type = env_gemm_type

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

72 + gemm_type = "grouped_contiguous"

73 + else:

74 + gemm_type = "indexed"

76 logger.info_once(f"Using {gemm_type} gemm for humming moe") # noqa

77 return gemm_type

@@ -89,6 +103,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

103 self._permute_scratch: MoEPermuteScratch | None = None

104

105 def init_humming_moe(self):

106 + from vllm.utils.humming import HummingMethod

107 +

108 self.compute_config = {

109 "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,

110 "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,

@@ -141,7 +157,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

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

158

159 @staticmethod

144 - def humming_gemm_type() -> HummingGemmType:

160 + def humming_gemm_type() -> "HummingGemmType":

161 raise NotImplementedError

162

163 @classmethod

@@ -153,12 +169,51 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

169 weight_key: QuantKey | None,

170 activation_key: QuantKey | None,

171 ) -> bool:

172 + SUPPORTED_W_A = [

173 + (kMxfp4Static, None),

174 + (kMxfp4Static, kMxfp4Dynamic),

175 + (kMxfp4Static, kMxfp8Dynamic),

176 + (kMxfp4Static, kFp8DynamicTokenSym),

177 + (kNvfp4Static, None),

178 + (kNvfp4Static, kFp8DynamicTokenSym),

179 + (kMxfp8Static, None),

180 + (kMxfp8Static, kFp8DynamicTokenSym),

181 + (kFp8StaticChannelSym, None),

182 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),

183 + (kFp8Static128BlockSym, None),

184 + (kFp8Static128BlockSym, kFp8DynamicTokenSym),

185 + (kInt4Static, None),

186 + (kInt4Static, kFp8DynamicTokenSym),

187 + (kInt8Static, None),

188 + (kInt8Static, kFp8DynamicTokenSym),

189 + ]

190 + return (weight_key, activation_key) in SUPPORTED_W_A

191 +

192 + @property

193 + def expects_unquantized_inputs(self) -> bool:

194 + """

195 + Humming kernels handle input quantization internally via

196 + HummingMethod.may_quant_input() in the apply() method.

197 +

198 + This property tells the prepare/finalize step to skip input

199 + quantization (by setting defer_input_quant=True) and pass

200 + unquantized inputs to the experts. This prevents double

201 + quantization: once in prepare and once in Humming's apply().

202 +

203 + Returns:

204 + True to indicate that this expert expects unquantized inputs

205 + and will handle quantization internally.

206 + """

207 return True

208

209 @staticmethod

210 def _supports_current_device() -> bool:

211 platform = current_platform

161 - return platform.is_cuda() and platform.has_device_capability((7, 5))

212 + return (

213 + has_humming()

214 + and platform.is_cuda()

215 + and platform.has_device_capability((7, 5))

216 + )

217

218 @staticmethod

219 def _supports_no_act_and_mul() -> bool:

@@ -182,10 +237,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

237

238 @staticmethod

239 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

185 - return not (

186 - moe_parallel_config.use_fi_nvl_two_sided_kernels

187 - or moe_parallel_config.use_fi_nvl_one_sided_kernels

188 - )

240 + return True

241

242 def moe_problem_size(

243 self,

@@ -194,6 +246,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

246 w2: torch.Tensor,

247 topk_ids: torch.Tensor,

248 ) -> tuple[int, int, int, int, int]:

249 + from vllm.utils.humming import HummingLayerMeta

250 +

251 meta1: HummingLayerMeta = self.layer.humming_metas["w13"]

252 meta2: HummingLayerMeta = self.layer.humming_metas["w2"]

253

@@ -215,6 +269,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

269 return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k

270

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

272 + from vllm.utils.humming import GemmType as HummingGemmType

273 + from vllm.utils.humming import dtypes

274 +

275 num_experts = self.num_experts

276 N = self.layer.intermediate_size_per_partition

277 K = self.layer.hidden_size

@@ -254,7 +311,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

311 torch_dtype_map = {

312 dtypes.float16: torch.float16,

313 dtypes.bfloat16: torch.bfloat16,

314 + dtypes.float32: torch.float32,

315 dtypes.float8e4m3: torch.float8_e4m3fn,

316 + dtypes.float8e5m2: torch.float8_e5m2,

317 dtypes.int8: torch.int8,

318 dtypes.int4: torch.uint8,

319 }

@@ -289,7 +348,13 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

348 for key in buffer_metas:

349 meta = buffer_metas[key]

350 if "quanted" in key and a_dtype.num_bits == 4:

292 - meta["shape"] = meta["shape"][:-1] + (meta["shape"][-1] // 2,)

351 + last_dim = meta["shape"][-1]

352 + if last_dim % 2 != 0:

353 + raise ValueError(

354 + f"Int4 packing requires last dimension to be even, "

355 + f"got {last_dim} for buffer '{key}'"

356 + )

357 + meta["shape"] = meta["shape"][:-1] + (last_dim // 2,)

358

359 if num_bits == 16:

360 required_buffers = ["gate_up_output", "activation_output", "down_output"]

@@ -325,8 +390,13 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

390

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

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

393 + elem_size = self.layer.params_dtype.itemsize

394

329 - return (workspace1_nbytes // 2,), (workspace2_nbytes // 2,), output_shape

395 + return (

396 + (workspace1_nbytes // elem_size,),

397 + (workspace2_nbytes // elem_size,),

398 + output_shape,

399 + )

400

401 def workspace_shapes(

402 self,

@@ -344,7 +414,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

414 def make_workspaces(self, M: int, topk: int, activation: MoEActivation):

415 shapes = self._workspace_shapes(M, topk, activation)

416 workspace1_shape, workspace2_shape, output_shape = shapes

347 - torch_dtype = self.layer.param_dtype

417 + torch_dtype = self.layer.params_dtype

418 workspace1, workspace2 = current_workspace_manager().get_simultaneous(

419 (workspace1_shape, torch_dtype),

420 (workspace2_shape, torch_dtype),

@@ -370,45 +440,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

440

441 return buffers

442

373 - def apply(

374 - self,

375 - output: torch.Tensor,

376 - hidden_states: torch.Tensor,

377 - w1: torch.Tensor,

378 - w2: torch.Tensor,

379 - topk_weights: torch.Tensor,

380 - topk_ids: torch.Tensor,

381 - activation: MoEActivation,

382 - global_num_experts: int,

383 - expert_map: torch.Tensor | None,

384 - a1q_scale: torch.Tensor | None,

385 - a2_scale: torch.Tensor | None,

386 - workspace13: torch.Tensor,

387 - workspace2: torch.Tensor,

388 - expert_tokens_meta: mk.ExpertTokensMetadata | None,

389 - apply_router_weight_on_input: bool,

390 - ):

391 - assert not apply_router_weight_on_input

392 -

393 - self.main_apply(

394 - hidden_states=hidden_states,

395 - topk_weights=topk_weights,

396 - topk_ids=topk_ids,

397 - workspace1=workspace13,

398 - workspace2=workspace2,

399 - expert_tokens_meta=expert_tokens_meta,

400 - )

401 -

402 - def main_apply(

403 - self,

404 - hidden_states: torch.Tensor,

405 - topk_weights: torch.Tensor,

406 - topk_ids: torch.Tensor,

407 - workspace1: torch.Tensor,

408 - workspace2: torch.Tensor,

409 - expert_tokens_meta: mk.ExpertTokensMetadata | None,

410 - ):

411 - raise NotImplementedError

443 + # Note: apply method is implemented by subclasses following the

444 + # standard FusedMoEExpertsModular.apply signature

445

446 @staticmethod

447 def is_supported_config(

@@ -418,24 +451,27 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

451 activation_key: QuantKey | None,

452 activation_format: mk.FusedMoEActivationFormat,

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

421 - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

422 - supported = cls.activation_format() == activation_format

423 - reason = "activation_format mismatched"

424 - elif activation_format == mk.FusedMoEActivationFormat.Standard:

425 - if cls.activation_format() != mk.FusedMoEActivationFormat.Standard:

426 - supported = False

427 - reason = "activation_format mismatched"

428 - else:

429 - assert hasattr(cls, "humming_gemm_type")

430 - gemm_type = cls.humming_gemm_type().value.lower()

431 - preferred_gemm_type = get_humming_moe_gemm_type().lower()

432 - supported = preferred_gemm_type == gemm_type

433 - reason = "preferred gemm type mismatched"

434 - else:

435 - supported = False

436 - reason = "unsupported activation_format"

454 + supported, reason = mk.FusedMoEExpertsModular.is_supported_config(

455 + cls,

456 + moe_config,

457 + weight_key,

458 + activation_key,

459 + activation_format,

460 + )

461

438 - return supported, None if supported else reason

462 + if supported:

463 + assert hasattr(cls, "humming_gemm_type")

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

465 + preferred_gemm_type = get_humming_moe_gemm_type()

466 + if preferred_gemm_type is not None:

467 + supported = preferred_gemm_type.lower() == gemm_type

468 + if not supported:

469 + reason = (

470 + f"preferred gemm type {preferred_gemm_type} != "

471 + f"supported gemm type {gemm_type}"

472 + )

473 +

474 + return supported, reason

475

476 def apply_activation(

477 self,

@@ -459,7 +495,9 @@ class HummingIndexedExperts(HummingExpertsBase):

495 return mk.FusedMoEActivationFormat.Standard

496

497 @staticmethod

462 - def humming_gemm_type() -> HummingGemmType:

498 + def humming_gemm_type() -> "HummingGemmType":

499 + from vllm.utils.humming import GemmType as HummingGemmType

500 +

501 return HummingGemmType.INDEXED

502

503 def prepare_humming_moe_kwargs(

@@ -470,12 +508,18 @@ class HummingIndexedExperts(HummingExpertsBase):

508 ) -> tuple[dict[str, Any], dict[str, Any]]:

509 valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

510

511 + moe_block_size = None

512 for min_shape_m, max_shape_m, config in self.w13_tuning_config:

513 if valid_shape_m > min_shape_m and valid_shape_m <= max_shape_m:

514 moe_block_size = config["block_shape"][0]

515 break

477 - else:

478 - raise ValueError(f"cannot found moe_block_size for shape {valid_shape_m}")

516 +

517 + if moe_block_size is None:

518 + logger.warning_once(

519 + "No tuning config found for shape %s, using default block_size=64",

520 + valid_shape_m,

521 + )

522 + moe_block_size = 64

523

524 sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size(

525 topk_ids=topk_ids,

@@ -501,27 +545,47 @@ class HummingIndexedExperts(HummingExpertsBase):

545

546 return moe_kwargs1, moe_kwargs2

547

504 - def main_apply(

548 + def apply(

549 self,

550 + output: torch.Tensor,

551 hidden_states: torch.Tensor,

552 + w1: torch.Tensor,

553 + w2: torch.Tensor,

554 topk_weights: torch.Tensor,

555 topk_ids: torch.Tensor,

509 - workspace1: torch.Tensor,

556 + activation: MoEActivation,

557 + global_num_experts: int,

558 + expert_map: torch.Tensor | None,

559 + a1q_scale: torch.Tensor | None,

560 + a2_scale: torch.Tensor | None,

561 + workspace13: torch.Tensor,

562 workspace2: torch.Tensor,

563 expert_tokens_meta: mk.ExpertTokensMetadata | None,

512 - ):

564 + apply_router_weight_on_input: bool,

565 + ) -> None:

566 + """

567 + Standard apply implementation for Humming indexed experts.

568 +

569 + Note: Humming kernels handle weights and quantization internally through

570 + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used.

571 + The output is written into workspace13 via the buffer management.

572 + """

573 + from vllm.utils.humming import HummingMethod

574 +

575 + assert not apply_router_weight_on_input

576 +

577 hidden_states = hidden_states.view(-1, hidden_states.size(-1))

578 buffers = self.prepare_buffers(

515 - workspace1,

579 + workspace13,

580 workspace2,

581 topk_ids.size(0),

582 topk_ids.size(1),

519 - self.layer.activation,

583 + activation,

584 )

585

586 moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs(

587 topk_ids=topk_ids,

524 - expert_map=self.layer.expert_map,

588 + expert_map=expert_map,

589 expert_tokens_meta=expert_tokens_meta,

590 )

591

@@ -542,7 +606,7 @@ class HummingIndexedExperts(HummingExpertsBase):

606 )

607

608 self.apply_activation(

545 - activation=self.layer.activation,

609 + activation=activation,

610 input=buffers["gate_up_output"],

611 output=buffers["activation_output"],

612 )

@@ -567,10 +631,13 @@ class HummingIndexedExperts(HummingExpertsBase):

631 inputs=buffers["down_output"].view(*topk_ids.shape, -1),

632 topk_weights=topk_weights,

633 topk_ids=topk_ids,

570 - expert_map=self.layer.expert_map,

634 + expert_map=expert_map,

635 outputs=buffers["output"],

636 )

637

638 + # Note: output is already written to buffers["output"]

639 + # which aliases workspace13/output

640 +

641

642 class HummingGroupedExperts(HummingExpertsBase):

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

@@ -581,35 +648,57 @@ class HummingGroupedExperts(HummingExpertsBase):

648 return mk.FusedMoEActivationFormat.Standard

649

650 @staticmethod

584 - def humming_gemm_type() -> HummingGemmType:

651 + def humming_gemm_type() -> "HummingGemmType":

652 + from vllm.utils.humming import GemmType as HummingGemmType

653 +

654 return HummingGemmType.GROUPED_CONTIGUOUS

655

587 - def main_apply(

656 + def apply(

657 self,

658 + output: torch.Tensor,

659 hidden_states: torch.Tensor,

660 + w1: torch.Tensor,

661 + w2: torch.Tensor,

662 topk_weights: torch.Tensor,

663 topk_ids: torch.Tensor,

592 - workspace1: torch.Tensor,

664 + activation: MoEActivation,

665 + global_num_experts: int,

666 + expert_map: torch.Tensor | None,

667 + a1q_scale: torch.Tensor | None,

668 + a2_scale: torch.Tensor | None,

669 + workspace13: torch.Tensor,

670 workspace2: torch.Tensor,

671 expert_tokens_meta: mk.ExpertTokensMetadata | None,

595 - ):

672 + apply_router_weight_on_input: bool,

673 + ) -> None:

674 + """

675 + Standard apply implementation for Humming grouped experts.

676 +

677 + Note: Humming kernels handle weights and quantization internally through

678 + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used.

679 + The output is written into workspace13 via the buffer management.

680 + """

681 + from vllm.utils.humming import HummingMethod

682 +

683 + assert not apply_router_weight_on_input

684 +

685 valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

686

687 buffers = self.prepare_buffers(

599 - workspace1,

688 + workspace13,

689 workspace2,

690 topk_ids.size(0),

691 topk_ids.size(1),

603 - self.layer.activation,

692 + activation,

693 )

694

695 hidden_states, _, expert_first_token_offset, inv_perm, _ = moe_permute(

696 hidden_states=hidden_states,

697 a1q_scale=None,

698 topk_ids=topk_ids,

610 - n_expert=self.global_num_experts,

699 + n_expert=global_num_experts,

700 n_local_expert=self.num_experts,

612 - expert_map=self.layer.expert_map,

701 + expert_map=expert_map,

702 scratch=self._get_permute_scratch(),

703 )

704

@@ -633,7 +722,7 @@ class HummingGroupedExperts(HummingExpertsBase):

722 )

723

724 self.apply_activation(

636 - activation=self.layer.activation,

725 + activation=activation,

726 input=buffers["gate_up_output"],

727 output=buffers["activation_output"],

728 )

@@ -665,6 +754,9 @@ class HummingGroupedExperts(HummingExpertsBase):

754 expert_first_token_offset=expert_first_token_offset,

755 )

756

757 + # Note: output is already written to buffers["output"]

758 + # which aliases workspace13/output

759 +

760

761 class BatchedHummingGroupedExperts(HummingExpertsBase):

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

@@ -675,29 +767,51 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):

767 return mk.FusedMoEActivationFormat.BatchedExperts

768

769 @staticmethod

678 - def humming_gemm_type() -> HummingGemmType:

770 + def humming_gemm_type() -> "HummingGemmType":

771 + from vllm.utils.humming import GemmType as HummingGemmType

772 +

773 return HummingGemmType.GROUPED_MASKED

774

681 - def main_apply(

775 + def apply(

776 self,

777 + output: torch.Tensor,

778 hidden_states: torch.Tensor,

779 + w1: torch.Tensor,

780 + w2: torch.Tensor,

781 topk_weights: torch.Tensor,

782 topk_ids: torch.Tensor,

686 - workspace1: torch.Tensor,

783 + activation: MoEActivation,

784 + global_num_experts: int,

785 + expert_map: torch.Tensor | None,

786 + a1q_scale: torch.Tensor | None,

787 + a2_scale: torch.Tensor | None,

788 + workspace13: torch.Tensor,

789 workspace2: torch.Tensor,

790 expert_tokens_meta: mk.ExpertTokensMetadata | None,

689 - ):

791 + apply_router_weight_on_input: bool,

792 + ) -> None:

793 + """

794 + Standard apply implementation for Humming batched grouped experts.

795 +

796 + Note: Humming kernels handle weights and quantization internally through

797 + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used.

798 + The output is written into workspace13 via the buffer management.

799 + """

800 + from vllm.utils.humming import HummingMethod

801 +

802 + assert not apply_router_weight_on_input

803 assert expert_tokens_meta is not None

804 +

805 hidden_states = hidden_states.view(-1, hidden_states.size(-1))

806 valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

807 expert_num_tokens = expert_tokens_meta.expert_num_tokens

808

809 buffers = self.prepare_buffers(

696 - workspace1,

810 + workspace13,

811 workspace2,

812 topk_ids.size(0),

813 topk_ids.size(1),

700 - self.layer.activation,

814 + activation,

815 )

816

817 inputs, input_scale = HummingMethod.may_quant_input(

@@ -720,7 +834,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):

834 )

835

836 self.apply_activation(

723 - activation=self.layer.activation,

837 + activation=activation,

838 input=buffers["gate_up_output"],

839 output=buffers["activation_output"],

840 )

@@ -743,3 +857,6 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):

857 tuning_config=self.w2_tuning_config_str,

858 sublayer_name="w2",

859 )

860 +

861 + # Note: output is already written to buffers["down_output"]

862 + # which aliases workspace13/output

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

+11 -6

@@ -711,10 +711,12 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

711

712 if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

713 from vllm.model_executor.layers.quantization.utils.humming_utils import (

714 - prepare_humming_moe_layer,

714 + convert_to_humming_moe_kernel_format,

715 )

716

717 - prepare_humming_moe_layer(layer, {"quant_method": "gpt_oss_mxfp4"})

717 + convert_to_humming_moe_kernel_format(

718 + layer, quant_config={"quant_method": "gpt_oss_mxfp4"}

719 + )

720 return (

721 layer.w13_weight,

722 layer.w2_weight,

@@ -1277,10 +1279,12 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1279

1280 if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

1281 from vllm.model_executor.layers.quantization.utils.humming_utils import (

1280 - prepare_humming_moe_layer,

1282 + convert_to_humming_moe_kernel_format,

1283 )

1284

1283 - prepare_humming_moe_layer(layer, {"quant_method": "mxfp4"})

1285 + convert_to_humming_moe_kernel_format(

1286 + layer, quant_config={"quant_method": "mxfp4"}

1287 + )

1288 return (

1289 layer.w13_weight,

1290 layer.w2_weight,

@@ -1569,7 +1573,7 @@ def make_mxfp4_moe_quant_config(

1573 w2_bias: torch.Tensor | None = None,

1574 a1_scale: torch.Tensor | None = None,

1575 a2_scale: torch.Tensor | None = None,

1572 - layer: torch.nn.Module | None = None,

1576 + layer: "RoutedExperts | None" = None,

1577 ) -> FusedMoEQuantConfig | None:

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

1579 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:

@@ -1662,7 +1666,7 @@ def make_mxfp4_moe_quant_config(

1666 get_humming_moe_quant_config,

1667 )

1668

1665 - assert isinstance(layer, RoutedExperts)

1669 + assert layer is not None

1670 return get_humming_moe_quant_config(

1671 layer,

1672 gemm1_alpha=gemm1_alpha,

@@ -1703,6 +1707,7 @@ def make_mxfp4_moe_kernel(

1707 assert prepare_finalize is not None

1708

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

1710 + logger.info_once("Using %s", experts_cls.__name__)

1711

1712 extra_kwargs = {}

1713 if mxfp4_backend == Mxfp4MoeBackend.HUMMING:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py

@@ -143,6 +143,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

143 mxfp4_backend=self.mxfp4_backend,

144 w1_scale=layer.w13_weight_scale,

145 w2_scale=layer.w2_weight_scale,

146 + layer=layer,

147 )

148

149 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

vllm/model_executor/layers/quantization/fp8.py

+9 -9

@@ -706,15 +706,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):

706 layer.w2_weight.is_shuffled = True

707

708 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

709 - if self.moe_quant_config:

710 - assert self.experts_cls is not None

711 - self.moe_kernel = make_fp8_moe_kernel(

712 - moe_quant_config=self.moe_quant_config,

713 - moe_config=self.moe,

714 - fp8_backend=self.fp8_backend,

715 - experts_cls=self.experts_cls,

716 - routing_tables=layer._expert_routing_tables(),

717 - )

709 + assert self.moe_quant_config is not None

710 + assert self.experts_cls is not None

711 + self.moe_kernel = make_fp8_moe_kernel(

712 + moe_quant_config=self.moe_quant_config,

713 + moe_config=self.moe,

714 + fp8_backend=self.fp8_backend,

715 + experts_cls=self.experts_cls,

716 + routing_tables=layer._expert_routing_tables(),

717 + )

718

719 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

720 # Allow for accessing weights and scales in standard way.

vllm/model_executor/layers/quantization/humming.py

+75 -163

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

30 QuantizationConfig,

31 QuantizeMethodBase,

32 )

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

34 + convert_to_humming_moe_kernel_format,

35 + get_humming_moe_quant_config,

36 + input_schema_to_quant_key,

37 + make_humming_moe_kernel,

38 + select_humming_moe_experts,

39 + weight_schema_to_quant_key,

40 +)

41 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

42 from vllm.model_executor.parameter import (

43 BasevLLMParameter,

@@ -106,7 +114,7 @@ def prepare_param(tensor, name, extra_attrs):

114 return param

115

116

109 -def prepare_moe_param(tensor, name, extra_attrs):

117 +def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, Any]):

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

119 if "scale_type" in extra_attrs:

120 extra_attrs["quant_method"] = extra_attrs["scale_type"]

@@ -605,12 +613,27 @@ class HummingMoEMethod(FusedMoEMethodBase):

613 ) -> None:

614 super().__init__(moe)

615 self.quant_config = quant_config

608 - self.moe = moe

616 self.weight_schema = quant_config.weight_schema

617 self.input_schema = quant_config.input_schema

618 self.force_weight_schema = quant_config.force_weight_schema

619 self.force_input_schema = quant_config.force_input_schema

620

621 + # Derive QuantKeys from humming schemas.

622 + # Prefer force schemas (the final format after requant) over base.

623 + weight_key = weight_schema_to_quant_key(

624 + self.force_weight_schema or self.weight_schema

625 + )

626 + activation_key = input_schema_to_quant_key(

627 + self.force_input_schema or self.input_schema

628 + )

629 +

630 + # Select Humming MoE experts

631 + self.experts_cls = select_humming_moe_experts(

632 + config=self.moe,

633 + weight_key=weight_key,

634 + activation_key=activation_key,

635 + )

636 +

637 def prepare_weight_loader(self, layer, weight_loader):

638 def new_weight_loader(

639 param: torch.nn.Parameter,

@@ -647,7 +670,7 @@ class HummingMoEMethod(FusedMoEMethodBase):

670 sublayer_name = "w2" if shard_id == "w2" else "w13"

671

672 param = getattr(layer, sublayer_name + "_" + key)

650 - part_subccess = param.weight_loader(

673 + part_success = param.weight_loader(

674 param=param,

675 loaded_weight=tensor.cpu(),

676 weight_name=shard_id + "_" + key,

@@ -655,7 +678,7 @@ class HummingMoEMethod(FusedMoEMethodBase):

678 expert_id=expert_id,

679 return_success=return_success,

680 )

658 - success = success and part_subccess

681 + success = success and part_success

682

683 return success if return_success else None

684

@@ -677,7 +700,7 @@ class HummingMoEMethod(FusedMoEMethodBase):

700

701 def create_weights(

702 self,

680 - layer: torch.nn.Module,

703 + layer: RoutedExperts,

704 num_experts: int,

705 hidden_size: int,

706 intermediate_size_per_partition: int,

@@ -734,160 +757,35 @@ class HummingMoEMethod(FusedMoEMethodBase):

757 locks = torch.zeros(1024, dtype=torch.int32)

758 layer.register_buffer("locks", locks)

759

737 - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

738 - from vllm.model_executor.layers.quantization.utils.humming_utils import (

739 - get_humming_moe_quant_config,

740 - )

741 -

760 + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:

761 return get_humming_moe_quant_config(layer)

762

763 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

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

765 return

766 self.processed = True

748 - layer.weight_schemas = {}

749 - layer.input_schemas = {}

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

751 - input_schema = self.input_schema

752 - weight_schema = self.weight_schema

753 - # convert from checkpoint format to humming format

754 - if not isinstance(weight_schema, _hm.HummingWeightSchema):

755 - tensors: dict[str, torch.Tensor] = dict(

756 - (key.removeprefix(sublayer_name + "_"), value)

757 - for key, value in layer.state_dict().items()

758 - if key.startswith(sublayer_name + "_")

759 - )

760 -

761 - shape_k_stacks = [configs["shape_k"]]

762 - shape_n_stacks = [configs["shape_n"]]

763 - if sublayer_name == "w13":

764 - shape_n_stacks = [configs["shape_n"] // 2] * 2

765 -

766 - weight_schema, tensors = weight_schema.convert_humming(

767 - tensors=tensors,

768 - shape_n_stacks=shape_n_stacks,

769 - shape_k_stacks=shape_k_stacks,

770 - param_dtype=layer.param_dtype,

771 - num_experts=layer.num_experts,

772 - )

773 -

774 - input_schema, _ = input_schema.convert_humming(

775 - tensors=tensors,

776 - shape_n_stacks=shape_n_stacks,

777 - shape_k_stacks=shape_k_stacks,

778 - param_dtype=layer.param_dtype,

779 - num_experts=layer.num_experts,

780 - )

781 -

782 - for name, _ in list(layer.named_parameters()):

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

784 - continue

785 - delattr(layer, name)

786 -

787 - for name, tensor in tensors.items():

788 - name = f"{sublayer_name}_{name}"

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

790 - setattr(layer, name, param)

791 -

792 - layer.weight_schemas[sublayer_name] = weight_schema

793 - layer.input_schemas[sublayer_name] = input_schema

794 -

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

796 - assert isinstance(weight_schema, _hm.HummingWeightSchema)

797 - force_requant = self.force_weight_schema is not None

798 - if force_requant and weight_schema != self.force_weight_schema:

799 - tensors = dict(

800 - (key.removeprefix(sublayer_name + "_"), value)

801 - for key, value in layer.state_dict().items()

802 - if key.startswith(sublayer_name + "_")

803 - )

804 -

805 - tensors = weight_schema.requant_tensors(

806 - tensors=tensors,

807 - target_weight_schema=self.force_weight_schema,

808 - param_dtype=layer.param_dtype,

809 - )

810 -

811 - weight_schema = self.force_weight_schema

812 -

813 - for name, _ in list(layer.named_parameters()):

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

815 - continue

816 - if name == sublayer_name + "_bias":

817 - continue

818 - delattr(layer, name)

819 -

820 - for name, tensor in tensors.items():

821 - name = f"{sublayer_name}_{name}"

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

823 - setattr(layer, name, param)

767

825 - del tensors

826 -

827 - # prepare layer config from humming kernel

828 - _hm.HummingMethod.prepare_layer_meta(

829 - layer=layer,

830 - shape_n=configs["shape_n"],

831 - shape_k=configs["shape_k"],

832 - pad_n_to_multiple=256,

833 - pad_k_to_multiple=128,

834 - input_schema=input_schema,

835 - weight_schema=weight_schema,

836 - has_bias=self.moe.has_bias,

837 - num_experts=layer.num_experts,

838 - torch_dtype=layer.param_dtype,

839 - sublayer_name=sublayer_name,

840 - )

841 -

842 - # preprocess weight for inference

843 - _hm.HummingMethod.transform_humming_layer(

844 - layer, sublayer_name=sublayer_name

845 - )

846 -

847 - from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (

848 - HummingGroupedExperts,

849 - HummingIndexedExperts,

850 - get_humming_moe_gemm_type,

768 + # Convert weights to Humming kernel format

769 + convert_to_humming_moe_kernel_format(

770 + layer=layer,

771 + sublayer_configs=layer.sublayer_configs,

772 + weight_schema=self.weight_schema,

773 + input_schema=self.input_schema,

774 + force_weight_schema=self.force_weight_schema,

775 )

776

853 - # use moe modular

854 - experts: HummingIndexedExperts | HummingGroupedExperts

855 - layer._ensure_moe_quant_config_init()

777 + # Build the MoE kernel

778 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

779 assert self.moe_quant_config is not None

857 - if get_humming_moe_gemm_type() == "indexed":

858 - experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config)

859 - else:

860 - experts = HummingGroupedExperts(layer, self.moe, self.moe_quant_config)

861 - self.experts = experts

862 -

863 - def select_gemm_impl(

864 - self,

865 - prepare_finalize,

866 - layer: torch.nn.Module,

867 - ):

868 - from vllm.model_executor.layers.fused_moe import modular_kernel as mk

869 - from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (

870 - BatchedHummingGroupedExperts,

871 - HummingGroupedExperts,

872 - HummingIndexedExperts,

873 - get_humming_moe_gemm_type,

780 + assert self.experts_cls is not None

781 + self.moe_kernel = make_humming_moe_kernel(

782 + self.moe_quant_config,

783 + self.moe,

784 + self.experts_cls,

785 + layer=layer,

786 + routing_tables=layer._expert_routing_tables(),

787 )

788

876 - activation_format = prepare_finalize.activation_format

877 - assert self.moe_quant_config is not None

878 - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

879 - return BatchedHummingGroupedExperts(

880 - layer=layer,

881 - moe_config=self.moe,

882 - quant_config=self.moe_quant_config,

883 - max_num_tokens=prepare_finalize.max_num_tokens_per_rank(),

884 - num_dispatchers=prepare_finalize.num_dispatchers(),

885 - )

886 - elif get_humming_moe_gemm_type() == "indexed":

887 - return HummingIndexedExperts(layer, self.moe, self.moe_quant_config)

888 - else:

889 - return HummingGroupedExperts(layer, self.moe, self.moe_quant_config)

890 -

789 def apply(

790 self,

791 layer: RoutedExperts,

@@ -896,22 +794,36 @@ class HummingMoEMethod(FusedMoEMethodBase):

794 topk_ids: torch.Tensor,

795 shared_experts: SharedExperts | None,

796 shared_experts_input: torch.Tensor | None,

899 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

900 - workspace1, workspace2, output = self.experts.make_workspaces(

901 - M=topk_ids.size(0),

902 - topk=topk_ids.size(1),

903 - activation=layer.activation,

904 - )

905 -

906 - assert workspace1.data_ptr() == output.data_ptr()

907 -

908 - self.experts.main_apply(

797 + ) -> torch.Tensor:

798 + """

799 + Apply Humming-quantized MoE computation using the standard kernel flow.

800 +

801 + This method uses FusedMoEKernel.apply() which orchestrates:

802 + 1. Preparation (quantization if needed - skipped for Humming via

803 + expects_unquantized_inputs=True to prevent double quantization)

804 + 2. Expert computation (via experts.apply())

805 + 3. Finalization (weight application & reduction - no-op for Humming

806 + since it's already done internally)

807 +

808 + Humming handles all quantization, weight application, and reduction

809 + internally in the experts.apply() method via HummingMethod calls.

810 +

811 + Note: Although w1/w2 weights are passed to the kernel for interface

812 + consistency, Humming's experts.apply() reads weights directly from

813 + the layer object via HummingMethod.forward_layer() and ignores the

814 + w1/w2 parameters.

815 + """

816 + assert self.moe_kernel is not None

817 + return self.moe_kernel.apply(

818 hidden_states=x,

910 - topk_weights=topk_weights,

819 + w1=layer.w13_weight,

820 + w2=layer.w2_weight,

821 topk_ids=topk_ids,

912 - workspace1=workspace1,

913 - workspace2=workspace2,

914 - expert_tokens_meta=None,

822 + topk_weights=topk_weights,

823 + activation=layer.activation,

824 + global_num_experts=layer.global_num_experts,

825 + expert_map=layer.expert_map,

826 + apply_router_weight_on_input=False,

827 + shared_experts=shared_experts,

828 + shared_experts_input=shared_experts_input,

829 )

916 -

917 - return output

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -1280,6 +1280,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1280 gemm1_alpha=getattr(layer, "swiglu_alpha", None),

1281 gemm1_beta=getattr(layer, "swiglu_beta", None),

1282 swiglu_limit=getattr(layer, "swiglu_limit", None),

1283 + layer=layer,

1284 )

1285

1286 # Emulation and other schemes

vllm/model_executor/layers/quantization/utils/humming_utils.py

+847 -110

@@ -1,20 +1,376 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 import json

4 -from typing import Any

4 +from typing import TYPE_CHECKING, Any

6 import regex as re

7 import torch

9 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

10 from vllm import envs

10 -from vllm.model_executor.layers.fused_moe import RoutedExperts

11 +from vllm.logger import init_logger

12 +from vllm.model_executor.layers.fused_moe.all2all_utils import (

13 + maybe_make_prepare_finalize,

14 +)

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

16 + FusedMoEConfig,

17 FusedMoEQuantConfig,

18 FusedMoEQuantDesc,

19 )

20 +from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (

21 + BatchedHummingGroupedExperts,

22 + HummingGroupedExperts,

23 + HummingIndexedExperts,

24 +)

25 from vllm.model_executor.layers.linear import LinearBase

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

17 -from vllm.utils.humming import BaseWeightSchema, HummingInputSchema, HummingMethod

26 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

27 + FP4_DTYPE,

28 + FP8_DTYPE,

29 + INT4_DTYPE,

30 + INT8_DTYPE,

31 + MXFP_SCALE_DTYPE,

32 + GroupShape,

33 + QuantKey,

34 + ScaleDesc,

35 +)

36 +from vllm.utils.import_utils import has_humming

37 +

38 +if TYPE_CHECKING:

39 + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts

40 + from vllm.utils.humming import (

41 + AWQWeightSchema,

42 + BaseInputSchema,

43 + BaseWeightSchema,

44 + CompressedTensorsInputSchema,

45 + CompressedTensorsWeightSchema,

46 + Fp8WeightSchema,

47 + GPTQWeightSchema,

48 + HummingInputSchema,

49 + HummingWeightSchema,

50 + )

51 + from vllm.utils.humming import dtypes as humming_dtypes

52 +

53 +logger = init_logger(__name__)

54 +

55 +if has_humming():

56 + from vllm.utils.humming import dtypes as humming_dtypes

57 +

58 + _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {

59 + humming_dtypes.float4e2m1: FP4_DTYPE,

60 + humming_dtypes.float8e4m3: FP8_DTYPE,

61 + humming_dtypes.float8e5m2: torch.float8_e5m2,

62 + humming_dtypes.int8: torch.int8,

63 + humming_dtypes.uint4: INT4_DTYPE,

64 + humming_dtypes.uint8: INT8_DTYPE,

65 + humming_dtypes.uint2: torch.uint8,

66 + humming_dtypes.uint3: torch.uint8,

67 + }

68 +

69 + _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {

70 + humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,

71 + humming_dtypes.float8e4m3: FP8_DTYPE,

72 + humming_dtypes.float16: torch.float16,

73 + humming_dtypes.bfloat16: torch.bfloat16,

74 + humming_dtypes.float32: torch.float32,

75 + }

76 +

77 +

78 +def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:

79 + """

80 + Map humming group sizes to QuantKey GroupShape.

81 +

82 + group_size: elements per group along K (col); 0 means full dimension.

83 + group_size_n: elements per group along N (row); 0 means 1 (per-row).

84 +

85 + GroupShape convention: row = N dim, col = K dim.

86 + """

87 + if group_size == 0 and group_size_n == 0:

88 + return GroupShape.PER_CHANNEL

89 +

90 + row = group_size_n if group_size_n > 0 else 1

91 + col = group_size if group_size > 0 else -1

92 + return GroupShape(row=row, col=col)

93 +

94 +

95 +# ---- HummingWeightSchema (post-conversion) --------------------------------

96 +

97 +

98 +def _humming_weight_schema_to_quant_key(

99 + schema: "HummingWeightSchema",

100 +) -> QuantKey:

101 + from vllm.utils.humming import WeightScaleType

102 +

103 + """Convert a HummingWeightSchema to a QuantKey."""

104 + dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]

105 +

106 + if schema.bs_dtype is not None:

107 + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]

108 + else:

109 + scale_dtype = torch.float32

110 +

111 + group_shape = _group_shape(

112 + schema.weight_scale_group_size,

113 + schema.weight_scale_group_size_n,

114 + )

115 +

116 + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)

117 +

118 + scale2 = None

119 + if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR:

120 + scale2 = ScaleDesc(

121 + dtype=torch.float32,

122 + static=True,

123 + group_shape=GroupShape.PER_TENSOR,

124 + )

125 +

126 + return QuantKey(

127 + dtype=dtype,

128 + scale=scale,

129 + scale2=scale2,

130 + symmetric=not schema.has_zero_point,

131 + )

132 +

133 +

134 +# ---- Checkpoint-format weight schemas (pre-conversion) --------------------

135 +

136 +

137 +def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey:

138 + if schema.weight_block_size is not None:

139 + gs_n, gs_k = schema.weight_block_size

140 + group_shape = GroupShape(row=gs_n, col=gs_k)

141 + else:

142 + group_shape = GroupShape.PER_CHANNEL

143 +

144 + scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape)

145 + return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True)

146 +

147 +

148 +def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey:

149 + group_shape = _group_shape(schema.group_size)

150 + scale = ScaleDesc(

151 + dtype=torch.float16,

152 + static=True,

153 + group_shape=group_shape,

154 + )

155 + return QuantKey(

156 + dtype=INT4_DTYPE,

157 + scale=scale,

158 + symmetric=not schema.zero_point,

159 + )

160 +

161 +

162 +def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey:

163 + group_shape = _group_shape(schema.group_size)

164 + scale = ScaleDesc(

165 + dtype=torch.float16,

166 + static=True,

167 + group_shape=group_shape,

168 + )

169 + return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym)

170 +

171 +

172 +def _compressed_tensors_weight_schema_to_quant_key(

173 + schema: "CompressedTensorsWeightSchema",

174 +) -> QuantKey:

175 + # Determine dtype from format/type/num_bits

176 + fmt = schema.format

177 + if fmt in ("int-quantized", "float-quantized", "naive-quantized"):

178 + dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE

179 + elif "nvfp4" in fmt or "mxfp4" in fmt:

180 + dtype = FP4_DTYPE

181 + else:

182 + dtype = _HUMMING_TO_QUANT_DTYPE[

183 + humming_dtypes.DataType.from_str(f"uint{schema.num_bits}")

184 + ]

185 +

186 + # Determine group shape from strategy

187 + if schema.strategy in ("group", "tensor_group"):

188 + group_shape = _group_shape(schema.group_size or 0)

189 + elif schema.strategy == "block" and schema.block_structure is not None:

190 + group_shape = GroupShape(

191 + row=schema.block_structure[0],

192 + col=schema.block_structure[1],

193 + )

194 + else:

195 + group_shape = GroupShape.PER_CHANNEL

196 +

197 + # Determine scale dtype

198 + if "mxfp" in fmt:

199 + scale_dtype = MXFP_SCALE_DTYPE

200 + elif "nvfp4" in fmt:

201 + scale_dtype = FP8_DTYPE

202 + else:

203 + scale_dtype = torch.float32

204 +

205 + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)

206 +

207 + scale2 = None

208 + if "nvfp4" in fmt or schema.strategy == "tensor_group":

209 + scale2 = ScaleDesc(

210 + dtype=torch.float32,

211 + static=True,

212 + group_shape=GroupShape.PER_TENSOR,

213 + )

214 +

215 + return QuantKey(

216 + dtype=dtype,

217 + scale=scale,

218 + scale2=scale2,

219 + symmetric=schema.symmetric,

220 + )

221 +

222 +

223 +# ---- Dispatch for any BaseWeightSchema ------------------------------------

224 +

225 +

226 +def weight_schema_to_quant_key(

227 + schema: "BaseWeightSchema",

228 +) -> QuantKey:

229 + from vllm.utils.humming import (

230 + AWQWeightSchema,

231 + BitnetWeightSchema,

232 + CompressedTensorsWeightSchema,

233 + Fp8WeightSchema,

234 + GptOssMxfp4WeightSchema,

235 + GPTQWeightSchema,

236 + HummingWeightSchema,

237 + ModeloptMxfp8WeightSchema,

238 + ModeloptNvfp4WeightSchema,

239 + Mxfp4WeightSchema,

240 + )

241 +

242 + """Convert any BaseWeightSchema to a QuantKey."""

243 + if isinstance(schema, HummingWeightSchema):

244 + return _humming_weight_schema_to_quant_key(schema)

245 +

246 + # Schemas with fixed QuantKeys

247 + if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)):

248 + return QuantKey(

249 + dtype=FP4_DTYPE,

250 + scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),

251 + )

252 + if isinstance(schema, ModeloptMxfp8WeightSchema):

253 + return QuantKey(

254 + dtype=FP8_DTYPE,

255 + scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),

256 + )

257 + if isinstance(schema, ModeloptNvfp4WeightSchema):

258 + return QuantKey(

259 + dtype=FP4_DTYPE,

260 + scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)),

261 + scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR),

262 + )

263 + if isinstance(schema, BitnetWeightSchema):

264 + return QuantKey(

265 + dtype=torch.uint8,

266 + scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL),

267 + )

268 +

269 + # Schemas requiring config inspection

270 + if isinstance(schema, Fp8WeightSchema):

271 + return _fp8_weight_schema_to_quant_key(schema)

272 + if isinstance(schema, AWQWeightSchema):

273 + return _awq_weight_schema_to_quant_key(schema)

274 + if isinstance(schema, GPTQWeightSchema):

275 + return _gptq_weight_schema_to_quant_key(schema)

276 + if isinstance(schema, CompressedTensorsWeightSchema):

277 + return _compressed_tensors_weight_schema_to_quant_key(schema)

278 +

279 + raise TypeError(f"Unsupported weight schema type: {type(schema)}")

280 +

281 +

282 +# ---- HummingInputSchema (post-conversion) ----------------------------------

283 +

284 +

285 +def _humming_input_schema_to_quant_key(

286 + schema: "HummingInputSchema",

287 +) -> QuantKey | None:

288 + """Convert a HummingInputSchema to a QuantKey. Returns None if

289 + the schema represents unquantized (bf16/fp16) inputs."""

290 + if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:

291 + return None

292 +

293 + dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]

294 +

295 + gs = schema.input_scale_group_size

296 + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN

297 +

298 + scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32

299 +

300 + scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)

301 +

302 + return QuantKey(dtype=dtype, scale=scale, symmetric=True)

303 +

304 +

305 +# ---- Checkpoint-format input schemas (pre-conversion) ----------------------

306 +

307 +

308 +def _resolve_input_quant_key(

309 + origin_a_dtype: "humming_dtypes.DataType",

310 + group_size: int,

311 +) -> QuantKey | None:

312 + from vllm.utils.humming import HummingInputSchema

313 +

314 + """Resolve the actual activation QuantKey after platform fallback."""

315 + a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype)

316 + if a_dtype is None or a_dtype.num_bits >= 16:

317 + return None

318 +

319 + dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype]

320 + gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0

321 + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN

322 + scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32

323 +

324 + scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)

325 + return QuantKey(dtype=dtype, scale=scale, symmetric=True)

326 +

327 +

328 +def _compressed_tensors_input_schema_to_quant_key(

329 + schema: "CompressedTensorsInputSchema",

330 +) -> QuantKey | None:

331 + type_bits_to_dtype = {

332 + ("float", 8): humming_dtypes.float8e4m3,

333 + ("float", 4): humming_dtypes.float4e2m1,

334 + ("int", 8): humming_dtypes.int8,

335 + ("int", 4): humming_dtypes.int4,

336 + }

337 + origin = type_bits_to_dtype.get((schema.type, schema.num_bits))

338 + if origin is None:

339 + return None

340 + return _resolve_input_quant_key(origin, schema.group_size)

341 +

342 +

343 +# ---- Dispatch for any BaseInputSchema -------------------------------------

344 +

345 +

346 +def input_schema_to_quant_key(

347 + schema: "BaseInputSchema",

348 +) -> QuantKey | None:

349 + from vllm.utils.humming import (

350 + CompressedTensorsInputSchema,

351 + Fp8InputSchema,

352 + HummingInputSchema,

353 + ModeloptNvfp4InputSchema,

354 + )

355 +

356 + """Convert any BaseInputSchema to a QuantKey. Returns None if

357 + the schema represents unquantized (bf16/fp16) inputs."""

358 + if isinstance(schema, HummingInputSchema):

359 + return _humming_input_schema_to_quant_key(schema)

360 +

361 + if isinstance(schema, Fp8InputSchema):

362 + return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0)

363 +

364 + if isinstance(schema, ModeloptNvfp4InputSchema):

365 + return _resolve_input_quant_key(

366 + humming_dtypes.float8e4m3,

367 + schema.group_size,

368 + )

369 +

370 + if isinstance(schema, CompressedTensorsInputSchema):

371 + return _compressed_tensors_input_schema_to_quant_key(schema)

372 +

373 + raise TypeError(f"Unsupported input schema type: {type(schema)}")

374

375

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

@@ -24,8 +380,9 @@ def humming_is_layer_skipped(config: dict[str, Any], prefix: str):

380 keys = ["ignored_layers", "ignore", "modules_to_not_convert"]

381 ignored_layers: list[str] = []

382 for key in keys:

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

28 - if not ignored_layers:

383 + candidate = config.get(key, []) or []

384 + if candidate:

385 + ignored_layers = candidate

386 break

387

388 if any(module_name in prefix for module_name in ignored_layers):

@@ -79,6 +436,12 @@ def convert_linear_layer_to_humming_standard(

436

437

438 def prepare_humming_layer(layer: LinearBase, quant_config: dict):

439 + from vllm.utils.humming import (

440 + BaseWeightSchema,

441 + HummingInputSchema,

442 + HummingMethod,

443 + )

444 +

445 weight_schema = BaseWeightSchema.from_config(quant_config)

446 input_schema = HummingInputSchema()

447

@@ -140,90 +503,59 @@ def prepare_humming_layer(layer: LinearBase, quant_config: dict):

503 layer.compute_config = json.dumps(compute_config)

504

505

143 -def prepare_humming_moe_layer(layer: RoutedExperts, quant_config: dict):

144 - weight_schema = BaseWeightSchema.from_config(quant_config)

145 - input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}

146 - if humming_is_layer_skipped(input_quant_config, layer.layer_name):

147 - input_schema = HummingInputSchema()

506 +def make_humming_moe_quant_config(

507 + quant_dtype: torch.dtype | str | None,

508 + weight_dtype: torch.dtype | str | None,

509 + weight_group_shape: GroupShape | None = None,

510 + w1_scale: torch.Tensor | None = None,

511 + w2_scale: torch.Tensor | None = None,

512 + w1_zp: torch.Tensor | None = None,

513 + w2_zp: torch.Tensor | None = None,

514 + w1_bias: torch.Tensor | None = None,

515 + w2_bias: torch.Tensor | None = None,

516 + w1_gscale: torch.Tensor | None = None,

517 + w2_gscale: torch.Tensor | None = None,

518 + gemm1_alpha: float | None = None,

519 + gemm1_beta: float | None = None,

520 + gemm1_clamp_limit: float | None = None,

521 +) -> FusedMoEQuantConfig:

522 + if quant_dtype is None:

523 + a_quant_desc = FusedMoEQuantDesc(dtype=None)

524 else:

149 - # TODO: read input_quant_config from quant_config

150 - input_schema = HummingInputSchema.from_config(input_quant_config)

151 -

152 - is_gated = layer.activation.is_gated

153 - shape_config = {

154 - "w13": (

155 - layer.moe_config.intermediate_size_per_partition * 2,

156 - layer.moe_config.hidden_dim,

157 - ),

158 - "w2": (

159 - layer.moe_config.hidden_dim,

160 - layer.moe_config.intermediate_size_per_partition * (1 if is_gated else 2),

161 - ),

162 - }

163 -

164 - layer.weight_schemas = {}

165 - layer.input_schemas = {}

166 -

167 - for sublayer_name in shape_config:

168 - # Step 1: convert weight to humming standard format

169 - tensors: dict[str, torch.Tensor] = dict(

170 - (key.removeprefix(sublayer_name + "_"), value)

171 - for key, value in layer.state_dict().items()

172 - if key.startswith(sublayer_name + "_")

173 - )

174 -

175 - shape_n, shape_k = shape_config[sublayer_name]

176 - shape_n_stacks = [shape_n]

177 - shape_k_stacks = [shape_k]

178 - if sublayer_name == "w13":

179 - shape_n_stacks = [shape_n // 2] * 2

180 -

181 - weight_schema_new, tensors = weight_schema.convert_humming(

182 - tensors=tensors,

183 - shape_n_stacks=shape_n_stacks,

184 - shape_k_stacks=shape_k_stacks,

185 - num_experts=layer.local_num_experts,

186 - param_dtype=layer.params_dtype,

187 - )

188 -

189 - layer.weight_schemas[sublayer_name] = weight_schema_new

190 - layer.input_schemas[sublayer_name] = input_schema

191 -

192 - for name, _ in list(layer.named_parameters()):

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

194 - continue

195 - delattr(layer, name)

196 -

197 - for name, tensor in tensors.items():

198 - name = f"{sublayer_name}_{name}"

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

200 - setattr(layer, name, param)

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

526 + a_quant_desc = FusedMoEQuantDesc(dtype=quant_dtype, shape=shape)

527

202 - # Step 2: transform weight (humming standard format) for forwarding

203 - HummingMethod.prepare_layer_meta(

204 - layer=layer,

205 - shape_n=shape_n,

206 - shape_k=shape_k,

207 - pad_n_to_multiple=256,

208 - pad_k_to_multiple=128,

209 - input_schema=input_schema,

210 - weight_schema=weight_schema_new,

211 - has_bias=layer.moe_config.has_bias,

212 - num_experts=layer.num_experts,

213 - torch_dtype=layer.params_dtype,

214 - sublayer_name=sublayer_name,

215 - )

528 + w1_quant_desc = FusedMoEQuantDesc(

529 + dtype=weight_dtype,

530 + shape=weight_group_shape,

531 + scale=w1_scale,

532 + alpha_or_gscale=w1_gscale,

533 + zp=w1_zp,

534 + bias=w1_bias,

535 + )

536

217 - HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name)

537 + w2_quant_desc = FusedMoEQuantDesc(

538 + dtype=weight_dtype,

539 + shape=weight_group_shape,

540 + scale=w2_scale,

541 + alpha_or_gscale=w2_gscale,

542 + zp=w2_zp,

543 + bias=w2_bias,

544 + )

545

219 - if not hasattr(layer, "locks"):

220 - device = layer.w13_weight.device

221 - locks = torch.zeros(1024, dtype=torch.int32, device=device)

222 - layer.register_buffer("locks", locks)

546 + return FusedMoEQuantConfig(

547 + _a1=a_quant_desc,

548 + _a2=a_quant_desc,

549 + _w1=w1_quant_desc,

550 + _w2=w2_quant_desc,

551 + gemm1_alpha=gemm1_alpha,

552 + gemm1_beta=gemm1_beta,

553 + gemm1_clamp_limit=gemm1_clamp_limit,

554 + )

555

556

557 def get_humming_moe_quant_config(

226 - layer: RoutedExperts,

558 + layer: "RoutedExperts",

559 gemm1_alpha: float | None = None,

560 gemm1_beta: float | None = None,

561 gemm1_clamp_limit: float | None = None,

@@ -231,12 +563,10 @@ def get_humming_moe_quant_config(

563 input_schema = layer.input_schemas["w13"]

564 weight_schema = layer.weight_schemas["w13"]

565

234 - a_dtype = input_schema.a_dtype

235 - if a_dtype is None or a_dtype.num_bits == 16:

236 - a_quant_desc = FusedMoEQuantDesc(dtype=None)

566 + if input_schema.a_dtype is None or input_schema.a_dtype.num_bits == 16:

567 + q_dtype = None

568 else:

238 - shape = GroupShape(row=1, col=-1)

239 - a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape)

569 + q_dtype = str(input_schema.a_dtype)

570

571 weight_scale_group_size = weight_schema.weight_scale_group_size

572 weight_scale_group_size_n = weight_schema.weight_scale_group_size_n

@@ -251,30 +581,437 @@ def get_humming_moe_quant_config(

581 else:

582 weight_group_shape = GroupShape(row=weight_scale_group_size, col=1)

583

254 - w1_quant_desc = FusedMoEQuantDesc(

255 - dtype=str(weight_schema.b_dtype),

256 - shape=weight_group_shape,

257 - scale=getattr(layer, "w13_weight_scale", None),

258 - alpha_or_gscale=getattr(layer, "w13_global_scale", None),

259 - zp=getattr(layer, "w13_zero_point", None),

260 - bias=getattr(layer, "w13_bias", None),

584 + return make_humming_moe_quant_config(

585 + quant_dtype=q_dtype,

586 + weight_dtype=str(weight_schema.b_dtype),

587 + weight_group_shape=weight_group_shape,

588 + w1_scale=getattr(layer, "w13_weight_scale", None),

589 + w1_gscale=getattr(layer, "w13_global_scale", None),

590 + w1_zp=getattr(layer, "w13_zero_point", None),

591 + w1_bias=getattr(layer, "w13_bias", None),

592 + w2_scale=getattr(layer, "w2_weight_scale", None),

593 + w2_gscale=getattr(layer, "w2_global_scale", None),

594 + w2_zp=getattr(layer, "w2_zero_point", None),

595 + w2_bias=getattr(layer, "w2_bias", None),

596 )

597

263 - w2_quant_desc = FusedMoEQuantDesc(

264 - dtype=str(weight_schema.b_dtype),

265 - shape=weight_group_shape,

266 - scale=getattr(layer, "w2_weight_scale", None),

267 - alpha_or_gscale=getattr(layer, "w2_global_scale", None),

268 - zp=getattr(layer, "w2_zero_point", None),

269 - bias=getattr(layer, "w2_bias", None),

598 +

599 +def select_humming_moe_experts(

600 + config: FusedMoEConfig,

601 + weight_key: QuantKey | None,

602 + activation_key: QuantKey | None,

603 +) -> type[mk.FusedMoEExperts] | None:

604 + """

605 + Select the primary Humming MoE Experts class

606 + Note: Shape-specific fallbacks may still occur at runtime.

607 + """

608 +

609 + if not has_humming():

610 + return None

611 +

612 + # NOTE: the kernels are selected in the following order.

613 + AVAILABLE_EXPERTS: list[type[mk.FusedMoEExperts]] = [

614 + BatchedHummingGroupedExperts,

615 + HummingGroupedExperts,

616 + HummingIndexedExperts,

617 + ]

618 +

619 + # NOTE(rob): We need to peak into the P/F selection to determine

620 + # if we are using the batched or standard expert format, which

621 + # if not ideal. Once we unify TP + DP/EP, we can select P/F first.

622 + activation_format = (

623 + mk.FusedMoEActivationFormat.BatchedExperts

624 + if config.moe_parallel_config.use_batched_activation_format

625 + else mk.FusedMoEActivationFormat.Standard

626 )

627

272 - return FusedMoEQuantConfig(

273 - _a1=a_quant_desc,

274 - _a2=a_quant_desc,

275 - _w1=w1_quant_desc,

276 - _w2=w2_quant_desc,

277 - gemm1_alpha=gemm1_alpha,

278 - gemm1_beta=gemm1_beta,

279 - gemm1_clamp_limit=gemm1_clamp_limit,

628 + def _make_log_backend(experts_cls: type[mk.FusedMoEExperts]):

629 + return f"Using {experts_cls.__name__} Humming MoE backend."

630 +

631 + def _make_log_unsupported(

632 + experts_cls: type[mk.FusedMoEExperts], reason: str | None

633 + ) -> str:

634 + if reason:

635 + return (

636 + f"Humming MoE experts {experts_cls.__name__} does not support the "

637 + f"deployment configuration since {reason}."

638 + )

639 + else:

640 + return (

641 + f"Humming MoE experts '{experts_cls.__name__}' does not support the "

642 + "deployment configuration."

643 + )

644 +

645 + for k_cls in AVAILABLE_EXPERTS:

646 + supported, reason = k_cls.is_supported_config(

647 + k_cls,

648 + config,

649 + weight_key,

650 + activation_key,

651 + activation_format,

652 + )

653 + if supported:

654 + logger.info_once(_make_log_backend(k_cls))

655 + return k_cls

656 + else:

657 + logger.debug_once(_make_log_unsupported(k_cls, reason))

658 +

659 + return None

660 +

661 +

662 +def make_humming_moe_kernel(

663 + moe_quant_config: FusedMoEQuantConfig,

664 + moe_config: FusedMoEConfig,

665 + experts_cls: type[mk.FusedMoEExperts],

666 + layer: "RoutedExperts",

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

668 +) -> mk.FusedMoEKernel:

669 + # Create Prepare/Finalize.

670 + prepare_finalize = maybe_make_prepare_finalize(

671 + moe=moe_config,

672 + quant_config=moe_quant_config,

673 + routing_tables=routing_tables,

674 + allow_new_interface=True,

675 + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),

676 + )

677 + assert prepare_finalize is not None

678 +

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

680 +

681 + extra_args: dict[str, Any] = {"layer": layer}

682 +

683 + # Create Experts.

684 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

685 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()

686 + assert max_num_tokens is not None

687 + experts = experts_cls(

688 + moe_config=moe_config,

689 + quant_config=moe_quant_config,

690 + max_num_tokens=max_num_tokens,

691 + num_dispatchers=prepare_finalize.num_dispatchers(),

692 + **extra_args,

693 + )

694 + else:

695 + experts = experts_cls(

696 + moe_config=moe_config,

697 + quant_config=moe_quant_config,

698 + **extra_args,

699 + )

700 +

701 + kernel = mk.FusedMoEKernel(

702 + prepare_finalize,

703 + experts,

704 + )

705 +

706 + return kernel

707 +

708 +

709 +def _extract_sublayer_tensors(

710 + layer: "RoutedExperts",

711 + sublayer_name: str,

712 +) -> dict[str, torch.Tensor]:

713 + """Extract tensors for a specific sublayer from the layer's state dict."""

714 + return dict(

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

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

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

718 + )

719 +

720 +

721 +def _replace_layer_parameters(

722 + layer: "RoutedExperts",

723 + sublayer_name: str,

724 + tensors: dict[str, torch.Tensor],

725 + preserve_bias: bool = False,

726 +) -> None:

727 + """

728 + Replace layer parameters for a sublayer with new tensors.

729 +

730 + Args:

731 + layer: The RoutedExperts layer

732 + sublayer_name: Name of the sublayer (e.g., "w13", "w2")

733 + tensors: Dict of parameter name to tensor

734 + preserve_bias: If True, don't delete bias parameters

735 + """

736 + # Delete old parameters

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

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

739 + continue

740 + if preserve_bias and name == sublayer_name + "_bias":

741 + continue

742 + delattr(layer, name)

743 +

744 + # Set new parameters

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

746 + param_name = f"{sublayer_name}_{name}"

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

748 + setattr(layer, param_name, param)

749 +

750 +

751 +def _convert_sublayer_to_humming(

752 + layer: "RoutedExperts",

753 + sublayer_name: str,

754 + shape_n: int,

755 + shape_k: int,

756 + weight_schema: Any,

757 + input_schema: Any,

758 + num_experts: int,

759 + param_dtype: torch.dtype,

760 +) -> tuple[Any, Any]:

761 + """

762 + Convert a sublayer's weights from checkpoint format to Humming format.

763 +

764 + Returns:

765 + Tuple of (converted_weight_schema, converted_input_schema)

766 + """

767 + from humming.schema import HummingWeightSchema

768 +

769 + if isinstance(weight_schema, HummingWeightSchema):

770 + # Already in Humming format

771 + return weight_schema, input_schema

772 +

773 + tensors = _extract_sublayer_tensors(layer, sublayer_name)

774 +

775 + shape_k_stacks = [shape_k]

776 + shape_n_stacks = [shape_n]

777 + if sublayer_name == "w13":

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

779 +

780 + converted_weight_schema, converted_tensors = weight_schema.convert_humming(

781 + tensors=tensors,

782 + shape_n_stacks=shape_n_stacks,

783 + shape_k_stacks=shape_k_stacks,

784 + param_dtype=param_dtype,

785 + num_experts=num_experts,

786 + )

787 +

788 + converted_input_schema, _ = input_schema.convert_humming(

789 + tensors=converted_tensors,

790 + shape_n_stacks=shape_n_stacks,

791 + shape_k_stacks=shape_k_stacks,

792 + param_dtype=param_dtype,

793 + num_experts=num_experts,

794 + )

795 +

796 + _replace_layer_parameters(layer, sublayer_name, converted_tensors)

797 +

798 + return converted_weight_schema, converted_input_schema

799 +

800 +

801 +def _prepare_and_transform_sublayer(

802 + layer: "RoutedExperts",

803 + sublayer_name: str,

804 + shape_n: int,

805 + shape_k: int,

806 + weight_schema: Any,

807 + input_schema: Any,

808 + has_bias: bool,

809 + num_experts: int,

810 + param_dtype: torch.dtype,

811 +) -> None:

812 + """

813 + Prepare layer metadata and transform weights for a sublayer.

814 +

815 + This calls Humming's prepare_layer_meta and transform_humming_layer.

816 + """

817 + from humming.layer import HummingMethod

818 +

819 + HummingMethod.prepare_layer_meta(

820 + layer=layer,

821 + shape_n=shape_n,

822 + shape_k=shape_k,

823 + pad_n_to_multiple=256,

824 + pad_k_to_multiple=128,

825 + input_schema=input_schema,

826 + weight_schema=weight_schema,

827 + has_bias=has_bias,

828 + num_experts=num_experts,

829 + torch_dtype=param_dtype,

830 + sublayer_name=sublayer_name,

831 + )

832 +

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

834 +

835 +

836 +def _process_single_sublayer(

837 + layer: "RoutedExperts",

838 + sublayer_name: str,

839 + shape_n: int,

840 + shape_k: int,

841 + weight_schema: Any,

842 + input_schema: Any,

843 + has_bias: bool,

844 + num_experts: int,

845 + param_dtype: torch.dtype,

846 + force_weight_schema: Any | None = None,

847 +) -> tuple[Any, Any]:

848 + """

849 + Process a single sublayer: convert, optionally requant, prepare, and transform.

850 +

851 + This combines the common logic from convert_to_humming_moe_kernel_format

852 + for processing a single sublayer.

853 +

854 + Args:

855 + layer: The RoutedExperts layer

856 + sublayer_name: Name of the sublayer (e.g., "w13", "w2")

857 + shape_n: Output dimension size

858 + shape_k: Input dimension size

859 + weight_schema: Initial weight quantization schema

860 + input_schema: Initial input quantization schema

861 + has_bias: Whether the layer has bias terms

862 + num_experts: Number of experts

863 + param_dtype: Parameter data type

864 + force_weight_schema: Optional schema to force requantization to

865 +

866 + Returns:

867 + Tuple of (final_weight_schema, final_input_schema)

868 + """

869 + from humming.schema import HummingWeightSchema

870 +

871 + # Step 1: Convert from checkpoint format to humming format if needed

872 + current_weight_schema, current_input_schema = _convert_sublayer_to_humming(

873 + layer=layer,

874 + sublayer_name=sublayer_name,

875 + shape_n=shape_n,

876 + shape_k=shape_k,

877 + weight_schema=weight_schema,

878 + input_schema=input_schema,

879 + num_experts=num_experts,

880 + param_dtype=param_dtype,

881 + )

882 +

883 + # Step 2: Force requant if needed

884 + assert isinstance(current_weight_schema, HummingWeightSchema)

885 + if force_weight_schema is not None and current_weight_schema != force_weight_schema:

886 + tensors = _extract_sublayer_tensors(layer, sublayer_name)

887 +

888 + tensors = current_weight_schema.requant_tensors(

889 + tensors=tensors,

890 + target_weight_schema=force_weight_schema,

891 + param_dtype=param_dtype,

892 + )

893 +

894 + current_weight_schema = force_weight_schema

895 + _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True)

896 + del tensors

897 +

898 + # Step 3: Prepare layer metadata and transform weights

899 + _prepare_and_transform_sublayer(

900 + layer=layer,

901 + sublayer_name=sublayer_name,

902 + shape_n=shape_n,

903 + shape_k=shape_k,

904 + weight_schema=current_weight_schema,

905 + input_schema=current_input_schema,

906 + has_bias=has_bias,

907 + num_experts=num_experts,

908 + param_dtype=param_dtype,

909 )

910 +

911 + return current_weight_schema, current_input_schema

912 +

913 +

914 +def convert_to_humming_moe_kernel_format(

915 + layer: "RoutedExperts",

916 + quant_config: dict | None = None,

917 + sublayer_configs: dict[str, Any] | None = None,

918 + weight_schema: Any | None = None,

919 + input_schema: Any | None = None,

920 + force_weight_schema: Any | None = None,

921 +) -> None:

922 + """

923 + Convert MoE weights from checkpoint format to Humming kernel format.

924 +

925 + This function processes weights for each sublayer (w13, w2) by:

926 + 1. Converting from checkpoint format to humming format if needed

927 + 2. Force requanting if a different quantization schema is specified

928 + 3. Preparing layer metadata for the Humming kernel

929 + 4. Transforming weights for inference

930 +

931 + Args:

932 + layer: The RoutedExperts layer containing weights to process

933 + quant_config: Optional quantization config dict. Required if weight_schema

934 + or input_schema are None. Used to build schemas via

935 + BaseWeightSchema.from_config().

936 + sublayer_configs: Optional configuration dict for each sublayer (w13, w2).

937 + Each config must have "shape_n" and "shape_k" keys.

938 + If None, configs are built from layer.moe_config properties.

939 + weight_schema: Optional initial weight quantization schema.

940 + If None, built from quant_config.

941 + input_schema: Optional initial input quantization schema.

942 + If None, built from quant_config or env vars.

943 + force_weight_schema: Optional schema to force requantization to

944 +

945 + Side effects:

946 + - Modifies layer parameters in place

947 + - Sets layer.weight_schemas and layer.input_schemas

948 + """

949 +

950 + # Build schemas from quant_config if not provided

951 + has_bias = layer.moe_config.has_bias

952 + num_experts = layer.moe_config.num_local_experts

953 + param_dtype = layer.params_dtype

954 +

955 + if weight_schema is None or input_schema is None:

956 + if quant_config is None:

957 + raise ValueError(

958 + "Must provide either weight_schema/input_schema or quant_config"

959 + )

960 +

961 + from humming.layer import HummingInputSchema

962 + from humming.schema import BaseWeightSchema

963 +

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

965 + humming_is_layer_skipped,

966 + )

967 +

968 + if weight_schema is None:

969 + weight_schema = BaseWeightSchema.from_config(quant_config)

970 +

971 + if input_schema is None:

972 + input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}

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

974 + input_schema = HummingInputSchema()

975 + else:

976 + # TODO: read input_quant_config from quant_config

977 + input_schema = HummingInputSchema.from_config(input_quant_config)

978 +

979 + # Build sublayer configs from layer properties if not provided

980 + if sublayer_configs is None:

981 + is_gated = layer.moe_config.activation.is_gated

982 + sublayer_configs = {

983 + "w13": {

984 + "shape_n": layer.moe_config.intermediate_size_per_partition * 2,

985 + "shape_k": layer.moe_config.hidden_dim,

986 + },

987 + "w2": {

988 + "shape_n": layer.moe_config.hidden_dim,

989 + "shape_k": layer.moe_config.intermediate_size_per_partition

990 + * (1 if is_gated else 2),

991 + },

992 + }

993 +

994 + layer.weight_schemas = {}

995 + layer.input_schemas = {}

996 +

997 + for sublayer_name, configs in sublayer_configs.items():

998 + final_weight_schema, final_input_schema = _process_single_sublayer(

999 + layer=layer,

1000 + sublayer_name=sublayer_name,

1001 + shape_n=configs["shape_n"],

1002 + shape_k=configs["shape_k"],

1003 + weight_schema=weight_schema,

1004 + input_schema=input_schema,

1005 + has_bias=has_bias,

1006 + num_experts=num_experts,

1007 + param_dtype=param_dtype,

1008 + force_weight_schema=force_weight_schema,

1009 + )

1010 +

1011 + layer.weight_schemas[sublayer_name] = final_weight_schema

1012 + layer.input_schemas[sublayer_name] = final_input_schema

1013 +

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

1015 + device = layer.w13_weight.device

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

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

vllm/utils/humming.py

+13

@@ -15,6 +15,7 @@ _EXPORTS: dict[str, str] = {

15 "dtypes": "humming.dtypes",

16 "DataType": "humming.dtypes:DataType",

17 "GemmType": "humming.config:GemmType",

18 + "WeightScaleType": "humming.config:WeightScaleType",

19 "HummingMethod": "humming.layer:HummingMethod",

20 "HummingLayerMeta": "humming.layer:HummingLayerMeta",

21 "BaseInputSchema": "humming.schema:BaseInputSchema",

@@ -22,6 +23,18 @@ _EXPORTS: dict[str, str] = {

23 "HummingInputSchema": "humming.schema:HummingInputSchema",

24 "HummingWeightSchema": "humming.schema:HummingWeightSchema",

25 "quantize_weight": "humming.utils.weight:quantize_weight",

26 + "AWQWeightSchema": "humming.schema:AWQWeightSchema",

27 + "BitnetWeightSchema": "humming.schema:BitnetWeightSchema",

28 + "ModeloptMxfp8WeightSchema": "humming.schema.modelopt:ModeloptMxfp8WeightSchema",

29 + "ModeloptNvfp4InputSchema": "humming.schema.modelopt:ModeloptNvfp4InputSchema",

30 + "ModeloptNvfp4WeightSchema": "humming.schema.modelopt:ModeloptNvfp4WeightSchema",

31 + "CompressedTensorsInputSchema": "humming.schema:CompressedTensorsInputSchema",

32 + "CompressedTensorsWeightSchema": "humming.schema:CompressedTensorsWeightSchema",

33 + "Fp8InputSchema": "humming.schema:Fp8InputSchema",

34 + "Fp8WeightSchema": "humming.schema.fp8:Fp8WeightSchema",

35 + "Mxfp4WeightSchema": "humming.schema:Mxfp4WeightSchema",

36 + "GptOssMxfp4WeightSchema": "humming.schema:GptOssMxfp4WeightSchema",

37 + "GPTQWeightSchema": "humming.schema:GPTQWeightSchema",

38 }