| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Device-side request mapping and sparse indexer metadata.""" |
| 4 | |
| 5 | from vllm.triton_utils import tl, triton |
| 6 | |
| 7 | |
| 8 | @triton.jit |
| 9 | def _token_request(query_start_loc, token, num_reqs): |
| 10 | lo = tl.full(token.shape, 0, tl.int32) |
| 11 | hi = tl.full(token.shape, num_reqs, tl.int32) |
| 12 | while tl.sum((lo < hi).to(tl.int32), 0) > 0: |
| 13 | mid = (lo + hi) // 2 |
| 14 | end = tl.load(query_start_loc + mid + 1, mid < num_reqs, other=0x7FFFFFFF) |
| 15 | right = (lo < hi) & (end <= token) |
| 16 | lo = tl.where(right, mid + 1, lo) |
| 17 | hi = tl.where((lo < hi) & ~right, mid, hi) |
| 18 | return tl.minimum(lo, num_reqs - 1) |
| 19 | |
| 20 | |
| 21 | @triton.jit(do_not_specialize=["num_reqs", "num_mapped", "num_tokens"]) |
| 22 | def _token_request_mapping_kernel( |
| 23 | query_start_loc, output, num_reqs, num_mapped, num_tokens |
| 24 | ): |
| 25 | token = tl.program_id(0) * 256 + tl.arange(0, 256) |
| 26 | req = _token_request(query_start_loc, token, num_reqs) |
| 27 | tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens) |
| 28 | |
| 29 | |
| 30 | @triton.jit( |
| 31 | do_not_specialize=[ |
| 32 | "num_reqs", |
| 33 | "num_tokens", |
| 34 | "capacity", |
| 35 | "input_stride", |
| 36 | "output_stride", |
| 37 | ] |
| 38 | ) |
| 39 | def _indexer_decode_metadata_kernel( |
| 40 | query_start_loc, |
| 41 | seq_lens, |
| 42 | block_table, |
| 43 | output_seq_lens, |
| 44 | output_block_table, |
| 45 | decode_lens, |
| 46 | indices, |
| 47 | per_req_lens, |
| 48 | num_reqs, |
| 49 | num_tokens, |
| 50 | capacity, |
| 51 | input_stride, |
| 52 | output_stride, |
| 53 | BLOCK_COLS: tl.constexpr, |
| 54 | ): |
| 55 | token = tl.program_id(0) |
| 56 | cols = tl.arange(0, triton.next_power_of_2(BLOCK_COLS)) |
| 57 | actual_tokens = tl.load(query_start_loc + num_reqs) |
| 58 | if token < num_tokens: |
| 59 | req = _token_request(query_start_loc, tl.full((1,), token, tl.int32), num_reqs) |
| 60 | req = tl.sum(req, 0) |
| 61 | valid = token < actual_tokens |
| 62 | start = tl.load(query_start_loc + req) |
| 63 | end = tl.load(query_start_loc + req + 1) |
| 64 | seq = tl.load(seq_lens + req) |
| 65 | length = tl.where(valid, seq - (end - start) + token - start + 1, 0) |
| 66 | tl.store(output_seq_lens + token, length) |
| 67 | tl.store(decode_lens + token, 1) |
| 68 | tl.store( |
| 69 | indices + token, tl.where(valid, req, num_reqs + token - actual_tokens) |
| 70 | ) |
| 71 | blocks = tl.load( |
| 72 | block_table + req.to(tl.int64) * input_stride + cols, |
| 73 | valid & (cols < BLOCK_COLS), |
| 74 | other=0, |
| 75 | ) |
| 76 | tl.store( |
| 77 | output_block_table + token.to(tl.int64) * output_stride + cols, |
| 78 | blocks, |
| 79 | cols < BLOCK_COLS, |
| 80 | ) |
| 81 | else: |
| 82 | tail = num_tokens + (token - num_tokens) * 256 + tl.arange(0, 256) |
| 83 | tl.store(output_seq_lens + tail, 0, tail < capacity) |
| 84 | if token < num_reqs: |
| 85 | start = tl.load(query_start_loc + token) |
| 86 | end = tl.load(query_start_loc + token + 1) |
| 87 | tl.store(per_req_lens + token, end - start) |