| 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 | |
| 8 | @triton.jit |
| 9 | def _min_p_kernel( |
| 10 | logits_ptr, |
| 11 | logits_stride, |
| 12 | expanded_idx_mapping_ptr, |
| 13 | min_p_ptr, |
| 14 | vocab_size, |
| 15 | BLOCK_SIZE: tl.constexpr, |
| 16 | ): |
| 17 | token_idx = tl.program_id(0).to(tl.int64) |
| 18 | req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) |
| 19 | min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32) |
| 20 | if min_p == 0.0: |
| 21 | return |
| 22 | |
| 23 | max_val = float("-inf") |
| 24 | for i in range(0, vocab_size, BLOCK_SIZE): |
| 25 | block = i + tl.arange(0, BLOCK_SIZE) |
| 26 | mask = block < vocab_size |
| 27 | logits = tl.load( |
| 28 | logits_ptr + token_idx * logits_stride + block, |
| 29 | mask=mask, |
| 30 | other=float("-inf"), |
| 31 | ) |
| 32 | max_val = tl.max(tl.maximum(logits, max_val)) |
| 33 | max_val = max_val.to(tl.float32) # type: ignore |
| 34 | |
| 35 | threshold = max_val + tl.log(min_p) |
| 36 | for i in range(0, vocab_size, BLOCK_SIZE): |
| 37 | block = i + tl.arange(0, BLOCK_SIZE) |
| 38 | mask = block < vocab_size |
| 39 | logits = tl.load( |
| 40 | logits_ptr + token_idx * logits_stride + block, |
| 41 | mask=mask, |
| 42 | other=float("-inf"), |
| 43 | ) |
| 44 | logits = tl.where(logits < threshold, float("-inf"), logits) |
| 45 | tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask) |
| 46 | |
| 47 | |
| 48 | def apply_min_p( |
| 49 | logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, min_p: torch.Tensor |
| 50 | ) -> None: |
| 51 | num_tokens, vocab_size = logits.shape |
| 52 | BLOCK_SIZE = 1024 |
| 53 | _min_p_kernel[(num_tokens,)]( |
| 54 | logits, |
| 55 | logits.stride(0), |
| 56 | expanded_idx_mapping, |
| 57 | min_p, |
| 58 | vocab_size, |
| 59 | BLOCK_SIZE=BLOCK_SIZE, |
| 60 | ) |