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