| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Set |
| 5 | from typing import TypeAlias |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.config.pooler import SequencePoolingType |
| 11 | from vllm.model_executor.layers.pooler import PoolingParamsUpdate |
| 12 | from vllm.tasks import PoolingTask |
| 13 | from vllm.utils.torch_utils import async_tensor_h2d |
| 14 | from vllm.v1.pool.metadata import PoolingMetadata |
| 15 | |
| 16 | SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor] |
| 17 | |
| 18 | _MEAN_POOL_ACCUMULATION_CHUNK_BYTES = 16 * 1024 * 1024 # 16MB |
| 19 | |
| 20 | |
| 21 | class SequencePoolingMethod(nn.Module, ABC): |
| 22 | def get_supported_tasks(self) -> Set[PoolingTask]: |
| 23 | return {"token_embed", "token_classify", "embed", "classify"} |
| 24 | |
| 25 | def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: |
| 26 | return PoolingParamsUpdate() |
| 27 | |
| 28 | @abstractmethod |
| 29 | def forward( |
| 30 | self, |
| 31 | hidden_states: torch.Tensor, |
| 32 | pooling_metadata: PoolingMetadata, |
| 33 | ) -> SequencePoolingMethodOutput: |
| 34 | raise NotImplementedError |
| 35 | |
| 36 | |
| 37 | class CLSPool(SequencePoolingMethod): |
| 38 | def forward( |
| 39 | self, |
| 40 | hidden_states: torch.Tensor, |
| 41 | pooling_metadata: PoolingMetadata, |
| 42 | ) -> SequencePoolingMethodOutput: |
| 43 | pooling_cursor = pooling_metadata.get_pooling_cursor() |
| 44 | if pooling_cursor.is_partial_prefill(): |
| 45 | raise RuntimeError("partial prefill is not supported with CLS pooling") |
| 46 | |
| 47 | return hidden_states[pooling_cursor.first_token_indices_gpu] |
| 48 | |
| 49 | |
| 50 | class LastPool(SequencePoolingMethod): |
| 51 | def forward( |
| 52 | self, |
| 53 | hidden_states: torch.Tensor, |
| 54 | pooling_metadata: PoolingMetadata, |
| 55 | ) -> SequencePoolingMethodOutput: |
| 56 | pooling_cursor = pooling_metadata.get_pooling_cursor() |
| 57 | return hidden_states[pooling_cursor.last_token_indices_gpu] |
| 58 | |
| 59 | |
| 60 | class MeanPool(SequencePoolingMethod): |
| 61 | def forward( |
| 62 | self, |
| 63 | hidden_states: torch.Tensor, |
| 64 | pooling_metadata: PoolingMetadata, |
| 65 | ) -> SequencePoolingMethodOutput: |
| 66 | pooling_cursor = pooling_metadata.get_pooling_cursor() |
| 67 | if pooling_cursor.is_partial_prefill(): |
| 68 | raise RuntimeError("partial prefill is not supported with MEAN pooling") |
| 69 | |
| 70 | prompt_lens_cpu = pooling_cursor.prompt_lens_cpu |
| 71 | num_seqs = prompt_lens_cpu.numel() |
| 72 | hidden_size = hidden_states.shape[-1] |
| 73 | |
| 74 | if num_seqs == 0: |
| 75 | # early return for empty batch |
| 76 | return hidden_states.new_empty((0, hidden_size), dtype=torch.float32) |
| 77 | |
| 78 | prompt_lens = async_tensor_h2d( |
| 79 | prompt_lens_cpu, device=hidden_states.device, dtype=torch.int64 |
| 80 | ) |
| 81 | # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] |
| 82 | segment_ids = torch.repeat_interleave( |
| 83 | torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long), |
| 84 | prompt_lens, |
| 85 | output_size=int(prompt_lens_cpu.sum()), |
| 86 | ) |
| 87 | segment_sums = torch.zeros( |
| 88 | (num_seqs, hidden_size), |
| 89 | dtype=torch.float32, |
| 90 | device=hidden_states.device, |
| 91 | ) |
| 92 | |
| 93 | bytes_per_token = hidden_size * torch.finfo(torch.float32).bits // 8 |
| 94 | chunk_size = max(1, _MEAN_POOL_ACCUMULATION_CHUNK_BYTES // bytes_per_token) |
| 95 | |
| 96 | # iterate over the batch in chunks |
| 97 | for start in range(0, hidden_states.shape[0], chunk_size): |
| 98 | end = min(start + chunk_size, hidden_states.shape[0]) |
| 99 | # using index_add_ to accumulate for each segment |
| 100 | segment_sums.index_add_( |
| 101 | 0, |
| 102 | segment_ids[start:end], |
| 103 | hidden_states[start:end].to(dtype=torch.float32), |
| 104 | ) |
| 105 | |
| 106 | return segment_sums / prompt_lens.unsqueeze(1) |
| 107 | |
| 108 | |
| 109 | def get_seq_pooling_method( |
| 110 | pooling_type: SequencePoolingType | str, |
| 111 | ) -> SequencePoolingMethod: |
| 112 | if pooling_type == "CLS": |
| 113 | return CLSPool() |
| 114 | if pooling_type == "LAST": |
| 115 | return LastPool() |
| 116 | if pooling_type == "MEAN": |
| 117 | return MeanPool() |
| 118 | |
| 119 | raise NotImplementedError(f"Unknown sequence pooling type: {pooling_type!r}") |