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
[Bugfix][v1] Fix step pooler implementation and step pooling usage in v1 (#19956)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Jun 24, 2025 at 02:38 UTC
61f4fc5dc6f165a2d2c38c277458d04fdd5c8bf5
14 files changed
+164
-40
tests/conftest.py
+11
-7
@@ -1027,13 +1027,13 @@ class VllmRunner:
1027
req_outputs = self.model.classify(prompts)
1028
return [req_output.outputs.probs for req_output in req_outputs]
1029
1030
- def encode(self,
1031
- prompts: list[str],
1032
- images: Optional[PromptImageInput] = None,
1033
- videos: Optional[PromptVideoInput] = None,
1034
- audios: Optional[PromptAudioInput] = None,
1035
- *args,
1036
- **kwargs) -> list[list[float]]:
1030
+ def embed(self,
1031
+ prompts: list[str],
1032
+ images: Optional[PromptImageInput] = None,
1033
+ videos: Optional[PromptVideoInput] = None,
1034
+ audios: Optional[PromptAudioInput] = None,
1035
+ *args,
1036
+ **kwargs) -> list[list[float]]:
1037
inputs = self.get_inputs(prompts,
1038
images=images,
1039
videos=videos,
@@ -1042,6 +1042,10 @@ class VllmRunner:
1042
req_outputs = self.model.embed(inputs, *args, **kwargs)
1043
return [req_output.outputs.embedding for req_output in req_outputs]
1044
1045
+ def encode(self, prompts: list[str]) -> list[list[float]]:
1046
+ req_outputs = self.model.encode(prompts)
1047
+ return [req_output.outputs.data for req_output in req_outputs]
1048
+
1049
def score(
1050
self,
1051
text_1: Union[str, list[str]],
tests/model_executor/test_model_load_with_params.py
+6
-6
@@ -29,8 +29,8 @@ def test_model_loading_with_params(vllm_runner):
29
revision=REVISION,
30
dtype="float16",
31
max_model_len=MAX_MODEL_LEN) as vllm_model:
32
- output = vllm_model.encode("Write a short story about a robot that"
33
- " dreams for the first time.\n")
32
+ output = vllm_model.embed("Write a short story about a robot that"
33
+ " dreams for the first time.\n")
34
35
model_config = vllm_model.model.llm_engine.model_config
36
model_tokenizer = vllm_model.model.llm_engine.tokenizer
@@ -67,8 +67,8 @@ def test_roberta_model_loading_with_params(vllm_runner):
67
revision=REVISION_ROBERTA,
68
dtype="float16",
69
max_model_len=MAX_MODEL_LEN) as vllm_model:
70
- output = vllm_model.encode("Write a short story about a robot that"
71
- " dreams for the first time.\n")
70
+ output = vllm_model.embed("Write a short story about a robot that"
71
+ " dreams for the first time.\n")
72
73
model_config = vllm_model.model.llm_engine.model_config
74
model_tokenizer = vllm_model.model.llm_engine.tokenizer
@@ -105,8 +105,8 @@ def test_facebook_roberta_model_loading_with_params(vllm_runner):
105
with vllm_runner(model_name=model_name,
106
dtype="float16",
107
max_model_len=MAX_MODEL_LEN) as vllm_model:
108
- output = vllm_model.encode("Write a short story about a robot that"
109
- " dreams for the first time.\n")
108
+ output = vllm_model.embed("Write a short story about a robot that"
109
+ " dreams for the first time.\n")
110
111
model_tokenizer = vllm_model.model.llm_engine.tokenizer
112
assert model_tokenizer.tokenizer_id == model_name
tests/models/language/pooling/embed_utils.py
+1
-1
@@ -55,7 +55,7 @@ def correctness_test_embed_models(hf_runner,
55
task="embed",
56
max_model_len=None,
57
**vllm_extra_kwargs) as vllm_model:
58
- vllm_outputs = vllm_model.encode(example_prompts)
58
+ vllm_outputs = vllm_model.embed(example_prompts)
59
60
with hf_runner(
61
model_info.name,
tests/models/language/pooling/test_embedding.py
+1
-1
@@ -89,7 +89,7 @@ def test_models(
89
task="embed",
90
max_model_len=512,
91
**vllm_extra_kwargs) as vllm_model:
92
- vllm_outputs = vllm_model.encode(example_prompts)
92
+ vllm_outputs = vllm_model.embed(example_prompts)
93
94
check_embeddings_close(
95
embeddings_0_lst=hf_outputs,
tests/models/language/pooling/test_jina.py
+2
-2
@@ -98,11 +98,11 @@ def test_matryoshka(
98
99
if dimensions not in matryoshka_dimensions:
100
with pytest.raises(ValueError):
101
- vllm_model.encode(
101
+ vllm_model.embed(
102
example_prompts,
103
pooling_params=PoolingParams(dimensions=dimensions))
104
else:
105
- vllm_outputs = vllm_model.encode(
105
+ vllm_outputs = vllm_model.embed(
106
example_prompts,
107
pooling_params=PoolingParams(dimensions=dimensions))
108
tests/models/language/pooling/test_reward.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+import torch.nn.functional as F
6
+from transformers import AutoModel
7
+
8
+from vllm.platforms import current_platform
9
+
10
+from ....conftest import HfRunner
11
+
12
+
13
+@pytest.fixture(autouse=True)
14
+def v1(run_with_both_engines):
15
+ # Simple autouse wrapper to run both engines for each test
16
+ # This can be promoted up to conftest.py to run for every
17
+ # test in a package
18
+ pass
19
+
20
+
21
+@pytest.fixture
22
+def math_step_prompts():
23
+ # ruff: noqa: E501
24
+ data = {
25
+ "system":
26
+ "Please reason step by step, and put your final answer within \\boxed{}. ",
27
+ "query":
28
+ "Sue lives in a fun neighborhood. One weekend, the neighbors decided to play a prank on Sue. On Friday morning, the neighbors placed 18 pink plastic flamingos out on Sue's front yard. On Saturday morning, the neighbors took back one third of the flamingos, painted them white, and put these newly painted white flamingos back out on Sue's front yard. Then, on Sunday morning, they added another 18 pink plastic flamingos to the collection. At noon on Sunday, how many more pink plastic flamingos were out than white plastic flamingos?",
29
+ "response": [
30
+ "To find out how many more pink plastic flamingos were out than white plastic flamingos at noon on Sunday, we can break down the problem into steps. First, on Friday, the neighbors start with 18 pink plastic flamingos.",
31
+ "On Saturday, they take back one third of the flamingos. Since there were 18 flamingos, (1/3 \\times 18 = 6) flamingos are taken back. So, they have (18 - 6 = 12) flamingos left in their possession. Then, they paint these 6 flamingos white and put them back out on Sue's front yard. Now, Sue has the original 12 pink flamingos plus the 6 new white ones. Thus, by the end of Saturday, Sue has (12 + 6 = 18) pink flamingos and 6 white flamingos.",
32
+ "On Sunday, the neighbors add another 18 pink plastic flamingos to Sue's front yard. By the end of Sunday morning, Sue has (18 + 18 = 36) pink flamingos and still 6 white flamingos.",
33
+ "To find the difference, subtract the number of white flamingos from the number of pink flamingos: (36 - 6 = 30). Therefore, at noon on Sunday, there were 30 more pink plastic flamingos out than white plastic flamingos. The answer is (\\boxed{30}).",
34
+ ],
35
+ }
36
+ answer = "<extra_0>".join(data['response']) + "<extra_0>"
37
+ prompt = f"<im_start>system\n{data['system']}<im_end>\n<im_start>user\n{data['query']}<im_end>\n<im_start>assistant\n{answer}<im_end><|endoftext|>"
38
+ return [prompt]
39
+
40
+
41
+def step_reward_patch_hf_model(hf_model: HfRunner):
42
+
43
+ # Patch the hf_runner to use the step reward function
44
+ def make_step_rewards(logits: torch.Tensor,
45
+ token_masks: torch.Tensor) -> list[list[float]]:
46
+ probabilities = F.softmax(logits, dim=-1)
47
+ probabilities = probabilities * token_masks.unsqueeze(-1)
48
+
49
+ all_scores_res: list[list[float]] = []
50
+ for i in range(probabilities.size(0)):
51
+ sample = probabilities[i] # seq_len, num_labels
52
+ positive_probs = sample[sample != 0].view(-1, 2)
53
+ non_zero_elements_list = positive_probs.cpu().tolist()
54
+ all_scores_res.append(non_zero_elements_list)
55
+ return all_scores_res
56
+
57
+ def reward(prompts: list[str]) -> list[list[float]]:
58
+ input_ids = hf_model.tokenizer(prompts, return_tensors="pt").input_ids
59
+ input_ids = hf_model.wrap_device(input_ids)
60
+ outputs = hf_model.model(input_ids=input_ids)
61
+
62
+ step_sep_id = hf_model.tokenizer.encode("<extra_0>")[0]
63
+ token_masks = (input_ids == step_sep_id)
64
+ return make_step_rewards(outputs[0], token_masks)
65
+
66
+ hf_model.reward = reward # type: ignore[attr-defined]
67
+
68
+ return hf_model
69
+
70
+
71
+@pytest.mark.parametrize(
72
+ "model",
73
+ [
74
+ pytest.param("Qwen/Qwen2.5-Math-PRM-7B",
75
+ marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
76
+ ],
77
+)
78
+@pytest.mark.parametrize("dtype", ["half"])
79
+def test_prm_models(
80
+ hf_runner,
81
+ vllm_runner,
82
+ math_step_prompts,
83
+ model: str,
84
+ dtype: str,
85
+ monkeypatch,
86
+) -> None:
87
+ if current_platform.is_rocm():
88
+ # ROCm Triton FA does not currently support sliding window attention
89
+ # switch to use ROCm CK FA backend
90
+ monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
91
+
92
+ with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
93
+ vllm_outputs = vllm_model.encode(math_step_prompts)
94
+
95
+ with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
96
+ hf_model = step_reward_patch_hf_model(hf_model)
97
+ hf_outputs = hf_model.reward(math_step_prompts)
98
+
99
+ # check logits difference
100
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
101
+ hf_output = torch.tensor(hf_output)
102
+ vllm_output = torch.tensor(vllm_output)
103
+
104
+ assert torch.allclose(hf_output, vllm_output, 1e-2)
tests/models/multimodal/pooling/test_dse_qwen2_vl.py
+2
-2
@@ -98,7 +98,7 @@ def _run_test(
98
max_model_len=8192) as vllm_model:
99
tokenizer = vllm_model.model.get_tokenizer()
100
texts = [
101
- # this is necessary because vllm_model.encode will not apply any
101
+ # this is necessary because vllm_model.embed will not apply any
102
# templating to the prompt, and therefore lacks an image_pad
103
# token unless one is inserted beforehand (the (28,28) image
104
# above is converted to an image pad token by the chat template).
@@ -109,7 +109,7 @@ def _run_test(
109
# vllm will replace the pad token with the actual image,
110
# which may be a placeholder image, later.
111
]
112
- vllm_outputs = vllm_model.encode(texts, images=input_images)
112
+ vllm_outputs = vllm_model.embed(texts, images=input_images)
113
114
hf_outputs = []
115
with hf_runner(model,
tests/models/multimodal/pooling/test_llava_next.py
+1
-1
@@ -68,7 +68,7 @@ def _run_test(
68
dtype=dtype,
69
max_model_len=4096,
70
enforce_eager=True) as vllm_model:
71
- vllm_outputs = vllm_model.encode(input_texts, images=input_images)
71
+ vllm_outputs = vllm_model.embed(input_texts, images=input_images)
72
73
with hf_runner(model, dtype=dtype,
74
auto_cls=AutoModelForImageTextToText) as hf_model:
tests/models/multimodal/pooling/test_phi3v.py
+1
-1
@@ -46,7 +46,7 @@ def _run_test(
46
# will hurt multiprocessing backend with fork method (the default method).
47
with vllm_runner(model, task="embed", dtype=dtype,
48
enforce_eager=True) as vllm_model:
49
- vllm_outputs = vllm_model.encode(input_texts, images=input_images)
49
+ vllm_outputs = vllm_model.embed(input_texts, images=input_images)
50
51
# use eager mode for hf runner, since phi3_v didn't work with flash_attn
52
hf_model_kwargs = {"_attn_implementation": "eager"}
tests/quantization/test_bitsandbytes.py
+1
-1
@@ -161,7 +161,7 @@ def test_4bit_bnb_embedding_model(
161
dtype=dtype,
162
gpu_memory_utilization=0.5,
163
quantization="bitsandbytes") as vllm_model:
164
- vllm_outputs = vllm_model.encode(example_prompts)
164
+ vllm_outputs = vllm_model.embed(example_prompts)
165
check_embeddings_close(
166
embeddings_0_lst=hf_outputs,
167
embeddings_1_lst=vllm_outputs,
vllm/model_executor/layers/pooler.py
+6
-7
@@ -239,25 +239,24 @@ class StepPool(SimplePooler):
239
prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
240
prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)
241
242
- pooled_data: list[torch.Tensor] = []
243
-
242
+ pooled_data_lst = list[torch.Tensor]()
243
if isinstance(hidden_states, list):
244
for req_state, prompt_len in zip(hidden_states, prompt_lens):
245
assert prompt_len == req_state.shape[0], \
247
- "partial prefill not supported with mean pooling"
248
- pooled_data = hidden_states
246
+ "partial prefill not supported with step pooling"
247
+ pooled_data_lst = hidden_states
248
else:
249
offset = 0
250
for prompt_len in prompt_lens:
251
pooled_data_i = hidden_states[offset:offset + prompt_len]
252
offset += prompt_len
254
- pooled_data.append(pooled_data_i)
253
+ pooled_data_lst.append(pooled_data_i)
254
256
- pooled_data = []
255
+ pooled_data = list[torch.Tensor]()
256
returned_token_ids = self.returned_token_ids
257
step_tag_id = self.step_tag_id
258
260
- for data, token_id in zip(pooled_data, prompt_token_ids):
259
+ for data, token_id in zip(pooled_data_lst, prompt_token_ids):
260
if returned_token_ids is not None and len(returned_token_ids) > 0:
261
data = data[:, returned_token_ids]
262
vllm/model_executor/models/interfaces.py
+6
@@ -489,6 +489,12 @@ def supports_cross_encoding(
489
return is_pooling_model(model) and _supports_cross_encoding(model)
490
491
492
+def has_step_pooler(model: Union[type[object], object]) -> bool:
493
+ """Check if the model uses step pooler."""
494
+ return is_pooling_model(model) and any(
495
+ type(module).__name__ == "StepPool" for module in model.modules())
496
+
497
+
498
class SupportsQuant:
499
"""The interface required for all models that support quantization."""
500
vllm/v1/worker/gpu_input_batch.py
+19
-11
@@ -59,14 +59,15 @@ class CachedRequestState:
59
class InputBatch:
60
61
def __init__(
62
- self,
63
- max_num_reqs: int,
64
- max_model_len: int,
65
- max_num_batched_tokens: int,
66
- device: torch.device,
67
- pin_memory: bool,
68
- vocab_size: int,
69
- block_sizes: list[int], # The block_size of each kv cache group
62
+ self,
63
+ max_num_reqs: int,
64
+ max_model_len: int,
65
+ max_num_batched_tokens: int,
66
+ device: torch.device,
67
+ pin_memory: bool,
68
+ vocab_size: int,
69
+ block_sizes: list[int], # The block_size of each kv cache group
70
+ logits_processing_needs_token_ids: bool = False,
71
):
72
self.max_num_reqs = max_num_reqs
73
self.max_model_len = max_model_len
@@ -74,6 +75,8 @@ class InputBatch:
75
self.device = device
76
self.pin_memory = pin_memory
77
self.vocab_size = vocab_size
78
+ self.logits_processing_needs_token_ids = (
79
+ logits_processing_needs_token_ids)
80
81
self._req_ids: list[Optional[str]] = []
82
self.req_id_to_index: dict[str, int] = {}
@@ -579,9 +582,14 @@ class InputBatch:
582
copy_slice(self.repetition_penalties_cpu_tensor,
583
self.repetition_penalties, num_reqs)
584
582
- # The prompt tokens are used only for applying penalties during
583
- # the sampling process. Hence copy these tensors only when
584
- # there are requests which need penalties to be applied.
585
+ needs_prompt_token_ids = (not self.no_penalties or
586
+ (self.num_reqs > 0
587
+ and self.logits_processing_needs_token_ids))
588
+ if needs_prompt_token_ids:
589
+ # The prompt tokens are used only for applying penalties or
590
+ # step pooling during the sampling/pooling process.
591
+ # Hence copy these tensors only when there are requests which
592
+ # need penalties/step_pooler to be applied.
593
prompt_token_ids = self._make_prompt_token_ids_tensor()
594
else:
595
prompt_token_ids = None
vllm/v1/worker/gpu_model_runner.py
+3
@@ -33,6 +33,7 @@ from vllm.logger import init_logger
33
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
34
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
35
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
36
+from vllm.model_executor.models.interfaces import has_step_pooler
37
from vllm.multimodal import MULTIMODAL_REGISTRY
38
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
39
from vllm.multimodal.utils import group_mm_inputs_by_modality
@@ -1708,6 +1709,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1709
)
1710
model_loader.load_weights(self.model,
1711
model_config=self.model_config)
1712
+ if has_step_pooler(self.model):
1713
+ self.input_batch.logits_processing_needs_token_ids = True
1714
if self.lora_config:
1715
self.model = self.load_lora_model(self.model,
1716
self.model_config,