diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index fd7c6ff0bb..e7ec663020 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -1092,7 +1092,7 @@ def test_thinking_budget_holder_spec_mode_tensor_layout(): False, ) assert h.in_spec_mode - assert h.mask.shape[0] == 8 * (2 + 1) + assert h._mask_capacity == 8 * (2 + 1) def test_thinking_budget_holder_empty_end_tokens_disables_row(): diff --git a/vllm/lora/ops/triton_ops/utils.py b/vllm/lora/ops/triton_ops/utils.py index af0a3157f0..dd4c0d2097 100644 --- a/vllm/lora/ops/triton_ops/utils.py +++ b/vllm/lora/ops/triton_ops/utils.py @@ -13,6 +13,7 @@ from vllm import envs from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import async_tensor_h2d logger = init_logger(__name__) is_batch_invariant = envs.VLLM_BATCH_INVARIANT @@ -49,7 +50,9 @@ def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device): lora_strides_d1.append(lora_a_weight.stride(1)) lora_strides_d2.append(lora_a_weight.stride(2)) if len(lora_a_weights) > 1: - lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64) + lora_ptr_tensor = async_tensor_h2d( + tensor_ptrs, dtype=torch.uint64, device=device + ) else: lora_ptr_tensor = lora_a_weights[0] @@ -106,10 +109,11 @@ def _get_lora_b_ptr( hidden_sizes.append(lora_b_weight.size(1)) if len(lora_weights) > 1: - # note these are device tensors - lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64) - slice_start_tensor = torch.tensor( - slice_offset_lst, device=device, dtype=torch.uint64 + lora_ptr_tensor = async_tensor_h2d( + tensor_ptrs, dtype=torch.uint64, device=device + ) + slice_start_tensor = async_tensor_h2d( + slice_offset_lst, dtype=torch.uint64, device=device ) else: slice_start_tensor = slice_offset_lst[0] @@ -129,10 +133,18 @@ def _get_lora_b_ptr( same_stride = True else: - lora_strides_d0_tensor = torch.tensor(lora_strides_d0, device=device) - lora_strides_d1_tensor = torch.tensor(lora_strides_d1, device=device) - lora_strides_d2_tensor = torch.tensor(lora_strides_d2, device=device) - hidden_sizes_tensor = torch.tensor(hidden_sizes, device=device) + lora_strides_d0_tensor = async_tensor_h2d( + lora_strides_d0, dtype=torch.int64, device=device + ) + lora_strides_d1_tensor = async_tensor_h2d( + lora_strides_d1, dtype=torch.int64, device=device + ) + lora_strides_d2_tensor = async_tensor_h2d( + lora_strides_d2, dtype=torch.int64, device=device + ) + hidden_sizes_tensor = async_tensor_h2d( + hidden_sizes, dtype=torch.int64, device=device + ) same_stride = False # MAX_N is the maximum hidden size among all the lora_b weights MAX_N = max(hidden_sizes) diff --git a/vllm/lora/punica_wrapper/utils.py b/vllm/lora/punica_wrapper/utils.py index 584745f86b..8cf5f1a176 100644 --- a/vllm/lora/punica_wrapper/utils.py +++ b/vllm/lora/punica_wrapper/utils.py @@ -5,6 +5,8 @@ from typing import TYPE_CHECKING import torch +from vllm.utils.torch_utils import async_tensor_h2d + if TYPE_CHECKING: # avoid circuit import from vllm.lora.layers import LoRAMapping @@ -110,8 +112,8 @@ def convert_mapping( embedding_indices, ] - indices = torch.tensor(indices_list, dtype=torch.long, device=device) - prompt_mapping_tensor = torch.tensor( + indices = async_tensor_h2d(indices_list, dtype=torch.long, device=device) + prompt_mapping_tensor = async_tensor_h2d( prompt_mapping, dtype=torch.long, device=device ) embeddings_indices = torch.stack( diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 01854b96d5..c9ea736329 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -559,13 +559,10 @@ def _encode_token_type_ids( def _decode_token_type_ids(input_ids: torch.Tensor) -> torch.Tensor: - ids_mask = ( - torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) - << TOKEN_TYPE_SHIFT - ) - tokens_mask = ids_mask.bitwise_not() + ids_mask = 1 << TOKEN_TYPE_SHIFT + tokens_mask = ~ids_mask - token_type_ids = input_ids.bitwise_and(ids_mask) >> TOKEN_TYPE_SHIFT + token_type_ids = (input_ids & ids_mask) >> TOKEN_TYPE_SHIFT input_ids.bitwise_and_(tokens_mask) diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 0f059b6d13..6ecadbcd67 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -305,7 +305,7 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]): return dict( pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches), - num_patches=MultiModalFieldConfig.batched("image"), + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 85f422342a..20caa8672d 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -723,6 +723,7 @@ class Glm4vVisionTransformer(nn.Module): # Use pre-computed cos_sin_cache from RotaryEmbedding cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) + pos_ids = pos_ids.to(cos.device, non_blocking=True) cos_combined = cos[pos_ids].flatten(1) sin_combined = sin[pos_ids].flatten(1) return cos_combined, sin_combined, pos_ids diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 2b9a03be08..d2a1cb5755 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -143,7 +143,7 @@ class GraniteSpeechMultiModalProcessor( ) -> Mapping[str, MultiModalFieldConfig]: return dict( input_features=MultiModalFieldConfig.batched("audio"), - audio_embed_sizes=MultiModalFieldConfig.batched("audio"), + audio_embed_sizes=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( @@ -717,13 +717,13 @@ class GraniteSpeechForConditionalGeneration( torch.Tensor: Mask of shape (bsz, num_features) to be applied to the audio features prior to splitting the audio embeddings. """ - most_audio_features = torch.max(audio_embed_sizes).item() - mask_indices = torch.arange( - most_audio_features, - device=audio_embed_sizes.device, - ).view(1, -1) + most_audio_features = int(torch.max(audio_embed_sizes)) + mask_indices = torch.arange(most_audio_features).view(1, -1) input_features_mask = mask_indices < audio_embed_sizes.view(-1, 1) - return input_features_mask + target_device = self.encoder.input_linear.weight.device + if target_device == input_features_mask.device: + return input_features_mask + return input_features_mask.pin_memory().to(target_device, non_blocking=True) def _pad_and_stack_input_features( self, diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index d3ffdd4cf2..cf7d51c86c 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -355,7 +355,7 @@ class Idefics3MultiModalProcessor(BaseMultiModalProcessor[Idefics3ProcessingInfo "image", num_patches ), image_embeds=MultiModalFieldConfig.batched("image"), - num_patches=MultiModalFieldConfig.batched("image"), + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index c8611a4993..f3918e302b 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -247,9 +247,11 @@ class BaseInternVLMultiModalProcessor(BaseMultiModalProcessor[_I]): pixel_values_flat=MultiModalFieldConfig.flat_from_sizes( "image", image_num_patches ), - image_num_patches=MultiModalFieldConfig.batched("image"), + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), - image_token_id=MultiModalFieldConfig.shared("image", num_images), + image_token_id=MultiModalFieldConfig.shared( + "image", num_images, keep_on_cpu=True + ), ) def _get_mm_fields_config( @@ -475,7 +477,7 @@ class InternVLMultiModalProcessor( pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes( "video", video_num_patches ), - video_num_patches=MultiModalFieldConfig.batched("video"), + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True), video_token_id=MultiModalFieldConfig.shared("video", num_videos), ) diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py index 739c90a429..fdd1d41af5 100644 --- a/vllm/model_executor/models/llava_next.py +++ b/vllm/model_executor/models/llava_next.py @@ -214,7 +214,7 @@ class LlavaNextMultiModalProcessor( ) -> Mapping[str, MultiModalFieldConfig]: return dict( pixel_values=MultiModalFieldConfig.batched("image"), - image_sizes=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), ) diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py index 638d9ba9d8..1beec4207d 100644 --- a/vllm/model_executor/models/llava_onevision.py +++ b/vllm/model_executor/models/llava_onevision.py @@ -316,7 +316,7 @@ class LlavaOnevisionMultiModalProcessor( ) -> Mapping[str, MultiModalFieldConfig]: return dict( pixel_values=MultiModalFieldConfig.batched("image"), - image_sizes=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), pixel_values_videos=MultiModalFieldConfig.batched("video"), ) diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index c3b09ed590..cd97522705 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -586,7 +586,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module): seq_len, batch_size, self.chunk_size, self.left_chunk ) device = xs_pad.device - enc_streaming_mask = enc_streaming_mask.to(device) + enc_streaming_mask = enc_streaming_mask.contiguous().to( + device, non_blocking=True + ) xs_pad = xs_pad.to(device) input_tensor = xs_pad @@ -605,7 +607,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module): seq_len, batch_size, chunk_size_nc, left_chunk_nc ) if device.type != "cpu": - enc_streaming_mask_nc = enc_streaming_mask_nc.to(device) + enc_streaming_mask_nc = enc_streaming_mask_nc.contiguous().to( + device, non_blocking=True + ) if masks is not None: hs_mask_nc = masks & enc_streaming_mask_nc else: @@ -917,7 +921,9 @@ class ConformerEncoder(TransformerEncoderBase): enc_streaming_mask = self._streaming_mask( max_audio_length, batch_size, self.chunk_size, self.left_chunk ) - enc_streaming_mask = enc_streaming_mask.to(device) + enc_streaming_mask = enc_streaming_mask.contiguous().to( + device, non_blocking=True + ) if mask is None: return enc_streaming_mask diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 14f4c424bb..ec74cb0f25 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -278,15 +278,15 @@ def create_qwen2_5_omni_thinker_field_factory( image_embeds=MultiModalFieldConfig.flat_from_sizes( "image", image_embed_grid_sizes ), - image_grid_thw=MultiModalFieldConfig.batched("image"), + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_grid_sizes ), video_embeds=MultiModalFieldConfig.flat_from_sizes( "video", video_embed_grid_sizes ), - video_grid_thw=MultiModalFieldConfig.batched("video"), - second_per_grid_ts=MultiModalFieldConfig.batched("video"), + video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True), use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos), ) diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 54334c91bf..9e1f394f53 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -84,6 +84,7 @@ from vllm.multimodal.processing import PromptReplacement, PromptUpdate from vllm.sequence import IntermediateTensors from vllm.utils.platform_utils import is_pin_memory_available 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 @@ -679,6 +680,7 @@ class Qwen2_5_VisionTransformer(nn.Module): # Use pre-computed cos_sin_cache from RotaryEmbedding cos, sin = self.rotary_pos_emb.get_cos_sin(max_size) + pos_ids = pos_ids.to(cos.device, non_blocking=True) cos_combined = cos[pos_ids].flatten(1) sin_combined = sin[pos_ids].flatten(1) @@ -737,9 +739,10 @@ class Qwen2_5_VisionTransformer(nn.Module): window_index_thw, cu_seqlens_window_thw = self.get_window_index_thw(t, h, w) cos_thw, sin_thw = self.rotary_pos_emb_thw(t, h, w) - cos_thw = cos_thw[window_index_thw, :, :] + window_index_thw_dev = window_index_thw.to(cos_thw.device, non_blocking=True) + cos_thw = cos_thw[window_index_thw_dev, :, :] cos_thw = cos_thw.flatten(start_dim=0, end_dim=1) - sin_thw = sin_thw[window_index_thw, :, :] + sin_thw = sin_thw[window_index_thw_dev, :, :] sin_thw = sin_thw.flatten(start_dim=0, end_dim=1) cu_seqlens_thw = torch.repeat_interleave( @@ -1036,7 +1039,7 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor): ) -> Mapping[str, MultiModalFieldConfig]: return dict( **super()._get_mm_fields_config(hf_inputs, hf_processor_mm_kwargs), - second_per_grid_ts=MultiModalFieldConfig.batched("video"), + second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True), ) def _call_hf_processor( @@ -1361,7 +1364,9 @@ class Qwen2_5_VLForConditionalGeneration( grid_thw_list = grid_thw.tolist() image_embeds_out = [] for emb, size in zip(image_embeds_split, grid_thw_list): - positions = compute_mrope_for_media(size, merge_size).to(emb.device) + positions = compute_mrope_for_media(size, merge_size).to( + emb.device, non_blocking=True + ) emb = torch.cat([emb, positions], dim=1) image_embeds_out.append(emb) image_embeds_split = image_embeds_out @@ -1443,7 +1448,7 @@ class Qwen2_5_VLForConditionalGeneration( merge_size, tokens_per_second=tokens_per_second, video_second_per_grid=video_second_per_grid_t.item(), - ).to(emb.device) + ).to(emb.device, non_blocking=True) emb = emb[retention_mask] positions = positions[retention_mask] @@ -1488,8 +1493,8 @@ class Qwen2_5_VLForConditionalGeneration( else mrope_positions.device ) - # Tensors - input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long) + # Tensors. + input_ids_t = async_tensor_h2d(input_ids, dtype=torch.long, device=device) mm_embeddings_out = [mm[:, :-4] for mm in multimodal_embeddings] mm_embeddings_pos = [ diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index e7e8d74714..9692d9615b 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -439,9 +439,9 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports num_audios, max_audio_tokens, embed_dim = audio_features.shape audio_output_lengths = audio_output_lengths.unsqueeze(1) audio_features_mask = ( - torch.arange(max_audio_tokens) - .expand(num_audios, max_audio_tokens) - .to(audio_output_lengths.device) + torch.arange(max_audio_tokens, device=audio_output_lengths.device).expand( + num_audios, max_audio_tokens + ) < audio_output_lengths ) masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim) diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index b05a73fcff..d0c059cdaf 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -634,6 +634,7 @@ class Qwen2VisionTransformer(nn.Module): # Use pre-computed cos_sin_cache from RotaryEmbedding cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) + pos_ids = pos_ids.to(cos.device, non_blocking=True) cos_combined = cos[pos_ids].flatten(1) sin_combined = sin[pos_ids].flatten(1) return cos_combined, sin_combined diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 44e656ab82..6d5123effa 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -80,6 +80,7 @@ from vllm.multimodal.processing.processor import ( ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interfaces import ( @@ -502,9 +503,9 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): cu_chunk_lens.extend([window_aftercnn] * num_full_chunks) if remainder: cu_chunk_lens.append(remainder) - cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum( - -1, dtype=torch.int32 - ) + cu_seqlens = async_tensor_h2d( + cu_chunk_lens, dtype=torch.int32, device=aftercnn_lens.device + ).cumsum(-1, dtype=torch.int32) max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) @@ -861,6 +862,7 @@ class Qwen3Omni_VisionTransformer(nn.Module): # Use pre-computed cos_sin_cache from RotaryEmbedding cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) + pos_ids = pos_ids.to(cos.device, non_blocking=True) cos_combined = cos[pos_ids].flatten(1) sin_combined = sin[pos_ids].flatten(1) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 7df44dda59..6f10259d13 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -101,6 +101,7 @@ from vllm.utils.collection_utils import is_list_of from vllm.utils.math_utils import round_up from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers +from ...utils.torch_utils import async_tensor_h2d from .interfaces import ( MultiModalEmbeddings, SupportsEagle, @@ -2157,7 +2158,9 @@ class Qwen3VLForConditionalGeneration( grid_thw_list = grid_thw.tolist() image_embeds_out = [] for emb, size in zip(image_embeds_split, grid_thw_list): - positions = compute_mrope_for_media(size, merge_size).to(emb.device) + positions = compute_mrope_for_media(size, merge_size).to( + emb.device, non_blocking=True + ) positions = torch.cat( [ positions, @@ -2388,7 +2391,7 @@ class Qwen3VLForConditionalGeneration( input_tokens=unpruned_token_ids, mm_features=[mm_feature], )[0] - .to(device) + .to(device, non_blocking=True) .permute(1, 0) ) full_is_video_embed = unpruned_token_ids_tensor == embed_token_id @@ -2636,7 +2639,7 @@ class Qwen3VLForConditionalGeneration( ) # Tensors - input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long) + input_ids_t = async_tensor_h2d(input_ids, device=device, dtype=torch.long) mm_embeddings_out = [] mm_embeddings_pos = [] diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index ab6ba91d24..4d900b5dde 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -152,7 +152,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): # Keep these as batched, as they always have batch size as first dim mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") - mm_fields["num_image_patches"] = MultiModalFieldConfig.batched("image") + mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) return mm_fields def _get_hf_mm_data( diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 48ea18ee19..095d0e363d 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -30,10 +30,8 @@ from vllm.model_executor.models.interfaces import supports_any_eagle from vllm.multimodal import NestedTensors from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv -from vllm.utils.platform_utils import ( - is_pin_memory_available, -) from vllm.utils.torch_utils import ( + async_tensor_h2d, direct_register_custom_op, ) @@ -498,10 +496,9 @@ def isin_list( elements: torch.Tensor, test_elements_list: list[int], ) -> torch.Tensor: - test_elements = torch.tensor( - test_elements_list, - pin_memory=is_pin_memory_available(), - ).to(device=elements.device, non_blocking=True) + test_elements = async_tensor_h2d( + test_elements_list, dtype=torch.int64, device=elements.device + ) return torch.isin(elements, test_elements) @@ -772,14 +769,9 @@ def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int: return layer_index -def cast_overflow_tensors( - tensors: torch.Tensor, - offset: float = 1000, -) -> torch.Tensor: - if tensors.isinf().any() or tensors.isnan().any(): - clamp_value = torch.finfo(tensors.dtype).max - offset - tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value) - return tensors +def cast_overflow_tensors(tensors: torch.Tensor, offset: float = 1000) -> torch.Tensor: + clamp_value = torch.finfo(tensors.dtype).max - offset + return torch.clamp(tensors, min=-clamp_value, max=clamp_value) def fast_topk( diff --git a/vllm/multimodal/evs.py b/vllm/multimodal/evs.py index 62611c8971..8584c08299 100644 --- a/vllm/multimodal/evs.py +++ b/vllm/multimodal/evs.py @@ -85,7 +85,7 @@ def compute_retention_mask( topk_indices = order[:retain_num_tokens] retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool) - retention_mask[topk_indices] = True + retention_mask.index_fill_(0, topk_indices, True) retention_mask = retention_mask.reshape(dissimilarity.size()) mask = retention_mask.view(-1) # "T H W -> (T H W)" diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index 74599a1e8c..ca5e2b66e0 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any import torch +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, MoveDirectionality, @@ -65,22 +66,9 @@ class ThinkingBudgetStateHolder: self.cu_num_tokens: dict[int, int] = {} if self.num_spec_tokens > 0: - self.mask = torch.zeros( - max_num_reqs * (self.num_spec_tokens + 1), - dtype=torch.bool, - device=device, - ) - self.force_token_ids = torch.full( - (max_num_reqs * (self.num_spec_tokens + 1),), - -1, - dtype=torch.long, - device=device, - ) + self._mask_capacity = max_num_reqs * (self.num_spec_tokens + 1) else: - self.mask = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) - self.force_token_ids = torch.full( - (max_num_reqs,), -1, dtype=torch.long, device=device - ) + self._mask_capacity = max_num_reqs def has_tracked_requests(self) -> bool: """True when ``sync_batch`` has state for a ``thinking_token_budget`` row. @@ -454,7 +442,6 @@ class ThinkingBudgetStateHolder: predict_bonus_token: bool, spec_token_ids_for_layout: list[list[int]], ) -> torch.Tensor: - self.mask[:] = False cumulative_total = 0 self.cu_num_tokens.clear() @@ -474,6 +461,11 @@ class ThinkingBudgetStateHolder: else: cumulative_total += 1 + # Build the active index / forced-token lists entirely on CPU so we + # avoid per-iteration scalar sync writes to GPU tensors. + active_indices_cpu: list[int] = [] + force_tokens_cpu: list[int] = [] + for seq_idx in sorted(self._state.keys()): if seq_idx not in self.cu_num_tokens: continue @@ -502,9 +494,12 @@ class ThinkingBudgetStateHolder: for force_idx in force_index: if end_count < len(self.think_end_token_ids): mask_idx = self.cu_num_tokens[seq_idx] + force_idx - if mask_idx < len(self.mask) and mask_idx < logits.shape[0]: - self.mask[mask_idx] = True - self.force_token_ids[mask_idx] = ( + if ( + mask_idx < self._mask_capacity + and mask_idx < logits.shape[0] + ): + active_indices_cpu.append(mask_idx) + force_tokens_cpu.append( self.think_end_token_ids[end_count] ) if predict_bonus_token: @@ -514,15 +509,16 @@ class ThinkingBudgetStateHolder: else: state["bonus_token_forced"] = True - has_active_thinking = any( - state.get("in_end", False) for state in self._state.values() - ) - - if has_active_thinking: - active_indices = self.mask.nonzero(as_tuple=False).view(-1) - - if len(active_indices) > 0: - force_tokens = self.force_token_ids[active_indices] - logits[active_indices, force_tokens] = 1e9 + if active_indices_cpu: + device = logits.device + active_indices = async_tensor_h2d( + active_indices_cpu, dtype=torch.long, device=device + ) + force_tokens = async_tensor_h2d( + force_tokens_cpu, dtype=torch.long, device=device + ) + # Avoid CPU->GPU sync. + fill = logits.new_full((len(active_indices_cpu),), 1e9) + logits.index_put_((active_indices, force_tokens), fill) return logits diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 529bbc2d4a..c4871d90be 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1765,8 +1765,6 @@ class GPUModelRunner( self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0], non_blocking=True, ) - if self.enable_prompt_embeds: - self.is_token_ids.gpu[:num_common_tokens] = True return # Upload the index tensors asynchronously so the scatter can be non-blocking. sampled_tokens_index_tensor = torch.tensor( @@ -3408,13 +3406,12 @@ class GPUModelRunner( # If a batch only has token ids, then including the embedding layer # in the CUDA graph will be more performant (like in the else case # below). - token_ids_idx = ( - self.is_token_ids.gpu[:num_scheduled_tokens] - .nonzero(as_tuple=False) - .squeeze(1) - ) + is_token_ids = self.is_token_ids.np[:num_scheduled_tokens] + token_ids_idx_np = np.nonzero(is_token_ids)[0] # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: + if token_ids_idx_np.size > 0: + token_ids_idx = torch.from_numpy(token_ids_idx_np) + token_ids_idx = token_ids_idx.to(self.device, non_blocking=True) token_ids = self.input_ids.gpu[token_ids_idx] tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds @@ -4601,6 +4598,9 @@ class GPUModelRunner( # appending a placeholder (-1) token id. if (req_state := self.requests.get(req_id)) is not None: req_state.output_token_ids.append(-1) + pos = self.input_batch.num_tokens_no_spec[i] + self.input_batch.is_token_ids[i, pos] = True + self.input_batch.num_tokens_no_spec[i] = pos + 1 self.input_batch.prev_req_id_to_index = prev_req_id_to_index def take_draft_token_ids(self) -> DraftTokenIds | None: