| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | import json |
| 6 | from functools import lru_cache |
| 7 | from pathlib import Path |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm import envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.math_utils import next_power_of_2 |
| 16 | from vllm.utils.torch_utils import async_tensor_h2d |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | is_batch_invariant = envs.VLLM_BATCH_INVARIANT |
| 20 | |
| 21 | _LORA_A_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {} |
| 22 | _LORA_B_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {} |
| 23 | |
| 24 | |
| 25 | def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device): |
| 26 | """`_LORA_A_PTR_DICT` collects the required information during `profile_run`, |
| 27 | After this, it remains constant and subsequent usage is through LUT. |
| 28 | Refer to: |
| 29 | https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py |
| 30 | """ |
| 31 | key = tuple(lora_weight.data_ptr() for lora_weight in lora_a_weights) |
| 32 | |
| 33 | if values := _LORA_A_PTR_DICT.get(key): |
| 34 | return values |
| 35 | |
| 36 | lora_strides_d0 = [] |
| 37 | lora_strides_d1 = [] |
| 38 | lora_strides_d2 = [] |
| 39 | tensor_ptrs = [] |
| 40 | for lora_a_weight in lora_a_weights: |
| 41 | if lora_a_weight.ndim == 4: # shape:(lora_num,1,size,rank) |
| 42 | assert lora_a_weight.size(1) == 1 |
| 43 | lora_a_weight = lora_a_weight.squeeze(dim=1) |
| 44 | else: |
| 45 | assert lora_a_weight.ndim == 3 # shape:(lora_num,size,rank) |
| 46 | assert lora_a_weight.is_contiguous() |
| 47 | tensor_ptrs.append(lora_a_weight.data_ptr()) |
| 48 | lora_strides_d0.append(lora_a_weight.stride(0)) |
| 49 | lora_strides_d1.append(lora_a_weight.stride(1)) |
| 50 | lora_strides_d2.append(lora_a_weight.stride(2)) |
| 51 | if len(lora_a_weights) > 1: |
| 52 | lora_ptr_tensor = async_tensor_h2d( |
| 53 | tensor_ptrs, dtype=torch.uint64, device=device |
| 54 | ) |
| 55 | else: |
| 56 | lora_ptr_tensor = lora_a_weights[0] |
| 57 | |
| 58 | if ( |
| 59 | len(set(lora_strides_d0)) > 1 |
| 60 | or len(set(lora_strides_d1)) > 1 |
| 61 | or len(set(lora_strides_d2)) > 1 |
| 62 | ): |
| 63 | raise ValueError("All LoRA weights must have the same stride.") |
| 64 | |
| 65 | _LORA_A_PTR_DICT[key] = ( |
| 66 | lora_ptr_tensor, |
| 67 | lora_strides_d0[0], |
| 68 | lora_strides_d1[0], |
| 69 | lora_strides_d2[0], |
| 70 | ) |
| 71 | return _LORA_A_PTR_DICT.get(key) |
| 72 | |
| 73 | |
| 74 | def _get_lora_b_ptr( |
| 75 | lora_weights: list[torch.Tensor], offset_start: int, device: torch.device |
| 76 | ): |
| 77 | """`_LORA_B_PTR_DICT` collects the required information during `profile_run`, |
| 78 | After this, it remains constant and subsequent usage is through LUT. |
| 79 | Refer to: |
| 80 | https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py |
| 81 | |
| 82 | """ |
| 83 | key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights) |
| 84 | if values := _LORA_B_PTR_DICT.get(key): |
| 85 | return values |
| 86 | slice_offset_lst = [] |
| 87 | tensor_ptrs = [] |
| 88 | lora_strides_d0 = [] |
| 89 | lora_strides_d1 = [] |
| 90 | lora_strides_d2 = [] |
| 91 | hidden_sizes = [] |
| 92 | slice_offset = offset_start |
| 93 | for lora_b_weight in lora_weights: |
| 94 | if lora_b_weight.ndim == 4: # shape:(lora_num,1,size,rank) |
| 95 | assert lora_b_weight.size(1) == 1 |
| 96 | lora_b_weight = lora_b_weight.squeeze(dim=1) |
| 97 | else: |
| 98 | assert lora_b_weight.ndim == 3 # shape:(lora_num,size,rank) |
| 99 | assert lora_b_weight.is_contiguous() |
| 100 | tensor_ptrs.append(lora_b_weight.data_ptr()) |
| 101 | lora_strides_d0.append(lora_b_weight.stride(0)) |
| 102 | lora_strides_d1.append(lora_b_weight.stride(1)) |
| 103 | lora_strides_d2.append(lora_b_weight.stride(2)) |
| 104 | slice_offset_lst.append(slice_offset) |
| 105 | slice_offset += lora_b_weight.size(1) |
| 106 | hidden_sizes.append(lora_b_weight.size(1)) |
| 107 | |
| 108 | if len(lora_weights) > 1: |
| 109 | lora_ptr_tensor = async_tensor_h2d( |
| 110 | tensor_ptrs, dtype=torch.uint64, device=device |
| 111 | ) |
| 112 | slice_start_tensor = async_tensor_h2d( |
| 113 | slice_offset_lst, dtype=torch.uint64, device=device |
| 114 | ) |
| 115 | else: |
| 116 | slice_start_tensor = slice_offset_lst[0] |
| 117 | lora_ptr_tensor = lora_b_weight[0] |
| 118 | |
| 119 | # If each lora has the same stride, there's no need to use a |
| 120 | # tensor for storage. |
| 121 | if ( |
| 122 | len(set(lora_strides_d0)) == 1 |
| 123 | and len(set(lora_strides_d1)) == 1 |
| 124 | and len(set(lora_strides_d2)) == 1 |
| 125 | ) and len(set(hidden_sizes)) == 1: |
| 126 | lora_strides_d0_tensor = lora_strides_d0[0] |
| 127 | lora_strides_d1_tensor = lora_strides_d1[0] |
| 128 | lora_strides_d2_tensor = lora_strides_d2[0] |
| 129 | hidden_sizes_tensor = hidden_sizes[0] |
| 130 | same_stride = True |
| 131 | |
| 132 | else: |
| 133 | lora_strides_d0_tensor = async_tensor_h2d( |
| 134 | lora_strides_d0, dtype=torch.int64, device=device |
| 135 | ) |
| 136 | lora_strides_d1_tensor = async_tensor_h2d( |
| 137 | lora_strides_d1, dtype=torch.int64, device=device |
| 138 | ) |
| 139 | lora_strides_d2_tensor = async_tensor_h2d( |
| 140 | lora_strides_d2, dtype=torch.int64, device=device |
| 141 | ) |
| 142 | hidden_sizes_tensor = async_tensor_h2d( |
| 143 | hidden_sizes, dtype=torch.int64, device=device |
| 144 | ) |
| 145 | same_stride = False |
| 146 | # MAX_N is the maximum hidden size among all the lora_b weights |
| 147 | MAX_N = max(hidden_sizes) |
| 148 | _LORA_B_PTR_DICT[key] = ( |
| 149 | slice_start_tensor, |
| 150 | lora_ptr_tensor, |
| 151 | lora_strides_d0_tensor, |
| 152 | lora_strides_d1_tensor, |
| 153 | lora_strides_d2_tensor, |
| 154 | hidden_sizes_tensor, |
| 155 | same_stride, |
| 156 | MAX_N, |
| 157 | ) |
| 158 | return _LORA_B_PTR_DICT.get(key) |
| 159 | |
| 160 | |
| 161 | @functools.lru_cache |
| 162 | def load_lora_op_config(op_type: str, add_inputs: bool | None) -> dict | None: |
| 163 | user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER |
| 164 | # Avoid optimizing for the batch invariant case. Use default config |
| 165 | if user_defined_config_folder is not None and not is_batch_invariant: |
| 166 | gpu_name = torch.cuda.get_device_name() |
| 167 | gpu_name = gpu_name.replace(" ", "_") |
| 168 | gpu_name = gpu_name.replace("-", "_") |
| 169 | |
| 170 | config_fname = None |
| 171 | # only expand op needs to consider add_inputs |
| 172 | if op_type == "expand": |
| 173 | config_fname = ( |
| 174 | f"{gpu_name}_{op_type.upper()}_{str(add_inputs).upper()}.json" |
| 175 | ) |
| 176 | else: |
| 177 | config_fname = f"{gpu_name}_{op_type.upper()}.json" |
| 178 | |
| 179 | config_path = Path(f"{user_defined_config_folder}/{config_fname}") |
| 180 | if not config_path.exists(): |
| 181 | logger.warning_once(f"No LoRA kernel configs found in {config_path}") |
| 182 | return None |
| 183 | |
| 184 | # Load json |
| 185 | logger.info_once(f"Using tuned LoRA kernel configs from {config_path}.") |
| 186 | with open(str(config_path)) as f: |
| 187 | config_data = json.load(f) |
| 188 | else: |
| 189 | config_data = None |
| 190 | |
| 191 | return config_data |
| 192 | |
| 193 | |
| 194 | @functools.lru_cache |
| 195 | def get_lora_op_configs( |
| 196 | op_type: str, |
| 197 | max_loras: int, |
| 198 | batch: int, |
| 199 | hidden_size: int, |
| 200 | rank: int, |
| 201 | num_slices: int, |
| 202 | add_inputs: bool | None = None, |
| 203 | moe_intermediate_size: int | None = None, |
| 204 | ) -> dict[str, int | None]: |
| 205 | # Add support for fused_moe_lora ops |
| 206 | assert op_type in [ |
| 207 | "shrink", |
| 208 | "expand", |
| 209 | "fused_moe_lora_w13_shrink", |
| 210 | "fused_moe_lora_w13_expand", |
| 211 | "fused_moe_lora_w2_shrink", |
| 212 | "fused_moe_lora_w2_expand", |
| 213 | ] |
| 214 | |
| 215 | # default config |
| 216 | default = {} |
| 217 | if op_type == "shrink": |
| 218 | split_k = 64 if batch < 128 else 8 |
| 219 | if is_batch_invariant: |
| 220 | split_k = 1 |
| 221 | default = { |
| 222 | "block_m": 32, |
| 223 | "block_n": 16, |
| 224 | "block_k": 256 if batch < 128 else 32, |
| 225 | "split_k": split_k, |
| 226 | "num_warps": 4, |
| 227 | "num_ctas": 1, |
| 228 | "group_size_m": 8, |
| 229 | "num_stages": 2, |
| 230 | "max_nreg": None, |
| 231 | } |
| 232 | # The default config for fused_moe_lora ops |
| 233 | elif op_type in [ |
| 234 | "fused_moe_lora_w13_shrink", |
| 235 | "fused_moe_lora_w2_shrink", |
| 236 | ]: |
| 237 | default = { |
| 238 | "block_m": 64, |
| 239 | "block_n": min(64, next_power_of_2(rank)), |
| 240 | "block_k": 32, |
| 241 | "num_warps": 4, |
| 242 | "num_stages": 3, |
| 243 | "group_size_m": 8, |
| 244 | "split_k": 1, |
| 245 | } |
| 246 | elif op_type in [ |
| 247 | "fused_moe_lora_w13_expand", |
| 248 | "fused_moe_lora_w2_expand", |
| 249 | ]: |
| 250 | default = { |
| 251 | "block_m": 64, |
| 252 | "block_n": 64, |
| 253 | "block_k": max(16, min(32, next_power_of_2(rank))), |
| 254 | "num_warps": 4, |
| 255 | "num_stages": 3, |
| 256 | "group_size_m": 8, |
| 257 | "split_k": 1, |
| 258 | } |
| 259 | else: |
| 260 | default = { |
| 261 | "block_m": 64, |
| 262 | "block_n": 64 if num_slices > 1 else 128, |
| 263 | "block_k": 32, |
| 264 | "num_warps": 4, |
| 265 | "num_ctas": 1, |
| 266 | "num_stages": 2, |
| 267 | "max_nreg": None, |
| 268 | } |
| 269 | m = batch |
| 270 | |
| 271 | k, n = (hidden_size, rank) if op_type == "shrink" else (rank, hidden_size) |
| 272 | |
| 273 | config_data: Any |
| 274 | config_data = load_lora_op_config(op_type, add_inputs) |
| 275 | if not config_data: |
| 276 | logger.warning_once("Using default LoRA kernel configs") |
| 277 | return default |
| 278 | |
| 279 | # config is structured as config_data[max_loras][num_slices][m][k][n] = {} |
| 280 | # slice by max_loras |
| 281 | config_data = ( |
| 282 | config_data.get(str(max_loras)) |
| 283 | or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - max_loras))] |
| 284 | ) |
| 285 | # slice by num_slices |
| 286 | config_data = config_data[str(num_slices)] |
| 287 | # slice by m |
| 288 | config_data = ( |
| 289 | config_data.get(str(m)) |
| 290 | or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - m))] |
| 291 | ) |
| 292 | # slice by k |
| 293 | config_data = ( |
| 294 | config_data.get(str(k)) |
| 295 | or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - k))] |
| 296 | ) |
| 297 | # slice by n |
| 298 | config_data = ( |
| 299 | config_data.get(str(n)) |
| 300 | or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - n))] |
| 301 | ) |
| 302 | |
| 303 | # slice by moe-intermediate-size if applicable |
| 304 | if moe_intermediate_size is not None: |
| 305 | i = moe_intermediate_size |
| 306 | config_data = ( |
| 307 | config_data.get(str(i)) |
| 308 | or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - i))] |
| 309 | ) |
| 310 | |
| 311 | assert config_data is not None |
| 312 | return config_data |
| 313 | |
| 314 | |
| 315 | @lru_cache |
| 316 | def supports_pdl(device: torch.device | None = None) -> bool: |
| 317 | """Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py.""" |
| 318 | # PDL requires compute capability SM90 or above |
| 319 | |
| 320 | return ( |
| 321 | current_platform.is_cuda() |
| 322 | and current_platform.has_device_capability(90) |
| 323 | and not envs.VLLM_LORA_DISABLE_PDL |
| 324 | ) |
| 325 | |
| 326 | |
| 327 | @lru_cache |
| 328 | def supports_tma(device: torch.device | None = None) -> bool: |
| 329 | # TMA requires compute capability SM90 or above |
| 330 | return current_platform.is_cuda() and current_platform.has_device_capability(90) |
| 331 | |
| 332 | |
| 333 | def _normalize_lora_config_keys( |
| 334 | config: dict[str, int | None], |
| 335 | ) -> dict[str, int | None]: |
| 336 | """Normalize Triton config dict keys to uppercase BLOCK_SIZE_* format.""" |
| 337 | out: dict[str, int | None] = {} |
| 338 | for key, val in config.items(): |
| 339 | if key.islower(): |
| 340 | if key.startswith("block_"): |
| 341 | nk = "BLOCK_SIZE_" + key.split("_")[-1].upper() |
| 342 | else: |
| 343 | nk = key.upper() |
| 344 | else: |
| 345 | nk = key |
| 346 | out[nk] = val |
| 347 | return out |