| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import torch |
| 7 | from torch.nn import Module |
| 8 | |
| 9 | if TYPE_CHECKING: |
| 10 | from vllm.model_executor.layers.fused_moe.config import ( |
| 11 | FusedMoEQuantConfig, |
| 12 | ) |
| 13 | |
| 14 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 15 | from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig |
| 16 | from vllm.model_executor.layers.fused_moe.oracle.int8 import ( |
| 17 | convert_to_int8_moe_kernel_format, |
| 18 | make_int8_moe_kernel, |
| 19 | make_int8_moe_quant_config, |
| 20 | select_int8_moe_backend, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization.online.moe_base import ( |
| 23 | OnlineMoEMethodBase, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 26 | amax_for_moe_weight_quant, |
| 27 | kInt8DynamicTokenSym, |
| 28 | kInt8StaticChannelSym, |
| 29 | weight_amax, |
| 30 | ) |
| 31 | from vllm.model_executor.utils import replace_parameter |
| 32 | |
| 33 | |
| 34 | class Int8OnlineMoEMethod(OnlineMoEMethodBase): |
| 35 | """Online per-channel INT8 MoE quantization. |
| 36 | Loads fp16/bf16 weights and quantizes them per-row to int8 during loading. |
| 37 | """ |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | *, |
| 42 | moe: FusedMoEConfig, |
| 43 | ): |
| 44 | super().__init__(moe) |
| 45 | self.int8_backend, self.experts_cls = select_int8_moe_backend( |
| 46 | config=self.moe, |
| 47 | weight_key=kInt8StaticChannelSym, |
| 48 | activation_key=kInt8DynamicTokenSym, |
| 49 | ) |
| 50 | |
| 51 | def process_weights_after_loading(self, layer: Module) -> None: |
| 52 | if getattr(layer, "_already_called_process_weights_after_loading", False): |
| 53 | return |
| 54 | |
| 55 | self._quantize_weights(layer) |
| 56 | self._setup_kernel(layer) |
| 57 | |
| 58 | layer._already_called_process_weights_after_loading = True |
| 59 | |
| 60 | def _quantize_weights(self, layer: Module) -> None: |
| 61 | vmax = torch.iinfo(torch.int8).max |
| 62 | |
| 63 | w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8) |
| 64 | w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8) |
| 65 | w13_scale = torch.zeros( |
| 66 | layer.num_experts, |
| 67 | layer.w13_weight.shape[1], |
| 68 | device=w13.device, |
| 69 | dtype=torch.float32, |
| 70 | ) |
| 71 | w2_scale = torch.zeros( |
| 72 | layer.num_experts, |
| 73 | layer.w2_weight.shape[1], |
| 74 | device=w2.device, |
| 75 | dtype=torch.float32, |
| 76 | ) |
| 77 | |
| 78 | w2_amax = weight_amax(layer.w2_weight, dim=-1) |
| 79 | w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size) |
| 80 | |
| 81 | for expert in range(layer.local_num_experts): |
| 82 | # w13: per-row quantization over hidden_size dim |
| 83 | w = layer.w13_weight[expert, :, :] |
| 84 | scales = w.abs().amax(dim=1) / vmax |
| 85 | q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) |
| 86 | w13[expert, :, :] = q.to(torch.int8) |
| 87 | w13_scale[expert, :] = scales |
| 88 | |
| 89 | # w2: per-row quantization over intermediate_size dim |
| 90 | w = layer.w2_weight[expert, :, :] |
| 91 | scales = w2_amax[expert] / vmax |
| 92 | q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) |
| 93 | w2[expert, :, :] = q.to(torch.int8) |
| 94 | w2_scale[expert, :] = scales |
| 95 | |
| 96 | replace_parameter(layer, "w13_weight", w13) |
| 97 | replace_parameter(layer, "w2_weight", w2) |
| 98 | replace_parameter(layer, "w13_scale", w13_scale) |
| 99 | replace_parameter(layer, "w2_scale", w2_scale) |
| 100 | |
| 101 | def _setup_kernel(self, layer: RoutedExperts) -> None: |
| 102 | w13, w2 = convert_to_int8_moe_kernel_format( |
| 103 | int8_backend=self.int8_backend, |
| 104 | w13=layer.w13_weight, |
| 105 | w2=layer.w2_weight, |
| 106 | layer=layer, |
| 107 | w13_scale=layer.w13_scale, |
| 108 | ) |
| 109 | replace_parameter(layer, "w13_weight", w13) |
| 110 | replace_parameter(layer, "w2_weight", w2) |
| 111 | |
| 112 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 113 | assert self.moe_quant_config is not None |
| 114 | assert self.experts_cls is not None |
| 115 | self.moe_kernel = make_int8_moe_kernel( |
| 116 | int8_backend=self.int8_backend, |
| 117 | moe_quant_config=self.moe_quant_config, |
| 118 | moe_config=self.moe, |
| 119 | experts_cls=self.experts_cls, |
| 120 | routing_tables=layer._expert_routing_tables(), |
| 121 | ) |
| 122 | self.moe_kernel.fused_experts.process_weights_after_loading(layer) |
| 123 | |
| 124 | def get_fused_moe_quant_config( |
| 125 | self, layer: torch.nn.Module |
| 126 | ) -> "FusedMoEQuantConfig | None": |
| 127 | return make_int8_moe_quant_config( |
| 128 | int8_backend=self.int8_backend, |
| 129 | w1_scale=getattr(layer, "w13_scale", None), |
| 130 | w2_scale=getattr(layer, "w2_scale", None), |
| 131 | w1_bias=getattr(layer, "w13_bias", None), |
| 132 | w2_bias=getattr(layer, "w2_bias", None), |
| 133 | layer=layer, |
| 134 | ) |