[Perf][LoRA] Replace O(n) list.index() with a dict in convert_mapping (#46542)
Signed-off-by: Lynn <lynnhe02@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Lynn committed
Jun 24, 2026 at 08:41 UTC
0bc479e6eb2b9fcdd27bb28a7ce09b88347949f1
1 file changed
+11
-3
vllm/lora/punica_wrapper/utils.py
+11
-3
@@ -92,14 +92,22 @@ def convert_mapping(
92
embedding_indices = index_mapping_indices.copy()
93
lora_indices = index_mapping_indices.copy()
94
95
+ # Build a reverse lookup (LoRA id -> index) once instead of repeatedly
96
+ # calling list.index(), which is an O(num_loras) linear scan performed for
97
+ # every prompt token and every batch token on each step.
98
+ lora_id_to_index = {
99
+ lora_id: index
100
+ for index, lora_id in enumerate(lora_index_to_id)
101
+ if lora_id is not None
102
+ }
103
+
104
prompt_mapping: list[int] = [
96
- lora_index_to_id.index(x) if x > 0 else -1 for x in mapping.prompt_mapping
105
+ lora_id_to_index[x] if x > 0 else -1 for x in mapping.prompt_mapping
106
]
107
lora_idx = None
108
for i in range(len(index_mapping_indices)):
100
- # TODO index can be slow. optimize
109
lora_idx = (
102
- lora_index_to_id.index(index_mapping_indices[i])
110
+ lora_id_to_index[index_mapping_indices[i]]
111
if index_mapping_indices[i] > 0
112
else -1
113
)