diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 2620b679b6e..de906a861b7 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -10,6 +10,7 @@ from unittest.mock import Mock import pytest import torch from compressed_tensors.quantization import ( + ActivationOrdering, QuantizationArgs, QuantizationStrategy, QuantizationType, @@ -683,6 +684,35 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner): assert output +@pytest.mark.parametrize( + "actorder,group_size,part,full,expected", + [ + # actorder="group" with real grouping: must load full-K w2 scales and, + # when sharded (part != full), report is_k_full=False. + (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)), + # actorder="group" but unsharded (part == full): full scales, k_full. + (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)), + # actorder="group" with channel-wise (group_size == -1): no full load. + (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)), + # "static"/"weight" reorder at quant time -> shard normally + k_full. + # Regression: static actorder under TP must keep is_k_full=True so the + # Marlin kernel never gets the invalid (group_size=16, is_k_full=0). + ("static", 32, 64, 128, (False, 64, True)), + ("weight", 32, 64, 128, (False, 64, True)), + (None, 32, 64, 128, (False, 64, True)), + ], +) +def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected): + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501 + CompressedTensorsWNA16MarlinMoEMethod, + ) + + result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding( + actorder, group_size, part, full + ) + assert result == expected + + @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(80), reason="MXFP4 requires ampere or newer", diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 88303f189f5..cfeacc902f4 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -97,6 +97,21 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): num_groups_w2 = num_groups_w13 = 1 self.group_size = -1 else: + if hidden_size % self.group_size != 0: + raise ValueError( + "CompressedTensors WNA16 MoE requires hidden_size " + f"({hidden_size}) to be divisible by group_size " + f"({self.group_size})." + ) + if intermediate_size_per_partition % self.group_size != 0: + raise ValueError( + "CompressedTensors WNA16 MoE with static group scales " + "requires the MoE intermediate size per tensor-parallel " + f"partition ({intermediate_size_per_partition}) to be " + f"divisible by group_size ({self.group_size}). Scale " + "groups would otherwise cross TP shard boundaries; use a " + "compatible TP size or enable expert parallelism." + ) num_groups_w2 = w2_scales_size // self.group_size num_groups_w13 = hidden_size // self.group_size diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 82734103917..0401a5b6e73 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -181,6 +181,29 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): backend_key = "Flashinfer" if is_flashinfer else "Marlin" return shape_map[weight_name][backend_key] + @staticmethod + def _w2_scale_sharding( + actorder, + group_size: int, + intermediate_size_per_partition: int, + intermediate_size_full: int, + ) -> tuple[bool, int, bool]: + """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE. + + Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime + and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``. + ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at + quantization time, so scales shard normally per TP rank. + """ + load_full_w2 = (actorder == "group") and group_size != -1 + w2_scales_size = ( + intermediate_size_full if load_full_w2 else intermediate_size_per_partition + ) + is_k_full = (actorder != "group") or ( + intermediate_size_per_partition == intermediate_size_full + ) + return load_full_w2, w2_scales_size, is_k_full + def create_weights( self, layer: torch.nn.Module, @@ -230,21 +253,36 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): layer.register_parameter("w2_weight_packed", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) - # In the case where we have actorder/g_idx, - # we do not partition the w2 scales - load_full_w2 = self.actorder and self.group_size != -1 - w2_scales_size = ( - intermediate_size_full if load_full_w2 else intermediate_size_per_partition - ) - - self.is_k_full = (not self.actorder) or ( - intermediate_size_per_partition == intermediate_size_full + load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding( + self.actorder, + self.group_size, + intermediate_size_per_partition, + intermediate_size_full, ) if self.strategy == "channel": num_groups_w2 = num_groups_w13 = 1 self.group_size = -1 else: + if hidden_size % self.group_size != 0: + raise ValueError( + "CompressedTensors WNA16 Marlin MoE requires hidden_size " + f"({hidden_size}) to be divisible by group_size " + f"({self.group_size})." + ) + if ( + not load_full_w2 + and intermediate_size_per_partition % self.group_size != 0 + ): + raise ValueError( + "CompressedTensors WNA16 Marlin MoE with static group " + "scales requires the MoE intermediate size per " + "tensor-parallel partition " + f"({intermediate_size_per_partition}) to be divisible by " + f"group_size ({self.group_size}). Scale groups would " + "otherwise cross TP shard boundaries; use a compatible TP " + "size or enable expert parallelism." + ) num_groups_w2 = w2_scales_size // self.group_size num_groups_w13 = hidden_size // self.group_size