| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Based on: |
| 4 | Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). |
| 5 | Punica: Multi-Tenant LoRA Serving. |
| 6 | https://arxiv.org/abs/2310.18547 |
| 7 | """ |
| 8 | |
| 9 | from typing import Any, final |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm import _custom_ops as ops |
| 14 | from vllm.lora.layers import LoRAMapping |
| 15 | from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink |
| 16 | from vllm.lora.utils import get_captured_lora_counts |
| 17 | from vllm.triton_utils import HAS_TRITON, triton |
| 18 | from vllm.utils.math_utils import round_up |
| 19 | |
| 20 | if HAS_TRITON: |
| 21 | from vllm.lora.ops.triton_ops import ( |
| 22 | LoRAKernelMeta, |
| 23 | fused_moe_lora, |
| 24 | ) |
| 25 | |
| 26 | from .punica_base import PunicaWrapperBase |
| 27 | |
| 28 | |
| 29 | @final |
| 30 | class PunicaWrapperXPU(PunicaWrapperBase): |
| 31 | """PunicaWrapperXPU is designed to manage and provide metadata for the punica |
| 32 | kernel. The main function is to maintain the state information for |
| 33 | Multi-LoRA, and to provide the interface for the punica ipex kernel. |
| 34 | """ |
| 35 | |
| 36 | def __init__( |
| 37 | self, |
| 38 | max_num_batched_tokens: int, |
| 39 | max_batches: int, |
| 40 | device: torch.device | str, |
| 41 | **kwargs, |
| 42 | ): |
| 43 | PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device) |
| 44 | torch._dynamo.mark_dynamic(self._token_lora_indices, 0) |
| 45 | torch._dynamo.mark_dynamic(self._embeddings_indices, 0) |
| 46 | torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0) |
| 47 | |
| 48 | self.lora_config = kwargs["lora_config"] |
| 49 | self.max_loras = self.lora_config.max_loras |
| 50 | |
| 51 | # Compute captured LoRA counts for cudagraph specialization. |
| 52 | captured_lora_counts = get_captured_lora_counts( |
| 53 | self.max_loras, self.lora_config.specialize_active_lora |
| 54 | ) |
| 55 | |
| 56 | self.token_mapping_meta = LoRAKernelMeta.make( |
| 57 | self.max_loras, |
| 58 | max_num_batched_tokens, |
| 59 | device=device, |
| 60 | captured_lora_counts=captured_lora_counts, |
| 61 | ) |
| 62 | |
| 63 | # When speculative decoding is enabled, max_num_samples is |
| 64 | # max_batches * (num_speculative_decoding_tokens + 1). |
| 65 | # This line can be optimized by replacing max_num_batched_tokens |
| 66 | # to max_batches * (num_speculative_decoding_tokens + 1). |
| 67 | self.prompt_mapping_meta = LoRAKernelMeta.make( |
| 68 | self.max_loras, |
| 69 | max_num_batched_tokens, |
| 70 | device=device, |
| 71 | captured_lora_counts=captured_lora_counts, |
| 72 | ) |
| 73 | |
| 74 | def update_metadata( |
| 75 | self, |
| 76 | mapping: LoRAMapping, |
| 77 | lora_index_to_id: list[int | None], |
| 78 | max_loras: int, |
| 79 | vocab_size: int, |
| 80 | **kwargs, |
| 81 | ): |
| 82 | self.is_prefill = mapping.is_prefill |
| 83 | self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size) |
| 84 | |
| 85 | # Prepare kernel metadata tensors |
| 86 | self.token_mapping_meta.prepare_tensors(self.token_lora_indices) |
| 87 | self.prompt_mapping_meta.prepare_tensors(self.sampler_indices) |
| 88 | |
| 89 | def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor: |
| 90 | return torch.narrow(self._token_lora_indices, 0, 0, x.size(0)) |
| 91 | |
| 92 | def _apply_shrink( |
| 93 | self, |
| 94 | y: torch.Tensor, |
| 95 | x: torch.Tensor, |
| 96 | w_t_all: torch.Tensor, |
| 97 | scale: float, |
| 98 | ): |
| 99 | buf = torch.zeros( |
| 100 | x.size(0), |
| 101 | w_t_all.size(-2), |
| 102 | dtype=x.dtype, |
| 103 | device=x.device, |
| 104 | ) |
| 105 | bgmv_shrink(x, w_t_all, buf, self._get_token_lora_indices(x), scale) |
| 106 | y.copy_(buf) |
| 107 | |
| 108 | def _apply_expand( |
| 109 | self, |
| 110 | y: torch.Tensor, |
| 111 | x: torch.Tensor, |
| 112 | w_t_all: torch.Tensor, |
| 113 | y_offset: int, |
| 114 | y_slice_size: int, |
| 115 | add_inputs: bool, |
| 116 | ): |
| 117 | token_lora_indices = self._get_token_lora_indices(x) |
| 118 | # After tensor-parallel all-gather (non-fully-sharded LoRA), x may |
| 119 | # have been gathered along the rank dim so x.size(1) == max_lora_rank |
| 120 | # * tp_size, while lora_b only uses max_lora_rank elements. The XPU |
| 121 | # C++ kernel requires inputs.size(1) == lora_b.size(-1), so truncate |
| 122 | # to the actual rank. x[:, :rank] is non-contiguous, hence the copy. |
| 123 | rank = w_t_all.size(-1) |
| 124 | if x.size(1) != rank: |
| 125 | x = x[:, :rank].contiguous() |
| 126 | bgmv_expand_slice( |
| 127 | x, w_t_all, y, token_lora_indices, y_offset, y_slice_size, add_inputs |
| 128 | ) |
| 129 | |
| 130 | def add_shrink( |
| 131 | self, |
| 132 | y: torch.Tensor, |
| 133 | x: torch.Tensor, |
| 134 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 135 | scale: float, |
| 136 | **kwargs: Any, |
| 137 | ): |
| 138 | """Performs GEMM for multiple slices of lora_a. |
| 139 | |
| 140 | Semantics: |
| 141 | for i in range(len(lora_a_stacked)): |
| 142 | y[i] += (x @ lora_a_stacked[i]) * scale |
| 143 | |
| 144 | Args: |
| 145 | y (torch.Tensor): Output tensors |
| 146 | x (torch.Tensor): Input tensor |
| 147 | lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights |
| 148 | scale (float): Scaling factor for the operation |
| 149 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 150 | |
| 151 | """ |
| 152 | x = x.view(-1, x.shape[-1]) |
| 153 | for slice_idx in range(len(lora_a_stacked)): |
| 154 | self._apply_shrink(y[slice_idx], x, lora_a_stacked[slice_idx], scale) |
| 155 | |
| 156 | def add_expand( |
| 157 | self, |
| 158 | y: torch.Tensor, |
| 159 | x: torch.Tensor, |
| 160 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 161 | output_slices: tuple[int, ...], |
| 162 | offset_start: int = 0, |
| 163 | add_inputs=True, |
| 164 | **kwargs: Any, |
| 165 | ) -> None: |
| 166 | """Performs GEMM for multiple slices of lora_b. |
| 167 | |
| 168 | Semantics: |
| 169 | for i in range(len(lora_b_stacked)): |
| 170 | slice = output_slices[i] |
| 171 | y[:, offset:offset+slice] += x[i] @ lora_b_stacked[i] |
| 172 | offset += slice |
| 173 | |
| 174 | Args: |
| 175 | y (torch.Tensor): Output tensor. |
| 176 | x (torch.Tensor): Input tensors |
| 177 | lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight |
| 178 | output_slices (tuple[int, ...]): Every slice's size |
| 179 | offset_start (int): The starting position of y, defaults to 0. |
| 180 | add_inputs (bool): Defaults to True. |
| 181 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 182 | |
| 183 | """ |
| 184 | y_org = y |
| 185 | y = y.view(-1, y.shape[-1]) |
| 186 | |
| 187 | assert x.ndim == 3 |
| 188 | assert x.size(0) == len(output_slices) |
| 189 | |
| 190 | # TODO fuse these kernels |
| 191 | for slice_idx in range(len(lora_b_stacked)): |
| 192 | self._apply_expand( |
| 193 | y, |
| 194 | x[slice_idx], |
| 195 | lora_b_stacked[slice_idx], |
| 196 | offset_start, |
| 197 | output_slices[slice_idx], |
| 198 | add_inputs=add_inputs, |
| 199 | ) |
| 200 | offset_start += output_slices[slice_idx] |
| 201 | y = y.view_as(y_org) |
| 202 | |
| 203 | def add_lora_embedding( |
| 204 | self, |
| 205 | y: torch.Tensor, |
| 206 | x: torch.Tensor, |
| 207 | lora_b_stacked: torch.Tensor, |
| 208 | add_inputs: bool = True, |
| 209 | **kwargs: Any, |
| 210 | ) -> None: |
| 211 | """Applies lora specifically for VocabParallelEmbeddingWithLoRA. |
| 212 | |
| 213 | Semantics: |
| 214 | y += x @ lora_b_stacked |
| 215 | |
| 216 | Args: |
| 217 | y (torch.Tensor): Output tensor. |
| 218 | x (torch.Tensor): Input tensor. |
| 219 | lora_b_stacked (torch.Tensor): lora_b's weights. |
| 220 | add_inputs (bool): Default to True. |
| 221 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 222 | |
| 223 | """ |
| 224 | token_lora_indices = self._get_token_lora_indices(x) |
| 225 | bgmv_expand(x, lora_b_stacked, y, token_lora_indices, add_inputs) |
| 226 | |
| 227 | def add_lora_linear( |
| 228 | self, |
| 229 | y: torch.Tensor, |
| 230 | x: torch.Tensor, |
| 231 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 232 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 233 | scale: float, |
| 234 | output_slices: tuple[int, ...], |
| 235 | *, |
| 236 | buffer: torch.Tensor | None = None, |
| 237 | **kwargs: Any, |
| 238 | ) -> None: |
| 239 | """Applicable to linear-related lora. |
| 240 | |
| 241 | Semantics: |
| 242 | for i in range(len(lora_a_stacked)): |
| 243 | y[i] += ( |
| 244 | x[i].unsqueeze(0) |
| 245 | @ lora_a_stacked[indices[i], layer_idx, :, :] |
| 246 | @ lora_b_stacked[indices[i], layer_idx, :, :] |
| 247 | * scale |
| 248 | ).squeeze(0) |
| 249 | |
| 250 | Args: |
| 251 | y (torch.Tensor): Output tensor. Will be changed in-place. |
| 252 | x (torch.Tensor): Input tensor |
| 253 | lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight. |
| 254 | lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight. |
| 255 | scale (float): Scaling factor. |
| 256 | output_slices (tuple[int, ...]): Every slice's size. |
| 257 | buffer (Optional[torch.Tensor]): Defaults to None. |
| 258 | **kwargs: Accepts `add_inputs` (bool, defaults to True); any other |
| 259 | keys are ignored. |
| 260 | |
| 261 | """ |
| 262 | assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices) |
| 263 | |
| 264 | assert buffer is None, ( |
| 265 | "To minimize overhead, the buffer should be created by " |
| 266 | ".add_lora_linear() instead of being passed in." |
| 267 | ) |
| 268 | r = lora_b_stacked[0].size(-1) |
| 269 | buffer = torch.zeros( # type: ignore |
| 270 | (len(output_slices), x.size(0), r), |
| 271 | dtype=x.dtype, |
| 272 | device=x.device, |
| 273 | ) |
| 274 | add_inputs = kwargs.pop("add_inputs", True) |
| 275 | self.add_shrink( |
| 276 | buffer, # type: ignore |
| 277 | x, |
| 278 | lora_a_stacked, |
| 279 | scale, |
| 280 | **kwargs, |
| 281 | ) |
| 282 | self.add_expand( |
| 283 | y, |
| 284 | buffer, # type: ignore |
| 285 | lora_b_stacked, |
| 286 | output_slices, |
| 287 | add_inputs=add_inputs, |
| 288 | **kwargs, |
| 289 | ) |
| 290 | |
| 291 | @property |
| 292 | def sampler_indices_padded(self) -> torch.Tensor: |
| 293 | """This property provides access to padded sampler indices.""" |
| 294 | return self._sampler_indices_padded[:] |
| 295 | |
| 296 | def add_lora_logits( |
| 297 | self, |
| 298 | y: torch.Tensor, |
| 299 | x: torch.Tensor, |
| 300 | lora_a_stacked: torch.Tensor, |
| 301 | lora_b_stacked: torch.Tensor, |
| 302 | scale, |
| 303 | *, |
| 304 | buffer: torch.Tensor | None = None, |
| 305 | **kwargs: Any, |
| 306 | ) -> None: |
| 307 | """Applies lora specifically for LogitsProcessorWithLoRA. |
| 308 | |
| 309 | Semantics: |
| 310 | buffer = (x @ lora_a_stacked) * scale |
| 311 | y += buffer @ lora_b_stacked |
| 312 | |
| 313 | Args: |
| 314 | y (torch.Tensor): Output tensor. |
| 315 | x (torch.Tensor): Input tensor. |
| 316 | lora_a_stacked (torch.Tensor): lora_a's weights. |
| 317 | lora_b_stacked (torch.Tensor): lora_b's weights. |
| 318 | scale (float): Scaling factor. |
| 319 | buffer (Optional[torch.Tensor]): Default to None. |
| 320 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 321 | |
| 322 | """ |
| 323 | y_org = y |
| 324 | y = y.view(-1, y.shape[-1]) |
| 325 | x = x.view(-1, x.shape[-1]) |
| 326 | r = lora_b_stacked.size(-1) |
| 327 | |
| 328 | assert buffer is None, ( |
| 329 | "To minimize overhead, the buffer should be created by " |
| 330 | ".add_lora_linear() instead of being passed in." |
| 331 | ) |
| 332 | buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device) |
| 333 | sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0)) |
| 334 | bgmv_shrink(x, lora_a_stacked, buffer, sampler_indices, scale) |
| 335 | bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True) |
| 336 | y = y.view_as(y_org) |
| 337 | |
| 338 | def apply_lora_full_linear( |
| 339 | self, |
| 340 | y: torch.Tensor, |
| 341 | x: torch.Tensor, |
| 342 | weight_stacked: torch.Tensor, |
| 343 | bias_stacked: torch.Tensor, |
| 344 | module_enabled: torch.Tensor, |
| 345 | ) -> None: |
| 346 | indices = self.sampler_indices |
| 347 | adapter_y = torch.zeros( |
| 348 | (x.size(0), weight_stacked.size(-2)), |
| 349 | dtype=torch.float32, |
| 350 | device=x.device, |
| 351 | ) |
| 352 | bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) |
| 353 | result = self._select_full_linear_output( |
| 354 | y, adapter_y, bias_stacked, module_enabled |
| 355 | ) |
| 356 | y.copy_(result) |
| 357 | |
| 358 | def moe_lora_align_block_size( |
| 359 | self, |
| 360 | topk_ids: torch.Tensor, |
| 361 | num_tokens: int, |
| 362 | block_size: int, |
| 363 | num_experts: int, |
| 364 | max_loras: int, |
| 365 | adapter_enabled: torch.Tensor, |
| 366 | expert_map: torch.Tensor | None = None, |
| 367 | pad_sorted_ids: bool = False, |
| 368 | naive_block_assignment: bool = False, |
| 369 | token_lora_mapping: torch.Tensor | None = None, |
| 370 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 371 | """Aligns tokens and experts into block-sized chunks for LoRA-based |
| 372 | mixture-of-experts (MoE) execution. |
| 373 | |
| 374 | When `token_lora_mapping` is provided, it overrides the global mapping |
| 375 | read from `self.token_mapping_meta`. This is how EP+LoRA injects the |
| 376 | per-rank-local token→LoRA map after all-to-all dispatch. |
| 377 | """ |
| 378 | ( |
| 379 | token_lora_mapping_meta, |
| 380 | _, |
| 381 | _, |
| 382 | _, |
| 383 | lora_ids, |
| 384 | _, |
| 385 | _, |
| 386 | ) = self.token_mapping_meta.meta_args( |
| 387 | num_tokens, self.lora_config.specialize_active_lora |
| 388 | ) |
| 389 | if token_lora_mapping is None: |
| 390 | token_lora_mapping = token_lora_mapping_meta |
| 391 | # Under EP the caller passes local_num_experts but topk_ids carries |
| 392 | # GLOBAL expert indices. The CUDA kernel uses num_experts to size |
| 393 | # its bucketing table; with EP we must size by global_num_experts |
| 394 | # so global topk_ids don't overflow. expert_map inside the kernel |
| 395 | # then translates global→local so the output expert_ids are local |
| 396 | # (mirrors the non-LoRA moe_align_block_size behavior). |
| 397 | kernel_num_experts = ( |
| 398 | expert_map.numel() if expert_map is not None else num_experts |
| 399 | ) |
| 400 | if naive_block_assignment: |
| 401 | expert_ids = topk_ids.reshape(-1) |
| 402 | sorted_ids = None |
| 403 | num_tokens_post_pad = None |
| 404 | else: |
| 405 | max_num_tokens_padded = topk_ids.numel() + kernel_num_experts * ( |
| 406 | block_size - 1 |
| 407 | ) |
| 408 | if pad_sorted_ids: |
| 409 | max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) |
| 410 | if topk_ids.numel() < kernel_num_experts: |
| 411 | max_num_tokens_padded = topk_ids.numel() * block_size |
| 412 | sorted_ids = torch.empty( |
| 413 | (max_loras * max_num_tokens_padded,), |
| 414 | dtype=torch.int32, |
| 415 | device=topk_ids.device, |
| 416 | ) |
| 417 | max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) |
| 418 | # Expert ids are initialized to -1 so unused (lora, expert) |
| 419 | # slots don't drive the LoRA Triton kernel into the wrong bucket. |
| 420 | # The kernel overwrites only active slots. |
| 421 | expert_ids = torch.full( |
| 422 | (max_loras * max_num_m_blocks,), |
| 423 | -1, |
| 424 | dtype=torch.int32, |
| 425 | device=topk_ids.device, |
| 426 | ) |
| 427 | num_tokens_post_pad = torch.empty( |
| 428 | (max_loras), dtype=torch.int32, device=topk_ids.device |
| 429 | ) |
| 430 | |
| 431 | ops.moe_lora_align_block_size( |
| 432 | topk_ids, |
| 433 | token_lora_mapping, |
| 434 | kernel_num_experts, |
| 435 | block_size, |
| 436 | max_loras, |
| 437 | max_num_tokens_padded, |
| 438 | max_num_m_blocks, |
| 439 | sorted_ids, |
| 440 | expert_ids, |
| 441 | num_tokens_post_pad, |
| 442 | adapter_enabled, |
| 443 | lora_ids, |
| 444 | expert_map, |
| 445 | ) |
| 446 | |
| 447 | return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad |
| 448 | |
| 449 | def add_lora_fused_moe( |
| 450 | self, |
| 451 | y: torch.Tensor, |
| 452 | x: torch.Tensor, |
| 453 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 454 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 455 | topk_weights: torch.Tensor, |
| 456 | sorted_token_ids: torch.Tensor | None, |
| 457 | expert_ids: torch.Tensor, |
| 458 | num_tokens_post_padded: torch.Tensor | None, |
| 459 | max_lora_rank: int, |
| 460 | top_k_num: int, |
| 461 | shrink_config, |
| 462 | expand_config, |
| 463 | adapter_enabled: torch.Tensor, |
| 464 | mul_routed_weight=False, |
| 465 | fully_sharded: bool = False, |
| 466 | offset: int = 0, |
| 467 | token_lora_mapping: torch.Tensor | None = None, |
| 468 | add_inputs: bool = True, |
| 469 | ): |
| 470 | """Perform a fused forward computation for a LoRA MoE layer.""" |
| 471 | ( |
| 472 | token_lora_mapping_meta, |
| 473 | _, |
| 474 | _, |
| 475 | _, |
| 476 | lora_ids, |
| 477 | no_lora_flag, |
| 478 | num_active_loras, |
| 479 | ) = self.token_mapping_meta.meta_args( |
| 480 | x.size(0), self.lora_config.specialize_active_lora |
| 481 | ) |
| 482 | |
| 483 | assert no_lora_flag.numel() == 1 |
| 484 | if no_lora_flag.item(): |
| 485 | # None of the inputs require LoRA. |
| 486 | return |
| 487 | |
| 488 | if token_lora_mapping is None: |
| 489 | token_lora_mapping = token_lora_mapping_meta |
| 490 | fused_moe_lora( |
| 491 | y, |
| 492 | x, |
| 493 | lora_a_stacked, |
| 494 | lora_b_stacked, |
| 495 | topk_weights, |
| 496 | sorted_token_ids, |
| 497 | expert_ids, |
| 498 | num_tokens_post_padded, |
| 499 | token_lora_mapping, |
| 500 | max_lora_rank, |
| 501 | top_k_num, |
| 502 | lora_ids, |
| 503 | num_active_loras, |
| 504 | adapter_enabled, |
| 505 | shrink_config.get("BLOCK_SIZE_M", 64), |
| 506 | shrink_config.get("BLOCK_SIZE_N", 64), |
| 507 | shrink_config.get("BLOCK_SIZE_K", 32), |
| 508 | shrink_config.get("GROUP_SIZE_M", 8), |
| 509 | shrink_config.get("NUM_WARPS", 4), |
| 510 | shrink_config.get("NUM_STAGES", 3), |
| 511 | shrink_config.get("SPLIT_K", 1), |
| 512 | expand_config.get("BLOCK_SIZE_M", 64), |
| 513 | expand_config.get("BLOCK_SIZE_N", 64), |
| 514 | expand_config.get("BLOCK_SIZE_K", 32), |
| 515 | expand_config.get("GROUP_SIZE_M", 8), |
| 516 | expand_config.get("NUM_WARPS", 4), |
| 517 | expand_config.get("NUM_STAGES", 3), |
| 518 | expand_config.get("SPLIT_K", 1), |
| 519 | mul_routed_weight, |
| 520 | fully_sharded, |
| 521 | offset, |
| 522 | add_inputs, |
| 523 | ) |
| 524 | |
| 525 | def add_lora_w13( |
| 526 | self, |
| 527 | y: torch.Tensor, |
| 528 | x: torch.Tensor, |
| 529 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 530 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 531 | topk_ids: torch.Tensor, |
| 532 | topk_weights: torch.Tensor, |
| 533 | expert_map: torch.Tensor | None, |
| 534 | w1: torch.Tensor, |
| 535 | w2: torch.Tensor, |
| 536 | num_tokens: int, |
| 537 | top_k_num: int, |
| 538 | max_loras: int, |
| 539 | adapter_enabled: torch.Tensor, |
| 540 | local_num_experts: int, |
| 541 | top_k: int, |
| 542 | num_slices: int, |
| 543 | fully_sharded: bool, |
| 544 | use_tuned_config: bool, |
| 545 | add_inputs: bool = True, |
| 546 | token_lora_mapping: torch.Tensor | None = None, |
| 547 | ) -> tuple[ |
| 548 | torch.Tensor | None, |
| 549 | torch.Tensor | None, |
| 550 | torch.Tensor | None, |
| 551 | torch.Tensor | None, |
| 552 | ]: |
| 553 | import functools |
| 554 | |
| 555 | from vllm.lora.layers.utils import try_get_optimal_moe_lora_config |
| 556 | from vllm.lora.ops.triton_ops.utils import ( |
| 557 | _normalize_lora_config_keys, |
| 558 | get_lora_op_configs, |
| 559 | ) |
| 560 | from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str |
| 561 | |
| 562 | config_dtype = _get_config_dtype_str( |
| 563 | dtype=x.dtype, |
| 564 | use_fp8_w8a8=False, |
| 565 | use_int8_w8a16=False, |
| 566 | use_int4_w4a16=False, |
| 567 | ) |
| 568 | max_lora_rank = lora_a_stacked[0].shape[-2] |
| 569 | |
| 570 | if use_tuned_config: |
| 571 | shrink_config = get_lora_op_configs( |
| 572 | op_type="fused_moe_lora_w13_shrink", |
| 573 | max_loras=max_loras, |
| 574 | batch=num_tokens, |
| 575 | hidden_size=x.shape[-1], |
| 576 | rank=max_lora_rank, |
| 577 | num_slices=num_slices, |
| 578 | moe_intermediate_size=lora_b_stacked[0].shape[-2], |
| 579 | ) |
| 580 | expand_config = get_lora_op_configs( |
| 581 | op_type="fused_moe_lora_w13_expand", |
| 582 | max_loras=max_loras, |
| 583 | batch=num_tokens, |
| 584 | hidden_size=x.shape[-1], |
| 585 | rank=max_lora_rank, |
| 586 | num_slices=num_slices, |
| 587 | moe_intermediate_size=lora_b_stacked[0].shape[-2], |
| 588 | ) |
| 589 | else: |
| 590 | get_config = functools.partial( |
| 591 | try_get_optimal_moe_lora_config, |
| 592 | w1_shape=w1.shape, |
| 593 | w2_shape=w2.shape, |
| 594 | rank=max_lora_rank, |
| 595 | top_k=top_k, |
| 596 | dtype=config_dtype, |
| 597 | M=num_tokens, |
| 598 | ) |
| 599 | shrink_config = get_config(op_type="fused_moe_lora_w13_shrink") |
| 600 | expand_config = get_config(op_type="fused_moe_lora_w13_expand") |
| 601 | |
| 602 | shrink_config = _normalize_lora_config_keys(shrink_config) |
| 603 | expand_config = _normalize_lora_config_keys(expand_config) |
| 604 | |
| 605 | SPARSITY_FACTOR = 8 |
| 606 | naive_block_assignment = ( |
| 607 | not fully_sharded |
| 608 | and expert_map is None |
| 609 | and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras |
| 610 | ) |
| 611 | |
| 612 | ( |
| 613 | token_lora_mapping, |
| 614 | sorted_token_ids_lora, |
| 615 | expert_ids_lora, |
| 616 | num_tokens_post_padded_lora, |
| 617 | ) = self.moe_lora_align_block_size( |
| 618 | topk_ids, |
| 619 | num_tokens, |
| 620 | int(shrink_config.get("BLOCK_SIZE_M") or 64), |
| 621 | local_num_experts, |
| 622 | max_loras, |
| 623 | adapter_enabled, |
| 624 | expert_map, |
| 625 | naive_block_assignment=naive_block_assignment, |
| 626 | token_lora_mapping=token_lora_mapping, |
| 627 | ) |
| 628 | |
| 629 | _sorted = sorted_token_ids_lora |
| 630 | _eids = expert_ids_lora |
| 631 | if _sorted is not None: |
| 632 | _eids = _eids.view(max_loras, -1) |
| 633 | _sorted = _sorted.view(max_loras, -1) |
| 634 | |
| 635 | self.add_lora_fused_moe( |
| 636 | y.view(-1, top_k_num, y.shape[-1]), |
| 637 | x, |
| 638 | lora_a_stacked, |
| 639 | lora_b_stacked, |
| 640 | topk_weights, |
| 641 | _sorted, |
| 642 | _eids, |
| 643 | num_tokens_post_padded_lora, |
| 644 | max_lora_rank, |
| 645 | top_k, |
| 646 | shrink_config, |
| 647 | expand_config, |
| 648 | adapter_enabled, |
| 649 | fully_sharded=fully_sharded, |
| 650 | token_lora_mapping=token_lora_mapping, |
| 651 | add_inputs=add_inputs, |
| 652 | ) |
| 653 | |
| 654 | return ( |
| 655 | sorted_token_ids_lora, |
| 656 | expert_ids_lora, |
| 657 | num_tokens_post_padded_lora, |
| 658 | token_lora_mapping, |
| 659 | ) |
| 660 | |
| 661 | def add_lora_w2( |
| 662 | self, |
| 663 | y: torch.Tensor, |
| 664 | x: torch.Tensor, |
| 665 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 666 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 667 | topk_weights: torch.Tensor, |
| 668 | sorted_token_ids_lora: torch.Tensor | None, |
| 669 | expert_ids_lora: torch.Tensor | None, |
| 670 | num_tokens_post_padded_lora: torch.Tensor | None, |
| 671 | token_lora_mapping: torch.Tensor | None, |
| 672 | num_tokens: int, |
| 673 | w1: torch.Tensor, |
| 674 | w2: torch.Tensor, |
| 675 | top_k_num: int, |
| 676 | max_loras: int, |
| 677 | adapter_enabled: torch.Tensor, |
| 678 | top_k: int, |
| 679 | fully_sharded: bool, |
| 680 | tp_rank: int, |
| 681 | use_tuned_config: bool, |
| 682 | add_inputs: bool = True, |
| 683 | ) -> None: |
| 684 | import functools |
| 685 | |
| 686 | from vllm.lora.layers.utils import try_get_optimal_moe_lora_config |
| 687 | from vllm.lora.ops.triton_ops.utils import ( |
| 688 | _normalize_lora_config_keys, |
| 689 | get_lora_op_configs, |
| 690 | ) |
| 691 | from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str |
| 692 | |
| 693 | config_dtype = _get_config_dtype_str( |
| 694 | dtype=x.dtype, |
| 695 | use_fp8_w8a8=False, |
| 696 | use_int8_w8a16=False, |
| 697 | use_int4_w4a16=False, |
| 698 | ) |
| 699 | max_lora_rank = lora_a_stacked[0].shape[-2] |
| 700 | |
| 701 | if use_tuned_config: |
| 702 | shrink_config = get_lora_op_configs( |
| 703 | op_type="fused_moe_lora_w2_shrink", |
| 704 | max_loras=max_loras, |
| 705 | batch=num_tokens, |
| 706 | hidden_size=y.shape[-1], |
| 707 | rank=max_lora_rank, |
| 708 | num_slices=1, |
| 709 | moe_intermediate_size=lora_a_stacked[0].shape[-1], |
| 710 | ) |
| 711 | expand_config = get_lora_op_configs( |
| 712 | op_type="fused_moe_lora_w2_expand", |
| 713 | max_loras=max_loras, |
| 714 | batch=num_tokens, |
| 715 | hidden_size=y.shape[-1], |
| 716 | rank=max_lora_rank, |
| 717 | num_slices=1, |
| 718 | moe_intermediate_size=lora_a_stacked[0].shape[-1], |
| 719 | ) |
| 720 | else: |
| 721 | get_config = functools.partial( |
| 722 | try_get_optimal_moe_lora_config, |
| 723 | w1_shape=w1.shape, |
| 724 | w2_shape=w2.shape, |
| 725 | rank=max_lora_rank, |
| 726 | top_k=top_k, |
| 727 | dtype=config_dtype, |
| 728 | M=num_tokens, |
| 729 | ) |
| 730 | shrink_config = get_config(op_type="fused_moe_lora_w2_shrink") |
| 731 | expand_config = get_config(op_type="fused_moe_lora_w2_expand") |
| 732 | |
| 733 | shrink_config = _normalize_lora_config_keys(shrink_config) |
| 734 | expand_config = _normalize_lora_config_keys(expand_config) |
| 735 | |
| 736 | _sorted = sorted_token_ids_lora |
| 737 | _eids = expert_ids_lora |
| 738 | if _sorted is not None: |
| 739 | assert _eids is not None |
| 740 | _eids = _eids.view(max_loras, -1) |
| 741 | _sorted = _sorted.view(max_loras, -1) |
| 742 | |
| 743 | # w2_lora_b shape[-2] is hidden_size // tp_size when fully_sharded |
| 744 | shard_size = lora_b_stacked[0].shape[-2] |
| 745 | offset = shard_size * tp_rank if fully_sharded else 0 |
| 746 | |
| 747 | self.add_lora_fused_moe( |
| 748 | y, |
| 749 | x, |
| 750 | lora_a_stacked, |
| 751 | lora_b_stacked, |
| 752 | topk_weights, |
| 753 | _sorted, |
| 754 | _eids, |
| 755 | num_tokens_post_padded_lora, |
| 756 | max_lora_rank, |
| 757 | top_k, |
| 758 | shrink_config, |
| 759 | expand_config, |
| 760 | adapter_enabled, |
| 761 | True, # mul_routed_weight |
| 762 | fully_sharded=fully_sharded, |
| 763 | offset=offset, |
| 764 | token_lora_mapping=token_lora_mapping, |
| 765 | add_inputs=add_inputs, |
| 766 | ) |