# SPDX-License-Identifier: MIT AND Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # # Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # TokenSpeed-specific top-p-only rejection/repair layout. from __future__ import annotations import torch from tokenspeed_kernel._triton import tl, triton _TOP_P_PARALLEL_BLOCK_SIZE = 1024 _TOP_P_PARALLEL_NUM_ATTEMPTS = 3 _TOP_P_REPAIR_NUM_ATTEMPTS = 8 @triton.jit def _top_p_parallel_stage1_kernel( logits_ptr, req_pool_indices_ptr, temperature_pool_ptr, seed_pool_ptr, offsets_pool_ptr, local_max_ptr, local_sum_ptr, local_argmax_ptr, local_scores_ptr, local_logits_ptr, local_ids_ptr, logits_row_stride: tl.constexpr, vocab_size: tl.constexpr, num_blocks: tl.constexpr, BLOCK_SIZE: tl.constexpr, NUM_ATTEMPTS: tl.constexpr, NUM_TOKENS_PER_REQ: tl.constexpr, ): row = tl.program_id(0) block = tl.program_id(1) req_row = row // NUM_TOKENS_PER_REQ spec_pos = row - req_row * NUM_TOKENS_PER_REQ pool_idx = tl.load(req_pool_indices_ptr + req_row) token_offsets = tl.arange(0, BLOCK_SIZE) cols = block * BLOCK_SIZE + token_offsets mask = cols < vocab_size temperature = tl.maximum( tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20 ) seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64) offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos vals = tl.load( logits_ptr + row * logits_row_stride + cols, mask=mask, other=float("-inf"), ).to(tl.float32) vals = vals / temperature block_max = tl.max(vals, axis=0) safe_block_max = tl.where(block_max > -float("inf"), block_max, 0.0) block_sum = tl.sum( tl.where(mask & (vals > -float("inf")), tl.exp(vals - safe_block_max), 0.0), axis=0, ) block_argmax = tl.min(tl.where(vals == block_max, cols, 2147483647), axis=0) block_base = row * num_blocks + block tl.store(local_max_ptr + block_base, block_max) tl.store(local_sum_ptr + block_base, block_sum) tl.store(local_argmax_ptr + block_base, block_argmax) for attempt in tl.static_range(0, NUM_ATTEMPTS): attempt_seed = tl.randint(seed, offset + attempt) uniform = tl.maximum(tl.rand(attempt_seed, cols), 1.0e-7) gumbel = -tl.log(-tl.log(uniform)) scores = tl.where(mask, vals + gumbel, float("-inf")) best_score = tl.max(scores, axis=0) best_id = tl.min(tl.where(scores == best_score, cols, 2147483647), axis=0) best_logit = tl.max(tl.where(cols == best_id, vals, float("-inf")), axis=0) out_offset = block_base * NUM_ATTEMPTS + attempt tl.store(local_scores_ptr + out_offset, best_score) tl.store(local_logits_ptr + out_offset, best_logit) tl.store(local_ids_ptr + out_offset, best_id) @triton.jit def _top_p_parallel_stage2_kernel( local_max_ptr, local_sum_ptr, local_argmax_ptr, local_scores_ptr, local_logits_ptr, local_ids_ptr, row_max_ptr, row_total_ptr, row_argmax_ptr, row_candidate_logits_ptr, row_candidate_ids_ptr, num_blocks: tl.constexpr, NUM_BLOCKS_PAD: tl.constexpr, NUM_ATTEMPTS: tl.constexpr, ): row = tl.program_id(0) block_offsets = tl.arange(0, NUM_BLOCKS_PAD) block_mask = block_offsets < num_blocks block_base = row * num_blocks + block_offsets local_max = tl.load( local_max_ptr + block_base, mask=block_mask, other=-float("inf") ) local_sum = tl.load(local_sum_ptr + block_base, mask=block_mask, other=0.0) row_max = tl.max(local_max, axis=0) safe_row_max = tl.where(row_max > -float("inf"), row_max, 0.0) total = tl.sum( tl.where(block_mask, local_sum * tl.exp(local_max - safe_row_max), 0.0), axis=0, ) local_argmax = tl.load( local_argmax_ptr + block_base, mask=block_mask, other=2147483647 ) row_argmax = tl.min( tl.where((local_max == row_max) & block_mask, local_argmax, 2147483647), axis=0, ) tl.store(row_max_ptr + row, row_max) tl.store(row_total_ptr + row, total) tl.store(row_argmax_ptr + row, row_argmax) for attempt in tl.static_range(0, NUM_ATTEMPTS): candidate_base = (row * num_blocks + block_offsets) * NUM_ATTEMPTS + attempt scores = tl.load( local_scores_ptr + candidate_base, mask=block_mask, other=-float("inf") ) ids = tl.load(local_ids_ptr + candidate_base, mask=block_mask, other=2147483647) logits = tl.load( local_logits_ptr + candidate_base, mask=block_mask, other=-float("inf") ) best_score = tl.max(scores, axis=0) best_id = tl.min( tl.where((scores == best_score) & block_mask, ids, 2147483647), axis=0, ) best_logit = tl.max(tl.where(ids == best_id, logits, -float("inf")), axis=0) row_candidate_offset = row * NUM_ATTEMPTS + attempt tl.store(row_candidate_ids_ptr + row_candidate_offset, best_id) tl.store(row_candidate_logits_ptr + row_candidate_offset, best_logit) @triton.jit def _top_p_parallel_stage3_kernel( logits_ptr, req_pool_indices_ptr, temperature_pool_ptr, row_max_ptr, row_candidate_logits_ptr, row_candidate_ids_ptr, partial_before_ptr, logits_row_stride: tl.constexpr, vocab_size: tl.constexpr, num_blocks: tl.constexpr, BLOCK_SIZE: tl.constexpr, NUM_ATTEMPTS: tl.constexpr, NUM_TOKENS_PER_REQ: tl.constexpr, ): row = tl.program_id(0) block = tl.program_id(1) req_row = row // NUM_TOKENS_PER_REQ pool_idx = tl.load(req_pool_indices_ptr + req_row) token_offsets = tl.arange(0, BLOCK_SIZE) cols = block * BLOCK_SIZE + token_offsets mask = cols < vocab_size row_max = tl.load(row_max_ptr + row) temperature = tl.maximum( tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20 ) vals = tl.load( logits_ptr + row * logits_row_stride + cols, mask=mask, other=float("-inf"), ).to(tl.float32) vals = vals / temperature weights = tl.exp(vals - row_max) for attempt in tl.static_range(0, NUM_ATTEMPTS): row_candidate_offset = row * NUM_ATTEMPTS + attempt candidate_logit = tl.load(row_candidate_logits_ptr + row_candidate_offset) candidate_id = tl.load(row_candidate_ids_ptr + row_candidate_offset) before_mask = (vals > candidate_logit) | ( (vals == candidate_logit) & (cols < candidate_id) ) before = tl.sum(tl.where(mask & before_mask, weights, 0.0), axis=0) out_offset = (row * num_blocks + block) * NUM_ATTEMPTS + attempt tl.store(partial_before_ptr + out_offset, before) @triton.jit def _top_p_parallel_stage4_kernel( top_p_pool_ptr, req_pool_indices_ptr, row_total_ptr, row_argmax_ptr, row_candidate_ids_ptr, partial_before_ptr, accepted_ptr, out_ptr, num_blocks: tl.constexpr, NUM_BLOCKS_PAD: tl.constexpr, NUM_ATTEMPTS: tl.constexpr, NUM_TOKENS_PER_REQ: tl.constexpr, ): row = tl.program_id(0) req_row = row // NUM_TOKENS_PER_REQ pool_idx = tl.load(req_pool_indices_ptr + req_row) top_p = tl.load(top_p_pool_ptr + pool_idx).to(tl.float32) target_mass = top_p * tl.load(row_total_ptr + row) block_offsets = tl.arange(0, NUM_BLOCKS_PAD) block_mask = block_offsets < num_blocks token = tl.load(row_argmax_ptr + row) found = tl.full((), 0, tl.int32) for attempt in tl.static_range(0, NUM_ATTEMPTS): before_base = (row * num_blocks + block_offsets) * NUM_ATTEMPTS + attempt before = tl.sum( tl.load(partial_before_ptr + before_base, mask=block_mask, other=0.0), axis=0, ) accepted = before < target_mass candidate_id = tl.load(row_candidate_ids_ptr + row * NUM_ATTEMPTS + attempt) take = (found == 0) & accepted token = tl.where(take, candidate_id, token) found = tl.where(take, 1, found) tl.store(accepted_ptr + row, found) tl.store(out_ptr + row, token) @triton.jit def _top_p_parallel_repair_kernel( logits_ptr, req_pool_indices_ptr, temperature_pool_ptr, top_p_pool_ptr, seed_pool_ptr, offsets_pool_ptr, row_max_ptr, row_total_ptr, row_argmax_ptr, accepted_ptr, out_ptr, logits_row_stride: tl.constexpr, vocab_size: tl.constexpr, BLOCK_SIZE: tl.constexpr, START_ATTEMPT: tl.constexpr, NUM_ATTEMPTS_TOTAL: tl.constexpr, NUM_TOKENS_PER_REQ: tl.constexpr, ): row = tl.program_id(0) accepted_found = tl.load(accepted_ptr + row) accepted_token = tl.load(out_ptr + row) req_row = row // NUM_TOKENS_PER_REQ spec_pos = row - req_row * NUM_TOKENS_PER_REQ pool_idx = tl.load(req_pool_indices_ptr + req_row) token_offsets = tl.arange(0, BLOCK_SIZE) temperature = tl.maximum( tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20 ) top_p = tl.load(top_p_pool_ptr + pool_idx).to(tl.float32) seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64) offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos row_max = tl.load(row_max_ptr + row) total = tl.load(row_total_ptr + row) target_mass = top_p * total row_argmax = tl.load(row_argmax_ptr + row) attempt = tl.full((), START_ATTEMPT, tl.int32) while (attempt < NUM_ATTEMPTS_TOTAL) & (accepted_found == 0): attempt_seed = tl.randint(seed, offset + attempt) best_score = tl.full((), float("-inf"), tl.float32) best_id = tl.full((), 2147483647, tl.int32) best_logit = tl.full((), float("-inf"), tl.float32) for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3): cols = start + token_offsets mask = cols < vocab_size vals = tl.load( logits_ptr + row * logits_row_stride + cols, mask=mask, other=float("-inf"), ).to(tl.float32) vals = vals / temperature uniform = tl.maximum(tl.rand(attempt_seed, cols), 1.0e-7) gumbel = -tl.log(-tl.log(uniform)) scores = tl.where(mask, vals + gumbel, float("-inf")) block_score = tl.max(scores, axis=0) block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0) block_logit = tl.max( tl.where(cols == block_id, vals, float("-inf")), axis=0 ) better = (block_score > best_score) | ( (block_score == best_score) & (block_id < best_id) ) best_score = tl.where(better, block_score, best_score) best_id = tl.where(better, block_id, best_id) best_logit = tl.where(better, block_logit, best_logit) before = tl.full((), 0.0, tl.float32) for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3): cols = start + token_offsets mask = cols < vocab_size vals = tl.load( logits_ptr + row * logits_row_stride + cols, mask=mask, other=float("-inf"), ).to(tl.float32) vals = vals / temperature weights = tl.exp(vals - row_max) before_mask = (vals > best_logit) | ( (vals == best_logit) & (cols < best_id) ) before += tl.sum(tl.where(mask & before_mask, weights, 0.0), axis=0) accepted = before < target_mass accepted_token = tl.where(accepted, best_id, accepted_token) accepted_found = tl.where(accepted, 1, accepted_found) attempt += 1 token = tl.where(accepted_found != 0, accepted_token, row_argmax) tl.store(out_ptr + row, token) tl.store(accepted_ptr + row, accepted_found) def gumbel_sample_top_p_parallel_from_pools( logits: torch.Tensor, req_pool_indices: torch.Tensor, temperature_pool: torch.Tensor, top_p_pool: torch.Tensor, seed_pool: torch.Tensor, offsets_pool: torch.Tensor, local_max: torch.Tensor, local_sum: torch.Tensor, local_argmax: torch.Tensor, local_scores: torch.Tensor, local_logits: torch.Tensor, local_ids: torch.Tensor, row_max: torch.Tensor, row_total: torch.Tensor, row_argmax: torch.Tensor, row_candidate_logits: torch.Tensor, row_candidate_ids: torch.Tensor, accepted: torch.Tensor, out: torch.Tensor, *, block_size: int = _TOP_P_PARALLEL_BLOCK_SIZE, num_attempts: int = _TOP_P_PARALLEL_NUM_ATTEMPTS, num_tokens_per_req: int = 1, ) -> torch.Tensor: """Block-parallel top-p-only Gumbel sampler.""" if logits.ndim != 2: raise ValueError("gumbel_sample_top_p_parallel_from_pools expects 2D logits") if logits.device.type != "cuda": raise ValueError("gumbel_sample_top_p_parallel_from_pools requires CUDA logits") if logits.stride(-1) != 1: raise ValueError( "gumbel_sample_top_p_parallel_from_pools requires stride-1 vocab dimension, " f"got stride={logits.stride()}" ) rows, vocab_size = logits.shape if vocab_size <= 0: raise ValueError( "gumbel_sample_top_p_parallel_from_pools requires non-empty vocab" ) if num_tokens_per_req <= 0: raise ValueError("num_tokens_per_req must be positive") if rows % num_tokens_per_req != 0: raise ValueError( "logits rows must be divisible by num_tokens_per_req, " f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}" ) request_rows = rows // num_tokens_per_req if req_pool_indices.shape[0] != request_rows: raise ValueError( "req_pool_indices length must match request rows for parallel top-p sample, " f"got {req_pool_indices.shape[0]} and {request_rows}" ) if num_attempts <= 0: raise ValueError("num_attempts must be positive") num_blocks = triton.cdiv(vocab_size, block_size) num_blocks_pad = triton.next_power_of_2(num_blocks) for name, tensor, dtype in ( ("req_pool_indices", req_pool_indices, torch.int32), ("seed_pool", seed_pool, torch.int64), ("local_argmax", local_argmax, torch.int32), ("local_ids", local_ids, torch.int32), ("row_argmax", row_argmax, torch.int32), ("row_candidate_ids", row_candidate_ids, torch.int32), ("accepted", accepted, torch.int32), ("out", out, torch.int32), ): if tensor.device.type != "cuda": raise ValueError(f"{name} must be CUDA") if tensor.dtype != dtype: raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") for name, tensor in ( ("temperature_pool", temperature_pool), ("top_p_pool", top_p_pool), ("offsets_pool", offsets_pool), ("local_max", local_max), ("local_sum", local_sum), ("local_scores", local_scores), ("local_logits", local_logits), ("row_max", row_max), ("row_total", row_total), ("row_candidate_logits", row_candidate_logits), ): if tensor.device.type != "cuda": raise ValueError(f"{name} must be CUDA") if out.shape[0] < rows: raise ValueError(f"out is too small: {out.shape[0]} < {rows}") if rows == 0: return out[:0] local_shape = (rows, num_blocks) candidate_shape = (rows, num_blocks, num_attempts) row_candidate_shape = (rows, num_attempts) if local_max.shape[0] < rows or local_max.shape[1] < num_blocks: raise ValueError(f"local_max must cover {local_shape}, got {local_max.shape}") if local_sum.shape[0] < rows or local_sum.shape[1] < num_blocks: raise ValueError(f"local_sum must cover {local_shape}, got {local_sum.shape}") if local_argmax.shape[0] < rows or local_argmax.shape[1] < num_blocks: raise ValueError( f"local_argmax must cover {local_shape}, got {local_argmax.shape}" ) for name, tensor in ( ("local_scores", local_scores), ("local_logits", local_logits), ("local_ids", local_ids), ): if ( tensor.shape[0] < rows or tensor.shape[1] < num_blocks or tensor.shape[2] < num_attempts ): raise ValueError(f"{name} must cover {candidate_shape}, got {tensor.shape}") for name, tensor in ( ("row_max", row_max), ("row_total", row_total), ("row_argmax", row_argmax), ("accepted", accepted), ): if tensor.shape[0] < rows: raise ValueError(f"{name} is too small: {tensor.shape[0]} < {rows}") for name, tensor in ( ("row_candidate_logits", row_candidate_logits), ("row_candidate_ids", row_candidate_ids), ): if tensor.shape[0] < rows or tensor.shape[1] < num_attempts: raise ValueError( f"{name} must cover {row_candidate_shape}, got {tensor.shape}" ) _top_p_parallel_stage1_kernel[(rows, num_blocks)]( logits, req_pool_indices, temperature_pool, seed_pool, offsets_pool, local_max, local_sum, local_argmax, local_scores, local_logits, local_ids, logits_row_stride=logits.stride(0), vocab_size=vocab_size, num_blocks=num_blocks, BLOCK_SIZE=block_size, NUM_ATTEMPTS=num_attempts, NUM_TOKENS_PER_REQ=num_tokens_per_req, num_warps=4, num_stages=3, ) _top_p_parallel_stage2_kernel[(rows,)]( local_max, local_sum, local_argmax, local_scores, local_logits, local_ids, row_max, row_total, row_argmax, row_candidate_logits, row_candidate_ids, num_blocks=num_blocks, NUM_BLOCKS_PAD=num_blocks_pad, NUM_ATTEMPTS=num_attempts, num_warps=8, num_stages=3, ) _top_p_parallel_stage3_kernel[(rows, num_blocks)]( logits, req_pool_indices, temperature_pool, row_max, row_candidate_logits, row_candidate_ids, local_scores, logits_row_stride=logits.stride(0), vocab_size=vocab_size, num_blocks=num_blocks, BLOCK_SIZE=block_size, NUM_ATTEMPTS=num_attempts, NUM_TOKENS_PER_REQ=num_tokens_per_req, num_warps=4, num_stages=3, ) _top_p_parallel_stage4_kernel[(rows,)]( top_p_pool, req_pool_indices, row_total, row_argmax, row_candidate_ids, local_scores, accepted, out, num_blocks=num_blocks, NUM_BLOCKS_PAD=num_blocks_pad, NUM_ATTEMPTS=num_attempts, NUM_TOKENS_PER_REQ=num_tokens_per_req, num_warps=8, num_stages=3, ) if num_attempts < _TOP_P_REPAIR_NUM_ATTEMPTS: _top_p_parallel_repair_kernel[(rows,)]( logits, req_pool_indices, temperature_pool, top_p_pool, seed_pool, offsets_pool, row_max, row_total, row_argmax, accepted, out, logits_row_stride=logits.stride(0), vocab_size=vocab_size, BLOCK_SIZE=block_size, START_ATTEMPT=num_attempts, NUM_ATTEMPTS_TOTAL=_TOP_P_REPAIR_NUM_ATTEMPTS, NUM_TOKENS_PER_REQ=num_tokens_per_req, num_warps=4, num_stages=3, ) return out[:rows]