[Perf] Batched projector for pooling model embed, 1.8% throughput improvement

tests/models/language/pooling/test_tokwise_pooler_batching.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import vllm.model_executor.layers.pooler.tokwise.methods as tokwise_methods
+from vllm.model_executor.layers.pooler.tokwise.heads import (
+    TokenEmbeddingPoolerHead,
+)
+from vllm.model_executor.layers.pooler.tokwise.methods import AllPool
+from vllm.model_executor.layers.pooler.tokwise.poolers import TokenPooler
+from vllm.pooling_params import PoolingParams
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+
+
+class CountingLinear(nn.Module):
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=False)
+        self.call_count = 0
+        self.input_shapes: list[tuple[int, ...]] = []
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.call_count += 1
+        self.input_shapes.append(tuple(x.shape))
+        return self.linear(x)
+
+
+def _patch_chunked_prefill(monkeypatch, enabled: bool) -> None:
+    monkeypatch.setattr(
+        tokwise_methods,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            scheduler_config=SimpleNamespace(enable_chunked_prefill=enabled)
+        ),
+    )
+
+
+def _build_pooling_metadata(
+    *,
+    prompt_lens: list[int],
+    pooling_params: list[PoolingParams],
+) -> PoolingMetadata:
+    prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.int64)
+    seq_lens_cpu = torch.tensor(prompt_lens, dtype=torch.int64)
+    query_start_loc_cpu = torch.tensor(
+        [0, *np.cumsum(prompt_lens, dtype=np.int64)],
+        dtype=torch.int64,
+    )
+
+    metadata = PoolingMetadata(
+        prompt_lens=prompt_lens_cpu,
+        prompt_token_ids=None,
+        prompt_token_ids_cpu=None,
+        pooling_params=pooling_params,
+        pooling_states=[PoolingStates() for _ in pooling_params],
+    )
+    metadata.build_pooling_cursor(
+        num_scheduled_tokens_np=np.asarray(prompt_lens, dtype=np.int64),
+        seq_lens_cpu=seq_lens_cpu,
+        device=torch.device("cpu"),
+        query_start_loc_gpu=query_start_loc_cpu,
+    )
+    return metadata
+
+
+@torch.inference_mode()
+def test_token_embed_pooler_projects_flat_batch_once(monkeypatch):
+    _patch_chunked_prefill(monkeypatch, enabled=False)
+
+    hidden_size = 4
+    lengths = [2, 3, 1]
+    hidden_states = torch.randn(sum(lengths), hidden_size)
+    projector = CountingLinear(hidden_size, 5)
+
+    pooling_params = [
+        PoolingParams(task="token_embed", dimensions=5, use_activation=False),
+        PoolingParams(task="token_embed", dimensions=3, use_activation=True),
+        PoolingParams(task="token_embed", dimensions=4, use_activation=False),
+    ]
+    pooling_metadata = _build_pooling_metadata(
+        prompt_lens=lengths,
+        pooling_params=pooling_params,
+    )
+    pooler = TokenPooler(
+        pooling=AllPool(),
+        head=TokenEmbeddingPoolerHead(
+            projector=projector,
+            activation=torch.tanh,
+        ),
+    )
+
+    outputs = pooler(hidden_states, pooling_metadata)
+
+    assert projector.call_count == 1
+    assert projector.input_shapes == [(sum(lengths), hidden_size)]
+
+    expected_outputs = []
+    offset = 0
+    for length, pooling_param in zip(lengths, pooling_params):
+        chunk = hidden_states[offset : offset + length]
+        embeddings = projector.linear(chunk)
+        embeddings = embeddings[..., : pooling_param.dimensions]
+        if pooling_param.use_activation:
+            embeddings = torch.tanh(embeddings)
+        expected_outputs.append(embeddings)
+        offset += length
+
+    assert len(outputs) == len(expected_outputs)
+    for output, expected in zip(outputs, expected_outputs):
+        assert output is not None
+        torch.testing.assert_close(output, expected)
+
+
+@torch.inference_mode()
+def test_token_embed_pooler_projects_uniform_postprocess_once(monkeypatch):
+    _patch_chunked_prefill(monkeypatch, enabled=False)
+
+    hidden_size = 4
+    lengths = [2, 2]
+    hidden_states = torch.randn(sum(lengths), hidden_size)
+    projector = CountingLinear(hidden_size, 6)
+
+    pooling_params = [
+        PoolingParams(task="token_embed", dimensions=4, use_activation=True),
+        PoolingParams(task="token_embed", dimensions=4, use_activation=True),
+    ]
+    pooling_metadata = _build_pooling_metadata(
+        prompt_lens=lengths,
+        pooling_params=pooling_params,
+    )
+    pooler = TokenPooler(
+        pooling=AllPool(),
+        head=TokenEmbeddingPoolerHead(
+            projector=projector,
+            activation=torch.tanh,
+        ),
+    )
+
+    outputs = pooler(hidden_states, pooling_metadata)
+
+    assert projector.call_count == 1
+    assert projector.input_shapes == [(sum(lengths), hidden_size)]
+
+    projected = torch.tanh(projector.linear(hidden_states)[..., :4])
+    expected_outputs = [projected[:2], projected[2:]]
+    assert len(outputs) == len(expected_outputs)
+    for output, expected in zip(outputs, expected_outputs):
+        assert output is not None
+        torch.testing.assert_close(output, expected)

vllm/model_executor/layers/pooler/tokwise/heads.py

@@ -12,7 +12,7 @@
 from vllm.tasks import PoolingTask
 from vllm.v1.pool.metadata import PoolingMetadata
 
-from .methods import TokenPoolingMethodOutputItem
+from .methods import RaggedTokenBatch, TokenPoolingMethodOutputItem
 
 TokenPoolerHeadOutputItem: TypeAlias = torch.Tensor | None
 
@@ -66,17 +66,48 @@ def forward_chunk(
         if pooled_data is None:
             return None
 
+        embeddings = self._project_batch(pooled_data)
+        return self._postprocess_embeddings(embeddings, pooling_param)
+
+    def forward_ragged(
+        self,
+        pooled_data: RaggedTokenBatch,
+        pooling_params: list[PoolingParams],
+    ) -> list[TokenPoolerHeadOutputItem]:
+        if pooled_data.num_items != len(pooling_params):
+            raise ValueError(
+                "pooled_data and pooling_params must have the same length: "
+                f"{pooled_data.num_items} != {len(pooling_params)}."
+            )
+
+        # doing projection for all tokens in the batch
+        embeddings = self._project_batch(pooled_data.values)
+        if self._has_uniform_postprocess(pooling_params):
+            embeddings = self._postprocess_embeddings(embeddings, pooling_params[0])
+            return pooled_data.with_values(embeddings).split()
+
+        # can't apply the same postprocess, doing it separately
+        pooled_outputs = pooled_data.with_values(embeddings).split()
+        return [
+            self._postprocess_embeddings(output, pooling_param)
+            for output, pooling_param in zip(pooled_outputs, pooling_params)
+        ]
+
+    def _project_batch(self, pooled_data: torch.Tensor) -> torch.Tensor:
         if self.head_dtype is not None:
             pooled_data = pooled_data.to(self.head_dtype)
         # pooled_data shape: [n_tokens, hidden_size]
 
         # Apply ST projector
         if self.projector is not None:
-            embeddings = self.projector(pooled_data)
-        else:
-            embeddings = pooled_data
-        # embeddings shape: [n_tokens, embedding_size]
+            return self.projector(pooled_data)
+        return pooled_data
 
+    def _postprocess_embeddings(
+        self,
+        embeddings: torch.Tensor,
+        pooling_param: PoolingParams,
+    ) -> torch.Tensor:
         # for matryoshka representation
         embeddings = embeddings[..., : pooling_param.dimensions]
 
@@ -87,6 +118,20 @@ def forward_chunk(
         # embeddings shape: [n_tokens, embedding_size]
         return embeddings
 
+    def _has_uniform_postprocess(self, pooling_params: list[PoolingParams]) -> bool:
+        # check if we can apply the same postprocess to all tokens in the batch
+        if not pooling_params:
+            return True
+
+        first_param = pooling_params[0]
+        first_dimensions = first_param.dimensions
+        first_use_activation = bool(first_param.use_activation)
+        return all(
+            param.dimensions == first_dimensions
+            and bool(param.use_activation) == first_use_activation
+            for param in pooling_params[1:]
+        )
+
 
 class TokenClassifierPoolerHead(TokenPoolerHead):
     def __init__(

vllm/model_executor/layers/pooler/tokwise/methods.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from abc import ABC, abstractmethod
 from collections.abc import Set
+from dataclasses import dataclass
 from typing import TypeAlias
 
 import torch
@@ -16,6 +17,59 @@
 TokenPoolingMethodOutputItem: TypeAlias = torch.Tensor | None
 
 
+@dataclass
+class RaggedTokenBatch:
+    values: torch.Tensor
+    cu_lengths_cpu: torch.Tensor
+
+    @classmethod
+    def from_lengths(
+        cls,
+        values: torch.Tensor,
+        lengths_cpu: torch.Tensor,
+    ) -> "RaggedTokenBatch":
+        return cls(
+            values=values,
+            cu_lengths_cpu=_make_cu_lengths_cpu(lengths_cpu),
+        )
+
+    @property
+    def num_items(self) -> int:
+        return self.cu_lengths_cpu.shape[0] - 1
+
+    def with_values(self, values: torch.Tensor) -> "RaggedTokenBatch":
+        return RaggedTokenBatch(
+            values=values,
+            cu_lengths_cpu=self.cu_lengths_cpu,
+        )
+
+    def split(self) -> list[TokenPoolingMethodOutputItem]:
+        outputs = list[TokenPoolingMethodOutputItem]()
+        cu_lengths_cpu = self.cu_lengths_cpu
+
+        for i in range(self.num_items):
+            start = int(cu_lengths_cpu[i])
+            end = int(cu_lengths_cpu[i + 1])
+            outputs.append(self.values[start:end])
+
+        return outputs
+
+
+def _make_cu_lengths_cpu(lengths_cpu: torch.Tensor) -> torch.Tensor:
+    # [1, 2, 3, 4] -> [0, 1, 3, 6, 10]
+    lengths_cpu = lengths_cpu.to(device="cpu", dtype=torch.int64)
+    cu_lengths_cpu = torch.zeros(
+        lengths_cpu.shape[0] + 1, dtype=torch.int64, device="cpu"
+    )
+    torch.cumsum(lengths_cpu, dim=0, out=cu_lengths_cpu[1:])
+    return cu_lengths_cpu
+
+
+TokenPoolingMethodOutput: TypeAlias = (
+    RaggedTokenBatch | list[TokenPoolingMethodOutputItem]
+)
+
+
 class TokenPoolingMethod(nn.Module, ABC):
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"token_embed", "token_classify"}
@@ -28,7 +82,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
-    ) -> list[TokenPoolingMethodOutputItem]:
+    ) -> TokenPoolingMethodOutput:
         raise NotImplementedError
 
 
@@ -45,8 +99,14 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
-    ) -> list[TokenPoolingMethodOutputItem]:
+    ) -> TokenPoolingMethodOutput:
         pooling_cursor = pooling_metadata.get_pooling_cursor()
+        if not self.enable_chunked_prefill:
+            return RaggedTokenBatch.from_lengths(
+                values=hidden_states,
+                lengths_cpu=pooling_cursor.num_scheduled_tokens_cpu,
+            )
+
         hidden_states_lst = [
             hidden_states[first : last + 1]
             for first, last in zip(
@@ -55,9 +115,6 @@ def forward(
             )
         ]
 
-        if not self.enable_chunked_prefill:
-            return hidden_states_lst
-
         pooling_states = pooling_metadata.pooling_states
 
         # If chunked_prefill is enabled
@@ -90,7 +147,12 @@ def forward(
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> list[TokenPoolingMethodOutputItem]:
-        pooled_data_lst = super().forward(hidden_states, pooling_metadata)
+        pooled_data = super().forward(hidden_states, pooling_metadata)
+        pooled_data_lst = (
+            pooled_data.split()
+            if isinstance(pooled_data, RaggedTokenBatch)
+            else pooled_data
+        )
         prompt_token_ids = pooling_metadata.get_prompt_token_ids()
         pooling_params = pooling_metadata.pooling_params
 

vllm/model_executor/layers/pooler/tokwise/poolers.py

@@ -28,14 +28,16 @@
     TokenPoolerHeadOutputItem,
 )
 from .methods import (
+    RaggedTokenBatch,
     TokenPoolingMethod,
+    TokenPoolingMethodOutput,
     TokenPoolingMethodOutputItem,
     get_tok_pooling_method,
 )
 
 TokenPoolingFn: TypeAlias = Callable[
     [torch.Tensor, PoolingMetadata],
-    list[TokenPoolingMethodOutputItem],
+    TokenPoolingMethodOutput,
 ]
 TokenPoolingHeadFn: TypeAlias = Callable[
     [list[TokenPoolingMethodOutputItem], PoolingMetadata],
@@ -89,6 +91,12 @@ def forward(
         pooling_metadata: PoolingMetadata,
     ) -> TokenPoolerOutput:
         pooled_data = self.pooling(hidden_states, pooling_metadata)
+        if isinstance(pooled_data, RaggedTokenBatch):
+            if isinstance(self.head, TokenEmbeddingPoolerHead):
+                return self.head.forward_ragged(
+                    pooled_data, pooling_metadata.pooling_params
+                )
+            pooled_data = pooled_data.split()
         if self.head is not None:
             pooled_data = self.head(pooled_data, pooling_metadata)
         return pooled_data