| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer. |
| 4 | |
| 5 | The candidate-source indexer publishes the top candidate blocks (in units of |
| 6 | ``candidate_block_size`` compressed positions). Instead of scoring every KV |
| 7 | position and masking the dense logits down to those blocks, this path scores |
| 8 | only the candidate tokens with DeepGEMM's sparse MQA logits kernels |
| 9 | (``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs |
| 10 | DeepSelect's top-k directly on the bf16 logits they produce, and remaps the |
| 11 | selected columns back to request-local positions. |
| 12 | |
| 13 | Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed |
| 14 | (granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The |
| 15 | FP8 indexer cache stores one fp32 scale per quant block instead. |
| 16 | |
| 17 | The functions here take caller-owned output buffers (``sparse_indices``, |
| 18 | ``end``, ``col_indices``) so the hot path allocates nothing; the indexer |
| 19 | metadata builder owns them. |
| 20 | """ |
| 21 | |
| 22 | import importlib.util |
| 23 | |
| 24 | import torch |
| 25 | |
| 26 | from vllm.model_executor.layers.indexer_topk import ( |
| 27 | deep_select_topk, |
| 28 | get_deep_select_stride_requirement, |
| 29 | ) |
| 30 | from vllm.platforms import current_platform |
| 31 | from vllm.triton_utils import tl, triton |
| 32 | from vllm.utils.deep_gemm import ( |
| 33 | fp8_fp4_paged_sparse_mqa_logits, |
| 34 | fp8_fp4_sparse_mqa_logits, |
| 35 | get_paged_sparse_mqa_logits_metadata, |
| 36 | get_sparse_mqa_logits_metadata, |
| 37 | ) |
| 38 | |
| 39 | # Sparse-block sizes (in tokens) supported by the DeepGEMM kernels. |
| 40 | SPARSE_BLOCK_KV_CHOICES = (16, 8) |
| 41 | |
| 42 | |
| 43 | def has_deep_select() -> bool: |
| 44 | """Whether the DeepSelect top-k extension is built and this GPU runs it.""" |
| 45 | return ( |
| 46 | importlib.util.find_spec("vllm._deepselect_C") is not None |
| 47 | and current_platform.is_cuda() |
| 48 | and current_platform.is_device_capability_family(100) |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: |
| 53 | """Raise unless DeepSelect accepts the sparse logits layout. |
| 54 | |
| 55 | The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]`` |
| 56 | logits and the top-k writes int32 ``[rows, topk_tokens]`` indices; |
| 57 | DeepSelect needs both row strides aligned (1024B / 32B on SM100). |
| 58 | """ |
| 59 | in_align, out_align = get_deep_select_stride_requirement() |
| 60 | logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize |
| 61 | if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0: |
| 62 | raise ValueError( |
| 63 | f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 " |
| 64 | f"columns, {logits_row_bytes} B) to be {in_align}B-aligned." |
| 65 | ) |
| 66 | if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0: |
| 67 | raise ValueError( |
| 68 | f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 " |
| 69 | f"output row, got topk={topk_tokens}." |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | _INT32_BIG = tl.constexpr(2**31 - 1) |
| 74 | |
| 75 | |
| 76 | @triton.jit(do_not_specialize_on_alignment=["end_ptr"]) |
| 77 | def _expand_candidates_kernel( |
| 78 | cand_ptr, |
| 79 | cand_stride, |
| 80 | ks_ptr, |
| 81 | ks_stride, |
| 82 | ke_ptr, |
| 83 | ke_stride, |
| 84 | si_ptr, |
| 85 | si_stride, |
| 86 | end_ptr, |
| 87 | end_stride, |
| 88 | K: tl.constexpr, |
| 89 | S: tl.constexpr, |
| 90 | RATIO: tl.constexpr, |
| 91 | SBK: tl.constexpr, |
| 92 | ): |
| 93 | """One program per row: expand + filter + sort candidate blocks. |
| 94 | |
| 95 | Each row's valid sparse blocks come out sorted ascending (no dedup pass: |
| 96 | production candidates are unique, and the kernel tolerates repeats), |
| 97 | padded by repeating the last valid block — DeepGEMM's padding convention. |
| 98 | ``end`` counts valid sparse-token columns; valid columns form a prefix |
| 99 | because block ids are sorted. |
| 100 | """ |
| 101 | row = tl.program_id(0) |
| 102 | ks = tl.load(ks_ptr + row * ks_stride) |
| 103 | ke = tl.load(ke_ptr + row * ke_stride) |
| 104 | cols = tl.arange(0, K) |
| 105 | cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask |
| 106 | if RATIO == 2: |
| 107 | cand = tl.interleave(cand * RATIO, cand * RATIO + 1) |
| 108 | |
| 109 | # Keep only blocks inside the row's KV range [0, ke - ks); the lower |
| 110 | # bound is free since candidates are non-negative. |
| 111 | kv_len = tl.maximum(ke - ks, 0) |
| 112 | max_block = (kv_len + SBK - 1) // SBK |
| 113 | valid = (cand >= 0) & (cand < max_block) |
| 114 | vals = tl.where(valid, cand, _INT32_BIG) |
| 115 | vals = tl.sort(vals) |
| 116 | |
| 117 | # Anchor at the row's ks rounded down to a sparse-block boundary |
| 118 | # (identity for the paged path, where ks == 0). |
| 119 | base = ks // SBK |
| 120 | si = tl.where(vals != _INT32_BIG, vals, 0) + base |
| 121 | n_valid = tl.sum(valid.to(tl.int32)) |
| 122 | last_rel = tl.max(tl.where(valid, cand, -1)) |
| 123 | pad_val = tl.where(n_valid > 0, last_rel + base, base) |
| 124 | tl.store( |
| 125 | si_ptr + row * si_stride + tl.arange(0, S), |
| 126 | tl.where(tl.arange(0, S) < n_valid, si, pad_val), |
| 127 | ) |
| 128 | |
| 129 | # Valid sparse-token column count; only the last block can be partial. |
| 130 | # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK, |
| 131 | # which is exactly last_rel * SBK + ks. |
| 132 | last_start = last_rel * SBK + ks |
| 133 | fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK) |
| 134 | end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0) |
| 135 | tl.store(end_ptr + row * end_stride, end) |
| 136 | |
| 137 | |
| 138 | @triton.jit |
| 139 | def _sparse_topk_remap_kernel( |
| 140 | col_ptr, |
| 141 | col_stride, |
| 142 | si_ptr, |
| 143 | si_stride, |
| 144 | ks_ptr, |
| 145 | ks_stride, |
| 146 | out_ptr, |
| 147 | out_stride, |
| 148 | k, |
| 149 | width, |
| 150 | SBK: tl.constexpr, |
| 151 | K_POW2: tl.constexpr, |
| 152 | ): |
| 153 | """Remap sparse top-k columns to request-local positions. |
| 154 | |
| 155 | Column ``j * SBK + o`` scores the token at |
| 156 | ``(si[row, j] - ks // SBK) * SBK + o``; columns outside ``[0, width)`` |
| 157 | (including DeepSelect's sentinel for NaN rows) become -1. |
| 158 | """ |
| 159 | row = tl.program_id(0) |
| 160 | ks_base = tl.load(ks_ptr + row * ks_stride) // SBK |
| 161 | cols = tl.arange(0, K_POW2) |
| 162 | c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1) |
| 163 | valid = (c >= 0) & (c < width) |
| 164 | cc = tl.where(valid, c, 0) |
| 165 | blocks = tl.load(si_ptr + row * si_stride + cc // SBK) |
| 166 | pos = (blocks - ks_base) * SBK + cc % SBK |
| 167 | tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k) |
| 168 | |
| 169 | |
| 170 | def pick_sparse_block_kv(candidate_block_size: int) -> int: |
| 171 | """Largest sparse-block size (tokens) dividing the candidate block size.""" |
| 172 | for sbk in SPARSE_BLOCK_KV_CHOICES: |
| 173 | if candidate_block_size % sbk == 0: |
| 174 | return sbk |
| 175 | raise ValueError( |
| 176 | f"candidate_block_size={candidate_block_size} is not a multiple of " |
| 177 | f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}." |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | def candidate_blocks_to_sparse_indices( |
| 182 | candidate_blocks: torch.Tensor, |
| 183 | row_ks: torch.Tensor, |
| 184 | row_ke: torch.Tensor, |
| 185 | candidate_block_size: int, |
| 186 | sparse_block_kv: int, |
| 187 | out: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 188 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 189 | """Expand request-local candidate blocks into DeepGEMM sparse block ids. |
| 190 | |
| 191 | Args: |
| 192 | candidate_blocks: [rows, K] int32 request-local candidate block ids |
| 193 | (-1 padded), in units of ``candidate_block_size`` positions. K |
| 194 | must be a power of two (the in-kernel sort); rows may be strided. |
| 195 | row_ks: [rows] int32 per-row K range start; bounds are in the same |
| 196 | (packed-workspace) coordinates the sparse kernel iterates over. |
| 197 | Pass zeros for the paged path, whose blocks are context-relative. |
| 198 | row_ke: [rows] int32 per-row K range end, in the same coordinates as |
| 199 | ``row_ks``. |
| 200 | candidate_block_size: Positions per candidate block. |
| 201 | sparse_block_kv: Positions per sparse block (8 or 16). |
| 202 | out: Optional ``(sparse_indices, end)`` buffers to write into, |
| 203 | shaped [rows, K * ratio] and [rows], both int32. |
| 204 | |
| 205 | Returns: |
| 206 | sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM, |
| 207 | sorted ascending per row with the valid prefix first (unique in |
| 208 | production, where candidates are unique) and the remaining slots |
| 209 | repeating the last valid block (the kernel's padding convention). |
| 210 | With the unaligned-ks variant, block ``i`` starts at |
| 211 | ``i * sparse_block_kv + ks % sparse_block_kv``. |
| 212 | end: [rows] int32 count of valid sparse-token columns; valid columns |
| 213 | always form a prefix because the block ids are sorted. |
| 214 | |
| 215 | """ |
| 216 | rows, num_candidates = candidate_blocks.shape |
| 217 | ratio = candidate_block_size // sparse_block_kv |
| 218 | assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), ( |
| 219 | candidate_block_size, |
| 220 | sparse_block_kv, |
| 221 | ) |
| 222 | assert num_candidates & (num_candidates - 1) == 0, num_candidates |
| 223 | num_sparse = num_candidates * ratio |
| 224 | if out is None: |
| 225 | device = candidate_blocks.device |
| 226 | out = ( |
| 227 | torch.empty(rows, num_sparse, dtype=torch.int32, device=device), |
| 228 | torch.empty(rows, dtype=torch.int32, device=device), |
| 229 | ) |
| 230 | sparse_indices, end = out |
| 231 | assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,) |
| 232 | assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1 |
| 233 | for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end): |
| 234 | assert t.dtype == torch.int32, t.dtype |
| 235 | |
| 236 | _expand_candidates_kernel[(rows,)]( |
| 237 | candidate_blocks, |
| 238 | candidate_blocks.stride(0), |
| 239 | row_ks, |
| 240 | row_ks.stride(0), |
| 241 | row_ke, |
| 242 | row_ke.stride(0), |
| 243 | sparse_indices, |
| 244 | sparse_indices.stride(0), |
| 245 | end, |
| 246 | end.stride(0), |
| 247 | K=num_candidates, |
| 248 | S=num_sparse, |
| 249 | RATIO=ratio, |
| 250 | SBK=sparse_block_kv, |
| 251 | num_warps=16, |
| 252 | ) |
| 253 | return sparse_indices, end |
| 254 | |
| 255 | |
| 256 | def sparse_topk_remap( |
| 257 | logits: torch.Tensor, |
| 258 | sparse_indices: torch.Tensor, |
| 259 | end: torch.Tensor, |
| 260 | row_ks: torch.Tensor, |
| 261 | sparse_block_kv: int, |
| 262 | topk_tokens: int, |
| 263 | topk_indices: torch.Tensor, |
| 264 | *, |
| 265 | col_indices: torch.Tensor, |
| 266 | ) -> None: |
| 267 | """Row top-k over sparse logits, remapped to request-local positions. |
| 268 | |
| 269 | Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at |
| 270 | request-local position |
| 271 | ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``. |
| 272 | |
| 273 | DeepSelect consumes the bf16 logits as they come out of the sparse |
| 274 | kernels; the valid columns form a prefix (block ids are sorted), so the |
| 275 | per-row ``end`` bounds drive it directly. Slots beyond a row's valid |
| 276 | count come back as -1. |
| 277 | |
| 278 | Args: |
| 279 | logits: [rows, width] bf16 sparse logits, as produced by the sparse |
| 280 | kernels. |
| 281 | sparse_indices: [rows, S] int32 sparse block ids backing ``logits``. |
| 282 | end: [rows] int32 per-row count of valid logit columns. |
| 283 | row_ks: [rows] int32 per-row K range start. |
| 284 | sparse_block_kv: Positions per sparse block (8 or 16). |
| 285 | topk_tokens: Number of tokens to keep per row. |
| 286 | topk_indices: [rows, topk_tokens] int32 output buffer. |
| 287 | col_indices: [rows, topk_tokens] int32 scratch for the sparse-column |
| 288 | top-k result (row stride 32B-aligned for DeepSelect). |
| 289 | |
| 290 | """ |
| 291 | rows, width = logits.shape |
| 292 | assert logits.dtype == torch.bfloat16 and width >= topk_tokens |
| 293 | assert end.dtype == torch.int32 and end.shape == (rows,) |
| 294 | cols = col_indices[:rows] |
| 295 | deep_select_topk(logits, topk_tokens, end=end, output_idx=cols) |
| 296 | _sparse_topk_remap_kernel[(rows,)]( |
| 297 | cols, |
| 298 | cols.stride(0), |
| 299 | sparse_indices, |
| 300 | sparse_indices.stride(0), |
| 301 | row_ks, |
| 302 | row_ks.stride(0), |
| 303 | topk_indices, |
| 304 | topk_indices.stride(0), |
| 305 | topk_tokens, |
| 306 | width, |
| 307 | SBK=sparse_block_kv, |
| 308 | K_POW2=triton.next_power_of_2(topk_tokens), |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | def sparse_mqa_logits_prefill_chunk( |
| 313 | q: torch.Tensor, |
| 314 | q_scale: torch.Tensor, |
| 315 | k_quant: torch.Tensor, |
| 316 | k_scale: torch.Tensor, |
| 317 | weights: torch.Tensor, |
| 318 | cu_seqlen_ks: torch.Tensor, |
| 319 | cu_seqlen_ke: torch.Tensor, |
| 320 | candidate_blocks: torch.Tensor, |
| 321 | candidate_block_size: int, |
| 322 | sparse_block_kv: int, |
| 323 | topk_tokens: int, |
| 324 | topk_indices: torch.Tensor, |
| 325 | *, |
| 326 | sparse_indices: torch.Tensor, |
| 327 | end: torch.Tensor, |
| 328 | col_indices: torch.Tensor, |
| 329 | kernel_metadata: torch.Tensor | None = None, |
| 330 | ) -> torch.Tensor: |
| 331 | """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace). |
| 332 | |
| 333 | Args: |
| 334 | q: [rows, H, D] packed Q values (MXFP4 viewed as int8). |
| 335 | q_scale: [rows, H] int32 packed UE8M0 Q scales. |
| 336 | k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8). |
| 337 | k_scale: [total_kv] int32 packed UE8M0 K scales. |
| 338 | weights: [rows, H] bf16 per-head weights; the sparse kernels take |
| 339 | bf16 and do not fold the Q scale in. |
| 340 | cu_seqlen_ks: [rows] int32 per-token K start bounds in the packed |
| 341 | workspace. |
| 342 | cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed |
| 343 | workspace. |
| 344 | candidate_blocks: [rows, K] int32 request-local candidate block ids. |
| 345 | candidate_block_size: Positions per candidate block. |
| 346 | sparse_block_kv: Positions per sparse block (8 or 16). |
| 347 | topk_tokens: Number of tokens to keep per row. |
| 348 | topk_indices: [rows, topk_tokens] output buffer. |
| 349 | sparse_indices: Caller-owned scratch, see |
| 350 | `candidate_blocks_to_sparse_indices`. |
| 351 | end: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`. |
| 352 | col_indices: Caller-owned scratch, see `sparse_topk_remap`. |
| 353 | kernel_metadata: DeepGEMM schedule from a previous call with the same |
| 354 | candidates and bounds (i.e. another indexer layer in the same |
| 355 | step). When given, the candidate expansion is skipped and |
| 356 | ``sparse_indices``/``end`` are assumed to be up to date. |
| 357 | |
| 358 | Returns: |
| 359 | The DeepGEMM schedule metadata used, for reuse by later layers. |
| 360 | |
| 361 | """ |
| 362 | assert weights.dtype == torch.bfloat16, weights.dtype |
| 363 | if kernel_metadata is None: |
| 364 | candidate_blocks_to_sparse_indices( |
| 365 | candidate_blocks, |
| 366 | cu_seqlen_ks, |
| 367 | cu_seqlen_ke, |
| 368 | candidate_block_size, |
| 369 | sparse_block_kv, |
| 370 | out=(sparse_indices, end), |
| 371 | ) |
| 372 | kernel_metadata = get_sparse_mqa_logits_metadata( |
| 373 | cu_seqlen_ks, |
| 374 | cu_seqlen_ke, |
| 375 | k_quant.shape[0], |
| 376 | sparse_indices, |
| 377 | q.dtype, |
| 378 | sparse_block_kv, |
| 379 | ) |
| 380 | logits = fp8_fp4_sparse_mqa_logits( |
| 381 | (q, q_scale), |
| 382 | (k_quant, k_scale), |
| 383 | weights, |
| 384 | kernel_metadata, |
| 385 | sparse_indices.shape[1], |
| 386 | sparse_block_kv, |
| 387 | ) |
| 388 | sparse_topk_remap( |
| 389 | logits, |
| 390 | sparse_indices, |
| 391 | end, |
| 392 | cu_seqlen_ks, |
| 393 | sparse_block_kv, |
| 394 | topk_tokens, |
| 395 | topk_indices, |
| 396 | col_indices=col_indices, |
| 397 | ) |
| 398 | return kernel_metadata |
| 399 | |
| 400 | |
| 401 | def sparse_mqa_logits_paged_decode( |
| 402 | q: torch.Tensor, |
| 403 | q_scale: torch.Tensor, |
| 404 | kv_cache: torch.Tensor, |
| 405 | weights: torch.Tensor, |
| 406 | context_lens: torch.Tensor, |
| 407 | block_table: torch.Tensor, |
| 408 | row_indices: torch.Tensor, |
| 409 | candidate_blocks: torch.Tensor, |
| 410 | candidate_block_size: int, |
| 411 | sparse_block_kv: int, |
| 412 | topk_tokens: int, |
| 413 | topk_indices: torch.Tensor, |
| 414 | *, |
| 415 | row_ks: torch.Tensor, |
| 416 | sparse_indices: torch.Tensor, |
| 417 | end: torch.Tensor, |
| 418 | col_indices: torch.Tensor, |
| 419 | kernel_metadata: torch.Tensor | None = None, |
| 420 | ) -> torch.Tensor: |
| 421 | """Sparse-MQA logits + top-k for paged decode rows. |
| 422 | |
| 423 | Every query is one row (spec-decode rows are flattened by the metadata |
| 424 | builder), so ``block_table`` is already per row. |
| 425 | |
| 426 | Args: |
| 427 | q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). |
| 428 | q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. |
| 429 | kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache |
| 430 | view; the page stride must be 512B-aligned. |
| 431 | weights: [rows, H] bf16 per-head weights. |
| 432 | context_lens: [rows] int32 per-row (compressed) context lengths. |
| 433 | block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. |
| 434 | row_indices: [rows] int32 row -> request map (pairing only affects |
| 435 | scheduling). |
| 436 | candidate_blocks: [rows, K] int32 candidate block ids. |
| 437 | candidate_block_size: Positions per candidate block. |
| 438 | sparse_block_kv: Positions per sparse block (8 or 16). |
| 439 | topk_tokens: Number of tokens to keep per row. |
| 440 | topk_indices: [rows, topk_tokens] output buffer. |
| 441 | row_ks: [rows] int32 zeros (paged blocks are context-relative). |
| 442 | sparse_indices: Caller-owned scratch. |
| 443 | end: Caller-owned scratch. |
| 444 | col_indices: Caller-owned scratch. |
| 445 | kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. |
| 446 | |
| 447 | Returns: |
| 448 | The DeepGEMM schedule metadata used, for reuse by later layers. |
| 449 | |
| 450 | """ |
| 451 | rows = q.shape[0] |
| 452 | page_kv = kv_cache.shape[1] |
| 453 | assert weights.dtype == torch.bfloat16, weights.dtype |
| 454 | assert kv_cache.stride(0) % 512 == 0, ( |
| 455 | "paged cache page stride must be 512B-aligned: " |
| 456 | f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" |
| 457 | ) |
| 458 | assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) |
| 459 | assert block_table.shape[0] == rows and block_table.stride(-1) == 1 |
| 460 | assert context_lens.shape == (rows,) and context_lens.is_contiguous() |
| 461 | |
| 462 | if kernel_metadata is None: |
| 463 | candidate_blocks_to_sparse_indices( |
| 464 | candidate_blocks, |
| 465 | row_ks, |
| 466 | context_lens, |
| 467 | candidate_block_size, |
| 468 | sparse_block_kv, |
| 469 | out=(sparse_indices, end), |
| 470 | ) |
| 471 | kernel_metadata = get_paged_sparse_mqa_logits_metadata( |
| 472 | context_lens, |
| 473 | block_table, |
| 474 | row_indices, |
| 475 | page_kv, |
| 476 | sparse_indices, |
| 477 | q.dtype, |
| 478 | sparse_block_kv, |
| 479 | ) |
| 480 | logits = fp8_fp4_paged_sparse_mqa_logits( |
| 481 | (q, q_scale), |
| 482 | kv_cache, |
| 483 | weights, |
| 484 | kernel_metadata, |
| 485 | sparse_indices.shape[1], |
| 486 | sparse_block_kv, |
| 487 | ) |
| 488 | sparse_topk_remap( |
| 489 | logits, |
| 490 | sparse_indices, |
| 491 | end, |
| 492 | row_ks, |
| 493 | sparse_block_kv, |
| 494 | topk_tokens, |
| 495 | topk_indices, |
| 496 | col_indices=col_indices, |
| 497 | ) |
| 498 | return kernel_metadata |