@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 454ea5b526

[MoE Refactor] Combine CompressedTensorsWNA16MarlinMoEMethod with CompressedTensorsWNA16MoEMethod (#44570)

Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Signed-off-by: Dobrzyniewicz, Agata <agata.dobrzyniewicz@intel.com> Signed-off-by: Tushar Jain <tushar00jain@users.noreply.github.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Agata Dobrzyniewicz <160237065+adobrzyn@users.noreply.github.com> Co-authored-by: Tushar Jain <8455015+tushar00jain@users.noreply.github.com> Co-authored-by: Tushar Jain <tushar00jain@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

bnellnm committed Jul 31, 2026 at 17:13 UTC 454ea5b52611c933e00581723e2db56f0144cea7
14 files changed +545 -855
tests/kernels/quantization/test_marlin_tile_padding.py
+1 -1
@@ -678,7 +678,7 @@ def test_check_moe_marlin_supports_layer_padding():
678 layer.hidden_size = hidden
679 layer.apply_router_weight_on_input = False
680 layer.moe_config = SimpleNamespace(
681 - intermediate_size_per_partition_unpadded=intermediate
681 + hidden_dim=hidden, intermediate_size_per_partition_unpadded=intermediate
682 )
683 return layer
684
tests/quantization/test_compressed_tensors.py
+4 -4
@@ -928,12 +928,12 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
928 (None, 32, 64, 128, (False, 64, True)),
929 ],
930 )
931 -def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected):
932 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501
933 - CompressedTensorsWNA16MarlinMoEMethod,
931 +def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected):
932 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
933 + CompressedTensorsWNA16MoEMethod,
934 )
935
936 - result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding(
936 + result = CompressedTensorsWNA16MoEMethod._w2_scale_sharding(
937 actorder, group_size, part, full
938 )
939 assert result == expected
tests/quantization/test_moe_wna16.py
+6
@@ -91,11 +91,17 @@ def test_map_wna16_backend_supports_triton():
91 def test_wna16_oracle_rejects_incompatible_quant_structures(
92 backend, quant_config, may_have_zp, may_have_bias, expected
93 ):
94 + from tests.kernels.moe.utils import make_dummy_moe_config
95 +
96 + moe_config = make_dummy_moe_config()
97 +
98 reason = _backend_incompatibility_reason(
99 backend=backend,
100 + moe_config=moe_config,
101 quant_config=quant_config,
102 may_have_zp=may_have_zp,
103 may_have_bias=may_have_bias,
104 + allow_tile_padding=True,
105 )
106
107 assert reason is not None
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
-1
@@ -593,7 +593,6 @@ class TritonWNA16Experts(TritonExperts):
593
594 @staticmethod
595 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
596 - # Why?
596 return not (
597 moe_parallel_config.use_fi_nvl_two_sided_kernels
598 or moe_parallel_config.use_fi_nvl_one_sided_kernels
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+34 -3
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
32 )
33 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
34 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
35 + check_moe_marlin_supports_config,
36 marlin_act_int8_process_scales,
37 marlin_moe_padded_intermediate,
38 marlin_moe_permute_scales,
@@ -125,9 +126,11 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
126
127 def _backend_incompatibility_reason(
128 backend: WNA16MoEBackend,
129 + moe_config: FusedMoEConfig,
130 quant_config: QuantizationConfig | QuantizationArgs,
131 may_have_zp: bool,
132 may_have_bias: bool,
133 + allow_tile_padding: bool,
134 ) -> str | None:
135 if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias):
136 return "zero points and bias are not supported"
@@ -149,7 +152,24 @@ def _backend_incompatibility_reason(
152 ):
153 return "group activation ordering is not supported"
154
152 - if isinstance(quant_config, MoeWNA16Config) and backend in (
155 + # Marlin only supports certain problem/group sizes.
156 + allow_marlin = not isinstance(quant_config, MoeWNA16Config)
157 +
158 + if allow_marlin and backend in (
159 + WNA16MoEBackend.MARLIN,
160 + WNA16MoEBackend.BATCHED_MARLIN,
161 + ):
162 + if isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, QuantizationArgs)):
163 + group_size = quant_config.group_size
164 + else:
165 + return "Marlin not supported for this layer"
166 +
167 + if not check_moe_marlin_supports_config(
168 + moe_config, group_size, allow_tile_padding
169 + ):
170 + return "Marlin not supported for this layer"
171 +
172 + if not allow_marlin and backend in (
173 WNA16MoEBackend.MARLIN,
174 WNA16MoEBackend.BATCHED_MARLIN,
175 WNA16MoEBackend.EMULATION,
@@ -182,6 +202,7 @@ def select_wna16_moe_backend(
202 quant_config: QuantizationConfig | QuantizationArgs,
203 may_have_zp: bool,
204 may_have_bias: bool,
205 + allow_tile_padding: bool = False,
206 ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
207 """Select the WNA16 MoE backend.
208
@@ -239,7 +260,12 @@ def select_wna16_moe_backend(
260 if runner_backend != "auto":
261 requested_backend = map_wna16_backend(runner_backend)
262 reason = _backend_incompatibility_reason(
242 - requested_backend, quant_config, may_have_zp, may_have_bias
263 + requested_backend,
264 + config,
265 + quant_config,
266 + may_have_zp,
267 + may_have_bias,
268 + allow_tile_padding,
269 )
270 if reason is not None:
271 raise ValueError(_make_log_unsupported(requested_backend, reason))
@@ -252,7 +278,12 @@ def select_wna16_moe_backend(
278
279 for backend in AVAILABLE_BACKENDS:
280 reason = _backend_incompatibility_reason(
255 - backend, quant_config, may_have_zp, may_have_bias
281 + backend,
282 + config,
283 + quant_config,
284 + may_have_zp,
285 + may_have_bias,
286 + allow_tile_padding,
287 )
288 if reason is not None:
289 logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
vllm/model_executor/layers/fused_moe/routed_experts.py
-2
@@ -206,7 +206,6 @@ class RoutedExperts(PluggableLayer):
206 def _needs_intermediate_size_param(self, quant_method: FusedMoEMethodBase) -> bool:
207 return quant_method.__class__.__name__ in (
208 "AutoGPTQMoEMethod",
209 - "CompressedTensorsWNA16MarlinMoEMethod",
209 "CompressedTensorsWNA16MoEMethod",
210 "CompressedTensorsW4A16FlydslMoEMethod",
211 )
@@ -659,7 +658,6 @@ class RoutedExperts(PluggableLayer):
658 # TODO (mgoin): check self.quant_method.quant_config.quant_format
659 # against known CompressionFormat enum values that have this quality
660 if quant_method_name in (
662 - "CompressedTensorsWNA16MarlinMoEMethod",
661 "CompressedTensorsWNA16MoEMethod",
662 "CompressedTensorsWNA16RDNA3MoEMethod",
663 "CompressedTensorsW4A16FlydslMoEMethod",
vllm/model_executor/layers/quantization/auto_awq.py
+1 -10
@@ -563,6 +563,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
563 quant_config=self.quant_config,
564 may_have_zp=self.quant_config.zero_point,
565 may_have_bias=True,
566 + allow_tile_padding=True,
567 )
568
569 def create_weights(
@@ -782,16 +783,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
783 gemm1_beta=getattr(layer, "swiglu_beta", None),
784 )
785
785 - def select_gemm_impl(
786 - self,
787 - prepare_finalize,
788 - layer: RoutedExperts,
789 - ):
790 - raise ValueError(
791 - f"{self.__class__.__name__} uses the new modular kernel "
792 - "initialization logic. This function should not be called."
793 - )
794 -
786 def apply(
787 self,
788 layer: RoutedExperts,
vllm/model_executor/layers/quantization/auto_gptq.py
+2 -11
@@ -490,8 +490,9 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
490 moe,
491 weight_key,
492 quant_config=self.quant_config,
493 - may_have_zp=True,
493 + may_have_zp=not self.quant_config.is_sym,
494 may_have_bias=True,
495 + allow_tile_padding=not self.quant_config.desc_act,
496 )
497
498 def create_weights(
@@ -806,16 +807,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
807 w2_bias=getattr(layer, "w2_bias", None),
808 )
809
809 - def select_gemm_impl(
810 - self,
811 - prepare_finalize,
812 - layer: RoutedExperts,
813 - ):
814 - raise ValueError(
815 - f"{self.__class__.__name__} uses the new modular kernel "
816 - "initialization logic. This function should not be called."
817 - )
818 -
810 def apply(
811 self,
812 layer: RoutedExperts,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+43 -84
@@ -5,7 +5,6 @@
5 import torch
6 from compressed_tensors import CompressionFormat
7 from compressed_tensors.quantization import (
8 - ActivationOrdering,
8 QuantizationStrategy,
9 QuantizationType,
10 )
@@ -19,9 +18,6 @@ from vllm.model_executor.layers.fused_moe import (
18 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
19 WNA16_SUPPORTED_BITS,
20 )
22 -from vllm.model_executor.layers.quantization.utils.marlin_utils import (
23 - check_moe_marlin_supports_layer,
24 -)
21 from vllm.platforms import current_platform
22
23 logger = init_logger(__name__)
@@ -78,9 +74,6 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
74 return CompressedTensorsW8A8Mxfp8MoEMethod(layer.moe_config)
75
76 if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
81 - # group_size=None means channelwise
82 - group_size = weight_quant.group_size or -1
83 -
77 valid_format_and_bits = (
78 weight_quant.num_bits in WNA16_SUPPORTED_BITS
79 and format == CompressionFormat.pack_quantized.value
@@ -95,87 +88,53 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
88 f" and bits: {weight_quant.num_bits}",
89 )
90
98 - # Prefer to use the MarlinMoE kernel when it is supported.
99 - is_actorder = (
100 - weight_quant.strategy == QuantizationStrategy.GROUP
101 - and weight_quant.actorder
102 - in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
103 - )
104 - if (
105 - not check_moe_marlin_supports_layer(
106 - layer, group_size, allow_tile_padding=not is_actorder
107 - )
108 - or current_platform.is_rocm()
109 - ):
110 - if is_actorder:
111 - raise ValueError(
112 - "WNA16MoE is not supported with actorder=group/dynamic."
91 + # Native ROCm HIP kernels (RDNA3, etc.)
92 + if current_platform.is_rocm():
93 + from . import rocm_moe_rdna
94 +
95 + if rocm_moe_rdna.is_supported(weight_quant):
96 + return rocm_moe_rdna.make_method(
97 + weight_quant, input_quant, layer.moe_config
98 + )
99 + from vllm.platforms.rocm import on_gfx950
100 +
101 + vllm_config = get_current_vllm_config()
102 + is_lora_disabled = vllm_config.lora_config is None
103 + moe_backend = vllm_config.kernel_config.moe_backend
104 + group_size = weight_quant.group_size or -1
105 + if (
106 + weight_quant.strategy == QuantizationStrategy.GROUP
107 + and weight_quant.type == QuantizationType.INT
108 + and group_size == 32
109 + and weight_quant.num_bits == 4
110 + and is_lora_disabled
111 + and on_gfx950()
112 + and moe_backend == "flydsl"
113 + ):
114 + from .compressed_tensors_moe_w4a16_flydsl import (
115 + CompressedTensorsW4A16FlydslMoEMethod,
116 )
117
115 - # Native ROCm HIP kernels (RDNA3, etc.)
116 - if current_platform.is_rocm():
117 - from . import rocm_moe_rdna
118 -
119 - if rocm_moe_rdna.is_supported(weight_quant):
120 - return rocm_moe_rdna.make_method(
121 - weight_quant, input_quant, layer.moe_config
122 - )
123 - from vllm.platforms.rocm import on_gfx950
124 -
125 - vllm_config = get_current_vllm_config()
126 - is_lora_disabled = vllm_config.lora_config is None
127 - moe_backend = vllm_config.kernel_config.moe_backend
128 - if (
129 - weight_quant.strategy == QuantizationStrategy.GROUP
130 - and weight_quant.type == QuantizationType.INT
131 - and group_size == 32
132 - and weight_quant.num_bits == 4
133 - and is_lora_disabled
134 - and on_gfx950()
135 - and moe_backend == "flydsl"
136 - ):
137 - from .compressed_tensors_moe_w4a16_flydsl import (
138 - CompressedTensorsW4A16FlydslMoEMethod,
139 - )
140 -
141 - logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod")
142 - return CompressedTensorsW4A16FlydslMoEMethod(
143 - weight_quant, input_quant, layer.moe_config
144 - )
145 - elif moe_backend == "emulation":
146 - # Although this is called 'Marlin', actually it selects
147 - # emulation backend by calling select_wna16_moe_backend.
148 - # TODO: we need to update CompressedTensorsWNA16MoeMethod
149 - # to honor "--moe-backend" option
150 - from .compressed_tensors_moe_wna16_marlin import (
151 - CompressedTensorsWNA16MarlinMoEMethod,
152 - )
153 -
154 - logger.info_once(
155 - "Using CompressedTensorsWNA16MarlinMoEMethod "
156 - "(emulation backend requested)"
157 - )
158 - return CompressedTensorsWNA16MarlinMoEMethod(
159 - weight_quant, input_quant, layer.moe_config, layer_name
160 - )
161 -
162 - from .compressed_tensors_moe_wna16 import (
163 - CompressedTensorsWNA16MoEMethod,
164 - )
118 + logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod")
119 + return CompressedTensorsW4A16FlydslMoEMethod(
120 + weight_quant, input_quant, layer.moe_config
121 + )
122 + elif moe_backend == "emulation":
123 + logger.info_once(
124 + "Using CompressedTensorsWNA16MoEMethod "
125 + "(emulation backend requested)"
126 + )
127
166 - logger.info_once("Using CompressedTensorsWNA16MoEMethod")
167 - return CompressedTensorsWNA16MoEMethod(
168 - weight_quant, input_quant, layer.moe_config
169 - )
170 - else:
171 - from .compressed_tensors_moe_wna16_marlin import (
172 - CompressedTensorsWNA16MarlinMoEMethod,
173 - )
128 + from .compressed_tensors_moe_wna16 import (
129 + CompressedTensorsWNA16MoEMethod,
130 + )
131
175 - logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
176 - return CompressedTensorsWNA16MarlinMoEMethod(
177 - weight_quant, input_quant, layer.moe_config
178 - )
132 + logger.info_once("Using CompressedTensorsWNA16MoEMethod")
133 + return CompressedTensorsWNA16MoEMethod(
134 + weight_quant,
135 + input_quant,
136 + layer.moe_config,
137 + )
138 elif quant_config._is_nvfp4_format(weight_quant):
139 from .compressed_tensors_moe_w4a4_nvfp4 import (
140 CompressedTensorsW4A4Nvfp4MoEMethod,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+435 -88
@@ -1,28 +1,51 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from typing import Any
5
6 import torch
7 from compressed_tensors.quantization import (
8 + ActivationOrdering,
9 QuantizationArgs,
10 + QuantizationStrategy,
11 )
12
10 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.fused_moe import (
15 + FusedMoEExpertsModular,
16 RoutedExperts,
17 SharedExperts,
18 )
19 from vllm.model_executor.layers.fused_moe.config import (
20 FusedMoEConfig,
21 FusedMoEQuantConfig,
19 - int4_w4a16_moe_quant_config,
20 - int8_w8a16_moe_quant_config,
22 +)
23 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
24 + WNA16MoEBackend,
25 + convert_to_wna16_moe_kernel_format,
26 + make_wna16_moe_kernel,
27 + make_wna16_moe_quant_config,
28 + select_wna16_moe_backend,
29 )
30 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
31 CompressedTensorsMoEMethod,
32 )
25 -from vllm.model_executor.utils import set_weight_attrs
33 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
34 + WNA16_SUPPORTED_TYPES_MAP,
35 + WNA16_ZP_SUPPORTED_TYPES_MAP,
36 +)
37 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
38 + check_moe_marlin_supports_config,
39 + get_marlin_input_dtype,
40 + marlin_make_workspace_new,
41 +)
42 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
43 + QuantKey,
44 + kInt4Static32GroupScale,
45 + kInt4StaticGroupScale,
46 + kInt8StaticGroupScale,
47 +)
48 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
49
50 logger = init_logger(__name__)
51
@@ -39,19 +62,180 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
62 self.weight_quant = weight_quant
63 self.input_quant = input_quant
64 # Extract properties from weight_quant
65 + self.symmetric = weight_quant.symmetric
66 self.num_bits = weight_quant.num_bits
67 self.packed_factor = 32 // weight_quant.num_bits
68 self.strategy = weight_quant.strategy
45 - # channelwise is not supported by this kernel
46 - assert weight_quant.strategy == "group"
69 self.group_size = weight_quant.group_size
48 - # grouped actorder isn't supported by this kernel
49 - assert weight_quant.actorder != "group"
50 - assert weight_quant.symmetric, (
51 - "Only symmetric quantization is supported for MoE. "
52 - "Try --moe-backend emulation."
70 + self.actorder = weight_quant.actorder
71 +
72 + # Extract quant_type and create weight key for oracle selection
73 + self.quant_type = (
74 + WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
75 + if self.symmetric
76 + else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]
77 + )
78 +
79 + if self.num_bits == 4:
80 + if self.group_size == 32:
81 + scale = kInt4Static32GroupScale
82 + else:
83 + scale = kInt4StaticGroupScale
84 + elif self.num_bits == 8:
85 + assert self.group_size == -1
86 + scale = kInt8StaticGroupScale
87 + else:
88 + raise ValueError(
89 + "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now."
90 + )
91 +
92 + weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
93 +
94 + is_actorder = self.strategy == QuantizationStrategy.GROUP and self.actorder in (
95 + ActivationOrdering.GROUP,
96 + ActivationOrdering.DYNAMIC,
97 )
98
99 + # Select WNA16 MoE backend via oracle.
100 + self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
101 + config=self.moe,
102 + weight_key=weight_key,
103 + quant_config=self.weight_quant,
104 + may_have_zp=not self.symmetric,
105 + may_have_bias=False,
106 + allow_tile_padding=not is_actorder,
107 + )
108 +
109 + self.is_marlin = self.wna16_backend in [
110 + WNA16MoEBackend.MARLIN,
111 + WNA16MoEBackend.BATCHED_MARLIN,
112 + ]
113 + self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
114 +
115 + if self.is_marlin:
116 + assert check_moe_marlin_supports_config(
117 + self.moe, self.group_size, allow_tile_padding=not is_actorder
118 + )
119 + self.input_dtype = get_marlin_input_dtype(layer_name)
120 + else:
121 + # channelwise is not supported by this kernel
122 + assert weight_quant.strategy == "group"
123 + # grouped actorder isn't supported by this kernel
124 + assert weight_quant.actorder != "group"
125 +
126 + assert self.symmetric, "Only symmetric quantization is supported for MoE"
127 +
128 + # Non-Marlin WNA16 always uses bf16/fp16 inputs
129 + self.input_dtype = torch.bfloat16
130 +
131 + def get_weight_shape(
132 + self,
133 + weight_name: str,
134 + num_experts: int,
135 + hidden_size: int,
136 + intermediate_size_per_partition: int,
137 + num_groups_w2: int | None = None,
138 + num_groups_w13: int | None = None,
139 + ) -> tuple[int, int, int]:
140 + """
141 + Get the shape of the weight based on the weight name, number of experts
142 + hidden size, intermediate size per partition, number of groups for w2,
143 + and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
144 + for weight scales/zero_points.
145 + """
146 + if weight_name in ("w13_scale", "w13_zp"):
147 + assert num_groups_w13 is not None, (
148 + "num_groups_w13 must be provided for weight scales/zero_points"
149 + )
150 + if weight_name in ("w2_scale", "w2_zp"):
151 + assert num_groups_w2 is not None, (
152 + "num_groups_w2 must be provided for weight scales/zero_points"
153 + )
154 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
155 + shape_map = {
156 + "w13_weight": {
157 + "Flashinfer": (
158 + num_experts,
159 + w13_num_shards * intermediate_size_per_partition,
160 + hidden_size // self.packed_factor,
161 + ),
162 + "Marlin": (
163 + num_experts,
164 + hidden_size // self.packed_factor,
165 + w13_num_shards * intermediate_size_per_partition,
166 + ),
167 + },
168 + "w13_scale": {
169 + "Flashinfer": (
170 + num_experts,
171 + w13_num_shards * intermediate_size_per_partition,
172 + num_groups_w13,
173 + ),
174 + "Marlin": (
175 + num_experts,
176 + num_groups_w13,
177 + w13_num_shards * intermediate_size_per_partition,
178 + ),
179 + },
180 + "w13_zp": {
181 + "Marlin": (
182 + num_experts,
183 + num_groups_w13,
184 + w13_num_shards
185 + * intermediate_size_per_partition
186 + // self.packed_factor,
187 + ),
188 + },
189 + "w2_weight": {
190 + "Flashinfer": (
191 + num_experts,
192 + hidden_size,
193 + intermediate_size_per_partition // self.packed_factor,
194 + ),
195 + "Marlin": (
196 + num_experts,
197 + intermediate_size_per_partition // self.packed_factor,
198 + hidden_size,
199 + ),
200 + },
201 + "w2_scale": {
202 + "Flashinfer": (num_experts, hidden_size, num_groups_w2),
203 + "Marlin": (num_experts, num_groups_w2, hidden_size),
204 + },
205 + "w2_zp": {
206 + "Marlin": (
207 + num_experts,
208 + num_groups_w2,
209 + hidden_size // self.packed_factor,
210 + ),
211 + },
212 + }
213 + backend_key = "Marlin" if self.is_transposed else "Flashinfer"
214 + return shape_map[weight_name][backend_key]
215 +
216 + @staticmethod
217 + def _w2_scale_sharding(
218 + actorder,
219 + group_size: int,
220 + intermediate_size_per_partition: int,
221 + intermediate_size_full: int,
222 + ) -> tuple[bool, int, bool]:
223 + """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE.
224 +
225 + Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime
226 + and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``.
227 + ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at
228 + quantization time, so scales shard normally per TP rank.
229 + """
230 + load_full_w2 = (actorder == "group") and group_size != -1
231 + w2_scales_size = (
232 + intermediate_size_full if load_full_w2 else intermediate_size_per_partition
233 + )
234 + is_k_full = (actorder != "group") or (
235 + intermediate_size_per_partition == intermediate_size_full
236 + )
237 + return load_full_w2, w2_scales_size, is_k_full
238 +
239 def create_weights(
240 self,
241 layer: torch.nn.Module,
@@ -61,18 +245,23 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
245 params_dtype: torch.dtype,
246 **extra_weight_attrs,
247 ):
248 + intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
249 +
250 # Will transpose the loaded weight along the
251 # intermediate and hidden dim sizes. Will
252 # shard for TP along the transposed dims
253 extra_weight_attrs.update(
68 - {"is_transposed": True, "quant_method": self.strategy}
254 + {"is_transposed": self.is_transposed, "quant_method": self.strategy}
255 )
70 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
256 +
257 w13_weight = torch.nn.Parameter(
258 torch.empty(
73 - num_experts,
74 - hidden_size // self.packed_factor,
75 - w13_num_shards * intermediate_size_per_partition,
259 + *self.get_weight_shape(
260 + "w13_weight",
261 + num_experts,
262 + hidden_size,
263 + intermediate_size_per_partition,
264 + ),
265 dtype=torch.int32,
266 ),
267 requires_grad=False,
@@ -82,9 +271,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
271
272 w2_weight = torch.nn.Parameter(
273 torch.empty(
85 - num_experts,
86 - intermediate_size_per_partition // self.packed_factor,
87 - hidden_size,
274 + *self.get_weight_shape(
275 + "w2_weight",
276 + num_experts,
277 + hidden_size,
278 + intermediate_size_per_partition,
279 + ),
280 dtype=torch.int32,
281 ),
282 requires_grad=False,
@@ -92,7 +284,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
284 layer.register_parameter("w2_weight_packed", w2_weight)
285 set_weight_attrs(w2_weight, extra_weight_attrs)
286
95 - w2_scales_size = intermediate_size_per_partition
287 + load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding(
288 + self.actorder,
289 + self.group_size,
290 + intermediate_size_per_partition,
291 + intermediate_size_full,
292 + )
293
294 if self.strategy == "channel":
295 num_groups_w2 = num_groups_w13 = 1
@@ -104,23 +301,34 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
301 f"({hidden_size}) to be divisible by group_size "
302 f"({self.group_size})."
303 )
107 - if intermediate_size_per_partition % self.group_size != 0:
304 + if (
305 + not load_full_w2
306 + and intermediate_size_per_partition % self.group_size != 0
307 + ):
308 raise ValueError(
109 - "CompressedTensors WNA16 MoE with static group scales "
110 - "requires the MoE intermediate size per tensor-parallel "
111 - f"partition ({intermediate_size_per_partition}) to be "
112 - f"divisible by group_size ({self.group_size}). Scale "
113 - "groups would otherwise cross TP shard boundaries; use a "
114 - "compatible TP size or enable expert parallelism."
309 + "CompressedTensors WNA16 MoE with static group "
310 + "scales requires the MoE intermediate size per "
311 + "tensor-parallel partition "
312 + f"({intermediate_size_per_partition}) to be divisible by "
313 + f"group_size ({self.group_size}). Scale groups would "
314 + "otherwise cross TP shard boundaries; use a compatible TP "
315 + "size or enable expert parallelism."
316 )
317 num_groups_w2 = w2_scales_size // self.group_size
318 num_groups_w13 = hidden_size // self.group_size
319
320 + layer.num_groups_w13 = num_groups_w13
321 + layer.num_groups_w2 = num_groups_w2
322 +
323 w13_scale = torch.nn.Parameter(
324 torch.ones(
121 - num_experts,
122 - num_groups_w13,
123 - w13_num_shards * intermediate_size_per_partition,
325 + *self.get_weight_shape(
326 + "w13_scale",
327 + num_experts,
328 + hidden_size,
329 + intermediate_size_per_partition,
330 + num_groups_w13=num_groups_w13,
331 + ),
332 dtype=params_dtype,
333 ),
334 requires_grad=False,
@@ -129,12 +337,54 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
337 set_weight_attrs(w13_scale, extra_weight_attrs)
338
339 w2_scale = torch.nn.Parameter(
132 - torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
340 + torch.ones(
341 + *self.get_weight_shape(
342 + "w2_scale",
343 + num_experts,
344 + hidden_size,
345 + intermediate_size_per_partition,
346 + num_groups_w2=num_groups_w2,
347 + ),
348 + dtype=params_dtype,
349 + ),
350 requires_grad=False,
351 )
352 layer.register_parameter("w2_weight_scale", w2_scale)
353 set_weight_attrs(w2_scale, extra_weight_attrs)
137 - set_weight_attrs(w2_scale, {"load_full_w2": False})
354 + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
355 +
356 + if not self.symmetric:
357 + w13_zp = torch.nn.Parameter(
358 + torch.zeros(
359 + *self.get_weight_shape(
360 + "w13_zp",
361 + num_experts,
362 + hidden_size,
363 + intermediate_size_per_partition,
364 + num_groups_w13=num_groups_w13,
365 + ),
366 + dtype=torch.int32,
367 + ),
368 + requires_grad=False,
369 + )
370 + layer.register_parameter("w13_weight_zero_point", w13_zp)
371 + set_weight_attrs(w13_zp, extra_weight_attrs)
372 +
373 + w2_zp = torch.nn.Parameter(
374 + torch.zeros(
375 + *self.get_weight_shape(
376 + "w2_zp",
377 + num_experts,
378 + hidden_size,
379 + intermediate_size_per_partition,
380 + num_groups_w2=num_groups_w2,
381 + ),
382 + dtype=torch.int32,
383 + ),
384 + requires_grad=False,
385 + )
386 + layer.register_parameter("w2_weight_zero_point", w2_zp)
387 + set_weight_attrs(w2_zp, extra_weight_attrs)
388
389 w2_weight_shape = torch.nn.Parameter(
390 torch.empty(num_experts, 2), requires_grad=False
@@ -195,65 +445,161 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
445 layer.a13_scale = None
446 layer.a2_scale = None
447
198 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
199 - # Reconfigure packed weights and scales to match moe_wna16 format
200 - layer.w13_weight_packed = torch.nn.Parameter(
201 - layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
202 - requires_grad=False,
203 - )
204 - layer.w2_weight_packed = torch.nn.Parameter(
205 - layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
206 - requires_grad=False,
207 - )
208 - layer.w13_weight_scale = torch.nn.Parameter(
209 - layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
448 + def _setup_kernel(self, layer: RoutedExperts):
449 + assert self.experts_cls is not None
450 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
451 + assert self.moe_quant_config is not None
452 +
453 + # Add Marlin-specific arguments
454 + marlin_args: dict[str, Any] = {}
455 + if self.is_marlin:
456 + marlin_args = {
457 + "w13_g_idx": layer.w13_weight_g_idx,
458 + "w2_g_idx": layer.w2_weight_g_idx,
459 + "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
460 + "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
461 + "is_k_full": self.is_k_full,
462 + }
463 +
464 + self.moe_kernel = make_wna16_moe_kernel(
465 + moe_quant_config=self.moe_quant_config,
466 + moe_config=self.moe,
467 + experts_cls=self.experts_cls,
468 + routing_tables=layer._expert_routing_tables(),
469 + **marlin_args,
470 )
211 - layer.w2_weight_scale = torch.nn.Parameter(
212 - layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
471 +
472 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
473 + # Process weights using the shared oracle infrastructure
474 + converted = convert_to_wna16_moe_kernel_format(
475 + backend=self.wna16_backend,
476 + layer=layer,
477 + quant_config=self.weight_quant,
478 + input_dtype=self.input_dtype,
479 + w13=layer.w13_weight_packed,
480 + w2=layer.w2_weight_packed,
481 + w13_scale=layer.w13_weight_scale,
482 + w2_scale=layer.w2_weight_scale,
483 + w13_g_idx=layer.w13_weight_g_idx,
484 + w2_g_idx=layer.w2_weight_g_idx,
485 + w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
486 + w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
487 )
488
489 + if converted is None:
490 + self._setup_kernel(layer)
491 + return
492 +
493 + (
494 + w13_qweight,
495 + w2_qweight,
496 + w13_scales,
497 + w2_scales,
498 + w13_g_idx_processed,
499 + w2_g_idx_processed,
500 + w13_g_idx_sort_indices,
501 + w2_g_idx_sort_indices,
502 + w13_qzeros,
503 + w2_qzeros,
504 + w13_input_global_scale,
505 + w2_input_global_scale,
506 + _, # w13_bias
507 + _, # w2_bias
508 + ) = converted
509 +
510 + # Replace common parameters
511 + replace_parameter(layer, "w13_weight_packed", w13_qweight)
512 + replace_parameter(layer, "w2_weight_packed", w2_qweight)
513 + replace_parameter(layer, "w13_weight_scale", w13_scales)
514 + replace_parameter(layer, "w2_weight_scale", w2_scales)
515 +
516 + # CPU fused_experts_cpu requires zero points even for symmetric quant
517 + if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU:
518 + assert w13_qzeros is not None and w2_qzeros is not None
519 + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
520 + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
521 +
522 + # Marlin-specific parameters (not needed for Flashinfer)
523 + if self.is_marlin:
524 + if w13_g_idx_processed is not None:
525 + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
526 + if w2_g_idx_processed is not None:
527 + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
528 + if w13_g_idx_sort_indices is not None:
529 + replace_parameter(
530 + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices
531 + )
532 + if w2_g_idx_sort_indices is not None:
533 + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
534 +
535 + # Register input global scales if present
536 + if w13_input_global_scale is not None:
537 + layer.register_parameter(
538 + "w13_input_global_scale",
539 + torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
540 + )
541 + if w2_input_global_scale is not None:
542 + layer.register_parameter(
543 + "w2_input_global_scale",
544 + torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
545 + )
546 +
547 + # Marlin workspace — only needed for Marlin-family backends, not emulation.
548 + if (
549 + self.experts_cls is not None
550 + and issubclass(self.experts_cls, FusedMoEExpertsModular)
551 + and self.wna16_backend != WNA16MoEBackend.EMULATION
552 + ):
553 + layer.workspace = marlin_make_workspace_new(
554 + layer.w13_weight_g_idx.device,
555 + 4,
556 + existing=getattr(layer, "workspace", None),
557 + )
558 +
559 + # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
560 + layer.w13_weight = layer.w13_weight_packed
561 + layer.w2_weight = layer.w2_weight_packed
562 +
563 + self._setup_kernel(layer)
564 +
565 def get_fused_moe_quant_config(
566 self, layer: torch.nn.Module
567 ) -> FusedMoEQuantConfig | None:
218 - assert self.num_bits == 4 or self.num_bits == 8
219 - config_builder = (
220 - int4_w4a16_moe_quant_config
221 - if self.num_bits == 4
222 - else int8_w8a16_moe_quant_config
223 - )
224 -
225 - return config_builder(
568 + return make_wna16_moe_quant_config(
569 w1_scale=layer.w13_weight_scale,
570 w2_scale=layer.w2_weight_scale,
228 - w1_zp=None,
229 - w2_zp=None,
230 - block_shape=[0, self.group_size],
571 + group_size=self.group_size,
572 + num_bits=self.num_bits,
573 + w1_zp=getattr(layer, "w13_weight_zero_point", None),
574 + w2_zp=getattr(layer, "w2_weight_zero_point", None),
575 + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
576 + gemm1_alpha=getattr(layer, "swiglu_alpha", None),
577 + gemm1_beta=getattr(layer, "swiglu_beta", None),
578 )
579
233 - def select_gemm_impl(
580 + def apply_monolithic(
581 self,
235 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
236 - layer: torch.nn.Module,
237 - ) -> mk.FusedMoEExpertsModular:
238 - if self.moe.is_lora_enabled:
239 - assert self.moe_quant_config is not None
240 - from vllm.triton_utils import HAS_TRITON
241 -
242 - if HAS_TRITON:
243 - from vllm.model_executor.layers.fused_moe import TritonWNA16Experts
244 -
245 - layer.w13_weight = layer.w13_weight_packed
246 - layer.w2_weight = layer.w2_weight_packed
247 - return TritonWNA16Experts(
248 - moe_config=self.moe, quant_config=self.moe_quant_config
249 - )
250 - else:
251 - raise NotImplementedError(
252 - "TritonExperts requires Triton. "
253 - "Install triton or disable LoRA for MoE."
254 - )
255 -
256 - raise NotImplementedError
582 + layer: RoutedExperts,
583 + x: torch.Tensor,
584 + router_logits: torch.Tensor,
585 + input_ids: torch.Tensor | None = None,
586 + ) -> torch.Tensor:
587 + assert self.is_monolithic
588 + assert self.moe_kernel is not None
589 + return self.moe_kernel.apply_monolithic(
590 + x,
591 + layer.w13_weight,
592 + layer.w2_weight,
593 + router_logits,
594 + activation=layer.activation,
595 + global_num_experts=layer.global_num_experts,
596 + expert_map=layer.expert_map,
597 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
598 + num_expert_group=layer.num_expert_group,
599 + topk_group=layer.topk_group,
600 + e_score_correction_bias=layer.e_score_correction_bias,
601 + routed_scaling_factor=layer.routed_scaling_factor,
602 + )
603
604 def apply(
605 self,
@@ -264,21 +610,22 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
610 shared_experts: SharedExperts | None,
611 shared_experts_input: torch.Tensor | None,
612 ) -> torch.Tensor:
267 - from vllm.model_executor.layers.fused_moe import fused_experts
268 -
269 - return fused_experts(
613 + assert not self.is_monolithic
614 + assert self.moe_kernel is not None
615 + return self.moe_kernel.apply(
616 x,
271 - layer.w13_weight_packed,
272 - layer.w2_weight_packed,
617 + layer.w13_weight,
618 + layer.w2_weight,
619 topk_weights=topk_weights,
620 topk_ids=topk_ids,
621 activation=layer.activation,
276 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
622 global_num_experts=layer.global_num_experts,
623 expert_map=layer.expert_map,
279 - quant_config=self.moe_quant_config,
624 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
625 + shared_experts=shared_experts,
626 + shared_experts_input=shared_experts_input,
627 )
628
629 @property
630 def supports_eplb(self) -> bool:
284 - return True
631 + return self.wna16_backend == WNA16MoEBackend.TRITON
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py deleted
-601
@@ -1,601 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Any
5 -
6 -import torch
7 -from compressed_tensors.quantization import (
8 - QuantizationArgs,
9 -)
10 -
11 -from vllm.logger import init_logger
12 -from vllm.model_executor.layers.fused_moe import (
13 - FusedMoEExpertsModular,
14 - RoutedExperts,
15 - SharedExperts,
16 -)
17 -from vllm.model_executor.layers.fused_moe.config import (
18 - FusedMoEConfig,
19 - FusedMoEQuantConfig,
20 -)
21 -from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
22 - WNA16MoEBackend,
23 - convert_to_wna16_moe_kernel_format,
24 - make_wna16_moe_kernel,
25 - make_wna16_moe_quant_config,
26 - select_wna16_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.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
32 - WNA16_SUPPORTED_TYPES_MAP,
33 - WNA16_ZP_SUPPORTED_TYPES_MAP,
34 -)
35 -from vllm.model_executor.layers.quantization.utils.marlin_utils import (
36 - get_marlin_input_dtype,
37 - marlin_make_workspace_new,
38 -)
39 -from vllm.model_executor.layers.quantization.utils.quant_utils import (
40 - QuantKey,
41 - kInt4Static32GroupScale,
42 - kInt4StaticGroupScale,
43 - kInt8StaticGroupScale,
44 -)
45 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs
46 -
47 -logger = init_logger(__name__)
48 -
49 -
50 -class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
51 - def __init__(
52 - self,
53 - weight_quant: QuantizationArgs,
54 - input_quant: QuantizationArgs | None,
55 - moe: FusedMoEConfig,
56 - layer_name: str | None = None,
57 - ):
58 - super().__init__(moe)
59 - self.weight_quant = weight_quant
60 - self.input_quant = input_quant
61 - self.symmetric = weight_quant.symmetric
62 - # Extract properties from weight_quant
63 - self.num_bits = weight_quant.num_bits
64 - self.packed_factor = 32 // weight_quant.num_bits
65 - self.strategy = weight_quant.strategy
66 - self.group_size = weight_quant.group_size
67 - self.actorder = weight_quant.actorder
68 -
69 - self.quant_type = (
70 - WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
71 - if self.symmetric
72 - else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]
73 - )
74 -
75 - self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
76 -
77 - if self.num_bits == 4:
78 - if self.group_size == 32:
79 - scale = kInt4Static32GroupScale
80 - else:
81 - scale = kInt4StaticGroupScale
82 - elif self.num_bits == 8:
83 - scale = kInt8StaticGroupScale
84 - else:
85 - raise ValueError(
86 - "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now."
87 - )
88 -
89 - weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
90 -
91 - # Select WNA16 MoE backend via oracle.
92 - self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
93 - config=self.moe,
94 - weight_key=weight_key,
95 - quant_config=self.weight_quant,
96 - may_have_zp=not self.symmetric,
97 - may_have_bias=False,
98 - )
99 - self.is_marlin = self.wna16_backend in [
100 - WNA16MoEBackend.MARLIN,
101 - WNA16MoEBackend.BATCHED_MARLIN,
102 - ]
103 - self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
104 -
105 - def get_weight_shape(
106 - self,
107 - weight_name: str,
108 - num_experts: int,
109 - hidden_size: int,
110 - intermediate_size_per_partition: int,
111 - num_groups_w2: int | None = None,
112 - num_groups_w13: int | None = None,
113 - ) -> tuple[int, int, int]:
114 - """
115 - Get the shape of the weight based on the weight name, number of experts
116 - hidden size, intermediate size per partition, number of groups for w2,
117 - and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
118 - for weight scales/zero_points.
119 - """
120 - if weight_name in ("w13_scale", "w13_zp"):
121 - assert num_groups_w13 is not None, (
122 - "num_groups_w13 must be provided for weight scales/zero_points"
123 - )
124 - if weight_name in ("w2_scale", "w2_zp"):
125 - assert num_groups_w2 is not None, (
126 - "num_groups_w2 must be provided for weight scales/zero_points"
127 - )
128 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1
129 - shape_map = {
130 - "w13_weight": {
131 - "Flashinfer": (
132 - num_experts,
133 - w13_num_shards * intermediate_size_per_partition,
134 - hidden_size // self.packed_factor,
135 - ),
136 - "Marlin": (
137 - num_experts,
138 - hidden_size // self.packed_factor,
139 - w13_num_shards * intermediate_size_per_partition,
140 - ),
141 - },
142 - "w13_scale": {
143 - "Flashinfer": (
144 - num_experts,
145 - w13_num_shards * intermediate_size_per_partition,
146 - num_groups_w13,
147 - ),
148 - "Marlin": (
149 - num_experts,
150 - num_groups_w13,
151 - w13_num_shards * intermediate_size_per_partition,
152 - ),
153 - },
154 - "w13_zp": {
155 - "Marlin": (
156 - num_experts,
157 - num_groups_w13,
158 - w13_num_shards
159 - * intermediate_size_per_partition
160 - // self.packed_factor,
161 - ),
162 - },
163 - "w2_weight": {
164 - "Flashinfer": (
165 - num_experts,
166 - hidden_size,
167 - intermediate_size_per_partition // self.packed_factor,
168 - ),
169 - "Marlin": (
170 - num_experts,
171 - intermediate_size_per_partition // self.packed_factor,
172 - hidden_size,
173 - ),
174 - },
175 - "w2_scale": {
176 - "Flashinfer": (num_experts, hidden_size, num_groups_w2),
177 - "Marlin": (num_experts, num_groups_w2, hidden_size),
178 - },
179 - "w2_zp": {
180 - "Marlin": (
181 - num_experts,
182 - num_groups_w2,
183 - hidden_size // self.packed_factor,
184 - ),
185 - },
186 - }
187 - backend_key = "Marlin" if self.is_transposed else "Flashinfer"
188 - return shape_map[weight_name][backend_key]
189 -
190 - @staticmethod
191 - def _w2_scale_sharding(
192 - actorder,
193 - group_size: int,
194 - intermediate_size_per_partition: int,
195 - intermediate_size_full: int,
196 - ) -> tuple[bool, int, bool]:
197 - """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE.
198 -
199 - Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime
200 - and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``.
201 - ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at
202 - quantization time, so scales shard normally per TP rank.
203 - """
204 - load_full_w2 = (actorder == "group") and group_size != -1
205 - w2_scales_size = (
206 - intermediate_size_full if load_full_w2 else intermediate_size_per_partition
207 - )
208 - is_k_full = (actorder != "group") or (
209 - intermediate_size_per_partition == intermediate_size_full
210 - )
211 - return load_full_w2, w2_scales_size, is_k_full
212 -
213 - def create_weights(
214 - self,
215 - layer: torch.nn.Module,
216 - num_experts: int,
217 - hidden_size: int,
218 - intermediate_size_per_partition: int,
219 - params_dtype: torch.dtype,
220 - **extra_weight_attrs,
221 - ):
222 - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
223 -
224 - # Will transpose the loaded weight along the
225 - # intermediate and hidden dim sizes. Will
226 - # shard for TP along the transposed dims
227 - extra_weight_attrs.update(
228 - {"is_transposed": self.is_transposed, "quant_method": self.strategy}
229 - )
230 -
231 - w13_weight = torch.nn.Parameter(
232 - torch.empty(
233 - *self.get_weight_shape(
234 - "w13_weight",
235 - num_experts,
236 - hidden_size,
237 - intermediate_size_per_partition,
238 - ),
239 - dtype=torch.int32,
240 - ),
241 - requires_grad=False,
242 - )
243 - layer.register_parameter("w13_weight_packed", w13_weight)
244 - set_weight_attrs(w13_weight, extra_weight_attrs)
245 -
246 - w2_weight = torch.nn.Parameter(
247 - torch.empty(
248 - *self.get_weight_shape(
249 - "w2_weight",
250 - num_experts,
251 - hidden_size,
252 - intermediate_size_per_partition,
253 - ),
254 - dtype=torch.int32,
255 - ),
256 - requires_grad=False,
257 - )
258 - layer.register_parameter("w2_weight_packed", w2_weight)
259 - set_weight_attrs(w2_weight, extra_weight_attrs)
260 -
261 - load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding(
262 - self.actorder,
263 - self.group_size,
264 - intermediate_size_per_partition,
265 - intermediate_size_full,
266 - )
267 -
268 - if self.strategy == "channel":
269 - num_groups_w2 = num_groups_w13 = 1
270 - self.group_size = -1
271 - else:
272 - if hidden_size % self.group_size != 0:
273 - raise ValueError(
274 - "CompressedTensors WNA16 Marlin MoE requires hidden_size "
275 - f"({hidden_size}) to be divisible by group_size "
276 - f"({self.group_size})."
277 - )
278 - if (
279 - not load_full_w2
280 - and intermediate_size_per_partition % self.group_size != 0
281 - ):
282 - raise ValueError(
283 - "CompressedTensors WNA16 Marlin MoE with static group "
284 - "scales requires the MoE intermediate size per "
285 - "tensor-parallel partition "
286 - f"({intermediate_size_per_partition}) to be divisible by "
287 - f"group_size ({self.group_size}). Scale groups would "
288 - "otherwise cross TP shard boundaries; use a compatible TP "
289 - "size or enable expert parallelism."
290 - )
291 - num_groups_w2 = w2_scales_size // self.group_size
292 - num_groups_w13 = hidden_size // self.group_size
293 -
294 - layer.num_groups_w13 = num_groups_w13
295 - layer.num_groups_w2 = num_groups_w2
296 -
297 - w13_scale = torch.nn.Parameter(
298 - torch.ones(
299 - *self.get_weight_shape(
300 - "w13_scale",
301 - num_experts,
302 - hidden_size,
303 - intermediate_size_per_partition,
304 - num_groups_w13=num_groups_w13,
305 - ),
306 - dtype=params_dtype,
307 - ),
308 - requires_grad=False,
309 - )
310 - layer.register_parameter("w13_weight_scale", w13_scale)
311 - set_weight_attrs(w13_scale, extra_weight_attrs)
312 -
313 - w2_scale = torch.nn.Parameter(
314 - torch.ones(
315 - *self.get_weight_shape(
316 - "w2_scale",
317 - num_experts,
318 - hidden_size,
319 - intermediate_size_per_partition,
320 - num_groups_w2=num_groups_w2,
321 - ),
322 - dtype=params_dtype,
323 - ),
324 - requires_grad=False,
325 - )
326 - layer.register_parameter("w2_weight_scale", w2_scale)
327 - set_weight_attrs(w2_scale, extra_weight_attrs)
328 - set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
329 -
330 - if not self.symmetric:
331 - w13_zp = torch.nn.Parameter(
332 - torch.zeros(
333 - *self.get_weight_shape(
334 - "w13_zp",
335 - num_experts,
336 - hidden_size,
337 - intermediate_size_per_partition,
338 - num_groups_w13=num_groups_w13,
339 - ),
340 - dtype=torch.int32,
341 - ),
342 - requires_grad=False,
343 - )
344 - layer.register_parameter("w13_weight_zero_point", w13_zp)
345 - set_weight_attrs(w13_zp, extra_weight_attrs)
346 -
347 - w2_zp = torch.nn.Parameter(
348 - torch.zeros(
349 - *self.get_weight_shape(
350 - "w2_zp",
351 - num_experts,
352 - hidden_size,
353 - intermediate_size_per_partition,
354 - num_groups_w2=num_groups_w2,
355 - ),
356 - dtype=torch.int32,
357 - ),
358 - requires_grad=False,
359 - )
360 - layer.register_parameter("w2_weight_zero_point", w2_zp)
361 - set_weight_attrs(w2_zp, extra_weight_attrs)
362 -
363 - w2_weight_shape = torch.nn.Parameter(
364 - torch.empty(num_experts, 2), requires_grad=False
365 - )
366 - layer.register_parameter("w2_weight_shape", w2_weight_shape)
367 - set_weight_attrs(w2_weight_shape, extra_weight_attrs)
368 - w13_weight_shape = torch.nn.Parameter(
369 - torch.empty(num_experts, 2), requires_grad=False
370 - )
371 -
372 - layer.register_parameter("w13_weight_shape", w13_weight_shape)
373 - set_weight_attrs(w13_weight_shape, extra_weight_attrs)
374 -
375 - w13_g_idx = torch.nn.Parameter(
376 - torch.empty(
377 - num_experts,
378 - hidden_size,
379 - dtype=torch.int32,
380 - ),
381 - requires_grad=False,
382 - )
383 - layer.register_parameter("w13_weight_g_idx", w13_g_idx)
384 - set_weight_attrs(w13_g_idx, extra_weight_attrs)
385 -
386 - w2_g_idx = torch.nn.Parameter(
387 - torch.empty(
388 - num_experts,
389 - intermediate_size_per_partition,
390 - dtype=torch.int32,
391 - ),
392 - requires_grad=False,
393 - )
394 - layer.register_parameter("w2_weight_g_idx", w2_g_idx)
395 - set_weight_attrs(w2_g_idx, extra_weight_attrs)
396 -
397 - w13_g_idx_sort_indices = torch.nn.Parameter(
398 - torch.empty(
399 - num_experts,
400 - hidden_size,
401 - dtype=torch.int32,
402 - ),
403 - requires_grad=False,
404 - )
405 - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
406 - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
407 -
408 - w2_g_idx_sort_indices = torch.nn.Parameter(
409 - torch.empty(
410 - num_experts,
411 - intermediate_size_per_partition,
412 - dtype=torch.int32,
413 - ),
414 - requires_grad=False,
415 - )
416 - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
417 - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
418 -
419 - layer.a13_scale = None
420 - layer.a2_scale = None
421 -
422 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
423 - # Process weights using the shared oracle infrastructure
424 - converted = convert_to_wna16_moe_kernel_format(
425 - backend=self.wna16_backend,
426 - layer=layer,
427 - quant_config=self.weight_quant,
428 - input_dtype=self.marlin_input_dtype,
429 - w13=layer.w13_weight_packed,
430 - w2=layer.w2_weight_packed,
431 - w13_scale=layer.w13_weight_scale,
432 - w2_scale=layer.w2_weight_scale,
433 - w13_g_idx=layer.w13_weight_g_idx,
434 - w2_g_idx=layer.w2_weight_g_idx,
435 - w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
436 - w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
437 - )
438 -
439 - if converted is None:
440 - # In-place backends (e.g. Humming) are not wired through this
441 - # marlin-only method; fail clearly rather than unpacking None.
442 - raise NotImplementedError(
443 - f"{type(self).__name__} does not support the "
444 - f"{self.wna16_backend.value} MoE backend."
445 - )
446 -
447 - (
448 - w13_qweight,
449 - w2_qweight,
450 - w13_scales,
451 - w2_scales,
452 - w13_g_idx_processed,
453 - w2_g_idx_processed,
454 - w13_g_idx_sort_indices,
455 - w2_g_idx_sort_indices,
456 - w13_qzeros,
457 - w2_qzeros,
458 - w13_input_global_scale,
459 - w2_input_global_scale,
460 - _, # w13_bias
461 - _, # w2_bias
462 - ) = converted
463 -
464 - # Replace common parameters
465 - replace_parameter(layer, "w13_weight_packed", w13_qweight)
466 - replace_parameter(layer, "w2_weight_packed", w2_qweight)
467 - replace_parameter(layer, "w13_weight_scale", w13_scales)
468 - replace_parameter(layer, "w2_weight_scale", w2_scales)
469 -
470 - # CPU fused_experts_cpu requires zero points even for symmetric quant
471 - if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU:
472 - replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
473 - replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
474 -
475 - # Marlin-specific parameters (not needed for Flashinfer)
476 - if self.is_marlin:
477 - if w13_g_idx_processed is not None:
478 - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
479 - if w2_g_idx_processed is not None:
480 - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
481 - if w13_g_idx_sort_indices is not None:
482 - replace_parameter(
483 - layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices
484 - )
485 - if w2_g_idx_sort_indices is not None:
486 - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
487 -
488 - # Register input global scales if present
489 - if w13_input_global_scale is not None:
490 - layer.register_parameter(
491 - "w13_input_global_scale",
492 - torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
493 - )
494 - if w2_input_global_scale is not None:
495 - layer.register_parameter(
496 - "w2_input_global_scale",
497 - torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
498 - )
499 -
500 - # Marlin workspace — only needed for Marlin-family backends, not emulation.
501 - if (
502 - self.experts_cls is not None
503 - and issubclass(self.experts_cls, FusedMoEExpertsModular)
504 - and self.wna16_backend != WNA16MoEBackend.EMULATION
505 - ):
506 - layer.workspace = marlin_make_workspace_new(
507 - layer.w13_weight_g_idx.device,
508 - 4,
509 - existing=getattr(layer, "workspace", None),
510 - )
511 -
512 - # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
513 - layer.w13_weight = layer.w13_weight_packed
514 - layer.w2_weight = layer.w2_weight_packed
515 -
516 - assert self.experts_cls is not None
517 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)
518 - assert self.moe_quant_config is not None
519 -
520 - # Add Marlin-specific arguments
521 - marlin_args: dict[str, Any] = {}
522 - if self.is_marlin:
523 - marlin_args = {
524 - "w13_g_idx": layer.w13_weight_g_idx,
525 - "w2_g_idx": layer.w2_weight_g_idx,
526 - "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
527 - "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
528 - "is_k_full": self.is_k_full,
529 - }
530 -
531 - self.moe_kernel = make_wna16_moe_kernel(
532 - moe_quant_config=self.moe_quant_config,
533 - moe_config=self.moe,
534 - experts_cls=self.experts_cls,
535 - routing_tables=layer._expert_routing_tables(),
536 - **marlin_args,
537 - )
538 -
539 - def get_fused_moe_quant_config(
540 - self, layer: torch.nn.Module
541 - ) -> FusedMoEQuantConfig | None:
542 - return make_wna16_moe_quant_config(
543 - w1_scale=layer.w13_weight_scale,
544 - w2_scale=layer.w2_weight_scale,
545 - group_size=self.group_size,
546 - num_bits=self.num_bits,
547 - w1_zp=getattr(layer, "w13_weight_zero_point", None),
548 - w2_zp=getattr(layer, "w2_weight_zero_point", None),
549 - gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
550 - gemm1_alpha=getattr(layer, "swiglu_alpha", None),
551 - gemm1_beta=getattr(layer, "swiglu_beta", None),
552 - )
553 -
554 - def apply_monolithic(
555 - self,
556 - layer: RoutedExperts,
557 - x: torch.Tensor,
558 - router_logits: torch.Tensor,
559 - input_ids: torch.Tensor | None = None,
560 - ) -> torch.Tensor:
561 - assert self.is_monolithic
562 - assert self.moe_kernel is not None
563 - return self.moe_kernel.apply_monolithic(
564 - x,
565 - layer.w13_weight,
566 - layer.w2_weight,
567 - router_logits,
568 - activation=layer.activation,
569 - global_num_experts=layer.global_num_experts,
570 - expert_map=layer.expert_map,
571 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
572 - num_expert_group=layer.num_expert_group,
573 - topk_group=layer.topk_group,
574 - e_score_correction_bias=layer.e_score_correction_bias,
575 - routed_scaling_factor=layer.routed_scaling_factor,
576 - )
577 -
578 - def apply(
579 - self,
580 - layer: RoutedExperts,
581 - x: torch.Tensor,
582 - topk_weights: torch.Tensor,
583 - topk_ids: torch.Tensor,
584 - shared_experts: SharedExperts | None,
585 - shared_experts_input: torch.Tensor | None,
586 - ) -> torch.Tensor:
587 - assert not self.is_monolithic
588 - assert self.moe_kernel is not None
589 - return self.moe_kernel.apply(
590 - x,
591 - layer.w13_weight,
592 - layer.w2_weight,
593 - topk_weights,
594 - topk_ids,
595 - activation=layer.activation,
596 - global_num_experts=layer.global_num_experts,
597 - expert_map=layer.expert_map,
598 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
599 - shared_experts=shared_experts,
600 - shared_experts_input=shared_experts_input,
601 - )
vllm/model_executor/layers/quantization/modelopt.py
-20
@@ -778,16 +778,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
778 "logic. This function should not be called."
779 )
780
781 - def select_gemm_impl(
782 - self,
783 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
784 - layer: RoutedExperts,
785 - ) -> mk.FusedMoEExpertsModular:
786 - raise ValueError(
787 - f"{self.__class__.__name__} uses the new modular kernel initialization "
788 - "logic. This function should not be called."
789 - )
790 -
781 def create_weights(
782 self,
783 layer: RoutedExperts,
@@ -2104,16 +2094,6 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2094 "logic. This function should not be called."
2095 )
2096
2107 - def select_gemm_impl(
2108 - self,
2109 - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
2110 - layer: RoutedExperts,
2111 - ) -> mk.FusedMoEExpertsModular:
2112 - raise ValueError(
2113 - f"{self.__class__.__name__} uses the new modular kernel initialization "
2114 - "logic. This function should not be called."
2115 - )
2116 -
2097 def get_fused_moe_quant_config(
2098 self, layer: RoutedExperts
2099 ) -> FusedMoEQuantConfig | None:
vllm/model_executor/layers/quantization/mxfp4.py
-20
@@ -430,16 +430,6 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
430 layer=layer,
431 )
432
433 - def select_gemm_impl(
434 - self,
435 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,
436 - layer: RoutedExperts,
437 - ) -> mk.FusedMoEExpertsModular:
438 - raise ValueError(
439 - f"{self.__class__.__name__} uses the new modular kernel "
440 - "initialization logic. This function should not be called."
441 - )
442 -
433 def apply(
434 self,
435 layer: RoutedExperts,
@@ -893,16 +883,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
883 layer=layer,
884 )
885
896 - def select_gemm_impl(
897 - self,
898 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,
899 - layer: RoutedExperts,
900 - ) -> mk.FusedMoEExpertsModular:
901 - raise ValueError(
902 - f"{self.__class__.__name__} uses the new modular kernel "
903 - "initialization logic. This function should not be called."
904 - )
905 -
886 def apply(
887 self,
888 layer: RoutedExperts,
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+19 -10
@@ -12,6 +12,7 @@ from vllm import _custom_ops as ops
12 from vllm.distributed.utils import verify_group_size_divides_partition
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.fused_moe import RoutedExperts
15 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
16 from vllm.model_executor.layers.linear import LinearBase
17 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
18 from vllm.model_executor.layers.quantization.utils.int8_utils import (
@@ -351,10 +352,12 @@ def marlin_moe_padded_intermediate(intermediate_size: int, group_size: int = -1)
352 return padded
353
354
354 -def check_moe_marlin_supports_layer(
355 - layer: RoutedExperts, group_size: int, allow_tile_padding: bool = False
355 +def check_moe_marlin_supports_config(
356 + config: FusedMoEConfig,
357 + group_size: int,
358 + allow_tile_padding: bool = False,
359 ) -> bool:
357 - """Whether the fused MoE Marlin kernel supports ``layer``.
360 + """Whether the fused MoE Marlin kernel supports ``config``.
361
362 Callers without act-order may pass ``allow_tile_padding=True``: a
363 tile-misaligned intermediate size is then zero-padded to a valid thread
@@ -364,15 +367,11 @@ def check_moe_marlin_supports_layer(
367 """
368 if current_platform.is_rocm():
369 return False
367 - hidden_size = layer.hidden_size
370 + hidden_size = config.hidden_dim
371 # The layer has not rounded intermediate_size yet; use the stable unpadded
372 # size. gate-up needs n=2*intermediate % 128, down needs k=intermediate % 64.
370 - intermediate_size_per_partition = (
371 - layer.moe_config.intermediate_size_per_partition_unpadded
372 - )
373 + intermediate_size_per_partition = config.intermediate_size_per_partition_unpadded
374 assert intermediate_size_per_partition is not None
374 - # apply_router_weight_on_input is not supported for moe marlin
375 - supports_router_weight = not layer.apply_router_weight_on_input
375
376 if allow_tile_padding:
377 supports_shape = hidden_size % 128 == 0 and (
@@ -384,7 +383,17 @@ def check_moe_marlin_supports_layer(
383 and intermediate_size_per_partition % max(64, group_size) == 0
384 )
385 supports_group_size = group_size in [-1, 32, 64, 128]
387 - return supports_shape and supports_group_size and supports_router_weight
386 + return supports_shape and supports_group_size
387 +
388 +
389 +def check_moe_marlin_supports_layer(
390 + layer: RoutedExperts,
391 + group_size: int,
392 + allow_tile_padding: bool = False,
393 +) -> bool:
394 + return check_moe_marlin_supports_config(
395 + layer.moe_config, group_size, allow_tile_padding
396 + )
397
398
399 def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tensor):