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] Remove UvaBufferPool for cpu->gpu copy (#33055)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Nick Hill <nhill@redhat.com>
Woosuk Kwon committed
Jan 26, 2026 at 16:47 UTC
6d86fde09c8b8ef10f3b085a1be93ecf6b7c0411
3 files changed
+27
-17
vllm/v1/worker/gpu/buffer_utils.py
+20
@@ -11,6 +11,26 @@ from vllm.utils.platform_utils import is_uva_available
11
from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
12
13
14
+def async_copy_to_gpu(
15
+ x: torch.Tensor | np.ndarray,
16
+ out: torch.Tensor | None = None,
17
+ device: torch.device | None = None,
18
+) -> torch.Tensor:
19
+ if isinstance(x, np.ndarray):
20
+ x = torch.from_numpy(x)
21
+ assert x.is_cpu
22
+ assert not x.is_pinned()
23
+
24
+ if out is None:
25
+ assert device is not None
26
+ out = torch.empty_like(x, device=device)
27
+
28
+ # CPU-to-CPU copy
29
+ tmp = x.pin_memory()
30
+ # CPU-to-GPU copy
31
+ return out.copy_(tmp, non_blocking=True)
32
+
33
+
34
class UvaBuffer:
35
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
36
if not is_uva_available():
vllm/v1/worker/gpu/mm/encoder_runner.py
+2
-5
@@ -6,7 +6,6 @@ import torch
6
from vllm.model_executor.models.interfaces import SupportsMultiModal
7
from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem
8
from vllm.multimodal.utils import group_mm_kwargs_by_modality
9
-from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
9
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
10
11
@@ -32,8 +31,6 @@ class EncoderRunner:
31
self.req_id_to_mm_features: dict[str, list[MultiModalFeatureSpec]] = {}
32
self.encoder_cache: dict[str, torch.Tensor] = {}
33
35
- self.tmp_is_mm_embed = UvaBufferPool(max_num_tokens, torch.bool)
36
-
34
def add_request(self, req_id: str, mm_features: list[MultiModalFeatureSpec]):
35
self.req_id_to_mm_features[req_id] = mm_features
36
@@ -114,7 +111,7 @@ class EncoderRunner:
111
total_num_scheduled_tokens,
112
dtype=torch.bool,
113
device="cpu",
117
- pin_memory=False,
114
+ pin_memory=True,
115
)
116
for i, req_id in enumerate(req_ids):
117
if not is_prefilling[i]:
@@ -163,7 +160,7 @@ class EncoderRunner:
160
mm_embeds.append(mm_embeds_item)
161
162
# Copy the is_mm_embed tensor to the GPU.
166
- is_mm_embed = self.tmp_is_mm_embed.copy_to_gpu(is_mm_embed)
163
+ is_mm_embed = is_mm_embed.to(device=self.device, non_blocking=True)
164
return mm_embeds, is_mm_embed
165
166
@torch.inference_mode()
vllm/v1/worker/gpu/model_runner.py
+5
-12
@@ -30,7 +30,7 @@ from vllm.v1.worker.gpu.attn_utils import (
30
init_kv_cache,
31
)
32
from vllm.v1.worker.gpu.block_table import BlockTables
33
-from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
33
+from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
34
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
35
from vllm.v1.worker.gpu.dp_utils import (
36
get_cudagraph_and_dp_padding,
@@ -172,11 +172,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
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)
177
- self.tmp_cu_num_logits = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
178
- self.tmp_query_start_loc = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
179
-
175
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
176
177
def update_max_model_len(self, max_model_len: int) -> None:
@@ -518,7 +513,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
513
self.req_states.req_id_to_index[req_id] for req_id in req_ids
514
]
515
idx_mapping_np = np.array(idx_mapping_list, dtype=np.int32)
521
- idx_mapping = self.tmp_idx_mapping.copy_to_gpu(idx_mapping_np)
516
+ idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
517
518
# Get the number of draft tokens for each request.
519
if not scheduler_output.scheduled_spec_decode_tokens:
@@ -546,7 +541,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
541
cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
542
cu_num_logits_np[0] = 0
543
np.cumsum(num_logits, out=cu_num_logits_np[1:])
549
- cu_num_logits = self.tmp_cu_num_logits.copy_to_gpu(cu_num_logits_np)
544
+ cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
545
546
expanded_idx_mapping = expand_idx_mapping(
547
idx_mapping,
@@ -565,10 +560,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
560
# Pad for full CUDA graph mode.
561
# Some attention backends like FA3 require query_start_loc to be non-decreasing.
562
query_start_loc_np[num_reqs + 1 :] = num_tokens
568
- self.tmp_query_start_loc.copy_to_gpu(
569
- query_start_loc_np,
570
- out=self.input_buffers.query_start_loc,
571
- )
563
+ async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
564
+
565
query_start_loc_np = query_start_loc_np[: num_reqs + 1]
566
query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
567
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]