[ModelRunner V2] Enable sequence pooling for embedding and classification models (#48791)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jul 29, 2026 at 09:30 UTC
43eaefba5a5dbccb71d2404e4bf28e21bb74fce6
10 files changed
+509
-37
tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py
+44
-1
@@ -15,7 +15,9 @@ from vllm.config import PoolerConfig
15
["Qwen/Qwen3-Embedding-0.6B"],
16
)
17
@torch.inference_mode
18
-def test_embed_models(hf_runner, vllm_runner, model: str):
18
+def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):
19
+ # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default.
20
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
21
chunk_size = 10
22
n_prompt_tokens = [55, 56, 57]
23
token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
@@ -104,3 +106,44 @@ def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: s
106
assert chunked == pytest.approx(unchunked, abs=5e-2), (
107
f"chunked score {chunked} diverged from unchunked {unchunked}"
108
)
109
+
110
+
111
+@torch.inference_mode
112
+def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
113
+ model = "Qwen/Qwen3-Embedding-0.6B"
114
+ chunk_size = 10
115
+ token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)]
116
+ prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts]
117
+
118
+ with hf_runner(model, auto_cls=AutoModel) as hf_model:
119
+ hf_outputs = []
120
+ for token_prompt in token_prompts:
121
+ inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
122
+ output = hf_model.model(inputs["input_ids"])
123
+ embedding = torch.nn.functional.normalize(
124
+ output.last_hidden_state.float()[0, -1], dim=0
125
+ )
126
+ hf_outputs.append(embedding.cpu().tolist())
127
+
128
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
129
+ with vllm_runner(
130
+ model,
131
+ runner="pooling",
132
+ pooler_config=PoolerConfig(task="embed"),
133
+ max_model_len=64,
134
+ max_num_batched_tokens=chunk_size,
135
+ max_num_seqs=2,
136
+ gpu_memory_utilization=0.25,
137
+ enforce_eager=True,
138
+ enable_chunked_prefill=True,
139
+ ) as vllm_model:
140
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
141
+ vllm_outputs = vllm_model.embed(prompts)
142
+
143
+ check_embeddings_close(
144
+ embeddings_0_lst=hf_outputs,
145
+ embeddings_1_lst=vllm_outputs,
146
+ name_0="hf",
147
+ name_1="vllm",
148
+ tol=1e-2,
149
+ )
tests/models/language/pooling/test_classification.py
+57
@@ -49,3 +49,60 @@ def test_models(
49
vllm_output,
50
rtol=2e-3 if dtype == "float" else 1e-2,
51
)
52
+
53
+
54
+@pytest.mark.core_model
55
+def test_bert_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
56
+ model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
57
+ score_inputs = (
58
+ "What is the capital of France?",
59
+ [
60
+ "Paris.",
61
+ "Paris is the capital and largest city of France.",
62
+ "William Shakespeare wrote Hamlet in the early seventeenth century.",
63
+ ],
64
+ )
65
+ prompt_batches = [
66
+ ["short input"],
67
+ [
68
+ "short input",
69
+ "a longer input that exercises mixed sequence lengths",
70
+ ],
71
+ ]
72
+
73
+ with hf_runner(
74
+ model, dtype="half", auto_cls=AutoModelForSequenceClassification
75
+ ) as hf_model:
76
+ # HfRunner uses problem_type to preserve the model's
77
+ # sbert_ce_default_activation_function=Identity raw logits.
78
+ hf_model.config.problem_type = "regression"
79
+ hf_outputs = [hf_model.classify(prompts) for prompts in prompt_batches]
80
+
81
+ text_1, text_2 = score_inputs
82
+ text_pairs = [[text_1, document] for document in text_2]
83
+ with hf_runner(model, dtype="half", is_cross_encoder=True) as hf_model:
84
+ hf_scores = hf_model.predict(text_pairs).tolist()
85
+
86
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
87
+ with vllm_runner(
88
+ model,
89
+ runner="pooling",
90
+ dtype="half",
91
+ max_model_len=64,
92
+ ) as vllm_model:
93
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
94
+ vllm_outputs = [vllm_model.classify(prompts) for prompts in prompt_batches]
95
+ vllm_scores = vllm_model.score(*score_inputs)
96
+
97
+ for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
98
+ hf_tensor = torch.tensor(hf_batch)
99
+ vllm_tensor = torch.tensor(vllm_batch)
100
+ assert vllm_tensor.shape == hf_tensor.shape
101
+ assert torch.allclose(vllm_tensor, hf_tensor, rtol=1e-2, atol=1e-4)
102
+
103
+ assert torch.allclose(
104
+ torch.tensor(vllm_scores),
105
+ torch.tensor(hf_scores),
106
+ rtol=1e-2,
107
+ atol=1e-4,
108
+ )
tests/models/language/pooling/test_embedding.py
+76
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
from transformers import AutoModel
7
8
+from vllm import PoolingParams
9
from vllm.config import PoolerConfig
10
11
from ...utils import check_embeddings_close
@@ -132,6 +133,7 @@ def test_encoder_only_model_runner_v2_attention(
133
task="embed", seq_pooling_type="LAST", use_activation=True
134
),
135
) as vllm_model:
136
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
137
vllm_outputs = vllm_model.embed(prompts)
138
139
check_embeddings_close(
@@ -141,3 +143,77 @@ def test_encoder_only_model_runner_v2_attention(
143
name_1="vllm",
144
tol=1e-2,
145
)
146
+
147
+
148
+@pytest.mark.core_model
149
+def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
150
+ model = "sentence-transformers/all-MiniLM-L6-v2"
151
+ prompt_batches = [
152
+ ["short input"],
153
+ [
154
+ "short input",
155
+ "a longer input that exercises mixed sequence lengths",
156
+ ],
157
+ ]
158
+
159
+ with hf_runner(model, is_sentence_transformer=True) as hf_model:
160
+ hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches]
161
+
162
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
163
+ with vllm_runner(
164
+ model,
165
+ runner="pooling",
166
+ max_model_len=64,
167
+ ) as vllm_model:
168
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
169
+ vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches]
170
+
171
+ for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
172
+ check_embeddings_close(
173
+ embeddings_0_lst=hf_batch,
174
+ embeddings_1_lst=vllm_batch,
175
+ name_0="hf",
176
+ name_1="vllm",
177
+ tol=1e-2,
178
+ )
179
+
180
+
181
+@pytest.mark.core_model
182
+def test_matryoshka_dimensions_model_runner_v2(
183
+ hf_runner, vllm_runner, monkeypatch
184
+) -> None:
185
+ model = "Snowflake/snowflake-arctic-embed-m-v1.5"
186
+ prompts = ["short input", "a longer input for a different output width"]
187
+ dimensions = [None, 256]
188
+
189
+ with hf_runner(model, is_sentence_transformer=True) as hf_model:
190
+ hf_outputs = hf_model.encode(prompts)
191
+
192
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
193
+ with vllm_runner(
194
+ model,
195
+ runner="pooling",
196
+ max_model_len=64,
197
+ gpu_memory_utilization=0.25,
198
+ ) as vllm_model:
199
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
200
+ vllm_outputs = vllm_model.embed(
201
+ prompts,
202
+ pooling_params=[PoolingParams(dimensions=d) for d in dimensions],
203
+ )
204
+
205
+ expected_outputs = []
206
+ for output, dimension in zip(hf_outputs, dimensions):
207
+ output = torch.as_tensor(output)
208
+ if dimension is not None:
209
+ output = torch.nn.functional.normalize(output[:dimension], dim=0)
210
+ expected_outputs.append(output.tolist())
211
+
212
+ assert [len(output) for output in vllm_outputs] == [768, 256]
213
+ check_embeddings_close(
214
+ embeddings_0_lst=expected_outputs,
215
+ embeddings_1_lst=vllm_outputs,
216
+ name_0="hf",
217
+ name_1="vllm",
218
+ tol=1e-2,
219
+ )
tests/models/language/pooling/test_splade_sparse_pooler.py
+75
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from unittest.mock import MagicMock, patch
5
+
6
+import numpy as np
7
import pytest
8
import torch
9
import torch.nn as nn
@@ -9,8 +12,13 @@ from vllm.model_executor.models.bert import (
12
BertMLMHead,
13
SPLADESparsePooler,
14
)
15
+from vllm.platforms import current_platform
16
from vllm.pooling_params import PoolingParams
17
+from vllm.utils.torch_utils import PIN_MEMORY
18
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
19
+from vllm.v1.worker.gpu.input_batch import InputBatch
20
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
21
+from vllm.v1.worker.gpu.states import RequestState
22
23
# ---------------------------------------------------------------------
24
# Functional test: SPLADE formula correctness (no HF download needed)
@@ -91,3 +99,70 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
99
rtol=1e-4,
100
atol=1e-4,
101
)
102
+
103
+
104
+def test_pooling_runner_gathers_required_token_ids() -> None:
105
+ runner = PoolingRunner.__new__(PoolingRunner)
106
+ pooling_params = PoolingParams(task="embed", requires_token_ids=True)
107
+ runner.pooling_params = {1: pooling_params, 3: pooling_params}
108
+ runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()}
109
+ runner.prompt_token_ids = {
110
+ 1: torch.tensor([101, 102]),
111
+ 3: torch.tensor([101, 11, 102]),
112
+ }
113
+
114
+ input_batch = MagicMock(spec=InputBatch)
115
+ input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32)
116
+ input_batch.num_reqs = 2
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(
120
+ input_batch, req_states, torch.device(current_platform.device_type)
121
+ )
122
+
123
+ expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
124
+ assert metadata.prompt_token_ids_cpu is not None
125
+ assert metadata.prompt_token_ids is not None
126
+ assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
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)
131
+ torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
132
+
133
+
134
+def test_pooling_runner_stores_only_required_token_ids() -> None:
135
+ runner = PoolingRunner.__new__(PoolingRunner)
136
+ runner.model = MagicMock()
137
+ runner.supported_tasks = frozenset({"embed"})
138
+ runner.pooling_params = {}
139
+ runner.pooling_states = {}
140
+ runner.prompt_token_ids = {}
141
+
142
+ runner.add_request(1, PoolingParams(task="embed"), [101, 102])
143
+ runner.add_request(
144
+ 2,
145
+ PoolingParams(task="embed", requires_token_ids=True),
146
+ [101, 11, 102],
147
+ )
148
+
149
+ assert 1 not in runner.prompt_token_ids
150
+ torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102]))
151
+
152
+
153
+def test_pooling_runner_rejects_unsupported_selected_task() -> None:
154
+ model = MagicMock()
155
+ model.pooler.get_supported_tasks.return_value = {
156
+ "embed",
157
+ "embed&token_classify",
158
+ "token_classify",
159
+ }
160
+ vllm_config = MagicMock()
161
+ vllm_config.scheduler_config.max_num_seqs = 2
162
+ vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
163
+
164
+ with (
165
+ patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]),
166
+ pytest.raises(ValueError, match="selects 'embed&token_classify'"),
167
+ ):
168
+ PoolingRunner(model, vllm_config)
tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
+1
@@ -41,6 +41,7 @@ def mock_model_runner_with_req_states():
41
runner.sampler = None
42
runner.prompt_logprobs_worker = None
43
runner.is_last_pp_rank = False
44
+ runner.pooling_runner = None
45
46
# Mock staged writes — they use Triton kernels that require GPU
47
runner.req_states.apply_staged_writes = Mock()
vllm/v1/worker/gpu/async_utils.py
+27
-14
@@ -6,7 +6,12 @@ import numpy as np
6
import torch
7
8
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
9
-from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput
9
+from vllm.v1.outputs import (
10
+ AsyncModelRunnerOutput,
11
+ LogprobsTensors,
12
+ ModelRunnerOutput,
13
+ PoolerOutput,
14
+)
15
from vllm.v1.worker.gpu.sample.output import SamplerOutput
16
17
@@ -89,34 +94,42 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
94
def __init__(
95
self,
96
model_runner_output: ModelRunnerOutput,
92
- pooler_output: torch.Tensor,
93
- is_valid: torch.Tensor | None,
97
+ pooler_output: PoolerOutput,
98
+ finished_mask: list[bool],
99
main_stream: torch.cuda.Stream,
100
copy_stream: torch.cuda.Stream,
101
):
102
self.model_runner_output = model_runner_output
103
self.pooler_output = pooler_output
99
- self.is_valid = is_valid
104
# Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
105
self.copy_event = torch.cuda.Event(blocking=True)
106
107
with stream(copy_stream, main_stream):
108
copy_stream.wait_stream(main_stream)
105
- self.pooler_output_cpu = self.pooler_output.to("cpu", non_blocking=True)
106
- if self.is_valid is not None:
107
- self.is_valid_cpu = self.is_valid.to("cpu", non_blocking=True)
109
+ if isinstance(self.pooler_output, torch.Tensor) and all(finished_mask):
110
+ self.pooler_output_cpu: PoolerOutput = self.pooler_output.to(
111
+ "cpu", non_blocking=True
112
+ )
113
else:
109
- self.is_valid_cpu = None
114
+ outputs = (
115
+ self.pooler_output.unbind()
116
+ if isinstance(self.pooler_output, torch.Tensor)
117
+ else self.pooler_output
118
+ )
119
+ self.pooler_output_cpu = [
120
+ None
121
+ if output is None or not is_finished
122
+ else output.to("cpu", non_blocking=True)
123
+ for output, is_finished in zip(outputs, finished_mask, strict=True)
124
+ ]
125
self.copy_event.record(copy_stream)
126
127
def get_output(self) -> ModelRunnerOutput:
113
- pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
128
+ if isinstance(self.pooler_output_cpu, torch.Tensor):
129
+ pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
130
+ else:
131
+ pooler_output = self.pooler_output_cpu
132
self.copy_event.synchronize()
115
- if self.is_valid_cpu is not None:
116
- is_valid_cpu = self.is_valid_cpu.tolist()
117
- for i, is_valid in enumerate(is_valid_cpu):
118
- if not is_valid:
119
- pooler_output[i] = None
133
self.model_runner_output.pooler_output = pooler_output
134
return self.model_runner_output
135
vllm/v1/worker/gpu/model_runner.py
+13
-3
@@ -366,7 +366,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
366
)
367
368
if self.is_pooling_model and self.is_last_pp_rank:
369
- self.pooling_runner = PoolingRunner(self.model)
369
+ self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
370
eplb_models_added |= self.eplb.maybe_register_model(
371
self.model,
372
self.model_config,
@@ -787,6 +787,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
787
req_idx = self.req_states.remove_request(req_id)
788
if req_idx is None:
789
return False
790
+ if self.pooling_runner is not None:
791
+ self.pooling_runner.remove_request(req_idx)
792
if self.pp_handler is not None:
793
self.pp_handler.on_req_idx_freed(req_idx)
794
if self.encoder_cache is not None:
@@ -839,6 +841,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
841
)
842
req_index = self.req_states.req_id_to_index[req_id]
843
844
+ if self.pooling_runner is not None:
845
+ assert new_req_data.pooling_params is not None
846
+ self.pooling_runner.add_request(
847
+ req_index,
848
+ new_req_data.pooling_params,
849
+ new_req_data.prompt_token_ids,
850
+ )
851
+
852
if self.encoder_cache is not None:
853
self.encoder_cache.add_request(req_id, new_req_data.mm_features)
854
@@ -1591,7 +1601,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1601
return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1602
1603
assert self.pooling_runner is not None
1594
- pooler_output, is_valid = self.pooling_runner.pool(
1604
+ pooler_output, finished_mask = self.pooling_runner.pool(
1605
hidden_states, input_batch, self.req_states
1606
)
1607
@@ -1604,7 +1614,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1614
async_output = AsyncPoolingOutput(
1615
model_runner_output=model_runner_output,
1616
pooler_output=pooler_output,
1607
- is_valid=is_valid,
1617
+ finished_mask=finished_mask,
1618
main_stream=self.main_stream,
1619
copy_stream=self.output_copy_stream,
1620
)
vllm/v1/worker/gpu/model_states/encoder_only.py
+49
-1
@@ -8,12 +8,13 @@ import torch.nn as nn
8
from vllm.config import VllmConfig, get_layers_from_vllm_config
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.model_executor.layers.attention import Attention
11
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
11
+from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
12
from vllm.v1.attention.backend import (
13
AttentionCGSupport,
14
AttentionType,
15
CommonAttentionMetadata,
16
)
17
+from vllm.v1.core.sched.output import NewRequestData
18
from vllm.v1.kv_cache_interface import (
19
AttentionSpec,
20
EncoderOnlyAttentionSpec,
@@ -22,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23
from vllm.v1.worker.gpu.input_batch import InputBatch
24
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
25
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
26
+from vllm.v1.worker.gpu.states import RequestState
27
from vllm.v1.worker.utils import AttentionGroup
28
29
@@ -95,6 +97,52 @@ class EncoderOnlyModelState(DefaultModelState):
97
self._dummy_slot_mapping = torch.zeros(
98
self.max_num_tokens, dtype=torch.int64, device=device
99
)
100
+ self.token_type_ids: dict[str, torch.Tensor] = {}
101
+
102
+ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
103
+ super().add_request(req_index, new_req_data)
104
+ pooling_params = new_req_data.pooling_params
105
+ if pooling_params is None or pooling_params.extra_kwargs is None:
106
+ return
107
+
108
+ token_type_start = pooling_params.extra_kwargs.get("compressed_token_type_ids")
109
+ if token_type_start is not None:
110
+ assert new_req_data.prompt_token_ids is not None
111
+ self.token_type_ids[new_req_data.req_id] = (
112
+ torch.arange(len(new_req_data.prompt_token_ids), dtype=torch.int32)
113
+ >= token_type_start
114
+ ).to(torch.int32)
115
+
116
+ def remove_request(self, req_id: str) -> None:
117
+ super().remove_request(req_id)
118
+ self.token_type_ids.pop(req_id, None)
119
+
120
+ def prepare_inputs(
121
+ self, input_batch: InputBatch, req_states: RequestState
122
+ ) -> dict[str, torch.Tensor | None]:
123
+ model_inputs = super().prepare_inputs(input_batch, req_states)
124
+ if not self.token_type_ids:
125
+ return model_inputs
126
+
127
+ token_type_ids_cpu = torch.zeros(
128
+ input_batch.num_tokens_after_padding,
129
+ dtype=torch.int32,
130
+ pin_memory=PIN_MEMORY,
131
+ )
132
+ offset = 0
133
+ for i, req_id in enumerate(input_batch.req_ids):
134
+ num_tokens = int(input_batch.num_scheduled_tokens[i])
135
+ request_token_type_ids = self.token_type_ids.get(req_id)
136
+ if request_token_type_ids is not None:
137
+ start = int(input_batch.num_computed_tokens_np[i])
138
+ token_type_ids_cpu[offset : offset + num_tokens].copy_(
139
+ request_token_type_ids[start : start + num_tokens]
140
+ )
141
+ offset += num_tokens
142
+ model_inputs["token_type_ids"] = token_type_ids_cpu.to(
143
+ self.device, non_blocking=True
144
+ )
145
+ return model_inputs
146
147
def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
148
# Encoder groups are built here rather than in init_attn_backend, so
vllm/v1/worker/gpu/pool/pooling_runner.py
+162
-17
@@ -2,45 +2,190 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import cast
4
5
+import numpy as np
6
import torch
7
import torch.nn as nn
7
-import torch.nn.functional as F
8
9
+from vllm.config import VllmConfig
10
from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
11
+from vllm.pooling_params import PoolingParams
12
from vllm.tasks import PoolingTask
13
+from vllm.utils.torch_utils import PIN_MEMORY
14
+from vllm.v1.outputs import PoolerOutput
15
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
16
from vllm.v1.worker.gpu.input_batch import InputBatch
17
from vllm.v1.worker.gpu.states import RequestState
18
19
+_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset({"embed", "classify"})
20
+
21
15
-# NOTE(woosuk): Currently, this class only supports the "LAST" pooling task
16
-# on decoder-only models. How to support other pooling tasks and models
17
-# is to be determined.
22
class PoolingRunner:
19
- def __init__(self, model: nn.Module):
23
+ def __init__(self, model: nn.Module, vllm_config: VllmConfig):
24
self.model = cast(VllmModelForPooling, model)
25
+ self.model_config = vllm_config.model_config
26
+ self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
27
+ model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))
28
+ selected_task = self.model_config.get_pooling_task(model_tasks)
29
+ if selected_task not in _SUPPORTED_TASKS:
30
+ raise ValueError(
31
+ "Model Runner V2 supports only sequence-level pooling tasks "
32
+ f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
33
+ f"{selected_task!r} from {list(model_tasks)}. Set an explicitly "
34
+ "supported task or VLLM_USE_V2_MODEL_RUNNER=0."
35
+ )
36
+ self.supported_tasks = frozenset(self.get_supported_tasks(model))
37
+ if not self.supported_tasks:
38
+ raise ValueError(
39
+ "Model Runner V2 supports only sequence-level pooling tasks "
40
+ f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
41
+ f"{list(model_tasks)}. "
42
+ "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
43
+ )
44
+ self.pooling_params: dict[int, PoolingParams] = {}
45
+ self.pooling_states: dict[int, PoolingStates] = {}
46
+ self.prompt_token_ids: dict[int, torch.Tensor] = {}
47
48
@staticmethod
49
def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
50
if not is_pooling_model(model):
51
return []
26
- assert "embed" in model.pooler.get_supported_tasks()
27
- return ["embed"]
52
+ return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
53
+
54
+ def add_request(
55
+ self,
56
+ req_index: int,
57
+ pooling_params: PoolingParams,
58
+ prompt_token_ids: list[int],
59
+ ) -> None:
60
+ task = pooling_params.task
61
+ if task not in self.supported_tasks:
62
+ raise ValueError(
63
+ f"Unsupported task: {task!r}. "
64
+ f"Supported tasks: {sorted(self.supported_tasks)}"
65
+ )
66
+ self.model.pooler.get_pooling_updates(task).apply(pooling_params)
67
+ self.pooling_params[req_index] = pooling_params
68
+ self.pooling_states[req_index] = PoolingStates()
69
+ if pooling_params.requires_token_ids:
70
+ self.prompt_token_ids[req_index] = torch.tensor(
71
+ prompt_token_ids, dtype=torch.int64
72
+ )
73
+
74
+ def remove_request(self, req_index: int) -> None:
75
+ self.pooling_params.pop(req_index, None)
76
+ if state := self.pooling_states.pop(req_index, None):
77
+ state.clean()
78
+ self.prompt_token_ids.pop(req_index, None)
79
+
80
+ def _get_pooling_metadata(
81
+ self,
82
+ input_batch: InputBatch,
83
+ req_states: RequestState,
84
+ device: torch.device,
85
+ ) -> PoolingMetadata:
86
+ req_indices = input_batch.idx_mapping_np.tolist()
87
+ pooling_params = [self.pooling_params[i] for i in req_indices]
88
+ pooling_states = [self.pooling_states[i] for i in req_indices]
89
+ prompt_lens = torch.from_numpy(
90
+ req_states.prompt_len.np[input_batch.idx_mapping_np].copy()
91
+ )
92
+
93
+ prompt_token_ids_cpu = None
94
+ prompt_token_ids = None
95
+ if any(params.requires_token_ids for params in pooling_params):
96
+ max_prompt_len = int(prompt_lens.max())
97
+ prompt_token_ids_cpu = torch.zeros(
98
+ (input_batch.num_reqs, max_prompt_len),
99
+ dtype=torch.int64,
100
+ pin_memory=PIN_MEMORY,
101
+ )
102
+ for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
103
+ if not params.requires_token_ids:
104
+ continue
105
+ token_ids = self.prompt_token_ids[req_index]
106
+ prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
107
+ prompt_token_ids = prompt_token_ids_cpu.to(device, non_blocking=True)
108
+
109
+ return PoolingMetadata(
110
+ prompt_lens=prompt_lens,
111
+ prompt_token_ids=prompt_token_ids,
112
+ prompt_token_ids_cpu=prompt_token_ids_cpu,
113
+ pooling_params=pooling_params,
114
+ pooling_states=pooling_states,
115
+ )
116
117
def pool(
118
self,
119
hidden_states: torch.Tensor,
120
input_batch: InputBatch,
121
req_states: RequestState,
34
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
35
- # TODO(woosuk): Support different types of pooling tasks.
36
- last_hidden_states = hidden_states[input_batch.logits_indices]
37
- # TODO(woosuk): Make normalization optional.
38
- last_hidden_states = F.normalize(last_hidden_states, p=2, dim=-1)
122
+ ) -> tuple[PoolerOutput, list[bool]]:
123
+ hidden_states = hidden_states[: input_batch.num_tokens]
124
+ pooling_metadata = self._get_pooling_metadata(
125
+ input_batch, req_states, hidden_states.device
126
+ )
127
+ num_reqs = input_batch.num_reqs
128
+ # Pooling has no speculative tokens, so this CPU upper bound is exact.
129
+ seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs]
130
+ pooling_metadata.build_pooling_cursor(
131
+ input_batch.num_scheduled_tokens,
132
+ seq_lens_cpu,
133
+ device=hidden_states.device,
134
+ query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
135
+ )
136
+ pooler_output = self.model.pooler(hidden_states, pooling_metadata)
137
+
138
+ finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()
139
+ return pooler_output, finished_mask
140
+
141
+ def _dummy_pooler_run_task(
142
+ self, hidden_states: torch.Tensor, task: PoolingTask
143
+ ) -> PoolerOutput:
144
+ num_tokens = hidden_states.shape[0]
145
+ num_reqs = min(num_tokens, self.max_num_reqs)
146
+ base_tokens = num_tokens // num_reqs
147
+ num_extra = num_tokens % num_reqs
148
+ num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
149
+ if num_extra > 0:
150
+ num_scheduled_tokens[-num_extra:] += 1
151
+ prompt_lens = torch.from_numpy(num_scheduled_tokens)
152
40
- prompt_len = req_states.prompt_len.gpu[input_batch.idx_mapping]
41
- is_valid = input_batch.seq_lens == prompt_len
42
- return last_hidden_states, is_valid
153
+ pooling_params = PoolingParams(task=task)
154
+ pooling_params.verify(self.model_config)
155
+ self.model.pooler.get_pooling_updates(task).apply(pooling_params)
156
+ prompt_token_ids = None
157
+ if pooling_params.requires_token_ids:
158
+ prompt_token_ids = torch.zeros(
159
+ (num_reqs, int(prompt_lens.max())),
160
+ dtype=torch.int64,
161
+ device=hidden_states.device,
162
+ )
163
+ pooling_metadata = PoolingMetadata(
164
+ prompt_lens=prompt_lens,
165
+ prompt_token_ids=prompt_token_ids,
166
+ prompt_token_ids_cpu=None
167
+ if prompt_token_ids is None
168
+ else prompt_token_ids.cpu(),
169
+ pooling_params=[pooling_params] * num_reqs,
170
+ pooling_states=[PoolingStates() for _ in range(num_reqs)],
171
+ )
172
+ pooling_metadata.build_pooling_cursor(
173
+ num_scheduled_tokens,
174
+ seq_lens_cpu=prompt_lens,
175
+ device=hidden_states.device,
176
+ )
177
+ try:
178
+ return self.model.pooler(hidden_states, pooling_metadata)
179
+ except RuntimeError as e:
180
+ if "out of memory" not in str(e):
181
+ raise
182
+ raise RuntimeError(
183
+ "CUDA out of memory occurred when warming up pooler "
184
+ f"({task=}) with {num_reqs} dummy requests. Please try "
185
+ "lowering `max_num_seqs` or `gpu_memory_utilization` when "
186
+ "initializing the engine."
187
+ ) from e
188
189
def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
45
- F.normalize(hidden_states, p=2, dim=-1)
46
- return
190
+ for task in sorted(self.supported_tasks):
191
+ self._dummy_pooler_run_task(hidden_states, task)
vllm/v1/worker/gpu/warmup.py
+5
-1
@@ -230,7 +230,11 @@ def warmup_kernels(
230
# SamplingParams exercising all sampling features.
231
if model_runner.is_pooling_model:
232
sampling_params = None
233
- pooling_params = PoolingParams()
233
+ pooling_task = model_runner.model_config.get_pooling_task(
234
+ model_runner.get_supported_tasks()
235
+ )
236
+ pooling_params = PoolingParams(task=pooling_task)
237
+ pooling_params.verify(model_runner.model_config)
238
else:
239
sampling_params = SamplingParams.for_sampler_warmup()
240
pooling_params = None