flashinfer.qsa¶
Packed BF16 query sparse attention (QSA), backed by cuDNN Frontend 1.32 or newer. These APIs compute attention over caller-selected four-token blocks. They do not generate selections or register an autograd function.
API |
Selection |
Query / KV heads |
GPU |
|---|---|---|---|
|
uint8 block bitmap |
12 / 1 or 24 / 2 |
SM100, SM103, SM107 |
|
int32 token indices |
12 / 1 |
SM100, SM103 |
Both require head dimension 256. Outputs are BF16 [T, Hq, 256], FP32
output before casting with the same shape, and FP32 natural-log LSE
[Hq, T]. Empty selections produce zero output and negative-infinity LSE.
All tensors must be contiguous on one CUDA device. cuDNN validates tensor
metadata and compiler support; packed sequence boundaries and selection values
are caller contracts and are not copied to the host for validation.
The indexed API consumes distinct four-token groups, with -1 padding and
an optional partial causal tail. Indices are absolute packed token IDs from
the query’s own sequence. Width must be divisible by four and at most 2052.
It avoids the bitmap API’s int32 [T, NB] compact-index scratch, which alone
uses 16 GiB at 128K tokens and 32K blocks.
Indexed forward is not bitwise repeatable: cuDNN’s parallel index compaction can change the attention accumulation order, including BF16 output values. Sorting caller indices does not remove this internal ordering variation. This API has no deterministic mode; account for this in reproducibility-sensitive workloads, including reinforcement learning.
cuDNN allocates outputs and scratch and uses the input device’s current CUDA
stream. Warm up each required configuration before CUDA graph capture.
Replay reads current selection tensors; no implicit selection cache is kept.
The cuDNN package requires CuTe DSL 4.8+; SM107 also needs a compiler build
with Rubin support. These APIs are separate from the NVFP4 paged-cache
flashinfer.sparse.qsa_sparse_paged_attention interface.
import torch
import flashinfer
q = torch.randn(8, 12, 256, device="cuda", dtype=torch.bfloat16)
k = torch.randn(8, 1, 256, device="cuda", dtype=torch.bfloat16)
v = torch.randn_like(k)
token = torch.arange(8, device="cuda", dtype=torch.int32)
indices = torch.where(token[None, :] <= token[:, None], token[None, :], -1)
out, out_fp32, lse = flashinfer.qsa_indexed_forward(q, k, v, indices)
|
Compute packed query sparse attention from a block-selection bitmap. |
|
Compute packed query sparse attention from selected token indices. |