| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utility functions for sparse MLA backends.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.distributed import get_dcp_group |
| 13 | from vllm.model_executor.warmup.jit_warmup import ( |
| 14 | kernel_launcher, |
| 15 | zip_inputs, |
| 16 | ) |
| 17 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 18 | LaunchSpec, |
| 19 | TritonWarmupTensor, |
| 20 | VllmTritonJitKernel, |
| 21 | triton_scalar_specialization_rep, |
| 22 | ) |
| 23 | from vllm.triton_utils import tl, triton |
| 24 | from vllm.utils.math_utils import cdiv |
| 25 | from vllm.v1.worker.block_table import get_block_table_width |
| 26 | |
| 27 | |
| 28 | def request_row_bounds(req_idx: np.ndarray) -> np.ndarray: |
| 29 | """Bounds of the runs of adjacent rows that belong to one request: run |
| 30 | ``r`` is rows ``[bounds[r], bounds[r + 1])``. |
| 31 | |
| 32 | Under PCP a rank holds two adjacent chunk rows of a split prefill; the |
| 33 | sparse backends give such a run one KV region. |
| 34 | """ |
| 35 | assert req_idx.size > 0 |
| 36 | bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True]) |
| 37 | assert bounds.size - 1 == np.unique(req_idx).size, ( |
| 38 | "rows of one request must be adjacent" |
| 39 | ) |
| 40 | return bounds |
| 41 | |
| 42 | |
| 43 | def flat_kv_row_view( |
| 44 | kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] |
| 45 | block_size: int, |
| 46 | ) -> tuple[torch.Tensor, int]: |
| 47 | """Flat [row, head_dim] view of a paged cache and its physical rows per block. |
| 48 | |
| 49 | Token offset is ``block_idx * block_stride_rows + offset_in_block``. |
| 50 | When other layers' pages sit between consecutive blocks of this cache, |
| 51 | ``block_stride_rows`` exceeds ``block_size``; those in-between rows are never |
| 52 | indexed (`triton_convert_req_index_to_global_index` ensures this). |
| 53 | """ |
| 54 | num_blocks, _, head_dim = kv_cache.shape |
| 55 | assert kv_cache.stride(0) % head_dim == 0, ( |
| 56 | "block stride is not a whole number of rows; flat row indexing would " |
| 57 | "silently misaddress" |
| 58 | ) |
| 59 | block_stride_rows = kv_cache.stride(0) // head_dim |
| 60 | num_rows = (num_blocks - 1) * block_stride_rows + block_size |
| 61 | rows = kv_cache.as_strided((num_rows, head_dim), (head_dim, 1)) |
| 62 | return rows, block_stride_rows |
| 63 | |
| 64 | |
| 65 | # Kernel with prefill workspace support and valid count tracking |
| 66 | |
| 67 | |
| 68 | class ConvertReqIndexToGlobalIndexKernel( |
| 69 | VllmTritonJitKernel["ConvertReqIndexToGlobalIndexKernel.CompileKey"] |
| 70 | ): |
| 71 | @dataclass(frozen=True) |
| 72 | class CompileKey: |
| 73 | block_size: int |
| 74 | block_stride_rows: int |
| 75 | block_n: int |
| 76 | has_prefill_workspace: bool |
| 77 | count_valid: bool |
| 78 | single_tile: bool |
| 79 | compact_to_front: bool |
| 80 | dcp_size: int |
| 81 | dcp_rank: int |
| 82 | dcp_interleave: int |
| 83 | num_warps: int |
| 84 | num_topk_tokens: int |
| 85 | block_table_stride: int |
| 86 | |
| 87 | @staticmethod |
| 88 | @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"]) |
| 89 | def sparse_mla_index_remap_kernel( |
| 90 | req_id_ptr, # int32 [num_tokens] |
| 91 | block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] |
| 92 | token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 93 | out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 94 | valid_count_ptr, # int32 [num_tokens] - output valid count per row |
| 95 | prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill |
| 96 | workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr |
| 97 | workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace |
| 98 | # shapes (compile-time where possible) |
| 99 | max_num_blocks_per_req, |
| 100 | BLOCK_SIZE: tl.constexpr, |
| 101 | BLOCK_STRIDE_ROWS: tl.constexpr, |
| 102 | BLOCK_N: tl.constexpr, # tile width along columns |
| 103 | NUM_TOPK_TOKENS: tl.constexpr, |
| 104 | HAS_PREFILL: tl.constexpr, |
| 105 | COUNT_VALID: tl.constexpr, # whether to count valid indices |
| 106 | # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count |
| 107 | # is an in-register reduction and needs no atomic. |
| 108 | SINGLE_TILE: tl.constexpr, |
| 109 | # When set, scatter valid slots to a contiguous prefix [0, valid_count) using |
| 110 | # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior |
| 111 | # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries). |
| 112 | # Requires COUNT_VALID; multi-tile output must be pre-filled with -1. The |
| 113 | # prefix is unspecified (only the selected set matters). |
| 114 | COMPACT_TO_FRONT: tl.constexpr, |
| 115 | # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op |
| 116 | DCP_SIZE: tl.constexpr, |
| 117 | DCP_RANK: tl.constexpr, |
| 118 | DCP_INTERLEAVE: tl.constexpr, |
| 119 | # strides (in elements) |
| 120 | bt_stride0, |
| 121 | bt_stride1, |
| 122 | ti_stride0, |
| 123 | ti_stride1, |
| 124 | out_stride0, |
| 125 | out_stride1, |
| 126 | ): |
| 127 | # program_id(0) -> token_id (row) |
| 128 | # program_id(1) -> tile index along columns |
| 129 | token_id = tl.program_id(0) |
| 130 | tile_id = tl.program_id(1) |
| 131 | |
| 132 | # Each program covers BLOCK_N consecutive columns |
| 133 | indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) |
| 134 | |
| 135 | # Load request id for this token (no mask: grid is exact) |
| 136 | req = tl.load(req_id_ptr + token_id) |
| 137 | |
| 138 | # Load token indices for this tile |
| 139 | ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 |
| 140 | tok = tl.load(ti_ptr, mask=indice_id < NUM_TOPK_TOKENS, other=-1) |
| 141 | |
| 142 | # Only token == -1 should propagate as -1 |
| 143 | is_invalid_tok = tok < 0 |
| 144 | is_prefill = False |
| 145 | if HAS_PREFILL: |
| 146 | prefill_req_id = tl.load(prefill_request_id_ptr + token_id) |
| 147 | is_prefill = prefill_req_id >= 0 |
| 148 | |
| 149 | # DCP de-interleave the global token id into this rank's local slot. |
| 150 | # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With |
| 151 | # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never |
| 152 | # remote) and local_idx == tok, so this reduces to the non-DCP path; with |
| 153 | # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE). |
| 154 | owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE |
| 155 | is_remote = owning_rank != DCP_RANK |
| 156 | local_idx = ( |
| 157 | tok // (DCP_SIZE * DCP_INTERLEAVE) |
| 158 | ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE |
| 159 | |
| 160 | # Compute block id and in-block offset |
| 161 | block_id = local_idx // BLOCK_SIZE |
| 162 | inblock_off = local_idx % BLOCK_SIZE |
| 163 | |
| 164 | # Guard block_table access |
| 165 | valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) |
| 166 | bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 |
| 167 | is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill) |
| 168 | base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0) |
| 169 | out_val = base * BLOCK_STRIDE_ROWS + inblock_off |
| 170 | |
| 171 | # Override with prefill output if prefill is enabled |
| 172 | if HAS_PREFILL: |
| 173 | workspace_start = tl.load( |
| 174 | workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 |
| 175 | ) |
| 176 | # Under DCP the prefill workspace is the all-gather of every rank's |
| 177 | # shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok. |
| 178 | prefill_out = ( |
| 179 | owning_rank * workspace_rank_stride + workspace_start + local_idx |
| 180 | ) |
| 181 | out_val = tl.where(is_prefill, prefill_out, out_val) |
| 182 | out_val = tl.where(is_invalid_tok, -1, out_val) |
| 183 | |
| 184 | if COMPACT_TO_FRONT: |
| 185 | # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix |
| 186 | # sum gives each valid lane a distinct local offset; one atomic add of the |
| 187 | # tile's valid count reserves a contiguous base across racing tiles. The |
| 188 | # out buffer is pre-filled with -1, so unwritten tail slots stay -1. |
| 189 | # With no racing tiles the base is 0 and the allocator becomes a store. |
| 190 | is_valid = (~is_invalid_tok).to(tl.int32) |
| 191 | local_offset = tl.cumsum(is_valid) - is_valid |
| 192 | tile_valid_count = tl.sum(is_valid) |
| 193 | if SINGLE_TILE: |
| 194 | base = 0 |
| 195 | tl.store(valid_count_ptr + token_id, tile_valid_count) |
| 196 | else: |
| 197 | base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) |
| 198 | dest = base + local_offset |
| 199 | out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1 |
| 200 | tl.store(out_ptr_dest, out_val, mask=is_valid == 1) |
| 201 | if SINGLE_TILE: |
| 202 | # This program owns both the compacted prefix and its padding. |
| 203 | # The stores are disjoint, so no separate initialization is needed. |
| 204 | tl.store( |
| 205 | out_ptr + token_id * out_stride0 + indice_id * out_stride1, |
| 206 | -1, |
| 207 | mask=(indice_id >= tile_valid_count) |
| 208 | & (indice_id < NUM_TOPK_TOKENS), |
| 209 | ) |
| 210 | else: |
| 211 | # Store results in place (input column == output column). |
| 212 | out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 |
| 213 | tl.store(out_ptr_ij, out_val, mask=indice_id < NUM_TOPK_TOKENS) |
| 214 | |
| 215 | # Accumulate the tile's valid count into the row total; a single tile's |
| 216 | # reduction *is* the total. |
| 217 | if COUNT_VALID: |
| 218 | tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32)) |
| 219 | if SINGLE_TILE: |
| 220 | tl.store(valid_count_ptr + token_id, tile_valid_count) |
| 221 | else: |
| 222 | tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) |
| 223 | |
| 224 | # Keep the warmup wrapper interface while exposing a descriptive CUDA symbol. |
| 225 | kernel = sparse_mla_index_remap_kernel |
| 226 | |
| 227 | def dispatch( # type: ignore[override] |
| 228 | self, |
| 229 | *, |
| 230 | BLOCK_SIZE: int, |
| 231 | BLOCK_STRIDE_ROWS: int, |
| 232 | BLOCK_N: int, |
| 233 | NUM_TOPK_TOKENS: int, |
| 234 | HAS_PREFILL_WORKSPACE: bool, |
| 235 | COUNT_VALID: bool, |
| 236 | COMPACT_TO_FRONT: bool, |
| 237 | DCP_SIZE: int, |
| 238 | DCP_RANK: int, |
| 239 | DCP_INTERLEAVE: int, |
| 240 | max_num_blocks_per_req: int, |
| 241 | ) -> CompileKey: |
| 242 | tiling = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID) |
| 243 | return self.CompileKey( |
| 244 | block_size=BLOCK_SIZE, |
| 245 | block_stride_rows=BLOCK_STRIDE_ROWS, |
| 246 | block_n=tiling[1], |
| 247 | has_prefill_workspace=HAS_PREFILL_WORKSPACE, |
| 248 | count_valid=COUNT_VALID, |
| 249 | single_tile=tiling[0], |
| 250 | compact_to_front=COMPACT_TO_FRONT, |
| 251 | dcp_size=DCP_SIZE, |
| 252 | dcp_rank=DCP_RANK, |
| 253 | dcp_interleave=DCP_INTERLEAVE, |
| 254 | num_warps=tiling[3], |
| 255 | num_topk_tokens=NUM_TOPK_TOKENS, |
| 256 | block_table_stride=triton_scalar_specialization_rep(max_num_blocks_per_req), |
| 257 | ) |
| 258 | |
| 259 | def get_warmup_keys( |
| 260 | self, |
| 261 | vllm_config: VllmConfig, |
| 262 | *, |
| 263 | block_stride_rows: int, |
| 264 | ) -> list[CompileKey]: |
| 265 | block_size = vllm_config.cache_config.block_size |
| 266 | dcp_size = vllm_config.parallel_config.decode_context_parallel_size |
| 267 | dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size |
| 268 | dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0 |
| 269 | num_topk_tokens = vllm_config.model_config.hf_text_config.index_topk |
| 270 | max_num_blocks = cdiv( |
| 271 | vllm_config.model_config.max_model_len, |
| 272 | block_size * dcp_size, |
| 273 | ) |
| 274 | max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size) |
| 275 | pcp_dcp_combos = [] |
| 276 | if ( |
| 277 | vllm_config.parallel_config.prefill_context_parallel_size > 1 |
| 278 | and dcp_size > 1 |
| 279 | ): |
| 280 | # PCP + DCP attends prefill rows over a DCP-gathered workspace. |
| 281 | pcp_dcp_combos.append( |
| 282 | dict( |
| 283 | HAS_PREFILL_WORKSPACE=True, |
| 284 | COUNT_VALID=True, |
| 285 | COMPACT_TO_FRONT=True, |
| 286 | DCP_SIZE=dcp_size, |
| 287 | DCP_RANK=dcp_rank, |
| 288 | DCP_INTERLEAVE=dcp_interleave, |
| 289 | ) |
| 290 | ) |
| 291 | return self._trace_dispatch(self.dispatch)( |
| 292 | zip_inputs( |
| 293 | dict( |
| 294 | HAS_PREFILL_WORKSPACE=False, |
| 295 | COUNT_VALID=False, |
| 296 | COMPACT_TO_FRONT=False, |
| 297 | DCP_SIZE=1, |
| 298 | DCP_RANK=0, |
| 299 | DCP_INTERLEAVE=1, |
| 300 | ), |
| 301 | dict( |
| 302 | HAS_PREFILL_WORKSPACE=False, |
| 303 | COUNT_VALID=True, |
| 304 | COMPACT_TO_FRONT=True, |
| 305 | DCP_SIZE=1, |
| 306 | DCP_RANK=0, |
| 307 | DCP_INTERLEAVE=1, |
| 308 | ), |
| 309 | dict( |
| 310 | HAS_PREFILL_WORKSPACE=True, |
| 311 | COUNT_VALID=True, |
| 312 | COMPACT_TO_FRONT=True, |
| 313 | DCP_SIZE=1, |
| 314 | DCP_RANK=0, |
| 315 | DCP_INTERLEAVE=1, |
| 316 | ), |
| 317 | dict( |
| 318 | HAS_PREFILL_WORKSPACE=False, |
| 319 | COUNT_VALID=True, |
| 320 | COMPACT_TO_FRONT=False, |
| 321 | DCP_SIZE=dcp_size, |
| 322 | DCP_RANK=dcp_rank, |
| 323 | DCP_INTERLEAVE=dcp_interleave, |
| 324 | ), |
| 325 | dict( |
| 326 | HAS_PREFILL_WORKSPACE=False, |
| 327 | COUNT_VALID=True, |
| 328 | COMPACT_TO_FRONT=dcp_size > 1, |
| 329 | DCP_SIZE=dcp_size, |
| 330 | DCP_RANK=dcp_rank, |
| 331 | DCP_INTERLEAVE=dcp_interleave, |
| 332 | ), |
| 333 | *pcp_dcp_combos, |
| 334 | ), |
| 335 | BLOCK_SIZE=block_size, |
| 336 | BLOCK_STRIDE_ROWS=block_stride_rows, |
| 337 | BLOCK_N=128, |
| 338 | NUM_TOPK_TOKENS=num_topk_tokens, |
| 339 | max_num_blocks_per_req=max_num_blocks_per_req, |
| 340 | ) |
| 341 | |
| 342 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 343 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 344 | token_shape = (1, compile_key.num_topk_tokens) |
| 345 | token_strides = (16, 1) |
| 346 | return dict( |
| 347 | req_id=int32_ptr, |
| 348 | block_table=TritonWarmupTensor( |
| 349 | torch.int32, |
| 350 | shape=(1, compile_key.block_table_stride), |
| 351 | ), |
| 352 | token_indices=TritonWarmupTensor( |
| 353 | torch.int32, |
| 354 | shape=token_shape, |
| 355 | strides=token_strides, |
| 356 | ), |
| 357 | out=TritonWarmupTensor( |
| 358 | torch.int32, |
| 359 | shape=token_shape, |
| 360 | strides=token_strides, |
| 361 | ), |
| 362 | valid_counts=int32_ptr if compile_key.count_valid else None, |
| 363 | prefill_workspace_request_ids=( |
| 364 | int32_ptr if compile_key.has_prefill_workspace else None |
| 365 | ), |
| 366 | prefill_workspace_starts=( |
| 367 | int32_ptr if compile_key.has_prefill_workspace else None |
| 368 | ), |
| 369 | max_num_blocks_per_req=1, |
| 370 | workspace_rank_stride=1, |
| 371 | BLOCK_SIZE=compile_key.block_size, |
| 372 | BLOCK_STRIDE_ROWS=compile_key.block_stride_rows, |
| 373 | BLOCK_N=compile_key.block_n, |
| 374 | NUM_TOPK_TOKENS=compile_key.num_topk_tokens, |
| 375 | HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace, |
| 376 | COUNT_VALID=compile_key.count_valid, |
| 377 | COMPACT_TO_FRONT=compile_key.compact_to_front, |
| 378 | DCP_SIZE=compile_key.dcp_size, |
| 379 | DCP_RANK=compile_key.dcp_rank, |
| 380 | DCP_INTERLEAVE=compile_key.dcp_interleave, |
| 381 | ) |
| 382 | |
| 383 | @kernel_launcher |
| 384 | def __call__( |
| 385 | self, |
| 386 | req_id: torch.Tensor, |
| 387 | block_table: torch.Tensor, |
| 388 | token_indices: torch.Tensor, |
| 389 | out: torch.Tensor, |
| 390 | valid_counts: torch.Tensor | None, |
| 391 | prefill_workspace_request_ids: torch.Tensor | None, |
| 392 | prefill_workspace_starts: torch.Tensor | None, |
| 393 | *, |
| 394 | workspace_rank_stride: int, |
| 395 | max_num_blocks_per_req: int, |
| 396 | BLOCK_SIZE: int, |
| 397 | BLOCK_STRIDE_ROWS: int, |
| 398 | BLOCK_N: int, |
| 399 | NUM_TOPK_TOKENS: int, |
| 400 | HAS_PREFILL_WORKSPACE: bool, |
| 401 | COUNT_VALID: bool, |
| 402 | COMPACT_TO_FRONT: bool, |
| 403 | DCP_SIZE: int, |
| 404 | DCP_RANK: int, |
| 405 | DCP_INTERLEAVE: int, |
| 406 | ) -> LaunchSpec: |
| 407 | single_tile, block_n, tiles_per_row, num_warps = _remap_tiling( |
| 408 | NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID |
| 409 | ) |
| 410 | # Exact 2D grid: tokens × column tiles |
| 411 | return (req_id.shape[0], tiles_per_row), dict( |
| 412 | valid_count_ptr=valid_counts, |
| 413 | prefill_request_id_ptr=prefill_workspace_request_ids, |
| 414 | workspace_starts_ptr=prefill_workspace_starts, |
| 415 | BLOCK_N=block_n, |
| 416 | HAS_PREFILL=HAS_PREFILL_WORKSPACE, |
| 417 | SINGLE_TILE=single_tile, |
| 418 | # Strides in elements |
| 419 | bt_stride0=block_table.stride(0), |
| 420 | bt_stride1=block_table.stride(1), |
| 421 | ti_stride0=token_indices.stride(0), |
| 422 | ti_stride1=token_indices.stride(1), |
| 423 | out_stride0=out.stride(0), |
| 424 | out_stride1=out.stride(1), |
| 425 | num_warps=num_warps, |
| 426 | ) |
| 427 | |
| 428 | |
| 429 | def _remap_tiling( |
| 430 | NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool |
| 431 | ) -> tuple[bool, int, int, int]: |
| 432 | """Pick the column tiling for the index remap kernel. |
| 433 | |
| 434 | Counting the valid slots per row is the only reason the column tiles have to |
| 435 | talk to each other, so when counting give one program the whole row: the |
| 436 | count becomes an in-register reduction plus a plain store, needing neither |
| 437 | atomics nor a zero-initialized counter. Pad modest non-power-of-two widths |
| 438 | (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic. |
| 439 | |
| 440 | Returns: |
| 441 | (single_tile, block_n, tiles_per_row, num_warps) |
| 442 | |
| 443 | """ |
| 444 | padded_width = triton.next_power_of_2(NUM_TOPK_TOKENS) |
| 445 | single_tile = count_valid and ( |
| 446 | padded_width == NUM_TOPK_TOKENS or padded_width <= 4096 |
| 447 | ) |
| 448 | if single_tile: |
| 449 | return True, padded_width, 1, 8 |
| 450 | return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4 |
| 451 | |
| 452 | |
| 453 | def triton_convert_req_index_to_global_index( |
| 454 | req_id: torch.Tensor, # int32 [num_tokens] |
| 455 | block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] |
| 456 | token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 457 | BLOCK_SIZE: int = 64, |
| 458 | BLOCK_STRIDE_ROWS: int | None = None, |
| 459 | NUM_TOPK_TOKENS: int = 2048, |
| 460 | BLOCK_N: int = 128, # tile width along columns |
| 461 | HAS_PREFILL_WORKSPACE: bool = False, |
| 462 | prefill_workspace_request_ids: torch.Tensor | None = None, |
| 463 | prefill_workspace_starts: torch.Tensor | None = None, |
| 464 | prefill_workspace_rank_stride: int | None = None, |
| 465 | dcp_size: int = 1, |
| 466 | dcp_rank: int = 0, |
| 467 | cp_kv_cache_interleave_size: int = 1, |
| 468 | return_valid_counts: bool = False, |
| 469 | out: torch.Tensor | None = None, |
| 470 | valid_counts_out: torch.Tensor | None = None, |
| 471 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 472 | """out[token_id, indice_id] = |
| 473 | block_table[req_id[token_id], |
| 474 | token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE |
| 475 | + token_indices[token_id, indice_id] % BLOCK_SIZE |
| 476 | |
| 477 | Only when token_indices[token_id, indice_id] == -1 do we output -1. |
| 478 | For safety, we also output -1 if the derived block_id would be |
| 479 | out-of-bounds. |
| 480 | |
| 481 | When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets |
| 482 | instead of global cache slots. prefill_workspace_request_ids and |
| 483 | prefill_workspace_starts must be provided. |
| 484 | |
| 485 | prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else |
| 486 | prefill request index (maps to prefill_workspace_starts) |
| 487 | prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace |
| 488 | starts for each prefill request |
| 489 | |
| 490 | When return_valid_counts is True, also returns the count of valid (non -1) |
| 491 | indices per row, computed during the same kernel pass (no extra overhead). |
| 492 | """ |
| 493 | assert req_id.dtype == torch.int32 |
| 494 | assert block_table.dtype == torch.int32 |
| 495 | assert token_indices.dtype == torch.int32 |
| 496 | assert req_id.shape[0] == token_indices.shape[0], ( |
| 497 | f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) " |
| 498 | "must cover the same tokens; the grid is sized by req_id but the output " |
| 499 | "is allocated like token_indices, so a longer req_id writes out of bounds" |
| 500 | ) |
| 501 | assert token_indices.shape[1] == NUM_TOPK_TOKENS |
| 502 | assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( |
| 503 | f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" |
| 504 | ) |
| 505 | |
| 506 | if HAS_PREFILL_WORKSPACE: |
| 507 | assert prefill_workspace_request_ids is not None |
| 508 | assert prefill_workspace_starts is not None |
| 509 | assert prefill_workspace_request_ids.dtype == torch.int32 |
| 510 | assert prefill_workspace_starts.dtype == torch.int32 |
| 511 | |
| 512 | num_tokens = req_id.shape[0] |
| 513 | max_num_blocks_per_req = block_table.shape[1] |
| 514 | |
| 515 | single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts) |
| 516 | |
| 517 | # Ensure contiguous tensors on the same device |
| 518 | req_id_c = req_id.contiguous() |
| 519 | block_table_c = block_table.contiguous() |
| 520 | token_indices_c = token_indices.contiguous() |
| 521 | # Only multi-tile compaction needs separate padding initialization. |
| 522 | if out is None: |
| 523 | out = ( |
| 524 | torch.full_like(token_indices_c, -1) |
| 525 | if return_valid_counts and not single_tile |
| 526 | else torch.empty_like(token_indices_c) |
| 527 | ) |
| 528 | else: |
| 529 | assert out.dtype == token_indices_c.dtype |
| 530 | assert out.device == token_indices_c.device |
| 531 | assert out.shape == token_indices_c.shape |
| 532 | assert out.is_contiguous() |
| 533 | if return_valid_counts and not single_tile: |
| 534 | out.fill_(-1) |
| 535 | |
| 536 | valid_counts: torch.Tensor | None = None |
| 537 | if return_valid_counts: |
| 538 | if valid_counts_out is None: |
| 539 | # Zero-init only matters for the atomic accumulation path. |
| 540 | alloc = torch.empty if single_tile else torch.zeros |
| 541 | valid_counts = alloc( |
| 542 | num_tokens, dtype=torch.int32, device=token_indices.device |
| 543 | ) |
| 544 | else: |
| 545 | assert valid_counts_out.dtype == torch.int32 |
| 546 | assert valid_counts_out.device == token_indices.device |
| 547 | assert valid_counts_out.shape == (num_tokens,) |
| 548 | assert valid_counts_out.is_contiguous() |
| 549 | valid_counts = valid_counts_out |
| 550 | if not single_tile: |
| 551 | valid_counts.zero_() |
| 552 | |
| 553 | # Prepare prefill pointers |
| 554 | if HAS_PREFILL_WORKSPACE: |
| 555 | assert prefill_workspace_request_ids is not None # for mypy |
| 556 | assert prefill_workspace_starts is not None # for mypy |
| 557 | assert prefill_workspace_request_ids.is_contiguous() |
| 558 | assert prefill_workspace_starts.is_contiguous() |
| 559 | |
| 560 | _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL( |
| 561 | req_id_c, |
| 562 | block_table_c, |
| 563 | token_indices_c, |
| 564 | out, |
| 565 | valid_counts, |
| 566 | prefill_workspace_request_ids, |
| 567 | prefill_workspace_starts, |
| 568 | # shapes / constexprs |
| 569 | workspace_rank_stride=prefill_workspace_rank_stride or 0, |
| 570 | max_num_blocks_per_req=max_num_blocks_per_req, |
| 571 | BLOCK_SIZE=BLOCK_SIZE, |
| 572 | BLOCK_STRIDE_ROWS=( |
| 573 | BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE |
| 574 | ), |
| 575 | BLOCK_N=BLOCK_N, |
| 576 | NUM_TOPK_TOKENS=NUM_TOPK_TOKENS, |
| 577 | HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE, |
| 578 | COUNT_VALID=return_valid_counts, |
| 579 | COMPACT_TO_FRONT=return_valid_counts, |
| 580 | DCP_SIZE=dcp_size, |
| 581 | DCP_RANK=dcp_rank, |
| 582 | DCP_INTERLEAVE=cp_kv_cache_interleave_size, |
| 583 | ) |
| 584 | |
| 585 | if return_valid_counts: |
| 586 | assert valid_counts is not None |
| 587 | return out, valid_counts |
| 588 | return out |
| 589 | |
| 590 | |
| 591 | def triton_filter_and_convert_dcp_index( |
| 592 | req_id: torch.Tensor, |
| 593 | block_table: torch.Tensor, |
| 594 | token_indices: torch.Tensor, |
| 595 | dcp_size: int, |
| 596 | dcp_rank: int, |
| 597 | cp_kv_cache_interleave_size: int = 1, |
| 598 | BLOCK_SIZE: int = 64, |
| 599 | BLOCK_STRIDE_ROWS: int | None = None, |
| 600 | NUM_TOPK_TOKENS: int = 2048, |
| 601 | BLOCK_N: int = 128, |
| 602 | return_valid_counts: bool = False, |
| 603 | compact_valid_to_front: bool = True, |
| 604 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 605 | """Filter global per-request indices to this DCP rank's local slots. |
| 606 | |
| 607 | With ``compact_valid_to_front`` (default), the conversion kernel scatters |
| 608 | this rank's owned slots to a contiguous prefix ``[0, valid_count)`` and |
| 609 | leaves the rest ``-1``. DCP filtering marks non-owned slots ``-1`` and so |
| 610 | creates interior gaps; the trtllm-gen sparse kernel reads the first |
| 611 | ``valid_count`` entries of each row, so they must be a contiguous prefix. |
| 612 | Compaction is fused into the kernel (atomic slot allocator) rather than a |
| 613 | separate sort/gather pass. Prefix order is unspecified (only the set matters). |
| 614 | """ |
| 615 | assert dcp_size >= 1 |
| 616 | assert 0 <= dcp_rank < dcp_size |
| 617 | # Interleave groups must align to KV blocks (globally enforced by |
| 618 | # VllmConfig: block_size % cp_kv_cache_interleave_size == 0); assert the |
| 619 | # local invariant so local_idx // BLOCK_SIZE never straddles a group. |
| 620 | assert BLOCK_SIZE % cp_kv_cache_interleave_size == 0, ( |
| 621 | f"BLOCK_SIZE ({BLOCK_SIZE}) must be divisible by " |
| 622 | f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})." |
| 623 | ) |
| 624 | assert req_id.dtype == torch.int32 |
| 625 | assert block_table.dtype == torch.int32 |
| 626 | assert token_indices.dtype == torch.int32 |
| 627 | assert token_indices.shape[1] == NUM_TOPK_TOKENS |
| 628 | assert NUM_TOPK_TOKENS % BLOCK_N == 0 |
| 629 | |
| 630 | if dcp_size == 1: |
| 631 | return triton_convert_req_index_to_global_index( |
| 632 | req_id, |
| 633 | block_table, |
| 634 | token_indices, |
| 635 | BLOCK_SIZE=BLOCK_SIZE, |
| 636 | BLOCK_STRIDE_ROWS=BLOCK_STRIDE_ROWS, |
| 637 | NUM_TOPK_TOKENS=NUM_TOPK_TOKENS, |
| 638 | BLOCK_N=BLOCK_N, |
| 639 | return_valid_counts=return_valid_counts, |
| 640 | ) |
| 641 | |
| 642 | num_tokens = req_id.shape[0] |
| 643 | max_num_blocks_per_req = block_table.shape[1] |
| 644 | |
| 645 | req_id_c = req_id.contiguous() |
| 646 | block_table_c = block_table.contiguous() |
| 647 | token_indices_c = token_indices.contiguous() |
| 648 | |
| 649 | # The compaction uses the valid-count buffer as a slot allocator, so it |
| 650 | # requires counting. Only the multi-tile path needs pre-filled padding. |
| 651 | count_valid = return_valid_counts or compact_valid_to_front |
| 652 | |
| 653 | # The compaction builds on the counting, so it shares the tiling. |
| 654 | single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid) |
| 655 | |
| 656 | if compact_valid_to_front and not single_tile: |
| 657 | out = torch.full_like(token_indices_c, -1) |
| 658 | else: |
| 659 | out = torch.empty_like(token_indices_c) |
| 660 | |
| 661 | valid_counts: torch.Tensor | None = None |
| 662 | if count_valid: |
| 663 | # Zero-init only matters for the atomic accumulation path. |
| 664 | alloc = torch.empty if single_tile else torch.zeros |
| 665 | valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) |
| 666 | |
| 667 | _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL( |
| 668 | req_id_c, |
| 669 | block_table_c, |
| 670 | token_indices_c, |
| 671 | out, |
| 672 | valid_counts, |
| 673 | # No prefill workspace on the DCP decode path. |
| 674 | None, |
| 675 | None, |
| 676 | workspace_rank_stride=0, |
| 677 | max_num_blocks_per_req=max_num_blocks_per_req, |
| 678 | BLOCK_SIZE=BLOCK_SIZE, |
| 679 | BLOCK_STRIDE_ROWS=( |
| 680 | BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE |
| 681 | ), |
| 682 | BLOCK_N=BLOCK_N, |
| 683 | NUM_TOPK_TOKENS=NUM_TOPK_TOKENS, |
| 684 | HAS_PREFILL_WORKSPACE=False, |
| 685 | COUNT_VALID=count_valid, |
| 686 | COMPACT_TO_FRONT=compact_valid_to_front, |
| 687 | DCP_SIZE=dcp_size, |
| 688 | DCP_RANK=dcp_rank, |
| 689 | DCP_INTERLEAVE=cp_kv_cache_interleave_size, |
| 690 | ) |
| 691 | |
| 692 | if return_valid_counts: |
| 693 | assert valid_counts is not None |
| 694 | return out, valid_counts |
| 695 | return out |
| 696 | |
| 697 | |
| 698 | _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel() |
| 699 | |
| 700 | |
| 701 | @triton.jit(do_not_specialize=["num_tokens"]) |
| 702 | def sparse_mla_prepare_safe_lengths_kernel( |
| 703 | indices_ptr, |
| 704 | counts_ptr, |
| 705 | safe_lengths_ptr, |
| 706 | num_tokens, |
| 707 | index_stride0: tl.constexpr, |
| 708 | count_stride: tl.constexpr, |
| 709 | BLOCK: tl.constexpr, |
| 710 | ): |
| 711 | token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| 712 | count = tl.load(counts_ptr + token * count_stride, token < num_tokens, other=1) |
| 713 | tl.store(safe_lengths_ptr + token, tl.maximum(count, 1), token < num_tokens) |
| 714 | # TRTLLM's NoPE sparse kernel requires at least one valid KV-cache entry. |
| 715 | tl.store( |
| 716 | indices_ptr + token * index_stride0, 0, (token < num_tokens) & (count == 0) |
| 717 | ) |
| 718 | |
| 719 | |
| 720 | def prepare_sparse_mla_safe_lengths( |
| 721 | physical_indices: torch.Tensor, valid_counts: torch.Tensor |
| 722 | ) -> torch.Tensor: |
| 723 | """Install dummy slots for empty queries and return nonzero kernel lengths. |
| 724 | |
| 725 | Preserve the raw counts so empty outputs and LSE can be neutralized later. |
| 726 | """ |
| 727 | num_tokens = valid_counts.numel() |
| 728 | safe_lengths = torch.empty( |
| 729 | (num_tokens,), dtype=valid_counts.dtype, device=valid_counts.device |
| 730 | ) |
| 731 | if num_tokens: |
| 732 | sparse_mla_prepare_safe_lengths_kernel[(triton.cdiv(num_tokens, 256),)]( |
| 733 | physical_indices, |
| 734 | valid_counts, |
| 735 | safe_lengths, |
| 736 | num_tokens, |
| 737 | physical_indices.stride(0), |
| 738 | valid_counts.stride(0), |
| 739 | BLOCK=256, |
| 740 | ) |
| 741 | return safe_lengths |