| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | |
| 5 | import numpy as np |
| 6 | import torch |
| 7 | |
| 8 | from vllm.pooling_params import PoolingParams |
| 9 | from vllm.tasks import PoolingTask |
| 10 | from vllm.utils.torch_utils import PIN_MEMORY |
| 11 | |
| 12 | |
| 13 | @dataclass |
| 14 | class PoolingCursor: |
| 15 | first_token_indices_gpu: torch.Tensor |
| 16 | last_token_indices_gpu: torch.Tensor |
| 17 | prompt_lens_cpu: torch.Tensor |
| 18 | seq_lens_cpu: torch.Tensor |
| 19 | num_scheduled_tokens_cpu: torch.Tensor |
| 20 | partial_prefill: bool |
| 21 | finished_mask: list[bool] |
| 22 | |
| 23 | def __getitem__(self, indices: slice) -> "PoolingCursor": |
| 24 | prompt_lens_cpu = self.prompt_lens_cpu[indices] |
| 25 | seq_lens_cpu = self.seq_lens_cpu[indices] |
| 26 | num_scheduled_tokens_cpu = self.num_scheduled_tokens_cpu[indices] |
| 27 | |
| 28 | return PoolingCursor( |
| 29 | first_token_indices_gpu=self.first_token_indices_gpu[indices], |
| 30 | last_token_indices_gpu=self.last_token_indices_gpu[indices], |
| 31 | prompt_lens_cpu=prompt_lens_cpu, |
| 32 | seq_lens_cpu=seq_lens_cpu, |
| 33 | num_scheduled_tokens_cpu=num_scheduled_tokens_cpu, |
| 34 | partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu), |
| 35 | finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(), |
| 36 | ) |
| 37 | |
| 38 | def is_partial_prefill(self) -> bool: |
| 39 | return self.partial_prefill |
| 40 | |
| 41 | def is_finished(self) -> torch.Tensor: |
| 42 | return self.prompt_lens_cpu == self.seq_lens_cpu |
| 43 | |
| 44 | def get_finished_mask(self) -> list[bool]: |
| 45 | return self.finished_mask |
| 46 | |
| 47 | |
| 48 | class PoolingStates: |
| 49 | def __init__(self) -> None: |
| 50 | # for chunked prefill with ALL pooling |
| 51 | self.hidden_states_cache: list[torch.Tensor] = [] |
| 52 | |
| 53 | def clean(self) -> None: |
| 54 | self.hidden_states_cache.clear() |
| 55 | |
| 56 | |
| 57 | @dataclass |
| 58 | class PoolingMetadata: |
| 59 | """Tensors for pooling.""" |
| 60 | |
| 61 | prompt_lens: torch.Tensor # CPU Tensor |
| 62 | prompt_token_ids: torch.Tensor | None # Model-device tensor |
| 63 | prompt_token_ids_cpu: torch.Tensor | None # CPU tensor |
| 64 | pooling_params: list[PoolingParams] |
| 65 | pooling_states: list[PoolingStates] |
| 66 | pooling_cursor: PoolingCursor | None = None |
| 67 | |
| 68 | def __post_init__(self) -> None: |
| 69 | pooling_params = self.pooling_params |
| 70 | |
| 71 | tasks: list[PoolingTask] = [ |
| 72 | task |
| 73 | for pooling_param in pooling_params |
| 74 | if (task := pooling_param.task) is not None |
| 75 | ] |
| 76 | if len(pooling_params) != len(tasks): |
| 77 | raise ValueError( |
| 78 | "Every pooling param must have a task set, but got " |
| 79 | f"{len(tasks)} tasks for {len(pooling_params)} pooling params" |
| 80 | ) |
| 81 | |
| 82 | self.tasks = tasks |
| 83 | |
| 84 | def __getitem__(self, indices: slice) -> "PoolingMetadata": |
| 85 | return PoolingMetadata( |
| 86 | prompt_lens=self.prompt_lens[indices], |
| 87 | prompt_token_ids=None |
| 88 | if self.prompt_token_ids is None |
| 89 | else self.prompt_token_ids[indices], |
| 90 | prompt_token_ids_cpu=None |
| 91 | if self.prompt_token_ids_cpu is None |
| 92 | else self.prompt_token_ids_cpu[indices], |
| 93 | pooling_params=self.pooling_params[indices], |
| 94 | pooling_states=self.pooling_states[indices], |
| 95 | pooling_cursor=None |
| 96 | if self.pooling_cursor is None |
| 97 | else self.pooling_cursor[indices], |
| 98 | ) |
| 99 | |
| 100 | def _get_prompt_token_ids( |
| 101 | self, |
| 102 | prompt_token_ids: torch.Tensor | None, |
| 103 | ) -> list[torch.Tensor]: |
| 104 | if prompt_token_ids is None: |
| 105 | raise ValueError( |
| 106 | "prompt_token_ids is required but was not set. " |
| 107 | "Please set `requires_token_ids=True` in `get_pooling_updates`" |
| 108 | ) |
| 109 | return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)] |
| 110 | |
| 111 | def get_prompt_token_ids(self) -> list[torch.Tensor]: |
| 112 | return self._get_prompt_token_ids(self.prompt_token_ids) |
| 113 | |
| 114 | def get_prompt_token_ids_cpu(self) -> list[torch.Tensor]: |
| 115 | return self._get_prompt_token_ids(self.prompt_token_ids_cpu) |
| 116 | |
| 117 | def get_pooling_cursor(self) -> PoolingCursor: |
| 118 | pooling_cursor = self.pooling_cursor |
| 119 | if pooling_cursor is None: |
| 120 | raise RuntimeError( |
| 121 | "pooling_cursor has not been initialized. " |
| 122 | "Call `build_pooling_cursor` before accessing it" |
| 123 | ) |
| 124 | |
| 125 | return pooling_cursor |
| 126 | |
| 127 | def build_pooling_cursor( |
| 128 | self, |
| 129 | num_scheduled_tokens_np: np.ndarray, |
| 130 | seq_lens_cpu: torch.Tensor, |
| 131 | device: torch.device, |
| 132 | query_start_loc_gpu: torch.Tensor | None = None, |
| 133 | ) -> None: |
| 134 | n_seq = len(num_scheduled_tokens_np) |
| 135 | prompt_lens = self.prompt_lens |
| 136 | |
| 137 | if len(prompt_lens) != n_seq: |
| 138 | raise ValueError( |
| 139 | f"prompt_lens length ({len(prompt_lens)}) does not match " |
| 140 | f"the number of sequences ({n_seq})" |
| 141 | ) |
| 142 | |
| 143 | num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np) |
| 144 | prompt_lens_np = prompt_lens.numpy() |
| 145 | seq_lens_np = seq_lens_cpu.numpy() |
| 146 | if query_start_loc_gpu is None: |
| 147 | cumsum = torch.zeros( |
| 148 | n_seq + 1, |
| 149 | dtype=torch.int64, |
| 150 | pin_memory=PIN_MEMORY, |
| 151 | device="cpu", |
| 152 | ) |
| 153 | torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:]) |
| 154 | cumsum = cumsum.to(device, non_blocking=True) |
| 155 | else: |
| 156 | if query_start_loc_gpu.shape[0] != n_seq + 1: |
| 157 | raise ValueError( |
| 158 | "query_start_loc_gpu length does not match " |
| 159 | f"the number of sequences: {query_start_loc_gpu.shape[0]} " |
| 160 | f"!= {n_seq + 1}." |
| 161 | ) |
| 162 | if query_start_loc_gpu.device != device: |
| 163 | raise ValueError( |
| 164 | "query_start_loc_gpu must be on the same device as the " |
| 165 | f"hidden states: {query_start_loc_gpu.device} != {device}." |
| 166 | ) |
| 167 | cumsum = query_start_loc_gpu |
| 168 | self.pooling_cursor = PoolingCursor( |
| 169 | first_token_indices_gpu=cumsum[:n_seq], |
| 170 | last_token_indices_gpu=cumsum[1:] - 1, |
| 171 | prompt_lens_cpu=prompt_lens, |
| 172 | seq_lens_cpu=seq_lens_cpu, |
| 173 | num_scheduled_tokens_cpu=num_scheduled_tokens_cpu, |
| 174 | partial_prefill=not np.array_equal(prompt_lens_np, num_scheduled_tokens_np), |
| 175 | finished_mask=np.equal(prompt_lens_np, seq_lens_np).tolist(), |
| 176 | ) |