flashinfer.cute_dsl.sparse.bsa_attn_sm120.bsa_attn_sm120_blk64_sage_fwd

flashinfer.cute_dsl.sparse.bsa_attn_sm120.bsa_attn_sm120_blk64_sage_fwd(q_int8: Tensor, k_int8: Tensor, v_fp8: Tensor, q_scale: Tensor, k_scale: Tensor, v_scale: Tensor, q2k_block_index: Tensor, block_sparse_num: int, block_sizes: Tensor | None = None, q2k_block_nums: Tensor | None = None, softmax_scale: float | None = None, *, out: Tensor, tma_descriptor_workspace: Tensor, uniform_block_count: bool = False, contiguous_block_indices: bool = False, backend: str = 'cake') → Tensor

Run prequantized SM120 Sage block-sparse attention.

All CUDA storage is caller-owned. out is contiguous BF16 BHSD and tma_descriptor_workspace is contiguous, 128-byte-aligned CUDA uint8 storage. Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS, and the operation supports MHA, head dimension 128, non-causal forward without LSE.

Parameters:

  • q_int8 (torch.Tensor) – Contiguous INT8 queries with shape [B, H, Sq, 128].

  • k_int8 (torch.Tensor) – Contiguous INT8 keys with shape [B, H, Sk, 128].

  • v_fp8 (torch.Tensor) – Contiguous FP8 E4M3 values in HDS layout with shape [B, H, 128, ceil(Sk / 64) * 64].

  • q_scale (torch.Tensor) – Contiguous FP32 query scales with shape [B, H, ceil(Sq / 128) * 4].

  • k_scale (torch.Tensor) – Contiguous FP32 key scales with shape [B, H, ceil(Sk / 64)].

  • v_scale (torch.Tensor) – Contiguous FP32 value scales with shape [B, H, 128].

  • q2k_block_index (torch.Tensor) – Contiguous INT32 selected KV block indices with shape [B, H, ceil(Sq / 64), capacity].

  • block_sparse_num (int) – Uniform number of selected KV blocks per query block when q2k_block_nums is omitted. Must be between zero and capacity.

  • block_sizes (torch.Tensor, optional) – Contiguous INT32 valid-token counts for the KV blocks. Supported shapes are [num_kv_blocks], [B, num_kv_blocks], and [B, H, num_kv_blocks].

  • q2k_block_nums (torch.Tensor, optional) – Contiguous INT32 selected-block counts with shape [B, H, ceil(Sq / 64)].

  • softmax_scale (float, optional) – Positive finite softmax scale. None selects 1 / sqrt(128).

  • out (torch.Tensor) – Caller-owned contiguous BF16 output with shape [B, H, Sq, 128].

  • tma_descriptor_workspace (torch.Tensor) – Caller-owned contiguous CUDA uint8 workspace, aligned to 128 bytes and large enough for the selected generated kernel.

  • uniform_block_count (bool) – Whether every query block uses block_sparse_num selected blocks.

  • contiguous_block_indices (bool) – Whether selected block indices are contiguous. This optimization requires uniform_block_count=True.

  • backend (str) – Backend name. The only supported value is "cake".

Returns:

The caller-owned out tensor after attention output is written.

Return type:

torch.Tensor