@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8277c42e4c

[Perf] Ensure async h2d copies are pinned in more places (#55202)

Nick Hill committed Sep 4, 2026 at 17:17 UTC 8277c42e4c74c5dd604f29b6d938a3456658f247

21 files changed +206 -73

tests/kernels/attention/test_flashinfer.py

@@ -190,6 +190,7 @@ def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None:

190 indptr_cpu=kv_indptr,

191 indices=kv_indices,

192 last_page_len_cpu=kv_last_page_lens,

193 + seq_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32, device="cpu"),

194 num_qo_heads=num_query_heads,

195 num_kv_heads=num_kv_heads,

196 head_dim=head_size,

@@ -268,6 +269,7 @@ def test_fast_plan_decode_matches_full_plan(

269 indptr_cpu=kv_indptr,

270 indices=kv_indices_buf,

271 last_page_len_cpu=kv_last_page_lens,

272 + seq_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32, device="cpu"),

273 num_qo_heads=num_query_heads,

274 num_kv_heads=num_kv_heads,

275 head_dim=head_size,

vllm/distributed/eplb/eplb_state.py

+10 -1

@@ -49,6 +49,7 @@ from vllm.logger import init_logger

49 from vllm.model_executor.models.interfaces import MixtureOfExperts

50 from vllm.platforms import current_platform

51 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

52 +from vllm.utils.torch_utils import PIN_MEMORY

54 from .async_worker import start_async_worker

55 from .eplb_communicator import EplbCommunicator, create_eplb_communicator

@@ -1298,6 +1299,8 @@ def _commit_eplb_maps_for_layer(

1299 f"Current number of physical experts: {dst.shape[0]}. New number of physical "

1300 f"experts {src.shape[0]}."

1301 )

1302 + if PIN_MEMORY and src.is_cpu:

1303 + src = src.new_empty(src.shape, pin_memory=True).copy_(src)

1304 dst.copy_(src, non_blocking=True)

1305

1306 num_logical_experts = model_state.logical_to_physical_map.shape[1]

@@ -1309,6 +1312,8 @@ def _commit_eplb_maps_for_layer(

1312 src = new_replica_count

1313 dst = model_state.logical_replica_count[layer]

1314 assert src.shape == dst.shape

1315 + if PIN_MEMORY:

1316 + src = src.pin_memory()

1317 dst.copy_(src, non_blocking=True)

1318

1319

@@ -1326,12 +1331,14 @@ def _commit_eplb_maps(

1331 src = new_physical_to_logical_map

1332 dst = model_state.physical_to_logical_map

1333

1334 + if PIN_MEMORY and src.is_cpu:

1335 + src = src.new_empty(src.shape, pin_memory=True).copy_(src)

1336 # Rare Case: When the number of physical experts has changed, discard the old

1337 # physical to logical expert map and use the new one. This only happens when the

1338 # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the

1339 # new map into the old one.

1340 if src.shape[1] != dst.shape[1]:

1334 - model_state.physical_to_logical_map = src.to(dst.device)

1341 + model_state.physical_to_logical_map = src.to(dst.device, non_blocking=True)

1342 else:

1343 dst.copy_(src, non_blocking=True)

1344

@@ -1346,6 +1353,8 @@ def _commit_eplb_maps(

1353 # Commit logical_replica_count

1354 src = new_replica_count

1355 dst = model_state.logical_replica_count

1356 + if PIN_MEMORY:

1357 + src = src.pin_memory()

1358 dst.copy_(src, non_blocking=True)

1359

1360

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

+6 -4

@@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

21 )

22 from vllm.distributed.parallel_state import get_tensor_model_parallel_rank

23 from vllm.logger import init_logger

24 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

25 from vllm.v1.attention.backend import AttentionMetadata

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

@@ -377,10 +378,11 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

378 with torch.cuda.stream(copy_stream):

379 # Move the CPU slot_mapping to GPU on the copy stream so the

380 # implicit H2D inside fancy indexing doesn't sync the default

380 - # stream.

381 - slot_mapping_gpu = slot_mapping.to(

382 - device=self._kv_cache.device, non_blocking=True

383 - )

381 + # stream. Deliberate bulk transfer in this reference connector.

382 + with gpu_sync_allowed():

383 + slot_mapping_gpu = slot_mapping.to(

384 + device=self._kv_cache.device, non_blocking=True

385 + )

386 hidden_states_gpu = extract_from_kv_cache(

387 self._kv_cache, slot_mapping_gpu, num_tokens

388 )

vllm/model_executor/layers/attention/mla_attention.py

+8 -2

@@ -265,6 +265,7 @@ from vllm.platforms import current_platform

265 from vllm.utils.flashinfer import has_flashinfer

266 from vllm.utils.math_utils import cdiv, round_down, round_up

267 from vllm.utils.torch_utils import (

268 + PIN_MEMORY,

269 LayerNameType,

270 _encode_layer_name,

271 _resolve_layer_name,

@@ -2510,8 +2511,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

2511 query_start_loc_cpu[reqs_start + 1 : num_reqs + 1]

2512 - query_start_loc_cpu[reqs_start:num_reqs]

2513 )

2513 - context_lens_cpu = (

2514 - seq_lens_cpu[reqs_start:num_reqs] - prefill_query_lens_cpu

2514 + context_lens_cpu = torch.empty(

2515 + num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY

2516 + )

2517 + torch.subtract(

2518 + seq_lens_cpu[reqs_start:num_reqs],

2519 + prefill_query_lens_cpu,

2520 + out=context_lens_cpu,

2521 )

2522 prefill_query_start_loc = (

2523 query_start_loc[reqs_start:] - query_start_loc[reqs_start]

vllm/model_executor/layers/fused_moe/modular_kernel.py

+5 -1

@@ -34,6 +34,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

34 QuantKey,

35 )

36 from vllm.platforms import current_platform

37 +from vllm.utils.torch_utils import PIN_MEMORY

38 from vllm.v1.worker.ubatching import (

39 dbo_enabled,

40 dbo_maybe_run_recv_hook,

@@ -109,7 +110,10 @@ class ExpertTokensMetadata:

110 expert_num_tokens_list: list[int], device: str

111 ) -> "ExpertTokensMetadata":

112 expert_num_tokens_cpu = torch.tensor(

112 - expert_num_tokens_list, device="cpu", dtype=torch.int32

113 + expert_num_tokens_list,

114 + device="cpu",

115 + dtype=torch.int32,

116 + pin_memory=PIN_MEMORY,

117 )

118 return ExpertTokensMetadata(

119 expert_num_tokens=expert_num_tokens_cpu.to(device, non_blocking=True),

vllm/model_executor/models/cosmos3_edge.py

+2 -4

@@ -21,6 +21,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

21 from vllm.multimodal.inputs import MultiModalFeatureSpec

22 from vllm.sequence import IntermediateTensors

23 from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig

24 +from vllm.utils.torch_utils import async_tensor_h2d

26 from .interfaces import (

27 MultiModalEmbeddings,

@@ -80,10 +81,7 @@ class Cosmos3EdgeVisionEncoder(Siglip2VisionTransformer):

81 dim=0,

82 )

83 lengths_cpu = spatial_shapes.prod(dim=-1).to(torch.int32)

83 - lengths = lengths_cpu.to(

84 - device=pixel_values.device,

85 - non_blocking=True,

86 - )

84 + lengths = async_tensor_h2d(lengths_cpu, pixel_values.device)

86 cu_seqlens = torch.zeros(

87 lengths.numel() + 1,

vllm/model_executor/models/ernie45_vl.py

+4 -4

@@ -538,7 +538,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

538 else:

539 max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

540

541 - cu_seqlens = cu_seqlens.to(device, non_blocking=True)

541 + cu_seqlens = async_tensor_h2d(cu_seqlens, device)

542

543 return {

544 "rotary_pos_emb": rotary_pos_emb,

@@ -1746,8 +1746,8 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1746 # Eager fallback: run the full pipeline (ViT + resampler). The result

1747 # is scattered directly, so it must be the post-merge embeddings.

1748 pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype)

1749 - grid_thw = mm_kwargs["image_grid_thw"].to(

1750 - self.vision_model.device, non_blocking=True

1749 + grid_thw = async_tensor_h2d(

1750 + mm_kwargs["image_grid_thw"], self.vision_model.device

1751 )

1752 image_features = self.vision_model(pixel_values, grid_thw)

1753 return self.resampler_model(image_features, grid_thw)

@@ -1768,7 +1768,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1768 output = outputs["default"]

1769 assert batch_mm_kwargs is not None

1770 grid_thw_cpu = batch_mm_kwargs["image_grid_thw"]

1771 - grid_thw = grid_thw_cpu.to(output.device, non_blocking=True)

1771 + grid_thw = async_tensor_h2d(grid_thw_cpu, output.device)

1772 # The valid token count slices the graph output for the eager

1773 # resampler call, so it has to come back to the host.

1774 num_valid = int(

vllm/model_executor/models/glm_ocr.py

+6 -2

@@ -70,6 +70,7 @@ from vllm.model_executor.models.glm4_1v import (

70 Glm4vVisionTransformer,

71 )

72 from vllm.multimodal import MULTIMODAL_REGISTRY

73 +from vllm.utils.torch_utils import PIN_MEMORY

75 from .utils import (

76 maybe_prefix,

@@ -344,8 +345,11 @@ class GlmOcrVisionTransformer(Glm4vVisionTransformer):

345 cu_seqlens = torch.repeat_interleave(

346 grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

347 ).cumsum(dim=0, dtype=torch.int32)

347 - cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])

348 - cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)

348 + # cat straight into a pinned buffer so the H2D copy stays non-blocking.

349 + cu_seqlens = torch.cat(

350 + [cu_seqlens.new_zeros(1), cu_seqlens],

351 + out=cu_seqlens.new_empty(cu_seqlens.numel() + 1, pin_memory=PIN_MEMORY),

352 + ).to(self.device, non_blocking=True)

353

354 # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations

355 max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

vllm/model_executor/models/isaac.py

+9 -3

@@ -69,6 +69,7 @@ from vllm.transformers_utils.processors.isaac import (

69 get_image_size_for_max_num_patches,

70 )

71 from vllm.utils.tensor_schema import TensorSchema, TensorShape

72 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

74 from .vision import is_vit_use_data_parallel

@@ -970,9 +971,14 @@ class IsaacForConditionalGeneration(

971 device = next(self.language_model.parameters()).device

972 dtype = self.vision_embedding.linear_fc1.weight.dtype

973 pixel_values = pixel_values.to(device=device, dtype=dtype)

973 - spatial_grids = image_grid_thw[:, 1:3].to(

974 - device, dtype=torch.int32, non_blocking=True

975 - )

974 + # The [:, 1:3] column slice isn't densely laid out, so stage it with a

975 + # single copy straight into pinned memory to keep the H2D non-blocking.

976 + spatial_grids_cpu = torch.empty(

977 + (image_grid_thw.shape[0], 2),

978 + dtype=image_grid_thw.dtype,

979 + pin_memory=PIN_MEMORY,

980 + ).copy_(image_grid_thw[:, 1:3])

981 + spatial_grids = async_tensor_h2d(spatial_grids_cpu, device, dtype=torch.int32)

982

983 vision_embeddings = self.vision_embedding((pixel_values, spatial_grids))

984 merge_size = self.config.vision_config.pixel_shuffle_scale_factor

vllm/model_executor/models/keye.py

+41 -16

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

59 )

60 from vllm.sequence import IntermediateTensors

61 from vllm.utils.tensor_schema import TensorSchema, TensorShape

62 -from vllm.utils.torch_utils import async_tensor_h2d

62 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

64 from .interfaces import (

65 MultiModalEmbeddings,

@@ -1305,17 +1305,29 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):

1305 )

1306 else:

1307 pixel_values = image_input["pixel_values"].type(self.visual.dtype)

1308 - # These are all built on the host, so stage them over

1309 - # non-blocking rather than forcing a sync per transfer.

1310 - siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1311 - pixel_values.device, non_blocking=True

1312 - )

1308 + # These are all built on the host; concat straight into pinned

1309 + # buffers so the H2D copies stay non-blocking.

1310 + siglip_position_ids = torch.concat(

1311 + siglip_position_ids,

1312 + dim=0,

1313 + out=torch.empty(

1314 + sum(t.numel() for t in siglip_position_ids),

1315 + dtype=torch.int64,

1316 + pin_memory=PIN_MEMORY,

1317 + ),

1318 + ).to(pixel_values.device, non_blocking=True)

1319 cu_seqlens = async_tensor_h2d(

1320 cu_seqlens, dtype=torch.int32, device=pixel_values.device

1321 )

1316 - sample_indices = torch.concat(sample_indices, dim=0).to(

1317 - pixel_values.device, non_blocking=True

1318 - )

1322 + sample_indices = torch.concat(

1323 + sample_indices,

1324 + dim=0,

1325 + out=torch.empty(

1326 + sum(t.numel() for t in sample_indices),

1327 + dtype=torch.int64,

1328 + pin_memory=PIN_MEMORY,

1329 + ),

1330 + ).to(pixel_values.device, non_blocking=True)

1331

1332 image_embeds = self.visual(

1333 pixel_values=pixel_values,

@@ -1359,16 +1371,29 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):

1371 )

1372 else:

1373 pixel_values_videos = pixel_values_videos.type(self.visual.dtype)

1362 - # These are host-built; stage them over non-blocking.

1363 - siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1364 - pixel_values_videos.device, non_blocking=True

1365 - )

1374 + # These are host-built; concat straight into pinned buffers so

1375 + # the H2D copies stay non-blocking.

1376 + siglip_position_ids = torch.concat(

1377 + siglip_position_ids,

1378 + dim=0,

1379 + out=torch.empty(

1380 + sum(t.numel() for t in siglip_position_ids),

1381 + dtype=torch.int64,

1382 + pin_memory=PIN_MEMORY,

1383 + ),

1384 + ).to(pixel_values_videos.device, non_blocking=True)

1385 cu_seqlens = async_tensor_h2d(

1386 cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device

1387 )

1369 - sample_indices = torch.concat(sample_indices, dim=0).to(

1370 - pixel_values_videos.device, non_blocking=True

1371 - )

1388 + sample_indices = torch.concat(

1389 + sample_indices,

1390 + dim=0,

1391 + out=torch.empty(

1392 + sum(t.numel() for t in sample_indices),

1393 + dtype=torch.int64,

1394 + pin_memory=PIN_MEMORY,

1395 + ),

1396 + ).to(pixel_values_videos.device, non_blocking=True)

1397

1398 video_embeds = self.visual(

1399 pixel_values=pixel_values_videos,

vllm/model_executor/models/paddleocr_vl.py

+12 -5

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

68 )

69 from vllm.sequence import IntermediateTensors

70 from vllm.utils.tensor_schema import TensorSchema, TensorShape

71 -from vllm.utils.torch_utils import async_tensor_h2d

71 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

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

74 from .ernie45 import Ernie4_5ForCausalLM

@@ -1142,10 +1142,17 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support

1142 siglip_position_ids.append(image_position_ids)

1143 cu_seqlens.append(cu_seqlens[-1] + numel)

1144

1145 - # Both are built on the host; stage them over non-blocking.

1146 - siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1147 - pixel_values.device, non_blocking=True

1148 - )

1145 + # Both are built on the host; concat straight into a pinned buffer

1146 + # so the H2D copy stays non-blocking.

1147 + siglip_position_ids = torch.concat(

1148 + siglip_position_ids,

1149 + dim=0,

1150 + out=torch.empty(

1151 + sum(t.numel() for t in siglip_position_ids),

1152 + dtype=torch.int64,

1153 + pin_memory=PIN_MEMORY,

1154 + ),

1155 + ).to(pixel_values.device, non_blocking=True)

1156 cu_seqlens = async_tensor_h2d(

1157 cu_seqlens, dtype=torch.int32, device=pixel_values.device

1158 )

vllm/model_executor/models/phi4mm_audio.py

+2 -2

@@ -643,8 +643,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

643 seq_len, batch_size, chunk_size_nc, left_chunk_nc

644 )

645 if device.type != "cpu":

646 - enc_streaming_mask_nc = enc_streaming_mask_nc.contiguous().to(

647 - device, non_blocking=True

646 + enc_streaming_mask_nc = async_tensor_h2d(

647 + enc_streaming_mask_nc.contiguous(), device

648 )

649 if masks is not None:

650 hs_mask_nc = masks & enc_streaming_mask_nc

vllm/model_executor/models/qwen2_5_omni_thinker.py

+3 -2

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

91 from vllm.sequence import IntermediateTensors

92 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

93 from vllm.utils.tensor_schema import TensorSchema, TensorShape

94 +from vllm.utils.torch_utils import async_tensor_h2d

96 from .interfaces import (

97 MultiModalEmbeddings,

@@ -985,8 +986,8 @@ class Qwen2_5OmniConditionalGenerationMixin:

986 input_features = audio_input["input_features"]

987 # audio_feature_lengths is keep_on_cpu; the audio tower derives

988 # device placement from feature_lens, so move it explicitly.

988 - audio_feature_lengths = audio_input["audio_feature_lengths"].to(

989 - input_features.device, non_blocking=True

989 + audio_feature_lengths = async_tensor_h2d(

990 + audio_input["audio_feature_lengths"], input_features.device

991 )

992

993 audio_feat_lengths, audio_output_lengths = (

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+3 -3

@@ -1001,7 +1001,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

1001

1002 # Move cu_seqlens to GPU; grid_thw may be on CPU during profile_run

1003 # and FA3 vit attention requires cu_seqlens on CUDA.

1004 - cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)

1004 + cu_seqlens = async_tensor_h2d(cu_seqlens, self.device)

1005 hidden_states = hidden_states.unsqueeze(1)

1006 rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)

1007 rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)

@@ -1574,8 +1574,8 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix

1574 input_features = audio_input["input_features"]

1575 # audio_feature_lengths is keep_on_cpu; the audio tower derives

1576 # device placement from feature_lens, so move it explicitly.

1577 - audio_feature_lengths = audio_input["audio_feature_lengths"].to(

1578 - input_features.device, non_blocking=True

1577 + audio_feature_lengths = async_tensor_h2d(

1578 + audio_input["audio_feature_lengths"], input_features.device

1579 )

1580

1581 audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)

vllm/model_executor/models/qwen3_vl.py

+7 -10

@@ -2367,8 +2367,8 @@ class Qwen3VLForConditionalGeneration(

2367 grid_thw_list = grid_thw.tolist()

2368 image_embeds_out = []

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

2370 - positions = compute_mrope_for_media(size, merge_size).to(

2371 - emb.device, non_blocking=True

2370 + positions = async_tensor_h2d(

2371 + compute_mrope_for_media(size, merge_size), emb.device

2372 )

2373 positions = torch.cat(

2374 [

@@ -2596,14 +2596,11 @@ class Qwen3VLForConditionalGeneration(

2596 identifier="DUMMY",

2597 mm_position=PlaceholderRange(offset=0, length=len(unpruned_token_ids)),

2598 )

2599 - original_mrope = (

2600 - self.get_mrope_input_positions(

2601 - input_tokens=unpruned_token_ids,

2602 - mm_features=[mm_feature],

2603 - )[0]

2604 - .to(device, non_blocking=True)

2605 - .permute(1, 0)

2606 - )

2599 + original_mrope_cpu = self.get_mrope_input_positions(

2600 + input_tokens=unpruned_token_ids,

2601 + mm_features=[mm_feature],

2602 + )[0]

2603 + original_mrope = async_tensor_h2d(original_mrope_cpu, device).permute(1, 0)

2604 full_is_video_embed = unpruned_token_ids_tensor == embed_token_id

2605

2606 with gpu_sync_allowed():

vllm/utils/flashinfer.py

+26

@@ -21,6 +21,7 @@ import vllm.envs as envs

21 from vllm.logger import init_logger

22 from vllm.platforms import current_platform

23 from vllm.utils.math_utils import cdiv

24 +from vllm.utils.torch_utils import PIN_MEMORY

26 logger = init_logger(__name__)

@@ -181,6 +182,31 @@ def _lazy_import_wrapper(

182 return wrapper

183

184

185 +def pin_host_range_buf(length: int) -> None:

186 + """Pin flashinfer's cached host-side arange buffer covering `length`.

187 +

188 + flashinfer's decode plan() builds its host qo_indptr from the cached

189 + `_get_range_buf` entry and copies it to the GPU with non_blocking=True on

190 + every call; the cache entry is unpinned, so the copy silently blocks the

191 + host. Replace the entry with a pinned copy so the copies stay async.

192 + """

193 + if not PIN_MEMORY:

194 + return

195 + mod = _get_submodule("flashinfer.utils")

196 + if mod is None:

197 + return

198 + get_range_buf = getattr(mod, "_get_range_buf", None)

199 + cache_buf = getattr(mod, "_cache_buf", None)

200 + ceil_pow2 = getattr(mod, "_ceil_pow2", None)

201 + if get_range_buf is None or cache_buf is None or ceil_pow2 is None:

202 + return

203 + get_range_buf(length, "cpu") # Ensure the cache entry exists.

204 + key = (f"range_{ceil_pow2(length)}", "cpu")

205 + buf = cache_buf[key]

206 + if not buf.is_pinned():

207 + cache_buf[key] = buf.pin_memory()

208 +

209 +

210 # Create lazy wrappers for each function

211 flashinfer_trtllm_bf16_moe = _lazy_import_wrapper(

212 "flashinfer.fused_moe", "trtllm_bf16_moe"

vllm/v1/attention/backends/flashinfer.py

+34 -1

@@ -15,6 +15,7 @@ from flashinfer import (

15 BatchPrefillWithPagedKVCacheWrapper,

16 BatchPrefillWithRaggedKVCacheWrapper,

17 MultiLevelCascadeAttentionWrapper,

18 + get_seq_lens,

19 )

20 from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache

21 from flashinfer.prefill import trtllm_batch_context_with_kv_cache

@@ -43,6 +44,7 @@ from vllm.utils.flashinfer import (

44 can_use_trtllm_attention,

45 flashinfer_xqa_batch_decode_with_kv_cache,

46 force_use_trtllm_attention,

47 + pin_host_range_buf,

48 supports_trtllm_attention,

49 use_trtllm_attention,

50 )

@@ -299,6 +301,7 @@ class BatchDCPPrefillWrapper:

301 paged_kv_indptr_cpu: torch.Tensor,

302 paged_kv_indices: torch.Tensor,

303 paged_kv_last_page_len_cpu: torch.Tensor,

304 + kv_lens_cpu: torch.Tensor,

305 page_size: int,

306 num_qo_heads: int,

307 dcp_world_size: int,

@@ -318,6 +321,7 @@ class BatchDCPPrefillWrapper:

321 paged_kv_indptr=paged_kv_indptr_cpu,

322 paged_kv_indices=paged_kv_indices,

323 paged_kv_last_page_len=paged_kv_last_page_len_cpu,

324 + seq_lens=kv_lens_cpu,

325 num_qo_heads=num_qo_heads * dcp_world_size,

326 num_kv_heads=num_kv_heads,

327 head_dim_qk=head_dim,

@@ -1580,6 +1584,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1584 paged_kv_last_page_len_prefill_cpu = (

1585 paged_kv_last_page_len_prefill_cpu.pin_memory()

1586 )

1587 + # plan() also copies kv_lens to the GPU with non_blocking=True,

1588 + # deriving them via get_seq_lens() into unpinned memory when

1589 + # seq_lens is not passed; pass an explicit pinned copy of the

1590 + # same values so that copy stays async too.

1591 + kv_lens_prefill_cpu = get_seq_lens(

1592 + paged_kv_indptr_prefill_cpu,

1593 + paged_kv_last_page_len_prefill_cpu,

1594 + page_size,

1595 + )

1596 + if PIN_MEMORY:

1597 + kv_lens_prefill_cpu = kv_lens_prefill_cpu.pin_memory()

1598 if self.use_dcp:

1599 assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)

1600 prefill_wrapper.plan(

@@ -1587,6 +1602,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1602 paged_kv_indptr_cpu=paged_kv_indptr_prefill_cpu,

1603 paged_kv_indices=paged_kv_indices,

1604 paged_kv_last_page_len_cpu=paged_kv_last_page_len_prefill_cpu,

1605 + kv_lens_cpu=kv_lens_prefill_cpu,

1606 page_size=self.page_size,

1607 num_qo_heads=self.num_qo_heads,

1608 dcp_world_size=self.dcp_world_size,

@@ -1616,6 +1632,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1632 paged_kv_indptr=paged_kv_indptr_prefill_cpu,

1633 paged_kv_indices=paged_kv_indices,

1634 paged_kv_last_page_len=paged_kv_last_page_len_prefill_cpu,

1635 + seq_lens=kv_lens_prefill_cpu,

1636 num_qo_heads=self.num_qo_heads,

1637 num_kv_heads=self.num_kv_heads,

1638 head_dim_qk=self.head_dim,

@@ -1716,11 +1733,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1733 if PIN_MEMORY:

1734 paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()

1735 paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()

1736 + # plan() also copies kv_lens to the GPU with non_blocking=True

1737 + # (trtllm-gen/cute-dsl backends), deriving them via

1738 + # get_seq_lens() into unpinned memory when seq_lens is not

1739 + # passed; pass an explicit pinned copy of the same values

1740 + # so that copy stays async too.

1741 + kv_lens_decode_cpu = get_seq_lens(

1742 + paged_kv_indptr_cpu, paged_kv_last_page_len_cpu, page_size

1743 + )

1744 + if PIN_MEMORY:

1745 + kv_lens_decode_cpu = kv_lens_decode_cpu.pin_memory()

1746 fast_plan_decode(

1747 decode_wrapper,

1748 indptr_cpu=paged_kv_indptr_cpu,

1749 indices=paged_kv_indices,

1750 last_page_len_cpu=paged_kv_last_page_len_cpu,

1751 + seq_lens_cpu=kv_lens_decode_cpu,

1752 num_qo_heads=self.num_qo_heads * self.dcp_world_size,

1753 num_kv_heads=self.num_kv_heads,

1754 head_dim=self.head_dim,

@@ -2579,6 +2607,7 @@ def fast_plan_decode(

2607 indptr_cpu: torch.Tensor,

2608 indices: torch.Tensor,

2609 last_page_len_cpu: torch.Tensor,

2610 + seq_lens_cpu: torch.Tensor,

2611 num_qo_heads: int,

2612 num_kv_heads: int,

2613 head_dim: int,

@@ -2614,6 +2643,10 @@ def fast_plan_decode(

2643 # original plan if we run for dynamic shape. For fixed shape (cudagraph),

2644 # this warm up is to generate the _cached_module for the decode wrapper.

2645 if not self.is_cuda_graph_enabled or getattr(self, "vllm_first_call", True):

2646 + # plan() copies qo_indptr from flashinfer's internal cached host

2647 + # arange buffer (_get_range_buf) with non_blocking=True; pin the

2648 + # cache entry so the copy stays async.

2649 + pin_host_range_buf(len(last_page_len_cpu) + 1)

2650 self.plan(

2651 indptr=indptr_cpu,

2652 indices=indices,

@@ -2634,7 +2667,7 @@ def fast_plan_decode(

2667 rope_theta=rope_theta,

2668 non_blocking=non_blocking,

2669 block_tables=None,

2637 - seq_lens=None,

2670 + seq_lens=seq_lens_cpu,

2671 fixed_split_size=fixed_split_size,

2672 disable_split_kv=disable_split_kv,

2673 )

vllm/v1/attention/backends/hpc_attn.py

+8 -2

@@ -19,6 +19,7 @@ from vllm.config.cache import CacheDType

19 from vllm.logger import init_logger

20 from vllm.platforms import current_platform

21 from vllm.platforms.interface import DeviceCapability

22 +from vllm.utils.torch_utils import PIN_MEMORY

23 from vllm.v1.attention.backend import (

24 AttentionBackend,

25 AttentionCGSupport,

@@ -216,8 +217,13 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):

217 if num_prefills > 0:

218 qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu

219 prefill_start = num_decodes

219 - qo_indptr_prefill_cpu = (

220 - qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]

220 + qo_indptr_prefill_cpu = torch.empty(

221 + num_prefills + 1, dtype=torch.int32, pin_memory=PIN_MEMORY

222 + )

223 + torch.subtract(

224 + qo_indptr_cpu[prefill_start:],

225 + qo_indptr_cpu[prefill_start],

226 + out=qo_indptr_prefill_cpu,

227 )

228 qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True)

229

vllm/v1/attention/backends/rocm_aiter_fa.py

+10 -5

@@ -21,7 +21,11 @@ from vllm.platforms.interface import DeviceCapability

21 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

22 from vllm.utils.math_utils import cdiv

23 from vllm.utils.platform_utils import num_compute_units

24 -from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache

24 +from vllm.utils.torch_utils import (

25 + async_tensor_h2d,

26 + get_dtype_size,

27 + is_quantized_kv_cache,

28 +)

29 from vllm.v1.attention.backend import (

30 AttentionBackend,

31 AttentionCGSupport,

@@ -593,6 +597,7 @@ class AiterFlashAttentionMetadataBuilder(

597 num_extends + 1,

598 dtype=torch.int32,

599 device=seq_lens_for_extend.device,

600 + pin_memory=True,

601 )

602 torch.cumsum(

603 swa_seqlen_for_extend,

@@ -623,8 +628,8 @@ class AiterFlashAttentionMetadataBuilder(

628

629 swa_metadata = AiterChunkSlidingWindowMetadata(

630 swa_cu_seqlens=cu_seq_lens.to(self.device, non_blocking=True),

626 - swa_seq_starts=seq_starts.to(self.device, non_blocking=True),

627 - swa_token_to_batch=token_to_seq.to(self.device, non_blocking=True),

631 + swa_seq_starts=async_tensor_h2d(seq_starts, self.device),

632 + swa_token_to_batch=async_tensor_h2d(token_to_seq, self.device),

633 swa_max_seqlens=max_seqlen_k,

634 swa_total_tokens=total_tokens,

635 swa_workspace=swa_workspace,

@@ -668,8 +673,8 @@ class AiterFlashAttentionMetadataBuilder(

673 chunk_context_metadata = AiterChunkContextMetadata(

674 workspace=self.extend_workspace,

675 cu_seq_lens_chunk=cu_seq_lens_cpu.to(self.device, non_blocking=True),

671 - chunk_starts=chunk_starts.to(self.device, non_blocking=True),

672 - token_to_batch=token_to_batch_tensor.to(self.device, non_blocking=True),

676 + chunk_starts=async_tensor_h2d(chunk_starts, self.device),

677 + token_to_batch=async_tensor_h2d(token_to_batch_tensor, self.device),

678 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

679 num_chunks=num_chunks,

680 total_token_per_batch=cu_seq_lens_cpu[:, -1].tolist(),

vllm/v1/worker/dp_utils.py

+3 -1

@@ -10,6 +10,7 @@ from vllm.config import ParallelConfig

10 from vllm.distributed.parallel_state import get_dp_group

11 from vllm.logger import init_logger

12 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

13 +from vllm.utils.torch_utils import PIN_MEMORY

14 from vllm.v1.worker.ubatch_utils import (

15 check_ubatch_thresholds,

16 is_last_ubatch_empty,

@@ -47,7 +48,8 @@ def _run_ar(

48 dp_rank = parallel_config.data_parallel_rank

49 device, group = _get_device_and_group(parallel_config)

50 # Populate this rank's contribution on CPU to reduce GPU syncs.

50 - tensor_cpu = torch.zeros(4, dp_size, dtype=torch.int32)

51 + pin_memory = PIN_MEMORY and torch.device(device).type == "cuda"

52 + tensor_cpu = torch.zeros(4, dp_size, dtype=torch.int32, pin_memory=pin_memory)

53 tensor_cpu[0][dp_rank] = orig_num_tokens_per_ubatch

54 tensor_cpu[1][dp_rank] = padded_num_tokens_per_ubatch

55 tensor_cpu[2][dp_rank] = 1 if should_ubatch else 0

vllm/v1/worker/gpu_model_runner.py

+5 -5

@@ -2258,8 +2258,10 @@ class GPUModelRunner(

2258 if self.use_async_spec_decode and (self.uses_mrope or self.uses_xdrope_dim > 0):

2259 drift = self.num_computed_tokens[req_indices_gpu].to(

2260 torch.int64

2261 - ) - self.input_batch.num_computed_tokens_cpu_tensor[req_indices].to(

2262 - device=self.device, dtype=torch.int64, non_blocking=True

2261 + ) - async_tensor_h2d(

2262 + self.input_batch.num_computed_tokens_cpu_tensor[req_indices],

2263 + device=self.device,

2264 + dtype=torch.int64,

2265 )

2266 target = self.mrope_positions if self.uses_mrope else self.xdrope_positions

2267 target.gpu[:, :total_num_scheduled_tokens] += drift

@@ -6324,9 +6326,7 @@ class GPUModelRunner(

6326 self.eplb_step(is_dummy=True, is_profile=is_profile)

6327

6328 logit_indices = np.cumsum(num_scheduled_tokens) - 1

6327 - logit_indices_device = torch.from_numpy(logit_indices).to(

6328 - self.device, non_blocking=True

6329 - )

6329 + logit_indices_device = async_tensor_h2d(logit_indices, device=self.device)

6330 return hidden_states, hidden_states[logit_indices_device]

6331

6332 @torch.inference_mode()