@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 989c176c0a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 03:45 UTC

[Perf][3/n] Eliminate GPU<->CPU syncs in attention impls (#41434)

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

Nick Hill committed May 7, 2026 at 19:44 UTC 989c176c0a14e1adc5a9ba33cb5c3a39fceec3d3

10 files changed +169 -68

vllm/utils/torch_utils.py

+9 -3

@@ -3,6 +3,7 @@

3 import contextlib

4 import importlib.metadata

5 import os

6 +import platform

7 import random

8 import threading

9 from collections.abc import Callable, Collection

@@ -67,6 +68,11 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {

68 T = TypeVar("T")

71 +# Pin memory in non-WSL case.

72 +# Logic duplicated here for now to avoid circular import.

73 +PIN_MEMORY = "microsoft" not in " ".join(platform.uname()).lower()

74 +

75 +

76 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:

77 return (

78 kv_cache_dtype.startswith("fp8")

@@ -602,12 +608,12 @@ def create_kv_caches_with_random(

608 def async_tensor_h2d(

609 data: list,

610 dtype: torch.dtype,

605 - target_device: str | torch.device,

606 - pin_memory: bool,

611 + device: str | torch.device,

612 + pin_memory: bool = PIN_MEMORY,

613 ) -> torch.Tensor:

614 """Asynchronously create a tensor and copy it from host to device."""

615 t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")

610 - return t.to(device=target_device, non_blocking=True)

616 + return t.to(device=device, non_blocking=True)

617

618

619 def make_ndarray_with_pad(

vllm/v1/attention/backends/flashinfer.py

+2 -1

@@ -1105,7 +1105,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1105 paged_kv_indptr_prefill_gpu = self.paged_kv_indptr.gpu[

1106 prefill_start : num_reqs + 1

1107 ]

1108 - paged_kv_indptr_prefill_gpu[0] = 0

1108 + # Assign to slice to avoid cpu sync.

1109 + paged_kv_indptr_prefill_gpu[:1] = 0

1110 torch.cumsum(

1111 num_blocks_per_req,

1112 dim=0,

vllm/v1/attention/backends/flex_attention.py

+30 -18

@@ -48,12 +48,16 @@ create_block_mask_compiled = torch.compile(

48 flex_attention_compiled = torch.compile(flex_attention, fullgraph=True)

51 -def _offsets_to_doc_ids_tensor(offsets: torch.Tensor) -> torch.Tensor:

52 - device = offsets.device

53 - counts = offsets[1:] - offsets[:-1]

54 - return torch.repeat_interleave(

55 - torch.arange(len(counts), device=device, dtype=torch.int32), counts

51 +def _offsets_to_doc_ids_tensor(

52 + offsets_cpu: torch.Tensor, device: torch.device

53 +) -> torch.Tensor:

54 + # Build on CPU (so `repeat_interleave` doesn't force a GPU->CPU sync to

55 + # learn the data-dependent output length) and upload non-blocking.

56 + counts = offsets_cpu[1:] - offsets_cpu[:-1]

57 + doc_ids = torch.repeat_interleave(

58 + torch.arange(len(counts), dtype=torch.int32), counts

59 )

60 + return doc_ids.to(device, non_blocking=True)

63 def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):

@@ -290,11 +294,13 @@ def unique_static_unsorted(

294 keep = (x_flat != ignored_val) & (idx == first_idx.gather(1, x_flat)) # [B, N]

295

296 # ── left-pack uniques into a fresh tensor ───────────────────────────

297 + # Route non-kept entries to a garbage slot at column N so we can do a

298 + # single scatter rather than using torch.nonzero (which would force a

299 + # GPU->CPU sync to enumerate kept positions).

300 dest_pos = torch.cumsum(keep.to(torch.long), dim=1) - 1 # where to go

294 - packed_flat = torch.full_like(x_flat, pad_val)

295 -

296 - rows, src_cols = torch.nonzero(keep, as_tuple=True)

297 - packed_flat[rows, dest_pos[rows, src_cols]] = x_flat[rows, src_cols]

301 + dest_pos = torch.where(keep, dest_pos, N)

302 + packed_extended = torch.full((B, N + 1), pad_val, device=device, dtype=x_flat.dtype)

303 + packed_flat = packed_extended.scatter_(1, dest_pos, x_flat)[:, :N]

304

305 # ── restore original layout ─────────────────────────────────────────

306 packed = packed_flat.reshape(x_perm.shape).movedim(-1, dim)

@@ -346,6 +352,9 @@ class FlexAttentionMetadata:

352 num_actual_tokens: int # Number of tokens excluding padding.

353 max_query_len: int

354 query_start_loc: torch.Tensor

355 + # CPU-resident copy of query_start_loc used to derive doc_ids without a

356 + # GPU->CPU sync from repeat_interleave's data-dependent output size.

357 + query_start_loc_cpu: torch.Tensor

358 max_seq_len: int

359 seq_lens: torch.Tensor

360 block_table: torch.Tensor

@@ -452,12 +461,7 @@ class FlexAttentionMetadata:

461 (is_valid, logical_q_idx, logical_kv_idx) = (

462 self._convert_physical_to_logical(self.doc_ids, q_idx, physical_kv_idx)

463 )

455 - # Apply mask modification only for valid indices

456 - return torch.where(

457 - is_valid,

458 - self.logical_mask_mod(b, h, logical_q_idx, logical_kv_idx),

459 - False,

460 - )

464 + return is_valid & self.logical_mask_mod(b, h, logical_q_idx, logical_kv_idx)

465

466 return final_mask_mod

467

@@ -469,7 +473,9 @@ class FlexAttentionMetadata:

473 packed query sequences.

474 """

475 # Create a lookup mapping from query indices -> request number

472 - request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)

476 + request_lookup = _offsets_to_doc_ids_tensor(

477 + self.query_start_loc_cpu, self.query_start_loc.device

478 + )

479

480 def final_mask_mod(

481 b: torch.Tensor,

@@ -581,7 +587,9 @@ class FlexAttentionMetadata:

587 return None

588

589 # Create a lookup mapping from query indices -> request number

584 - request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)

590 + request_lookup = _offsets_to_doc_ids_tensor(

591 + self.query_start_loc_cpu, self.query_start_loc.device

592 + )

593 user_score_mod = self.score_mod

594

595 def transformed_score_mod(

@@ -726,7 +734,9 @@ class FlexAttentionMetadata:

734 assert self.prefix_kv_lens is None, "Not implemented yet."

735 assert self.suffix_kv_lens is None, "Not implemented yet."

736 # Create a lookup mapping from query indices -> request number

729 - self.doc_ids = _offsets_to_doc_ids_tensor(self.query_start_loc)

737 + self.doc_ids = _offsets_to_doc_ids_tensor(

738 + self.query_start_loc_cpu, self.query_start_loc.device

739 + )

740 self.doc_ids = copy_to_persistent(self.persistent_doc_ids, self.doc_ids)

741 self.num_blocks = self.total_cache_tokens // self.block_size

742

@@ -807,6 +817,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

817

818 max_seq_len = common_attn_metadata.max_seq_len

819 query_start_loc = common_attn_metadata.query_start_loc

820 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

821 seq_lens = common_attn_metadata.seq_lens

822 block_table_tensor = common_attn_metadata.block_table_tensor

823 slot_mapping = common_attn_metadata.slot_mapping

@@ -871,6 +882,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

882 num_actual_tokens=num_actual_tokens,

883 max_query_len=max_query_len,

884 query_start_loc=query_start_loc,

885 + query_start_loc_cpu=query_start_loc_cpu,

886 max_seq_len=max_seq_len,

887 seq_lens=seq_lens,

888 block_table=block_table_tensor,

vllm/v1/attention/backends/mamba_attn.py

+18 -19

@@ -9,6 +9,7 @@ import torch

10 from vllm.config import VllmConfig

11 from vllm.utils.math_utils import cdiv

12 +from vllm.utils.torch_utils import async_tensor_h2d

13 from vllm.v1.attention.backend import (

14 AttentionCGSupport,

15 AttentionMetadataBuilder,

@@ -270,16 +271,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

271 num_prefills = common.num_prefills

272 num_decode_tokens = common.num_decode_tokens

273

273 - num_computed_tokens_cpu = (

274 - common_attn_metadata.compute_num_computed_tokens().cpu()

275 - )

276 - num_computed_tokens_p_cpu = num_computed_tokens_cpu[

277 - num_reqs - num_prefills : num_reqs

278 - ]

274 + # Derive prefill context lengths from CPU data only.

275 + # `seq_lens_cpu_upper_bound` is precise for prefill rows in all modes

276 + # (including async spec decode), so this avoids the D2H sync that

277 + # `compute_num_computed_tokens().cpu()` would force.

278 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

279 + assert seq_lens_cpu is not None

280 query_start_loc_p_cpu = (

281 common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]

282 - num_decode_tokens

283 )

284 + prefill_query_lens_cpu = query_start_loc_p_cpu[1:] - query_start_loc_p_cpu[:-1]

285 + num_computed_tokens_p_cpu = (

286 + seq_lens_cpu[num_reqs - num_prefills : num_reqs] - prefill_query_lens_cpu

287 + )

288

289 cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(

290 chunk_size,

@@ -289,20 +294,14 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

294 )

295

296 device = common_attn_metadata.query_start_loc.device

292 - cu_chunk_seqlen_p = torch.as_tensor(

293 - cu_chunk_seqlen,

294 - device=device,

295 - dtype=torch.int32,

296 - )

297 - seq_idx_p = torch.as_tensor(

298 - seq_idx,

299 - device=device,

300 - dtype=torch.int32,

297 + # Build on pinned CPU and upload non-blocking to avoid the synchronous

298 + # H2D copy that `torch.as_tensor(list, device=cuda)` would force.

299 + cu_chunk_seqlen_p = async_tensor_h2d(

300 + cu_chunk_seqlen, dtype=torch.int32, device=device

301 )

302 - last_chunk_indices_p = torch.as_tensor(

303 - last_chunk_indices,

304 - device=device,

305 - dtype=torch.int32,

302 + seq_idx_p = async_tensor_h2d(seq_idx, dtype=torch.int32, device=device)

303 + last_chunk_indices_p = async_tensor_h2d(

304 + last_chunk_indices, dtype=torch.int32, device=device

305 )

306 return cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p

307

vllm/v1/attention/backends/tree_attn.py

+52 -6

@@ -90,6 +90,14 @@ class TreeAttentionMetadata:

90 num_prefills: int = 0

91 num_decodes: int = 0

93 + # Precomputed (on CPU in the builder) max_query_len and max_seq_len for

94 + # the prefill-only and decode-only sub-batches. Used by the properties

95 + # below to avoid a GPU->CPU sync via `.max().item()` on every forward.

96 + max_query_len_prefill: int = 0

97 + max_seq_len_prefill: int = 0

98 + max_query_len_decode: int = 0

99 + max_seq_len_decode: int = 0

100 +

101 tree_attn_bias: torch.Tensor | None = None

102

103 # Cached Prefill/decode metadata.

@@ -107,14 +115,13 @@ class TreeAttentionMetadata:

115 return self._cached_prefill_metadata

116

117 q_start_loc = self.query_start_loc[self.num_decodes :]

110 - q_seqlens = torch.diff(q_start_loc)

118 kv_seqlens = self.seq_lens[self.num_decodes :]

119 # Construct & cache prefill-phase attention metadata structure

120 self._cached_prefill_metadata = TreeAttentionMetadata(

121 num_actual_tokens=self.num_prefill_tokens,

115 - max_query_len=int(q_seqlens.max().item()),

122 + max_query_len=self.max_query_len_prefill,

123 query_start_loc=q_start_loc - q_start_loc[0],

117 - max_seq_len=int(kv_seqlens.max().item()),

124 + max_seq_len=self.max_seq_len_prefill,

125 seq_lens=kv_seqlens,

126 block_table=self.block_table[self.num_decodes :],

127 slot_mapping=self.slot_mapping[self.num_decode_tokens :],

@@ -132,14 +139,13 @@ class TreeAttentionMetadata:

139 return self._cached_decode_metadata

140

141 q_start_loc = self.query_start_loc[: self.num_decodes + 1]

135 - q_seqlens = torch.diff(q_start_loc)

142 kv_seqlens = self.seq_lens[: self.num_decodes]

143 # Construct & cache decode-phase attention metadata structure

144 self._cached_decode_metadata = TreeAttentionMetadata(

145 num_actual_tokens=self.num_decode_tokens,

140 - max_query_len=int(q_seqlens.max().item()),

146 + max_query_len=self.max_query_len_decode,

147 query_start_loc=q_start_loc,

142 - max_seq_len=int(kv_seqlens.max().item()),

148 + max_seq_len=self.max_seq_len_decode,

149 seq_lens=kv_seqlens,

150 block_table=self.block_table[: self.num_decodes],

151 slot_mapping=self.slot_mapping[: self.num_decode_tokens],

@@ -199,6 +205,42 @@ class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadat

205 block_table = common_attn_metadata.block_table_tensor

206 slot_mapping = common_attn_metadata.slot_mapping

207

208 + # Precompute prefill/decode sub-batch max_query_len / max_seq_len on

209 + # CPU so the prefill_metadata / decode_metadata properties don't need

210 + # a GPU->CPU sync via `.max().item()` on every forward.

211 + # Prefer `seq_lens_cpu_upper_bound` over the (deprecated)

212 + # `seq_lens_cpu` property: the upper bound is precise for prefill

213 + # rows and optimistic-but-safe for decode rows (workspace sizing

214 + # from `max()` is fine with an over-estimate), and avoids the

215 + # `seq_lens.to("cpu")` sync the property would fall through to in

216 + # async-spec-decode mode. The draft-attention path (eagle

217 + # speculator) doesn't populate it; fall back to the batch-wide

218 + # `max_seq_len` as a safe upper bound for both sub-batches.

219 + q_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

220 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

221 + if num_prefills > 0:

222 + q_seqlens_p = torch.diff(q_start_loc_cpu[num_decodes:])

223 + max_query_len_prefill = int(q_seqlens_p.max())

224 + max_seq_len_prefill = (

225 + int(seq_lens_cpu[num_decodes:].max())

226 + if seq_lens_cpu is not None

227 + else max_seq_len

228 + )

229 + else:

230 + max_query_len_prefill = 0

231 + max_seq_len_prefill = 0

232 + if num_decodes > 0:

233 + q_seqlens_d = torch.diff(q_start_loc_cpu[: num_decodes + 1])

234 + max_query_len_decode = int(q_seqlens_d.max())

235 + max_seq_len_decode = (

236 + int(seq_lens_cpu[:num_decodes].max())

237 + if seq_lens_cpu is not None

238 + else max_seq_len

239 + )

240 + else:

241 + max_query_len_decode = 0

242 + max_seq_len_decode = 0

243 +

244 return TreeAttentionMetadata(

245 num_actual_tokens=num_actual_tokens,

246 num_prefill_tokens=num_prefill_tokens,

@@ -211,6 +253,10 @@ class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadat

253 seq_lens=kv_seqlens,

254 block_table=block_table,

255 slot_mapping=slot_mapping,

256 + max_query_len_prefill=max_query_len_prefill,

257 + max_seq_len_prefill=max_seq_len_prefill,

258 + max_query_len_decode=max_query_len_decode,

259 + max_seq_len_decode=max_seq_len_decode,

260 tree_attn_bias=self.tree_attn_bias,

261 )

262

vllm/v1/attention/backends/triton_attn.py

+4 -5

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

18 from vllm.platforms import current_platform

19 from vllm.platforms.interface import DeviceCapability

20 from vllm.utils.math_utils import next_power_of_2

21 -from vllm.utils.torch_utils import is_quantized_kv_cache

21 +from vllm.utils.torch_utils import async_tensor_h2d, is_quantized_kv_cache

22 from vllm.v1.attention.backend import (

23 AttentionBackend,

24 AttentionCGSupport,

@@ -117,10 +117,9 @@ class TritonAttentionMetadata:

117 for r in range_lists:

118 padded_r = list(r) + [(0, 0)] * (max_ranges - len(r))

119 padded.append(padded_r)

120 - # Create tensor with efficient H2D transfer

121 - return torch.tensor(padded, dtype=torch.int32, device=device).view(

122 - num_seqs, max_ranges, 2

123 - )

120 + # Build on pinned CPU memory so the H2D transfer is non-blocking.

121 + padded = async_tensor_h2d(padded, dtype=torch.int32, device=device)

122 + return padded.view(num_seqs, max_ranges, 2)

123

124

125 class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):

vllm/v1/attention/backends/turboquant_attn.py

+37 -9

@@ -187,6 +187,10 @@ class TurboQuantMetadata(AttentionMetadata):

187 is_prefill: bool = False

188 num_decodes: int = 0 # number of decode requests (first in batch)

189 num_decode_tokens: int = 0 # tokens from decode requests

190 + # CPU-resident copies used by the prefill path for per-request iteration

191 + # without per-step D2H syncs.

192 + query_start_loc_cpu: torch.Tensor | None = None

193 + seq_lens_cpu: torch.Tensor | None = None

194

195

196 class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):

@@ -230,6 +234,8 @@ class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):

234 is_prefill=(cam.max_query_len > 1),

235 num_decodes=num_decodes,

236 num_decode_tokens=num_decode_tokens,

237 + query_start_loc_cpu=cam.query_start_loc_cpu,

238 + seq_lens_cpu=cam.seq_lens_cpu_upper_bound,

239 )

240

241

@@ -474,11 +480,21 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

480 # first-chunk prefills. Using full-batch max_seq_len breaks

481 # this because decode requests inflate max_seq_len.

482 prefill_seq_lens = attn_metadata.seq_lens[num_decodes:]

477 - # Use CPU-side max to avoid GPU→CPU sync from .item()

478 - prefill_max_seq = max(attn_metadata.seq_lens[num_decodes:].tolist())

483 + # Use the CPU-resident `seq_lens` upper-bound from the metadata

484 + # (populated in the builder) to compute the prefill sub-batch

485 + # max without a GPU→CPU sync.

486 + if attn_metadata.seq_lens_cpu is not None:

487 + prefill_max_seq = int(attn_metadata.seq_lens_cpu[num_decodes:].max())

488 + else:

489 + prefill_max_seq = attn_metadata.max_seq_len

490 prefill_qsl = (

491 attn_metadata.query_start_loc[num_decodes:] - num_decode_tokens

492 )

493 + prefill_qsl_cpu = None

494 + if attn_metadata.query_start_loc_cpu is not None:

495 + prefill_qsl_cpu = (

496 + attn_metadata.query_start_loc_cpu[num_decodes:] - num_decode_tokens

497 + )

498 prefill_meta = TurboQuantMetadata(

499 seq_lens=prefill_seq_lens,

500 slot_mapping=attn_metadata.slot_mapping[num_decode_tokens:N],

@@ -488,6 +504,10 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

504 max_query_len=attn_metadata.max_query_len,

505 max_seq_len=prefill_max_seq,

506 is_prefill=True,

507 + query_start_loc_cpu=prefill_qsl_cpu,

508 + seq_lens_cpu=attn_metadata.seq_lens_cpu[num_decodes:]

509 + if attn_metadata.seq_lens_cpu is not None

510 + else None,

511 )

512 k = key[:N].view(N, self.num_kv_heads, self.head_size)

513 v = value[:N].view(N, self.num_kv_heads, self.head_size)

@@ -578,10 +598,16 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

598

599 output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype)

600

581 - # Convert to Python lists once (single CPU-GPU sync) instead of

582 - # per-request .item() calls that each force a sync.

583 - qsl = query_start_loc.tolist()

584 - seq_lens_list = attn_metadata.seq_lens.tolist()

601 + # Prefer the CPU-resident copies from the metadata if populated —

602 + # otherwise `.tolist()` on GPU tensors forces a synchronizing copy.

603 + if attn_metadata.query_start_loc_cpu is not None:

604 + qsl = attn_metadata.query_start_loc_cpu.tolist()

605 + else:

606 + qsl = query_start_loc.tolist()

607 + if attn_metadata.seq_lens_cpu is not None:

608 + seq_lens_list = attn_metadata.seq_lens_cpu.tolist()

609 + else:

610 + seq_lens_list = attn_metadata.seq_lens.tolist()

611

612 # Pre-allocate cu_seqlens for single-request flash_attn calls

613 # to avoid per-request host→device tensor creation.

@@ -612,7 +638,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

638 if q_len == seq_len:

639 # First-chunk prefill: all K/V are in the current batch.

640 if _HAS_FLASH_ATTN:

615 - self._cu_2[1] = q_len

641 + # Assign to slice to avoid gpu/cpu sync.

642 + self._cu_2[1:2] = q_len

643 cu = self._cu_2

644 out = self._flash_attn_varlen(

645 q=q_seq,

@@ -791,8 +818,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

818 if not hasattr(self, "_cu_2_q"):

819 self._cu_2_q = torch.zeros(2, device=device, dtype=torch.int32)

820 self._cu_2_k = torch.zeros(2, device=device, dtype=torch.int32)

794 - self._cu_2_q[1] = q_len

795 - self._cu_2_k[1] = seq_len

821 + # Assigning to slice uses fill_ which avoids cpu/gpu sync.

822 + self._cu_2_q[1:2] = q_len

823 + self._cu_2_k[1:2] = seq_len

824 cu_seqlens_q = self._cu_2_q

825 cu_seqlens_k = self._cu_2_k

826 return self._flash_attn_varlen(

vllm/v1/attention/backends/utils.py

+15 -4

@@ -332,8 +332,10 @@ def make_local_attention_virtual_batches(

332 # regression when using numpy arrays (batch and block indices) to index into

333 # torch tensor (block_table). As a workaround, convert numpy arrays to torch

334 # tensor first, which recovers perf.

335 - batch_indices_torch = torch.from_numpy(batch_indices)

336 - block_indices_torch = torch.from_numpy(block_indices)

335 + # Upload the index tensors to the block_table's device up-front so that the

336 + # fancy indexing below doesn't implicitly force a synchronous H2D copy.

337 + batch_indices_torch = torch.from_numpy(batch_indices).to(device, non_blocking=True)

338 + block_indices_torch = torch.from_numpy(block_indices).to(device, non_blocking=True)

339

340 # Save as a lambda so we can return this for update_block_table

341 make_block_table = lambda block_table: block_table[

@@ -391,7 +393,16 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(

393

394 # Figure out how many tokens are in each request

395 # num_decode_tokens: [1, 2, 1]

394 - num_decode_tokens = torch.bincount(request_ids, minlength=num_reqs)

396 + # Avoid `torch.bincount` here — on CUDA it forces a sync to determine

397 + # the output size (even with `minlength`, the kernel must confirm no

398 + # value exceeds the bound). `scatter_add_` into a preallocated buffer

399 + # is equivalent and stays async.

400 + num_decode_tokens = torch.zeros(

401 + num_reqs, dtype=request_ids.dtype, device=request_ids.device

402 + )

403 + num_decode_tokens.scatter_add_(

404 + 0, request_ids.to(num_decode_tokens.dtype), torch.ones_like(request_ids)

405 + )

406

407 # Calculate new query_start_loc with tokens in generation_indices

408 # decode_query_start_loc: [0, 1, 3, 4]

@@ -399,7 +410,7 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(

410 num_reqs + 1, device=query_start_loc.device, dtype=query_start_loc.dtype

411 )

412

402 - decode_query_start_loc[0] = 0

413 + decode_query_start_loc[:1].fill_(0) # Avoid sync from scalar assignment.

414 decode_query_start_loc[1:] = torch.cumsum(num_decode_tokens, dim=0)

415 decode_max_query_len = int(num_decode_tokens.max().item())

416 total_num_decode_tokens = int(num_decode_tokens.sum().item())

vllm/v1/worker/gpu/buffer_utils.py

+1 -1

@@ -167,7 +167,7 @@ class StagedWriteTensor:

167

168 # Special handling for write_contents

169 write_contents = async_tensor_h2d(

170 - self._staged_write_contents, self.dtype, self.device, pin_memory=True

170 + self._staged_write_contents, self.dtype, self.device

171 )

172

173 # Write diffs to the GPU buffer

vllm/v1/worker/gpu/sample/penalties.py

+1 -2

@@ -58,8 +58,7 @@ class PenaltiesState:

58 idx_mapping = async_tensor_h2d(

59 self._new_penalties_reqs,

60 dtype=torch.int32,

61 - target_device=self.device,

62 - pin_memory=True,

61 + device=self.device,

62 )

64 prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs]