| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """LoRA utilities for the Model Runner V2 and cudagraph.""" |
| 4 | |
| 5 | from collections.abc import Callable |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import numpy as np |
| 9 | |
| 10 | from vllm.lora.request import LoRARequest |
| 11 | from vllm.lora.utils import get_captured_lora_counts |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.config.compilation import CompilationConfig |
| 15 | from vllm.config.lora import LoRAConfig |
| 16 | |
| 17 | NO_LORA_ID = 0 |
| 18 | |
| 19 | |
| 20 | def get_lora_capture_cases( |
| 21 | lora_config: "LoRAConfig | None", |
| 22 | compilation_config: "CompilationConfig", |
| 23 | ) -> list[int]: |
| 24 | """Return num_active_loras values for cudagraph capture. |
| 25 | |
| 26 | When cudagraph_specialize_lora=True: powers of 2 up to max_loras, plus |
| 27 | max_loras+1. When False: [0, max_loras+1]. When LoRA disabled: [0]. |
| 28 | """ |
| 29 | if lora_config is None: |
| 30 | return [0] |
| 31 | if compilation_config.cudagraph_specialize_lora: |
| 32 | specialize = getattr(lora_config, "specialize_active_lora", False) |
| 33 | captured = get_captured_lora_counts(lora_config.max_loras, specialize) |
| 34 | return [0] + [c for c in captured if c > 0] |
| 35 | return [0, lora_config.max_loras + 1] |
| 36 | |
| 37 | |
| 38 | def get_num_active_loras_for_dispatch( |
| 39 | lora_config: "LoRAConfig | None", |
| 40 | lora_state: "LoraState", |
| 41 | req_ids: list[str], |
| 42 | dummy_run: bool, |
| 43 | ) -> int: |
| 44 | """Compute num_active_loras for cudagraph dispatch.""" |
| 45 | if lora_config and not dummy_run: |
| 46 | return len(lora_state.get_activate_loras(req_ids)) |
| 47 | if dummy_run and lora_config: |
| 48 | return lora_config.max_loras + 1 |
| 49 | return 0 |
| 50 | |
| 51 | |
| 52 | def create_lora_capture_hook( |
| 53 | lora_config: "LoRAConfig | None", |
| 54 | runner: Any, |
| 55 | ) -> Callable[[int, int, int], None] | None: |
| 56 | """Create a hook to set up LoRA state before each cudagraph capture.""" |
| 57 | if lora_config is None: |
| 58 | return None |
| 59 | |
| 60 | def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None: |
| 61 | # Match InputBatch.make_dummy: distribute the remainder evenly so no |
| 62 | # dummy request exceeds ceil(num_tokens / num_reqs) tokens. |
| 63 | num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32) |
| 64 | num_extra = num_tokens % num_reqs |
| 65 | if num_extra > 0: |
| 66 | num_scheduled[-num_extra:] += 1 |
| 67 | with runner.maybe_select_dummy_loras( |
| 68 | lora_config, num_scheduled, num_active_loras=num_active_loras |
| 69 | ): |
| 70 | pass |
| 71 | |
| 72 | return hook |
| 73 | |
| 74 | |
| 75 | class LoraState: |
| 76 | def __init__(self, max_num_reqs: int): |
| 77 | self.lora_ids = np.zeros(max_num_reqs, dtype=np.int32) |
| 78 | self.lora_ids.fill(NO_LORA_ID) |
| 79 | # req_id -> lora_request |
| 80 | self.lora_requests: dict[str, LoRARequest] = {} |
| 81 | |
| 82 | def add_request( |
| 83 | self, req_id: str, req_index: int, lora_request: LoRARequest | None |
| 84 | ) -> None: |
| 85 | if lora_request is not None: |
| 86 | self.lora_requests[req_id] = lora_request |
| 87 | self.lora_ids[req_index] = lora_request.lora_int_id |
| 88 | else: |
| 89 | self.lora_ids[req_index] = NO_LORA_ID |
| 90 | |
| 91 | def remove_request(self, req_id: str) -> None: |
| 92 | self.lora_requests.pop(req_id, None) |
| 93 | |
| 94 | def make_lora_inputs( |
| 95 | self, |
| 96 | req_ids: list[str], |
| 97 | idx_mapping: np.ndarray, |
| 98 | num_scheduled_tokens: np.ndarray, |
| 99 | ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: |
| 100 | lora_ids = self.lora_ids[idx_mapping] |
| 101 | prompt_lora_mapping = tuple(lora_ids) |
| 102 | token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens)) |
| 103 | active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids) |
| 104 | return prompt_lora_mapping, token_lora_mapping, active_lora_requests |
| 105 | |
| 106 | def get_activate_loras(self, req_ids: list[str]) -> set[LoRARequest]: |
| 107 | active_lora_requests: set[LoRARequest] = set() |
| 108 | for req_id in req_ids: |
| 109 | lora_request = self.lora_requests.get(req_id) |
| 110 | if lora_request is not None: |
| 111 | active_lora_requests.add(lora_request) |
| 112 | return active_lora_requests |