| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.triton_utils import tl, triton |
| 7 | |
| 8 | |
| 9 | @triton.jit |
| 10 | def _max_with_nan(a, b): |
| 11 | return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL) |
| 12 | |
| 13 | |
| 14 | @triton.jit(do_not_specialize=["width", "nblocks"]) |
| 15 | def _block_scores_kernel( |
| 16 | logits, |
| 17 | starts, |
| 18 | ends, |
| 19 | scores, |
| 20 | stride_row, |
| 21 | stride_col, |
| 22 | stride_start, |
| 23 | stride_end, |
| 24 | width, |
| 25 | nblocks, |
| 26 | BLOCK_SIZE: tl.constexpr, |
| 27 | HAS_STARTS: tl.constexpr, |
| 28 | ROW_REPEAT: tl.constexpr, |
| 29 | TILE: tl.constexpr, |
| 30 | ): |
| 31 | row = tl.program_id(0).to(tl.int64) |
| 32 | blocks = tl.program_id(1) * TILE + tl.arange(0, TILE) |
| 33 | start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0 |
| 34 | end = tl.load(ends + row // ROW_REPEAT * stride_end) |
| 35 | offsets = tl.arange(0, triton.next_power_of_2(BLOCK_SIZE)) |
| 36 | cols = start + blocks[:, None] * BLOCK_SIZE + offsets[None, :] |
| 37 | values = tl.load( |
| 38 | logits + row * stride_row + cols * stride_col, |
| 39 | (blocks[:, None] < nblocks) |
| 40 | & (offsets[None, :] < BLOCK_SIZE) |
| 41 | & (cols < end) |
| 42 | & (cols < width), |
| 43 | other=-float("inf"), |
| 44 | ) |
| 45 | reduced = tl.reduce(values, 1, _max_with_nan) |
| 46 | reduced = tl.where( |
| 47 | (end > start) & (blocks == (end - start - 1) // BLOCK_SIZE), |
| 48 | float("inf"), |
| 49 | reduced, |
| 50 | ) |
| 51 | tl.store(scores + row * nblocks + blocks, reduced, blocks < nblocks) |
| 52 | |
| 53 | |
| 54 | @triton.jit(do_not_specialize=["k"]) |
| 55 | def _store_candidates_kernel( |
| 56 | values, |
| 57 | indices, |
| 58 | output, |
| 59 | out_stride_row, |
| 60 | out_stride_col, |
| 61 | k, |
| 62 | OUT_K: tl.constexpr, |
| 63 | TILE: tl.constexpr, |
| 64 | ): |
| 65 | row = tl.program_id(0).to(tl.int64) |
| 66 | cols = tl.program_id(1) * TILE + tl.arange(0, TILE) |
| 67 | value = tl.load(values + row * k + cols, cols < k, other=-float("inf")) |
| 68 | index = tl.load(indices + row * k + cols, cols < k, other=-1) |
| 69 | # NaN scores can occur during warmup; only -inf denotes padding. |
| 70 | tl.store( |
| 71 | output + row * out_stride_row + cols * out_stride_col, |
| 72 | tl.where(value != -float("inf"), index, -1), |
| 73 | cols < OUT_K, |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | @triton.jit(do_not_specialize=["width", "nblocks"]) |
| 78 | def _candidate_flags_kernel( |
| 79 | candidates, |
| 80 | starts, |
| 81 | flags, |
| 82 | stride_row, |
| 83 | stride_col, |
| 84 | stride_start, |
| 85 | width, |
| 86 | nblocks, |
| 87 | BLOCK_SIZE: tl.constexpr, |
| 88 | K: tl.constexpr, |
| 89 | HAS_STARTS: tl.constexpr, |
| 90 | ROW_REPEAT: tl.constexpr, |
| 91 | ): |
| 92 | row = tl.program_id(0).to(tl.int64) |
| 93 | offsets = tl.arange(0, 1024) |
| 94 | for tile in range(tl.cdiv(nblocks + 1, 1024)): |
| 95 | slots = tile * 1024 + offsets |
| 96 | tl.store(flags + row * (nblocks + 1) + slots, 0, slots <= nblocks) |
| 97 | start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0 |
| 98 | cols = tl.arange(0, triton.next_power_of_2(K)) |
| 99 | block = tl.load( |
| 100 | candidates + row * stride_row + cols * stride_col, cols < K, other=-1 |
| 101 | ).to(tl.int64) |
| 102 | # Preserve the packed-column clamp for candidates beyond the logits width. |
| 103 | block = tl.where(start + block * BLOCK_SIZE >= width, nblocks, block) |
| 104 | tl.debug_barrier() |
| 105 | tl.store(flags + row * (nblocks + 1) + block, 1, (cols < K) & (block >= 0)) |
| 106 | |
| 107 | |
| 108 | @triton.jit(do_not_specialize=["width", "nblocks"]) |
| 109 | def _mask_candidates_kernel( |
| 110 | logits, |
| 111 | starts, |
| 112 | ends, |
| 113 | flags, |
| 114 | stride_row, |
| 115 | stride_col, |
| 116 | stride_start, |
| 117 | stride_end, |
| 118 | width, |
| 119 | nblocks, |
| 120 | BLOCK_SIZE: tl.constexpr, |
| 121 | HAS_STARTS: tl.constexpr, |
| 122 | ROW_REPEAT: tl.constexpr, |
| 123 | TILE: tl.constexpr, |
| 124 | ): |
| 125 | row = tl.program_id(0).to(tl.int64) |
| 126 | cols = tl.program_id(1) * TILE + tl.arange(0, TILE) |
| 127 | start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0 |
| 128 | end = tl.load(ends + row // ROW_REPEAT * stride_end) |
| 129 | valid = (cols >= start) & (cols < end) & (cols < width) |
| 130 | block = (cols - start) // BLOCK_SIZE |
| 131 | keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0) |
| 132 | edge = tl.load(flags + row * (nblocks + 1) + nblocks) |
| 133 | keep = (keep != 0) | ((cols == width - 1) & (edge != 0)) |
| 134 | tl.store( |
| 135 | logits + row * stride_row + cols * stride_col, |
| 136 | -float("inf"), |
| 137 | (cols < width) & ~(valid & keep), |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | def select_candidate_blocks( |
| 142 | logits: torch.Tensor, |
| 143 | row_ks: torch.Tensor | None, |
| 144 | row_ke: torch.Tensor, |
| 145 | topk_blocks: int, |
| 146 | block_size: int, |
| 147 | out: torch.Tensor, |
| 148 | row_repeat: int = 1, |
| 149 | ) -> None: |
| 150 | """Select local block IDs by maximum score, pinning each row's newest block. |
| 151 | |
| 152 | Row bounds are in packed column space; absent starts mean zero. |
| 153 | Decode rows share bounds in groups of ``row_repeat``. Output is -1 padded. |
| 154 | """ |
| 155 | assert logits.is_cuda |
| 156 | rows, width = logits.shape |
| 157 | if not rows: |
| 158 | return |
| 159 | if not width: |
| 160 | out.fill_(-1) |
| 161 | return |
| 162 | nblocks = triton.cdiv(width, block_size) |
| 163 | scores = logits.new_empty((rows, nblocks)) |
| 164 | _block_scores_kernel[(rows, triton.cdiv(nblocks, 128))]( |
| 165 | logits, |
| 166 | row_ks, |
| 167 | row_ke, |
| 168 | scores, |
| 169 | *logits.stride(), |
| 170 | row_ks.stride(0) if row_ks is not None else 0, |
| 171 | row_ke.stride(0), |
| 172 | width, |
| 173 | nblocks, |
| 174 | block_size, |
| 175 | row_ks is not None, |
| 176 | row_repeat, |
| 177 | 128, |
| 178 | ) |
| 179 | # Keep the existing top-k tie behavior. |
| 180 | top = scores.topk(min(topk_blocks, nblocks), dim=-1) |
| 181 | _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))]( |
| 182 | top.values, |
| 183 | top.indices, |
| 184 | out, |
| 185 | *out.stride(), |
| 186 | top.values.shape[1], |
| 187 | topk_blocks, |
| 188 | 256, |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | def apply_candidate_mask( |
| 193 | logits: torch.Tensor, |
| 194 | row_ks: torch.Tensor | None, |
| 195 | row_ke: torch.Tensor, |
| 196 | candidate_blocks: torch.Tensor, |
| 197 | block_size: int, |
| 198 | row_repeat: int = 1, |
| 199 | ) -> None: |
| 200 | """Mask packed logits outside causal bounds and request-local candidates.""" |
| 201 | assert logits.is_cuda |
| 202 | rows, width = logits.shape |
| 203 | if not rows or not width: |
| 204 | return |
| 205 | nblocks = triton.cdiv(width, block_size) |
| 206 | flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8) |
| 207 | start_stride = row_ks.stride(0) if row_ks is not None else 0 |
| 208 | _candidate_flags_kernel[(rows,)]( |
| 209 | candidate_blocks, |
| 210 | row_ks, |
| 211 | flags, |
| 212 | *candidate_blocks.stride(), |
| 213 | start_stride, |
| 214 | width, |
| 215 | nblocks, |
| 216 | block_size, |
| 217 | candidate_blocks.shape[1], |
| 218 | row_ks is not None, |
| 219 | row_repeat, |
| 220 | ) |
| 221 | _mask_candidates_kernel[(rows, triton.cdiv(width, 1024))]( |
| 222 | logits, |
| 223 | row_ks, |
| 224 | row_ke, |
| 225 | flags, |
| 226 | *logits.stride(), |
| 227 | start_stride, |
| 228 | row_ke.stride(0), |
| 229 | width, |
| 230 | nblocks, |
| 231 | block_size, |
| 232 | row_ks is not None, |
| 233 | row_repeat, |
| 234 | 1024, |
| 235 | ) |