iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Nick Hill
2026-05-19 10:35:54 -04:00
committed by GitHub
parent a78b842d0e
commit b82e908b4c
23 changed files with 129 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/logits_processors/test_correctness.py
+1 -1
View File
@@ -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():
vllm/lora/ops/triton_ops/utils.py
+21 -9
View File
@@ -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)
vllm/lora/punica_wrapper/utils.py
+4 -2
View File
@@ -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(
vllm/model_executor/models/bert.py
+3 -6
View File
@@ -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)
vllm/model_executor/models/gemma3_mm.py
+1 -1
View File
@@ -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(
vllm/model_executor/models/glm4_1v.py
+1
View File
@@ -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
vllm/model_executor/models/granite_speech.py
+7 -7
View File
@@ -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,
vllm/model_executor/models/idefics3.py
+1 -1
View File
@@ -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(
vllm/model_executor/models/internvl.py
+5 -3
View File
@@ -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),
)
vllm/model_executor/models/llava_next.py
+1 -1
View File
@@ -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"),
)
vllm/model_executor/models/llava_onevision.py
+1 -1
View File
@@ -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"),
)
vllm/model_executor/models/phi4mm_audio.py
+9 -3
View File
@@ -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
vllm/model_executor/models/qwen2_5_omni_thinker.py
+3 -3
View File
@@ -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),
)
vllm/model_executor/models/qwen2_5_vl.py
+12 -7
View File
@@ -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 = [
vllm/model_executor/models/qwen2_audio.py
+3 -3
View File
@@ -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)
vllm/model_executor/models/qwen2_vl.py
+1
View File
@@ -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
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+5 -3
View File
@@ -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)
vllm/model_executor/models/qwen3_vl.py
+6 -3
View File
@@ -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 = []
vllm/model_executor/models/transformers/multimodal.py
+3 -1
View File
@@ -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(
vllm/model_executor/models/utils.py
+7 -15
View File
@@ -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(
vllm/multimodal/evs.py
+1 -1
View File
@@ -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)"
vllm/v1/sample/thinking_budget_state.py
+25 -29
View File
@@ -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
vllm/v1/worker/gpu_model_runner.py
+8 -8
View File
@@ -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: