b82e908b4c · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/logits_processors/test_correctness.py

+1 -1

@@ -1092,7 +1092,7 @@ def test_thinking_budget_holder_spec_mode_tensor_layout():

1092 False,

1093 )

1094 assert h.in_spec_mode

1095 - assert h.mask.shape[0] == 8 * (2 + 1)

1095 + assert h._mask_capacity == 8 * (2 + 1)

1096

1097

1098 def test_thinking_budget_holder_empty_end_tokens_disables_row():

vllm/lora/ops/triton_ops/utils.py

+21 -9

@@ -13,6 +13,7 @@ from vllm import envs

13 from vllm.logger import init_logger

14 from vllm.platforms import current_platform

15 from vllm.utils.math_utils import next_power_of_2

16 +from vllm.utils.torch_utils import async_tensor_h2d

17

18 logger = init_logger(__name__)

19 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):

50 lora_strides_d1.append(lora_a_weight.stride(1))

51 lora_strides_d2.append(lora_a_weight.stride(2))

52 if len(lora_a_weights) > 1:

52 - lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)

53 + lora_ptr_tensor = async_tensor_h2d(

54 + tensor_ptrs, dtype=torch.uint64, device=device

55 + )

56 else:

57 lora_ptr_tensor = lora_a_weights[0]

58

@@ -106,10 +109,11 @@ def _get_lora_b_ptr(

109 hidden_sizes.append(lora_b_weight.size(1))

110

111 if len(lora_weights) > 1:

109 - # note these are device tensors

110 - lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)

111 - slice_start_tensor = torch.tensor(

112 - slice_offset_lst, device=device, dtype=torch.uint64

112 + lora_ptr_tensor = async_tensor_h2d(

113 + tensor_ptrs, dtype=torch.uint64, device=device

114 + )

115 + slice_start_tensor = async_tensor_h2d(

116 + slice_offset_lst, dtype=torch.uint64, device=device

117 )

118 else:

119 slice_start_tensor = slice_offset_lst[0]

@@ -129,10 +133,18 @@ def _get_lora_b_ptr(

133 same_stride = True

134

135 else:

132 - lora_strides_d0_tensor = torch.tensor(lora_strides_d0, device=device)

133 - lora_strides_d1_tensor = torch.tensor(lora_strides_d1, device=device)

134 - lora_strides_d2_tensor = torch.tensor(lora_strides_d2, device=device)

135 - hidden_sizes_tensor = torch.tensor(hidden_sizes, device=device)

136 + lora_strides_d0_tensor = async_tensor_h2d(

137 + lora_strides_d0, dtype=torch.int64, device=device

138 + )

139 + lora_strides_d1_tensor = async_tensor_h2d(

140 + lora_strides_d1, dtype=torch.int64, device=device

141 + )

142 + lora_strides_d2_tensor = async_tensor_h2d(

143 + lora_strides_d2, dtype=torch.int64, device=device

144 + )

145 + hidden_sizes_tensor = async_tensor_h2d(

146 + hidden_sizes, dtype=torch.int64, device=device

147 + )

148 same_stride = False

149 # MAX_N is the maximum hidden size among all the lora_b weights

150 MAX_N = max(hidden_sizes)

vllm/lora/punica_wrapper/utils.py

+4 -2

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING

5

6 import torch

7

8 +from vllm.utils.torch_utils import async_tensor_h2d

9 +

10 if TYPE_CHECKING:

11 # avoid circuit import

12 from vllm.lora.layers import LoRAMapping

@@ -110,8 +112,8 @@ def convert_mapping(

112 embedding_indices,

113 ]

114

113 - indices = torch.tensor(indices_list, dtype=torch.long, device=device)

114 - prompt_mapping_tensor = torch.tensor(

115 + indices = async_tensor_h2d(indices_list, dtype=torch.long, device=device)

116 + prompt_mapping_tensor = async_tensor_h2d(

117 prompt_mapping, dtype=torch.long, device=device

118 )

119 embeddings_indices = torch.stack(

vllm/model_executor/models/bert.py

+3 -6

@@ -559,13 +559,10 @@ def _encode_token_type_ids(

559

560

561 def _decode_token_type_ids(input_ids: torch.Tensor) -> torch.Tensor:

562 - ids_mask = (

563 - torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)

564 - << TOKEN_TYPE_SHIFT

565 - )

566 - tokens_mask = ids_mask.bitwise_not()

562 + ids_mask = 1 << TOKEN_TYPE_SHIFT

563 + tokens_mask = ~ids_mask

564

568 - token_type_ids = input_ids.bitwise_and(ids_mask) >> TOKEN_TYPE_SHIFT

565 + token_type_ids = (input_ids & ids_mask) >> TOKEN_TYPE_SHIFT

566

567 input_ids.bitwise_and_(tokens_mask)

568

vllm/model_executor/models/gemma3_mm.py

+1 -1

@@ -305,7 +305,7 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):

305

306 return dict(

307 pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),

308 - num_patches=MultiModalFieldConfig.batched("image"),

308 + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

309 )

310

311 def _get_prompt_updates(

vllm/model_executor/models/glm4_1v.py

+1

@@ -723,6 +723,7 @@ class Glm4vVisionTransformer(nn.Module):

723 # Use pre-computed cos_sin_cache from RotaryEmbedding

724 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

725

726 + pos_ids = pos_ids.to(cos.device, non_blocking=True)

727 cos_combined = cos[pos_ids].flatten(1)

728 sin_combined = sin[pos_ids].flatten(1)

729 return cos_combined, sin_combined, pos_ids

vllm/model_executor/models/granite_speech.py

+7 -7

@@ -143,7 +143,7 @@ class GraniteSpeechMultiModalProcessor(

143 ) -> Mapping[str, MultiModalFieldConfig]:

144 return dict(

145 input_features=MultiModalFieldConfig.batched("audio"),

146 - audio_embed_sizes=MultiModalFieldConfig.batched("audio"),

146 + audio_embed_sizes=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),

147 )

148

149 def _get_prompt_updates(

@@ -717,13 +717,13 @@ class GraniteSpeechForConditionalGeneration(

717 torch.Tensor: Mask of shape (bsz, num_features) to be applied to

718 the audio features prior to splitting the audio embeddings.

719 """

720 - most_audio_features = torch.max(audio_embed_sizes).item()

721 - mask_indices = torch.arange(

722 - most_audio_features,

723 - device=audio_embed_sizes.device,

724 - ).view(1, -1)

720 + most_audio_features = int(torch.max(audio_embed_sizes))

721 + mask_indices = torch.arange(most_audio_features).view(1, -1)

722 input_features_mask = mask_indices < audio_embed_sizes.view(-1, 1)

726 - return input_features_mask

723 + target_device = self.encoder.input_linear.weight.device

724 + if target_device == input_features_mask.device:

725 + return input_features_mask

726 + return input_features_mask.pin_memory().to(target_device, non_blocking=True)

727

728 def _pad_and_stack_input_features(

729 self,

vllm/model_executor/models/idefics3.py

+1 -1

@@ -355,7 +355,7 @@ class Idefics3MultiModalProcessor(BaseMultiModalProcessor[Idefics3ProcessingInfo

355 "image", num_patches

356 ),

357 image_embeds=MultiModalFieldConfig.batched("image"),

358 - num_patches=MultiModalFieldConfig.batched("image"),

358 + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

359 )

360

361 def _get_prompt_updates(

vllm/model_executor/models/internvl.py

+5 -3

@@ -247,9 +247,11 @@ class BaseInternVLMultiModalProcessor(BaseMultiModalProcessor[_I]):

247 pixel_values_flat=MultiModalFieldConfig.flat_from_sizes(

248 "image", image_num_patches

249 ),

250 - image_num_patches=MultiModalFieldConfig.batched("image"),

250 + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

251 image_embeds=MultiModalFieldConfig.batched("image"),

252 - image_token_id=MultiModalFieldConfig.shared("image", num_images),

252 + image_token_id=MultiModalFieldConfig.shared(

253 + "image", num_images, keep_on_cpu=True

254 + ),

255 )

256

257 def _get_mm_fields_config(

@@ -475,7 +477,7 @@ class InternVLMultiModalProcessor(

477 pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes(

478 "video", video_num_patches

479 ),

478 - video_num_patches=MultiModalFieldConfig.batched("video"),

480 + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

481 video_token_id=MultiModalFieldConfig.shared("video", num_videos),

482 )

483

vllm/model_executor/models/llava_next.py

+1 -1

@@ -214,7 +214,7 @@ class LlavaNextMultiModalProcessor(

214 ) -> Mapping[str, MultiModalFieldConfig]:

215 return dict(

216 pixel_values=MultiModalFieldConfig.batched("image"),

217 - image_sizes=MultiModalFieldConfig.batched("image"),

217 + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

218 image_embeds=MultiModalFieldConfig.batched("image"),

219 )

220

vllm/model_executor/models/llava_onevision.py

+1 -1

@@ -316,7 +316,7 @@ class LlavaOnevisionMultiModalProcessor(

316 ) -> Mapping[str, MultiModalFieldConfig]:

317 return dict(

318 pixel_values=MultiModalFieldConfig.batched("image"),

319 - image_sizes=MultiModalFieldConfig.batched("image"),

319 + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

320 image_embeds=MultiModalFieldConfig.batched("image"),

321 pixel_values_videos=MultiModalFieldConfig.batched("video"),

322 )

vllm/model_executor/models/phi4mm_audio.py

+9 -3

@@ -586,7 +586,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

586 seq_len, batch_size, self.chunk_size, self.left_chunk

587 )

588 device = xs_pad.device

589 - enc_streaming_mask = enc_streaming_mask.to(device)

589 + enc_streaming_mask = enc_streaming_mask.contiguous().to(

590 + device, non_blocking=True

591 + )

592 xs_pad = xs_pad.to(device)

593

594 input_tensor = xs_pad

@@ -605,7 +607,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

607 seq_len, batch_size, chunk_size_nc, left_chunk_nc

608 )

609 if device.type != "cpu":

608 - enc_streaming_mask_nc = enc_streaming_mask_nc.to(device)

610 + enc_streaming_mask_nc = enc_streaming_mask_nc.contiguous().to(

611 + device, non_blocking=True

612 + )

613 if masks is not None:

614 hs_mask_nc = masks & enc_streaming_mask_nc

615 else:

@@ -917,7 +921,9 @@ class ConformerEncoder(TransformerEncoderBase):

921 enc_streaming_mask = self._streaming_mask(

922 max_audio_length, batch_size, self.chunk_size, self.left_chunk

923 )

920 - enc_streaming_mask = enc_streaming_mask.to(device)

924 + enc_streaming_mask = enc_streaming_mask.contiguous().to(

925 + device, non_blocking=True

926 + )

927 if mask is None:

928 return enc_streaming_mask

929

vllm/model_executor/models/qwen2_5_omni_thinker.py

+3 -3

@@ -278,15 +278,15 @@ def create_qwen2_5_omni_thinker_field_factory(

278 image_embeds=MultiModalFieldConfig.flat_from_sizes(

279 "image", image_embed_grid_sizes

280 ),

281 - image_grid_thw=MultiModalFieldConfig.batched("image"),

281 + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

282 pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

283 "video", video_grid_sizes

284 ),

285 video_embeds=MultiModalFieldConfig.flat_from_sizes(

286 "video", video_embed_grid_sizes

287 ),

288 - video_grid_thw=MultiModalFieldConfig.batched("video"),

289 - second_per_grid_ts=MultiModalFieldConfig.batched("video"),

288 + video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

289 + second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

290 use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos),

291 )

292

vllm/model_executor/models/qwen2_5_vl.py

+12 -7

@@ -84,6 +84,7 @@ from vllm.multimodal.processing import PromptReplacement, PromptUpdate

84 from vllm.sequence import IntermediateTensors

85 from vllm.utils.platform_utils import is_pin_memory_available

86 from vllm.utils.tensor_schema import TensorSchema, TensorShape

87 +from vllm.utils.torch_utils import async_tensor_h2d

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

89 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

90

@@ -679,6 +680,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

680 # Use pre-computed cos_sin_cache from RotaryEmbedding

681 cos, sin = self.rotary_pos_emb.get_cos_sin(max_size)

682

683 + pos_ids = pos_ids.to(cos.device, non_blocking=True)

684 cos_combined = cos[pos_ids].flatten(1)

685 sin_combined = sin[pos_ids].flatten(1)

686

@@ -737,9 +739,10 @@ class Qwen2_5_VisionTransformer(nn.Module):

739 window_index_thw, cu_seqlens_window_thw = self.get_window_index_thw(t, h, w)

740 cos_thw, sin_thw = self.rotary_pos_emb_thw(t, h, w)

741

740 - cos_thw = cos_thw[window_index_thw, :, :]

742 + window_index_thw_dev = window_index_thw.to(cos_thw.device, non_blocking=True)

743 + cos_thw = cos_thw[window_index_thw_dev, :, :]

744 cos_thw = cos_thw.flatten(start_dim=0, end_dim=1)

742 - sin_thw = sin_thw[window_index_thw, :, :]

745 + sin_thw = sin_thw[window_index_thw_dev, :, :]

746 sin_thw = sin_thw.flatten(start_dim=0, end_dim=1)

747

748 cu_seqlens_thw = torch.repeat_interleave(

@@ -1036,7 +1039,7 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):

1039 ) -> Mapping[str, MultiModalFieldConfig]:

1040 return dict(

1041 **super()._get_mm_fields_config(hf_inputs, hf_processor_mm_kwargs),

1039 - second_per_grid_ts=MultiModalFieldConfig.batched("video"),

1042 + second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

1043 )

1044

1045 def _call_hf_processor(

@@ -1361,7 +1364,9 @@ class Qwen2_5_VLForConditionalGeneration(

1364 grid_thw_list = grid_thw.tolist()

1365 image_embeds_out = []

1366 for emb, size in zip(image_embeds_split, grid_thw_list):

1364 - positions = compute_mrope_for_media(size, merge_size).to(emb.device)

1367 + positions = compute_mrope_for_media(size, merge_size).to(

1368 + emb.device, non_blocking=True

1369 + )

1370 emb = torch.cat([emb, positions], dim=1)

1371 image_embeds_out.append(emb)

1372 image_embeds_split = image_embeds_out

@@ -1443,7 +1448,7 @@ class Qwen2_5_VLForConditionalGeneration(

1448 merge_size,

1449 tokens_per_second=tokens_per_second,

1450 video_second_per_grid=video_second_per_grid_t.item(),

1446 - ).to(emb.device)

1451 + ).to(emb.device, non_blocking=True)

1452

1453 emb = emb[retention_mask]

1454 positions = positions[retention_mask]

@@ -1488,8 +1493,8 @@ class Qwen2_5_VLForConditionalGeneration(

1493 else mrope_positions.device

1494 )

1495

1491 - # Tensors

1492 - input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long)

1496 + # Tensors.

1497 + input_ids_t = async_tensor_h2d(input_ids, dtype=torch.long, device=device)

1498

1499 mm_embeddings_out = [mm[:, :-4] for mm in multimodal_embeddings]

1500 mm_embeddings_pos = [

vllm/model_executor/models/qwen2_audio.py

+3 -3

@@ -439,9 +439,9 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

439 num_audios, max_audio_tokens, embed_dim = audio_features.shape

440 audio_output_lengths = audio_output_lengths.unsqueeze(1)

441 audio_features_mask = (

442 - torch.arange(max_audio_tokens)

443 - .expand(num_audios, max_audio_tokens)

444 - .to(audio_output_lengths.device)

442 + torch.arange(max_audio_tokens, device=audio_output_lengths.device).expand(

443 + num_audios, max_audio_tokens

444 + )

445 < audio_output_lengths

446 )

447 masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim)

vllm/model_executor/models/qwen2_vl.py

+1

@@ -634,6 +634,7 @@ class Qwen2VisionTransformer(nn.Module):

634 # Use pre-computed cos_sin_cache from RotaryEmbedding

635 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

636

637 + pos_ids = pos_ids.to(cos.device, non_blocking=True)

638 cos_combined = cos[pos_ids].flatten(1)

639 sin_combined = sin[pos_ids].flatten(1)

640 return cos_combined, sin_combined

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+5 -3

@@ -80,6 +80,7 @@ from vllm.multimodal.processing.processor import (

80 )

81 from vllm.sequence import IntermediateTensors

82 from vllm.transformers_utils.processor import cached_processor_from_config

83 +from vllm.utils.torch_utils import async_tensor_h2d

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

85

86 from .interfaces import (

@@ -502,9 +503,9 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

503 cu_chunk_lens.extend([window_aftercnn] * num_full_chunks)

504 if remainder:

505 cu_chunk_lens.append(remainder)

505 - cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(

506 - -1, dtype=torch.int32

507 - )

506 + cu_seqlens = async_tensor_h2d(

507 + cu_chunk_lens, dtype=torch.int32, device=aftercnn_lens.device

508 + ).cumsum(-1, dtype=torch.int32)

509

510 max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

511

@@ -861,6 +862,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

862 # Use pre-computed cos_sin_cache from RotaryEmbedding

863 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

864

865 + pos_ids = pos_ids.to(cos.device, non_blocking=True)

866 cos_combined = cos[pos_ids].flatten(1)

867 sin_combined = sin[pos_ids].flatten(1)

868

vllm/model_executor/models/qwen3_vl.py

+6 -3

@@ -101,6 +101,7 @@ from vllm.utils.collection_utils import is_list_of

101 from vllm.utils.math_utils import round_up

102 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

103

104 +from ...utils.torch_utils import async_tensor_h2d

105 from .interfaces import (

106 MultiModalEmbeddings,

107 SupportsEagle,

@@ -2157,7 +2158,9 @@ class Qwen3VLForConditionalGeneration(

2158 grid_thw_list = grid_thw.tolist()

2159 image_embeds_out = []

2160 for emb, size in zip(image_embeds_split, grid_thw_list):

2160 - positions = compute_mrope_for_media(size, merge_size).to(emb.device)

2161 + positions = compute_mrope_for_media(size, merge_size).to(

2162 + emb.device, non_blocking=True

2163 + )

2164 positions = torch.cat(

2165 [

2166 positions,

@@ -2388,7 +2391,7 @@ class Qwen3VLForConditionalGeneration(

2391 input_tokens=unpruned_token_ids,

2392 mm_features=[mm_feature],

2393 )[0]

2391 - .to(device)

2394 + .to(device, non_blocking=True)

2395 .permute(1, 0)

2396 )

2397 full_is_video_embed = unpruned_token_ids_tensor == embed_token_id

@@ -2636,7 +2639,7 @@ class Qwen3VLForConditionalGeneration(

2639 )

2640

2641 # Tensors

2639 - input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long)

2642 + input_ids_t = async_tensor_h2d(input_ids, device=device, dtype=torch.long)

2643

2644 mm_embeddings_out = []

2645 mm_embeddings_pos = []

vllm/model_executor/models/transformers/multimodal.py

+3 -1

@@ -152,7 +152,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):

152 # Keep these as batched, as they always have batch size as first dim

153 mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image")

154 mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image")

155 - mm_fields["num_image_patches"] = MultiModalFieldConfig.batched("image")

155 + mm_fields["num_image_patches"] = MultiModalFieldConfig.batched(

156 + "image", keep_on_cpu=True

157 + )

158 return mm_fields

159

160 def _get_hf_mm_data(

vllm/model_executor/models/utils.py

+7 -15

@@ -30,10 +30,8 @@ from vllm.model_executor.models.interfaces import supports_any_eagle

30 from vllm.multimodal import NestedTensors

31 from vllm.sequence import IntermediateTensors

32 from vllm.utils.math_utils import cdiv

33 -from vllm.utils.platform_utils import (

34 - is_pin_memory_available,

35 -)

33 from vllm.utils.torch_utils import (

34 + async_tensor_h2d,

35 direct_register_custom_op,

36 )

37

@@ -498,10 +496,9 @@ def isin_list(

496 elements: torch.Tensor,

497 test_elements_list: list[int],

498 ) -> torch.Tensor:

501 - test_elements = torch.tensor(

502 - test_elements_list,

503 - pin_memory=is_pin_memory_available(),

504 - ).to(device=elements.device, non_blocking=True)

499 + test_elements = async_tensor_h2d(

500 + test_elements_list, dtype=torch.int64, device=elements.device

501 + )

502

503 return torch.isin(elements, test_elements)

504

@@ -772,14 +769,9 @@ def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:

769 return layer_index

770

771

775 -def cast_overflow_tensors(

776 - tensors: torch.Tensor,

777 - offset: float = 1000,

778 -) -> torch.Tensor:

779 - if tensors.isinf().any() or tensors.isnan().any():

780 - clamp_value = torch.finfo(tensors.dtype).max - offset

781 - tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value)

782 - return tensors

772 +def cast_overflow_tensors(tensors: torch.Tensor, offset: float = 1000) -> torch.Tensor:

773 + clamp_value = torch.finfo(tensors.dtype).max - offset

774 + return torch.clamp(tensors, min=-clamp_value, max=clamp_value)

775

776

777 def fast_topk(

vllm/multimodal/evs.py

+1 -1

@@ -85,7 +85,7 @@ def compute_retention_mask(

85 topk_indices = order[:retain_num_tokens]

86

87 retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool)

88 - retention_mask[topk_indices] = True

88 + retention_mask.index_fill_(0, topk_indices, True)

89 retention_mask = retention_mask.reshape(dissimilarity.size())

90

91 mask = retention_mask.view(-1) # "T H W -> (T H W)"

vllm/v1/sample/thinking_budget_state.py

+25 -29

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

6

7 import torch

8

9 +from vllm.utils.torch_utils import async_tensor_h2d

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

11 BatchUpdate,

12 MoveDirectionality,

@@ -65,22 +66,9 @@ class ThinkingBudgetStateHolder:

66 self.cu_num_tokens: dict[int, int] = {}

67

68 if self.num_spec_tokens > 0:

68 - self.mask = torch.zeros(

69 - max_num_reqs * (self.num_spec_tokens + 1),

70 - dtype=torch.bool,

71 - device=device,

72 - )

73 - self.force_token_ids = torch.full(

74 - (max_num_reqs * (self.num_spec_tokens + 1),),

75 - -1,

76 - dtype=torch.long,

77 - device=device,

78 - )

69 + self._mask_capacity = max_num_reqs * (self.num_spec_tokens + 1)

70 else:

80 - self.mask = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)

81 - self.force_token_ids = torch.full(

82 - (max_num_reqs,), -1, dtype=torch.long, device=device

83 - )

71 + self._mask_capacity = max_num_reqs

72

73 def has_tracked_requests(self) -> bool:

74 """True when ``sync_batch`` has state for a ``thinking_token_budget`` row.

@@ -454,7 +442,6 @@ class ThinkingBudgetStateHolder:

442 predict_bonus_token: bool,

443 spec_token_ids_for_layout: list[list[int]],

444 ) -> torch.Tensor:

457 - self.mask[:] = False

445 cumulative_total = 0

446 self.cu_num_tokens.clear()

447

@@ -474,6 +461,11 @@ class ThinkingBudgetStateHolder:

461 else:

462 cumulative_total += 1

463

464 + # Build the active index / forced-token lists entirely on CPU so we

465 + # avoid per-iteration scalar sync writes to GPU tensors.

466 + active_indices_cpu: list[int] = []

467 + force_tokens_cpu: list[int] = []

468 +

469 for seq_idx in sorted(self._state.keys()):

470 if seq_idx not in self.cu_num_tokens:

471 continue

@@ -502,9 +494,12 @@ class ThinkingBudgetStateHolder:

494 for force_idx in force_index:

495 if end_count < len(self.think_end_token_ids):

496 mask_idx = self.cu_num_tokens[seq_idx] + force_idx

505 - if mask_idx < len(self.mask) and mask_idx < logits.shape[0]:

506 - self.mask[mask_idx] = True

507 - self.force_token_ids[mask_idx] = (

497 + if (

498 + mask_idx < self._mask_capacity

499 + and mask_idx < logits.shape[0]

500 + ):

501 + active_indices_cpu.append(mask_idx)

502 + force_tokens_cpu.append(

503 self.think_end_token_ids[end_count]

504 )

505 if predict_bonus_token:

@@ -514,15 +509,16 @@ class ThinkingBudgetStateHolder:

509 else:

510 state["bonus_token_forced"] = True

511

517 - has_active_thinking = any(

518 - state.get("in_end", False) for state in self._state.values()

519 - )

520 -

521 - if has_active_thinking:

522 - active_indices = self.mask.nonzero(as_tuple=False).view(-1)

523 -

524 - if len(active_indices) > 0:

525 - force_tokens = self.force_token_ids[active_indices]

526 - logits[active_indices, force_tokens] = 1e9

512 + if active_indices_cpu:

513 + device = logits.device

514 + active_indices = async_tensor_h2d(

515 + active_indices_cpu, dtype=torch.long, device=device

516 + )

517 + force_tokens = async_tensor_h2d(

518 + force_tokens_cpu, dtype=torch.long, device=device

519 + )

520 + # Avoid CPU->GPU sync.

521 + fill = logits.new_full((len(active_indices_cpu),), 1e9)

522 + logits.index_put_((active_indices, force_tokens), fill)

523

524 return logits

vllm/v1/worker/gpu_model_runner.py

+8 -8

@@ -1765,8 +1765,6 @@ class GPUModelRunner(

1765 self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0],

1766 non_blocking=True,

1767 )

1768 - if self.enable_prompt_embeds:

1769 - self.is_token_ids.gpu[:num_common_tokens] = True

1768 return

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

1770 sampled_tokens_index_tensor = torch.tensor(

@@ -3408,13 +3406,12 @@ class GPUModelRunner(

3406 # If a batch only has token ids, then including the embedding layer

3407 # in the CUDA graph will be more performant (like in the else case

3408 # below).

3411 - token_ids_idx = (

3412 - self.is_token_ids.gpu[:num_scheduled_tokens]

3413 - .nonzero(as_tuple=False)

3414 - .squeeze(1)

3415 - )

3409 + is_token_ids = self.is_token_ids.np[:num_scheduled_tokens]

3410 + token_ids_idx_np = np.nonzero(is_token_ids)[0]

3411 # Some tokens ids may need to become embeds

3417 - if token_ids_idx.numel() > 0:

3412 + if token_ids_idx_np.size > 0:

3413 + token_ids_idx = torch.from_numpy(token_ids_idx_np)

3414 + token_ids_idx = token_ids_idx.to(self.device, non_blocking=True)

3415 token_ids = self.input_ids.gpu[token_ids_idx]

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

3417 self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds

@@ -4601,6 +4598,9 @@ class GPUModelRunner(

4598 # appending a placeholder (-1) token id.

4599 if (req_state := self.requests.get(req_id)) is not None:

4600 req_state.output_token_ids.append(-1)

4601 + pos = self.input_batch.num_tokens_no_spec[i]

4602 + self.input_batch.is_token_ids[i, pos] = True

4603 + self.input_batch.num_tokens_no_spec[i] = pos + 1

4604 self.input_batch.prev_req_id_to_index = prev_req_id_to_index

4605

4606 def take_draft_token_ids(self) -> DraftTokenIds | None:

[Perf][4/n] Eliminate various GPU<->CPU syncs (#42347)