Fix static actorder handling for compressed-tensors WNA16 MoE (#41161)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
ZewenShen-Cohere committed
Jun 22, 2026 at 18:46 UTC
ca5b24695bfae510ccfa30bbbe1dfc006ef49711
3 files changed
+92
-9
tests/quantization/test_compressed_tensors.py
+30
@@ -10,6 +10,7 @@ from unittest.mock import Mock
10
import pytest
11
import torch
12
from compressed_tensors.quantization import (
13
+ ActivationOrdering,
14
QuantizationArgs,
15
QuantizationStrategy,
16
QuantizationType,
@@ -683,6 +684,35 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
684
assert output
685
686
687
+@pytest.mark.parametrize(
688
+ "actorder,group_size,part,full,expected",
689
+ [
690
+ # actorder="group" with real grouping: must load full-K w2 scales and,
691
+ # when sharded (part != full), report is_k_full=False.
692
+ (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)),
693
+ # actorder="group" but unsharded (part == full): full scales, k_full.
694
+ (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)),
695
+ # actorder="group" with channel-wise (group_size == -1): no full load.
696
+ (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)),
697
+ # "static"/"weight" reorder at quant time -> shard normally + k_full.
698
+ # Regression: static actorder under TP must keep is_k_full=True so the
699
+ # Marlin kernel never gets the invalid (group_size=16, is_k_full=0).
700
+ ("static", 32, 64, 128, (False, 64, True)),
701
+ ("weight", 32, 64, 128, (False, 64, True)),
702
+ (None, 32, 64, 128, (False, 64, True)),
703
+ ],
704
+)
705
+def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected):
706
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501
707
+ CompressedTensorsWNA16MarlinMoEMethod,
708
+ )
709
+
710
+ result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding(
711
+ actorder, group_size, part, full
712
+ )
713
+ assert result == expected
714
+
715
+
716
@pytest.mark.skipif(
717
not current_platform.is_cuda() or not current_platform.has_device_capability(80),
718
reason="MXFP4 requires ampere or newer",
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+15
@@ -97,6 +97,21 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
97
num_groups_w2 = num_groups_w13 = 1
98
self.group_size = -1
99
else:
100
+ if hidden_size % self.group_size != 0:
101
+ raise ValueError(
102
+ "CompressedTensors WNA16 MoE requires hidden_size "
103
+ f"({hidden_size}) to be divisible by group_size "
104
+ f"({self.group_size})."
105
+ )
106
+ if intermediate_size_per_partition % self.group_size != 0:
107
+ raise ValueError(
108
+ "CompressedTensors WNA16 MoE with static group scales "
109
+ "requires the MoE intermediate size per tensor-parallel "
110
+ f"partition ({intermediate_size_per_partition}) to be "
111
+ f"divisible by group_size ({self.group_size}). Scale "
112
+ "groups would otherwise cross TP shard boundaries; use a "
113
+ "compatible TP size or enable expert parallelism."
114
+ )
115
num_groups_w2 = w2_scales_size // self.group_size
116
num_groups_w13 = hidden_size // self.group_size
117
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+47
-9
@@ -181,6 +181,29 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
181
backend_key = "Flashinfer" if is_flashinfer else "Marlin"
182
return shape_map[weight_name][backend_key]
183
184
+ @staticmethod
185
+ def _w2_scale_sharding(
186
+ actorder,
187
+ group_size: int,
188
+ intermediate_size_per_partition: int,
189
+ intermediate_size_full: int,
190
+ ) -> tuple[bool, int, bool]:
191
+ """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE.
192
+
193
+ Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime
194
+ and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``.
195
+ ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at
196
+ quantization time, so scales shard normally per TP rank.
197
+ """
198
+ load_full_w2 = (actorder == "group") and group_size != -1
199
+ w2_scales_size = (
200
+ intermediate_size_full if load_full_w2 else intermediate_size_per_partition
201
+ )
202
+ is_k_full = (actorder != "group") or (
203
+ intermediate_size_per_partition == intermediate_size_full
204
+ )
205
+ return load_full_w2, w2_scales_size, is_k_full
206
+
207
def create_weights(
208
self,
209
layer: torch.nn.Module,
@@ -230,21 +253,36 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
253
layer.register_parameter("w2_weight_packed", w2_weight)
254
set_weight_attrs(w2_weight, extra_weight_attrs)
255
233
- # In the case where we have actorder/g_idx,
234
- # we do not partition the w2 scales
235
- load_full_w2 = self.actorder and self.group_size != -1
236
- w2_scales_size = (
237
- intermediate_size_full if load_full_w2 else intermediate_size_per_partition
238
- )
239
-
240
- self.is_k_full = (not self.actorder) or (
241
- intermediate_size_per_partition == intermediate_size_full
256
+ load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding(
257
+ self.actorder,
258
+ self.group_size,
259
+ intermediate_size_per_partition,
260
+ intermediate_size_full,
261
)
262
263
if self.strategy == "channel":
264
num_groups_w2 = num_groups_w13 = 1
265
self.group_size = -1
266
else:
267
+ if hidden_size % self.group_size != 0:
268
+ raise ValueError(
269
+ "CompressedTensors WNA16 Marlin MoE requires hidden_size "
270
+ f"({hidden_size}) to be divisible by group_size "
271
+ f"({self.group_size})."
272
+ )
273
+ if (
274
+ not load_full_w2
275
+ and intermediate_size_per_partition % self.group_size != 0
276
+ ):
277
+ raise ValueError(
278
+ "CompressedTensors WNA16 Marlin MoE with static group "
279
+ "scales requires the MoE intermediate size per "
280
+ "tensor-parallel partition "
281
+ f"({intermediate_size_per_partition}) to be divisible by "
282
+ f"group_size ({self.group_size}). Scale groups would "
283
+ "otherwise cross TP shard boundaries; use a compatible TP "
284
+ "size or enable expert parallelism."
285
+ )
286
num_groups_w2 = w2_scales_size // self.group_size
287
num_groups_w13 = hidden_size // self.group_size
288