| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 9 | from vllm.model_executor.layers.fused_moe.config import ( |
| 10 | FusedMoEParallelConfig, |
| 11 | FusedMoEQuantConfig, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 14 | TopKWeightAndReduceNoOP, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 17 | QuantKey, |
| 18 | kFp8Dynamic128Sym, |
| 19 | kFp8Static128BlockSym, |
| 20 | kFp8StaticTensorSym, |
| 21 | ) |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.utils.hpc import has_hpc, hpc_fuse_moe, hpc_fuse_moe_blockwise |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | |
| 28 | class HPCExperts(mk.FusedMoEExpertsModular): |
| 29 | """MoE implementation powered by [HPC](https://github.com/Tencent/hpc-ops). |
| 30 | |
| 31 | Only supported on NVIDIA Hopper GPUs (e.g. H20, H200), and currently limited to |
| 32 | FP8 models such as Hy3-FP8, Qwen3-235B-A22B-FP8, etc. |
| 33 | """ |
| 34 | |
| 35 | def __init__( |
| 36 | self, |
| 37 | moe_config: mk.FusedMoEConfig, |
| 38 | quant_config: FusedMoEQuantConfig, |
| 39 | ): |
| 40 | super().__init__(moe_config, quant_config) |
| 41 | |
| 42 | assert quant_config.weight_quant_dtype in (torch.float8_e4m3fn,), ( |
| 43 | "Only fp8 quantization is currently supported." |
| 44 | ) |
| 45 | |
| 46 | self.device = moe_config.device |
| 47 | self.num_experts = moe_config.num_local_experts |
| 48 | self.ep_rank = moe_config.moe_parallel_config.ep_rank |
| 49 | self.ep_size = moe_config.moe_parallel_config.ep_size |
| 50 | self.tp_rank = moe_config.moe_parallel_config.tp_rank |
| 51 | self.tp_size = moe_config.moe_parallel_config.tp_size |
| 52 | self.out_dtype = moe_config.in_dtype |
| 53 | |
| 54 | @property |
| 55 | def expects_unquantized_inputs(self) -> bool: |
| 56 | return False |
| 57 | |
| 58 | @staticmethod |
| 59 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 60 | return mk.FusedMoEActivationFormat.Standard |
| 61 | |
| 62 | @staticmethod |
| 63 | def _supports_current_device() -> bool: |
| 64 | p = current_platform |
| 65 | return p.is_cuda() and p.is_device_capability(90) and has_hpc() |
| 66 | |
| 67 | @staticmethod |
| 68 | def _supports_no_act_and_mul() -> bool: |
| 69 | return False |
| 70 | |
| 71 | @staticmethod |
| 72 | def _supports_quant_scheme( |
| 73 | weight_key: QuantKey | None, |
| 74 | activation_key: QuantKey | None, |
| 75 | ) -> bool: |
| 76 | scheme = (weight_key, activation_key) |
| 77 | # The following are supported by HPCExperts: |
| 78 | return scheme in [ |
| 79 | # fp8 static per-tensor on 9.0+ |
| 80 | (kFp8StaticTensorSym, kFp8StaticTensorSym), |
| 81 | (kFp8Static128BlockSym, kFp8Dynamic128Sym), |
| 82 | ] |
| 83 | |
| 84 | @staticmethod |
| 85 | def _supports_activation(activation: MoEActivation) -> bool: |
| 86 | return activation in [ |
| 87 | MoEActivation.SILU, |
| 88 | ] |
| 89 | |
| 90 | @staticmethod |
| 91 | def _supports_shape(hidden_dim: int) -> bool: |
| 92 | # HPC fused MoE kernels process hidden_size in blocks of 128: |
| 93 | # block-wise fp8 requires hidden_size % 128 == 0 (per-128 quant), and |
| 94 | # the group GEMM tiles N by 128. Require 128-alignment to cover all |
| 95 | # code paths. |
| 96 | return hidden_dim % 128 == 0 |
| 97 | |
| 98 | @staticmethod |
| 99 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 100 | return True |
| 101 | |
| 102 | def supports_expert_map(self) -> bool: |
| 103 | return False |
| 104 | |
| 105 | def supports_chunking(self) -> bool: |
| 106 | # This refers to TP chunking; DP chunking is handled separately. |
| 107 | return True |
| 108 | |
| 109 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 110 | return TopKWeightAndReduceNoOP() |
| 111 | |
| 112 | def workspace_shapes( |
| 113 | self, |
| 114 | M: int, |
| 115 | N: int, |
| 116 | K: int, |
| 117 | topk: int, |
| 118 | global_num_experts: int, |
| 119 | local_num_experts: int, |
| 120 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 121 | activation: MoEActivation, |
| 122 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 123 | # We use global_num_experts due to how moe_align_block_size handles |
| 124 | # expert_maps. |
| 125 | """Compute the shapes for the temporary and final outputs of the two gemms |
| 126 | and activation in the fused expert function. Since the gemms are |
| 127 | independent, the workspace for the first gemm can be shared with the |
| 128 | workspace for the last gemm. |
| 129 | |
| 130 | Returns a tuple of: |
| 131 | - workspace13 shape tuple: must be large enough to hold the |
| 132 | result of either expert gemm. |
| 133 | - workspace2 shape tuple: must be large enough to hold the |
| 134 | result of the activation function. |
| 135 | - output shape tuple: must be exact size of the final gemm output. |
| 136 | - Workspace type: The dtype to use for the workspace tensors. |
| 137 | - Note: in order for activation chunking to work, the first dimension |
| 138 | of each tuple must be the number of tokens. |
| 139 | """ |
| 140 | workspace1 = (M, K) |
| 141 | workspace2 = (0,) |
| 142 | output_shape = (M, K) |
| 143 | # The workspace is determined by `aq`, since it comes after any |
| 144 | # potential communication op and is involved in the expert computation. |
| 145 | return (workspace1, workspace2, output_shape) |
| 146 | |
| 147 | def apply( |
| 148 | self, |
| 149 | output: torch.Tensor, |
| 150 | hidden_states: torch.Tensor, |
| 151 | w1: torch.Tensor, |
| 152 | w2: torch.Tensor, |
| 153 | topk_weights: torch.Tensor, |
| 154 | topk_ids: torch.Tensor, |
| 155 | activation: MoEActivation, |
| 156 | global_num_experts: int, |
| 157 | expert_map: torch.Tensor | None, |
| 158 | a1q_scale: torch.Tensor | None, |
| 159 | a2_scale: torch.Tensor | None, |
| 160 | workspace13: torch.Tensor | None, |
| 161 | workspace2: torch.Tensor | None, |
| 162 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 163 | apply_router_weight_on_input: bool | None, |
| 164 | ): |
| 165 | assert self._supports_activation(activation), f"{activation=} not supported" |
| 166 | assert self.quant_config.w1_scale is not None, ( |
| 167 | "w13_weight_scale must be provided" |
| 168 | ) |
| 169 | assert self.quant_config.w2_scale is not None, ( |
| 170 | "w2_weight_scale must be provided" |
| 171 | ) |
| 172 | |
| 173 | if self.quant_config.is_block_quantized: |
| 174 | hpc_fuse_moe_blockwise( |
| 175 | x=hidden_states, |
| 176 | x_scale=a1q_scale, |
| 177 | gate_up_weight=w1, |
| 178 | gate_up_weight_scale=self.quant_config.w1_scale, |
| 179 | down_weight=w2, |
| 180 | down_weight_scale=self.quant_config.w2_scale, |
| 181 | topk_ids=topk_ids, |
| 182 | topk_scale=topk_weights, |
| 183 | rank_ep=self.ep_rank, |
| 184 | num_expert_total=global_num_experts, |
| 185 | output=output, |
| 186 | ) |
| 187 | else: |
| 188 | assert self.quant_config.a1_scale is not None, ( |
| 189 | "w13_input_scale must be provided" |
| 190 | ) |
| 191 | assert self.quant_config.a2_scale is not None, ( |
| 192 | "w2_input_scale must be provided" |
| 193 | ) |
| 194 | hpc_fuse_moe( |
| 195 | x=hidden_states, |
| 196 | gate_up_weight=w1, |
| 197 | down_weight=w2, |
| 198 | gate_up_scale=self.quant_config.g1_alphas, |
| 199 | down_scale=self.quant_config.g2_alphas, |
| 200 | act_and_mul_scale=self.quant_config.a2_gscale, |
| 201 | topk_ids=topk_ids, |
| 202 | topk_scale=topk_weights, |
| 203 | rank_ep=self.ep_rank, |
| 204 | num_expert_total=global_num_experts, |
| 205 | output=output, |
| 206 | ) |