@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 271 lines 9.88 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import numpy as np
5 import torch
6
7 from vllm.sampling_params import MAX_LOGPROB_TOKEN_IDS, SamplingParams
8 from vllm.triton_utils import tl, triton
9 from vllm.v1.outputs import LogprobsTensors
10 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11
12 # Upper bound on the topk kernel's per-iteration gather width.
13 _MAX_TOPK_BLOCK = 1024
14
15
16 @triton.jit
17 def _topk_log_softmax_kernel(
18 output_ptr,
19 logits_ptr,
20 logits_stride,
21 topk_ids_ptr,
22 topk_ids_stride,
23 topk,
24 vocab_size,
25 BLOCK_SIZE: tl.constexpr,
26 TOPK_BLOCK_SIZE: tl.constexpr,
27 ):
28 req_idx = tl.program_id(0).to(tl.int64)
29 row_ptr = logits_ptr + req_idx * logits_stride
30
31 max_val = float("-inf")
32 for i in range(0, vocab_size, BLOCK_SIZE):
33 block = i + tl.arange(0, BLOCK_SIZE)
34 logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
35 max_val = tl.max(tl.maximum(logits, max_val))
36 max_val = max_val.to(tl.float32) # type: ignore
37
38 se = 0.0
39 for i in range(0, vocab_size, BLOCK_SIZE):
40 block = i + tl.arange(0, BLOCK_SIZE)
41 logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0)
42 # NOTE(woosuk): Make sure that logits and all following operations use FP32.
43 logits = logits.to(tl.float32)
44 e = tl.exp(logits - max_val)
45 e = tl.where(block < vocab_size, e, 0.0)
46 se += tl.sum(e)
47 lse = tl.log(se)
48
49 for j in range(0, topk, TOPK_BLOCK_SIZE):
50 k_offset = j + tl.arange(0, TOPK_BLOCK_SIZE)
51 k_mask = k_offset < topk
52 topk_ids = tl.load(
53 topk_ids_ptr + req_idx * topk_ids_stride + k_offset, mask=k_mask, other=0
54 )
55 logits = tl.load(row_ptr + topk_ids, mask=k_mask)
56 logits = logits.to(tl.float32)
57 o = logits - max_val - lse
58 tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
59
60
61 @triton.jit
62 def _ranks_kernel(
63 output_ptr,
64 logits_ptr,
65 logits_stride,
66 token_ids_ptr,
67 vocab_size,
68 BLOCK_SIZE: tl.constexpr,
69 ):
70 req_idx = tl.program_id(0).to(tl.int64)
71 row_ptr = logits_ptr + req_idx * logits_stride
72
73 token_id = tl.load(token_ids_ptr + req_idx)
74 x = tl.load(row_ptr + token_id)
75
76 n = 0
77 for i in range(0, vocab_size, BLOCK_SIZE):
78 block = i + tl.arange(0, BLOCK_SIZE)
79 logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
80 n += tl.sum((logits >= x).to(tl.int32))
81 tl.store(output_ptr + req_idx, n)
82
83
84 def compute_token_logprobs(
85 logits: torch.Tensor, token_ids: torch.Tensor
86 ) -> torch.Tensor:
87 # NOTE(woosuk): To save GPU memory, we do not materialize the full
88 # [batch_size, vocab_size] logprobs tensor. The kernel computes
89 # max + logsumexp per row and only emits logprobs at `token_ids`.
90 batch_size, vocab_size = logits.shape
91 token_ids = token_ids.to(torch.int64)
92 num_logprobs = token_ids.shape[1]
93 logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32)
94 # Cap the kernel's per-iteration width so very large num_logprobs requests
95 # stream the gather in bounded-size chunks, avoiding excessive mem use.
96 topk_block_size = min(triton.next_power_of_2(num_logprobs), _MAX_TOPK_BLOCK)
97 _topk_log_softmax_kernel[(batch_size,)](
98 logprobs,
99 logits,
100 logits.stride(0),
101 token_ids,
102 token_ids.stride(0),
103 num_logprobs,
104 vocab_size,
105 BLOCK_SIZE=1024, # type: ignore
106 TOPK_BLOCK_SIZE=topk_block_size,
107 )
108 return logprobs
109
110
111 def compute_topk_scores(
112 logits: torch.Tensor,
113 num_logprobs: int,
114 sampled_token_ids: torch.Tensor,
115 cu_num_logits: list[int] | torch.Tensor | None = None,
116 logprob_token_ids_state: "LogprobTokenIdsState | None" = None,
117 expanded_idx_mapping: torch.Tensor | None = None,
118 max_per_req_token_ids: int = 0,
119 logits_mode: bool = False,
120 ) -> LogprobsTensors:
121 assert num_logprobs >= 0
122 batch_size, vocab_size = logits.shape
123
124 if max_per_req_token_ids == 0:
125 # Fast path: no request asked for custom logprob_token_ids.
126 logprob_token_ids = sampled_token_ids.unsqueeze(-1)
127 if num_logprobs > 0:
128 topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
129 logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)
130 if logits_mode:
131 scores = logits.gather(-1, logprob_token_ids).to(torch.float32)
132 else:
133 scores = compute_token_logprobs(logits, logprob_token_ids)
134 else:
135 # Some requests specified logprob_token_ids. Build the [batch_size,
136 # 1 + max_cols] token_ids matrix and validity mask on the GPU via a
137 # single triton kernel, overriding the topk columns with per-request
138 # tokens where applicable.
139 assert logprob_token_ids_state is not None
140 assert expanded_idx_mapping is not None
141
142 if num_logprobs > 0:
143 topk_token_ids = torch.topk(logits, num_logprobs, dim=-1).indices
144 topk_token_ids = topk_token_ids.to(torch.int32)
145 else:
146 # This tensor just used as an int32 pointer, data not accessed.
147 topk_token_ids = logprob_token_ids_state.token_ids.gpu
148
149 num_cols = max(num_logprobs, max_per_req_token_ids)
150 logprob_token_ids = sampled_token_ids.new_zeros((batch_size, 1 + num_cols))
151 valid_mask = torch.zeros_like(logprob_token_ids, dtype=torch.bool)
152 _fill_logprob_token_ids_kernel[(batch_size,)](
153 logprob_token_ids,
154 logprob_token_ids.stride(0),
155 valid_mask,
156 valid_mask.stride(0),
157 sampled_token_ids,
158 topk_token_ids,
159 topk_token_ids.stride(0),
160 expanded_idx_mapping,
161 logprob_token_ids_state.num_token_ids.gpu,
162 logprob_token_ids_state.token_ids.gpu,
163 logprob_token_ids_state.token_ids.gpu.stride(0),
164 NUM_TOPK=num_logprobs,
165 PADDED_COLS=triton.next_power_of_2(num_cols),
166 )
167 if logits_mode:
168 scores = logits.gather(-1, logprob_token_ids).to(torch.float32)
169 else:
170 scores = compute_token_logprobs(logits, logprob_token_ids)
171 scores = scores.masked_fill(~valid_mask, float("-inf"))
172
173 token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device)
174 _ranks_kernel[(batch_size,)](
175 token_ranks,
176 logits,
177 logits.stride(0),
178 sampled_token_ids,
179 vocab_size,
180 BLOCK_SIZE=8192, # type: ignore
181 )
182 is_tensor = isinstance(cu_num_logits, torch.Tensor)
183 return LogprobsTensors(
184 logprob_token_ids=logprob_token_ids,
185 logprobs=scores,
186 selected_token_ranks=token_ranks,
187 cu_num_generated_tokens=None if is_tensor else cu_num_logits,
188 cu_num_generated_tokens_tensor=cu_num_logits if is_tensor else None,
189 )
190
191
192 @triton.jit
193 def _fill_logprob_token_ids_kernel(
194 # [batch_size, 1 + num_cols]
195 out_token_ids_ptr,
196 out_token_ids_stride,
197 # [batch_size, 1 + num_cols]
198 out_valid_mask_ptr,
199 out_valid_mask_stride,
200 sampled_token_ids_ptr, # [batch_size]
201 topk_indices_ptr, # [batch_size, NUM_TOPK] (unused when NUM_TOPK == 0)
202 topk_indices_stride,
203 expanded_idx_mapping_ptr, # [batch_size] -> req_state_idx
204 num_per_req_token_ids_ptr, # [max_num_reqs]
205 per_req_token_ids_ptr, # [max_num_reqs, MAX_LOGPROB_TOKEN_IDS]
206 per_req_token_ids_stride,
207 NUM_TOPK: tl.constexpr,
208 PADDED_COLS: tl.constexpr,
209 ):
210 batch_idx = tl.program_id(0)
211
212 # Column 0: always the sampled token, always valid.
213 sampled = tl.load(sampled_token_ids_ptr + batch_idx)
214 tl.store(out_token_ids_ptr + batch_idx * out_token_ids_stride, sampled)
215 tl.store(out_valid_mask_ptr + batch_idx * out_valid_mask_stride, 1)
216
217 req_state_idx = tl.load(expanded_idx_mapping_ptr + batch_idx)
218 num_custom = tl.load(num_per_req_token_ids_ptr + req_state_idx)
219
220 col = tl.arange(0, PADDED_COLS)
221 tid_base = out_token_ids_ptr + batch_idx * out_token_ids_stride + 1
222 mask_base = out_valid_mask_ptr + batch_idx * out_valid_mask_stride + 1
223
224 if num_custom > 0:
225 # Override topk with per-request custom tokens.
226 src = per_req_token_ids_ptr + req_state_idx * per_req_token_ids_stride
227 valid = col < num_custom
228 else:
229 # Fill with topk indices (no-op when NUM_TOPK == 0).
230 src = topk_indices_ptr + batch_idx * topk_indices_stride
231 valid = col < NUM_TOPK
232
233 tokens = tl.load(src + col, mask=valid, other=0).to(tl.int64)
234 tl.store(tid_base + col, tokens, mask=valid)
235 tl.store(mask_base + col, tl.full([PADDED_COLS], 1, tl.int1), mask=valid)
236
237
238 class LogprobTokenIdsState:
239 """Per-request override of which token ids' logprobs to return.
240
241 See `SamplingParams.logprob_token_ids`.
242 """
243
244 def __init__(self, max_num_reqs: int, device: torch.device):
245 self.max_num_reqs = max_num_reqs
246 self.num_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
247 self.token_ids = StagedWriteTensor(
248 (max_num_reqs, MAX_LOGPROB_TOKEN_IDS),
249 dtype=torch.int32,
250 device=device,
251 )
252
253 def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
254 token_ids = sampling_params.logprob_token_ids
255 if not token_ids:
256 self.num_token_ids.np[req_idx] = 0
257 return
258 n = len(token_ids)
259 if n > MAX_LOGPROB_TOKEN_IDS:
260 raise ValueError(
261 f"Too many logprob_token_ids: {n}. The max is {MAX_LOGPROB_TOKEN_IDS}."
262 )
263 self.num_token_ids.np[req_idx] = n
264 self.token_ids.stage_write(req_idx, 0, token_ids)
265
266 def apply_staged_writes(self) -> None:
267 self.num_token_ids.copy_to_uva()
268 self.token_ids.apply_write()
269
270 def max_num_token_ids(self, idx_mapping_np: np.ndarray) -> int:
271 return int(self.num_token_ids.np[idx_mapping_np].max(initial=0))