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 12, 08:30 UTC
[Model Runner V2] Don't use UVA buffer for prefill_len (#29713)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 28, 2025 at 20:27 UTC
4a80ad0a2548cf0da02af843af346a783fa7754a
2 files changed
+6
-1
vllm/v1/worker/gpu/model_runner.py
+2
@@ -410,6 +410,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
410
cu_num_new_blocks[i].append(x + len(block_ids))
411
new_block_ids[i].extend(block_ids)
412
overwrite.append(True)
413
+ if scheduler_output.scheduled_new_reqs:
414
+ self.req_states.prefill_len.copy_to_gpu()
415
416
# Add new blocks for the existing requests.
417
cached_reqs = scheduler_output.scheduled_cached_reqs
vllm/v1/worker/gpu/states.py
+4
-1
@@ -117,7 +117,10 @@ class RequestState:
117
self.prefill_token_ids = UvaBuffer(
118
self.max_num_reqs, self.max_model_len, dtype=torch.int32
119
)
120
- self.prefill_len = UvaBuffer(self.max_num_reqs, dtype=torch.int32)
120
+ # NOTE(woosuk): We don't use UVA for prefill_len because its GPU view
121
+ # can be used outside of update_states and prepare_inputs.
122
+ # Without async barrier, using UVA can cause race conditions.
123
+ self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
124
# Number of computed tokens.
125
self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
126
self.num_computed_tokens = torch.zeros(