| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Combined Top-K and Top-P Triton kernels. |
| 4 | |
| 5 | Based on the paper "Qrita: High-performance Top-k and Top-p Algorithm for GPUs |
| 6 | using Pivot-based Truncation and Selection" By Park et al. |
| 7 | (https://arxiv.org/abs/2602.01518) |
| 8 | |
| 9 | """ |
| 10 | |
| 11 | from typing import Any |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.model_executor.warmup.jit_warmup import ( |
| 16 | WarmupChoices, |
| 17 | WarmupIntRange, |
| 18 | _when, |
| 19 | ) |
| 20 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 21 | DispatchSpec, |
| 22 | TritonWarmupTensor, |
| 23 | triton_kernel_dispatcher_with_warmup, |
| 24 | ) |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.triton_utils import tl, triton |
| 27 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 28 | from vllm.utils.math_utils import next_power_of_2 |
| 29 | from vllm.utils.platform_utils import num_compute_units |
| 30 | |
| 31 | _TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {} |
| 32 | _TRITON_BUFFER_CACHE: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {} |
| 33 | |
| 34 | # fmt: off |
| 35 | _NORMAL_CDF_TO_SIGMA_TABLE = [ |
| 36 | 3.656, 3.650, 3.650, 3.650, 3.626, 3.626, 3.626, 3.514, 3.514, 3.503, |
| 37 | 3.503, 3.434, 3.434, 3.428, 3.428, 3.387, 3.380, 3.380, 3.376, 3.373, |
| 38 | 3.373, 3.356, 3.354, 3.354, 3.291, 3.249, 3.234, 3.214, 3.198, 3.198, |
| 39 | 3.185, 3.177, 3.177, 3.165, 3.164, 3.161, 3.138, 3.120, 3.115, 3.113, |
| 40 | 3.093, 3.066, 3.054, 3.043, 3.037, 3.023, 2.993, 2.991, 2.976, 2.970, |
| 41 | 2.952, 2.946, 2.932, 2.908, 2.902, 2.895, 2.886, 2.874, 2.861, 2.844, |
| 42 | 2.836, 2.810, 2.801, 2.790, 2.784, 2.779, 2.767, 2.757, 2.745, 2.733, |
| 43 | 2.723, 2.716, 2.693, 2.678, 2.671, 2.656, 2.649, 2.629, 2.611, 2.595, |
| 44 | 2.592, 2.585, 2.574, 2.550, 2.543, 2.534, 2.521, 2.518, 2.497, 2.485, |
| 45 | 2.468, 2.450, 2.441, 2.430, 2.412, 2.402, 2.389, 2.383, 2.377, 2.364, |
| 46 | 2.349, 2.338, 2.332, 2.319, 2.310, 2.301, 2.282, 2.274, 2.266, 2.250, |
| 47 | 2.242, 2.236, 2.226, 2.215, 2.207, 2.196, 2.179, 2.171, 2.162, 2.147, |
| 48 | 2.135, 2.121, 2.109, 2.095, 2.085, 2.073, 2.063, 2.045, 2.030, 2.016, |
| 49 | 2.003, 1.992, 1.983, 1.972, 1.960, 1.949, 1.940, 1.928, 1.912, 1.897, |
| 50 | 1.881, 1.869, 1.854, 1.838, 1.824, 1.807, 1.792, 1.779, 1.764, 1.751, |
| 51 | 1.739, 1.726, 1.711, 1.697, 1.685, 1.668, 1.652, 1.636, 1.622, 1.603, |
| 52 | 1.585, 1.568, 1.551, 1.534, 1.513, 1.499, 1.480, 1.464, 1.441, 1.422, |
| 53 | 1.394, 1.373, 1.347, 1.320, 1.296, 1.270, 1.246, 1.219, 1.190, 1.163, |
| 54 | 1.135, 1.104, 1.073, 1.041, 1.006, 0.969, 0.931, 0.894, 0.851, 0.806, |
| 55 | 0.757, 0.702, 0.643, 0.574, 0.498, 0.405, 0.288, 0.134, -0.110, -3.813 |
| 56 | ] |
| 57 | |
| 58 | _PERCENTILE_TO_STD_TABLE = [ |
| 59 | 2.576, 2.319, 2.178, 2.064, 1.968, 1.892, 1.819, 1.757, 1.708, 1.659, |
| 60 | 1.616, 1.568, 1.526, 1.492, 1.456, 1.420, 1.382, 1.342, 1.309, 1.280, |
| 61 | 1.249, 1.221, 1.193, 1.169, 1.145, 1.121, 1.095, 1.073, 1.050, 1.030, |
| 62 | 1.008, 0.987, 0.966, 0.945, 0.926, 0.910, 0.891, 0.871, 0.854, 0.837, |
| 63 | 0.819, 0.803, 0.784, 0.767, 0.753, 0.734, 0.719, 0.702, 0.690, 0.675, |
| 64 | 0.658, 0.640, 0.625, 0.609, 0.595, 0.578, 0.564, 0.550, 0.537, 0.521, |
| 65 | 0.509, 0.495, 0.481, 0.466, 0.453, 0.439, 0.424, 0.410, 0.397, 0.383, |
| 66 | 0.370, 0.356, 0.343, 0.330, 0.316, 0.302, 0.289, 0.274, 0.261, 0.247, |
| 67 | 0.235, 0.223, 0.209, 0.196, 0.184, 0.172, 0.159, 0.149, 0.137, 0.124, |
| 68 | 0.112, 0.100, 0.086, 0.074, 0.062, 0.050, 0.035, 0.023, 0.009, -0.003, |
| 69 | -0.015, -0.027, -0.039, -0.052, -0.063, -0.074, -0.085, -0.097, -0.109, -0.122, |
| 70 | -0.134, -0.147, -0.158, -0.171, -0.184, -0.196, -0.210, -0.223, -0.235, -0.248, |
| 71 | -0.261, -0.275, -0.289, -0.302, -0.317, -0.328, -0.341, -0.353, -0.368, -0.382, |
| 72 | -0.396, -0.410, -0.426, -0.439, -0.452, -0.465, -0.480, -0.493, -0.507, -0.521, |
| 73 | -0.537, -0.551, -0.568, -0.582, -0.597, -0.614, -0.628, -0.643, -0.658, -0.673, |
| 74 | -0.691, -0.706, -0.721, -0.738, -0.754, -0.769, -0.789, -0.808, -0.824, -0.838, |
| 75 | -0.857, -0.877, -0.893, -0.912, -0.929, -0.947, -0.965, -0.983, -1.003, -1.027, |
| 76 | -1.050, -1.070, -1.092, -1.117, -1.139, -1.162, -1.189, -1.216, -1.241, -1.272, |
| 77 | -1.300, -1.330, -1.367, -1.404, -1.441, -1.485, -1.523, -1.564, -1.607, -1.658, |
| 78 | -1.710, -1.778, -1.832, -1.901, -1.978, -2.068, -2.174, -2.325, -2.577, -3.813 |
| 79 | ] |
| 80 | # fmt: on |
| 81 | |
| 82 | |
| 83 | @triton.jit |
| 84 | def _update_min_larger_stats(data, above_mask, min_larger, num_min_larger, sentinel): |
| 85 | """Update running (min, count) of values above a pivot across tiles. |
| 86 | |
| 87 | Tracks the smallest value strictly above a pivot and how many times |
| 88 | it occurs. Called once per tile per pivot; the running state is |
| 89 | carried across tiles via `min_larger` / `num_min_larger`. |
| 90 | |
| 91 | Merge rule: |
| 92 | - tile min < running min → replace both |
| 93 | - tile min == running min → accumulate count |
| 94 | - tile min > running min → keep running values |
| 95 | """ |
| 96 | tile_min = tl.min(tl.where(above_mask, data, sentinel)) |
| 97 | tile_eq = above_mask & (tl.abs(data - tile_min) < 1e-9) |
| 98 | tile_cnt = tl.sum(tile_eq) |
| 99 | is_new = tile_min < min_larger |
| 100 | is_same = tl.abs(tile_min - min_larger) < 1e-9 |
| 101 | num_min_larger = tl.where(is_new, tile_cnt, num_min_larger + tile_cnt * is_same) |
| 102 | min_larger = tl.minimum(min_larger, tile_min) |
| 103 | return min_larger, num_min_larger |
| 104 | |
| 105 | |
| 106 | @triton.jit(do_not_specialize_on_alignment=["BATCH_SIZE"]) |
| 107 | def _topk_topp_kernel( |
| 108 | LOGITS, |
| 109 | LOGITS_STRIDE_0, |
| 110 | BUFFER, |
| 111 | PERCENTILE_TO_STD_TABLE, |
| 112 | NORMAL_CDF_TO_SIGMA_TABLE, |
| 113 | K, |
| 114 | P, |
| 115 | BATCH_SIZE, |
| 116 | VOCAB_SIZE: tl.constexpr, |
| 117 | MASK_VALUE: tl.constexpr, |
| 118 | BLOCK_SIZE: tl.constexpr, |
| 119 | BLOCK_SIZE_TRUNC: tl.constexpr, |
| 120 | TOPK_ENABLED: tl.constexpr, |
| 121 | TOPP_ENABLED: tl.constexpr, |
| 122 | SPLIT_COVERS_PONLY: tl.constexpr, |
| 123 | ): |
| 124 | NUM_TILES: tl.constexpr = (VOCAB_SIZE + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 125 | pid = tl.program_id(0) |
| 126 | num_programs = tl.num_programs(0) |
| 127 | for row_id in tl.range(pid, BATCH_SIZE, num_programs): |
| 128 | LOGITS_ROW = LOGITS + row_id.to(tl.int64) * LOGITS_STRIDE_0 |
| 129 | BUFFER_ROW = BUFFER + pid * VOCAB_SIZE |
| 130 | |
| 131 | final_pivot = -float("inf") |
| 132 | duplicate_logit = float("inf") |
| 133 | num_duplicate_logit = tl.zeros((), dtype=tl.uint32) |
| 134 | num_keep = tl.zeros((), dtype=tl.uint32) |
| 135 | num_kept = tl.zeros((), dtype=tl.uint32) |
| 136 | |
| 137 | max_logit = -float("inf") |
| 138 | min_logit = float("inf") |
| 139 | |
| 140 | if TOPK_ENABLED: |
| 141 | k = tl.load(K + row_id) |
| 142 | if k < VOCAB_SIZE: |
| 143 | # Zeroth pass: Compute avg and std from a sample block |
| 144 | offs = tl.arange(0, BLOCK_SIZE) |
| 145 | mask_n = offs < VOCAB_SIZE |
| 146 | logits_blk0 = tl.load( |
| 147 | LOGITS_ROW + offs, mask=mask_n, other=-float("inf") |
| 148 | ) |
| 149 | # Exclude -inf values (e.g. from grammar bitmasks) from |
| 150 | # statistics to avoid NaN in pivot computation. |
| 151 | finite_mask = (logits_blk0 > -float("inf")) & mask_n |
| 152 | num_finite = tl.sum(finite_mask) |
| 153 | finite_logits = tl.where(finite_mask, logits_blk0, 0.0) |
| 154 | avg_logit = tl.where( |
| 155 | num_finite > 0, tl.sum(finite_logits) / num_finite, 0.0 |
| 156 | ) |
| 157 | sq_avg_logit = tl.where( |
| 158 | num_finite > 0, |
| 159 | tl.sum(finite_logits * finite_logits) / num_finite, |
| 160 | 0.0, |
| 161 | ) |
| 162 | std_logit = tl.sqrt( |
| 163 | tl.maximum(sq_avg_logit - avg_logit * avg_logit, 0.0) |
| 164 | ) |
| 165 | |
| 166 | # Calculate outlier pivot t for Gaussian sigma-truncation |
| 167 | percentile = tl.cast(k / VOCAB_SIZE * 200, tl.uint32) |
| 168 | percentile = tl.minimum(percentile, 199) |
| 169 | sigma = tl.load(PERCENTILE_TO_STD_TABLE + percentile) |
| 170 | sigma = sigma + tl.abs(sigma) * -0.15 |
| 171 | outlier_pivot = avg_logit + std_logit * sigma |
| 172 | num_outliers = tl.zeros((), dtype=tl.uint32) |
| 173 | |
| 174 | # First pass: compute max and min logits and gather outliers |
| 175 | num_finite_total = tl.zeros((), dtype=tl.uint32) |
| 176 | for i in range(0, NUM_TILES): |
| 177 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 178 | mask_n = offs_n < VOCAB_SIZE |
| 179 | logits_blk = tl.load( |
| 180 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 181 | ) |
| 182 | |
| 183 | max_logit = tl.maximum(max_logit, tl.max(logits_blk)) |
| 184 | # Exclude -inf from min to keep binary search bounds |
| 185 | # finite (avoids NaN pivots). |
| 186 | finite_blk_mask = logits_blk > -float("inf") |
| 187 | finite_blk = tl.where(finite_blk_mask, logits_blk, float("inf")) |
| 188 | min_logit = tl.minimum(min_logit, tl.min(finite_blk)) |
| 189 | num_finite_total += tl.sum(finite_blk_mask & mask_n) |
| 190 | |
| 191 | outlier_mask = (logits_blk > outlier_pivot) & mask_n |
| 192 | cumulative_pos = tl.cast( |
| 193 | tl.cumsum(outlier_mask) - 1 + num_outliers, tl.int32 |
| 194 | ) |
| 195 | num_outliers += tl.sum(outlier_mask) |
| 196 | write_pos = tl.where(outlier_mask, cumulative_pos, -1) |
| 197 | tl.store(BUFFER_ROW + write_pos, logits_blk, mask=outlier_mask) |
| 198 | |
| 199 | # If no finite logits exist (all -inf), clamp min to |
| 200 | # max so the search converges to -inf (no masking). |
| 201 | min_logit = tl.minimum(min_logit, max_logit) |
| 202 | |
| 203 | # Second passes: Ternary search for pivot |
| 204 | num_iters = 0 |
| 205 | k_pivot = float("inf") |
| 206 | k_pivots_num = tl.zeros((), dtype=tl.uint32) |
| 207 | min_larger = float("inf") |
| 208 | num_min_larger = tl.zeros((), dtype=tl.uint32) |
| 209 | if num_outliers > k: |
| 210 | max_range = max_logit |
| 211 | min_range = outlier_pivot |
| 212 | search_range = tl.cast(num_outliers, tl.int32) |
| 213 | search_iters = tl.cast( |
| 214 | (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC, |
| 215 | tl.int32, |
| 216 | ) |
| 217 | found_pivot = 0 |
| 218 | while found_pivot == 0: |
| 219 | k_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range |
| 220 | k_pivots_num_0 = tl.zeros((), dtype=tl.uint32) |
| 221 | min_larger_0 = float("inf") |
| 222 | num_min_larger_0 = tl.zeros((), dtype=tl.uint32) |
| 223 | |
| 224 | k_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range |
| 225 | k_pivots_num_1 = tl.zeros((), dtype=tl.uint32) |
| 226 | min_larger_1 = float("inf") |
| 227 | num_min_larger_1 = tl.zeros((), dtype=tl.uint32) |
| 228 | |
| 229 | # Single fused pass: compute k_pivots_num, |
| 230 | # min_larger, and num_min_larger together to avoid |
| 231 | # a second data scan. See _update_min_larger_stats |
| 232 | # for the tile-level merge logic. |
| 233 | for i in range(0, search_iters): |
| 234 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 235 | 0, BLOCK_SIZE_TRUNC |
| 236 | ) |
| 237 | mask_n_2 = offs_n < search_range |
| 238 | logits_blk2 = tl.load( |
| 239 | BUFFER_ROW + offs_n, mask=mask_n_2, other=-float("inf") |
| 240 | ) |
| 241 | |
| 242 | above_0 = logits_blk2 > k_pivot_0 |
| 243 | above_1 = logits_blk2 > k_pivot_1 |
| 244 | k_pivots_num_0 += tl.sum(above_0) |
| 245 | k_pivots_num_1 += tl.sum(above_1) |
| 246 | |
| 247 | min_larger_0, num_min_larger_0 = _update_min_larger_stats( |
| 248 | logits_blk2, |
| 249 | above_0, |
| 250 | min_larger_0, |
| 251 | num_min_larger_0, |
| 252 | float("inf"), |
| 253 | ) |
| 254 | min_larger_1, num_min_larger_1 = _update_min_larger_stats( |
| 255 | logits_blk2, |
| 256 | above_1, |
| 257 | min_larger_1, |
| 258 | num_min_larger_1, |
| 259 | float("inf"), |
| 260 | ) |
| 261 | |
| 262 | # Check if any of the pivots satisfy termination condition |
| 263 | if ( |
| 264 | k_pivots_num_0 >= k |
| 265 | and k_pivots_num_0 - num_min_larger_0 < k |
| 266 | ): |
| 267 | k_pivot = k_pivot_0 |
| 268 | k_pivots_num = k_pivots_num_0 |
| 269 | min_larger = min_larger_0 |
| 270 | num_min_larger = num_min_larger_0 |
| 271 | found_pivot = 1 |
| 272 | if ( |
| 273 | k_pivots_num_1 >= k |
| 274 | and k_pivots_num_1 - num_min_larger_1 < k |
| 275 | ): |
| 276 | k_pivot = k_pivot_1 |
| 277 | k_pivots_num = k_pivots_num_1 |
| 278 | min_larger = min_larger_1 |
| 279 | num_min_larger = num_min_larger_1 |
| 280 | found_pivot = 1 |
| 281 | |
| 282 | # Update range |
| 283 | if k_pivots_num_1 > k: |
| 284 | min_range = k_pivot_1 |
| 285 | elif k_pivots_num_0 > k: |
| 286 | min_range = k_pivot_0 |
| 287 | |
| 288 | if k_pivots_num_0 < k: |
| 289 | max_range = k_pivot_0 |
| 290 | elif k_pivots_num_1 < k: |
| 291 | max_range = k_pivot_1 |
| 292 | |
| 293 | num_iters += 1 |
| 294 | if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: |
| 295 | k_pivot = (max_range + min_range) / 2.0 |
| 296 | min_larger = min_larger_0 |
| 297 | num_min_larger = num_min_larger_0 |
| 298 | found_pivot = 1 |
| 299 | else: |
| 300 | # If top-k outlier gathering failed, search whole logit space |
| 301 | max_range = max_logit |
| 302 | min_range = min_logit |
| 303 | found_pivot = 0 |
| 304 | while found_pivot == 0: |
| 305 | k_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range |
| 306 | k_pivots_num_0 = tl.zeros((), dtype=tl.uint32) |
| 307 | min_larger_0 = float("inf") |
| 308 | num_min_larger_0 = tl.zeros((), dtype=tl.uint32) |
| 309 | |
| 310 | k_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range |
| 311 | k_pivots_num_1 = tl.zeros((), dtype=tl.uint32) |
| 312 | min_larger_1 = float("inf") |
| 313 | num_min_larger_1 = tl.zeros((), dtype=tl.uint32) |
| 314 | |
| 315 | # Single fused pass over full vocab (same approach |
| 316 | # as the buffer path above). |
| 317 | for i in range(0, NUM_TILES): |
| 318 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 319 | mask_n = offs_n < VOCAB_SIZE |
| 320 | logits_blk2 = tl.load( |
| 321 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 322 | ) |
| 323 | |
| 324 | above_0 = logits_blk2 > k_pivot_0 |
| 325 | above_1 = logits_blk2 > k_pivot_1 |
| 326 | k_pivots_num_0 += tl.sum(above_0) |
| 327 | k_pivots_num_1 += tl.sum(above_1) |
| 328 | |
| 329 | min_larger_0, num_min_larger_0 = _update_min_larger_stats( |
| 330 | logits_blk2, |
| 331 | above_0, |
| 332 | min_larger_0, |
| 333 | num_min_larger_0, |
| 334 | float("inf"), |
| 335 | ) |
| 336 | min_larger_1, num_min_larger_1 = _update_min_larger_stats( |
| 337 | logits_blk2, |
| 338 | above_1, |
| 339 | min_larger_1, |
| 340 | num_min_larger_1, |
| 341 | float("inf"), |
| 342 | ) |
| 343 | |
| 344 | # Check if any of the pivots satisfy termination condition |
| 345 | if ( |
| 346 | k_pivots_num_0 >= k |
| 347 | and k_pivots_num_0 - num_min_larger_0 < k |
| 348 | ): |
| 349 | k_pivot = k_pivot_0 |
| 350 | k_pivots_num = k_pivots_num_0 |
| 351 | min_larger = min_larger_0 |
| 352 | num_min_larger = num_min_larger_0 |
| 353 | found_pivot = 1 |
| 354 | if ( |
| 355 | k_pivots_num_1 >= k |
| 356 | and k_pivots_num_1 - num_min_larger_1 < k |
| 357 | ): |
| 358 | k_pivot = k_pivot_1 |
| 359 | k_pivots_num = k_pivots_num_1 |
| 360 | min_larger = min_larger_1 |
| 361 | num_min_larger = num_min_larger_1 |
| 362 | found_pivot = 1 |
| 363 | |
| 364 | # Update range |
| 365 | if k_pivots_num_1 > k: |
| 366 | min_range = k_pivot_1 |
| 367 | elif k_pivots_num_0 > k: |
| 368 | min_range = k_pivot_0 |
| 369 | |
| 370 | if k_pivots_num_0 < k: |
| 371 | max_range = k_pivot_0 |
| 372 | elif k_pivots_num_1 < k: |
| 373 | max_range = k_pivot_1 |
| 374 | |
| 375 | num_iters += 1 |
| 376 | if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: |
| 377 | k_pivot = (max_range + min_range) / 2.0 |
| 378 | min_larger = min_larger_0 |
| 379 | num_min_larger = num_min_larger_0 |
| 380 | found_pivot = 1 |
| 381 | |
| 382 | duplicate_logit = min_larger |
| 383 | num_duplicate_logit = num_min_larger |
| 384 | num_keep = num_duplicate_logit - (k_pivots_num - k) |
| 385 | num_kept = tl.zeros((), dtype=tl.uint32) |
| 386 | |
| 387 | # Top-k only path. If there are fewer finite values |
| 388 | # than k (e.g. grammar mask), keep everything. |
| 389 | final_pivot = k_pivot if num_finite_total > k else -float("inf") |
| 390 | |
| 391 | if TOPP_ENABLED and num_finite_total > k: |
| 392 | #### TOP-P SAMPLING AFTER TOP-K #### |
| 393 | p = tl.load(P + row_id) |
| 394 | if p < 1.0: |
| 395 | min_logit = k_pivot |
| 396 | sum_exp_logits = 0.0 |
| 397 | num_outliers_2 = tl.zeros((), dtype=tl.uint32) |
| 398 | search_range = tl.cast(num_outliers, tl.int32) |
| 399 | search_iters = tl.cast( |
| 400 | (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC, |
| 401 | tl.int32, |
| 402 | ) |
| 403 | |
| 404 | # Third pass: Calculate exp logits and sum, gather outliers |
| 405 | if num_outliers > k: |
| 406 | for i in range(0, search_iters): |
| 407 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 408 | 0, BLOCK_SIZE_TRUNC |
| 409 | ) |
| 410 | mask_n_2 = offs_n < search_range |
| 411 | |
| 412 | probs_blk = tl.load( |
| 413 | BUFFER_ROW + offs_n, |
| 414 | mask=mask_n_2, |
| 415 | other=-float("inf"), |
| 416 | ) |
| 417 | |
| 418 | outlier_mask = (probs_blk > min_logit) & mask_n_2 |
| 419 | |
| 420 | # Duplicate logit handling for Top-k |
| 421 | if num_keep < num_duplicate_logit: |
| 422 | duplicate_mask = ( |
| 423 | tl.abs(probs_blk - duplicate_logit) < 1e-9 |
| 424 | ) |
| 425 | duplicate_count = ( |
| 426 | tl.cumsum(duplicate_mask) + num_kept |
| 427 | ) |
| 428 | duplicate_keep_mask = ( |
| 429 | duplicate_count <= num_keep |
| 430 | ) & duplicate_mask |
| 431 | duplicate_remove_mask = ( |
| 432 | duplicate_mask & ~duplicate_keep_mask |
| 433 | ) |
| 434 | outlier_mask = outlier_mask & ( |
| 435 | ~duplicate_remove_mask |
| 436 | ) |
| 437 | num_kept += tl.sum(duplicate_keep_mask) |
| 438 | |
| 439 | probs_blk = tl.where( |
| 440 | outlier_mask, probs_blk, -float("inf") |
| 441 | ) |
| 442 | probs_blk = probs_blk - max_logit |
| 443 | probs_blk = tl.exp(probs_blk) |
| 444 | sum_exp_logits += tl.sum(probs_blk) |
| 445 | |
| 446 | # Fourth pass: Calculate BUFFER and get outliers |
| 447 | for i in range(0, search_iters): |
| 448 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 449 | 0, BLOCK_SIZE_TRUNC |
| 450 | ) |
| 451 | mask_n_2 = offs_n < search_range |
| 452 | |
| 453 | probs_blk = tl.load( |
| 454 | BUFFER_ROW + offs_n, |
| 455 | mask=mask_n_2, |
| 456 | other=-float("inf"), |
| 457 | ) |
| 458 | |
| 459 | probs_blk = probs_blk - max_logit |
| 460 | probs_blk = tl.exp(probs_blk) |
| 461 | probs_blk = probs_blk / sum_exp_logits |
| 462 | tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n_2) |
| 463 | else: |
| 464 | # If top-k outlier gathering failed, |
| 465 | # retry gathering using top-k pivot |
| 466 | for i in range(0, NUM_TILES): |
| 467 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 468 | mask_n = offs_n < VOCAB_SIZE |
| 469 | |
| 470 | probs_blk = tl.load( |
| 471 | LOGITS_ROW + offs_n, |
| 472 | mask=mask_n, |
| 473 | other=-float("inf"), |
| 474 | ) |
| 475 | |
| 476 | outlier_mask = (probs_blk > min_logit) & mask_n |
| 477 | |
| 478 | # Duplicate logit handling for Top-k |
| 479 | duplicate_mask = ( |
| 480 | tl.abs(probs_blk - duplicate_logit) < 1e-9 |
| 481 | ) |
| 482 | duplicate_count = tl.cumsum(duplicate_mask) + num_kept |
| 483 | duplicate_keep_mask = ( |
| 484 | duplicate_count <= num_keep |
| 485 | ) & duplicate_mask |
| 486 | duplicate_remove_mask = ( |
| 487 | duplicate_mask & ~duplicate_keep_mask |
| 488 | ) |
| 489 | outlier_mask = outlier_mask & (~duplicate_remove_mask) |
| 490 | num_kept += tl.sum(duplicate_keep_mask) |
| 491 | |
| 492 | probs_blk = tl.where( |
| 493 | outlier_mask, probs_blk, -float("inf") |
| 494 | ) |
| 495 | probs_blk = probs_blk - max_logit |
| 496 | probs_blk = tl.exp(probs_blk) |
| 497 | sum_exp_logits += tl.sum(probs_blk) |
| 498 | |
| 499 | cumulative_pos = tl.cast( |
| 500 | tl.cumsum(outlier_mask) - 1 + num_outliers_2, |
| 501 | tl.int32, |
| 502 | ) |
| 503 | num_outliers_2 += tl.sum(outlier_mask) |
| 504 | write_pos = tl.where(outlier_mask, cumulative_pos, -1) |
| 505 | tl.store( |
| 506 | BUFFER_ROW + write_pos, probs_blk, mask=outlier_mask |
| 507 | ) |
| 508 | |
| 509 | search_range = tl.cast(num_outliers_2, tl.int32) |
| 510 | search_iters = tl.cast( |
| 511 | (num_outliers_2 + BLOCK_SIZE_TRUNC - 1) |
| 512 | // BLOCK_SIZE_TRUNC, |
| 513 | tl.int32, |
| 514 | ) |
| 515 | |
| 516 | # Fourth pass: Calculate BUFFER and get outliers |
| 517 | for i in range(0, search_iters): |
| 518 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 519 | 0, BLOCK_SIZE_TRUNC |
| 520 | ) |
| 521 | mask_n_2 = offs_n < search_range |
| 522 | |
| 523 | probs_blk = tl.load( |
| 524 | BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 |
| 525 | ) |
| 526 | probs_blk = probs_blk / sum_exp_logits |
| 527 | tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n_2) |
| 528 | |
| 529 | max_range = tl.exp(max_logit - max_logit) / sum_exp_logits |
| 530 | min_range = tl.exp(min_logit - max_logit) / sum_exp_logits |
| 531 | |
| 532 | p_pivot = 1.0 |
| 533 | num_iters = 0 |
| 534 | min_larger_prob = 1.0 |
| 535 | num_min_larger = tl.zeros((), dtype=tl.uint32) |
| 536 | p_pivots_sum = 0.0 |
| 537 | |
| 538 | # Fifth passes: Search for p_pivot |
| 539 | found_pivot = 0 |
| 540 | while found_pivot == 0: |
| 541 | p_pivot_0 = (max_range - min_range) * 0.5 + min_range |
| 542 | p_pivots_sum_0 = 0.0 |
| 543 | min_larger_0 = 1.0 |
| 544 | num_min_larger_0 = tl.zeros((), dtype=tl.uint32) |
| 545 | |
| 546 | # Single fused pass: compute p_pivots_sum, |
| 547 | # min_larger, and num_min_larger together. |
| 548 | # See _update_min_larger_stats for the |
| 549 | # tile-level merge logic. |
| 550 | for i in range(0, search_iters): |
| 551 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 552 | 0, BLOCK_SIZE_TRUNC |
| 553 | ) |
| 554 | mask_n_2 = offs_n < search_range |
| 555 | probs_blk = tl.load( |
| 556 | BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 |
| 557 | ) |
| 558 | |
| 559 | above_0 = probs_blk > p_pivot_0 |
| 560 | p_pivots_sum_0 += tl.sum(probs_blk * above_0) |
| 561 | |
| 562 | min_larger_0, num_min_larger_0 = ( |
| 563 | _update_min_larger_stats( |
| 564 | probs_blk, |
| 565 | above_0, |
| 566 | min_larger_0, |
| 567 | num_min_larger_0, |
| 568 | 1.0, |
| 569 | ) |
| 570 | ) |
| 571 | |
| 572 | # Check if the pivot satisfies termination condition |
| 573 | if p_pivots_sum_0 >= p and ( |
| 574 | p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p |
| 575 | ): |
| 576 | p_pivot = p_pivot_0 |
| 577 | min_larger_prob = min_larger_0 |
| 578 | num_min_larger = num_min_larger_0 |
| 579 | p_pivots_sum = p_pivots_sum_0 |
| 580 | found_pivot = 1 |
| 581 | |
| 582 | # Update range |
| 583 | if p_pivots_sum_0 > p: |
| 584 | min_range = p_pivot_0 |
| 585 | elif p_pivots_sum_0 < p: |
| 586 | max_range = p_pivot_0 |
| 587 | |
| 588 | num_iters += 1 |
| 589 | if (max_range - min_range) < 1e-9 or num_iters >= 18: |
| 590 | p_pivot = (max_range + min_range) / 2.0 |
| 591 | min_larger_prob = min_larger_0 |
| 592 | num_min_larger = num_min_larger_0 |
| 593 | p_pivots_sum = p_pivots_sum_0 |
| 594 | found_pivot = 1 |
| 595 | |
| 596 | duplicate_logit = ( |
| 597 | tl.log(min_larger_prob * sum_exp_logits) + max_logit |
| 598 | ) |
| 599 | num_duplicate_logit = num_min_larger |
| 600 | num_keep = num_duplicate_logit - tl.cast( |
| 601 | (p_pivots_sum - p) / min_larger_prob, tl.uint32 |
| 602 | ) |
| 603 | num_kept = tl.zeros((), dtype=tl.uint32) |
| 604 | |
| 605 | # Top-k + Top-p path |
| 606 | final_pivot = tl.log(p_pivot * sum_exp_logits) + max_logit |
| 607 | |
| 608 | if TOPP_ENABLED and final_pivot == -float("inf"): |
| 609 | #### STANDALONE TOP-P SAMPLING #### |
| 610 | # When the split top-p pipeline co-runs (mixed top-k + top-p |
| 611 | # batches), it covers p-only rows (k >= VOCAB_SIZE); skip those |
| 612 | # here. Rows whose top-k was a no-op because they have <= k |
| 613 | # finite logits (e.g. grammar masks) are NOT covered by the |
| 614 | # split pipeline and must get standalone top-p here. |
| 615 | run_standalone = True |
| 616 | if SPLIT_COVERS_PONLY and TOPK_ENABLED: |
| 617 | run_standalone = tl.load(K + row_id) < VOCAB_SIZE |
| 618 | p = tl.load(P + row_id) |
| 619 | if run_standalone and p < 1.0: |
| 620 | # Zeroth pass: Compute avg and std from a sample block |
| 621 | offs = tl.arange(0, BLOCK_SIZE) |
| 622 | mask_n = offs < VOCAB_SIZE |
| 623 | logits_blk0 = tl.load( |
| 624 | LOGITS_ROW + offs, mask=mask_n, other=-float("inf") |
| 625 | ) |
| 626 | # Exclude -inf values (e.g. from grammar bitmasks) from |
| 627 | # statistics to avoid NaN in pivot computation. |
| 628 | finite_mask = (logits_blk0 > -float("inf")) & mask_n |
| 629 | num_finite = tl.sum(finite_mask) |
| 630 | finite_logits = tl.where(finite_mask, logits_blk0, 0.0) |
| 631 | avg_logit = tl.where( |
| 632 | num_finite > 0, tl.sum(finite_logits) / num_finite, 0.0 |
| 633 | ) |
| 634 | sq_avg_logit = tl.where( |
| 635 | num_finite > 0, |
| 636 | tl.sum(finite_logits * finite_logits) / num_finite, |
| 637 | 0.0, |
| 638 | ) |
| 639 | std_logit = tl.sqrt( |
| 640 | tl.maximum(sq_avg_logit - avg_logit * avg_logit, 0.0) |
| 641 | ) |
| 642 | max_sample = avg_logit + std_logit * 10.0 |
| 643 | sum_exp_logits = 0.0 |
| 644 | |
| 645 | # First pass: compute max and min logits and sum_exp_logits |
| 646 | for i in range(0, NUM_TILES): |
| 647 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 648 | mask_n = offs_n < VOCAB_SIZE |
| 649 | logits_blk = tl.load( |
| 650 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 651 | ) |
| 652 | max_logit = tl.maximum(max_logit, tl.max(logits_blk)) |
| 653 | # Exclude -inf from min to keep binary search bounds |
| 654 | # finite (avoids NaN pivots). |
| 655 | finite_blk = tl.where( |
| 656 | logits_blk > -float("inf"), logits_blk, float("inf") |
| 657 | ) |
| 658 | min_logit = tl.minimum(min_logit, tl.min(finite_blk)) |
| 659 | |
| 660 | probs_blk = tl.exp(logits_blk - max_sample) |
| 661 | probs_blk = tl.where(mask_n, probs_blk, 0.0) |
| 662 | sum_exp_logits += tl.sum(probs_blk) |
| 663 | |
| 664 | # If no finite logits exist (all -inf), clamp min to |
| 665 | # max so the search converges to -inf (no masking). |
| 666 | min_logit = tl.minimum(min_logit, max_logit) |
| 667 | |
| 668 | idx = tl.cast(p * 200, tl.int32) |
| 669 | idx = tl.maximum(0, tl.minimum(idx, 199)) |
| 670 | sigma = tl.load(NORMAL_CDF_TO_SIGMA_TABLE + idx) |
| 671 | sigma = sigma + tl.abs(sigma) * -0.25 |
| 672 | outlier_pivot = avg_logit + std_logit * sigma |
| 673 | |
| 674 | outlier_prob = tl.exp(outlier_pivot - max_sample) / sum_exp_logits |
| 675 | sum_outlier_probs = 0.0 |
| 676 | num_outliers = tl.zeros((), dtype=tl.uint32) |
| 677 | |
| 678 | # Second pass: Calculate softmax and gather outliers |
| 679 | for i in range(0, NUM_TILES): |
| 680 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 681 | mask_n = offs_n < VOCAB_SIZE |
| 682 | |
| 683 | probs_blk = tl.load( |
| 684 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 685 | ) |
| 686 | probs_blk = tl.exp(probs_blk - max_sample) |
| 687 | probs_blk = probs_blk / sum_exp_logits |
| 688 | |
| 689 | outlier_mask = (probs_blk > outlier_prob) & mask_n |
| 690 | sum_outlier_probs += tl.sum(outlier_mask * probs_blk) |
| 691 | cumulative_pos = tl.cast( |
| 692 | tl.cumsum(outlier_mask) - 1 + num_outliers, tl.int32 |
| 693 | ) |
| 694 | num_outliers += tl.sum(outlier_mask) |
| 695 | write_pos = tl.where(outlier_mask, cumulative_pos, -1) |
| 696 | tl.store(BUFFER_ROW + write_pos, probs_blk, mask=outlier_mask) |
| 697 | |
| 698 | max_range = tl.exp(max_logit - max_sample) / sum_exp_logits |
| 699 | min_range = tl.exp(min_logit - max_sample) / sum_exp_logits |
| 700 | |
| 701 | p_pivot = 1.0 |
| 702 | num_iters = 0 |
| 703 | min_larger_prob = 1.0 |
| 704 | num_min_larger = tl.zeros((), dtype=tl.uint32) |
| 705 | p_pivots_sum = 0.0 |
| 706 | |
| 707 | # Third pass: Search for p_pivot |
| 708 | if sum_outlier_probs > p: |
| 709 | min_range = outlier_prob |
| 710 | search_range = tl.cast(num_outliers, tl.int32) |
| 711 | search_iters = tl.cast( |
| 712 | (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC, |
| 713 | tl.int32, |
| 714 | ) |
| 715 | |
| 716 | found_pivot = 0 |
| 717 | while found_pivot == 0: |
| 718 | p_pivot_0 = (max_range - min_range) * 0.5 + min_range |
| 719 | p_pivots_sum_0 = 0.0 |
| 720 | min_larger_0 = 1.0 |
| 721 | num_min_larger_0 = tl.zeros((), dtype=tl.uint32) |
| 722 | |
| 723 | # Single fused pass: compute p_pivots_sum, |
| 724 | # min_larger, and num_min_larger together. |
| 725 | # See _update_min_larger_stats for the |
| 726 | # tile-level merge logic. |
| 727 | for i in range(0, search_iters): |
| 728 | offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( |
| 729 | 0, BLOCK_SIZE_TRUNC |
| 730 | ) |
| 731 | mask_n_2 = offs_n < search_range |
| 732 | probs_blk = tl.load( |
| 733 | BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 |
| 734 | ) |
| 735 | |
| 736 | above_0 = probs_blk > p_pivot_0 |
| 737 | p_pivots_sum_0 += tl.sum(probs_blk * above_0) |
| 738 | |
| 739 | min_larger_0, num_min_larger_0 = _update_min_larger_stats( |
| 740 | probs_blk, |
| 741 | above_0, |
| 742 | min_larger_0, |
| 743 | num_min_larger_0, |
| 744 | 1.0, |
| 745 | ) |
| 746 | |
| 747 | # Check if the pivot satisfies termination condition |
| 748 | if ( |
| 749 | p_pivots_sum_0 >= p |
| 750 | and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p |
| 751 | ): |
| 752 | p_pivot = p_pivot_0 |
| 753 | min_larger_prob = min_larger_0 |
| 754 | num_min_larger = num_min_larger_0 |
| 755 | p_pivots_sum = p_pivots_sum_0 |
| 756 | found_pivot = 1 |
| 757 | |
| 758 | # Update range |
| 759 | if p_pivots_sum_0 > p: |
| 760 | min_range = p_pivot_0 |
| 761 | elif p_pivots_sum_0 < p: |
| 762 | max_range = p_pivot_0 |
| 763 | |
| 764 | num_iters += 1 |
| 765 | if (max_range - min_range) < 1e-9 or num_iters >= 18: |
| 766 | p_pivot = (max_range + min_range) / 2.0 |
| 767 | min_larger_prob = min_larger_0 |
| 768 | num_min_larger = num_min_larger_0 |
| 769 | p_pivots_sum = p_pivots_sum_0 |
| 770 | found_pivot = 1 |
| 771 | else: |
| 772 | # Re-populate the buffer with full softmax probabilities |
| 773 | for i in range(0, NUM_TILES): |
| 774 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 775 | mask_n = offs_n < VOCAB_SIZE |
| 776 | |
| 777 | probs_blk = tl.load( |
| 778 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 779 | ) |
| 780 | probs_blk = tl.exp(probs_blk - max_sample) |
| 781 | probs_blk = probs_blk / sum_exp_logits |
| 782 | tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n) |
| 783 | |
| 784 | found_pivot = 0 |
| 785 | while found_pivot == 0: |
| 786 | p_pivot_0 = (max_range - min_range) * 0.5 + min_range |
| 787 | p_pivots_sum_0 = 0.0 |
| 788 | min_larger_0 = 1.0 |
| 789 | num_min_larger_0 = tl.zeros((), dtype=tl.uint32) |
| 790 | |
| 791 | # Single fused pass: compute p_pivots_sum, |
| 792 | # min_larger, and num_min_larger together. |
| 793 | # See _update_min_larger_stats for the |
| 794 | # tile-level merge logic. |
| 795 | for i in range(0, NUM_TILES): |
| 796 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 797 | mask_n = offs_n < VOCAB_SIZE |
| 798 | probs_blk = tl.load( |
| 799 | BUFFER_ROW + offs_n, mask=mask_n, other=0.0 |
| 800 | ) |
| 801 | |
| 802 | above_0 = probs_blk > p_pivot_0 |
| 803 | p_pivots_sum_0 += tl.sum(probs_blk * above_0) |
| 804 | |
| 805 | min_larger_0, num_min_larger_0 = _update_min_larger_stats( |
| 806 | probs_blk, |
| 807 | above_0, |
| 808 | min_larger_0, |
| 809 | num_min_larger_0, |
| 810 | 1.0, |
| 811 | ) |
| 812 | |
| 813 | # Check if the pivot satisfies termination condition |
| 814 | if ( |
| 815 | p_pivots_sum_0 >= p |
| 816 | and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p |
| 817 | ): |
| 818 | p_pivot = p_pivot_0 |
| 819 | min_larger_prob = min_larger_0 |
| 820 | num_min_larger = num_min_larger_0 |
| 821 | p_pivots_sum = p_pivots_sum_0 |
| 822 | found_pivot = 1 |
| 823 | |
| 824 | # Update range |
| 825 | if p_pivots_sum_0 > p: |
| 826 | min_range = p_pivot_0 |
| 827 | elif p_pivots_sum_0 < p: |
| 828 | max_range = p_pivot_0 |
| 829 | |
| 830 | num_iters += 1 |
| 831 | if (max_range - min_range) < 1e-9 or num_iters >= 18: |
| 832 | p_pivot = (max_range + min_range) / 2.0 |
| 833 | min_larger_prob = min_larger_0 |
| 834 | num_min_larger = num_min_larger_0 |
| 835 | p_pivots_sum = p_pivots_sum_0 |
| 836 | found_pivot = 1 |
| 837 | |
| 838 | duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_sample |
| 839 | num_duplicate_logit = num_min_larger |
| 840 | num_keep = num_duplicate_logit - tl.cast( |
| 841 | (p_pivots_sum - p) / min_larger_prob, tl.uint32 |
| 842 | ) |
| 843 | num_kept = tl.zeros((), dtype=tl.uint32) |
| 844 | |
| 845 | # Top-p only path |
| 846 | final_pivot = tl.log(p_pivot * sum_exp_logits) + max_sample |
| 847 | |
| 848 | # Sixth pass: Apply mask and store final output. |
| 849 | # If the pivot >= max logit (or is NaN), no token would |
| 850 | # survive the strict `>` keep_mask. Skip masking. |
| 851 | # Using `not <` instead of `>=` so that NaN is also caught. |
| 852 | if not (final_pivot < max_logit): |
| 853 | final_pivot = -float("inf") |
| 854 | elif final_pivot != -float("inf"): |
| 855 | for i in range(0, NUM_TILES): |
| 856 | offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 857 | mask_n = offs_n < VOCAB_SIZE |
| 858 | logits_blk = tl.load( |
| 859 | LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") |
| 860 | ) |
| 861 | keep_mask = (logits_blk > final_pivot) & mask_n |
| 862 | |
| 863 | # Duplicate logit handling |
| 864 | if num_keep < num_duplicate_logit: |
| 865 | duplicate_mask = ( |
| 866 | tl.abs(logits_blk - duplicate_logit) < 1e-9 |
| 867 | ) & mask_n |
| 868 | duplicate_count = tl.cumsum(duplicate_mask) + num_kept |
| 869 | duplicate_keep_mask = (duplicate_count <= num_keep) & duplicate_mask |
| 870 | duplicate_remove_mask = duplicate_mask & ~duplicate_keep_mask |
| 871 | num_kept += tl.sum(duplicate_keep_mask) |
| 872 | keep_mask = keep_mask & (~duplicate_remove_mask) |
| 873 | |
| 874 | logits_blk = tl.where(keep_mask, logits_blk, MASK_VALUE) |
| 875 | tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n) |
| 876 | |
| 877 | |
| 878 | def _max_sampler_batch_size(vllm_config: Any) -> int: |
| 879 | return vllm_config.scheduler_config.max_num_seqs * ( |
| 880 | vllm_config.num_speculative_tokens + 1 |
| 881 | ) |
| 882 | |
| 883 | |
| 884 | def _topk_topp_warmup_inputs(vllm_config: Any) -> dict[str, Any]: |
| 885 | vocab_size = vllm_config.model_config.get_vocab_size() |
| 886 | split_enabled = current_platform.is_cuda_alike() |
| 887 | mode: Any = WarmupChoices((True, False), (True, True), (False, True)) |
| 888 | logits_stride: Any = WarmupChoices(16, 2) |
| 889 | batch_size: Any = WarmupIntRange(1, _max_sampler_batch_size(vllm_config) + 1) |
| 890 | topk_enabled = mode[0] |
| 891 | topp_enabled = mode[1] |
| 892 | _when(topk_enabled or not (split_enabled and batch_size <= _SPLIT_MAX_BATCH)) |
| 893 | logits = TritonWarmupTensor( |
| 894 | torch.float32, |
| 895 | shape=(batch_size, vocab_size), |
| 896 | strides=(logits_stride, 1), |
| 897 | ) |
| 898 | return dict( |
| 899 | logits=logits, |
| 900 | buffer=TritonWarmupTensor(torch.float32), |
| 901 | percentile_to_std_table=TritonWarmupTensor(torch.float32), |
| 902 | normal_cdf_to_sigma_table=TritonWarmupTensor(torch.float32), |
| 903 | k=TritonWarmupTensor(torch.int32) if topk_enabled else None, |
| 904 | p=TritonWarmupTensor(torch.float32) if topp_enabled else None, |
| 905 | mask_value=float("-inf"), |
| 906 | num_sm=num_compute_units(), |
| 907 | ) |
| 908 | |
| 909 | |
| 910 | @triton_kernel_dispatcher_with_warmup( |
| 911 | kernel=_topk_topp_kernel, |
| 912 | warmup_inputs=_topk_topp_warmup_inputs, |
| 913 | ) |
| 914 | def _topk_topp( |
| 915 | logits: torch.Tensor, |
| 916 | buffer: torch.Tensor, |
| 917 | percentile_to_std_table: torch.Tensor, |
| 918 | normal_cdf_to_sigma_table: torch.Tensor, |
| 919 | k: torch.Tensor | None, |
| 920 | p: torch.Tensor | None, |
| 921 | mask_value: float, |
| 922 | num_sm: int, |
| 923 | ) -> DispatchSpec: |
| 924 | batch_size, vocab_size = logits.shape |
| 925 | topk_enabled = k is not None |
| 926 | topp_enabled = p is not None |
| 927 | k_ptr = k if k is not None else logits |
| 928 | p_ptr = p if p is not None else logits |
| 929 | num_programs = min(num_sm, batch_size) |
| 930 | # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. |
| 931 | # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot |
| 932 | # approximation; use smaller tiles for accurate top-p results. |
| 933 | if logits.device.type == "cpu": |
| 934 | block_size, block_size_trunc, num_warps = 256, 128, 4 |
| 935 | elif logits.device.type == "xpu": |
| 936 | block_size, block_size_trunc, num_warps = 4096, 2048, 4 |
| 937 | else: |
| 938 | block_size, block_size_trunc, num_warps = 8192, 4096, 8 |
| 939 | # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so |
| 940 | # per-tile latency bounds kernel latency, and Triton's default of 4 |
| 941 | # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is |
| 942 | # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not. |
| 943 | split_covers_ponly = ( |
| 944 | current_platform.is_cuda_alike() |
| 945 | and topp_enabled |
| 946 | and batch_size <= _SPLIT_MAX_BATCH |
| 947 | ) |
| 948 | return (num_programs,), dict( |
| 949 | K=k_ptr, |
| 950 | P=p_ptr, |
| 951 | BATCH_SIZE=batch_size, |
| 952 | VOCAB_SIZE=vocab_size, |
| 953 | BLOCK_SIZE=block_size, |
| 954 | BLOCK_SIZE_TRUNC=block_size_trunc, |
| 955 | TOPK_ENABLED=topk_enabled, |
| 956 | TOPP_ENABLED=topp_enabled, |
| 957 | SPLIT_COVERS_PONLY=split_covers_ponly, |
| 958 | num_warps=num_warps, |
| 959 | ) |
| 960 | |
| 961 | |
| 962 | # --------------------------------------------------------------------------- |
| 963 | # Split-row pipeline for small batches. |
| 964 | # |
| 965 | # The monolithic kernel above assigns one program per row, so at small batch |
| 966 | # sizes nearly all SMs idle and every full-vocab sweep is latency-bound |
| 967 | # (~25us each on SM100). For standalone top-p (no top-k truncation) the pivot |
| 968 | # search can hit its iteration cap while sweeping ~the whole vocab per |
| 969 | # iteration, which is slower than the PyTorch sort path. |
| 970 | # |
| 971 | # The pipeline below instead splits each row across S programs and runs the |
| 972 | # pivot search as parallel partial reductions over a fixed number of rounds |
| 973 | # (no host synchronization, CUDA-graph safe). The search state is never |
| 974 | # stored: the per-round combine is a tiny pure function of the stored |
| 975 | # partials, so every program redundantly recomputes the bracket chain from |
| 976 | # them. This is deterministic (same inputs, same code) and far cheaper than |
| 977 | # a vocab sweep, and it saves a kernel launch per combine. The pipeline |
| 978 | # handles only "p-only" rows (p < 1 and no active top-k); rows with an |
| 979 | # active top-k are left to the monolithic kernel, whose truncation path is |
| 980 | # already fast for them. |
| 981 | # --------------------------------------------------------------------------- |
| 982 | |
| 983 | _SPLIT_MAX_BATCH = 64 |
| 984 | _SPLIT_MAX_SPLITS = 32 # power of two |
| 985 | _SPLIT_FANOUT = 8 |
| 986 | _SPLIT_ROUNDS = 5 |
| 987 | |
| 988 | _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {} |
| 989 | |
| 990 | |
| 991 | def _topp_split_count(batch_size: int, num_sm: int) -> int: |
| 992 | splits = 1 |
| 993 | while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm: |
| 994 | splits *= 2 |
| 995 | return splits |
| 996 | |
| 997 | |
| 998 | @triton.jit |
| 999 | def _topp_sb_stats_kernel( |
| 1000 | LOGITS, |
| 1001 | LOGITS_STRIDE_0, |
| 1002 | STATS, |
| 1003 | K, |
| 1004 | P, |
| 1005 | HAS_K: tl.constexpr, |
| 1006 | VOCAB_SIZE: tl.constexpr, |
| 1007 | S: tl.constexpr, |
| 1008 | BLOCK: tl.constexpr, |
| 1009 | ): |
| 1010 | # Per-slice online-softmax partials: (max, sum_exp, min, finite_count). |
| 1011 | pid = tl.program_id(0) |
| 1012 | row = pid // S |
| 1013 | slice_id = pid % S |
| 1014 | p = tl.load(P + row) |
| 1015 | if p >= 1.0: |
| 1016 | return |
| 1017 | if HAS_K and tl.load(K + row) < VOCAB_SIZE: |
| 1018 | return |
| 1019 | slice_len = (VOCAB_SIZE + S - 1) // S |
| 1020 | start = slice_id * slice_len |
| 1021 | end = tl.minimum(start + slice_len, VOCAB_SIZE) |
| 1022 | ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0 |
| 1023 | offs = tl.arange(0, BLOCK) |
| 1024 | m = -float("inf") |
| 1025 | exp_sum = 0.0 |
| 1026 | mn = float("inf") |
| 1027 | cnt = tl.zeros((), dtype=tl.float32) |
| 1028 | for i in range(start, end, BLOCK): |
| 1029 | x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf")) |
| 1030 | finite = x > -float("inf") |
| 1031 | nm = tl.maximum(m, tl.max(x)) |
| 1032 | rescale = tl.where(m == -float("inf"), 0.0, tl.exp(m - nm)) |
| 1033 | e = tl.where(finite, tl.exp(x - nm), 0.0) |
| 1034 | exp_sum = exp_sum * rescale + tl.sum(e) |
| 1035 | m = nm |
| 1036 | mn = tl.minimum(mn, tl.min(tl.where(finite, x, float("inf")))) |
| 1037 | cnt += tl.sum(finite.to(tl.float32)) |
| 1038 | base = (row * S + slice_id) * 4 |
| 1039 | tl.store(STATS + base + 0, m) |
| 1040 | tl.store(STATS + base + 1, exp_sum) |
| 1041 | tl.store(STATS + base + 2, mn) |
| 1042 | tl.store(STATS + base + 3, cnt) |
| 1043 | |
| 1044 | |
| 1045 | def _topp_split_stats_warmup_inputs(vllm_config: Any) -> dict[str, Any]: |
| 1046 | vocab_size = vllm_config.model_config.get_vocab_size() |
| 1047 | num_sm = num_compute_units() |
| 1048 | logits_stride: Any = WarmupChoices(16, 2) |
| 1049 | batch_size: Any = WarmupIntRange( |
| 1050 | 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 |
| 1051 | ) |
| 1052 | has_k: Any = WarmupChoices(False, True) |
| 1053 | logits = TritonWarmupTensor( |
| 1054 | torch.float32, |
| 1055 | shape=(batch_size, vocab_size), |
| 1056 | strides=(logits_stride, 1), |
| 1057 | ) |
| 1058 | return dict( |
| 1059 | logits=logits, |
| 1060 | stats=TritonWarmupTensor(torch.float32), |
| 1061 | k=TritonWarmupTensor(torch.int32) if has_k else None, |
| 1062 | p=TritonWarmupTensor(torch.float32), |
| 1063 | num_sm=num_sm, |
| 1064 | ) |
| 1065 | |
| 1066 | |
| 1067 | @triton_kernel_dispatcher_with_warmup( |
| 1068 | kernel=_topp_sb_stats_kernel, |
| 1069 | warmup_inputs=_topp_split_stats_warmup_inputs, |
| 1070 | ) |
| 1071 | def _topp_split_stats( |
| 1072 | logits: torch.Tensor, |
| 1073 | stats: torch.Tensor, |
| 1074 | k: torch.Tensor | None, |
| 1075 | p: torch.Tensor, |
| 1076 | num_sm: int, |
| 1077 | ) -> DispatchSpec: |
| 1078 | batch_size, vocab_size = logits.shape |
| 1079 | splits = _topp_split_count(batch_size, num_sm) |
| 1080 | return (logits.shape[0] * splits,), dict( |
| 1081 | K=k if k is not None else logits, |
| 1082 | HAS_K=k is not None, |
| 1083 | VOCAB_SIZE=vocab_size, |
| 1084 | S=splits, |
| 1085 | BLOCK=8192, |
| 1086 | num_warps=8, |
| 1087 | ) |
| 1088 | |
| 1089 | |
| 1090 | @triton.jit |
| 1091 | def _topp_sb_row_stats(STATS, row, S: tl.constexpr): |
| 1092 | """Combine per-slice partials into row (max, sum_exp, min, n_finite).""" |
| 1093 | sidx = tl.arange(0, S) |
| 1094 | sb = STATS + row * S * 4 |
| 1095 | m_s = tl.load(sb + sidx * 4 + 0) |
| 1096 | exp_s = tl.load(sb + sidx * 4 + 1) |
| 1097 | mn_s = tl.load(sb + sidx * 4 + 2) |
| 1098 | c_s = tl.load(sb + sidx * 4 + 3) |
| 1099 | M = tl.max(m_s) |
| 1100 | w = tl.where(m_s == -float("inf"), 0.0, tl.exp(m_s - M)) |
| 1101 | Z = tl.sum(exp_s * w) |
| 1102 | return M, Z, tl.min(mn_s), tl.sum(c_s) |
| 1103 | |
| 1104 | |
| 1105 | @triton.jit |
| 1106 | def _topp_sb_ladder(lo, hi, F: tl.constexpr): |
| 1107 | """Geometric ladder of F pivots inside (lo, hi).""" |
| 1108 | fidx = tl.arange(0, F).to(tl.float32) |
| 1109 | return lo * tl.exp(tl.log(hi / lo) * (fidx + 1.0) / (F + 1.0)) |
| 1110 | |
| 1111 | |
| 1112 | @triton.jit |
| 1113 | def _topp_sb_combine( |
| 1114 | PART_BASE, |
| 1115 | Z, |
| 1116 | p, |
| 1117 | lo, |
| 1118 | hi, |
| 1119 | best, |
| 1120 | best_mass, |
| 1121 | best_minl, |
| 1122 | best_nmin, |
| 1123 | S: tl.constexpr, |
| 1124 | F: tl.constexpr, |
| 1125 | IS_LAST: tl.constexpr, |
| 1126 | ): |
| 1127 | """Combine one round's partials and advance the search state. |
| 1128 | |
| 1129 | Returns (done, nomask, pivot, dup, numdup, numkeep, lo, hi, best, |
| 1130 | best_mass, best_minl, best_nmin, sel_fidx, best_sel). `best*` tracks the |
| 1131 | tightest fallback: the largest evaluated pivot whose above-mass is >= p, |
| 1132 | which guarantees the top-p mass invariant if the search fails to |
| 1133 | converge. Pivots are recomputed from (lo, hi) via the same ladder the |
| 1134 | sweep used, so partials line up with pivots without storing them. |
| 1135 | `sel_fidx` / `best_sel` report which ladder slot produced `dup` / |
| 1136 | updated `best` this round (-1 if none), so callers can recover the |
| 1137 | per-slice (minl, nmin) partials behind the boundary value for |
| 1138 | deterministic duplicate handling. |
| 1139 | """ |
| 1140 | fidx = tl.arange(0, F) |
| 1141 | sidx = tl.arange(0, S) |
| 1142 | offs2 = sidx[:, None] * (F * 3) + fidx[None, :] * 3 |
| 1143 | mass_s = tl.load(PART_BASE + offs2 + 0) |
| 1144 | minl_s = tl.load(PART_BASE + offs2 + 1) |
| 1145 | nmin_s = tl.load(PART_BASE + offs2 + 2) |
| 1146 | mass = tl.sum(mass_s, axis=0) |
| 1147 | gmin = tl.min(minl_s, axis=0) |
| 1148 | nmin = tl.sum(tl.where(tl.abs(minl_s - gmin[None, :]) < 1e-9, nmin_s, 0.0), axis=0) |
| 1149 | pivs = _topp_sb_ladder(lo, hi, F) |
| 1150 | |
| 1151 | ok = (mass >= p) & (mass - gmin * nmin < p) |
| 1152 | ok_idx = tl.max(tl.where(ok, fidx, -1).to(tl.float32)) |
| 1153 | ge = mass >= p |
| 1154 | ge_idx = tl.max(tl.where(ge, fidx, -1).to(tl.float32)) |
| 1155 | |
| 1156 | best_sel = -1.0 |
| 1157 | if ge_idx >= 0: |
| 1158 | bmask = fidx == ge_idx.to(tl.int32) |
| 1159 | bp = tl.sum(tl.where(bmask, pivs, 0.0)) |
| 1160 | if bp > best: |
| 1161 | best = bp |
| 1162 | best_mass = tl.sum(tl.where(bmask, mass, 0.0)) |
| 1163 | best_minl = tl.sum(tl.where(bmask, gmin, 0.0)) |
| 1164 | best_nmin = tl.sum(tl.where(bmask, nmin, 0.0)) |
| 1165 | best_sel = ge_idx |
| 1166 | |
| 1167 | done = 0.0 |
| 1168 | nomask = 0.0 |
| 1169 | pivot = 0.0 |
| 1170 | sel_mass = 0.0 |
| 1171 | dup = 1.0 |
| 1172 | numdup = 1.0 |
| 1173 | numkeep = 1.0 |
| 1174 | sel_fidx = -1.0 |
| 1175 | if ok_idx >= 0: |
| 1176 | bmask = fidx == ok_idx.to(tl.int32) |
| 1177 | pivot = tl.sum(tl.where(bmask, pivs, 0.0)) |
| 1178 | sel_mass = tl.sum(tl.where(bmask, mass, 0.0)) |
| 1179 | dup = tl.sum(tl.where(bmask, gmin, 0.0)) |
| 1180 | numdup = tl.sum(tl.where(bmask, nmin, 0.0)) |
| 1181 | sel_fidx = ok_idx |
| 1182 | done = 1.0 |
| 1183 | else: |
| 1184 | new_lo = tl.maximum(lo, tl.max(tl.where(ge, pivs, -float("inf")))) |
| 1185 | new_hi = tl.minimum(hi, tl.min(tl.where(~ge, pivs, float("inf")))) |
| 1186 | if IS_LAST or not (new_hi > new_lo * (1.0 + 1e-7)): |
| 1187 | # Out of rounds or the range collapsed: fall back to the |
| 1188 | # tightest pivot with mass >= p, or keep the whole row if no |
| 1189 | # pivot qualified. |
| 1190 | if best > 0.0: |
| 1191 | pivot = best |
| 1192 | sel_mass = best_mass |
| 1193 | dup = best_minl |
| 1194 | numdup = best_nmin |
| 1195 | else: |
| 1196 | nomask = 1.0 |
| 1197 | done = 1.0 |
| 1198 | lo = new_lo |
| 1199 | hi = new_hi |
| 1200 | if (done != 0.0) and (nomask == 0.0): |
| 1201 | # Mirror the monolithic guard: a pivot at/above the max logit would |
| 1202 | # mask everything under the strict `>` comparison. |
| 1203 | nomask = tl.where(pivot * Z >= 1.0, 1.0, 0.0) |
| 1204 | # Duplicate (boundary-value) handling, same formula as the |
| 1205 | # monolithic kernel: keep only some of the boundary duplicates. |
| 1206 | numkeep = numdup - tl.cast(tl.maximum(sel_mass - p, 0.0) / dup, tl.uint32).to( |
| 1207 | tl.float32 |
| 1208 | ) |
| 1209 | numkeep = tl.minimum(tl.maximum(numkeep, 1.0), numdup) |
| 1210 | return ( |
| 1211 | done, |
| 1212 | nomask, |
| 1213 | pivot, |
| 1214 | dup, |
| 1215 | numdup, |
| 1216 | numkeep, |
| 1217 | lo, |
| 1218 | hi, |
| 1219 | best, |
| 1220 | best_mass, |
| 1221 | best_minl, |
| 1222 | best_nmin, |
| 1223 | sel_fidx, |
| 1224 | best_sel, |
| 1225 | ) |
| 1226 | |
| 1227 | |
| 1228 | @triton.jit(do_not_specialize_on_alignment=["ROUND"]) |
| 1229 | def _topp_sb_step_kernel( |
| 1230 | LOGITS, |
| 1231 | LOGITS_STRIDE_0, |
| 1232 | STATS, |
| 1233 | PARTS, |
| 1234 | K, |
| 1235 | P, |
| 1236 | ROUND, |
| 1237 | HAS_K: tl.constexpr, |
| 1238 | S: tl.constexpr, |
| 1239 | F: tl.constexpr, |
| 1240 | NUM_ROUNDS: tl.constexpr, |
| 1241 | VOCAB_SIZE: tl.constexpr, |
| 1242 | BLOCK: tl.constexpr, |
| 1243 | ): |
| 1244 | pid = tl.program_id(0) |
| 1245 | row = pid // S |
| 1246 | slice_id = pid % S |
| 1247 | p = tl.load(P + row) |
| 1248 | if p >= 1.0: |
| 1249 | return |
| 1250 | if HAS_K and tl.load(K + row) < VOCAB_SIZE: |
| 1251 | return |
| 1252 | M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S) |
| 1253 | if (Z <= 0.0) or (num_finite < 1.0): |
| 1254 | return |
| 1255 | max_prob = 1.0 / Z |
| 1256 | min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30) |
| 1257 | if max_prob <= min_prob * (1.0 + 1e-6): |
| 1258 | # (Nearly) uniform row: no pivot can differentiate; keep everything. |
| 1259 | return |
| 1260 | # Replay the combine chain for rounds 0..ROUND-1 (deterministic; each |
| 1261 | # is a tiny pure function of the stored partials). Partials from rounds |
| 1262 | # past a row's convergence are stale, so the chain stops at `done`. |
| 1263 | lo = min_prob |
| 1264 | hi = max_prob |
| 1265 | best = -1.0 |
| 1266 | best_mass = 0.0 |
| 1267 | best_minl = 1.0 |
| 1268 | best_nmin = 0.0 |
| 1269 | done = 0.0 |
| 1270 | for j in range(ROUND): |
| 1271 | if done == 0.0: |
| 1272 | ( |
| 1273 | done, |
| 1274 | nomask, |
| 1275 | pivot, |
| 1276 | dup, |
| 1277 | numdup, |
| 1278 | numkeep, |
| 1279 | lo, |
| 1280 | hi, |
| 1281 | best, |
| 1282 | best_mass, |
| 1283 | best_minl, |
| 1284 | best_nmin, |
| 1285 | sel_fidx, |
| 1286 | best_sel, |
| 1287 | ) = _topp_sb_combine( |
| 1288 | PARTS + (row * NUM_ROUNDS + j) * S * F * 3, |
| 1289 | Z, |
| 1290 | p, |
| 1291 | lo, |
| 1292 | hi, |
| 1293 | best, |
| 1294 | best_mass, |
| 1295 | best_minl, |
| 1296 | best_nmin, |
| 1297 | S=S, |
| 1298 | F=F, |
| 1299 | IS_LAST=False, |
| 1300 | ) |
| 1301 | if done != 0.0: |
| 1302 | return |
| 1303 | # Sweep this slice against the ladder pivots derived from (lo, hi). |
| 1304 | pivs = _topp_sb_ladder(lo, hi, F) |
| 1305 | slice_len = (VOCAB_SIZE + S - 1) // S |
| 1306 | start = slice_id * slice_len |
| 1307 | end = tl.minimum(start + slice_len, VOCAB_SIZE) |
| 1308 | ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0 |
| 1309 | offs = tl.arange(0, BLOCK) |
| 1310 | fidx = tl.arange(0, F) |
| 1311 | mass = tl.zeros([F], dtype=tl.float32) |
| 1312 | minl = tl.full([F], 1.0, dtype=tl.float32) |
| 1313 | nmin = tl.zeros([F], dtype=tl.float32) |
| 1314 | for i in range(start, end, BLOCK): |
| 1315 | x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf")) |
| 1316 | pr = tl.exp(x - M) / Z |
| 1317 | pr2 = tl.broadcast_to(pr[None, :], (F, BLOCK)) |
| 1318 | above = pr2 > pivs[:, None] |
| 1319 | mass += tl.sum(tl.where(above, pr2, 0.0), axis=1) |
| 1320 | cand = tl.where(above, pr2, float("inf")) |
| 1321 | tmin = tl.min(cand, axis=1) |
| 1322 | teq = above & (tl.abs(pr2 - tmin[:, None]) < 1e-9) |
| 1323 | tcnt = tl.sum(teq.to(tl.float32), axis=1) |
| 1324 | is_new = tmin < minl |
| 1325 | is_same = tl.abs(tmin - minl) < 1e-9 |
| 1326 | nmin = tl.where(is_new, tcnt, nmin + tl.where(is_same, tcnt, 0.0)) |
| 1327 | minl = tl.minimum(minl, tmin) |
| 1328 | pbase = PARTS + ((row * NUM_ROUNDS + ROUND) * S + slice_id) * F * 3 |
| 1329 | tl.store(pbase + fidx * 3 + 0, mass) |
| 1330 | tl.store(pbase + fidx * 3 + 1, minl) |
| 1331 | tl.store(pbase + fidx * 3 + 2, nmin) |
| 1332 | |
| 1333 | |
| 1334 | def _topp_split_step_warmup_inputs(vllm_config: Any) -> dict[str, Any]: |
| 1335 | vocab_size = vllm_config.model_config.get_vocab_size() |
| 1336 | num_sm = num_compute_units() |
| 1337 | logits_stride: Any = WarmupChoices(16, 2) |
| 1338 | batch_size: Any = WarmupIntRange( |
| 1339 | 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 |
| 1340 | ) |
| 1341 | round: Any = WarmupIntRange(0, _SPLIT_ROUNDS) |
| 1342 | has_k: Any = WarmupChoices(False, True) |
| 1343 | logits = TritonWarmupTensor( |
| 1344 | torch.float32, |
| 1345 | shape=(batch_size, vocab_size), |
| 1346 | strides=(logits_stride, 1), |
| 1347 | ) |
| 1348 | return dict( |
| 1349 | logits=logits, |
| 1350 | stats=TritonWarmupTensor(torch.float32), |
| 1351 | parts=TritonWarmupTensor(torch.float32), |
| 1352 | k=TritonWarmupTensor(torch.int32) if has_k else None, |
| 1353 | p=TritonWarmupTensor(torch.float32), |
| 1354 | round=round, |
| 1355 | num_sm=num_sm, |
| 1356 | ) |
| 1357 | |
| 1358 | |
| 1359 | @triton_kernel_dispatcher_with_warmup( |
| 1360 | kernel=_topp_sb_step_kernel, |
| 1361 | warmup_inputs=_topp_split_step_warmup_inputs, |
| 1362 | ) |
| 1363 | def _topp_split_step( |
| 1364 | logits: torch.Tensor, |
| 1365 | stats: torch.Tensor, |
| 1366 | parts: torch.Tensor, |
| 1367 | k: torch.Tensor | None, |
| 1368 | p: torch.Tensor, |
| 1369 | round: int, |
| 1370 | num_sm: int, |
| 1371 | ) -> DispatchSpec: |
| 1372 | batch_size, vocab_size = logits.shape |
| 1373 | splits = _topp_split_count(batch_size, num_sm) |
| 1374 | return (logits.shape[0] * splits,), dict( |
| 1375 | K=k if k is not None else logits, |
| 1376 | HAS_K=k is not None, |
| 1377 | S=splits, |
| 1378 | F=_SPLIT_FANOUT, |
| 1379 | NUM_ROUNDS=_SPLIT_ROUNDS, |
| 1380 | VOCAB_SIZE=vocab_size, |
| 1381 | BLOCK=2048, |
| 1382 | num_warps=8, |
| 1383 | ) |
| 1384 | |
| 1385 | |
| 1386 | @triton.jit |
| 1387 | def _topp_sb_mask_kernel( |
| 1388 | LOGITS, |
| 1389 | LOGITS_STRIDE_0, |
| 1390 | STATS, |
| 1391 | PARTS, |
| 1392 | K, |
| 1393 | P, |
| 1394 | HAS_K: tl.constexpr, |
| 1395 | MASK_VALUE: tl.constexpr, |
| 1396 | S: tl.constexpr, |
| 1397 | F: tl.constexpr, |
| 1398 | NUM_ROUNDS: tl.constexpr, |
| 1399 | VOCAB_SIZE: tl.constexpr, |
| 1400 | BLOCK: tl.constexpr, |
| 1401 | ): |
| 1402 | pid = tl.program_id(0) |
| 1403 | row = pid // S |
| 1404 | slice_id = pid % S |
| 1405 | p = tl.load(P + row) |
| 1406 | if p >= 1.0: |
| 1407 | return |
| 1408 | if HAS_K and tl.load(K + row) < VOCAB_SIZE: |
| 1409 | return |
| 1410 | M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S) |
| 1411 | if (Z <= 0.0) or (num_finite < 1.0): |
| 1412 | return |
| 1413 | max_prob = 1.0 / Z |
| 1414 | min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30) |
| 1415 | if max_prob <= min_prob * (1.0 + 1e-6): |
| 1416 | return # (nearly) uniform row: keep everything |
| 1417 | # Replay the full combine chain; the last round finalizes with fallback |
| 1418 | # semantics. Rows that converged early stop the chain at that round. |
| 1419 | # `fin` tracks the (round, ladder slot) whose per-slice partials define |
| 1420 | # the boundary value, for deterministic duplicate handling below. |
| 1421 | lo = min_prob |
| 1422 | hi = max_prob |
| 1423 | best = -1.0 |
| 1424 | best_mass = 0.0 |
| 1425 | best_minl = 1.0 |
| 1426 | best_nmin = 0.0 |
| 1427 | done = 0.0 |
| 1428 | nomask = 0.0 |
| 1429 | pivot = 0.0 |
| 1430 | dup = 1.0 |
| 1431 | numdup = 1.0 |
| 1432 | numkeep = 1.0 |
| 1433 | fin = -1.0 |
| 1434 | best_fin = -1.0 |
| 1435 | for j in range(NUM_ROUNDS): |
| 1436 | if done == 0.0: |
| 1437 | ( |
| 1438 | done, |
| 1439 | nomask, |
| 1440 | pivot, |
| 1441 | dup, |
| 1442 | numdup, |
| 1443 | numkeep, |
| 1444 | lo, |
| 1445 | hi, |
| 1446 | best, |
| 1447 | best_mass, |
| 1448 | best_minl, |
| 1449 | best_nmin, |
| 1450 | sel_fidx, |
| 1451 | best_sel, |
| 1452 | ) = _topp_sb_combine( |
| 1453 | PARTS + (row * NUM_ROUNDS + j) * S * F * 3, |
| 1454 | Z, |
| 1455 | p, |
| 1456 | lo, |
| 1457 | hi, |
| 1458 | best, |
| 1459 | best_mass, |
| 1460 | best_minl, |
| 1461 | best_nmin, |
| 1462 | S=S, |
| 1463 | F=F, |
| 1464 | IS_LAST=j == NUM_ROUNDS - 1, |
| 1465 | ) |
| 1466 | if sel_fidx >= 0.0: |
| 1467 | fin = j * F + sel_fidx |
| 1468 | if best_sel >= 0.0: |
| 1469 | best_fin = j * F + best_sel |
| 1470 | if nomask != 0.0: |
| 1471 | return |
| 1472 | if fin < 0.0: |
| 1473 | fin = best_fin |
| 1474 | logZ = tl.log(Z) |
| 1475 | pivot_logit = tl.log(pivot) + logZ + M |
| 1476 | # numdup/numkeep are exact small integers held in fp32. |
| 1477 | ties = numkeep < numdup |
| 1478 | num_preceding = 0.0 |
| 1479 | if ties: |
| 1480 | # Deterministic, index-ordered duplicate budget: per-slice counts of |
| 1481 | # the boundary value come from the finalizing round's partials (a |
| 1482 | # slice holds copies of `dup` iff its minl matches `dup`, with the |
| 1483 | # count in nmin), so each slice knows exactly how many copies were |
| 1484 | # kept by lower-indexed slices. |
| 1485 | fin_round = tl.cast(fin / F, tl.int32) |
| 1486 | fin_slot = tl.cast(fin - tl.cast(fin_round, tl.float32) * F, tl.int32) |
| 1487 | pbase = PARTS + ((row * NUM_ROUNDS + fin_round) * S * F + fin_slot) * 3 |
| 1488 | sidx = tl.arange(0, S) |
| 1489 | minl_s = tl.load(pbase + sidx * (F * 3) + 1) |
| 1490 | nmin_s = tl.load(pbase + sidx * (F * 3) + 2) |
| 1491 | dupc_s = tl.where(tl.abs(minl_s - dup) < 1e-9, nmin_s, 0.0) |
| 1492 | num_preceding = tl.sum(tl.where(sidx < slice_id, dupc_s, 0.0)) |
| 1493 | slice_len = (VOCAB_SIZE + S - 1) // S |
| 1494 | start = slice_id * slice_len |
| 1495 | end = tl.minimum(start + slice_len, VOCAB_SIZE) |
| 1496 | ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0 |
| 1497 | offs = tl.arange(0, BLOCK) |
| 1498 | for i in range(start, end, BLOCK): |
| 1499 | mask_n = i + offs < end |
| 1500 | x = tl.load(ROW + i + offs, mask=mask_n, other=-float("inf")) |
| 1501 | keep = x > pivot_logit |
| 1502 | if ties: |
| 1503 | # Match the boundary value in probability space with exact |
| 1504 | # equality: `dup` is a probability computed by the step kernels |
| 1505 | # as exp(x - M) / Z, so recomputing it here with the same ops |
| 1506 | # reproduces it bitwise on any platform. (The log/exp round-trip |
| 1507 | # to logit space loses ~1e-6 to cancellation, breaking any |
| 1508 | # absolute tolerance.) |
| 1509 | dmask = tl.exp(x - M) / Z == dup |
| 1510 | cum = tl.cumsum(dmask.to(tl.int32)).to(tl.float32) + num_preceding |
| 1511 | keep_dup = dmask & (cum <= numkeep) |
| 1512 | keep = keep & (~dmask | keep_dup) |
| 1513 | tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n) |
| 1514 | |
| 1515 | |
| 1516 | def _topp_split_mask_warmup_inputs(vllm_config: Any) -> dict[str, Any]: |
| 1517 | vocab_size = vllm_config.model_config.get_vocab_size() |
| 1518 | num_sm = num_compute_units() |
| 1519 | logits_stride: Any = WarmupChoices(16, 2) |
| 1520 | batch_size: Any = WarmupIntRange( |
| 1521 | 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 |
| 1522 | ) |
| 1523 | has_k: Any = WarmupChoices(False, True) |
| 1524 | logits = TritonWarmupTensor( |
| 1525 | torch.float32, |
| 1526 | shape=(batch_size, vocab_size), |
| 1527 | strides=(logits_stride, 1), |
| 1528 | ) |
| 1529 | return dict( |
| 1530 | logits=logits, |
| 1531 | stats=TritonWarmupTensor(torch.float32), |
| 1532 | parts=TritonWarmupTensor(torch.float32), |
| 1533 | k=TritonWarmupTensor(torch.int32) if has_k else None, |
| 1534 | p=TritonWarmupTensor(torch.float32), |
| 1535 | mask_value=float("-inf"), |
| 1536 | num_sm=num_sm, |
| 1537 | ) |
| 1538 | |
| 1539 | |
| 1540 | @triton_kernel_dispatcher_with_warmup( |
| 1541 | kernel=_topp_sb_mask_kernel, |
| 1542 | warmup_inputs=_topp_split_mask_warmup_inputs, |
| 1543 | ) |
| 1544 | def _topp_split_mask( |
| 1545 | logits: torch.Tensor, |
| 1546 | stats: torch.Tensor, |
| 1547 | parts: torch.Tensor, |
| 1548 | k: torch.Tensor | None, |
| 1549 | p: torch.Tensor, |
| 1550 | mask_value: float, |
| 1551 | num_sm: int, |
| 1552 | ) -> DispatchSpec: |
| 1553 | batch_size, vocab_size = logits.shape |
| 1554 | splits = _topp_split_count(batch_size, num_sm) |
| 1555 | return (logits.shape[0] * splits,), dict( |
| 1556 | K=k if k is not None else logits, |
| 1557 | HAS_K=k is not None, |
| 1558 | S=splits, |
| 1559 | F=_SPLIT_FANOUT, |
| 1560 | NUM_ROUNDS=_SPLIT_ROUNDS, |
| 1561 | VOCAB_SIZE=vocab_size, |
| 1562 | BLOCK=8192, |
| 1563 | num_warps=8, |
| 1564 | ) |
| 1565 | |
| 1566 | |
| 1567 | def _apply_topp_split( |
| 1568 | logits: torch.Tensor, |
| 1569 | k: torch.Tensor | None, |
| 1570 | p: torch.Tensor, |
| 1571 | mask_value: float, |
| 1572 | num_sm: int, |
| 1573 | ) -> None: |
| 1574 | """Split-row top-p pipeline for small batches; masks logits in place.""" |
| 1575 | ws = _TRITON_SPLIT_CACHE.get(logits.device) |
| 1576 | if ws is None: |
| 1577 | ws = { |
| 1578 | "stats": logits.new_empty((_SPLIT_MAX_BATCH * _SPLIT_MAX_SPLITS, 4)), |
| 1579 | "parts": logits.new_empty( |
| 1580 | (_SPLIT_MAX_BATCH, _SPLIT_ROUNDS, _SPLIT_MAX_SPLITS, _SPLIT_FANOUT, 3) |
| 1581 | ), |
| 1582 | } |
| 1583 | _TRITON_SPLIT_CACHE[logits.device] = ws |
| 1584 | |
| 1585 | _topp_split_stats( |
| 1586 | logits, |
| 1587 | ws["stats"], |
| 1588 | k, |
| 1589 | p, |
| 1590 | num_sm, |
| 1591 | ) |
| 1592 | for round_i in range(_SPLIT_ROUNDS): |
| 1593 | _topp_split_step( |
| 1594 | logits, |
| 1595 | ws["stats"], |
| 1596 | ws["parts"], |
| 1597 | k, |
| 1598 | p, |
| 1599 | round_i, |
| 1600 | num_sm, |
| 1601 | ) |
| 1602 | _topp_split_mask( |
| 1603 | logits, |
| 1604 | ws["stats"], |
| 1605 | ws["parts"], |
| 1606 | k, |
| 1607 | p, |
| 1608 | mask_value, |
| 1609 | num_sm, |
| 1610 | ) |
| 1611 | |
| 1612 | |
| 1613 | def apply_top_k_top_p_triton( |
| 1614 | logits: torch.Tensor, |
| 1615 | k: torch.Tensor | None, |
| 1616 | p: torch.Tensor | None, |
| 1617 | mask_value: float = float("-inf"), |
| 1618 | ) -> torch.Tensor: |
| 1619 | """Apply combined top-k and top-p masking using Triton. |
| 1620 | |
| 1621 | Top-k is applied first (by logit value), then top-p is applied |
| 1622 | to the remaining k values (by probability). |
| 1623 | |
| 1624 | Args: |
| 1625 | logits: [batch_size, vocab_size] float32 tensor. The returned tensor |
| 1626 | may alias this input or be a new contiguous tensor for unsupported |
| 1627 | layouts. |
| 1628 | k: [batch_size] int32 tensor of top-k values per row, or None to disable top-k |
| 1629 | p: [batch_size] float32 tensor of top-p values per row (0 to 1), |
| 1630 | or None to disable top-p |
| 1631 | mask_value: Value for masked positions (default: -inf) |
| 1632 | |
| 1633 | Returns: |
| 1634 | The masked logits tensor. It may or may not be modified in-place. |
| 1635 | |
| 1636 | """ |
| 1637 | assert logits.ndim == 2 |
| 1638 | assert logits.dtype == torch.float32 |
| 1639 | batch_size, vocab_size = logits.shape |
| 1640 | topk_enabled = k is not None |
| 1641 | topp_enabled = p is not None |
| 1642 | |
| 1643 | if batch_size == 0 or not (topk_enabled or topp_enabled): |
| 1644 | return logits |
| 1645 | |
| 1646 | # The Triton kernel supports arbitrary row strides, but it still assumes |
| 1647 | # the vocab dimension is laid out contiguously within each row. |
| 1648 | if logits.stride(1) != 1: |
| 1649 | logits = logits.contiguous() |
| 1650 | |
| 1651 | if k is not None: |
| 1652 | assert k.ndim == 1 and k.shape[0] == batch_size |
| 1653 | k_ptr = k.to(torch.int32) |
| 1654 | else: |
| 1655 | k_ptr = logits # Dummy pointer (won't be read) |
| 1656 | |
| 1657 | if p is not None: |
| 1658 | assert p.ndim == 1 and p.shape[0] == batch_size |
| 1659 | p_ptr = p.to(torch.float32) |
| 1660 | else: |
| 1661 | p_ptr = logits # Dummy pointer (won't be read) |
| 1662 | |
| 1663 | num_sm = num_compute_units(logits.device.index) |
| 1664 | |
| 1665 | # At small batch sizes the monolithic kernel's standalone top-p path is |
| 1666 | # latency-bound (one program per row, serial vocab sweeps), so use the |
| 1667 | # split-row pipeline for rows without an active top-k. Rows with an |
| 1668 | # active top-k still go through the monolithic kernel (its truncation |
| 1669 | # path is already fast). The monolithic's standalone top-p branch then |
| 1670 | # covers only rows whose top-k is a no-op (k < vocab but <= k finite |
| 1671 | # logits, e.g. grammar masks), which the split pipeline skips. |
| 1672 | use_split = ( |
| 1673 | topp_enabled and batch_size <= _SPLIT_MAX_BATCH and logits.device.type == "cuda" |
| 1674 | ) |
| 1675 | if use_split and not topk_enabled: |
| 1676 | _apply_topp_split(logits, None, p_ptr, mask_value, num_sm) |
| 1677 | return logits |
| 1678 | |
| 1679 | NUM_PROGRAMS = min(num_sm, batch_size) |
| 1680 | |
| 1681 | # Cache per-Triton Program buffer on each device. |
| 1682 | buf_key = (logits.device, logits.dtype, vocab_size) |
| 1683 | buffer = _TRITON_BUFFER_CACHE.get(buf_key) |
| 1684 | if buffer is None or buffer.shape[0] < NUM_PROGRAMS: |
| 1685 | size = min(next_power_of_2(NUM_PROGRAMS), num_sm) |
| 1686 | buffer = logits.new_empty((size, vocab_size)) |
| 1687 | _TRITON_BUFFER_CACHE[buf_key] = buffer |
| 1688 | if buffer.shape[0] > NUM_PROGRAMS: |
| 1689 | buffer = buffer[:NUM_PROGRAMS] |
| 1690 | |
| 1691 | # Cache lookup table entries on each device. |
| 1692 | tables = _TRITON_TABLE_CACHE.get(logits.device) |
| 1693 | if tables is None: |
| 1694 | with gpu_sync_allowed(): |
| 1695 | normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) |
| 1696 | percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) |
| 1697 | _TRITON_TABLE_CACHE[logits.device] = ( |
| 1698 | normal_cdf_to_sigma_table, |
| 1699 | percentile_to_std_table, |
| 1700 | ) |
| 1701 | else: |
| 1702 | normal_cdf_to_sigma_table, percentile_to_std_table = tables |
| 1703 | |
| 1704 | _topk_topp( |
| 1705 | logits, |
| 1706 | buffer, |
| 1707 | percentile_to_std_table, |
| 1708 | normal_cdf_to_sigma_table, |
| 1709 | k_ptr if topk_enabled else None, |
| 1710 | p_ptr if topp_enabled else None, |
| 1711 | mask_value, |
| 1712 | num_sm, |
| 1713 | ) |
| 1714 | if use_split: |
| 1715 | _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm) |
| 1716 | |
| 1717 | return logits |
| 1718 | |
| 1719 | |
| 1720 | def reset_buffer_cache(): |
| 1721 | _TRITON_BUFFER_CACHE.clear() |
| 1722 | _TRITON_TABLE_CACHE.clear() |
| 1723 | _TRITON_SPLIT_CACHE.clear() |
| 1724 | torch.accelerator.empty_cache() |