| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from dataclasses import dataclass |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config.model import LogprobsMode |
| 10 | from vllm.sampling_params import SamplingParams |
| 11 | from vllm.triton_utils import tl, triton |
| 12 | from vllm.utils.torch_utils import async_tensor_h2d |
| 13 | from vllm.v1.outputs import LogprobsTensors |
| 14 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 15 | from vllm.v1.worker.gpu.sample.logprob import ( |
| 16 | compute_token_logprobs, |
| 17 | compute_topk_scores, |
| 18 | ) |
| 19 | |
| 20 | CHUNK_SIZE = 1024 |
| 21 | |
| 22 | |
| 23 | @dataclass |
| 24 | class _TokenIdScores: |
| 25 | token_ids: torch.Tensor # [num_ids], on device |
| 26 | start: int # first prompt row to score |
| 27 | scores: torch.Tensor | None = None # [num_rows, num_ids], filled per chunk |
| 28 | |
| 29 | |
| 30 | class PromptLogprobsWorker: |
| 31 | def __init__( |
| 32 | self, |
| 33 | max_num_reqs: int, |
| 34 | device: torch.device, |
| 35 | logprobs_mode: LogprobsMode = "raw_logprobs", |
| 36 | ): |
| 37 | self.max_num_reqs = max_num_reqs |
| 38 | self.device = device |
| 39 | self.logprobs_mode = logprobs_mode |
| 40 | |
| 41 | self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool) |
| 42 | self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32) |
| 43 | # req_idx -> list of in-progress LogprobsTensors |
| 44 | self.in_progress_prompt_logprobs: dict[str, list[LogprobsTensors]] = {} |
| 45 | # req_id -> fixed-ID scoring state |
| 46 | self.token_id_scores: dict[str, _TokenIdScores] = {} |
| 47 | |
| 48 | def add_request(self, req_id: str, req_idx: int, sampling_params: SamplingParams): |
| 49 | uses_prompt_logprobs = sampling_params.prompt_logprobs is not None |
| 50 | self.uses_prompt_logprobs[req_idx] = uses_prompt_logprobs |
| 51 | self.num_prompt_logprobs[req_idx] = sampling_params.prompt_logprobs or 0 |
| 52 | if uses_prompt_logprobs: |
| 53 | self.in_progress_prompt_logprobs[req_id] = [] |
| 54 | if sampling_params.prompt_logprob_token_ids is not None: |
| 55 | self.token_id_scores[req_id] = _TokenIdScores( |
| 56 | async_tensor_h2d( |
| 57 | sampling_params.prompt_logprob_token_ids, self.device, torch.int64 |
| 58 | ), |
| 59 | sampling_params.prompt_logprob_start or 0, |
| 60 | ) |
| 61 | |
| 62 | def remove_request(self, req_id: str) -> None: |
| 63 | self.in_progress_prompt_logprobs.pop(req_id, None) |
| 64 | self.token_id_scores.pop(req_id, None) |
| 65 | |
| 66 | def compute_prompt_token_id_logprobs( |
| 67 | self, |
| 68 | logits_fn: Callable[[torch.Tensor], torch.Tensor], |
| 69 | hidden_states: torch.Tensor, |
| 70 | input_batch: InputBatch, |
| 71 | prompt_lens: np.ndarray, |
| 72 | ) -> dict[str, torch.Tensor]: |
| 73 | """Fill scores row (prompt row - req.start) per chunk; emit on the last.""" |
| 74 | if not self.token_id_scores: |
| 75 | return {} |
| 76 | logits_mode = self.logprobs_mode in ("raw_logits", "processed_logits") |
| 77 | out: dict[str, torch.Tensor] = {} |
| 78 | for i, req_id in enumerate(input_batch.req_ids): |
| 79 | req = self.token_id_scores.get(req_id) |
| 80 | if req is None: |
| 81 | continue |
| 82 | prompt_len = int(prompt_lens[input_batch.idx_mapping_np[i]]) |
| 83 | chunk_start = int(input_batch.num_computed_prefill_tokens_np[i]) |
| 84 | chunk_end = chunk_start + int(input_batch.num_scheduled_tokens[i]) |
| 85 | # Skip decode steps and, as compute_prompt_logprobs does, requests |
| 86 | # resumed after preemption: their scores were already emitted. |
| 87 | if chunk_start >= prompt_len or prompt_len < input_batch.prefill_len_np[i]: |
| 88 | continue |
| 89 | if req.scores is None: |
| 90 | if chunk_start > req.start: |
| 91 | # This prefill starts past the first scored row, so those |
| 92 | # rows will never be written; drop the request instead of |
| 93 | # emitting a buffer with unwritten rows. |
| 94 | self.token_id_scores.pop(req_id, None) |
| 95 | continue |
| 96 | req.scores = hidden_states.new_empty( |
| 97 | (max(prompt_len - 1 - req.start, 0), len(req.token_ids)), |
| 98 | dtype=torch.float32, |
| 99 | ) |
| 100 | # The last prompt row predicts the first decode token; skip it. |
| 101 | lo = max(req.start, chunk_start) |
| 102 | hi = min(chunk_end, prompt_len - 1) |
| 103 | base = int(input_batch.query_start_loc_np[i]) - chunk_start |
| 104 | for a in range(lo, hi, CHUNK_SIZE): |
| 105 | b = min(a + CHUNK_SIZE, hi) |
| 106 | logits = logits_fn(hidden_states[base + a : base + b]) |
| 107 | ids = req.token_ids.expand(b - a, -1) |
| 108 | req.scores[a - req.start : b - req.start] = ( |
| 109 | logits.gather(-1, ids) |
| 110 | if logits_mode |
| 111 | else compute_token_logprobs(logits, ids) |
| 112 | ) |
| 113 | if chunk_end >= prompt_len: |
| 114 | out[req_id] = req.scores |
| 115 | req.scores = None |
| 116 | return out |
| 117 | |
| 118 | def compute_prompt_logprobs( |
| 119 | self, |
| 120 | logits_fn: Callable[[torch.Tensor], torch.Tensor], |
| 121 | hidden_states: torch.Tensor, |
| 122 | input_batch: InputBatch, |
| 123 | # [max_num_reqs, max_model_len] |
| 124 | all_token_ids: torch.Tensor, |
| 125 | # [max_num_reqs] |
| 126 | num_computed_tokens: torch.Tensor, |
| 127 | # [max_num_reqs] |
| 128 | prompt_lens: np.ndarray, |
| 129 | ) -> dict[str, LogprobsTensors]: |
| 130 | idx_mapping_np = input_batch.idx_mapping_np |
| 131 | needs_prompt_logprobs = self.uses_prompt_logprobs[idx_mapping_np] |
| 132 | if not np.any(needs_prompt_logprobs): |
| 133 | # Common case: No request asks for prompt logprobs. |
| 134 | return {} |
| 135 | |
| 136 | num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np] |
| 137 | prompt_lens = prompt_lens[idx_mapping_np] |
| 138 | computed_prefill = input_batch.num_computed_prefill_tokens_np |
| 139 | includes_prompt = computed_prefill < prompt_lens |
| 140 | # NOTE(woosuk): If the request was resumed after preemption, its prompt |
| 141 | # logprobs must have been computed before preemption. Skip. |
| 142 | resumed_after_prompt = prompt_lens < input_batch.prefill_len_np |
| 143 | needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt |
| 144 | if not np.any(needs_prompt_logprobs): |
| 145 | return {} |
| 146 | |
| 147 | # get the maximum number in this batch |
| 148 | requested_num_prompt_logprobs = num_prompt_logprobs[needs_prompt_logprobs] |
| 149 | max_num_prompt_logprobs = ( |
| 150 | -1 |
| 151 | if np.any(requested_num_prompt_logprobs == -1) |
| 152 | else int(requested_num_prompt_logprobs.max()) |
| 153 | ) |
| 154 | |
| 155 | # Get the prompt logprobs token_ids. |
| 156 | prompt_logprobs_token_ids = get_prompt_logprobs_token_ids( |
| 157 | input_batch.num_tokens, |
| 158 | input_batch.query_start_loc, |
| 159 | input_batch.idx_mapping, |
| 160 | num_computed_tokens, |
| 161 | all_token_ids, |
| 162 | ) |
| 163 | prompt_token_ids, prompt_logprobs, prompt_ranks = ( |
| 164 | compute_prompt_logprobs_with_chunking( |
| 165 | prompt_logprobs_token_ids, |
| 166 | hidden_states[: input_batch.num_tokens], |
| 167 | logits_fn, |
| 168 | max_num_prompt_logprobs, |
| 169 | self.logprobs_mode, |
| 170 | ) |
| 171 | ) |
| 172 | |
| 173 | pos_after_step = computed_prefill + input_batch.num_scheduled_tokens |
| 174 | is_prompt_chunked = pos_after_step < prompt_lens |
| 175 | |
| 176 | query_start_loc_np = input_batch.query_start_loc_np |
| 177 | prompt_logprobs_dict: dict[str, LogprobsTensors] = {} |
| 178 | for i, req_id in enumerate(input_batch.req_ids): |
| 179 | if not needs_prompt_logprobs[i]: |
| 180 | continue |
| 181 | |
| 182 | req_is_prompt_chunked = is_prompt_chunked[i] |
| 183 | req_num_prompt_logprobs = int(num_prompt_logprobs[i]) |
| 184 | start_idx = query_start_loc_np[i] |
| 185 | end_idx = query_start_loc_np[i + 1] |
| 186 | assert start_idx < end_idx, ( |
| 187 | f"start_idx ({start_idx}) >= end_idx ({end_idx})" |
| 188 | ) |
| 189 | if not req_is_prompt_chunked: |
| 190 | end_idx -= 1 |
| 191 | |
| 192 | width = ( |
| 193 | prompt_logprobs.shape[1] |
| 194 | if req_num_prompt_logprobs == -1 |
| 195 | else req_num_prompt_logprobs + 1 |
| 196 | ) |
| 197 | # no logprobs if start_idx >= end_idx |
| 198 | logprobs = ( |
| 199 | None |
| 200 | if start_idx >= end_idx |
| 201 | else LogprobsTensors( |
| 202 | logprob_token_ids=prompt_token_ids[start_idx:end_idx, :width], |
| 203 | logprobs=prompt_logprobs[start_idx:end_idx, :width], |
| 204 | selected_token_ranks=prompt_ranks[start_idx:end_idx], |
| 205 | ) |
| 206 | ) |
| 207 | |
| 208 | prompt_logprobs_list = self.in_progress_prompt_logprobs[req_id] |
| 209 | if logprobs is not None and (req_is_prompt_chunked or prompt_logprobs_list): |
| 210 | prompt_logprobs_list.append(logprobs) |
| 211 | if req_is_prompt_chunked: |
| 212 | # Prompt is chunked. Do not return the logprobs yet. |
| 213 | continue |
| 214 | |
| 215 | if prompt_logprobs_list: |
| 216 | # Merge the in-progress logprobs. |
| 217 | logprobs = LogprobsTensors.cat(prompt_logprobs_list) |
| 218 | prompt_logprobs_list.clear() |
| 219 | |
| 220 | if logprobs is None: |
| 221 | continue |
| 222 | |
| 223 | prompt_logprobs_dict[req_id] = logprobs |
| 224 | return prompt_logprobs_dict |
| 225 | |
| 226 | |
| 227 | @triton.jit |
| 228 | def _prompt_logprobs_token_ids_kernel( |
| 229 | prompt_logprobs_token_ids_ptr, |
| 230 | query_start_loc_ptr, |
| 231 | idx_mapping_ptr, |
| 232 | num_computed_tokens_ptr, |
| 233 | all_token_ids_ptr, |
| 234 | all_token_ids_stride, |
| 235 | BLOCK_SIZE: tl.constexpr, |
| 236 | ): |
| 237 | batch_idx = tl.program_id(0) |
| 238 | req_state_idx = tl.load(idx_mapping_ptr + batch_idx) |
| 239 | |
| 240 | query_start = tl.load(query_start_loc_ptr + batch_idx) |
| 241 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) |
| 242 | query_len = query_end - query_start |
| 243 | |
| 244 | num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx) |
| 245 | for i in range(0, query_len, BLOCK_SIZE): |
| 246 | block = i + tl.arange(0, BLOCK_SIZE) |
| 247 | mask = block < query_len |
| 248 | # NOTE(woosuk): We should shift the pos by one |
| 249 | # because the logprob is computed for the next token. |
| 250 | target_pos = num_computed_tokens + 1 + block |
| 251 | token_ids = tl.load( |
| 252 | all_token_ids_ptr + req_state_idx * all_token_ids_stride + target_pos, |
| 253 | mask=mask, |
| 254 | ) |
| 255 | tl.store( |
| 256 | prompt_logprobs_token_ids_ptr + query_start + block, token_ids, mask=mask |
| 257 | ) |
| 258 | |
| 259 | |
| 260 | def get_prompt_logprobs_token_ids( |
| 261 | num_tokens: int, |
| 262 | query_start_loc: torch.Tensor, |
| 263 | idx_mapping: torch.Tensor, |
| 264 | num_computed_tokens: torch.Tensor, |
| 265 | all_token_ids: torch.Tensor, |
| 266 | ) -> torch.Tensor: |
| 267 | token_ids = torch.empty(num_tokens, dtype=torch.int64, device=idx_mapping.device) |
| 268 | num_reqs = idx_mapping.shape[0] |
| 269 | _prompt_logprobs_token_ids_kernel[(num_reqs,)]( |
| 270 | token_ids, |
| 271 | query_start_loc, |
| 272 | idx_mapping, |
| 273 | num_computed_tokens, |
| 274 | all_token_ids, |
| 275 | all_token_ids.stride(0), |
| 276 | BLOCK_SIZE=1024, |
| 277 | ) |
| 278 | return token_ids |
| 279 | |
| 280 | |
| 281 | def compute_prompt_logprobs_with_chunking( |
| 282 | prompt_token_ids: torch.Tensor, |
| 283 | prompt_hidden_states: torch.Tensor, |
| 284 | logits_fn: Callable[[torch.Tensor], torch.Tensor], |
| 285 | num_prompt_logprobs: int, |
| 286 | logprobs_mode: LogprobsMode = "raw_logprobs", |
| 287 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 288 | # Since materializing the full prompt logits can take too much memory, |
| 289 | # we compute it in chunks. |
| 290 | token_ids = [] |
| 291 | scores = [] |
| 292 | ranks = [] |
| 293 | logits_mode = logprobs_mode in ("raw_logits", "processed_logits") |
| 294 | prompt_token_ids = prompt_token_ids.to(torch.int64) |
| 295 | for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE): |
| 296 | end_idx = start_idx + CHUNK_SIZE |
| 297 | # NOTE(woosuk): logits_fn can be slow because it involves all-gather. |
| 298 | prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx]) |
| 299 | requested_num = ( |
| 300 | prompt_logits.shape[-1] |
| 301 | if num_prompt_logprobs == -1 |
| 302 | else num_prompt_logprobs |
| 303 | ) |
| 304 | result = compute_topk_scores( |
| 305 | prompt_logits, |
| 306 | requested_num, |
| 307 | prompt_token_ids[start_idx:end_idx], |
| 308 | logits_mode=logits_mode, |
| 309 | ) |
| 310 | token_ids.append(result.logprob_token_ids) |
| 311 | scores.append(result.logprobs) |
| 312 | ranks.append(result.selected_token_ranks) |
| 313 | |
| 314 | token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0] |
| 315 | scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0] |
| 316 | ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0] |
| 317 | return token_ids, scores, ranks |