| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused MoE W4A16 experts on the RDNA3 (gfx1100) HIP kernel. |
| 4 | |
| 5 | ``moe_gptq_gemm_rdna3`` is a single HIP kernel launch per GEMM that handles |
| 6 | expert routing + W4A16 dequant + dot product with atomic output accumulation. |
| 7 | |
| 8 | Weight format (per expert, same as the dense RDNA3 W4A16 kernel): |
| 9 | - Packed int32 ``[E, K/8, N]`` with exllama shuffle |
| 10 | - Scales ``[E, groups, N]`` in activation dtype |
| 11 | - Zero points ``[E, groups, N/8]`` packed int32 (synthesized, symmetric only) |
| 12 | """ |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | import vllm._custom_ops as ops |
| 17 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 18 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 19 | MoEActivation, |
| 20 | apply_moe_activation_supported, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.config import ( |
| 23 | FusedMoEConfig, |
| 24 | FusedMoEParallelConfig, |
| 25 | FusedMoEQuantConfig, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( |
| 28 | moe_align_block_size, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 31 | TopKWeightAndReduceNoOP, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 34 | QuantKey, |
| 35 | kInt4Static, |
| 36 | kInt4Static32, |
| 37 | ) |
| 38 | from vllm.platforms import current_platform |
| 39 | |
| 40 | |
| 41 | def rdna3_moe_kernel_available() -> bool: |
| 42 | """Whether the fused RDNA3 MoE HIP kernel is built into this binary.""" |
| 43 | if not current_platform.is_rocm(): |
| 44 | return False |
| 45 | |
| 46 | from vllm.platforms.rocm import on_gfx1100 |
| 47 | |
| 48 | return ( |
| 49 | on_gfx1100() |
| 50 | and hasattr(torch.ops, "_rocm_C") |
| 51 | and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3") |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | class Rdna3WNA16Experts(mk.FusedMoEExpertsModular): |
| 56 | """W4A16 experts backed by ``moe_gptq_gemm_rdna3`` (gfx1100). |
| 57 | |
| 58 | Both GEMMs accumulate atomically, so their destinations are zeroed first. |
| 59 | The second GEMM is given ``output_topk``, which makes it reduce over the |
| 60 | top-k dimension while it accumulates — the ``moe_sum`` launch and the |
| 61 | ``[M * top_k, K]`` intermediate are fused away, hence |
| 62 | ``TopKWeightAndReduceNoOP``. |
| 63 | """ |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | moe_config: FusedMoEConfig, |
| 68 | quant_config: FusedMoEQuantConfig, |
| 69 | max_num_tokens: int | None = None, |
| 70 | num_dispatchers: int | None = None, |
| 71 | ): |
| 72 | assert quant_config.use_int4_w4a16, "Supports only int4_w4a16" |
| 73 | super().__init__( |
| 74 | moe_config=moe_config, |
| 75 | quant_config=quant_config, |
| 76 | max_num_tokens=max_num_tokens, |
| 77 | num_dispatchers=num_dispatchers, |
| 78 | ) |
| 79 | self._empty_topk_weights: torch.Tensor | None = None |
| 80 | |
| 81 | @staticmethod |
| 82 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 83 | return mk.FusedMoEActivationFormat.Standard |
| 84 | |
| 85 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 86 | return TopKWeightAndReduceNoOP() |
| 87 | |
| 88 | @staticmethod |
| 89 | def _supports_current_device() -> bool: |
| 90 | return rdna3_moe_kernel_available() |
| 91 | |
| 92 | @staticmethod |
| 93 | def _supports_no_act_and_mul() -> bool: |
| 94 | return True |
| 95 | |
| 96 | @staticmethod |
| 97 | def _supports_quant_scheme( |
| 98 | weight_key: QuantKey | None, |
| 99 | activation_key: QuantKey | None, |
| 100 | ) -> bool: |
| 101 | # Symmetric int4 weights only: the kernel consumes synthesized zero |
| 102 | # points and has no path for checkpoint zero points. |
| 103 | return activation_key is None and weight_key in (kInt4Static, kInt4Static32) |
| 104 | |
| 105 | @staticmethod |
| 106 | def _supports_activation(activation: MoEActivation) -> bool: |
| 107 | return apply_moe_activation_supported(activation) |
| 108 | |
| 109 | @staticmethod |
| 110 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 111 | return not ( |
| 112 | moe_parallel_config.use_fi_nvl_two_sided_kernels |
| 113 | or moe_parallel_config.use_fi_nvl_one_sided_kernels |
| 114 | ) |
| 115 | |
| 116 | def moe_problem_size( |
| 117 | self, |
| 118 | a1: torch.Tensor, |
| 119 | w1: torch.Tensor, |
| 120 | w2: torch.Tensor, |
| 121 | topk_ids: torch.Tensor, |
| 122 | ) -> tuple[int, int, int, int, int]: |
| 123 | # Weights are packed K-first: w1 is [E, K // 8, w13_shards * N] and |
| 124 | # w2 is [E, N // 8, K], so neither N nor K can be read off a trailing |
| 125 | # dimension the way the base implementation does. |
| 126 | assert w1.dim() == 3 and w2.dim() == 3 |
| 127 | assert a1.dim() == 2 |
| 128 | assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" |
| 129 | return w1.size(0), a1.size(0), w2.size(1) * 8, a1.size(-1), topk_ids.size(1) |
| 130 | |
| 131 | def workspace_shapes( |
| 132 | self, |
| 133 | M: int, |
| 134 | N: int, |
| 135 | K: int, |
| 136 | topk: int, |
| 137 | global_num_experts: int, |
| 138 | local_num_experts: int, |
| 139 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 140 | activation: MoEActivation, |
| 141 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 142 | # The modular kernel provisions the output buffer out of workspace13, |
| 143 | # so both GEMM-1 buffers have to live in workspace2 — the second GEMM |
| 144 | # reads the activation output while accumulating into the output. |
| 145 | gate_up = N if not activation.is_gated else 2 * N |
| 146 | return ((M, K), (M * topk * (gate_up + N),), (M, K)) |
| 147 | |
| 148 | def apply( |
| 149 | self, |
| 150 | output: torch.Tensor, |
| 151 | hidden_states: torch.Tensor, |
| 152 | w1: torch.Tensor, |
| 153 | w2: torch.Tensor, |
| 154 | topk_weights: torch.Tensor, |
| 155 | topk_ids: torch.Tensor, |
| 156 | activation: MoEActivation, |
| 157 | global_num_experts: int, |
| 158 | expert_map: torch.Tensor | None, |
| 159 | a1q_scale: torch.Tensor | None, |
| 160 | a2_scale: torch.Tensor | None, |
| 161 | workspace13: torch.Tensor, |
| 162 | workspace2: torch.Tensor, |
| 163 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 164 | apply_router_weight_on_input: bool, |
| 165 | ) -> None: |
| 166 | assert self.w1_scale is not None and self.w2_scale is not None |
| 167 | assert self.w1_zp is not None and self.w2_zp is not None |
| 168 | |
| 169 | E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids) |
| 170 | if global_num_experts == -1: |
| 171 | global_num_experts = E |
| 172 | |
| 173 | gate_up = w1.size(2) |
| 174 | act_n = self.adjust_N_for_activation(gate_up, activation) |
| 175 | rows = M * top_k |
| 176 | |
| 177 | scratch = workspace2.view(-1) |
| 178 | gate_up_out = scratch[: rows * gate_up].view(rows, gate_up) |
| 179 | act_out = scratch[rows * gate_up : rows * (gate_up + act_n)].view(rows, act_n) |
| 180 | |
| 181 | # BLOCK_SIZE_M=1 for decode (no padding waste), 4 for prefill. |
| 182 | block_size_m = 1 if M <= 4 else 4 |
| 183 | sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( |
| 184 | topk_ids, |
| 185 | block_size_m, |
| 186 | global_num_experts, |
| 187 | expert_map, |
| 188 | ) |
| 189 | |
| 190 | if ( |
| 191 | self._empty_topk_weights is None |
| 192 | or self._empty_topk_weights.device != hidden_states.device |
| 193 | ): |
| 194 | self._empty_topk_weights = torch.empty( |
| 195 | 0, dtype=torch.float32, device=hidden_states.device |
| 196 | ) |
| 197 | topk_weights_f32 = topk_weights.reshape(-1).float() |
| 198 | no_topk_weights = self._empty_topk_weights |
| 199 | |
| 200 | gate_up_out.zero_() |
| 201 | ops.moe_gptq_gemm_rdna3( |
| 202 | hidden_states, |
| 203 | gate_up_out, |
| 204 | w1, |
| 205 | self.w1_scale, |
| 206 | self.w1_zp, |
| 207 | topk_weights_f32 if apply_router_weight_on_input else no_topk_weights, |
| 208 | sorted_token_ids, |
| 209 | expert_ids, |
| 210 | num_tokens_post_padded, |
| 211 | top_k, |
| 212 | block_size_m, |
| 213 | apply_router_weight_on_input, |
| 214 | ) |
| 215 | |
| 216 | self.activation(activation, act_out, gate_up_out) |
| 217 | |
| 218 | # output_topk=top_k makes the kernel accumulate into out[token_id], |
| 219 | # fusing the top-k reduction into the atomic write-back. |
| 220 | output.zero_() |
| 221 | ops.moe_gptq_gemm_rdna3( |
| 222 | act_out, |
| 223 | output, |
| 224 | w2, |
| 225 | self.w2_scale, |
| 226 | self.w2_zp, |
| 227 | no_topk_weights if apply_router_weight_on_input else topk_weights_f32, |
| 228 | sorted_token_ids, |
| 229 | expert_ids, |
| 230 | num_tokens_post_padded, |
| 231 | 1, |
| 232 | block_size_m, |
| 233 | not apply_router_weight_on_input, |
| 234 | output_topk=top_k, |
| 235 | ) |