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 9, 14:45 UTC
[Model Runner V2] Add LoRAState to consolidate lora logic (#33062)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jan 25, 2026 at 22:21 UTC
a9b53dd435bb82f311f340ebebc15e62b9624a9d
3 files changed
+53
-42
vllm/v1/worker/gpu/lora_utils.py
new
+47
@@ -0,0 +1,47 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+
5
+from vllm.lora.request import LoRARequest
6
+
7
+NO_LORA_ID = 0
8
+
9
+
10
+class LoraState:
11
+ def __init__(self, max_num_reqs: int):
12
+ self.lora_ids = np.zeros(max_num_reqs, dtype=np.int32)
13
+ self.lora_ids.fill(NO_LORA_ID)
14
+ # req_id -> lora_request
15
+ self.lora_requests: dict[str, LoRARequest] = {}
16
+
17
+ def add_request(
18
+ self,
19
+ req_id: str,
20
+ req_index: int,
21
+ lora_request: LoRARequest | None,
22
+ ) -> None:
23
+ if lora_request is not None:
24
+ self.lora_requests[req_id] = lora_request
25
+ self.lora_ids[req_index] = lora_request.lora_int_id
26
+ else:
27
+ self.lora_ids[req_index] = NO_LORA_ID
28
+
29
+ def remove_request(self, req_id: str) -> None:
30
+ self.lora_requests.pop(req_id, None)
31
+
32
+ def make_lora_inputs(
33
+ self,
34
+ req_ids: list[str],
35
+ idx_mapping: np.ndarray,
36
+ num_scheduled_tokens: np.ndarray,
37
+ ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
38
+ lora_ids = self.lora_ids[idx_mapping]
39
+ prompt_lora_mapping = tuple(lora_ids)
40
+ token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
41
+
42
+ active_lora_requests: set[LoRARequest] = set()
43
+ for req_id in req_ids:
44
+ lora_request = self.lora_requests.get(req_id, None)
45
+ if lora_request is not None:
46
+ active_lora_requests.add(lora_request)
47
+ return prompt_lora_mapping, token_lora_mapping, active_lora_requests
vllm/v1/worker/gpu/model_runner.py
+6
-2
@@ -51,6 +51,7 @@ from vllm.v1.worker.gpu.kv_connector import (
51
KVConnector,
52
get_kv_connector,
53
)
54
+from vllm.v1.worker.gpu.lora_utils import LoraState
55
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
56
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
57
from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -168,6 +169,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
169
max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
170
vocab_size=self.vocab_size,
171
)
172
+ # LoRA-related workers.
173
+ self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
174
175
# Buffers for CPU-to-GPU copies.
176
self.tmp_idx_mapping = UvaBufferPool(self.max_num_reqs, torch.int32)
@@ -426,6 +429,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
429
if self.supports_mm_inputs:
430
self.encoder_runner.remove_request(req_id)
431
self.prompt_logprobs_worker.remove_request(req_id)
432
+ self.lora_state.remove_request(req_id)
433
434
def free_states(self, scheduler_output: SchedulerOutput) -> None:
435
if self.supports_mm_inputs:
@@ -444,7 +448,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
448
prompt_len=prompt_len,
449
prefill_token_ids=new_req_data.prefill_token_ids,
450
num_computed_tokens=new_req_data.num_computed_tokens,
447
- lora_request=new_req_data.lora_request,
451
)
452
req_index = self.req_states.req_id_to_index[req_id]
453
@@ -469,6 +472,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
472
self.prompt_logprobs_worker.add_request(
473
req_id, req_index, new_req_data.sampling_params
474
)
475
+ self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
476
477
if scheduler_output.scheduled_new_reqs:
478
self.req_states.apply_staged_writes()
@@ -841,7 +845,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
845
)
846
if self.lora_config:
847
# Activate LoRA adapters.
844
- lora_inputs = self.req_states.make_lora_inputs(
848
+ lora_inputs = self.lora_state.make_lora_inputs(
849
input_batch.req_ids,
850
input_batch.idx_mapping_np,
851
input_batch.num_scheduled_tokens,
vllm/v1/worker/gpu/states.py
-40
@@ -1,15 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
4
-
3
import numpy as np
4
import torch
5
8
-from vllm.lora.request import LoRARequest
6
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
7
11
-NO_LORA_ID = 0
12
-
8
9
class RequestState:
10
def __init__(
@@ -31,7 +26,6 @@ class RequestState:
26
self.req_id_to_index: dict[str, int] = {}
27
self.index_to_req_id: dict[int, str] = {}
28
self.free_indices = list(range(max_num_reqs))
34
- self.extra_data: dict[str, ExtraData] = {}
29
30
self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)
31
# NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
@@ -70,10 +64,6 @@ class RequestState:
64
self.max_num_reqs, dtype=torch.int32, device=device
65
)
66
73
- # LoRA.
74
- self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
75
- self.lora_ids.fill(NO_LORA_ID)
76
-
67
@property
68
def num_reqs(self) -> int:
69
return len(self.req_id_to_index)
@@ -84,13 +74,11 @@ class RequestState:
74
prompt_len: int,
75
prefill_token_ids: list[int],
76
num_computed_tokens: int,
87
- lora_request: LoRARequest | None,
77
) -> None:
78
assert len(self.free_indices) > 0, "No free indices"
79
req_idx = self.free_indices.pop()
80
self.req_id_to_index[req_id] = req_idx
81
self.index_to_req_id[req_idx] = req_id
93
- self.extra_data[req_id] = ExtraData(lora_request)
82
83
self.prompt_len[req_idx] = prompt_len
84
prefill_len = len(prefill_token_ids)
@@ -102,43 +90,15 @@ class RequestState:
90
self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
91
self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
92
105
- if lora_request is not None:
106
- self.lora_ids[req_idx] = lora_request.lora_int_id
107
- else:
108
- self.lora_ids[req_idx] = NO_LORA_ID
109
-
93
def apply_staged_writes(self) -> None:
94
self.prefill_len.copy_to_uva()
95
self.prefill_token_ids.apply_write()
96
self.num_computed_tokens.apply_write()
97
98
def remove_request(self, req_id: str) -> None:
116
- self.extra_data.pop(req_id, None)
99
req_idx = self.req_id_to_index.pop(req_id, None)
100
if req_idx is None:
101
# Request not found.
102
return
103
self.index_to_req_id.pop(req_idx, None)
104
self.free_indices.append(req_idx)
123
-
124
- def make_lora_inputs(
125
- self,
126
- req_ids: list[str],
127
- idx_mapping: np.ndarray,
128
- num_scheduled_tokens: np.ndarray,
129
- ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
130
- lora_ids = self.lora_ids[idx_mapping]
131
- prompt_lora_mapping = tuple(lora_ids)
132
- token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
133
-
134
- active_lora_requests: set[LoRARequest] = set()
135
- for req_id in req_ids:
136
- lora_request = self.extra_data[req_id].lora_request
137
- if lora_request is not None:
138
- active_lora_requests.add(lora_request)
139
- return prompt_lora_mapping, token_lora_mapping, active_lora_requests
140
-
141
-
142
-@dataclass
143
-class ExtraData:
144
- lora_request: LoRARequest | None