@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 799397ee4f
Fork 0

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 15, 04:45 UTC

Support embedding models in V1 (#16188)

Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Signed-off-by: Max de Bayser <maxdebayser@gmail.com> Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com>

Maximilien de Bayser committed Jun 19, 2025 at 01:36 UTC 799397ee4f57b90ee1b5f12f88b12f4de0de0d1d
56 files changed +895 -287
examples/offline_inference/basic/embed.py
+4 -1
@@ -12,7 +12,10 @@ def parse_args():
12 parser = EngineArgs.add_cli_args(parser)
13 # Set example specific arguments
14 parser.set_defaults(
15 - model="intfloat/e5-mistral-7b-instruct", task="embed", enforce_eager=True
15 + model="intfloat/e5-mistral-7b-instruct",
16 + task="embed",
17 + enforce_eager=True,
18 + max_model_len=1024,
19 )
20 return parser.parse_args()
21
examples/offline_inference/vision_language_embedding.py
+1
@@ -94,6 +94,7 @@ def run_vlm2vec(query: Query) -> ModelRequestData:
94 engine_args = EngineArgs(
95 model="TIGER-Lab/VLM2Vec-Full",
96 task="embed",
97 + max_model_len=4096,
98 trust_remote_code=True,
99 mm_processor_kwargs={"num_crops": 4},
100 limit_mm_per_prompt={"image": 1},
tests/compile/test_basic_correctness.py
+18 -14
@@ -31,7 +31,7 @@ class TestSetting:
31 # basic llama model
32 TestSetting(
33 model="meta-llama/Llama-3.2-1B-Instruct",
34 - model_args=[],
34 + model_args=["--max-model-len", "2048"],
35 pp_size=2,
36 tp_size=2,
37 attn_backend="FLASHINFER",
@@ -41,7 +41,7 @@ class TestSetting:
41 # llama model with quantization
42 TestSetting(
43 model="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ",
44 - model_args=["--quantization", "gptq"],
44 + model_args=["--quantization", "gptq", "--max-model-len", "2048"],
45 pp_size=1,
46 tp_size=1,
47 attn_backend="FLASH_ATTN",
@@ -51,7 +51,7 @@ class TestSetting:
51 # MoE model
52 TestSetting(
53 model="ibm/PowerMoE-3b",
54 - model_args=[],
54 + model_args=["--max-model-len", "2048"],
55 pp_size=1,
56 tp_size=2,
57 attn_backend="FLASH_ATTN",
@@ -61,23 +61,27 @@ class TestSetting:
61 # embedding model
62 TestSetting(
63 model="BAAI/bge-multilingual-gemma2",
64 - model_args=["--task", "embed", "--dtype", "bfloat16"],
64 + model_args=[
65 + "--task", "embed", "--dtype", "bfloat16", "--max-model-len",
66 + "2048"
67 + ],
68 pp_size=1,
69 tp_size=1,
70 attn_backend="FLASH_ATTN",
71 method="encode",
72 fullgraph=True,
73 ),
71 - # encoder-based embedding model (BERT)
72 - TestSetting(
73 - model="BAAI/bge-base-en-v1.5",
74 - model_args=["--task", "embed"],
75 - pp_size=1,
76 - tp_size=1,
77 - attn_backend="XFORMERS",
78 - method="encode",
79 - fullgraph=True,
80 - ),
74 + # TODO: bert models are not supported in V1 yet
75 + # # encoder-based embedding model (BERT)
76 + # TestSetting(
77 + # model="BAAI/bge-base-en-v1.5",
78 + # model_args=["--task", "embed"],
79 + # pp_size=1,
80 + # tp_size=1,
81 + # attn_backend="XFORMERS",
82 + # method="encode",
83 + # fullgraph=True,
84 + # ),
85 # vision language model
86 TestSetting(
87 model="microsoft/Phi-3.5-vision-instruct",
tests/conftest.py
+3
@@ -145,6 +145,7 @@ def run_with_both_engines(request, monkeypatch):
145 # Automatically runs tests twice, once with V1 and once without
146 use_v1 = request.param
147 # Tests decorated with `@skip_v1` are only run without v1
148 + skip_v0 = request.node.get_closest_marker("skip_v0")
149 skip_v1 = request.node.get_closest_marker("skip_v1")
150
151 if use_v1:
@@ -152,6 +153,8 @@ def run_with_both_engines(request, monkeypatch):
153 pytest.skip("Skipping test on vllm V1")
154 monkeypatch.setenv('VLLM_USE_V1', '1')
155 else:
156 + if skip_v0:
157 + pytest.skip("Skipping test on vllm V0")
158 monkeypatch.setenv('VLLM_USE_V1', '0')
159
160 yield
tests/entrypoints/llm/test_encode.py
+20 -4
@@ -8,6 +8,8 @@ import pytest
8 from vllm import LLM, PoolingParams, PoolingRequestOutput
9 from vllm.distributed import cleanup_dist_env_and_memory
10
11 +from ...models.utils import check_embeddings_close
12 +
13 MODEL_NAME = "intfloat/multilingual-e5-small"
14
15 PROMPTS = [
@@ -27,6 +29,14 @@ TOKEN_IDS = [
29 ]
30
31
32 +@pytest.fixture(autouse=True)
33 +def v1(run_with_both_engines):
34 + # Simple autouse wrapper to run both engines for each test
35 + # This can be promoted up to conftest.py to run for every
36 + # test in a package
37 + pass
38 +
39 +
40 @pytest.fixture(scope="module")
41 def llm():
42 # pytest caches the fixture so we use weakref.proxy to
@@ -46,9 +56,15 @@ def llm():
56 cleanup_dist_env_and_memory()
57
58
49 -def assert_outputs_equal(o1: list[PoolingRequestOutput],
59 +def assert_outputs_match(o1: list[PoolingRequestOutput],
60 o2: list[PoolingRequestOutput]):
51 - assert [o.outputs for o in o1] == [o.outputs for o in o2]
61 + check_embeddings_close(
62 + embeddings_0_lst=[o.outputs.data for o in o1],
63 + embeddings_1_lst=[o.outputs.data for o in o2],
64 + name_0="hf",
65 + name_1="vllm",
66 + tol=1e-2,
67 + )
68
69
70 @pytest.mark.skip_global_cleanup
@@ -63,7 +79,7 @@ def test_v1_v2_api_consistency_single_prompt_tokens(llm: LLM,
79
80 v2_output = llm.encode({"prompt_token_ids": prompt_token_ids},
81 pooling_params=pooling_params)
66 - assert_outputs_equal(v1_output, v2_output)
82 + assert_outputs_match(v1_output, v2_output)
83
84
85 @pytest.mark.skip_global_cleanup
@@ -80,7 +96,7 @@ def test_v1_v2_api_consistency_multi_prompt_tokens(llm: LLM):
96 } for p in TOKEN_IDS],
97 pooling_params=pooling_params,
98 )
83 - assert_outputs_equal(v1_output, v2_output)
99 + assert_outputs_match(v1_output, v2_output)
100
101
102 @pytest.mark.skip_global_cleanup
tests/entrypoints/openai/test_embedding.py
+8
@@ -21,6 +21,14 @@ DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' +
21 DTYPE = "bfloat16"
22
23
24 +@pytest.fixture(autouse=True)
25 +def v1(run_with_both_engines):
26 + # Simple autouse wrapper to run both engines for each test
27 + # This can be promoted up to conftest.py to run for every
28 + # test in a package
29 + pass
30 +
31 +
32 @pytest.fixture(scope="module")
33 def server():
34 args = [
tests/entrypoints/openai/test_pooling.py
+11 -4
@@ -7,6 +7,7 @@ import numpy as np
7 import pytest
8 import requests
9
10 +from tests.models.utils import check_embeddings_close
11 from vllm.entrypoints.openai.protocol import PoolingResponse
12 from vllm.transformers_utils.tokenizer import get_tokenizer
13
@@ -223,8 +224,11 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer,
224 np.frombuffer(base64.b64decode(data.data),
225 dtype="float32").tolist())
226
226 - assert responses_float.data[0].data == decoded_responses_base64_data[0]
227 - assert responses_float.data[1].data == decoded_responses_base64_data[1]
227 + check_embeddings_close(
228 + embeddings_0_lst=[d.data for d in responses_float.data],
229 + embeddings_1_lst=decoded_responses_base64_data,
230 + name_0="float32",
231 + name_1="base64")
232
233 # Default response is float32 decoded from base64 by OpenAI Client
234 default_response = requests.post(
@@ -237,5 +241,8 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer,
241 default_response.raise_for_status()
242 responses_default = PoolingResponse.model_validate(default_response.json())
243
240 - assert responses_float.data[0].data == responses_default.data[0].data
241 - assert responses_float.data[1].data == responses_default.data[1].data
244 + check_embeddings_close(
245 + embeddings_0_lst=[d.data for d in responses_default.data],
246 + embeddings_1_lst=[d.data for d in responses_default.data],
247 + name_0="float32",
248 + name_1="base64")
tests/entrypoints/openai/test_rerank.py
+8
@@ -12,6 +12,14 @@ MODEL_NAME = "BAAI/bge-reranker-base"
12 DTYPE = "bfloat16"
13
14
15 +@pytest.fixture(autouse=True)
16 +def v1(run_with_both_engines):
17 + # Simple autouse wrapper to run both engines for each test
18 + # This can be promoted up to conftest.py to run for every
19 + # test in a package
20 + pass
21 +
22 +
23 @pytest.fixture(scope="module")
24 def server():
25 args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
tests/entrypoints/openai/test_score.py
+9
@@ -11,6 +11,15 @@ from vllm.entrypoints.openai.protocol import ScoreResponse
11
12 from ...utils import RemoteOpenAIServer
13
14 +
15 +@pytest.fixture(autouse=True)
16 +def v1(run_with_both_engines):
17 + # Simple autouse wrapper to run both engines for each test
18 + # This can be promoted up to conftest.py to run for every
19 + # test in a package
20 + pass
21 +
22 +
23 MODELS = [
24 {
25 "name": "BAAI/bge-reranker-v2-m3",
tests/models/language/pooling/test_classification.py
+9 -1
@@ -6,6 +6,14 @@ from transformers import AutoModelForSequenceClassification
6
7 from vllm.platforms import current_platform
8
9 +# TODO: enable when float32 is supported by V1
10 +# @pytest.fixture(autouse=True)
11 +# def v1(run_with_both_engines):
12 +# # Simple autouse wrapper to run both engines for each test
13 +# # This can be promoted up to conftest.py to run for every
14 +# # test in a package
15 +# pass
16 +
17
18 @pytest.mark.parametrize(
19 "model",
@@ -29,7 +37,7 @@ def test_models(
37 # switch to use ROCm CK FA backend
38 monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
39
32 - with vllm_runner(model, dtype=dtype) as vllm_model:
40 + with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:
41 vllm_outputs = vllm_model.classify(example_prompts)
42
43 with hf_runner(model,
tests/models/language/pooling/test_embedding.py
+27 -7
@@ -8,6 +8,14 @@ from vllm.platforms import current_platform
8 from ...utils import check_embeddings_close
9
10
11 +@pytest.fixture(autouse=True)
12 +def v1(run_with_both_engines):
13 + # Simple autouse wrapper to run both engines for each test
14 + # This can be promoted up to conftest.py to run for every
15 + # test in a package
16 + pass
17 +
18 +
19 @pytest.mark.parametrize(
20 "model",
21 [
@@ -20,15 +28,27 @@ from ...utils import check_embeddings_close
28 marks=[pytest.mark.core_model]),
29 pytest.param("intfloat/e5-mistral-7b-instruct",
30 marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
23 - pytest.param("ssmits/Qwen2-7B-Instruct-embed-base"),
31 + # the qwen models interfere with each other (see PR
32 + # https://github.com/vllm-project/vllm/pull/18720).
33 + # To avoid this problem, for now we skip v0 since it will be
34 + # deprecated anyway.
35 + pytest.param("ssmits/Qwen2-7B-Instruct-embed-base",
36 + marks=[pytest.mark.skip_v0]),
37 # [Encoder-only]
38 pytest.param("BAAI/bge-base-en-v1.5",
26 - marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
27 - pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
28 - pytest.param("intfloat/multilingual-e5-small"),
29 - pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct"),
39 + marks=[
40 + pytest.mark.core_model, pytest.mark.cpu_model,
41 + pytest.mark.skip_v1
42 + ]),
43 + pytest.param("sentence-transformers/all-MiniLM-L12-v2",
44 + marks=[pytest.mark.skip_v1]),
45 + pytest.param("intfloat/multilingual-e5-small",
46 + marks=[pytest.mark.skip_v1]),
47 + pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
48 + marks=[pytest.mark.skip_v1]),
49 # [Cross-Encoder]
31 - pytest.param("sentence-transformers/stsb-roberta-base-v2"),
50 + pytest.param("sentence-transformers/stsb-roberta-base-v2",
51 + marks=[pytest.mark.skip_v1]),
52 ],
53 )
54 def test_models(
@@ -62,7 +82,7 @@ def test_models(
82
83 with vllm_runner(model,
84 task="embed",
65 - max_model_len=None,
85 + max_model_len=512,
86 **vllm_extra_kwargs) as vllm_model:
87 vllm_outputs = vllm_model.encode(example_prompts)
88
tests/models/registry.py
+11 -11
@@ -265,8 +265,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
265
266 _EMBEDDING_EXAMPLE_MODELS = {
267 # [Text-only]
268 - "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
269 - "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
268 + "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5", v0_only=True),
269 + "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2", v0_only=True), # noqa: E501
270 "GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
271 "GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
272 trust_remote_code=True),
@@ -279,16 +279,16 @@ _EMBEDDING_EXAMPLE_MODELS = {
279 "LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
280 "MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
281 "ModernBertModel": _HfExamplesInfo("Alibaba-NLP/gte-modernbert-base",
282 - trust_remote_code=True),
282 + trust_remote_code=True, v0_only=True),
283 "NomicBertModel": _HfExamplesInfo("nomic-ai/nomic-embed-text-v2-moe",
284 - trust_remote_code=True),
284 + trust_remote_code=True, v0_only=True), # noqa: E501
285 "Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
286 "Qwen2ForRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-RM-72B"),
287 "Qwen2ForProcessRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-PRM-7B"),
288 "Qwen2ForSequenceClassification": _HfExamplesInfo("jason9693/Qwen2.5-1.5B-apeach"), # noqa: E501
289 - "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"), # noqa: E501
290 - "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"), # noqa: E501
291 - "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"),
289 + "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2", v0_only=True), # noqa: E501
290 + "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1", v0_only=True), # noqa: E501
291 + "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small", v0_only=True), # noqa: E501
292 # [Multimodal]
293 "LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
294 "Phi3VForCausalLM": _HfExamplesInfo("TIGER-Lab/VLM2Vec-Full",
@@ -300,10 +300,10 @@ _EMBEDDING_EXAMPLE_MODELS = {
300
301 _CROSS_ENCODER_EXAMPLE_MODELS = {
302 # [Text-only]
303 - "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2"), # noqa: E501
304 - "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base"), # noqa: E501
305 - "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3"), # noqa: E501
306 - "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base"), # noqa: E501
303 + "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2", v0_only=True), # noqa: E501
304 + "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base", v0_only=True), # noqa: E501
305 + "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3", v0_only=True), # noqa: E501
306 + "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base", v0_only=True), # noqa: E501
307 }
308
309 _MULTIMODAL_EXAMPLE_MODELS = {
tests/tokenization/test_detokenize.py
+1
@@ -68,6 +68,7 @@ def _run_incremental_decode(tokenizer,
68 None,
69 params,
70 None,
71 + None,
72 0.0,
73 None,
74 cache_salt=None,
tests/v1/core/test_kv_cache_utils.py
+1
@@ -43,6 +43,7 @@ def make_request(request_id,
43 multi_modal_hashes=mm_hashes,
44 multi_modal_placeholders=mm_positions,
45 sampling_params=SamplingParams(max_tokens=17),
46 + pooling_params=None,
47 eos_token_id=100,
48 lora_request=None,
49 cache_salt=cache_salt,
tests/v1/core/test_prefix_caching.py
+1
@@ -39,6 +39,7 @@ def make_request(request_id,
39 multi_modal_placeholders=mm_positions,
40 sampling_params=SamplingParams(max_tokens=17,
41 prompt_logprobs=prompt_logprobs),
42 + pooling_params=None,
43 eos_token_id=100,
44 lora_request=None,
45 cache_salt=cache_salt,
tests/v1/core/test_scheduler.py
+22 -4
@@ -135,6 +135,7 @@ def create_requests(num_requests: int,
135 request_id=f"{i}",
136 prompt_token_ids=[i] * num_tokens,
137 sampling_params=sampling_params,
138 + pooling_params=None,
139 multi_modal_inputs=mm_inputs,
140 multi_modal_placeholders=mm_position,
141 multi_modal_hashes=None,
@@ -283,6 +284,7 @@ def test_schedule_partial_requests():
284 spec_token_ids=None,
285 logprobs=None,
286 prompt_logprobs_dict={},
287 + pooler_output=[],
288 )
289 scheduler.update_from_output(output, model_runner_output)
290
@@ -333,6 +335,7 @@ def test_no_mm_input_chunking():
335 spec_token_ids=None,
336 logprobs=None,
337 prompt_logprobs_dict={},
338 + pooler_output=[],
339 )
340 scheduler.update_from_output(output, model_runner_output)
341
@@ -396,6 +399,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
399 spec_token_ids=None,
400 logprobs=None,
401 prompt_logprobs_dict={},
402 + pooler_output=[],
403 )
404 scheduler.update_from_output(output, model_runner_output)
405
@@ -420,6 +424,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
424 spec_token_ids=None,
425 logprobs=None,
426 prompt_logprobs_dict={},
427 + pooler_output=[],
428 )
429 scheduler.update_from_output(output1, model_runner_output)
430 output2 = scheduler.schedule()
@@ -473,7 +478,8 @@ def test_stop_via_update_from_output():
478 11]], # First request hits EOS, second continues
479 spec_token_ids=None,
480 logprobs=None,
476 - prompt_logprobs_dict={})
481 + prompt_logprobs_dict={},
482 + pooler_output=[])
483
484 scheduler.update_from_output(scheduler_output, model_output)
485
@@ -523,7 +529,8 @@ def test_stop_via_update_from_output():
529 [13, 14]], # First request hits stop token
530 spec_token_ids=None,
531 logprobs=None,
526 - prompt_logprobs_dict={})
532 + prompt_logprobs_dict={},
533 + pooler_output=[])
534
535 scheduler.update_from_output(scheduler_output, model_output)
536
@@ -572,7 +579,8 @@ def test_stop_via_update_from_output():
579 [13]], # First request exceeds max_tokens
580 spec_token_ids=None,
581 logprobs=None,
575 - prompt_logprobs_dict={})
582 + prompt_logprobs_dict={},
583 + pooler_output=[])
584
585 scheduler.update_from_output(scheduler_output, model_output)
586
@@ -614,7 +622,8 @@ def test_stop_via_update_from_output():
622 sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
623 spec_token_ids=None,
624 logprobs=None,
617 - prompt_logprobs_dict={})
625 + prompt_logprobs_dict={},
626 + pooler_output=[])
627
628 scheduler.update_from_output(scheduler_output, model_output)
629
@@ -663,6 +672,7 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],
672 spec_token_ids=None,
673 logprobs=None,
674 prompt_logprobs_dict={},
675 + pooler_output=[],
676 )
677 scheduler.update_from_output(scheduler_output0, model_runner_output)
678
@@ -680,6 +690,7 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],
690 spec_token_ids=None,
691 logprobs=None,
692 prompt_logprobs_dict={},
693 + pooler_output=[],
694 )
695 scheduler.update_from_output(scheduler_output1, model_runner_output)
696
@@ -730,6 +741,7 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
741 spec_token_ids=spec_tokens,
742 logprobs=None,
743 prompt_logprobs_dict={},
744 + pooler_output=[],
745 )
746 engine_core_outputs = scheduler.update_from_output(output,
747 model_runner_output)
@@ -769,6 +781,7 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
781 spec_token_ids=None,
782 logprobs=None,
783 prompt_logprobs_dict={},
784 + pooler_output=[],
785 )
786 engine_core_outputs = scheduler.update_from_output(output,
787 model_runner_output)
@@ -896,6 +909,7 @@ def test_kv_connector_basic():
909 spec_token_ids=None,
910 logprobs=None,
911 prompt_logprobs_dict={},
912 + pooler_output=[],
913 )
914
915 # Ensure ScheduleOutput is correct.
@@ -941,6 +955,7 @@ def test_kv_connector_basic():
955 spec_token_ids=None,
956 logprobs=None,
957 prompt_logprobs_dict={},
958 + pooler_output=[],
959 )
960
961 # We should get a local cache hit of NUM_TOKENS_PREFIX and
@@ -1007,6 +1022,7 @@ def test_kv_connector_unable_to_allocate():
1022 spec_token_ids=None,
1023 logprobs=None,
1024 prompt_logprobs_dict={},
1025 + pooler_output=[],
1026 )
1027
1028 # Just one request should be running.
@@ -1087,6 +1103,7 @@ def test_kv_connector_handles_preemption():
1103 spec_token_ids=None,
1104 logprobs=None,
1105 prompt_logprobs_dict={},
1106 + pooler_output=[],
1107 )
1108
1109 # All can be scheduled - 1st token.
@@ -1181,6 +1198,7 @@ def make_output(scheduler: Scheduler):
1198 spec_token_ids=None,
1199 logprobs=None,
1200 prompt_logprobs_dict={},
1201 + pooler_output=[],
1202 )
1203
1204
tests/v1/engine/test_engine_core.py
+1
@@ -39,6 +39,7 @@ def make_request() -> EngineCoreRequest:
39 mm_hashes=None,
40 mm_placeholders=None,
41 sampling_params=SamplingParams(),
42 + pooling_params=None,
43 eos_token_id=None,
44 arrival_time=time.time(),
45 lora_request=None,
tests/v1/engine/test_engine_core_client.py
+1
@@ -53,6 +53,7 @@ def make_request(
53 mm_hashes=None,
54 mm_placeholders=None,
55 sampling_params=params,
56 + pooling_params=None,
57 eos_token_id=None,
58 arrival_time=time.time(),
59 lora_request=None,
tests/v1/engine/test_fast_incdec_prefix_err.py
+1
@@ -33,6 +33,7 @@ def test_fast_inc_detok_invalid_utf8_err_case():
33 None,
34 params,
35 None,
36 + None,
37 0.0,
38 None,
39 cache_salt=None,
tests/v1/engine/test_output_processor.py
+9 -4
@@ -66,7 +66,8 @@ def test_incremental_detokenization(request_output_kind: RequestOutputKind,
66 output_kind=request_output_kind,
67 stop=[],
68 include_stop_str_in_output=False,
69 - ))
69 + ),
70 + pooling_params=None)
71 for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
72 ]
73
@@ -416,7 +417,8 @@ def test_logprobs_processor(request_output_kind: RequestOutputKind,
417 include_stop_str_in_output=False,
418 logprobs=num_sample_logprobs,
419 prompt_logprobs=num_prompt_logprobs,
419 - ))
420 + ),
421 + pooling_params=None)
422 for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
423 ]
424
@@ -582,7 +584,8 @@ def test_stop_token(include_stop_str_in_output: bool,
584 logprobs=num_sample_logprobs,
585 prompt_logprobs=None,
586 ignore_eos=ignore_eos,
585 - ))
587 + ),
588 + pooling_params=None)
589
590 # Add request to the detokenizer.
591 output_processor.add_request(request, prompt_string)
@@ -678,7 +681,8 @@ def test_stop_string(include_stop_str_in_output: bool,
681 include_stop_str_in_output=include_stop_str_in_output,
682 logprobs=num_sample_logprobs,
683 prompt_logprobs=None,
681 - ))
684 + ),
685 + pooling_params=None)
686 for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
687 ]
688
@@ -786,6 +790,7 @@ def test_iteration_stats(dummy_test_vectors):
790 cache_salt=None,
791 data_parallel_rank=None,
792 sampling_params=SamplingParams(),
793 + pooling_params=None,
794 ) for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
795 ]
796
tests/v1/kv_connector/unit/utils.py
+2
@@ -150,6 +150,7 @@ def create_request(
150 request_id=f"id-{request_id}",
151 prompt_token_ids=prompt_token_ids,
152 sampling_params=sampling_params,
153 + pooling_params=None,
154 multi_modal_inputs=None,
155 multi_modal_placeholders=None,
156 multi_modal_hashes=None,
@@ -183,6 +184,7 @@ def create_model_runner_output(
184 spec_token_ids=None,
185 logprobs=None,
186 prompt_logprobs_dict={},
187 + pooler_output=None,
188 finished_sending=finished_sending,
189 finished_recving=finished_recving,
190 )
tests/v1/worker/test_gpu_input_batch.py
+3 -1
@@ -10,6 +10,7 @@ import torch
10
11 from vllm.sampling_params import SamplingParams
12 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
13 +from vllm.v1.pool.metadata import PoolingMetadata
14 from vllm.v1.sample.metadata import SamplingMetadata
15 from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable
16 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -46,7 +47,7 @@ def _compare_objs(obj1, obj2):
47 for a_i, b_i in zip(a.block_tables, b.block_tables):
48 _compare_objs(a_i, b_i)
49 is_same = True
49 - elif isinstance(a, (BlockTable, SamplingMetadata)):
50 + elif isinstance(a, (BlockTable, SamplingMetadata, PoolingMetadata)):
51 _compare_objs(a, b)
52 is_same = True # if we make it here must be same
53 elif a == b:
@@ -201,6 +202,7 @@ def _construct_cached_request_state(req_id_suffix: int):
202 req_id=f"req_id_{req_id_suffix}",
203 prompt_token_ids=prompt_token_ids,
204 sampling_params=_create_sampling_params(),
205 + pooling_params=None,
206 mm_inputs=[],
207 mm_positions=[],
208 block_ids=([], ),
tests/v1/worker/test_gpu_model_runner.py
+1
@@ -122,6 +122,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
122 mm_hashes=[],
123 mm_positions=[],
124 sampling_params=SamplingParams(),
125 + pooling_params=None,
126 block_ids=([0], ),
127 num_computed_tokens=0,
128 lora_request=None,
vllm/config.py
+23 -3
@@ -4496,11 +4496,31 @@ class VllmConfig:
4496
4497 if self.compilation_config.full_cuda_graph and \
4498 not self.model_config.disable_cascade_attn:
4499 - logger.warning_once(
4500 - "full_cuda_graph is not supported with "
4501 - "cascade attention. Disabling cascade attention.")
4499 + logger.info("full_cuda_graph is not supported with "
4500 + "cascade attention. Disabling cascade attention.")
4501 self.model_config.disable_cascade_attn = True
4502
4503 + disable_chunked_prefill_reasons: list[str] = []
4504 +
4505 + if self.model_config and self.model_config.pooler_config:
4506 + pooling_type = self.model_config.pooler_config.pooling_type
4507 + if pooling_type is None or pooling_type.lower() != "last":
4508 + disable_chunked_prefill_reasons.append(
4509 + "Only \"last\" pooling supports chunked "
4510 + "prefill and prefix caching; disabling both.")
4511 +
4512 + if disable_chunked_prefill_reasons:
4513 + for reason in disable_chunked_prefill_reasons:
4514 + logger.info(reason)
4515 + self.scheduler_config.chunked_prefill_enabled = False
4516 + self.scheduler_config.long_prefill_token_threshold = 0
4517 + self.scheduler_config.max_num_batched_tokens = max(
4518 + self.scheduler_config.max_model_len,
4519 + DEFAULT_MAX_NUM_BATCHED_TOKENS)
4520 +
4521 + if self.cache_config is not None:
4522 + self.cache_config.enable_prefix_caching = False
4523 +
4524 if (self.kv_events_config is not None
4525 and self.kv_events_config.enable_kv_cache_events
4526 and not self.cache_config.enable_prefix_caching):
vllm/engine/arg_utils.py
+31 -14
@@ -1041,7 +1041,7 @@ class EngineArgs:
1041
1042 # Set default arguments for V0 or V1 Engine.
1043 if use_v1:
1044 - self._set_default_args_v1(usage_context)
1044 + self._set_default_args_v1(usage_context, model_config)
1045 else:
1046 self._set_default_args_v0(model_config)
1047
@@ -1349,13 +1349,7 @@ class EngineArgs:
1349 recommend_to_remove=False)
1350 return False
1351
1352 - # No Embedding Models so far.
1353 - if model_config.task not in ["generate"]:
1354 - _raise_or_fallback(feature_name=f"--task {model_config.task}",
1355 - recommend_to_remove=False)
1356 - return False
1357 -
1358 - # No Encoder-Decoder, not all Mamba so far.
1352 + # No Mamba or Encoder-Decoder so far.
1353 if not model_config.is_v1_compatible:
1354 _raise_or_fallback(feature_name=model_config.architectures,
1355 recommend_to_remove=False)
@@ -1523,15 +1517,38 @@ class EngineArgs:
1517 if self.max_num_seqs is None:
1518 self.max_num_seqs = 256
1519
1526 - def _set_default_args_v1(self, usage_context: UsageContext) -> None:
1520 + def _set_default_args_v1(self, usage_context: UsageContext,
1521 + model_config: ModelConfig) -> None:
1522 """Set Default Arguments for V1 Engine."""
1523
1529 - # V1 always uses chunked prefills.
1530 - self.enable_chunked_prefill = True
1524 + # V1 always uses chunked prefills and prefix caching
1525 + # for non-pooling tasks.
1526 + # For pooling tasks the default is False
1527 + if model_config.runner_type != "pooling":
1528 + self.enable_chunked_prefill = True
1529 + if self.enable_prefix_caching is None:
1530 + self.enable_prefix_caching = True
1531 + else:
1532 +
1533 + pooling_type = model_config.pooler_config.pooling_type
1534 +
1535 + # TODO: when encoder models are supported we'll have to
1536 + # check for causal attention here.
1537 + incremental_prefill_supported = (pooling_type is not None and
1538 + pooling_type.lower() == "last")
1539
1532 - # V1 enables prefix caching by default.
1533 - if self.enable_prefix_caching is None:
1534 - self.enable_prefix_caching = True
1540 + action = "Enabling" if \
1541 + incremental_prefill_supported else "Disabling"
1542 +
1543 + if self.enable_chunked_prefill is None:
1544 + self.enable_chunked_prefill = incremental_prefill_supported
1545 + logger.info("(%s) chunked prefill by default", action)
1546 + if self.enable_prefix_caching is None:
1547 + self.enable_prefix_caching = incremental_prefill_supported
1548 + logger.info("(%s) prefix caching by default", action)
1549 +
1550 + if not self.enable_chunked_prefill:
1551 + self.max_num_batched_tokens = model_config.max_model_len
1552
1553 # V1 should use the new scheduler by default.
1554 # Swap it only if this arg is set to the original V0 default
vllm/entrypoints/llm.py
+1 -1
@@ -1266,7 +1266,7 @@ class LLM:
1266 # the tokenizer for models such as
1267 # "cross-encoder/ms-marco-MiniLM-L-6-v2" doesn't support passing
1268 # lists of tokens to the `text` and `text_pair` kwargs
1269 - tokenizer = self.llm_engine.get_tokenizer()
1269 + tokenizer = self.get_tokenizer()
1270
1271 def ensure_str(prompt: SingletonPrompt):
1272 if isinstance(prompt, dict):
vllm/entrypoints/openai/serving_pooling.py
+3 -1
@@ -9,6 +9,7 @@ from typing import Final, Literal, Optional, Union, cast
9
10 import jinja2
11 import numpy as np
12 +import torch
13 from fastapi import Request
14 from typing_extensions import assert_never
15
@@ -39,7 +40,8 @@ def _get_data(
40 elif encoding_format == "base64":
41 # Force to use float32 for base64 encoding
42 # to match the OpenAI python client behavior
42 - pooling_bytes = np.array(output.data, dtype="float32").tobytes()
43 + pt_float32 = output.data.to(dtype=torch.float32)
44 + pooling_bytes = np.array(pt_float32, dtype="float32").tobytes()
45 return base64.b64encode(pooling_bytes).decode("utf-8")
46
47 assert_never(encoding_format)
vllm/model_executor/layers/pooler.py
+116 -33
@@ -10,11 +10,15 @@ import torch.nn.functional as F
10 from typing_extensions import assert_never
11
12 from vllm.config import ModelConfig, PoolerConfig
13 -from vllm.model_executor.pooling_metadata import (PoolingMetadata,
14 - PoolingTensors)
13 +from vllm.model_executor.pooling_metadata import ( # noqa: E501
14 + PoolingMetadata as V0PoolingMetadata)
15 +from vllm.model_executor.pooling_metadata import PoolingTensors
16 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
17 from vllm.transformers_utils.config import (
18 get_cross_encoder_activation_function)
19 +from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
20 +
21 +PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
22
23
24 class PoolingType(IntEnum):
@@ -75,15 +79,18 @@ class SimplePooler(nn.Module):
79
80 def get_prompt_lens(
81 self,
78 - hidden_states: torch.Tensor,
82 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
83 pooling_metadata: PoolingMetadata,
84 ) -> torch.Tensor:
85 + if isinstance(pooling_metadata, V1PoolingMetadata):
86 + return pooling_metadata.prompt_lens
87 + assert isinstance(hidden_states, torch.Tensor)
88 return PoolingTensors.from_pooling_metadata(
89 pooling_metadata, hidden_states.device).prompt_lens
90
91 def extract_states(
92 self,
86 - hidden_states: torch.Tensor,
93 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
94 pooling_metadata: PoolingMetadata,
95 ) -> Union[list[torch.Tensor], torch.Tensor]:
96 raise NotImplementedError
@@ -93,7 +100,7 @@ class SimplePooler(nn.Module):
100
101 def forward(
102 self,
96 - hidden_states: torch.Tensor,
103 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
104 pooling_metadata: PoolingMetadata,
105 ) -> PoolerOutput:
106 pooled_data = self.extract_states(hidden_states, pooling_metadata)
@@ -106,11 +113,19 @@ class CLSPool(SimplePooler):
113
114 def extract_states(
115 self,
109 - hidden_states: torch.Tensor,
116 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
117 pooling_metadata: PoolingMetadata,
118 ) -> Union[list[torch.Tensor], torch.Tensor]:
119 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
120
121 + if isinstance(hidden_states, list):
122 + result = []
123 + for req_state, prompt_len in zip(hidden_states, prompt_lens):
124 + assert prompt_len == req_state.shape[0], \
125 + "partial prefill not supported with CLS pooling"
126 + result.append(req_state[0])
127 + return result
128 +
129 first_token_flat_indices = torch.zeros_like(prompt_lens)
130 first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]
131 return hidden_states[first_token_flat_indices]
@@ -120,9 +135,12 @@ class LastPool(SimplePooler):
135
136 def extract_states(
137 self,
123 - hidden_states: torch.Tensor,
138 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
139 pooling_metadata: PoolingMetadata,
140 ) -> Union[list[torch.Tensor], torch.Tensor]:
141 + if isinstance(hidden_states, list):
142 + return [h[-1] for h in hidden_states]
143 +
144 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
145
146 last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1
@@ -133,11 +151,17 @@ class AllPool(SimplePooler):
151
152 def extract_states(
153 self,
136 - hidden_states: torch.Tensor,
154 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
155 pooling_metadata: PoolingMetadata,
156 ) -> Union[list[torch.Tensor], torch.Tensor]:
157 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
158
159 + if isinstance(hidden_states, list):
160 + for req_state, prompt_len in zip(hidden_states, prompt_lens):
161 + assert prompt_len == req_state.shape[0], \
162 + "partial prefill not supported with ALL pooling"
163 + return hidden_states
164 +
165 offset = 0
166 pooled_data = list[torch.Tensor]()
167 for prompt_len in prompt_lens:
@@ -151,11 +175,20 @@ class MeanPool(SimplePooler):
175
176 def extract_states(
177 self,
154 - hidden_states: torch.Tensor,
178 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
179 pooling_metadata: PoolingMetadata,
180 ) -> Union[list[torch.Tensor], torch.Tensor]:
181 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
182
183 + if isinstance(hidden_states, list):
184 + result = []
185 + for req_state, prompt_len in zip(hidden_states, prompt_lens):
186 + assert prompt_len == req_state.shape[0], \
187 + "partial prefill not supported with mean pooling"
188 + result.append(torch.mean(req_state, dim=0,
189 + dtype=torch.float32))
190 + return result
191 +
192 # Use float32 for torch.cumsum in MeanPool,
193 # otherwise precision will be lost significantly.
194 cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
@@ -184,30 +217,53 @@ class StepPool(SimplePooler):
217 self.step_tag_id = step_tag_id
218 self.returned_token_ids = returned_token_ids
219
220 + def get_prompt_token_ids(
221 + self,
222 + pooling_metadata: PoolingMetadata,
223 + ) -> list[torch.Tensor]:
224 + if isinstance(pooling_metadata, V1PoolingMetadata):
225 + return [
226 + pooling_metadata.prompt_token_ids[i, :num]
227 + for i, num in enumerate(pooling_metadata.prompt_lens)
228 + ]
229 + return [
230 + torch.tensor(seq_data_i.prompt_token_ids)
231 + for seq_data_i in pooling_metadata.seq_data.values()
232 + ]
233 +
234 def extract_states(
235 self,
189 - hidden_states: torch.Tensor,
236 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
237 pooling_metadata: PoolingMetadata,
238 ) -> Union[list[torch.Tensor], torch.Tensor]:
239 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
240 + prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)
241
194 - returned_token_ids = self.returned_token_ids
195 - if returned_token_ids is not None and len(returned_token_ids) > 0:
196 - hidden_states = hidden_states[:, returned_token_ids]
242 + pooled_data: list[torch.Tensor] = []
243
244 + if isinstance(hidden_states, list):
245 + for req_state, prompt_len in zip(hidden_states, prompt_lens):
246 + assert prompt_len == req_state.shape[0], \
247 + "partial prefill not supported with mean pooling"
248 + pooled_data = hidden_states
249 + else:
250 + offset = 0
251 + for prompt_len in prompt_lens:
252 + pooled_data_i = hidden_states[offset:offset + prompt_len]
253 + offset += prompt_len
254 + pooled_data.append(pooled_data_i)
255 +
256 + pooled_data = []
257 + returned_token_ids = self.returned_token_ids
258 step_tag_id = self.step_tag_id
259
200 - offset = 0
201 - pooled_data = list[torch.Tensor]()
202 - for prompt_len, seq_data_i in zip(prompt_lens,
203 - pooling_metadata.seq_data.values()):
204 - pooled_data_i = hidden_states[offset:offset + prompt_len]
205 - if step_tag_id is not None:
206 - token_ids = torch.tensor(seq_data_i.prompt_token_ids)
207 - pooled_data_i = pooled_data_i[token_ids == step_tag_id]
260 + for data, token_id in zip(pooled_data, prompt_token_ids):
261 + if returned_token_ids is not None and len(returned_token_ids) > 0:
262 + data = data[:, returned_token_ids]
263
209 - offset += prompt_len
210 - pooled_data.append(pooled_data_i)
264 + if step_tag_id is not None:
265 + data = data[token_id == step_tag_id]
266 + pooled_data.append(data)
267
268 return pooled_data
269
@@ -230,10 +286,17 @@ class PoolerHead(nn.Module):
286 else:
287 pooled_data = pooled_data.to(torch.float32)
288
233 - dimensions_list = [
234 - pooling_param.dimensions
235 - for _, pooling_param in pooling_metadata.seq_groups
236 - ]
289 + if isinstance(pooling_metadata, V0PoolingMetadata):
290 + dimensions_list = [
291 + pooling_param.dimensions
292 + for _, pooling_param in pooling_metadata.seq_groups
293 + ]
294 + else:
295 + assert isinstance(pooled_data, list)
296 + dimensions_list = [
297 + pooling_param.dimensions
298 + for pooling_param in pooling_metadata.pooling_params
299 + ]
300 if any(d is not None for d in dimensions_list):
301 # change the output dimension
302 assert len(pooled_data) == len(dimensions_list)
@@ -325,20 +388,41 @@ class ClassifierPooler(nn.Module):
388 raise NotImplementedError(f"task={config.task!r} is not supported"
389 " with the classification pooler")
390
391 + def get_prompt_lens(
392 + self,
393 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
394 + pooling_metadata: PoolingMetadata,
395 + ) -> torch.Tensor:
396 + if isinstance(pooling_metadata, V1PoolingMetadata):
397 + return pooling_metadata.prompt_lens
398 + assert isinstance(hidden_states, torch.Tensor)
399 + return PoolingTensors.from_pooling_metadata(
400 + pooling_metadata, hidden_states.device).prompt_lens
401 +
402 def forward(
403 self,
330 - hidden_states: torch.Tensor,
404 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
405 pooling_metadata: PoolingMetadata,
406 ) -> PoolerOutput:
407 """Pools sentence pair scores from the hidden_states."""
408 + prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
409
335 - prompt_lens = PoolingTensors.from_pooling_metadata(
336 - pooling_metadata, hidden_states.device).prompt_lens
410 + pooled_data = list[torch.Tensor]()
411 + if isinstance(hidden_states, list):
412 + for req_state, prompt_len in zip(hidden_states, prompt_lens):
413 + assert prompt_len == req_state.shape[0], \
414 + "partial prefill not supported with classifier"
415 + pooled_data = hidden_states
416 + else:
417 + offset = 0
418 + for prompt_len in prompt_lens:
419 + pooled_data_i = hidden_states[offset:offset + prompt_len]
420 + offset += prompt_len
421 + pooled_data.append(pooled_data_i)
422
423 offset = 0
424 pooled_data_lst = []
340 - for prompt_len in prompt_lens:
341 - pooled_data_i = hidden_states[offset:offset + prompt_len]
425 + for pooled_data_i in pooled_data:
426
427 if self.pooler is not None:
428 final_shape_tensor = self.pooler(pooled_data_i)
@@ -346,7 +430,6 @@ class ClassifierPooler(nn.Module):
430 final_shape_tensor = self.classifier(pooled_data_i)
431
432 pooled_data_lst.append(final_shape_tensor)
349 - offset += prompt_len
433
434 pooled_output = torch.stack(pooled_data_lst)
435
vllm/model_executor/models/bert.py
+2 -2
@@ -446,8 +446,8 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
446 softmax=False)
447
448
449 -class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
450 - SupportsQuant):
449 +class BertForSequenceClassification(nn.Module, SupportsV0Only,
450 + SupportsCrossEncoding, SupportsQuant):
451 """A model that uses Bert to provide embedding functionalities.
452
453 This class encapsulates the BertModel and provides an interface for
vllm/model_executor/models/modernbert.py
+3 -2
@@ -21,7 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21 from vllm.model_executor.pooling_metadata import PoolingMetadata
22 from vllm.sequence import IntermediateTensors, PoolerOutput
23
24 -from .interfaces import SupportsCrossEncoding
24 +from .interfaces import SupportsCrossEncoding, SupportsV0Only
25 from .utils import WeightsMapper, maybe_prefix
26
27
@@ -270,7 +270,8 @@ class ModernBertPooler(nn.Module):
270 return pooled_output
271
272
273 -class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
273 +class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
274 + SupportsCrossEncoding):
275
276 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
277 super().__init__()
vllm/model_executor/models/qwen3.py
+6 -1
@@ -375,7 +375,12 @@ class Qwen3ForSequenceClassification(nn.Module, SupportsLoRA,
375 ) -> Optional[PoolerOutput]:
376 hidden_states = self._pooler.extract_states(hidden_states,
377 pooling_metadata)
378 - logits, _ = self.score(hidden_states)
378 +
379 + if isinstance(hidden_states, list):
380 + logits = [self.score(state)[0] for state in hidden_states]
381 + else:
382 + logits, _ = self.score(hidden_states)
383 +
384 pooled_data = self._pooler.head(logits, pooling_metadata)
385 pooled_outputs = [
386 self._pooler.build_output(data.squeeze(-1)) for data in pooled_data
vllm/pooling_params.py
+7
@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Any, Optional
5
6 import msgspec
7
8 +from vllm.sampling_params import RequestOutputKind
9 +
10 if TYPE_CHECKING:
11 from vllm.config import ModelConfig
12
@@ -23,6 +25,7 @@ class PoolingParams(
25
26 dimensions: Optional[int] = None
27 additional_data: Optional[Any] = None
28 + output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
29
30 def clone(self) -> "PoolingParams":
31 """Returns a deep copy of the PoolingParams instance."""
@@ -52,3 +55,7 @@ class PoolingParams(
55 return (f"PoolingParams("
56 f"dimensions={self.dimensions}, "
57 f"additional_metadata={self.additional_data})")
58 +
59 + def __post_init__(self) -> None:
60 + assert self.output_kind == RequestOutputKind.FINAL_ONLY,\
61 + "For pooling output_kind has to be FINAL_ONLY"
vllm/v1/core/kv_cache_manager.py
+2 -1
@@ -146,7 +146,8 @@ class KVCacheManager:
146 # Prefix caching is disabled or
147 # When the request requires prompt logprobs, we skip prefix caching.
148 if (not self.enable_caching
149 - or request.sampling_params.prompt_logprobs is not None):
149 + or (request.sampling_params is not None
150 + and request.sampling_params.prompt_logprobs is not None)):
151 return self.create_empty_block_list(), 0
152
153 # The block hashes for the request may already be computed
vllm/v1/core/sched/output.py
+4 -1
@@ -14,6 +14,7 @@ if TYPE_CHECKING:
14 KVConnectorMetadata)
15 from vllm.lora.request import LoRARequest
16 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
17 + from vllm.pooling_params import PoolingParams
18 from vllm.sampling_params import SamplingParams
19 from vllm.v1.request import Request
20
@@ -26,7 +27,8 @@ class NewRequestData:
27 mm_inputs: list[MultiModalKwargs]
28 mm_hashes: list[str]
29 mm_positions: list[PlaceholderRange]
29 - sampling_params: SamplingParams
30 + sampling_params: Optional[SamplingParams]
31 + pooling_params: Optional[PoolingParams]
32 block_ids: tuple[list[int], ...]
33 num_computed_tokens: int
34 lora_request: Optional[LoRARequest]
@@ -44,6 +46,7 @@ class NewRequestData:
46 mm_hashes=request.mm_hashes,
47 mm_positions=request.mm_positions,
48 sampling_params=request.sampling_params,
49 + pooling_params=request.pooling_params,
50 block_ids=block_ids,
51 num_computed_tokens=request.num_computed_tokens,
52 lora_request=request.lora_request,
vllm/v1/core/sched/scheduler.py
+25 -3
@@ -402,6 +402,15 @@ class Scheduler(SchedulerInterface):
402 < num_new_tokens):
403 num_new_tokens = (
404 self.scheduler_config.long_prefill_token_threshold)
405 +
406 + # chunked prefill has to be enabled explicitly to allow
407 + # pooling requests to be chunked
408 + if not self.scheduler_config.chunked_prefill_enabled and \
409 + num_new_tokens > token_budget:
410 + self.waiting.popleft()
411 + skipped_waiting_requests.appendleft(request)
412 + continue
413 +
414 num_new_tokens = min(num_new_tokens, token_budget)
415 assert num_new_tokens > 0
416
@@ -707,6 +716,7 @@ class Scheduler(SchedulerInterface):
716 logprobs = model_runner_output.logprobs
717 prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
718 num_scheduled_tokens = scheduler_output.num_scheduled_tokens
719 + pooler_outputs = model_runner_output.pooler_output
720
721 new_running: list[Request] = []
722 outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
@@ -724,7 +734,8 @@ class Scheduler(SchedulerInterface):
734 continue
735
736 req_index = model_runner_output.req_id_to_index[req_id]
727 - generated_token_ids = sampled_token_ids[req_index]
737 + generated_token_ids = sampled_token_ids[
738 + req_index] if sampled_token_ids else []
739
740 scheduled_spec_token_ids = (
741 scheduler_output.scheduled_spec_decode_tokens.get(req_id))
@@ -776,8 +787,17 @@ class Scheduler(SchedulerInterface):
787 del new_token_ids[num_new:] # Trim new tokens if needed.
788 break
789
790 + pooler_output = None
791 + if pooler_outputs:
792 + pooler_output = pooler_outputs[req_index]
793 + stopped = check_stop(request, self.max_model_len,
794 + pooler_output)
795 + if stopped:
796 + kv_transfer_params = self._free_request(request)
797 +
798 # Extract sample logprobs if needed.
780 - if request.sampling_params.logprobs is not None and logprobs:
799 + if request.sampling_params is not None \
800 + and request.sampling_params.logprobs is not None and logprobs:
801 # NOTE: once we support N tokens per step (spec decode),
802 # the outer lists can be of length > 1.
803 new_logprobs = logprobs.slice(req_index, req_index + 1)
@@ -802,7 +822,8 @@ class Scheduler(SchedulerInterface):
822
823 # Get prompt logprobs for this request.
824 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
805 - if new_token_ids or kv_transfer_params:
825 + if new_token_ids or pooler_output is not None \
826 + or kv_transfer_params:
827
828 # Add EngineCoreOutput for this Request.
829 outputs[request.client_index].append(
@@ -812,6 +833,7 @@ class Scheduler(SchedulerInterface):
833 finish_reason=request.get_finished_reason(),
834 new_logprobs=new_logprobs,
835 new_prompt_logprobs_tensors=prompt_logprobs_tensors,
836 + pooling_output=pooler_output,
837 stop_reason=request.stop_reason,
838 events=request.take_events(),
839 kv_transfer_params=kv_transfer_params,
vllm/v1/core/sched/utils.py
+14 -1
@@ -1,15 +1,28 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Optional
4 +
5 +import torch
6 +
7 from vllm.v1.request import Request, RequestStatus
8
9
6 -def check_stop(request: Request, max_model_len: int) -> bool:
10 +def check_stop(request: Request,
11 + max_model_len: int,
12 + pooler_output: Optional[torch.Tensor] = None) -> bool:
13 if (request.num_tokens >= max_model_len
14 or request.num_output_tokens >= request.max_tokens):
15 request.status = RequestStatus.FINISHED_LENGTH_CAPPED
16 return True
17
18 + if request.pooling_params:
19 + if pooler_output is not None:
20 + request.status = RequestStatus.FINISHED_STOPPED
21 + return True
22 + return False
23 +
24 sampling_params = request.sampling_params
25 + assert sampling_params is not None
26 last_token_id = request.output_token_ids[-1]
27 if (not sampling_params.ignore_eos
28 and last_token_id == request.eos_token_id):
vllm/v1/engine/__init__.py
+6 -1
@@ -7,10 +7,12 @@ from collections.abc import Sequence
7 from typing import Any, Optional, Union
8
9 import msgspec
10 +import torch
11
12 from vllm.lora.request import LoRARequest
13 from vllm.multimodal import MultiModalKwargs
14 from vllm.multimodal.inputs import PlaceholderRange
15 +from vllm.pooling_params import PoolingParams
16 from vllm.sampling_params import SamplingParams
17 from vllm.v1.metrics.stats import SchedulerStats
18 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
@@ -50,7 +52,8 @@ class EngineCoreRequest(
52 mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]]
53 mm_hashes: Optional[list[str]]
54 mm_placeholders: Optional[list[PlaceholderRange]]
53 - sampling_params: SamplingParams
55 + sampling_params: Optional[SamplingParams]
56 + pooling_params: Optional[PoolingParams]
57 eos_token_id: Optional[int]
58 arrival_time: float
59 lora_request: Optional[LoRARequest]
@@ -104,6 +107,8 @@ class EngineCoreOutput(
107 new_logprobs: Optional[LogprobsLists] = None
108 new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None
109
110 + pooling_output: Optional[torch.Tensor] = None
111 +
112 finish_reason: Optional[FinishReason] = None
113 stop_reason: Union[int, str, None] = None
114 events: Optional[list[EngineCoreEvent]] = None
vllm/v1/engine/async_llm.py
+73 -7
@@ -17,7 +17,7 @@ from vllm.inputs.preprocess import InputPreprocessor
17 from vllm.logger import init_logger
18 from vllm.lora.request import LoRARequest
19 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
20 -from vllm.outputs import RequestOutput
20 +from vllm.outputs import PoolingRequestOutput, RequestOutput
21 from vllm.pooling_params import PoolingParams
22 from vllm.prompt_adapter.request import PromptAdapterRequest
23 from vllm.sampling_params import SamplingParams
@@ -228,8 +228,7 @@ class AsyncLLM(EngineClient):
228 if self.errored:
229 raise EngineDeadError()
230
231 - assert isinstance(params, SamplingParams), \
232 - "Pooling is not supported in V1"
231 + is_pooling = isinstance(params, PoolingParams)
232
233 # Create a new output collector for the request.
234 queue = RequestOutputCollector(output_kind=params.output_kind)
@@ -240,7 +239,7 @@ class AsyncLLM(EngineClient):
239 tokenization_kwargs, trace_headers, prompt_adapter_request,
240 priority, data_parallel_rank)
241
243 - if params.n == 1:
242 + if is_pooling or params.n == 1:
243 await self._add_request(request, prompt_str, None, 0, queue)
244 return queue
245
@@ -443,7 +442,7 @@ class AsyncLLM(EngineClient):
442 stat_logger.record(scheduler_stats=scheduler_stats,
443 iteration_stats=iteration_stats)
444
446 - def encode(
445 + async def encode(
446 self,
447 prompt: PromptType,
448 pooling_params: PoolingParams,
@@ -451,8 +450,75 @@ class AsyncLLM(EngineClient):
450 lora_request: Optional[LoRARequest] = None,
451 trace_headers: Optional[Mapping[str, str]] = None,
452 priority: int = 0,
454 - ):
455 - raise ValueError("Not Supported on V1 yet.")
453 + ) -> AsyncGenerator[PoolingRequestOutput, None]:
454 + """
455 + Main function called by the API server to kick off a request
456 + * 1) Making an AsyncStream corresponding to the Request.
457 + * 2) Processing the Input.
458 + * 3) Adding the Request to the EngineCore (separate process).
459 +
460 + A separate output_handler loop runs in a background AsyncIO task,
461 + pulling outputs from EngineCore and putting them into the
462 + per-request AsyncStream.
463 +
464 + The caller of generate() iterates the returned AsyncGenerator,
465 + returning the RequestOutput back to the caller.
466 + """
467 +
468 + try:
469 + # We start the output_handler on the first call to generate() so
470 + # we can call __init__ before the event loop, which enables us
471 + # to handle startup failure gracefully in the OpenAI server.
472 + self._run_output_handler()
473 +
474 + q = await self.add_request(
475 + request_id,
476 + prompt,
477 + pooling_params,
478 + lora_request=lora_request,
479 + trace_headers=trace_headers,
480 + priority=priority,
481 + )
482 +
483 + # The output_handler task pushes items into the queue.
484 + # This task pulls from the queue and yields to caller.
485 + finished = False
486 + while not finished:
487 + # Note: drain queue without await if possible (avoids
488 + # task switching under load which helps performance).
489 + out = q.get_nowait() or await q.get()
490 + assert isinstance(out, PoolingRequestOutput)
491 + # Note: both OutputProcessor and EngineCore handle their
492 + # own request cleanup based on finished.
493 + finished = out.finished
494 + yield out
495 +
496 + # If the request is disconnected by the client, generate()
497 + # is cancelled. So, we abort the request if we end up here.
498 + except asyncio.CancelledError:
499 + await self.abort(request_id)
500 + if self.log_requests:
501 + logger.info("Request %s aborted.", request_id)
502 + raise
503 +
504 + # Engine is dead. Do not abort since we shut down.
505 + except EngineDeadError:
506 + if self.log_requests:
507 + logger.info("Request %s failed (engine dead).", request_id)
508 + raise
509 +
510 + # Request validation error.
511 + except ValueError:
512 + if self.log_requests:
513 + logger.info("Request %s failed (bad request).", request_id)
514 + raise
515 +
516 + # Unexpected error in the generate() task (possibly recoverable).
517 + except Exception as e:
518 + await self.abort(request_id)
519 + if self.log_requests:
520 + logger.info("Request %s failed.", request_id)
521 + raise EngineGenerateError() from e
522
523 async def get_vllm_config(self) -> VllmConfig:
524 return self.vllm_config
vllm/v1/engine/core.py
-1
@@ -60,7 +60,6 @@ class EngineCore:
60 executor_class: type[Executor],
61 log_stats: bool,
62 executor_fail_callback: Optional[Callable] = None):
63 - assert vllm_config.model_config.runner_type != "pooling"
63
64 # plugins need to be loaded at the engine/scheduler level too
65 from vllm.plugins import load_general_plugins
vllm/v1/engine/detokenizer.py
+7 -3
@@ -50,6 +50,8 @@ class IncrementalDetokenizer:
50 request: EngineCoreRequest,
51 ) -> "IncrementalDetokenizer":
52
53 + assert request.sampling_params is not None
54 +
55 if tokenizer is None:
56 # No tokenizer => skipping detokenization.
57 return IncrementalDetokenizer()
@@ -70,6 +72,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC):
72
73 # Stop strings
74 params = request.sampling_params
75 + assert params is not None
76 self.stop = stop = params.stop
77 self.include_stop_str_in_output = params.include_stop_str_in_output
78
@@ -164,6 +167,7 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
167 super().__init__(request)
168
169 sampling_params = request.sampling_params
170 + assert sampling_params is not None
171
172 self.request_id = request.request_id
173 self.skip_special_tokens = sampling_params.skip_special_tokens
@@ -245,20 +249,20 @@ class SlowIncrementalDetokenizer(BaseIncrementalDetokenizer):
249 super().__init__(request)
250
251 self.tokenizer = tokenizer
252 + params = request.sampling_params
253 + assert params is not None
254
255 # Metadata for incremental detokenization.
256 self.tokens, self.prefix_offset, self.read_offset = (
257 convert_prompt_ids_to_tokens(
258 tokenizer=tokenizer,
259 prompt_ids=request.prompt_token_ids,
254 - skip_special_tokens=request.sampling_params.
255 - skip_special_tokens,
260 + skip_special_tokens=params.skip_special_tokens,
261 ))
262
263 self.token_ids.extend(request.prompt_token_ids)
264 self.prompt_len = len(request.prompt_token_ids)
265
261 - params = request.sampling_params
266 self.skip_special_tokens = params.skip_special_tokens
267 self.spaces_between_special_tokens = (
268 params.spaces_between_special_tokens)
vllm/v1/engine/llm_engine.py
+2 -2
@@ -15,7 +15,7 @@ from vllm.inputs import PromptType
15 from vllm.logger import init_logger
16 from vllm.lora.request import LoRARequest
17 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
18 -from vllm.outputs import RequestOutput
18 +from vllm.outputs import PoolingRequestOutput, RequestOutput
19 from vllm.pooling_params import PoolingParams
20 from vllm.prompt_adapter.request import PromptAdapterRequest
21 from vllm.sampling_params import SamplingParams
@@ -221,7 +221,7 @@ class LLMEngine:
221 # Add the request to EngineCore.
222 self.engine_core.add_request(child_request)
223
224 - def step(self) -> list[RequestOutput]:
224 + def step(self) -> Union[list[RequestOutput], list[PoolingRequestOutput]]:
225
226 if self.should_execute_dummy_batch:
227 self.should_execute_dummy_batch = False
vllm/v1/engine/logprobs.py
+1
@@ -38,6 +38,7 @@ class LogprobsProcessor:
38 tokenizer: Optional[AnyTokenizer],
39 request: EngineCoreRequest,
40 ) -> "LogprobsProcessor":
41 + assert request.sampling_params is not None
42 num_logprobs = request.sampling_params.logprobs
43 num_prompt_logprobs = request.sampling_params.prompt_logprobs
44 return cls(
vllm/v1/engine/output_processor.py
+93 -44
@@ -4,9 +4,12 @@
4 import asyncio
5 from collections.abc import Iterable
6 from dataclasses import dataclass
7 -from typing import Any, Optional, Union
7 +from typing import Any, Optional, Union, cast
8
9 -from vllm.outputs import CompletionOutput, RequestOutput
9 +import torch
10 +
11 +from vllm.outputs import (CompletionOutput, PoolingOutput,
12 + PoolingRequestOutput, RequestOutput)
13 from vllm.sampling_params import RequestOutputKind
14 from vllm.transformers_utils.tokenizer import AnyTokenizer
15 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
@@ -29,20 +32,22 @@ class RequestOutputCollector:
32
33 def __init__(self, output_kind: RequestOutputKind):
34 self.aggregate = output_kind == RequestOutputKind.DELTA
32 - self.output: Optional[Union[RequestOutput, Exception]] = None
35 + self.output: Optional[Union[RequestOutput, PoolingRequestOutput,
36 + Exception]] = None
37 self.ready = asyncio.Event()
38
35 - def put(self, output: Union[RequestOutput, Exception]) -> None:
39 + def put(self, output: Union[RequestOutput, PoolingRequestOutput,
40 + Exception]) -> None:
41 """Non-blocking put operation."""
42 if self.output is None or isinstance(output, Exception):
43 self.output = output
44 self.ready.set()
40 - elif isinstance(self.output, RequestOutput):
45 + elif isinstance(self.output, (RequestOutput, PoolingRequestOutput)):
46 # This ensures that request outputs with different request indexes
47 # (if n > 1) do not override each other.
48 self.output.add(output, aggregate=self.aggregate)
49
45 - async def get(self) -> RequestOutput:
50 + async def get(self) -> Union[RequestOutput, PoolingRequestOutput]:
51 """Get operation blocks on put event."""
52 while (output := self.output) is None:
53 await self.ready.wait()
@@ -52,7 +57,8 @@ class RequestOutputCollector:
57 raise output
58 return output
59
55 - def get_nowait(self) -> Optional[RequestOutput]:
60 + def get_nowait(
61 + self) -> Optional[Union[RequestOutput, PoolingRequestOutput]]:
62 """Non-blocking get operation."""
63 output = self.output
64 if output is not None:
@@ -66,7 +72,7 @@ class RequestOutputCollector:
72 @dataclass
73 class OutputProcessorOutput:
74
69 - request_outputs: list[RequestOutput]
75 + request_outputs: list[Union[RequestOutput, PoolingRequestOutput]]
76 reqs_to_abort: list[str]
77
78
@@ -81,8 +87,8 @@ class RequestState:
87 output_kind: RequestOutputKind,
88 prompt: Optional[str],
89 prompt_token_ids: list[int],
84 - logprobs_processor: LogprobsProcessor,
85 - detokenizer: IncrementalDetokenizer,
90 + logprobs_processor: Optional[LogprobsProcessor],
91 + detokenizer: Optional[IncrementalDetokenizer],
92 max_tokens_param: Optional[int],
93 arrival_time: float,
94 queue: Optional[RequestOutputCollector],
@@ -116,27 +122,39 @@ class RequestState:
122 queue: Optional[RequestOutputCollector],
123 log_stats: bool,
124 ) -> "RequestState":
119 - if not request.sampling_params.detokenize:
120 - tokenizer = None
125 +
126 + if sampling_params := request.sampling_params:
127 + if not sampling_params.detokenize:
128 + tokenizer = None
129 + output_kind = sampling_params.output_kind
130 + logprobs_processor = LogprobsProcessor.from_new_request(
131 + tokenizer=tokenizer,
132 + request=request,
133 + )
134 + detokenizer = IncrementalDetokenizer.from_new_request(
135 + tokenizer=tokenizer,
136 + request=request,
137 + )
138 + max_tokens_param = sampling_params.max_tokens
139 + else:
140 + logprobs_processor = None
141 + detokenizer = None
142 + max_tokens_param = None
143 + assert request.pooling_params is not None
144 + output_kind = request.pooling_params.output_kind
145 +
146 return cls(
147 request_id=request.request_id,
148 parent_req=parent_req,
149 request_index=request_index,
150 lora_name=(request.lora_request.name
151 if request.lora_request is not None else None),
127 - output_kind=request.sampling_params.output_kind,
152 + output_kind=output_kind,
153 prompt=prompt,
154 prompt_token_ids=request.prompt_token_ids,
130 - logprobs_processor=LogprobsProcessor.from_new_request(
131 - tokenizer=tokenizer,
132 - request=request,
133 - ),
134 - detokenizer=IncrementalDetokenizer.from_new_request(
135 - tokenizer=tokenizer,
136 - request=request,
137 - ),
138 - max_tokens_param=(request.sampling_params.max_tokens if
139 - request.sampling_params is not None else None),
155 + logprobs_processor=logprobs_processor,
156 + detokenizer=detokenizer,
157 + max_tokens_param=max_tokens_param,
158 arrival_time=request.arrival_time,
159 queue=queue,
160 log_stats=log_stats,
@@ -145,11 +163,12 @@ class RequestState:
163 def make_request_output(
164 self,
165 new_token_ids: list[int],
166 + pooling_output: Optional[torch.Tensor],
167 finish_reason: Optional[FinishReason],
168 stop_reason: Union[int, str, None],
169 kv_transfer_params: Optional[dict[str, Any]] = None,
170 num_cached_tokens: int = 0,
152 - ) -> Optional[RequestOutput]:
171 + ) -> Optional[Union[RequestOutput, PoolingRequestOutput]]:
172
173 finished = finish_reason is not None
174 final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
@@ -158,15 +177,20 @@ class RequestState:
177 # Only the final output is required in FINAL_ONLY mode.
178 return None
179
161 - completion_output = self._new_completion_output(
162 - new_token_ids, finish_reason, stop_reason)
163 -
180 request_id = self.request_id
181 + if pooling_output is not None:
182 + return self._new_request_output(
183 + request_id, [self._new_pooling_output(pooling_output)],
184 + finished)
185 +
186 + output = self._new_completion_output(new_token_ids, finish_reason,
187 + stop_reason)
188 +
189 if self.parent_req is None:
166 - outputs = [completion_output]
190 + outputs = [output]
191 else:
192 request_id, outputs, finished = self.parent_req.get_outputs(
169 - request_id, completion_output)
193 + request_id, output)
194 if not outputs:
195 return None
196
@@ -176,12 +200,21 @@ class RequestState:
200 def _new_request_output(
201 self,
202 request_id: str,
179 - outputs: list[CompletionOutput],
203 + outputs: Union[list[CompletionOutput], list[PoolingOutput]],
204 finished: bool,
205 kv_transfer_params: Optional[dict[str, Any]] = None,
206 num_cached_tokens: int = 0,
183 - ) -> RequestOutput:
184 -
207 + ) -> Union[RequestOutput, PoolingRequestOutput]:
208 +
209 + if isinstance(outputs[0], PoolingOutput):
210 + assert len(outputs) == 1
211 + return PoolingRequestOutput(
212 + request_id=request_id,
213 + outputs=outputs[0],
214 + prompt_token_ids=self.prompt_token_ids,
215 + finished=finished,
216 + )
217 + assert self.logprobs_processor is not None
218 if self.output_kind == RequestOutputKind.DELTA:
219 # Side effect: logprobs processor forgets prompt logprobs
220 prompt_logprobs = self.logprobs_processor.pop_prompt_logprobs()
@@ -193,7 +226,7 @@ class RequestState:
226 prompt=self.prompt,
227 prompt_token_ids=self.prompt_token_ids,
228 prompt_logprobs=prompt_logprobs,
196 - outputs=outputs,
229 + outputs=cast(list[CompletionOutput], outputs),
230 finished=finished,
231 kv_transfer_params=kv_transfer_params,
232 num_cached_tokens=num_cached_tokens,
@@ -206,6 +239,8 @@ class RequestState:
239 stop_reason: Union[int, str, None],
240 ) -> CompletionOutput:
241
242 + assert self.detokenizer is not None
243 + assert self.logprobs_processor is not None
244 finished = finish_reason is not None
245 delta = self.output_kind == RequestOutputKind.DELTA
246
@@ -228,6 +263,13 @@ class RequestState:
263 finish_reason=str(finish_reason) if finished else None,
264 stop_reason=stop_reason if finished else None)
265
266 + def _new_pooling_output(
267 + self,
268 + pooling_output: torch.Tensor,
269 + ) -> PoolingOutput:
270 +
271 + return PoolingOutput(data=pooling_output)
272 +
273
274 class OutputProcessor:
275 """Process EngineCoreOutputs into RequestOutputs."""
@@ -326,7 +368,8 @@ class OutputProcessor:
368 within the loop below.
369 """
370
329 - request_outputs: list[RequestOutput] = []
371 + request_outputs: Union[list[RequestOutput],
372 + list[PoolingRequestOutput]] = []
373 reqs_to_abort: list[str] = []
374 for engine_core_output in engine_core_outputs:
375 req_id = engine_core_output.request_id
@@ -341,25 +384,31 @@ class OutputProcessor:
384 iteration_stats)
385
386 new_token_ids = engine_core_output.new_token_ids
387 + pooling_output = engine_core_output.pooling_output
388 finish_reason = engine_core_output.finish_reason
389 stop_reason = engine_core_output.stop_reason
390 kv_transfer_params = engine_core_output.kv_transfer_params
391 num_cached_tokens = engine_core_output.num_cached_tokens
392 req_state.is_prefilling = False
393
350 - # 2) Detokenize the token ids into text and perform stop checks.
351 - stop_string = req_state.detokenizer.update(
352 - new_token_ids, finish_reason == FinishReason.STOP)
353 - if stop_string:
354 - finish_reason = FinishReason.STOP
355 - stop_reason = stop_string
356 -
357 - # 3) Compute sample and prompt logprobs for request, if required.
358 - req_state.logprobs_processor.update_from_output(engine_core_output)
394 + if pooling_output is None:
395 + assert req_state.detokenizer is not None
396 + assert req_state.logprobs_processor is not None
397 + # 2) Detokenize the token ids into text and perform stop checks.
398 + stop_string = req_state.detokenizer.update(
399 + new_token_ids, finish_reason == FinishReason.STOP)
400 + if stop_string:
401 + finish_reason = FinishReason.STOP
402 + stop_reason = stop_string
403 +
404 + # 3) Compute sample and prompt logprobs for request,
405 + # if required.
406 + req_state.logprobs_processor.update_from_output(
407 + engine_core_output)
408
409 # 4) Create and handle RequestOutput objects.
410 if request_output := req_state.make_request_output(
362 - new_token_ids, finish_reason, stop_reason,
411 + new_token_ids, pooling_output, finish_reason, stop_reason,
412 kv_transfer_params, num_cached_tokens):
413 if req_state.queue is not None:
414 # AsyncLLM: put into queue for handling by generate().
vllm/v1/engine/processor.py
+19 -14
@@ -136,8 +136,8 @@ class Processor:
136 Should raise ValueError if unsupported for API Server.
137 """
138
139 - if not isinstance(params, SamplingParams):
140 - raise ValueError("V1 does not yet support Pooling models.")
139 + if isinstance(params, PoolingParams):
140 + return
141
142 self._validate_logprobs(params)
143 self._validate_sampling_params(params, lora_request)
@@ -263,18 +263,22 @@ class Processor:
263 if encoder_inputs is not None:
264 raise NotImplementedError
265
266 - assert isinstance(params, SamplingParams)
267 - # TODO: can we avoid cloning here in multiproc case?
268 - sampling_params = params.clone()
269 - # If unset max tokens, then generate up to the max_model_len.
270 - if sampling_params.max_tokens is None:
271 - sampling_params.max_tokens = (
272 - self.model_config.max_model_len -
273 - len(decoder_inputs["prompt_token_ids"]))
274 - sampling_params.update_from_generation_config(
275 - self.generation_config_fields, eos_token_id)
276 - sampling_params.update_from_tokenizer(
277 - self.tokenizer.get_lora_tokenizer(lora_request))
266 + sampling_params = None
267 + pooling_params = None
268 + if isinstance(params, SamplingParams):
269 + # TODO: can we avoid cloning here in multiproc case?
270 + sampling_params = params.clone()
271 + # If unset max tokens, then generate up to the max_model_len.
272 + if sampling_params.max_tokens is None:
273 + sampling_params.max_tokens = (
274 + self.model_config.max_model_len -
275 + len(decoder_inputs["prompt_token_ids"]))
276 + sampling_params.update_from_generation_config(
277 + self.generation_config_fields, eos_token_id)
278 + sampling_params.update_from_tokenizer(
279 + self.tokenizer.get_lora_tokenizer(lora_request))
280 + else:
281 + pooling_params = params.clone()
282
283 # Multimodal related.
284 sorted_mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]] = None
@@ -331,6 +335,7 @@ class Processor:
335 mm_hashes=sorted_mm_hashes,
336 mm_placeholders=sorted_mm_positions,
337 sampling_params=sampling_params,
338 + pooling_params=pooling_params,
339 eos_token_id=eos_token_id,
340 arrival_time=arrival_time,
341 lora_request=lora_request,
vllm/v1/metrics/loggers.py
+3 -2
@@ -481,8 +481,9 @@ class PrometheusStatLogger(StatLoggerBase):
481 finished_request.num_prompt_tokens)
482 self.histogram_num_generation_tokens_request.observe(
483 finished_request.num_generation_tokens)
484 - self.histogram_max_tokens_request.observe(
485 - finished_request.max_tokens_param)
484 + if finished_request.max_tokens_param:
485 + self.histogram_max_tokens_request.observe(
486 + finished_request.max_tokens_param)
487
488 if self.gauge_lora_info is not None:
489 running_lora_adapters = \
vllm/v1/metrics/stats.py
-1
@@ -106,7 +106,6 @@ class IterationStats:
106
107 self.num_generation_tokens += num_new_generation_tokens
108 if is_prefilling:
109 - assert num_new_generation_tokens > 0
109 self.num_prompt_tokens += prompt_len
110
111 first_token_latency = self._time_since(req_stats.arrival_time)
vllm/v1/outputs.py
+4
@@ -101,6 +101,9 @@ class ModelRunnerOutput:
101 # [prompt_len]
102 prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]]
103
104 + # [num_reqs, hidden_size]
105 + pooler_output: list[Optional[torch.Tensor]]
106 +
107 # [req_ids]
108 finished_sending: Optional[set[str]] = None
109 finished_recving: Optional[set[str]] = None
@@ -112,5 +115,6 @@ EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[],
115 spec_token_ids=None,
116 logprobs=None,
117 prompt_logprobs_dict={},
118 + pooler_output=[],
119 finished_sending=None,
120 finished_recving=None)
vllm/v1/pool/__init__.py
vllm/v1/pool/metadata.py new
+16
@@ -0,0 +1,16 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from dataclasses import dataclass
3 +from typing import Optional
4 +
5 +import torch
6 +
7 +from vllm.pooling_params import PoolingParams
8 +
9 +
10 +@dataclass
11 +class PoolingMetadata:
12 + """Tensors for pooling."""
13 +
14 + prompt_lens: torch.Tensor
15 + prompt_token_ids: Optional[torch.Tensor]
16 + pooling_params: list[PoolingParams]
vllm/v1/request.py
+30 -13
@@ -5,6 +5,7 @@ import enum
5 from typing import TYPE_CHECKING, Any, Optional, Union
6
7 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
8 +from vllm.pooling_params import PoolingParams
9 from vllm.sampling_params import SamplingParams
10 from vllm.utils import is_list_of
11 from vllm.v1.engine import (EngineCoreEvent, EngineCoreEventType,
@@ -25,7 +26,8 @@ class Request:
26 multi_modal_inputs: Optional[list[MultiModalKwargs]],
27 multi_modal_hashes: Optional[list[str]],
28 multi_modal_placeholders: Optional[list[PlaceholderRange]],
28 - sampling_params: SamplingParams,
29 + sampling_params: Optional[SamplingParams],
30 + pooling_params: Optional[PoolingParams],
31 eos_token_id: Optional[int],
32 client_index: int = 0,
33 lora_request: Optional["LoRARequest"] = None,
@@ -35,18 +37,35 @@ class Request:
37 self.request_id = request_id
38 self.client_index = client_index
39 self.sampling_params = sampling_params
40 + self.pooling_params = pooling_params
41 # Because of LoRA, the eos token id can be different for each request.
42 self.eos_token_id = eos_token_id
43 self.lora_request = lora_request
44 self.structured_output_request = structured_output_request
45
43 - self.status = (RequestStatus.WAITING_FOR_FSM
44 - if sampling_params.guided_decoding is not None else
45 - RequestStatus.WAITING)
46 + self.status = RequestStatus.WAITING
47 + if sampling_params and sampling_params.guided_decoding is not None:
48 + self.status = RequestStatus.WAITING_FOR_FSM
49 self.events: list[EngineCoreEvent] = []
50 self.stop_reason: Union[int, str, None] = None
48 - assert sampling_params.max_tokens is not None
49 - self.max_tokens = sampling_params.max_tokens
51 +
52 + # P/D: Connector-specific KV transfer parameters.
53 + self.kv_transfer_params: Optional[dict[str, Any]] = None
54 +
55 + if pooling_params is not None:
56 + self.max_tokens = 1
57 + elif sampling_params is not None:
58 + assert sampling_params.max_tokens is not None
59 + self.max_tokens = sampling_params.max_tokens
60 + if sampling_params.guided_decoding is not None:
61 + self.status = RequestStatus.WAITING_FOR_FSM
62 +
63 + if sampling_params.extra_args is not None:
64 + self.kv_transfer_params = \
65 + sampling_params.extra_args.get("kv_transfer_params")
66 + else:
67 + raise ValueError(
68 + "sampling_params and pooling_params can't both be unset")
69
70 self.prompt_token_ids = prompt_token_ids
71 self.num_prompt_tokens = len(self.prompt_token_ids)
@@ -63,11 +82,6 @@ class Request:
82 self.num_encoder_inputs = len(self.mm_inputs)
83 self.has_encoder_inputs = self.num_encoder_inputs > 0
84
66 - # P/D: Connector-specific KV transfer parameters.
67 - kv_params = (None if sampling_params.extra_args is None else
68 - sampling_params.extra_args.get("kv_transfer_params"))
69 - self.kv_transfer_params: Optional[dict[str, Any]] = kv_params
70 -
85 # Sanity check
86 assert len(self.mm_inputs) == len(self.mm_positions)
87 if self.mm_hashes:
@@ -98,10 +112,12 @@ class Request:
112 multi_modal_hashes=request.mm_hashes,
113 multi_modal_placeholders=request.mm_placeholders,
114 sampling_params=request.sampling_params,
115 + pooling_params=request.pooling_params,
116 eos_token_id=request.eos_token_id,
117 lora_request=request.lora_request,
118 structured_output_request=StructuredOutputRequest(
104 - sampling_params=request.sampling_params),
119 + sampling_params=request.sampling_params) \
120 + if request.sampling_params else None,
121 cache_salt=request.cache_salt,
122 )
123
@@ -141,7 +157,8 @@ class Request:
157
158 @property
159 def use_structured_output(self) -> bool:
144 - return self.sampling_params.guided_decoding is not None
160 + return self.sampling_params is not None and \
161 + self.sampling_params.guided_decoding is not None
162
163 def record_event(
164 self,
vllm/v1/structured_output/__init__.py
+3 -1
@@ -62,13 +62,15 @@ class StructuredOutputManager:
62 return
63
64 if TYPE_CHECKING:
65 - assert request.sampling_params.guided_decoding is not None
65 + assert request.sampling_params is not None and \
66 + request.sampling_params.guided_decoding is not None
67
68 # Initialize the backend the first time it is needed.
69 #
70 # NOTE: We only support a single backend. We do NOT support different
71 # backends on a per-request basis in V1 (for now, anyway...).
72 if self.backend is None:
73 + assert request.sampling_params is not None
74 backend = request.sampling_params.guided_decoding.backend
75 vocab_size = self.vllm_config.model_config.get_vocab_size()
76 if backend == "xgrammar":
vllm/v1/worker/gpu_input_batch.py
+101 -70
@@ -10,9 +10,11 @@ import torch
10
11 from vllm.lora.request import LoRARequest
12 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
13 +from vllm.pooling_params import PoolingParams
14 from vllm.sampling_params import SamplingParams, SamplingType
15 from vllm.utils import swap_dict_values
16 from vllm.v1.outputs import LogprobsTensors
17 +from vllm.v1.pool.metadata import PoolingMetadata
18 from vllm.v1.sample.metadata import SamplingMetadata
19 from vllm.v1.utils import copy_slice
20 from vllm.v1.worker.block_table import MultiGroupBlockTable
@@ -27,7 +29,8 @@ class CachedRequestState:
29 prompt_token_ids: list[int]
30 mm_inputs: list[MultiModalKwargs]
31 mm_positions: list[PlaceholderRange]
30 - sampling_params: SamplingParams
32 + sampling_params: Optional[SamplingParams]
33 + pooling_params: Optional[PoolingParams]
34 generator: Optional[torch.Generator]
35
36 block_ids: tuple[list[int], ...]
@@ -226,6 +229,8 @@ class InputBatch:
229 # This is updated each time the batch constituents change.
230 self.sampling_metadata = self._make_sampling_metadata()
231
232 + self.pooling_params: dict[str, PoolingParams] = {}
233 +
234 @property
235 def req_ids(self) -> list[str]:
236 # None elements should only be present transiently
@@ -269,77 +274,83 @@ class InputBatch:
274 self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
275 self.block_table.add_row(request.block_ids, req_index)
276
272 - sampling_params = request.sampling_params
273 - if sampling_params.sampling_type == SamplingType.GREEDY:
274 - # Avoid later division by zero.
275 - self.temperature_cpu[req_index] = -1.0
276 - self.greedy_reqs.add(req_id)
277 - else:
278 - self.temperature_cpu[req_index] = sampling_params.temperature
279 - self.random_reqs.add(req_id)
280 -
281 - self.top_p_cpu[req_index] = sampling_params.top_p
282 - if sampling_params.top_p < 1:
283 - self.top_p_reqs.add(req_id)
284 - top_k = sampling_params.top_k
285 - if 0 < top_k < self.vocab_size:
286 - self.top_k_reqs.add(req_id)
287 - else:
288 - top_k = self.vocab_size
289 - self.top_k_cpu[req_index] = top_k
290 - self.min_p_cpu[req_index] = sampling_params.min_p
291 - self.frequency_penalties_cpu[
292 - req_index] = sampling_params.frequency_penalty
293 - if sampling_params.min_p > _SAMPLING_EPS:
294 - self.min_p_reqs.add(req_id)
295 - if sampling_params.frequency_penalty != 0.0:
296 - self.frequency_penalties_reqs.add(req_id)
297 - self.presence_penalties_cpu[
298 - req_index] = sampling_params.presence_penalty
299 - if sampling_params.presence_penalty != 0.0:
300 - self.presence_penalties_reqs.add(req_id)
301 - self.repetition_penalties_cpu[
302 - req_index] = sampling_params.repetition_penalty
303 - if sampling_params.repetition_penalty != 1.0:
304 - self.repetition_penalties_reqs.add(req_id)
305 - if sampling_params.min_tokens:
306 - self.min_tokens[req_index] = (sampling_params.min_tokens,
307 - sampling_params.all_stop_token_ids)
308 -
309 - # NOTE(woosuk): self.generators should not include the requests that
310 - # do not have their own generator.
311 - if request.generator is not None:
312 - self.generators[req_index] = request.generator
313 -
314 - if sampling_params.logprobs is not None:
315 - self.num_logprobs[req_id] = sampling_params.logprobs
316 - if sampling_params.prompt_logprobs is not None:
317 - self.num_prompt_logprobs[req_id] = sampling_params.prompt_logprobs
318 - if sampling_params.logit_bias is not None:
319 - self.logit_bias[req_index] = sampling_params.logit_bias
320 -
321 - if sampling_params.allowed_token_ids:
322 - self.has_allowed_token_ids.add(req_id)
323 - if self.allowed_token_ids_mask_cpu_tensor is None:
324 - # Lazy allocation for this tensor, which can be large.
277 + if sampling_params := request.sampling_params:
278 + if sampling_params.sampling_type == SamplingType.GREEDY:
279 + # Avoid later division by zero.
280 + self.temperature_cpu[req_index] = -1.0
281 + self.greedy_reqs.add(req_id)
282 + else:
283 + self.temperature_cpu[req_index] = sampling_params.temperature
284 + self.random_reqs.add(req_id)
285 +
286 + self.top_p_cpu[req_index] = sampling_params.top_p
287 + if sampling_params.top_p < 1:
288 + self.top_p_reqs.add(req_id)
289 + top_k = sampling_params.top_k
290 + if 0 < top_k < self.vocab_size:
291 + self.top_k_reqs.add(req_id)
292 + else:
293 + top_k = self.vocab_size
294 + self.top_k_cpu[req_index] = top_k
295 + self.min_p_cpu[req_index] = sampling_params.min_p
296 + self.frequency_penalties_cpu[
297 + req_index] = sampling_params.frequency_penalty
298 + if sampling_params.min_p > _SAMPLING_EPS:
299 + self.min_p_reqs.add(req_id)
300 + if sampling_params.frequency_penalty != 0.0:
301 + self.frequency_penalties_reqs.add(req_id)
302 + self.presence_penalties_cpu[
303 + req_index] = sampling_params.presence_penalty
304 + if sampling_params.presence_penalty != 0.0:
305 + self.presence_penalties_reqs.add(req_id)
306 + self.repetition_penalties_cpu[
307 + req_index] = sampling_params.repetition_penalty
308 + if sampling_params.repetition_penalty != 1.0:
309 + self.repetition_penalties_reqs.add(req_id)
310 + if sampling_params.min_tokens:
311 + self.min_tokens[req_index] = (
312 + sampling_params.min_tokens,
313 + sampling_params.all_stop_token_ids)
314 +
315 + # NOTE(woosuk): self.generators should not include the requests that
316 + # do not have their own generator.
317 + if request.generator is not None:
318 + self.generators[req_index] = request.generator
319 +
320 + if sampling_params.logprobs is not None:
321 + self.num_logprobs[req_id] = sampling_params.logprobs
322 + if sampling_params.prompt_logprobs is not None:
323 + self.num_prompt_logprobs[
324 + req_id] = sampling_params.prompt_logprobs
325 + if sampling_params.logit_bias is not None:
326 + self.logit_bias[req_index] = sampling_params.logit_bias
327 +
328 + if sampling_params.allowed_token_ids:
329 + self.has_allowed_token_ids.add(req_id)
330 + if self.allowed_token_ids_mask_cpu_tensor is None:
331 + # Lazy allocation for this tensor, which can be large.
332 + # False means we don't fill with -inf.
333 + self.allowed_token_ids_mask = torch.zeros(
334 + self.max_num_reqs,
335 + self.vocab_size,
336 + dtype=torch.bool,
337 + device=self.device)
338 + self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
339 + self.max_num_reqs,
340 + self.vocab_size,
341 + dtype=torch.bool,
342 + device="cpu")
343 + self.allowed_token_ids_mask_cpu_tensor[req_index] = True
344 # False means we don't fill with -inf.
326 - self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,
327 - self.vocab_size,
328 - dtype=torch.bool,
329 - device=self.device)
330 - self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
331 - self.max_num_reqs,
332 - self.vocab_size,
333 - dtype=torch.bool,
334 - device="cpu")
335 - self.allowed_token_ids_mask_cpu_tensor[req_index] = True
336 - # False means we don't fill with -inf.
337 - self.allowed_token_ids_mask_cpu_tensor[req_index][
338 - sampling_params.allowed_token_ids] = False
345 + self.allowed_token_ids_mask_cpu_tensor[req_index][
346 + sampling_params.allowed_token_ids] = False
347
340 - if sampling_params.bad_words_token_ids:
341 - self.bad_words_token_ids[
342 - req_index] = sampling_params.bad_words_token_ids
348 + if sampling_params.bad_words_token_ids:
349 + self.bad_words_token_ids[
350 + req_index] = sampling_params.bad_words_token_ids
351 + else:
352 + assert request.pooling_params is not None
353 + self.pooling_params[req_id] = request.pooling_params
354
355 # Add request lora ID
356 if request.lora_request:
@@ -392,6 +403,7 @@ class InputBatch:
403 # False means we don't fill with -inf.
404 self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
405 self.bad_words_token_ids.pop(req_index, None)
406 + self.pooling_params.pop(req_id, None)
407 return req_index
408
409 def swap_states(self, i1: int, i2: int) -> None:
@@ -602,6 +614,25 @@ class InputBatch:
614 bad_words_token_ids=self.bad_words_token_ids,
615 )
616
617 + @property
618 + def pooling_metadata(self) -> PoolingMetadata:
619 + if len(self.pooling_params) == 0:
620 + pooling_params = []
621 + else:
622 + # Note, for now this assumes that all request in the batch
623 + # are either sampling or pooling requests
624 + assert len(self.req_ids) == len(self.pooling_params)
625 + pooling_params = [
626 + self.pooling_params[req_id] for req_id in self.req_ids
627 + ]
628 +
629 + return PoolingMetadata(
630 + prompt_lens=torch.from_numpy(
631 + self.num_prompt_tokens[:self.num_reqs]).to(self.device),
632 + prompt_token_ids=self.sampling_metadata.prompt_token_ids,
633 + pooling_params=pooling_params,
634 + )
635 +
636 def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
637 max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
638 prompt_token_ids_cpu_tensor = torch.empty(
vllm/v1/worker/gpu_model_runner.py
+115 -11
@@ -36,6 +36,7 @@ from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
36 from vllm.multimodal import MULTIMODAL_REGISTRY
37 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
38 from vllm.multimodal.utils import group_mm_inputs_by_modality
39 +from vllm.pooling_params import PoolingParams
40 from vllm.sampling_params import SamplingType
41 from vllm.sequence import IntermediateTensors
42 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
@@ -51,6 +52,7 @@ from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
52 SlidingWindowSpec)
53 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
54 ModelRunnerOutput)
55 +from vllm.v1.pool.metadata import PoolingMetadata
56 from vllm.v1.sample.metadata import SamplingMetadata
57 from vllm.v1.sample.rejection_sampler import RejectionSampler
58 from vllm.v1.sample.sampler import Sampler
@@ -119,6 +121,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
121 cache_config.cache_dtype]
122
123 self.is_multimodal_model = model_config.is_multimodal_model
124 + self.is_pooling_model = model_config.pooler_config is not None
125 self.max_model_len = model_config.max_model_len
126 self.max_num_tokens = scheduler_config.max_num_batched_tokens
127 self.max_num_reqs = scheduler_config.max_num_seqs
@@ -394,7 +397,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
397 for new_req_data in scheduler_output.scheduled_new_reqs:
398 req_id = new_req_data.req_id
399 sampling_params = new_req_data.sampling_params
397 - if sampling_params.sampling_type == SamplingType.RANDOM_SEED:
400 + pooling_params = new_req_data.pooling_params
401 + if sampling_params and \
402 + sampling_params.sampling_type == SamplingType.RANDOM_SEED:
403 generator = torch.Generator(device=self.device)
404 generator.manual_seed(sampling_params.seed)
405 else:
@@ -406,6 +411,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
411 mm_inputs=new_req_data.mm_inputs,
412 mm_positions=new_req_data.mm_positions,
413 sampling_params=sampling_params,
414 + pooling_params=pooling_params,
415 generator=generator,
416 block_ids=new_req_data.block_ids,
417 num_computed_tokens=new_req_data.num_computed_tokens,
@@ -563,7 +569,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
569 self,
570 scheduler_output: "SchedulerOutput",
571 ) -> tuple[dict[str, Any], bool, torch.Tensor,
566 - Optional[SpecDecodeMetadata]]:
572 + Optional[SpecDecodeMetadata], np.ndarray]:
573 """
574 :return: tuple[
575 attn_metadata: layer-to-attention_metadata mapping,
@@ -750,7 +756,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
756 self.set_active_loras(self.input_batch, num_scheduled_tokens)
757
758 return (attn_metadata, attention_cuda_graphs, logits_indices,
753 - spec_decode_metadata)
759 + spec_decode_metadata, num_scheduled_tokens)
760
761 def _compute_cascade_attn_prefix_len(
762 self,
@@ -1197,6 +1203,51 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1203 dtype=torch.int32)
1204 return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding
1205
1206 + def _pool(
1207 + self,
1208 + hidden_states: torch.Tensor,
1209 + num_scheduled_tokens: int,
1210 + num_scheduled_tokens_np: np.ndarray,
1211 + finished_sending: Optional[set[str]],
1212 + finished_recving: Optional[set[str]],
1213 + ) -> ModelRunnerOutput:
1214 + assert self.input_batch.num_reqs ==\
1215 + len(self.input_batch.pooling_params), \
1216 + "Either all or none of the requests in" \
1217 + " a batch must be pooling request"
1218 +
1219 + extracted_hidden_states = list(
1220 + torch.split(hidden_states[:num_scheduled_tokens],
1221 + num_scheduled_tokens_np.tolist()))
1222 +
1223 + pooling_metadata = self.input_batch.pooling_metadata
1224 +
1225 + raw_pooler_output = self.model.pooler(
1226 + hidden_states=extracted_hidden_states,
1227 + pooling_metadata=pooling_metadata)
1228 +
1229 + pooler_output: list[Optional[torch.Tensor]] = []
1230 + seq_lens = self.seq_lens[:self.input_batch.num_reqs]
1231 + for raw_output, seq_len, prompt_len in zip(
1232 + raw_pooler_output, seq_lens, pooling_metadata.prompt_lens):
1233 +
1234 + if seq_len == prompt_len:
1235 + pooler_output.append(raw_output.data.cpu())
1236 + else:
1237 + pooler_output.append(None)
1238 +
1239 + return ModelRunnerOutput(
1240 + req_ids=self.input_batch.req_ids,
1241 + req_id_to_index=self.input_batch.req_id_to_index,
1242 + sampled_token_ids=[],
1243 + spec_token_ids=None,
1244 + logprobs=None,
1245 + prompt_logprobs_dict={},
1246 + pooler_output=pooler_output,
1247 + finished_sending=finished_sending,
1248 + finished_recving=finished_recving,
1249 + )
1250 +
1251 @torch.inference_mode()
1252 def execute_model(
1253 self,
@@ -1214,7 +1265,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1265
1266 # Prepare the decoder inputs.
1267 (attn_metadata, attention_cuda_graphs, logits_indices,
1217 - spec_decode_metadata) = (self._prepare_inputs(scheduler_output))
1268 + spec_decode_metadata,
1269 + num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
1270 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1271 if (self.use_cuda_graph
1272 and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -1284,7 +1336,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1336 # compiled with full CUDA graphs, we have to skip them entirely.
1337 skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
1338
1287 - # Run the decoder.
1339 + # Run the model.
1340 # Use persistent buffers for CUDA graphs.
1341 with set_forward_context(
1342 attn_metadata,
@@ -1326,6 +1378,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1378 all_gather_group=get_tp_group())
1379 logits = None
1380 else:
1381 + if self.input_batch.pooling_params:
1382 + return self._pool(hidden_states, num_scheduled_tokens,
1383 + num_scheduled_tokens_np, finished_sending,
1384 + finished_recving)
1385 +
1386 sample_hidden_states = hidden_states[logits_indices]
1387 logits = self.model.compute_logits(sample_hidden_states, None)
1388 if broadcast_pp_output:
@@ -1541,6 +1598,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1598 spec_token_ids=spec_token_ids,
1599 logprobs=logprobs_lists,
1600 prompt_logprobs_dict=prompt_logprobs_dict,
1601 + pooler_output=[],
1602 finished_sending=finished_sending,
1603 finished_recving=finished_recving,
1604 )
@@ -1802,7 +1860,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1860 self,
1861 num_tokens: int,
1862 capture_attn_cudagraph: bool = False,
1805 - ) -> torch.Tensor:
1863 + ) -> tuple[torch.Tensor, torch.Tensor]:
1864
1865 # Padding for DP
1866 num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
@@ -1899,7 +1957,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1957 self.drafter.dummy_run(num_tokens)
1958
1959 logit_indices = np.cumsum(num_scheduled_tokens) - 1
1902 - return hidden_states[logit_indices]
1960 + return hidden_states, hidden_states[logit_indices]
1961
1962 @torch.inference_mode()
1963 def _dummy_sampler_run(
@@ -1978,6 +2036,48 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2036 )
2037 return sampler_output
2038
2039 + @torch.inference_mode()
2040 + def _dummy_pooler_run(
2041 + self,
2042 + hidden_states: torch.Tensor,
2043 + ) -> torch.Tensor:
2044 +
2045 + num_tokens = hidden_states.shape[0]
2046 + max_num_reqs = self.scheduler_config.max_num_seqs
2047 + num_reqs = min(num_tokens, max_num_reqs)
2048 + min_tokens_per_req = num_tokens // num_reqs
2049 + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
2050 + num_scheduled_tokens_list[-1] += num_tokens % num_reqs
2051 + assert sum(num_scheduled_tokens_list) == num_tokens
2052 + assert len(num_scheduled_tokens_list) == num_reqs
2053 +
2054 + hidden_states_list = list(
2055 + torch.split(hidden_states, num_scheduled_tokens_list))
2056 +
2057 + req_num_tokens = num_tokens // num_reqs
2058 +
2059 + dummy_metadata = PoolingMetadata(
2060 + prompt_lens=torch.tensor([h.shape[0] for h in hidden_states_list],
2061 + device=self.device),
2062 + prompt_token_ids=torch.zeros((num_reqs, req_num_tokens),
2063 + dtype=torch.int32,
2064 + device=self.device),
2065 + pooling_params=[PoolingParams()] * num_reqs)
2066 +
2067 + try:
2068 + pooler_output = self.model.pooler(hidden_states=hidden_states_list,
2069 + pooling_metadata=dummy_metadata)
2070 + except RuntimeError as e:
2071 + if 'out of memory' in str(e):
2072 + raise RuntimeError(
2073 + "CUDA out of memory occurred when warming up pooler with "
2074 + f"{num_reqs} dummy requests. Please try lowering "
2075 + "`max_num_seqs` or `gpu_memory_utilization` when "
2076 + "initializing the engine.") from e
2077 + else:
2078 + raise e
2079 + return pooler_output
2080 +
2081 def profile_run(self) -> None:
2082 # Profile with multimodal encoder & encoder cache.
2083 # TODO: handle encoder-decoder models once we support them.
@@ -2048,13 +2148,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2148 # Cache the dummy encoder outputs.
2149 self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
2150
2051 - hidden_states = self._dummy_run(self.max_num_tokens)
2151 + hidden_states, last_hidden_states \
2152 + = self._dummy_run(self.max_num_tokens)
2153 if get_pp_group().is_last_rank:
2053 - sampler_output = self._dummy_sampler_run(hidden_states)
2154 + if self.is_pooling_model:
2155 + output = self._dummy_pooler_run(hidden_states)
2156 + else:
2157 + output = self._dummy_sampler_run(last_hidden_states)
2158 else:
2055 - sampler_output = None
2159 + output = None
2160 self._sync_device()
2057 - del hidden_states, sampler_output
2161 + del hidden_states, output
2162 self.encoder_cache.clear()
2163 gc.collect()
2164
vllm/v1/worker/gpu_worker.py
+8 -3
@@ -273,9 +273,14 @@ class Worker(WorkerBase):
273 if get_pp_group().is_last_rank:
274 max_num_reqs = min(self.scheduler_config.max_num_seqs,
275 self.scheduler_config.max_num_batched_tokens)
276 - self.model_runner._dummy_sampler_run(
277 - hidden_states=self.model_runner._dummy_run(
278 - num_tokens=max_num_reqs))
276 +
277 + hidden_states, last_hidden_states = \
278 + self.model_runner._dummy_run(num_tokens=max_num_reqs)
279 + if self.model_runner.is_pooling_model:
280 + self.model_runner._dummy_pooler_run(hidden_states)
281 + else:
282 + self.model_runner._dummy_sampler_run(
283 + hidden_states=last_hidden_states)
284
285 # Reset the seed to ensure that the random state is not affected by
286 # the model initialization and profiling.
vllm/v1/worker/tpu_input_batch.py
+1
@@ -231,6 +231,7 @@ class InputBatch:
231 self.block_table.add_row(request.block_ids, req_index)
232
233 sampling_params = request.sampling_params
234 + assert sampling_params is not None, "pooling requests not supported yet"
235 if sampling_params.sampling_type == SamplingType.GREEDY:
236 # Avoid later division by zero.
237 self.temperature_cpu[req_index] = -1.0
vllm/v1/worker/tpu_model_runner.py
+4
@@ -386,6 +386,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
386 req_ids_to_add: list[str] = []
387 # Add new requests to the cached states.
388 for new_req_data in scheduler_output.scheduled_new_reqs:
389 + assert new_req_data.sampling_params is not None,\
390 + "Pooling is not supported in TPU yet"
391 req_id = new_req_data.req_id
392 sampling_params = new_req_data.sampling_params
393
@@ -395,6 +397,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
397 mm_inputs=new_req_data.mm_inputs,
398 mm_positions=new_req_data.mm_positions,
399 sampling_params=sampling_params,
400 + pooling_params=None,
401 generator=None,
402 block_ids=new_req_data.block_ids,
403 num_computed_tokens=new_req_data.num_computed_tokens,
@@ -956,6 +959,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
959 spec_token_ids=None,
960 logprobs=logprobs_lists,
961 prompt_logprobs_dict=prompt_logprobs_dict,
962 + pooler_output=[],
963 )
964
965 # Check there are no new graphs compiled - all the graphs should be