A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 01:15 UTC
[Feature] Batch invariant: Lora (#30097)
Signed-off-by: quanliu <18646313696@163.com>
quanliu committed
Dec 23, 2025 at 10:32 UTC
a37328fc5c6c281042b3bbf4cb067e8bd3cd2ce5
1 file changed
+8
-2
vllm/lora/ops/triton_ops/utils.py
+8
-2
@@ -11,9 +11,11 @@ import torch
11
12
from vllm import envs
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
15
from vllm.platforms import current_platform
16
17
logger = init_logger(__name__)
18
+is_batch_invariant = vllm_is_batch_invariant()
19
20
_LORA_A_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}
21
_LORA_B_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}
@@ -150,7 +152,8 @@ def _get_lora_b_ptr(
152
@functools.lru_cache
153
def load_lora_op_config(op_type: str, add_inputs: bool | None) -> dict | None:
154
user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER
153
- if user_defined_config_folder is not None:
155
+ # Avoid optimizing for the batch invariant case. Use default config
156
+ if user_defined_config_folder is not None and not is_batch_invariant:
157
gpu_name = torch.cuda.get_device_name()
158
gpu_name = gpu_name.replace(" ", "_")
159
gpu_name = gpu_name.replace("-", "_")
@@ -203,11 +206,14 @@ def get_lora_op_configs(
206
# default config
207
default = {}
208
if op_type == "shrink":
209
+ split_k = 64 if batch < 128 else 8
210
+ if is_batch_invariant:
211
+ split_k = 1
212
default = {
213
"block_m": 32,
214
"block_n": 16,
215
"block_k": 256 if batch < 128 else 32,
210
- "split_k": 64 if batch < 128 else 8,
216
+ "split_k": split_k,
217
"num_warps": 4,
218
"num_ctas": 1,
219
"group_size_m": 8,