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 8, 14:13 UTC
[Model Runner V2] Use pinned memory for write_contents (#34222)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 10, 2026 at 08:55 UTC
a2443de5fa4a0605607f6c3d9219022c7f6ac480
1 file changed
+9
-20
vllm/v1/worker/gpu/buffer_utils.py
+9
-20
@@ -7,9 +7,11 @@ import numpy as np
7
import torch
8
9
from vllm.triton_utils import tl, triton
10
-from vllm.utils.math_utils import next_power_of_2
10
from vllm.utils.platform_utils import is_uva_available
12
-from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
11
+from vllm.utils.torch_utils import (
12
+ async_tensor_h2d,
13
+ get_accelerator_view_from_cpu_tensor,
14
+)
15
16
17
def async_copy_to_gpu(
@@ -117,6 +119,7 @@ class StagedWriteTensor:
119
)
120
self.num_rows = size if isinstance(size, int) else size[0]
121
self.dtype = dtype
122
+ self.device = device
123
self.max_concurrency = max_concurrency
124
125
if not uva_instead_of_gpu:
@@ -137,8 +140,6 @@ class StagedWriteTensor:
140
141
self.write_indices = new_buffer(self.num_rows, dtype=torch.int32)
142
self.write_starts = new_buffer(self.num_rows, dtype=torch.int32)
140
- init_size = next_power_of_2(self.num_rows)
141
- self.write_contents = new_buffer(init_size, dtype=dtype)
143
self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32)
144
145
def stage_write(
@@ -170,21 +171,9 @@ class StagedWriteTensor:
171
cu_lens_uva = self.write_cu_lens.copy_to_uva(self._staged_write_cu_lens)
172
173
# Special handling for write_contents
173
- diff_len = len(self._staged_write_contents)
174
- assert isinstance(self.write_contents.size, int)
175
- if diff_len > self.write_contents.size:
176
- # Re-allocate a larger buffer for the write_contents
177
- new_size = next_power_of_2(diff_len)
178
- self.write_contents = UvaBufferPool(
179
- new_size, dtype=self.dtype, max_concurrency=self.max_concurrency
180
- )
181
- # NOTE(woosuk): Since the previous write_contents buffer is released,
182
- # we perform a synchronization here to ensure that all data transfers
183
- # involving the old buffer have finished before allocating a new one.
184
- # This prevents potential race conditions. The slight overhead is
185
- # negligible because the reallocations are infrequent in practice.
186
- torch.cuda.synchronize()
187
- contents_uva = self.write_contents.copy_to_uva(self._staged_write_contents)
174
+ write_contents = async_tensor_h2d(
175
+ self._staged_write_contents, self.dtype, self.device, pin_memory=True
176
+ )
177
178
# Write diffs to the GPU buffer
179
_apply_write_kernel[(n,)](
@@ -192,7 +181,7 @@ class StagedWriteTensor:
181
self.gpu.stride(0),
182
indices_uva,
183
starts_uva,
195
- contents_uva,
184
+ write_contents,
185
cu_lens_uva,
186
BLOCK_SIZE=1024,
187
)