@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9b0afeb4f6

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

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Aug 8, 2026 at 17:13 UTC 9b0afeb4f6c435a06d394f8a25209d4682a215af

13 files changed +99 -49

tests/basic_correctness/test_basic_correctness.py

+2 -1

@@ -307,7 +307,8 @@ def test_raise_on_logit_nans(

307

308 def compute_nan_logits(self, *args, **kwargs):

309 logits = original_compute_logits(self, *args, **kwargs)

310 - logits[0, 0] = float("nan")

310 + # `logits[0, 0] = ...` copies the scalar H2D and blocks.

311 + logits[0, 0].fill_(float("nan"))

312 return logits

313

314 monkeypatch.setattr(model_cls, "compute_logits", compute_nan_logits)

tests/v1/e2e/general/test_mamba_prefix_cache.py

+7 -7

@@ -19,6 +19,7 @@ from vllm.distributed import cleanup_dist_env_and_memory

19 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

20 from vllm.platforms import current_platform

21 from vllm.sequence import IntermediateTensors

22 +from vllm.utils.torch_utils import async_tensor_h2d

23 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

24 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager

25 from vllm.v1.core.sched.output import SchedulerOutput

@@ -77,7 +78,7 @@ def get_fake_sample_fn() -> SamplerOutput:

78 first_token_id_index = num_computed_tokens + 1

79 if spec_decode_metadata is None:

80 return SamplerOutput(

80 - sampled_token_ids=torch.tensor(

81 + sampled_token_ids=async_tensor_h2d(

82 [[prompt_token_ids[first_token_id_index]]],

83 device=DEVICE_TYPE,

84 dtype=torch.int32,

@@ -90,7 +91,7 @@ def get_fake_sample_fn() -> SamplerOutput:

91 ]

92 sampled_token_ids = accepted_tokens

93 return SamplerOutput(

93 - sampled_token_ids=torch.tensor(

94 + sampled_token_ids=async_tensor_h2d(

95 [sampled_token_ids],

96 device=DEVICE_TYPE,

97 dtype=torch.int32,

@@ -132,7 +133,7 @@ def get_fake_propose_draft_token_ids_fn():

133 ]

134 ]

135

135 - next_token_ids = torch.tensor(

136 + next_token_ids = async_tensor_h2d(

137 prompt_token_ids[

138 first_token_id_index - 1 : first_token_id_index

139 - 1

@@ -142,7 +143,7 @@ def get_fake_propose_draft_token_ids_fn():

143 dtype=torch.int32,

144 )

145

145 - valid_sampled_tokens_count = torch.tensor(

146 + valid_sampled_tokens_count = async_tensor_h2d(

147 [num_accepted_tokens],

148 device=DEVICE_TYPE,

149 dtype=torch.int32,

@@ -150,7 +151,7 @@ def get_fake_propose_draft_token_ids_fn():

151

152 self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)

153

153 - return torch.tensor(

154 + return async_tensor_h2d(

155 proposed_draft_token_ids,

156 device=DEVICE_TYPE,

157 dtype=torch.int32,

@@ -1088,11 +1089,10 @@ def _run_mamba_prefix_cache_mrv2(

1089 device=hidden_states.device,

1090 dtype=torch.int64,

1091 )

1091 - num_logits = torch.tensor(

1092 + num_logits = async_tensor_h2d(

1093 input_batch.cu_num_logits_np[1 : num_reqs + 1]

1094 - input_batch.cu_num_logits_np[:num_reqs],

1095 device=hidden_states.device,

1095 - dtype=torch.int32,

1096 )

1097 accepted = torch.full_like(num_logits, num_accepted_tokens)

1098 num_sampled = torch.minimum(accepted, num_logits)

tests/v1/logits_processors/utils.py

+14 -6

@@ -14,6 +14,7 @@ from vllm.config import VllmConfig

14 from vllm.exceptions import VLLMValidationError

15 from vllm.logger import init_logger

16 from vllm.sampling_params import SamplingParams

17 +from vllm.utils.torch_utils import async_tensor_h2d

18 from vllm.v1.sample.logits_processor import (

19 LOGITSPROCS_GROUP,

20 AdapterLogitsProcessor,

@@ -90,17 +91,24 @@ class DummyLogitsProcessor(LogitsProcessor):

91 if not self.req_info:

92 return logits

93 - # Save target values before modification

94 - cols = torch.tensor(

94 + # Save target values before modification.

95 + cols = async_tensor_h2d(

96 list(self.req_info.values()), dtype=torch.long, device=logits.device

97 )

97 - rows = torch.tensor(

98 + rows = async_tensor_h2d(

99 list(self.req_info.keys()), dtype=torch.long, device=logits.device

100 )

101 values_to_keep = logits[rows, cols].clone()

102

102 - # Mask all but target tokens

103 - logits[rows] = float("-inf")

103 + # Mask all but target tokens. Use an on-device fill tensor so the

104 + # scatter doesn't force a synchronizing scalar H2D.

105 + fill = torch.full(

106 + (rows.numel(), logits.size(-1)),

107 + float("-inf"),

108 + dtype=logits.dtype,

109 + device=logits.device,

110 + )

111 + logits[rows] = fill

112 logits[rows, cols] = values_to_keep

113

114 return logits

@@ -141,7 +149,7 @@ class DummyPerReqLogitsProcessor:

149 output_ids: list[int],

150 logits: torch.Tensor,

151 ) -> torch.Tensor:

144 - val_to_keep = logits[self.target_token].item()

152 + val_to_keep = logits[self.target_token].clone()

153 logits[:] = float("-inf")

154 logits[self.target_token] = val_to_keep

155 return logits

vllm/distributed/kv_transfer/kv_connector/utils.py

+8 -3

@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

20 from vllm.logger import init_logger

21 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

22 from vllm.platforms import current_platform

23 +from vllm.utils.torch_utils import async_tensor_h2d

24 from vllm.v1.attention.backend import AttentionBackend

25 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

@@ -179,9 +180,13 @@ def _make_src_and_dst_indices(

180 src_device: torch.device | str,

181 dst_device: torch.device | str,

182 ) -> tuple[torch.Tensor, torch.Tensor]:

182 - src_indices = torch.tensor(src_block_ids, device=src_device, dtype=torch.int64)

183 - dst_indices = torch.tensor(dst_block_ids, device=dst_device, dtype=torch.int64)

184 - return src_indices, dst_indices

183 + def _to(block_ids: list[int], device: torch.device | str) -> torch.Tensor:

184 + device = torch.device(device) if isinstance(device, str) else device

185 + if device.type == "cpu":

186 + return torch.tensor(block_ids, dtype=torch.int64, device=device)

187 + return async_tensor_h2d(block_ids, dtype=torch.int64, device=device)

188 +

189 + return _to(src_block_ids, src_device), _to(dst_block_ids, dst_device)

190

191

192 def copy_kv_blocks(

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

+4 -3

@@ -135,6 +135,7 @@ class ExampleConnector(KVConnectorBase_V1):

135 slot_mapping (torch.Tensor): the slot mapping. In shape

136 [num_tokens].

137 """

138 + slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True)

139 if isinstance(attn_metadata, MLACommonMetadata):

140 dst_kv_cache_layer_shape = dst_kv_cache_layer.shape

141 num_pages = dst_kv_cache_layer_shape[0]

@@ -178,9 +179,8 @@ class ExampleConnector(KVConnectorBase_V1):

179 filename = self._generate_filename_debug(

180 layer_name, request.token_ids, request.mm_hashes

181 )

181 - kv_cache = safetensors.torch.load_file(

182 - filename, device=str(kv_cache_layer.device)

183 - )["kv_cache"]

182 + kv_cache_cpu = safetensors.torch.load_file(filename)["kv_cache"]

183 + kv_cache = kv_cache_cpu.to("cuda", non_blocking=True)

184 if isinstance(attn_metadata, dict):

185 inject_kv_into_layer(

186 kv_cache_layer,

@@ -227,6 +227,7 @@ class ExampleConnector(KVConnectorBase_V1):

227 Assume the shape of the layer is (num_pages, page_size, xxx)

228 for MLA, and (num_pages, 2, page_size, xxx) otherwise.

229 """

230 + slot_mapping = slot_mapping.to(layer.device, non_blocking=True)

231 if isinstance(attn_metadata, MLACommonMetadata):

232 num_pages, page_size = layer.shape[0], layer.shape[1]

233 return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+2 -2

@@ -11,7 +11,7 @@ from vllm.platforms import current_platform

11 from vllm.triton_utils import tl, triton

12 from vllm.triton_utils.allocation import set_triton_allocator

13 from vllm.utils.mem_utils import get_max_shared_memory_bytes

14 -from vllm.utils.torch_utils import direct_register_custom_op

14 +from vllm.utils.torch_utils import async_tensor_h2d, direct_register_custom_op

16 from .utils import supports_pdl, supports_tma

@@ -866,7 +866,7 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):

866 tensor_ptrs = []

867 for lora_weight in lora_weights:

868 tensor_ptrs.append(lora_weight.data_ptr())

869 - ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)

869 + ptr_tensor = async_tensor_h2d(tensor_ptrs, dtype=torch.uint64, device=device)

870

871 _LORA_PTR_DICT[key] = ptr_tensor

872 return _LORA_PTR_DICT.get(key)

vllm/model_executor/models/chameleon.py

+32 -4

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

55 )

56 from vllm.sequence import IntermediateTensors

57 from vllm.utils.tensor_schema import TensorSchema, TensorShape

58 +from vllm.utils.torch_utils import async_tensor_h2d

60 from .interfaces import (

61 MultiModalEmbeddings,

@@ -819,8 +820,20 @@ class ChameleonImageVocabularyMapping:

820

821 def convert_img2bpe(self, img_batch: torch.Tensor) -> torch.Tensor:

822 device = img_batch.device

822 - img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]

823 - return img_tokens.to(device)

823 + # Cache a per-device copy of the (small, static) mapping tensor so we

824 + # can index entirely on `device` instead of forcing a D2H on

825 + # `img_batch` and an H2D on the result.

826 + cache = getattr(self, "_img2bpe_mapping_cache", None)

827 + if cache is None:

828 + cache = {}

829 + self._img2bpe_mapping_cache = cache

830 + mapping_on_device = cache.get(device)

831 + if mapping_on_device is None:

832 + mapping_on_device = async_tensor_h2d(

833 + self.img2bpe_mapping_tensor, device=device

834 + )

835 + cache[device] = mapping_on_device

836 + return mapping_on_device[img_batch]

837

838

839 class ChameleonModel(nn.Module):

@@ -1025,8 +1038,23 @@ class ChameleonForConditionalGeneration(

1038 # Disallow image tokens which does not include special

1039 # begin-image and end-image tokens

1040 if logits is not None:

1028 - image_tokens = self.model.vocabulary_mapping.image_tokens

1029 - logits[:, image_tokens] = torch.finfo(logits.dtype).min

1041 + # Cache a per-device index tensor for the (static) image-token

1042 + # set, and use `index_fill_` instead of advanced-index assign

1043 + # so the scatter runs entirely on device (no host roundtrip

1044 + # for the scalar fill value or for the Python-list indices).

1045 + cache = getattr(self, "_image_tokens_index_cache", None)

1046 + if cache is None:

1047 + cache = {}

1048 + self._image_tokens_index_cache = cache

1049 + image_tokens_idx = cache.get(logits.device)

1050 + if image_tokens_idx is None:

1051 + image_tokens_idx = async_tensor_h2d(

1052 + self.model.vocabulary_mapping.image_tokens,

1053 + dtype=torch.long,

1054 + device=logits.device,

1055 + )

1056 + cache[logits.device] = image_tokens_idx

1057 + logits.index_fill_(1, image_tokens_idx, torch.finfo(logits.dtype).min)

1058

1059 return logits

1060

vllm/model_executor/models/gemma3n_mm.py

+13 -4

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

54 )

55 from vllm.sequence import IntermediateTensors

56 from vllm.utils.tensor_schema import TensorSchema, TensorShape

57 +from vllm.utils.torch_utils import async_tensor_h2d

59 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription

60 from .utils import (

@@ -683,10 +684,18 @@ class Gemma3nForConditionalGeneration(

684 # We handle both cases:

685 # - If fewer tokens: pad with the embedding of the last vocab token

686 # - If more tokens: truncate to the expected count

686 - # TODO precompute and cache padding

687 - audio_padding_toks = torch.tensor(

688 - [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device

689 - )

687 + # Cache the single-scalar padding-token tensor per-device to avoid a

688 + # synchronous H2D tensor construction on every forward.

689 + cache = getattr(self, "_audio_padding_toks_cache", None)

690 + if cache is None:

691 + cache = {}

692 + self._audio_padding_toks_cache = cache

693 + audio_padding_toks = cache.get(audio_features.device)

694 + if audio_padding_toks is None:

695 + audio_padding_toks = async_tensor_h2d(

696 + [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device

697 + )

698 + cache[audio_features.device] = audio_padding_toks

699 audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)

700 audio_features = torch.where(

701 audio_mask.unsqueeze(-1), audio_padding_embs, audio_features

vllm/model_executor/models/glm4_1v.py

+11 -4

@@ -99,6 +99,7 @@ from vllm.sequence import IntermediateTensors

99 from vllm.transformers_utils.processor import get_processor_cls_name_from_config

100 from vllm.transformers_utils.utils import convert_model_repo_to_path

101 from vllm.utils.tensor_schema import TensorSchema, TensorShape

102 +from vllm.utils.torch_utils import async_tensor_h2d

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

104 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

105

@@ -535,7 +536,8 @@ class Glm4vVisionEmbeddings(nn.Module):

536 total_seq = h_coords.shape[0]

537 device = pos_embed_weight.device

538

538 - # Move coordinates to correct device

539 + # Move coordinates to correct device. The only caller already hands

540 + # these over pinned + non-blocking, so this is normally a no-op.

541 h_coords, w_coords = h_coords.to(device), w_coords.to(device)

542

543 # Handle empty sequence case

@@ -816,10 +818,15 @@ class Glm4vVisionTransformer(nn.Module):

818 )

819

820 lengths = [h * w] * t

819 - image_shapes = torch.tensor([[t, h, w]], device=device)

821 + image_shapes = async_tensor_h2d(

822 + [[t, h, w]], dtype=torch.long, device=device

823 + )

824

821 - h_coords_repeated = h_coords.repeat(t)

822 - w_coords_repeated = w_coords.repeat(t)

825 + # Build the coordinates on the host (cheap integer math) but move

826 + # them across pinned + non-blocking, so the consumer's

827 + # `.to(device)` below is a no-op rather than a blocking H2D.

828 + h_coords_repeated = async_tensor_h2d(h_coords.repeat(t), device=device)

829 + w_coords_repeated = async_tensor_h2d(w_coords.repeat(t), device=device)

830

831 embeds = self.embeddings(

832 embeddings=torch.zeros(

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+1 -1

@@ -437,7 +437,7 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

437 # Compute chunk information

438 chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()

439

440 - chunk_lengths = torch.tensor(

440 + chunk_lengths = async_tensor_h2d(

441 [self.n_window * 2] * chunk_num.sum(),

442 dtype=torch.long,

443 device=feature_lens.device,

vllm/model_executor/models/voxtral.py

+3 -7

@@ -773,15 +773,11 @@ class VoxtralEncoderModel(nn.Module):

773 if global_log_mel_max := self.config.global_log_mel_max:

774 if not isinstance(global_log_mel_max, float):

775 raise TypeError(f"{global_log_mel_max=} needs to be of type float.")

776 - log_spec_max = torch.tensor(

777 - global_log_mel_max,

778 - device=log_spec.device,

779 - dtype=log_spec.dtype,

780 - )

776 + # Use `clamp` to avoid gpu<->cpu sync.

777 + log_spec = torch.clamp(log_spec, min=global_log_mel_max - 8.0)

778 else:

779 log_spec_max = log_spec.max()

783 -

784 - log_spec = torch.maximum(log_spec, log_spec_max - 8.0)

780 + log_spec = torch.maximum(log_spec, log_spec_max - 8.0)

781 log_spec = (log_spec + 4.0) / 4.0

782 return log_spec.to(input_dtype)

783

vllm/v1/spec_decode/extract_hidden_states.py

+1 -6

@@ -13,7 +13,6 @@ from vllm.distributed.eplb.eplb_state import EplbState

13 from vllm.forward_context import set_forward_context

14 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

15 from vllm.model_executor.model_loader import get_model

16 -from vllm.utils.torch_utils import PIN_MEMORY

16 from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata

17 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

18 from vllm.v1.utils import CpuGpuBuffer

@@ -59,11 +58,7 @@ class ExtractHiddenStatesProposer:

58 )

60 self.backup_next_token_ids = CpuGpuBuffer(

62 - max_batch_size,

63 - dtype=torch.int32,

64 - pin_memory=PIN_MEMORY,

65 - device=device,

66 - with_numpy=True,

61 + max_batch_size, dtype=torch.int32, device=device

62 )

64 self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -3069,7 +3069,7 @@ class GPUModelRunner(

3069

3070 self._cache_encoder_output(

3071 mm_hashes[i],

3072 - pe_tensor.to(self.device),

3072 + async_tensor_h2d(pe_tensor, device=self.device),

3073 scheduler_output.ec_manager_metadata,

3074 scheduler_output.free_encoder_mm_hashes,

3075 )