From 54dd98be2830bbd7b70fd60a0e917b7d0df5adab Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Thu, 20 Aug 2026 00:31:32 +0800 Subject: [PATCH] [CT] Support Humming for WNA16 MoE (#48918) Signed-off-by: yiliu30 Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Misha Goin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/quantization/test_compressed_tensors.py | 44 +++++++ tests/quantization/test_moe_wna16.py | 110 ++++++++++++++++++ .../fused_moe/experts/fused_humming_moe.py | 32 ++++- .../layers/fused_moe/oracle/int_wna16.py | 16 ++- .../compressed_tensors_moe_wna16.py | 56 ++++++--- .../quantization/utils/humming_utils.py | 3 + .../layers/quantization/utils/quant_utils.py | 16 +-- 7 files changed, 253 insertions(+), 24 deletions(-) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 3f11cec1c1..a301e8d9cf 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -43,8 +43,14 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( find_matched_target, ) from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + QuantKey, + ScaleDesc, +) from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types from vllm.v1.attention.backends.fa_utils import get_flash_attn_version # AITER only supports per-channel-per-channel INT8 gemm @@ -971,6 +977,44 @@ def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected) assert result == expected +@pytest.mark.parametrize("num_bits", range(2, 9)) +def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits): + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingExpertsBase, + ) + from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa: E501 + WNA16_SUPPORTED_TYPES_MAP, + ) + + weight_key = QuantKey( + dtype=WNA16_SUPPORTED_TYPES_MAP[num_bits], + scale=ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=GroupShape(row=1, col=128), + ), + symmetric=True, + ) + + assert HummingExpertsBase._supports_quant_scheme(weight_key, None) + + +def test_quant_key_str_supports_scalar_type_dtypes(): + quant_key = QuantKey( + dtype=scalar_types.uint2b2, + scale=ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=GroupShape(row=1, col=128), + ), + symmetric=True, + ) + + assert str(quant_key) == ( + "QuantKey(uint2b2,scale(f16,static,GroupShape(row=1, col=128)),symmetric)" + ) + + @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/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 8624832e7b..49e224fef4 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.moe_wna16 import ( MoeWNA16Config, MoeWNA16Method, ) +from vllm.platforms import current_platform def test_map_wna16_backend_supports_triton(): @@ -203,3 +204,112 @@ def test_moe_wna16_uses_humming_quant_config(monkeypatch): ) assert method.get_fused_moe_quant_config(layer) is quant_config + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Compressed-tensors Humming WNA16 MoE requires CUDA", +) +@pytest.mark.parametrize("num_bits", [3, 5, 6, 7]) +def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes( + num_bits, +): + pytest.importorskip("humming") + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, + ) + + quant_args = QuantizationArgs( + num_bits=num_bits, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + ) + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=256, + intermediate_size=512, + ) + moe_config.moe_backend = "humming" + method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) + layer = torch.nn.Module() + + method.create_weights( + layer, + num_experts=2, + hidden_size=256, + intermediate_size_per_partition=512, + params_dtype=torch.float16, + intermediate_size_full=512, + ) + + packed_hidden = (256 * num_bits + 31) // 32 + packed_intermediate = (512 * num_bits + 31) // 32 + assert method.wna16_backend == WNA16MoEBackend.HUMMING + assert layer.w13_weight_packed.shape == (2, 1024, packed_hidden) + assert layer.w2_weight_packed.shape == (2, 256, packed_intermediate) + assert layer.w13_weight_scale.shape == (2, 1024, 2) + assert layer.w2_weight_scale.shape == (2, 256, 4) + assert layer.w13_weight_packed.dtype is torch.int32 + assert layer.w2_weight_scale.dtype is torch.float16 + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Compressed-tensors Humming WNA16 MoE requires CUDA", +) +def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel(): + pytest.importorskip("humming") + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, + ) + + quant_args = QuantizationArgs( + num_bits=3, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + ) + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=256, + intermediate_size=512, + ) + moe_config.moe_backend = "humming" + method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) + layer = torch.nn.Module() + layer.moe_config = moe_config + layer.params_dtype = torch.bfloat16 + layer.layer_name = "test.humming_moe" + layer._expert_routing_tables = lambda: (None, None, None) + + method.create_weights( + layer, + num_experts=2, + hidden_size=256, + intermediate_size_per_partition=512, + intermediate_size_full=512, + params_dtype=torch.bfloat16, + ) + layer.cuda() + for parameter in layer.parameters(): + parameter.data.zero_() + + method.process_weights_after_loading(layer) + + assert method.wna16_backend == WNA16MoEBackend.HUMMING + assert method.moe_kernel is not None + assert set(layer.weight_schemas) == {"w13", "w2"} + assert set(layer.humming_configs) == {"w13", "w2"} + assert not hasattr(layer, "w13_weight_packed") + assert not hasattr(layer, "w2_weight_packed") + assert layer.w13_weight.dtype is torch.int32 + assert layer.w2_weight.dtype is torch.int32 diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 96c42967e9..69843c228a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -40,6 +40,9 @@ from vllm.model_executor.layers.fused_moe.utils import ( swiglu_limit_func, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + INT4_DTYPE, + INT8_DTYPE, + GroupShape, QuantKey, kFp8Dynamic128Sym, kFp8DynamicTokenSym, @@ -58,6 +61,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Static, ) from vllm.platforms import current_platform +from vllm.scalar_type import ScalarType from vllm.utils.import_utils import has_humming from vllm.v1.worker.workspace import current_workspace_manager @@ -76,6 +80,29 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _is_supported_wna16_weight_key(weight_key: QuantKey | None) -> bool: + if weight_key is None or weight_key.scale2 is not None: + return False + + group_shape = weight_key.scale.group_shape + if not ( + group_shape == GroupShape.PER_CHANNEL + or (group_shape.row == 1 and group_shape.col > 0) + ): + return False + + dtype = weight_key.dtype + if dtype in (INT4_DTYPE, INT8_DTYPE, torch.uint8): + return True + + return ( + isinstance(dtype, ScalarType) + and dtype.is_integer() + and not dtype.is_signed() + and 2 <= dtype.size_bits <= 8 + ) + + def get_humming_moe_gemm_type() -> str: env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE gemm_type = "indexed" @@ -267,7 +294,10 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): # mxfp8 (compressed-tensors / modelopt / online) (kMxfp8Static, kMxfp8Dynamic), ] - return (weight_key, activation_key) in SUPPORTED_W_A + return (weight_key, activation_key) in SUPPORTED_W_A or ( + activation_key in (None, kFp8DynamicTokenSym) + and _is_supported_wna16_weight_key(weight_key) + ) @property def expects_unquantized_inputs(self) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 5e4f30b9d1..c56712ac58 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1104,9 +1104,21 @@ def _humming_wna16_weight_schema( "desc_act": quant_config.desc_act, "sym": quant_config.is_sym, } + if isinstance(quant_config, QuantizationArgs): + quant_type = getattr(quant_config.type, "value", quant_config.type) + quant_strategy = getattr(quant_config.strategy, "value", quant_config.strategy) + return { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "type": str(quant_type), + "num_bits": quant_config.num_bits, + "strategy": str(quant_strategy), + "group_size": quant_config.group_size, + "symmetric": quant_config.symmetric, + } raise TypeError( - "Humming WNA16 checkpoint schema requires AutoAWQConfig or " - "AutoGPTQConfig, " + "Humming WNA16 checkpoint schema requires AutoAWQConfig, " + "AutoGPTQConfig or QuantizationArgs, " f"got {type(quant_config).__name__}." ) 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 35894b3224..1ab5085145 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 @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math +from fractions import Fraction from typing import Any import torch @@ -40,7 +42,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_make_workspace_new, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, QuantKey, + ScaleDesc, kInt4Static32GroupScale, kInt4StaticGroupScale, kInt8StaticGroupScale, @@ -64,7 +68,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): # Extract properties from weight_quant self.symmetric = weight_quant.symmetric self.num_bits = weight_quant.num_bits - self.packed_factor = 32 // weight_quant.num_bits + self.packed_factor = Fraction(32, weight_quant.num_bits) self.strategy = weight_quant.strategy self.group_size = weight_quant.group_size self.actorder = weight_quant.actorder @@ -84,8 +88,14 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): elif self.num_bits == 8: scale = kInt8StaticGroupScale else: - raise ValueError( - "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now." + scale = ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=( + GroupShape.PER_CHANNEL + if self.group_size == -1 + else GroupShape(row=1, col=self.group_size) + ), ) weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) @@ -109,7 +119,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] - self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM + self.is_transposed = self.wna16_backend not in ( + WNA16MoEBackend.FLASHINFER_TRTLLM, + WNA16MoEBackend.HUMMING, + ) if self.is_marlin: assert check_moe_marlin_supports_config( @@ -125,6 +138,9 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): # Non-Marlin WNA16 always uses bf16/fp16 inputs self.input_dtype = torch.bfloat16 + def _packed_dim(self, dim: int) -> int: + return math.ceil(dim * self.num_bits / 32) + def get_weight_shape( self, weight_name: str, @@ -149,16 +165,16 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): "num_groups_w2 must be provided for weight scales/zero_points" ) w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - shape_map = { + shape_map: dict[str, dict[str, tuple[int, int | None, int | None]]] = { "w13_weight": { "Flashinfer": ( num_experts, w13_num_shards * intermediate_size_per_partition, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), ), "Marlin": ( num_experts, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), w13_num_shards * intermediate_size_per_partition, ), }, @@ -178,20 +194,18 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): "Marlin": ( num_experts, num_groups_w13, - w13_num_shards - * intermediate_size_per_partition - // self.packed_factor, + self._packed_dim(w13_num_shards * intermediate_size_per_partition), ), }, "w2_weight": { "Flashinfer": ( num_experts, hidden_size, - intermediate_size_per_partition // self.packed_factor, + self._packed_dim(intermediate_size_per_partition), ), "Marlin": ( num_experts, - intermediate_size_per_partition // self.packed_factor, + self._packed_dim(intermediate_size_per_partition), hidden_size, ), }, @@ -203,12 +217,14 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): "Marlin": ( num_experts, num_groups_w2, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), ), }, } backend_key = "Marlin" if self.is_transposed else "Flashinfer" - return shape_map[weight_name][backend_key] + shape = shape_map[weight_name][backend_key] + assert shape[1] is not None and shape[2] is not None + return shape[0], shape[1], shape[2] @staticmethod def _w2_scale_sharding( @@ -483,6 +499,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_backend, routing_tables=layer._expert_routing_tables(), **marlin_args, ) @@ -586,6 +603,17 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: + if self.wna16_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config( + layer, + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + ) return make_wna16_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 7447cbef12..1303a92d91 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -73,6 +73,9 @@ if has_humming(): humming_dtypes.uint8: INT8_DTYPE, humming_dtypes.uint2: torch.uint8, humming_dtypes.uint3: torch.uint8, + humming_dtypes.uint5: torch.uint8, + humming_dtypes.uint6: torch.uint8, + humming_dtypes.uint7: torch.uint8, } _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 7ec658a375..20f33be240 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -25,6 +25,13 @@ INT4_DTYPE = scalar_types.uint4b8 INT8_DTYPE = scalar_types.uint8b128 +def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str: + """Return a stable short name for torch and ScalarType dtypes.""" + if isinstance(dtype, ScalarType): + return str(dtype) + return fx.graph.dtype_abbrs[dtype] + + def weight_amax( weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False ) -> torch.Tensor: @@ -152,7 +159,7 @@ class ScaleDesc: group_shape = d.get(self.group_shape, str(self.group_shape)) return ( - f"{fx.graph.dtype_abbrs[self.dtype]}," + f"{_dtype_abbr(self.dtype)}," f"{'static' if self.static else 'dynamic'},{group_shape}" ) @@ -178,13 +185,8 @@ class QuantKey: def __str__(self): scale2_str = f"scale2({self.scale2})," if self.scale2 else "" - dtype_description = ( - fx.graph.dtype_abbrs[self.dtype] - if isinstance(self.dtype, torch.dtype) - else self.dtype - ) return ( - f"QuantKey({dtype_description}," + f"QuantKey({_dtype_abbr(self.dtype)}," f"scale({self.scale}),{scale2_str}" f"{'a' if not self.symmetric else ''}symmetric)" )