diff --git a/vllm/model_executor/layers/pooler/seqwise/methods.py b/vllm/model_executor/layers/pooler/seqwise/methods.py index b967ff4ede7b..82170b5fbdc4 100644 --- a/vllm/model_executor/layers/pooler/seqwise/methods.py +++ b/vllm/model_executor/layers/pooler/seqwise/methods.py @@ -68,21 +68,23 @@ def forward( "partial prefill not supported with MEAN pooling" ) - prompt_lens = pooling_cursor.prompt_lens_cpu.to( - hidden_states.device, dtype=torch.int64, non_blocking=True - ) - - num_seqs = prompt_lens.numel() + prompt_lens_cpu = pooling_cursor.prompt_lens_cpu + num_seqs = prompt_lens_cpu.numel() hidden_size = hidden_states.shape[-1] if num_seqs == 0: # early return for empty batch return hidden_states.new_empty((0, hidden_size), dtype=torch.float32) - # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] + # Build segment_ids on CPU so repeat_interleave doesn't need to sync + # GPU->CPU to learn its data-dependent output length, then upload + # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] segment_ids = torch.repeat_interleave( - torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long), - prompt_lens, + torch.arange(num_seqs, dtype=torch.long), + prompt_lens_cpu, + ).to(hidden_states.device, non_blocking=True) + prompt_lens = prompt_lens_cpu.to( + hidden_states.device, dtype=torch.int64, non_blocking=True ) segment_sums = torch.zeros( (num_seqs, hidden_size), diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index d06663b5b947..ae5926cd62ff 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses from collections.abc import Mapping, Set from itertools import groupby @@ -80,9 +81,11 @@ def forward( pooling_metadata: PoolingMetadata, ) -> PoolerOutput: poolers_by_task = self.poolers_by_task + cursor = pooling_metadata.pooling_cursor outputs = list[torch.Tensor | None]() offset = 0 + token_offset = 0 for task, group in groupby(pooling_metadata.tasks): if not (pooler := poolers_by_task.get(task)): raise ValueError( @@ -91,10 +94,37 @@ def forward( ) num_items = len(list(group)) - group_output: PoolerOutput = pooler( - hidden_states, - pooling_metadata[offset : offset + num_items], - ) + group_metadata = pooling_metadata[offset : offset + num_items] + if cursor is None: + group_hidden_states = hidden_states + else: + # Slice out this group's tokens so sub-poolers see only their + # portion of the batch. Token offset is computed from the CPU + # `num_scheduled_tokens_cpu` to avoid a GPU->CPU sync. + group_cursor = group_metadata.pooling_cursor + num_group_tokens = int(group_cursor.num_scheduled_tokens_cpu.sum()) + group_hidden_states = hidden_states[ + token_offset : token_offset + num_group_tokens + ] + if token_offset: + # Shift first/last indices to be relative to the slice + # so seqwise poolers (which index `hidden_states` directly) + # remain correct. + pooling_cursor = dataclasses.replace( + group_cursor, + first_token_indices_gpu=( + group_cursor.first_token_indices_gpu - token_offset + ), + last_token_indices_gpu=( + group_cursor.last_token_indices_gpu - token_offset + ), + ) + group_metadata = dataclasses.replace( + group_metadata, pooling_cursor=pooling_cursor + ) + token_offset += num_group_tokens + + group_output: PoolerOutput = pooler(group_hidden_states, group_metadata) outputs.extend(group_output) offset += num_items diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index d3fefb745cfe..59b7234661b5 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/vllm/model_executor/layers/pooler/tokwise/methods.py @@ -47,17 +47,12 @@ def forward( pooling_metadata: PoolingMetadata, ) -> list[TokenPoolingMethodOutputItem]: pooling_cursor = pooling_metadata.get_pooling_cursor() - split_sizes = pooling_cursor.num_scheduled_tokens_cpu.tolist() - if split_sizes: - # DispatchPooler passes the full hidden_states tensor. - # slice out the subgroup once, then split it by - # per-request token counts - group_start = int(pooling_cursor.first_token_indices_gpu[0].item()) - group_end = int(pooling_cursor.last_token_indices_gpu[-1].item()) + 1 - hidden_states_group = hidden_states[group_start:group_end] - hidden_states_lst = list(hidden_states_group.split(split_sizes)) - else: - hidden_states_lst = [] + # Use the already-CPU num_scheduled_tokens tensor so `.tolist()` + # doesn't trigger a GPU->CPU sync. torch.split produces the same + # consecutive slices as indexing with first/last per-sequence indices. + hidden_states_lst = list( + torch.split(hidden_states, pooling_cursor.num_scheduled_tokens_cpu.tolist()) + ) if not self.enable_chunked_prefill: return hidden_states_lst @@ -95,12 +90,14 @@ def forward( pooling_metadata: PoolingMetadata, ) -> list[TokenPoolingMethodOutputItem]: pooled_data_lst = super().forward(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids() + # Use the CPU copy of prompt_token_ids so the step_tag_id mask can be + # resolved to indices without a d2h sync from boolean indexing. + prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu() pooling_params = pooling_metadata.pooling_params pooled_data = list[torch.Tensor | None]() - for data, token_id, pooling_param in zip( - pooled_data_lst, prompt_token_ids, pooling_params + for data, token_id_cpu, pooling_param in zip( + pooled_data_lst, prompt_token_ids_cpu, pooling_params ): # for unfinished chunked prefill if data is None: @@ -113,7 +110,9 @@ def forward( data = data[:, returned_token_ids] if step_tag_id is not None: - data = data[token_id == step_tag_id] + idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0] + idx = idx_cpu.to(data.device, non_blocking=True) + data = data[idx] pooled_data.append(data)