@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 409 lines 13.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.triton_utils import tl, triton
6
7 # Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
8 # retained Philox helper for rejection sampling.
9 #
10 # Triton requires globals accessed from `@triton.jit` functions to be wrapped
11 # in `tl.constexpr(...)`.
12 _TL_RAND_MIN = tl.constexpr(4.6566127342e-10)
13
14 # Offset salt keeping the draft's Gumbel noise disjoint from the target's.
15 # Verification is a probability-ratio test, not a Gumbel coupling, so a proposal
16 # and the residual it is resampled from must not share a noise vector.
17 # Positions are int64 and never approach 2**30, so the streams cannot collide.
18 _DRAFT_NOISE_SALT = tl.constexpr(1 << 30)
19
20
21 @triton.jit
22 def _temperature_kernel(
23 logits_ptr,
24 logits_stride,
25 expanded_idx_mapping_ptr,
26 temperature_ptr,
27 vocab_size,
28 BLOCK_SIZE: tl.constexpr,
29 ):
30 token_idx = tl.program_id(0).to(tl.int64)
31 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
32 temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
33 if temperature == 0.0 or temperature == 1.0:
34 # Early return to avoid loading logits.
35 return
36
37 block_idx = tl.program_id(1)
38 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
39 mask = block < vocab_size
40
41 logits = tl.load(logits_ptr + token_idx * logits_stride + block, mask=mask)
42 logits = logits.to(tl.float32)
43 logits = logits / temperature
44 tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
45
46
47 def apply_temperature(
48 logits: torch.Tensor,
49 expanded_idx_mapping: torch.Tensor,
50 temperature: torch.Tensor,
51 ) -> None:
52 num_tokens, vocab_size = logits.shape
53 BLOCK_SIZE = 8192
54 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
55 _temperature_kernel[(num_tokens, num_blocks)](
56 logits,
57 logits.stride(0),
58 expanded_idx_mapping,
59 temperature,
60 vocab_size,
61 BLOCK_SIZE=BLOCK_SIZE,
62 )
63
64
65 @triton.jit
66 def tl_rand32(seed, offset, includes_zero: tl.constexpr):
67 u = tl.rand(seed, offset)
68 if not includes_zero:
69 u = tl.maximum(u, _TL_RAND_MIN)
70 return u
71
72
73 @triton.jit
74 def _murmur3_rotl32(value, shift: tl.constexpr):
75 return (value << shift) | (value >> (32 - shift))
76
77
78 @triton.jit
79 def _murmur3_mix(h, key):
80 key *= 0xCC9E2D51
81 key = _murmur3_rotl32(key, 15)
82 key *= 0x1B873593
83 h ^= key
84 h = _murmur3_rotl32(h, 13)
85 return h * 5 + 0xE6546B64
86
87
88 @triton.jit
89 def _murmur3_fmix32(h):
90 h ^= h >> 16
91 h *= 0x85EBCA6B
92 h ^= h >> 13
93 h *= 0xC2B2AE35
94 return h ^ (h >> 16)
95
96
97 @triton.jit
98 def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0):
99 seed = seed.to(tl.int64)
100 pos = pos.to(tl.int64)
101 offset = offset.to(tl.uint32)
102 # Keep the request-wide prefix scalar until the token offset is mixed in.
103 h = (seed ^ seed).to(tl.uint32)
104 h ^= domain
105 h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32))
106 h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32))
107 h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32))
108 h = _murmur3_mix(h, offset)
109 return _murmur3_fmix32(h ^ 16)
110
111
112 @triton.jit
113 def murmur3_uniform32(seed, pos, offset):
114 random32 = murmur3_hash32(seed, pos, offset)
115 # Split the uint32 before converting to fp32 so backends without a native
116 # uint32-to-float conversion can still use all 32 source bits. Both 16-bit
117 # halves convert exactly; their sum is the correctly rounded fp32 value of
118 # (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps
119 # the full 32-bit source resolution instead of being truncated to 24 bits.
120 hi16 = (random32 >> 16).to(tl.int32)
121 lo16 = (random32 & 0xFFFF).to(tl.int32)
122 return (
123 hi16.to(tl.float32) * 1.52587890625e-05
124 + (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10
125 )
126
127
128 @triton.jit
129 def _uniform64_from_random53(random53):
130 uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16
131 # The largest midpoint rounds to 1.0 in fp64; keep the uniform open without
132 # relying on a near-one literal that Triton may materialize in fp32.
133 return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform)
134
135
136 @triton.jit
137 def murmur3_uniform64(seed, pos, offset):
138 lo = murmur3_hash32(seed, pos, offset).to(tl.uint64)
139 hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64)
140 random53 = ((hi << 32) | lo) >> 11
141 return _uniform64_from_random53(random53)
142
143
144 @triton.jit
145 def _log1p_neg_stable(value):
146 # Preserve precision for the positive Gumbel tail without relying on a
147 # backend-specific libdevice log1p. The degree-8 series has absolute error
148 # below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for
149 # the part of the distribution that can win the argmax.
150 polynomial = 1.0 / 8.0
151 polynomial = 1.0 / 7.0 + value * polynomial
152 polynomial = 1.0 / 6.0 + value * polynomial
153 polynomial = 1.0 / 5.0 + value * polynomial
154 polynomial = 1.0 / 4.0 + value * polynomial
155 polynomial = 1.0 / 3.0 + value * polynomial
156 polynomial = 1.0 / 2.0 + value * polynomial
157 polynomial = 1.0 + value * polynomial
158 series = -value * polynomial
159
160 direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08))
161 return tl.where(value < 0.25, series, direct)
162
163
164 @triton.jit
165 def gumbel_noised_argmax(
166 logits,
167 keys,
168 mask,
169 seed,
170 pos,
171 temp,
172 IS_DRAFTING: tl.constexpr,
173 USE_FP64: tl.constexpr,
174 APPLY_TEMPERATURE: tl.constexpr = True,
175 ):
176 """Argmax of logits under Gumbel-max sampling, or plain argmax at temp 0.
177
178 `keys` indexes the noise, so the same token draws the same noise wherever it
179 appears; `pos` and `seed` place the draw in the request's stream, which is
180 what lets a draft and its verification agree.
181 """
182 if temp != 0.0 and APPLY_TEMPERATURE:
183 # Match the behavior of _temperature_kernel: if that kernel uses
184 # tl.div_rn, this must too.
185 logits = logits / temp
186
187 # fp32 is the default reduction dtype; fp64 is ~1/32-1/64x the throughput
188 # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.
189 if USE_FP64:
190 logits = logits.to(tl.float64)
191 if temp != 0.0:
192 if IS_DRAFTING:
193 pos = pos + _DRAFT_NOISE_SALT
194 if USE_FP64:
195 u = murmur3_uniform64(seed, pos, keys)
196 gumbel_noise = -tl.log(-tl.log(u))
197 else:
198 u = murmur3_uniform32(seed, pos, keys)
199 # Draw the large-noise tail (which decides the argmax winner) from
200 # u -> 0, where fp32 has fine resolution. Avoid backend-specific
201 # log1p while preserving precision in the winning tail.
202 gumbel_noise = -tl.log(-_log1p_neg_stable(u))
203 logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
204
205 return tl.max(logits, axis=0, return_indices=True)
206
207
208 @triton.jit
209 def gumbel_block_argmax(
210 logits,
211 block,
212 mask,
213 token_idx,
214 expanded_idx_mapping_ptr,
215 temp_ptr,
216 seeds_ptr,
217 pos_ptr,
218 # [max_num_reqs, num_cols, vocab_size]
219 logits_cache_ptr,
220 logits_cache_stride_0,
221 logits_cache_stride_1,
222 logits_cache_col_ptr,
223 logits_cache_source_ptr,
224 logits_cache_source_stride,
225 vocab_size,
226 IS_DRAFTING: tl.constexpr,
227 APPLY_TEMPERATURE: tl.constexpr,
228 USE_FP64: tl.constexpr,
229 PER_TOKEN_COL: tl.constexpr = False,
230 ):
231 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)
232 is_valid_req = req_state_idx >= 0
233 temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to(
234 tl.float32
235 )
236 if logits_cache_ptr is not None:
237 # Store the logits *before* temperature. Dividing first would produce a
238 # value that is generally not representable in the cache's dtype, forcing
239 # it to be fp32. Consumers (the rejection sampler) divide by the same
240 # temperature on load, which reproduces the value used below bitwise.
241 if PER_TOKEN_COL:
242 col = tl.load(logits_cache_col_ptr + token_idx)
243 else:
244 col = tl.load(logits_cache_col_ptr)
245 cached_logits = tl.load(
246 logits_cache_source_ptr + token_idx * logits_cache_source_stride + block,
247 mask=mask,
248 )
249 tl.store(
250 logits_cache_ptr
251 + req_state_idx * logits_cache_stride_0
252 + col * logits_cache_stride_1
253 + block,
254 cached_logits,
255 mask=mask & is_valid_req,
256 )
257
258 seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
259 pos = tl.load(pos_ptr + token_idx)
260 return gumbel_noised_argmax(
261 logits,
262 block,
263 mask,
264 seed,
265 pos,
266 temp,
267 IS_DRAFTING=IS_DRAFTING,
268 USE_FP64=USE_FP64,
269 APPLY_TEMPERATURE=APPLY_TEMPERATURE,
270 )
271
272
273 @triton.jit
274 def _gumbel_sample_kernel(
275 local_argmax_ptr,
276 local_argmax_stride,
277 local_max_ptr,
278 local_max_stride,
279 # [max_num_reqs, num_cols, vocab_size]
280 logits_cache_ptr,
281 logits_cache_stride_0,
282 logits_cache_stride_1,
283 logits_cache_col_ptr,
284 logits_cache_source_ptr,
285 logits_cache_source_stride,
286 logits_ptr,
287 logits_stride,
288 expanded_idx_mapping_ptr,
289 seeds_ptr,
290 pos_ptr,
291 temp_ptr,
292 vocab_size,
293 BLOCK_SIZE: tl.constexpr,
294 IS_DRAFTING: tl.constexpr,
295 APPLY_TEMPERATURE: tl.constexpr,
296 USE_FP64: tl.constexpr,
297 PER_TOKEN_COL: tl.constexpr,
298 ):
299 token_idx = tl.program_id(0).to(tl.int64)
300 block_idx = tl.program_id(1)
301 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
302 mask = block < vocab_size
303 logits = tl.load(
304 logits_ptr + token_idx * logits_stride + block,
305 mask=mask,
306 other=float("-inf"),
307 )
308 logits = logits.to(tl.float32)
309
310 value, idx = gumbel_block_argmax(
311 logits,
312 block,
313 mask,
314 token_idx,
315 expanded_idx_mapping_ptr,
316 temp_ptr,
317 seeds_ptr,
318 pos_ptr,
319 logits_cache_ptr,
320 logits_cache_stride_0,
321 logits_cache_stride_1,
322 logits_cache_col_ptr,
323 logits_cache_source_ptr,
324 logits_cache_source_stride,
325 vocab_size,
326 IS_DRAFTING=IS_DRAFTING,
327 APPLY_TEMPERATURE=APPLY_TEMPERATURE,
328 USE_FP64=USE_FP64,
329 PER_TOKEN_COL=PER_TOKEN_COL,
330 )
331 token_id = block_idx * BLOCK_SIZE + idx
332 tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
333 tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
334
335
336 def gumbel_sample(
337 logits: torch.Tensor, # [num_tokens, vocab_size]
338 expanded_idx_mapping: torch.Tensor, # [num_tokens]
339 temperature: torch.Tensor, # [max_num_reqs]
340 seed: torch.Tensor, # [max_num_reqs]
341 pos: torch.Tensor, # [num_tokens]
342 apply_temperature: bool,
343 is_drafting: bool,
344 logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size]
345 logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens]
346 use_fp64: bool = False,
347 logits_cache_source: torch.Tensor | None = None,
348 ) -> torch.Tensor:
349 # Enforce contiguity on non-strided input tensors
350 expanded_idx_mapping = expanded_idx_mapping.contiguous()
351 pos = pos.contiguous()
352 if logits_cache_col is not None:
353 logits_cache_col = logits_cache_col.contiguous()
354 num_tokens, vocab_size = logits.shape
355 if logits_cache is not None:
356 if logits_cache_source is None:
357 logits_cache_source = logits
358 assert logits_cache_source.shape == logits.shape, (
359 "logits cache source must match sampled logits shape"
360 )
361 assert logits_cache_source.device == logits.device, (
362 "logits cache source must be on the sampled logits device"
363 )
364 assert logits_cache_source.dtype == logits_cache.dtype, (
365 "logits cache source and destination must have the same dtype"
366 )
367 assert logits_cache.size(-1) >= vocab_size, (
368 f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "
369 f"than the sampled logits ({vocab_size}). Cached logits would be "
370 "truncated."
371 )
372 elif logits_cache_source is not None:
373 raise ValueError("logits_cache_source requires logits_cache")
374 if logits_cache_source is not None and logits_cache_source.stride(-1) != 1:
375 logits_cache_source = logits_cache_source.contiguous()
376 BLOCK_SIZE = 1024
377 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
378 local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
379 local_max_dtype = torch.float64 if use_fp64 else torch.float32
380 local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype)
381 per_token_col = logits_cache_col is not None and logits_cache_col.dim() > 0
382 _gumbel_sample_kernel[(num_tokens, num_blocks)](
383 local_argmax,
384 local_argmax.stride(0),
385 local_max,
386 local_max.stride(0),
387 logits_cache,
388 logits_cache.stride(0) if logits_cache is not None else 0,
389 logits_cache.stride(1) if logits_cache is not None else 0,
390 logits_cache_col,
391 logits_cache_source,
392 logits_cache_source.stride(0) if logits_cache_source is not None else 0,
393 logits,
394 logits.stride(0),
395 expanded_idx_mapping,
396 seed,
397 pos,
398 temperature,
399 vocab_size,
400 BLOCK_SIZE=BLOCK_SIZE,
401 IS_DRAFTING=is_drafting,
402 APPLY_TEMPERATURE=apply_temperature,
403 USE_FP64=use_fp64,
404 PER_TOKEN_COL=per_token_col,
405 )
406 # NOTE(woosuk): Use int64 for later indexing.
407 max_block_idx = local_max.argmax(dim=-1, keepdim=True)
408 sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
409 return sampled