flashinfer.qsa.qsa_indexed_forward

flashinfer.qsa.qsa_indexed_forward(q: Tensor, k: Tensor, v: Tensor, indices: Tensor, scale: float = 0.0625) → tuple[Tensor, Tensor, Tensor]

Compute packed query sparse attention from selected token indices.

Requires cuDNN Frontend 1.32+, SM100/SM103, BF16 12Q/1KV and dimension 256. All tensors must be contiguous on one CUDA device.

Parameters:

  • q (torch.Tensor) – BF16 queries, shape [T, 12, 256].

  • k (torch.Tensor) – BF16 keys and values, shape [T, 1, 256].

  • v (torch.Tensor) – BF16 keys and values, shape [T, 1, 256].

  • indices (torch.Tensor) – int32 absolute packed token IDs, shape [T, W]. W is a positive multiple of four, at most 2052. Each four-slot group is a distinct complete block, four -1 padding slots, or a partial causal tail followed by -1. Groups may appear in any order but must not overlap. Tokens must belong to the query’s sequence and cannot exceed its token position. Block origins restart at each sequence boundary. These device-value rules are a caller contract, not host-side checks.

  • scale (float) – Finite attention score scale, default 1 / sqrt(256).

Returns:

out, out_fp32, lse – BF16 output, FP32 output before casting (both [T, 12, 256]), and FP32 natural-log LSE ([12, T]). Empty rows produce zero outputs and negative-infinity LSE.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Notes

Consumes existing indices; no scoring or top-k selection is performed. Outputs and scratch are allocated by cuDNN on the input device’s current CUDA stream. Warm up before capture; replay rereads the indices. This explicit forward call does not register autograd.

Results are not bitwise repeatable. cuDNN’s parallel index compaction can change the accumulation order, including BF16 output values. Sorting caller indices does not remove this internal ordering variation.