@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b2b2c5239e
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, 12:15 UTC

[MoE Refactor] Split up compressed_tensors_moe.py (#38960)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Apr 6, 2026 at 20:07 UTC b2b2c5239ec65fc6ba1109b0d06cb7462d99b70e
12 files changed +2770 -2543
docs/design/moe_kernel_features.md
+2 -2
@@ -57,8 +57,8 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
57
58 - [`ModelOptFp8MoEMethod`][vllm.model_executor.layers.quantization.modelopt.ModelOptFp8MoEMethod]
59 - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]
60 -- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod]
61 -- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW8A8Fp8MoEMethod]
60 +- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
61 +- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
62 - [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
63 - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
64
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py deleted
-2541
@@ -1,2541 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import enum
5 -from enum import Enum
6 -
7 -import torch
8 -from compressed_tensors import CompressionFormat
9 -from compressed_tensors.quantization import (
10 - ActivationOrdering,
11 - QuantizationArgs,
12 - QuantizationStrategy,
13 -)
14 -
15 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16 -from vllm import _custom_ops as ops
17 -from vllm.distributed import get_tensor_model_parallel_world_size
18 -from vllm.logger import init_logger
19 -from vllm.model_executor.layers.fused_moe import (
20 - FusedMoE,
21 - FusedMoEActivationFormat,
22 - FusedMoEExpertsModular,
23 - FusedMoEMethodBase,
24 - FusedMoeWeightScaleSupported,
25 - UnquantizedFusedMoEMethod,
26 -)
27 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
28 -from vllm.model_executor.layers.fused_moe.config import (
29 - FusedMoEConfig,
30 - FusedMoEQuantConfig,
31 - int4_w4a16_moe_quant_config,
32 - int4_w4afp8_moe_quant_config,
33 - int8_w8a8_moe_quant_config,
34 - int8_w8a16_moe_quant_config,
35 -)
36 -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts
37 -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
38 - BatchedMarlinExperts,
39 - MarlinExperts,
40 - fused_marlin_moe,
41 -)
42 -from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
43 - convert_to_fp8_moe_kernel_format,
44 - make_fp8_moe_kernel,
45 - make_fp8_moe_quant_config,
46 - select_fp8_moe_backend,
47 -)
48 -from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
49 - Mxfp4MoeBackend,
50 - make_mxfp4_moe_kernel,
51 - make_mxfp4_moe_quant_config,
52 -)
53 -from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
54 - convert_to_nvfp4_moe_kernel_format,
55 - is_global_sf_supported_for_nvfp4_backend,
56 - make_nvfp4_moe_kernel,
57 - make_nvfp4_moe_quant_config,
58 - select_nvfp4_moe_backend,
59 -)
60 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
61 - WNA16_SUPPORTED_BITS,
62 - WNA16_SUPPORTED_TYPES_MAP,
63 -)
64 -from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
65 - flashinfer_trtllm_mxint4_moe,
66 - is_flashinfer_mxint4_moe_available,
67 - prepare_static_weights_for_trtllm_mxint4_moe,
68 -)
69 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (
70 - process_fp8_input_tensor_strategy_moe,
71 - process_fp8_weight_tensor_strategy_moe,
72 -)
73 -from vllm.model_executor.layers.quantization.utils.marlin_utils import (
74 - check_moe_marlin_supports_layer,
75 - get_marlin_input_dtype,
76 - marlin_act_int8_process_scales,
77 - marlin_make_workspace_new,
78 - marlin_moe_permute_scales,
79 -)
80 -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
81 - prepare_moe_fp4_layer_for_marlin,
82 -)
83 -from vllm.model_executor.layers.quantization.utils.quant_utils import (
84 - convert_bf16_scales_to_fp8,
85 - convert_packed_uint4b8_to_signed_int4_inplace,
86 - kFp8Dynamic128Sym,
87 - kFp8DynamicTokenSym,
88 - kFp8Static128BlockSym,
89 - kFp8StaticChannelSym,
90 - kFp8StaticTensorSym,
91 - kNvfp4Dynamic,
92 - kNvfp4Static,
93 -)
94 -from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
95 - normalize_e4m3fn_to_e4m3fnuz,
96 -)
97 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs
98 -from vllm.platforms import CpuArchEnum, current_platform
99 -
100 -logger = init_logger(__name__)
101 -
102 -
103 -class GPTQMarlinState(Enum):
104 - REPACK = enum.auto()
105 - READY = enum.auto()
106 -
107 -
108 -__all__ = [
109 - "CompressedTensorsMoEMethod",
110 - "CompressedTensorsW8A8Fp8MoEMethod",
111 - "CompressedTensorsW8A8Int8MoEMethod",
112 - "CompressedTensorsWNA16MarlinMoEMethod",
113 - "CompressedTensorsWNA16MoEMethod",
114 - "CompressedTensorsW4A4Nvfp4MoEMethod",
115 - "CompressedTensorsW4A8Int8MoEMethod",
116 -]
117 -
118 -
119 -class CompressedTensorsMoEMethod(FusedMoEMethodBase):
120 - @staticmethod
121 - def get_moe_method(
122 - quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
123 - layer: torch.nn.Module,
124 - layer_name: str,
125 - ) -> FusedMoEMethodBase:
126 - # FusedMoE was made by combining multiple Linears so need to
127 - # make sure quantization config for Linear can target it
128 - quant_config._add_fused_moe_to_target_scheme_map()
129 - unfused_names = [
130 - layer_name + proj_name
131 - for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
132 - ]
133 - # TODO: refactor this to use expert_mapping and check all layer numbers
134 - all_scheme_dicts = [
135 - quant_config.get_scheme_dict(layer, name) for name in unfused_names
136 - ]
137 - scheme_dict = all_scheme_dicts.pop()
138 -
139 - # multiple schemes found
140 - if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
141 - raise ValueError(
142 - "All MoE projections need to have same "
143 - "quantization scheme but found multiple"
144 - )
145 -
146 - if scheme_dict is None: # ignored layer
147 - return UnquantizedFusedMoEMethod(layer.moe_config)
148 -
149 - # TODO: @dsikka: refactor this to use schemes as other kernels
150 - # are supported + check if the layer is being ignored.
151 - weight_quant = scheme_dict.get("weights")
152 - input_quant = scheme_dict.get("input_activations")
153 - format = scheme_dict.get("format")
154 -
155 - if quant_config._is_mxfp4(weight_quant):
156 - return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
157 -
158 - if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
159 - # group_size=None means channelwise
160 - group_size = weight_quant.group_size or -1
161 -
162 - valid_format_and_bits = (
163 - weight_quant.num_bits in WNA16_SUPPORTED_BITS
164 - and format == CompressionFormat.pack_quantized.value
165 - )
166 -
167 - if not valid_format_and_bits:
168 - raise ValueError(
169 - "For Fused MoE layers, only format: ",
170 - f"{CompressionFormat.pack_quantized.value} ",
171 - f" and bits: {WNA16_SUPPORTED_BITS} is supported ",
172 - f"but got format: {CompressionFormat.pack_quantized.value} "
173 - f" and bits: {weight_quant.num_bits}",
174 - )
175 -
176 - # Prefer to use the MarlinMoE kernel when it is supported.
177 - if (
178 - not check_moe_marlin_supports_layer(layer, group_size)
179 - or current_platform.is_rocm()
180 - ):
181 - if (
182 - weight_quant.strategy == QuantizationStrategy.GROUP
183 - and weight_quant.actorder
184 - in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
185 - ):
186 - raise ValueError(
187 - "WNA16MoE is not supported with actorder=group/dynamic."
188 - )
189 - logger.info_once("Using CompressedTensorsWNA16MoEMethod")
190 - return CompressedTensorsWNA16MoEMethod(
191 - weight_quant, input_quant, layer.moe_config
192 - )
193 - else:
194 - logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
195 - return CompressedTensorsWNA16MarlinMoEMethod(
196 - weight_quant, input_quant, layer.moe_config
197 - )
198 - elif quant_config._is_nvfp4_format(weight_quant):
199 - _is_valid_nvfp4_activations = (
200 - quant_config._is_nvfp4_format(input_quant) or input_quant is None
201 - )
202 - if not _is_valid_nvfp4_activations:
203 - raise ValueError(
204 - "For NVFP4 weights, input quantization must also be NVFP4 format ",
205 - f"or None for NVFP4A16, found {input_quant}",
206 - )
207 - return CompressedTensorsW4A4Nvfp4MoEMethod(
208 - layer.moe_config, layer_name, use_a16=(input_quant is None)
209 - )
210 - elif (
211 - quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
212 - or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
213 - or quant_config._is_fp8_w8a8(weight_quant, input_quant)
214 - ):
215 - return CompressedTensorsW8A8Fp8MoEMethod(
216 - weight_quant, input_quant, layer.moe_config
217 - )
218 - elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
219 - return CompressedTensorsW8A8Int8MoEMethod(
220 - weight_quant, input_quant, layer.moe_config
221 - )
222 - elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
223 - logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod")
224 - return CompressedTensorsW4A8Fp8MoEMethod(
225 - weight_quant, input_quant, layer.moe_config
226 - )
227 - elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
228 - return CompressedTensorsW4A8Int8MoEMethod(
229 - weight_quant, input_quant, layer.moe_config
230 - )
231 - else:
232 - raise RuntimeError(
233 - f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
234 - )
235 -
236 -
237 -class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
238 - def __init__(self, moe):
239 - super().__init__(moe)
240 - self.group_size = 32
241 - self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
242 - self.experts_cls = MarlinExperts
243 -
244 - def create_weights(
245 - self,
246 - layer: torch.nn.Module,
247 - num_experts: int,
248 - hidden_size: int,
249 - intermediate_size_per_partition: int,
250 - params_dtype: torch.dtype,
251 - **extra_weight_attrs,
252 - ):
253 - layer.num_experts = num_experts
254 - layer.params_dtype = params_dtype
255 -
256 - w13_weight = torch.nn.Parameter(
257 - torch.empty(
258 - num_experts,
259 - 2 * intermediate_size_per_partition,
260 - # 2 fp4 items are packed in the input dimension
261 - hidden_size // 2,
262 - requires_grad=False,
263 - dtype=torch.uint8,
264 - ),
265 - requires_grad=False,
266 - )
267 - layer.register_parameter("w13_weight_packed", w13_weight)
268 - set_weight_attrs(w13_weight, extra_weight_attrs)
269 -
270 - w2_weight = torch.nn.Parameter(
271 - torch.empty(
272 - num_experts,
273 - hidden_size,
274 - # 2 fp4 items are packed in the input dimension
275 - intermediate_size_per_partition // 2,
276 - dtype=torch.uint8,
277 - ),
278 - requires_grad=False,
279 - )
280 - layer.register_parameter("w2_weight_packed", w2_weight)
281 - set_weight_attrs(w2_weight, extra_weight_attrs)
282 -
283 - w13_weight_scale = torch.nn.Parameter(
284 - torch.empty(
285 - num_experts,
286 - 2 * intermediate_size_per_partition,
287 - # 2 fp4 items are packed in the input dimension
288 - hidden_size // self.group_size,
289 - dtype=torch.uint8,
290 - ),
291 - requires_grad=False,
292 - )
293 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
294 - extra_weight_attrs.update(
295 - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
296 - )
297 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
298 -
299 - w2_weight_scale = torch.nn.Parameter(
300 - torch.empty(
301 - num_experts,
302 - hidden_size,
303 - # 2 fp4 items are packed in the input dimension
304 - intermediate_size_per_partition // self.group_size,
305 - dtype=torch.uint8,
306 - ),
307 - requires_grad=False,
308 - )
309 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
310 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
311 -
312 - def get_fused_moe_quant_config(
313 - self, layer: torch.nn.Module
314 - ) -> FusedMoEQuantConfig | None:
315 - return make_mxfp4_moe_quant_config(
316 - mxfp4_backend=self.mxfp4_backend,
317 - w1_scale=layer.w13_weight_scale,
318 - w2_scale=layer.w2_weight_scale,
319 - )
320 -
321 - def process_weights_after_loading(self, layer: FusedMoE) -> None:
322 - layer.w13_weight = torch.nn.Parameter(
323 - layer.w13_weight_packed.data, requires_grad=False
324 - )
325 - delattr(layer, "w13_weight_packed")
326 -
327 - layer.w2_weight = torch.nn.Parameter(
328 - layer.w2_weight_packed.data, requires_grad=False
329 - )
330 - delattr(layer, "w2_weight_packed")
331 -
332 - logger.warning_once(
333 - "Your GPU does not have native support for FP4 computation but "
334 - "FP4 quantization is being used. Weight-only FP4 compression "
335 - "will be used leveraging the Marlin kernel. This may degrade "
336 - "performance for compute-heavy workloads."
337 - )
338 - prepare_moe_fp4_layer_for_marlin(layer)
339 -
340 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)
341 - if self.moe_quant_config is not None:
342 - self.moe_kernel = make_mxfp4_moe_kernel(
343 - moe_quant_config=self.moe_quant_config,
344 - moe_config=self.moe,
345 - experts_cls=self.experts_cls,
346 - mxfp4_backend=self.mxfp4_backend,
347 - shared_experts=layer.shared_experts,
348 - routing_tables=layer._maybe_init_expert_routing_tables(),
349 - )
350 -
351 - def apply(
352 - self,
353 - layer: FusedMoE,
354 - x: torch.Tensor,
355 - topk_weights: torch.Tensor,
356 - topk_ids: torch.Tensor,
357 - shared_experts_input: torch.Tensor | None,
358 - ) -> torch.Tensor:
359 - assert self.moe_kernel is not None
360 - return self.moe_kernel.apply(
361 - x,
362 - layer.w13_weight,
363 - layer.w2_weight,
364 - topk_weights,
365 - topk_ids,
366 - activation=layer.activation,
367 - global_num_experts=layer.global_num_experts,
368 - expert_map=layer.expert_map,
369 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
370 - shared_experts_input=shared_experts_input,
371 - )
372 -
373 -
374 -class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
375 - def __init__(
376 - self,
377 - moe: FusedMoEConfig,
378 - layer_name: str | None = None,
379 - use_a16: bool = False,
380 - ):
381 - super().__init__(moe)
382 - self.group_size = 16
383 -
384 - # Select experts implementation.
385 - self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
386 - config=self.moe,
387 - weight_key=kNvfp4Static,
388 - activation_key=None if use_a16 else kNvfp4Dynamic,
389 - )
390 -
391 - self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
392 - self.nvfp4_backend
393 - )
394 -
395 - def create_weights(
396 - self,
397 - layer: torch.nn.Module,
398 - num_experts: int,
399 - hidden_size: int,
400 - intermediate_size_per_partition: int,
401 - params_dtype: torch.dtype,
402 - **extra_weight_attrs,
403 - ):
404 - layer.num_experts = num_experts
405 - layer.params_dtype = params_dtype
406 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
407 -
408 - w13_weight = torch.nn.Parameter(
409 - torch.empty(
410 - num_experts,
411 - w13_num_shards * intermediate_size_per_partition,
412 - # 2 fp4 items are packed in the input dimension
413 - hidden_size // 2,
414 - requires_grad=False,
415 - dtype=torch.uint8,
416 - ),
417 - requires_grad=False,
418 - )
419 - layer.register_parameter("w13_weight_packed", w13_weight)
420 - set_weight_attrs(w13_weight, extra_weight_attrs)
421 -
422 - w2_weight = torch.nn.Parameter(
423 - torch.empty(
424 - num_experts,
425 - hidden_size,
426 - # 2 fp4 items are packed in the input dimension
427 - intermediate_size_per_partition // 2,
428 - dtype=torch.uint8,
429 - ),
430 - requires_grad=False,
431 - )
432 - layer.register_parameter("w2_weight_packed", w2_weight)
433 - set_weight_attrs(w2_weight, extra_weight_attrs)
434 -
435 - # Weight Scales
436 - w13_weight_scale = torch.nn.Parameter(
437 - torch.empty(
438 - num_experts,
439 - w13_num_shards * intermediate_size_per_partition,
440 - # 2 fp4 items are packed in the input dimension
441 - hidden_size // self.group_size,
442 - dtype=torch.float8_e4m3fn,
443 - ),
444 - requires_grad=False,
445 - )
446 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
447 - extra_weight_attrs.update(
448 - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
449 - )
450 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
451 -
452 - w2_weight_scale = torch.nn.Parameter(
453 - torch.empty(
454 - num_experts,
455 - hidden_size,
456 - # 2 fp4 items are packed in the input dimension
457 - intermediate_size_per_partition // self.group_size,
458 - dtype=torch.float8_e4m3fn,
459 - ),
460 - requires_grad=False,
461 - )
462 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
463 - extra_weight_attrs.update(
464 - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
465 - )
466 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
467 -
468 - # Weight Global Scales
469 - w13_weight_scale_2 = torch.nn.Parameter(
470 - torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
471 - requires_grad=False,
472 - )
473 - layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)
474 - extra_weight_attrs.update(
475 - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
476 - )
477 - set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
478 -
479 - w2_weight_scale_2 = torch.nn.Parameter(
480 - torch.empty(num_experts, dtype=torch.float32), requires_grad=False
481 - )
482 - layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
483 - extra_weight_attrs.update(
484 - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
485 - )
486 - set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
487 -
488 - # Input Global Scales
489 - w13_input_scale = torch.nn.Parameter(
490 - torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
491 - requires_grad=False,
492 - )
493 - layer.register_parameter("w13_input_global_scale", w13_input_scale)
494 - extra_weight_attrs.update(
495 - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
496 - )
497 - set_weight_attrs(w13_input_scale, extra_weight_attrs)
498 -
499 - w2_input_scale = torch.nn.Parameter(
500 - torch.empty(num_experts, dtype=torch.float32), requires_grad=False
501 - )
502 - layer.register_parameter("w2_input_global_scale", w2_input_scale)
503 - extra_weight_attrs.update(
504 - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
505 - )
506 - set_weight_attrs(w2_input_scale, extra_weight_attrs)
507 -
508 - def process_weights_after_loading(self, layer: FusedMoE) -> None:
509 - """
510 - Convert NVFP4 MoE weights into kernel format and setup the kernel.
511 - """
512 - # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
513 - # requires this naming convention. However, the name change breaks
514 - # reloading because the state dict no longer matches disk. Once we
515 - # remove MKM, we should revert this change to ensure compatibility.
516 - layer.w13_weight = torch.nn.Parameter(
517 - layer.w13_weight_packed.data, requires_grad=False
518 - )
519 - delattr(layer, "w13_weight_packed")
520 -
521 - layer.w2_weight = torch.nn.Parameter(
522 - layer.w2_weight_packed.data, requires_grad=False
523 - )
524 - delattr(layer, "w2_weight_packed")
525 -
526 - # Use a single gscale for w13.
527 - if self.moe.is_act_and_mul and not torch.allclose(
528 - layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]
529 - ):
530 - logger.warning_once(
531 - "w1_weight_global_scale must match w3_weight_global_scale. "
532 - "Accuracy may be affected.",
533 - )
534 - w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()
535 -
536 - # Shuffle weights into the NvFp4 kernel format.
537 - (
538 - w13,
539 - w13_scale,
540 - w13_scale_2,
541 - a13_scale,
542 - w2,
543 - w2_scale,
544 - w2_scale_2,
545 - a2_scale,
546 - ) = convert_to_nvfp4_moe_kernel_format(
547 - nvfp4_backend=self.nvfp4_backend,
548 - layer=layer,
549 - w13=layer.w13_weight,
550 - w13_scale=layer.w13_weight_scale,
551 - w13_scale_2=(1.0 / w13_weight_global_scale),
552 - a13_scale=(1.0 / layer.w13_input_global_scale),
553 - w2=layer.w2_weight,
554 - w2_scale=layer.w2_weight_scale,
555 - w2_scale_2=(1.0 / layer.w2_weight_global_scale),
556 - a2_scale=(1.0 / layer.w2_input_global_scale),
557 - is_act_and_mul=self.moe.is_act_and_mul,
558 - )
559 -
560 - replace_parameter(layer, "w13_weight", w13)
561 - replace_parameter(layer, "w13_weight_scale", w13_scale)
562 - replace_parameter(layer, "w2_weight", w2)
563 - replace_parameter(layer, "w2_weight_scale", w2_scale)
564 - layer.w13_weight_scale_2 = w13_scale_2
565 - layer.w2_weight_scale_2 = w2_scale_2
566 - layer.w13_input_scale = a13_scale
567 - layer.w2_input_scale = a2_scale
568 -
569 - # Setup modular kernel.
570 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)
571 - assert self.experts_cls is not None
572 - self.moe_kernel = make_nvfp4_moe_kernel(
573 - moe_quant_config=self.moe_quant_config,
574 - moe_config=self.moe,
575 - experts_cls=self.experts_cls,
576 - shared_experts=layer.shared_experts,
577 - routing_tables=layer._maybe_init_expert_routing_tables(),
578 - )
579 - self.moe_kernel.fused_experts.process_weights_after_loading(layer)
580 -
581 - def maybe_make_prepare_finalize(
582 - self,
583 - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
584 - ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
585 - raise ValueError(
586 - f"{self.__class__.__name__} uses the new modular kernel initialization "
587 - "logic. This function should not be called."
588 - )
589 -
590 - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
591 - return make_nvfp4_moe_quant_config(
592 - backend=self.nvfp4_backend,
593 - w13_scale=layer.w13_weight_scale,
594 - w2_scale=layer.w2_weight_scale,
595 - w13_scale_2=layer.w13_weight_scale_2,
596 - w2_scale_2=layer.w2_weight_scale_2,
597 - a13_scale=layer.w13_input_scale,
598 - a2_scale=layer.w2_input_scale,
599 - )
600 -
601 - def apply_monolithic(
602 - self,
603 - layer: FusedMoE,
604 - x: torch.Tensor,
605 - router_logits: torch.Tensor,
606 - ) -> torch.Tensor:
607 - assert self.is_monolithic
608 - assert self.moe_kernel is not None
609 - return self.moe_kernel.apply_monolithic(
610 - x,
611 - layer.w13_weight,
612 - layer.w2_weight,
613 - router_logits,
614 - activation=layer.activation,
615 - global_num_experts=layer.global_num_experts,
616 - expert_map=layer.expert_map,
617 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
618 - num_expert_group=layer.num_expert_group,
619 - topk_group=layer.topk_group,
620 - e_score_correction_bias=layer.e_score_correction_bias,
621 - routed_scaling_factor=layer.routed_scaling_factor,
622 - )
623 -
624 - def apply(
625 - self,
626 - layer: FusedMoE,
627 - x: torch.Tensor,
628 - topk_weights: torch.Tensor,
629 - topk_ids: torch.Tensor,
630 - shared_experts_input: torch.Tensor | None,
631 - ) -> torch.Tensor:
632 - assert self.moe_kernel is not None
633 - return self.moe_kernel.apply(
634 - x,
635 - layer.w13_weight,
636 - layer.w2_weight,
637 - topk_weights,
638 - topk_ids,
639 - activation=layer.activation,
640 - global_num_experts=layer.global_num_experts,
641 - expert_map=layer.expert_map,
642 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
643 - shared_experts_input=shared_experts_input,
644 - )
645 -
646 -
647 -class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
648 - """W8A8 FP8 MoE quantization using compressed tensors."""
649 -
650 - def __init__(
651 - self,
652 - weight_quant: QuantizationArgs,
653 - input_quant: QuantizationArgs,
654 - moe: FusedMoEConfig,
655 - layer_name: str | None = None,
656 - ):
657 - super().__init__(moe)
658 - self.weight_quant = weight_quant
659 - self.input_quant = input_quant
660 -
661 - per_tensor = (
662 - self.weight_quant.strategy == QuantizationStrategy.TENSOR
663 - and self.input_quant.strategy == QuantizationStrategy.TENSOR
664 - )
665 - per_channel = (
666 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL
667 - and self.input_quant.strategy == QuantizationStrategy.TOKEN
668 - )
669 - if not (per_tensor or per_channel):
670 - assert self.weight_quant.strategy == QuantizationStrategy.BLOCK
671 - self.weight_block_size = self.weight_quant.block_structure
672 - assert self.weight_quant.dynamic is not None
673 - else:
674 - self.weight_block_size = None
675 - self.block_quant = self.weight_block_size is not None
676 -
677 - self.static_input_scales = not self.input_quant.dynamic
678 - if self.static_input_scales and per_channel:
679 - raise ValueError(
680 - "For FP8 Fused MoE layer, we require either per tensor or "
681 - "channelwise, dynamic per token quantization."
682 - )
683 -
684 - ct2vllm_weight = {
685 - QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
686 - QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
687 - QuantizationStrategy.BLOCK: kFp8Static128BlockSym,
688 - }
689 - ct2vllm_act = {
690 - QuantizationStrategy.TOKEN: kFp8DynamicTokenSym,
691 - QuantizationStrategy.TENSOR: (
692 - kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym
693 - ),
694 - }
695 - weight_key = ct2vllm_weight[self.weight_quant.strategy]
696 - if weight_key == kFp8Static128BlockSym:
697 - activation_key = kFp8Dynamic128Sym
698 - else:
699 - activation_key = ct2vllm_act[self.input_quant.strategy]
700 -
701 - # Select Fp8 MoE backend
702 - self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
703 - config=self.moe,
704 - weight_key=weight_key,
705 - activation_key=activation_key,
706 - allow_vllm_cutlass=True,
707 - )
708 -
709 - def create_weights(
710 - self,
711 - layer: torch.nn.Module,
712 - num_experts: int,
713 - hidden_size: int,
714 - intermediate_size_per_partition: int,
715 - params_dtype: torch.dtype,
716 - **extra_weight_attrs,
717 - ):
718 - layer.num_experts = num_experts
719 - layer.orig_dtype = params_dtype
720 - layer.weight_block_size = None
721 -
722 - params_dtype = torch.float8_e4m3fn
723 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
724 -
725 - if self.block_quant:
726 - assert self.weight_block_size is not None
727 - layer.weight_block_size = self.weight_block_size
728 - tp_size = get_tensor_model_parallel_world_size()
729 - block_n, block_k = (
730 - self.weight_block_size[0],
731 - self.weight_block_size[1],
732 - )
733 - # NOTE: To ensure proper alignment of the block-wise quantization
734 - # scales, the output_size of the weights for both the gate and up
735 - # layers must be divisible by block_n.
736 - # Required by column parallel or enabling merged weights
737 - if intermediate_size_per_partition % block_n != 0:
738 - raise ValueError(
739 - f"The output_size of gate's and up's weight = "
740 - f"{intermediate_size_per_partition} is not divisible by "
741 - f"weight quantization block_n = {block_n}."
742 - )
743 - if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
744 - # Required by row parallel
745 - raise ValueError(
746 - f"The input_size of down's weight = "
747 - f"{intermediate_size_per_partition} is not divisible by "
748 - f"weight quantization block_k = {block_k}."
749 - )
750 -
751 - # WEIGHTS
752 - w13_weight = torch.nn.Parameter(
753 - torch.empty(
754 - num_experts,
755 - w13_num_shards * intermediate_size_per_partition,
756 - hidden_size,
757 - dtype=params_dtype,
758 - ),
759 - requires_grad=False,
760 - )
761 - layer.register_parameter("w13_weight", w13_weight)
762 - set_weight_attrs(w13_weight, extra_weight_attrs)
763 -
764 - w2_weight = torch.nn.Parameter(
765 - torch.empty(
766 - num_experts,
767 - hidden_size,
768 - intermediate_size_per_partition,
769 - dtype=params_dtype,
770 - ),
771 - requires_grad=False,
772 - )
773 - layer.register_parameter("w2_weight", w2_weight)
774 - set_weight_attrs(w2_weight, extra_weight_attrs)
775 -
776 - # WEIGHT_SCALES
777 - if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
778 - # For gated MoE, allocate 2 scales for w1 and w3 respectively.
779 - # They will be combined to a single scale after weight loading.
780 - # For non-gated MoE, allocate 1 scale for w13.
781 - w13_weight_scale = torch.nn.Parameter(
782 - torch.ones(num_experts, w13_num_shards, dtype=torch.float32),
783 - requires_grad=False,
784 - )
785 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
786 - w2_weight_scale = torch.nn.Parameter(
787 - torch.ones(num_experts, dtype=torch.float32), requires_grad=False
788 - )
789 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
790 - # Add PER-TENSOR quantization for FusedMoE.weight_loader.
791 - extra_weight_attrs.update(
792 - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
793 - )
794 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
795 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
796 -
797 - elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL:
798 - w13_weight_scale = torch.nn.Parameter(
799 - torch.ones(
800 - num_experts,
801 - w13_num_shards * intermediate_size_per_partition,
802 - 1,
803 - dtype=torch.float32,
804 - ),
805 - requires_grad=False,
806 - )
807 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
808 - w2_weight_scale = torch.nn.Parameter(
809 - torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
810 - requires_grad=False,
811 - )
812 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
813 - # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
814 - extra_weight_attrs.update(
815 - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
816 - )
817 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
818 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
819 -
820 - elif self.weight_quant.strategy == QuantizationStrategy.BLOCK:
821 - w13_weight_scale = torch.nn.Parameter(
822 - torch.ones(
823 - num_experts,
824 - w13_num_shards
825 - * ((intermediate_size_per_partition + block_n - 1) // block_n),
826 - (hidden_size + block_k - 1) // block_k,
827 - dtype=torch.float32,
828 - ),
829 - requires_grad=False,
830 - )
831 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
832 - w2_weight_scale = torch.nn.Parameter(
833 - torch.ones(
834 - num_experts,
835 - (hidden_size + block_n - 1) // block_n,
836 - (intermediate_size_per_partition + block_k - 1) // block_k,
837 - dtype=torch.float32,
838 - ),
839 - requires_grad=False,
840 - )
841 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
842 - # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
843 - extra_weight_attrs.update(
844 - {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
845 - )
846 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
847 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
848 -
849 - # INPUT_SCALES
850 - if self.static_input_scales:
851 - w13_input_scale = torch.nn.Parameter(
852 - torch.ones(num_experts, dtype=torch.float32), requires_grad=False
853 - )
854 - layer.register_parameter("w13_input_scale", w13_input_scale)
855 - set_weight_attrs(w13_input_scale, extra_weight_attrs)
856 -
857 - w2_input_scale = torch.nn.Parameter(
858 - torch.ones(num_experts, dtype=torch.float32), requires_grad=False
859 - )
860 - layer.register_parameter("w2_input_scale", w2_input_scale)
861 - set_weight_attrs(w2_input_scale, extra_weight_attrs)
862 - else:
863 - layer.w13_input_scale = None
864 - layer.w2_input_scale = None
865 -
866 - def process_weights_after_loading(self, layer: FusedMoE) -> None:
867 - # Allow for accessing weights and scales in standard way.
868 - w13 = layer.w13_weight
869 - w2 = layer.w2_weight
870 - w13_scale = layer.w13_weight_scale
871 - w2_scale = layer.w2_weight_scale
872 - w13_input_scale = layer.w13_input_scale
873 - w2_input_scale = layer.w2_input_scale
874 -
875 - # MI300x and MI325x use FNUZ format for FP8. Convert if needed.
876 - if current_platform.is_fp8_fnuz():
877 - w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
878 - w13, w13_scale, w13_input_scale
879 - )
880 - w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
881 - w2, w2_scale, w2_input_scale
882 - )
883 -
884 - # Per tensor kernels require single activation scale. Use the max.
885 - if self.static_input_scales:
886 - assert self.input_quant.strategy == QuantizationStrategy.TENSOR
887 - assert w13_input_scale is not None and w2_input_scale is not None
888 - w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
889 - w13_input_scale, w2_input_scale
890 - )
891 - replace_parameter(layer, "w13_input_scale", w13_input_scale)
892 - replace_parameter(layer, "w2_input_scale", w2_input_scale)
893 -
894 - # Per-tensor kernels use a single scale, for W13, but on disk there
895 - # is a separate scale for W1 and W3. Requantize with the max scale.
896 - if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
897 - w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
898 - w13,
899 - w13_scale,
900 - shard_size=layer.intermediate_size_per_partition,
901 - num_experts=layer.local_num_experts,
902 - is_act_and_mul=self.moe.is_act_and_mul,
903 - )
904 -
905 - w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
906 - fp8_backend=self.fp8_backend,
907 - layer=layer,
908 - w13=w13,
909 - w2=w2,
910 - w13_scale=w13_scale,
911 - w2_scale=w2_scale,
912 - w13_input_scale=w13_input_scale,
913 - w2_input_scale=w2_input_scale,
914 - )
915 -
916 - # Replace parameters with updated versions. Note that this helper
917 - # function ensures the replacement is compatible with RL weight reloads.
918 - replace_parameter(layer, "w13_weight", w13)
919 - replace_parameter(layer, "w2_weight", w2)
920 - replace_parameter(layer, "w13_weight_scale", w13_scale)
921 - replace_parameter(layer, "w2_weight_scale", w2_scale)
922 -
923 - # Setup modular kernel for TP case and naive DP/EP case.
924 - # In non-naive DP/EP case, we will create a ModularKernelMethod.
925 - # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
926 - # in both cases.
927 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)
928 - if self.moe_quant_config:
929 - assert self.experts_cls is not None
930 - self.moe_kernel = make_fp8_moe_kernel(
931 - moe_quant_config=self.moe_quant_config,
932 - moe_config=self.moe,
933 - fp8_backend=self.fp8_backend,
934 - experts_cls=self.experts_cls,
935 - routing_tables=layer._maybe_init_expert_routing_tables(),
936 - shared_experts=layer.shared_experts,
937 - )
938 -
939 - def maybe_make_prepare_finalize(
940 - self,
941 - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
942 - ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
943 - raise ValueError(
944 - f"{self.__class__.__name__} uses the new modular kernel initialization "
945 - "logic. This function should not be called."
946 - )
947 -
948 - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
949 - is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
950 - return make_fp8_moe_quant_config(
951 - fp8_backend=self.fp8_backend,
952 - w1_scale=layer.w13_weight_scale,
953 - w2_scale=layer.w2_weight_scale,
954 - a1_scale=layer.w13_input_scale,
955 - a2_scale=layer.w2_input_scale,
956 - per_act_token_quant=is_per_token,
957 - per_out_ch_quant=is_per_token,
958 - block_shape=self.weight_block_size,
959 - )
960 -
961 - def apply_monolithic(
962 - self,
963 - layer: FusedMoE,
964 - x: torch.Tensor,
965 - router_logits: torch.Tensor,
966 - ) -> torch.Tensor:
967 - assert self.moe_kernel is not None
968 - return self.moe_kernel.apply_monolithic(
969 - x,
970 - layer.w13_weight,
971 - layer.w2_weight,
972 - router_logits,
973 - activation=layer.activation,
974 - global_num_experts=layer.global_num_experts,
975 - expert_map=layer.expert_map,
976 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
977 - num_expert_group=layer.num_expert_group,
978 - topk_group=layer.topk_group,
979 - e_score_correction_bias=layer.e_score_correction_bias,
980 - routed_scaling_factor=layer.routed_scaling_factor,
981 - )
982 -
983 - def apply(
984 - self,
985 - layer: FusedMoE,
986 - x: torch.Tensor,
987 - topk_weights: torch.Tensor,
988 - topk_ids: torch.Tensor,
989 - shared_experts_input: torch.Tensor | None,
990 - ) -> torch.Tensor:
991 - assert not self.is_monolithic
992 - assert self.moe_kernel is not None
993 - return self.moe_kernel.apply(
994 - x,
995 - layer.w13_weight,
996 - layer.w2_weight,
997 - topk_weights,
998 - topk_ids,
999 - activation=layer.activation,
1000 - global_num_experts=layer.global_num_experts,
1001 - # TODO(rob): investigate the disable_expert_map introduced by:
1002 - # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
1003 - expert_map=layer.expert_map,
1004 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1005 - shared_experts_input=shared_experts_input,
1006 - )
1007 -
1008 - @property
1009 - def supports_eplb(self) -> bool:
1010 - return True
1011 -
1012 -
1013 -class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1014 - def __init__(
1015 - self,
1016 - weight_quant: QuantizationArgs,
1017 - input_quant: QuantizationArgs,
1018 - moe: FusedMoEConfig,
1019 - layer_name: str | None = None,
1020 - ):
1021 - super().__init__(moe)
1022 - self.weight_quant = weight_quant
1023 - self.input_quant = input_quant
1024 -
1025 - per_channel = (
1026 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL
1027 - and self.input_quant.strategy == QuantizationStrategy.TOKEN
1028 - )
1029 - if not per_channel:
1030 - raise ValueError(
1031 - "For INT8 Fused MoE layers, we require channelwise, "
1032 - "dynamic per token quantization. Found "
1033 - f"{self.weight_quant}, {self.input_quant}"
1034 - )
1035 -
1036 - self.static_input_scales = not self.input_quant.dynamic
1037 - if self.static_input_scales:
1038 - raise ValueError(
1039 - "For INT8 Fused MoE layers, we require channelwise, "
1040 - "dynamic per token quantization. Found static input scales."
1041 - )
1042 -
1043 - def create_weights(
1044 - self,
1045 - layer: torch.nn.Module,
1046 - num_experts: int,
1047 - hidden_size: int,
1048 - intermediate_size_per_partition: int,
1049 - params_dtype: torch.dtype,
1050 - **extra_weight_attrs,
1051 - ):
1052 - params_dtype = torch.int8
1053 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1054 -
1055 - # WEIGHTS
1056 - w13_weight = torch.nn.Parameter(
1057 - torch.empty(
1058 - num_experts,
1059 - w13_num_shards * intermediate_size_per_partition,
1060 - hidden_size,
1061 - dtype=params_dtype,
1062 - ),
1063 - requires_grad=False,
1064 - )
1065 - layer.register_parameter("w13_weight", w13_weight)
1066 - set_weight_attrs(w13_weight, extra_weight_attrs)
1067 -
1068 - w2_weight = torch.nn.Parameter(
1069 - torch.empty(
1070 - num_experts,
1071 - hidden_size,
1072 - intermediate_size_per_partition,
1073 - dtype=params_dtype,
1074 - ),
1075 - requires_grad=False,
1076 - )
1077 - layer.register_parameter("w2_weight", w2_weight)
1078 - set_weight_attrs(w2_weight, extra_weight_attrs)
1079 -
1080 - # WEIGHT_SCALES
1081 - assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL
1082 - w13_weight_scale = torch.nn.Parameter(
1083 - torch.ones(
1084 - num_experts,
1085 - w13_num_shards * intermediate_size_per_partition,
1086 - 1,
1087 - dtype=torch.float32,
1088 - ),
1089 - requires_grad=False,
1090 - )
1091 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
1092 - w2_weight_scale = torch.nn.Parameter(
1093 - torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
1094 - requires_grad=False,
1095 - )
1096 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
1097 - # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
1098 - extra_weight_attrs.update(
1099 - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
1100 - )
1101 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
1102 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
1103 -
1104 - # INPUT_SCALES
1105 - assert not self.static_input_scales
1106 - layer.w13_input_scale = None
1107 - layer.w2_input_scale = None
1108 -
1109 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1110 - pass
1111 -
1112 - def get_fused_moe_quant_config(
1113 - self, layer: torch.nn.Module
1114 - ) -> FusedMoEQuantConfig | None:
1115 - return int8_w8a8_moe_quant_config(
1116 - w1_scale=layer.w13_weight_scale,
1117 - w2_scale=layer.w2_weight_scale,
1118 - a1_scale=layer.w13_input_scale,
1119 - a2_scale=layer.w2_input_scale,
1120 - per_act_token_quant=True,
1121 - )
1122 -
1123 - def apply(
1124 - self,
1125 - layer: FusedMoE,
1126 - x: torch.Tensor,
1127 - topk_weights: torch.Tensor,
1128 - topk_ids: torch.Tensor,
1129 - shared_experts_input: torch.Tensor | None,
1130 - ) -> torch.Tensor:
1131 - from vllm.model_executor.layers.fused_moe import fused_experts
1132 -
1133 - return fused_experts(
1134 - hidden_states=x,
1135 - w1=layer.w13_weight,
1136 - w2=layer.w2_weight,
1137 - topk_weights=topk_weights,
1138 - topk_ids=topk_ids,
1139 - inplace=not self.moe.disable_inplace,
1140 - activation=layer.activation,
1141 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1142 - global_num_experts=layer.global_num_experts,
1143 - expert_map=layer.expert_map,
1144 - quant_config=self.moe_quant_config,
1145 - )
1146 -
1147 -
1148 -class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1149 - def __init__(
1150 - self,
1151 - weight_quant: QuantizationArgs,
1152 - input_quant: QuantizationArgs | None,
1153 - moe: FusedMoEConfig,
1154 - layer_name: str | None = None,
1155 - ):
1156 - super().__init__(moe)
1157 - self.weight_quant = weight_quant
1158 - self.input_quant = input_quant
1159 - assert weight_quant.symmetric, (
1160 - "Only symmetric quantization is supported for MoE"
1161 - )
1162 - # Extract properties from weight_quant
1163 - self.num_bits = weight_quant.num_bits
1164 - self.packed_factor = 32 // weight_quant.num_bits
1165 - self.strategy = weight_quant.strategy
1166 - self.group_size = weight_quant.group_size
1167 - self.actorder = weight_quant.actorder
1168 -
1169 - self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
1170 -
1171 - self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
1172 - self.use_flashinfer_mxint4_moe = (
1173 - is_flashinfer_mxint4_moe_available()
1174 - and self.group_size == 32
1175 - and weight_quant.num_bits == 4
1176 - )
1177 - self.kernel_backend = (
1178 - "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin"
1179 - )
1180 - logger.info_once(
1181 - f"Using {self.kernel_backend} backend for WNA16 MoE "
1182 - f"(group_size={self.group_size}, num_bits={self.num_bits})",
1183 - scope="local",
1184 - )
1185 -
1186 - def get_weight_shape(
1187 - self,
1188 - weight_name: str,
1189 - num_experts: int,
1190 - hidden_size: int,
1191 - intermediate_size_per_partition: int,
1192 - num_groups_w2: int | None = None,
1193 - num_groups_w13: int | None = None,
1194 - ) -> tuple[int, int, int]:
1195 - """
1196 - Get the shape of the weight based on the weight name, number of experts
1197 - hidden size, intermediate size per partition, number of groups for w2,
1198 - and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
1199 - for weight scales.
1200 - """
1201 - if weight_name == "w13_scale":
1202 - assert num_groups_w13 is not None, (
1203 - "num_groups_w13 must be provided for weight scales"
1204 - )
1205 - if weight_name == "w2_scale":
1206 - assert num_groups_w2 is not None, (
1207 - "num_groups_w2 must be provided for weight scales"
1208 - )
1209 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1210 - shape_map = {
1211 - "w13_weight": {
1212 - "Flashinfer": (
1213 - num_experts,
1214 - w13_num_shards * intermediate_size_per_partition,
1215 - hidden_size // self.packed_factor,
1216 - ),
1217 - "Marlin": (
1218 - num_experts,
1219 - hidden_size // self.packed_factor,
1220 - w13_num_shards * intermediate_size_per_partition,
1221 - ),
1222 - },
1223 - "w13_scale": {
1224 - "Flashinfer": (
1225 - num_experts,
1226 - w13_num_shards * intermediate_size_per_partition,
1227 - num_groups_w13,
1228 - ),
1229 - "Marlin": (
1230 - num_experts,
1231 - num_groups_w13,
1232 - w13_num_shards * intermediate_size_per_partition,
1233 - ),
1234 - },
1235 - "w2_weight": {
1236 - "Flashinfer": (
1237 - num_experts,
1238 - hidden_size,
1239 - intermediate_size_per_partition // self.packed_factor,
1240 - ),
1241 - "Marlin": (
1242 - num_experts,
1243 - intermediate_size_per_partition // self.packed_factor,
1244 - hidden_size,
1245 - ),
1246 - },
1247 - "w2_scale": {
1248 - "Flashinfer": (num_experts, hidden_size, num_groups_w2),
1249 - "Marlin": (num_experts, num_groups_w2, hidden_size),
1250 - },
1251 - }
1252 - return shape_map[weight_name][self.kernel_backend]
1253 -
1254 - def create_weights(
1255 - self,
1256 - layer: torch.nn.Module,
1257 - num_experts: int,
1258 - hidden_size: int,
1259 - intermediate_size_per_partition: int,
1260 - params_dtype: torch.dtype,
1261 - **extra_weight_attrs,
1262 - ):
1263 - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
1264 -
1265 - # Will transpose the loaded weight along the
1266 - # intermediate and hidden dim sizes. Will
1267 - # shard for TP along the transposed dims
1268 - is_transposed = self.kernel_backend != "Flashinfer"
1269 - extra_weight_attrs.update(
1270 - {"is_transposed": is_transposed, "quant_method": self.strategy}
1271 - )
1272 -
1273 - w13_weight = torch.nn.Parameter(
1274 - torch.empty(
1275 - *self.get_weight_shape(
1276 - "w13_weight",
1277 - num_experts,
1278 - hidden_size,
1279 - intermediate_size_per_partition,
1280 - ),
1281 - dtype=torch.int32,
1282 - ),
1283 - requires_grad=False,
1284 - )
1285 - layer.register_parameter("w13_weight_packed", w13_weight)
1286 - set_weight_attrs(w13_weight, extra_weight_attrs)
1287 -
1288 - w2_weight = torch.nn.Parameter(
1289 - torch.empty(
1290 - *self.get_weight_shape(
1291 - "w2_weight",
1292 - num_experts,
1293 - hidden_size,
1294 - intermediate_size_per_partition,
1295 - ),
1296 - dtype=torch.int32,
1297 - ),
1298 - requires_grad=False,
1299 - )
1300 - layer.register_parameter("w2_weight_packed", w2_weight)
1301 - set_weight_attrs(w2_weight, extra_weight_attrs)
1302 -
1303 - # In the case where we have actorder/g_idx,
1304 - # we do not partition the w2 scales
1305 - load_full_w2 = self.actorder and self.group_size != -1
1306 - w2_scales_size = (
1307 - intermediate_size_full if load_full_w2 else intermediate_size_per_partition
1308 - )
1309 -
1310 - self.is_k_full = (not self.actorder) or (
1311 - intermediate_size_per_partition == intermediate_size_full
1312 - )
1313 -
1314 - if self.strategy == "channel":
1315 - num_groups_w2 = num_groups_w13 = 1
1316 - self.group_size = -1
1317 - else:
1318 - num_groups_w2 = w2_scales_size // self.group_size
1319 - num_groups_w13 = hidden_size // self.group_size
1320 -
1321 - layer.num_groups_w13 = num_groups_w13
1322 - layer.num_groups_w2 = num_groups_w2
1323 -
1324 - w13_scale = torch.nn.Parameter(
1325 - torch.ones(
1326 - *self.get_weight_shape(
1327 - "w13_scale",
1328 - num_experts,
1329 - hidden_size,
1330 - intermediate_size_per_partition,
1331 - num_groups_w13=num_groups_w13,
1332 - ),
1333 - dtype=params_dtype,
1334 - ),
1335 - requires_grad=False,
1336 - )
1337 - layer.register_parameter("w13_weight_scale", w13_scale)
1338 - set_weight_attrs(w13_scale, extra_weight_attrs)
1339 -
1340 - w2_scale = torch.nn.Parameter(
1341 - torch.ones(
1342 - *self.get_weight_shape(
1343 - "w2_scale",
1344 - num_experts,
1345 - hidden_size,
1346 - intermediate_size_per_partition,
1347 - num_groups_w2=num_groups_w2,
1348 - ),
1349 - dtype=params_dtype,
1350 - ),
1351 - requires_grad=False,
1352 - )
1353 - layer.register_parameter("w2_weight_scale", w2_scale)
1354 - set_weight_attrs(w2_scale, extra_weight_attrs)
1355 - set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
1356 -
1357 - w2_weight_shape = torch.nn.Parameter(
1358 - torch.empty(num_experts, 2), requires_grad=False
1359 - )
1360 - layer.register_parameter("w2_weight_shape", w2_weight_shape)
1361 - set_weight_attrs(w2_weight_shape, extra_weight_attrs)
1362 - w13_weight_shape = torch.nn.Parameter(
1363 - torch.empty(num_experts, 2), requires_grad=False
1364 - )
1365 -
1366 - layer.register_parameter("w13_weight_shape", w13_weight_shape)
1367 - set_weight_attrs(w13_weight_shape, extra_weight_attrs)
1368 -
1369 - w13_g_idx = torch.nn.Parameter(
1370 - torch.empty(
1371 - num_experts,
1372 - hidden_size,
1373 - dtype=torch.int32,
1374 - ),
1375 - requires_grad=False,
1376 - )
1377 - layer.register_parameter("w13_weight_g_idx", w13_g_idx)
1378 - set_weight_attrs(w13_g_idx, extra_weight_attrs)
1379 -
1380 - w2_g_idx = torch.nn.Parameter(
1381 - torch.empty(
1382 - num_experts,
1383 - intermediate_size_per_partition,
1384 - dtype=torch.int32,
1385 - ),
1386 - requires_grad=False,
1387 - )
1388 - layer.register_parameter("w2_weight_g_idx", w2_g_idx)
1389 - set_weight_attrs(w2_g_idx, extra_weight_attrs)
1390 -
1391 - w13_g_idx_sort_indices = torch.nn.Parameter(
1392 - torch.empty(
1393 - num_experts,
1394 - hidden_size,
1395 - dtype=torch.int32,
1396 - ),
1397 - requires_grad=False,
1398 - )
1399 - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
1400 - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
1401 -
1402 - w2_g_idx_sort_indices = torch.nn.Parameter(
1403 - torch.empty(
1404 - num_experts,
1405 - intermediate_size_per_partition,
1406 - dtype=torch.int32,
1407 - ),
1408 - requires_grad=False,
1409 - )
1410 - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
1411 - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
1412 -
1413 - layer.a13_scale = None
1414 - layer.a2_scale = None
1415 - layer.marlin_state = GPTQMarlinState.REPACK
1416 -
1417 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1418 - num_experts = layer.w13_weight_g_idx.shape[0]
1419 - device = layer.w13_weight_g_idx.device
1420 - if self.kernel_backend == "Flashinfer":
1421 - dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
1422 - layer.w13_weight_packed,
1423 - layer.w13_weight_scale,
1424 - layer.w2_weight_packed,
1425 - layer.w2_weight_scale,
1426 - )
1427 - replace_parameter(
1428 - layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"]
1429 - )
1430 - replace_parameter(
1431 - layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"]
1432 - )
1433 - replace_parameter(
1434 - layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"]
1435 - )
1436 - replace_parameter(
1437 - layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"]
1438 - )
1439 - return None
1440 -
1441 - is_a_8bit = (
1442 - self.marlin_input_dtype is not None
1443 - and self.marlin_input_dtype.itemsize == 1
1444 - )
1445 -
1446 - if self.marlin_input_dtype == torch.float8_e4m3fn:
1447 - # NOTE: for non-zp quantization format only
1448 - ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True)
1449 - ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True)
1450 - layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512
1451 - layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512
1452 -
1453 - # when running models with grouped act order,
1454 - # resort to g_idx values provided in checkpoint
1455 - if self.actorder == "group":
1456 - w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx)
1457 - w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx)
1458 - w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx)
1459 - w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx)
1460 -
1461 - for e in range(num_experts):
1462 - w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to(
1463 - torch.int32
1464 - )
1465 - w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to(
1466 - torch.int32
1467 - )
1468 - w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][
1469 - w13_g_idx_sort_indices[e]
1470 - ]
1471 - w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]]
1472 -
1473 - replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx)
1474 - replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx)
1475 - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
1476 - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
1477 -
1478 - else:
1479 - layer.w13_weight_g_idx = torch.nn.Parameter(
1480 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),
1481 - requires_grad=False,
1482 - )
1483 - layer.w2_weight_g_idx = torch.nn.Parameter(
1484 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),
1485 - requires_grad=False,
1486 - )
1487 - layer.w13_g_idx_sort_indices = torch.nn.Parameter(
1488 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),
1489 - requires_grad=False,
1490 - )
1491 - layer.w2_g_idx_sort_indices = torch.nn.Parameter(
1492 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),
1493 - requires_grad=False,
1494 - )
1495 -
1496 - marlin_w13_qweight = ops.gptq_marlin_moe_repack(
1497 - layer.w13_weight_packed,
1498 - layer.w13_g_idx_sort_indices,
1499 - layer.w13_weight_packed.shape[1] * self.packed_factor,
1500 - layer.w13_weight_packed.shape[2],
1501 - self.num_bits,
1502 - is_a_8bit=is_a_8bit,
1503 - )
1504 - replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight)
1505 -
1506 - marlin_w2_qweight = ops.gptq_marlin_moe_repack(
1507 - layer.w2_weight_packed,
1508 - layer.w2_g_idx_sort_indices,
1509 - layer.w2_weight_packed.shape[1] * self.packed_factor,
1510 - layer.w2_weight_packed.shape[2],
1511 - self.num_bits,
1512 - is_a_8bit=is_a_8bit,
1513 - )
1514 - replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight)
1515 -
1516 - # Repack scales
1517 - marlin_w13_scales = marlin_moe_permute_scales(
1518 - s=layer.w13_weight_scale,
1519 - size_k=layer.w13_weight_packed.shape[2],
1520 - size_n=layer.w13_weight_scale.shape[2],
1521 - group_size=self.group_size,
1522 - is_a_8bit=is_a_8bit,
1523 - )
1524 - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1:
1525 - marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
1526 - marlin_w13_scales
1527 - )
1528 - layer.register_parameter(
1529 - "w13_input_global_scale",
1530 - torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
1531 - )
1532 - replace_parameter(layer, "w13_weight_scale", marlin_w13_scales)
1533 -
1534 - marlin_w2_scales = marlin_moe_permute_scales(
1535 - s=layer.w2_weight_scale,
1536 - size_k=layer.w2_weight_scale.shape[1]
1537 - * (self.group_size if self.group_size != -1 else self.packed_factor),
1538 - size_n=layer.w2_weight_scale.shape[2],
1539 - group_size=self.group_size,
1540 - is_a_8bit=is_a_8bit,
1541 - )
1542 - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1:
1543 - marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
1544 - marlin_w2_scales
1545 - )
1546 - layer.register_parameter(
1547 - "w2_input_global_scale",
1548 - torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
1549 - )
1550 - replace_parameter(layer, "w2_weight_scale", marlin_w2_scales)
1551 -
1552 - layer.workspace = marlin_make_workspace_new(device, 4)
1553 -
1554 - def get_fused_moe_quant_config(
1555 - self, layer: torch.nn.Module
1556 - ) -> FusedMoEQuantConfig | None:
1557 - if self.num_bits != 4:
1558 - return None
1559 - return int4_w4a16_moe_quant_config(
1560 - w1_scale=layer.w13_weight_scale,
1561 - w2_scale=layer.w2_weight_scale,
1562 - w1_zp=None,
1563 - w2_zp=None,
1564 - block_shape=[0, self.group_size],
1565 - )
1566 -
1567 - def select_gemm_impl(
1568 - self,
1569 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
1570 - layer: torch.nn.Module,
1571 - ) -> mk.FusedMoEExpertsModular:
1572 - assert self.num_bits == 4, "only supporting w4"
1573 - layer.w13_weight = layer.w13_weight_packed
1574 - layer.w2_weight = layer.w2_weight_packed
1575 - assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])
1576 - assert self.moe_quant_config is not None
1577 - if (
1578 - prepare_finalize.activation_format
1579 - == mk.FusedMoEActivationFormat.BatchedExperts
1580 - ):
1581 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
1582 - assert max_num_tokens_per_rank is not None
1583 - return BatchedMarlinExperts(
1584 - max_num_tokens=max_num_tokens_per_rank,
1585 - num_dispatchers=prepare_finalize.num_dispatchers(),
1586 - moe_config=self.moe,
1587 - quant_config=self.moe_quant_config,
1588 - w13_g_idx=layer.w13_weight_g_idx,
1589 - w2_g_idx=layer.w2_weight_g_idx,
1590 - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
1591 - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
1592 - is_k_full=self.is_k_full,
1593 - )
1594 - else:
1595 - return MarlinExperts(
1596 - moe_config=self.moe,
1597 - quant_config=self.moe_quant_config,
1598 - w13_g_idx=layer.w13_weight_g_idx,
1599 - w2_g_idx=layer.w2_weight_g_idx,
1600 - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
1601 - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
1602 - is_k_full=self.is_k_full,
1603 - )
1604 -
1605 - @property
1606 - def is_monolithic(self) -> bool:
1607 - return self.kernel_backend == "Flashinfer"
1608 -
1609 - def apply_monolithic(
1610 - self,
1611 - layer: FusedMoE,
1612 - x: torch.Tensor,
1613 - router_logits: torch.Tensor,
1614 - ) -> torch.Tensor:
1615 - assert self.kernel_backend == "Flashinfer"
1616 - return flashinfer_trtllm_mxint4_moe(
1617 - x=x,
1618 - router_logits=router_logits,
1619 - w13_weight_packed=layer.w13_weight_packed,
1620 - w13_weight_scale=layer.w13_weight_scale,
1621 - w2_weight_packed=layer.w2_weight_packed,
1622 - w2_weight_scale=layer.w2_weight_scale,
1623 - global_num_experts=layer.global_num_experts,
1624 - top_k=layer.top_k,
1625 - intermediate_size_per_partition=layer.intermediate_size_per_partition,
1626 - local_num_experts=layer.local_num_experts,
1627 - ep_rank=layer.ep_rank,
1628 - num_expert_group=layer.num_expert_group,
1629 - topk_group=layer.topk_group,
1630 - e_score_correction_bias=layer.e_score_correction_bias,
1631 - routing_method_type=layer.routing_method_type,
1632 - )
1633 -
1634 - def apply(
1635 - self,
1636 - layer: FusedMoE,
1637 - x: torch.Tensor,
1638 - topk_weights: torch.Tensor,
1639 - topk_ids: torch.Tensor,
1640 - shared_experts_input: torch.Tensor | None,
1641 - ) -> torch.Tensor:
1642 - assert self.kernel_backend == "Marlin"
1643 - return fused_marlin_moe(
1644 - x,
1645 - layer.w13_weight_packed,
1646 - layer.w2_weight_packed,
1647 - None,
1648 - None,
1649 - layer.w13_weight_scale,
1650 - layer.w2_weight_scale,
1651 - topk_weights,
1652 - topk_ids,
1653 - input_global_scale1=getattr(layer, "w13_input_global_scale", None),
1654 - input_global_scale2=getattr(layer, "w2_input_global_scale", None),
1655 - quant_type_id=self.quant_type.id,
1656 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1657 - global_num_experts=layer.global_num_experts,
1658 - activation=layer.activation,
1659 - expert_map=layer.expert_map,
1660 - g_idx1=layer.w13_weight_g_idx,
1661 - g_idx2=layer.w2_weight_g_idx,
1662 - sort_indices1=layer.w13_g_idx_sort_indices,
1663 - sort_indices2=layer.w2_g_idx_sort_indices,
1664 - workspace=layer.workspace,
1665 - input_dtype=self.marlin_input_dtype,
1666 - is_k_full=self.is_k_full,
1667 - inplace=not self.moe.disable_inplace,
1668 - )
1669 -
1670 -
1671 -class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1672 - def __init__(
1673 - self,
1674 - weight_quant: QuantizationArgs,
1675 - input_quant: QuantizationArgs | None,
1676 - moe: FusedMoEConfig,
1677 - layer_name: str | None = None,
1678 - ):
1679 - super().__init__(moe)
1680 - self.weight_quant = weight_quant
1681 - self.input_quant = input_quant
1682 - # Extract properties from weight_quant
1683 - self.num_bits = weight_quant.num_bits
1684 - self.packed_factor = 32 // weight_quant.num_bits
1685 - self.strategy = weight_quant.strategy
1686 - # channelwise is not supported by this kernel
1687 - assert weight_quant.strategy == "group"
1688 - self.group_size = weight_quant.group_size
1689 - # grouped actorder isn't supported by this kernel
1690 - assert weight_quant.actorder != "group"
1691 - assert weight_quant.symmetric, (
1692 - "Only symmetric quantization is supported for MoE"
1693 - )
1694 -
1695 - def create_weights(
1696 - self,
1697 - layer: torch.nn.Module,
1698 - num_experts: int,
1699 - hidden_size: int,
1700 - intermediate_size_per_partition: int,
1701 - params_dtype: torch.dtype,
1702 - **extra_weight_attrs,
1703 - ):
1704 - # Will transpose the loaded weight along the
1705 - # intermediate and hidden dim sizes. Will
1706 - # shard for TP along the transposed dims
1707 - extra_weight_attrs.update(
1708 - {"is_transposed": True, "quant_method": self.strategy}
1709 - )
1710 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1711 - w13_weight = torch.nn.Parameter(
1712 - torch.empty(
1713 - num_experts,
1714 - hidden_size // self.packed_factor,
1715 - w13_num_shards * intermediate_size_per_partition,
1716 - dtype=torch.int32,
1717 - ),
1718 - requires_grad=False,
1719 - )
1720 - layer.register_parameter("w13_weight_packed", w13_weight)
1721 - set_weight_attrs(w13_weight, extra_weight_attrs)
1722 -
1723 - w2_weight = torch.nn.Parameter(
1724 - torch.empty(
1725 - num_experts,
1726 - intermediate_size_per_partition // self.packed_factor,
1727 - hidden_size,
1728 - dtype=torch.int32,
1729 - ),
1730 - requires_grad=False,
1731 - )
1732 - layer.register_parameter("w2_weight_packed", w2_weight)
1733 - set_weight_attrs(w2_weight, extra_weight_attrs)
1734 -
1735 - w2_scales_size = intermediate_size_per_partition
1736 -
1737 - if self.strategy == "channel":
1738 - num_groups_w2 = num_groups_w13 = 1
1739 - self.group_size = -1
1740 - else:
1741 - num_groups_w2 = w2_scales_size // self.group_size
1742 - num_groups_w13 = hidden_size // self.group_size
1743 -
1744 - w13_scale = torch.nn.Parameter(
1745 - torch.ones(
1746 - num_experts,
1747 - num_groups_w13,
1748 - w13_num_shards * intermediate_size_per_partition,
1749 - dtype=params_dtype,
1750 - ),
1751 - requires_grad=False,
1752 - )
1753 - layer.register_parameter("w13_weight_scale", w13_scale)
1754 - set_weight_attrs(w13_scale, extra_weight_attrs)
1755 -
1756 - w2_scale = torch.nn.Parameter(
1757 - torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
1758 - requires_grad=False,
1759 - )
1760 - layer.register_parameter("w2_weight_scale", w2_scale)
1761 - set_weight_attrs(w2_scale, extra_weight_attrs)
1762 - set_weight_attrs(w2_scale, {"load_full_w2": False})
1763 -
1764 - w2_weight_shape = torch.nn.Parameter(
1765 - torch.empty(num_experts, 2), requires_grad=False
1766 - )
1767 - layer.register_parameter("w2_weight_shape", w2_weight_shape)
1768 - set_weight_attrs(w2_weight_shape, extra_weight_attrs)
1769 - w13_weight_shape = torch.nn.Parameter(
1770 - torch.empty(num_experts, 2), requires_grad=False
1771 - )
1772 -
1773 - layer.register_parameter("w13_weight_shape", w13_weight_shape)
1774 - set_weight_attrs(w13_weight_shape, extra_weight_attrs)
1775 -
1776 - w13_g_idx = torch.nn.Parameter(
1777 - torch.empty(
1778 - num_experts,
1779 - hidden_size,
1780 - dtype=torch.int32,
1781 - ),
1782 - requires_grad=False,
1783 - )
1784 - layer.register_parameter("w13_weight_g_idx", w13_g_idx)
1785 - set_weight_attrs(w13_g_idx, extra_weight_attrs)
1786 -
1787 - w2_g_idx = torch.nn.Parameter(
1788 - torch.empty(
1789 - num_experts,
1790 - intermediate_size_per_partition,
1791 - dtype=torch.int32,
1792 - ),
1793 - requires_grad=False,
1794 - )
1795 - layer.register_parameter("w2_weight_g_idx", w2_g_idx)
1796 - set_weight_attrs(w2_g_idx, extra_weight_attrs)
1797 -
1798 - w13_g_idx_sort_indices = torch.nn.Parameter(
1799 - torch.empty(
1800 - num_experts,
1801 - hidden_size,
1802 - dtype=torch.int32,
1803 - ),
1804 - requires_grad=False,
1805 - )
1806 - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
1807 - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
1808 -
1809 - w2_g_idx_sort_indices = torch.nn.Parameter(
1810 - torch.empty(
1811 - num_experts,
1812 - intermediate_size_per_partition,
1813 - dtype=torch.int32,
1814 - ),
1815 - requires_grad=False,
1816 - )
1817 - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
1818 - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
1819 -
1820 - layer.a13_scale = None
1821 - layer.a2_scale = None
1822 -
1823 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1824 - # Reconfigure packed weights and scales to match moe_wna16 format
1825 - layer.w13_weight_packed = torch.nn.Parameter(
1826 - layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
1827 - requires_grad=False,
1828 - )
1829 - layer.w2_weight_packed = torch.nn.Parameter(
1830 - layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
1831 - requires_grad=False,
1832 - )
1833 - layer.w13_weight_scale = torch.nn.Parameter(
1834 - layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
1835 - )
1836 - layer.w2_weight_scale = torch.nn.Parameter(
1837 - layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
1838 - )
1839 -
1840 - def get_fused_moe_quant_config(
1841 - self, layer: torch.nn.Module
1842 - ) -> FusedMoEQuantConfig | None:
1843 - assert self.num_bits == 4 or self.num_bits == 8
1844 - config_builder = (
1845 - int4_w4a16_moe_quant_config
1846 - if self.num_bits == 4
1847 - else int8_w8a16_moe_quant_config
1848 - )
1849 -
1850 - return config_builder(
1851 - w1_scale=layer.w13_weight_scale,
1852 - w2_scale=layer.w2_weight_scale,
1853 - w1_zp=None,
1854 - w2_zp=None,
1855 - block_shape=[0, self.group_size],
1856 - )
1857 -
1858 - def select_gemm_impl(
1859 - self,
1860 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
1861 - layer: torch.nn.Module,
1862 - ) -> mk.FusedMoEExpertsModular:
1863 - if self.moe.is_lora_enabled:
1864 - assert self.moe_quant_config is not None
1865 - from vllm.triton_utils import HAS_TRITON
1866 -
1867 - if HAS_TRITON:
1868 - from vllm.model_executor.layers.fused_moe import TritonWNA16Experts
1869 -
1870 - layer.w13_weight = layer.w13_weight_packed
1871 - layer.w2_weight = layer.w2_weight_packed
1872 - return TritonWNA16Experts(
1873 - moe_config=self.moe, quant_config=self.moe_quant_config
1874 - )
1875 - else:
1876 - raise NotImplementedError(
1877 - "TritonExperts requires Triton. "
1878 - "Install triton or disable LoRA for MoE."
1879 - )
1880 -
1881 - raise NotImplementedError
1882 -
1883 - def apply(
1884 - self,
1885 - layer: FusedMoE,
1886 - x: torch.Tensor,
1887 - topk_weights: torch.Tensor,
1888 - topk_ids: torch.Tensor,
1889 - shared_experts_input: torch.Tensor | None,
1890 - ) -> torch.Tensor:
1891 - from vllm.model_executor.layers.fused_moe import fused_experts
1892 -
1893 - return fused_experts(
1894 - x,
1895 - layer.w13_weight_packed,
1896 - layer.w2_weight_packed,
1897 - topk_weights=topk_weights,
1898 - topk_ids=topk_ids,
1899 - inplace=not self.moe.disable_inplace,
1900 - activation=layer.activation,
1901 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1902 - global_num_experts=layer.global_num_experts,
1903 - expert_map=layer.expert_map,
1904 - quant_config=self.moe_quant_config,
1905 - )
1906 -
1907 - @property
1908 - def supports_eplb(self) -> bool:
1909 - return True
1910 -
1911 -
1912 -class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
1913 - """
1914 - CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
1915 - - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles)
1916 - - Scales: Fp32 for Channelwise , bf16 for groupwise quantization
1917 - - Bias: Same data type as original weights
1918 - - Activations: FP32/Bf16 dynamic per-token (A8 Int),
1919 - quantized inside the kernel
1920 - """
1921 -
1922 - def __init__(
1923 - self,
1924 - weight_quant: QuantizationArgs,
1925 - input_quant: QuantizationArgs,
1926 - moe: FusedMoEConfig,
1927 - layer_name: str | None = None,
1928 - ):
1929 - super().__init__(moe)
1930 - self.has_bias = self.moe.has_bias
1931 - self.weight_quant = weight_quant
1932 - self.input_quant = input_quant
1933 -
1934 - # Validate scheme: weights=W4 (channel or group),
1935 - # activations=dynamic TOKEN (A8)
1936 -
1937 - # Must be dynamic per-token activations
1938 - if (
1939 - input_quant.strategy != QuantizationStrategy.TOKEN
1940 - or not input_quant.dynamic
1941 - ):
1942 - raise ValueError(
1943 - "W4A8-int MoE needs dynamic per-token activation quantization."
1944 - )
1945 -
1946 - # Weight can be channel-wise (group_size=None) or group-wise
1947 - self.group_size = (
1948 - weight_quant.group_size if (weight_quant.group_size is not None) else -1
1949 - )
1950 - if weight_quant.num_bits != 4:
1951 - raise ValueError("This method only supports 4-bit weights (num_bits=4).")
1952 -
1953 - # CPU only
1954 - if not current_platform.is_cpu():
1955 - raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.")
1956 -
1957 - # Arm: check _dyn ops availability
1958 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
1959 - try:
1960 - _ = torch.ops.aten._dyn_quant_matmul_4bit
1961 - _ = torch.ops.aten._dyn_quant_pack_4bit_weight
1962 - except AttributeError as err:
1963 - raise RuntimeError(
1964 - f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
1965 - install a newer build."""
1966 - ) from err
1967 - self.static_input_scales = False # always dynamic per token
1968 -
1969 - # ---- parameter creation ----
1970 - def create_weights(
1971 - self,
1972 - layer: torch.nn.Module,
1973 - num_experts: int,
1974 - hidden_size: int,
1975 - intermediate_size_per_partition: int,
1976 - params_dtype: torch.dtype,
1977 - **extra_weight_attrs,
1978 - ):
1979 - # Shapes per local rank (TP/EP):
1980 - # w13: [E, 2*I_local, H] int8 (int4 values in [-8,7])
1981 - # w2 : [E, H, I_local] int8
1982 - # Scales:
1983 - # channel-wise: group_size=-1 -> per-output-row, single scale per row
1984 - # group-wise : group_size=g ->
1985 - # per-output-row, (in_features/g) scales
1986 -
1987 - E = num_experts
1988 - H = hidden_size
1989 - IN = intermediate_size_per_partition
1990 - g = self.group_size
1991 -
1992 - # Per-row scale columns
1993 - def _n_scale_cols(in_features: int) -> int:
1994 - return 1 if g == -1 else (in_features // g)
1995 -
1996 - # Register unpacked int4-as-int8 weights the loader will fill.
1997 - w13 = torch.nn.Parameter(
1998 - torch.empty(E, 2 * IN, H, dtype=torch.int8), requires_grad=False
1999 - )
2000 - set_weight_attrs(w13, extra_weight_attrs)
2001 - layer.register_parameter("w13_weight", w13)
2002 -
2003 - w2 = torch.nn.Parameter(
2004 - torch.empty(E, H, IN, dtype=torch.int8), requires_grad=False
2005 - )
2006 - set_weight_attrs(w2, extra_weight_attrs)
2007 - layer.register_parameter("w2_weight", w2)
2008 -
2009 - # Register scales
2010 - # KleidiAI groupwise kernels accepts float32 scales
2011 - # KleidiAI groupwise kernels accepts bfloat16 scales
2012 - scale_dtype = torch.float32 if g == -1 else torch.bfloat16
2013 -
2014 - w13_s = torch.nn.Parameter(
2015 - torch.ones(E, 2 * IN, _n_scale_cols(H), dtype=scale_dtype),
2016 - requires_grad=False,
2017 - )
2018 - set_weight_attrs(
2019 - w13_s,
2020 - {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
2021 - )
2022 - layer.register_parameter("w13_weight_scale", w13_s)
2023 -
2024 - w2_s = torch.nn.Parameter(
2025 - torch.ones(E, H, _n_scale_cols(IN), dtype=scale_dtype), requires_grad=False
2026 - )
2027 - set_weight_attrs(
2028 - w2_s,
2029 - {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
2030 - )
2031 - layer.register_parameter("w2_weight_scale", w2_s)
2032 -
2033 - if self.has_bias:
2034 - w13_bias = torch.nn.Parameter(
2035 - torch.zeros(E, 2 * IN, dtype=params_dtype), requires_grad=False
2036 - )
2037 - layer.register_parameter("w13_bias", w13_bias)
2038 - set_weight_attrs(w13_bias, extra_weight_attrs)
2039 -
2040 - w2_bias = torch.nn.Parameter(
2041 - torch.zeros(num_experts, hidden_size, dtype=params_dtype),
2042 - requires_grad=False,
2043 - )
2044 - layer.register_parameter("w2_bias", w2_bias)
2045 - set_weight_attrs(w2_bias, extra_weight_attrs)
2046 -
2047 - # Placeholders for packed weights (will be replaced after packing)
2048 - layer.register_parameter(
2049 - "w13_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
2050 - )
2051 - set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs)
2052 -
2053 - layer.register_parameter(
2054 - "w2_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
2055 - )
2056 - set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs)
2057 -
2058 - # dims for 4 bit fused matmuls
2059 - layer.w13_in_features = H
2060 - layer.w13_out_features = 2 * IN
2061 - layer.w2_in_features = IN
2062 - layer.w2_out_features = H
2063 - layer.group_size = g
2064 -
2065 - # post-load packing to dyn-4bit KleidiAI kernel's format
2066 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
2067 - E = layer.w13_weight.shape[0]
2068 - H = layer.w13_in_features
2069 - I2 = layer.w13_out_features
2070 - IN = layer.w2_in_features
2071 - g = layer.group_size
2072 -
2073 - def _pack_matrix(
2074 - int4_as_int8_2d: torch.Tensor,
2075 - scales_2d: torch.Tensor,
2076 - bias_1d: torch.Tensor | None,
2077 - in_features: int,
2078 - out_features: int,
2079 - ) -> torch.Tensor:
2080 - # int4 values are stored as int8 in [-8,7].
2081 - # Shift to unsigned nibble and pack pairs along input-dim.
2082 - tmp = int4_as_int8_2d.add(8) # [out, in]
2083 - uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(
2084 - torch.uint8
2085 - ) # [out, in//2]
2086 -
2087 - # KleidiAI groupwise kernels accepts float32 scales
2088 - # KleidiAI groupwise kernels accepts bfloat16 scales
2089 - scale_dtype = torch.float32 if g == -1 else torch.bfloat16
2090 - scales = scales_2d.to(scale_dtype)
2091 - bias = None if bias_1d is None else bias_1d.to(torch.float32)
2092 - return torch.ops.aten._dyn_quant_pack_4bit_weight(
2093 - uint8_nibbles,
2094 - scales,
2095 - bias,
2096 - g if g != -1 else in_features,
2097 - in_features,
2098 - out_features,
2099 - )
2100 -
2101 - # Pack per expert
2102 - w13_packed_list = []
2103 - w2_packed_list = []
2104 -
2105 - has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None
2106 - has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None
2107 -
2108 - for e in range(E):
2109 - w13_packed_list.append(
2110 - _pack_matrix(
2111 - layer.w13_weight[e], # [2I, H]
2112 - layer.w13_weight_scale[e], # [2I, H/g or 1]
2113 - layer.w13_bias[e] if has_w13_bias else None, # [2I]
2114 - H,
2115 - I2,
2116 - )
2117 - )
2118 - w2_packed_list.append(
2119 - _pack_matrix(
2120 - # w2 shape is [H, IN]; we need [out, in] == [H, IN].
2121 - layer.w2_weight[e], # [H, IN]
2122 - layer.w2_weight_scale[e], # [H, IN/g or 1]
2123 - layer.w2_bias[e] if has_w2_bias else None, # [H]
2124 - IN,
2125 - layer.w2_out_features, # in_features=IN, out_features=H
2126 - )
2127 - )
2128 -
2129 - # each packed tensor has identical shape per expert; stack on dim 0
2130 - w13_packed = torch.stack(w13_packed_list, dim=0)
2131 - w2_packed = torch.stack(w2_packed_list, dim=0)
2132 -
2133 - replace_parameter(
2134 - layer,
2135 - "w13_weight_packed",
2136 - torch.nn.Parameter(w13_packed, requires_grad=False),
2137 - )
2138 - replace_parameter(
2139 - layer,
2140 - "w2_weight_packed",
2141 - torch.nn.Parameter(w2_packed, requires_grad=False),
2142 - )
2143 -
2144 - # free raw tensors/scales/bias now that they're packed into the payload.
2145 - replace_parameter(
2146 - layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
2147 - )
2148 - replace_parameter(
2149 - layer, "w2_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
2150 - )
2151 - replace_parameter(
2152 - layer,
2153 - "w13_weight_scale",
2154 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
2155 - )
2156 - replace_parameter(
2157 - layer,
2158 - "w2_weight_scale",
2159 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
2160 - )
2161 - if has_w13_bias:
2162 - replace_parameter(
2163 - layer,
2164 - "w13_bias",
2165 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
2166 - )
2167 - if has_w2_bias:
2168 - replace_parameter(
2169 - layer,
2170 - "w2_bias",
2171 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
2172 - )
2173 -
2174 - def get_fused_moe_quant_config(
2175 - self, layer: torch.nn.Module
2176 - ) -> FusedMoEQuantConfig | None:
2177 - # CPU dynamic 4-bit MoE path does not use modular kernels or
2178 - # fused_experts; quant config is not needed.
2179 - return None
2180 -
2181 - @property
2182 - def is_monolithic(self) -> bool:
2183 - return True
2184 -
2185 - def apply_monolithic(
2186 - self,
2187 - layer: FusedMoE,
2188 - x: torch.Tensor,
2189 - router_logits: torch.Tensor,
2190 - ) -> torch.Tensor:
2191 - assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."
2192 - assert layer.activation in (
2193 - MoEActivation.SILU,
2194 - MoEActivation.SWIGLUOAI,
2195 - MoEActivation.SWIGLUSTEP,
2196 - ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."
2197 - assert layer.expert_map is None, """expert_map/EP not implemented
2198 - for CPU dyn-4bit MoE."""
2199 -
2200 - def _act_kind(s: MoEActivation) -> int:
2201 - # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU
2202 - if s == MoEActivation.SWIGLUSTEP:
2203 - return 0
2204 - if s == MoEActivation.SWIGLUOAI:
2205 - return 1
2206 - if s == MoEActivation.SILU:
2207 - return 2
2208 - raise ValueError(f"Unknown activation '{s}'")
2209 -
2210 - # Apply topk softmax on router output
2211 - topk_weights, topk_ids = select_experts(
2212 - hidden_states=x,
2213 - router_logits=router_logits,
2214 - top_k=layer.top_k,
2215 - use_grouped_topk=layer.use_grouped_topk,
2216 - renormalize=layer.renormalize,
2217 - )
2218 -
2219 - return torch.ops._C.dynamic_4bit_int_moe(
2220 - x,
2221 - topk_ids.to(torch.long),
2222 - topk_weights,
2223 - layer.w13_weight_packed,
2224 - layer.w2_weight_packed,
2225 - layer.w2_out_features,
2226 - layer.w2_in_features,
2227 - layer.w13_out_features,
2228 - layer.group_size,
2229 - layer.apply_router_weight_on_input,
2230 - int(_act_kind(layer.activation)),
2231 - )
2232 -
2233 -
2234 -class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2235 - def __init__(
2236 - self,
2237 - weight_quant: QuantizationArgs,
2238 - input_quant: QuantizationArgs,
2239 - moe: FusedMoEConfig,
2240 - layer_name: str | None = None,
2241 - ):
2242 - super().__init__(moe)
2243 - self.weight_quant = weight_quant
2244 - self.input_quant = input_quant
2245 -
2246 - self.group_size = self.weight_quant.group_size
2247 - self.num_bits = self.weight_quant.num_bits
2248 - self.packed_factor = 32 // self.num_bits
2249 -
2250 - assert self.weight_quant.symmetric, (
2251 - "Only symmetric quantization is supported for W4A8 MoE"
2252 - )
2253 - assert self.weight_quant.actorder != "group"
2254 - assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
2255 -
2256 - self.disable_expert_map = False
2257 - self.layer_name = layer_name
2258 -
2259 - from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
2260 - from vllm.model_executor.layers.quantization.utils.quant_utils import (
2261 - GroupShape,
2262 - )
2263 -
2264 - self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
2265 -
2266 - def create_weights(
2267 - self,
2268 - layer: torch.nn.Module,
2269 - num_experts: int,
2270 - hidden_size: int,
2271 - intermediate_size_per_partition: int,
2272 - params_dtype: torch.dtype,
2273 - **extra_weight_attrs,
2274 - ):
2275 - layer.num_experts = num_experts
2276 - layer.orig_dtype = params_dtype
2277 - layer.weight_block_size = None
2278 -
2279 - # requirement for CUTLASS reorder_tensor
2280 - assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256"
2281 - assert intermediate_size_per_partition % 256 == 0, (
2282 - f"{intermediate_size_per_partition=} must be divisible by 256"
2283 - )
2284 - # storage type, pack 8xint4 into int32
2285 - params_dtype = torch.int32
2286 -
2287 - # WEIGHTS
2288 - w13_weight_packed = torch.nn.Parameter(
2289 - torch.empty(
2290 - num_experts,
2291 - 2 * intermediate_size_per_partition,
2292 - hidden_size // self.packed_factor,
2293 - dtype=params_dtype,
2294 - ),
2295 - requires_grad=False,
2296 - )
2297 - layer.register_parameter("w13_weight_packed", w13_weight_packed)
2298 - set_weight_attrs(w13_weight_packed, extra_weight_attrs)
2299 -
2300 - w2_weight_packed = torch.nn.Parameter(
2301 - torch.empty(
2302 - num_experts,
2303 - hidden_size,
2304 - intermediate_size_per_partition // self.packed_factor,
2305 - dtype=params_dtype,
2306 - ),
2307 - requires_grad=False,
2308 - )
2309 - layer.register_parameter("w2_weight_packed", w2_weight_packed)
2310 - set_weight_attrs(w2_weight_packed, extra_weight_attrs)
2311 -
2312 - # SCALES
2313 - # weight_scale refers to the group-wise scales
2314 - # they are initially loaded as bf16, we will convert to fp8
2315 - # after loading
2316 - w13_weight_scale = torch.nn.Parameter(
2317 - torch.ones(
2318 - num_experts,
2319 - 2 * intermediate_size_per_partition,
2320 - hidden_size // self.group_size,
2321 - dtype=layer.orig_dtype,
2322 - ),
2323 - requires_grad=False,
2324 - )
2325 - layer.register_parameter("w13_weight_scale", w13_weight_scale)
2326 -
2327 - w2_weight_scale = torch.nn.Parameter(
2328 - torch.ones(
2329 - num_experts,
2330 - hidden_size,
2331 - intermediate_size_per_partition // self.group_size,
2332 - dtype=layer.orig_dtype,
2333 - ),
2334 - requires_grad=False,
2335 - )
2336 - layer.register_parameter("w2_weight_scale", w2_weight_scale)
2337 - # Add PER-GROUP quantization for FusedMoE.weight_loader.
2338 - extra_weight_attrs.update(
2339 - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
2340 - )
2341 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)
2342 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)
2343 -
2344 - # weight shapes
2345 - w2_weight_shape = torch.nn.Parameter(
2346 - torch.empty(num_experts, 2), requires_grad=False
2347 - )
2348 - layer.register_parameter("w2_weight_shape", w2_weight_shape)
2349 - set_weight_attrs(w2_weight_shape, extra_weight_attrs)
2350 - w13_weight_shape = torch.nn.Parameter(
2351 - torch.empty(num_experts, 2), requires_grad=False
2352 - )
2353 - layer.register_parameter("w13_weight_shape", w13_weight_shape)
2354 - set_weight_attrs(w13_weight_shape, extra_weight_attrs)
2355 -
2356 - # don't use input scales
2357 - layer.w13_input_scale = None
2358 - layer.w2_input_scale = None
2359 -
2360 - def process_weights_after_loading(self, layer):
2361 - device = layer.w13_weight_packed.device
2362 -
2363 - # STRIDES
2364 - # A, C
2365 - self.a_strides1_c_strides2 = torch.full(
2366 - (layer.local_num_experts,),
2367 - layer.hidden_size,
2368 - device=device,
2369 - dtype=torch.int64,
2370 - )
2371 - self.a_strides2 = torch.full(
2372 - (layer.local_num_experts,),
2373 - layer.intermediate_size_per_partition,
2374 - device=device,
2375 - dtype=torch.int64,
2376 - )
2377 - self.c_strides1 = torch.full(
2378 - (layer.local_num_experts,),
2379 - 2 * layer.intermediate_size_per_partition,
2380 - device=device,
2381 - dtype=torch.int64,
2382 - )
2383 -
2384 - # S (group-wise scales)
2385 - # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it
2386 - self.s_strides1 = torch.zeros(
2387 - (layer.local_num_experts, 2), device=device, dtype=torch.int64
2388 - )
2389 - self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition
2390 -
2391 - self.s_strides2 = torch.zeros(
2392 - (layer.local_num_experts, 2), device=device, dtype=torch.int64
2393 - )
2394 - self.s_strides2[:, 0] = layer.hidden_size
2395 -
2396 - # encode and reorder weight tensors, and get the layout to pass to
2397 - # the grouped gemm kernel. `b_strides1/2` specifies the entire layout
2398 - convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed)
2399 - w13_weight_shuffled, self.b_strides1 = (
2400 - ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed)
2401 - )
2402 - replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled)
2403 - convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed)
2404 - w2_weight_shuffled, self.b_strides2 = (
2405 - ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed)
2406 - )
2407 - replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled)
2408 -
2409 - # convert bf16 scales to (fp8_scales, channel_scales)
2410 - w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
2411 - self.quant_fp8, layer.w13_weight_scale
2412 - )
2413 - w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8(
2414 - self.quant_fp8, layer.w2_weight_scale
2415 - )
2416 -
2417 - # register channel scales
2418 - layer.register_parameter(
2419 - "w13_weight_chan_scale",
2420 - torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False),
2421 - )
2422 - layer.register_parameter(
2423 - "w2_weight_chan_scale",
2424 - torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False),
2425 - )
2426 -
2427 - # The scales are stored as (E, N, K // 128) but the kernel expects
2428 - # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims
2429 - # and make it contiguous
2430 - w13_weight_scale_packed = ops.cutlass_pack_scale_fp8(
2431 - w13_weight_scale.permute(0, 2, 1).contiguous()
2432 - )
2433 - replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed)
2434 - w2_weight_scale_packed = ops.cutlass_pack_scale_fp8(
2435 - w2_weight_scale.permute(0, 2, 1).contiguous()
2436 - )
2437 - replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed)
2438 -
2439 - def maybe_make_prepare_finalize(
2440 - self,
2441 - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
2442 - ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
2443 - return super().maybe_make_prepare_finalize(routing_tables)
2444 -
2445 - def get_fused_moe_quant_config(
2446 - self, layer: torch.nn.Module
2447 - ) -> FusedMoEQuantConfig | None:
2448 - # Store quantization scales; both per-group and per-channel
2449 - # Note we haven't specified the group size here because
2450 - # the quant config logic assumes group-wise scaling
2451 - # and channel-wise scaling are exclusive.
2452 - return int4_w4afp8_moe_quant_config(
2453 - w1_scale=layer.w13_weight_scale, # group scale
2454 - w2_scale=layer.w2_weight_scale, # group scale
2455 - g1_alphas=layer.w13_weight_chan_scale,
2456 - g2_alphas=layer.w2_weight_chan_scale,
2457 - per_act_token_quant=True, # always use dynamic per-token
2458 - per_out_ch_quant=True, # always use per-channel
2459 - )
2460 -
2461 - def select_gemm_impl(
2462 - self,
2463 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
2464 - layer: torch.nn.Module,
2465 - ) -> mk.FusedMoEExpertsModular:
2466 - assert self.moe_quant_config is not None
2467 - assert (
2468 - prepare_finalize.activation_format == FusedMoEActivationFormat.Standard
2469 - ), "BatchedExperts not supported"
2470 -
2471 - from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8
2472 -
2473 - experts: FusedMoEExpertsModular
2474 -
2475 - logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__)
2476 - experts = CutlassExpertsW4A8Fp8(
2477 - out_dtype=self.moe.in_dtype,
2478 - a_strides1=self.a_strides1_c_strides2,
2479 - a_strides2=self.a_strides2,
2480 - b_strides1=self.b_strides1,
2481 - b_strides2=self.b_strides2,
2482 - c_strides1=self.c_strides1,
2483 - c_strides2=self.a_strides1_c_strides2,
2484 - s_strides1=self.s_strides1,
2485 - s_strides2=self.s_strides2,
2486 - moe_config=self.moe,
2487 - quant_config=self.moe_quant_config,
2488 - group_size=self.group_size,
2489 - )
2490 -
2491 - num_dispatchers = prepare_finalize.num_dispatchers()
2492 - self.disable_expert_map = (
2493 - num_dispatchers > 1 or not experts.supports_expert_map()
2494 - )
2495 -
2496 - return experts
2497 -
2498 - def apply(
2499 - self,
2500 - layer: FusedMoE,
2501 - x: torch.Tensor,
2502 - topk_weights: torch.Tensor,
2503 - topk_ids: torch.Tensor,
2504 - shared_experts_input: torch.Tensor | None,
2505 - ) -> torch.Tensor:
2506 - if layer.enable_eplb:
2507 - raise NotImplementedError(
2508 - "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
2509 - )
2510 - assert self.moe_quant_config is not None
2511 -
2512 - from vllm.model_executor.layers.fused_moe.cutlass_moe import (
2513 - cutlass_moe_w4a8_fp8,
2514 - )
2515 -
2516 - return cutlass_moe_w4a8_fp8(
2517 - x,
2518 - layer.w13_weight_packed,
2519 - layer.w2_weight_packed,
2520 - topk_weights,
2521 - topk_ids,
2522 - moe_config=self.moe,
2523 - quant_config=self.moe_quant_config,
2524 - activation=layer.activation,
2525 - global_num_experts=layer.global_num_experts,
2526 - expert_map=None if self.disable_expert_map else layer.expert_map,
2527 - a_strides1=self.a_strides1_c_strides2,
2528 - a_strides2=self.a_strides2,
2529 - b_strides1=self.b_strides1,
2530 - b_strides2=self.b_strides2,
2531 - c_strides1=self.c_strides1,
2532 - c_strides2=self.a_strides1_c_strides2,
2533 - s_strides1=self.s_strides1,
2534 - s_strides2=self.s_strides2,
2535 - group_size=self.group_size,
2536 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
2537 - )
2538 -
2539 - @property
2540 - def supports_eplb(self) -> bool:
2541 - return False
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/__init__.py new
+10
@@ -0,0 +1,10 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501
5 + CompressedTensorsMoEMethod,
6 +)
7 +
8 +__all__ = [
9 + "CompressedTensorsMoEMethod",
10 +]
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py new
+175
@@ -0,0 +1,175 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors import CompressionFormat
7 +from compressed_tensors.quantization import (
8 + ActivationOrdering,
9 + QuantizationStrategy,
10 +)
11 +
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe import (
14 + FusedMoEMethodBase,
15 + UnquantizedFusedMoEMethod,
16 +)
17 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
18 + WNA16_SUPPORTED_BITS,
19 +)
20 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
21 + check_moe_marlin_supports_layer,
22 +)
23 +from vllm.platforms import current_platform
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +class CompressedTensorsMoEMethod(FusedMoEMethodBase):
29 + @staticmethod
30 + def get_moe_method(
31 + quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
32 + layer: torch.nn.Module,
33 + layer_name: str,
34 + ) -> FusedMoEMethodBase:
35 + # FusedMoE was made by combining multiple Linears so need to
36 + # make sure quantization config for Linear can target it
37 + quant_config._add_fused_moe_to_target_scheme_map()
38 + unfused_names = [
39 + layer_name + proj_name
40 + for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
41 + ]
42 + # TODO: refactor this to use expert_mapping and check all layer numbers
43 + all_scheme_dicts = [
44 + quant_config.get_scheme_dict(layer, name) for name in unfused_names
45 + ]
46 + scheme_dict = all_scheme_dicts.pop()
47 +
48 + # multiple schemes found
49 + if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
50 + raise ValueError(
51 + "All MoE projections need to have same "
52 + "quantization scheme but found multiple"
53 + )
54 +
55 + if scheme_dict is None: # ignored layer
56 + return UnquantizedFusedMoEMethod(layer.moe_config)
57 +
58 + # TODO: @dsikka: refactor this to use schemes as other kernels
59 + # are supported + check if the layer is being ignored.
60 + weight_quant = scheme_dict.get("weights")
61 + input_quant = scheme_dict.get("input_activations")
62 + format = scheme_dict.get("format")
63 +
64 + if quant_config._is_mxfp4(weight_quant):
65 + from .compressed_tensors_moe_w4a4_mxfp4 import (
66 + CompressedTensorsW4A4Mxfp4MoEMethod,
67 + )
68 +
69 + return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
70 +
71 + if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
72 + # group_size=None means channelwise
73 + group_size = weight_quant.group_size or -1
74 +
75 + valid_format_and_bits = (
76 + weight_quant.num_bits in WNA16_SUPPORTED_BITS
77 + and format == CompressionFormat.pack_quantized.value
78 + )
79 +
80 + if not valid_format_and_bits:
81 + raise ValueError(
82 + "For Fused MoE layers, only format: ",
83 + f"{CompressionFormat.pack_quantized.value} ",
84 + f" and bits: {WNA16_SUPPORTED_BITS} is supported ",
85 + f"but got format: {CompressionFormat.pack_quantized.value} "
86 + f" and bits: {weight_quant.num_bits}",
87 + )
88 +
89 + # Prefer to use the MarlinMoE kernel when it is supported.
90 + if (
91 + not check_moe_marlin_supports_layer(layer, group_size)
92 + or current_platform.is_rocm()
93 + ):
94 + from .compressed_tensors_moe_wna16 import (
95 + CompressedTensorsWNA16MoEMethod,
96 + )
97 +
98 + if (
99 + weight_quant.strategy == QuantizationStrategy.GROUP
100 + and weight_quant.actorder
101 + in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
102 + ):
103 + raise ValueError(
104 + "WNA16MoE is not supported with actorder=group/dynamic."
105 + )
106 + logger.info_once("Using CompressedTensorsWNA16MoEMethod")
107 + return CompressedTensorsWNA16MoEMethod(
108 + weight_quant, input_quant, layer.moe_config
109 + )
110 + else:
111 + from .compressed_tensors_moe_wna16_marlin import (
112 + CompressedTensorsWNA16MarlinMoEMethod,
113 + )
114 +
115 + logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
116 + return CompressedTensorsWNA16MarlinMoEMethod(
117 + weight_quant, input_quant, layer.moe_config
118 + )
119 + elif quant_config._is_nvfp4_format(weight_quant):
120 + from .compressed_tensors_moe_w4a4_nvfp4 import (
121 + CompressedTensorsW4A4Nvfp4MoEMethod,
122 + )
123 +
124 + _is_valid_nvfp4_activations = (
125 + quant_config._is_nvfp4_format(input_quant) or input_quant is None
126 + )
127 + if not _is_valid_nvfp4_activations:
128 + raise ValueError(
129 + "For NVFP4 weights, input quantization must also be NVFP4 format ",
130 + f"or None for NVFP4A16, found {input_quant}",
131 + )
132 + return CompressedTensorsW4A4Nvfp4MoEMethod(
133 + layer.moe_config, layer_name, use_a16=(input_quant is None)
134 + )
135 + elif (
136 + quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
137 + or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
138 + or quant_config._is_fp8_w8a8(weight_quant, input_quant)
139 + ):
140 + from .compressed_tensors_moe_w8a8_fp8 import (
141 + CompressedTensorsW8A8Fp8MoEMethod,
142 + )
143 +
144 + return CompressedTensorsW8A8Fp8MoEMethod(
145 + weight_quant, input_quant, layer.moe_config
146 + )
147 + elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
148 + from .compressed_tensors_moe_w8a8_int8 import (
149 + CompressedTensorsW8A8Int8MoEMethod,
150 + )
151 +
152 + return CompressedTensorsW8A8Int8MoEMethod(
153 + weight_quant, input_quant, layer.moe_config
154 + )
155 + elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
156 + from .compressed_tensors_moe_w4a8_fp8 import (
157 + CompressedTensorsW4A8Fp8MoEMethod,
158 + )
159 +
160 + logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod")
161 + return CompressedTensorsW4A8Fp8MoEMethod(
162 + weight_quant, input_quant, layer.moe_config
163 + )
164 + elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
165 + from .compressed_tensors_moe_w4a8_int8 import (
166 + CompressedTensorsW4A8Int8MoEMethod,
167 + )
168 +
169 + return CompressedTensorsW4A8Int8MoEMethod(
170 + weight_quant, input_quant, layer.moe_config
171 + )
172 + else:
173 + raise RuntimeError(
174 + f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
175 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py new
+168
@@ -0,0 +1,168 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +
7 +from vllm.logger import init_logger
8 +from vllm.model_executor.layers.fused_moe import (
9 + FusedMoE,
10 + FusedMoeWeightScaleSupported,
11 +)
12 +from vllm.model_executor.layers.fused_moe.config import (
13 + FusedMoEQuantConfig,
14 +)
15 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
16 + MarlinExperts,
17 +)
18 +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
19 + Mxfp4MoeBackend,
20 + make_mxfp4_moe_kernel,
21 + make_mxfp4_moe_quant_config,
22 +)
23 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
24 + CompressedTensorsMoEMethod,
25 +)
26 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
27 + prepare_moe_fp4_layer_for_marlin,
28 +)
29 +from vllm.model_executor.utils import set_weight_attrs
30 +
31 +logger = init_logger(__name__)
32 +
33 +
34 +class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
35 + def __init__(self, moe):
36 + super().__init__(moe)
37 + self.group_size = 32
38 + self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
39 + self.experts_cls = MarlinExperts
40 +
41 + def create_weights(
42 + self,
43 + layer: torch.nn.Module,
44 + num_experts: int,
45 + hidden_size: int,
46 + intermediate_size_per_partition: int,
47 + params_dtype: torch.dtype,
48 + **extra_weight_attrs,
49 + ):
50 + layer.num_experts = num_experts
51 + layer.params_dtype = params_dtype
52 +
53 + w13_weight = torch.nn.Parameter(
54 + torch.empty(
55 + num_experts,
56 + 2 * intermediate_size_per_partition,
57 + # 2 fp4 items are packed in the input dimension
58 + hidden_size // 2,
59 + requires_grad=False,
60 + dtype=torch.uint8,
61 + ),
62 + requires_grad=False,
63 + )
64 + layer.register_parameter("w13_weight_packed", w13_weight)
65 + set_weight_attrs(w13_weight, extra_weight_attrs)
66 +
67 + w2_weight = torch.nn.Parameter(
68 + torch.empty(
69 + num_experts,
70 + hidden_size,
71 + # 2 fp4 items are packed in the input dimension
72 + intermediate_size_per_partition // 2,
73 + dtype=torch.uint8,
74 + ),
75 + requires_grad=False,
76 + )
77 + layer.register_parameter("w2_weight_packed", w2_weight)
78 + set_weight_attrs(w2_weight, extra_weight_attrs)
79 +
80 + w13_weight_scale = torch.nn.Parameter(
81 + torch.empty(
82 + num_experts,
83 + 2 * intermediate_size_per_partition,
84 + # 2 fp4 items are packed in the input dimension
85 + hidden_size // self.group_size,
86 + dtype=torch.uint8,
87 + ),
88 + requires_grad=False,
89 + )
90 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
91 + extra_weight_attrs.update(
92 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
93 + )
94 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
95 +
96 + w2_weight_scale = torch.nn.Parameter(
97 + torch.empty(
98 + num_experts,
99 + hidden_size,
100 + # 2 fp4 items are packed in the input dimension
101 + intermediate_size_per_partition // self.group_size,
102 + dtype=torch.uint8,
103 + ),
104 + requires_grad=False,
105 + )
106 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
107 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
108 +
109 + def get_fused_moe_quant_config(
110 + self, layer: torch.nn.Module
111 + ) -> FusedMoEQuantConfig | None:
112 + return make_mxfp4_moe_quant_config(
113 + mxfp4_backend=self.mxfp4_backend,
114 + w1_scale=layer.w13_weight_scale,
115 + w2_scale=layer.w2_weight_scale,
116 + )
117 +
118 + def process_weights_after_loading(self, layer: FusedMoE) -> None:
119 + layer.w13_weight = torch.nn.Parameter(
120 + layer.w13_weight_packed.data, requires_grad=False
121 + )
122 + delattr(layer, "w13_weight_packed")
123 +
124 + layer.w2_weight = torch.nn.Parameter(
125 + layer.w2_weight_packed.data, requires_grad=False
126 + )
127 + delattr(layer, "w2_weight_packed")
128 +
129 + logger.warning_once(
130 + "Your GPU does not have native support for FP4 computation but "
131 + "FP4 quantization is being used. Weight-only FP4 compression "
132 + "will be used leveraging the Marlin kernel. This may degrade "
133 + "performance for compute-heavy workloads."
134 + )
135 + prepare_moe_fp4_layer_for_marlin(layer)
136 +
137 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
138 + if self.moe_quant_config is not None:
139 + self.moe_kernel = make_mxfp4_moe_kernel(
140 + moe_quant_config=self.moe_quant_config,
141 + moe_config=self.moe,
142 + experts_cls=self.experts_cls,
143 + mxfp4_backend=self.mxfp4_backend,
144 + shared_experts=layer.shared_experts,
145 + routing_tables=layer._maybe_init_expert_routing_tables(),
146 + )
147 +
148 + def apply(
149 + self,
150 + layer: FusedMoE,
151 + x: torch.Tensor,
152 + topk_weights: torch.Tensor,
153 + topk_ids: torch.Tensor,
154 + shared_experts_input: torch.Tensor | None,
155 + ) -> torch.Tensor:
156 + assert self.moe_kernel is not None
157 + return self.moe_kernel.apply(
158 + x,
159 + layer.w13_weight,
160 + layer.w2_weight,
161 + topk_weights,
162 + topk_ids,
163 + activation=layer.activation,
164 + global_num_experts=layer.global_num_experts,
165 + expert_map=layer.expert_map,
166 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
167 + shared_experts_input=shared_experts_input,
168 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py new
+306
@@ -0,0 +1,306 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.layers.fused_moe import (
10 + FusedMoE,
11 + FusedMoeWeightScaleSupported,
12 +)
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FusedMoEConfig,
15 + FusedMoEQuantConfig,
16 +)
17 +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
18 + convert_to_nvfp4_moe_kernel_format,
19 + is_global_sf_supported_for_nvfp4_backend,
20 + make_nvfp4_moe_kernel,
21 + make_nvfp4_moe_quant_config,
22 + select_nvfp4_moe_backend,
23 +)
24 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
25 + CompressedTensorsMoEMethod,
26 +)
27 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 + kNvfp4Dynamic,
29 + kNvfp4Static,
30 +)
31 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
32 +
33 +logger = init_logger(__name__)
34 +
35 +
36 +class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
37 + def __init__(
38 + self,
39 + moe: FusedMoEConfig,
40 + layer_name: str | None = None,
41 + use_a16: bool = False,
42 + ):
43 + super().__init__(moe)
44 + self.group_size = 16
45 +
46 + # Select experts implementation.
47 + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
48 + config=self.moe,
49 + weight_key=kNvfp4Static,
50 + activation_key=None if use_a16 else kNvfp4Dynamic,
51 + )
52 +
53 + self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
54 + self.nvfp4_backend
55 + )
56 +
57 + def create_weights(
58 + self,
59 + layer: torch.nn.Module,
60 + num_experts: int,
61 + hidden_size: int,
62 + intermediate_size_per_partition: int,
63 + params_dtype: torch.dtype,
64 + **extra_weight_attrs,
65 + ):
66 + layer.num_experts = num_experts
67 + layer.params_dtype = params_dtype
68 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
69 +
70 + w13_weight = torch.nn.Parameter(
71 + torch.empty(
72 + num_experts,
73 + w13_num_shards * intermediate_size_per_partition,
74 + # 2 fp4 items are packed in the input dimension
75 + hidden_size // 2,
76 + requires_grad=False,
77 + dtype=torch.uint8,
78 + ),
79 + requires_grad=False,
80 + )
81 + layer.register_parameter("w13_weight_packed", w13_weight)
82 + set_weight_attrs(w13_weight, extra_weight_attrs)
83 +
84 + w2_weight = torch.nn.Parameter(
85 + torch.empty(
86 + num_experts,
87 + hidden_size,
88 + # 2 fp4 items are packed in the input dimension
89 + intermediate_size_per_partition // 2,
90 + dtype=torch.uint8,
91 + ),
92 + requires_grad=False,
93 + )
94 + layer.register_parameter("w2_weight_packed", w2_weight)
95 + set_weight_attrs(w2_weight, extra_weight_attrs)
96 +
97 + # Weight Scales
98 + w13_weight_scale = torch.nn.Parameter(
99 + torch.empty(
100 + num_experts,
101 + w13_num_shards * intermediate_size_per_partition,
102 + # 2 fp4 items are packed in the input dimension
103 + hidden_size // self.group_size,
104 + dtype=torch.float8_e4m3fn,
105 + ),
106 + requires_grad=False,
107 + )
108 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
109 + extra_weight_attrs.update(
110 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
111 + )
112 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
113 +
114 + w2_weight_scale = torch.nn.Parameter(
115 + torch.empty(
116 + num_experts,
117 + hidden_size,
118 + # 2 fp4 items are packed in the input dimension
119 + intermediate_size_per_partition // self.group_size,
120 + dtype=torch.float8_e4m3fn,
121 + ),
122 + requires_grad=False,
123 + )
124 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
125 + extra_weight_attrs.update(
126 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
127 + )
128 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
129 +
130 + # Weight Global Scales
131 + w13_weight_scale_2 = torch.nn.Parameter(
132 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
133 + requires_grad=False,
134 + )
135 + layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)
136 + extra_weight_attrs.update(
137 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
138 + )
139 + set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
140 +
141 + w2_weight_scale_2 = torch.nn.Parameter(
142 + torch.empty(num_experts, dtype=torch.float32), requires_grad=False
143 + )
144 + layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
145 + extra_weight_attrs.update(
146 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
147 + )
148 + set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
149 +
150 + # Input Global Scales
151 + w13_input_scale = torch.nn.Parameter(
152 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
153 + requires_grad=False,
154 + )
155 + layer.register_parameter("w13_input_global_scale", w13_input_scale)
156 + extra_weight_attrs.update(
157 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
158 + )
159 + set_weight_attrs(w13_input_scale, extra_weight_attrs)
160 +
161 + w2_input_scale = torch.nn.Parameter(
162 + torch.empty(num_experts, dtype=torch.float32), requires_grad=False
163 + )
164 + layer.register_parameter("w2_input_global_scale", w2_input_scale)
165 + extra_weight_attrs.update(
166 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
167 + )
168 + set_weight_attrs(w2_input_scale, extra_weight_attrs)
169 +
170 + def process_weights_after_loading(self, layer: FusedMoE) -> None:
171 + """
172 + Convert NVFP4 MoE weights into kernel format and setup the kernel.
173 + """
174 + # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
175 + # requires this naming convention. However, the name change breaks
176 + # reloading because the state dict no longer matches disk. Once we
177 + # remove MKM, we should revert this change to ensure compatibility.
178 + layer.w13_weight = torch.nn.Parameter(
179 + layer.w13_weight_packed.data, requires_grad=False
180 + )
181 + delattr(layer, "w13_weight_packed")
182 +
183 + layer.w2_weight = torch.nn.Parameter(
184 + layer.w2_weight_packed.data, requires_grad=False
185 + )
186 + delattr(layer, "w2_weight_packed")
187 +
188 + # Use a single gscale for w13.
189 + if self.moe.is_act_and_mul and not torch.allclose(
190 + layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]
191 + ):
192 + logger.warning_once(
193 + "w1_weight_global_scale must match w3_weight_global_scale. "
194 + "Accuracy may be affected.",
195 + )
196 + w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()
197 +
198 + # Shuffle weights into the NvFp4 kernel format.
199 + (
200 + w13,
201 + w13_scale,
202 + w13_scale_2,
203 + a13_scale,
204 + w2,
205 + w2_scale,
206 + w2_scale_2,
207 + a2_scale,
208 + ) = convert_to_nvfp4_moe_kernel_format(
209 + nvfp4_backend=self.nvfp4_backend,
210 + layer=layer,
211 + w13=layer.w13_weight,
212 + w13_scale=layer.w13_weight_scale,
213 + w13_scale_2=(1.0 / w13_weight_global_scale),
214 + a13_scale=(1.0 / layer.w13_input_global_scale),
215 + w2=layer.w2_weight,
216 + w2_scale=layer.w2_weight_scale,
217 + w2_scale_2=(1.0 / layer.w2_weight_global_scale),
218 + a2_scale=(1.0 / layer.w2_input_global_scale),
219 + is_act_and_mul=self.moe.is_act_and_mul,
220 + )
221 +
222 + replace_parameter(layer, "w13_weight", w13)
223 + replace_parameter(layer, "w13_weight_scale", w13_scale)
224 + replace_parameter(layer, "w2_weight", w2)
225 + replace_parameter(layer, "w2_weight_scale", w2_scale)
226 + layer.w13_weight_scale_2 = w13_scale_2
227 + layer.w2_weight_scale_2 = w2_scale_2
228 + layer.w13_input_scale = a13_scale
229 + layer.w2_input_scale = a2_scale
230 +
231 + # Setup modular kernel.
232 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
233 + assert self.experts_cls is not None
234 + self.moe_kernel = make_nvfp4_moe_kernel(
235 + moe_quant_config=self.moe_quant_config,
236 + moe_config=self.moe,
237 + experts_cls=self.experts_cls,
238 + shared_experts=layer.shared_experts,
239 + routing_tables=layer._maybe_init_expert_routing_tables(),
240 + )
241 + self.moe_kernel.fused_experts.process_weights_after_loading(layer)
242 +
243 + def maybe_make_prepare_finalize(
244 + self,
245 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
246 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
247 + raise ValueError(
248 + f"{self.__class__.__name__} uses the new modular kernel initialization "
249 + "logic. This function should not be called."
250 + )
251 +
252 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
253 + return make_nvfp4_moe_quant_config(
254 + backend=self.nvfp4_backend,
255 + w13_scale=layer.w13_weight_scale,
256 + w2_scale=layer.w2_weight_scale,
257 + w13_scale_2=layer.w13_weight_scale_2,
258 + w2_scale_2=layer.w2_weight_scale_2,
259 + a13_scale=layer.w13_input_scale,
260 + a2_scale=layer.w2_input_scale,
261 + )
262 +
263 + def apply_monolithic(
264 + self,
265 + layer: FusedMoE,
266 + x: torch.Tensor,
267 + router_logits: torch.Tensor,
268 + ) -> torch.Tensor:
269 + assert self.is_monolithic
270 + assert self.moe_kernel is not None
271 + return self.moe_kernel.apply_monolithic(
272 + x,
273 + layer.w13_weight,
274 + layer.w2_weight,
275 + router_logits,
276 + activation=layer.activation,
277 + global_num_experts=layer.global_num_experts,
278 + expert_map=layer.expert_map,
279 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
280 + num_expert_group=layer.num_expert_group,
281 + topk_group=layer.topk_group,
282 + e_score_correction_bias=layer.e_score_correction_bias,
283 + routed_scaling_factor=layer.routed_scaling_factor,
284 + )
285 +
286 + def apply(
287 + self,
288 + layer: FusedMoE,
289 + x: torch.Tensor,
290 + topk_weights: torch.Tensor,
291 + topk_ids: torch.Tensor,
292 + shared_experts_input: torch.Tensor | None,
293 + ) -> torch.Tensor:
294 + assert self.moe_kernel is not None
295 + return self.moe_kernel.apply(
296 + x,
297 + layer.w13_weight,
298 + layer.w2_weight,
299 + topk_weights,
300 + topk_ids,
301 + activation=layer.activation,
302 + global_num_experts=layer.global_num_experts,
303 + expert_map=layer.expert_map,
304 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
305 + shared_experts_input=shared_experts_input,
306 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py new
+343
@@ -0,0 +1,343 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors.quantization import (
7 + QuantizationArgs,
8 +)
9 +
10 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 +from vllm import _custom_ops as ops
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe import (
14 + FusedMoE,
15 + FusedMoEActivationFormat,
16 + FusedMoEExpertsModular,
17 + FusedMoeWeightScaleSupported,
18 +)
19 +from vllm.model_executor.layers.fused_moe.config import (
20 + FusedMoEConfig,
21 + FusedMoEQuantConfig,
22 + int4_w4afp8_moe_quant_config,
23 +)
24 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
25 + CompressedTensorsMoEMethod,
26 +)
27 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 + convert_bf16_scales_to_fp8,
29 + convert_packed_uint4b8_to_signed_int4_inplace,
30 +)
31 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
32 +
33 +logger = init_logger(__name__)
34 +
35 +
36 +class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
37 + def __init__(
38 + self,
39 + weight_quant: QuantizationArgs,
40 + input_quant: QuantizationArgs,
41 + moe: FusedMoEConfig,
42 + layer_name: str | None = None,
43 + ):
44 + super().__init__(moe)
45 + self.weight_quant = weight_quant
46 + self.input_quant = input_quant
47 +
48 + self.group_size = self.weight_quant.group_size
49 + self.num_bits = self.weight_quant.num_bits
50 + self.packed_factor = 32 // self.num_bits
51 +
52 + assert self.weight_quant.symmetric, (
53 + "Only symmetric quantization is supported for W4A8 MoE"
54 + )
55 + assert self.weight_quant.actorder != "group"
56 + assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
57 +
58 + self.disable_expert_map = False
59 + self.layer_name = layer_name
60 +
61 + from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
62 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
63 + GroupShape,
64 + )
65 +
66 + self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
67 +
68 + def create_weights(
69 + self,
70 + layer: torch.nn.Module,
71 + num_experts: int,
72 + hidden_size: int,
73 + intermediate_size_per_partition: int,
74 + params_dtype: torch.dtype,
75 + **extra_weight_attrs,
76 + ):
77 + layer.num_experts = num_experts
78 + layer.orig_dtype = params_dtype
79 + layer.weight_block_size = None
80 +
81 + # requirement for CUTLASS reorder_tensor
82 + assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256"
83 + assert intermediate_size_per_partition % 256 == 0, (
84 + f"{intermediate_size_per_partition=} must be divisible by 256"
85 + )
86 + # storage type, pack 8xint4 into int32
87 + params_dtype = torch.int32
88 +
89 + # WEIGHTS
90 + w13_weight_packed = torch.nn.Parameter(
91 + torch.empty(
92 + num_experts,
93 + 2 * intermediate_size_per_partition,
94 + hidden_size // self.packed_factor,
95 + dtype=params_dtype,
96 + ),
97 + requires_grad=False,
98 + )
99 + layer.register_parameter("w13_weight_packed", w13_weight_packed)
100 + set_weight_attrs(w13_weight_packed, extra_weight_attrs)
101 +
102 + w2_weight_packed = torch.nn.Parameter(
103 + torch.empty(
104 + num_experts,
105 + hidden_size,
106 + intermediate_size_per_partition // self.packed_factor,
107 + dtype=params_dtype,
108 + ),
109 + requires_grad=False,
110 + )
111 + layer.register_parameter("w2_weight_packed", w2_weight_packed)
112 + set_weight_attrs(w2_weight_packed, extra_weight_attrs)
113 +
114 + # SCALES
115 + # weight_scale refers to the group-wise scales
116 + # they are initially loaded as bf16, we will convert to fp8
117 + # after loading
118 + w13_weight_scale = torch.nn.Parameter(
119 + torch.ones(
120 + num_experts,
121 + 2 * intermediate_size_per_partition,
122 + hidden_size // self.group_size,
123 + dtype=layer.orig_dtype,
124 + ),
125 + requires_grad=False,
126 + )
127 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
128 +
129 + w2_weight_scale = torch.nn.Parameter(
130 + torch.ones(
131 + num_experts,
132 + hidden_size,
133 + intermediate_size_per_partition // self.group_size,
134 + dtype=layer.orig_dtype,
135 + ),
136 + requires_grad=False,
137 + )
138 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
139 + # Add PER-GROUP quantization for FusedMoE.weight_loader.
140 + extra_weight_attrs.update(
141 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
142 + )
143 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
144 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
145 +
146 + # weight shapes
147 + w2_weight_shape = torch.nn.Parameter(
148 + torch.empty(num_experts, 2), requires_grad=False
149 + )
150 + layer.register_parameter("w2_weight_shape", w2_weight_shape)
151 + set_weight_attrs(w2_weight_shape, extra_weight_attrs)
152 + w13_weight_shape = torch.nn.Parameter(
153 + torch.empty(num_experts, 2), requires_grad=False
154 + )
155 + layer.register_parameter("w13_weight_shape", w13_weight_shape)
156 + set_weight_attrs(w13_weight_shape, extra_weight_attrs)
157 +
158 + # don't use input scales
159 + layer.w13_input_scale = None
160 + layer.w2_input_scale = None
161 +
162 + def process_weights_after_loading(self, layer):
163 + device = layer.w13_weight_packed.device
164 +
165 + # STRIDES
166 + # A, C
167 + self.a_strides1_c_strides2 = torch.full(
168 + (layer.local_num_experts,),
169 + layer.hidden_size,
170 + device=device,
171 + dtype=torch.int64,
172 + )
173 + self.a_strides2 = torch.full(
174 + (layer.local_num_experts,),
175 + layer.intermediate_size_per_partition,
176 + device=device,
177 + dtype=torch.int64,
178 + )
179 + self.c_strides1 = torch.full(
180 + (layer.local_num_experts,),
181 + 2 * layer.intermediate_size_per_partition,
182 + device=device,
183 + dtype=torch.int64,
184 + )
185 +
186 + # S (group-wise scales)
187 + # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it
188 + self.s_strides1 = torch.zeros(
189 + (layer.local_num_experts, 2), device=device, dtype=torch.int64
190 + )
191 + self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition
192 +
193 + self.s_strides2 = torch.zeros(
194 + (layer.local_num_experts, 2), device=device, dtype=torch.int64
195 + )
196 + self.s_strides2[:, 0] = layer.hidden_size
197 +
198 + # encode and reorder weight tensors, and get the layout to pass to
199 + # the grouped gemm kernel. `b_strides1/2` specifies the entire layout
200 + convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed)
201 + w13_weight_shuffled, self.b_strides1 = (
202 + ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed)
203 + )
204 + replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled)
205 + convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed)
206 + w2_weight_shuffled, self.b_strides2 = (
207 + ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed)
208 + )
209 + replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled)
210 +
211 + # convert bf16 scales to (fp8_scales, channel_scales)
212 + w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
213 + self.quant_fp8, layer.w13_weight_scale
214 + )
215 + w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8(
216 + self.quant_fp8, layer.w2_weight_scale
217 + )
218 +
219 + # register channel scales
220 + layer.register_parameter(
221 + "w13_weight_chan_scale",
222 + torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False),
223 + )
224 + layer.register_parameter(
225 + "w2_weight_chan_scale",
226 + torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False),
227 + )
228 +
229 + # The scales are stored as (E, N, K // 128) but the kernel expects
230 + # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims
231 + # and make it contiguous
232 + w13_weight_scale_packed = ops.cutlass_pack_scale_fp8(
233 + w13_weight_scale.permute(0, 2, 1).contiguous()
234 + )
235 + replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed)
236 + w2_weight_scale_packed = ops.cutlass_pack_scale_fp8(
237 + w2_weight_scale.permute(0, 2, 1).contiguous()
238 + )
239 + replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed)
240 +
241 + def maybe_make_prepare_finalize(
242 + self,
243 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
244 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
245 + return super().maybe_make_prepare_finalize(routing_tables)
246 +
247 + def get_fused_moe_quant_config(
248 + self, layer: torch.nn.Module
249 + ) -> FusedMoEQuantConfig | None:
250 + # Store quantization scales; both per-group and per-channel
251 + # Note we haven't specified the group size here because
252 + # the quant config logic assumes group-wise scaling
253 + # and channel-wise scaling are exclusive.
254 + return int4_w4afp8_moe_quant_config(
255 + w1_scale=layer.w13_weight_scale, # group scale
256 + w2_scale=layer.w2_weight_scale, # group scale
257 + g1_alphas=layer.w13_weight_chan_scale,
258 + g2_alphas=layer.w2_weight_chan_scale,
259 + per_act_token_quant=True, # always use dynamic per-token
260 + per_out_ch_quant=True, # always use per-channel
261 + )
262 +
263 + def select_gemm_impl(
264 + self,
265 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
266 + layer: torch.nn.Module,
267 + ) -> mk.FusedMoEExpertsModular:
268 + assert self.moe_quant_config is not None
269 + assert (
270 + prepare_finalize.activation_format == FusedMoEActivationFormat.Standard
271 + ), "BatchedExperts not supported"
272 +
273 + from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8
274 +
275 + experts: FusedMoEExpertsModular
276 +
277 + logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__)
278 + experts = CutlassExpertsW4A8Fp8(
279 + out_dtype=self.moe.in_dtype,
280 + a_strides1=self.a_strides1_c_strides2,
281 + a_strides2=self.a_strides2,
282 + b_strides1=self.b_strides1,
283 + b_strides2=self.b_strides2,
284 + c_strides1=self.c_strides1,
285 + c_strides2=self.a_strides1_c_strides2,
286 + s_strides1=self.s_strides1,
287 + s_strides2=self.s_strides2,
288 + moe_config=self.moe,
289 + quant_config=self.moe_quant_config,
290 + group_size=self.group_size,
291 + )
292 +
293 + num_dispatchers = prepare_finalize.num_dispatchers()
294 + self.disable_expert_map = (
295 + num_dispatchers > 1 or not experts.supports_expert_map()
296 + )
297 +
298 + return experts
299 +
300 + def apply(
301 + self,
302 + layer: FusedMoE,
303 + x: torch.Tensor,
304 + topk_weights: torch.Tensor,
305 + topk_ids: torch.Tensor,
306 + shared_experts_input: torch.Tensor | None,
307 + ) -> torch.Tensor:
308 + if layer.enable_eplb:
309 + raise NotImplementedError(
310 + "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
311 + )
312 + assert self.moe_quant_config is not None
313 +
314 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
315 + cutlass_moe_w4a8_fp8,
316 + )
317 +
318 + return cutlass_moe_w4a8_fp8(
319 + x,
320 + layer.w13_weight_packed,
321 + layer.w2_weight_packed,
322 + topk_weights,
323 + topk_ids,
324 + moe_config=self.moe,
325 + quant_config=self.moe_quant_config,
326 + activation=layer.activation,
327 + global_num_experts=layer.global_num_experts,
328 + expert_map=None if self.disable_expert_map else layer.expert_map,
329 + a_strides1=self.a_strides1_c_strides2,
330 + a_strides2=self.a_strides2,
331 + b_strides1=self.b_strides1,
332 + b_strides2=self.b_strides2,
333 + c_strides1=self.c_strides1,
334 + c_strides2=self.a_strides1_c_strides2,
335 + s_strides1=self.s_strides1,
336 + s_strides2=self.s_strides2,
337 + group_size=self.group_size,
338 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
339 + )
340 +
341 + @property
342 + def supports_eplb(self) -> bool:
343 + return False
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py new
+349
@@ -0,0 +1,349 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors.quantization import (
7 + QuantizationArgs,
8 + QuantizationStrategy,
9 +)
10 +
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe import (
13 + FusedMoE,
14 +)
15 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 +from vllm.model_executor.layers.fused_moe.config import (
17 + FusedMoEConfig,
18 + FusedMoEQuantConfig,
19 +)
20 +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts
21 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
22 + CompressedTensorsMoEMethod,
23 +)
24 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
25 +from vllm.platforms import CpuArchEnum, current_platform
26 +
27 +logger = init_logger(__name__)
28 +
29 +
30 +class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
31 + """
32 + CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
33 + - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles)
34 + - Scales: Fp32 for Channelwise , bf16 for groupwise quantization
35 + - Bias: Same data type as original weights
36 + - Activations: FP32/Bf16 dynamic per-token (A8 Int),
37 + quantized inside the kernel
38 + """
39 +
40 + def __init__(
41 + self,
42 + weight_quant: QuantizationArgs,
43 + input_quant: QuantizationArgs,
44 + moe: FusedMoEConfig,
45 + layer_name: str | None = None,
46 + ):
47 + super().__init__(moe)
48 + self.has_bias = self.moe.has_bias
49 + self.weight_quant = weight_quant
50 + self.input_quant = input_quant
51 +
52 + # Validate scheme: weights=W4 (channel or group),
53 + # activations=dynamic TOKEN (A8)
54 +
55 + # Must be dynamic per-token activations
56 + if (
57 + input_quant.strategy != QuantizationStrategy.TOKEN
58 + or not input_quant.dynamic
59 + ):
60 + raise ValueError(
61 + "W4A8-int MoE needs dynamic per-token activation quantization."
62 + )
63 +
64 + # Weight can be channel-wise (group_size=None) or group-wise
65 + self.group_size = (
66 + weight_quant.group_size if (weight_quant.group_size is not None) else -1
67 + )
68 + if weight_quant.num_bits != 4:
69 + raise ValueError("This method only supports 4-bit weights (num_bits=4).")
70 +
71 + # CPU only
72 + if not current_platform.is_cpu():
73 + raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.")
74 +
75 + # Arm: check _dyn ops availability
76 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
77 + try:
78 + _ = torch.ops.aten._dyn_quant_matmul_4bit
79 + _ = torch.ops.aten._dyn_quant_pack_4bit_weight
80 + except AttributeError as err:
81 + raise RuntimeError(
82 + f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
83 + install a newer build."""
84 + ) from err
85 + self.static_input_scales = False # always dynamic per token
86 +
87 + # ---- parameter creation ----
88 + def create_weights(
89 + self,
90 + layer: torch.nn.Module,
91 + num_experts: int,
92 + hidden_size: int,
93 + intermediate_size_per_partition: int,
94 + params_dtype: torch.dtype,
95 + **extra_weight_attrs,
96 + ):
97 + # Shapes per local rank (TP/EP):
98 + # w13: [E, 2*I_local, H] int8 (int4 values in [-8,7])
99 + # w2 : [E, H, I_local] int8
100 + # Scales:
101 + # channel-wise: group_size=-1 -> per-output-row, single scale per row
102 + # group-wise : group_size=g ->
103 + # per-output-row, (in_features/g) scales
104 +
105 + E = num_experts
106 + H = hidden_size
107 + IN = intermediate_size_per_partition
108 + g = self.group_size
109 +
110 + # Per-row scale columns
111 + def _n_scale_cols(in_features: int) -> int:
112 + return 1 if g == -1 else (in_features // g)
113 +
114 + # Register unpacked int4-as-int8 weights the loader will fill.
115 + w13 = torch.nn.Parameter(
116 + torch.empty(E, 2 * IN, H, dtype=torch.int8), requires_grad=False
117 + )
118 + set_weight_attrs(w13, extra_weight_attrs)
119 + layer.register_parameter("w13_weight", w13)
120 +
121 + w2 = torch.nn.Parameter(
122 + torch.empty(E, H, IN, dtype=torch.int8), requires_grad=False
123 + )
124 + set_weight_attrs(w2, extra_weight_attrs)
125 + layer.register_parameter("w2_weight", w2)
126 +
127 + # Register scales
128 + # KleidiAI groupwise kernels accepts float32 scales
129 + # KleidiAI groupwise kernels accepts bfloat16 scales
130 + scale_dtype = torch.float32 if g == -1 else torch.bfloat16
131 +
132 + w13_s = torch.nn.Parameter(
133 + torch.ones(E, 2 * IN, _n_scale_cols(H), dtype=scale_dtype),
134 + requires_grad=False,
135 + )
136 + set_weight_attrs(
137 + w13_s,
138 + {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
139 + )
140 + layer.register_parameter("w13_weight_scale", w13_s)
141 +
142 + w2_s = torch.nn.Parameter(
143 + torch.ones(E, H, _n_scale_cols(IN), dtype=scale_dtype), requires_grad=False
144 + )
145 + set_weight_attrs(
146 + w2_s,
147 + {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
148 + )
149 + layer.register_parameter("w2_weight_scale", w2_s)
150 +
151 + if self.has_bias:
152 + w13_bias = torch.nn.Parameter(
153 + torch.zeros(E, 2 * IN, dtype=params_dtype), requires_grad=False
154 + )
155 + layer.register_parameter("w13_bias", w13_bias)
156 + set_weight_attrs(w13_bias, extra_weight_attrs)
157 +
158 + w2_bias = torch.nn.Parameter(
159 + torch.zeros(num_experts, hidden_size, dtype=params_dtype),
160 + requires_grad=False,
161 + )
162 + layer.register_parameter("w2_bias", w2_bias)
163 + set_weight_attrs(w2_bias, extra_weight_attrs)
164 +
165 + # Placeholders for packed weights (will be replaced after packing)
166 + layer.register_parameter(
167 + "w13_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
168 + )
169 + set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs)
170 +
171 + layer.register_parameter(
172 + "w2_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
173 + )
174 + set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs)
175 +
176 + # dims for 4 bit fused matmuls
177 + layer.w13_in_features = H
178 + layer.w13_out_features = 2 * IN
179 + layer.w2_in_features = IN
180 + layer.w2_out_features = H
181 + layer.group_size = g
182 +
183 + # post-load packing to dyn-4bit KleidiAI kernel's format
184 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
185 + E = layer.w13_weight.shape[0]
186 + H = layer.w13_in_features
187 + I2 = layer.w13_out_features
188 + IN = layer.w2_in_features
189 + g = layer.group_size
190 +
191 + def _pack_matrix(
192 + int4_as_int8_2d: torch.Tensor,
193 + scales_2d: torch.Tensor,
194 + bias_1d: torch.Tensor | None,
195 + in_features: int,
196 + out_features: int,
197 + ) -> torch.Tensor:
198 + # int4 values are stored as int8 in [-8,7].
199 + # Shift to unsigned nibble and pack pairs along input-dim.
200 + tmp = int4_as_int8_2d.add(8) # [out, in]
201 + uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(
202 + torch.uint8
203 + ) # [out, in//2]
204 +
205 + # KleidiAI groupwise kernels accepts float32 scales
206 + # KleidiAI groupwise kernels accepts bfloat16 scales
207 + scale_dtype = torch.float32 if g == -1 else torch.bfloat16
208 + scales = scales_2d.to(scale_dtype)
209 + bias = None if bias_1d is None else bias_1d.to(torch.float32)
210 + return torch.ops.aten._dyn_quant_pack_4bit_weight(
211 + uint8_nibbles,
212 + scales,
213 + bias,
214 + g if g != -1 else in_features,
215 + in_features,
216 + out_features,
217 + )
218 +
219 + # Pack per expert
220 + w13_packed_list = []
221 + w2_packed_list = []
222 +
223 + has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None
224 + has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None
225 +
226 + for e in range(E):
227 + w13_packed_list.append(
228 + _pack_matrix(
229 + layer.w13_weight[e], # [2I, H]
230 + layer.w13_weight_scale[e], # [2I, H/g or 1]
231 + layer.w13_bias[e] if has_w13_bias else None, # [2I]
232 + H,
233 + I2,
234 + )
235 + )
236 + w2_packed_list.append(
237 + _pack_matrix(
238 + # w2 shape is [H, IN]; we need [out, in] == [H, IN].
239 + layer.w2_weight[e], # [H, IN]
240 + layer.w2_weight_scale[e], # [H, IN/g or 1]
241 + layer.w2_bias[e] if has_w2_bias else None, # [H]
242 + IN,
243 + layer.w2_out_features, # in_features=IN, out_features=H
244 + )
245 + )
246 +
247 + # each packed tensor has identical shape per expert; stack on dim 0
248 + w13_packed = torch.stack(w13_packed_list, dim=0)
249 + w2_packed = torch.stack(w2_packed_list, dim=0)
250 +
251 + replace_parameter(
252 + layer,
253 + "w13_weight_packed",
254 + torch.nn.Parameter(w13_packed, requires_grad=False),
255 + )
256 + replace_parameter(
257 + layer,
258 + "w2_weight_packed",
259 + torch.nn.Parameter(w2_packed, requires_grad=False),
260 + )
261 +
262 + # free raw tensors/scales/bias now that they're packed into the payload.
263 + replace_parameter(
264 + layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
265 + )
266 + replace_parameter(
267 + layer, "w2_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
268 + )
269 + replace_parameter(
270 + layer,
271 + "w13_weight_scale",
272 + torch.nn.Parameter(torch.empty(0), requires_grad=False),
273 + )
274 + replace_parameter(
275 + layer,
276 + "w2_weight_scale",
277 + torch.nn.Parameter(torch.empty(0), requires_grad=False),
278 + )
279 + if has_w13_bias:
280 + replace_parameter(
281 + layer,
282 + "w13_bias",
283 + torch.nn.Parameter(torch.empty(0), requires_grad=False),
284 + )
285 + if has_w2_bias:
286 + replace_parameter(
287 + layer,
288 + "w2_bias",
289 + torch.nn.Parameter(torch.empty(0), requires_grad=False),
290 + )
291 +
292 + def get_fused_moe_quant_config(
293 + self, layer: torch.nn.Module
294 + ) -> FusedMoEQuantConfig | None:
295 + # CPU dynamic 4-bit MoE path does not use modular kernels or
296 + # fused_experts; quant config is not needed.
297 + return None
298 +
299 + @property
300 + def is_monolithic(self) -> bool:
301 + return True
302 +
303 + def apply_monolithic(
304 + self,
305 + layer: FusedMoE,
306 + x: torch.Tensor,
307 + router_logits: torch.Tensor,
308 + ) -> torch.Tensor:
309 + assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."
310 + assert layer.activation in (
311 + MoEActivation.SILU,
312 + MoEActivation.SWIGLUOAI,
313 + MoEActivation.SWIGLUSTEP,
314 + ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."
315 + assert layer.expert_map is None, """expert_map/EP not implemented
316 + for CPU dyn-4bit MoE."""
317 +
318 + def _act_kind(s: MoEActivation) -> int:
319 + # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU
320 + if s == MoEActivation.SWIGLUSTEP:
321 + return 0
322 + if s == MoEActivation.SWIGLUOAI:
323 + return 1
324 + if s == MoEActivation.SILU:
325 + return 2
326 + raise ValueError(f"Unknown activation '{s}'")
327 +
328 + # Apply topk softmax on router output
329 + topk_weights, topk_ids = select_experts(
330 + hidden_states=x,
331 + router_logits=router_logits,
332 + top_k=layer.top_k,
333 + use_grouped_topk=layer.use_grouped_topk,
334 + renormalize=layer.renormalize,
335 + )
336 +
337 + return torch.ops._C.dynamic_4bit_int_moe(
338 + x,
339 + topk_ids.to(torch.long),
340 + topk_weights,
341 + layer.w13_weight_packed,
342 + layer.w2_weight_packed,
343 + layer.w2_out_features,
344 + layer.w2_in_features,
345 + layer.w13_out_features,
346 + layer.group_size,
347 + layer.apply_router_weight_on_input,
348 + int(_act_kind(layer.activation)),
349 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py new
+414
@@ -0,0 +1,414 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors.quantization import (
7 + QuantizationArgs,
8 + QuantizationStrategy,
9 +)
10 +
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 +from vllm.distributed import get_tensor_model_parallel_world_size
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.layers.fused_moe import (
15 + FusedMoE,
16 + FusedMoeWeightScaleSupported,
17 +)
18 +from vllm.model_executor.layers.fused_moe.config import (
19 + FusedMoEConfig,
20 + FusedMoEQuantConfig,
21 +)
22 +from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
23 + convert_to_fp8_moe_kernel_format,
24 + make_fp8_moe_kernel,
25 + make_fp8_moe_quant_config,
26 + select_fp8_moe_backend,
27 +)
28 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
29 + CompressedTensorsMoEMethod,
30 +)
31 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
32 + process_fp8_input_tensor_strategy_moe,
33 + process_fp8_weight_tensor_strategy_moe,
34 +)
35 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
36 + kFp8Dynamic128Sym,
37 + kFp8DynamicTokenSym,
38 + kFp8Static128BlockSym,
39 + kFp8StaticChannelSym,
40 + kFp8StaticTensorSym,
41 +)
42 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
43 + normalize_e4m3fn_to_e4m3fnuz,
44 +)
45 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
46 +from vllm.platforms import current_platform
47 +
48 +logger = init_logger(__name__)
49 +
50 +
51 +class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
52 + """W8A8 FP8 MoE quantization using compressed tensors."""
53 +
54 + def __init__(
55 + self,
56 + weight_quant: QuantizationArgs,
57 + input_quant: QuantizationArgs,
58 + moe: FusedMoEConfig,
59 + layer_name: str | None = None,
60 + ):
61 + super().__init__(moe)
62 + self.weight_quant = weight_quant
63 + self.input_quant = input_quant
64 +
65 + per_tensor = (
66 + self.weight_quant.strategy == QuantizationStrategy.TENSOR
67 + and self.input_quant.strategy == QuantizationStrategy.TENSOR
68 + )
69 + per_channel = (
70 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL
71 + and self.input_quant.strategy == QuantizationStrategy.TOKEN
72 + )
73 + if not (per_tensor or per_channel):
74 + assert self.weight_quant.strategy == QuantizationStrategy.BLOCK
75 + self.weight_block_size = self.weight_quant.block_structure
76 + assert self.weight_quant.dynamic is not None
77 + else:
78 + self.weight_block_size = None
79 + self.block_quant = self.weight_block_size is not None
80 +
81 + self.static_input_scales = not self.input_quant.dynamic
82 + if self.static_input_scales and per_channel:
83 + raise ValueError(
84 + "For FP8 Fused MoE layer, we require either per tensor or "
85 + "channelwise, dynamic per token quantization."
86 + )
87 +
88 + ct2vllm_weight = {
89 + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
90 + QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
91 + QuantizationStrategy.BLOCK: kFp8Static128BlockSym,
92 + }
93 + ct2vllm_act = {
94 + QuantizationStrategy.TOKEN: kFp8DynamicTokenSym,
95 + QuantizationStrategy.TENSOR: (
96 + kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym
97 + ),
98 + }
99 + weight_key = ct2vllm_weight[self.weight_quant.strategy]
100 + if weight_key == kFp8Static128BlockSym:
101 + activation_key = kFp8Dynamic128Sym
102 + else:
103 + activation_key = ct2vllm_act[self.input_quant.strategy]
104 +
105 + # Select Fp8 MoE backend
106 + self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
107 + config=self.moe,
108 + weight_key=weight_key,
109 + activation_key=activation_key,
110 + allow_vllm_cutlass=True,
111 + )
112 +
113 + def create_weights(
114 + self,
115 + layer: torch.nn.Module,
116 + num_experts: int,
117 + hidden_size: int,
118 + intermediate_size_per_partition: int,
119 + params_dtype: torch.dtype,
120 + **extra_weight_attrs,
121 + ):
122 + layer.num_experts = num_experts
123 + layer.orig_dtype = params_dtype
124 + layer.weight_block_size = None
125 +
126 + params_dtype = torch.float8_e4m3fn
127 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
128 +
129 + if self.block_quant:
130 + assert self.weight_block_size is not None
131 + layer.weight_block_size = self.weight_block_size
132 + tp_size = get_tensor_model_parallel_world_size()
133 + block_n, block_k = (
134 + self.weight_block_size[0],
135 + self.weight_block_size[1],
136 + )
137 + # NOTE: To ensure proper alignment of the block-wise quantization
138 + # scales, the output_size of the weights for both the gate and up
139 + # layers must be divisible by block_n.
140 + # Required by column parallel or enabling merged weights
141 + if intermediate_size_per_partition % block_n != 0:
142 + raise ValueError(
143 + f"The output_size of gate's and up's weight = "
144 + f"{intermediate_size_per_partition} is not divisible by "
145 + f"weight quantization block_n = {block_n}."
146 + )
147 + if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
148 + # Required by row parallel
149 + raise ValueError(
150 + f"The input_size of down's weight = "
151 + f"{intermediate_size_per_partition} is not divisible by "
152 + f"weight quantization block_k = {block_k}."
153 + )
154 +
155 + # WEIGHTS
156 + w13_weight = torch.nn.Parameter(
157 + torch.empty(
158 + num_experts,
159 + w13_num_shards * intermediate_size_per_partition,
160 + hidden_size,
161 + dtype=params_dtype,
162 + ),
163 + requires_grad=False,
164 + )
165 + layer.register_parameter("w13_weight", w13_weight)
166 + set_weight_attrs(w13_weight, extra_weight_attrs)
167 +
168 + w2_weight = torch.nn.Parameter(
169 + torch.empty(
170 + num_experts,
171 + hidden_size,
172 + intermediate_size_per_partition,
173 + dtype=params_dtype,
174 + ),
175 + requires_grad=False,
176 + )
177 + layer.register_parameter("w2_weight", w2_weight)
178 + set_weight_attrs(w2_weight, extra_weight_attrs)
179 +
180 + # WEIGHT_SCALES
181 + if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
182 + # For gated MoE, allocate 2 scales for w1 and w3 respectively.
183 + # They will be combined to a single scale after weight loading.
184 + # For non-gated MoE, allocate 1 scale for w13.
185 + w13_weight_scale = torch.nn.Parameter(
186 + torch.ones(num_experts, w13_num_shards, dtype=torch.float32),
187 + requires_grad=False,
188 + )
189 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
190 + w2_weight_scale = torch.nn.Parameter(
191 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False
192 + )
193 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
194 + # Add PER-TENSOR quantization for FusedMoE.weight_loader.
195 + extra_weight_attrs.update(
196 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
197 + )
198 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
199 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
200 +
201 + elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL:
202 + w13_weight_scale = torch.nn.Parameter(
203 + torch.ones(
204 + num_experts,
205 + w13_num_shards * intermediate_size_per_partition,
206 + 1,
207 + dtype=torch.float32,
208 + ),
209 + requires_grad=False,
210 + )
211 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
212 + w2_weight_scale = torch.nn.Parameter(
213 + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
214 + requires_grad=False,
215 + )
216 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
217 + # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
218 + extra_weight_attrs.update(
219 + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
220 + )
221 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
222 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
223 +
224 + elif self.weight_quant.strategy == QuantizationStrategy.BLOCK:
225 + w13_weight_scale = torch.nn.Parameter(
226 + torch.ones(
227 + num_experts,
228 + w13_num_shards
229 + * ((intermediate_size_per_partition + block_n - 1) // block_n),
230 + (hidden_size + block_k - 1) // block_k,
231 + dtype=torch.float32,
232 + ),
233 + requires_grad=False,
234 + )
235 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
236 + w2_weight_scale = torch.nn.Parameter(
237 + torch.ones(
238 + num_experts,
239 + (hidden_size + block_n - 1) // block_n,
240 + (intermediate_size_per_partition + block_k - 1) // block_k,
241 + dtype=torch.float32,
242 + ),
243 + requires_grad=False,
244 + )
245 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
246 + # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
247 + extra_weight_attrs.update(
248 + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
249 + )
250 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
251 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
252 +
253 + # INPUT_SCALES
254 + if self.static_input_scales:
255 + w13_input_scale = torch.nn.Parameter(
256 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False
257 + )
258 + layer.register_parameter("w13_input_scale", w13_input_scale)
259 + set_weight_attrs(w13_input_scale, extra_weight_attrs)
260 +
261 + w2_input_scale = torch.nn.Parameter(
262 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False
263 + )
264 + layer.register_parameter("w2_input_scale", w2_input_scale)
265 + set_weight_attrs(w2_input_scale, extra_weight_attrs)
266 + else:
267 + layer.w13_input_scale = None
268 + layer.w2_input_scale = None
269 +
270 + def process_weights_after_loading(self, layer: FusedMoE) -> None:
271 + # Allow for accessing weights and scales in standard way.
272 + w13 = layer.w13_weight
273 + w2 = layer.w2_weight
274 + w13_scale = layer.w13_weight_scale
275 + w2_scale = layer.w2_weight_scale
276 + w13_input_scale = layer.w13_input_scale
277 + w2_input_scale = layer.w2_input_scale
278 +
279 + # MI300x and MI325x use FNUZ format for FP8. Convert if needed.
280 + if current_platform.is_fp8_fnuz():
281 + w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
282 + w13, w13_scale, w13_input_scale
283 + )
284 + w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
285 + w2, w2_scale, w2_input_scale
286 + )
287 +
288 + # Per tensor kernels require single activation scale. Use the max.
289 + if self.static_input_scales:
290 + assert self.input_quant.strategy == QuantizationStrategy.TENSOR
291 + assert w13_input_scale is not None and w2_input_scale is not None
292 + w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
293 + w13_input_scale, w2_input_scale
294 + )
295 + replace_parameter(layer, "w13_input_scale", w13_input_scale)
296 + replace_parameter(layer, "w2_input_scale", w2_input_scale)
297 +
298 + # Per-tensor kernels use a single scale, for W13, but on disk there
299 + # is a separate scale for W1 and W3. Requantize with the max scale.
300 + if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
301 + w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
302 + w13,
303 + w13_scale,
304 + shard_size=layer.intermediate_size_per_partition,
305 + num_experts=layer.local_num_experts,
306 + is_act_and_mul=self.moe.is_act_and_mul,
307 + )
308 +
309 + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
310 + fp8_backend=self.fp8_backend,
311 + layer=layer,
312 + w13=w13,
313 + w2=w2,
314 + w13_scale=w13_scale,
315 + w2_scale=w2_scale,
316 + w13_input_scale=w13_input_scale,
317 + w2_input_scale=w2_input_scale,
318 + )
319 +
320 + # Replace parameters with updated versions. Note that this helper
321 + # function ensures the replacement is compatible with RL weight reloads.
322 + replace_parameter(layer, "w13_weight", w13)
323 + replace_parameter(layer, "w2_weight", w2)
324 + replace_parameter(layer, "w13_weight_scale", w13_scale)
325 + replace_parameter(layer, "w2_weight_scale", w2_scale)
326 +
327 + # Setup modular kernel for TP case and naive DP/EP case.
328 + # In non-naive DP/EP case, we will create a ModularKernelMethod.
329 + # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
330 + # in both cases.
331 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
332 + if self.moe_quant_config:
333 + assert self.experts_cls is not None
334 + self.moe_kernel = make_fp8_moe_kernel(
335 + moe_quant_config=self.moe_quant_config,
336 + moe_config=self.moe,
337 + fp8_backend=self.fp8_backend,
338 + experts_cls=self.experts_cls,
339 + routing_tables=layer._maybe_init_expert_routing_tables(),
340 + shared_experts=layer.shared_experts,
341 + )
342 +
343 + def maybe_make_prepare_finalize(
344 + self,
345 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
346 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
347 + raise ValueError(
348 + f"{self.__class__.__name__} uses the new modular kernel initialization "
349 + "logic. This function should not be called."
350 + )
351 +
352 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
353 + is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
354 + return make_fp8_moe_quant_config(
355 + fp8_backend=self.fp8_backend,
356 + w1_scale=layer.w13_weight_scale,
357 + w2_scale=layer.w2_weight_scale,
358 + a1_scale=layer.w13_input_scale,
359 + a2_scale=layer.w2_input_scale,
360 + per_act_token_quant=is_per_token,
361 + per_out_ch_quant=is_per_token,
362 + block_shape=self.weight_block_size,
363 + )
364 +
365 + def apply_monolithic(
366 + self,
367 + layer: FusedMoE,
368 + x: torch.Tensor,
369 + router_logits: torch.Tensor,
370 + ) -> torch.Tensor:
371 + assert self.moe_kernel is not None
372 + return self.moe_kernel.apply_monolithic(
373 + x,
374 + layer.w13_weight,
375 + layer.w2_weight,
376 + router_logits,
377 + activation=layer.activation,
378 + global_num_experts=layer.global_num_experts,
379 + expert_map=layer.expert_map,
380 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
381 + num_expert_group=layer.num_expert_group,
382 + topk_group=layer.topk_group,
383 + e_score_correction_bias=layer.e_score_correction_bias,
384 + routed_scaling_factor=layer.routed_scaling_factor,
385 + )
386 +
387 + def apply(
388 + self,
389 + layer: FusedMoE,
390 + x: torch.Tensor,
391 + topk_weights: torch.Tensor,
392 + topk_ids: torch.Tensor,
393 + shared_experts_input: torch.Tensor | None,
394 + ) -> torch.Tensor:
395 + assert not self.is_monolithic
396 + assert self.moe_kernel is not None
397 + return self.moe_kernel.apply(
398 + x,
399 + layer.w13_weight,
400 + layer.w2_weight,
401 + topk_weights,
402 + topk_ids,
403 + activation=layer.activation,
404 + global_num_experts=layer.global_num_experts,
405 + # TODO(rob): investigate the disable_expert_map introduced by:
406 + # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
407 + expert_map=layer.expert_map,
408 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
409 + shared_experts_input=shared_experts_input,
410 + )
411 +
412 + @property
413 + def supports_eplb(self) -> bool:
414 + return True
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py new
+161
@@ -0,0 +1,161 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors.quantization import (
7 + QuantizationArgs,
8 + QuantizationStrategy,
9 +)
10 +
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe import (
13 + FusedMoE,
14 + FusedMoeWeightScaleSupported,
15 +)
16 +from vllm.model_executor.layers.fused_moe.config import (
17 + FusedMoEConfig,
18 + FusedMoEQuantConfig,
19 + int8_w8a8_moe_quant_config,
20 +)
21 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
22 + CompressedTensorsMoEMethod,
23 +)
24 +from vllm.model_executor.utils import set_weight_attrs
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
30 + def __init__(
31 + self,
32 + weight_quant: QuantizationArgs,
33 + input_quant: QuantizationArgs,
34 + moe: FusedMoEConfig,
35 + layer_name: str | None = None,
36 + ):
37 + super().__init__(moe)
38 + self.weight_quant = weight_quant
39 + self.input_quant = input_quant
40 +
41 + per_channel = (
42 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL
43 + and self.input_quant.strategy == QuantizationStrategy.TOKEN
44 + )
45 + if not per_channel:
46 + raise ValueError(
47 + "For INT8 Fused MoE layers, we require channelwise, "
48 + "dynamic per token quantization. Found "
49 + f"{self.weight_quant}, {self.input_quant}"
50 + )
51 +
52 + self.static_input_scales = not self.input_quant.dynamic
53 + if self.static_input_scales:
54 + raise ValueError(
55 + "For INT8 Fused MoE layers, we require channelwise, "
56 + "dynamic per token quantization. Found static input scales."
57 + )
58 +
59 + def create_weights(
60 + self,
61 + layer: torch.nn.Module,
62 + num_experts: int,
63 + hidden_size: int,
64 + intermediate_size_per_partition: int,
65 + params_dtype: torch.dtype,
66 + **extra_weight_attrs,
67 + ):
68 + params_dtype = torch.int8
69 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
70 +
71 + # WEIGHTS
72 + w13_weight = torch.nn.Parameter(
73 + torch.empty(
74 + num_experts,
75 + w13_num_shards * intermediate_size_per_partition,
76 + hidden_size,
77 + dtype=params_dtype,
78 + ),
79 + requires_grad=False,
80 + )
81 + layer.register_parameter("w13_weight", w13_weight)
82 + set_weight_attrs(w13_weight, extra_weight_attrs)
83 +
84 + w2_weight = torch.nn.Parameter(
85 + torch.empty(
86 + num_experts,
87 + hidden_size,
88 + intermediate_size_per_partition,
89 + dtype=params_dtype,
90 + ),
91 + requires_grad=False,
92 + )
93 + layer.register_parameter("w2_weight", w2_weight)
94 + set_weight_attrs(w2_weight, extra_weight_attrs)
95 +
96 + # WEIGHT_SCALES
97 + assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL
98 + w13_weight_scale = torch.nn.Parameter(
99 + torch.ones(
100 + num_experts,
101 + w13_num_shards * intermediate_size_per_partition,
102 + 1,
103 + dtype=torch.float32,
104 + ),
105 + requires_grad=False,
106 + )
107 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
108 + w2_weight_scale = torch.nn.Parameter(
109 + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
110 + requires_grad=False,
111 + )
112 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
113 + # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
114 + extra_weight_attrs.update(
115 + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
116 + )
117 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
118 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
119 +
120 + # INPUT_SCALES
121 + assert not self.static_input_scales
122 + layer.w13_input_scale = None
123 + layer.w2_input_scale = None
124 +
125 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
126 + pass
127 +
128 + def get_fused_moe_quant_config(
129 + self, layer: torch.nn.Module
130 + ) -> FusedMoEQuantConfig | None:
131 + return int8_w8a8_moe_quant_config(
132 + w1_scale=layer.w13_weight_scale,
133 + w2_scale=layer.w2_weight_scale,
134 + a1_scale=layer.w13_input_scale,
135 + a2_scale=layer.w2_input_scale,
136 + per_act_token_quant=True,
137 + )
138 +
139 + def apply(
140 + self,
141 + layer: FusedMoE,
142 + x: torch.Tensor,
143 + topk_weights: torch.Tensor,
144 + topk_ids: torch.Tensor,
145 + shared_experts_input: torch.Tensor | None,
146 + ) -> torch.Tensor:
147 + from vllm.model_executor.layers.fused_moe import fused_experts
148 +
149 + return fused_experts(
150 + hidden_states=x,
151 + w1=layer.w13_weight,
152 + w2=layer.w2_weight,
153 + topk_weights=topk_weights,
154 + topk_ids=topk_ids,
155 + inplace=not self.moe.disable_inplace,
156 + activation=layer.activation,
157 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
158 + global_num_experts=layer.global_num_experts,
159 + expert_map=layer.expert_map,
160 + quant_config=self.moe_quant_config,
161 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py new
+267
@@ -0,0 +1,267 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from compressed_tensors.quantization import (
7 + QuantizationArgs,
8 +)
9 +
10 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe import (
13 + FusedMoE,
14 +)
15 +from vllm.model_executor.layers.fused_moe.config import (
16 + FusedMoEConfig,
17 + FusedMoEQuantConfig,
18 + int4_w4a16_moe_quant_config,
19 + int8_w8a16_moe_quant_config,
20 +)
21 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
22 + CompressedTensorsMoEMethod,
23 +)
24 +from vllm.model_executor.utils import set_weight_attrs
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
30 + def __init__(
31 + self,
32 + weight_quant: QuantizationArgs,
33 + input_quant: QuantizationArgs | None,
34 + moe: FusedMoEConfig,
35 + layer_name: str | None = None,
36 + ):
37 + super().__init__(moe)
38 + self.weight_quant = weight_quant
39 + self.input_quant = input_quant
40 + # Extract properties from weight_quant
41 + self.num_bits = weight_quant.num_bits
42 + self.packed_factor = 32 // weight_quant.num_bits
43 + self.strategy = weight_quant.strategy
44 + # channelwise is not supported by this kernel
45 + assert weight_quant.strategy == "group"
46 + self.group_size = weight_quant.group_size
47 + # grouped actorder isn't supported by this kernel
48 + assert weight_quant.actorder != "group"
49 + assert weight_quant.symmetric, (
50 + "Only symmetric quantization is supported for MoE"
51 + )
52 +
53 + def create_weights(
54 + self,
55 + layer: torch.nn.Module,
56 + num_experts: int,
57 + hidden_size: int,
58 + intermediate_size_per_partition: int,
59 + params_dtype: torch.dtype,
60 + **extra_weight_attrs,
61 + ):
62 + # Will transpose the loaded weight along the
63 + # intermediate and hidden dim sizes. Will
64 + # shard for TP along the transposed dims
65 + extra_weight_attrs.update(
66 + {"is_transposed": True, "quant_method": self.strategy}
67 + )
68 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
69 + w13_weight = torch.nn.Parameter(
70 + torch.empty(
71 + num_experts,
72 + hidden_size // self.packed_factor,
73 + w13_num_shards * intermediate_size_per_partition,
74 + dtype=torch.int32,
75 + ),
76 + requires_grad=False,
77 + )
78 + layer.register_parameter("w13_weight_packed", w13_weight)
79 + set_weight_attrs(w13_weight, extra_weight_attrs)
80 +
81 + w2_weight = torch.nn.Parameter(
82 + torch.empty(
83 + num_experts,
84 + intermediate_size_per_partition // self.packed_factor,
85 + hidden_size,
86 + dtype=torch.int32,
87 + ),
88 + requires_grad=False,
89 + )
90 + layer.register_parameter("w2_weight_packed", w2_weight)
91 + set_weight_attrs(w2_weight, extra_weight_attrs)
92 +
93 + w2_scales_size = intermediate_size_per_partition
94 +
95 + if self.strategy == "channel":
96 + num_groups_w2 = num_groups_w13 = 1
97 + self.group_size = -1
98 + else:
99 + num_groups_w2 = w2_scales_size // self.group_size
100 + num_groups_w13 = hidden_size // self.group_size
101 +
102 + w13_scale = torch.nn.Parameter(
103 + torch.ones(
104 + num_experts,
105 + num_groups_w13,
106 + w13_num_shards * intermediate_size_per_partition,
107 + dtype=params_dtype,
108 + ),
109 + requires_grad=False,
110 + )
111 + layer.register_parameter("w13_weight_scale", w13_scale)
112 + set_weight_attrs(w13_scale, extra_weight_attrs)
113 +
114 + w2_scale = torch.nn.Parameter(
115 + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
116 + requires_grad=False,
117 + )
118 + layer.register_parameter("w2_weight_scale", w2_scale)
119 + set_weight_attrs(w2_scale, extra_weight_attrs)
120 + set_weight_attrs(w2_scale, {"load_full_w2": False})
121 +
122 + w2_weight_shape = torch.nn.Parameter(
123 + torch.empty(num_experts, 2), requires_grad=False
124 + )
125 + layer.register_parameter("w2_weight_shape", w2_weight_shape)
126 + set_weight_attrs(w2_weight_shape, extra_weight_attrs)
127 + w13_weight_shape = torch.nn.Parameter(
128 + torch.empty(num_experts, 2), requires_grad=False
129 + )
130 +
131 + layer.register_parameter("w13_weight_shape", w13_weight_shape)
132 + set_weight_attrs(w13_weight_shape, extra_weight_attrs)
133 +
134 + w13_g_idx = torch.nn.Parameter(
135 + torch.empty(
136 + num_experts,
137 + hidden_size,
138 + dtype=torch.int32,
139 + ),
140 + requires_grad=False,
141 + )
142 + layer.register_parameter("w13_weight_g_idx", w13_g_idx)
143 + set_weight_attrs(w13_g_idx, extra_weight_attrs)
144 +
145 + w2_g_idx = torch.nn.Parameter(
146 + torch.empty(
147 + num_experts,
148 + intermediate_size_per_partition,
149 + dtype=torch.int32,
150 + ),
151 + requires_grad=False,
152 + )
153 + layer.register_parameter("w2_weight_g_idx", w2_g_idx)
154 + set_weight_attrs(w2_g_idx, extra_weight_attrs)
155 +
156 + w13_g_idx_sort_indices = torch.nn.Parameter(
157 + torch.empty(
158 + num_experts,
159 + hidden_size,
160 + dtype=torch.int32,
161 + ),
162 + requires_grad=False,
163 + )
164 + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
165 + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
166 +
167 + w2_g_idx_sort_indices = torch.nn.Parameter(
168 + torch.empty(
169 + num_experts,
170 + intermediate_size_per_partition,
171 + dtype=torch.int32,
172 + ),
173 + requires_grad=False,
174 + )
175 + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
176 + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
177 +
178 + layer.a13_scale = None
179 + layer.a2_scale = None
180 +
181 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
182 + # Reconfigure packed weights and scales to match moe_wna16 format
183 + layer.w13_weight_packed = torch.nn.Parameter(
184 + layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
185 + requires_grad=False,
186 + )
187 + layer.w2_weight_packed = torch.nn.Parameter(
188 + layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
189 + requires_grad=False,
190 + )
191 + layer.w13_weight_scale = torch.nn.Parameter(
192 + layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
193 + )
194 + layer.w2_weight_scale = torch.nn.Parameter(
195 + layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
196 + )
197 +
198 + def get_fused_moe_quant_config(
199 + self, layer: torch.nn.Module
200 + ) -> FusedMoEQuantConfig | None:
201 + assert self.num_bits == 4 or self.num_bits == 8
202 + config_builder = (
203 + int4_w4a16_moe_quant_config
204 + if self.num_bits == 4
205 + else int8_w8a16_moe_quant_config
206 + )
207 +
208 + return config_builder(
209 + w1_scale=layer.w13_weight_scale,
210 + w2_scale=layer.w2_weight_scale,
211 + w1_zp=None,
212 + w2_zp=None,
213 + block_shape=[0, self.group_size],
214 + )
215 +
216 + def select_gemm_impl(
217 + self,
218 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
219 + layer: torch.nn.Module,
220 + ) -> mk.FusedMoEExpertsModular:
221 + if self.moe.is_lora_enabled:
222 + assert self.moe_quant_config is not None
223 + from vllm.triton_utils import HAS_TRITON
224 +
225 + if HAS_TRITON:
226 + from vllm.model_executor.layers.fused_moe import TritonWNA16Experts
227 +
228 + layer.w13_weight = layer.w13_weight_packed
229 + layer.w2_weight = layer.w2_weight_packed
230 + return TritonWNA16Experts(
231 + moe_config=self.moe, quant_config=self.moe_quant_config
232 + )
233 + else:
234 + raise NotImplementedError(
235 + "TritonExperts requires Triton. "
236 + "Install triton or disable LoRA for MoE."
237 + )
238 +
239 + raise NotImplementedError
240 +
241 + def apply(
242 + self,
243 + layer: FusedMoE,
244 + x: torch.Tensor,
245 + topk_weights: torch.Tensor,
246 + topk_ids: torch.Tensor,
247 + shared_experts_input: torch.Tensor | None,
248 + ) -> torch.Tensor:
249 + from vllm.model_executor.layers.fused_moe import fused_experts
250 +
251 + return fused_experts(
252 + x,
253 + layer.w13_weight_packed,
254 + layer.w2_weight_packed,
255 + topk_weights=topk_weights,
256 + topk_ids=topk_ids,
257 + inplace=not self.moe.disable_inplace,
258 + activation=layer.activation,
259 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
260 + global_num_experts=layer.global_num_experts,
261 + expert_map=layer.expert_map,
262 + quant_config=self.moe_quant_config,
263 + )
264 +
265 + @property
266 + def supports_eplb(self) -> bool:
267 + return True
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py new
+575
@@ -0,0 +1,575 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import enum
5 +from enum import Enum
6 +
7 +import torch
8 +from compressed_tensors.quantization import (
9 + QuantizationArgs,
10 +)
11 +
12 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 +from vllm import _custom_ops as ops
14 +from vllm.logger import init_logger
15 +from vllm.model_executor.layers.fused_moe import (
16 + FusedMoE,
17 +)
18 +from vllm.model_executor.layers.fused_moe.config import (
19 + FusedMoEConfig,
20 + FusedMoEQuantConfig,
21 + int4_w4a16_moe_quant_config,
22 +)
23 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
24 + BatchedMarlinExperts,
25 + MarlinExperts,
26 + fused_marlin_moe,
27 +)
28 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
29 + CompressedTensorsMoEMethod,
30 +)
31 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
32 + WNA16_SUPPORTED_TYPES_MAP,
33 +)
34 +from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
35 + flashinfer_trtllm_mxint4_moe,
36 + is_flashinfer_mxint4_moe_available,
37 + prepare_static_weights_for_trtllm_mxint4_moe,
38 +)
39 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
40 + get_marlin_input_dtype,
41 + marlin_act_int8_process_scales,
42 + marlin_make_workspace_new,
43 + marlin_moe_permute_scales,
44 +)
45 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
46 +
47 +logger = init_logger(__name__)
48 +
49 +
50 +class GPTQMarlinState(Enum):
51 + REPACK = enum.auto()
52 + READY = enum.auto()
53 +
54 +
55 +class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
56 + def __init__(
57 + self,
58 + weight_quant: QuantizationArgs,
59 + input_quant: QuantizationArgs | None,
60 + moe: FusedMoEConfig,
61 + layer_name: str | None = None,
62 + ):
63 + super().__init__(moe)
64 + self.weight_quant = weight_quant
65 + self.input_quant = input_quant
66 + assert weight_quant.symmetric, (
67 + "Only symmetric quantization is supported for MoE"
68 + )
69 + # Extract properties from weight_quant
70 + self.num_bits = weight_quant.num_bits
71 + self.packed_factor = 32 // weight_quant.num_bits
72 + self.strategy = weight_quant.strategy
73 + self.group_size = weight_quant.group_size
74 + self.actorder = weight_quant.actorder
75 +
76 + self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
77 +
78 + self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
79 + self.use_flashinfer_mxint4_moe = (
80 + is_flashinfer_mxint4_moe_available()
81 + and self.group_size == 32
82 + and weight_quant.num_bits == 4
83 + )
84 + self.kernel_backend = (
85 + "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin"
86 + )
87 + logger.info_once(
88 + f"Using {self.kernel_backend} backend for WNA16 MoE "
89 + f"(group_size={self.group_size}, num_bits={self.num_bits})",
90 + scope="local",
91 + )
92 +
93 + def get_weight_shape(
94 + self,
95 + weight_name: str,
96 + num_experts: int,
97 + hidden_size: int,
98 + intermediate_size_per_partition: int,
99 + num_groups_w2: int | None = None,
100 + num_groups_w13: int | None = None,
101 + ) -> tuple[int, int, int]:
102 + """
103 + Get the shape of the weight based on the weight name, number of experts
104 + hidden size, intermediate size per partition, number of groups for w2,
105 + and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
106 + for weight scales.
107 + """
108 + if weight_name == "w13_scale":
109 + assert num_groups_w13 is not None, (
110 + "num_groups_w13 must be provided for weight scales"
111 + )
112 + if weight_name == "w2_scale":
113 + assert num_groups_w2 is not None, (
114 + "num_groups_w2 must be provided for weight scales"
115 + )
116 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
117 + shape_map = {
118 + "w13_weight": {
119 + "Flashinfer": (
120 + num_experts,
121 + w13_num_shards * intermediate_size_per_partition,
122 + hidden_size // self.packed_factor,
123 + ),
124 + "Marlin": (
125 + num_experts,
126 + hidden_size // self.packed_factor,
127 + w13_num_shards * intermediate_size_per_partition,
128 + ),
129 + },
130 + "w13_scale": {
131 + "Flashinfer": (
132 + num_experts,
133 + w13_num_shards * intermediate_size_per_partition,
134 + num_groups_w13,
135 + ),
136 + "Marlin": (
137 + num_experts,
138 + num_groups_w13,
139 + w13_num_shards * intermediate_size_per_partition,
140 + ),
141 + },
142 + "w2_weight": {
143 + "Flashinfer": (
144 + num_experts,
145 + hidden_size,
146 + intermediate_size_per_partition // self.packed_factor,
147 + ),
148 + "Marlin": (
149 + num_experts,
150 + intermediate_size_per_partition // self.packed_factor,
151 + hidden_size,
152 + ),
153 + },
154 + "w2_scale": {
155 + "Flashinfer": (num_experts, hidden_size, num_groups_w2),
156 + "Marlin": (num_experts, num_groups_w2, hidden_size),
157 + },
158 + }
159 + return shape_map[weight_name][self.kernel_backend]
160 +
161 + def create_weights(
162 + self,
163 + layer: torch.nn.Module,
164 + num_experts: int,
165 + hidden_size: int,
166 + intermediate_size_per_partition: int,
167 + params_dtype: torch.dtype,
168 + **extra_weight_attrs,
169 + ):
170 + intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
171 +
172 + # Will transpose the loaded weight along the
173 + # intermediate and hidden dim sizes. Will
174 + # shard for TP along the transposed dims
175 + is_transposed = self.kernel_backend != "Flashinfer"
176 + extra_weight_attrs.update(
177 + {"is_transposed": is_transposed, "quant_method": self.strategy}
178 + )
179 +
180 + w13_weight = torch.nn.Parameter(
181 + torch.empty(
182 + *self.get_weight_shape(
183 + "w13_weight",
184 + num_experts,
185 + hidden_size,
186 + intermediate_size_per_partition,
187 + ),
188 + dtype=torch.int32,
189 + ),
190 + requires_grad=False,
191 + )
192 + layer.register_parameter("w13_weight_packed", w13_weight)
193 + set_weight_attrs(w13_weight, extra_weight_attrs)
194 +
195 + w2_weight = torch.nn.Parameter(
196 + torch.empty(
197 + *self.get_weight_shape(
198 + "w2_weight",
199 + num_experts,
200 + hidden_size,
201 + intermediate_size_per_partition,
202 + ),
203 + dtype=torch.int32,
204 + ),
205 + requires_grad=False,
206 + )
207 + layer.register_parameter("w2_weight_packed", w2_weight)
208 + set_weight_attrs(w2_weight, extra_weight_attrs)
209 +
210 + # In the case where we have actorder/g_idx,
211 + # we do not partition the w2 scales
212 + load_full_w2 = self.actorder and self.group_size != -1
213 + w2_scales_size = (
214 + intermediate_size_full if load_full_w2 else intermediate_size_per_partition
215 + )
216 +
217 + self.is_k_full = (not self.actorder) or (
218 + intermediate_size_per_partition == intermediate_size_full
219 + )
220 +
221 + if self.strategy == "channel":
222 + num_groups_w2 = num_groups_w13 = 1
223 + self.group_size = -1
224 + else:
225 + num_groups_w2 = w2_scales_size // self.group_size
226 + num_groups_w13 = hidden_size // self.group_size
227 +
228 + layer.num_groups_w13 = num_groups_w13
229 + layer.num_groups_w2 = num_groups_w2
230 +
231 + w13_scale = torch.nn.Parameter(
232 + torch.ones(
233 + *self.get_weight_shape(
234 + "w13_scale",
235 + num_experts,
236 + hidden_size,
237 + intermediate_size_per_partition,
238 + num_groups_w13=num_groups_w13,
239 + ),
240 + dtype=params_dtype,
241 + ),
242 + requires_grad=False,
243 + )
244 + layer.register_parameter("w13_weight_scale", w13_scale)
245 + set_weight_attrs(w13_scale, extra_weight_attrs)
246 +
247 + w2_scale = torch.nn.Parameter(
248 + torch.ones(
249 + *self.get_weight_shape(
250 + "w2_scale",
251 + num_experts,
252 + hidden_size,
253 + intermediate_size_per_partition,
254 + num_groups_w2=num_groups_w2,
255 + ),
256 + dtype=params_dtype,
257 + ),
258 + requires_grad=False,
259 + )
260 + layer.register_parameter("w2_weight_scale", w2_scale)
261 + set_weight_attrs(w2_scale, extra_weight_attrs)
262 + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
263 +
264 + w2_weight_shape = torch.nn.Parameter(
265 + torch.empty(num_experts, 2), requires_grad=False
266 + )
267 + layer.register_parameter("w2_weight_shape", w2_weight_shape)
268 + set_weight_attrs(w2_weight_shape, extra_weight_attrs)
269 + w13_weight_shape = torch.nn.Parameter(
270 + torch.empty(num_experts, 2), requires_grad=False
271 + )
272 +
273 + layer.register_parameter("w13_weight_shape", w13_weight_shape)
274 + set_weight_attrs(w13_weight_shape, extra_weight_attrs)
275 +
276 + w13_g_idx = torch.nn.Parameter(
277 + torch.empty(
278 + num_experts,
279 + hidden_size,
280 + dtype=torch.int32,
281 + ),
282 + requires_grad=False,
283 + )
284 + layer.register_parameter("w13_weight_g_idx", w13_g_idx)
285 + set_weight_attrs(w13_g_idx, extra_weight_attrs)
286 +
287 + w2_g_idx = torch.nn.Parameter(
288 + torch.empty(
289 + num_experts,
290 + intermediate_size_per_partition,
291 + dtype=torch.int32,
292 + ),
293 + requires_grad=False,
294 + )
295 + layer.register_parameter("w2_weight_g_idx", w2_g_idx)
296 + set_weight_attrs(w2_g_idx, extra_weight_attrs)
297 +
298 + w13_g_idx_sort_indices = torch.nn.Parameter(
299 + torch.empty(
300 + num_experts,
301 + hidden_size,
302 + dtype=torch.int32,
303 + ),
304 + requires_grad=False,
305 + )
306 + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
307 + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
308 +
309 + w2_g_idx_sort_indices = torch.nn.Parameter(
310 + torch.empty(
311 + num_experts,
312 + intermediate_size_per_partition,
313 + dtype=torch.int32,
314 + ),
315 + requires_grad=False,
316 + )
317 + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
318 + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
319 +
320 + layer.a13_scale = None
321 + layer.a2_scale = None
322 + layer.marlin_state = GPTQMarlinState.REPACK
323 +
324 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
325 + num_experts = layer.w13_weight_g_idx.shape[0]
326 + device = layer.w13_weight_g_idx.device
327 + if self.kernel_backend == "Flashinfer":
328 + dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
329 + layer.w13_weight_packed,
330 + layer.w13_weight_scale,
331 + layer.w2_weight_packed,
332 + layer.w2_weight_scale,
333 + )
334 + replace_parameter(
335 + layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"]
336 + )
337 + replace_parameter(
338 + layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"]
339 + )
340 + replace_parameter(
341 + layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"]
342 + )
343 + replace_parameter(
344 + layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"]
345 + )
346 + return None
347 +
348 + is_a_8bit = (
349 + self.marlin_input_dtype is not None
350 + and self.marlin_input_dtype.itemsize == 1
351 + )
352 +
353 + if self.marlin_input_dtype == torch.float8_e4m3fn:
354 + # NOTE: for non-zp quantization format only
355 + ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True)
356 + ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True)
357 + layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512
358 + layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512
359 +
360 + # when running models with grouped act order,
361 + # resort to g_idx values provided in checkpoint
362 + if self.actorder == "group":
363 + w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx)
364 + w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx)
365 + w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx)
366 + w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx)
367 +
368 + for e in range(num_experts):
369 + w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to(
370 + torch.int32
371 + )
372 + w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to(
373 + torch.int32
374 + )
375 + w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][
376 + w13_g_idx_sort_indices[e]
377 + ]
378 + w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]]
379 +
380 + replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx)
381 + replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx)
382 + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
383 + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
384 +
385 + else:
386 + layer.w13_weight_g_idx = torch.nn.Parameter(
387 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),
388 + requires_grad=False,
389 + )
390 + layer.w2_weight_g_idx = torch.nn.Parameter(
391 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),
392 + requires_grad=False,
393 + )
394 + layer.w13_g_idx_sort_indices = torch.nn.Parameter(
395 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),
396 + requires_grad=False,
397 + )
398 + layer.w2_g_idx_sort_indices = torch.nn.Parameter(
399 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),
400 + requires_grad=False,
401 + )
402 +
403 + marlin_w13_qweight = ops.gptq_marlin_moe_repack(
404 + layer.w13_weight_packed,
405 + layer.w13_g_idx_sort_indices,
406 + layer.w13_weight_packed.shape[1] * self.packed_factor,
407 + layer.w13_weight_packed.shape[2],
408 + self.num_bits,
409 + is_a_8bit=is_a_8bit,
410 + )
411 + replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight)
412 +
413 + marlin_w2_qweight = ops.gptq_marlin_moe_repack(
414 + layer.w2_weight_packed,
415 + layer.w2_g_idx_sort_indices,
416 + layer.w2_weight_packed.shape[1] * self.packed_factor,
417 + layer.w2_weight_packed.shape[2],
418 + self.num_bits,
419 + is_a_8bit=is_a_8bit,
420 + )
421 + replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight)
422 +
423 + # Repack scales
424 + marlin_w13_scales = marlin_moe_permute_scales(
425 + s=layer.w13_weight_scale,
426 + size_k=layer.w13_weight_packed.shape[2],
427 + size_n=layer.w13_weight_scale.shape[2],
428 + group_size=self.group_size,
429 + is_a_8bit=is_a_8bit,
430 + )
431 + if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1:
432 + marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
433 + marlin_w13_scales
434 + )
435 + layer.register_parameter(
436 + "w13_input_global_scale",
437 + torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
438 + )
439 + replace_parameter(layer, "w13_weight_scale", marlin_w13_scales)
440 +
441 + marlin_w2_scales = marlin_moe_permute_scales(
442 + s=layer.w2_weight_scale,
443 + size_k=layer.w2_weight_scale.shape[1]
444 + * (self.group_size if self.group_size != -1 else self.packed_factor),
445 + size_n=layer.w2_weight_scale.shape[2],
446 + group_size=self.group_size,
447 + is_a_8bit=is_a_8bit,
448 + )
449 + if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1:
450 + marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
451 + marlin_w2_scales
452 + )
453 + layer.register_parameter(
454 + "w2_input_global_scale",
455 + torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
456 + )
457 + replace_parameter(layer, "w2_weight_scale", marlin_w2_scales)
458 +
459 + layer.workspace = marlin_make_workspace_new(device, 4)
460 +
461 + def get_fused_moe_quant_config(
462 + self, layer: torch.nn.Module
463 + ) -> FusedMoEQuantConfig | None:
464 + if self.num_bits != 4:
465 + return None
466 + return int4_w4a16_moe_quant_config(
467 + w1_scale=layer.w13_weight_scale,
468 + w2_scale=layer.w2_weight_scale,
469 + w1_zp=None,
470 + w2_zp=None,
471 + block_shape=[0, self.group_size],
472 + )
473 +
474 + def select_gemm_impl(
475 + self,
476 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
477 + layer: torch.nn.Module,
478 + ) -> mk.FusedMoEExpertsModular:
479 + assert self.num_bits == 4, "only supporting w4"
480 + layer.w13_weight = layer.w13_weight_packed
481 + layer.w2_weight = layer.w2_weight_packed
482 + assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])
483 + assert self.moe_quant_config is not None
484 + if (
485 + prepare_finalize.activation_format
486 + == mk.FusedMoEActivationFormat.BatchedExperts
487 + ):
488 + max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
489 + assert max_num_tokens_per_rank is not None
490 + return BatchedMarlinExperts(
491 + max_num_tokens=max_num_tokens_per_rank,
492 + num_dispatchers=prepare_finalize.num_dispatchers(),
493 + moe_config=self.moe,
494 + quant_config=self.moe_quant_config,
495 + w13_g_idx=layer.w13_weight_g_idx,
496 + w2_g_idx=layer.w2_weight_g_idx,
497 + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
498 + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
499 + is_k_full=self.is_k_full,
500 + )
501 + else:
502 + return MarlinExperts(
503 + moe_config=self.moe,
504 + quant_config=self.moe_quant_config,
505 + w13_g_idx=layer.w13_weight_g_idx,
506 + w2_g_idx=layer.w2_weight_g_idx,
507 + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
508 + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
509 + is_k_full=self.is_k_full,
510 + )
511 +
512 + @property
513 + def is_monolithic(self) -> bool:
514 + return self.kernel_backend == "Flashinfer"
515 +
516 + def apply_monolithic(
517 + self,
518 + layer: FusedMoE,
519 + x: torch.Tensor,
520 + router_logits: torch.Tensor,
521 + ) -> torch.Tensor:
522 + assert self.kernel_backend == "Flashinfer"
523 + return flashinfer_trtllm_mxint4_moe(
524 + x=x,
525 + router_logits=router_logits,
526 + w13_weight_packed=layer.w13_weight_packed,
527 + w13_weight_scale=layer.w13_weight_scale,
528 + w2_weight_packed=layer.w2_weight_packed,
529 + w2_weight_scale=layer.w2_weight_scale,
530 + global_num_experts=layer.global_num_experts,
531 + top_k=layer.top_k,
532 + intermediate_size_per_partition=layer.intermediate_size_per_partition,
533 + local_num_experts=layer.local_num_experts,
534 + ep_rank=layer.ep_rank,
535 + num_expert_group=layer.num_expert_group,
536 + topk_group=layer.topk_group,
537 + e_score_correction_bias=layer.e_score_correction_bias,
538 + routing_method_type=layer.routing_method_type,
539 + )
540 +
541 + def apply(
542 + self,
543 + layer: FusedMoE,
544 + x: torch.Tensor,
545 + topk_weights: torch.Tensor,
546 + topk_ids: torch.Tensor,
547 + shared_experts_input: torch.Tensor | None,
548 + ) -> torch.Tensor:
549 + assert self.kernel_backend == "Marlin"
550 + return fused_marlin_moe(
551 + x,
552 + layer.w13_weight_packed,
553 + layer.w2_weight_packed,
554 + None,
555 + None,
556 + layer.w13_weight_scale,
557 + layer.w2_weight_scale,
558 + topk_weights,
559 + topk_ids,
560 + input_global_scale1=getattr(layer, "w13_input_global_scale", None),
561 + input_global_scale2=getattr(layer, "w2_input_global_scale", None),
562 + quant_type_id=self.quant_type.id,
563 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
564 + global_num_experts=layer.global_num_experts,
565 + activation=layer.activation,
566 + expert_map=layer.expert_map,
567 + g_idx1=layer.w13_weight_g_idx,
568 + g_idx2=layer.w2_weight_g_idx,
569 + sort_indices1=layer.w13_g_idx_sort_indices,
570 + sort_indices2=layer.w2_g_idx_sort_indices,
571 + workspace=layer.workspace,
572 + input_dtype=self.marlin_input_dtype,
573 + is_k_full=self.is_k_full,
574 + inplace=not self.moe.disable_inplace,
575 + )