@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 74c4d80c6c

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 12, 01:00 UTC

[Model][6/N] Improve all pooling task | Support chunked prefill with ALL pooling (#27145)

Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

wang.yuqi committed Dec 4, 2025 at 21:44 UTC 74c4d80c6ca8160578b6e812079cb11dfd8b3d22

15 files changed +224 -93

docs/features/README.md

+1 -1

@@ -54,7 +54,7 @@ th:not(:first-child) {

54 | beam-search | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ❔ | ✅ | ✅ | |

55 | [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ❌ | ❔ | ❔ | ✅ |

57 -\* Chunked prefill and prefix caching are only applicable to last-token pooling.

57 +\* Chunked prefill and prefix caching are only applicable to last-token or all pooling with causal attention.

58 <sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.

60 ### Feature x Hardware

tests/entrypoints/pooling/classify/test_offline.py

+2 -5

@@ -61,11 +61,8 @@ def test_pooling_params(llm: LLM):

63 @pytest.mark.skip_global_cleanup

64 -def test_encode_api(llm: LLM):

65 - # chunked prefill does not support all pooling

66 - err_msg = "pooling_task must be one of.+"

67 - with pytest.raises(ValueError, match=err_msg):

68 - llm.encode(prompts, pooling_task="token_classify", use_tqdm=False)

64 +def test_token_classify(llm: LLM):

65 + llm.encode(prompts, pooling_task="token_classify", use_tqdm=False)

68 def test_score_api(llm: LLM):

tests/entrypoints/pooling/classify/test_online.py

+6 -6

@@ -255,21 +255,21 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):

255 @pytest.mark.asyncio

256 @pytest.mark.parametrize("model_name", [MODEL_NAME])

257 async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):

258 - # token_classify uses ALL pooling, which does not support chunked prefill.

258 task = "token_classify"

259 + input_text = ["This product was excellent and exceeded my expectations"]

260 response = requests.post(

261 server.url_for("pooling"),

262 json={

263 "model": model_name,

264 - "input": "test",

264 + "input": input_text,

265 "encoding_format": "float",

266 "task": task,

267 },

268 )

269 - assert response.json()["error"]["type"] == "BadRequestError"

270 - assert response.json()["error"]["message"].startswith(

271 - f"Task {task} is not supported"

272 - )

269 + poolings = PoolingResponse.model_validate(response.json())

270 + assert len(poolings.data) == 1

271 + assert len(poolings.data[0].data) == 8

272 + assert len(poolings.data[0].data[0]) == 2

273

274

275 @pytest.mark.asyncio

tests/entrypoints/pooling/embed/test_offline.py

+1 -1

@@ -42,7 +42,7 @@ def llm():

44 @pytest.mark.skip_global_cleanup

45 -def test_encode_api(llm: LLM):

45 +def test_token_embed(llm: LLM):

46 outputs = llm.encode(prompts, pooling_task="token_embed", use_tqdm=False)

47 multi_vector = outputs[0].outputs.data

48 assert multi_vector.shape == (11, 384)

tests/entrypoints/pooling/reward/test_offline.py

@@ -36,6 +36,13 @@ def llm():

36 cleanup_dist_env_and_memory()

39 +@pytest.mark.skip_global_cleanup

40 +def test_config(llm: LLM):

41 + vllm_config = llm.llm_engine.vllm_config

42 + assert vllm_config.cache_config.enable_prefix_caching

43 + assert vllm_config.scheduler_config.enable_chunked_prefill

44 +

45 +

46 def test_pooling_params(llm: LLM):

47 def get_outputs(use_activation):

48 outputs = llm.reward(

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py new

+53

@@ -0,0 +1,53 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import pytest

4 +import torch

5 +from transformers import AutoModel

6 +

7 +from tests.models.utils import check_embeddings_close

8 +from vllm import TokensPrompt

9 +

10 +

11 +@pytest.mark.parametrize(

12 + "model",

13 + ["Qwen/Qwen3-Embedding-0.6B"],

14 +)

15 +@torch.inference_mode

16 +def test_embed_models(hf_runner, vllm_runner, model: str):

17 + chunk_size = 10

18 + n_prompt_tokens = [55, 56, 57]

19 + token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]

20 +

21 + with vllm_runner(

22 + model,

23 + runner="pooling",

24 + max_model_len=128,

25 + max_num_batched_tokens=chunk_size,

26 + enforce_eager=True,

27 + # `enable_chunked_prefill`: Set to `False` instead of `None` in VllmRunner

28 + enable_chunked_prefill=True,

29 + enable_prefix_caching=True,

30 + ) as vllm_model:

31 + vllm_outputs = vllm_model.token_embed(

32 + [TokensPrompt(prompt_token_ids=t) for t in token_prompts],

33 + )

34 +

35 + with hf_runner(

36 + model,

37 + auto_cls=AutoModel,

38 + ) as hf_model:

39 + hf_outputs = []

40 + for token_prompt in token_prompts:

41 + inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})

42 + input_ids = inputs["input_ids"]

43 + output = hf_model.model(input_ids)

44 + hf_outputs.append(output.last_hidden_state.cpu().float()[0])

45 +

46 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):

47 + check_embeddings_close(

48 + embeddings_0_lst=hf_output,

49 + embeddings_1_lst=vllm_output,

50 + name_0="hf",

51 + name_1="vllm",

52 + tol=1e-2,

53 + )

tests/models/language/pooling/test_extract_hidden_states.py

-1

@@ -20,7 +20,6 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):

20 max_model_len=128,

21 enforce_eager=True,

22 runner="pooling",

23 - enable_chunked_prefill=False,

23 enable_prefix_caching=True,

24 ) as vllm_model:

25 pooling_outputs = vllm_model.llm.encode(

tests/test_config.py

+4 -4

@@ -629,8 +629,8 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files):

629 (

630 "internlm/internlm2-1_8b-reward",

631 "decoder",

632 - False,

633 - "Pooling models with all pooling does not support chunked prefill.",

632 + True,

633 + "Pooling models with causal attn and all pooling support chunked prefill.",

634 ),

635 (

636 "BAAI/bge-base-en",

@@ -748,8 +748,8 @@ def test_is_chunked_prefill_supported(

748 (

749 "internlm/internlm2-1_8b-reward",

750 "decoder",

751 - False,

752 - "Pooling models with all pooling does not support prefix caching.",

751 + True,

752 + "Pooling models with causal attn and all pooling support prefix caching.",

753 ),

754 (

755 "BAAI/bge-base-en",

vllm/config/model.py

+14 -10

@@ -1780,20 +1780,22 @@ class ModelConfig:

1780 return False

1781 elif attn_type == "decoder":

1782 pooling_type = self.pooler_config.pooling_type.lower()

1783 - if pooling_type in ["all", "mean", "step", "cls"]:

1783 + if pooling_type in ["mean", "step", "cls"]:

1784 logger.debug(

1785 "Pooling models with %s pooling does not "

1786 "support chunked prefill.",

1787 pooling_type,

1788 )

1789 return False

1790 - else:

1791 - # pooling_type == "last"

1790 + elif pooling_type in ["all", "last"]:

1791 logger.debug(

1793 - "Pooling models with causal attn and last pooling support "

1794 - "chunked prefill."

1792 + "Pooling models with causal attn and %s pooling support "

1793 + "chunked prefill.",

1794 + pooling_type,

1795 )

1796 return True

1797 + else:

1798 + raise ValueError(f"{pooling_type=} not supported.")

1799 # vllm currently does not have pooling models using hybrid,

1800 # attention_free or encoder_decoder attn types.

1801 return attn_type != "encoder_decoder"

@@ -1817,20 +1819,22 @@ class ModelConfig:

1819 return False

1820 elif attn_type == "decoder":

1821 pooling_type = self.pooler_config.pooling_type.lower()

1820 - if pooling_type in ["all", "mean", "step", "cls"]:

1822 + if pooling_type in ["mean", "step", "cls"]:

1823 logger.debug(

1824 "Pooling models with %s pooling does not "

1825 "support prefix caching.",

1826 pooling_type,

1827 )

1828 return False

1827 - else:

1828 - # pooling_type == "last"

1829 + elif pooling_type in ["all", "last"]:

1830 logger.debug(

1830 - "Pooling models with causal attn and last pooling support "

1831 - "prefix caching."

1831 + "Pooling models with causal attn and %s pooling support "

1832 + "prefix caching.",

1833 + pooling_type,

1834 )

1835 return True

1836 + else:

1837 + raise ValueError(f"{pooling_type=} not supported.")

1838 # vllm currently does not have pooling models using hybrid,

1839 # attention_free or encoder_decoder attn types.

1840 return False

vllm/model_executor/layers/pooler.py

+71 -22

@@ -127,14 +127,14 @@ class PoolingMethod(nn.Module, ABC):

127 self,

128 hidden_states: torch.Tensor,

129 pooling_cursor: PoolingCursor,

130 - ) -> list[torch.Tensor] | torch.Tensor:

130 + ) -> PoolerOutput:

131 raise NotImplementedError

132

133 def forward(

134 self,

135 hidden_states: torch.Tensor,

136 pooling_metadata: PoolingMetadata,

137 - ) -> list[torch.Tensor] | torch.Tensor:

137 + ) -> PoolerOutput:

138 pooling_cursor = pooling_metadata.pooling_cursor

139 return self.forward_all(hidden_states, pooling_cursor)

140

@@ -147,7 +147,7 @@ class CLSPool(PoolingMethod):

147 self,

148 hidden_states: torch.Tensor,

149 pooling_cursor: PoolingCursor,

150 - ) -> list[torch.Tensor] | torch.Tensor:

150 + ) -> PoolerOutput:

151 assert not pooling_cursor.is_partial_prefill(), (

152 "partial prefill not supported with CLS pooling"

153 )

@@ -163,27 +163,65 @@ class LastPool(PoolingMethod):

163 self,

164 hidden_states: torch.Tensor,

165 pooling_cursor: PoolingCursor,

166 - ) -> list[torch.Tensor] | torch.Tensor:

166 + ) -> PoolerOutput:

167 return hidden_states[pooling_cursor.last_token_indices_gpu]

168

169

170 class AllPool(PoolingMethod):

171 + def __init__(self):

172 + super().__init__()

173 +

174 + vllm_config = get_current_vllm_config()

175 + self.enable_chunked_prefill = (

176 + vllm_config.scheduler_config.enable_chunked_prefill

177 + )

178 +

179 def get_supported_tasks(self) -> Set[PoolingTask]:

180 return {"token_embed", "token_classify"}

181

182 def forward_all(

175 - self,

176 - hidden_states: torch.Tensor,

177 - pooling_cursor: PoolingCursor,

178 - ) -> list[torch.Tensor] | torch.Tensor:

179 - assert not pooling_cursor.is_partial_prefill(), (

180 - "partial prefill not supported with ALL pooling"

183 + self, hidden_states: torch.Tensor, pooling_cursor: PoolingCursor

184 + ) -> PoolerOutput:

185 + raise NotImplementedError(

186 + "forward_all is not implemented for AllPool. Use forward instead."

187 )

188

189 + def forward(

190 + self,

191 + hidden_states: torch.Tensor,

192 + pooling_metadata: PoolingMetadata,

193 + ) -> PoolerOutput:

194 + pooling_cursor = pooling_metadata.pooling_cursor

195 + is_finished = pooling_cursor.is_finished()

196 hidden_states_lst = list(

197 hidden_states.split(pooling_cursor.num_scheduled_tokens_cpu.tolist())

198 )

186 - return [hidden_states_lst[i] for i in pooling_cursor.index]

199 + hidden_states_lst = [hidden_states_lst[i] for i in pooling_cursor.index]

200 +

201 + if not self.enable_chunked_prefill:

202 + return hidden_states_lst

203 +

204 + pooling_states = pooling_metadata.pooling_states

205 +

206 + # If chunked_prefill is enabled

207 + # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache

208 + for p, hs_chunk in zip(pooling_states, hidden_states_lst):

209 + p.hidden_states_cache.append(hs_chunk)

210 +

211 + # 2. Once prefill is finished, send hidden_states_cache to PoolerHead

212 + output_list: PoolerOutput = []

213 + for p, finished in zip(pooling_states, is_finished):

214 + if finished:

215 + hidden_states_cache = p.hidden_states_cache

216 + if len(hidden_states_cache) == 1:

217 + output_list.append(hidden_states_cache[0])

218 + else:

219 + output_list.append(torch.concat(hidden_states_cache, dim=0))

220 + p.clean()

221 + else:

222 + output_list.append(None)

223 +

224 + return output_list

225

226

227 class MeanPool(PoolingMethod):

@@ -194,7 +232,7 @@ class MeanPool(PoolingMethod):

232 self,

233 hidden_states: torch.Tensor,

234 pooling_cursor: PoolingCursor,

197 - ) -> list[torch.Tensor] | torch.Tensor:

235 + ) -> PoolerOutput:

236 assert not pooling_cursor.is_partial_prefill(), (

237 "partial prefill not supported with MEAN pooling"

238 )

@@ -399,7 +437,7 @@ class PoolerHead(nn.Module):

437 self,

438 pooled_data: list[torch.Tensor] | torch.Tensor,

439 pooling_metadata: PoolingMetadata,

402 - ):

440 + ) -> PoolerOutput:

441 return self.activation(pooled_data)

442

443

@@ -418,7 +456,7 @@ class EmbeddingPoolerHead(PoolerHead):

456 self,

457 pooled_data: list[torch.Tensor] | torch.Tensor,

458 pooling_metadata: PoolingMetadata,

421 - ):

459 + ) -> PoolerOutput:

460 if isinstance(pooled_data, list):

461 pooled_data = torch.stack(pooled_data)

462 # pooled_data shape: [batchsize, hidden_dimension]

@@ -586,8 +624,12 @@ class ClassifierPooler(Pooler):

624

625 class TokenEmbeddingPoolerHead(EmbeddingPoolerHead):

626 def forward(

589 - self, pooled_data: torch.Tensor, pooling_param: PoolingParams

590 - ) -> torch.Tensor:

627 + self, pooled_data: torch.Tensor | None, pooling_param: PoolingParams

628 + ) -> PoolerOutput:

629 + # for unfinished chunked prefill

630 + if pooled_data is None:

631 + return None

632 +

633 pooled_data = pooled_data.to(self.head_dtype)

634 # pooled_data shape: [n_tokens, hidden_dimension]

635

@@ -630,9 +672,13 @@ class TokenClassifierPoolerHead(nn.Module):

672

673 def forward(

674 self,

633 - hidden_states: torch.Tensor,

675 + hidden_states: torch.Tensor | None,

676 pooling_param: PoolingParams,

635 - ) -> torch.Tensor:

677 + ) -> PoolerOutput:

678 + # for unfinished chunked prefill

679 + if hidden_states is None:

680 + return None

681 +

682 hidden_states = hidden_states.to(self.head_dtype)

683 # hidden_states shape: [n_token, hidden_size]

684

@@ -686,17 +732,20 @@ class StepPooler(Pooler):

732 self,

733 hidden_states: torch.Tensor | list[torch.Tensor],

734 pooling_metadata: PoolingMetadata,

689 - ) -> torch.Tensor | list[torch.Tensor]:

735 + ) -> PoolerOutput:

736 pooled_data_lst = self.pooling(hidden_states, pooling_metadata)

737 prompt_token_ids = pooling_metadata.get_prompt_token_ids()

692 -

693 - pooled_data = list[torch.Tensor]()

694 -

738 pooling_params = pooling_metadata.pooling_params

739

740 + pooled_data: PoolerOutput = []

741 for data, token_id, pooling_param in zip(

742 pooled_data_lst, prompt_token_ids, pooling_params

743 ):

744 + # for unfinished chunked prefill

745 + if data is None:

746 + pooled_data.append(data)

747 + continue

748 +

749 step_tag_id = pooling_param.step_tag_id

750 returned_token_ids = pooling_param.returned_token_ids

751

vllm/model_executor/models/terratorch.py

+2 -2

@@ -64,7 +64,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder

64 from vllm.sequence import IntermediateTensors

66 from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal

67 -from .interfaces_base import default_pooling_type

67 +from .interfaces_base import attn_type

69 logger = init_logger(__name__)

@@ -220,7 +220,7 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor):

220 )

221

222

223 -@default_pooling_type("All")

223 +@attn_type("attention_free")

224 @MULTIMODAL_REGISTRY.register_processor(

225 TerratorchMultiModalProcessor,

226 info=TerratorchProcessingInfo,

vllm/v1/outputs.py

+1 -1

@@ -89,7 +89,7 @@ class LogprobsTensors(NamedTuple):

90 # [num_reqs, <dynamic>]

91 # The shape of each element depends on the pooler used

92 -PoolerOutput = torch.Tensor | list[torch.Tensor]

92 +PoolerOutput = list[torch.Tensor | None] | torch.Tensor | None

95 @dataclass

vllm/v1/pool/metadata.py

+26 -3

@@ -17,6 +17,7 @@ class PoolingCursor:

17 first_token_indices_gpu: torch.Tensor

18 last_token_indices_gpu: torch.Tensor

19 prompt_lens_cpu: torch.Tensor

20 + seq_lens_cpu: torch.Tensor

21 num_scheduled_tokens_cpu: torch.Tensor

23 def __getitem__(self, indices: slice):

@@ -25,12 +26,25 @@ class PoolingCursor:

26 first_token_indices_gpu=self.first_token_indices_gpu[indices],

27 last_token_indices_gpu=self.last_token_indices_gpu[indices],

28 prompt_lens_cpu=self.prompt_lens_cpu[indices],

29 + seq_lens_cpu=self.seq_lens_cpu[indices],

30 num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],

31 )

33 def is_partial_prefill(self):

34 return not torch.all(self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)

36 + def is_finished(self):

37 + return self.prompt_lens_cpu == self.seq_lens_cpu

38 +

39 +

40 +class PoolingStates:

41 + def __init__(self):

42 + # for chunked prefill with ALL pooling

43 + self.hidden_states_cache: list[torch.Tensor] = []

44 +

45 + def clean(self):

46 + self.hidden_states_cache.clear()

47 +

49 @dataclass

50 class PoolingMetadata:

@@ -39,6 +53,7 @@ class PoolingMetadata:

53 prompt_lens: torch.Tensor # CPU Tensor

54 prompt_token_ids: torch.Tensor | None

55 pooling_params: list[PoolingParams]

56 + pooling_states: list[PoolingStates]

57 pooling_cursor: PoolingCursor | None = None

59 def __post_init__(self) -> None:

@@ -60,6 +75,7 @@ class PoolingMetadata:

75 if self.prompt_token_ids is None

76 else self.prompt_token_ids[indices],

77 pooling_params=self.pooling_params[indices],

78 + pooling_states=self.pooling_states[indices],

79 pooling_cursor=None

80 if self.pooling_cursor is None

81 else self.pooling_cursor[indices],

@@ -74,15 +90,21 @@ class PoolingMetadata:

90 return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]

92 def build_pooling_cursor(

77 - self, num_scheduled_tokens: list[int], device: torch.device

93 + self,

94 + num_scheduled_tokens: list[int],

95 + seq_lens_cpu: torch.Tensor,

96 + device: torch.device,

97 ):

98 self.pooling_cursor = build_pooling_cursor(

80 - num_scheduled_tokens, self.prompt_lens, device

99 + num_scheduled_tokens, seq_lens_cpu, self.prompt_lens, device

100 )

101

102

103 def build_pooling_cursor(

85 - num_scheduled_tokens: list[int], prompt_lens: torch.Tensor, device: torch.device

104 + num_scheduled_tokens: list[int],

105 + seq_lens_cpu: torch.Tensor,

106 + prompt_lens: torch.Tensor,

107 + device: torch.device,

108 ):

109 assert len(prompt_lens) == len(num_scheduled_tokens)

110

@@ -99,5 +121,6 @@ def build_pooling_cursor(

121 first_token_indices_gpu=cumsum[:n_seq],

122 last_token_indices_gpu=cumsum[1:] - 1,

123 prompt_lens_cpu=prompt_lens,

124 + seq_lens_cpu=seq_lens_cpu,

125 num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,

126 )

vllm/v1/worker/gpu_input_batch.py

+21 -2

@@ -15,7 +15,7 @@ from vllm.sampling_params import SamplingParams, SamplingType

15 from vllm.utils import length_from_prompt_token_ids_or_embeds

16 from vllm.utils.collection_utils import swap_dict_values

17 from vllm.v1.outputs import LogprobsTensors

18 -from vllm.v1.pool.metadata import PoolingMetadata

18 +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

19 from vllm.v1.sample.logits_processor import (

20 BatchUpdateBuilder,

21 LogitsProcessors,

@@ -33,7 +33,6 @@ class CachedRequestState:

33 prompt_token_ids: list[int] | None

34 mm_features: list[MultiModalFeatureSpec]

35 sampling_params: SamplingParams | None

36 - pooling_params: PoolingParams | None

36 generator: torch.Generator | None

38 block_ids: tuple[list[int], ...]

@@ -51,11 +50,18 @@ class CachedRequestState:

50 # Used when both async_scheduling and spec_decode are enabled.

51 prev_num_draft_len: int = 0

53 + # for pooling models

54 + pooling_params: PoolingParams | None = None

55 + pooling_states: PoolingStates | None = None

56 +

57 def __post_init__(self):

58 self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

59 self.prompt_token_ids, self.prompt_embeds

60 )

62 + if self.pooling_params is not None:

63 + self.pooling_states = PoolingStates()

64 +

65 @property

66 def num_tokens(self) -> int:

67 return self.num_prompt_tokens + len(self.output_token_ids)

@@ -255,7 +261,9 @@ class InputBatch:

261 # This is updated each time the batch constituents change.

262 self.sampling_metadata = self._make_sampling_metadata()

263

264 + # for pooling models

265 self.pooling_params: dict[str, PoolingParams] = {}

266 + self.pooling_states: dict[str, PoolingStates] = {}

267

268 # Cached reference to the GPU tensor of previously sampled tokens

269 self.prev_sampled_token_ids: torch.Tensor | None = None

@@ -413,7 +421,11 @@ class InputBatch:

421 sampling_params.bad_words_token_ids

422 )

423 elif pooling_params := request.pooling_params:

424 + pooling_states = request.pooling_states

425 + assert pooling_states is not None

426 +

427 self.pooling_params[req_id] = pooling_params

428 + self.pooling_states[req_id] = pooling_states

429 self.logits_processing_needs_token_ids[req_index] = (

430 pooling_params.requires_token_ids

431 )

@@ -469,6 +481,7 @@ class InputBatch:

481

482 if self.is_pooling_model:

483 self.pooling_params.pop(req_id, None)

484 + self.pooling_states.pop(req_id, None)

485 return req_index

486

487 self.greedy_reqs.discard(req_id)

@@ -837,13 +850,19 @@ class InputBatch:

850 assert len(self.req_ids) == len(self.pooling_params)

851 return [self.pooling_params[req_id] for req_id in self.req_ids]

852

853 + def get_pooling_states(self) -> list[PoolingStates]:

854 + assert len(self.req_ids) == len(self.pooling_states)

855 + return [self.pooling_states[req_id] for req_id in self.req_ids]

856 +

857 def get_pooling_metadata(self) -> PoolingMetadata:

858 pooling_params = self.get_pooling_params()

859 + pooling_states = self.get_pooling_states()

860

861 return PoolingMetadata(

862 prompt_lens=torch.from_numpy(self.num_prompt_tokens[: self.num_reqs]),

863 prompt_token_ids=self.sampling_metadata.prompt_token_ids,

864 pooling_params=pooling_params,

865 + pooling_states=pooling_states,

866 )

867

868 def _make_prompt_token_ids_tensor(self) -> torch.Tensor:

vllm/v1/worker/gpu_model_runner.py

+15 -35

@@ -131,7 +131,7 @@ from vllm.v1.outputs import (

131 SamplerOutput,

132 make_empty_encoder_model_runner_output,

133 )

134 -from vllm.v1.pool.metadata import PoolingMetadata

134 +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

135 from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs

136 from vllm.v1.sample.logits_processor.interface import LogitsProcessor

137 from vllm.v1.sample.metadata import SamplingMetadata

@@ -2291,20 +2291,6 @@ class GPUModelRunner(

2291

2292 supported_tasks = list(model.pooler.get_supported_tasks())

2293

2294 - if self.scheduler_config.enable_chunked_prefill:

2295 - if "token_embed" in supported_tasks:

2296 - supported_tasks.remove("token_embed")

2297 - if "token_classify" in supported_tasks:

2298 - supported_tasks.remove("token_classify")

2299 -

2300 - logger.debug_once(

2301 - "Chunked prefill is not supported with "

2302 - "token_embed and token_classify tasks "

2303 - "which using ALL pooling. "

2304 - "Please turn off chunked prefill by "

2305 - "`--no-enable-chunked-prefill` before using it."

2306 - )

2307 -

2294 if "score" in supported_tasks:

2295 num_labels = getattr(self.model_config.hf_config, "num_labels", 0)

2296 if num_labels != 1:

@@ -2381,11 +2367,12 @@ class GPUModelRunner(

2367 )

2368

2369 hidden_states = hidden_states[:num_scheduled_tokens]

2370 + seq_lens_cpu = self.seq_lens.cpu[: self.input_batch.num_reqs]

2371 +

2372 pooling_metadata = self.input_batch.get_pooling_metadata()

2373 pooling_metadata.build_pooling_cursor(

2386 - num_scheduled_tokens_np.tolist(), device=hidden_states.device

2374 + num_scheduled_tokens_np.tolist(), seq_lens_cpu, device=hidden_states.device

2375 )

2388 - seq_lens_cpu = self.seq_lens.cpu[: self.input_batch.num_reqs]

2376

2377 model = cast(VllmModelForPooling, self.model)

2378 raw_pooler_output: PoolerOutput = model.pooler(

@@ -2393,7 +2380,7 @@ class GPUModelRunner(

2380 pooling_metadata=pooling_metadata,

2381 )

2382 raw_pooler_output = json_map_leaves(

2396 - lambda x: x.to("cpu", non_blocking=True),

2383 + lambda x: x.to("cpu", non_blocking=True) if x is not None else x,

2384 raw_pooler_output,

2385 )

2386 self._sync_device()

@@ -4248,10 +4235,13 @@ class GPUModelRunner(

4235 prompt_lens=dummy_prompt_lens,

4236 prompt_token_ids=dummy_token_ids,

4237 pooling_params=[dummy_pooling_params] * num_reqs,

4238 + pooling_states=[PoolingStates() for i in range(num_reqs)],

4239 )

4240

4241 dummy_metadata.build_pooling_cursor(

4254 - num_scheduled_tokens_list, device=hidden_states.device

4242 + num_scheduled_tokens_list,

4243 + seq_lens_cpu=dummy_prompt_lens,

4244 + device=hidden_states.device,

4245 )

4246

4247 try:

@@ -4278,22 +4268,12 @@ class GPUModelRunner(

4268 supported_pooling_tasks = self.get_supported_pooling_tasks()

4269

4270 if not supported_pooling_tasks:

4281 - if self.scheduler_config.enable_chunked_prefill:

4282 - raise RuntimeError(

4283 - f"Model {self.model_config.model} does not support "

4284 - "any pooling tasks with chunked prefill enabled. "

4285 - "Please add --no-enable-chunked-prefill to your "

4286 - "config or CLI args. See "

4287 - "https://docs.vllm.ai/en/latest/models/pooling_models.html "

4288 - "to learn more."

4289 - )

4290 - else:

4291 - raise RuntimeError(

4292 - f"Model {self.model_config.model} does not support "

4293 - "any pooling tasks. See "

4294 - "https://docs.vllm.ai/en/latest/models/pooling_models.html "

4295 - "to learn more."

4296 - )

4271 + raise RuntimeError(

4272 + f"Model {self.model_config.model} does not support "

4273 + "any pooling tasks. See "

4274 + "https://docs.vllm.ai/en/latest/models/pooling_models.html "

4275 + "to learn more."

4276 + )

4277

4278 output_size = dict[PoolingTask, float]()

4279 for task in supported_pooling_tasks: