diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 583c87eb5ba..3412b30814e 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/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) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index db0cab471b4..4381e2c8d1f 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/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) diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index f57ea285eb7..a2181d8e188 100644 --- a/tests/v1/logits_processors/utils.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index eff9bec8ee9..7b068adf8b6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/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( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index 58195d13560..5e2371dea6b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/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, ...] diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py index 6ddaaf08df0..3133ea62289 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py +++ b/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) diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py index fec2fa64097..94b6b4d3458 100644 --- a/vllm/model_executor/models/chameleon.py +++ b/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 diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 7762ddd3a3e..2f5688853d6 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/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 diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 19389d84463..456a38dc36c 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/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( diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index a5b759e7f7a..bb938167aec 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/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, diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index f15ec491af1..5be82632f47 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/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) diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index de7a075e2f7..6325c1bf5da 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/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 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d5b23a24383..6d637b326d3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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, )