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
[Misc] Simplify PoolerOutput and move to `v1/outputs` (#25629)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 25, 2025 at 14:47 UTC
755ed7b05be4743237d3339c4ff8c22bcaae04f4
6 files changed
+34
-82
vllm/executor/executor_base.py
+2
-2
@@ -15,10 +15,10 @@ from vllm.config import VllmConfig
15
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
16
from vllm.logger import init_logger
17
from vllm.lora.request import LoRARequest
18
-from vllm.sequence import ExecuteModelRequest, PoolerOutput
18
+from vllm.sequence import ExecuteModelRequest
19
from vllm.tasks import SupportedTask
20
from vllm.utils import make_async
21
-from vllm.v1.outputs import SamplerOutput
21
+from vllm.v1.outputs import PoolerOutput, SamplerOutput
22
from vllm.worker.worker_base import WorkerBase
23
24
logger = init_logger(__name__)
vllm/model_executor/layers/pooler.py
+8
-21
@@ -16,9 +16,9 @@ from vllm.config import ModelConfig, PoolerConfig, get_current_vllm_config
16
from vllm.logger import init_logger
17
from vllm.model_executor.models.adapters import _load_st_projector
18
from vllm.pooling_params import PoolingParams
19
-from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
19
from vllm.tasks import PoolingTask
21
-from vllm.utils import current_stream, resolve_obj_by_qualname
20
+from vllm.utils import resolve_obj_by_qualname
21
+from vllm.v1.outputs import PoolerOutput
22
from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata
23
24
logger = init_logger(__name__)
@@ -190,19 +190,6 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
190
return PoolerClassify()
191
192
193
-def build_output(
194
- all_data: Union[torch.Tensor, list[torch.Tensor]], ) -> PoolerOutput:
195
- # Pooling models D2H & synchronize occurs here
196
- if isinstance(all_data, list):
197
- all_data = [d.to("cpu", non_blocking=True) for d in all_data]
198
- else:
199
- all_data = all_data.to("cpu", non_blocking=True)
200
- current_stream().synchronize()
201
-
202
- all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]
203
- return PoolerOutput(outputs=all_outputs)
204
-
205
-
193
class PoolingMethod(nn.Module, ABC):
194
195
@staticmethod
@@ -556,7 +543,7 @@ class SimplePooler(Pooler):
543
) -> PoolerOutput:
544
pooled_data = self.pooling(hidden_states, pooling_metadata)
545
pooled_data = self.head(pooled_data, pooling_metadata)
559
- return build_output(pooled_data)
546
+ return pooled_data
547
548
549
class StepPooler(Pooler):
@@ -607,7 +594,7 @@ class StepPooler(Pooler):
594
) -> PoolerOutput:
595
pooled_data = self.extract_states(hidden_states, pooling_metadata)
596
pooled_data = self.head(pooled_data, pooling_metadata)
610
- return build_output(pooled_data)
597
+ return pooled_data
598
599
600
class ClassifierPooler(Pooler):
@@ -678,7 +665,7 @@ class ClassifierPooler(Pooler):
665
]
666
667
# scores shape: [batchsize, num_labels]
681
- return build_output(scores)
668
+ return scores
669
670
671
class DispatchPooler(Pooler):
@@ -708,7 +695,7 @@ class DispatchPooler(Pooler):
695
) -> PoolerOutput:
696
poolers_by_task = self.poolers_by_task
697
711
- outputs = list[PoolingSequenceGroupOutput]()
698
+ outputs = list[torch.Tensor]()
699
offset = 0
700
for task, group in groupby(get_tasks(pooling_metadata)):
701
if not (pooler := poolers_by_task.get(task)):
@@ -722,10 +709,10 @@ class DispatchPooler(Pooler):
709
pooling_metadata[offset:offset + num_items],
710
)
711
725
- outputs.extend(group_output.outputs)
712
+ outputs.extend(group_output)
713
offset += num_items
714
728
- return PoolerOutput(outputs)
715
+ return outputs
716
717
def extra_repr(self) -> str:
718
s = f"supported_task={self.get_supported_tasks()}"
vllm/model_executor/models/gritlm.py
+3
-3
@@ -12,12 +12,12 @@ from vllm.logger import init_logger
12
from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
13
PoolerHead, PoolerNormalize,
14
PoolingParamsUpdate,
15
- build_output, get_prompt_lens,
15
+ get_prompt_lens,
16
get_prompt_token_ids)
17
from vllm.model_executor.models.llama import LlamaForCausalLM
18
-from vllm.sequence import PoolerOutput
18
from vllm.tasks import PoolingTask
19
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
20
+from vllm.v1.outputs import PoolerOutput
21
from vllm.v1.pool.metadata import PoolingMetadata
22
23
from .interfaces_base import default_pooling_type
@@ -212,7 +212,7 @@ class GritLMPooler(Pooler):
212
) -> PoolerOutput:
213
pooled_data = self.pooling(hidden_states, pooling_metadata)
214
pooled_data = self.head(pooled_data, pooling_metadata)
215
- return build_output(pooled_data)
215
+ return pooled_data
216
217
218
@default_pooling_type("MEAN")
vllm/sequence.py
-48
@@ -11,7 +11,6 @@ if TYPE_CHECKING:
11
from vllm.v1.worker.kv_connector_model_runner_mixin import (
12
KVConnectorOutput)
13
else:
14
- LoRARequest = Any
14
KVConnectorOutput = Any
15
16
VLLM_TOKEN_ID_ARRAY_TYPE = "l"
@@ -48,29 +47,6 @@ class RequestMetrics:
47
model_execute_time: Optional[float] = None
48
49
51
-class PoolingSequenceGroupOutput(
52
- msgspec.Struct,
53
- omit_defaults=True, # type: ignore[call-arg]
54
- array_like=True, # type: ignore[call-arg]
55
-):
56
- """The model output associated with a pooling sequence group."""
57
- # Annotated as Any to be compatible with msgspec
58
- # The actual type is in SequenceGroup.pooled_data
59
- data: Any
60
-
61
- def get_data_nbytes(self) -> int:
62
- data: torch.Tensor = self.data
63
- return data.nbytes
64
-
65
- def __repr__(self) -> str:
66
- return f"PoolingSequenceGroupOutput(data={self.data}"
67
-
68
- def __eq__(self, other: object) -> bool:
69
- if not isinstance(other, PoolingSequenceGroupOutput):
70
- raise NotImplementedError()
71
- return self.data == other.data
72
-
73
-
50
# cannot use msgspec.Struct here because Dynamo does not support it
51
@dataclass
52
class IntermediateTensors:
@@ -119,30 +95,6 @@ class IntermediateTensors:
95
return f"IntermediateTensors(tensors={self.tensors})"
96
97
122
-class PoolerOutput(
123
- msgspec.Struct,
124
- omit_defaults=True, # type: ignore[call-arg]
125
- array_like=True): # type: ignore[call-arg]
126
- """The output from a pooling operation in the pooling model."""
127
- outputs: list[PoolingSequenceGroupOutput]
128
-
129
- def get_data_nbytes(self) -> int:
130
- return sum(o.get_data_nbytes() for o in self.outputs)
131
-
132
- def __getitem__(self, idx: int) -> PoolingSequenceGroupOutput:
133
- return self.outputs[idx]
134
-
135
- def __setitem__(self, idx: int, value: PoolingSequenceGroupOutput):
136
- self.outputs[idx] = value
137
-
138
- def __len__(self):
139
- return len(self.outputs)
140
-
141
- def __eq__(self, other: object):
142
- return isinstance(other,
143
- self.__class__) and self.outputs == other.outputs
144
-
145
-
98
class ExecuteModelRequest(
99
msgspec.Struct,
100
array_like=True, # type: ignore[call-arg]
vllm/v1/outputs.py
+6
-1
@@ -3,7 +3,7 @@
3
4
from abc import ABC, abstractmethod
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, NamedTuple, Optional
6
+from typing import TYPE_CHECKING, NamedTuple, Optional, Union
7
8
import torch
9
@@ -65,6 +65,11 @@ class LogprobsTensors(NamedTuple):
65
)
66
67
68
+# [num_reqs, <dynamic>]
69
+# The shape of each element depends on the pooler used
70
+PoolerOutput = Union[torch.Tensor, list[torch.Tensor]]
71
+
72
+
73
@dataclass
74
class SamplerOutput:
75
vllm/v1/worker/gpu_model_runner.py
+15
-7
@@ -52,13 +52,14 @@ from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargsItem,
52
from vllm.multimodal.utils import group_mm_kwargs_by_modality
53
from vllm.pooling_params import PoolingParams
54
from vllm.sampling_params import SamplingType
55
-from vllm.sequence import IntermediateTensors, PoolerOutput
55
+from vllm.sequence import IntermediateTensors
56
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
57
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
58
GiB_bytes, cdiv, check_use_alibi, get_dtype_size,
59
is_pin_memory_available,
60
length_from_prompt_token_ids_or_embeds, round_up,
61
supports_dynamo)
62
+from vllm.utils.jsontree import json_map_leaves
63
from vllm.v1.attention.backends.flash_attn import AttentionMetadata
64
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
65
from vllm.v1.attention.backends.utils import (
@@ -79,7 +80,7 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,
80
# yapf: enable
81
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
82
DraftTokenIds, LogprobsLists, LogprobsTensors,
82
- ModelRunnerOutput, SamplerOutput)
83
+ ModelRunnerOutput, PoolerOutput, SamplerOutput)
84
from vllm.v1.pool.metadata import PoolingMetadata
85
from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs
86
from vllm.v1.sample.metadata import SamplingMetadata
@@ -1823,15 +1824,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1824
device=hidden_states.device)
1825
seq_lens_cpu = self.seq_lens.cpu[:self.input_batch.num_reqs]
1826
1826
- # Pooling models D2H & synchronize occurs in pooler.py:build_output
1827
- raw_pooler_output = self.model.pooler(
1828
- hidden_states=hidden_states, pooling_metadata=pooling_metadata)
1827
+ model = cast(VllmModelForPooling, self.model)
1828
+ raw_pooler_output: PoolerOutput = model.pooler(
1829
+ hidden_states=hidden_states,
1830
+ pooling_metadata=pooling_metadata,
1831
+ )
1832
+ raw_pooler_output = json_map_leaves(
1833
+ lambda x: x.to("cpu", non_blocking=True),
1834
+ raw_pooler_output,
1835
+ )
1836
+ self._sync_device()
1837
1838
pooler_output: list[Optional[torch.Tensor]] = []
1839
for raw_output, seq_len, prompt_len in zip(
1840
raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens):
1841
1834
- output = raw_output.data if seq_len == prompt_len else None
1842
+ output = raw_output if seq_len == prompt_len else None
1843
pooler_output.append(output)
1844
1845
return ModelRunnerOutput(
@@ -3233,7 +3241,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3241
for task in self.get_supported_pooling_tasks():
3242
# Run a full batch with each task to ensure none of them OOMs
3243
output = self._dummy_pooler_run_task(hidden_states, task)
3236
- output_size[task] = output.get_data_nbytes()
3244
+ output_size[task] = sum(o.nbytes for o in output)
3245
del output # Allow GC
3246
3247
max_task = max(output_size.items(), key=lambda x: x[1])[0]