[Pooling] Improve BGE-M3 sync pooling throughput by up to 3.13% (#53464)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

tests/model_executor/layers/test_pooler_methods.py

@@ -5,9 +5,11 @@
 from dataclasses import dataclass
 from unittest.mock import patch
 
+import numpy as np
 import pytest
 import torch
 
+from vllm.model_executor.layers.pooler import PoolingParamsUpdate
 from vllm.model_executor.layers.pooler.seqwise.methods import (
     CLSPool,
     LastPool,
@@ -52,6 +54,8 @@ def _make_pooling_cursor(
         prompt_lens_cpu=prompt_lens_cpu,
         seq_lens_cpu=seq_lens_cpu,
         num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
+        partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu),
+        finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(),
     )
 
 
@@ -101,6 +105,38 @@ def _make_metadata(
     )
 
 
+def test_build_pooling_cursor_caches_cpu_status() -> None:
+    metadata = PoolingMetadata(
+        prompt_lens=torch.tensor([3, 4, 2], dtype=torch.long),
+        prompt_token_ids=None,
+        prompt_token_ids_cpu=None,
+        pooling_params=[PoolingParams(task="embed") for _ in range(3)],
+        pooling_states=[PoolingStates() for _ in range(3)],
+    )
+    metadata.build_pooling_cursor(
+        np.array([3, 2, 2], dtype=np.int32),
+        seq_lens_cpu=torch.tensor([3, 2, 1], dtype=torch.long),
+        device=_CPU,
+    )
+
+    cursor = metadata.get_pooling_cursor()
+
+    assert cursor.is_partial_prefill() is True
+    assert cursor.get_finished_mask() == [True, False, False]
+    assert cursor[:1].is_partial_prefill() is False
+    assert cursor[1:2].is_partial_prefill() is True
+    assert cursor[1:].get_finished_mask() == [False, False]
+
+
+def test_pooling_params_update_combines_token_id_requirements() -> None:
+    update = PoolingParamsUpdate() | PoolingParamsUpdate(requires_token_ids=True)
+    params = PoolingParams(task="embed")
+
+    update.apply(params)
+
+    assert params.requires_token_ids is True
+
+
 # ---------------------------------------------------------------------------
 # CLSPool
 # ---------------------------------------------------------------------------

tests/models/language/pooling/test_splade_sparse_pooler.py

@@ -12,9 +12,7 @@ from vllm.model_executor.models.bert import (
     BertMLMHead,
     SPLADESparsePooler,
 )
-from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
-from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.pool.late_interaction_runner import LateInteractionRunner
 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
 from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -119,18 +117,17 @@ def test_pooling_runner_gathers_required_token_ids() -> None:
     req_states = MagicMock(spec=RequestState)
     req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32))
     metadata = runner._get_pooling_metadata(
-        input_batch, req_states, torch.device(current_platform.device_type)
+        input_batch, req_states, torch.device("cpu")
     )
 
     expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
     assert metadata.prompt_token_ids_cpu is not None
-    assert metadata.prompt_token_ids is not None
-    assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
+    assert metadata.prompt_token_ids is None
+    assert metadata.prompt_token_ids_cpu.is_pinned() is False
     torch.testing.assert_close(
         metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32)
     )
     torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected)
-    torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
 
 
 def test_pooling_runner_stores_only_required_token_ids() -> None:

tests/v1/worker/test_gpu_input_batch.py

@@ -475,15 +475,14 @@ def test_placeholder_spec_token_ids_written_verbatim():
 
 
 @pytest.mark.parametrize(
-    ("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
+    ("pooling_params", "expect_cpu_prompt_token_ids"),
     [
-        ({"task": "classify"}, False, False),
-        ({"task": "classify", "requires_token_ids": True}, True, True),
+        ({"task": "classify"}, False),
+        ({"task": "classify", "requires_token_ids": True}, True),
     ],
 )
 def test_pooling_metadata_token_id_buffers(
     pooling_params: dict[str, object],
-    expect_device_prompt_token_ids: bool,
     expect_cpu_prompt_token_ids: bool,
 ):
     from vllm.pooling_params import PoolingParams
@@ -504,13 +503,8 @@ def test_pooling_metadata_token_id_buffers(
     input_batch.refresh_metadata()
 
     metadata = input_batch.get_pooling_metadata()
-    if expect_device_prompt_token_ids:
-        assert input_batch.sampling_metadata.prompt_token_ids is not None
-        assert metadata.prompt_token_ids is not None
-        assert metadata.get_prompt_token_ids()[0].tolist() == req.prompt_token_ids
-    else:
-        assert input_batch.sampling_metadata.prompt_token_ids is None
-        assert metadata.prompt_token_ids is None
+    assert input_batch.sampling_metadata.prompt_token_ids is None
+    assert metadata.prompt_token_ids is None
 
     if expect_cpu_prompt_token_ids:
         assert metadata.prompt_token_ids_cpu is not None

vllm/model_executor/layers/pooler/common.py

@@ -18,7 +18,7 @@ ActivationFn = Callable[[_T], _T]
 @dataclass(frozen=True)
 class PoolingParamsUpdate:
     requires_token_ids: bool = False
-    """Set this flag to enable prompt token IDs for your pooler."""
+    """Set this flag to enable CPU prompt token IDs for your pooler."""
 
     def __or__(self, other: "PoolingParamsUpdate") -> "PoolingParamsUpdate":
         return PoolingParamsUpdate(

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

@@ -69,7 +69,7 @@ class AllPool(TokenPoolingMethod):
             return hidden_states_lst
 
         pooling_states = pooling_metadata.pooling_states
-        finished_mask = pooling_cursor.is_finished().tolist()
+        finished_mask = pooling_cursor.get_finished_mask()
 
         # If chunked_prefill is enabled
         # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache

vllm/model_executor/models/jina.py

@@ -12,11 +12,12 @@ from torch import nn
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.pooler import PoolingParamsUpdate
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import PoolingTask
 from vllm.transformers_utils.repo_utils import get_hf_file_bytes
-from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.pool.metadata import PoolingMetadata
 
 from ..layers.pooler import DispatchPooler
@@ -104,33 +105,37 @@ class JinaForRankingPool(StepPool):
     def get_supported_tasks(self) -> set[PoolingTask]:
         return {"token_embed"}
 
+    def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
+        return PoolingParamsUpdate(requires_token_ids=True)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> list[TokenPoolingMethodOutputItem]:
         pooled_data_lst = super().forward(hidden_states, pooling_metadata)
-        prompt_token_ids = pooling_metadata.get_prompt_token_ids()
+        prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu()
 
         embeds_list = list[torch.Tensor | None]()
-        # `torch.where` resolves the match count on the host.
-        with gpu_sync_allowed():
-            for data, token_ids in zip(pooled_data_lst, prompt_token_ids):
-                # for unfinished chunked prefill
-                if data is None:
-                    embeds_list.append(None)
-                else:
-                    doc_match = torch.eq(token_ids, self.doc_token_id)
-                    docs_indexes = torch.where(doc_match)[0]
-                    query_indexes = torch.where(
-                        torch.eq(token_ids, self.query_token_id)
-                    )[0]
+        for data, token_ids_cpu in zip(pooled_data_lst, prompt_token_ids_cpu):
+            # for unfinished chunked prefill
+            if data is None:
+                embeds_list.append(None)
+            else:
+                doc_match = torch.eq(token_ids_cpu, self.doc_token_id)
+                docs_indexes_cpu = torch.where(doc_match)[0]
+                query_indexes_cpu = torch.where(
+                    torch.eq(token_ids_cpu, self.query_token_id)
+                )[0]
 
-                    # The JinaForRanking model concatenates docs first, then query.
-                    # Let's stay consistent with this novel design.
-                    indexes = torch.cat([docs_indexes, query_indexes])
-                    embeds = self.projector(data[indexes])
-                    embeds_list.append(embeds)
+                # The JinaForRanking model concatenates docs first, then query.
+                # Let's stay consistent with this novel design.
+                indexes_cpu = torch.cat([docs_indexes_cpu, query_indexes_cpu])
+                indexes = async_tensor_h2d(
+                    indexes_cpu, device=data.device, dtype=torch.long
+                )
+                embeds = self.projector(data[indexes])
+                embeds_list.append(embeds)
 
         return embeds_list
 

vllm/v1/pool/metadata.py

@@ -17,22 +17,33 @@ class PoolingCursor:
     prompt_lens_cpu: torch.Tensor
     seq_lens_cpu: torch.Tensor
     num_scheduled_tokens_cpu: torch.Tensor
+    partial_prefill: bool
+    finished_mask: list[bool]
 
     def __getitem__(self, indices: slice) -> "PoolingCursor":
+        prompt_lens_cpu = self.prompt_lens_cpu[indices]
+        seq_lens_cpu = self.seq_lens_cpu[indices]
+        num_scheduled_tokens_cpu = self.num_scheduled_tokens_cpu[indices]
+
         return PoolingCursor(
             first_token_indices_gpu=self.first_token_indices_gpu[indices],
             last_token_indices_gpu=self.last_token_indices_gpu[indices],
-            prompt_lens_cpu=self.prompt_lens_cpu[indices],
-            seq_lens_cpu=self.seq_lens_cpu[indices],
-            num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],
+            prompt_lens_cpu=prompt_lens_cpu,
+            seq_lens_cpu=seq_lens_cpu,
+            num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
+            partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu),
+            finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(),
         )
 
     def is_partial_prefill(self) -> bool:
-        return not torch.all(self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)
+        return self.partial_prefill
 
     def is_finished(self) -> torch.Tensor:
         return self.prompt_lens_cpu == self.seq_lens_cpu
 
+    def get_finished_mask(self) -> list[bool]:
+        return self.finished_mask
+
 
 class PoolingStates:
     def __init__(self) -> None:
@@ -130,9 +141,14 @@ class PoolingMetadata:
             )
 
         num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
+        prompt_lens_np = prompt_lens.numpy()
+        seq_lens_np = seq_lens_cpu.numpy()
         if query_start_loc_gpu is None:
             cumsum = torch.zeros(
-                n_seq + 1, dtype=torch.int64, pin_memory=PIN_MEMORY, device="cpu"
+                n_seq + 1,
+                dtype=torch.int64,
+                pin_memory=PIN_MEMORY,
+                device="cpu",
             )
             torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
             cumsum = cumsum.to(device, non_blocking=True)
@@ -155,4 +171,6 @@ class PoolingMetadata:
             prompt_lens_cpu=prompt_lens,
             seq_lens_cpu=seq_lens_cpu,
             num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
+            partial_prefill=not np.array_equal(prompt_lens_np, num_scheduled_tokens_np),
+            finished_mask=np.equal(prompt_lens_np, seq_lens_np).tolist(),
         )

vllm/v1/worker/gpu/pool/pooling_runner.py

@@ -103,28 +103,26 @@ class PoolingRunner:
         pooling_params = [self.pooling_params[i] for i in req_indices]
         pooling_states = [self.pooling_states[i] for i in req_indices]
         prompt_lens = torch.from_numpy(
-            req_states.prompt_len.np[input_batch.idx_mapping_np].copy()
+            req_states.prompt_len.np[input_batch.idx_mapping_np]
         )
 
         prompt_token_ids_cpu = None
-        prompt_token_ids = None
         if any(params.requires_token_ids for params in pooling_params):
             max_prompt_len = int(prompt_lens.max())
             prompt_token_ids_cpu = torch.zeros(
                 (input_batch.num_reqs, max_prompt_len),
                 dtype=torch.int64,
-                pin_memory=PIN_MEMORY,
+                pin_memory=PIN_MEMORY and device.type != "cpu",
             )
             for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
                 if not params.requires_token_ids:
                     continue
                 token_ids = self.prompt_token_ids[req_index]
                 prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
-            prompt_token_ids = prompt_token_ids_cpu.to(device, non_blocking=True)
 
         return PoolingMetadata(
             prompt_lens=prompt_lens,
-            prompt_token_ids=prompt_token_ids,
+            prompt_token_ids=None,
             prompt_token_ids_cpu=prompt_token_ids_cpu,
             pooling_params=pooling_params,
             pooling_states=pooling_states,
@@ -150,7 +148,7 @@ class PoolingRunner:
             query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
         )
         pooler_output = self.model.pooler(hidden_states, pooling_metadata)
-        finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()
+        finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask()
         pooler_output = self.late_interaction_runner.postprocess_pooler_output(
             raw_pooler_output=pooler_output,
             pooling_params=pooling_metadata.pooling_params,
@@ -174,19 +172,16 @@ class PoolingRunner:
         pooling_params = PoolingParams(task=task)
         pooling_params.verify(self.model_config)
         self.model.pooler.get_pooling_updates(task).apply(pooling_params)
-        prompt_token_ids = None
+        prompt_token_ids_cpu = None
         if pooling_params.requires_token_ids:
-            prompt_token_ids = torch.zeros(
+            prompt_token_ids_cpu = torch.zeros(
                 (num_reqs, int(prompt_lens.max())),
                 dtype=torch.int64,
-                device=hidden_states.device,
             )
         pooling_metadata = PoolingMetadata(
             prompt_lens=prompt_lens,
-            prompt_token_ids=prompt_token_ids,
-            prompt_token_ids_cpu=None
-            if prompt_token_ids is None
-            else prompt_token_ids.cpu(),
+            prompt_token_ids=None,
+            prompt_token_ids_cpu=prompt_token_ids_cpu,
             pooling_params=[pooling_params] * num_reqs,
             pooling_states=[PoolingStates() for _ in range(num_reqs)],
         )

vllm/v1/worker/gpu_input_batch.py

@@ -476,9 +476,7 @@ class InputBatch:
 
             self.pooling_params[req_id] = pooling_params
             self.pooling_states[req_id] = pooling_states
-            self.logits_processing_needs_token_ids[req_index] = (
-                pooling_params.requires_token_ids
-            )
+            self.logits_processing_needs_token_ids[req_index] = False
         else:
             raise NotImplementedError("Unrecognized request type")
 
@@ -890,8 +888,8 @@ class InputBatch:
             not self.no_penalties
             or self.logits_processing_needs_token_ids[:num_reqs].any()
         )
-        # The prompt tokens are used only for applying penalties or
-        # step pooling during the sampling/pooling process.
+        # The device prompt tokens are used only for applying penalties or
+        # pooling methods that explicitly request GPU token IDs.
         # Hence copy these tensors only when there are requests which
         # need penalties/step_pooler to be applied.
         prompt_token_ids_cpu = (
@@ -979,7 +977,7 @@ class InputBatch:
 
         return PoolingMetadata(
             prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(),
-            prompt_token_ids=self.sampling_metadata.prompt_token_ids,
+            prompt_token_ids=None,
             prompt_token_ids_cpu=prompt_token_ids_cpu,
             pooling_params=pooling_params,
             pooling_states=pooling_states,

vllm/v1/worker/gpu_model_runner.py

@@ -3550,10 +3550,7 @@ class GPUModelRunner(
             hidden_states=hidden_states, pooling_metadata=pooling_metadata
         )
 
-        finished_mask = [
-            seq_len == prompt_len
-            for seq_len, prompt_len in zip(seq_lens_cpu, pooling_metadata.prompt_lens)
-        ]
+        finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask()
         raw_pooler_output = self.late_interaction_runner.postprocess_pooler_output(
             raw_pooler_output=raw_pooler_output,
             pooling_params=pooling_metadata.pooling_params,