flashinfer.page.get_batch_indices_positions¶
- flashinfer.page.get_batch_indices_positions(append_indptr: Tensor, seq_lens: Tensor, nnz: int, batch_indices: Tensor | None = None, positions: Tensor | None = None) Tuple[Tensor, Tensor]¶
将 append indptr 和序列长度转换为批次索引和位置。
- 参数:
append_indptr (torch.Tensor) – 不规则张量的 indptr，形状为：
[batch_size + 1]。seq_lens (torch.Tensor) – KV-Cache 中每个请求的序列长度，形状为：
[batch_size]。nnz (int) – 不规则张量中的条目数。
batch_indices (Optional[torch.Tensor]) – 用于 batch_indices 的预分配输出张量。如果为
None，则自动分配。positions (Optional[torch.Tensor]) – 用于 positions 的预分配输出张量。如果为
None，则自动分配。
- 返回值:
batch_indices (torch.Tensor) – 不规则张量中每个条目的批次索引，形状为：
[nnz]。positions (torch.Tensor) – 不规则张量中每个条目的位置，形状为：
[nnz]。
示例
>>> import torch >>> import flashinfer >>> nnz_kv = 10 >>> append_indptr = torch.tensor([0, 1, 3, 6, 10], dtype=torch.int32, device="cuda:0") >>> seq_lens = torch.tensor([5, 5, 5, 5]) >>> batch_indices, positions = flashinfer.get_batch_indices_positions(append_indptr, seq_lens, nnz_kv) >>> batch_indices tensor([0, 1, 1, 2, 2, 2, 3, 3, 3, 3], device='cuda:0', dtype=torch.int32) >>> positions # the rightmost column index of each row tensor([4, 3, 4, 2, 3, 4, 1, 2, 3, 4], device='cuda:0', dtype=torch.int32)
注意
此函数类似于 cuSPARSE 库中的 CSR2COO 转换，区别在于我们将从不规则张量（不需要列索引数组）转换为 COO 格式。