iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CT] Support Humming for WNA16 MoE (#48918)

Browse Source 
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Yi Liu
2026-08-19 12:31:32 -04:00
committed by GitHub
co-authored by OpenAI Codex OpenAI Codex Misha Goin mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent c676232313
commit 54dd98be28
7 changed files with 253 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/quantization/test_compressed_tensors.py
+44
View File
@@ -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",
tests/quantization/test_moe_wna16.py
+110
View File
@@ -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
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+31 -1
View File
@@ -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:
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+14 -2
View File
@@ -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__}."
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+42 -14
View File
@@ -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,
vllm/model_executor/layers/quantization/utils/humming_utils.py
+3
View File
@@ -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] = {
vllm/model_executor/layers/quantization/utils/quant_utils.py
+9 -7
View File
@@ -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)"
)