@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7df3d7dada

[Core] Ensure memory is pinned prior to async h2d copy (#45424)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jun 20, 2026 at 20:02 UTC 7df3d7dada840c68b85b26b79de7f59f676d58e3

49 files changed +254 -264

tests/v1/logits_processors/test_correctness.py

-2

@@ -145,7 +145,6 @@ def _generate_fake_sampling_metadata(

145 vllm_config.scheduler_config.max_num_seqs,

146 num_spec,

147 device,

148 - PIN_MEMORY_AVAILABLE,

148 )

149 fake_sampling_metadata = SamplingMetadata(

150 temperature=torch.full((batch_size,), 0.0),

@@ -880,7 +879,6 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():

879 cfg.scheduler_config.max_num_seqs,

880 0,

881 torch.device("cpu"),

883 - False,

882 )

883 is None

884 )

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

-1

@@ -35,7 +35,6 @@ def mock_model_runner_with_input_batch():

35 max_model_len=1024,

36 max_num_batched_tokens=1024,

37 device="cpu",

38 - pin_memory=False,

38 vocab_size=32000,

39 block_sizes=[16],

40 kernel_block_sizes=[16],

tests/v1/worker/test_gpu_input_batch.py

-6

@@ -10,7 +10,6 @@ import torch

11 from vllm.platforms import current_platform

12 from vllm.sampling_params import SamplingParams

13 -from vllm.utils.platform_utils import is_pin_memory_available

13 from vllm.utils.torch_utils import make_tensor_with_pad

14 from vllm.v1.pool.metadata import PoolingMetadata

15 from vllm.v1.sample.logits_processor import LogitsProcessors

@@ -236,7 +235,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

235 max_model_len=1024,

236 max_num_batched_tokens=1024,

237 device=torch.device(device),

239 - pin_memory=is_pin_memory_available(),

238 vocab_size=1024,

239 block_sizes=[1],

240 kernel_block_sizes=[1],

@@ -331,7 +329,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis

329 max_model_len=1024,

330 max_num_batched_tokens=1024,

331 device=torch.device(device),

334 - pin_memory=is_pin_memory_available(),

332 vocab_size=1024,

333 block_sizes=[1],

334 kernel_block_sizes=[1],

@@ -341,7 +338,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis

338 max_model_len=1024,

339 max_num_batched_tokens=1024,

340 device=torch.device(device),

344 - pin_memory=is_pin_memory_available(),

341 vocab_size=1024,

342 block_sizes=[1],

343 kernel_block_sizes=[1],

@@ -410,7 +406,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):

406 max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,

407 max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),

408 device=torch.device(device),

413 - pin_memory=is_pin_memory_available(),

409 vocab_size=VOCAB_SIZE,

410 block_sizes=[16],

411 kernel_block_sizes=[16],

@@ -459,7 +454,6 @@ def test_pooling_metadata_token_id_buffers(

454 max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,

455 max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,

456 device=torch.device("cpu"),

462 - pin_memory=False,

457 vocab_size=VOCAB_SIZE,

458 block_sizes=[16],

459 kernel_block_sizes=[16],

tests/v1/worker/test_gpu_model_runner.py

-4

@@ -85,7 +85,6 @@ def initialize_kv_cache(runner: GPUModelRunner):

85 max_model_len=runner.max_model_len,

86 max_num_batched_tokens=runner.max_num_tokens,

87 device=runner.device,

88 - pin_memory=runner.pin_memory,

88 vocab_size=runner.model_config.get_vocab_size(),

89 block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],

90 kernel_block_sizes=[

@@ -1405,7 +1404,6 @@ def test_input_batch_with_kernel_block_sizes():

1404 max_model_len = 512

1405 max_num_batched_tokens = 512

1406 device = torch.device(DEVICE_TYPE)

1408 - pin_memory = False

1407 vocab_size = 50272

1408

1409 # Test with different kernel block sizes

@@ -1417,7 +1415,6 @@ def test_input_batch_with_kernel_block_sizes():

1415 max_model_len=max_model_len,

1416 max_num_batched_tokens=max_num_batched_tokens,

1417 device=device,

1420 - pin_memory=pin_memory,

1418 vocab_size=vocab_size,

1419 block_sizes=block_sizes,

1420 kernel_block_sizes=kernel_block_sizes,

@@ -1478,7 +1475,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):

1475 max_model_len=runner.max_model_len,

1476 max_num_batched_tokens=runner.max_num_tokens,

1477 device=runner.device,

1481 - pin_memory=runner.pin_memory,

1478 vocab_size=runner.model_config.get_vocab_size(),

1479 block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],

1480 kernel_block_sizes=[16],

vllm/device_allocator/cumem.py

+2 -2

@@ -18,8 +18,8 @@ import torch

19 from vllm.device_allocator import AllocationData, HandleType

20 from vllm.logger import init_logger

21 -from vllm.utils.platform_utils import is_pin_memory_available

21 from vllm.utils.system_utils import find_loaded_library

22 +from vllm.utils.torch_utils import PIN_MEMORY

24 logger = init_logger(__name__)

@@ -196,7 +196,7 @@ class CuMemAllocator:

196 size_in_bytes,

197 dtype=torch.uint8,

198 device="cpu",

199 - pin_memory=is_pin_memory_available(),

199 + pin_memory=PIN_MEMORY,

200 )

201 cpu_ptr = cpu_backup_tensor.data_ptr()

202 libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)

vllm/device_allocator/xpumem.py

+2 -2

@@ -11,7 +11,7 @@ import torch

12 from vllm.device_allocator import AllocationData, HandleType

13 from vllm.logger import init_logger

14 -from vllm.utils.platform_utils import is_pin_memory_available

14 +from vllm.utils.torch_utils import PIN_MEMORY

16 logger = init_logger(__name__)

@@ -188,7 +188,7 @@ class XpuMemAllocator:

188 size_in_bytes,

189 dtype=torch.uint8,

190 device="cpu",

191 - pin_memory=is_pin_memory_available(),

191 + pin_memory=PIN_MEMORY,

192 )

193 cpu_ptr = cpu_backup_tensor.data_ptr()

194 _xpu_memcpy_sync(

vllm/lora/lora_model.py

+2 -2

@@ -17,7 +17,7 @@ from vllm.lora.utils import (

17 )

18 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

19 from vllm.model_executor.models.utils import WeightsMapper

20 -from vllm.utils.platform_utils import is_pin_memory_available

20 +from vllm.utils.torch_utils import PIN_MEMORY

22 logger = init_logger(__name__)

@@ -126,7 +126,7 @@ class LoRAModel:

126 skip_prefixes: list[str] | None = None,

127 ) -> "LoRAModel":

128 """Create a LoRAModel from a dictionary of tensors."""

129 - pin_memory = str(device) == "cpu" and is_pin_memory_available()

129 + pin_memory = str(device) == "cpu" and PIN_MEMORY

130 loras: dict[str, LoRALayerWeights] = {}

131 for tensor_name, tensor in tensors.items():

132 if is_base_embedding_weights(tensor_name):

vllm/lora/lora_weights.py

+2 -2

@@ -7,7 +7,7 @@ import torch

7 import torch.types

9 from vllm.lora.peft_helper import PEFTHelper

10 -from vllm.utils.platform_utils import is_pin_memory_available

10 +from vllm.utils.torch_utils import PIN_MEMORY

13 class LoRALayerWeights:

@@ -79,7 +79,7 @@ class LoRALayerWeights:

79 dtype: torch.dtype,

80 device: torch.types.Device,

81 ) -> "LoRALayerWeights":

82 - pin_memory = str(device) == "cpu" and is_pin_memory_available()

82 + pin_memory = str(device) == "cpu" and PIN_MEMORY

83 lora_a = torch.zeros(

84 [rank, input_dim], dtype=dtype, device=device, pin_memory=pin_memory

85 )

vllm/lora/model_manager.py

+2 -2

@@ -42,7 +42,7 @@ from vllm.model_executor.models.utils import PPMissingLayer

42 from vllm.multimodal import MULTIMODAL_REGISTRY

43 from vllm.multimodal.encoder_budget import MultiModalBudget

44 from vllm.utils.cache import LRUCache

45 -from vllm.utils.platform_utils import is_pin_memory_available

45 +from vllm.utils.torch_utils import PIN_MEMORY

47 logger = init_logger(__name__)

@@ -801,7 +801,7 @@ class LoRAModelManager:

801 # 2. The weight packing above (e.g., pack_moe) may invalidate the

802 # pin_memory allocation, so we execute it after packing.

803

804 - pin_memory = str(lora_device) == "cpu" and is_pin_memory_available()

804 + pin_memory = str(lora_device) == "cpu" and PIN_MEMORY

805 if pin_memory:

806 for lora in lora_model.loras.values():

807 if isinstance(lora.lora_a, list):

vllm/model_executor/layers/attention/mla_attention.py

+10 -8

@@ -1684,12 +1684,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1684 # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

1685 # Note(simon): this is done in CPU because of downstream's

1686 # of `to_list`.

1687 - chunk_starts = (

1687 + chunk_starts = torch.empty(

1688 + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1689 + ).copy_(

1690 torch.arange(num_chunks, dtype=torch.int32)

1691 + .multiply_(max_context_chunk)

1692 .unsqueeze(1)

1690 - .expand(-1, num_prefills)

1691 - * max_context_chunk

1692 - ).pin_memory()

1693 + )

1694 chunk_ends = torch.min(

1695 context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk

1696 )

@@ -1746,12 +1747,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1747 )

1748 * self.dcp_local_block_size

1749 )

1749 - local_chunk_starts = (

1750 + local_chunk_starts = torch.empty(

1751 + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1752 + ).copy_(

1753 torch.arange(num_chunks, dtype=torch.int32)

1754 + .multiply_(padded_local_max_context_chunk_across_ranks)

1755 .unsqueeze(1)

1752 - .expand(-1, num_prefills)

1753 - * padded_local_max_context_chunk_across_ranks

1754 - ).pin_memory()

1756 + )

1757 local_chunk_ends = torch.min(

1758 padded_local_context_lens_cpu.unsqueeze(0),

1759 local_chunk_starts

vllm/model_executor/layers/attention/mm_encoder_attention.py

+2 -1

@@ -28,6 +28,7 @@ from vllm.utils.flashinfer import (

28 is_flashinfer_cudnn_fp8_prefill_attn_supported,

29 )

30 from vllm.utils.math_utils import round_up

31 +from vllm.utils.torch_utils import async_tensor_h2d

32 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

33 from vllm.v1.attention.backends.registry import AttentionBackendEnum

34 from vllm.v1.attention.ops.vit_attn_wrappers import (

@@ -311,7 +312,7 @@ class MMEncoderAttention(CustomOp):

312 )

313 cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])

314

314 - cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True)

315 + cu_seqlens = async_tensor_h2d(cu_seqlens, device=device)

316 return cu_seqlens

317

318 def __init__(

vllm/model_executor/layers/pooler/seqwise/methods.py

+8 -8

@@ -10,6 +10,7 @@ import torch.nn as nn

10 from vllm.config.pooler import SequencePoolingType

11 from vllm.model_executor.layers.pooler import PoolingParamsUpdate

12 from vllm.tasks import PoolingTask

13 +from vllm.utils.torch_utils import async_tensor_h2d

14 from vllm.v1.pool.metadata import PoolingMetadata

16 SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]

@@ -74,15 +75,14 @@ class MeanPool(SequencePoolingMethod):

75 # early return for empty batch

76 return hidden_states.new_empty((0, hidden_size), dtype=torch.float32)

77 - # Build segment_ids on CPU so repeat_interleave doesn't need to sync

78 - # GPU->CPU to learn its data-dependent output length, then upload

79 - # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]

78 + prompt_lens = async_tensor_h2d(

79 + prompt_lens_cpu, device=hidden_states.device, dtype=torch.int64

80 + )

81 + # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]

82 segment_ids = torch.repeat_interleave(

81 - torch.arange(num_seqs, dtype=torch.long),

82 - prompt_lens_cpu,

83 - ).to(hidden_states.device, non_blocking=True)

84 - prompt_lens = prompt_lens_cpu.to(

85 - hidden_states.device, dtype=torch.int64, non_blocking=True

83 + torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long),

84 + prompt_lens,

85 + output_size=int(prompt_lens_cpu.sum()),

86 )

87 segment_sums = torch.zeros(

88 (num_seqs, hidden_size),

vllm/model_executor/models/moonvit.py

+3 -2

@@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import maybe_prefix

66 from vllm.model_executor.models.vision import is_vit_use_data_parallel

67 from vllm.platforms import current_platform

68 from vllm.transformers_utils.configs.moonvit import MoonViTConfig

69 +from vllm.utils.torch_utils import async_tensor_h2d

72 def _apply_rope_input_validation(x, freqs_cis):

@@ -758,7 +759,7 @@ class MoonVitPretrainedModel(PreTrainedModel):

759 ),

760 ]

761 )

761 - metadata["cu_seqlens"] = torch.from_numpy(cu_seqlens_np).to(device)

762 + metadata["cu_seqlens"] = async_tensor_h2d(cu_seqlens_np, device=device)

763

764 if max_seqlen_override is not None:

765 max_seqlen_val = int(max_seqlen_override)

@@ -770,7 +771,7 @@ class MoonVitPretrainedModel(PreTrainedModel):

771 metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)

772

773 gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size)

773 - metadata["merge_gather_idx"] = torch.from_numpy(gather_idx_np).to(device)

774 + metadata["merge_gather_idx"] = async_tensor_h2d(gather_idx_np, device=device)

775

776 return metadata

777

vllm/model_executor/models/qwen2_5_vl.py

+2 -3

@@ -83,9 +83,8 @@ from vllm.multimodal.parse import MultiModalDataItems

83 from vllm.multimodal.processing import PromptReplacement, PromptUpdate

84 from vllm.platforms import current_platform

85 from vllm.sequence import IntermediateTensors

86 -from vllm.utils.platform_utils import is_pin_memory_available

86 from vllm.utils.tensor_schema import TensorSchema, TensorShape

88 -from vllm.utils.torch_utils import async_tensor_h2d

87 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

88 from vllm.v1.attention.backends.registry import AttentionBackendEnum

89 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

@@ -825,7 +824,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

824 @staticmethod

825 def invert_permutation(perm: torch.Tensor) -> torch.Tensor:

826 # building the inverse permutation in O(n) time

828 - inv = torch.empty_like(perm, pin_memory=is_pin_memory_available())

827 + inv = torch.empty_like(perm, pin_memory=PIN_MEMORY)

828 inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)

829 return inv

830

vllm/models/deepseek_v4/sparse_mla.py

+2 -1

@@ -13,6 +13,7 @@ from vllm.config.cache import CacheDType

13 from vllm.platforms.interface import DeviceCapability

14 from vllm.triton_utils import tl, triton

15 from vllm.utils.math_utils import cdiv

16 +from vllm.utils.torch_utils import np_to_pinned_tensor

17 from vllm.v1.attention.backend import (

18 AttentionBackend,

19 AttentionCGSupport,

@@ -207,7 +208,7 @@ class DeepseekV4FlashMLAMetadataBuilder(

208 # Zero-fill for cudagraphs

209 self.req_id_per_token_buffer.fill_(0)

210 self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

210 - torch.from_numpy(req_id_per_token), non_blocking=True

211 + np_to_pinned_tensor(req_id_per_token), non_blocking=True

212 )

213 req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

214

vllm/multimodal/inputs.py

+14 -2

@@ -488,7 +488,13 @@ class MultiModalBatchedField(BaseMultiModalField):

488 # An optimization when `batch` contains only one tensor:

489 # - produce exactly same result as `torch.stack(batch)`

490 # - will achieve zero-copy if the tensor is contiguous

491 - return batch[0].unsqueeze(0).contiguous()

491 + out = batch[0].unsqueeze(0)

492 + if not pin_memory:

493 + return out.contiguous()

494 + # Avoid extra copy - pinning unpinned memory will make it contiguous

495 + if not out.is_contiguous() and out.is_pinned():

496 + out = out.contiguous()

497 + return out.pin_memory()

498 first_shape = batch[0].shape

499 if all(elem.shape == first_shape for elem in batch):

500 out = torch.empty(

@@ -538,7 +544,13 @@ class MultiModalFlatField(BaseMultiModalField):

544 # An optimization when `batch` contains only one tensor:

545 # - produce exactly same result as `torch.concat(batch)`

546 # - will achieve zero-copy if the tensor is contiguous

541 - return batch[0].contiguous()

547 + out = batch[0]

548 + if not pin_memory:

549 + return out.contiguous()

550 + # Avoid extra copy - pinning unpinned memory will make it contiguous

551 + if not out.is_contiguous() and out.is_pinned():

552 + out = out.contiguous()

553 + return out.pin_memory()

554

555 dim = self.dim + (self.dim < 0) * len(batch[0].shape)

556

vllm/platforms/__init__.py

+1 -2

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING

9 from vllm import envs

10 from vllm.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group

11 from vllm.utils.import_utils import resolve_obj_by_qualname

12 -from vllm.utils.torch_utils import supports_xccl

13 from .interface import CpuArchEnum, Platform, PlatformEnum

@@ -135,7 +134,7 @@ def xpu_platform_plugin() -> str | None:

134 try:

135 import torch

136

138 - if supports_xccl():

137 + if torch.distributed.is_xccl_available():

138 dist_backend = "xccl"

139 from vllm.platforms.xpu import XPUPlatform

140

vllm/platforms/cuda.py

+2 -1

@@ -23,7 +23,6 @@ import vllm._C_stable_libtorch # noqa

23 import vllm.envs as envs

24 from vllm.logger import init_logger

25 from vllm.utils.import_utils import import_pynvml

26 -from vllm.utils.torch_utils import is_quantized_kv_cache

26 from vllm.v1.attention.backends.registry import AttentionBackendEnum

28 from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl

@@ -88,6 +87,8 @@ def _get_backend_priorities(

87 kv_cache_dtype: CacheDType | None = None,

88 ) -> list[AttentionBackendEnum]:

89 """Get backend priorities with lazy import to avoid circular dependency."""

90 + from vllm.utils.torch_utils import is_quantized_kv_cache

91 +

92 if use_mla:

93 if device_capability.major == 10:

94 # Sparse MLA backend priorities

vllm/platforms/xpu.py

+6 -6

@@ -14,7 +14,6 @@ import vllm_xpu_kernels._xpu_C # noqa

15 import vllm.envs as envs

16 from vllm.logger import init_logger

17 -from vllm.utils.torch_utils import supports_xpu_graph

17 from vllm.v1.attention.backends.registry import AttentionBackendEnum

19 from .interface import DeviceCapability, Platform, PlatformEnum

@@ -178,8 +177,6 @@ class XPUPlatform(Platform):

177

178 @classmethod

179 def check_and_update_config(cls, vllm_config: VllmConfig) -> None:

181 - parallel_config = vllm_config.parallel_config

182 -

180 # lazy import to avoid circular import

181 from vllm.config import CUDAGraphMode

182

@@ -190,6 +187,10 @@ class XPUPlatform(Platform):

187 attention_config = vllm_config.attention_config

188 if attention_config.backend is None:

189 attention_config.backend = AttentionBackendEnum.FLASH_ATTN

190 +

191 + # lazy import to avoid circular import

192 + from vllm.utils.torch_utils import supports_xpu_graph

193 +

194 if not supports_xpu_graph():

195 compilation_config.cudagraph_mode = CUDAGraphMode.NONE

196 logger.warning(

@@ -324,9 +325,8 @@ class XPUPlatform(Platform):

325

326 @classmethod

327 def get_device_communicator_cls(cls) -> str:

327 - from vllm.utils.torch_utils import supports_xccl

328 -

329 - if not supports_xccl():

328 + if not torch.distributed.is_xccl_available():

329 + # Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform

330 logger.warning(

331 "xccl is not enabled in this torch build, communication"

332 " is not available."

vllm/utils/torch_utils.py

+18 -15

@@ -3,7 +3,6 @@

3 import contextlib

4 import importlib.metadata

5 import os

6 -import platform

6 import random

7 import threading

8 from collections.abc import Callable, Collection

@@ -18,6 +17,7 @@ from torch.library import Library, infer_schema

18 import vllm.envs as envs

19 from vllm.logger import init_logger

20 +from vllm.utils.platform_utils import is_pin_memory_available

22 if TYPE_CHECKING:

23 from vllm.config import ModelConfig

@@ -68,9 +68,7 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {

68 T = TypeVar("T")

71 -# Pin memory in non-WSL case.

72 -# Logic duplicated here for now to avoid circular import.

73 -PIN_MEMORY = "microsoft" not in " ".join(platform.uname()).lower()

71 +PIN_MEMORY = is_pin_memory_available()

74 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:

@@ -606,14 +604,24 @@ def create_kv_caches_with_random(

604

605

606 def async_tensor_h2d(

609 - data: list,

610 - dtype: torch.dtype,

607 + data: list | np.ndarray | torch.Tensor,

608 device: str | torch.device,

612 - pin_memory: bool = PIN_MEMORY,

609 + dtype: torch.dtype | None = None,

610 ) -> torch.Tensor:

614 - """Asynchronously create a tensor and copy it from host to device."""

615 - t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")

616 - return t.to(device=device, non_blocking=True)

611 + """Copy list/numpy array/tensor async from host to device."""

612 + if isinstance(data, np.ndarray):

613 + data = torch.from_numpy(data)

614 + if isinstance(data, torch.Tensor):

615 + t = data.pin_memory() if PIN_MEMORY else data

616 + else:

617 + t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu")

618 + assert t.is_cpu

619 + return t.to(device=device, dtype=dtype, non_blocking=True)

620 +

621 +

622 +def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor:

623 + t = torch.from_numpy(array)

624 + return t.pin_memory() if PIN_MEMORY else t

625

626

627 def make_ndarray_with_pad(

@@ -914,11 +922,6 @@ def _encode_layer_name(layer_name: str) -> str | LayerName:

922 return LayerName(layer_name) if _USE_LAYERNAME else layer_name

923

924

917 -# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform

918 -def supports_xccl() -> bool:

919 - return torch.distributed.is_xccl_available()

920 -

921 -

925 # Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform

926 def supports_xpu_graph() -> bool:

927 return is_torch_equal_or_newer("2.11.0.dev")

vllm/v1/attention/backends/flashinfer.py

+2 -4

@@ -41,8 +41,8 @@ from vllm.utils.flashinfer import (

41 use_trtllm_attention,

42 )

43 from vllm.utils.math_utils import cdiv

44 -from vllm.utils.platform_utils import is_pin_memory_available

44 from vllm.utils.torch_utils import (

45 + PIN_MEMORY,

46 canonicalize_singleton_dim_strides,

47 is_quantized_kv_cache,

48 is_strictly_contiguous,

@@ -708,9 +708,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

708 # Since we do not have explicit synchronization in ModelRunnerV2, we do not pin

709 # reused CPU buffers to avoid a race condition between step N async copies to

710 # GPU and step N+1 buffer updates.

711 - self.pin_memory = (

712 - not vllm_config.use_v2_model_runner and is_pin_memory_available()

713 - )

711 + self.pin_memory = not vllm_config.use_v2_model_runner and PIN_MEMORY

712 self.paged_kv_indptr = self._make_buffer(max_num_reqs + 1)

713 self.paged_kv_indptr_cpu_buffer = torch.zeros_like(

714 self.paged_kv_indptr.cpu, pin_memory=self.pin_memory

vllm/v1/attention/backends/flex_attention.py

+6 -2

@@ -28,7 +28,11 @@ from vllm.logger import init_logger

28 from vllm.model_executor.layers.attention import Attention

29 from vllm.platforms import current_platform

30 from vllm.utils.math_utils import cdiv

31 -from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer

31 +from vllm.utils.torch_utils import (

32 + async_tensor_h2d,

33 + is_quantized_kv_cache,

34 + is_torch_equal_or_newer,

35 +)

36 from vllm.v1.attention.backend import (

37 AttentionBackend,

38 AttentionCGSupport,

@@ -58,7 +62,7 @@ def _offsets_to_doc_ids_tensor(

62 doc_ids = torch.repeat_interleave(

63 torch.arange(len(counts), dtype=torch.int32), counts

64 )

61 - return doc_ids.to(device, non_blocking=True)

65 + return async_tensor_h2d(doc_ids, device=device)

68 def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):

vllm/v1/attention/backends/gdn_attn.py

+11 -8

@@ -8,6 +8,7 @@ from typing import Literal

8 import torch

10 from vllm.config import VllmConfig

11 +from vllm.utils.torch_utils import async_tensor_h2d

12 from vllm.v1.attention.backend import (

13 AttentionBackend,

14 AttentionCGSupport,

@@ -203,8 +204,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

204 spec_sequence_masks = None

205 spec_sequence_masks_cpu = None

206 else:

206 - spec_sequence_masks = spec_sequence_masks_cpu.to(

207 - query_start_loc.device, non_blocking=True

207 + spec_sequence_masks = async_tensor_h2d(

208 + spec_sequence_masks_cpu, device=query_start_loc.device

209 )

210

211 if spec_sequence_masks is None:

@@ -376,12 +377,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

377 )

378

379 assert prefill_query_start_loc_cpu is not None

379 - chunk_indices = prepare_chunk_indices(

380 - prefill_query_start_loc_cpu, FLA_CHUNK_SIZE

381 - ).to(device=gpu_device, non_blocking=True)

382 - chunk_offsets = prepare_chunk_offsets(

383 - prefill_query_start_loc_cpu, FLA_CHUNK_SIZE

384 - ).to(device=gpu_device, non_blocking=True)

380 + chunk_indices = async_tensor_h2d(

381 + prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),

382 + device=gpu_device,

383 + )

384 + chunk_offsets = async_tensor_h2d(

385 + prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),

386 + device=gpu_device,

387 + )

388

389 if num_prefills > 0:

390 has_initial_state = context_lens_tensor > 0

vllm/v1/attention/backends/mamba2_attn.py

+14 -13

@@ -7,6 +7,7 @@ from typing import Any

7 import torch

9 from vllm.config import VllmConfig

10 +from vllm.utils.torch_utils import async_tensor_h2d

11 from vllm.v1.attention.backend import (

12 AttentionBackend,

13 CommonAttentionMetadata,

@@ -68,22 +69,22 @@ def compute_varlen_chunk_metadata(

70 # Exclusive prefix sum over logical-chunk lengths

71 if chunk_lens:

71 - cu_chunk_seqlens = torch.tensor(

72 - [0] + list(itertools.accumulate(chunk_lens)),

73 - device=device,

74 - dtype=torch.int32,

75 - )

76 - # Final boundary must equal total tokens

77 - assert int(cu_chunk_seqlens[-1].item()) == total

72 + cu_chunk_seqlens_list = [0] + list(itertools.accumulate(chunk_lens))

73 + # Final boundary must equal total tokens (check on host to avoid a sync)

74 + assert cu_chunk_seqlens_list[-1] == total

75 else:

79 - cu_chunk_seqlens = torch.tensor([0], device=device, dtype=torch.int32)

76 + cu_chunk_seqlens_list = [0]

77 + cu_chunk_seqlens = async_tensor_h2d(

78 + cu_chunk_seqlens_list, dtype=torch.int32, device=device

79 + )

81 - last_chunk_indices_t = (

82 - torch.tensor(last_chunk_indices, device=device, dtype=torch.int32)

83 - if len(starts) > 0

84 - else torch.empty((0,), device=device, dtype=torch.int32)

81 + # last_chunk_indices is empty when there are no sequences (len(starts) == 0).

82 + last_chunk_indices_t = async_tensor_h2d(

83 + last_chunk_indices, dtype=torch.int32, device=device

84 + )

85 + seq_idx_chunks_t = async_tensor_h2d(

86 + seq_idx_chunks, dtype=torch.int32, device=device

87 )

86 - seq_idx_chunks_t = torch.tensor(seq_idx_chunks, device=device, dtype=torch.int32)

88 return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

+2 -2

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

26 get_mla_dims,

27 )

28 from vllm.platforms.interface import DeviceCapability

29 -from vllm.utils.torch_utils import is_quantized_kv_cache

29 +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

30 from vllm.v1.attention.backend import (

31 AttentionBackend,

32 AttentionCGSupport,

@@ -217,7 +217,7 @@ class FlashInferMLASparseMetadataBuilder(

217 # Zero-fill for cudagraphs

218 self.req_id_per_token_buffer.fill_(0)

219 self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

220 - torch.from_numpy(req_id_per_token), non_blocking=True

220 + np_to_pinned_tensor(req_id_per_token), non_blocking=True

221 )

222 req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens]

223

vllm/v1/attention/backends/mla/flashmla_sparse.py

+2 -2

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

16 from vllm.platforms import current_platform

17 from vllm.platforms.interface import DeviceCapability

18 from vllm.utils.platform_utils import num_compute_units

19 -from vllm.utils.torch_utils import is_quantized_kv_cache

19 +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

20 from vllm.v1.attention.backend import (

21 AttentionBackend,

22 AttentionCGSupport,

@@ -503,7 +503,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

503 # Zero-fill for cudagraphs

504 self.req_id_per_token_buffer.fill_(0)

505 self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

506 - torch.from_numpy(req_id_per_token), non_blocking=True

506 + np_to_pinned_tensor(req_id_per_token), non_blocking=True

507 )

508 req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

509

vllm/v1/attention/backends/utils.py

+16 -19

@@ -17,7 +17,7 @@ from typing_extensions import runtime_checkable

18 from vllm.config import VllmConfig, get_layers_from_vllm_config

19 from vllm.utils.math_utils import cdiv

20 -from vllm.utils.torch_utils import async_tensor_h2d

20 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor

21 from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec

23 if TYPE_CHECKING:

@@ -364,8 +364,8 @@ def make_local_attention_virtual_batches(

364 # tensor first, which recovers perf.

365 # Upload the index tensors to the block_table's device up-front so that the

366 # fancy indexing below doesn't implicitly force a synchronous H2D copy.

367 - batch_indices_torch = torch.from_numpy(batch_indices).to(device, non_blocking=True)

368 - block_indices_torch = torch.from_numpy(block_indices).to(device, non_blocking=True)

367 + batch_indices_torch = async_tensor_h2d(batch_indices, device=device)

368 + block_indices_torch = async_tensor_h2d(block_indices, device=device)

369

370 # Save as a lambda so we can return this for update_block_table

371 make_block_table = lambda block_table: block_table[

@@ -379,8 +379,8 @@ def make_local_attention_virtual_batches(

379

380 return CommonAttentionMetadata(

381 query_start_loc_cpu=query_start_loc_cpu,

382 - query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),

383 - seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),

382 + query_start_loc=async_tensor_h2d(query_start_loc_cpu, device=device),

383 + seq_lens=async_tensor_h2d(seq_lens_cpu, device=device),

384 num_reqs=len(seq_lens_cpu),

385 num_actual_tokens=common_attn_metadata.num_actual_tokens,

386 max_query_len=seqlens_q_local.max(),

@@ -808,14 +808,12 @@ def create_fast_prefill_custom_backend(

808

809

810 def compute_causal_conv1d_metadata(

811 - query_start_loc_p_cpu: torch.Tensor,

812 - *,

813 - device: torch.device,

814 -):

811 + query_start_loc_p_cpu: torch.Tensor, *, device: torch.device

812 +) -> tuple[dict[int, dict[str, Any]], torch.Tensor, torch.Tensor]:

813 # Needed for causal_conv1d. Use the CPU query_start_loc to avoid DtoH sync.

814 assert query_start_loc_p_cpu.device.type == "cpu"

815 seqlens = query_start_loc_p_cpu.diff()

818 - nums_dict = {} # type: ignore

816 + nums_dict: dict[int, dict[str, Any]] = {}

817 batch_ptr = None

818 token_chunk_offset_ptr = None

819 for BLOCK_M in [8]: # cover all BLOCK_M values

@@ -823,7 +821,7 @@ def compute_causal_conv1d_metadata(

821 nums_dict[BLOCK_M] = {}

822 nums_dict[BLOCK_M]["nums"] = nums

823 nums_dict[BLOCK_M]["tot"] = nums.sum().item()

826 - mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums))

824 + mlist = np_to_pinned_tensor(np.repeat(np.arange(len(nums)), nums))

825 nums_dict[BLOCK_M]["mlist"] = mlist

826 mlist_len = len(nums_dict[BLOCK_M]["mlist"])

827 nums_dict[BLOCK_M]["mlist_len"] = mlist_len

@@ -831,7 +829,7 @@ def compute_causal_conv1d_metadata(

829 offsetlist = [] # type: ignore

830 for idx, num in enumerate(nums):

831 offsetlist.extend(range(num))

834 - offsetlist = torch.tensor(offsetlist, dtype=torch.int32)

832 + offsetlist = torch.tensor(offsetlist, dtype=torch.int32, pin_memory=PIN_MEMORY)

833 nums_dict[BLOCK_M]["offsetlist"] = offsetlist

834

835 if batch_ptr is None:

@@ -845,16 +843,15 @@ def compute_causal_conv1d_metadata(

843 else:

844 if batch_ptr.nelement() < MAX_NUM_PROGRAMS:

845 batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)

848 - token_chunk_offset_ptr.resize_( # type: ignore

849 - MAX_NUM_PROGRAMS

850 - ).fill_(PAD_SLOT_ID)

846 + assert token_chunk_offset_ptr is not None

847 + token_chunk_offset_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)

848

849 + assert batch_ptr is not None

850 batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True)

853 - token_chunk_offset_ptr[ # type: ignore

854 - 0:mlist_len

855 - ].copy_(offsetlist, non_blocking=True)

851 + assert token_chunk_offset_ptr is not None

852 + token_chunk_offset_ptr[0:mlist_len].copy_(offsetlist, non_blocking=True)

853 nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr

857 - nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr # type: ignore

854 + nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr

855

856 return nums_dict, batch_ptr, token_chunk_offset_ptr

857

vllm/v1/kv_offload/cpu/gpu_worker.py

+3 -3

@@ -14,7 +14,7 @@ from vllm.logger import init_logger

14 from vllm.platforms import current_platform

15 from vllm.triton_utils import HAS_TRITON, triton

16 from vllm.utils.math_utils import cdiv

17 -from vllm.utils.platform_utils import is_pin_memory_available

17 +from vllm.utils.torch_utils import PIN_MEMORY

18 from vllm.v1.kv_offload.base import (

19 BlockIDsLoadStoreSpec,

20 CanonicalKVCacheRef,

@@ -156,7 +156,7 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:

156 def _new_descriptor_buffers(

157 num_copy_ops: int,

158 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

159 - pin = is_pin_memory_available()

159 + pin = PIN_MEMORY

160 # CUDA cache_kernels.cu requires int64; XPU DMA engine requires uint64.

161 ptr_dtype = torch.uint64 if current_platform.is_xpu() else torch.int64

162 return (

@@ -482,7 +482,7 @@ class CpuGpuOffloadingHandlers:

482 num_cpu_blocks: int,

483 mmap_region: SharedOffloadRegion | None = None,

484 ):

485 - pin_memory = is_pin_memory_available()

485 + pin_memory = PIN_MEMORY

486 logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))

487 self._mmap_region = mmap_region

488 if mmap_region is not None and pin_memory:

vllm/v1/pool/metadata.py

+2 -4

@@ -7,9 +7,7 @@ import torch

8 from vllm.pooling_params import PoolingParams

9 from vllm.tasks import PoolingTask

10 -from vllm.utils.platform_utils import is_pin_memory_available

11 -

12 -pin_memory = is_pin_memory_available()

10 +from vllm.utils.torch_utils import PIN_MEMORY

13 @dataclass

@@ -134,7 +132,7 @@ class PoolingMetadata:

132 num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)

133 if query_start_loc_gpu is None:

134 cumsum = torch.zeros(

137 - n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"

135 + n_seq + 1, dtype=torch.int64, pin_memory=PIN_MEMORY, device="cpu"

136 )

137 torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])

138 cumsum = cumsum.to(device, non_blocking=True)

vllm/v1/sample/logits_processor/builtin.py

+5 -10

@@ -7,6 +7,7 @@ import numpy as np

7 import torch

9 from vllm import SamplingParams

10 +from vllm.utils.torch_utils import async_tensor_h2d

11 from vllm.v1.sample.logits_processor.interface import (

12 BatchUpdate,

13 LogitsProcessor,

@@ -118,7 +119,6 @@ class MinPLogitsProcessor(LogitsProcessor):

119 class LogitBiasLogitsProcessor(LogitsProcessor):

120 def __init__(self, _, device: torch.device, is_pin_memory: bool):

121 self.device = device

121 - self.pin_memory = is_pin_memory

122 self.biases: dict[int, dict[int, float]] = {}

123

124 self.bias_tensor: torch.Tensor = torch.tensor(())

@@ -154,9 +154,7 @@ class LogitBiasLogitsProcessor(LogitsProcessor):

154 )

155

156 def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:

157 - return torch.tensor(

158 - data, device="cpu", dtype=dtype, pin_memory=self.pin_memory

159 - ).to(device=self.device, non_blocking=True)

157 + return async_tensor_h2d(data, device=self.device, dtype=dtype)

158

159 def apply(self, logits: torch.Tensor) -> torch.Tensor:

160 if self.biases:

@@ -170,7 +168,6 @@ class MinTokensLogitsProcessor(LogitsProcessor):

168 ):

169 # index -> (min_toks, output_token_ids, stop_token_ids)

170 self.device = device

173 - self.pin_memory = is_pin_memory

171 self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}

172

173 # (req_idx_tensor,eos_tok_id_tensor)

@@ -227,9 +224,7 @@ class MinTokensLogitsProcessor(LogitsProcessor):

224 )

225

226 def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:

230 - return torch.tensor(

231 - data, device="cpu", dtype=dtype, pin_memory=self.pin_memory

232 - ).to(device=self.device, non_blocking=True)

227 + return async_tensor_h2d(data, device=self.device, dtype=dtype)

228

229 def apply(self, logits: torch.Tensor) -> torch.Tensor:

230 if self.min_toks:

@@ -283,8 +278,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):

278 toks_arr = np.concatenate(all_toks)

279 # (row_indices, token_indices) for index_put_ to set -inf.

280 logits_slice = (

286 - torch.from_numpy(rows_arr).to(self.device, non_blocking=True),

287 - torch.from_numpy(toks_arr).to(self.device, non_blocking=True),

281 + async_tensor_h2d(rows_arr, device=self.device),

282 + async_tensor_h2d(toks_arr, device=self.device),

283 )

284 logits.index_put_(logits_slice, self.neg_inf_tensor)

285

vllm/v1/sample/ops/penalties.py

+2 -3

@@ -4,8 +4,7 @@

4 import torch

6 from vllm.model_executor.layers.utils import apply_penalties

7 -from vllm.utils.platform_utils import is_pin_memory_available

8 -from vllm.utils.torch_utils import make_tensor_with_pad

7 +from vllm.utils.torch_utils import PIN_MEMORY, make_tensor_with_pad

10 def apply_all_penalties(

@@ -52,6 +51,6 @@ def _convert_to_tensors(

51 pad=vocab_size,

52 device="cpu",

53 dtype=torch.int64,

55 - pin_memory=is_pin_memory_available(),

54 + pin_memory=PIN_MEMORY,

55 )

56 return output_tokens_tensor.to(device, non_blocking=True)

vllm/v1/sample/sampler.py

+2 -2

@@ -6,7 +6,7 @@ import torch

6 import torch.nn as nn

8 from vllm.config.model import LogprobsMode

9 -from vllm.utils.platform_utils import is_pin_memory_available

9 +from vllm.utils.torch_utils import PIN_MEMORY

10 from vllm.v1.outputs import LogprobsTensors, SamplerOutput

11 from vllm.v1.sample.metadata import SamplingMetadata

12 from vllm.v1.sample.ops.bad_words import apply_bad_words

@@ -65,7 +65,7 @@ class Sampler(nn.Module):

65 ):

66 super().__init__()

67 self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel)

68 - self.pin_memory = is_pin_memory_available()

68 + self.pin_memory = PIN_MEMORY

69 self.logprobs_mode = logprobs_mode

70 self.use_fp64_gumbel = use_fp64_gumbel

vllm/v1/sample/thinking_budget_state.py

+2 -3

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any

7 import torch

9 from vllm.platforms import current_platform

10 -from vllm.utils.torch_utils import async_tensor_h2d

10 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

11 from vllm.v1.sample.logits_processor.interface import (

12 BatchUpdate,

13 MoveDirectionality,

@@ -22,12 +22,11 @@ def maybe_create_thinking_budget_state_holder(

22 max_num_seqs: int,

23 num_spec_tokens: int,

24 device: torch.device,

25 - is_pin_memory: bool,

25 ) -> "ThinkingBudgetStateHolder | None":

26 if reasoning_config is None:

27 return None

28 return ThinkingBudgetStateHolder(

30 - reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory

29 + reasoning_config, max_num_seqs, num_spec_tokens, device, PIN_MEMORY

30 )

vllm/v1/serial_utils.py

+2 -2

@@ -33,7 +33,7 @@ from vllm.multimodal.inputs import (

33 MultiModalSharedField,

34 NestedTensors,

35 )

36 -from vllm.utils.platform_utils import is_pin_memory_available

36 +from vllm.utils.torch_utils import PIN_MEMORY

37 from vllm.v1.utils import tensor_data

39 logger = init_logger(__name__)

@@ -327,7 +327,7 @@ class MsgpackDecoder:

327 oob_tensor_provider: OOBTensorProvider | None = None,

328 ):

329 self.share_mem = share_mem

330 - self.pin_tensors = is_pin_memory_available()

330 + self.pin_tensors = PIN_MEMORY

331 args = () if t is None else (t,)

332 self.decoder = msgpack.Decoder(

333 *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook

vllm/v1/simple_kv_offload/worker.py

+2 -2

@@ -8,7 +8,7 @@ import torch

9 from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 -from vllm.utils.platform_utils import is_pin_memory_available

11 +from vllm.utils.torch_utils import PIN_MEMORY

12 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend

13 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor

14 from vllm.v1.simple_kv_offload.metadata import (

@@ -149,7 +149,7 @@ class SimpleCPUOffloadWorker:

149 (self.num_cpu_blocks * total_bytes_per_block) / (1024**3),

150 )

151

152 - pin_memory = is_pin_memory_available()

152 + pin_memory = PIN_MEMORY

153 if not pin_memory:

154 logger.warning(

155 "Pinned memory not available. CPU offload performance may be degraded."

vllm/v1/spec_decode/extract_hidden_states.py

+2 -3

@@ -12,7 +12,7 @@ from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config

12 from vllm.forward_context import set_forward_context

13 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

14 from vllm.model_executor.model_loader import get_model

15 -from vllm.utils.platform_utils import is_pin_memory_available

15 +from vllm.utils.torch_utils import PIN_MEMORY

16 from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata

17 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

18 from vllm.v1.utils import CpuGpuBuffer

@@ -58,7 +58,7 @@ class ExtractHiddenStatesProposer:

58 self.backup_next_token_ids = CpuGpuBuffer(

59 max_batch_size,

60 dtype=torch.int32,

61 - pin_memory=is_pin_memory_available(),

61 + pin_memory=PIN_MEMORY,

62 device=device,

63 with_numpy=True,

64 )

@@ -317,7 +317,6 @@ class ExtractHiddenStatesProposer:

317 (batch_size, 1). For each request we either use the sampled token

318 (if valid and not discarded) or a backup token from the request state.

319 """

320 - num_reqs = gpu_input_batch.num_reqs

320

321 # Precompute backup token IDs for discarded requests.

322 num_reqs = gpu_input_batch.num_reqs

vllm/v1/spec_decode/llm_base_proposer.py

+7 -9

@@ -26,7 +26,7 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

26 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM

27 from vllm.multimodal import MULTIMODAL_REGISTRY

28 from vllm.platforms import current_platform

29 -from vllm.utils.platform_utils import is_pin_memory_available

29 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

30 from vllm.v1.attention.backend import CommonAttentionMetadata

31 from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

@@ -228,7 +228,7 @@ class SpecDecodeBaseProposer:

228 self.backup_next_token_ids = CpuGpuBuffer(

229 self.max_batch_size,

230 dtype=torch.int32,

231 - pin_memory=is_pin_memory_available(),

231 + pin_memory=PIN_MEMORY,

232 device=device,

233 with_numpy=True,

234 )

@@ -239,9 +239,7 @@ class SpecDecodeBaseProposer:

239 self._last_draft_probs: torch.Tensor | None = None

240

241 self._slot_mapping_buffer = torch.zeros(

242 - self.max_positions,

243 - dtype=torch.int64,

244 - device=device,

242 + self.max_positions, dtype=torch.int64, device=device

243 )

244

245 # Determine allowed attention backends once during initialization.

@@ -1127,7 +1125,7 @@ class SpecDecodeBaseProposer:

1125 new_query_start_loc_cpu = torch.zeros(

1126 query_start_loc_cpu.shape,

1127 dtype=torch.int32,

1130 - pin_memory=is_pin_memory_available(),

1128 + pin_memory=PIN_MEMORY,

1129 )

1130 new_query_start_loc_np = new_query_start_loc_cpu.numpy()

1131 np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])

@@ -1160,11 +1158,11 @@ class SpecDecodeBaseProposer:

1158 # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2

1159 # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3

1160 token_indices_np = token_offsets + old_query_start_locs_expanded

1163 - token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)

1161 + token_indices = async_tensor_h2d(token_indices_np, device=device)

1162

1163 spec_common_attn_metadata = CommonAttentionMetadata(

1166 - query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),

1167 - seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),

1164 + query_start_loc=async_tensor_h2d(new_query_start_loc_cpu, device=device),

1165 + seq_lens=async_tensor_h2d(new_seq_lens_cpu, device=device),

1166 query_start_loc_cpu=new_query_start_loc_cpu,

1167 _seq_lens_cpu=new_seq_lens_cpu,

1168 _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

vllm/v1/spec_decode/ngram_proposer_gpu.py

+2 -2

@@ -545,7 +545,7 @@ def update_ngram_gpu_tensors_incremental(

545 num_tokens = input_batch.num_tokens_no_spec[idx]

546 if num_tokens > 0:

547 token_ids_gpu_tensor[idx, :num_tokens].copy_(

548 - input_batch.token_ids_cpu_tensor[idx, :num_tokens],

548 + input_batch.token_ids_cpu_tensor[idx, :num_tokens].pin_memory(),

549 non_blocking=True,

550 )

551

@@ -591,7 +591,7 @@ def update_ngram_gpu_tensors_incremental(

591 num_tokens = input_batch.num_tokens_no_spec[new_req_idx]

592 if num_tokens > 0:

593 token_ids_gpu_tensor[new_req_idx, :num_tokens].copy_(

594 - input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens],

594 + input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens].pin_memory(),

595 non_blocking=True,

596 )

597

vllm/v1/structured_output/backend_lm_format_enforcer.py

+2 -2

@@ -11,7 +11,7 @@ from transformers import PreTrainedTokenizerBase

12 from vllm.sampling_params import SamplingParams

13 from vllm.utils.import_utils import LazyLoader

14 -from vllm.utils.platform_utils import is_pin_memory_available

14 +from vllm.utils.torch_utils import PIN_MEMORY

15 from vllm.v1.structured_output.backend_types import (

16 StructuredOutputBackend,

17 StructuredOutputGrammar,

@@ -139,7 +139,7 @@ class LMFormatEnforcerBackend(StructuredOutputBackend):

139 (max_num_seqs, (self.vocab_size + 31) // 32),

140 -1,

141 dtype=torch.int32,

142 - pin_memory=is_pin_memory_available(),

142 + pin_memory=PIN_MEMORY,

143 )

144

145 def destroy(self):

vllm/v1/structured_output/backend_outlines.py

+2 -2

@@ -15,7 +15,7 @@ from regex import escape as regex_escape

16 from vllm.sampling_params import SamplingParams

17 from vllm.utils.import_utils import LazyLoader

18 -from vllm.utils.platform_utils import is_pin_memory_available

18 +from vllm.utils.torch_utils import PIN_MEMORY

19 from vllm.v1.structured_output.backend_types import (

20 StructuredOutputBackend,

21 StructuredOutputGrammar,

@@ -101,7 +101,7 @@ class OutlinesBackend(StructuredOutputBackend):

101 (max_num_seqs, (self.vocab_size + 31) // 32),

102 -1,

103 dtype=torch.int32,

104 - pin_memory=is_pin_memory_available(),

104 + pin_memory=PIN_MEMORY,

105 )

106

107 def destroy(self):

vllm/v1/structured_output/utils.py

+11 -14

@@ -10,7 +10,6 @@ from collections.abc import Callable

10 from concurrent.futures import ThreadPoolExecutor, TimeoutError

11 from typing import TYPE_CHECKING, TypeVar

13 -import numpy as np

13 import regex as re

14 import torch

15 from cachetools import LRUCache

@@ -18,7 +17,7 @@ from cachetools import LRUCache

17 import vllm.envs as envs

18 from vllm.logger import init_logger

19 from vllm.utils.import_utils import LazyLoader

21 -from vllm.utils.platform_utils import is_pin_memory_available

20 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

21 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

23 if TYPE_CHECKING:

@@ -123,11 +122,13 @@ def apply_grammar_bitmask(

122 out_indices = []

123

124 # Reorder the bitmask to match the order of the requests in the batch.

126 - sorted_bitmask = np.full(

127 - shape=(logits.shape[0], grammar_bitmask.shape[1]),

128 - fill_value=-1,

129 - dtype=grammar_bitmask.dtype,

125 + sorted_bitmask_tensor = torch.full(

126 + (logits.shape[0], grammar_bitmask.shape[1]),

127 + -1,

128 + dtype=torch.from_numpy(grammar_bitmask[:0]).dtype,

129 + pin_memory=PIN_MEMORY,

130 )

131 + sorted_bitmask = sorted_bitmask_tensor.numpy()

132 cumulative_index = 0

133 for req_id in grammar_output.structured_output_request_ids:

134 num_spec_tokens = len(spec_tokens.get(req_id, ()))

@@ -138,10 +139,8 @@ def apply_grammar_bitmask(

139 out_indices.append(bitmask_index)

140 cumulative_index += 1 + num_spec_tokens

141

141 - # Copy async to device as tensor.

142 - grammar_bitmask = torch.from_numpy(sorted_bitmask).to(

143 - logits.device, non_blocking=True

144 - )

142 + # Copy async to device.

143 + grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True)

144

145 # If the length of out indices and the logits have the same shape

146 # we don't need to pass indices to the kernel,

@@ -154,11 +153,9 @@ def apply_grammar_bitmask(

153 # xgrammar expects a python list of indices but it will actually work with

154 # a tensor. If we copy the tensor ourselves here we can do it in a

155 # non_blocking manner and there should be no cpu sync within xgrammar.

157 - pin_memory = is_pin_memory_available()

158 - index_tensor = torch.tensor(

159 - out_indices, dtype=torch.int32, device="cpu", pin_memory=pin_memory

156 + index_tensor = async_tensor_h2d(

157 + out_indices, dtype=torch.int32, device=logits.device

158 )

161 - index_tensor = index_tensor.to(logits.device, non_blocking=True)

159

160 xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)

161 return

vllm/v1/utils.py

+2 -1

@@ -31,6 +31,7 @@ from vllm.logger import init_logger

31 from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message

32 from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri

33 from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title

34 +from vllm.utils.torch_utils import PIN_MEMORY

35 from vllm.v1.core.sched.output import SchedulerOutput

37 if TYPE_CHECKING:

@@ -114,7 +115,7 @@ class CpuGpuBuffer:

115 *size: int | torch.SymInt,

116 dtype: torch.dtype,

117 device: torch.device,

117 - pin_memory: bool,

118 + pin_memory: bool = PIN_MEMORY,

119 with_numpy: bool = True,

120 ) -> None:

121 # these buffers are mutable runtime state, so allocate them as normal

vllm/v1/worker/cpu/shm.py

+8 -3

@@ -7,6 +7,8 @@

8 from typing import Any

10 +import numpy as np

11 +

12 # Patch torch APIs

13 import torch

@@ -45,11 +47,14 @@ import vllm.utils.torch_utils as torch_utils

49 def async_tensor_h2d(

48 - data: list,

49 - dtype: torch.dtype,

50 + data: list | np.ndarray | torch.Tensor,

51 device: str | torch.device,

51 - pin_memory: bool = False,

52 + dtype: torch.dtype | None = None,

53 ) -> torch.Tensor:

54 + if isinstance(data, np.ndarray):

55 + data = torch.from_numpy(data)

56 + if isinstance(data, torch.Tensor):

57 + return data.to(dtype=dtype)

58 return torch.tensor(data, dtype=dtype, device="cpu")

vllm/v1/worker/gpu/buffer_utils.py

+5 -6

@@ -36,10 +36,9 @@ def async_copy_to_gpu(

36 assert device is not None

37 out = torch.empty_like(x, device=device)

39 - # Copy directly to GPU — explicit pin_memory() causes sporadic stalls

40 - # under high concurrency due to CUDA driver contention. The driver

41 - # handles the transfer efficiently without manual pinning.

42 - return out.copy_(x, non_blocking=True)

39 + # pin_memory() is no-op if the memory is already pinned.

40 + pinned = x.pin_memory()

41 + return out.copy_(pinned, non_blocking=True)

44 class UvaBuffer:

@@ -183,7 +182,7 @@ class StagedWriteTensor:

182

183 # Special handling for write_contents

184 write_contents = async_tensor_h2d(

186 - self._staged_write_contents, self.dtype, self.device

185 + self._staged_write_contents, device=self.device, dtype=self.dtype

186 )

187

188 # Write diffs to the GPU buffer

@@ -255,7 +254,7 @@ class FusedStagedWriter:

254 indices_uva = self.indices.copy_to_uva(indices)

255 starts_uva = self.starts.copy_to_uva(starts)

256 cu_lens_uva = self.cu_lens.copy_to_uva(cu_lens)

258 - contents_gpu = async_tensor_h2d(contents, torch.int32, self.device)

257 + contents_gpu = async_tensor_h2d(contents, device=self.device, dtype=torch.int32)

258

259 _apply_write_kernel[(len(group_ids),)](

260 output_ptrs,

vllm/v1/worker/gpu/mm/encoder_runner.py

+2 -3

@@ -49,12 +49,11 @@ class EncoderRunner:

50 @torch.inference_mode()

51 def execute_mm_encoder(

52 - self,

53 - mm_kwargs: list[tuple[str, MultiModalKwargsItem]],

52 + self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]

53 ) -> list[torch.Tensor]:

54 encoder_outputs: list[torch.Tensor] = []

55 for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(

57 - mm_kwargs, device=self.device, pin_memory=False

56 + mm_kwargs, device=self.device, pin_memory=True

57 ):

58 batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch)

59 sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)

vllm/v1/worker/gpu/model_runner.py

+2 -3

@@ -46,8 +46,7 @@ from vllm.sequence import IntermediateTensors

46 from vllm.tasks import SupportedTask

47 from vllm.utils.math_utils import cdiv

48 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib

49 -from vllm.utils.platform_utils import is_pin_memory_available

50 -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

49 +from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE

50 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

51 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

52 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

@@ -498,7 +497,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

497 """Build KV-block zeroing metadata; invoked from gpu_worker."""

498 self.kv_block_zeroer = KVBlockZeroer(

499 self.device,

501 - is_pin_memory_available(),

500 + pin_memory=PIN_MEMORY,

501 attn_groups_iter=(g for groups in self.attn_groups for g in groups),

502 kernel_block_sizes=self.kernel_block_sizes,

503 cache_dtype=self.cache_config.cache_dtype,

vllm/v1/worker/gpu/model_states/whisper.py

+3 -4

@@ -170,7 +170,8 @@ class WhisperModelState(ModelState):

170 for_capture: bool,

171 num_reqs: int,

172 ) -> dict[int, tuple[torch.Tensor, np.ndarray]]:

173 - encoder_seq_lens_np = np.zeros(num_reqs, dtype=np.int32)

173 + encoder_seq_lens = torch.zeros(num_reqs, dtype=torch.int32, pin_memory=True)

174 + encoder_seq_lens_np = encoder_seq_lens.numpy()

175 if not for_capture:

176 # During normal execution, use actual encoder lengths.

177 for i, req_id in enumerate(req_ids):

@@ -183,9 +184,7 @@ class WhisperModelState(ModelState):

184 # is captured with the correct value for cross-attention.

185 encoder_seq_lens_np[:] = self.max_encoder_len

186

186 - self.encoder_seq_lens_gpu[:num_reqs].copy_(

187 - torch.from_numpy(encoder_seq_lens_np), non_blocking=True

188 - )

187 + self.encoder_seq_lens_gpu[:num_reqs].copy_(encoder_seq_lens, non_blocking=True)

188 self.encoder_seq_lens_gpu[num_reqs:].fill_(0)

189 encoder_seq_lens_gpu = self.encoder_seq_lens_gpu[:num_reqs]

190

vllm/v1/worker/gpu_input_batch.py

+17 -16

@@ -15,6 +15,7 @@ from vllm.pooling_params import PoolingParams

15 from vllm.sampling_params import SamplingParams, SamplingType

16 from vllm.utils import length_from_prompt_token_ids_or_embeds

17 from vllm.utils.collection_utils import swap_dict_values

18 +from vllm.utils.torch_utils import PIN_MEMORY

19 from vllm.v1.outputs import LogprobsTensors

20 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

21 from vllm.v1.sample.logits_processor import (

@@ -95,7 +96,6 @@ class InputBatch:

96 max_model_len: int,

97 max_num_batched_tokens: int,

98 device: torch.device,

98 - pin_memory: bool,

99 vocab_size: int,

100 block_sizes: list[int], # The block_size of each kv cache group

101 kernel_block_sizes: list[int],

@@ -112,7 +112,6 @@ class InputBatch:

112 max_num_reqs,

113 num_spec_tokens,

114 device,

115 - pin_memory,

115 )

116 self.thinking_token_budget_reqs: set[str] = set()

117 self.is_pooling_model = is_pooling_model

@@ -120,7 +119,6 @@ class InputBatch:

119 self.max_model_len = max_model_len

120 self.max_num_batched_tokens = max_num_batched_tokens

121 self.device = device

123 - self.pin_memory = pin_memory

122 self.vocab_size = vocab_size

123

124 self._req_ids: list[str | None] = []

@@ -138,7 +136,10 @@ class InputBatch:

136 )

137 self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()

138 self.is_token_ids_tensor = torch.zeros(

141 - (max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False

139 + (max_num_reqs, max_model_len),

140 + device="cpu",

141 + dtype=bool,

142 + pin_memory=False,

143 )

144 self.is_token_ids = self.is_token_ids_tensor.numpy()

145 # Store prompt embeddings per request to avoid OOM from large upfront

@@ -149,21 +150,21 @@ class InputBatch:

150 (max_num_reqs,),

151 device="cpu",

152 dtype=torch.int32,

152 - pin_memory=pin_memory,

153 + pin_memory=PIN_MEMORY,

154 )

155 self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy()

156 self.num_prompt_tokens_cpu_tensor = torch.zeros(

157 (max_num_reqs,),

158 device="cpu",

159 dtype=torch.int32,

159 - pin_memory=pin_memory,

160 + pin_memory=PIN_MEMORY,

161 )

162 self.num_prompt_tokens = self.num_prompt_tokens_cpu_tensor.numpy()

163 self.num_computed_tokens_cpu_tensor = torch.zeros(

164 (max_num_reqs,),

165 device="cpu",

166 dtype=torch.int32,

166 - pin_memory=pin_memory,

167 + pin_memory=PIN_MEMORY,

168 )

169 self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy()

170

@@ -172,7 +173,7 @@ class InputBatch:

173 max_num_reqs=max_num_reqs,

174 max_model_len=max_model_len,

175 max_num_batched_tokens=max_num_batched_tokens,

175 - pin_memory=pin_memory,

176 + pin_memory=PIN_MEMORY,

177 device=device,

178 block_sizes=block_sizes,

179 kernel_block_sizes=kernel_block_sizes,

@@ -185,7 +186,7 @@ class InputBatch:

186 (max_num_reqs,), dtype=torch.float32, device=device

187 )

188 self.temperature_cpu_tensor = torch.empty(

188 - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory

189 + (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY

190 )

191 self.temperature_cpu = self.temperature_cpu_tensor.numpy()

192 self.greedy_reqs: set[str] = set()

@@ -193,14 +194,14 @@ class InputBatch:

194

195 self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device)

196 self.top_p_cpu_tensor = torch.empty(

196 - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory

197 + (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY

198 )

199 self.top_p_cpu = self.top_p_cpu_tensor.numpy()

200 self.top_p_reqs: set[str] = set()

201

202 self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device)

203 self.top_k_cpu_tensor = torch.empty(

203 - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory

204 + (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY

205 )

206 self.top_k_cpu = self.top_k_cpu_tensor.numpy()

207 self.top_k_reqs: set[str] = set()

@@ -210,7 +211,7 @@ class InputBatch:

211 (max_num_reqs,), dtype=torch.float, device=device

212 )

213 self.frequency_penalties_cpu_tensor = torch.empty(

213 - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory

214 + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY

215 )

216 self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy()

217 self.frequency_penalties_reqs: set[str] = set()

@@ -220,7 +221,7 @@ class InputBatch:

221 (max_num_reqs,), dtype=torch.float, device=device

222 )

223 self.presence_penalties_cpu_tensor = torch.empty(

223 - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory

224 + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY

225 )

226 self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy()

227 self.presence_penalties_reqs: set[str] = set()

@@ -230,14 +231,14 @@ class InputBatch:

231 (max_num_reqs,), dtype=torch.float, device=device

232 )

233 self.repetition_penalties_cpu_tensor = torch.empty(

233 - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory

234 + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY

235 )

236 self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy()

237 self.repetition_penalties_reqs: set[str] = set()

238

239 # Speculative decoding

240 self.num_accepted_tokens_cpu_tensor = torch.ones(

240 - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory

241 + (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY

242 )

243 self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy()

244

@@ -963,7 +964,7 @@ class InputBatch:

964 (self.num_reqs, max_prompt_len),

965 device="cpu",

966 dtype=torch.int64,

966 - pin_memory=self.pin_memory,

967 + pin_memory=PIN_MEMORY,

968 )

969 prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()

970 prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len]

vllm/v1/worker/gpu_model_runner.py

+38 -45

@@ -116,8 +116,10 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds

116 from vllm.utils.math_utils import cdiv, round_up

117 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib

118 from vllm.utils.nvtx_pytorch_hooks import PytHooks

119 -from vllm.utils.platform_utils import is_pin_memory_available, num_compute_units

119 +from vllm.utils.platform_utils import num_compute_units

120 from vllm.utils.torch_utils import (

121 + PIN_MEMORY,

122 + async_tensor_h2d,

123 get_dtype_size,

124 is_quantized_kv_cache,

125 kv_cache_dtype_str_to_dtype,

@@ -441,7 +443,6 @@ class GPUModelRunner(

443 scheduler_config = self.scheduler_config

444 parallel_config = self.parallel_config

445 self.device = device

444 - self.pin_memory = is_pin_memory_available()

446 self.dtype = self.model_config.dtype

447

448 self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(

@@ -666,7 +667,6 @@ class GPUModelRunner(

667 max_model_len=max(self.max_model_len, self.max_encoder_len),

668 max_num_batched_tokens=self.max_num_tokens,

669 device=self.device,

669 - pin_memory=self.pin_memory,

670 vocab_size=self.model_config.get_vocab_size(),

671 block_sizes=[placeholder_block_size],

672 kernel_block_sizes=[placeholder_block_size],

@@ -674,7 +674,7 @@ class GPUModelRunner(

674 logitsprocs=build_logitsprocs(

675 self.vllm_config,

676 self.device,

677 - self.pin_memory,

677 + PIN_MEMORY,

678 self.is_pooling_model,

679 custom_logitsprocs,

680 ),

@@ -729,7 +729,7 @@ class GPUModelRunner(

729 self.max_num_reqs, dtype=torch.int32, device=self.device

730 )

731 self.optimistic_seq_lens_cpu = torch.zeros(

732 - self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory

732 + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY

733 )

734 self.num_computed_tokens = torch.zeros(

735 self.max_num_reqs, dtype=torch.int32, device=self.device

@@ -846,7 +846,7 @@ class GPUModelRunner(

846 and self.speculative_config.use_ngram_gpu()

847 ):

848 self._num_valid_draft_tokens_cpu = torch.empty(

849 - self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory

849 + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY

850 )

851 self._num_valid_draft_tokens_event = torch.cuda.Event()

852 self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream()

@@ -857,7 +857,7 @@ class GPUModelRunner(

857 (self.max_num_reqs, 1),

858 dtype=torch.int64,

859 device="cpu",

860 - pin_memory=self.pin_memory,

860 + pin_memory=PIN_MEMORY,

861 )

862

863 # Pre-allocated tensor for copying valid sampled token counts to CPU,

@@ -879,7 +879,7 @@ class GPUModelRunner(

879 (self.max_num_reqs, self.num_spec_tokens),

880 dtype=torch.int64,

881 device="cpu",

882 - pin_memory=self.pin_memory,

882 + pin_memory=PIN_MEMORY,

883 )

884 if self.use_async_scheduling:

885 self.valid_sampled_token_count_event = torch.Event()

@@ -888,7 +888,7 @@ class GPUModelRunner(

888 self.max_num_reqs,

889 dtype=torch.int32,

890 device="cpu",

891 - pin_memory=self.pin_memory,

891 + pin_memory=PIN_MEMORY,

892 )

893

894 # Model weight offloader

@@ -1000,7 +1000,6 @@ class GPUModelRunner(

1000 *size,

1001 dtype=dtype,

1002 device=self.device,

1003 - pin_memory=self.pin_memory,

1003 with_numpy=numpy,

1004 )

1005

@@ -1055,7 +1054,7 @@ class GPUModelRunner(

1054 token_type_ids.append(ids)

1055

1056 token_type_ids_cpu = torch.empty(

1058 - sum(seq_lens_cpu), dtype=torch.int32, pin_memory=self.pin_memory

1057 + sum(seq_lens_cpu), dtype=torch.int32, pin_memory=PIN_MEMORY

1058 )

1059 torch.cat(token_type_ids, out=token_type_ids_cpu)

1060 model_kwargs["token_type_ids"] = token_type_ids_cpu.to(

@@ -1095,12 +1094,12 @@ class GPUModelRunner(

1094 """

1095 self._kv_block_zeroer = KVBlockZeroer(

1096 self.device,

1098 - self.pin_memory,

1097 + pin_memory=PIN_MEMORY,

1098 attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),

1099 kernel_block_sizes=self._kernel_block_sizes,

1100 cache_dtype=self.cache_config.cache_dtype,

1101 runner_only_attn_layers=self.runner_only_attn_layers,

1103 - static_forward_context=(self.compilation_config.static_forward_context),

1102 + static_forward_context=self.compilation_config.static_forward_context,

1103 )

1104

1105 def _zero_block_ids(self, block_ids: list[int]) -> None:

@@ -1651,7 +1650,7 @@ class GPUModelRunner(

1650 for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs(

1651 mm_kwargs,

1652 device=self.device,

1654 - pin_memory=self.pin_memory,

1653 + pin_memory=PIN_MEMORY,

1654 ):

1655 mm_kwargs_combined.update(mm_kwargs_batch)

1656

@@ -1807,10 +1806,10 @@ class GPUModelRunner(

1806 return

1807 # Upload the index tensors asynchronously so the scatter can be non-blocking.

1808 sampled_tokens_index_tensor = torch.tensor(

1810 - sample_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1809 + sample_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY

1810 ).to(self.device, non_blocking=True)

1811 prev_common_req_indices_tensor = torch.tensor(

1813 - prev_indices, dtype=torch.int64, pin_memory=self.pin_memory

1812 + prev_indices, dtype=torch.int64, pin_memory=PIN_MEMORY

1813 ).to(self.device, non_blocking=True)

1814 self.input_ids.gpu.scatter_(

1815 dim=0,

@@ -1826,10 +1825,10 @@ class GPUModelRunner(

1825

1826 assert isinstance(self._draft_token_ids, torch.Tensor)

1827 draft_tokens_index_tensor = torch.tensor(

1829 - spec_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1828 + spec_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY

1829 ).to(self.device, non_blocking=True)

1830 prev_draft_token_indices_tensor = torch.tensor(

1832 - prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory

1831 + prev_draft_token_indices, dtype=torch.int64, pin_memory=PIN_MEMORY

1832 ).to(self.device, non_blocking=True)

1833

1834 # because input_ids dtype is torch.int32,

@@ -2788,21 +2787,16 @@ class GPUModelRunner(

2787 # [0, 1, 2, 5, 6, 9]

2788 target_logits_indices += self._arange_scratch[: cu_num_draft_tokens[-1]]

2789

2791 - # TODO: Optimize the CPU -> GPU copy.

2792 - cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(

2793 - self.device, non_blocking=True

2794 - )

2795 - cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to(

2796 - self.device, non_blocking=True

2790 + cu_num_draft_tokens = async_tensor_h2d(cu_num_draft_tokens, device=self.device)

2791 + cu_num_sampled_tokens = async_tensor_h2d(

2792 + cu_num_sampled_tokens, device=self.device

2793 )

2798 - logits_indices = torch.from_numpy(logits_indices).to(

2799 - self.device, non_blocking=True

2794 + logits_indices = async_tensor_h2d(logits_indices, device=self.device)

2795 + target_logits_indices = async_tensor_h2d(

2796 + target_logits_indices, device=self.device

2797 )

2801 - target_logits_indices = torch.from_numpy(target_logits_indices).to(

2802 - self.device, non_blocking=True

2803 - )

2804 - bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(

2805 - self.device, non_blocking=True

2798 + bonus_logits_indices = async_tensor_h2d(

2799 + bonus_logits_indices, device=self.device

2800 )

2801

2802 # Compute the draft token ids.

@@ -3012,9 +3006,7 @@ class GPUModelRunner(

3006 # Track the current index in mm_kwargs/mm_lora_refs to map groups to request IDs

3007 current_item_idx = 0

3008 for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(

3015 - mm_kwargs,

3016 - device=self.device,

3017 - pin_memory=self.pin_memory,

3009 + mm_kwargs, device=self.device, pin_memory=PIN_MEMORY

3010 ):

3011 batch_outputs: MultiModalEmbeddings

3012

@@ -3048,7 +3040,7 @@ class GPUModelRunner(

3040 group_and_batch_mm_kwargs(

3041 [video_mm_kwargs_item],

3042 device=self.device,

3051 - pin_memory=self.pin_memory,

3043 + pin_memory=PIN_MEMORY,

3044 )

3045 )

3046

@@ -3107,7 +3099,10 @@ class GPUModelRunner(

3099

3100 mm_embeds = list[torch.Tensor]()

3101 is_mm_embed = torch.zeros(

3110 - total_num_scheduled_tokens, dtype=torch.bool, device="cpu"

3102 + total_num_scheduled_tokens,

3103 + dtype=torch.bool,

3104 + device="cpu",

3105 + pin_memory=PIN_MEMORY,

3106 )

3107

3108 req_start_idx = 0

@@ -3515,8 +3510,7 @@ class GPUModelRunner(

3510 token_ids_idx_np = np.nonzero(is_token_ids)[0]

3511 # Some tokens ids may need to become embeds

3512 if token_ids_idx_np.size > 0:

3518 - token_ids_idx = torch.from_numpy(token_ids_idx_np)

3519 - token_ids_idx = token_ids_idx.to(self.device, non_blocking=True)

3513 + token_ids_idx = async_tensor_h2d(token_ids_idx_np, device=self.device)

3514 token_ids = self.input_ids.gpu[token_ids_idx]

3515 tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)

3516 self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds

@@ -4953,7 +4947,7 @@ class GPUModelRunner(

4947 ):

4948 indices.append(offset + len(tokens) - 1)

4949 offset += num_draft + 1

4956 - indices = torch.tensor(indices, device=self.device)

4950 + indices = async_tensor_h2d(indices, device=self.device)

4951 hidden_states = sample_hidden_states[indices]

4952

4953 draft_token_ids = self.drafter.propose(

@@ -5483,8 +5477,8 @@ class GPUModelRunner(

5477 continue

5478

5479 num_prompt_tokens = len(request.prompt_token_ids)

5486 - prompt_token_ids = torch.tensor(request.prompt_token_ids).to(

5487 - self.device, non_blocking=True

5480 + prompt_token_ids = async_tensor_h2d(

5481 + request.prompt_token_ids, device=self.device

5482 )

5483

5484 # Set up target LogprobsTensors object.

@@ -5651,7 +5645,7 @@ class GPUModelRunner(

5645 for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs(

5646 [(modality, dummy_mm_item)] * max_items_per_batch,

5647 device=self.device,

5654 - pin_memory=self.pin_memory,

5648 + pin_memory=PIN_MEMORY,

5649 )

5650 )

5651

@@ -6995,7 +6989,6 @@ class GPUModelRunner(

6989 max_model_len=max_model_len,

6990 max_num_batched_tokens=self.max_num_tokens,

6991 device=self.device,

6998 - pin_memory=self.pin_memory,

6992 vocab_size=self.model_config.get_vocab_size(),

6993 block_sizes=block_sizes,

6994 kernel_block_sizes=kernel_block_sizes,

@@ -7430,7 +7423,7 @@ class GPUModelRunner(

7423 self.routed_experts_capturer.device_buffer.shape,

7424 dtype=self.routed_experts_capturer.device_buffer.dtype,

7425 device="cpu",

7433 - pin_memory=self.pin_memory,

7426 + pin_memory=PIN_MEMORY,

7427 )

7428 # ``slot_mapping`` dtype is fixed to int64 by

7429 # ``block_table.slot_mapping``; we mirror that here.

@@ -7439,7 +7432,7 @@ class GPUModelRunner(

7432 (max_tokens,),

7433 dtype=torch.int64,

7434 device="cpu",

7442 - pin_memory=self.pin_memory,

7435 + pin_memory=PIN_MEMORY,

7436 )

7437 # Private device buffer so the shared ``block_table.slot_mapping``

7438 # can be overwritten by the next ``_prepare_inputs`` while the