iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][2/n] Eliminate GPU<->CPU syncs in pooling code (#41433)

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Nick Hill
2026-05-05 02:43:25 +00:00
committed by GitHub
parent 685bf811d6
commit 416f9cdede
3 changed files with 58 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/pooler/seqwise/methods.py
+10 -8
View File
@@ -68,21 +68,23 @@ class MeanPool(SequencePoolingMethod):
"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),
vllm/model_executor/layers/pooler/special.py
+34 -4
View File
@@ -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 @@ class DispatchPooler(Pooler):
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 @@ class DispatchPooler(Pooler):
)
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
vllm/model_executor/layers/pooler/tokwise/methods.py
+14 -15
View File
@@ -47,17 +47,12 @@ class AllPool(TokenPoolingMethod):
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 @@ class StepPool(AllPool):
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 @@ class StepPool(AllPool):
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)