Skip to content

vllm.model_executor.layers.pooler.tokwise.heads

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)
    ]

_mean_pool_chunks(data, chunk_size)

Mean-pool complete contextual states, accumulating in FP32.

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def _mean_pool_chunks(data: torch.Tensor, chunk_size: int) -> torch.Tensor:
    """Mean-pool complete contextual states, accumulating in FP32."""
    chunk_size = min(chunk_size, len(data))
    num_full, remainder = divmod(len(data), chunk_size)
    full_end = num_full * chunk_size
    full_chunks = data[:full_end].reshape(num_full, chunk_size, data.shape[-1])
    if not remainder:
        return full_chunks.mean(dim=1, dtype=torch.float32)

    return torch.cat(
        [
            full_chunks.mean(dim=1, dtype=torch.float32),
            data[full_end:].mean(dim=0, keepdim=True, dtype=torch.float32),
        ]
    )