iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
382970ee6ca490aeaaaf4e32c53695b581ff61ba
vLLM/vllm/v1/worker/gpu/sample/gumbel.py
T

397 lines
13 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.triton_utils import tl, triton
# Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
# retained Philox helper for rejection sampling.
#
# Triton requires globals accessed from `@triton.jit` functions to be wrapped
# in `tl.constexpr(...)`.
_TL_RAND_MIN = tl.constexpr(4.6566127342e-10)
# Offset salt keeping the draft's Gumbel noise disjoint from the target's.
# Verification is a probability-ratio test, not a Gumbel coupling, so a proposal
# and the residual it is resampled from must not share a noise vector.
# Positions are int64 and never approach 2**30, so the streams cannot collide.
_DRAFT_NOISE_SALT = tl.constexpr(1 << 30)
@triton.jit
def _temperature_kernel(
logits_ptr,
logits_stride,
expanded_idx_mapping_ptr,
temperature_ptr,
vocab_size,
BLOCK_SIZE: tl.constexpr,
):
token_idx = tl.program_id(0).to(tl.int64)
req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
if temperature == 0.0 or temperature == 1.0:
# Early return to avoid loading logits.
return
block_idx = tl.program_id(1)
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = block < vocab_size
logits = tl.load(logits_ptr + token_idx * logits_stride + block, mask=mask)
logits = logits.to(tl.float32)
logits = logits / temperature
tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
def apply_temperature(
logits: torch.Tensor,
expanded_idx_mapping: torch.Tensor,
temperature: torch.Tensor,
) -> None:
num_tokens, vocab_size = logits.shape
BLOCK_SIZE = 8192
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
_temperature_kernel[(num_tokens, num_blocks)](
logits,
logits.stride(0),
expanded_idx_mapping,
temperature,
vocab_size,
BLOCK_SIZE=BLOCK_SIZE,
)
@triton.jit
def tl_rand64(seed, offset, includes_zero: tl.constexpr):
lo, hi, _, _ = tl.randint4x(seed, offset)
lo = lo.to(tl.uint32, bitcast=True).to(tl.uint64)
hi = hi.to(tl.uint32, bitcast=True).to(tl.uint64)
r = (hi << 32) | lo
# 1 / 2**64
scale = 5.421010862427522170037e-20
u = r.to(tl.float64) * scale
if not includes_zero:
u = tl.maximum(u, 2.2250738585072014e-308) # float64 tiny
return u
@triton.jit
def tl_rand32(seed, offset, includes_zero: tl.constexpr):
u = tl.rand(seed, offset)
if not includes_zero:
u = tl.maximum(u, _TL_RAND_MIN)
return u
@triton.jit
def _murmur3_rotl32(value, shift: tl.constexpr):
return (value << shift) | (value >> (32 - shift))
@triton.jit
def _murmur3_mix(h, key):
key *= 0xCC9E2D51
key = _murmur3_rotl32(key, 15)
key *= 0x1B873593
h ^= key
h = _murmur3_rotl32(h, 13)
return h * 5 + 0xE6546B64
@triton.jit
def _murmur3_fmix32(h):
h ^= h >> 16
h *= 0x85EBCA6B
h ^= h >> 13
h *= 0xC2B2AE35
return h ^ (h >> 16)
@triton.jit
def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0):
seed = seed.to(tl.int64)
pos = pos.to(tl.int64)
offset = offset.to(tl.uint32)
# Keep the request-wide prefix scalar until the token offset is mixed in.
h = (seed ^ seed).to(tl.uint32)
h ^= domain
h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32))
h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32))
h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32))
h = _murmur3_mix(h, offset)
return _murmur3_fmix32(h ^ 16)
@triton.jit
def murmur3_uniform32(seed, pos, offset):
random32 = murmur3_hash32(seed, pos, offset)
# Split the uint32 before converting to fp32 so backends without a native
# uint32-to-float conversion can still use all 32 source bits. Both 16-bit
# halves convert exactly; their sum is the correctly rounded fp32 value of
# (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps
# the full 32-bit source resolution instead of being truncated to 24 bits.
hi16 = (random32 >> 16).to(tl.int32)
lo16 = (random32 & 0xFFFF).to(tl.int32)
return (
hi16.to(tl.float32) * 1.52587890625e-05
+ (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10
)
@triton.jit
def _uniform64_from_random53(random53):
uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16
# The largest midpoint rounds to 1.0 in fp64; keep the uniform open without
# relying on a near-one literal that Triton may materialize in fp32.
return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform)
@triton.jit
def murmur3_uniform64(seed, pos, offset):
lo = murmur3_hash32(seed, pos, offset).to(tl.uint64)
hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64)
random53 = ((hi << 32) | lo) >> 11
return _uniform64_from_random53(random53)
@triton.jit
def _log1p_neg_stable(value):
# Preserve precision for the positive Gumbel tail without relying on a
# backend-specific libdevice log1p. The degree-8 series has absolute error
# below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for
# the part of the distribution that can win the argmax.
polynomial = 1.0 / 8.0
polynomial = 1.0 / 7.0 + value * polynomial
polynomial = 1.0 / 6.0 + value * polynomial
polynomial = 1.0 / 5.0 + value * polynomial
polynomial = 1.0 / 4.0 + value * polynomial
polynomial = 1.0 / 3.0 + value * polynomial
polynomial = 1.0 / 2.0 + value * polynomial
polynomial = 1.0 + value * polynomial
series = -value * polynomial
direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08))
return tl.where(value < 0.25, series, direct)
@triton.jit
def gumbel_noised_argmax(
logits,
keys,
mask,
seed,
pos,
temp,
IS_DRAFTING: tl.constexpr,
USE_FP64: tl.constexpr,
APPLY_TEMPERATURE: tl.constexpr = True,
):
"""Argmax of logits under Gumbel-max sampling, or plain argmax at temp 0.
`keys` indexes the noise, so the same token draws the same noise wherever it
appears; `pos` and `seed` place the draw in the request's stream, which is
what lets a draft and its verification agree.
"""
if temp != 0.0 and APPLY_TEMPERATURE:
# Match the behavior of _temperature_kernel: if that kernel uses
# tl.div_rn, this must too.
logits = logits / temp
# fp32 is the default reduction dtype; fp64 is ~1/32-1/64x the throughput
# on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.
if USE_FP64:
logits = logits.to(tl.float64)
if temp != 0.0:
if IS_DRAFTING:
pos = pos + _DRAFT_NOISE_SALT
if USE_FP64:
u = murmur3_uniform64(seed, pos, keys)
gumbel_noise = -tl.log(-tl.log(u))
else:
u = murmur3_uniform32(seed, pos, keys)
# Draw the large-noise tail (which decides the argmax winner) from
# u -> 0, where fp32 has fine resolution. Avoid backend-specific
# log1p while preserving precision in the winning tail.
gumbel_noise = -tl.log(-_log1p_neg_stable(u))
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
return tl.max(logits, axis=0, return_indices=True)
@triton.jit
def gumbel_block_argmax(
logits,
block,
mask,
token_idx,
expanded_idx_mapping_ptr,
temp_ptr,
seeds_ptr,
pos_ptr,
# [max_num_reqs, num_cols, vocab_size]
logits_cache_ptr,
logits_cache_stride_0,
logits_cache_stride_1,
logits_cache_col_ptr,
vocab_size,
IS_DRAFTING: tl.constexpr,
APPLY_TEMPERATURE: tl.constexpr,
USE_FP64: tl.constexpr,
PER_TOKEN_COL: tl.constexpr = False,
):
req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)
is_valid_req = req_state_idx >= 0
temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to(
tl.float32
)
if logits_cache_ptr is not None:
# Store the logits *before* temperature. Dividing first would produce a
# value that is generally not representable in the cache's dtype, forcing
# it to be fp32. Consumers (the rejection sampler) divide by the same
# temperature on load, which reproduces the value used below bitwise.
if PER_TOKEN_COL:
col = tl.load(logits_cache_col_ptr + token_idx)
else:
col = tl.load(logits_cache_col_ptr)
tl.store(
logits_cache_ptr
+ req_state_idx * logits_cache_stride_0
+ col * logits_cache_stride_1
+ block,
logits,
mask=mask & is_valid_req,
)
seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
pos = tl.load(pos_ptr + token_idx)
return gumbel_noised_argmax(
logits,
block,
mask,
seed,
pos,
temp,
IS_DRAFTING=IS_DRAFTING,
USE_FP64=USE_FP64,
APPLY_TEMPERATURE=APPLY_TEMPERATURE,
)
@triton.jit
def _gumbel_sample_kernel(
local_argmax_ptr,
local_argmax_stride,
local_max_ptr,
local_max_stride,
# [max_num_reqs, num_cols, vocab_size]
logits_cache_ptr,
logits_cache_stride_0,
logits_cache_stride_1,
logits_cache_col_ptr,
logits_ptr,
logits_stride,
expanded_idx_mapping_ptr,
seeds_ptr,
pos_ptr,
temp_ptr,
vocab_size,
BLOCK_SIZE: tl.constexpr,
IS_DRAFTING: tl.constexpr,
APPLY_TEMPERATURE: tl.constexpr,
USE_FP64: tl.constexpr,
PER_TOKEN_COL: tl.constexpr,
):
token_idx = tl.program_id(0).to(tl.int64)
block_idx = tl.program_id(1)
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = block < vocab_size
logits = tl.load(
logits_ptr + token_idx * logits_stride + block,
mask=mask,
other=float("-inf"),
)
logits = logits.to(tl.float32)
value, idx = gumbel_block_argmax(
logits,
block,
mask,
token_idx,
expanded_idx_mapping_ptr,
temp_ptr,
seeds_ptr,
pos_ptr,
logits_cache_ptr,
logits_cache_stride_0,
logits_cache_stride_1,
logits_cache_col_ptr,
vocab_size,
IS_DRAFTING=IS_DRAFTING,
APPLY_TEMPERATURE=APPLY_TEMPERATURE,
USE_FP64=USE_FP64,
PER_TOKEN_COL=PER_TOKEN_COL,
)
token_id = block_idx * BLOCK_SIZE + idx
tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
def gumbel_sample(
logits: torch.Tensor, # [num_tokens, vocab_size]
expanded_idx_mapping: torch.Tensor, # [num_tokens]
temperature: torch.Tensor, # [max_num_reqs]
seed: torch.Tensor, # [max_num_reqs]
pos: torch.Tensor, # [num_tokens]
apply_temperature: bool,
is_drafting: bool,
logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size]
logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens]
use_fp64: bool = False,
) -> torch.Tensor:
# Enforce contiguity on non-strided input tensors
expanded_idx_mapping = expanded_idx_mapping.contiguous()
pos = pos.contiguous()
if logits_cache_col is not None:
logits_cache_col = logits_cache_col.contiguous()
num_tokens, vocab_size = logits.shape
if logits_cache is not None:
assert logits_cache.size(-1) >= vocab_size, (
f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "
f"than the sampled logits ({vocab_size}). Cached logits would be "
"truncated."
)
BLOCK_SIZE = 1024
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
local_max_dtype = torch.float64 if use_fp64 else torch.float32
local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype)
per_token_col = logits_cache_col is not None and logits_cache_col.dim() > 0
_gumbel_sample_kernel[(num_tokens, num_blocks)](
local_argmax,
local_argmax.stride(0),
local_max,
local_max.stride(0),
logits_cache,
logits_cache.stride(0) if logits_cache is not None else 0,
logits_cache.stride(1) if logits_cache is not None else 0,
logits_cache_col,
logits,
logits.stride(0),
expanded_idx_mapping,
seed,
pos,
temperature,
vocab_size,
BLOCK_SIZE=BLOCK_SIZE,
IS_DRAFTING=is_drafting,
APPLY_TEMPERATURE=apply_temperature,
USE_FP64=use_fp64,
PER_TOKEN_COL=per_token_col,
)
# NOTE(woosuk): Use int64 for later indexing.
max_block_idx = local_max.argmax(dim=-1, keepdim=True)
sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
return sampled
Reference in New Issue View Git Blame Copy Permalink