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 | None = None, tma_descriptor_workspace: Tensor | None = None, uniform_block_count: bool = False, contiguous_block_indices: bool = False, backend: str = 'cake') → Tensor

Run prequantized SM120 Sage block-sparse attention.

Supports two backends selected via backend:

"cake" (default)

tvm-ffi generated kernel. All CUDA storage is caller-owned: out is a required keyword argument. Tensor-map descriptors are passed by value as grid-constant kernel parameters; tma_descriptor_workspace is accepted and ignored. uniform_block_count and contiguous_block_indices tuning flags are only honoured by this backend.

"cute_dsl"

CuTe-DSL port of upstream Block-Sparse-Attention. Uses BHSD tensor layout throughout and allocates out internally when not provided. Supports SM120 only (exact compute capability 12.0), MHA only, head dimension 128, non-causal forward without LSE. uniform_block_count and contiguous_block_indices are silently ignored by this backend.

Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS.

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, optional) – Pre-allocated contiguous BF16 output with shape [B, H, Sq, 128]. Required for backend="cake". For backend="cute_dsl", allocated internally when not provided.

  • tma_descriptor_workspace (torch.Tensor, optional) – Accepted for compatibility and ignored by both backends. Tensor-map descriptors are passed by value in the kernel launch, so no caller storage is read or written.

  • uniform_block_count (bool) – Whether every query block uses block_sparse_num selected blocks. Only honoured by backend="cake", where it is a caller-provided guarantee: when true, q2k_block_nums is ignored and its values are not checked. Silently ignored by backend="cute_dsl".

  • contiguous_block_indices (bool) – Whether selected block indices are contiguous. Requires uniform_block_count=True. Only honoured by backend="cake", which trusts the caller’s guarantee without checking the indices; non-contiguous top-k indices must use False, even when sorted. Incorrectly setting this flag to True under backend="cake" can produce incorrect output without raising an error. Silently ignored by backend="cute_dsl".

  • backend (str) – Backend name. "cake" (default) or "cute_dsl".

Returns:

The output tensor after attention output is written. For backend="cake" this is always the caller-owned out. For backend="cute_dsl", out is returned directly when provided; otherwise a newly allocated tensor is returned.

Return type:

torch.Tensor

Notes

Non-contiguous indices support both per-row q2k_block_nums and the scalar block_sparse_num count, with every supported block_sizes layout and with aligned or partial final KV blocks. Omitting q2k_block_nums uses the scalar count even when uniform_block_count=False. Under backend="cake", setting uniform_block_count=True uses the scalar count and ignores any supplied per-row counts; backend="cute_dsl" accepts the flag but silently ignores it. A row selecting zero blocks produces an exactly zero output.

Initialize the JIT module before CUDA Graph capture. Keep captured input and output tensors alive while their graph is in use. Each captured launch retains its own tensor-map descriptors, so launches with different tensor bindings can be interleaved without shared descriptor storage.