| 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 import get_current_vllm_config |
| 11 | from vllm.config.pooler import TokenPoolingType |
| 12 | from vllm.model_executor.layers.pooler import PoolingParamsUpdate |
| 13 | from vllm.tasks import PoolingTask |
| 14 | from vllm.utils.torch_utils import async_tensor_h2d |
| 15 | from vllm.v1.pool.metadata import PoolingMetadata |
| 16 | |
| 17 | TokenPoolingMethodOutputItem: TypeAlias = torch.Tensor | None |
| 18 | |
| 19 | |
| 20 | class TokenPoolingMethod(nn.Module, ABC): |
| 21 | def get_supported_tasks(self) -> Set[PoolingTask]: |
| 22 | return {"token_embed", "token_classify"} |
| 23 | |
| 24 | def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: |
| 25 | return PoolingParamsUpdate() |
| 26 | |
| 27 | @abstractmethod |
| 28 | def forward( |
| 29 | self, |
| 30 | hidden_states: torch.Tensor, |
| 31 | pooling_metadata: PoolingMetadata, |
| 32 | ) -> list[TokenPoolingMethodOutputItem]: |
| 33 | raise NotImplementedError |
| 34 | |
| 35 | |
| 36 | class AllPool(TokenPoolingMethod): |
| 37 | def __init__(self): |
| 38 | super().__init__() |
| 39 | |
| 40 | vllm_config = get_current_vllm_config() |
| 41 | scheduler_config = vllm_config.scheduler_config |
| 42 | |
| 43 | self.enable_chunked_prefill = scheduler_config.enable_chunked_prefill |
| 44 | # With async scheduling, the next step may overwrite the hidden states |
| 45 | # buffer before outputs are copied to CPU, so they must own their storage. |
| 46 | self.clone_finished = bool(scheduler_config.async_scheduling) |
| 47 | |
| 48 | def extra_repr(self) -> str: |
| 49 | return ( |
| 50 | f"enable_chunked_prefill={self.enable_chunked_prefill}, " |
| 51 | f"clone_finished={self.clone_finished}" |
| 52 | ) |
| 53 | |
| 54 | def forward( |
| 55 | self, |
| 56 | hidden_states: torch.Tensor, |
| 57 | pooling_metadata: PoolingMetadata, |
| 58 | ) -> list[TokenPoolingMethodOutputItem]: |
| 59 | pooling_cursor = pooling_metadata.get_pooling_cursor() |
| 60 | # Use the already-CPU num_scheduled_tokens tensor so `.tolist()` |
| 61 | # doesn't trigger a GPU->CPU sync. torch.split produces the same |
| 62 | # consecutive slices as indexing with first/last per-sequence indices. |
| 63 | hidden_states_lst = list( |
| 64 | torch.split(hidden_states, pooling_cursor.num_scheduled_tokens_cpu.tolist()) |
| 65 | ) |
| 66 | |
| 67 | if not self.enable_chunked_prefill: |
| 68 | if self.clone_finished: |
| 69 | return [hs.clone() for hs in hidden_states_lst] |
| 70 | return hidden_states_lst |
| 71 | |
| 72 | pooling_states = pooling_metadata.pooling_states |
| 73 | finished_mask = pooling_cursor.get_finished_mask() |
| 74 | |
| 75 | # If chunked_prefill is enabled |
| 76 | # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache |
| 77 | for p, hs_chunk, finished in zip( |
| 78 | pooling_states, hidden_states_lst, finished_mask |
| 79 | ): |
| 80 | # Own data retained across steps or returned without concatenation. |
| 81 | needs_owned_storage = not finished or ( |
| 82 | self.clone_finished and not p.hidden_states_cache |
| 83 | ) |
| 84 | p.hidden_states_cache.append( |
| 85 | hs_chunk.clone() if needs_owned_storage else hs_chunk |
| 86 | ) |
| 87 | |
| 88 | # 2. Once prefill is finished, send hidden_states_cache to PoolerHead |
| 89 | output_list = list[TokenPoolingMethodOutputItem]() |
| 90 | for p, finished in zip(pooling_states, finished_mask): |
| 91 | if finished: |
| 92 | hidden_states_cache = p.hidden_states_cache |
| 93 | if len(hidden_states_cache) == 1: |
| 94 | output_list.append(hidden_states_cache[0]) |
| 95 | else: |
| 96 | output_list.append(torch.concat(hidden_states_cache, dim=0)) |
| 97 | p.clean() |
| 98 | else: |
| 99 | output_list.append(None) |
| 100 | |
| 101 | return output_list |
| 102 | |
| 103 | |
| 104 | class StepPool(AllPool): |
| 105 | def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: |
| 106 | return PoolingParamsUpdate(requires_token_ids=True) |
| 107 | |
| 108 | def forward( |
| 109 | self, |
| 110 | hidden_states: torch.Tensor, |
| 111 | pooling_metadata: PoolingMetadata, |
| 112 | ) -> list[TokenPoolingMethodOutputItem]: |
| 113 | pooled_data_lst = super().forward(hidden_states, pooling_metadata) |
| 114 | # Use the CPU copy of prompt_token_ids so the step_tag_id mask can be |
| 115 | # resolved to indices without a d2h sync from boolean indexing. |
| 116 | prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu() |
| 117 | pooling_params = pooling_metadata.pooling_params |
| 118 | |
| 119 | pooled_data = list[torch.Tensor | None]() |
| 120 | for data, token_id_cpu, pooling_param in zip( |
| 121 | pooled_data_lst, prompt_token_ids_cpu, pooling_params |
| 122 | ): |
| 123 | # for unfinished chunked prefill |
| 124 | if data is None: |
| 125 | pooled_data.append(None) |
| 126 | else: |
| 127 | step_tag_id = pooling_param.step_tag_id |
| 128 | returned_token_ids = pooling_param.returned_token_ids |
| 129 | |
| 130 | if returned_token_ids is not None and len(returned_token_ids) > 0: |
| 131 | data = data[:, returned_token_ids] |
| 132 | |
| 133 | if step_tag_id is not None: |
| 134 | idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0] |
| 135 | idx = async_tensor_h2d(idx_cpu, data.device) |
| 136 | data = data[idx] |
| 137 | |
| 138 | pooled_data.append(data) |
| 139 | |
| 140 | return pooled_data |
| 141 | |
| 142 | |
| 143 | def get_tok_pooling_method(pooling_type: TokenPoolingType | str): |
| 144 | if pooling_type == "ALL": |
| 145 | return AllPool() |
| 146 | if pooling_type == "STEP": |
| 147 | return StepPool() |
| 148 | |
| 149 | raise NotImplementedError(f"Unknown tokenwise pooling type: {pooling_type!r}") |