f97e502969 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/models/audioflamingo3.py

+2 -2

@@ -248,13 +248,13 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):

248 feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(

249 "audio", chunk_counts, dim=0

250 ),

251 - chunk_counts=MultiModalFieldConfig.batched("audio"),

251 + chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),

252 )

253 return dict(

254 audio_embeds=MultiModalFieldConfig.batched("audio"),

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

256 feature_attention_mask=MultiModalFieldConfig.batched("audio"),

257 - chunk_counts=MultiModalFieldConfig.batched("audio"),

257 + chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),

258 )

259

260

vllm/model_executor/models/diffusion_gemma.py

+20 -12

@@ -48,6 +48,7 @@ from vllm.model_executor.models.transformers.utils import recursive_replace_line

48 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix

49 from vllm.multimodal import MULTIMODAL_REGISTRY

50 from vllm.platforms import current_platform

51 +from vllm.utils.torch_utils import async_tensor_h2d

52 from vllm.v1.outputs import LogprobsTensors

53 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p

54 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata

@@ -728,10 +729,13 @@ class DiffusionGemmaRequestStates:

729 max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device

730 )

731

731 - def init_canvas(self, slot_indices_np: np.ndarray) -> None:

732 - """Initialize canvas with random tokens for the given slots."""

733 - n = slot_indices_np.shape[0]

734 - self.canvas[slot_indices_np] = torch.randint(

732 + def init_canvas(self, slot_indices: torch.Tensor) -> None:

733 + """Initialize canvas with random tokens for the given slots.

734 +

735 + `slot_indices` must already be on device to avoid a cpu->gpu sync.

736 + """

737 + n = slot_indices.shape[0]

738 + self.canvas[slot_indices] = torch.randint(

739 0,

740 self.vocab_size,

741 (n, self.canvas_length),

@@ -740,15 +744,15 @@ class DiffusionGemmaRequestStates:

744 )

745

746 def add_request(self, slot_idx: int) -> None:

743 - self.is_encoder_phase[slot_idx] = True

744 - self.init_canvas(torch.tensor([slot_idx], device=self.device))

745 - self.step[slot_idx] = 0

746 - self.accepted_canvas_history_len[slot_idx] = 0

747 + self.is_encoder_phase[slot_idx].fill_(True)

748 + self.init_canvas(async_tensor_h2d([slot_idx], device=self.device))

749 + self.step[slot_idx].fill_(0)

750 + self.accepted_canvas_history_len[slot_idx].fill_(0)

751 self.self_conditioning_embeds[slot_idx] = 0

752

753 def remove_request(self, slot_idx: int) -> None:

750 - self.is_encoder_phase[slot_idx] = False

751 - self.accepted_canvas_history_len[slot_idx] = 0

754 + self.is_encoder_phase[slot_idx].fill_(False)

755 + self.accepted_canvas_history_len[slot_idx].fill_(0)

756 self.self_conditioning_embeds[slot_idx] = 0

757

758

@@ -1148,11 +1152,15 @@ class DiffusionSampler:

1152 ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]]

1153 if len(ps) == 0:

1154 return

1151 - states.init_canvas(ps)

1152 - self.req_states.draft_tokens[ps, : self.canvas_length] = states.canvas[ps]

1155 + # Move the slot indices across once, up front: indexing a device

1156 + # tensor with a numpy array copies them over synchronously each time.

1157 ps_gpu = async_copy_to_gpu(

1158 ps.astype(np.int64), device=states.is_encoder_phase.device

1159 )

1160 + states.init_canvas(ps_gpu)

1161 + self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[

1162 + ps_gpu

1163 + ]

1164 states.is_encoder_phase.index_fill_(0, ps_gpu, False)

1165

1166 def _handle_prefill(

vllm/model_executor/models/gemma4_mm.py

+15 -1

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

69 )

70 from vllm.sequence import IntermediateTensors

71 from vllm.utils.tensor_schema import TensorSchema, TensorShape

72 +from vllm.utils.torch_utils import async_tensor_h2d

73

74 from .interfaces import (

75 MultiModalEmbeddings,

@@ -2039,7 +2040,20 @@ class Gemma4ForConditionalGeneration(

2040 ) -> torch.Tensor | None:

2041 logits = self.language_model.compute_logits(hidden_states)

2042 if logits is not None and self._suppress_token_ids:

2042 - logits[:, self._suppress_token_ids] = -float("inf")

2043 + # Cache a per-device index tensor for the (static) suppressed-token

2044 + # set and use `index_fill_`, so neither the Python-list indices nor

2045 + # the scalar fill value take a host roundtrip per call.

2046 + cache = getattr(self, "_suppress_token_ids_cache", None)

2047 + if cache is None:

2048 + cache = {}

2049 + self._suppress_token_ids_cache = cache

2050 + suppress_idx = cache.get(logits.device)

2051 + if suppress_idx is None:

2052 + suppress_idx = async_tensor_h2d(

2053 + self._suppress_token_ids, dtype=torch.long, device=logits.device

2054 + )

2055 + cache[logits.device] = suppress_idx

2056 + logits.index_fill_(1, suppress_idx, -float("inf"))

2057 return logits

2058

2059 # ------------------------------------------------------------------ #

vllm/model_executor/models/glm4_1v.py

+2 -2

@@ -548,9 +548,9 @@ class Glm4vVisionEmbeddings(nn.Module):

548 else:

549 # Convert inputs to tensors if needed

550 if isinstance(lengths, list):

551 - lengths = torch.tensor(lengths, device=device, dtype=torch.long)

551 + lengths = async_tensor_h2d(lengths, device=device, dtype=torch.long)

552 if not isinstance(image_shapes, torch.Tensor):

553 - image_shapes = torch.tensor(

553 + image_shapes = async_tensor_h2d(

554 image_shapes, device=device, dtype=torch.long

555 )

556

vllm/model_executor/models/interns1.py

+6 -2

@@ -426,12 +426,16 @@ class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo

426 ),

427 image_num_patches=MultiModalFieldConfig.batched("image"),

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

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

429 + image_token_id=MultiModalFieldConfig.shared(

430 + "image", num_images, keep_on_cpu=True

431 + ),

432 pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

433 "video", video_num_patches

434 ),

435 video_num_patches=MultiModalFieldConfig.batched("video"),

434 - video_token_id=MultiModalFieldConfig.shared("video", num_videos),

436 + video_token_id=MultiModalFieldConfig.shared(

437 + "video", num_videos, keep_on_cpu=True

438 + ),

439 )

440

441 def _get_prompt_updates(

vllm/model_executor/models/internvl.py

+3 -1

@@ -478,7 +478,9 @@ class InternVLMultiModalProcessor(

478 "video", video_num_patches

479 ),

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

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

481 + video_token_id=MultiModalFieldConfig.shared(

482 + "video", num_videos, keep_on_cpu=True

483 + ),

484 )

485

486 def _get_mm_fields_config(

vllm/model_executor/models/keye.py

+5 -4

@@ -871,12 +871,12 @@ def _keye_field_config(

871 return dict(

872 pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),

873 image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),

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

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

875 pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

876 "video", video_grid_sizes

877 ),

878 video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),

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

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

880 )

881

882

@@ -1291,8 +1291,9 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):

1291 image_grid_thw = image_input["image_grid_thw"]

1292 assert image_grid_thw.ndim == 2

1293

1294 - for idx, thaw in enumerate(image_grid_thw):

1295 - thw_tuple = tuple(thaw.detach().cpu().numpy().tolist())

1294 + image_grid_thw_list = image_grid_thw.tolist()

1295 + for idx, thw in enumerate(image_grid_thw_list):

1296 + thw_tuple = tuple(thw)

1297 numel = np.prod(thw_tuple)

1298 image_grid_hws.append(thw_tuple)

1299 image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])

vllm/model_executor/models/phi3v.py

+3 -2

@@ -266,9 +266,10 @@ class Phi3HDImageEmbedding(nn.Module):

266 )

267

268 batch_image_features_proj = []

269 + image_sizes_list = image_sizes.tolist()

270 # need a for loop to process each image because of different image sizes

271 # (patch arrangement is different for each image)

271 - for i, img_size in enumerate(image_sizes):

272 + for i, img_size in enumerate(image_sizes_list):

273 h, w = img_size

274 h_crop = h // 336

275 w_crop = w // 336

@@ -426,7 +427,7 @@ class Phi3VMultiModalProcessor(BaseMultiModalProcessor[Phi3VProcessingInfo]):

427 ) -> Mapping[str, MultiModalFieldConfig]:

428 return dict(

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

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

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

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

432 )

433

vllm/model_executor/models/phi4mm.py

+1 -2

@@ -270,8 +270,7 @@ class Phi4MMImageEncoder(nn.Module):

270 pixel_values = pixel_values.flatten(0, 1)

271

272 img_features = self.get_img_features(

273 - pixel_values,

274 - image_attention_mask.type(torch.BoolTensor).flatten(0, 1).to(target_device),

273 + pixel_values, image_attention_mask.bool().flatten(0, 1)

274 )

275

276 base_feat_height_target = self.base_feat_height_target

vllm/model_executor/models/step3_vl.py

+3 -2

@@ -200,7 +200,7 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo])

200 patch_pixel_values=MultiModalFieldConfig.flat_from_sizes(

201 "image", num_patches

202 ),

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

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

204 patch_newline_mask=MultiModalFieldConfig.flat_from_sizes(

205 "image", num_patches

206 ),

@@ -682,7 +682,8 @@ class Step3VLForConditionalGeneration(

682

683 merged_image_features = []

684 cur_patch_idx = 0

685 - for i, num_patch in enumerate(num_patches):

685 + num_patches_list = num_patches.tolist()

686 + for i, num_patch in enumerate(num_patches_list):

687 cur_feature = []

688 if num_patch > 0:

689 patch_slice = patch_image_features[

vllm/v1/attention/backend.py

+2

@@ -442,6 +442,7 @@ class CommonAttentionMetadata:

442 # Needed by FastPrefillAttentionBuilder

443 logits_indices_padded: torch.Tensor | None = None

444 num_logits_indices: int | None = None

445 + max_logits_per_req: int | None = None

446

447 # Needed by CrossAttentionBuilder

448 encoder_seq_lens: torch.Tensor | None = None

@@ -590,6 +591,7 @@ class CommonAttentionMetadata:

591 else self.causal,

592 logits_indices_padded=self.logits_indices_padded,

593 num_logits_indices=self.num_logits_indices,

594 + max_logits_per_req=self.max_logits_per_req,

595 encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),

596 encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),

597 dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),

vllm/v1/attention/backends/mamba2_attn.py

+9 -5

@@ -150,11 +150,15 @@ class Mamba2AttentionMetadataBuilder(

150

151 # Compute seq_idx for prefill only

152 if common.num_prefills > 0:

153 - prep_initial_states = (

154 - torch.any(common.has_initial_states_p).item()

155 - if common.has_initial_states_p is not None

156 - else False

157 - )

153 + prep_initial_states = False

154 + if common.has_initial_states_p is not None:

155 + # Same condition as `has_initial_states_p`, but derived from CPU

156 + # data so it needs no D2H. `seq_lens_cpu_upper_bound` is precise

157 + # for prefill rows, which is all this slice covers.

158 + num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata(

159 + common, common_attn_metadata

160 + )

161 + prep_initial_states = bool((num_computed_tokens_p_cpu > 0).any())

162

163 cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p = (

164 self._build_chunk_metadata_tensors(

vllm/v1/attention/backends/mamba_attn.py

+29 -15

@@ -324,34 +324,48 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

324

325 return cu_chunk_seqlen, seq_idx, last_chunk_indices

326

327 - def _build_chunk_metadata_tensors(

327 + def _prefill_cpu_metadata(

328 self,

329 - chunk_size: int,

329 common: M,

330 common_attn_metadata: CommonAttentionMetadata,

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

333 - """

334 - Compute chunk metadata and return as device tensors.

335 - Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p).

336 - """

337 - num_reqs = common.num_reqs

338 - num_prefills = common.num_prefills

339 - num_decode_tokens = common.num_decode_tokens

331 + ) -> tuple[torch.Tensor, torch.Tensor]:

332 + """Prefill context lengths and query offsets, from CPU data only.

333

341 - # Derive prefill context lengths from CPU data only.

342 - # `seq_lens_cpu_upper_bound` is precise for prefill rows in all modes

343 - # (including async spec decode), so this avoids the D2H sync that

344 - # `compute_num_computed_tokens().cpu()` would force.

334 + `seq_lens_cpu_upper_bound` is precise for prefill rows in all modes

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

336 + `compute_num_computed_tokens().cpu()` would force.

337 +

338 + Returns (num_computed_tokens_p_cpu, query_start_loc_p_cpu).

339 + """

340 seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

341 assert seq_lens_cpu is not None

342 + num_reqs = common.num_reqs

343 + num_prefills = common.num_prefills

344 query_start_loc_p_cpu = (

345 common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]

349 - - num_decode_tokens

346 + - common.num_decode_tokens

347 )

348 prefill_query_lens_cpu = query_start_loc_p_cpu[1:] - query_start_loc_p_cpu[:-1]

349 num_computed_tokens_p_cpu = (

350 seq_lens_cpu[num_reqs - num_prefills : num_reqs] - prefill_query_lens_cpu

351 )

352 + return num_computed_tokens_p_cpu, query_start_loc_p_cpu

353 +

354 + def _build_chunk_metadata_tensors(

355 + self,

356 + chunk_size: int,

357 + common: M,

358 + common_attn_metadata: CommonAttentionMetadata,

359 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

360 + """

361 + Compute chunk metadata and return as device tensors.

362 + Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p).

363 + """

364 + num_prefills = common.num_prefills

365 +

366 + num_computed_tokens_p_cpu, query_start_loc_p_cpu = self._prefill_cpu_metadata(

367 + common, common_attn_metadata

368 + )

369

370 cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(

371 chunk_size,

vllm/v1/attention/backends/utils.py

+8 -2

@@ -539,8 +539,14 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(

539

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

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

542 - decode_max_query_len = int(num_decode_tokens.max().item())

543 - total_num_decode_tokens = int(num_decode_tokens.sum().item())

542 +

543 + # `num_decode_tokens` is a histogram over `logits_indices`, so its total is

544 + # just how many there were -- already known as a Python int.

545 + total_num_decode_tokens = num_logits_indices

546 +

547 + # Largest per-request logits count.

548 + decode_max_query_len = common_attn_metadata.max_logits_per_req

549 + assert decode_max_query_len is not None

550

551 common_attn_metadata = CommonAttentionMetadata(

552 query_start_loc=decode_query_start_loc,

vllm/v1/worker/gpu_model_runner.py

+8 -4

@@ -2009,10 +2009,11 @@ class GPUModelRunner(

2009 ) -> tuple[

2010 torch.Tensor,

2011 SpecDecodeMetadata | None,

2012 + int,

2013 ]:

2014 """

2015 Returns:

2015 - tuple[logits_indices, spec_decode_metadata]

2016 + tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens]

2017 """

2018 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2019 assert total_num_scheduled_tokens > 0

@@ -2324,6 +2325,7 @@ class GPUModelRunner(

2325 return (

2326 logits_indices,

2327 spec_decode_metadata,

2328 + int(num_sampled_tokens.max()),

2329 )

2330

2331 def _build_attention_metadata(

@@ -2335,6 +2337,7 @@ class GPUModelRunner(

2337 num_reqs_padded: int | None = None,

2338 ubatch_slices: UBatchSlices | None = None,

2339 logits_indices: torch.Tensor | None = None,

2340 + max_num_sampled_tokens: int | None = None,

2341 use_spec_decode: bool = False,

2342 for_cudagraph_capture: bool = False,

2343 num_scheduled_tokens: dict[str, int] | None = None,

@@ -2509,6 +2512,7 @@ class GPUModelRunner(

2512

2513 if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill:

2514 cm_base.num_logits_indices = logits_indices.size(0)

2515 + cm_base.max_logits_per_req = max_num_sampled_tokens

2516 cm_base.logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

2517 logits_indices

2518 )

@@ -4341,9 +4345,8 @@ class GPUModelRunner(

4345 max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())

4346 num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

4347

4344 - logits_indices, spec_decode_metadata = self._prepare_inputs(

4345 - scheduler_output,

4346 - num_scheduled_tokens_np,

4348 + logits_indices, spec_decode_metadata, max_num_sampled_tokens = (

4349 + self._prepare_inputs(scheduler_output, num_scheduled_tokens_np)

4350 )

4351

4352 cascade_attn_prefix_lens = None

@@ -4481,6 +4484,7 @@ class GPUModelRunner(

4484 max_query_len=max_num_scheduled_tokens,

4485 ubatch_slices=ubatch_slices_attn,

4486 logits_indices=logits_indices,

4487 + max_num_sampled_tokens=max_num_sampled_tokens,

4488 use_spec_decode=use_spec_decode,

4489 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,

4490 cascade_attn_prefix_lens=cascade_attn_prefix_lens,

[Perf] Avoid more GPU<->CPU syncs on the model execution path (#51738)