| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from enum import Enum |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 10 | from vllm.config.kernel import MoEBackend |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 13 | maybe_make_prepare_finalize, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.fused_moe.config import ( |
| 16 | FusedMoEConfig, |
| 17 | FusedMoEQuantConfig, |
| 18 | int8_w8a8_moe_quant_config, |
| 19 | int8_w8a16_moe_quant_config, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 22 | QuantKey, |
| 23 | kInt8DynamicTokenSym, |
| 24 | kInt8StaticChannelSym, |
| 25 | ) |
| 26 | from vllm.model_executor.utils import replace_parameter |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | class Int8MoeBackend(Enum): |
| 32 | TRITON = "TRITON" |
| 33 | HUMMING = "HUMMING" |
| 34 | CPU = "CPU" |
| 35 | |
| 36 | |
| 37 | def _get_priority_backends( |
| 38 | moe_config: FusedMoEConfig, |
| 39 | ) -> list[Int8MoeBackend]: |
| 40 | """Get available backends in priority order based on platform and config.""" |
| 41 | return [ |
| 42 | Int8MoeBackend.TRITON, |
| 43 | Int8MoeBackend.HUMMING, |
| 44 | Int8MoeBackend.CPU, |
| 45 | ] |
| 46 | |
| 47 | |
| 48 | def backend_to_kernel_cls( |
| 49 | backend: Int8MoeBackend, |
| 50 | ) -> list[type[mk.FusedMoEExperts]]: |
| 51 | if backend == Int8MoeBackend.TRITON: |
| 52 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( |
| 53 | TritonExperts, |
| 54 | ) |
| 55 | |
| 56 | return [TritonExperts] |
| 57 | |
| 58 | elif backend == Int8MoeBackend.HUMMING: |
| 59 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 60 | BatchedHummingGroupedExperts, |
| 61 | HummingGroupedExperts, |
| 62 | HummingIndexedExperts, |
| 63 | ) |
| 64 | |
| 65 | return [ |
| 66 | BatchedHummingGroupedExperts, |
| 67 | HummingGroupedExperts, |
| 68 | HummingIndexedExperts, |
| 69 | ] |
| 70 | elif backend == Int8MoeBackend.CPU: |
| 71 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 72 | ArmCPUExpertsInt8, |
| 73 | CPUExpertsInt8, |
| 74 | PowerCPUExpertsInt8, |
| 75 | ZenCPUExpertsInt8, |
| 76 | ) |
| 77 | |
| 78 | return [ |
| 79 | ZenCPUExpertsInt8, |
| 80 | PowerCPUExpertsInt8, |
| 81 | ArmCPUExpertsInt8, |
| 82 | CPUExpertsInt8, |
| 83 | ] |
| 84 | else: |
| 85 | raise ValueError(f"Unknown Int8 MoE backend: {backend.value}") |
| 86 | |
| 87 | |
| 88 | def map_int8_backend(runner_backend: MoEBackend) -> Int8MoeBackend: |
| 89 | """Map user's MoEBackend to Int8MoeBackend.""" |
| 90 | mapping = { |
| 91 | "triton": Int8MoeBackend.TRITON, |
| 92 | "humming": Int8MoeBackend.HUMMING, |
| 93 | } |
| 94 | if backend := mapping.get(runner_backend): |
| 95 | return backend |
| 96 | raise ValueError( |
| 97 | f"moe_backend='{runner_backend}' is not supported for Int8 MoE. " |
| 98 | f"Expected one of {list(mapping.keys())}." |
| 99 | ) |
| 100 | |
| 101 | |
| 102 | def select_int8_moe_backend( |
| 103 | config: FusedMoEConfig, |
| 104 | weight_key: QuantKey | None = kInt8StaticChannelSym, |
| 105 | activation_key: QuantKey | None = kInt8DynamicTokenSym, |
| 106 | ) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]: |
| 107 | """Select the primary Int8 MoE backend. |
| 108 | Note: Shape-specific fallbacks may still occur at runtime. |
| 109 | """ |
| 110 | AVAILABLE_BACKENDS = _get_priority_backends(config) |
| 111 | |
| 112 | activation_format = ( |
| 113 | mk.FusedMoEActivationFormat.BatchedExperts |
| 114 | if config.moe_parallel_config.use_batched_activation_format |
| 115 | else mk.FusedMoEActivationFormat.Standard |
| 116 | ) |
| 117 | |
| 118 | def _make_log_backend(backend: Int8MoeBackend) -> str: |
| 119 | available_backend_strs = [b.value for b in AVAILABLE_BACKENDS] |
| 120 | return ( |
| 121 | f"Using {backend.value} Int8 MoE backend out " |
| 122 | f"of potential backends: {available_backend_strs}." |
| 123 | ) |
| 124 | |
| 125 | def _make_log_unsupported(backend: Int8MoeBackend, reason: str | None) -> str: |
| 126 | if reason: |
| 127 | return ( |
| 128 | f"Int8 MoE backend {backend.value} does not support the " |
| 129 | f"deployment configuration since {reason}." |
| 130 | ) |
| 131 | else: |
| 132 | return ( |
| 133 | f"Int8 MoE backend '{backend.value}' does not support the " |
| 134 | "deployment configuration." |
| 135 | ) |
| 136 | |
| 137 | def _return_or_raise( |
| 138 | backend: Int8MoeBackend, |
| 139 | ) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]: |
| 140 | for k_cls in backend_to_kernel_cls(backend): |
| 141 | supported, reason = k_cls.is_supported_config( |
| 142 | k_cls, config, weight_key, activation_key, activation_format |
| 143 | ) |
| 144 | if supported: |
| 145 | logger.info_once(_make_log_backend(backend)) |
| 146 | return backend, k_cls |
| 147 | raise ValueError(_make_log_unsupported(backend, reason)) |
| 148 | |
| 149 | # Handle explicit moe_backend from user. |
| 150 | runner_backend = config.moe_backend |
| 151 | if runner_backend != "auto": |
| 152 | requested_backend = map_int8_backend(runner_backend) |
| 153 | return _return_or_raise(requested_backend) |
| 154 | |
| 155 | # Select kernels in order of backend. |
| 156 | for backend in AVAILABLE_BACKENDS: |
| 157 | for k_cls in backend_to_kernel_cls(backend): |
| 158 | supported, reason = k_cls.is_supported_config( |
| 159 | k_cls, |
| 160 | config, |
| 161 | weight_key, |
| 162 | activation_key, |
| 163 | activation_format, |
| 164 | ) |
| 165 | if supported: |
| 166 | logger.info_once(_make_log_backend(backend)) |
| 167 | return backend, k_cls |
| 168 | else: |
| 169 | logger.debug_once(_make_log_unsupported(backend, reason)) |
| 170 | |
| 171 | raise NotImplementedError( |
| 172 | "No Int8 MoE backend supports the deployment configuration " |
| 173 | f"(weight_key={weight_key}, activation_key={activation_key}). " |
| 174 | "Set `VLLM_LOGGING_LEVEL=DEBUG` to see per-backend unsupported reasons." |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | def make_int8_moe_quant_config( |
| 179 | int8_backend: Int8MoeBackend, |
| 180 | w1_scale: torch.Tensor, |
| 181 | w2_scale: torch.Tensor, |
| 182 | a1_scale: torch.Tensor | None = None, |
| 183 | a2_scale: torch.Tensor | None = None, |
| 184 | w1_bias: torch.Tensor | None = None, |
| 185 | w2_bias: torch.Tensor | None = None, |
| 186 | per_act_token_quant: bool = False, |
| 187 | layer: torch.nn.Module | None = None, |
| 188 | ) -> FusedMoEQuantConfig: |
| 189 | if (a1_scale is None) != (a2_scale is None): |
| 190 | raise ValueError("a1_scale and a2_scale must both be provided or both be None") |
| 191 | |
| 192 | scales_absent = a1_scale is None and a2_scale is None |
| 193 | |
| 194 | if int8_backend == Int8MoeBackend.HUMMING: |
| 195 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 196 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 197 | get_humming_moe_quant_config, |
| 198 | ) |
| 199 | |
| 200 | assert isinstance(layer, RoutedExperts) |
| 201 | return get_humming_moe_quant_config( |
| 202 | layer, |
| 203 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 204 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 205 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 206 | ) |
| 207 | |
| 208 | if scales_absent and not per_act_token_quant: |
| 209 | return int8_w8a16_moe_quant_config( |
| 210 | w1_scale=w1_scale, |
| 211 | w2_scale=w2_scale, |
| 212 | w1_zp=None, |
| 213 | w2_zp=None, |
| 214 | w1_bias=w1_bias, |
| 215 | w2_bias=w2_bias, |
| 216 | ) |
| 217 | |
| 218 | return int8_w8a8_moe_quant_config( |
| 219 | w1_scale=w1_scale, |
| 220 | w2_scale=w2_scale, |
| 221 | a1_scale=a1_scale, |
| 222 | a2_scale=a2_scale, |
| 223 | w1_bias=w1_bias, |
| 224 | w2_bias=w2_bias, |
| 225 | per_act_token_quant=per_act_token_quant, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def _humming_int8_weight_schema( |
| 230 | weight: torch.Tensor, weight_scale: torch.Tensor |
| 231 | ) -> dict[str, Any]: |
| 232 | """Build the humming compressed-tensors int8 schema from the canonical |
| 233 | on-device tensors; humming does the signed-int8 -> native conversion.""" |
| 234 | config: dict[str, Any] = { |
| 235 | "quant_method": "compressed-tensors", |
| 236 | "format": "int-quantized", |
| 237 | "type": "int", |
| 238 | "num_bits": 8, |
| 239 | "symmetric": True, |
| 240 | "strategy": "channel", |
| 241 | } |
| 242 | num_experts, num_output = weight.shape[0], weight.shape[-2] |
| 243 | if weight_scale.numel() < num_experts * num_output: |
| 244 | config["strategy"] = "tensor" |
| 245 | return config |
| 246 | |
| 247 | |
| 248 | def convert_to_int8_moe_kernel_format( |
| 249 | int8_backend: Int8MoeBackend, |
| 250 | w13: torch.Tensor, |
| 251 | w2: torch.Tensor, |
| 252 | layer: torch.nn.Module | None = None, |
| 253 | w13_scale: torch.Tensor | None = None, |
| 254 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 255 | """Convert INT8 MoE weights to backend-specific kernel format.""" |
| 256 | if int8_backend == Int8MoeBackend.HUMMING: |
| 257 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 258 | convert_to_humming_moe_kernel_format, |
| 259 | ) |
| 260 | |
| 261 | assert layer is not None |
| 262 | # The Humming conversion expects canonical CT scale parameter names. |
| 263 | # Online int8 produces per-channel (E, N) w*_scale; expose them as |
| 264 | # (E, N, 1) w*_weight_scale tensors before conversion. |
| 265 | for sub in ("w13", "w2"): |
| 266 | if hasattr(layer, f"{sub}_weight_scale"): |
| 267 | continue |
| 268 | scale = getattr(layer, f"{sub}_scale").data |
| 269 | if scale.dim() < 3: |
| 270 | scale = scale.unsqueeze(-1) |
| 271 | replace_parameter(layer, f"{sub}_weight_scale", scale) |
| 272 | delattr(layer, f"{sub}_scale") |
| 273 | convert_to_humming_moe_kernel_format( |
| 274 | layer, |
| 275 | quant_config=_humming_int8_weight_schema(w13, layer.w13_weight_scale), |
| 276 | ) |
| 277 | return layer.w13_weight, layer.w2_weight |
| 278 | elif int8_backend not in (Int8MoeBackend.TRITON, Int8MoeBackend.CPU): |
| 279 | raise ValueError(f"Unsupported Int8 MoE backend: {int8_backend.value}") |
| 280 | |
| 281 | return w13, w2 |
| 282 | |
| 283 | |
| 284 | def make_int8_moe_kernel( |
| 285 | int8_backend: Int8MoeBackend, |
| 286 | moe_quant_config: FusedMoEQuantConfig, |
| 287 | moe_config: FusedMoEConfig, |
| 288 | experts_cls: type[mk.FusedMoEExperts], |
| 289 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 290 | ) -> mk.FusedMoEKernel: |
| 291 | # Create Prepare/Finalize. |
| 292 | prepare_finalize = maybe_make_prepare_finalize( |
| 293 | moe=moe_config, |
| 294 | quant_config=moe_quant_config, |
| 295 | routing_tables=routing_tables, |
| 296 | allow_new_interface=True, |
| 297 | use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), |
| 298 | ) |
| 299 | assert prepare_finalize is not None |
| 300 | |
| 301 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 302 | |
| 303 | # Create Experts. |
| 304 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 305 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 306 | assert max_num_tokens is not None |
| 307 | experts = experts_cls( |
| 308 | moe_config=moe_config, |
| 309 | quant_config=moe_quant_config, |
| 310 | max_num_tokens=max_num_tokens, |
| 311 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 312 | ) |
| 313 | else: |
| 314 | experts = experts_cls( |
| 315 | moe_config=moe_config, |
| 316 | quant_config=moe_quant_config, |
| 317 | ) |
| 318 | |
| 319 | kernel = mk.FusedMoEKernel( |
| 320 | prepare_finalize, |
| 321 | experts, |
| 322 | ) |
| 323 | |
| 324 | return kernel |