| 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.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 8 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 9 | TopKWeightAndReduceDelegate, |
| 10 | TopKWeightAndReduceNaiveBatched, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 13 | moe_kernel_quantize_input, |
| 14 | normalize_scales_shape, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): |
| 19 | """A reference prepare/finalize class that reorganizes the tokens into |
| 20 | expert batched format, i.e. E x max_num_tokens x K. This is the format |
| 21 | that the batched dispatch/combine kernels use. |
| 22 | """ |
| 23 | |
| 24 | def __init__( |
| 25 | self, |
| 26 | max_num_tokens: int, |
| 27 | num_local_experts: int, |
| 28 | num_dispatchers: int, |
| 29 | rank: int, |
| 30 | ): |
| 31 | super().__init__() |
| 32 | self.max_num_tokens = max_num_tokens |
| 33 | self.num_local_experts = num_local_experts |
| 34 | self.rank = rank |
| 35 | self.num_dispatchers_ = num_dispatchers |
| 36 | |
| 37 | @property |
| 38 | def activation_format(self) -> mk.FusedMoEActivationFormat: |
| 39 | return mk.FusedMoEActivationFormat.BatchedExperts |
| 40 | |
| 41 | def max_num_tokens_per_rank(self) -> int | None: |
| 42 | return self.max_num_tokens |
| 43 | |
| 44 | def topk_indices_dtype(self) -> torch.dtype | None: |
| 45 | return None |
| 46 | |
| 47 | def num_dispatchers(self) -> int: |
| 48 | return self.num_dispatchers_ |
| 49 | |
| 50 | def output_is_reduced(self) -> bool: |
| 51 | return False |
| 52 | |
| 53 | def prepare( |
| 54 | self, |
| 55 | a1: torch.Tensor, |
| 56 | topk_weights: torch.Tensor, |
| 57 | topk_ids: torch.Tensor, |
| 58 | num_experts: int, |
| 59 | expert_map: torch.Tensor | None, |
| 60 | apply_router_weight_on_input: bool, |
| 61 | quant_config: FusedMoEQuantConfig, |
| 62 | defer_input_quant: bool = False, |
| 63 | ) -> mk.PrepareResultType: |
| 64 | if defer_input_quant: |
| 65 | raise NotImplementedError( |
| 66 | f"{self.__class__.__name__} does not support defer_input_quant=True. " |
| 67 | "Please select an MoE kernel that accepts quantized inputs." |
| 68 | ) |
| 69 | assert a1.dim() == 2 |
| 70 | assert topk_ids.dim() == 2 |
| 71 | assert topk_ids.size(0) == a1.size(0) |
| 72 | |
| 73 | if apply_router_weight_on_input: |
| 74 | topk = topk_ids.size(1) |
| 75 | # TODO: this only works for topK=1, will need to update for topK>1 |
| 76 | assert topk == 1, ( |
| 77 | "apply_router_weight_on_input is only implemented for topk=1" |
| 78 | ) |
| 79 | a1.mul_(topk_weights.to(a1.dtype)) |
| 80 | |
| 81 | num_tokens, hidden_dim = a1.size() |
| 82 | topk = topk_ids.size(1) |
| 83 | |
| 84 | tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device) |
| 85 | |
| 86 | num_local_experts = self.num_local_experts |
| 87 | |
| 88 | if quant_config.quant_dtype is None: |
| 89 | b_type = a1.dtype |
| 90 | else: |
| 91 | b_type = quant_config.quant_dtype |
| 92 | |
| 93 | b_a1 = torch.zeros( |
| 94 | (num_local_experts, self.max_num_tokens, hidden_dim), |
| 95 | dtype=b_type, |
| 96 | device=a1.device, |
| 97 | ) |
| 98 | |
| 99 | if quant_config.is_quantized: |
| 100 | scale_shape = quant_config.batched_scale_shape( |
| 101 | num_local_experts, self.max_num_tokens, hidden_dim |
| 102 | ) |
| 103 | |
| 104 | b_a1_scale = torch.zeros(scale_shape, dtype=torch.float32, device=a1.device) |
| 105 | else: |
| 106 | assert quant_config.a1_scale is None |
| 107 | b_a1_scale = None |
| 108 | |
| 109 | first_expert = num_local_experts * self.rank |
| 110 | last_expert = first_expert + num_local_experts |
| 111 | |
| 112 | a1_scale = normalize_scales_shape(quant_config.a1_scale) |
| 113 | |
| 114 | if quant_config.quant_dtype is None: |
| 115 | # Vectorized dispatch: compute the [E_local, T] hit mask and |
| 116 | # per-expert slot offsets in one shot (single nonzero sync instead |
| 117 | # of a launch-bound per-expert Python loop). |
| 118 | local_ids = torch.arange(first_expert, last_expert, device=a1.device).view( |
| 119 | -1, 1, 1 |
| 120 | ) |
| 121 | hits = (topk_ids.unsqueeze(0) == local_ids).any(dim=2) # [E_local, T] |
| 122 | tokens_per_expert[:num_local_experts] = hits.sum(dim=1).to(torch.int32) |
| 123 | # slot of each token within its expert batch, preserving token order |
| 124 | slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T] |
| 125 | e_idx, t_idx = hits.nonzero(as_tuple=True) |
| 126 | b_a1[e_idx, slots[e_idx, t_idx]] = a1[t_idx].to(b_type) |
| 127 | else: |
| 128 | for expert_id in range(first_expert, last_expert): |
| 129 | topks = torch.any(topk_ids == expert_id, dim=1).flatten() |
| 130 | rows = torch.count_nonzero(topks.flatten()) |
| 131 | if rows == 0: |
| 132 | continue |
| 133 | idx = expert_id - first_expert |
| 134 | tokens_per_expert[idx] = rows |
| 135 | rhs = a1[: topks.numel()][topks] |
| 136 | if a1_scale is not None: |
| 137 | if quant_config.is_per_act_token: |
| 138 | rhs_a1_scale = a1_scale[: topks.numel()][topks] |
| 139 | else: |
| 140 | rhs_a1_scale = a1_scale |
| 141 | else: |
| 142 | rhs_a1_scale = None |
| 143 | b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input( |
| 144 | rhs, |
| 145 | rhs_a1_scale, |
| 146 | quant_config.quant_dtype, |
| 147 | quant_config.per_act_token_quant, |
| 148 | quant_config.block_shape, |
| 149 | ) |
| 150 | assert b_s is not None |
| 151 | if quant_config.is_per_act_token: |
| 152 | b_a1_scale[idx, :rows] = b_s[:rows] |
| 153 | else: |
| 154 | b_a1_scale[idx, : b_s.shape[0]] = b_s |
| 155 | |
| 156 | assert b_a1_scale is None or b_a1_scale.ndim == 3 |
| 157 | |
| 158 | expert_tokens_meta = mk.ExpertTokensMetadata( |
| 159 | expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None |
| 160 | ) |
| 161 | |
| 162 | return b_a1, b_a1_scale, expert_tokens_meta, None, None |
| 163 | |
| 164 | def finalize( |
| 165 | self, |
| 166 | output: torch.Tensor, |
| 167 | fused_expert_output: torch.Tensor, |
| 168 | topk_weights: torch.Tensor, |
| 169 | topk_ids: torch.Tensor, |
| 170 | apply_router_weight_on_input: bool, |
| 171 | weight_and_reduce_impl: mk.TopKWeightAndReduce, |
| 172 | ) -> None: |
| 173 | if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): |
| 174 | weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank) |
| 175 | weight_and_reduce_impl.apply( |
| 176 | output=output, |
| 177 | fused_expert_output=fused_expert_output, |
| 178 | topk_weights=topk_weights, |
| 179 | topk_ids=topk_ids, |
| 180 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 181 | ) |