diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py index 8460d7dfe8..53ab98ff44 100644 --- a/tests/model_executor/layers/test_pooler_methods.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index 0218d3694f..0b3561f90f 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/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: diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index e7aa2e03e0..3ef4e0dabe 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/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 diff --git a/vllm/model_executor/layers/pooler/common.py b/vllm/model_executor/layers/pooler/common.py index 55fc4b4576..9a43432c2d 100644 --- a/vllm/model_executor/layers/pooler/common.py +++ b/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( diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index 98b58c286c..1eb8a8ba3d 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/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 diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 456a51b3f3..1a4d2b0441 100644 --- a/vllm/model_executor/models/jina.py +++ b/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 diff --git a/vllm/v1/pool/metadata.py b/vllm/v1/pool/metadata.py index 9a9bb2b0e7..ffa14fbb01 100644 --- a/vllm/v1/pool/metadata.py +++ b/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(), ) diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index 29fae75273..839edc8b8c 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/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)], ) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 0b805fca76..7585f4d892 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/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, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8a0e27d3c8..93e19fca21 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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,