[Pooling] Improve BGE-M3 sync pooling throughput by up to 3.13% (#53464)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Aug 24, 2026 at 08:33 UTC
79bb395eea64dbfef99a55f010d2854db71f8571
10 files changed
+107
-67
tests/model_executor/layers/test_pooler_methods.py
+36
@@ -5,9 +5,11 @@
5
from dataclasses import dataclass
6
from unittest.mock import patch
7
8
+import numpy as np
9
import pytest
10
import torch
11
12
+from vllm.model_executor.layers.pooler import PoolingParamsUpdate
13
from vllm.model_executor.layers.pooler.seqwise.methods import (
14
CLSPool,
15
LastPool,
@@ -52,6 +54,8 @@ def _make_pooling_cursor(
54
prompt_lens_cpu=prompt_lens_cpu,
55
seq_lens_cpu=seq_lens_cpu,
56
num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
57
+ partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu),
58
+ finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(),
59
)
60
61
@@ -101,6 +105,38 @@ def _make_metadata(
105
)
106
107
108
+def test_build_pooling_cursor_caches_cpu_status() -> None:
109
+ metadata = PoolingMetadata(
110
+ prompt_lens=torch.tensor([3, 4, 2], dtype=torch.long),
111
+ prompt_token_ids=None,
112
+ prompt_token_ids_cpu=None,
113
+ pooling_params=[PoolingParams(task="embed") for _ in range(3)],
114
+ pooling_states=[PoolingStates() for _ in range(3)],
115
+ )
116
+ metadata.build_pooling_cursor(
117
+ np.array([3, 2, 2], dtype=np.int32),
118
+ seq_lens_cpu=torch.tensor([3, 2, 1], dtype=torch.long),
119
+ device=_CPU,
120
+ )
121
+
122
+ cursor = metadata.get_pooling_cursor()
123
+
124
+ assert cursor.is_partial_prefill() is True
125
+ assert cursor.get_finished_mask() == [True, False, False]
126
+ assert cursor[:1].is_partial_prefill() is False
127
+ assert cursor[1:2].is_partial_prefill() is True
128
+ assert cursor[1:].get_finished_mask() == [False, False]
129
+
130
+
131
+def test_pooling_params_update_combines_token_id_requirements() -> None:
132
+ update = PoolingParamsUpdate() | PoolingParamsUpdate(requires_token_ids=True)
133
+ params = PoolingParams(task="embed")
134
+
135
+ update.apply(params)
136
+
137
+ assert params.requires_token_ids is True
138
+
139
+
140
# ---------------------------------------------------------------------------
141
# CLSPool
142
# ---------------------------------------------------------------------------
tests/models/language/pooling/test_splade_sparse_pooler.py
+3
-6
@@ -12,9 +12,7 @@ from vllm.model_executor.models.bert import (
12
BertMLMHead,
13
SPLADESparsePooler,
14
)
15
-from vllm.platforms import current_platform
15
from vllm.pooling_params import PoolingParams
17
-from vllm.utils.torch_utils import PIN_MEMORY
16
from vllm.v1.pool.late_interaction_runner import LateInteractionRunner
17
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
18
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -119,18 +117,17 @@ def test_pooling_runner_gathers_required_token_ids() -> None:
117
req_states = MagicMock(spec=RequestState)
118
req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32))
119
metadata = runner._get_pooling_metadata(
122
- input_batch, req_states, torch.device(current_platform.device_type)
120
+ input_batch, req_states, torch.device("cpu")
121
)
122
123
expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
124
assert metadata.prompt_token_ids_cpu is not None
127
- assert metadata.prompt_token_ids is not None
128
- assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
125
+ assert metadata.prompt_token_ids is None
126
+ assert metadata.prompt_token_ids_cpu.is_pinned() is False
127
torch.testing.assert_close(
128
metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32)
129
)
130
torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected)
133
- torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
131
132
133
def test_pooling_runner_stores_only_required_token_ids() -> None:
tests/v1/worker/test_gpu_input_batch.py
+5
-11
@@ -475,15 +475,14 @@ def test_placeholder_spec_token_ids_written_verbatim():
475
476
477
@pytest.mark.parametrize(
478
- ("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
478
+ ("pooling_params", "expect_cpu_prompt_token_ids"),
479
[
480
- ({"task": "classify"}, False, False),
481
- ({"task": "classify", "requires_token_ids": True}, True, True),
480
+ ({"task": "classify"}, False),
481
+ ({"task": "classify", "requires_token_ids": True}, True),
482
],
483
)
484
def test_pooling_metadata_token_id_buffers(
485
pooling_params: dict[str, object],
486
- expect_device_prompt_token_ids: bool,
486
expect_cpu_prompt_token_ids: bool,
487
):
488
from vllm.pooling_params import PoolingParams
@@ -504,13 +503,8 @@ def test_pooling_metadata_token_id_buffers(
503
input_batch.refresh_metadata()
504
505
metadata = input_batch.get_pooling_metadata()
507
- if expect_device_prompt_token_ids:
508
- assert input_batch.sampling_metadata.prompt_token_ids is not None
509
- assert metadata.prompt_token_ids is not None
510
- assert metadata.get_prompt_token_ids()[0].tolist() == req.prompt_token_ids
511
- else:
512
- assert input_batch.sampling_metadata.prompt_token_ids is None
513
- assert metadata.prompt_token_ids is None
506
+ assert input_batch.sampling_metadata.prompt_token_ids is None
507
+ assert metadata.prompt_token_ids is None
508
509
if expect_cpu_prompt_token_ids:
510
assert metadata.prompt_token_ids_cpu is not None
vllm/model_executor/layers/pooler/common.py
+1
-1
@@ -18,7 +18,7 @@ ActivationFn = Callable[[_T], _T]
18
@dataclass(frozen=True)
19
class PoolingParamsUpdate:
20
requires_token_ids: bool = False
21
- """Set this flag to enable prompt token IDs for your pooler."""
21
+ """Set this flag to enable CPU prompt token IDs for your pooler."""
22
23
def __or__(self, other: "PoolingParamsUpdate") -> "PoolingParamsUpdate":
24
return PoolingParamsUpdate(
vllm/model_executor/layers/pooler/tokwise/methods.py
+1
-1
@@ -69,7 +69,7 @@ class AllPool(TokenPoolingMethod):
69
return hidden_states_lst
70
71
pooling_states = pooling_metadata.pooling_states
72
- finished_mask = pooling_cursor.is_finished().tolist()
72
+ finished_mask = pooling_cursor.get_finished_mask()
73
74
# If chunked_prefill is enabled
75
# 1. first store the chunked hidden_states in pooling_states.hidden_states_cache
vllm/model_executor/models/jina.py
+25
-20
@@ -12,11 +12,12 @@ from torch import nn
12
13
from vllm.config import VllmConfig
14
from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.pooler import PoolingParamsUpdate
16
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
17
from vllm.sequence import IntermediateTensors
18
from vllm.tasks import PoolingTask
19
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
19
-from vllm.utils.gpu_sync_debug import gpu_sync_allowed
20
+from vllm.utils.torch_utils import async_tensor_h2d
21
from vllm.v1.pool.metadata import PoolingMetadata
22
23
from ..layers.pooler import DispatchPooler
@@ -104,33 +105,37 @@ class JinaForRankingPool(StepPool):
105
def get_supported_tasks(self) -> set[PoolingTask]:
106
return {"token_embed"}
107
108
+ def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
109
+ return PoolingParamsUpdate(requires_token_ids=True)
110
+
111
def forward(
112
self,
113
hidden_states: torch.Tensor,
114
pooling_metadata: PoolingMetadata,
115
) -> list[TokenPoolingMethodOutputItem]:
116
pooled_data_lst = super().forward(hidden_states, pooling_metadata)
113
- prompt_token_ids = pooling_metadata.get_prompt_token_ids()
117
+ prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu()
118
119
embeds_list = list[torch.Tensor | None]()
116
- # `torch.where` resolves the match count on the host.
117
- with gpu_sync_allowed():
118
- for data, token_ids in zip(pooled_data_lst, prompt_token_ids):
119
- # for unfinished chunked prefill
120
- if data is None:
121
- embeds_list.append(None)
122
- else:
123
- doc_match = torch.eq(token_ids, self.doc_token_id)
124
- docs_indexes = torch.where(doc_match)[0]
125
- query_indexes = torch.where(
126
- torch.eq(token_ids, self.query_token_id)
127
- )[0]
128
-
129
- # The JinaForRanking model concatenates docs first, then query.
130
- # Let's stay consistent with this novel design.
131
- indexes = torch.cat([docs_indexes, query_indexes])
132
- embeds = self.projector(data[indexes])
133
- embeds_list.append(embeds)
120
+ for data, token_ids_cpu in zip(pooled_data_lst, prompt_token_ids_cpu):
121
+ # for unfinished chunked prefill
122
+ if data is None:
123
+ embeds_list.append(None)
124
+ else:
125
+ doc_match = torch.eq(token_ids_cpu, self.doc_token_id)
126
+ docs_indexes_cpu = torch.where(doc_match)[0]
127
+ query_indexes_cpu = torch.where(
128
+ torch.eq(token_ids_cpu, self.query_token_id)
129
+ )[0]
130
+
131
+ # The JinaForRanking model concatenates docs first, then query.
132
+ # Let's stay consistent with this novel design.
133
+ indexes_cpu = torch.cat([docs_indexes_cpu, query_indexes_cpu])
134
+ indexes = async_tensor_h2d(
135
+ indexes_cpu, device=data.device, dtype=torch.long
136
+ )
137
+ embeds = self.projector(data[indexes])
138
+ embeds_list.append(embeds)
139
140
return embeds_list
141
vllm/v1/pool/metadata.py
+23
-5
@@ -17,22 +17,33 @@ class PoolingCursor:
17
prompt_lens_cpu: torch.Tensor
18
seq_lens_cpu: torch.Tensor
19
num_scheduled_tokens_cpu: torch.Tensor
20
+ partial_prefill: bool
21
+ finished_mask: list[bool]
22
23
def __getitem__(self, indices: slice) -> "PoolingCursor":
24
+ prompt_lens_cpu = self.prompt_lens_cpu[indices]
25
+ seq_lens_cpu = self.seq_lens_cpu[indices]
26
+ num_scheduled_tokens_cpu = self.num_scheduled_tokens_cpu[indices]
27
+
28
return PoolingCursor(
29
first_token_indices_gpu=self.first_token_indices_gpu[indices],
30
last_token_indices_gpu=self.last_token_indices_gpu[indices],
25
- prompt_lens_cpu=self.prompt_lens_cpu[indices],
26
- seq_lens_cpu=self.seq_lens_cpu[indices],
27
- num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],
31
+ prompt_lens_cpu=prompt_lens_cpu,
32
+ seq_lens_cpu=seq_lens_cpu,
33
+ num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
34
+ partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu),
35
+ finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(),
36
)
37
38
def is_partial_prefill(self) -> bool:
31
- return not torch.all(self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)
39
+ return self.partial_prefill
40
41
def is_finished(self) -> torch.Tensor:
42
return self.prompt_lens_cpu == self.seq_lens_cpu
43
44
+ def get_finished_mask(self) -> list[bool]:
45
+ return self.finished_mask
46
+
47
48
class PoolingStates:
49
def __init__(self) -> None:
@@ -130,9 +141,14 @@ class PoolingMetadata:
141
)
142
143
num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
144
+ prompt_lens_np = prompt_lens.numpy()
145
+ seq_lens_np = seq_lens_cpu.numpy()
146
if query_start_loc_gpu is None:
147
cumsum = torch.zeros(
135
- n_seq + 1, dtype=torch.int64, pin_memory=PIN_MEMORY, device="cpu"
148
+ n_seq + 1,
149
+ dtype=torch.int64,
150
+ pin_memory=PIN_MEMORY,
151
+ device="cpu",
152
)
153
torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
154
cumsum = cumsum.to(device, non_blocking=True)
@@ -155,4 +171,6 @@ class PoolingMetadata:
171
prompt_lens_cpu=prompt_lens,
172
seq_lens_cpu=seq_lens_cpu,
173
num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
174
+ partial_prefill=not np.array_equal(prompt_lens_np, num_scheduled_tokens_np),
175
+ finished_mask=np.equal(prompt_lens_np, seq_lens_np).tolist(),
176
)
vllm/v1/worker/gpu/pool/pooling_runner.py
+8
-13
@@ -103,28 +103,26 @@ class PoolingRunner:
103
pooling_params = [self.pooling_params[i] for i in req_indices]
104
pooling_states = [self.pooling_states[i] for i in req_indices]
105
prompt_lens = torch.from_numpy(
106
- req_states.prompt_len.np[input_batch.idx_mapping_np].copy()
106
+ req_states.prompt_len.np[input_batch.idx_mapping_np]
107
)
108
109
prompt_token_ids_cpu = None
110
- prompt_token_ids = None
110
if any(params.requires_token_ids for params in pooling_params):
111
max_prompt_len = int(prompt_lens.max())
112
prompt_token_ids_cpu = torch.zeros(
113
(input_batch.num_reqs, max_prompt_len),
114
dtype=torch.int64,
116
- pin_memory=PIN_MEMORY,
115
+ pin_memory=PIN_MEMORY and device.type != "cpu",
116
)
117
for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
118
if not params.requires_token_ids:
119
continue
120
token_ids = self.prompt_token_ids[req_index]
121
prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
123
- prompt_token_ids = prompt_token_ids_cpu.to(device, non_blocking=True)
122
123
return PoolingMetadata(
124
prompt_lens=prompt_lens,
127
- prompt_token_ids=prompt_token_ids,
125
+ prompt_token_ids=None,
126
prompt_token_ids_cpu=prompt_token_ids_cpu,
127
pooling_params=pooling_params,
128
pooling_states=pooling_states,
@@ -150,7 +148,7 @@ class PoolingRunner:
148
query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
149
)
150
pooler_output = self.model.pooler(hidden_states, pooling_metadata)
153
- finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()
151
+ finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask()
152
pooler_output = self.late_interaction_runner.postprocess_pooler_output(
153
raw_pooler_output=pooler_output,
154
pooling_params=pooling_metadata.pooling_params,
@@ -174,19 +172,16 @@ class PoolingRunner:
172
pooling_params = PoolingParams(task=task)
173
pooling_params.verify(self.model_config)
174
self.model.pooler.get_pooling_updates(task).apply(pooling_params)
177
- prompt_token_ids = None
175
+ prompt_token_ids_cpu = None
176
if pooling_params.requires_token_ids:
179
- prompt_token_ids = torch.zeros(
177
+ prompt_token_ids_cpu = torch.zeros(
178
(num_reqs, int(prompt_lens.max())),
179
dtype=torch.int64,
182
- device=hidden_states.device,
180
)
181
pooling_metadata = PoolingMetadata(
182
prompt_lens=prompt_lens,
186
- prompt_token_ids=prompt_token_ids,
187
- prompt_token_ids_cpu=None
188
- if prompt_token_ids is None
189
- else prompt_token_ids.cpu(),
183
+ prompt_token_ids=None,
184
+ prompt_token_ids_cpu=prompt_token_ids_cpu,
185
pooling_params=[pooling_params] * num_reqs,
186
pooling_states=[PoolingStates() for _ in range(num_reqs)],
187
)
vllm/v1/worker/gpu_input_batch.py
+4
-6
@@ -476,9 +476,7 @@ class InputBatch:
476
477
self.pooling_params[req_id] = pooling_params
478
self.pooling_states[req_id] = pooling_states
479
- self.logits_processing_needs_token_ids[req_index] = (
480
- pooling_params.requires_token_ids
481
- )
479
+ self.logits_processing_needs_token_ids[req_index] = False
480
else:
481
raise NotImplementedError("Unrecognized request type")
482
@@ -890,8 +888,8 @@ class InputBatch:
888
not self.no_penalties
889
or self.logits_processing_needs_token_ids[:num_reqs].any()
890
)
893
- # The prompt tokens are used only for applying penalties or
894
- # step pooling during the sampling/pooling process.
891
+ # The device prompt tokens are used only for applying penalties or
892
+ # pooling methods that explicitly request GPU token IDs.
893
# Hence copy these tensors only when there are requests which
894
# need penalties/step_pooler to be applied.
895
prompt_token_ids_cpu = (
@@ -979,7 +977,7 @@ class InputBatch:
977
978
return PoolingMetadata(
979
prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(),
982
- prompt_token_ids=self.sampling_metadata.prompt_token_ids,
980
+ prompt_token_ids=None,
981
prompt_token_ids_cpu=prompt_token_ids_cpu,
982
pooling_params=pooling_params,
983
pooling_states=pooling_states,
vllm/v1/worker/gpu_model_runner.py
+1
-4
@@ -3550,10 +3550,7 @@ class GPUModelRunner(
3550
hidden_states=hidden_states, pooling_metadata=pooling_metadata
3551
)
3552
3553
- finished_mask = [
3554
- seq_len == prompt_len
3555
- for seq_len, prompt_len in zip(seq_lens_cpu, pooling_metadata.prompt_lens)
3556
- ]
3553
+ finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask()
3554
raw_pooler_output = self.late_interaction_runner.postprocess_pooler_output(
3555
raw_pooler_output=raw_pooler_output,
3556
pooling_params=pooling_metadata.pooling_params,