| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Top-k kernels for the DSA sparse attention indexer.""" |
| 4 | |
| 5 | import functools |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm import _custom_ops as ops |
| 10 | from vllm.config import get_current_vllm_config |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.utils.flashinfer import has_flashinfer |
| 13 | from vllm.v1.worker.workspace import current_workspace_manager |
| 14 | |
| 15 | RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 |
| 16 | |
| 17 | # cooperative_topk's hard row limit: one cluster wave covers at most this many |
| 18 | # rows. "auto" uses cooperative_topk for every batch within the limit and falls |
| 19 | # back to persistent_topk past it; explicit ``cooperative`` requests are |
| 20 | # validated against the same bound. |
| 21 | AUTO_COOPERATIVE_MAX_ROWS = 64 |
| 22 | |
| 23 | # --------------------------------------------------------------------------- |
| 24 | # DeepSelect (vllm._deepselect_C) |
| 25 | # --------------------------------------------------------------------------- |
| 26 | |
| 27 | # Matches the -1 fill convention used for topk_indices_buffer elsewhere. |
| 28 | IDX_OOB_FILL_VALUE = -1 |
| 29 | |
| 30 | try: |
| 31 | import vllm._deepselect_C # noqa: F401 (registers torch.ops.deep_select) |
| 32 | except ImportError as e: |
| 33 | from vllm.logger import init_logger |
| 34 | |
| 35 | init_logger(__name__).warning( |
| 36 | "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e |
| 37 | ) |
| 38 | |
| 39 | |
| 40 | @functools.lru_cache(maxsize=1) |
| 41 | def get_deep_select_stride_requirement() -> tuple[int, int]: |
| 42 | """Stride alignment requirement (input, output) in bytes.""" |
| 43 | return torch.ops.deep_select.get_alignment_requirement() |
| 44 | |
| 45 | |
| 46 | def is_deep_select_supported(input: torch.Tensor, topk: int) -> bool: |
| 47 | """Whether the kernel accepts this input (dtype/stride/topk constraints).""" |
| 48 | return ( |
| 49 | topk <= 4096 |
| 50 | and input.shape[1] < 2**23 |
| 51 | and input.dtype in (torch.float32, torch.bfloat16) |
| 52 | and input.stride(1) == 1 |
| 53 | and input.stride(0) |
| 54 | * input.element_size() |
| 55 | % get_deep_select_stride_requirement()[0] |
| 56 | == 0 |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | def _get_empty_and_aligned_tensor( |
| 61 | dim0: int, dim1: int, device: torch.device, dtype: torch.dtype |
| 62 | ) -> torch.Tensor: |
| 63 | """Tensor with shape (dim0, dim1) whose stride(0) is 32B-aligned.""" |
| 64 | output_stride_requirement = get_deep_select_stride_requirement()[1] // ( |
| 65 | dtype.itemsize |
| 66 | ) |
| 67 | assert output_stride_requirement > 0 |
| 68 | dim1_rounded = ( |
| 69 | (dim1 + output_stride_requirement - 1) // output_stride_requirement |
| 70 | ) * output_stride_requirement |
| 71 | return torch.empty((dim0, dim1_rounded), device=device, dtype=dtype)[:, :dim1] |
| 72 | |
| 73 | |
| 74 | def deep_select_topk( |
| 75 | input: torch.Tensor, |
| 76 | topk: int, |
| 77 | end: torch.Tensor | None = None, |
| 78 | output_idx: torch.Tensor | None = None, |
| 79 | indices_dtype: torch.dtype = torch.int32, |
| 80 | ) -> torch.Tensor: |
| 81 | """Select the top-k indices per row of `input` with DeepSelect. |
| 82 | |
| 83 | Args: |
| 84 | input: (num_rows, vocab_size), bf16 or fp32. stride(1) must be 1 and |
| 85 | stride(0) must be 1024B-aligned. |
| 86 | topk: Number of elements to select per row; must be <= 4096. |
| 87 | end: Optional (num_rows,) int32 tensor with the exclusive right |
| 88 | boundary of each row. Rows with `end[i] < topk` get their |
| 89 | remaining indices filled with -1. |
| 90 | output_idx: Optional preallocated (num_rows, topk) output tensor whose |
| 91 | stride(0) is 32B-aligned (e.g. a slice of a wider buffer). |
| 92 | indices_dtype: Output dtype when `output_idx` is not provided. |
| 93 | |
| 94 | Returns: |
| 95 | The (num_rows, topk) indices tensor. |
| 96 | |
| 97 | """ |
| 98 | assert input.dim() == 2 and input.stride(1) == 1 |
| 99 | assert ( |
| 100 | input.stride(0) * input.element_size() % get_deep_select_stride_requirement()[0] |
| 101 | == 0 |
| 102 | ) |
| 103 | |
| 104 | num_rows = input.shape[0] |
| 105 | if output_idx is None: |
| 106 | output_idx = _get_empty_and_aligned_tensor( |
| 107 | num_rows, topk, input.device, indices_dtype |
| 108 | ) |
| 109 | else: |
| 110 | assert output_idx.dtype == indices_dtype |
| 111 | assert output_idx.shape[0] >= num_rows and output_idx.shape[1] >= topk |
| 112 | |
| 113 | torch.ops.deep_select.topk( |
| 114 | input, |
| 115 | topk, |
| 116 | None, # begin is not supported |
| 117 | end, |
| 118 | False, # sorted_value |
| 119 | False, # sorted_index |
| 120 | None, # output_value |
| 121 | output_idx, |
| 122 | None, # output_idx_offset |
| 123 | IDX_OOB_FILL_VALUE, |
| 124 | float("-inf"), # value_oob_fill_value |
| 125 | False, # return_value |
| 126 | False, # abort_when_nan_found: dummy/capture passes can feed NaN |
| 127 | # logits from uninitialized KV cache; the reference topk kernels do |
| 128 | # not trap on NaN, and trapping would abort CUDA graph capture. |
| 129 | ) |
| 130 | return output_idx |
| 131 | |
| 132 | |
| 133 | # --------------------------------------------------------------------------- |
| 134 | # Backend selection |
| 135 | # --------------------------------------------------------------------------- |
| 136 | |
| 137 | |
| 138 | @functools.cache |
| 139 | def get_indexer_topk(backend: str) -> "SparseIndexerTopk": |
| 140 | return SparseIndexerTopk(backend) |
| 141 | |
| 142 | |
| 143 | class SparseIndexerTopk(torch.nn.Module): |
| 144 | """The sparse indexer's decode top-k stage. |
| 145 | |
| 146 | Selects among the available top-k kernels (see |
| 147 | kernel_config.sparse_indexer_topk_backend) and runs the chosen one. |
| 148 | """ |
| 149 | |
| 150 | def __init__(self, backend: str | None = None) -> None: |
| 151 | super().__init__() |
| 152 | if backend is None: |
| 153 | backend = ( |
| 154 | get_current_vllm_config().kernel_config.sparse_indexer_topk_backend |
| 155 | ) |
| 156 | self._backend = backend |
| 157 | self._is_cuda = current_platform.is_cuda() |
| 158 | self._has_deep_select = self._is_cuda and ( |
| 159 | current_platform.is_device_capability_family(100) |
| 160 | ) |
| 161 | self._has_flashinfer_topk = has_flashinfer() |
| 162 | self._cooperative_capable = self._is_cuda and ( |
| 163 | current_platform.has_device_capability(90) |
| 164 | and not current_platform.is_device_capability_family(120) |
| 165 | ) |
| 166 | self._is_rocm = current_platform.is_rocm() |
| 167 | self._aiter_enabled = False |
| 168 | self._aiter_capable = False |
| 169 | if self._is_rocm: |
| 170 | from vllm._aiter_ops import rocm_aiter_ops |
| 171 | |
| 172 | self._aiter_enabled = bool(rocm_aiter_ops.is_enabled()) |
| 173 | self._aiter_capable = bool(rocm_aiter_ops.is_indexer_top_k_enabled()) |
| 174 | |
| 175 | def resolve_backend( |
| 176 | self, logits: torch.Tensor, topk_tokens: int, num_rows: int |
| 177 | ) -> str: |
| 178 | """Resolve the decode top-k implementation from the configured |
| 179 | backend ("auto" = the pre-existing chain, or a validated explicit |
| 180 | value).""" |
| 181 | if self._backend == "auto": |
| 182 | return self._resolve_auto(logits, topk_tokens, num_rows) |
| 183 | |
| 184 | failures: list[str] = [] |
| 185 | if self._backend == "cooperative": |
| 186 | failures = self._cooperative_constraints(logits, topk_tokens, num_rows) |
| 187 | elif self._backend == "persistent": |
| 188 | if not self._is_cuda: |
| 189 | failures.append("requires a CUDA platform") |
| 190 | if topk_tokens not in (512, 1024, 2048): |
| 191 | failures.append( |
| 192 | f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" |
| 193 | ) |
| 194 | elif self._backend == "deep_select": |
| 195 | if not self._is_cuda: |
| 196 | failures.append("requires a CUDA platform") |
| 197 | elif not self._has_deep_select: |
| 198 | failures.append("requires SM100a/SM103a (10.x device family)") |
| 199 | elif not is_deep_select_supported(logits, topk_tokens): |
| 200 | failures.append( |
| 201 | f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," |
| 202 | f" stride={logits.stride()}, topk={topk_tokens}" |
| 203 | ) |
| 204 | elif self._backend == "aiter": |
| 205 | if not self._is_rocm: |
| 206 | failures.append("requires a ROCm platform") |
| 207 | elif not self._aiter_enabled: |
| 208 | failures.append("requires AITER (VLLM_ROCM_USE_AITER=1)") |
| 209 | elif not self._aiter_capable: |
| 210 | failures.append( |
| 211 | "AITER has no tuned indexer top-k kernels for this GPU arch" |
| 212 | ) |
| 213 | if topk_tokens not in (512, 1024, 2048, 4096): |
| 214 | failures.append( |
| 215 | f"topk_tokens must be in (512, 1024, 2048, 4096), got {topk_tokens}" |
| 216 | ) |
| 217 | elif self._backend == "flashinfer": |
| 218 | if not self._is_cuda: |
| 219 | failures.append("requires a CUDA platform") |
| 220 | if not self._has_flashinfer_topk: |
| 221 | failures.append( |
| 222 | "flashinfer.topk.top_k_ragged_transform is not importable" |
| 223 | ) |
| 224 | if logits.dtype != torch.float32 or logits.stride(1) != 1: |
| 225 | failures.append( |
| 226 | f"requires fp32 logits with stride(1) == 1, got " |
| 227 | f"dtype={logits.dtype}, stride={logits.stride()}" |
| 228 | ) |
| 229 | if failures: |
| 230 | raise RuntimeError( |
| 231 | f"sparse_indexer_topk_backend='{self._backend}' was requested, but: " |
| 232 | + "; ".join(failures) |
| 233 | ) |
| 234 | return self._backend |
| 235 | |
| 236 | def _resolve_auto( |
| 237 | self, logits: torch.Tensor, topk_tokens: int, num_rows: int |
| 238 | ) -> str: |
| 239 | """The priority chain: cooperative -> persistent -> per_row. |
| 240 | aiter/deep_select/flashinfer/torch are opt-in only. |
| 241 | |
| 242 | cooperative_topk is preferred whenever it is applicable, i.e. within |
| 243 | its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to |
| 244 | persistent_topk. |
| 245 | |
| 246 | aiter not auto-resolved as it does not outperform the per-row on every shape, |
| 247 | hence we delegate to the caller to decide when to use it. |
| 248 | """ |
| 249 | if not self._cooperative_constraints(logits, topk_tokens, num_rows): |
| 250 | return "cooperative" |
| 251 | if self._is_cuda and topk_tokens in (512, 1024, 2048): |
| 252 | return "persistent" |
| 253 | return "per_row" |
| 254 | |
| 255 | def _cooperative_constraints( |
| 256 | self, logits: torch.Tensor, topk_tokens: int, num_rows: int |
| 257 | ) -> list[str]: |
| 258 | """Unmet constraints of cooperative_topk (empty when applicable).""" |
| 259 | failures = [] |
| 260 | if not self._is_cuda: |
| 261 | failures.append("requires a CUDA platform") |
| 262 | if topk_tokens not in (512, 1024, 2048): |
| 263 | failures.append( |
| 264 | f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" |
| 265 | ) |
| 266 | if num_rows > AUTO_COOPERATIVE_MAX_ROWS: |
| 267 | failures.append( |
| 268 | f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}" |
| 269 | ) |
| 270 | if logits.stride(0) % 4 != 0: |
| 271 | failures.append( |
| 272 | f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}" |
| 273 | ) |
| 274 | if self._is_cuda and not self._cooperative_capable: |
| 275 | failures.append("requires SM90+ and is not supported on the SM12x family") |
| 276 | return failures |
| 277 | |
| 278 | @staticmethod |
| 279 | def _row_ends(seq_lens: torch.Tensor, next_n: int, num_rows: int) -> torch.Tensor: |
| 280 | """Per-row exclusive end offsets (int32, (num_rows,)) for top-k |
| 281 | kernels that take ragged lengths (DeepSelect, FlashInfer, torch |
| 282 | reference). |
| 283 | |
| 284 | seq_lens is (B, next_n) per-row effective lens for native spec decode |
| 285 | and (B, 1) otherwise, in which case per-row lens are derived the same |
| 286 | way as the other decode top-k kernels. |
| 287 | """ |
| 288 | if seq_lens.numel() == num_rows: |
| 289 | row_ends = seq_lens.reshape(-1) |
| 290 | else: |
| 291 | next_n_offsets = torch.arange( |
| 292 | next_n, dtype=torch.int32, device=seq_lens.device |
| 293 | ) |
| 294 | row_ends = ( |
| 295 | (seq_lens.reshape(-1, 1) - next_n + 1 + next_n_offsets) |
| 296 | .clamp_(min=0) |
| 297 | .reshape(-1) |
| 298 | ) |
| 299 | assert row_ends.dtype == torch.int32 |
| 300 | return row_ends |
| 301 | |
| 302 | def forward( |
| 303 | self, |
| 304 | logits: torch.Tensor, |
| 305 | seq_lens: torch.Tensor, |
| 306 | next_n: int, |
| 307 | topk_indices: torch.Tensor, |
| 308 | topk_tokens: int, |
| 309 | max_seq_len: int, |
| 310 | ) -> None: |
| 311 | """Run the resolved decode top-k implementation, writing into |
| 312 | topk_indices (int32, -1 fill for rows shorter than topk_tokens). |
| 313 | """ |
| 314 | backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) |
| 315 | if backend == "deep_select": |
| 316 | row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) |
| 317 | deep_select_topk(logits, topk_tokens, end=row_ends, output_idx=topk_indices) |
| 318 | elif backend == "cooperative": |
| 319 | (topk_workspace,) = current_workspace_manager().get_simultaneous( |
| 320 | ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), |
| 321 | ) |
| 322 | torch.ops._C.cooperative_topk( |
| 323 | logits, |
| 324 | seq_lens, |
| 325 | topk_indices, |
| 326 | topk_workspace, |
| 327 | topk_tokens, |
| 328 | max_seq_len, |
| 329 | ) |
| 330 | elif backend == "persistent": |
| 331 | (topk_workspace,) = current_workspace_manager().get_simultaneous( |
| 332 | ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), |
| 333 | ) |
| 334 | # persistent_topk's last argument is the host-side "max seq_len |
| 335 | # across rows" bound (it gates the sampled_topk path and clamps |
| 336 | # per-row lengths), not the padded row width: logits is sized to |
| 337 | # max_model_len, so passing its width would incorrectly enable the |
| 338 | # long-row kernels on short-context batches. |
| 339 | torch.ops._C.persistent_topk( |
| 340 | logits, |
| 341 | seq_lens, |
| 342 | topk_indices, |
| 343 | topk_workspace, |
| 344 | topk_tokens, |
| 345 | max_seq_len, |
| 346 | ) |
| 347 | elif backend == "flashinfer": |
| 348 | # Deferred: importing flashinfer initializes CUDA at import time. |
| 349 | from flashinfer.topk import top_k_ragged_transform |
| 350 | |
| 351 | row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) |
| 352 | offsets = torch.zeros( |
| 353 | logits.shape[0], dtype=torch.int32, device=logits.device |
| 354 | ) |
| 355 | # top_k_ragged_transform selects within [0, row_ends[i]) per |
| 356 | # row and -1-fills past the row length. |
| 357 | indices = top_k_ragged_transform(logits, offsets, row_ends, topk_tokens) |
| 358 | topk_indices.copy_(indices) |
| 359 | elif backend == "torch": |
| 360 | # Debug reference: mask everything past each row's end, then topk. |
| 361 | row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) |
| 362 | cols = torch.arange(logits.shape[1], device=logits.device) |
| 363 | masked = logits.masked_fill( |
| 364 | cols.unsqueeze(0) >= row_ends.unsqueeze(1), float("-inf") |
| 365 | ) |
| 366 | indices = masked.topk(topk_tokens, dim=-1).indices |
| 367 | in_range = torch.arange(topk_tokens, device=logits.device).unsqueeze( |
| 368 | 0 |
| 369 | ) < row_ends.unsqueeze(1) |
| 370 | indices = torch.where(in_range, indices, -1) |
| 371 | topk_indices.copy_(indices) |
| 372 | elif backend == "aiter": |
| 373 | from vllm._aiter_ops import rocm_aiter_ops |
| 374 | |
| 375 | # AITER's decode kernel only implements the 1D/next_n == 1 form, |
| 376 | # which is equivalent once the ragged ends are expanded. |
| 377 | rocm_aiter_ops.indexer_top_k_decode( |
| 378 | logits, |
| 379 | 1, |
| 380 | self._row_ends(seq_lens, next_n, logits.shape[0]), |
| 381 | topk_indices, |
| 382 | topk_tokens, |
| 383 | ) |
| 384 | else: |
| 385 | assert backend == "per_row", f"unknown topk backend: {backend}" |
| 386 | self._per_row(logits, seq_lens, next_n, topk_indices, topk_tokens) |
| 387 | |
| 388 | @staticmethod |
| 389 | def _per_row( |
| 390 | logits: torch.Tensor, |
| 391 | seq_lens: torch.Tensor, |
| 392 | next_n: int, |
| 393 | topk_indices: torch.Tensor, |
| 394 | topk_tokens: int, |
| 395 | ) -> None: |
| 396 | ops.top_k_per_row_decode( |
| 397 | logits, |
| 398 | next_n, |
| 399 | seq_lens, |
| 400 | topk_indices, |
| 401 | logits.shape[0], |
| 402 | logits.stride(0), |
| 403 | logits.stride(1), |
| 404 | topk_tokens, |
| 405 | ) |