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

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 14, 05:30 UTC

[Performance] V1 Pooling Models E2E Performance Optimization (#23162)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Aug 21, 2025 at 21:26 UTC d70a16625dc74d9517641aa82f4ae7367854da96

8 files changed +161 -167

vllm/model_executor/layers/pooler.py

+48 -83

@@ -19,7 +19,8 @@ from vllm.model_executor.pooling_metadata import PoolingTensors

19 from vllm.pooling_params import PoolingParams

20 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput

21 from vllm.tasks import PoolingTask

22 -from vllm.utils import resolve_obj_by_qualname

22 +from vllm.utils import current_stream, resolve_obj_by_qualname

23 +from vllm.v1.pool.metadata import PoolingCursor

24 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata

26 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]

@@ -205,6 +206,13 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):

206

207 def build_output(

208 all_data: Union[torch.Tensor, list[torch.Tensor]], ) -> PoolerOutput:

209 + # Pooling models D2H & synchronize occurs here

210 + if isinstance(all_data, list):

211 + all_data = [d.to("cpu", non_blocking=True) for d in all_data]

212 + else:

213 + all_data = all_data.to("cpu", non_blocking=True)

214 + current_stream().synchronize()

215 +

216 all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]

217 return PoolerOutput(outputs=all_outputs)

218

@@ -231,40 +239,21 @@ class PoolingMethod(nn.Module, ABC):

239 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

240 return PoolingParamsUpdate()

241

234 - @abstractmethod

235 - def forward_one(

236 - self,

237 - hidden_states: torch.Tensor,

238 - prompt_len: Optional[torch.Tensor] = None,

239 - ) -> torch.Tensor:

240 - """

241 - Note:

242 - `prompt_len=None` means `prompt_len=len(hidden_states)`.

243 - """

244 - raise NotImplementedError

245 -

242 @abstractmethod

243 def forward_all(

244 self,

245 hidden_states: torch.Tensor,

250 - prompt_lens: torch.Tensor,

246 + pooling_cursor: PoolingCursor,

247 ) -> Union[list[torch.Tensor], torch.Tensor]:

248 raise NotImplementedError

249

250 def forward(

251 self,

256 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

252 + hidden_states: torch.Tensor,

253 pooling_metadata: PoolingMetadata,

254 ) -> Union[list[torch.Tensor], torch.Tensor]:

259 - prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)

260 -

261 - if isinstance(hidden_states, list):

262 - return [

263 - self.forward_one(h, prompt_len)

264 - for h, prompt_len in zip(hidden_states, prompt_lens)

265 - ]

266 -

267 - return self.forward_all(hidden_states, prompt_lens)

255 + pooling_cursor = pooling_metadata.pooling_cursor

256 + return self.forward_all(hidden_states, pooling_cursor)

257

258

259 class CLSPool(PoolingMethod):

@@ -272,24 +261,15 @@ class CLSPool(PoolingMethod):

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

262 return {"encode", "embed", "classify", "score"}

263

275 - def forward_one(

276 - self,

277 - hidden_states: torch.Tensor,

278 - prompt_len: Optional[torch.Tensor] = None,

279 - ) -> torch.Tensor:

280 - assert prompt_len is None or prompt_len == hidden_states.shape[0], \

281 - "partial prefill not supported with CLS pooling"

282 -

283 - return hidden_states[0]

284 -

264 def forward_all(

265 self,

266 hidden_states: torch.Tensor,

288 - prompt_lens: torch.Tensor,

267 + pooling_cursor: PoolingCursor,

268 ) -> Union[list[torch.Tensor], torch.Tensor]:

290 - first_token_flat_indices = torch.zeros_like(prompt_lens)

291 - first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]

292 - return hidden_states[first_token_flat_indices]

269 + assert not pooling_cursor.is_partial_prefill(), \

270 + "partial prefill not supported with CLS pooling"

271 +

272 + return hidden_states[pooling_cursor.first_token_indices_gpu]

273

274

275 class LastPool(PoolingMethod):

@@ -297,20 +277,12 @@ class LastPool(PoolingMethod):

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

278 return {"encode", "embed", "classify", "score"}

279

300 - def forward_one(

301 - self,

302 - hidden_states: torch.Tensor,

303 - prompt_len: Optional[torch.Tensor] = None,

304 - ) -> torch.Tensor:

305 - return hidden_states[-1]

306 -

280 def forward_all(

281 self,

282 hidden_states: torch.Tensor,

310 - prompt_lens: torch.Tensor,

283 + pooling_cursor: PoolingCursor,

284 ) -> Union[list[torch.Tensor], torch.Tensor]:

312 - last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1

313 - return hidden_states[last_token_flat_indices]

285 + return hidden_states[pooling_cursor.last_token_indices_gpu]

286

287

288 class AllPool(PoolingMethod):

@@ -318,22 +290,19 @@ class AllPool(PoolingMethod):

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

291 return {"encode"}

292

321 - def forward_one(

322 - self,

323 - hidden_states: torch.Tensor,

324 - prompt_len: Optional[torch.Tensor] = None,

325 - ) -> torch.Tensor:

326 - assert prompt_len is None or prompt_len == hidden_states.shape[0], \

327 - "partial prefill not supported with ALL pooling"

328 -

329 - return hidden_states

330 -

293 def forward_all(

294 self,

295 hidden_states: torch.Tensor,

334 - prompt_lens: torch.Tensor,

296 + pooling_cursor: PoolingCursor,

297 ) -> Union[list[torch.Tensor], torch.Tensor]:

336 - return list(hidden_states.split_with_sizes(prompt_lens.tolist()))

298 +

299 + assert not pooling_cursor.is_partial_prefill(), \

300 + "partial prefill not supported with ALL pooling"

301 +

302 + hidden_states_lst = list(

303 + hidden_states.split(

304 + pooling_cursor.num_scheduled_tokens_cpu.tolist()))

305 + return [hidden_states_lst[i] for i in pooling_cursor.index]

306

307

308 class MeanPool(PoolingMethod):

@@ -341,31 +310,25 @@ class MeanPool(PoolingMethod):

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

311 return {"encode", "embed", "classify", "score"}

312

344 - def forward_one(

313 + def forward_all(

314 self,

315 hidden_states: torch.Tensor,

347 - prompt_len: Optional[torch.Tensor] = None,

348 - ) -> torch.Tensor:

349 - assert prompt_len is None or prompt_len == hidden_states.shape[0], \

316 + pooling_cursor: PoolingCursor,

317 + ) -> Union[list[torch.Tensor], torch.Tensor]:

318 +

319 + assert not pooling_cursor.is_partial_prefill(), \

320 "partial prefill not supported with MEAN pooling"

321

352 - return hidden_states.mean(dim=0, dtype=torch.float32)

322 + prompt_lens = pooling_cursor.prompt_lens_cpu.to(hidden_states.device,

323 + non_blocking=True)

324

354 - def forward_all(

355 - self,

356 - hidden_states: torch.Tensor,

357 - prompt_lens: torch.Tensor,

358 - ) -> Union[list[torch.Tensor], torch.Tensor]:

325 # Use float32 for torch.cumsum in MeanPool,

326 # otherwise precision will be lost significantly.

327 cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)

328

363 - start_indices = torch.cat([

364 - torch.tensor([0], device=hidden_states.device),

365 - torch.cumsum(prompt_lens[:-1], dim=0)

366 - ])

367 - end_indices = torch.cumsum(prompt_lens, dim=0)

368 - return (cumsum[end_indices - 1] - cumsum[start_indices] +

329 + start_indices = pooling_cursor.first_token_indices_gpu

330 + end_indices = pooling_cursor.last_token_indices_gpu

331 + return (cumsum[end_indices] - cumsum[start_indices] +

332 hidden_states[start_indices]) / prompt_lens.unsqueeze(1)

333

334

@@ -477,6 +440,10 @@ class EmbeddingPoolerHead(PoolerHead):

440

441 pooling_params = get_pooling_params(pooling_metadata)

442

443 + if isinstance(pooled_data, list):

444 + pooled_data = torch.stack(pooled_data)

445 + # pooled_data shape: [batchsize, embedding_dimension]

446 +

447 # for matryoshka representation

448 dimensions_list = [

449 pooling_param.dimensions for pooling_param in pooling_params

@@ -667,6 +634,10 @@ class ClassifierPooler(Pooler):

634 ) -> PoolerOutput:

635 pooled_data = self.pooling(hidden_states, pooling_metadata)

636

637 + if isinstance(pooled_data, list):

638 + pooled_data = torch.stack(pooled_data)

639 + # pooled_data shape: [batchsize, hidden_size]

640 +

641 if self.classifier is not None:

642 # apply classifier once on the full batch if possible

643 if isinstance(pooled_data, torch.Tensor):

@@ -717,12 +688,6 @@ class DispatchPooler(Pooler):

688 ) -> PoolerOutput:

689 poolers_by_task = self.poolers_by_task

690

720 - if isinstance(hidden_states, list):

721 - hidden_states_lst = hidden_states

722 - else:

723 - prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)

724 - hidden_states_lst = list(hidden_states.split(prompt_lens.tolist()))

725 -

691 outputs = list[PoolingSequenceGroupOutput]()

692 offset = 0

693 for task, group in groupby(get_tasks(pooling_metadata)):

@@ -733,7 +698,7 @@ class DispatchPooler(Pooler):

698

699 num_items = len(list(group))

700 group_output: PoolerOutput = pooler(

736 - hidden_states_lst[offset:offset + num_items],

701 + hidden_states,

702 pooling_metadata[offset:offset + num_items],

703 )

704

vllm/model_executor/models/bert.py

+3 -3

@@ -528,9 +528,9 @@ def _encode_token_type_ids(input_ids: torch.Tensor,

528

529 def _decode_token_type_ids(input_ids: torch.Tensor) -> torch.Tensor:

530

531 - ids_mask = torch.ones(input_ids.shape,

532 - dtype=torch.int32,

533 - device=input_ids.device) << TOKEN_TYPE_SHIFT

531 + ids_mask = torch.ones_like(input_ids,

532 + dtype=torch.int32,

533 + device=input_ids.device) << TOKEN_TYPE_SHIFT

534 tokens_mask = ids_mask.bitwise_not()

535

536 token_type_ids = input_ids.bitwise_and(ids_mask) >> TOKEN_TYPE_SHIFT

vllm/model_executor/models/roberta.py

+10 -55

@@ -9,7 +9,6 @@ from torch import nn

9 from transformers import RobertaConfig

11 from vllm.config import VllmConfig

12 -from vllm.forward_context import get_forward_context

12 from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,

13 DispatchPooler, Pooler)

14 from vllm.model_executor.layers.vocab_parallel_embedding import (

@@ -100,7 +99,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

100 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

101 super().__init__(vllm_config=vllm_config, prefix=prefix)

103 - self.padding_idx = vllm_config.model_config.hf_config.pad_token_id

102 + self.padding_idx: int = vllm_config.model_config.hf_config.pad_token_id

103

104 def forward(

105 self,

@@ -178,7 +177,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

177 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

178 super().__init__()

179 config = vllm_config.model_config.hf_config

181 - self.padding_idx = vllm_config.model_config.hf_config.pad_token_id

180 + self.padding_idx: int = vllm_config.model_config.hf_config.pad_token_id

181

182 self.num_labels = config.num_labels

183 self.roberta = BertModel(vllm_config=vllm_config,

@@ -233,58 +232,14 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

232 intermediate_tensors=intermediate_tensors)

233

234

236 -# Adapted from transformers

237 -def create_position_ids_from_input_ids(input_ids,

238 - padding_idx,

239 - past_key_values_length=0):

240 - """

241 - Replace non-padding symbols with their position numbers.

242 - Position numbers begin at padding_idx+1. Padding symbols

243 - are ignored. This is modified from fairseq's `utils.make_positions`.

244 -

245 - Args:

246 - x: torch.Tensor x:

247 -

248 - Returns: torch.Tensor

249 - """

250 - # The series of casts and type-conversions here are carefully

251 - # balanced to both work with ONNX export and XLA.

252 - mask = input_ids.ne(padding_idx).int()

253 -

254 - incremental_indices = (torch.cumsum(mask, dim=0).type_as(mask) +

255 - past_key_values_length) * mask

256 -

257 - return incremental_indices.long() + padding_idx

258 -

259 -

235 def replace_roberta_positions(input_ids: torch.Tensor,

236 position_ids: torch.Tensor,

237 padding_idx: int) -> None:

263 -

264 - seq_lens: Optional[torch.Tensor] = None

265 - attn_metadata = get_forward_context().attn_metadata

266 - if attn_metadata is not None: # can be None during warmup

267 - if isinstance(attn_metadata, dict):

268 - attn_metadata = next(iter(attn_metadata.values()))

269 - # TODO: remove "seq_lens_tensor" after V0 is removed

270 - seq_lens = getattr(attn_metadata, "seq_lens_tensor",

271 - getattr(attn_metadata, "seq_lens", None))

272 -

273 - if seq_lens is not None:

274 - assert isinstance(seq_lens, torch.Tensor)

275 -

276 - # Replace position ids because in RoBERTa models

277 - # they have to start at padding_idx + 1 and ignore

278 - # existing padding tokens

279 - # References:

280 - # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133

281 - # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669

282 - token_list = torch.split(input_ids[:torch.sum(seq_lens)],

283 - seq_lens.tolist())

284 -

285 - offset = 0

286 - for tokens in token_list:

287 - length = tokens.shape[0]

288 - position_ids[offset:offset+length] = \

289 - create_position_ids_from_input_ids(tokens, padding_idx)

290 - offset = offset + length

238 + # Replace position ids because in RoBERTa models

239 + # they have to start at padding_idx + 1 and ignore

240 + # existing padding tokens

241 + # References:

242 + # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133

243 + # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669

244 + # vllm does not use padding tokens, let's make things simpler

245 + position_ids += padding_idx + 1

vllm/model_executor/pooling_metadata.py

+17 -6

@@ -2,12 +2,13 @@

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

4 from dataclasses import dataclass

5 -from typing import Any

5 +from typing import Any, Optional

7 import torch

9 from vllm.pooling_params import PoolingParams

10 from vllm.utils import is_pin_memory_available

11 +from vllm.v1.pool.metadata import PoolingCursor, build_pooling_cursor

14 class PoolingMetadata:

@@ -23,14 +24,15 @@ class PoolingMetadata:

24 """

26 def __init__(

26 - self,

27 - seq_groups: list[tuple[list[int], PoolingParams]],

28 - seq_data: dict[int, Any], # Specific data related to sequences

29 - prompt_lens: list[int],

30 - ) -> None:

27 + self,

28 + seq_groups: list[tuple[list[int], PoolingParams]],

29 + seq_data: dict[int, Any], # Specific data related to sequences

30 + prompt_lens: list[int],

31 + pooling_cursor: Optional[PoolingCursor] = None) -> None:

32 self.seq_groups = seq_groups

33 self.seq_data = seq_data

34 self.prompt_lens = prompt_lens

35 + self.pooling_cursor: Optional[PoolingCursor] = pooling_cursor

37 def __repr__(self) -> str:

38 return ("PoolingMetadata("

@@ -43,8 +45,17 @@ class PoolingMetadata:

45 seq_groups=self.seq_groups[indices],

46 seq_data=dict(list(self.seq_data.items())[indices]),

47 prompt_lens=self.prompt_lens[indices],

48 + pooling_cursor=None

49 + if self.pooling_cursor is None else self.pooling_cursor[indices],

50 )

52 + def build_pooling_cursor(self, num_scheduled_tokens: list[int],

53 + device: torch.device):

54 + prompt_lens = torch.tensor(self.prompt_lens, device="cpu")

55 + self.pooling_cursor = build_pooling_cursor(num_scheduled_tokens,

56 + prompt_lens,

57 + device=device)

58 +

60 @dataclass

61 class PoolingTensors:

vllm/v1/pool/metadata.py

+54 -2

@@ -6,15 +6,40 @@ from typing import Optional

6 import torch

8 from vllm.pooling_params import PoolingParams

9 +from vllm.utils import is_pin_memory_available

10 +

11 +pin_memory = is_pin_memory_available()

12 +

13 +

14 +@dataclass

15 +class PoolingCursor:

16 + index: list[int]

17 + first_token_indices_gpu: torch.Tensor

18 + last_token_indices_gpu: torch.Tensor

19 + prompt_lens_cpu: torch.Tensor

20 + num_scheduled_tokens_cpu: torch.Tensor

21 +

22 + def __getitem__(self, indices: slice):

23 + return PoolingCursor(

24 + index=self.index[indices],

25 + first_token_indices_gpu=self.first_token_indices_gpu[indices],

26 + last_token_indices_gpu=self.last_token_indices_gpu[indices],

27 + prompt_lens_cpu=self.prompt_lens_cpu[indices],

28 + num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],

29 + )

30 +

31 + def is_partial_prefill(self):

32 + return not torch.all(

33 + self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)

36 @dataclass

37 class PoolingMetadata:

38 """Tensors for pooling."""

14 -

15 - prompt_lens: torch.Tensor

39 + prompt_lens: torch.Tensor # CPU Tensor

40 prompt_token_ids: Optional[torch.Tensor]

41 pooling_params: list[PoolingParams]

42 + pooling_cursor: Optional[PoolingCursor] = None

44 def __getitem__(self, indices: slice):

45 return PoolingMetadata(

@@ -22,4 +47,31 @@ class PoolingMetadata:

47 prompt_token_ids=None if self.prompt_token_ids is None else

48 self.prompt_token_ids[indices],

49 pooling_params=self.pooling_params[indices],

50 + pooling_cursor=None

51 + if self.pooling_cursor is None else self.pooling_cursor[indices],

52 )

53 +

54 + def build_pooling_cursor(self, num_scheduled_tokens: list[int],

55 + device: torch.device):

56 + self.pooling_cursor = build_pooling_cursor(num_scheduled_tokens,

57 + self.prompt_lens, device)

58 +

59 +

60 +def build_pooling_cursor(num_scheduled_tokens: list[int],

61 + prompt_lens: torch.Tensor, device: torch.device):

62 + assert len(prompt_lens) == len(num_scheduled_tokens)

63 +

64 + n_seq = len(num_scheduled_tokens)

65 + index = list(range(n_seq))

66 + num_scheduled_tokens = torch.tensor(num_scheduled_tokens, device="cpu")

67 + cumsum = torch.zeros(n_seq + 1,

68 + dtype=torch.int64,

69 + pin_memory=pin_memory,

70 + device="cpu")

71 + torch.cumsum(num_scheduled_tokens, dim=0, out=cumsum[1:])

72 + cumsum = cumsum.to(device, non_blocking=True)

73 + return PoolingCursor(index=index,

74 + first_token_indices_gpu=cumsum[:n_seq],

75 + last_token_indices_gpu=cumsum[1:] - 1,

76 + prompt_lens_cpu=prompt_lens,

77 + num_scheduled_tokens_cpu=num_scheduled_tokens)

vllm/v1/worker/gpu_input_batch.py

+1 -1

@@ -713,7 +713,7 @@ class InputBatch:

713

714 return PoolingMetadata(

715 prompt_lens=torch.from_numpy(

716 - self.num_prompt_tokens[:self.num_reqs]).to(self.device),

716 + self.num_prompt_tokens[:self.num_reqs]),

717 prompt_token_ids=self.sampling_metadata.prompt_token_ids,

718 pooling_params=pooling_params,

719 )

vllm/v1/worker/gpu_model_runner.py

+20 -16

@@ -1476,23 +1476,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1476 "Either all or none of the requests in" \

1477 " a batch must be pooling request"

1478

1479 - extracted_hidden_states = list(

1480 - torch.split(hidden_states[:num_scheduled_tokens],

1481 - num_scheduled_tokens_np.tolist()))

1482 -

1479 + hidden_states = hidden_states[:num_scheduled_tokens]

1480 pooling_metadata = self.input_batch.pooling_metadata

1481 + pooling_metadata.build_pooling_cursor(num_scheduled_tokens_np.tolist(),

1482 + device=hidden_states.device)

1483 + seq_lens_cpu = self.seq_lens_cpu[:self.input_batch.num_reqs]

1484

1485 + # Pooling models D2H & synchronize occurs in pooler.py:build_output

1486 raw_pooler_output = self.model.pooler(

1486 - hidden_states=extracted_hidden_states,

1487 - pooling_metadata=pooling_metadata)

1487 + hidden_states=hidden_states, pooling_metadata=pooling_metadata)

1488

1489 pooler_output: list[Optional[torch.Tensor]] = []

1490 - seq_lens = self.seq_lens[:self.input_batch.num_reqs]

1490 for raw_output, seq_len, prompt_len in zip(

1492 - raw_pooler_output, seq_lens, pooling_metadata.prompt_lens):

1491 + raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens):

1492

1493 if seq_len == prompt_len:

1495 - pooler_output.append(raw_output.data.cpu())

1494 + pooler_output.append(raw_output.data)

1495 else:

1496 pooler_output.append(None)

1497

@@ -2524,13 +2523,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2523 assert sum(num_scheduled_tokens_list) == num_tokens

2524 assert len(num_scheduled_tokens_list) == num_reqs

2525

2527 - hidden_states_list = list(

2528 - torch.split(hidden_states, num_scheduled_tokens_list))

2526 req_num_tokens = num_tokens // num_reqs

2527

2528 dummy_prompt_lens = torch.tensor(

2532 - [h.shape[0] for h in hidden_states_list],

2533 - device=self.device,

2529 + num_scheduled_tokens_list,

2530 + device="cpu",

2531 )

2532 dummy_token_ids = torch.zeros((num_reqs, req_num_tokens),

2533 dtype=torch.int32,

@@ -2547,8 +2544,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2544 pooling_params=[dummy_pooling_params] * num_reqs,

2545 )

2546

2547 + dummy_metadata.build_pooling_cursor(num_scheduled_tokens_list,

2548 + device=hidden_states.device)

2549 +

2550 try:

2551 - return model.pooler(hidden_states=hidden_states_list,

2551 + return model.pooler(hidden_states=hidden_states,

2552 pooling_metadata=dummy_metadata)

2553 except RuntimeError as e:

2554 if 'out of memory' in str(e):

@@ -3316,10 +3316,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3316

3317 dummy_block_table = torch.zeros((num_reqs, 1),

3318 dtype=torch.int32,

3319 - device=self.device)

3319 + pin_memory=self.pin_memory,

3320 + device="cpu").to(self.device,

3321 + non_blocking=True)

3322 dummy_slot_mapping = torch.zeros((total_num_scheduled_tokens, ),

3323 dtype=torch.int32,

3322 - device=self.device)

3324 + pin_memory=self.pin_memory,

3325 + device="cpu").to(self.device,

3326 + non_blocking=True)

3327

3328 group_metadata = dict[str, tuple[CommonAttentionMetadata, Any]]()

3329

vllm/worker/pooling_model_runner.py

+8 -1

@@ -149,9 +149,16 @@ class PoolingModelRunner(

149 if not self.is_driver_worker:

150 return []

151

152 + pooling_metadata = model_input.pooling_metadata

153 + assert pooling_metadata is not None

154 +

155 + pooling_metadata.build_pooling_cursor(

156 + num_scheduled_tokens=pooling_metadata.prompt_lens,

157 + device=hidden_or_intermediate_states.device)

158 +

159 return [

160 self.model.pooler(hidden_states=hidden_or_intermediate_states,

154 - pooling_metadata=model_input.pooling_metadata)

161 + pooling_metadata=pooling_metadata)

162 ]

163

164 def make_model_input_from_broadcasted_tensor_dict(