| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Oracle for W4A8 (INT4, FP8) MoE |
| 4 | from enum import Enum |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 12 | maybe_make_prepare_finalize, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.fused_moe.config import ( |
| 15 | FusedMoEConfig, |
| 16 | FusedMoEQuantConfig, |
| 17 | int4_w4afp8_moe_quant_config, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 20 | QuantKey, |
| 21 | kFp8DynamicTokenSym, |
| 22 | kInt4Static, |
| 23 | ) |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | class W4A8MoeBackend(Enum): |
| 32 | CUTLASS = "CUTLASS" |
| 33 | HUMMING = "HUMMING" |
| 34 | |
| 35 | |
| 36 | def backend_to_kernel_cls( |
| 37 | backend: W4A8MoeBackend, |
| 38 | ) -> list[type[mk.FusedMoEExperts]]: |
| 39 | if backend == W4A8MoeBackend.CUTLASS: |
| 40 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 41 | CutlassExpertsW4A8Fp8, |
| 42 | ) |
| 43 | |
| 44 | return [CutlassExpertsW4A8Fp8] |
| 45 | elif backend == W4A8MoeBackend.HUMMING: |
| 46 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 47 | BatchedHummingGroupedExperts, |
| 48 | HummingGroupedExperts, |
| 49 | HummingIndexedExperts, |
| 50 | ) |
| 51 | |
| 52 | return [ |
| 53 | BatchedHummingGroupedExperts, |
| 54 | HummingGroupedExperts, |
| 55 | HummingIndexedExperts, |
| 56 | ] |
| 57 | else: |
| 58 | raise ValueError(f"Unknown W4A8 MoE backend: {backend.value}") |
| 59 | |
| 60 | |
| 61 | def select_w4a8_moe_backend( |
| 62 | config: FusedMoEConfig, |
| 63 | weight_key: QuantKey | None = kInt4Static, |
| 64 | activation_key: QuantKey | None = kFp8DynamicTokenSym, |
| 65 | ) -> tuple[W4A8MoeBackend, type[mk.FusedMoEExperts]]: |
| 66 | backends = [W4A8MoeBackend.CUTLASS, W4A8MoeBackend.HUMMING] |
| 67 | if config.moe_backend != "auto": |
| 68 | backend_map = { |
| 69 | "cutlass": W4A8MoeBackend.CUTLASS, |
| 70 | "humming": W4A8MoeBackend.HUMMING, |
| 71 | } |
| 72 | if config.moe_backend not in backend_map: |
| 73 | raise ValueError( |
| 74 | f"moe_backend='{config.moe_backend}' is not supported for W4A8 MoE. " |
| 75 | f"Expected one of {list(backend_map)}." |
| 76 | ) |
| 77 | backends = [backend_map[config.moe_backend]] |
| 78 | |
| 79 | activation_format = ( |
| 80 | mk.FusedMoEActivationFormat.BatchedExperts |
| 81 | if config.moe_parallel_config.use_batched_activation_format |
| 82 | else mk.FusedMoEActivationFormat.Standard |
| 83 | ) |
| 84 | |
| 85 | reasons = [] |
| 86 | for backend in backends: |
| 87 | for kernel_cls in backend_to_kernel_cls(backend): |
| 88 | supported, reason = kernel_cls.is_supported_config( |
| 89 | kernel_cls, |
| 90 | config, |
| 91 | weight_key, |
| 92 | activation_key, |
| 93 | activation_format, |
| 94 | ) |
| 95 | if supported: |
| 96 | logger.info_once("Using %s W4A8 MoE backend.", backend.value) |
| 97 | return backend, kernel_cls |
| 98 | reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 99 | |
| 100 | raise NotImplementedError( |
| 101 | "No W4A8 MoE backend supports the deployment configuration: " |
| 102 | + "; ".join(reasons) |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | def convert_to_w4a8_moe_kernel_format( |
| 107 | backend: W4A8MoeBackend, |
| 108 | layer: "RoutedExperts", |
| 109 | group_size: int, |
| 110 | w13_weight: torch.Tensor, |
| 111 | w2_weight: torch.Tensor, |
| 112 | w13_weight_scale: torch.Tensor, |
| 113 | w2_weight_scale: torch.Tensor, |
| 114 | ) -> tuple[ |
| 115 | torch.Tensor, |
| 116 | torch.Tensor, |
| 117 | torch.Tensor, |
| 118 | torch.Tensor, |
| 119 | torch.Tensor | None, |
| 120 | torch.Tensor | None, |
| 121 | torch.Tensor | None, |
| 122 | torch.Tensor | None, |
| 123 | ]: |
| 124 | if backend == W4A8MoeBackend.HUMMING: |
| 125 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 126 | convert_to_humming_moe_kernel_format, |
| 127 | ) |
| 128 | from vllm.utils.humming import ( |
| 129 | HummingInputSchema, |
| 130 | HummingWeightSchema, |
| 131 | dtypes, |
| 132 | ) |
| 133 | |
| 134 | convert_to_humming_moe_kernel_format( |
| 135 | layer, |
| 136 | weight_schema=HummingWeightSchema( |
| 137 | b_dtype=dtypes.uint4, weight_scale_group_size=group_size |
| 138 | ), |
| 139 | input_schema=HummingInputSchema(a_dtype=dtypes.float8e4m3), |
| 140 | allow_input_schema_fallback=False, |
| 141 | ) |
| 142 | return ( |
| 143 | layer.w13_weight, |
| 144 | layer.w2_weight, |
| 145 | layer.w13_weight_scale, |
| 146 | layer.w2_weight_scale, |
| 147 | None, |
| 148 | None, |
| 149 | None, |
| 150 | None, |
| 151 | ) |
| 152 | |
| 153 | from vllm import _custom_ops as ops |
| 154 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 155 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 156 | GroupShape, |
| 157 | convert_bf16_scales_to_fp8, |
| 158 | convert_packed_uint4b8_to_signed_int4_inplace, |
| 159 | ) |
| 160 | |
| 161 | quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN) |
| 162 | |
| 163 | convert_packed_uint4b8_to_signed_int4_inplace(w13_weight) |
| 164 | # Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness. |
| 165 | torch.accelerator.synchronize() |
| 166 | w13_weight_shuffled, b_strides1 = ops.cutlass_encode_and_reorder_int4b_grouped( |
| 167 | w13_weight |
| 168 | ) |
| 169 | |
| 170 | convert_packed_uint4b8_to_signed_int4_inplace(w2_weight) |
| 171 | # Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness. |
| 172 | torch.accelerator.synchronize() |
| 173 | w2_weight_shuffled, b_strides2 = ops.cutlass_encode_and_reorder_int4b_grouped( |
| 174 | w2_weight |
| 175 | ) |
| 176 | |
| 177 | w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8( |
| 178 | quant_fp8, w13_weight_scale |
| 179 | ) |
| 180 | w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8( |
| 181 | quant_fp8, w2_weight_scale |
| 182 | ) |
| 183 | |
| 184 | # Scales are stored as (E, N, K // 128), but the kernel expects |
| 185 | # (E, K // 128, N) in row-major format. |
| 186 | w13_weight_scale_packed = ops.cutlass_pack_scale_fp8( |
| 187 | w13_weight_scale.permute(0, 2, 1).contiguous() |
| 188 | ) |
| 189 | w2_weight_scale_packed = ops.cutlass_pack_scale_fp8( |
| 190 | w2_weight_scale.permute(0, 2, 1).contiguous() |
| 191 | ) |
| 192 | |
| 193 | return ( |
| 194 | w13_weight_shuffled, |
| 195 | w2_weight_shuffled, |
| 196 | w13_weight_scale_packed, |
| 197 | w2_weight_scale_packed, |
| 198 | w13_weight_chan_scale, |
| 199 | w2_weight_chan_scale, |
| 200 | b_strides1, |
| 201 | b_strides2, |
| 202 | ) |
| 203 | |
| 204 | |
| 205 | def make_w4a8_moe_quant_config( |
| 206 | backend: W4A8MoeBackend, |
| 207 | w1_scale: torch.Tensor, |
| 208 | w2_scale: torch.Tensor, |
| 209 | g1_alphas: torch.Tensor | None, |
| 210 | g2_alphas: torch.Tensor | None, |
| 211 | layer: torch.nn.Module | None = None, |
| 212 | ) -> FusedMoEQuantConfig: |
| 213 | if backend == W4A8MoeBackend.HUMMING: |
| 214 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 215 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 216 | get_humming_moe_quant_config, |
| 217 | ) |
| 218 | |
| 219 | assert isinstance(layer, RoutedExperts) |
| 220 | return get_humming_moe_quant_config(layer) |
| 221 | |
| 222 | assert g1_alphas is not None and g2_alphas is not None |
| 223 | return int4_w4afp8_moe_quant_config( |
| 224 | w1_scale=w1_scale, |
| 225 | w2_scale=w2_scale, |
| 226 | g1_alphas=g1_alphas, |
| 227 | g2_alphas=g2_alphas, |
| 228 | per_act_token_quant=True, |
| 229 | per_out_ch_quant=True, |
| 230 | ) |
| 231 | |
| 232 | |
| 233 | def make_w4a8_moe_kernel( |
| 234 | backend: W4A8MoeBackend, |
| 235 | moe_quant_config: FusedMoEQuantConfig, |
| 236 | moe_config: FusedMoEConfig, |
| 237 | experts_cls: type[mk.FusedMoEExperts], |
| 238 | b_strides1: torch.Tensor | None, |
| 239 | b_strides2: torch.Tensor | None, |
| 240 | group_size: int, |
| 241 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 242 | ) -> mk.FusedMoEKernel: |
| 243 | if backend == W4A8MoeBackend.HUMMING: |
| 244 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 245 | make_humming_moe_kernel, |
| 246 | ) |
| 247 | |
| 248 | return make_humming_moe_kernel( |
| 249 | moe_quant_config, moe_config, experts_cls, routing_tables=routing_tables |
| 250 | ) |
| 251 | |
| 252 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 253 | CutlassExpertsW4A8Fp8, |
| 254 | ) |
| 255 | |
| 256 | assert issubclass(experts_cls, CutlassExpertsW4A8Fp8) |
| 257 | assert b_strides1 is not None and b_strides2 is not None |
| 258 | prepare_finalize = maybe_make_prepare_finalize( |
| 259 | moe=moe_config, |
| 260 | quant_config=moe_quant_config, |
| 261 | routing_tables=routing_tables, |
| 262 | allow_new_interface=True, |
| 263 | ) |
| 264 | assert prepare_finalize is not None |
| 265 | |
| 266 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 267 | |
| 268 | experts = experts_cls( |
| 269 | moe_config=moe_config, |
| 270 | quant_config=moe_quant_config, |
| 271 | b_strides1=b_strides1, |
| 272 | b_strides2=b_strides2, |
| 273 | group_size=group_size, |
| 274 | ) |
| 275 | |
| 276 | return mk.FusedMoEKernel( |
| 277 | prepare_finalize, |
| 278 | experts, |
| 279 | ) |