flashinfer.token_sparse.token_sparse_kv_indices

flashinfer.token_sparse.token_sparse_kv_indices(positions: Tensor, requests: Tensor, block_table: Tensor, *, block_tables: Tensor, seq_lens: Tensor, requests_per_token: int) → None

Build the K/V slot lists of token-sparse attention from key positions.

Each token becomes requests_per_token consecutive attention requests, for example one per group of its query heads, each with the token’s K and V slot lists and count. These are the block_tables and seq_lens that trtllm_batch_decode_with_kv_cache(..., sparse_top_k=K, uses_shared_paged_kv_idx=False) takes over a cache of one-token pages. It runs as one launch without host synchronization, so it can be captured in a CUDA graph (run it once before capture to compile it).

Parameters:

  • positions (torch.Tensor) – int32 [tokens, K], K a multiple of 512: each token’s key positions, valid ones first, then -1.

  • requests (torch.Tensor) – int32 [tokens]: each token’s request, a row of block_table.

  • block_table (torch.Tensor) – int32 [requests, blocks]: each request’s pages of 64 keys. Key k is in slot block_table[request, k // 64] * 64 + k % 64.

  • block_tables (torch.Tensor) – int32 [rows, 2, K], output: request r gets the slots in both rows (K, then V), -1 past its count. Rows past tokens * requests_per_token are left as they are.

  • seq_lens (torch.Tensor) – int32 [rows], output: each request’s count of valid keys.

  • requests_per_token (int) – Attention requests per token.