| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k. |
| 4 | |
| 5 | Index queries score each 128-token block of index keys (max over the block), |
| 6 | then the top-k blocks (plus forced init/local blocks) are selected per query |
| 7 | token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the |
| 8 | sparse block size (128), so one sparse block maps to exactly one page. |
| 9 | |
| 10 | Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head). |
| 11 | |
| 12 | Only the paths MiniMax M3 uses are implemented: score_type="max", index value |
| 13 | disabled (score-only indexer), single shared index head. The selected block ids |
| 14 | feed the block-sparse attention kernels in ``sparse_attn``. |
| 15 | """ |
| 16 | |
| 17 | import torch |
| 18 | |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.triton_utils import tl, triton |
| 21 | from vllm.utils.math_utils import round_up |
| 22 | |
| 23 | # One sparse block == one KV page. |
| 24 | SPARSE_BLOCK_SIZE = 128 |
| 25 | |
| 26 | |
| 27 | # --------------------------------------------------------------------------- |
| 28 | # Bitonic top-k helpers (layout-agnostic). |
| 29 | # --------------------------------------------------------------------------- |
| 30 | @triton.jit |
| 31 | def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr): |
| 32 | n_outer: tl.constexpr = x.numel >> n_dims |
| 33 | shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)] |
| 34 | y = tl.reshape(x, shape) |
| 35 | mask = tl.arange(0, 2)[None, :, None] |
| 36 | left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype) |
| 37 | right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype) |
| 38 | left = tl.reshape(left, x.shape) |
| 39 | right = tl.reshape(right, x.shape) |
| 40 | y_idx = tl.reshape(ids, shape) |
| 41 | left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape) |
| 42 | right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape) |
| 43 | left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype) |
| 44 | right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype) |
| 45 | idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) |
| 46 | ileft = left.to(idtype, bitcast=True) |
| 47 | iright = right.to(idtype, bitcast=True) |
| 48 | ix = x.to(idtype, bitcast=True) |
| 49 | cond = (left > right) != flip |
| 50 | ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix)) |
| 51 | new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids)) |
| 52 | return ret.to(x.dtype, bitcast=True), new_ids |
| 53 | |
| 54 | |
| 55 | @triton.jit |
| 56 | def _bitonic_merge( |
| 57 | x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr |
| 58 | ): |
| 59 | n_outer: tl.constexpr = x.numel >> n_dims |
| 60 | tl.static_assert(stage <= n_dims) |
| 61 | if order == 2: |
| 62 | shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage] |
| 63 | flip = tl.reshape( |
| 64 | tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape |
| 65 | ) |
| 66 | else: |
| 67 | flip = order |
| 68 | for i in tl.static_range(stage): |
| 69 | x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims) |
| 70 | return x, ids |
| 71 | |
| 72 | |
| 73 | # --------------------------------------------------------------------------- |
| 74 | # Index block-score kernel (paged). score[h, token, block] = max over the |
| 75 | # 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so |
| 76 | # each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1). |
| 77 | # --------------------------------------------------------------------------- |
| 78 | # since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens |
| 79 | # might lose pointer alignment, which trigger Triton recompiles. we don't actually |
| 80 | # need pointer alignment for those tensors anyway because we do scalar load. |
| 81 | @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) |
| 82 | def _index_block_score_kernel( |
| 83 | q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] |
| 84 | ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] |
| 85 | score_ptr, # [num_idx_heads, total_q, max_block] |
| 86 | block_table_ptr, # [num_reqs, max_blocks] |
| 87 | cu_seqlens, # [batch+1] query start offsets |
| 88 | seq_lens, # [batch] total K length |
| 89 | prefix_lens, # [batch] context length before this chunk's queries |
| 90 | num_idx_heads, |
| 91 | head_dim: tl.constexpr, |
| 92 | stride_q_n, |
| 93 | stride_q_h, |
| 94 | stride_q_d, |
| 95 | stride_ik_blk, |
| 96 | stride_ik_pos, |
| 97 | stride_ik_d, |
| 98 | stride_s_h, |
| 99 | stride_s_n, |
| 100 | stride_s_k, |
| 101 | stride_bt_b, |
| 102 | BLOCK_SIZE_Q: tl.constexpr, |
| 103 | BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) |
| 104 | USE_SPLIT_K: tl.constexpr, |
| 105 | ): |
| 106 | pid_q = tl.program_id(0) |
| 107 | pid_bh = tl.program_id(1) |
| 108 | pid_b = pid_bh // num_idx_heads |
| 109 | pid_h = pid_bh % num_idx_heads |
| 110 | |
| 111 | seq_start = tl.load(cu_seqlens + pid_b) |
| 112 | q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start |
| 113 | seq_len = tl.load(seq_lens + pid_b) |
| 114 | prefix_len = tl.load(prefix_lens + pid_b) |
| 115 | if BLOCK_SIZE_Q * pid_q >= q_len: |
| 116 | return |
| 117 | |
| 118 | q_ptrs = tl.make_block_ptr( |
| 119 | base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, |
| 120 | shape=(q_len, head_dim), |
| 121 | strides=(stride_q_n, stride_q_d), |
| 122 | offsets=(pid_q * BLOCK_SIZE_Q, 0), |
| 123 | block_shape=(BLOCK_SIZE_Q, head_dim), |
| 124 | order=(1, 0), |
| 125 | ) |
| 126 | q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") |
| 127 | # FP8 has no mixed-dtype dot with bf16/fp32. Leave bf16 and fp32 loads |
| 128 | # in their stored dtype so fp32 pipeline tests keep full precision. |
| 129 | if q.dtype.is_fp8(): |
| 130 | q = q.to(tl.bfloat16) |
| 131 | q_start = prefix_len + pid_q * BLOCK_SIZE_Q |
| 132 | |
| 133 | off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len |
| 134 | off_k = tl.arange(0, BLOCK_SIZE_K) |
| 135 | off_d = tl.arange(0, head_dim) |
| 136 | # Block table row for this request. |
| 137 | bt_row = block_table_ptr + pid_b * stride_bt_b |
| 138 | # Causal window: only blocks up to the last query token's position. |
| 139 | hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q) |
| 140 | num_blocks = tl.cdiv(hi, BLOCK_SIZE_K) |
| 141 | if USE_SPLIT_K: |
| 142 | pid_k = tl.program_id(2) |
| 143 | blocks_per_split = tl.cdiv(num_blocks, tl.num_programs(2)) |
| 144 | block_start = pid_k * blocks_per_split |
| 145 | block_end = tl.minimum(block_start + blocks_per_split, num_blocks) |
| 146 | if block_start >= block_end: |
| 147 | return |
| 148 | else: |
| 149 | block_start = 0 |
| 150 | block_end = num_blocks |
| 151 | for blk in tl.range(block_start, block_end): |
| 152 | i = blk * BLOCK_SIZE_K |
| 153 | page = tl.load(bt_row + blk).to(tl.int64) |
| 154 | pos = i + off_k |
| 155 | # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed) |
| 156 | # we don't need masked load for K, because KV cache ensures |
| 157 | # allocation is multiple of BLOCK_SIZE_K. |
| 158 | # for tokens beyond seqlen, they will be masked in qk later. |
| 159 | k = tl.load( |
| 160 | ik_cache_ptr |
| 161 | + page * stride_ik_blk |
| 162 | + off_k[None, :] * stride_ik_pos |
| 163 | + off_d[:, None] * stride_ik_d, |
| 164 | ) |
| 165 | if k.dtype.is_fp8(): |
| 166 | k = k.to(tl.bfloat16) |
| 167 | qk = tl.dot(q, k, out_dtype=tl.float32) |
| 168 | # apply causal mask as needed |
| 169 | if q_start < i + BLOCK_SIZE_K: |
| 170 | qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf")) |
| 171 | # one sparse block per K-tile -> max over the 128 positions |
| 172 | score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q] |
| 173 | s_ptrs = ( |
| 174 | score_ptr |
| 175 | + pid_h * stride_s_h |
| 176 | + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) |
| 177 | * stride_s_n |
| 178 | + blk * stride_s_k |
| 179 | ) |
| 180 | q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len |
| 181 | tl.store(s_ptrs, score, mask=q_store_mask) |
| 182 | |
| 183 | |
| 184 | # --------------------------------------------------------------------------- |
| 185 | # Top-k selection over per-token block scores (layout-agnostic). block_size_q |
| 186 | # is 1 for M3, so top-k is computed per query token. |
| 187 | # --------------------------------------------------------------------------- |
| 188 | # since prefill metadata is sliced from mixed batch metadata, prefix_lens |
| 189 | # might lose pointer alignment, which trigger Triton recompiles. we don't actually |
| 190 | # need pointer alignment for those tensors anyway because we do scalar load. |
| 191 | @triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) |
| 192 | @triton.autotune( |
| 193 | configs=[ |
| 194 | triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2), |
| 195 | triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2), |
| 196 | triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2), |
| 197 | triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), |
| 198 | triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), |
| 199 | triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), |
| 200 | ], |
| 201 | key=["BLOCK_SIZE_T"], |
| 202 | ) |
| 203 | @triton.jit(do_not_specialize_on_alignment=["prefix_lens"]) |
| 204 | def _topk_index_kernel( |
| 205 | s_ptr, # [num_heads, total_q, max_block] |
| 206 | ti_ptr, # [num_heads, total_q, topk] |
| 207 | sample_interval: tl.constexpr, # block_size_q (1 for M3) |
| 208 | block_size: tl.constexpr, # sparse block size (128) |
| 209 | cu_seqlens, |
| 210 | cu_seqblocks_q, |
| 211 | prefix_lens, |
| 212 | topk, |
| 213 | init_blocks: tl.constexpr, |
| 214 | local_blocks: tl.constexpr, |
| 215 | stride_s_h, |
| 216 | stride_s_n, |
| 217 | stride_s_k, |
| 218 | stride_ti_h, |
| 219 | stride_ti_n, |
| 220 | stride_ti_t, |
| 221 | BLOCK_SIZE_K: tl.constexpr, |
| 222 | BLOCK_SIZE_T: tl.constexpr, |
| 223 | MASK_INIT: tl.constexpr, |
| 224 | MASK_LOCAL: tl.constexpr, |
| 225 | ): |
| 226 | tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T) |
| 227 | pid_q = tl.program_id(0) |
| 228 | pid_b = tl.program_id(1) |
| 229 | pid_h = tl.program_id(2) |
| 230 | seq_start = tl.load(cu_seqlens + pid_b) |
| 231 | block_start = tl.load(cu_seqblocks_q + pid_b) |
| 232 | block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start |
| 233 | prefix_len = tl.load(prefix_lens + pid_b) |
| 234 | if pid_q >= block_num: |
| 235 | return |
| 236 | off_k = tl.arange(0, BLOCK_SIZE_K) |
| 237 | off_t = tl.arange(0, BLOCK_SIZE_T) |
| 238 | s_ptrs = ( |
| 239 | s_ptr |
| 240 | + (seq_start + pid_q * sample_interval) * stride_s_n |
| 241 | + pid_h * stride_s_h |
| 242 | + off_k * stride_s_k |
| 243 | ) |
| 244 | topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) |
| 245 | topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) |
| 246 | left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 |
| 247 | valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size |
| 248 | for i in tl.range(0, valid_blocks, BLOCK_SIZE_K): |
| 249 | causal_mask = i + off_k < valid_blocks |
| 250 | local_mask = i + off_k >= max(0, valid_blocks - local_blocks) |
| 251 | init_mask = i + off_k < init_blocks |
| 252 | score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32) |
| 253 | score = tl.where(score != score, -1e30, score) |
| 254 | s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K |
| 255 | if MASK_INIT: |
| 256 | score = tl.where(causal_mask & init_mask, score - 1e29, score) |
| 257 | else: |
| 258 | score = tl.where(causal_mask & init_mask, 1e30, score) |
| 259 | if MASK_LOCAL: |
| 260 | score = tl.where(causal_mask & local_mask, score - 1e28, score) |
| 261 | else: |
| 262 | score = tl.where(causal_mask & local_mask, 1e29, score) |
| 263 | topk_score, last_topk_score = score, topk_score |
| 264 | topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx) |
| 265 | n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) |
| 266 | for j in tl.static_range(1, n_dims): |
| 267 | topk_score, topk_idx = _bitonic_merge( |
| 268 | topk_score, topk_idx.to(tl.int32), j, 2, n_dims |
| 269 | ) |
| 270 | if i != 0: |
| 271 | topk_score, topk_idx = _bitonic_merge( |
| 272 | topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims |
| 273 | ) |
| 274 | topk_score_new = last_topk_score * left_half_mask + topk_score * ( |
| 275 | 1 - left_half_mask |
| 276 | ) |
| 277 | topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( |
| 278 | 1 - left_half_mask |
| 279 | ) |
| 280 | topk_score, topk_idx = _bitonic_merge( |
| 281 | topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims |
| 282 | ) |
| 283 | else: |
| 284 | topk_score, topk_idx = _bitonic_merge( |
| 285 | topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims |
| 286 | ) |
| 287 | topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 |
| 288 | topk_idx = tl.sum( |
| 289 | topk_mask[:, None] |
| 290 | * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), |
| 291 | axis=0, |
| 292 | ) |
| 293 | ti_ptrs = ( |
| 294 | ti_ptr |
| 295 | + (block_start + pid_q) * stride_ti_n |
| 296 | + pid_h * stride_ti_h |
| 297 | + off_t * stride_ti_t |
| 298 | ) |
| 299 | store_mask = off_t < topk |
| 300 | valid_mask = off_t < valid_blocks |
| 301 | topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1) |
| 302 | tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask) |
| 303 | |
| 304 | |
| 305 | # --------------------------------------------------------------------------- |
| 306 | # Decode index-score kernel (split-K over seq blocks). Decode batches are |
| 307 | # flattened request-major, with a runtime query length used to map each query |
| 308 | # token back to its request metadata. Chunk counts depend only on shape |
| 309 | # constants so the grid is fixed within a cuda graph. The score scale is omitted |
| 310 | # because decode only consumes block ordering. |
| 311 | # --------------------------------------------------------------------------- |
| 312 | @triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"]) |
| 313 | def _decode_index_score_kernel( |
| 314 | q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] |
| 315 | ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] |
| 316 | score_ptr, # [num_idx_heads, total_q, max_block] |
| 317 | block_table_ptr, # [num_reqs, max_blocks] |
| 318 | seq_lens, # [num_reqs] |
| 319 | num_idx_heads: tl.constexpr, |
| 320 | head_dim: tl.constexpr, |
| 321 | init_blocks, |
| 322 | local_blocks, |
| 323 | decode_query_len, |
| 324 | stride_q_n, |
| 325 | stride_q_h, |
| 326 | stride_q_d, |
| 327 | stride_ik_blk, |
| 328 | stride_ik_pos, |
| 329 | stride_ik_d, |
| 330 | stride_s_h, |
| 331 | stride_s_n, |
| 332 | stride_s_k, |
| 333 | stride_bt_b, |
| 334 | BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) |
| 335 | BLOCK_SIZE_Q: tl.constexpr, |
| 336 | num_kv_chunks, |
| 337 | USE_PDL: tl.constexpr, |
| 338 | ): |
| 339 | BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q |
| 340 | pid_r = tl.program_id(0) |
| 341 | pid_c = tl.program_id(1) |
| 342 | hq_offsets = tl.arange(0, BLOCK_SIZE_HQ) |
| 343 | h_offsets = hq_offsets // BLOCK_SIZE_Q |
| 344 | q_offsets = hq_offsets % BLOCK_SIZE_Q |
| 345 | q_mask = q_offsets < decode_query_len |
| 346 | q_ids = pid_r * decode_query_len + q_offsets |
| 347 | |
| 348 | if USE_PDL: |
| 349 | tl.extra.cuda.gdc_wait() |
| 350 | tl.extra.cuda.gdc_launch_dependents() |
| 351 | |
| 352 | seq_len = tl.load(seq_lens + pid_r) |
| 353 | query_pos = seq_len - decode_query_len + q_offsets |
| 354 | # Full-CG padding uses zero-length request rows. Clamp to an empty |
| 355 | # attention range instead of letting padded rows produce negative lengths. |
| 356 | kv_len = tl.maximum(query_pos + 1, 0) |
| 357 | num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K |
| 358 | kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0) |
| 359 | num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K |
| 360 | |
| 361 | # block-aligned fixed-count split: grid independent of seq_len (cuda graph). |
| 362 | chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks |
| 363 | chunk_start_block = pid_c * chunk_size_blocks |
| 364 | chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks) |
| 365 | if chunk_start_block >= chunk_end_block: |
| 366 | return |
| 367 | off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block |
| 368 | off_d = tl.arange(0, head_dim) |
| 369 | bt_row = block_table_ptr + pid_r * stride_bt_b |
| 370 | # Force-select init (1e30) and local (1e29, higher priority) blocks. |
| 371 | local_start = tl.maximum(0, num_blocks_q - local_blocks) |
| 372 | # Query vectors for all index heads in a small spec-decode block. |
| 373 | q = tl.load( |
| 374 | q_ptr |
| 375 | + q_ids[None, :] * stride_q_n |
| 376 | + h_offsets[None, :] * stride_q_h |
| 377 | + off_d[:, None] * stride_q_d, |
| 378 | mask=q_mask[None, :], |
| 379 | other=0.0, |
| 380 | ) # [D,HQ] |
| 381 | if q.dtype.is_fp8(): |
| 382 | q = q.to(tl.bfloat16) |
| 383 | for blk in tl.range(chunk_start_block, chunk_end_block): |
| 384 | page = tl.load(bt_row + blk).to(tl.int64) |
| 385 | pos = blk * BLOCK_SIZE_K + off_k |
| 386 | pos_mask = pos[:, None] < kv_len[None, :] |
| 387 | # we don't need masked load for K, because KV cache ensures |
| 388 | # allocation is multiple of BLOCK_SIZE_K. |
| 389 | # for tokens beyond seqlen, they will be masked in qk later. |
| 390 | k = tl.load( |
| 391 | ik_cache_ptr |
| 392 | + page * stride_ik_blk |
| 393 | + off_k[:, None] * stride_ik_pos |
| 394 | + off_d * stride_ik_d, |
| 395 | ) # [N,D] |
| 396 | # Upcast only FP8 cache loads. BF16 operands keep FP8 dot |
| 397 | # instructions out of the fallback path; FP32 accumulation |
| 398 | # preserves score accuracy. BF16/FP32 loads stay as stored. |
| 399 | if k.dtype.is_fp8(): |
| 400 | k = k.to(tl.bfloat16) |
| 401 | kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ] |
| 402 | kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf")) |
| 403 | score = tl.max(kq, axis=0) # [HQ] |
| 404 | is_visible_block = blk < num_blocks_q |
| 405 | is_init = (blk < init_blocks) & is_visible_block |
| 406 | is_local = (blk >= local_start) & is_visible_block |
| 407 | score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score)) |
| 408 | tl.store( |
| 409 | score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k, |
| 410 | score, |
| 411 | mask=q_mask, |
| 412 | ) |
| 413 | |
| 414 | |
| 415 | # --------------------------------------------------------------------------- |
| 416 | # Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local |
| 417 | # blocks are already encoded in the scores. |
| 418 | # --------------------------------------------------------------------------- |
| 419 | @triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) |
| 420 | @triton.autotune( |
| 421 | configs=[ |
| 422 | triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), |
| 423 | triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2), |
| 424 | triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), |
| 425 | triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3), |
| 426 | triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), |
| 427 | ], |
| 428 | key=["topk"], |
| 429 | ) |
| 430 | @triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"]) |
| 431 | def _topk_index_partial_kernel( |
| 432 | s_ptr, # score: [num_idx_heads, total_q, max_block] |
| 433 | ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] |
| 434 | ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape |
| 435 | seq_lens, # [num_reqs] |
| 436 | block_size: tl.constexpr, # sparse block size (128) |
| 437 | topk: tl.constexpr, |
| 438 | chunk_blocks, # how many score-blocks each chunk owns |
| 439 | decode_query_len, |
| 440 | stride_s_h, |
| 441 | stride_s_b, |
| 442 | stride_s_k, |
| 443 | stride_ts_c, |
| 444 | stride_ts_h, |
| 445 | stride_ts_b, |
| 446 | stride_ts_t, |
| 447 | stride_ti_c, |
| 448 | stride_ti_h, |
| 449 | stride_ti_b, |
| 450 | stride_ti_t, |
| 451 | BLOCK_SIZE_K: tl.constexpr, |
| 452 | BLOCK_SIZE_T: tl.constexpr, |
| 453 | USE_PDL: tl.constexpr, |
| 454 | ): |
| 455 | tl.static_assert(topk < BLOCK_SIZE_K) |
| 456 | pid_b = tl.program_id(0) # flattened query-token id |
| 457 | pid_h = tl.program_id(1) |
| 458 | pid_chunk = tl.program_id(2) |
| 459 | req_id = pid_b // decode_query_len |
| 460 | q_offset = pid_b - req_id * decode_query_len |
| 461 | |
| 462 | if USE_PDL: |
| 463 | tl.extra.cuda.gdc_wait() |
| 464 | |
| 465 | seq_len = tl.load(seq_lens + req_id) |
| 466 | query_pos = seq_len - decode_query_len + q_offset |
| 467 | # Full-CG padding uses zero-length request rows. Clamp to an empty |
| 468 | # attention range instead of letting padded rows produce negative lengths. |
| 469 | kv_len = tl.maximum(query_pos + 1, 0) |
| 470 | num_blocks = (kv_len + block_size - 1) // block_size |
| 471 | |
| 472 | # Slice this chunk owns within [0, num_blocks). |
| 473 | chunk_start = pid_chunk * chunk_blocks |
| 474 | chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks) |
| 475 | chunk_actual = tl.maximum(chunk_end - chunk_start, 0) |
| 476 | |
| 477 | off_k = tl.arange(0, BLOCK_SIZE_K) |
| 478 | off_t = tl.arange(0, BLOCK_SIZE_T) |
| 479 | |
| 480 | s_ptrs = ( |
| 481 | s_ptr |
| 482 | + pid_b * stride_s_b |
| 483 | + pid_h * stride_s_h |
| 484 | + (chunk_start + off_k) * stride_s_k |
| 485 | ) |
| 486 | |
| 487 | topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) |
| 488 | topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) |
| 489 | left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 |
| 490 | |
| 491 | # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty |
| 492 | # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0. |
| 493 | for i in tl.range(0, chunk_actual, BLOCK_SIZE_K): |
| 494 | mask = off_k < chunk_actual - i |
| 495 | score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32) |
| 496 | score = tl.where(score != score, -1e30, score) |
| 497 | s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K |
| 498 | topk_score, last_topk_score = score, topk_score |
| 499 | topk_idx, last_topk_idx = ( |
| 500 | tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global |
| 501 | topk_idx, |
| 502 | ) |
| 503 | n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) |
| 504 | for j in tl.static_range(1, n_dims): |
| 505 | topk_score, topk_idx = _bitonic_merge( |
| 506 | topk_score, topk_idx.to(tl.int32), j, 2, n_dims |
| 507 | ) |
| 508 | if i != 0: |
| 509 | topk_score, topk_idx = _bitonic_merge( |
| 510 | topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims |
| 511 | ) |
| 512 | topk_score_new = last_topk_score * left_half_mask + topk_score * ( |
| 513 | 1 - left_half_mask |
| 514 | ) |
| 515 | topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( |
| 516 | 1 - left_half_mask |
| 517 | ) |
| 518 | topk_score, topk_idx = _bitonic_merge( |
| 519 | topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims |
| 520 | ) |
| 521 | else: |
| 522 | topk_score, topk_idx = _bitonic_merge( |
| 523 | topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims |
| 524 | ) |
| 525 | |
| 526 | if USE_PDL: |
| 527 | tl.extra.cuda.gdc_launch_dependents() |
| 528 | |
| 529 | # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort). |
| 530 | topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 |
| 531 | final_score = tl.sum( |
| 532 | topk_mask_extract[:, None] |
| 533 | * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), |
| 534 | axis=0, |
| 535 | ) |
| 536 | final_idx = tl.sum( |
| 537 | topk_mask_extract[:, None] |
| 538 | * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), |
| 539 | axis=0, |
| 540 | ) |
| 541 | |
| 542 | # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0 |
| 543 | # sentinels and lose to real scores in the merge stage. |
| 544 | ts_ptrs = ( |
| 545 | ts_partial_ptr |
| 546 | + pid_chunk * stride_ts_c |
| 547 | + pid_b * stride_ts_b |
| 548 | + pid_h * stride_ts_h |
| 549 | + off_t * stride_ts_t |
| 550 | ) |
| 551 | ti_ptrs = ( |
| 552 | ti_partial_ptr |
| 553 | + pid_chunk * stride_ti_c |
| 554 | + pid_b * stride_ti_b |
| 555 | + pid_h * stride_ti_h |
| 556 | + off_t * stride_ti_t |
| 557 | ) |
| 558 | tl.store(ts_ptrs, final_score) |
| 559 | tl.store(ti_ptrs, final_idx) |
| 560 | |
| 561 | |
| 562 | @triton.heuristics( |
| 563 | { |
| 564 | "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]), |
| 565 | "BLOCK_SIZE_K": lambda args: triton.next_power_of_2( |
| 566 | args["num_topk_chunks"] * triton.next_power_of_2(args["topk"]) |
| 567 | ), |
| 568 | } |
| 569 | ) |
| 570 | @triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"]) |
| 571 | def _topk_index_merge_kernel( |
| 572 | ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] |
| 573 | ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape |
| 574 | ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk] |
| 575 | seq_lens, # [num_reqs] |
| 576 | block_size: tl.constexpr, # sparse block size (128) |
| 577 | topk: tl.constexpr, |
| 578 | decode_query_len, |
| 579 | stride_ts_c, |
| 580 | stride_ts_h, |
| 581 | stride_ts_b, |
| 582 | stride_ts_t, |
| 583 | stride_ti_c, |
| 584 | stride_ti_h, |
| 585 | stride_ti_b, |
| 586 | stride_ti_t, |
| 587 | stride_tif_h, |
| 588 | stride_tif_b, |
| 589 | stride_tif_t, |
| 590 | num_topk_chunks, |
| 591 | BLOCK_SIZE_K: tl.constexpr, |
| 592 | BLOCK_SIZE_T: tl.constexpr, |
| 593 | USE_PDL: tl.constexpr, |
| 594 | ): |
| 595 | pid_b = tl.program_id(0) # flattened query-token id |
| 596 | pid_h = tl.program_id(1) |
| 597 | req_id = pid_b // decode_query_len |
| 598 | q_offset = pid_b - req_id * decode_query_len |
| 599 | |
| 600 | if USE_PDL: |
| 601 | tl.extra.cuda.gdc_wait() |
| 602 | tl.extra.cuda.gdc_launch_dependents() |
| 603 | |
| 604 | seq_len = tl.load(seq_lens + req_id) |
| 605 | query_pos = seq_len - decode_query_len + q_offset |
| 606 | # Full-CG padding uses zero-length request rows. Clamp to an empty |
| 607 | # attention range instead of letting padded rows produce negative lengths. |
| 608 | kv_len = tl.maximum(query_pos + 1, 0) |
| 609 | num_blocks = (kv_len + block_size - 1) // block_size |
| 610 | |
| 611 | # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K. |
| 612 | # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T, |
| 613 | # in_chunk = p % BLOCK_SIZE_T. |
| 614 | off = tl.arange(0, BLOCK_SIZE_K) |
| 615 | chunk_idx = off // BLOCK_SIZE_T |
| 616 | in_chunk_idx = off % BLOCK_SIZE_T |
| 617 | valid = chunk_idx < num_topk_chunks |
| 618 | |
| 619 | score_offset = ( |
| 620 | chunk_idx * stride_ts_c |
| 621 | + pid_h * stride_ts_h |
| 622 | + pid_b * stride_ts_b |
| 623 | + in_chunk_idx * stride_ts_t |
| 624 | ) |
| 625 | idx_offset = ( |
| 626 | chunk_idx * stride_ti_c |
| 627 | + pid_h * stride_ti_h |
| 628 | + pid_b * stride_ti_b |
| 629 | + in_chunk_idx * stride_ti_t |
| 630 | ) |
| 631 | |
| 632 | score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to( |
| 633 | tl.float32 |
| 634 | ) |
| 635 | score = tl.where(score != score, -1e30, score) |
| 636 | idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32) |
| 637 | |
| 638 | # Full bitonic descending sort of BLOCK_SIZE_K items. |
| 639 | n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) |
| 640 | for j in tl.static_range(1, n_dims): |
| 641 | score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims) |
| 642 | score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims) |
| 643 | |
| 644 | # Extract first BLOCK_SIZE_T positions — these are the global top-K. |
| 645 | extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 |
| 646 | topk_idx_final = tl.sum( |
| 647 | extract_mask[:, None] |
| 648 | * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), |
| 649 | axis=0, |
| 650 | ) |
| 651 | |
| 652 | off_t = tl.arange(0, BLOCK_SIZE_T) |
| 653 | tif_ptrs = ( |
| 654 | ti_final_ptr |
| 655 | + pid_h * stride_tif_h |
| 656 | + pid_b * stride_tif_b |
| 657 | + off_t * stride_tif_t |
| 658 | ) |
| 659 | store_mask = off_t < topk |
| 660 | topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1) |
| 661 | tl.store( |
| 662 | tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask |
| 663 | ) |
| 664 | |
| 665 | |
| 666 | # --------------------------------------------------------------------------- |
| 667 | # Python wrappers |
| 668 | # --------------------------------------------------------------------------- |
| 669 | @torch.no_grad() |
| 670 | def minimax_m3_index_score( |
| 671 | idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] |
| 672 | index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] |
| 673 | block_table: torch.Tensor, # [batch, max_blocks] |
| 674 | cu_seqlens_q: torch.Tensor, # [batch+1] int32 |
| 675 | seq_lens: torch.Tensor, # [batch] int32 |
| 676 | prefix_lens: torch.Tensor, # [batch] int32 |
| 677 | max_query_len: int, |
| 678 | max_seq_len: int, |
| 679 | num_kv_heads: int, |
| 680 | ) -> torch.Tensor: |
| 681 | """Compute per-token index scores for each visible sparse block. |
| 682 | |
| 683 | Returns score [num_kv_heads, total_q, max_block], where each score is the |
| 684 | max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads. |
| 685 | """ |
| 686 | total_q, num_idx_heads, head_dim = idx_q.shape |
| 687 | assert num_idx_heads == num_kv_heads, ( |
| 688 | "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" |
| 689 | ) |
| 690 | batch = cu_seqlens_q.shape[0] - 1 |
| 691 | max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) |
| 692 | |
| 693 | # Keep score strides 16-divisible to avoid Triton recompiles. |
| 694 | score_block_stride = round_up(max_block, 16) |
| 695 | score = torch.empty( |
| 696 | (num_idx_heads, total_q, score_block_stride), |
| 697 | dtype=torch.float32, |
| 698 | device=idx_q.device, |
| 699 | ) |
| 700 | BLOCK_SIZE_Q = 64 |
| 701 | n_q_tiles = triton.cdiv(max_query_len, BLOCK_SIZE_Q) |
| 702 | SCORE_TARGET_GRID = 48 |
| 703 | split_k = max( |
| 704 | 1, |
| 705 | min(max_block, SCORE_TARGET_GRID // max(1, n_q_tiles * batch * num_idx_heads)), |
| 706 | ) |
| 707 | if not ( |
| 708 | current_platform.is_cuda() and current_platform.is_device_capability((12, 0)) |
| 709 | ): |
| 710 | split_k = 1 |
| 711 | grid_score = (n_q_tiles, batch * num_idx_heads, split_k) |
| 712 | _index_block_score_kernel[grid_score]( |
| 713 | idx_q, |
| 714 | index_kv_cache, |
| 715 | score, |
| 716 | block_table, |
| 717 | cu_seqlens_q, |
| 718 | seq_lens, |
| 719 | prefix_lens, |
| 720 | num_idx_heads, |
| 721 | head_dim, |
| 722 | idx_q.stride(0), |
| 723 | idx_q.stride(1), |
| 724 | idx_q.stride(2), |
| 725 | index_kv_cache.stride(0), |
| 726 | index_kv_cache.stride(1), |
| 727 | index_kv_cache.stride(2), |
| 728 | score.stride(0), |
| 729 | score.stride(1), |
| 730 | score.stride(2), |
| 731 | block_table.stride(0), |
| 732 | BLOCK_SIZE_Q=BLOCK_SIZE_Q, |
| 733 | BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, |
| 734 | USE_SPLIT_K=split_k > 1, |
| 735 | ) |
| 736 | return score |
| 737 | |
| 738 | |
| 739 | @torch.no_grad() |
| 740 | def minimax_m3_index_topk( |
| 741 | score: torch.Tensor, # [num_idx_heads, total_q, max_block] |
| 742 | cu_seqlens_q: torch.Tensor, # [batch+1] int32 |
| 743 | prefix_lens: torch.Tensor, # [batch] int32 |
| 744 | max_query_len: int, |
| 745 | topk: int, |
| 746 | init_blocks: int, |
| 747 | local_blocks: int, |
| 748 | out: torch.Tensor | None = None, |
| 749 | ) -> torch.Tensor: |
| 750 | """Select index top-k from a precomputed score tensor. |
| 751 | |
| 752 | When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the |
| 753 | result is written into ``out[:, :total_q, :]`` instead of a fresh tensor -- |
| 754 | used to keep the top-k output at a stable address for cudagraph capture. |
| 755 | """ |
| 756 | num_idx_heads = score.shape[0] |
| 757 | batch = cu_seqlens_q.shape[0] - 1 |
| 758 | total_q = score.shape[1] |
| 759 | if out is not None: |
| 760 | topk_idx = out[:, :total_q, :] |
| 761 | else: |
| 762 | topk_idx = torch.empty( |
| 763 | (num_idx_heads, total_q, topk), |
| 764 | dtype=torch.int32, |
| 765 | device=score.device, |
| 766 | ) |
| 767 | # block_size_q == 1 -> query blocks coincide with query tokens. |
| 768 | grid_topk = (max_query_len, batch, num_idx_heads) |
| 769 | _topk_index_kernel[grid_topk]( |
| 770 | score, |
| 771 | topk_idx, |
| 772 | 1, # sample_interval (block_size_q) |
| 773 | SPARSE_BLOCK_SIZE, |
| 774 | cu_seqlens_q, |
| 775 | cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1 |
| 776 | prefix_lens, |
| 777 | topk, |
| 778 | init_blocks, |
| 779 | local_blocks, |
| 780 | score.stride(0), |
| 781 | score.stride(1), |
| 782 | score.stride(2), |
| 783 | topk_idx.stride(0), |
| 784 | topk_idx.stride(1), |
| 785 | topk_idx.stride(2), |
| 786 | MASK_INIT=False, |
| 787 | MASK_LOCAL=False, |
| 788 | ) |
| 789 | return topk_idx |
| 790 | |
| 791 | |
| 792 | @torch.no_grad() |
| 793 | def minimax_m3_index_decode_score( |
| 794 | idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] |
| 795 | index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] |
| 796 | block_table: torch.Tensor, # [num_reqs, max_blocks] |
| 797 | seq_lens: torch.Tensor, # [num_reqs] int32 |
| 798 | max_seq_len: int, |
| 799 | init_blocks: int, |
| 800 | local_blocks: int, |
| 801 | num_kv_heads: int, |
| 802 | decode_query_len: int, |
| 803 | max_decode_query_len: int, |
| 804 | score_out: torch.Tensor | None = None, |
| 805 | ) -> torch.Tensor: |
| 806 | """Decode index block-score (split-K, cudagraph-safe); no top-k. |
| 807 | |
| 808 | Returns score [num_kv_heads, total_q, >=max_block] (fp32; init/local blocks |
| 809 | forced to 1e30/1e29). When ``score_out`` is given the scores are written into |
| 810 | it (read/written by strides, so a transposed view of a unified buffer is |
| 811 | accepted) instead of a fresh tensor -- used to share a unified score buffer |
| 812 | with the prefill side and run a single top-k over both. |
| 813 | """ |
| 814 | total_q, num_idx_heads, head_dim = idx_q.shape |
| 815 | assert num_idx_heads == num_kv_heads, ( |
| 816 | "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" |
| 817 | ) |
| 818 | assert decode_query_len <= max_decode_query_len |
| 819 | assert total_q == seq_lens.shape[0] * decode_query_len |
| 820 | max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) |
| 821 | use_pdl = current_platform.is_arch_support_pdl() |
| 822 | # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA |
| 823 | # SM9+); this ROCm Triton rejects it even when False ("Keyword argument |
| 824 | # launch_pdl was specified but unrecognised"). Only pass it when PDL is |
| 825 | # actually supported -- on ROCm use_pdl is always False, so it's omitted. |
| 826 | pdl_kwargs: dict[str, bool | int] = {} |
| 827 | if use_pdl: |
| 828 | pdl_kwargs.update({"launch_pdl": True}) |
| 829 | # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block; |
| 830 | # reduce stages to ease memory/register pressure. Keep no-spec and TP=4 |
| 831 | # single-head codegen unchanged. |
| 832 | score_kwargs = pdl_kwargs.copy() |
| 833 | if num_idx_heads > 1 and max_decode_query_len > 1: |
| 834 | score_kwargs.update({"num_warps": 4, "num_stages": 2}) |
| 835 | |
| 836 | if score_out is not None: |
| 837 | score = score_out |
| 838 | else: |
| 839 | # Keep score strides 16-divisible to avoid Triton recompiles. |
| 840 | score_block_stride = round_up(max_block, 16) |
| 841 | score = torch.empty( |
| 842 | (num_idx_heads, total_q, score_block_stride), |
| 843 | dtype=torch.float32, |
| 844 | device=idx_q.device, |
| 845 | ) |
| 846 | # split-K over seq blocks; chunk count depends only on shape constants so |
| 847 | # the grid is fixed within a cuda graph. |
| 848 | TARGET_GRID = 4096 |
| 849 | MAX_NUM_KV_CHUNKS = 256 |
| 850 | # Use the configured max decode length to avoid Triton recompiles when |
| 851 | # switching between qlen=1 and spec-decode verification batches. |
| 852 | BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len) |
| 853 | score_ctas_per_chunk = seq_lens.shape[0] |
| 854 | target = max( |
| 855 | 1, |
| 856 | min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)), |
| 857 | ) |
| 858 | num_kv_chunks = 1 << (target.bit_length() - 1) |
| 859 | grid_score = (seq_lens.shape[0], num_kv_chunks) |
| 860 | _decode_index_score_kernel[grid_score]( |
| 861 | idx_q, |
| 862 | index_kv_cache, |
| 863 | score, |
| 864 | block_table, |
| 865 | seq_lens, |
| 866 | num_idx_heads, |
| 867 | head_dim, |
| 868 | init_blocks, |
| 869 | local_blocks, |
| 870 | decode_query_len, |
| 871 | idx_q.stride(0), |
| 872 | idx_q.stride(1), |
| 873 | idx_q.stride(2), |
| 874 | index_kv_cache.stride(0), |
| 875 | index_kv_cache.stride(1), |
| 876 | index_kv_cache.stride(2), |
| 877 | score.stride(0), |
| 878 | score.stride(1), |
| 879 | score.stride(2), |
| 880 | block_table.stride(0), |
| 881 | BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, |
| 882 | BLOCK_SIZE_Q=BLOCK_SIZE_Q, |
| 883 | num_kv_chunks=num_kv_chunks, |
| 884 | USE_PDL=use_pdl, |
| 885 | **score_kwargs, |
| 886 | ) |
| 887 | return score |
| 888 | |
| 889 | |
| 890 | @torch.no_grad() |
| 891 | def minimax_m3_index_decode( |
| 892 | idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] |
| 893 | index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] |
| 894 | block_table: torch.Tensor, # [num_reqs, max_blocks] |
| 895 | seq_lens: torch.Tensor, # [num_reqs] int32 |
| 896 | max_seq_len: int, |
| 897 | topk: int, |
| 898 | init_blocks: int, |
| 899 | local_blocks: int, |
| 900 | num_kv_heads: int, |
| 901 | decode_query_len: int, |
| 902 | max_decode_query_len: int, |
| 903 | out: torch.Tensor | None = None, |
| 904 | score_out: torch.Tensor | None = None, |
| 905 | ) -> torch.Tensor: |
| 906 | """Decode index block-score + top-k, both split-K (cudagraph-safe). |
| 907 | |
| 908 | Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). |
| 909 | When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into |
| 910 | ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. |
| 911 | When ``score_out`` ([num_kv_heads, total_q, >=max_block]) is given, the block |
| 912 | scores are written into it (read back by the top-k) instead of a fresh |
| 913 | tensor -- used to share a unified score buffer with the prefill side. Reads |
| 914 | via strides, so a transposed view of a block-major buffer is accepted. |
| 915 | """ |
| 916 | total_q, num_idx_heads, _ = idx_q.shape |
| 917 | batch = total_q |
| 918 | max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) |
| 919 | use_pdl = current_platform.is_arch_support_pdl() |
| 920 | pdl_kwargs: dict[str, bool | int] = {} |
| 921 | if use_pdl: |
| 922 | pdl_kwargs.update({"launch_pdl": True}) |
| 923 | score = minimax_m3_index_decode_score( |
| 924 | idx_q, |
| 925 | index_kv_cache, |
| 926 | block_table, |
| 927 | seq_lens, |
| 928 | max_seq_len, |
| 929 | init_blocks, |
| 930 | local_blocks, |
| 931 | num_kv_heads, |
| 932 | decode_query_len, |
| 933 | max_decode_query_len, |
| 934 | score_out=score_out, |
| 935 | ) |
| 936 | |
| 937 | if out is not None: |
| 938 | topk_idx = out[:, :total_q, :] |
| 939 | else: |
| 940 | topk_idx = torch.empty( |
| 941 | (num_idx_heads, total_q, topk), |
| 942 | dtype=torch.int32, |
| 943 | device=idx_q.device, |
| 944 | ) |
| 945 | # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts |
| 946 | # pow2(num_topk_chunks * pow2(topk)) candidates. |
| 947 | TOPK_TARGET_GRID = 512 |
| 948 | MAX_NUM_TOPK_CHUNKS = 16 |
| 949 | topk_target = max( |
| 950 | 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads)) |
| 951 | ) |
| 952 | num_topk_chunks = 1 << (topk_target.bit_length() - 1) |
| 953 | block_size_t = triton.next_power_of_2(topk) |
| 954 | chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks |
| 955 | topk_score_partial = torch.empty( |
| 956 | num_topk_chunks, |
| 957 | num_idx_heads, |
| 958 | batch, |
| 959 | block_size_t, |
| 960 | dtype=torch.float32, |
| 961 | device=idx_q.device, |
| 962 | ) |
| 963 | topk_idx_partial = torch.empty( |
| 964 | num_topk_chunks, |
| 965 | num_idx_heads, |
| 966 | batch, |
| 967 | block_size_t, |
| 968 | dtype=torch.int32, |
| 969 | device=idx_q.device, |
| 970 | ) |
| 971 | _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)]( |
| 972 | score, |
| 973 | topk_score_partial, |
| 974 | topk_idx_partial, |
| 975 | seq_lens, |
| 976 | SPARSE_BLOCK_SIZE, |
| 977 | topk, |
| 978 | chunk_blocks, |
| 979 | decode_query_len, |
| 980 | score.stride(0), |
| 981 | score.stride(1), |
| 982 | score.stride(2), |
| 983 | topk_score_partial.stride(0), |
| 984 | topk_score_partial.stride(1), |
| 985 | topk_score_partial.stride(2), |
| 986 | topk_score_partial.stride(3), |
| 987 | topk_idx_partial.stride(0), |
| 988 | topk_idx_partial.stride(1), |
| 989 | topk_idx_partial.stride(2), |
| 990 | topk_idx_partial.stride(3), |
| 991 | USE_PDL=use_pdl, |
| 992 | **pdl_kwargs, |
| 993 | ) |
| 994 | _topk_index_merge_kernel[(batch, num_idx_heads)]( |
| 995 | topk_score_partial, |
| 996 | topk_idx_partial, |
| 997 | topk_idx, |
| 998 | seq_lens, |
| 999 | SPARSE_BLOCK_SIZE, |
| 1000 | topk, |
| 1001 | decode_query_len, |
| 1002 | topk_score_partial.stride(0), |
| 1003 | topk_score_partial.stride(1), |
| 1004 | topk_score_partial.stride(2), |
| 1005 | topk_score_partial.stride(3), |
| 1006 | topk_idx_partial.stride(0), |
| 1007 | topk_idx_partial.stride(1), |
| 1008 | topk_idx_partial.stride(2), |
| 1009 | topk_idx_partial.stride(3), |
| 1010 | topk_idx.stride(0), |
| 1011 | topk_idx.stride(1), |
| 1012 | topk_idx.stride(2), |
| 1013 | num_topk_chunks=num_topk_chunks, |
| 1014 | USE_PDL=use_pdl, |
| 1015 | **pdl_kwargs, |
| 1016 | ) |
| 1017 | return topk_idx |