Skip to content

vllm.model_executor.layers.pooler.tokwise

Poolers that produce an output for each token in the sequence.

Modules:

Classes:

TokenEmbeddingPoolerHead

Bases: TokenPoolerHead

Methods:

  • forward_batched –

    Batch projection, splitting only when finishing needs per-request options.

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
class TokenEmbeddingPoolerHead(TokenPoolerHead):
    def __init__(
        self,
        head_dtype: torch.dtype | str | None = None,
        projector: ProjectorFn | None = None,
        activation: ActivationFn | None = None,
    ) -> None:
        super().__init__()

        self.head_dtype = head_dtype
        self.projector = projector
        self.activation = activation

    def extra_repr(self) -> str:
        attrs = []
        if self.head_dtype is not None:
            attrs.append(f"head_dtype={self.head_dtype}")
        if self.projector is not None:
            attrs.append("projector=True")
        if self.activation is not None:
            attrs.append(f"activation={self.activation.__class__.__name__}")
        return ", ".join(attrs)

    def get_supported_tasks(self) -> Set[PoolingTask]:
        return {"token_embed"}

    def forward_batched(
        self,
        hidden_states: torch.Tensor,
        pooling_params: list[PoolingParams],
        split: Callable[[torch.Tensor], list[TokenPoolingMethodOutputItem]],
    ) -> list[TokenPoolerHeadOutputItem] | None:
        """Batch projection, splitting only when finishing needs per-request options.

        Keep eligibility checks in sync with PoolingParams fields read by forward_chunk.
        """
        if envs.VLLM_BATCH_INVARIANT or any(
            p.late_chunking_params is not None for p in pooling_params
        ):
            return None
        # [total_tokens, hidden_size] -> [total_tokens, embedding_size]
        embeddings = self._project(hidden_states)
        first = pooling_params[0]
        if all(
            p.dimensions == first.dimensions
            and p.use_activation == first.use_activation
            for p in pooling_params
        ):
            embeddings = self._finish(embeddings, first)
            # Each output has shape [request_tokens, output_size].
            return split(embeddings)
        # Each part has shape [request_tokens, embedding_size].
        pooled_data = split(embeddings)
        del embeddings
        return [
            self._finish(part, param) if part is not None else None
            for part, param in zip(pooled_data, pooling_params, strict=True)
        ]

    def _project(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Project [num_rows, hidden_size] to [num_rows, embedding_size]."""
        if self.head_dtype is not None:
            hidden_states = hidden_states.to(self.head_dtype)
        if self.projector is not None:
            return self.projector(hidden_states)
        return hidden_states

    def _finish(
        self, embeddings: torch.Tensor, pooling_param: PoolingParams
    ) -> torch.Tensor:
        """Slice/normalize [num_rows, embedding_size] to [num_rows, output_size]."""
        if pooling_param.dimensions is not None:
            embeddings = embeddings[..., : pooling_param.dimensions]
        if self.activation is not None and pooling_param.use_activation:
            embeddings = self.activation(embeddings)
        return embeddings

    def forward_chunk(
        self,
        pooled_data: TokenPoolingMethodOutputItem,
        pooling_param: PoolingParams,
    ) -> TokenPoolerHeadOutputItem:
        # for unfinished chunked prefill
        if pooled_data is None:
            return None

        # pooled_data: [num_tokens, hidden_size]
        if (late_chunking := pooling_param.late_chunking_params) is not None:
            means = _mean_pool_chunks(pooled_data, late_chunking.chunk_size)
            pooled_data = (
                means.to(pooled_data.dtype) if self.head_dtype is None else means
            )
            # pooled_data: [num_chunks, hidden_size]

        # Output: [num_rows, output_size], one row per token or late chunk.
        return self._finish(self._project(pooled_data), pooling_param)

_finish(embeddings, pooling_param)

Slice/normalize [num_rows, embedding_size] to [num_rows, output_size].

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
133
134
135
136
137
138
139
140
141
def _finish(
    self, embeddings: torch.Tensor, pooling_param: PoolingParams
) -> torch.Tensor:
    """Slice/normalize [num_rows, embedding_size] to [num_rows, output_size]."""
    if pooling_param.dimensions is not None:
        embeddings = embeddings[..., : pooling_param.dimensions]
    if self.activation is not None and pooling_param.use_activation:
        embeddings = self.activation(embeddings)
    return embeddings

_project(hidden_states)

Project [num_rows, hidden_size] to [num_rows, embedding_size].

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
125
126
127
128
129
130
131
def _project(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Project [num_rows, hidden_size] to [num_rows, embedding_size]."""
    if self.head_dtype is not None:
        hidden_states = hidden_states.to(self.head_dtype)
    if self.projector is not None:
        return self.projector(hidden_states)
    return hidden_states

forward_batched(hidden_states, pooling_params, split)

Batch projection, splitting only when finishing needs per-request options.

Keep eligibility checks in sync with PoolingParams fields read by forward_chunk.

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def forward_batched(
    self,
    hidden_states: torch.Tensor,
    pooling_params: list[PoolingParams],
    split: Callable[[torch.Tensor], list[TokenPoolingMethodOutputItem]],
) -> list[TokenPoolerHeadOutputItem] | None:
    """Batch projection, splitting only when finishing needs per-request options.

    Keep eligibility checks in sync with PoolingParams fields read by forward_chunk.
    """
    if envs.VLLM_BATCH_INVARIANT or any(
        p.late_chunking_params is not None for p in pooling_params
    ):
        return None
    # [total_tokens, hidden_size] -> [total_tokens, embedding_size]
    embeddings = self._project(hidden_states)
    first = pooling_params[0]
    if all(
        p.dimensions == first.dimensions
        and p.use_activation == first.use_activation
        for p in pooling_params
    ):
        embeddings = self._finish(embeddings, first)
        # Each output has shape [request_tokens, output_size].
        return split(embeddings)
    # Each part has shape [request_tokens, embedding_size].
    pooled_data = split(embeddings)
    del embeddings
    return [
        self._finish(part, param) if part is not None else None
        for part, param in zip(pooled_data, pooling_params, strict=True)
    ]

TokenPooler

Bases: Pooler

A layer that pools specific information from hidden states.

This layer does the following: 1. Extracts specific tokens or aggregates data based on pooling method. 2. Postprocesses the output based on pooling head. 3. Returns structured results as PoolerOutput.

Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class TokenPooler(Pooler):
    """A layer that pools specific information from hidden states.

    This layer does the following:
    1. Extracts specific tokens or aggregates data based on pooling method.
    2. Postprocesses the output based on pooling head.
    3. Returns structured results as `PoolerOutput`.
    """

    def __init__(
        self,
        pooling: TokenPoolingMethod | TokenPoolingFn,
        head: TokenPoolerHead | TokenPoolingHeadFn | None = None,
    ) -> None:
        super().__init__()

        self.pooling = pooling
        self.head = head

    def extra_repr(self) -> str:
        head_name = self.head.__class__.__name__ if self.head is not None else None
        return f"pooling={self.pooling.__class__.__name__}, head={head_name}"

    def get_supported_tasks(self) -> Set[PoolingTask]:
        tasks = set(POOLING_TASKS)

        if isinstance(self.pooling, TokenPoolingMethod):
            tasks &= self.pooling.get_supported_tasks()
        if isinstance(self.head, TokenPoolerHead):
            tasks &= self.head.get_supported_tasks()

        return tasks

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

        if isinstance(self.pooling, TokenPoolingMethod):
            updates |= self.pooling.get_pooling_updates(task)

        return updates

    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> TokenPoolerOutput:
        params = pooling_metadata.pooling_params
        if (
            params
            and type(self.pooling) is AllPool
            and not self.pooling.enable_chunked_prefill
            and type(self.head) is TokenEmbeddingPoolerHead
            and (
                outputs := self.head.forward_batched(
                    hidden_states,
                    params,
                    lambda embeddings: self.pooling(embeddings, pooling_metadata),
                )
            )
            is not None
        ):
            return outputs
        pooled_data = self.pooling(hidden_states, pooling_metadata)
        if self.head is not None:
            pooled_data = self.head(pooled_data, pooling_metadata)
        return pooled_data