[MRV2] Buffer util simplifications (#56888)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Sep 15, 2026 at 13:36 UTC
f2aad6aa70749ccc4604d31d6f8a04c6728c3361
21 files changed
+125
-112
tests/utils_/test_torch_utils.py
+25
@@ -1,10 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
import pytest
5
import torch
6
7
from vllm.utils.torch_utils import (
8
OMP_NUM_THREADS_SET_BY_VLLM,
9
+ async_tensor_h2d,
10
available_cpu_count,
11
common_broadcastable_dtype,
12
current_stream,
@@ -192,3 +194,26 @@ def test_runtime_threads_override_vllm_set_omp_num_threads(
194
torch.set_num_threads(3)
195
set_torch_threads_for_runtime()
196
assert torch.get_num_threads() == 1
197
+
198
+
199
+@pytest.mark.parametrize(
200
+ "device", ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
201
+)
202
+def test_async_tensor_h2d_staging(device):
203
+ """Staging must preserve values while normalizing dtype and layout."""
204
+ # Non-contiguous source with a dtype conversion: would previously sync.
205
+ src = torch.arange(12, dtype=torch.int64).reshape(3, 4).T
206
+ result = async_tensor_h2d(src, device=device, dtype=torch.int32)
207
+ assert result.dtype == torch.int32
208
+ assert result.is_contiguous()
209
+ assert torch.equal(result.cpu(), src.to(torch.int32))
210
+
211
+ # numpy and list sources still work.
212
+ assert torch.equal(
213
+ async_tensor_h2d(np.arange(4, dtype=np.int64), device=device),
214
+ torch.arange(4, device=device),
215
+ )
216
+ assert torch.equal(
217
+ async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32),
218
+ torch.tensor([1, 2, 3], dtype=torch.int32, device=device),
219
+ )
tests/v1/attention/test_indexer_dcp_localize.py
+1
-1
@@ -996,7 +996,7 @@ def test_pcp_plan_deinterleave_restores_global_order(
996
997
monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False)
998
monkeypatch.setattr(
999
- "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu",
999
+ "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d",
1000
lambda x, device: x.to(device),
1001
)
1002
scheduled = np.array(req_lens, dtype=np.int64)
tests/v1/worker/test_gpu_pcp_manager.py
+3
-3
@@ -17,7 +17,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_c
17
from vllm.v1.worker.gpu.pcp_manager import PCPManager
18
19
20
-def _copy_to_cpu(value, out=None, device=None):
20
+def _copy_to_cpu(value, device=None, dtype=None, out=None):
21
tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value
22
if out is not None:
23
return out.copy_(tensor)
@@ -83,7 +83,7 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch):
83
device=torch.device("cpu"),
84
dcp_world_size=1,
85
)
86
- monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
86
+ monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
87
88
segments_by_rank, per_rank_num_tokens = manager._build_batch_layout(
89
num_scheduled_tokens=np.ones(3, dtype=np.int32),
@@ -158,7 +158,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
158
device=torch.device("cpu"),
159
dcp_world_size=1,
160
)
161
- monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
161
+ monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
162
163
with pytest.raises(ValueError, match="smaller than the largest rank-local batch"):
164
manager._build_batch_layout(
vllm/model_executor/layers/attention/mla_attention.py
+17
-18
@@ -269,11 +269,11 @@ from vllm.utils.torch_utils import (
269
LayerNameType,
270
_encode_layer_name,
271
_resolve_layer_name,
272
+ async_tensor_h2d,
273
direct_register_custom_op,
274
get_dtype_size,
275
is_quantized_kv_cache,
276
kv_cache_dtype_str_to_dtype,
276
- np_to_pinned_tensor,
277
)
278
from vllm.v1.attention.backend import (
279
AttentionBackend,
@@ -1980,11 +1980,6 @@ def plan_mla_context_chunks(
1980
return plans
1981
1982
1983
-def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor:
1984
- """Pinned int32 CPU tensor backing one concatenated per-chunk field."""
1985
- return np_to_pinned_tensor(np.asarray(values, dtype=np.int32))
1986
-
1987
-
1983
def align_mla_chunked_context_workspace_size(
1984
vllm_config: VllmConfig,
1985
workspace_size: int,
@@ -2161,24 +2156,28 @@ def build_mla_chunked_context_metadata(
2156
token_offset += num_tokens
2157
local_token_offset += num_local_tokens
2158
2164
- seq_lens_cpu = _flat_int32(seq_lens_flat)
2165
- starts_and_context_lens = _flat_int32(starts_flat + context_lens).to(
2166
- device, non_blocking=True
2159
+ seq_lens_cpu = torch.tensor(
2160
+ seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
2161
+ )
2162
+ starts_and_context_lens = async_tensor_h2d(
2163
+ starts_flat + context_lens, device=device, dtype=torch.int32
2164
)
2165
starts = starts_and_context_lens[: len(starts_flat)]
2166
context_lens_gpu = starts_and_context_lens[len(starts_flat) :]
2170
- cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True)
2171
- cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True)
2172
- token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to(
2173
- device, non_blocking=True
2167
+ cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32)
2168
+ cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32)
2169
+ token_to_seq = async_tensor_h2d(
2170
+ np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32
2171
)
2172
if use_dcp:
2176
- padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to(
2177
- device, non_blocking=True
2173
+ padded_local_cu_seq_lens = async_tensor_h2d(
2174
+ padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32
2175
+ )
2176
+ padded_local_token_to_seq = async_tensor_h2d(
2177
+ np.concatenate(padded_local_token_to_seq_parts),
2178
+ device=device,
2179
+ dtype=torch.int32,
2180
)
2179
- padded_local_token_to_seq = _flat_int32(
2180
- np.concatenate(padded_local_token_to_seq_parts)
2181
- ).to(device, non_blocking=True)
2181
2182
chunks: list[MLACommonPrefillMetadata.ContextChunk] = []
2183
for index, (plan, layout) in enumerate(zip(plans, layouts)):
vllm/model_executor/models/diffusion_gemma.py
+3
-3
@@ -53,7 +53,7 @@ from vllm.utils.torch_utils import async_tensor_h2d
53
from vllm.v1.outputs import LogprobsTensors
54
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
55
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
56
-from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu
56
+from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
57
from vllm.v1.worker.gpu.input_batch import InputBatch
58
from vllm.v1.worker.gpu.model_states.interface import ModelState
59
from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
@@ -1159,7 +1159,7 @@ class DiffusionSampler:
1159
return
1160
# Move the slot indices across once, up front: indexing a device
1161
# tensor with a numpy array copies them over synchronously each time.
1162
- ps_gpu = async_copy_to_gpu(
1162
+ ps_gpu = async_tensor_h2d(
1163
ps.astype(np.int64), device=states.is_encoder_phase.device
1164
)
1165
states.init_canvas(ps_gpu)
@@ -1266,7 +1266,7 @@ class DiffusionSampler:
1266
# was truncated near max_model_len, in which case the scheduler gave us
1267
# fewer than CL logits for that request.
1268
valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0]
1269
- valid_canvas_len = async_copy_to_gpu(
1269
+ valid_canvas_len = async_tensor_h2d(
1270
valid_canvas_len_np.astype(np.int64), device=device
1271
)
1272
vllm/models/dots3_note/nvidia/attention.py
+15
-7
@@ -19,7 +19,11 @@ from vllm.model_executor.layers.attention.mla_attention import (
19
)
20
from vllm.platforms import current_platform
21
from vllm.triton_utils import tl, triton
22
-from vllm.utils.torch_utils import is_quantized_kv_cache
22
+from vllm.utils.torch_utils import (
23
+ PIN_MEMORY,
24
+ async_tensor_h2d,
25
+ is_quantized_kv_cache,
26
+)
27
from vllm.v1.attention.backend import (
28
AttentionCGSupport,
29
AttentionLayer,
@@ -231,8 +235,12 @@ def _build_sliding_window_metadata(
235
query_lens = query_lens_cpu[req_start:req_end]
236
kv_lens = kv_lens_cpu[req_start:req_end]
237
num_reqs = req_end - req_start
234
- cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
235
- cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
238
+ cu_seq_lens_q_cpu = torch.zeros(
239
+ num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY
240
+ )
241
+ cu_seq_lens_k_cpu = torch.zeros(
242
+ num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY
243
+ )
244
torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:])
245
torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:])
246
token_to_seq_cpu = torch.repeat_interleave(
@@ -248,8 +256,8 @@ def _build_sliding_window_metadata(
256
query_end=query_end,
257
cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True),
258
cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True),
251
- starts=starts_cpu[req_start:req_end].to(device, non_blocking=True),
252
- token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
259
+ starts=async_tensor_h2d(starts_cpu[req_start:req_end], device),
260
+ token_to_seq=async_tensor_h2d(token_to_seq_cpu, device),
261
num_kv_tokens=num_kv_tokens,
262
max_seq_len_q=int(query_lens.max().item()),
263
max_seq_len_k=int(kv_lens.max().item()),
@@ -396,8 +404,8 @@ class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder):
404
seq_lens = common_attn_metadata.seq_lens[
405
reqs_start + req_start : reqs_start + req_end
406
]
399
- query_lens = query_lens_cpu[req_start:req_end].to(
400
- self.device, non_blocking=True
407
+ query_lens = async_tensor_h2d(
408
+ query_lens_cpu[req_start:req_end], self.device
409
)
410
kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1)
411
chunk.cu_seq_lens_k.zero_()
vllm/third_party/flash_linear_attention/ops/index.py
+3
-2
@@ -10,6 +10,7 @@
10
import torch
11
12
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
13
+from vllm.utils.torch_utils import async_tensor_h2d
14
from vllm.triton_utils import triton
15
16
from .utils import tensor_cache
@@ -27,8 +28,8 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te
28
chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
29
indices = torch.cat([torch.arange(n) for n in chunk_counts])
30
chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1)
30
- return chunk_indices.to(
31
- device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True
31
+ return async_tensor_h2d(
32
+ chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype
33
)
34
35
vllm/utils/torch_utils.py
+15
-2
@@ -705,17 +705,30 @@ def create_kv_caches_with_random(
705
706
def async_tensor_h2d(
707
data: list | np.ndarray | torch.Tensor,
708
- device: str | torch.device,
708
+ device: str | torch.device | None = None,
709
dtype: torch.dtype | None = None,
710
+ out: torch.Tensor | None = None,
711
) -> torch.Tensor:
712
"""Copy list/numpy array/tensor async from host to device."""
713
+ if dtype is None and out is not None:
714
+ dtype = out.dtype
715
if isinstance(data, np.ndarray):
716
data = torch.from_numpy(data)
717
if isinstance(data, torch.Tensor):
715
- t = data.pin_memory() if PIN_MEMORY else data
718
+ t = data
719
+ if PIN_MEMORY and not t.is_pinned():
720
+ # Stage in pinned, contiguous buffer to ensure fully async copy.
721
+ t = torch.empty(
722
+ t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True
723
+ ).copy_(t)
724
else:
725
t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu")
726
assert t.is_cpu
727
+
728
+ if out is not None:
729
+ assert out.dtype == dtype
730
+ return out.copy_(t, non_blocking=True)
731
+ assert device is not None, "must provide destination tensor or device"
732
return t.to(device=device, dtype=dtype, non_blocking=True)
733
734
vllm/v1/attention/backends/flashinfer.py
+3
-14
@@ -1264,15 +1264,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1264
"""
1265
# write self.paged_kv_indptr_cpu inplace (0-index is always 0)
1266
np.cumsum(
1267
- num_blocks_np,
1268
- dtype=np.int32,
1269
- out=self.paged_kv_indptr.np[1 : num_reqs + 1],
1267
+ num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1]
1268
)
1271
- paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1]
1272
- paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1]
1273
- if PIN_MEMORY:
1274
- paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()
1275
- paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True)
1269
+ paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1)
1270
1271
# write self.paged_kv_indices inplace
1272
num_actual_pages = self.paged_kv_indptr.np[num_reqs]
@@ -1292,12 +1286,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1286
page_size,
1287
paged_kv_last_page_len_np,
1288
)
1295
- paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]
1296
- if PIN_MEMORY:
1297
- paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()
1298
- self.paged_kv_last_page_len.gpu[:num_reqs].copy_(
1299
- paged_kv_last_page_len_cpu, non_blocking=True
1300
- )
1289
+ self.paged_kv_last_page_len.copy_to_gpu(num_reqs)
1290
return paged_kv_indices
1291
1292
def build(
vllm/v1/attention/backends/mla/flashmla_sparse.py
+6
-3
@@ -20,7 +20,11 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
20
from vllm.platforms import current_platform
21
from vllm.platforms.interface import DeviceCapability
22
from vllm.utils.platform_utils import num_compute_units
23
-from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache
23
+from vllm.utils.torch_utils import (
24
+ async_tensor_h2d,
25
+ current_stream,
26
+ is_quantized_kv_cache,
27
+)
28
from vllm.v1.attention.backend import (
29
AttentionBackend,
30
AttentionCGSupport,
@@ -48,7 +52,6 @@ from vllm.v1.attention.ops.flashmla import (
52
get_mla_metadata,
53
)
54
from vllm.v1.kv_cache_interface import AttentionSpec
51
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
55
from vllm.v1.worker.workspace import current_workspace_manager
56
57
if TYPE_CHECKING:
@@ -519,7 +522,7 @@ class FlashMLASparseMetadataBuilder(
522
)
523
workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32))
524
max_prefill_buffer_size //= self.dcp_world_size
522
- entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device)
525
+ entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device)
526
prefill_block_table = prefill_block_table.index_select(0, entry_rows)
527
prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows)
528
num_entries = len(workspace_rows)
vllm/v1/attention/backends/mla/indexer.py
+3
-4
@@ -27,7 +27,7 @@ from vllm.utils.deep_gemm import (
27
)
28
from vllm.utils.math_utils import round_down
29
from vllm.utils.platform_utils import num_compute_units
30
-from vllm.utils.torch_utils import PIN_MEMORY
30
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
31
from vllm.v1.attention.backend import (
32
AttentionBackend,
33
AttentionCGSupport,
@@ -47,7 +47,6 @@ from vllm.v1.kv_cache_interface import (
47
KVCacheSpec,
48
MLAAttentionSpec,
49
)
50
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
50
51
logger = init_logger(__name__)
52
@@ -338,14 +337,14 @@ def build_pcp_global_chunk_plan(
337
+ local
338
)
339
341
- cu = async_copy_to_gpu(cu, device=device)
340
+ cu = async_tensor_h2d(cu, device=device)
341
return PCPGlobalChunkPlan(
342
row_start_cu=cu[0],
343
global_cu=cu[1],
344
padded_local_cu=cu[2],
345
padded_local_total=padded_total,
346
total=total,
348
- deinterleave_idx=async_copy_to_gpu(idx, device=device),
347
+ deinterleave_idx=async_tensor_h2d(idx, device=device),
348
)
349
350
vllm/v1/sample/rejection_sampler.py
+2
-3
@@ -19,6 +19,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
19
triton_kernel_dispatcher_with_warmup,
20
)
21
from vllm.triton_utils import tl, triton
22
+from vllm.utils.torch_utils import async_tensor_h2d
23
from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
24
from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor
25
from vllm.v1.sample.metadata import SamplingMetadata
@@ -327,9 +328,7 @@ class RejectionSampler(nn.Module):
328
num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")
329
original_indices = torch.arange(num_requests, device="cpu")
330
repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens)
330
- repeat_indices = repeat_indices_cpu.to(
331
- device=logits.device, non_blocking=True
332
- )
331
+ repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device)
332
logits = self.apply_penalties(
333
logits, sampling_metadata, metadata, repeat_indices, output_token_ids
334
)
vllm/v1/spec_decode/llm_base_proposer.py
-1
@@ -255,7 +255,6 @@ class SpecDecodeBaseProposer:
255
self.backup_next_token_ids = CpuGpuBuffer(
256
self.max_batch_size,
257
dtype=torch.int32,
258
- pin_memory=PIN_MEMORY,
258
device=device,
259
with_numpy=True,
260
)
vllm/v1/utils.py
+4
-3
@@ -137,9 +137,10 @@ class CpuGpuBuffer:
137
self.np = self.cpu.numpy()
138
139
def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
140
- if n is None:
141
- return self.gpu.copy_(self.cpu, non_blocking=True)
142
- return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
140
+ cpu, gpu = self.cpu, self.gpu
141
+ if n is not None:
142
+ cpu, gpu = cpu[:n], gpu[:n]
143
+ return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True)
144
145
def copy_to_cpu(self, n: int | None = None) -> torch.Tensor:
146
"""NOTE: Because this method is non-blocking, explicit synchronization
vllm/v1/worker/cpu/shm.py
+6
-3
@@ -76,14 +76,17 @@ import vllm.utils.torch_utils as torch_utils
76
77
def async_tensor_h2d(
78
data: list | np.ndarray | torch.Tensor,
79
- device: str | torch.device,
79
+ device: str | torch.device | None = None,
80
dtype: torch.dtype | None = None,
81
+ out: torch.Tensor | None = None,
82
) -> torch.Tensor:
83
if isinstance(data, np.ndarray):
84
data = torch.from_numpy(data)
84
- if isinstance(data, torch.Tensor):
85
+ if not isinstance(data, torch.Tensor):
86
+ data = torch.tensor(data, dtype=dtype, device="cpu")
87
+ elif out is None:
88
return data.to(dtype=dtype)
86
- return torch.tensor(data, dtype=dtype, device="cpu")
89
+ return data if out is None else out.copy_(data)
90
91
92
torch_utils.async_tensor_h2d = async_tensor_h2d
vllm/v1/worker/gpu/buffer_utils.py
-19
@@ -10,7 +10,6 @@ from vllm.logger import init_logger
10
from vllm.triton_utils import tl, triton
11
from vllm.utils.platform_utils import is_uva_available
12
from vllm.utils.torch_utils import (
13
- PIN_MEMORY,
13
async_tensor_h2d,
14
get_accelerator_view_from_cpu_tensor,
15
)
@@ -27,24 +26,6 @@ def set_default_max_concurrency(n: int) -> None:
26
_DEFAULT_MAX_CONCURRENCY = max(2, n)
27
28
30
-def async_copy_to_gpu(
31
- x: torch.Tensor | np.ndarray,
32
- out: torch.Tensor | None = None,
33
- device: torch.device | None = None,
34
-) -> torch.Tensor:
35
- if isinstance(x, np.ndarray):
36
- x = torch.from_numpy(x)
37
- assert x.is_cpu
38
-
39
- if out is None:
40
- assert device is not None
41
- out = torch.empty_like(x, device=device)
42
-
43
- # pin_memory() is no-op if the memory is already pinned.
44
- pinned = x.pin_memory() if PIN_MEMORY else x
45
- return out.copy_(pinned, non_blocking=True)
46
-
47
-
29
class UvaBuffer:
30
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
31
if not is_uva_available():
vllm/v1/worker/gpu/model_runner.py
+5
-8
@@ -64,7 +64,7 @@ from vllm.sequence import IntermediateTensors
64
from vllm.tasks import SupportedTask
65
from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
66
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
67
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
67
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d
68
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
69
from vllm.v1.kv_cache_interface import (
70
KVCacheConfig,
@@ -99,10 +99,7 @@ from vllm.v1.worker.gpu.attn_utils import (
99
init_kv_cache,
100
)
101
from vllm.v1.worker.gpu.block_table import BlockTables
102
-from vllm.v1.worker.gpu.buffer_utils import (
103
- async_copy_to_gpu,
104
- set_default_max_concurrency,
105
-)
102
+from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
103
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
104
from vllm.v1.worker.gpu.cudagraph_utils import (
105
BatchExecutionDescriptor,
@@ -1289,7 +1286,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1286
req_ids = batch_req_state.req_ids
1287
num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
1288
idx_mapping_np = batch_req_state.idx_mapping_np
1292
- idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
1289
+ idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
1290
num_reqs = len(req_ids)
1291
1292
# Get the number of draft tokens for each request.
@@ -1324,7 +1321,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1321
cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
1322
cu_num_logits_np[0] = 0
1323
np.cumsum(num_logits, out=cu_num_logits_np[1:])
1327
- cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
1324
+ cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)
1325
1326
adaptive_verification = (
1327
self.adaptive_verification if num_draft_tokens_per_req is not None else None
@@ -1352,7 +1349,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1349
# Some attention backends like FA3 require query_start_loc to be non-decreasing.
1350
query_start_loc_np[num_reqs + 1 :] = num_tokens
1351
query_start_loc = self.input_buffers.query_start_loc
1355
- async_copy_to_gpu(query_start_loc_np, out=query_start_loc)
1352
+ async_tensor_h2d(query_start_loc_np, out=query_start_loc)
1353
if adaptive_verification is not None:
1354
cu_num_logits, query_start_loc, total_num_draft_tokens = (
1355
adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
vllm/v1/worker/gpu/pcp_manager.py
+6
-6
@@ -10,9 +10,9 @@ import torch
10
from vllm.config import CUDAGraphMode, VllmConfig
11
from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group
12
from vllm.logger import init_logger
13
+from vllm.utils.torch_utils import async_tensor_h2d
14
from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
15
from vllm.v1.worker.gpu.block_table import BlockTables
15
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
16
from vllm.v1.worker.gpu.input_batch import (
17
InputBatch,
18
InputBuffers,
@@ -358,13 +358,13 @@ class PCPManager:
358
dtype=np.int64,
359
)
360
361
- self._hidden_restore_idx = async_copy_to_gpu(
361
+ self._hidden_restore_idx = async_tensor_h2d(
362
hidden_restore_idx, device=self.device
363
)
364
- self._padded_gather_idx = async_copy_to_gpu(
364
+ self._padded_gather_idx = async_tensor_h2d(
365
padded_gather_idx, device=self.device
366
)
367
- self._gathered_kv_write_mask = async_copy_to_gpu(
367
+ self._gathered_kv_write_mask = async_tensor_h2d(
368
gathered_kv_write_mask, device=self.device
369
)
370
return segments_by_rank, per_rank_num_tokens
@@ -536,12 +536,12 @@ class PCPManager:
536
local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1]
537
np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out)
538
local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens
539
- async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc)
539
+ async_tensor_h2d(local_query_start_loc_np, out=input_buffers.query_start_loc)
540
local_query_start_loc = input_buffers.query_start_loc[
541
: num_reqs_after_padding + 1
542
]
543
544
- local_to_global_req_idx = async_copy_to_gpu(
544
+ local_to_global_req_idx = async_tensor_h2d(
545
local_to_global_req_idx_np, device=self.device
546
)
547
seq_lens = input_buffers.seq_lens[:num_reqs_after_padding]
vllm/v1/worker/gpu/pp_utils.py
+3
-3
@@ -11,7 +11,7 @@ import torch
11
from vllm.distributed.parallel_state import get_pp_group
12
from vllm.platforms import current_platform
13
from vllm.sequence import IntermediateTensors
14
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
14
+from vllm.utils.torch_utils import async_tensor_h2d
15
from vllm.v1.worker.gpu.input_batch import InputBatch
16
17
@@ -138,7 +138,7 @@ class PPHandler:
138
return None
139
# Filter excluded request indices.
140
idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
141
- idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
141
+ idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
142
143
self.main_stream.wait_event(slot.event)
144
if slot.draft_tokens is not None and draft_tokens_to_update is not None:
@@ -148,7 +148,7 @@ class PPHandler:
148
keep = ~exclude_mask
149
keep_t = torch.as_tensor(keep, device=self.device)
150
draft_tokens = draft_tokens[keep_t]
151
- draft_idx_mapping = async_copy_to_gpu(
151
+ draft_idx_mapping = async_tensor_h2d(
152
slot.idx_mapping_np[keep], device=self.device
153
)
154
draft_tokens_to_update[draft_idx_mapping] = draft_tokens
vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+3
-6
@@ -14,6 +14,7 @@ from vllm.config.compilation import CUDAGraphMode
14
from vllm.distributed.parallel_state import get_tp_group
15
from vllm.logger import init_logger
16
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
17
+from vllm.utils.torch_utils import async_tensor_h2d
18
from vllm.v1.attention.backend import AttentionCGSupport
19
from vllm.v1.utils import CpuGpuBuffer
20
from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream
@@ -21,7 +22,6 @@ from vllm.v1.worker.gpu.attn_utils import (
22
get_attn_cg_support,
23
get_query_lens_mismatch_unsupported_backend,
24
)
24
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
25
26
logger = init_logger(__name__)
27
_PROFILE_REPLAYS = 5
@@ -413,7 +413,7 @@ class AdaptiveVerificationManager:
413
if draft_budget == 0:
414
capacities.zero_()
415
else:
416
- async_copy_to_gpu(scheduled_drafts, out=capacities)
416
+ async_tensor_h2d(scheduled_drafts, out=capacities)
417
if draft_budget < int(scheduled_drafts.sum()):
418
_assign_draft_token_budget_compiled(
419
self._confidence_probs,
@@ -424,10 +424,7 @@ class AdaptiveVerificationManager:
424
)
425
426
num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs]
427
- async_copy_to_gpu(
428
- num_non_draft_tokens,
429
- out=num_non_draft_tokens_gpu,
430
- )
427
+ async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu)
428
self._cu_num_logits[:1].zero_()
429
torch.cumsum(
430
capacities + self.num_bonus_tokens,
vllm/v1/worker/gpu/structured_outputs.py
+2
-3
@@ -5,8 +5,7 @@ import torch
5
6
from vllm.triton_utils import tl, triton
7
from vllm.utils.math_utils import cdiv
8
-from vllm.utils.torch_utils import PIN_MEMORY
9
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
8
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
9
from vllm.v1.worker.gpu.input_batch import InputBatch
10
11
@@ -68,7 +67,7 @@ class StructuredOutputsWorker:
67
68
# Asynchronously copy the bitmask to GPU.
69
with torch.cuda.stream(self.copy_stream):
71
- bitmask = async_copy_to_gpu(
70
+ bitmask = async_tensor_h2d(
71
grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]]
72
)
73