flashinfer.sdsa.indexer_scores

flashinfer.sdsa.indexer_scores(query: Tensor, key_cache: Tensor, block_table: Tensor, context_lengths: Tensor, *, requests: Tensor, scores: Tensor, query_scales: Tensor | None = None, key_scales: Tensor | None = None) → None

Compute FP32 scores of query rows against their request’s paged keys.

This is fp8_paged_mqa_logits() for an indexer with one query head. With one head there is no head weight and no ReLU, so a score is the plain dot product of a query row and a key. Query rows come in groups of 1 to 8 consecutive rows, such as a decode step’s token and its draft tokens. 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). Queries, keys, pages and lengths can change between replays. On SM10x it runs a CuTe DSL kernel; on other GPUs (e.g. Hopper) it scores BF16 keys with a Triton kernel, and NVFP4 keys are not supported.

Parameters:

  • query (torch.Tensor) – BF16 [rows, 128], or NVFP4 uint8 [rows, 64] (E2M1 values, two to a byte).

  • key_cache (torch.Tensor) – BF16 [pages, 64, 128], or NVFP4 uint8 [pages, 64, 64]. Pages can be any multiple of 16 bytes apart, such as views of one cache that holds each page’s keys and then their scales.

  • block_table (torch.Tensor) – int32 [requests, blocks]: each request’s pages of 64 keys. Only pages below a group’s length are read.

  • context_lengths (torch.Tensor) – int32 [groups]: each group’s number of keys.

  • requests (torch.Tensor) – int32 [groups]: each group’s request, a row of block_table. A group of length 0 reads nothing, so its entry can hold any value.

  • scores (torch.Tensor) – float32 [rows, max_context], output: scores[q, k] is the dot product of query row q with key k, for k below its group’s length. Later entries are not written. max_context sizes the grid, and the block table must cover it.

  • query_scales (torch.Tensor, optional) – E4M3 [rows, 8] for NVFP4 queries, a scale per 16 values. None for BF16.

  • key_scales (torch.Tensor, optional) – E4M3 [pages, 64, 8] for NVFP4 keys, paged like the keys. Pages can be any multiple of 8 bytes apart. None for BF16.