diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index cd7da3e1e3..a03651c0a9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -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: diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8.py index d8e026ba31..5df5b099af 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8.py @@ -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, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index a4cd5d4d13..cb69a1d665 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -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 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index f2d2a9c945..d965127b81 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -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, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py index 8f98cb4813..4a2223aab8 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py @@ -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,