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

[Bugfix][Quantization] Add Humming to the W4A8 (INT4xFP8) MoE oracle (#58427)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Misha Goin
2026-09-24 12:37:33 -07:00
committed by GitHub
parent e773247f8a
commit cd06e81c41
5 changed files with 163 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+3 -1
View File
@@ -1326,7 +1326,9 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
@staticmethod
def _supports_current_device() -> bool:
return cutlass_group_gemm_supported()
return (
current_platform.is_device_capability(90) and cutlass_group_gemm_supported()
)
@staticmethod
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/oracle/w4a8.py
+120 -36
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Oracle for W4A8 (INT4, FP8) MoE
from enum import Enum
from typing import TYPE_CHECKING
@@ -22,26 +23,37 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
)
if TYPE_CHECKING:
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
CutlassExpertsW4A8Fp8,
)
from vllm.model_executor.layers.fused_moe import RoutedExperts
logger = init_logger(__name__)
class W4A8MoeBackend(Enum):
CUTLASS = "CUTLASS"
HUMMING = "HUMMING"
def backend_to_kernel_cls(
backend: W4A8MoeBackend,
) -> list[type["CutlassExpertsW4A8Fp8"]]:
) -> list[type[mk.FusedMoEExperts]]:
if backend == W4A8MoeBackend.CUTLASS:
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
CutlassExpertsW4A8Fp8,
)
return [CutlassExpertsW4A8Fp8]
elif backend == W4A8MoeBackend.HUMMING:
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
BatchedHummingGroupedExperts,
HummingGroupedExperts,
HummingIndexedExperts,
)
return [
BatchedHummingGroupedExperts,
HummingGroupedExperts,
HummingIndexedExperts,
]
else:
raise ValueError(f"Unknown W4A8 MoE backend: {backend.value}")
@@ -50,8 +62,19 @@ def select_w4a8_moe_backend(
config: FusedMoEConfig,
weight_key: QuantKey | None = kInt4Static,
activation_key: QuantKey | None = kFp8DynamicTokenSym,
) -> tuple[W4A8MoeBackend, type["CutlassExpertsW4A8Fp8"]]:
backend = W4A8MoeBackend.CUTLASS
) -> tuple[W4A8MoeBackend, type[mk.FusedMoEExperts]]:
backends = [W4A8MoeBackend.CUTLASS, W4A8MoeBackend.HUMMING]
if config.moe_backend != "auto":
backend_map = {
"cutlass": W4A8MoeBackend.CUTLASS,
"humming": W4A8MoeBackend.HUMMING,
}
if config.moe_backend not in backend_map:
raise ValueError(
f"moe_backend='{config.moe_backend}' is not supported for W4A8 MoE. "
f"Expected one of {list(backend_map)}."
)
backends = [backend_map[config.moe_backend]]
activation_format = (
mk.FusedMoEActivationFormat.BatchedExperts
@@ -59,29 +82,33 @@ def select_w4a8_moe_backend(
else mk.FusedMoEActivationFormat.Standard
)
last_reason: str | None = None
for kernel_cls in backend_to_kernel_cls(backend):
supported, reason = kernel_cls.is_supported_config(
kernel_cls,
config,
weight_key,
activation_key,
activation_format,
)
if supported:
logger.info_once("Using %s W4A8 MoE backend.", backend.value)
return backend, kernel_cls
last_reason = reason
reasons = []
for backend in backends:
for kernel_cls in backend_to_kernel_cls(backend):
supported, reason = kernel_cls.is_supported_config(
kernel_cls,
config,
weight_key,
activation_key,
activation_format,
)
if supported:
logger.info_once("Using %s W4A8 MoE backend.", backend.value)
return backend, kernel_cls
reasons.append(f"{kernel_cls.__name__}: {reason}")
raise NotImplementedError(
f"W4A8 MoE backend {backend.value} does not support the "
f"deployment configuration: {last_reason}."
"No W4A8 MoE backend supports the deployment configuration: "
+ "; ".join(reasons)
)
def convert_to_w4a8_moe_kernel_format(
w13_weight_packed: torch.Tensor,
w2_weight_packed: torch.Tensor,
backend: W4A8MoeBackend,
layer: "RoutedExperts",
group_size: int,
w13_weight: torch.Tensor,
w2_weight: torch.Tensor,
w13_weight_scale: torch.Tensor,
w2_weight_scale: torch.Tensor,
) -> tuple[
@@ -89,11 +116,40 @@ def convert_to_w4a8_moe_kernel_format(
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
if backend == W4A8MoeBackend.HUMMING:
from vllm.model_executor.layers.quantization.utils.humming import (
convert_to_humming_moe_kernel_format,
)
from vllm.utils.humming import (
HummingInputSchema,
HummingWeightSchema,
dtypes,
)
convert_to_humming_moe_kernel_format(
layer,
weight_schema=HummingWeightSchema(
b_dtype=dtypes.uint4, weight_scale_group_size=group_size
),
input_schema=HummingInputSchema(a_dtype=dtypes.float8e4m3),
allow_input_schema_fallback=False,
)
return (
layer.w13_weight,
layer.w2_weight,
layer.w13_weight_scale,
layer.w2_weight_scale,
None,
None,
None,
None,
)
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -104,18 +160,18 @@ def convert_to_w4a8_moe_kernel_format(
quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
convert_packed_uint4b8_to_signed_int4_inplace(w13_weight_packed)
convert_packed_uint4b8_to_signed_int4_inplace(w13_weight)
# Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
torch.accelerator.synchronize()
w13_weight_shuffled, b_strides1 = ops.cutlass_encode_and_reorder_int4b_grouped(
w13_weight_packed
w13_weight
)
convert_packed_uint4b8_to_signed_int4_inplace(w2_weight_packed)
convert_packed_uint4b8_to_signed_int4_inplace(w2_weight)
# Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
torch.accelerator.synchronize()
w2_weight_shuffled, b_strides2 = ops.cutlass_encode_and_reorder_int4b_grouped(
w2_weight_packed
w2_weight
)
w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
@@ -147,11 +203,23 @@ def convert_to_w4a8_moe_kernel_format(
def make_w4a8_moe_quant_config(
backend: W4A8MoeBackend,
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
g1_alphas: torch.Tensor,
g2_alphas: torch.Tensor,
g1_alphas: torch.Tensor | None,
g2_alphas: torch.Tensor | None,
layer: torch.nn.Module | None = None,
) -> FusedMoEQuantConfig:
if backend == W4A8MoeBackend.HUMMING:
from vllm.model_executor.layers.fused_moe import RoutedExperts
from vllm.model_executor.layers.quantization.utils.humming import (
get_humming_moe_quant_config,
)
assert isinstance(layer, RoutedExperts)
return get_humming_moe_quant_config(layer)
assert g1_alphas is not None and g2_alphas is not None
return int4_w4afp8_moe_quant_config(
w1_scale=w1_scale,
w2_scale=w2_scale,
@@ -163,14 +231,30 @@ def make_w4a8_moe_quant_config(
def make_w4a8_moe_kernel(
backend: W4A8MoeBackend,
moe_quant_config: FusedMoEQuantConfig,
moe_config: FusedMoEConfig,
experts_cls: type["CutlassExpertsW4A8Fp8"],
b_strides1: torch.Tensor,
b_strides2: torch.Tensor,
experts_cls: type[mk.FusedMoEExperts],
b_strides1: torch.Tensor | None,
b_strides2: torch.Tensor | None,
group_size: int,
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEKernel:
if backend == W4A8MoeBackend.HUMMING:
from vllm.model_executor.layers.quantization.utils.humming import (
make_humming_moe_kernel,
)
return make_humming_moe_kernel(
moe_quant_config, moe_config, experts_cls, routing_tables=routing_tables
)
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
CutlassExpertsW4A8Fp8,
)
assert issubclass(experts_cls, CutlassExpertsW4A8Fp8)
assert b_strides1 is not None and b_strides2 is not None
prepare_finalize = maybe_make_prepare_finalize(
moe=moe_config,
quant_config=moe_quant_config,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+12
View File
@@ -511,6 +511,13 @@ class CompressedTensorsConfig(QuantizationConfig):
def _is_dynamic_token_w4a8_int(
weight_quant: QuantizationArgs, input_quant: QuantizationArgs
) -> bool:
if not weight_quant or not input_quant:
return False
if (
weight_quant.type != QuantizationType.INT
or input_quant.type != QuantizationType.INT
):
return False
is_weight_4_bits = weight_quant.num_bits == 4
is_activation_8_bits = input_quant.num_bits == 8
weight_strategy = (
@@ -575,6 +582,11 @@ class CompressedTensorsConfig(QuantizationConfig):
) -> bool:
if not weight_quant or not input_quant:
return False
if (
weight_quant.type != QuantizationType.INT
or input_quant.type != QuantizationType.FLOAT
):
return False
is_weight_4_bits = weight_quant.num_bits == 4
is_activation_8_bits = input_quant.num_bits == 8
weight_strategy = weight_quant.strategy == QuantizationStrategy.GROUP.value
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+1 -1
View File
@@ -164,7 +164,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
return CompressedTensorsW8A8Int8MoEMethod(
weight_quant, input_quant, layer.moe_config
)
elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
elif quant_config._is_fp8_w4a8(weight_quant, input_quant):
from .compressed_tensors_moe_w4a8_fp8 import (
CompressedTensorsW4A8Fp8MoEMethod,
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
+27 -33
View File
@@ -143,57 +143,49 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
layer.register_parameter("w13_weight_shape", w13_weight_shape)
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
w13_weight_chan_scale = torch.nn.Parameter(
torch.ones(
num_experts,
self.moe.w13_num_shards * intermediate_size_per_partition,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_chan_scale", w13_weight_chan_scale)
w2_weight_chan_scale = torch.nn.Parameter(
torch.ones(num_experts, hidden_size, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_chan_scale", w2_weight_chan_scale)
# don't use input scales
layer.w13_input_scale = None
layer.w2_input_scale = None
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
for name in ("w13", "w2"):
setattr(layer, f"{name}_weight", getattr(layer, f"{name}_weight_packed"))
delattr(layer, f"{name}_weight_packed")
(
w13_weight_packed,
w2_weight_packed,
w13_weight,
w2_weight,
w13_weight_scale,
w2_weight_scale,
w13_weight_chan_scale,
w2_weight_chan_scale,
b_strides1,
b_strides2,
self.b_strides1,
self.b_strides2,
) = convert_to_w4a8_moe_kernel_format(
w13_weight_packed=layer.w13_weight_packed,
w2_weight_packed=layer.w2_weight_packed,
backend=self.w4a8_backend,
layer=layer,
group_size=self.group_size,
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
)
replace_parameter(layer, "w13_weight_packed", w13_weight_packed)
replace_parameter(layer, "w2_weight_packed", w2_weight_packed)
replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
replace_parameter(layer, "w13_weight_chan_scale", w13_weight_chan_scale)
replace_parameter(layer, "w2_weight_chan_scale", w2_weight_chan_scale)
self.b_strides1 = b_strides1
self.b_strides2 = b_strides2
for name, tensor in (
("w13_weight", w13_weight),
("w2_weight", w2_weight),
("w13_weight_scale", w13_weight_scale),
("w2_weight_scale", w2_weight_scale),
("w13_weight_chan_scale", w13_weight_chan_scale),
("w2_weight_chan_scale", w2_weight_chan_scale),
):
replace_parameter(layer, name, tensor)
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
if self.moe_quant_config is not None:
assert self.experts_cls is not None
self.moe_kernel = make_w4a8_moe_kernel(
backend=self.w4a8_backend,
moe_quant_config=self.moe_quant_config,
moe_config=self.moe,
experts_cls=self.experts_cls,
@@ -205,10 +197,12 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
return make_w4a8_moe_quant_config(
backend=self.w4a8_backend,
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
g1_alphas=layer.w13_weight_chan_scale,
g2_alphas=layer.w2_weight_chan_scale,
layer=layer,
)
def apply(
@@ -224,8 +218,8 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
assert self.moe_kernel is not None
return self.moe_kernel.apply(
hidden_states=x,
w1=layer.w13_weight_packed,
w2=layer.w2_weight_packed,
w1=layer.w13_weight,
w2=layer.w2_weight,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=layer.activation,