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 14, 05:30 UTC
[V0 Deprecation] Remove pooling model support in V0 (#23434)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Maximilien de Bayser committed
Aug 29, 2025 at 04:04 UTC
2554b27baa58b15843367f92d7f73d71bb89033d
38 files changed
+99
-808
tests/distributed/test_pipeline_parallel.py
+6
-2
@@ -118,6 +118,8 @@ class PPTestSettings:
118
multi_node_only: bool = False,
119
load_format: Optional[str] = None,
120
):
121
+ vllm_major_versions = ["1"] if runner == "pooling" else ["0"]
122
+
123
return PPTestSettings(
124
parallel_setups=[
125
ParallelSetup(tp_size=tp_base,
@@ -126,7 +128,7 @@ class PPTestSettings:
128
chunked_prefill=False),
129
],
130
distributed_backends=["mp"],
129
- vllm_major_versions=["0"],
131
+ vllm_major_versions=vllm_major_versions,
132
runner=runner,
133
test_options=PPTestOptions(multi_node_only=multi_node_only,
134
load_format=load_format),
@@ -213,7 +215,9 @@ TEXT_GENERATION_MODELS = {
215
EMBEDDING_MODELS = { # type: ignore[var-annotated]
216
# [Text-only]
217
"intfloat/e5-mistral-7b-instruct": PPTestSettings.fast(runner="pooling"),
216
- "BAAI/bge-multilingual-gemma2": PPTestSettings.fast(runner="pooling"),
218
+ # TODO: re-enable when https://github.com/vllm-project/vllm/issues/23883
219
+ # is fixed
220
+ #"BAAI/bge-multilingual-gemma2": PPTestSettings.fast(runner="pooling"),
221
"Qwen/Qwen2.5-Math-RM-72B": PPTestSettings.fast(
222
load_format="dummy", runner="pooling"
223
),
tests/entrypoints/llm/test_classify.py
-8
@@ -16,14 +16,6 @@ MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
16
prompts = ["The chef prepared a delicious meal."]
17
18
19
-@pytest.fixture(autouse=True)
20
-def v1(run_with_both_engines):
21
- # Simple autouse wrapper to run both engines for each test
22
- # This can be promoted up to conftest.py to run for every
23
- # test in a package
24
- pass
25
-
26
-
19
@pytest.fixture(scope="module")
20
def llm():
21
# pytest caches the fixture so we use weakref.proxy to
tests/entrypoints/llm/test_encode.py
-8
@@ -27,14 +27,6 @@ TOKEN_IDS = [
27
]
28
29
30
-@pytest.fixture(autouse=True)
31
-def v1(run_with_both_engines):
32
- # Simple autouse wrapper to run both engines for each test
33
- # This can be promoted up to conftest.py to run for every
34
- # test in a package
35
- pass
36
-
37
-
30
@pytest.fixture(scope="module")
31
def llm():
32
# pytest caches the fixture so we use weakref.proxy to
tests/entrypoints/llm/test_reward.py
-8
@@ -16,14 +16,6 @@ MODEL_NAME = "internlm/internlm2-1_8b-reward"
16
prompts = ["The chef prepared a delicious meal."]
17
18
19
-@pytest.fixture(autouse=True)
20
-def v1(run_with_both_engines):
21
- # Simple autouse wrapper to run both engines for each test
22
- # This can be promoted up to conftest.py to run for every
23
- # test in a package
24
- pass
25
-
26
-
19
@pytest.fixture(scope="module")
20
def llm():
21
# pytest caches the fixture so we use weakref.proxy to
tests/entrypoints/llm/test_score.py
-8
@@ -14,14 +14,6 @@ from ...models.utils import softmax
14
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
15
16
17
-@pytest.fixture(autouse=True)
18
-def v1(run_with_both_engines):
19
- # Simple autouse wrapper to run both engines for each test
20
- # This can be promoted up to conftest.py to run for every
21
- # test in a package
22
- pass
23
-
24
-
17
@pytest.fixture(scope="module")
18
def llm():
19
# pytest caches the fixture so we use weakref.proxy to
tests/entrypoints/offline_mode/test_offline_mode.py
+10
-9
@@ -32,15 +32,16 @@ MODEL_CONFIGS = [
32
"tensor_parallel_size": 1,
33
"tokenizer_mode": "mistral",
34
},
35
- {
36
- "model": "sentence-transformers/all-MiniLM-L12-v2",
37
- "enforce_eager": True,
38
- "gpu_memory_utilization": 0.20,
39
- "max_model_len": 64,
40
- "max_num_batched_tokens": 64,
41
- "max_num_seqs": 64,
42
- "tensor_parallel_size": 1,
43
- },
35
+ # TODO: re-enable once these tests are run with V1
36
+ # {
37
+ # "model": "sentence-transformers/all-MiniLM-L12-v2",
38
+ # "enforce_eager": True,
39
+ # "gpu_memory_utilization": 0.20,
40
+ # "max_model_len": 64,
41
+ # "max_num_batched_tokens": 64,
42
+ # "max_num_seqs": 64,
43
+ # "tensor_parallel_size": 1,
44
+ # },
45
]
46
47
tests/entrypoints/openai/test_embedding.py
-8
@@ -24,14 +24,6 @@ DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' +
24
DTYPE = "bfloat16"
25
26
27
-@pytest.fixture(autouse=True)
28
-def v1(run_with_both_engines):
29
- # Simple autouse wrapper to run both engines for each test
30
- # This can be promoted up to conftest.py to run for every
31
- # test in a package
32
- pass
33
-
34
-
27
@pytest.fixture(scope="module")
28
def server():
29
args = [
tests/entrypoints/openai/test_rerank.py
-8
@@ -14,14 +14,6 @@ MODEL_NAME = "BAAI/bge-reranker-base"
14
DTYPE = "bfloat16"
15
16
17
-@pytest.fixture(autouse=True)
18
-def v1(run_with_both_engines):
19
- # Simple autouse wrapper to run both engines for each test
20
- # This can be promoted up to conftest.py to run for every
21
- # test in a package
22
- pass
23
-
24
-
17
@pytest.fixture(scope="module")
18
def server():
19
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
tests/entrypoints/openai/test_score.py
-9
@@ -12,15 +12,6 @@ from vllm.entrypoints.openai.protocol import ScoreResponse
12
13
from ...utils import RemoteOpenAIServer
14
15
-
16
-@pytest.fixture(autouse=True)
17
-def v1(run_with_both_engines):
18
- # Simple autouse wrapper to run both engines for each test
19
- # This can be promoted up to conftest.py to run for every
20
- # test in a package
21
- pass
22
-
23
-
15
MODELS = [
16
{
17
"name": "BAAI/bge-reranker-v2-m3",
tests/models/language/pooling/test_embedding.py
+3
-17
@@ -10,14 +10,6 @@ from vllm.platforms import current_platform
10
from ...utils import check_embeddings_close, check_transformers_version
11
12
13
-@pytest.fixture(autouse=True)
14
-def v1(run_with_both_engines):
15
- # Simple autouse wrapper to run both engines for each test
16
- # This can be promoted up to conftest.py to run for every
17
- # test in a package
18
- pass
19
-
20
-
13
@pytest.mark.parametrize(
14
"model",
15
[
@@ -32,21 +24,15 @@ def v1(run_with_both_engines):
24
"intfloat/e5-mistral-7b-instruct",
25
# CPU v1 doesn't support sliding window
26
marks=[pytest.mark.core_model]),
35
- # the qwen models interfere with each other (see PR
36
- # https://github.com/vllm-project/vllm/pull/18720).
37
- # To avoid this problem, for now we skip v0 since it will be
38
- # deprecated anyway.
27
pytest.param("ssmits/Qwen2-7B-Instruct-embed-base",
40
- marks=[pytest.mark.skip_v0, pytest.mark.cpu_model]),
28
+ marks=[pytest.mark.cpu_model]),
29
# [Encoder-only]
30
pytest.param("BAAI/bge-base-en-v1.5", marks=[pytest.mark.core_model]),
31
pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
32
pytest.param("intfloat/multilingual-e5-small"),
45
- pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
46
- marks=[pytest.mark.skip_v1]),
33
+ pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct"),
34
# [Cross-Encoder]
48
- pytest.param("sentence-transformers/stsb-roberta-base-v2",
49
- marks=[pytest.mark.skip_v1]),
35
+ pytest.param("sentence-transformers/stsb-roberta-base-v2"),
36
],
37
)
38
def test_models(
tests/models/language/pooling/test_reward.py
-8
@@ -13,14 +13,6 @@ from ....conftest import HfRunner
13
from ...utils import check_transformers_version
14
15
16
-@pytest.fixture(autouse=True)
17
-def v1(run_with_both_engines):
18
- # Simple autouse wrapper to run both engines for each test
19
- # This can be promoted up to conftest.py to run for every
20
- # test in a package
21
- pass
22
-
23
-
16
@pytest.fixture
17
def math_step_prompts():
18
# ruff: noqa: E501
tests/models/language/pooling/test_scoring.py
-9
@@ -23,15 +23,6 @@ TEXTS_2 = [
23
"The capital of Germany is Berlin.",
24
]
25
26
-
27
-@pytest.fixture(autouse=True)
28
-def v1(run_with_both_engines):
29
- # Simple autouse wrapper to run both engines for each test
30
- # This can be promoted up to conftest.py to run for every
31
- # test in a package
32
- pass
33
-
34
-
26
DTYPE = "half"
27
28
tests/models/registry.py
+11
-12
@@ -323,8 +323,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
323
324
_EMBEDDING_EXAMPLE_MODELS = {
325
# [Text-only]
326
- "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5", v0_only=True),
327
- "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2", v0_only=True), # noqa: E501
326
+ "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
327
+ "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), # noqa: E501
328
"GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
329
"GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
330
trust_remote_code=True),
@@ -337,9 +337,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
337
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
338
"MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
339
"ModernBertModel": _HfExamplesInfo("Alibaba-NLP/gte-modernbert-base",
340
- trust_remote_code=True, v0_only=True),
340
+ trust_remote_code=True),
341
"NomicBertModel": _HfExamplesInfo("nomic-ai/nomic-embed-text-v2-moe",
342
- trust_remote_code=True, v0_only=True), # noqa: E501
342
+ trust_remote_code=True), # noqa: E501
343
"Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
344
"Qwen2ForRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-RM-72B",
345
max_transformers_version="4.53",
@@ -347,9 +347,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
347
"Qwen2ForProcessRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-PRM-7B",
348
max_transformers_version="4.53",
349
transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers"), # noqa: E501
350
- "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2", v0_only=True), # noqa: E501
351
- "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1", v0_only=True), # noqa: E501
352
- "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small", v0_only=True), # noqa: E501
350
+ "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"), # noqa: E501
351
+ "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"), # noqa: E501
352
+ "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"), # noqa: E501
353
# [Multimodal]
354
"LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
355
"Phi3VForCausalLM": _HfExamplesInfo("TIGER-Lab/VLM2Vec-Full",
@@ -364,20 +364,19 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
364
"GPT2ForSequenceClassification": _HfExamplesInfo("nie3e/sentiment-polish-gpt2-small"), # noqa: E501
365
366
# [Cross-encoder]
367
- "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2", v0_only=True), # noqa: E501
367
+ "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2"), # noqa: E501
368
"GteNewForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-multilingual-reranker-base", # noqa: E501
369
trust_remote_code=True,
370
hf_overrides={
371
"architectures": ["GteNewForSequenceClassification"]}),# noqa: E501
372
- "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base", v0_only=True), # noqa: E501
373
- "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base", v0_only=True), # noqa: E501
374
- "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3", v0_only=True), # noqa: E501
372
+ "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base"), # noqa: E501
373
+ "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base"), # noqa: E501
374
+ "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3"), # noqa: E501
375
}
376
377
_AUTOMATIC_CONVERTED_MODELS = {
378
# Use as_seq_cls_model for automatic conversion
379
"GemmaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-gemma", # noqa: E501
380
- v0_only=True,
380
hf_overrides={"architectures": ["GemmaForSequenceClassification"], # noqa: E501
381
"classifier_from_token": ["Yes"], # noqa: E501
382
"method": "no_post_processing"}), # noqa: E501
tests/worker/test_model_input.py
-54
@@ -9,10 +9,7 @@ from vllm.attention import AttentionMetadata, AttentionMetadataBuilder
9
from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.attention.backends.utils import CommonAttentionState
11
from vllm.model_executor import SamplingMetadata
12
-from vllm.model_executor.pooling_metadata import PoolingMetadata
12
from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
14
-from vllm.worker.pooling_model_runner import (
15
- ModelInputForGPUWithPoolingMetadata)
13
14
15
class MockAttentionBackend(AttentionBackend):
@@ -114,54 +111,3 @@ def test_model_runner_input():
111
assert (received_model_input.sampling_metadata.selected_token_indices ==
112
sampling_metadata.selected_token_indices)
113
assert received_model_input.sampling_metadata.seq_groups is None
117
-
118
-
119
-def test_embedding_model_runner_input():
120
- pooling_metadata = PoolingMetadata(
121
- seq_groups=[[0]],
122
- seq_data={},
123
- prompt_lens=[1],
124
- )
125
- attn_metadata = AttentionMetadata(
126
- num_prefills=1,
127
- num_prefill_tokens=2,
128
- num_decode_tokens=3,
129
- slot_mapping=torch.zeros(1),
130
- multi_modal_placeholder_index_maps=None,
131
- enable_kv_scales_calculation=True,
132
- )
133
- model_input = ModelInputForGPUWithPoolingMetadata(
134
- input_tokens=torch.ones(10),
135
- input_positions=torch.ones(10),
136
- pooling_metadata=pooling_metadata,
137
- attn_metadata=attn_metadata)
138
-
139
- assert isinstance(model_input, ModelInputForGPUWithPoolingMetadata)
140
-
141
- # Test round trip serialization.
142
- tensor_dict = model_input.as_broadcastable_tensor_dict()
143
- attn_backend = MockAttentionBackend()
144
- received_model_input = (
145
- ModelInputForGPUWithPoolingMetadata.from_broadcasted_tensor_dict(
146
- tensor_dict, attn_backend=attn_backend))
147
- # Check that received copy has correct values.
148
- assert isinstance(received_model_input,
149
- ModelInputForGPUWithPoolingMetadata)
150
- assert received_model_input.input_tokens is not None
151
- assert (
152
- received_model_input.input_tokens == model_input.input_tokens).all()
153
- assert received_model_input.input_positions is not None
154
- assert (received_model_input.input_positions == model_input.input_positions
155
- ).all()
156
- assert received_model_input.multi_modal_kwargs is None
157
- assert (received_model_input.multi_modal_kwargs ==
158
- model_input.multi_modal_kwargs)
159
- assert received_model_input.lora_requests is None
160
- assert received_model_input.lora_requests == model_input.lora_requests
161
- assert received_model_input.lora_mapping is None
162
- assert received_model_input.lora_mapping == model_input.lora_mapping
163
- for field in dataclasses.fields(AttentionMetadata):
164
- assert getattr(received_model_input.attn_metadata, field.name,
165
- None) == getattr(attn_metadata, field.name, None)
166
- # Pooling metadata is not broadcast.
167
- assert received_model_input.pooling_metadata is None
vllm/core/scheduler.py
-1
@@ -1591,7 +1591,6 @@ class Scheduler:
1591
encoder_seq_data=encoder_seq_data,
1592
cross_block_table=cross_block_table,
1593
state=seq_group.state,
1594
- token_type_ids=seq_group.token_type_ids,
1594
# `multi_modal_data` will only be present for the 1st comm
1595
# between engine and worker.
1596
# the subsequent comms can still use delta, but
vllm/engine/arg_utils.py
+1
-6
@@ -1566,8 +1566,7 @@ class EngineArgs:
1566
use_spec_decode = self.speculative_config is not None
1567
1568
if (is_gpu and not use_sliding_window and not use_spec_decode
1569
- and not self.enable_lora
1570
- and model_config.runner_type != "pooling"):
1569
+ and not self.enable_lora):
1570
self.enable_chunked_prefill = True
1571
logger.warning(
1572
"Chunked prefill is enabled by default for models "
@@ -1585,10 +1584,6 @@ class EngineArgs:
1584
"OOM during the initial memory profiling phase, or result "
1585
"in low performance due to small KV cache size. Consider "
1586
"setting --max-model-len to a smaller value.", max_model_len)
1588
- elif (self.enable_chunked_prefill
1589
- and model_config.runner_type == "pooling"):
1590
- msg = "Chunked prefill is not supported for pooling models"
1591
- raise ValueError(msg)
1587
1588
# if using prefix caching, we must set a hash algo
1589
if self.enable_prefix_caching:
vllm/engine/async_llm_engine.py
+14
-97
@@ -72,8 +72,8 @@ STOP_ITERATION = Exception() # Sentinel
72
73
74
class AsyncStream:
75
- """A stream of RequestOutputs or PoolingRequestOutputs for a request
76
- that can be iterated over asynchronously via an async generator."""
75
+ """A stream of RequestOutputs for a request that can be iterated over
76
+ asynchronously via an async generator."""
77
78
def __init__(self, request_id: str, cancel: Callable[[str], None]) -> None:
79
self.request_id = request_id
@@ -81,8 +81,7 @@ class AsyncStream:
81
self._queue: asyncio.Queue = asyncio.Queue()
82
self._finished = False
83
84
- def put(self, item: Union[RequestOutput, PoolingRequestOutput,
85
- Exception]) -> None:
84
+ def put(self, item: Union[RequestOutput, Exception]) -> None:
85
if not self._finished:
86
self._queue.put_nowait(item)
87
@@ -99,9 +98,7 @@ class AsyncStream:
98
def finished(self) -> bool:
99
return self._finished
100
102
- async def generator(
103
- self
104
- ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:
101
+ async def generator(self) -> AsyncGenerator[RequestOutput, None]:
102
try:
103
while True:
104
result = await self._queue.get()
@@ -151,8 +148,7 @@ class RequestTracker:
148
self.abort_request(rid, exception=exc)
149
150
def process_request_output(self,
154
- request_output: Union[RequestOutput,
155
- PoolingRequestOutput],
151
+ request_output: RequestOutput,
152
*,
153
verbose: bool = False) -> None:
154
"""Process a request output from the engine."""
@@ -261,9 +257,7 @@ class _AsyncLLMEngine(LLMEngine):
257
def __init__(self, *args, **kwargs):
258
super().__init__(*args, **kwargs)
259
264
- async def step_async(
265
- self, virtual_engine: int
266
- ) -> List[Union[RequestOutput, PoolingRequestOutput]]:
260
+ async def step_async(self, virtual_engine: int) -> List[RequestOutput]:
261
"""Performs one decoding iteration and returns newly generated results.
262
The workers are ran asynchronously if possible.
263
@@ -405,7 +399,7 @@ class _AsyncLLMEngine(LLMEngine):
399
self,
400
request_id: str,
401
prompt: PromptType,
408
- params: Union[SamplingParams, PoolingParams],
402
+ params: SamplingParams,
403
arrival_time: Optional[float] = None,
404
lora_request: Optional[LoRARequest] = None,
405
trace_headers: Optional[Mapping[str, str]] = None,
@@ -779,14 +773,14 @@ class AsyncLLMEngine(EngineClient):
773
self,
774
request_id: str,
775
prompt: PromptType,
782
- params: Union[SamplingParams, PoolingParams],
776
+ params: SamplingParams,
777
arrival_time: Optional[float] = None,
778
lora_request: Optional[LoRARequest] = None,
779
trace_headers: Optional[Mapping[str, str]] = None,
780
priority: int = 0,
781
data_parallel_rank: Optional[int] = None,
782
tokenization_kwargs: Optional[dict[str, Any]] = None,
789
- ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:
783
+ ) -> AsyncGenerator[RequestOutput, None]:
784
if not self.is_running:
785
if self.start_engine_loop:
786
self.start_background_loop()
@@ -908,7 +902,7 @@ class AsyncLLMEngine(EngineClient):
902
await self.abort(request_id)
903
raise
904
911
- async def encode(
905
+ def encode(
906
self,
907
prompt: PromptType,
908
pooling_params: PoolingParams,
@@ -918,85 +912,8 @@ class AsyncLLMEngine(EngineClient):
912
priority: int = 0,
913
tokenization_kwargs: Optional[dict[str, Any]] = None,
914
) -> AsyncGenerator[PoolingRequestOutput, None]:
921
- """Generate outputs for a request from a pooling model.
922
-
923
- Generate outputs for a request. This method is a coroutine. It adds the
924
- request into the waiting queue of the LLMEngine and streams the outputs
925
- from the LLMEngine to the caller.
926
-
927
- Args:
928
- prompt: The prompt to the LLM. See
929
- [`PromptType`][vllm.inputs.PromptType] for more details about
930
- the format of each input.
931
- pooling_params: The pooling parameters of the request.
932
- request_id: The unique id of the request.
933
- lora_request: LoRA request to use for generation, if any.
934
- trace_headers: OpenTelemetry trace headers.
935
- priority: The priority of the request.
936
- Only applicable with priority scheduling.
937
-
938
- Yields:
939
- The output `PoolingRequestOutput` objects from the LLMEngine
940
- for the request.
941
-
942
- Details:
943
- - If the engine is not running, start the background loop,
944
- which iteratively invokes
945
- [`vllm.engine.async_llm_engine.AsyncLLMEngine.engine_step`][]
946
- to process the waiting requests.
947
- - Add the request to the engine's `RequestTracker`.
948
- On the next background loop, this request will be sent to
949
- the underlying engine.
950
- Also, a corresponding `AsyncStream` will be created.
951
- - Wait for the request outputs from `AsyncStream` and yield them.
952
-
953
- Example:
954
- ```
955
- # Please refer to entrypoints/api_server.py for
956
- # the complete example.
957
-
958
- # initialize the engine and the example input
959
- # note that engine_args here is AsyncEngineArgs instance
960
- engine = AsyncLLMEngine.from_engine_args(engine_args)
961
- example_input = {
962
- "input": "What is LLM?",
963
- "request_id": 0,
964
- }
965
-
966
- # start the generation
967
- results_generator = engine.encode(
968
- example_input["input"],
969
- PoolingParams(),
970
- example_input["request_id"])
971
-
972
- # get the results
973
- final_output = None
974
- async for request_output in results_generator:
975
- if await request.is_disconnected():
976
- # Abort the request if the client disconnects.
977
- await engine.abort(request_id)
978
- # Return or raise an error
979
- ...
980
- final_output = request_output
981
-
982
- # Process and return the final output
983
- ...
984
- ```
985
- """
986
- try:
987
- async for output in await self.add_request(
988
- request_id,
989
- prompt,
990
- pooling_params,
991
- lora_request=lora_request,
992
- trace_headers=trace_headers,
993
- priority=priority,
994
- tokenization_kwargs=tokenization_kwargs,
995
- ):
996
- yield LLMEngine.validate_output(output, PoolingRequestOutput)
997
- except asyncio.CancelledError:
998
- await self.abort(request_id)
999
- raise
915
+ raise NotImplementedError(
916
+ "Pooling models are not supported in vLLM V0")
917
918
async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
919
"""Abort a request.
@@ -1104,8 +1021,8 @@ class AsyncLLMEngine(EngineClient):
1021
async def is_sleeping(self) -> bool:
1022
return self.engine.is_sleeping()
1023
1107
- async def add_lora(self, lora_request: LoRARequest) -> None:
1108
- self.engine.add_lora(lora_request)
1024
+ async def add_lora(self, lora_request: LoRARequest) -> bool:
1025
+ return self.engine.add_lora(lora_request)
1026
1027
async def collective_rpc(self,
1028
method: str,
vllm/engine/llm_engine.py
+15
-67
@@ -40,12 +40,11 @@ from vllm.multimodal.cache import processor_only_cache_from_config
40
from vllm.multimodal.processing import EncDecMultiModalProcessor
41
from vllm.outputs import (PoolingRequestOutput, RequestOutput,
42
RequestOutputFactory)
43
-from vllm.pooling_params import PoolingParams
43
from vllm.sampling_params import RequestOutputKind, SamplingParams
44
from vllm.sequence import (ExecuteModelRequest, ParallelSampleSequenceGroup,
46
- PoolingSequenceGroupOutput, Sequence, SequenceGroup,
47
- SequenceGroupBase, SequenceGroupMetadata,
48
- SequenceGroupOutput, SequenceStatus)
45
+ Sequence, SequenceGroup, SequenceGroupBase,
46
+ SequenceGroupMetadata, SequenceGroupOutput,
47
+ SequenceStatus)
48
from vllm.tracing import (SpanAttributes, SpanKind, extract_trace_context,
49
init_tracer)
50
from vllm.transformers_utils.detokenizer import Detokenizer
@@ -93,8 +92,7 @@ class SchedulerContext:
92
93
def __init__(self) -> None:
94
self.output_queue: Deque[OutputData] = deque()
96
- self.request_outputs: List[Union[RequestOutput,
97
- PoolingRequestOutput]] = []
95
+ self.request_outputs: List[RequestOutput] = []
96
self.seq_group_metadata_list: Optional[
97
List[SequenceGroupMetadata]] = None
98
self.scheduler_outputs: Optional[SchedulerOutputs] = None
@@ -261,8 +259,7 @@ class LLMEngine:
259
260
self.model_executor = executor_class(vllm_config=vllm_config)
261
264
- if self.model_config.runner_type != "pooling":
265
- self._initialize_kv_caches()
262
+ self._initialize_kv_caches()
263
264
# If usage stat is enabled, collect relevant info.
265
if is_usage_stats_enabled():
@@ -541,7 +538,7 @@ class LLMEngine:
538
self,
539
request_id: str,
540
processed_inputs: ProcessorInputs,
544
- params: Union[SamplingParams, PoolingParams],
541
+ params: SamplingParams,
542
arrival_time: float,
543
lora_request: Optional[LoRARequest],
544
trace_headers: Optional[Mapping[str, str]] = None,
@@ -577,7 +574,7 @@ class LLMEngine:
574
encoder_seq = (None if encoder_inputs is None else Sequence(
575
seq_id, encoder_inputs, block_size, eos_token_id, lora_request))
576
580
- # Create a SequenceGroup based on SamplingParams or PoolingParams
577
+ # Create a SequenceGroup based on SamplingParams
578
if isinstance(params, SamplingParams):
579
seq_group = self._create_sequence_group_with_sampling(
580
request_id,
@@ -588,18 +585,8 @@ class LLMEngine:
585
trace_headers=trace_headers,
586
encoder_seq=encoder_seq,
587
priority=priority)
591
- elif isinstance(params, PoolingParams):
592
- seq_group = self._create_sequence_group_with_pooling(
593
- request_id,
594
- seq,
595
- params,
596
- arrival_time=arrival_time,
597
- lora_request=lora_request,
598
- encoder_seq=encoder_seq,
599
- priority=priority)
588
else:
601
- raise ValueError(
602
- "Either SamplingParams or PoolingParams must be provided.")
589
+ raise ValueError("SamplingParams must be provided.")
590
591
# Add the sequence group to the scheduler with least unfinished seqs.
592
costs = [
@@ -618,7 +605,7 @@ class LLMEngine:
605
self,
606
request_id: str,
607
prompt: PromptType,
621
- params: Union[SamplingParams, PoolingParams],
608
+ params: SamplingParams,
609
arrival_time: Optional[float] = None,
610
lora_request: Optional[LoRARequest] = None,
611
tokenization_kwargs: Optional[dict[str, Any]] = None,
@@ -636,9 +623,8 @@ class LLMEngine:
623
prompt: The prompt to the LLM. See
624
[PromptType][vllm.inputs.PromptType]
625
for more details about the format of each input.
639
- params: Parameters for sampling or pooling.
626
+ params: Parameters for sampling.
627
[SamplingParams][vllm.SamplingParams] for text generation.
641
- [PoolingParams][vllm.PoolingParams] for pooling.
628
arrival_time: The arrival time of the request. If None, we use
629
the current monotonic time.
630
lora_request: The LoRA request to add.
@@ -760,29 +746,6 @@ class LLMEngine:
746
747
return seq_group
748
763
- def _create_sequence_group_with_pooling(
764
- self,
765
- request_id: str,
766
- seq: Sequence,
767
- pooling_params: PoolingParams,
768
- arrival_time: float,
769
- lora_request: Optional[LoRARequest],
770
- encoder_seq: Optional[Sequence] = None,
771
- priority: int = 0,
772
- ) -> SequenceGroup:
773
- """Creates a SequenceGroup with PoolingParams."""
774
- # Defensive copy of PoolingParams, which are used by the pooler
775
- pooling_params = pooling_params.clone()
776
- # Create the sequence group.
777
- seq_group = SequenceGroup(request_id=request_id,
778
- seqs=[seq],
779
- arrival_time=arrival_time,
780
- lora_request=lora_request,
781
- pooling_params=pooling_params,
782
- encoder_seq=encoder_seq,
783
- priority=priority)
784
- return seq_group
785
-
749
def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
750
"""Aborts a request(s) with the given ID.
751
@@ -856,18 +819,6 @@ class LLMEngine:
819
success = success and scheduler.reset_prefix_cache(device)
820
return success
821
859
- @staticmethod
860
- def _process_sequence_group_outputs(
861
- seq_group: SequenceGroup,
862
- outputs: List[PoolingSequenceGroupOutput],
863
- ) -> None:
864
- seq_group.pooled_data = outputs[0].data
865
-
866
- for seq in seq_group.get_seqs():
867
- seq.status = SequenceStatus.FINISHED_STOPPED
868
-
869
- return
870
-
822
def _process_model_outputs(self,
823
ctx: SchedulerContext,
824
request_id: Optional[str] = None) -> None:
@@ -962,13 +913,10 @@ class LLMEngine:
913
seq_group.metrics.model_execute_time = (
914
o.model_execute_time)
915
965
- if self.model_config.runner_type == "pooling":
966
- self._process_sequence_group_outputs(seq_group, output)
967
- else:
968
- self.output_processor.process_prompt_logprob(seq_group, output)
969
- if seq_group_meta.do_sample:
970
- self.output_processor.process_outputs(
971
- seq_group, output, is_async)
916
+ self.output_processor.process_prompt_logprob(seq_group, output)
917
+ if seq_group_meta.do_sample:
918
+ self.output_processor.process_outputs(seq_group, output,
919
+ is_async)
920
921
if seq_group.is_finished():
922
finished_now.append(i)
@@ -1090,7 +1038,7 @@ class LLMEngine:
1038
seq.append_token_id(sample.output_token, sample.logprobs,
1039
sample.output_embed)
1040
1093
- def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
1041
+ def step(self) -> List[RequestOutput]:
1042
"""Performs one decoding iteration and returns newly generated results.
1043
1044
<figure markdown="span">
vllm/engine/multiprocessing/__init__.py
+1
@@ -120,6 +120,7 @@ class RPCLoadAdapterRequest:
120
@dataclass
121
class RPCAdapterLoadedResponse:
122
request_id: str
123
+ lora_loaded: bool
124
125
126
RPC_REQUEST_T = Union[RPCProcessRequest, RPCAbortRequest, RPCStartupRequest,
vllm/engine/multiprocessing/client.py
+13
-38
@@ -6,7 +6,7 @@ import copy
6
import pickle
7
from contextlib import contextmanager, suppress
8
from typing import (Any, AsyncGenerator, Dict, Iterable, Iterator, List,
9
- Mapping, Optional, Union, cast)
9
+ Mapping, Optional, Union)
10
11
import cloudpickle
12
import psutil
@@ -477,10 +477,8 @@ class MQLLMEngineClient(EngineClient):
477
Any priority other than 0 will lead to an error if the
478
scheduling policy is not "priority".
479
"""
480
- return cast(
481
- AsyncGenerator[RequestOutput, None],
482
- self._process_request(prompt, sampling_params, request_id,
483
- lora_request, trace_headers, priority))
480
+ return self._process_request(prompt, sampling_params, request_id,
481
+ lora_request, trace_headers, priority)
482
483
def encode(
484
self,
@@ -490,45 +488,20 @@ class MQLLMEngineClient(EngineClient):
488
lora_request: Optional[LoRARequest] = None,
489
trace_headers: Optional[Mapping[str, str]] = None,
490
priority: int = 0,
491
+ tokenization_kwargs: Optional[dict[str, Any]] = None,
492
) -> AsyncGenerator[PoolingRequestOutput, None]:
494
- """Generate outputs for a request from a pooling model.
495
-
496
- Generate outputs for a request. This method is a coroutine. It adds the
497
- request into the waiting queue of the LLMEngine and streams the outputs
498
- from the LLMEngine to the caller.
499
-
500
- Args:
501
- prompt: The prompt to the LLM. See
502
- [`PromptType`][vllm.inputs.PromptType] for more details about
503
- the format of each input.
504
- pooling_params: The pooling parameters of the request.
505
- request_id: The unique id of the request.
506
- lora_request: LoRA request to use for generation, if any.
507
- trace_headers: OpenTelemetry trace headers.
508
-
509
- Yields:
510
- The output `PoolingRequestOutput` objects from the LLMEngine
511
- for the request.
512
- """
513
- return cast(
514
- AsyncGenerator[PoolingRequestOutput, None],
515
- self._process_request(prompt,
516
- pooling_params,
517
- request_id,
518
- lora_request,
519
- trace_headers,
520
- priority=priority))
493
+ raise NotImplementedError(
494
+ "Pooling models are not supported in vLLM V0")
495
496
async def _process_request(
497
self,
498
prompt: PromptType,
525
- params: Union[SamplingParams, PoolingParams],
499
+ params: SamplingParams,
500
request_id: str,
501
lora_request: Optional[LoRARequest] = None,
502
trace_headers: Optional[Mapping[str, str]] = None,
503
priority: int = 0,
530
- ) -> Union[AsyncGenerator[RequestOutput, None], AsyncGenerator[
531
- PoolingRequestOutput, None]]:
504
+ ) -> AsyncGenerator[RequestOutput, None]:
505
"""Send an RPCGenerateRequest to the RPCServer and stream responses."""
506
507
# If already dead, error out.
@@ -547,7 +520,7 @@ class MQLLMEngineClient(EngineClient):
520
try:
521
# 2) Detach logits processors so that they can be pickled
522
# separately (may require cloudpickle which is slower)
550
- if isinstance(params, SamplingParams) and params.logits_processors:
523
+ if params.logits_processors:
524
# Defensive shallow copy
525
params = copy.copy(params)
526
logits_processors = params.logits_processors
@@ -646,13 +619,14 @@ class MQLLMEngineClient(EngineClient):
619
raise request_output
620
return request_output.is_sleeping
621
649
- async def add_lora(self, lora_request: LoRARequest) -> None:
622
+ async def add_lora(self, lora_request: LoRARequest) -> bool:
623
"""Load a new LoRA adapter into the engine for future requests."""
624
# Uses the same I/O as generate requests
625
request = RPCLoadAdapterRequest(lora_request)
626
627
# Create output queue for this request.
655
- queue: asyncio.Queue[Union[None, BaseException]] = asyncio.Queue()
628
+ queue: asyncio.Queue[Union[
629
+ BaseException, RPCAdapterLoadedResponse]] = asyncio.Queue()
630
self.output_queues[request.request_id] = queue
631
632
# Send the request
@@ -666,3 +640,4 @@ class MQLLMEngineClient(EngineClient):
640
# Raise on error, otherwise happily return None
641
if isinstance(request_output, BaseException):
642
raise request_output
643
+ return request_output.lora_loaded
vllm/engine/multiprocessing/engine.py
+3
-2
@@ -347,7 +347,7 @@ class MQLLMEngine:
347
348
def _handle_load_adapter_request(self, request: RPCLoadAdapterRequest):
349
try:
350
- self.engine.add_lora(request.lora_request)
350
+ lora_loaded = self.engine.add_lora(request.lora_request)
351
except BaseException as e:
352
# Send back an error if the adater fails to load
353
rpc_err = RPCError(request_id=request.request_id,
@@ -357,7 +357,8 @@ class MQLLMEngine:
357
return
358
# Otherwise, send back the successful load message
359
self._send_outputs(
360
- RPCAdapterLoadedResponse(request_id=request.request_id))
360
+ RPCAdapterLoadedResponse(request_id=request.request_id,
361
+ lora_loaded=lora_loaded))
362
363
def _handle_is_sleeping_request(self, request: RPCIsSleepingRequest):
364
is_sleeping = self.is_sleeping()
vllm/engine/protocol.py
+3
-2
@@ -3,7 +3,7 @@
3
4
import asyncio
5
from abc import ABC, abstractmethod
6
-from typing import AsyncGenerator, Iterable, Mapping, Optional, Union
6
+from typing import Any, AsyncGenerator, Iterable, Mapping, Optional, Union
7
8
from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
9
from vllm.config import DecodingConfig, ModelConfig, VllmConfig
@@ -224,6 +224,7 @@ class EngineClient(ABC):
224
lora_request: Optional[LoRARequest] = None,
225
trace_headers: Optional[Mapping[str, str]] = None,
226
priority: int = 0,
227
+ tokenization_kwargs: Optional[dict[str, Any]] = None,
228
) -> AsyncGenerator[PoolingRequestOutput, None]:
229
"""Generate outputs for a request from a pooling model."""
230
...
@@ -320,7 +321,7 @@ class EngineClient(ABC):
321
...
322
323
@abstractmethod
323
- async def add_lora(self, lora_request: LoRARequest) -> None:
324
+ async def add_lora(self, lora_request: LoRARequest) -> bool:
325
"""Load a new LoRA adapter into the engine for future requests."""
326
...
327
vllm/entrypoints/llm.py
+1
-2
@@ -1156,8 +1156,7 @@ class LLM:
1156
tokenization_kwargs=tokenization_kwargs,
1157
)
1158
1159
- if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
1160
- "token_type_ids", None)):
1159
+ if (token_type_ids := engine_prompt.pop("token_type_ids", None)):
1160
params = pooling_params.clone()
1161
compressed = compress_token_type_ids(token_type_ids)
1162
params.extra_kwargs = {"compressed_token_type_ids": compressed}
vllm/entrypoints/openai/serving_score.py
+1
-3
@@ -7,7 +7,6 @@ from typing import Any, Optional, Union
7
8
from fastapi import Request
9
10
-from vllm import envs
10
from vllm.config import ModelConfig
11
from vllm.engine.protocol import EngineClient
12
from vllm.entrypoints.logger import RequestLogger
@@ -229,8 +228,7 @@ class ServingScores(OpenAIServing):
228
params=default_pooling_params,
229
lora_request=lora_request)
230
232
- if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
233
- "token_type_ids", None)):
231
+ if (token_type_ids := engine_prompt.pop("token_type_ids", None)):
232
pooling_params = default_pooling_params.clone()
233
compressed = compress_token_type_ids(token_type_ids)
234
pooling_params.extra_kwargs = {
vllm/inputs/data.py
-6
@@ -174,9 +174,6 @@ class TokenInputs(TypedDict):
174
prompt_token_ids: list[int]
175
"""The token IDs of the prompt."""
176
177
- token_type_ids: NotRequired[list[int]]
178
- """The token type IDs of the prompt."""
179
-
177
prompt: NotRequired[str]
178
"""
179
The original prompt text corresponding to the token IDs, if available.
@@ -190,7 +187,6 @@ class TokenInputs(TypedDict):
187
188
def token_inputs(
189
prompt_token_ids: list[int],
193
- token_type_ids: Optional[list[int]] = None,
190
prompt: Optional[str] = None,
191
cache_salt: Optional[str] = None,
192
) -> TokenInputs:
@@ -200,8 +196,6 @@ def token_inputs(
196
197
if prompt is not None:
198
inputs["prompt"] = prompt
203
- if token_type_ids is not None:
204
- inputs["token_type_ids"] = token_type_ids
199
if cache_salt is not None:
200
inputs["cache_salt"] = cache_salt
201
vllm/inputs/preprocess.py
+2
-10
@@ -355,7 +355,6 @@ class InputPreprocessor:
355
mm_hash_overrides: Optional[dict[str, list[str]]] = None,
356
) -> Union[TokenInputs, MultiModalInputs]:
357
prompt_token_ids = parsed_content["prompt_token_ids"]
358
- token_type_ids = parsed_content.get("token_type_ids")
358
359
inputs: Union[TokenInputs, MultiModalInputs]
360
if multi_modal_data := parsed_content.get("multi_modal_data"):
@@ -368,10 +367,7 @@ class InputPreprocessor:
367
mm_hash_overrides=mm_hash_overrides,
368
)
369
else:
371
- inputs = token_inputs(
372
- prompt_token_ids=prompt_token_ids,
373
- token_type_ids=token_type_ids,
374
- )
370
+ inputs = token_inputs(prompt_token_ids=prompt_token_ids)
371
372
if cache_salt := parsed_content.get("cache_salt"):
373
inputs["cache_salt"] = cache_salt
@@ -387,7 +383,6 @@ class InputPreprocessor:
383
mm_hash_overrides: Optional[dict[str, list[str]]] = None,
384
) -> Union[TokenInputs, MultiModalInputs]:
385
prompt_token_ids = parsed_content["prompt_token_ids"]
390
- token_type_ids = parsed_content.get("token_type_ids")
386
387
inputs: Union[TokenInputs, MultiModalInputs]
388
if multi_modal_data := parsed_content.get("multi_modal_data"):
@@ -400,10 +395,7 @@ class InputPreprocessor:
395
mm_hash_overrides=mm_hash_overrides,
396
)
397
else:
403
- inputs = token_inputs(
404
- prompt_token_ids=prompt_token_ids,
405
- token_type_ids=token_type_ids,
406
- )
398
+ inputs = token_inputs(prompt_token_ids=prompt_token_ids, )
399
400
if cache_salt := parsed_content.get("cache_salt"):
401
inputs["cache_salt"] = cache_salt
vllm/model_executor/layers/pooler.py
+7
-25
@@ -13,17 +13,12 @@ import torch.nn.functional as F
13
from transformers import PretrainedConfig
14
15
from vllm.config import ModelConfig, PoolerConfig
16
-from vllm.model_executor.pooling_metadata import ( # noqa: E501
17
- PoolingMetadata as V0PoolingMetadata)
18
-from vllm.model_executor.pooling_metadata import PoolingTensors
16
from vllm.pooling_params import PoolingParams
17
from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
18
from vllm.tasks import PoolingTask
19
from vllm.utils import current_stream, resolve_obj_by_qualname
23
-from vllm.v1.pool.metadata import PoolingCursor
24
-from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
20
+from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata
21
26
-PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
22
PoolingFn = Callable[
23
[Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
24
Union[torch.Tensor, list[torch.Tensor]]]
@@ -127,36 +122,23 @@ def get_prompt_lens(
122
hidden_states: Union[torch.Tensor, list[torch.Tensor]],
123
pooling_metadata: PoolingMetadata,
124
) -> torch.Tensor:
130
- if isinstance(pooling_metadata, V1PoolingMetadata):
131
- return pooling_metadata.prompt_lens
132
-
133
- return PoolingTensors.from_pooling_metadata(
134
- pooling_metadata, hidden_states[0].device).prompt_lens
125
+ return pooling_metadata.prompt_lens
126
127
128
def get_prompt_token_ids(
129
pooling_metadata: PoolingMetadata) -> list[torch.Tensor]:
139
- if isinstance(pooling_metadata, V1PoolingMetadata):
140
- assert pooling_metadata.prompt_token_ids is not None, (
141
- "Please set `requires_token_ids=True` in `get_pooling_updates`")
142
-
143
- return [
144
- pooling_metadata.prompt_token_ids[i, :num]
145
- for i, num in enumerate(pooling_metadata.prompt_lens)
146
- ]
130
+ assert pooling_metadata.prompt_token_ids is not None, (
131
+ "Please set `requires_token_ids=True` in `get_pooling_updates`")
132
133
return [
149
- torch.tensor(seq_data_i.prompt_token_ids)
150
- for seq_data_i in pooling_metadata.seq_data.values()
134
+ pooling_metadata.prompt_token_ids[i, :num]
135
+ for i, num in enumerate(pooling_metadata.prompt_lens)
136
]
137
138
139
def get_pooling_params(
140
pooling_metadata: PoolingMetadata) -> list[PoolingParams]:
156
- if isinstance(pooling_metadata, V0PoolingMetadata):
157
- pooling_params = [p for _, p in pooling_metadata.seq_groups]
158
- else:
159
- pooling_params = pooling_metadata.pooling_params
141
+ pooling_params = pooling_metadata.pooling_params
142
return pooling_params
143
144
vllm/model_executor/models/bert.py
+1
-1
@@ -24,9 +24,9 @@ from vllm.model_executor.layers.pooler import (ClassifierPooler,
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.layers.vocab_parallel_embedding import (
26
VocabParallelEmbedding)
27
-from vllm.model_executor.pooling_metadata import PoolingMetadata
27
from vllm.sequence import IntermediateTensors
28
from vllm.tasks import PoolingTask
29
+from vllm.v1.pool.metadata import PoolingMetadata
30
31
from .interfaces import SupportsCrossEncoding, SupportsQuant
32
from .interfaces_base import default_pooling_type
vllm/model_executor/models/gritlm.py
+1
-1
@@ -15,10 +15,10 @@ from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
15
build_output, get_prompt_lens,
16
get_prompt_token_ids)
17
from vllm.model_executor.models.llama import LlamaForCausalLM
18
-from vllm.model_executor.pooling_metadata import PoolingMetadata
18
from vllm.sequence import PoolerOutput
19
from vllm.tasks import PoolingTask
20
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
21
+from vllm.v1.pool.metadata import PoolingMetadata
22
23
from .interfaces_base import default_pooling_type
24
vllm/model_executor/models/modernbert.py
+1
-1
@@ -22,9 +22,9 @@ from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
22
from vllm.model_executor.layers.vocab_parallel_embedding import (
23
VocabParallelEmbedding)
24
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
-from vllm.model_executor.pooling_metadata import PoolingMetadata
25
from vllm.sequence import IntermediateTensors
26
from vllm.tasks import PoolingTask
27
+from vllm.v1.pool.metadata import PoolingMetadata
28
29
from .interfaces import SupportsCrossEncoding
30
from .interfaces_base import default_pooling_type
vllm/model_executor/pooling_metadata.py
deleted
-90
@@ -1,90 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass
5
-from typing import Any, Optional
6
-
7
-import torch
8
-
9
-from vllm.pooling_params import PoolingParams
10
-from vllm.utils import is_pin_memory_available
11
-from vllm.v1.pool.metadata import PoolingCursor, build_pooling_cursor
12
-
13
-
14
-class PoolingMetadata:
15
- """Metadata for pooling operations in the Pooler layer.
16
-
17
- This class holds the necessary information for pooling operations,
18
- providing context for how to perform pooling and other related operations.
19
-
20
- Attributes:
21
- seq_groups: List of (seq_ids, pooling_params).
22
- seq_data: A mapping of sequence ID to additional sequence data.
23
- prompt_lens: List of the lengths of each prompt.
24
- """
25
-
26
- def __init__(
27
- self,
28
- seq_groups: list[tuple[list[int], PoolingParams]],
29
- seq_data: dict[int, Any], # Specific data related to sequences
30
- prompt_lens: list[int],
31
- pooling_cursor: Optional[PoolingCursor] = None) -> None:
32
- self.seq_groups = seq_groups
33
- self.seq_data = seq_data
34
- self.prompt_lens = prompt_lens
35
- self.pooling_cursor: Optional[PoolingCursor] = pooling_cursor
36
-
37
- def __repr__(self) -> str:
38
- return ("PoolingMetadata("
39
- f"seq_groups={self.seq_groups}, "
40
- f"seq_data={self.seq_data}, "
41
- f"prompt_lens={self.prompt_lens})")
42
-
43
- def __getitem__(self, indices: slice):
44
- return PoolingMetadata(
45
- seq_groups=self.seq_groups[indices],
46
- seq_data=dict(list(self.seq_data.items())[indices]),
47
- prompt_lens=self.prompt_lens[indices],
48
- pooling_cursor=None
49
- if self.pooling_cursor is None else self.pooling_cursor[indices],
50
- )
51
-
52
- def build_pooling_cursor(self, num_scheduled_tokens: list[int],
53
- device: torch.device):
54
- prompt_lens = torch.tensor(self.prompt_lens, device="cpu")
55
- self.pooling_cursor = build_pooling_cursor(num_scheduled_tokens,
56
- prompt_lens,
57
- device=device)
58
-
59
-
60
-@dataclass
61
-class PoolingTensors:
62
- """Tensors for pooling."""
63
-
64
- prompt_lens: torch.Tensor
65
-
66
- @classmethod
67
- def from_pooling_metadata(
68
- cls,
69
- pooling_metadata: "PoolingMetadata",
70
- device: torch.device,
71
- ) -> "PoolingTensors":
72
- """
73
- Create PoolingTensors from PoolingMetadata.
74
-
75
- Args:
76
- pooling_metadata: PoolingMetadata instance to convert.
77
- device: Device to store the tensors.
78
- """
79
- # Convert prompt lengths to tensor
80
- pin_memory = is_pin_memory_available()
81
-
82
- prompt_lens_t = torch.tensor(
83
- pooling_metadata.prompt_lens,
84
- device="cpu",
85
- dtype=torch.long,
86
- pin_memory=pin_memory,
87
- )
88
-
89
- return cls(prompt_lens=prompt_lens_t.to(device=device,
90
- non_blocking=True), )
vllm/multimodal/inputs.py
-6
@@ -913,9 +913,6 @@ class MultiModalInputs(TypedDict):
913
prompt_token_ids: list[int]
914
"""The processed token IDs which includes placeholder tokens."""
915
916
- token_type_ids: NotRequired[list[int]]
917
- """The token type IDs of the prompt."""
918
-
916
mm_kwargs: MultiModalKwargsOptionalItems
917
"""Keyword arguments to be directly passed to the model after batching."""
918
@@ -946,6 +943,3 @@ class MultiModalEncDecInputs(MultiModalInputs):
943
944
encoder_prompt_token_ids: list[int]
945
"""The processed token IDs of the encoder prompt."""
949
-
950
- encoder_token_type_ids: NotRequired[list[int]]
951
- """The token type IDs of the encoder prompt."""
vllm/sequence.py
-11
@@ -508,12 +508,6 @@ class Sequence:
508
return [0] * len(self.inputs["prompt_embeds"])
509
return self.inputs["prompt_token_ids"]
510
511
- @property
512
- def token_type_ids(self) -> list[int]:
513
- if self.inputs["type"] == "embeds":
514
- return []
515
- return self.inputs.get("token_type_ids", [])
516
-
511
@property
512
def multi_modal_data(self) -> MultiModalKwargs:
513
if self.inputs["type"] == "multimodal":
@@ -765,10 +759,6 @@ class SequenceGroup:
759
return (self.encoder_seq.prompt_token_ids
760
if self.encoder_seq is not None else None)
761
768
- @property
769
- def token_type_ids(self) -> Optional[list[int]]:
770
- return self.first_seq.token_type_ids
771
-
762
@property
763
def multi_modal_data(self) -> MultiModalKwargs:
764
if self.first_seq.multi_modal_data:
@@ -972,7 +962,6 @@ class SequenceGroupMetadata(
962
computed_block_nums: Optional[list[int]] = None
963
state: Optional[SequenceGroupState] = msgspec.field(
964
default_factory=lambda: SequenceGroupState())
975
- token_type_ids: Optional[list[int]] = None
965
multi_modal_data: Optional[MultiModalKwargs] = None
966
multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
967
encoder_seq_data: Optional[SequenceData] = None
vllm/worker/enc_dec_model_runner.py
+2
-3
@@ -24,8 +24,7 @@ from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalKwargs,
24
MultiModalRegistry)
25
from vllm.platforms import _Backend
26
from vllm.sampling_params import SamplingParams
27
-from vllm.sequence import (IntermediateTensors, PoolerOutput,
28
- SequenceGroupMetadata)
27
+from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
28
from vllm.utils import STR_NOT_IMPL_ENC_DEC_BACKEND, make_tensor_with_pad
29
from vllm.worker.model_runner import (GPUModelRunnerBase,
30
ModelInputForGPUBuilder,
@@ -161,7 +160,7 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
160
kv_caches: List[torch.Tensor],
161
intermediate_tensors: Optional[IntermediateTensors] = None,
162
num_steps: int = 1,
164
- ) -> Optional[List[PoolerOutput]]:
163
+ ) -> Optional[List[SamplerOutput]]:
164
if num_steps > 1:
165
raise ValueError("num_steps > 1 is not supported in "
166
"EncoderDecoderModelRunner")
vllm/worker/model_runner.py
-29
@@ -86,7 +86,6 @@ class ModelInputForGPU(ModelRunnerInputBase):
86
input_tokens: Optional[torch.Tensor] = None
87
inputs_embeds: Optional[torch.Tensor] = None
88
input_positions: Optional[torch.Tensor] = None
89
- token_types: Optional[torch.Tensor] = None
89
seq_lens: Optional[List[int]] = None
90
query_lens: Optional[List[int]] = None
91
lora_mapping: Optional["LoRAMapping"] = None
@@ -192,7 +191,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
191
self.input_tokens[0].clear() # type: ignore
192
self.inputs_embeds = None # type: ignore
193
self.input_positions[0].clear() # type: ignore
195
- self.token_types[0].clear() # type: ignore
194
self.mrope_input_positions = None # type: ignore
195
self.seq_lens[0] = 0 # type: ignore
196
self.orig_seq_lens[0] = 0 # type: ignore
@@ -219,7 +217,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
217
input_tokens: Optional[List[List[int]]] = None,
218
inputs_embeds: Optional[torch.Tensor] = None,
219
input_positions: Optional[List[List[int]]] = None,
222
- token_types: Optional[List[List[int]]] = None,
220
mrope_input_positions: Optional[List[List[List[int]]]] = None,
221
222
# The sequence length (may be capped to the sliding window).
@@ -284,12 +281,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
281
for seq_id in range(len(self.seq_ids)):
282
self.input_positions[seq_id].clear()
283
287
- if token_types:
288
- self.token_types = token_types
289
- else:
290
- for seq_id in range(len(self.seq_ids)):
291
- self.token_types[seq_id].clear()
292
-
284
self.mrope_input_positions = None
285
286
if seq_lens:
@@ -348,7 +339,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
339
self.input_tokens = input_tokens or []
340
self.inputs_embeds = inputs_embeds
341
self.input_positions = input_positions or []
351
- self.token_types = token_types or []
342
self.mrope_input_positions = mrope_input_positions or None
343
self.seq_lens = seq_lens or []
344
self.orig_seq_lens = orig_seq_lens or []
@@ -376,7 +366,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
366
367
self.input_tokens = [[] for _ in range(self.n_seqs)]
368
self.input_positions = [[] for _ in range(self.n_seqs)]
379
- self.token_types = [[] for _ in range(self.n_seqs)]
369
self.mrope_input_positions = None
370
self.seq_lens = [0] * self.n_seqs
371
self.orig_seq_lens = [0] * self.n_seqs
@@ -400,7 +389,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
389
f"inputs_embeds.shape="
390
f"{getattr(self.inputs_embeds, 'shape', None)}, "
391
f"input_positions={self.input_positions}, "
403
- f"token_types={self.token_types}, "
392
f"mrope_input_positions={self.mrope_input_positions}, "
393
f"seq_lens={self.seq_lens}, "
394
f"orig_seq_lens={self.orig_seq_lens}, "
@@ -522,8 +510,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
510
prompt_embeds = seq_data.get_token_embeddings(
511
)[context_len:seq_len]
512
525
- token_types = seq_group_metadata.token_type_ids
526
-
513
inter_data.seq_lens[seq_idx] = seq_len
514
inter_data.orig_seq_lens[seq_idx] = seq_len
515
inter_data.prompt_lens[seq_idx] = seq_data.get_prompt_len()
@@ -531,8 +517,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
517
inter_data.input_tokens[seq_idx].extend(tokens)
518
inter_data.inputs_embeds = prompt_embeds
519
inter_data.input_positions[seq_idx].extend(range(context_len, seq_len))
534
- inter_data.token_types[seq_idx].extend(
535
- token_types if token_types else [])
520
inter_data.query_lens[seq_idx] = seq_len - context_len
521
522
if seq_data.mrope_position_delta is not None:
@@ -590,8 +574,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
574
seq_idx][uncomputed_start:]
575
inter_data.input_positions[seq_idx] = inter_data.input_positions[
576
seq_idx][uncomputed_start:]
593
- inter_data.token_types[seq_idx] = inter_data.token_types[seq_idx][
594
- uncomputed_start:]
577
context_len = prefix_cache_len
578
579
inter_data.context_lens[seq_idx] = context_len
@@ -606,8 +588,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
588
seq_idx][-1:]
589
inter_data.input_positions[seq_idx] = inter_data.input_positions[
590
seq_idx][-1:]
609
- inter_data.token_types[seq_idx] = inter_data.token_types[seq_idx][
610
- -1:]
591
inter_data.query_lens[seq_idx] = 1
592
inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1
593
@@ -802,12 +782,9 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
782
# Combine and flatten intermediate data.
783
input_tokens = list[int]()
784
inputs_embeds_list = list[torch.Tensor]()
805
- token_types = list[int]()
785
for inter_data in self.inter_data_list:
786
for cur_input_tokens in inter_data.input_tokens:
787
input_tokens.extend(cur_input_tokens)
809
- for cur_token_types in inter_data.token_types:
810
- token_types.extend(cur_token_types)
788
if inter_data.inputs_embeds is not None:
789
inputs_embeds_list.append(
790
inter_data.inputs_embeds.to(
@@ -890,11 +867,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
867
self.runner.device,
868
self.runner.pin_memory)
869
893
- token_types_tensor = async_tensor_h2d(token_types, torch.long,
894
- self.runner.device,
895
- self.runner.pin_memory) \
896
- if token_types else None
897
-
870
if mrope_input_positions is not None:
871
for idx in range(3):
872
mrope_input_positions[idx].extend(
@@ -951,7 +923,6 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
923
input_tokens=input_tokens_tensor,
924
inputs_embeds=inputs_embeds,
925
input_positions=input_positions_tensor,
954
- token_types=token_types_tensor,
926
attn_metadata=attn_metadata,
927
seq_lens=seq_lens,
928
query_lens=query_lens,
vllm/worker/model_runner_base.py
+2
-12
@@ -13,10 +13,9 @@ from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.sampler import SamplerOutput
15
from vllm.model_executor.models.interfaces import supports_transcription
16
-from vllm.model_executor.models.interfaces_base import (
17
- is_pooling_model, is_text_generation_model)
16
+from vllm.model_executor.models.interfaces_base import is_text_generation_model
17
from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
19
-from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
18
+from vllm.tasks import GenerationTask, SupportedTask
19
20
if TYPE_CHECKING:
21
from vllm.attention import AttentionMetadata
@@ -241,20 +240,11 @@ class ModelRunnerBase(ABC, Generic[T]):
240
241
return supported_tasks
242
244
- def get_supported_pooling_tasks(self) -> list[PoolingTask]:
245
- model = self.get_model()
246
- if not is_pooling_model(model):
247
- return []
248
-
249
- return list(model.pooler.get_supported_tasks())
250
-
243
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
244
tasks = list[SupportedTask]()
245
246
if self.model_config.runner_type == "generate":
247
tasks.extend(self.get_supported_generation_tasks())
256
- if self.model_config.runner_type == "pooling":
257
- tasks.extend(self.get_supported_pooling_tasks())
248
249
return tuple(tasks)
250
vllm/worker/pooling_model_runner.py
deleted
-222
@@ -1,222 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
6
-
7
-import torch
8
-
9
-from vllm.config import VllmConfig
10
-from vllm.distributed import get_pp_group
11
-from vllm.forward_context import set_forward_context
12
-from vllm.logger import init_logger
13
-from vllm.model_executor.models.interfaces_base import VllmModelForPooling
14
-from vllm.model_executor.pooling_metadata import PoolingMetadata
15
-from vllm.multimodal import MultiModalKwargs
16
-from vllm.pooling_params import PoolingParams
17
-from vllm.sequence import (IntermediateTensors, PoolerOutput, SequenceData,
18
- SequenceGroupMetadata)
19
-from vllm.worker.model_runner import (GPUModelRunnerBase, ModelInputForGPU,
20
- ModelInputForGPUBuilder)
21
-
22
-logger = init_logger(__name__)
23
-
24
-
25
-@dataclasses.dataclass(frozen=True)
26
-class ModelInputForGPUWithPoolingMetadata(ModelInputForGPU):
27
- """
28
- Used by the PoolingModelRunner.
29
- """
30
- pooling_metadata: Optional["PoolingMetadata"] = None
31
-
32
-
33
-class PoolingModelRunner(
34
- GPUModelRunnerBase[ModelInputForGPUWithPoolingMetadata]):
35
- _model_input_cls: Type[ModelInputForGPUWithPoolingMetadata] = (
36
- ModelInputForGPUWithPoolingMetadata)
37
- _builder_cls: Type[ModelInputForGPUBuilder] = ModelInputForGPUBuilder
38
-
39
- def __init__(
40
- self,
41
- vllm_config: VllmConfig,
42
- kv_cache_dtype: Optional[str] = "auto",
43
- is_driver_worker: bool = False,
44
- ):
45
- super().__init__(vllm_config=vllm_config,
46
- kv_cache_dtype=kv_cache_dtype,
47
- is_driver_worker=is_driver_worker)
48
-
49
- @torch.inference_mode()
50
- def execute_model(
51
- self,
52
- model_input: ModelInputForGPUWithPoolingMetadata,
53
- kv_caches: List[torch.Tensor],
54
- intermediate_tensors: Optional[IntermediateTensors] = None,
55
- num_steps: int = 1,
56
- ) -> Optional[Union[List[PoolerOutput], IntermediateTensors]]:
57
- if num_steps > 1:
58
- raise ValueError(
59
- "PoolingModelRunner does not support multi-step execution.")
60
-
61
- if self.lora_config:
62
- assert model_input.lora_requests is not None
63
- assert model_input.lora_mapping is not None
64
- self.set_active_loras(model_input.lora_requests,
65
- model_input.lora_mapping)
66
-
67
- # Currently cuda graph is only supported by the decode phase.
68
- assert model_input.attn_metadata is not None
69
- prefill_meta = model_input.attn_metadata.prefill_metadata
70
- decode_meta = model_input.attn_metadata.decode_metadata
71
- virtual_engine = model_input.virtual_engine
72
- # Pooling models are (ab-)used also to integrate non text models that
73
- # are not autoregressive (PrithviGeosaptialMAE).
74
- # These model might not use attention and do not really have a prefill
75
- # and decode phase. The model input is processed in one shot and both
76
- # decode_metadata and prefill_metadata would be None for such models.
77
- # See the PlaceholderAttentionMetadata class.
78
- # TODO: Figure out if cuda_graph is of any use for these models and
79
- # explore how to leverage it.
80
- if (prefill_meta is None and decode_meta is not None
81
- and decode_meta.use_cuda_graph):
82
- if model_input.inputs_embeds is None:
83
- assert model_input.input_tokens is not None
84
- graph_batch_size = model_input.input_tokens.shape[0]
85
- model_executable = (
86
- self.graph_runners[model_input.virtual_engine][(
87
- graph_batch_size, False)])
88
- else:
89
- graph_batch_size = model_input.inputs_embeds.shape[0]
90
- model_executable = (
91
- self.graph_runners[model_input.virtual_engine][(
92
- graph_batch_size, True)])
93
- else:
94
- model_executable = self.model
95
-
96
- multi_modal_kwargs = model_input.multi_modal_kwargs or {}
97
- seqlen_agnostic_kwargs = {
98
- "finished_requests_ids": model_input.finished_requests_ids,
99
- "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
100
- } if self.has_inner_state else {}
101
- if (self.observability_config is not None
102
- and self.observability_config.collect_model_forward_time):
103
- model_forward_start = torch.cuda.Event(enable_timing=True)
104
- model_forward_end = torch.cuda.Event(enable_timing=True)
105
- model_forward_start.record()
106
-
107
- cross_enc_kwargs = {}
108
- if model_input.token_types is not None:
109
- cross_enc_kwargs["token_type_ids"] = model_input.token_types
110
-
111
- with set_forward_context(model_input.attn_metadata, self.vllm_config,
112
- virtual_engine):
113
- hidden_or_intermediate_states = model_executable(
114
- input_ids=model_input.input_tokens,
115
- positions=model_input.input_positions,
116
- intermediate_tensors=intermediate_tensors,
117
- **MultiModalKwargs.as_kwargs(
118
- multi_modal_kwargs,
119
- device=self.device,
120
- ),
121
- **cross_enc_kwargs,
122
- **seqlen_agnostic_kwargs,
123
- )
124
-
125
- if (self.observability_config is not None
126
- and self.observability_config.collect_model_forward_time):
127
- model_forward_end.record()
128
-
129
- # Only perform pooling in the last pipeline stage.
130
- if not get_pp_group().is_last_rank:
131
- if (self.is_driver_worker
132
- and hidden_or_intermediate_states is not None
133
- and isinstance(hidden_or_intermediate_states,
134
- IntermediateTensors)
135
- and self.observability_config is not None
136
- and self.observability_config.collect_model_forward_time):
137
- model_forward_end.synchronize()
138
- model_forward_time = model_forward_start.elapsed_time(
139
- model_forward_end)
140
- orig_model_forward_time = 0.0
141
- if intermediate_tensors is not None:
142
- orig_model_forward_time = intermediate_tensors.tensors.get(
143
- "model_forward_time", torch.tensor(0.0)).item()
144
- hidden_or_intermediate_states.tensors["model_forward_time"] = (
145
- torch.tensor(model_forward_time + orig_model_forward_time))
146
- return hidden_or_intermediate_states
147
-
148
- # Only perform pooling in the driver worker.
149
- if not self.is_driver_worker:
150
- return []
151
-
152
- pooling_metadata = model_input.pooling_metadata
153
- assert pooling_metadata is not None
154
-
155
- pooling_metadata.build_pooling_cursor(
156
- num_scheduled_tokens=pooling_metadata.prompt_lens,
157
- device=hidden_or_intermediate_states.device)
158
-
159
- return [
160
- self.model.pooler(hidden_states=hidden_or_intermediate_states,
161
- pooling_metadata=pooling_metadata)
162
- ]
163
-
164
- def make_model_input_from_broadcasted_tensor_dict(
165
- self,
166
- tensor_dict: Dict[str,
167
- Any]) -> ModelInputForGPUWithPoolingMetadata:
168
- return ModelInputForGPUWithPoolingMetadata.from_broadcasted_tensor_dict(
169
- tensor_dict,
170
- attn_backend=self.attn_backend,
171
- )
172
-
173
- def prepare_model_input(
174
- self,
175
- seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
176
- virtual_engine: int = 0,
177
- finished_requests_ids: Optional[List[str]] = None
178
- ) -> ModelInputForGPUWithPoolingMetadata:
179
- assert seq_group_metadata_list is not None
180
- model_input = self._prepare_model_input_tensors(
181
- seq_group_metadata_list, finished_requests_ids)
182
- # Prepare PoolingMetadata.
183
- assert model_input.seq_lens is not None
184
- pooling_metadata = self._prepare_pooling(seq_group_metadata_list,
185
- model_input.seq_lens)
186
-
187
- return dataclasses.replace(model_input,
188
- pooling_metadata=pooling_metadata)
189
-
190
- def _prepare_pooling(
191
- self,
192
- seq_group_metadata_list: List[SequenceGroupMetadata],
193
- prompt_lens: List[int],
194
- ) -> PoolingMetadata:
195
- """Prepare PoolingMetadata for the sequence group metadata list."""
196
- seq_groups: List[Tuple[List[int], PoolingParams]] = []
197
- for i, seq_group_metadata in enumerate(seq_group_metadata_list):
198
- seq_ids = list(seq_group_metadata.seq_data.keys())
199
-
200
- pooling_params = seq_group_metadata.pooling_params
201
- assert pooling_params is not None
202
-
203
- task = pooling_params.task
204
- assert task is not None, "You did not set `task` in the API"
205
-
206
- model = cast(VllmModelForPooling, self.model)
207
- to_update = model.pooler.get_pooling_updates(task)
208
- to_update.apply(pooling_params)
209
-
210
- seq_groups.append((seq_ids, pooling_params))
211
-
212
- seq_data: Dict[int, SequenceData] = {}
213
- for seq_group_metadata in seq_group_metadata_list:
214
- seq_data.update(seq_group_metadata.seq_data)
215
-
216
- pooling_metadata = PoolingMetadata(
217
- seq_groups=seq_groups,
218
- seq_data=seq_data,
219
- prompt_lens=prompt_lens,
220
- )
221
-
222
- return pooling_metadata
vllm/worker/worker.py
+1
-5
@@ -30,7 +30,6 @@ from vllm.utils import (GiB_bytes, MemorySnapshot, bind_kv_cache,
30
from vllm.worker.cache_engine import CacheEngine
31
from vllm.worker.enc_dec_model_runner import EncoderDecoderModelRunner
32
from vllm.worker.model_runner import GPUModelRunnerBase, ModelRunner
33
-from vllm.worker.pooling_model_runner import PoolingModelRunner
33
from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,
34
WorkerInput)
35
@@ -83,9 +82,7 @@ class Worker(LocalOrDistributedWorkerBase):
82
else {"return_hidden_states": True}
83
84
ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner
86
- if model_config.runner_type == "pooling":
87
- ModelRunnerClass = PoolingModelRunner
88
- elif self.model_config.is_encoder_decoder:
85
+ if self.model_config.is_encoder_decoder:
86
ModelRunnerClass = EncoderDecoderModelRunner
87
self.model_runner: GPUModelRunnerBase = ModelRunnerClass(
88
vllm_config=self.vllm_config,
@@ -99,7 +96,6 @@ class Worker(LocalOrDistributedWorkerBase):
96
# Uninitialized cache engine. Will be initialized by
97
# initialize_cache.
98
self.cache_engine: List[CacheEngine]
102
- # Initialize gpu_cache as pooling models don't initialize kv_caches
99
self.gpu_cache: Optional[List[List[torch.Tensor]]] = None
100
self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}
101