[Perf] Avoid some more unnecessary GPU<->CPU syncs (#51458)

Signed-off-by: Nick Hill <[email protected]>

tests/basic_correctness/test_basic_correctness.py

@@ -307,7 +307,8 @@ def test_raise_on_logit_nans(
 
             def compute_nan_logits(self, *args, **kwargs):
                 logits = original_compute_logits(self, *args, **kwargs)
-                logits[0, 0] = float("nan")
+                # `logits[0, 0] = ...` copies the scalar H2D and blocks.
+                logits[0, 0].fill_(float("nan"))
                 return logits
 
             monkeypatch.setattr(model_cls, "compute_logits", compute_nan_logits)

tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -19,6 +19,7 @@ from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -77,7 +78,7 @@ def get_fake_sample_fn() -> SamplerOutput:
             first_token_id_index = num_computed_tokens + 1
         if spec_decode_metadata is None:
             return SamplerOutput(
-                sampled_token_ids=torch.tensor(
+                sampled_token_ids=async_tensor_h2d(
                     [[prompt_token_ids[first_token_id_index]]],
                     device=DEVICE_TYPE,
                     dtype=torch.int32,
@@ -90,7 +91,7 @@ def get_fake_sample_fn() -> SamplerOutput:
         ]
         sampled_token_ids = accepted_tokens
         return SamplerOutput(
-            sampled_token_ids=torch.tensor(
+            sampled_token_ids=async_tensor_h2d(
                 [sampled_token_ids],
                 device=DEVICE_TYPE,
                 dtype=torch.int32,
@@ -132,7 +133,7 @@ def get_fake_propose_draft_token_ids_fn():
             ]
         ]
 
-        next_token_ids = torch.tensor(
+        next_token_ids = async_tensor_h2d(
             prompt_token_ids[
                 first_token_id_index - 1 : first_token_id_index
                 - 1
@@ -142,7 +143,7 @@ def get_fake_propose_draft_token_ids_fn():
             dtype=torch.int32,
         )
 
-        valid_sampled_tokens_count = torch.tensor(
+        valid_sampled_tokens_count = async_tensor_h2d(
             [num_accepted_tokens],
             device=DEVICE_TYPE,
             dtype=torch.int32,
@@ -150,7 +151,7 @@ def get_fake_propose_draft_token_ids_fn():
 
         self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
 
-        return torch.tensor(
+        return async_tensor_h2d(
             proposed_draft_token_ids,
             device=DEVICE_TYPE,
             dtype=torch.int32,
@@ -1088,11 +1089,10 @@ def _run_mamba_prefix_cache_mrv2(
             device=hidden_states.device,
             dtype=torch.int64,
         )
-        num_logits = torch.tensor(
+        num_logits = async_tensor_h2d(
             input_batch.cu_num_logits_np[1 : num_reqs + 1]
             - input_batch.cu_num_logits_np[:num_reqs],
             device=hidden_states.device,
-            dtype=torch.int32,
         )
         accepted = torch.full_like(num_logits, num_accepted_tokens)
         num_sampled = torch.minimum(accepted, num_logits)

tests/v1/logits_processors/utils.py

@@ -14,6 +14,7 @@ from vllm.config import VllmConfig
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.sample.logits_processor import (
     LOGITSPROCS_GROUP,
     AdapterLogitsProcessor,
@@ -90,17 +91,24 @@ class DummyLogitsProcessor(LogitsProcessor):
         if not self.req_info:
             return logits
 
-        # Save target values before modification
-        cols = torch.tensor(
+        # Save target values before modification.
+        cols = async_tensor_h2d(
             list(self.req_info.values()), dtype=torch.long, device=logits.device
         )
-        rows = torch.tensor(
+        rows = async_tensor_h2d(
             list(self.req_info.keys()), dtype=torch.long, device=logits.device
         )
         values_to_keep = logits[rows, cols].clone()
 
-        # Mask all but target tokens
-        logits[rows] = float("-inf")
+        # Mask all but target tokens. Use an on-device fill tensor so the
+        # scatter doesn't force a synchronizing scalar H2D.
+        fill = torch.full(
+            (rows.numel(), logits.size(-1)),
+            float("-inf"),
+            dtype=logits.dtype,
+            device=logits.device,
+        )
+        logits[rows] = fill
         logits[rows, cols] = values_to_keep
 
         return logits
@@ -141,7 +149,7 @@ class DummyPerReqLogitsProcessor:
         output_ids: list[int],
         logits: torch.Tensor,
     ) -> torch.Tensor:
-        val_to_keep = logits[self.target_token].item()
+        val_to_keep = logits[self.target_token].clone()
         logits[:] = float("-inf")
         logits[self.target_token] = val_to_keep
         return logits

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
@@ -179,9 +180,13 @@ def _make_src_and_dst_indices(
     src_device: torch.device | str,
     dst_device: torch.device | str,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    src_indices = torch.tensor(src_block_ids, device=src_device, dtype=torch.int64)
-    dst_indices = torch.tensor(dst_block_ids, device=dst_device, dtype=torch.int64)
-    return src_indices, dst_indices
+    def _to(block_ids: list[int], device: torch.device | str) -> torch.Tensor:
+        device = torch.device(device) if isinstance(device, str) else device
+        if device.type == "cpu":
+            return torch.tensor(block_ids, dtype=torch.int64, device=device)
+        return async_tensor_h2d(block_ids, dtype=torch.int64, device=device)
+
+    return _to(src_block_ids, src_device), _to(dst_block_ids, dst_device)
 
 
 def copy_kv_blocks(

vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py

@@ -135,6 +135,7 @@ class ExampleConnector(KVConnectorBase_V1):
                 slot_mapping (torch.Tensor): the slot mapping. In shape
                     [num_tokens].
             """
+            slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True)
             if isinstance(attn_metadata, MLACommonMetadata):
                 dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
                 num_pages = dst_kv_cache_layer_shape[0]
@@ -178,9 +179,8 @@ class ExampleConnector(KVConnectorBase_V1):
                 filename = self._generate_filename_debug(
                     layer_name, request.token_ids, request.mm_hashes
                 )
-                kv_cache = safetensors.torch.load_file(
-                    filename, device=str(kv_cache_layer.device)
-                )["kv_cache"]
+                kv_cache_cpu = safetensors.torch.load_file(filename)["kv_cache"]
+                kv_cache = kv_cache_cpu.to("cuda", non_blocking=True)
                 if isinstance(attn_metadata, dict):
                     inject_kv_into_layer(
                         kv_cache_layer,
@@ -227,6 +227,7 @@ class ExampleConnector(KVConnectorBase_V1):
             Assume the shape of the layer is (num_pages, page_size, xxx)
             for MLA, and (num_pages, 2, page_size, xxx) otherwise.
             """
+            slot_mapping = slot_mapping.to(layer.device, non_blocking=True)
             if isinstance(attn_metadata, MLACommonMetadata):
                 num_pages, page_size = layer.shape[0], layer.shape[1]
                 return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -11,7 +11,7 @@ from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.mem_utils import get_max_shared_memory_bytes
-from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.utils.torch_utils import async_tensor_h2d, direct_register_custom_op
 
 from .utils import supports_pdl, supports_tma
 
@@ -866,7 +866,7 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
     tensor_ptrs = []
     for lora_weight in lora_weights:
         tensor_ptrs.append(lora_weight.data_ptr())
-    ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)
+    ptr_tensor = async_tensor_h2d(tensor_ptrs, dtype=torch.uint64, device=device)
 
     _LORA_PTR_DICT[key] = ptr_tensor
     return _LORA_PTR_DICT.get(key)

vllm/model_executor/models/chameleon.py

@@ -55,6 +55,7 @@ from vllm.multimodal.processing import (
 )
 from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
+from vllm.utils.torch_utils import async_tensor_h2d
 
 from .interfaces import (
     MultiModalEmbeddings,
@@ -819,8 +820,20 @@ class ChameleonImageVocabularyMapping:
 
     def convert_img2bpe(self, img_batch: torch.Tensor) -> torch.Tensor:
         device = img_batch.device
-        img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]
-        return img_tokens.to(device)
+        # Cache a per-device copy of the (small, static) mapping tensor so we
+        # can index entirely on `device` instead of forcing a D2H on
+        # `img_batch` and an H2D on the result.
+        cache = getattr(self, "_img2bpe_mapping_cache", None)
+        if cache is None:
+            cache = {}
+            self._img2bpe_mapping_cache = cache
+        mapping_on_device = cache.get(device)
+        if mapping_on_device is None:
+            mapping_on_device = async_tensor_h2d(
+                self.img2bpe_mapping_tensor, device=device
+            )
+            cache[device] = mapping_on_device
+        return mapping_on_device[img_batch]
 
 
 class ChameleonModel(nn.Module):
@@ -1025,8 +1038,23 @@ class ChameleonForConditionalGeneration(
         # Disallow image tokens which does not include special
         # begin-image and end-image tokens
         if logits is not None:
-            image_tokens = self.model.vocabulary_mapping.image_tokens
-            logits[:, image_tokens] = torch.finfo(logits.dtype).min
+            # Cache a per-device index tensor for the (static) image-token
+            # set, and use `index_fill_` instead of advanced-index assign
+            # so the scatter runs entirely on device (no host roundtrip
+            # for the scalar fill value or for the Python-list indices).
+            cache = getattr(self, "_image_tokens_index_cache", None)
+            if cache is None:
+                cache = {}
+                self._image_tokens_index_cache = cache
+            image_tokens_idx = cache.get(logits.device)
+            if image_tokens_idx is None:
+                image_tokens_idx = async_tensor_h2d(
+                    self.model.vocabulary_mapping.image_tokens,
+                    dtype=torch.long,
+                    device=logits.device,
+                )
+                cache[logits.device] = image_tokens_idx
+            logits.index_fill_(1, image_tokens_idx, torch.finfo(logits.dtype).min)
 
         return logits
 

vllm/model_executor/models/gemma3n_mm.py

@@ -54,6 +54,7 @@ from vllm.multimodal.processing.processor import (
 )
 from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
+from vllm.utils.torch_utils import async_tensor_h2d
 
 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription
 from .utils import (
@@ -683,10 +684,18 @@ class Gemma3nForConditionalGeneration(
         # We handle both cases:
         # - If fewer tokens: pad with the embedding of the last vocab token
         # - If more tokens: truncate to the expected count
-        # TODO precompute and cache padding
-        audio_padding_toks = torch.tensor(
-            [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device
-        )
+        # Cache the single-scalar padding-token tensor per-device to avoid a
+        # synchronous H2D tensor construction on every forward.
+        cache = getattr(self, "_audio_padding_toks_cache", None)
+        if cache is None:
+            cache = {}
+            self._audio_padding_toks_cache = cache
+        audio_padding_toks = cache.get(audio_features.device)
+        if audio_padding_toks is None:
+            audio_padding_toks = async_tensor_h2d(
+                [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device
+            )
+            cache[audio_features.device] = audio_padding_toks
         audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
         audio_features = torch.where(
             audio_mask.unsqueeze(-1), audio_padding_embs, audio_features

vllm/model_executor/models/glm4_1v.py

@@ -99,6 +99,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.processor import get_processor_cls_name_from_config
 from vllm.transformers_utils.utils import convert_model_repo_to_path
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
 
@@ -535,7 +536,8 @@ class Glm4vVisionEmbeddings(nn.Module):
         total_seq = h_coords.shape[0]
         device = pos_embed_weight.device
 
-        # Move coordinates to correct device
+        # Move coordinates to correct device. The only caller already hands
+        # these over pinned + non-blocking, so this is normally a no-op.
         h_coords, w_coords = h_coords.to(device), w_coords.to(device)
 
         # Handle empty sequence case
@@ -816,10 +818,15 @@ class Glm4vVisionTransformer(nn.Module):
             )
 
             lengths = [h * w] * t
-            image_shapes = torch.tensor([[t, h, w]], device=device)
+            image_shapes = async_tensor_h2d(
+                [[t, h, w]], dtype=torch.long, device=device
+            )
 
-            h_coords_repeated = h_coords.repeat(t)
-            w_coords_repeated = w_coords.repeat(t)
+            # Build the coordinates on the host (cheap integer math) but move
+            # them across pinned + non-blocking, so the consumer's
+            # `.to(device)` below is a no-op rather than a blocking H2D.
+            h_coords_repeated = async_tensor_h2d(h_coords.repeat(t), device=device)
+            w_coords_repeated = async_tensor_h2d(w_coords.repeat(t), device=device)
 
             embeds = self.embeddings(
                 embeddings=torch.zeros(

vllm/model_executor/models/qwen3_omni_moe_thinker.py

@@ -437,7 +437,7 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):
         # Compute chunk information
         chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
 
-        chunk_lengths = torch.tensor(
+        chunk_lengths = async_tensor_h2d(
             [self.n_window * 2] * chunk_num.sum(),
             dtype=torch.long,
             device=feature_lens.device,

vllm/model_executor/models/voxtral.py

@@ -773,15 +773,11 @@ class VoxtralEncoderModel(nn.Module):
         if global_log_mel_max := self.config.global_log_mel_max:
             if not isinstance(global_log_mel_max, float):
                 raise TypeError(f"{global_log_mel_max=} needs to be of type float.")
-            log_spec_max = torch.tensor(
-                global_log_mel_max,
-                device=log_spec.device,
-                dtype=log_spec.dtype,
-            )
+            # Use `clamp` to avoid gpu<->cpu sync.
+            log_spec = torch.clamp(log_spec, min=global_log_mel_max - 8.0)
         else:
             log_spec_max = log_spec.max()
-
-        log_spec = torch.maximum(log_spec, log_spec_max - 8.0)
+            log_spec = torch.maximum(log_spec, log_spec_max - 8.0)
         log_spec = (log_spec + 4.0) / 4.0
         return log_spec.to(input_dtype)
 

vllm/v1/spec_decode/extract_hidden_states.py

@@ -13,7 +13,6 @@ from vllm.distributed.eplb.eplb_state import EplbState
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.model_loader import get_model
-from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 from vllm.v1.utils import CpuGpuBuffer
@@ -59,11 +58,7 @@ class ExtractHiddenStatesProposer:
         )
 
         self.backup_next_token_ids = CpuGpuBuffer(
-            max_batch_size,
-            dtype=torch.int32,
-            pin_memory=PIN_MEMORY,
-            device=device,
-            with_numpy=True,
+            max_batch_size, dtype=torch.int32, device=device
         )
 
         self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config

vllm/v1/worker/gpu_model_runner.py

@@ -3069,7 +3069,7 @@ class GPUModelRunner(
 
                 self._cache_encoder_output(
                     mm_hashes[i],
-                    pe_tensor.to(self.device),
+                    async_tensor_h2d(pe_tensor, device=self.device),
                     scheduler_output.ec_manager_metadata,
                     scheduler_output.free_encoder_mm_hashes,
                 )