flashinfer.qsa.qsa_forward

flashinfer.qsa.qsa_forward(q: Tensor, k: Tensor, v: Tensor, sel_blocks: Tensor, lo: Tensor, hi: Tensor, blk_base: Tensor, tok_base: Tensor, block_first: Tensor, block_last: Tensor, scale: float = 0.0625, block_size: int = 4) → tuple[Tensor, Tensor, Tensor]

Compute packed query sparse attention from a block-selection bitmap.

Requires cuDNN Frontend 1.32+, BF16, head dimension 256, 12Q/1KV or 24Q/2KV, and SM100, SM103 or SM107. cuDNN checks geometry, storage and compiler support. All tensors must be contiguous on one CUDA device.

Parameters:

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

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

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

  • sel_blocks (torch.Tensor) – uint8 bitmap, shape [T, NB]. Nonzero selects a global block.

  • lo (torch.Tensor) – int32 inclusive key-token bounds, each shape [T].

  • hi (torch.Tensor) – int32 inclusive key-token bounds, each shape [T].

  • blk_base (torch.Tensor) – int32 first global block/token of each query’s sequence, shape [T].

  • tok_base (torch.Tensor) – int32 first global block/token of each query’s sequence, shape [T].

  • block_first (torch.Tensor) – int32 inclusive global token span of each block, shape [NB].

  • block_last (torch.Tensor) – int32 inclusive global token span of each block, shape [NB].

  • scale (float) – Finite attention score scale. Defaults to 1 / sqrt(256).

  • block_size (int) – Must be 4. Block numbering restarts at each packed sequence boundary.

Returns:

  • out (torch.Tensor) – BF16 attention output, shape [T, Hq, 256].

  • out_fp32 (torch.Tensor) – FP32 output before the BF16 cast, with the same shape.

  • lse (torch.Tensor) – FP32 natural-log softmax normalizer, shape [Hq, T].

Notes

Metadata must describe one consistent packed partition. Values remain on device and are a caller contract. No causal tail is added implicitly. Empty rows return zero outputs and negative-infinity LSE.

This explicit forward call does not register autograd. cuDNN allocates outputs and scratch on the input device and uses its current CUDA stream. Warm up before CUDA graph capture. Bitmap scratch includes int32 [T, NB] indices; prefer qsa_indexed_forward() for long sequences when token indices are already available.