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, 07:30 UTC
[Perf][2/n] Eliminate GPU<->CPU syncs in pooling code (#41433)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 4, 2026 at 19:43 UTC
416f9cdede967edbf712727fa8a510c70f18aacb
3 files changed
+58
-27
vllm/model_executor/layers/pooler/seqwise/methods.py
+10
-8
@@ -68,21 +68,23 @@ class MeanPool(SequencePoolingMethod):
68
"partial prefill not supported with MEAN pooling"
69
)
70
71
- prompt_lens = pooling_cursor.prompt_lens_cpu.to(
72
- hidden_states.device, dtype=torch.int64, non_blocking=True
73
- )
74
-
75
- num_seqs = prompt_lens.numel()
71
+ prompt_lens_cpu = pooling_cursor.prompt_lens_cpu
72
+ num_seqs = prompt_lens_cpu.numel()
73
hidden_size = hidden_states.shape[-1]
74
75
if num_seqs == 0:
76
# early return for empty batch
77
return hidden_states.new_empty((0, hidden_size), dtype=torch.float32)
78
82
- # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]
79
+ # Build segment_ids on CPU so repeat_interleave doesn't need to sync
80
+ # GPU->CPU to learn its data-dependent output length, then upload
81
+ # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]
82
segment_ids = torch.repeat_interleave(
84
- torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long),
85
- prompt_lens,
83
+ torch.arange(num_seqs, dtype=torch.long),
84
+ prompt_lens_cpu,
85
+ ).to(hidden_states.device, non_blocking=True)
86
+ prompt_lens = prompt_lens_cpu.to(
87
+ hidden_states.device, dtype=torch.int64, non_blocking=True
88
)
89
segment_sums = torch.zeros(
90
(num_seqs, hidden_size),
vllm/model_executor/layers/pooler/special.py
+34
-4
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import dataclasses
4
from collections.abc import Mapping, Set
5
from itertools import groupby
6
@@ -80,9 +81,11 @@ class DispatchPooler(Pooler):
81
pooling_metadata: PoolingMetadata,
82
) -> PoolerOutput:
83
poolers_by_task = self.poolers_by_task
84
+ cursor = pooling_metadata.pooling_cursor
85
86
outputs = list[torch.Tensor | None]()
87
offset = 0
88
+ token_offset = 0
89
for task, group in groupby(pooling_metadata.tasks):
90
if not (pooler := poolers_by_task.get(task)):
91
raise ValueError(
@@ -91,10 +94,37 @@ class DispatchPooler(Pooler):
94
)
95
96
num_items = len(list(group))
94
- group_output: PoolerOutput = pooler(
95
- hidden_states,
96
- pooling_metadata[offset : offset + num_items],
97
- )
97
+ group_metadata = pooling_metadata[offset : offset + num_items]
98
+ if cursor is None:
99
+ group_hidden_states = hidden_states
100
+ else:
101
+ # Slice out this group's tokens so sub-poolers see only their
102
+ # portion of the batch. Token offset is computed from the CPU
103
+ # `num_scheduled_tokens_cpu` to avoid a GPU->CPU sync.
104
+ group_cursor = group_metadata.pooling_cursor
105
+ num_group_tokens = int(group_cursor.num_scheduled_tokens_cpu.sum())
106
+ group_hidden_states = hidden_states[
107
+ token_offset : token_offset + num_group_tokens
108
+ ]
109
+ if token_offset:
110
+ # Shift first/last indices to be relative to the slice
111
+ # so seqwise poolers (which index `hidden_states` directly)
112
+ # remain correct.
113
+ pooling_cursor = dataclasses.replace(
114
+ group_cursor,
115
+ first_token_indices_gpu=(
116
+ group_cursor.first_token_indices_gpu - token_offset
117
+ ),
118
+ last_token_indices_gpu=(
119
+ group_cursor.last_token_indices_gpu - token_offset
120
+ ),
121
+ )
122
+ group_metadata = dataclasses.replace(
123
+ group_metadata, pooling_cursor=pooling_cursor
124
+ )
125
+ token_offset += num_group_tokens
126
+
127
+ group_output: PoolerOutput = pooler(group_hidden_states, group_metadata)
128
129
outputs.extend(group_output)
130
offset += num_items
vllm/model_executor/layers/pooler/tokwise/methods.py
+14
-15
@@ -47,17 +47,12 @@ class AllPool(TokenPoolingMethod):
47
pooling_metadata: PoolingMetadata,
48
) -> list[TokenPoolingMethodOutputItem]:
49
pooling_cursor = pooling_metadata.get_pooling_cursor()
50
- split_sizes = pooling_cursor.num_scheduled_tokens_cpu.tolist()
51
- if split_sizes:
52
- # DispatchPooler passes the full hidden_states tensor.
53
- # slice out the subgroup once, then split it by
54
- # per-request token counts
55
- group_start = int(pooling_cursor.first_token_indices_gpu[0].item())
56
- group_end = int(pooling_cursor.last_token_indices_gpu[-1].item()) + 1
57
- hidden_states_group = hidden_states[group_start:group_end]
58
- hidden_states_lst = list(hidden_states_group.split(split_sizes))
59
- else:
60
- hidden_states_lst = []
50
+ # Use the already-CPU num_scheduled_tokens tensor so `.tolist()`
51
+ # doesn't trigger a GPU->CPU sync. torch.split produces the same
52
+ # consecutive slices as indexing with first/last per-sequence indices.
53
+ hidden_states_lst = list(
54
+ torch.split(hidden_states, pooling_cursor.num_scheduled_tokens_cpu.tolist())
55
+ )
56
57
if not self.enable_chunked_prefill:
58
return hidden_states_lst
@@ -95,12 +90,14 @@ class StepPool(AllPool):
90
pooling_metadata: PoolingMetadata,
91
) -> list[TokenPoolingMethodOutputItem]:
92
pooled_data_lst = super().forward(hidden_states, pooling_metadata)
98
- prompt_token_ids = pooling_metadata.get_prompt_token_ids()
93
+ # Use the CPU copy of prompt_token_ids so the step_tag_id mask can be
94
+ # resolved to indices without a d2h sync from boolean indexing.
95
+ prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu()
96
pooling_params = pooling_metadata.pooling_params
97
98
pooled_data = list[torch.Tensor | None]()
102
- for data, token_id, pooling_param in zip(
103
- pooled_data_lst, prompt_token_ids, pooling_params
99
+ for data, token_id_cpu, pooling_param in zip(
100
+ pooled_data_lst, prompt_token_ids_cpu, pooling_params
101
):
102
# for unfinished chunked prefill
103
if data is None:
@@ -113,7 +110,9 @@ class StepPool(AllPool):
110
data = data[:, returned_token_ids]
111
112
if step_tag_id is not None:
116
- data = data[token_id == step_tag_id]
113
+ idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0]
114
+ idx = idx_cpu.to(data.device, non_blocking=True)
115
+ data = data[idx]
116
117
pooled_data.append(data)
118