| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass, field |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.v1.worker.workspace import ( |
| 9 | current_workspace_manager, |
| 10 | is_workspace_manager_initialized, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | @dataclass |
| 15 | class MoEPermuteScratch: |
| 16 | # Reused metadata buffers for repeated grouped-MoE permutes. |
| 17 | max_num_tokens: int |
| 18 | topk: int |
| 19 | num_experts: int |
| 20 | num_local_experts: int |
| 21 | device: torch.device |
| 22 | hidden_size: int | None = None |
| 23 | hidden_dtype: torch.dtype | None = None |
| 24 | token_expert_indices: torch.Tensor = field(init=False) |
| 25 | expert_first_token_offset: torch.Tensor = field(init=False) |
| 26 | permuted_idx: torch.Tensor = field(init=False) |
| 27 | inv_permuted_idx: torch.Tensor = field(init=False) |
| 28 | permuted_hidden_states: torch.Tensor | None = field(init=False, default=None) |
| 29 | sort_workspace: torch.Tensor = field(init=False) |
| 30 | permuted_experts_id: torch.Tensor = field(init=False) |
| 31 | sorted_row_idx: torch.Tensor = field(init=False) |
| 32 | topk_ids_int32: torch.Tensor = field(init=False) |
| 33 | topk_ids_for_sort: torch.Tensor = field(init=False) |
| 34 | max_expanded_rows: int = field(init=False) |
| 35 | |
| 36 | def __post_init__(self) -> None: |
| 37 | assert self.max_num_tokens > 0 |
| 38 | assert self.topk > 0 |
| 39 | assert self.num_experts > 0 |
| 40 | assert self.num_local_experts > 0 |
| 41 | if self.hidden_size is None: |
| 42 | assert self.hidden_dtype is None |
| 43 | else: |
| 44 | assert self.hidden_dtype is not None |
| 45 | |
| 46 | self.max_expanded_rows = self.max_num_tokens * self.topk |
| 47 | self.token_expert_indices = torch.arange( |
| 48 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 49 | ) |
| 50 | self.expert_first_token_offset = torch.empty( |
| 51 | self.num_local_experts + 1, dtype=torch.int64, device=self.device |
| 52 | ) |
| 53 | self.permuted_idx = torch.empty( |
| 54 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 55 | ) |
| 56 | self.inv_permuted_idx = torch.empty( |
| 57 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 58 | ) |
| 59 | if self.hidden_size is not None: |
| 60 | hidden_numel = self.max_expanded_rows * self.hidden_size |
| 61 | self.permuted_hidden_states = torch.empty( |
| 62 | hidden_numel, dtype=self.hidden_dtype, device=self.device |
| 63 | ) |
| 64 | self.permuted_experts_id = torch.empty( |
| 65 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 66 | ) |
| 67 | self.sorted_row_idx = torch.empty( |
| 68 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 69 | ) |
| 70 | self.topk_ids_int32 = torch.empty( |
| 71 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 72 | ) |
| 73 | self.topk_ids_for_sort = torch.empty( |
| 74 | self.max_expanded_rows, dtype=torch.int32, device=self.device |
| 75 | ) |
| 76 | sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size( |
| 77 | self.max_expanded_rows, self.num_experts |
| 78 | ) |
| 79 | self.sort_workspace = torch.empty( |
| 80 | sorter_size, dtype=torch.int8, device=self.device |
| 81 | ) |
| 82 | # torch.device("cuda") in config, after initialized, |
| 83 | # will be changed to cuda:{index}, so we need to refresh here. |
| 84 | self.device = self.token_expert_indices.device |
| 85 | |
| 86 | def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: |
| 87 | n_token, n_hidden = hidden_states.shape |
| 88 | assert hidden_states.device == self.device |
| 89 | assert topk_ids.device == self.device |
| 90 | assert n_token <= self.max_num_tokens |
| 91 | assert topk_ids.size(1) == self.topk |
| 92 | assert topk_ids.size(0) == n_token |
| 93 | if self.hidden_size is not None: |
| 94 | assert n_hidden == self.hidden_size |
| 95 | assert hidden_states.dtype == self.hidden_dtype |
| 96 | assert self.permuted_hidden_states is not None |
| 97 | |
| 98 | def token_expert_indices_view(self, n_token: int) -> torch.Tensor: |
| 99 | return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk) |
| 100 | |
| 101 | def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: |
| 102 | if topk_ids.dtype == torch.int32: |
| 103 | return topk_ids |
| 104 | numel = topk_ids.numel() |
| 105 | topk_ids_int32 = self.topk_ids_int32[:numel].view_as(topk_ids) |
| 106 | topk_ids_int32.copy_(topk_ids) |
| 107 | return topk_ids_int32 |
| 108 | |
| 109 | |
| 110 | def moe_prepare_scatter( |
| 111 | topk_ids: torch.Tensor, |
| 112 | expert_map: torch.Tensor | None, |
| 113 | scratch: MoEPermuteScratch, |
| 114 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 115 | """Generate expert offsets and shared scatter/unpermute destination indices.""" |
| 116 | assert topk_ids.device == scratch.device |
| 117 | assert topk_ids.size(0) <= scratch.max_num_tokens |
| 118 | assert topk_ids.size(1) == scratch.topk |
| 119 | n_token, topk = topk_ids.shape |
| 120 | expanded_rows = topk_ids.numel() |
| 121 | inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) |
| 122 | if expanded_rows == 0: |
| 123 | scratch.expert_first_token_offset.zero_() |
| 124 | return scratch.expert_first_token_offset, inverse |
| 125 | torch.ops._moe_C.moe_prepare_scatter( |
| 126 | scratch.prepare_topk_ids(topk_ids), |
| 127 | scratch.token_expert_indices_view(n_token), |
| 128 | expert_map, |
| 129 | scratch.num_experts, |
| 130 | scratch.num_local_experts, |
| 131 | scratch.expert_first_token_offset, |
| 132 | inverse, |
| 133 | scratch.sort_workspace, |
| 134 | scratch.permuted_experts_id[:expanded_rows].view(n_token, topk), |
| 135 | scratch.sorted_row_idx[:expanded_rows].view(n_token, topk), |
| 136 | scratch.topk_ids_for_sort[:expanded_rows].view(n_token, topk), |
| 137 | ) |
| 138 | return scratch.expert_first_token_offset, inverse |
| 139 | |
| 140 | |
| 141 | def get_moe_permute_scratch( |
| 142 | *, |
| 143 | max_num_tokens: int, |
| 144 | topk: int, |
| 145 | num_experts: int, |
| 146 | num_local_experts: int, |
| 147 | device: torch.device, |
| 148 | hidden_size: int | None = None, |
| 149 | hidden_dtype: torch.dtype | None = None, |
| 150 | ) -> MoEPermuteScratch: |
| 151 | """Share scratch across sequential layers in the current ubatch and lane. |
| 152 | |
| 153 | Without a workspace manager, allocate new scratch for each call. |
| 154 | """ |
| 155 | |
| 156 | def create_scratch() -> MoEPermuteScratch: |
| 157 | return MoEPermuteScratch( |
| 158 | max_num_tokens=max_num_tokens, |
| 159 | topk=topk, |
| 160 | num_experts=num_experts, |
| 161 | num_local_experts=num_local_experts, |
| 162 | device=device, |
| 163 | hidden_size=hidden_size, |
| 164 | hidden_dtype=hidden_dtype, |
| 165 | ) |
| 166 | |
| 167 | if not is_workspace_manager_initialized(): |
| 168 | return create_scratch() |
| 169 | |
| 170 | key = ( |
| 171 | MoEPermuteScratch, |
| 172 | max_num_tokens, |
| 173 | topk, |
| 174 | num_experts, |
| 175 | num_local_experts, |
| 176 | device, |
| 177 | hidden_size, |
| 178 | hidden_dtype, |
| 179 | ) |
| 180 | return current_workspace_manager().get_persistent_resource(key, create_scratch) |
| 181 | |
| 182 | |
| 183 | def moe_permute( |
| 184 | hidden_states: torch.Tensor, |
| 185 | a1q_scale: torch.Tensor | None, |
| 186 | topk_ids: torch.Tensor, |
| 187 | n_expert: int, |
| 188 | n_local_expert: int = -1, |
| 189 | expert_map: torch.Tensor | None = None, |
| 190 | permuted_hidden_states: torch.Tensor | None = None, |
| 191 | scratch: MoEPermuteScratch | None = None, |
| 192 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 193 | """This function expands and permutes activation to gather uncontinuous tokens |
| 194 | for each expert. |
| 195 | |
| 196 | Args: |
| 197 | hidden_states (torch.Tensor): The input tensor to the MoE layer. |
| 198 | a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states |
| 199 | topk_ids (torch.Tensor): topk expert route id for each token. |
| 200 | n_expert (int): The number of expert. |
| 201 | n_local_expert (int): The number of expert in current EP rank. |
| 202 | expert_map (Optional[torch.Tensor]): A tensor mapping expert indices |
| 203 | from the global expert space to the local expert space of the expert |
| 204 | parallel shard. |
| 205 | permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. |
| 206 | If None, the output tensor will be created in this function. |
| 207 | scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch |
| 208 | buffers. Validated against hidden_states and topk_ids when given, |
| 209 | otherwise the buffers are allocated in this function. |
| 210 | |
| 211 | Returns: |
| 212 | - permuted_hidden_states (torch.Tensor): permuted activation. |
| 213 | - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states |
| 214 | if original scale not per-tensor scaling |
| 215 | - expert_first_token_offset (torch.Tensor): offset of the first token |
| 216 | of each expert for standard grouped gemm. |
| 217 | - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute. |
| 218 | - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden. |
| 219 | |
| 220 | """ |
| 221 | n_token, n_hidden = hidden_states.size() |
| 222 | topk = topk_ids.size(1) |
| 223 | assert (n_hidden * hidden_states.element_size()) % 16 == 0, ( |
| 224 | "permue kernel need hidden dim align to 16B" |
| 225 | ) |
| 226 | permuted_row_size = n_token * topk |
| 227 | if n_local_expert == -1: |
| 228 | n_local_expert = n_expert |
| 229 | if permuted_hidden_states is None: |
| 230 | if scratch is None: |
| 231 | permuted_hidden_states = torch.empty( |
| 232 | (permuted_row_size, n_hidden), |
| 233 | dtype=hidden_states.dtype, |
| 234 | device=hidden_states.device, |
| 235 | ) |
| 236 | else: |
| 237 | scratch.validate(hidden_states, topk_ids) |
| 238 | hidden_numel = permuted_row_size * n_hidden |
| 239 | scratch_hidden_states = scratch.permuted_hidden_states |
| 240 | assert scratch_hidden_states is not None |
| 241 | permuted_hidden_states = scratch_hidden_states[:hidden_numel].view( |
| 242 | permuted_row_size, n_hidden |
| 243 | ) |
| 244 | assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), ( |
| 245 | f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}" |
| 246 | f" but got {permuted_hidden_states.size()}" |
| 247 | ) |
| 248 | |
| 249 | if scratch is None: |
| 250 | token_expert_indices = torch.arange( |
| 251 | 0, n_token * topk, dtype=torch.int32, device=hidden_states.device |
| 252 | ).reshape((n_token, topk)) |
| 253 | |
| 254 | expert_first_token_offset = torch.empty( |
| 255 | n_local_expert + 1, dtype=torch.int64, device=hidden_states.device |
| 256 | ) |
| 257 | permuted_idx = torch.full( |
| 258 | (permuted_row_size,), |
| 259 | n_token * topk, |
| 260 | dtype=torch.int32, |
| 261 | device=hidden_states.device, |
| 262 | ) |
| 263 | inv_permuted_idx = torch.empty( |
| 264 | (n_token, topk), dtype=torch.int32, device=hidden_states.device |
| 265 | ) |
| 266 | topk_ids_int32 = topk_ids.to(torch.int32) |
| 267 | torch.ops._moe_C.moe_permute( |
| 268 | hidden_states, |
| 269 | topk_ids_int32, |
| 270 | token_expert_indices, |
| 271 | expert_map, |
| 272 | n_expert, |
| 273 | n_local_expert, |
| 274 | topk, |
| 275 | permuted_hidden_states, |
| 276 | expert_first_token_offset, |
| 277 | inv_permuted_idx, |
| 278 | permuted_idx, |
| 279 | ) |
| 280 | else: |
| 281 | scratch.validate(hidden_states, topk_ids) |
| 282 | assert n_expert == scratch.num_experts |
| 283 | assert n_local_expert == scratch.num_local_experts |
| 284 | token_expert_indices = scratch.token_expert_indices_view(n_token) |
| 285 | expert_first_token_offset = scratch.expert_first_token_offset |
| 286 | permuted_idx = scratch.permuted_idx[:permuted_row_size] |
| 287 | permuted_idx.fill_(permuted_row_size) |
| 288 | inv_permuted_idx = scratch.inv_permuted_idx[:permuted_row_size].view( |
| 289 | n_token, topk |
| 290 | ) |
| 291 | permuted_experts_id = scratch.permuted_experts_id[:permuted_row_size].view( |
| 292 | n_token, topk |
| 293 | ) |
| 294 | sorted_row_idx = scratch.sorted_row_idx[:permuted_row_size].view(n_token, topk) |
| 295 | topk_ids_for_sort = scratch.topk_ids_for_sort[:permuted_row_size].view( |
| 296 | n_token, topk |
| 297 | ) |
| 298 | topk_ids_int32 = scratch.prepare_topk_ids(topk_ids) |
| 299 | torch.ops._moe_C.moe_permute_with_scratch( |
| 300 | hidden_states, |
| 301 | topk_ids_int32, |
| 302 | token_expert_indices, |
| 303 | expert_map, |
| 304 | n_expert, |
| 305 | n_local_expert, |
| 306 | topk, |
| 307 | permuted_hidden_states, |
| 308 | expert_first_token_offset, |
| 309 | inv_permuted_idx, |
| 310 | permuted_idx, |
| 311 | scratch.sort_workspace, |
| 312 | permuted_experts_id, |
| 313 | sorted_row_idx, |
| 314 | topk_ids_for_sort, |
| 315 | ) |
| 316 | |
| 317 | if a1q_scale is not None and a1q_scale.dim() > 1: |
| 318 | a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) // topk] |
| 319 | return ( |
| 320 | permuted_hidden_states, |
| 321 | a1q_scale, |
| 322 | expert_first_token_offset, |
| 323 | inv_permuted_idx.flatten(), |
| 324 | permuted_idx, |
| 325 | ) |
| 326 | |
| 327 | |
| 328 | def moe_unpermute( |
| 329 | out: torch.Tensor, |
| 330 | permuted_hidden_states: torch.Tensor, |
| 331 | topk_weights: torch.Tensor, |
| 332 | inv_permuted_idx: torch.Tensor, |
| 333 | expert_first_token_offset: torch.Tensor | None = None, |
| 334 | ) -> None: |
| 335 | """This function expands and permutes activation to gathering uncontinuous |
| 336 | tokens for each expert. |
| 337 | |
| 338 | Args: |
| 339 | out (torch.Tensor): output tensor |
| 340 | permuted_hidden_states (torch.Tensor): permuted activation. |
| 341 | topk_weights (torch.Tensor): topk expert route weight for each token. |
| 342 | inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute. |
| 343 | expert_first_token_offset (Optional[torch.Tensor]): offset of the first |
| 344 | token of each expert for grouped gemm. |
| 345 | |
| 346 | Returns: |
| 347 | - hidden_states (torch.Tensor): The reduced and unpermuted activation |
| 348 | tensor. |
| 349 | |
| 350 | """ |
| 351 | topk = topk_weights.size(1) |
| 352 | n_hidden = permuted_hidden_states.size(-1) |
| 353 | assert (n_hidden * permuted_hidden_states.element_size()) % 16 == 0, ( |
| 354 | "unpermue kernel need hidden dim align to 16B" |
| 355 | ) |
| 356 | |
| 357 | torch.ops._moe_C.moe_unpermute( |
| 358 | permuted_hidden_states, |
| 359 | topk_weights, |
| 360 | inv_permuted_idx, |
| 361 | expert_first_token_offset, |
| 362 | topk, |
| 363 | out, |
| 364 | ) |
| 365 | |
| 366 | |
| 367 | def moe_permute_unpermute_supported(): |
| 368 | return torch.ops._moe_C.moe_permute_unpermute_supported() |