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, 02:15 UTC
[Misc] Move functions into `PoolingMetadata` (#30027)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Dec 4, 2025 at 16:21 UTC
68eb5c8d970a453a440776211f8dbff215fb40c3
3 files changed
+30
-47
vllm/model_executor/layers/pooler.py
+7
-43
@@ -64,42 +64,6 @@ class PoolingParamsUpdate:
64
params.requires_token_ids = self.requires_token_ids
65
66
67
-def get_prompt_lens(
68
- hidden_states: torch.Tensor | list[torch.Tensor],
69
- pooling_metadata: PoolingMetadata,
70
-) -> torch.Tensor:
71
- return pooling_metadata.prompt_lens
72
-
73
-
74
-def get_prompt_token_ids(pooling_metadata: PoolingMetadata) -> list[torch.Tensor]:
75
- assert pooling_metadata.prompt_token_ids is not None, (
76
- "Please set `requires_token_ids=True` in `get_pooling_updates`"
77
- )
78
-
79
- return [
80
- pooling_metadata.prompt_token_ids[i, :num]
81
- for i, num in enumerate(pooling_metadata.prompt_lens)
82
- ]
83
-
84
-
85
-def get_pooling_params(pooling_metadata: PoolingMetadata) -> list[PoolingParams]:
86
- pooling_params = pooling_metadata.pooling_params
87
- return pooling_params
88
-
89
-
90
-def get_tasks(pooling_metadata: PoolingMetadata) -> list[PoolingTask]:
91
- pooling_params = get_pooling_params(pooling_metadata)
92
-
93
- tasks: list[PoolingTask] = [
94
- task
95
- for pooling_param in pooling_params
96
- if (task := pooling_param.task) is not None
97
- ]
98
- assert len(pooling_params) == len(tasks)
99
-
100
- return tasks
101
-
102
-
67
def get_classification_activation_function(config: PretrainedConfig):
68
# Implement alignment with transformers ForSequenceClassificationLoss
69
# https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92
@@ -466,7 +430,7 @@ class EmbeddingPoolerHead(PoolerHead):
430
pooled_data = self.projector(pooled_data)
431
# pooled_data shape: [batchsize, embedding_dimension]
432
469
- pooling_params = get_pooling_params(pooling_metadata)
433
+ pooling_params = pooling_metadata.pooling_params
434
435
# for matryoshka representation
436
dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params]
@@ -606,7 +570,7 @@ class ClassifierPooler(Pooler):
570
if self.logit_bias is not None:
571
pooled_data -= self.logit_bias
572
609
- pooling_params = get_pooling_params(pooling_metadata)
573
+ pooling_params = pooling_metadata.pooling_params
574
flags = [p.use_activation for p in pooling_params]
575
576
if len(set(flags)) == 1:
@@ -704,7 +668,7 @@ class AllPooler(Pooler):
668
pooling_metadata: PoolingMetadata,
669
) -> PoolerOutput:
670
pooled_data = self.pooling(hidden_states, pooling_metadata)
707
- pooling_params = get_pooling_params(pooling_metadata)
671
+ pooling_params = pooling_metadata.pooling_params
672
assert len(pooled_data) == len(pooling_params)
673
674
pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
@@ -724,11 +688,11 @@ class StepPooler(Pooler):
688
pooling_metadata: PoolingMetadata,
689
) -> torch.Tensor | list[torch.Tensor]:
690
pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
727
- prompt_token_ids = get_prompt_token_ids(pooling_metadata)
691
+ prompt_token_ids = pooling_metadata.get_prompt_token_ids()
692
693
pooled_data = list[torch.Tensor]()
694
731
- pooling_params = get_pooling_params(pooling_metadata)
695
+ pooling_params = pooling_metadata.pooling_params
696
697
for data, token_id, pooling_param in zip(
698
pooled_data_lst, prompt_token_ids, pooling_params
@@ -757,7 +721,7 @@ class StepPooler(Pooler):
721
pooling_metadata: PoolingMetadata,
722
) -> PoolerOutput:
723
pooled_data = self.extract_states(hidden_states, pooling_metadata)
760
- pooling_params = get_pooling_params(pooling_metadata)
724
+ pooling_params = pooling_metadata.pooling_params
725
assert len(pooled_data) == len(pooling_params)
726
727
pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
@@ -794,7 +758,7 @@ class DispatchPooler(Pooler):
758
759
outputs = list[torch.Tensor]()
760
offset = 0
797
- for task, group in groupby(get_tasks(pooling_metadata)):
761
+ for task, group in groupby(pooling_metadata.tasks):
762
if not (pooler := poolers_by_task.get(task)):
763
raise ValueError(
764
f"Unsupported task: {task} "
vllm/model_executor/models/gritlm.py
+2
-4
@@ -14,8 +14,6 @@ from vllm.model_executor.layers.pooler import (
14
PoolerHead,
15
PoolerNormalize,
16
PoolingParamsUpdate,
17
- get_prompt_lens,
18
- get_prompt_token_ids,
17
)
18
from vllm.model_executor.models.llama import LlamaForCausalLM
19
from vllm.tasks import PoolingTask
@@ -153,11 +151,11 @@ class GritLMMeanPool(nn.Module):
151
hidden_states: torch.Tensor | list[torch.Tensor],
152
pooling_metadata: PoolingMetadata,
153
) -> list[torch.Tensor] | torch.Tensor:
156
- prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)
154
+ prompt_lens = pooling_metadata.prompt_lens
155
instr_lens = torch.tensor(
156
[
157
self._get_instruction_len(token_ids.cpu().numpy())
160
- for token_ids in get_prompt_token_ids(pooling_metadata)
158
+ for token_ids in pooling_metadata.get_prompt_token_ids()
159
],
160
device="cpu",
161
)
vllm/v1/pool/metadata.py
+21
@@ -5,6 +5,7 @@ from dataclasses import dataclass
5
import torch
6
7
from vllm.pooling_params import PoolingParams
8
+from vllm.tasks import PoolingTask
9
from vllm.utils.platform_utils import is_pin_memory_available
10
11
pin_memory = is_pin_memory_available()
@@ -40,6 +41,18 @@ class PoolingMetadata:
41
pooling_params: list[PoolingParams]
42
pooling_cursor: PoolingCursor | None = None
43
44
+ def __post_init__(self) -> None:
45
+ pooling_params = self.pooling_params
46
+
47
+ tasks: list[PoolingTask] = [
48
+ task
49
+ for pooling_param in pooling_params
50
+ if (task := pooling_param.task) is not None
51
+ ]
52
+ assert len(pooling_params) == len(tasks)
53
+
54
+ self.tasks = tasks
55
+
56
def __getitem__(self, indices: slice):
57
return PoolingMetadata(
58
prompt_lens=self.prompt_lens[indices],
@@ -52,6 +65,14 @@ class PoolingMetadata:
65
else self.pooling_cursor[indices],
66
)
67
68
+ def get_prompt_token_ids(self) -> list[torch.Tensor]:
69
+ prompt_token_ids = self.prompt_token_ids
70
+ assert prompt_token_ids is not None, (
71
+ "Please set `requires_token_ids=True` in `get_pooling_updates`"
72
+ )
73
+
74
+ return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
75
+
76
def build_pooling_cursor(
77
self, num_scheduled_tokens: list[int], device: torch.device
78
):