@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0ae89f18fd

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 5, 22:15 UTC

[Refactor] Move FusedMoE hidden_size roundup to quant_method (#34285)

Signed-off-by: Bowen Bao <bowenbao@amd.com>

Bowen Bao committed Mar 26, 2026 at 23:38 UTC 0ae89f18fd75c5fcac48ff711ed84464c3da5d33

12 files changed +204 -222

vllm/model_executor/layers/fused_moe/config.py

+12 -3

@@ -229,9 +229,6 @@ class FusedMoEQuantConfig:

229 _w1: FusedMoEQuantDesc

230 _w2: FusedMoEQuantDesc

231 is_nvfp4_scale_swizzled: bool = True

232 - # CK MXFP4 (gfx950) padding info for rocm_aiter_ops.fused_moe()

233 - hidden_pad: int = 0

234 - intermediate_pad: int = 0

232

233 def __post_init__(self):

234 assert not self.per_act_token_quant or self.block_shape is None, (

@@ -1172,6 +1169,11 @@ class FusedMoEConfig:

1169 # Defaults to in_dtype if not specified.

1170 router_logits_dtype: torch.dtype | None = None

1171

1172 + # Defaults to hidden_dim if not specified.

1173 + hidden_dim_unpadded: int | None = None

1174 + # Defaults to intermediate_size_per_partition if not specified.

1175 + intermediate_size_per_partition_unpadded: int | None = None

1176 +

1177 moe_backend: str = "auto"

1178 max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE

1179 has_bias: bool = False

@@ -1195,6 +1197,13 @@ class FusedMoEConfig:

1197 if self.router_logits_dtype is None:

1198 self.router_logits_dtype = self.in_dtype

1199

1200 + if self.hidden_dim_unpadded is None:

1201 + self.hidden_dim_unpadded = self.hidden_dim

1202 + if self.intermediate_size_per_partition_unpadded is None:

1203 + self.intermediate_size_per_partition_unpadded = (

1204 + self.intermediate_size_per_partition

1205 + )

1206 +

1207 @property

1208 def tp_size(self):

1209 return self.moe_parallel_config.tp_size

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+33

@@ -9,6 +9,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.logger import init_logger

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

11 FusedMoEConfig,

12 + FusedMoEParallelConfig,

13 FusedMoEQuantConfig,

14 )

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

@@ -65,6 +66,38 @@ class FusedMoEMethodBase(QuantizeMethodBase):

66 """

67 return False

69 + def maybe_roundup_sizes(

70 + self,

71 + hidden_size: int,

72 + intermediate_size_per_partition: int,

73 + act_dtype: torch.dtype,

74 + moe_parallel_config: FusedMoEParallelConfig,

75 + ) -> tuple[int, int]:

76 + """

77 + Given layer hidden size and intermediate size per partition and MoE

78 + configurations, round up hidden_size and intermediate_size_per_partition

79 + if necessary.

80 +

81 + Args:

82 + hidden_size: Layer hidden-size

83 + intermediate_size_per_partition: Intermediate size per partition for

84 + the layer.

85 + act_dtype: Data type of the layer activations.

86 + moe_parallel_config: Fused MoE parallelization strategy configuration.

87 +

88 + Return:

89 + A tuple of (rounded_hidden_size, rounded_intermediate_size_per_partition),

90 + where:

91 + - rounded_hidden_size is the possibly rounded up hidden size.

92 + - rounded_intermediate_size_per_partition is the possibly rounded

93 + up intermediate size per partition.

94 + """

95 + from .all2all_utils import maybe_roundup_layer_hidden_size

96 +

97 + return maybe_roundup_layer_hidden_size(

98 + hidden_size, act_dtype, moe_parallel_config

99 + ), intermediate_size_per_partition

100 +

101 def maybe_make_prepare_finalize(

102 self,

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

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+3 -13

@@ -428,13 +428,9 @@ def triton_kernel_fused_mxfp4_w4a8_experts(

428 assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32

429 assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32

430

431 - # Shape check: when weights are padded (e.g. hidden_size padded for

432 - # GFX950 swizzle), unpadded_K_w1 carries the original dimension.

433 - expected_K_w1 = unpadded_K_w1 if unpadded_K_w1 is not None else w1.shape[-2]

434 - assert hidden_states.shape[-1] == expected_K_w1, (

435 - f"hidden_states K={hidden_states.shape[-1]} != "

436 - f"expected K={expected_K_w1} (w1 K={w1.shape[-2]})"

437 - )

431 + # Shape check: weights are padded (e.g. hidden_size padded for

432 + # GFX950 swizzle).

433 + assert hidden_states.shape[-1] == w1.shape[-2]

434 assert w2.shape[-1] == w1.shape[1]

435

436 E, _, N = w1.shape

@@ -494,12 +490,6 @@ def triton_kernel_fused_mxfp4_w4a8_experts(

490 unpadded_K=unpadded_K_w2,

491 )

492

497 - # When hidden_size was padded for alignment (e.g. GFX950 swizzle),

498 - # the kernel output has the padded dimension. Slice back to the

499 - # original hidden_size so downstream layers see the expected shape.

500 - if unpadded_N_w2 is not None and intermediate_cache3.shape[-1] != unpadded_N_w2:

501 - intermediate_cache3 = intermediate_cache3[..., :unpadded_N_w2].contiguous()

502 -

493 return intermediate_cache3

494

495

vllm/model_executor/layers/fused_moe/layer.py

+63 -71

@@ -210,42 +210,6 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:

210 )

211

212

213 -# TODO(rob): move this down to the kernel.

214 -def maybe_roundup_hidden_size(

215 - hidden_size: int,

216 - act_dtype: torch.dtype,

217 - moe_parallel_config: FusedMoEParallelConfig,

218 - is_lora_enabled: bool,

219 - model_type: str | None,

220 -) -> int:

221 - """

222 - Given layer hidden size and MoE configurations, round up hidden_size

223 - if necessary.

224 -

225 - Args:

226 - hidden_size: Layer hidden-size

227 - act_dtype: Data type of the layer activations.

228 - moe_parallel_config: Fused MoE parallelization strategy configuration.

229 - is_lora_enabled: True if the engine is enabled with LoRA. This

230 - is used in the case of mxfp4 quantization in selecting the

231 - MxFP4Backend.

232 - model_type: for checking if gpt-oss

233 -

234 - Return:

235 - Rounded up hidden_size if rounding up is required based on the configs.

236 - Original hidden size otherwise.

237 - """

238 - from vllm.model_executor.layers.fused_moe.all2all_utils import (

239 - maybe_roundup_layer_hidden_size,

240 - )

241 -

242 - hidden_size = maybe_roundup_layer_hidden_size(

243 - hidden_size, act_dtype, moe_parallel_config

244 - )

245 -

246 - return hidden_size

247 -

248 -

213 # --8<-- [start:fused_moe]

214 @CustomOp.register("fused_moe")

215 class FusedMoE(CustomOp):

@@ -459,7 +423,7 @@ class FusedMoE(CustomOp):

423 ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."

424

425 assert intermediate_size % self.tp_size == 0

462 - self.intermediate_size_per_partition = intermediate_size // self.tp_size

426 + intermediate_size_per_partition = intermediate_size // self.tp_size

427 self.reduce_results = reduce_results

428 self.renormalize = renormalize

429

@@ -501,28 +465,13 @@ class FusedMoE(CustomOp):

465 )

466 self.routing_method_type: RoutingMethodType = self.router.routing_method_type

467

504 - # Round up hidden size before creating moe_config.

505 - # This way moe_config is created with the correct hidden_size from the start.

506 - unpadded_hidden_size = hidden_size

507 - self.model_type = (

508 - self.vllm_config.model_config.hf_config.model_type

509 - if self.vllm_config.model_config is not None

510 - else None

511 - )

512 - hidden_size = maybe_roundup_hidden_size(

513 - hidden_size=hidden_size,

514 - act_dtype=moe_in_dtype,

515 - moe_parallel_config=self.moe_parallel_config,

516 - is_lora_enabled=vllm_config.lora_config is not None,

517 - model_type=self.model_type,

518 - )

519 - self.hidden_size = hidden_size

520 -

468 self.moe_config: FusedMoEConfig = FusedMoEConfig(

469 num_experts=self.global_num_experts,

470 experts_per_token=top_k,

471 hidden_dim=hidden_size,

525 - intermediate_size_per_partition=self.intermediate_size_per_partition,

472 + hidden_dim_unpadded=hidden_size,

473 + intermediate_size_per_partition=intermediate_size_per_partition,

474 + intermediate_size_per_partition_unpadded=intermediate_size_per_partition,

475 num_local_experts=self.local_num_experts,

476 num_logical_experts=self.logical_num_experts,

477 moe_parallel_config=self.moe_parallel_config,

@@ -567,13 +516,6 @@ class FusedMoE(CustomOp):

516 # for heuristic purposes, so it must be initialized first.

517 self.quant_method: FusedMoEMethodBase = _get_quant_method()

518

570 - # Quant methods (e.g. Mxfp4MoEMethod) may round up hidden_dim

571 - # and intermediate_size in moe_config during __init__. Sync

572 - # self.hidden_size so downstream consumers (e.g. LoRA) see the

573 - # padded value.

574 - if self.moe_config.hidden_dim != self.hidden_size:

575 - self.hidden_size = self.moe_config.hidden_dim

576 -

519 if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike():

520 raise NotImplementedError(

521 "is_act_and_mul=False is supported only for CUDA and ROCm for now"

@@ -591,11 +533,24 @@ class FusedMoE(CustomOp):

533 f"EPLB is not supported {self.quant_method.__class__.__name__}."

534 )

535

536 + # Round up hidden size and update moe_config.

537 + hidden_size, intermediate_size_per_partition = (

538 + self.quant_method.maybe_roundup_sizes(

539 + hidden_size,

540 + intermediate_size_per_partition,

541 + moe_in_dtype,

542 + self.moe_parallel_config,

543 + )

544 + )

545 + self.moe_config.hidden_dim = hidden_size

546 + self.moe_config.intermediate_size_per_partition = (

547 + intermediate_size_per_partition

548 + )

549 +

550 moe_quant_params = {

551 "num_experts": self.local_num_experts,

596 - "hidden_size": self.hidden_size,

597 - "unpadded_hidden_size": unpadded_hidden_size,

598 - "intermediate_size_per_partition": self.intermediate_size_per_partition,

552 + "hidden_size": hidden_size,

553 + "intermediate_size_per_partition": intermediate_size_per_partition,

554 "params_dtype": params_dtype,

555 "weight_loader": self.weight_loader,

556 "global_num_experts": self.global_num_experts,

@@ -933,9 +888,17 @@ class FusedMoE(CustomOp):

888 # Only narrow if the loaded_weight is not a scalar (0-dim tensor)

889 # and we're not loading the full weight

890 if not load_full and loaded_weight.ndim > 0:

936 - loaded_weight = loaded_weight.narrow(

937 - shard_dim, shard_size * tp_rank, shard_size

938 - )

891 + # Handle padding: loaded_weight might be smaller than shard_size on last

892 + # TP rank

893 + start_offset = shard_size * tp_rank

894 + available = loaded_weight.shape[shard_dim] - start_offset

895 + if available <= 0:

896 + # If there is no available weight to load for this TP rank

897 + # (can happen on last TP rank with padding), we can skip

898 + # loading and return early

899 + return

900 + narrow_size = min(shard_size, available)

901 + loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)

902 # Narrow parameter and load.

903 # w1, gate_proj: Load into first logical weight of w13.

904 if shard_id == "w1":

@@ -944,6 +907,13 @@ class FusedMoE(CustomOp):

907 else:

908 assert shard_id == "w3"

909 expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)

910 +

911 + # Handle padding: if loaded_weight is smaller than expert_data (can happen

912 + # on last TP shard with padding), copy to top-left corner

913 + if expert_data.shape != loaded_weight.shape:

914 + expert_data = expert_data[

915 + : loaded_weight.shape[0], : loaded_weight.shape[1]

916 + ]

917 expert_data.copy_(loaded_weight)

918

919 def _load_w2(

@@ -961,10 +931,24 @@ class FusedMoE(CustomOp):

931 # Only narrow if the loaded_weight is not a scalar (0-dim tensor)

932 # and we're not loading the full weight

933 if not load_full and loaded_weight.ndim > 0:

964 - loaded_weight = loaded_weight.narrow(

965 - shard_dim, shard_size * tp_rank, shard_size

966 - )

934 + # Handle padding: loaded_weight might be smaller than shard_size on last

935 + # TP rank

936 + start_offset = shard_size * tp_rank

937 + available = loaded_weight.shape[shard_dim] - start_offset

938 + if available <= 0:

939 + # If there is no available weight to load for this TP rank

940 + # (can happen on last TP rank with padding), we can skip

941 + # loading and return early

942 + return

943 + narrow_size = min(shard_size, available)

944 + loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)

945 # w2, down_proj: Load into only logical weight of w2.

946 + # Handle padding: if loaded_weight is smaller than expert_data (can happen

947 + # on last TP shard with padding), copy to top-left corner

948 + if expert_data.shape != loaded_weight.shape:

949 + expert_data = expert_data[

950 + : loaded_weight.shape[0], : loaded_weight.shape[1]

951 + ]

952 expert_data.copy_(loaded_weight)

953

954 def _load_single_value(

@@ -1549,6 +1533,14 @@ class FusedMoE(CustomOp):

1533 ]

1534 ]

1535

1536 + @property

1537 + def hidden_size(self) -> int:

1538 + return self.moe_config.hidden_dim

1539 +

1540 + @property

1541 + def intermediate_size_per_partition(self) -> int:

1542 + return self.moe_config.intermediate_size_per_partition

1543 +

1544 def extra_repr(self) -> str:

1545 s = (

1546 f"global_num_experts={self.global_num_experts}, "

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

+1 -7

@@ -779,8 +779,6 @@ def make_mxfp4_moe_quant_config(

779 w2_scale: Union[torch.Tensor, "PrecisionConfig"],

780 w1_bias: torch.Tensor | None = None,

781 w2_bias: torch.Tensor | None = None,

782 - hidden_pad: int = 0,

783 - intermediate_pad: int = 0,

782 ) -> FusedMoEQuantConfig | None:

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

784 if mxfp4_backend in (

@@ -802,16 +800,12 @@ def make_mxfp4_moe_quant_config(

800 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,

801 Mxfp4MoeBackend.CK,

802 ):

805 - config = mxfp4_w4a16_moe_quant_config(

803 + return mxfp4_w4a16_moe_quant_config(

804 w1_bias=w1_bias,

805 w2_bias=w2_bias,

806 w1_scale=w1_scale,

807 w2_scale=w2_scale,

808 )

811 - if mxfp4_backend == Mxfp4MoeBackend.CK:

812 - config.hidden_pad = hidden_pad

813 - config.intermediate_pad = intermediate_pad

814 - return config

809 else:

810 return ocp_mx_moe_quant_config(

811 quant_dtype="mxfp4",

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+16 -2

@@ -10,6 +10,7 @@ from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

12 FUSED_MOE_UNQUANTIZED_CONFIG,

13 + FusedMoEConfig,

14 FusedMoEParallelConfig,

15 FusedMoEQuantConfig,

16 )

@@ -186,6 +187,7 @@ def rocm_aiter_fused_experts(

187 w2: torch.Tensor,

188 topk_weights: torch.Tensor,

189 topk_ids: torch.Tensor,

190 + moe_config: FusedMoEConfig,

191 activation: MoEActivation = MoEActivation.SILU,

192 apply_router_weight_on_input: bool = False,

193 expert_map: torch.Tensor | None = None,

@@ -276,6 +278,17 @@ def rocm_aiter_fused_experts(

278 "Only support topk=1 when `apply_router_weight_on_input` is True"

279 )

280

281 + # Compute padding on-the-fly for CK MXFP4 kernels

282 + hidden_pad = 0

283 + intermediate_pad = 0

284 + assert moe_config.hidden_dim_unpadded is not None

285 + assert moe_config.intermediate_size_per_partition_unpadded is not None

286 + hidden_pad = hidden_states.shape[1] - moe_config.hidden_dim_unpadded

287 + intermediate_pad = (

288 + moe_config.intermediate_size_per_partition

289 + - moe_config.intermediate_size_per_partition_unpadded

290 + )

291 +

292 return rocm_aiter_ops.fused_moe(

293 hidden_states,

294 w1,

@@ -292,8 +305,8 @@ def rocm_aiter_fused_experts(

305 doweight_stage1=apply_router_weight_on_input,

306 num_local_tokens=num_local_tokens,

307 output_dtype=output_dtype,

295 - hidden_pad=quant_config.hidden_pad,

296 - intermediate_pad=quant_config.intermediate_pad,

308 + hidden_pad=hidden_pad,

309 + intermediate_pad=intermediate_pad,

310 bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,

311 bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,

312 )

@@ -419,6 +432,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):

432 apply_router_weight_on_input=apply_router_weight_on_input,

433 expert_map=expert_map,

434 quant_config=self.quant_config,

435 + moe_config=self.moe_config,

436 a1q_scale=a1q_scale,

437 num_local_tokens=num_local_tokens,

438 output_dtype=output.dtype,

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

-4

@@ -715,8 +715,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

715 params_dtype: torch.dtype,

716 **extra_weight_attrs,

717 ):

718 - layer.intermediate_size_per_partition = intermediate_size_per_partition

719 - layer.hidden_size = hidden_size

718 layer.num_experts = num_experts

719 layer.orig_dtype = params_dtype

720 layer.weight_block_size = None

@@ -2274,8 +2272,6 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2272 params_dtype: torch.dtype,

2273 **extra_weight_attrs,

2274 ):

2277 - layer.intermediate_size_per_partition = intermediate_size_per_partition

2278 - layer.hidden_size = hidden_size

2275 layer.num_experts = num_experts

2276 layer.orig_dtype = params_dtype

2277 layer.weight_block_size = None

vllm/model_executor/layers/quantization/fp8.py

-4

@@ -672,8 +672,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

672 params_dtype: torch.dtype,

673 **extra_weight_attrs,

674 ):

675 - layer.intermediate_size_per_partition = intermediate_size_per_partition

676 - layer.hidden_size = hidden_size

675 layer.num_experts = num_experts

676 layer.orig_dtype = params_dtype

677 layer.weight_block_size = None

@@ -1011,8 +1009,6 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1009 params_dtype: torch.dtype,

1010 **extra_weight_attrs,

1011 ):

1014 - layer.intermediate_size_per_partition = intermediate_size_per_partition

1015 - layer.hidden_size = hidden_size

1012 layer.num_experts = num_experts

1013 layer.orig_dtype = params_dtype

1014 layer.weight_block_size = None

vllm/model_executor/layers/quantization/mxfp4.py

+25 -37

@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe import (

13 )

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

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

16 + FusedMoEParallelConfig,

17 FusedMoEQuantConfig,

18 )

19 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

@@ -107,18 +108,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

108 self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

109 self.moe_kernel: mk.FusedMoEKernel | None = None

110

110 - # Round up dims once based on backend. This mutates the shared

111 - # FusedMoEConfig in-place so that create_weights() and all

112 - # downstream code see the padded dimensions. This must happen

113 - # before create_weights() is called.

114 - self.moe.hidden_dim, self.moe.intermediate_size_per_partition = (

115 - mxfp4_round_up_hidden_size_and_intermediate_size(

116 - self.mxfp4_backend,

117 - self.moe.hidden_dim,

118 - self.moe.intermediate_size_per_partition,

119 - )

120 - )

121 -

111 # Used for triton kernel precision configs

112 self.w13_precision_config = None

113 self.w2_precision_config = None

@@ -129,6 +118,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

118 # so can skip the padding in the forward before applying the moe method

119 return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8

120

121 + def maybe_roundup_sizes(

122 + self,

123 + hidden_size: int,

124 + intermediate_size_per_partition: int,

125 + act_dtype: torch.dtype,

126 + moe_parallel_config: FusedMoEParallelConfig,

127 + ) -> tuple[int, int]:

128 + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(

129 + hidden_size=hidden_size,

130 + intermediate_size_per_partition=intermediate_size_per_partition,

131 + act_dtype=act_dtype,

132 + moe_parallel_config=moe_parallel_config,

133 + )

134 + return mxfp4_round_up_hidden_size_and_intermediate_size(

135 + self.mxfp4_backend, hidden_size, intermediate_size_per_partition

136 + )

137 +

138 def create_weights(

139 self,

140 layer: torch.nn.Module,

@@ -143,32 +149,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

149 scale_dtype = torch.uint8

150 mxfp4_block = 32

151

146 - # Use pre-rounded sizes from config

147 - self.intermediate_size = intermediate_size_per_partition_after_pad = (

148 - self.moe.intermediate_size_per_partition

149 - )

150 - self.hidden_size = hidden_size = self.moe.hidden_dim

151 -

152 - # Expose padded dimensions on the layer for LoRA and Marlin code

153 - # that reads layer.hidden_size / layer.intermediate_size_per_partition.

152 layer.params_dtype = params_dtype

153 layer.num_experts = num_experts

156 - layer.hidden_size = hidden_size

157 - layer.intermediate_size_per_partition = (

158 - intermediate_size_per_partition_after_pad

159 - )

160 -

161 - # CK (gfx950) padding info for rocm_aiter_ops.fused_moe()

162 - self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)

163 - self.intermediate_pad = (

164 - intermediate_size_per_partition_after_pad - intermediate_size_per_partition

165 - )

154 + self.intermediate_size = intermediate_size_per_partition

155 + self.hidden_size = hidden_size

156

157 # Fused gate_up_proj (column parallel)

158 w13_weight = torch.nn.Parameter(

159 torch.zeros(

160 num_experts,

171 - 2 * intermediate_size_per_partition_after_pad,

161 + 2 * intermediate_size_per_partition,

162 hidden_size // 2,

163 dtype=weight_dtype,

164 ),

@@ -180,7 +170,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

170 w13_weight_scale = torch.nn.Parameter(

171 torch.zeros(

172 num_experts,

183 - 2 * intermediate_size_per_partition_after_pad,

173 + 2 * intermediate_size_per_partition,

174 hidden_size // mxfp4_block,

175 dtype=scale_dtype,

176 ),

@@ -194,7 +184,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

184 torch.zeros(

185 num_experts,

186 hidden_size,

197 - intermediate_size_per_partition_after_pad // 2,

187 + intermediate_size_per_partition // 2,

188 dtype=weight_dtype,

189 ),

190 requires_grad=False,

@@ -206,7 +196,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

196 torch.zeros(

197 num_experts,

198 hidden_size,

209 - intermediate_size_per_partition_after_pad // mxfp4_block,

199 + intermediate_size_per_partition // mxfp4_block,

200 dtype=scale_dtype,

201 ),

202 requires_grad=False,

@@ -218,7 +208,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

208 w13_bias = torch.nn.Parameter(

209 torch.zeros(

210 num_experts,

221 - 2 * intermediate_size_per_partition_after_pad,

211 + 2 * intermediate_size_per_partition,

212 dtype=torch.bfloat16,

213 ),

214 requires_grad=False,

@@ -368,8 +358,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

358 w2_scale=w2_scale,

359 w1_bias=w1_bias,

360 w2_bias=w2_bias,

371 - hidden_pad=self.hidden_pad,

372 - intermediate_pad=self.intermediate_pad,

361 )

362

363 def select_gemm_impl(

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

+51 -79

@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import (

18 MoEActivation,

19 )

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

21 + FusedMoEParallelConfig,

22 FusedMoEQuantConfig,

23 fp8_w8a8_moe_quant_config,

24 mxfp4_w4a8_moe_quant_config,

@@ -27,13 +28,13 @@ from vllm.model_executor.layers.fused_moe.config import (

28 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

29 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

30 Mxfp4MoeBackend,

31 + mxfp4_round_up_hidden_size_and_intermediate_size,

32 select_mxfp4_moe_backend,

33 )

34 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

35 prepare_fp8_moe_layer_for_marlin,

36 )

37 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (

36 - CK_MXFP4_MOE_DIM_ALIGNMENT,

38 _swizzle_mxfp4,

39 )

40 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

@@ -49,7 +50,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

50 from vllm.model_executor.utils import set_weight_attrs

51 from vllm.platforms import current_platform

52 from vllm.scalar_type import scalar_types

52 -from vllm.utils.math_utils import round_up

54 logger = init_logger(__name__)

@@ -173,8 +173,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

173 params_dtype: torch.dtype,

174 **extra_weight_attrs,

175 ):

176 - layer.intermediate_size_per_partition = intermediate_size_per_partition

177 - layer.hidden_size = hidden_size

176 layer.num_experts = num_experts

177 layer.orig_dtype = params_dtype

178 layer.weight_block_size = None

@@ -182,7 +180,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

180

181 # WEIGHTS

182 w13_weight = torch.nn.Parameter(

185 - torch.empty(

183 + torch.zeros(

184 num_experts,

185 2 * intermediate_size_per_partition,

186 hidden_size,

@@ -194,7 +192,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

192 set_weight_attrs(w13_weight, extra_weight_attrs)

193

194 w2_weight = torch.nn.Parameter(

197 - torch.empty(

195 + torch.zeros(

196 num_experts,

197 hidden_size,

198 intermediate_size_per_partition,

@@ -461,6 +459,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

459 activation=layer.activation,

460 apply_router_weight_on_input=layer.apply_router_weight_on_input,

461 quant_config=self.moe_quant_config,

462 + moe_config=layer.moe_config,

463 expert_map=layer.expert_map,

464 )

465 elif self.use_marlin:

@@ -527,7 +526,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

526 ):

527 params_dtype = torch.uint32

528 w13_weight = torch.nn.Parameter(

530 - torch.empty(

529 + torch.zeros(

530 num_experts,

531 2 * intermediate_size_per_partition,

532 hidden_size // 8, # INT32 packing for W4

@@ -536,7 +535,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

535 requires_grad=False,

536 )

537 w2_weight = torch.nn.Parameter(

539 - torch.empty(

538 + torch.zeros(

539 num_experts,

540 hidden_size,

541 intermediate_size_per_partition // 8, # INT32 packing for W4

@@ -649,6 +648,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

648 activation=layer.activation,

649 apply_router_weight_on_input=layer.apply_router_weight_on_input,

650 quant_config=self.moe_quant_config,

651 + moe_config=layer.moe_config,

652 expert_map=layer.expert_map,

653 )

654

@@ -702,6 +702,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

702 self.mxfp4_backend: Mxfp4MoeBackend | None = None

703 if self.ocp_mx_scheme == "w_mxfp4":

704 self.mxfp4_backend, _ = select_mxfp4_moe_backend(moe)

705 + elif self.ocp_mx_scheme.startswith("w_mxfp4"):

706 + # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes.

707 + self.mxfp4_backend = Mxfp4MoeBackend.NONE

708

709 if self.input_quant is not None:

710 self.static_input_scales = not self.input_quant.get("is_dynamic")

@@ -734,36 +737,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

737 self.emulate = (

738 not current_platform.supports_mx()

739 or not self.ocp_mx_scheme.startswith("w_mxfp4")

737 - ) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe)

738 -

739 - # CK's pre-compiled MXFP4 MoE GEMM kernel instances have dimension

740 - # alignment requirements. When violated (e.g. MiniMax-M2.1 with

741 - # TP=4 yields intermediate_size_per_partition=384), AITER raises:

742 - # "device_gemm ... does not support this GEMM problem".

743 - # Fall back to emulation in that case.

744 - # For gpt_oss models, create_weights rounds up the dimensions

745 - # internally, so the alignment check is skipped.

746 - if (

747 - not self.emulate

748 - and self.use_rocm_aiter_moe

749 - and self.ocp_mx_scheme is not None

750 - and self.ocp_mx_scheme.startswith("w_mxfp4")

751 - and self.model_type != "gpt_oss"

752 - and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0

753 - ):

754 - logger.warning_once(

755 - "AITER CK MXFP4 MoE GEMM does not support "

756 - "intermediate_size_per_partition=%d (not a multiple of %d). "

757 - "This typically happens when intermediate_size / "

758 - "tensor_parallel_size produces an incompatible dimension. "

759 - "Falling back to emulation mode. To avoid this overhead, "

760 - "use a compatible tensor_parallel_size or set "

761 - "VLLM_ROCM_USE_AITER_MOE=0.",

762 - moe.intermediate_size_per_partition,

763 - CK_MXFP4_MOE_DIM_ALIGNMENT,

764 - )

765 - self.use_rocm_aiter_moe = False

766 - self.emulate = True

740 + ) and (

741 + self.mxfp4_backend is None

742 + or self.mxfp4_backend is Mxfp4MoeBackend.NONE

743 + or not self.use_rocm_aiter_moe

744 + )

745

746 if self.emulate:

747 logger.warning_once(

@@ -780,6 +758,27 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

758 "The current mode supports native MoE MXFP4 computation"

759 )

760

761 + def maybe_roundup_sizes(

762 + self,

763 + hidden_size: int,

764 + intermediate_size_per_partition: int,

765 + act_dtype: torch.dtype,

766 + moe_parallel_config: FusedMoEParallelConfig,

767 + ) -> tuple[int, int]:

768 + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(

769 + hidden_size=hidden_size,

770 + intermediate_size_per_partition=intermediate_size_per_partition,

771 + act_dtype=act_dtype,

772 + moe_parallel_config=moe_parallel_config,

773 + )

774 + if self.mxfp4_backend is not None:

775 + hidden_size, intermediate_size_per_partition = (

776 + mxfp4_round_up_hidden_size_and_intermediate_size(

777 + self.mxfp4_backend, hidden_size, intermediate_size_per_partition

778 + )

779 + )

780 + return hidden_size, intermediate_size_per_partition

781 +

782 def get_packed_dim(self, dim: int, quant_dtype: str):

783 if quant_dtype == "mxfp4":

784 assert dim % 2 == 0

@@ -805,40 +804,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

804 )

805

806 params_dtype = torch.uint8

808 - self.intermediate_size_per_partition = intermediate_size_per_partition

809 - if self.model_type == "gpt_oss":

810 - if current_platform.is_rocm():

811 - intermediate_size_per_partition_after_pad = round_up(

812 - intermediate_size_per_partition, 256

813 - )

814 - else:

815 - intermediate_size_per_partition_after_pad = round_up(

816 - intermediate_size_per_partition, 64

817 - )

818 - else:

819 - intermediate_size_per_partition_after_pad = intermediate_size_per_partition

820 -

821 - self.unpadded_hidden_size = extra_weight_attrs.get(

822 - "unpadded_hidden_size", hidden_size

823 - )

824 -

825 - # On GFX950, the GFX950MXScaleLayout swizzle requires

826 - # hidden_size to be a multiple of 256 (SCALE_K = hidden_size / 32

827 - # must be divisible by 8). Pad hidden_size for weight/scale

828 - # allocation; the original value is preserved in unpadded_hidden_size.

829 - # Only applies to the native (non-emulated) CK path on GFX950.

830 - if (

831 - self.model_type == "gpt_oss"

832 - and current_platform.is_rocm()

833 - and not self.emulate

834 - ):

835 - hidden_size = round_up(hidden_size, 256)

807

808 # WEIGHTS

809 w13_weight = torch.nn.Parameter(

839 - torch.empty(

810 + torch.zeros(

811 num_experts,

841 - 2 * intermediate_size_per_partition_after_pad,

812 + 2 * intermediate_size_per_partition,

813 self.get_packed_dim(hidden_size, self.weight_dtype),

814 dtype=params_dtype,

815 ),

@@ -849,12 +820,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

820 set_weight_attrs(w13_weight, extra_weight_attrs)

821

822 w2_weight = torch.nn.Parameter(

852 - torch.empty(

823 + torch.zeros(

824 num_experts,

825 hidden_size,

855 - self.get_packed_dim(

856 - intermediate_size_per_partition_after_pad, self.weight_dtype

857 - ),

826 + self.get_packed_dim(intermediate_size_per_partition, self.weight_dtype),

827 dtype=params_dtype,

828 ),

829 requires_grad=False,

@@ -867,7 +836,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

836 w13_weight_scale = torch.nn.Parameter(

837 torch.ones(

838 num_experts,

870 - 2 * intermediate_size_per_partition_after_pad,

839 + 2 * intermediate_size_per_partition,

840 hidden_size // OCP_MX_BLOCK_SIZE,

841 dtype=params_dtype,

842 ),

@@ -877,7 +846,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

846 torch.ones(

847 num_experts,

848 hidden_size,

880 - intermediate_size_per_partition_after_pad // OCP_MX_BLOCK_SIZE,

849 + intermediate_size_per_partition // OCP_MX_BLOCK_SIZE,

850 dtype=params_dtype,

851 ),

852 requires_grad=False,

@@ -892,7 +861,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

861 w13_bias = torch.nn.Parameter(

862 torch.zeros(

863 num_experts,

895 - 2 * intermediate_size_per_partition_after_pad,

864 + 2 * intermediate_size_per_partition,

865 dtype=torch.float32,

866 ),

867 requires_grad=False,

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

1041 topk_ids=topk_ids,

1042 activation=layer.activation,

1043 quant_config=self.moe_quant_config,

1044 + moe_config=layer.moe_config,

1045 expert_map=layer.expert_map,

1046 )

1047 else:

@@ -1204,6 +1174,8 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):

1174 triton_kernel_moe_forward,

1175 )

1176

1177 + assert self.moe.hidden_dim_unpadded is not None

1178 + assert self.moe.intermediate_size_per_partition_unpadded is not None

1179 return triton_kernel_moe_forward(

1180 hidden_states=x,

1181 w1=self.w13_weight_triton_tensor,

@@ -1215,8 +1187,8 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):

1187 expert_map=expert_map,

1188 quant_config=self.moe_quant_config,

1189 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1218 - unpadded_N_w1=self.intermediate_size_per_partition * 2,

1219 - unpadded_K_w1=self.unpadded_hidden_size,

1220 - unpadded_N_w2=self.unpadded_hidden_size,

1221 - unpadded_K_w2=self.intermediate_size_per_partition,

1190 + unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2,

1191 + unpadded_K_w1=self.moe.hidden_dim_unpadded,

1192 + unpadded_N_w2=self.moe.hidden_dim_unpadded,

1193 + unpadded_K_w2=self.moe.intermediate_size_per_partition_unpadded,

1194 )

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

-1

@@ -254,7 +254,6 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

254 w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment

255 )

256 )

257 - layer.intermediate_size_per_partition = padded_intermediate

257 layer.moe_config.intermediate_size_per_partition = padded_intermediate

258

259 w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(

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

-1

@@ -439,7 +439,6 @@ def prepare_fp8_moe_layer_for_fi(

439 layer.moe_config.is_act_and_mul,

440 min_alignment,

441 )

442 - layer.intermediate_size_per_partition = new_intermediate

442 layer.moe_config.intermediate_size_per_partition = new_intermediate

443

444 # FI kernels require W31 layout rather than W13.