QSA (Qwen Sparse Attention)

QSA is the sparse attention layer used by Qwen’s hybrid attention stack. It has two stages. A lightweight MQA indexer scores compressed micro-blocks, the caller keeps the top-scoring blocks, and core attention then runs on the original tokens inside those selected blocks.

MQA names the indexer’s structure, not a second attention mechanism: several query heads share one compressed key head, which keeps block scoring cheap. Core attention is the main attention over the selected tokens.

MATE implements the two compute stages for prepared inputs:

  • qsa_indexer_mqa_score, qsa_mqa_prefill (also exported as qsa_mqa) and qsa_mqa_decode produce indexer scores.

  • qsa_sparse_attention runs sparse core attention over the selected tokens.

Projection, RMSNorm, RoPE, cache updates, top-k selection and index expansion belong to the caller. Native execution supports MP31 (S5000) with TileLang-MUSA 0.1.12+musa.4; other architectures are rejected before compilation.

MQA score

The canonical score input is BF16 or FP16 query [B,Q,4,128] and compressed key [B,N,128]. Each compressed key represents four tokens. Scores sum the four head-wise ReLU dot products, scaled by 1/sqrt(128). Incomplete or future blocks are masked.

qsa_mqa_prefill takes packed queries, shared keys and per-row key intervals. qsa_mqa_decode takes paged compressed keys [pages, page_size, 1, dim], a page table and compressed context lengths. Its native path supports only heads, dim == (4, 128) with BF16/FP16 and page_size in 8, 16, 32, 64; any other page size raises under backend="tilelang" and requires backend="reference". scale is a multiplier; score_scale is a denominator. Do not pass both. The caller supplies valid active pages; unused entries may be -1. Native decode masks missing or out-of-range pages to -inf.

Native score and prefill kernels read Q and K through vector loads, so the same three layout rules apply to query and block_keys in score and to query and keys in prefill:

  • the feature stride (stride(-1)) must be 1;

  • data_ptr() must be 16-byte aligned;

  • every outer axis with more than one element must have a positive stride whose byte span (stride * element_size) is a multiple of 16.

For BF16 and FP16 this means outer strides must be positive multiples of 8 elements; axes of size 1 are exempt. A sliced view that satisfies the shape rules but not these stride rules fails in the native path with vector loads require unit feature stride and 16-byte alignment; pass a contiguous copy or adjust the slice offsets instead. qsa_mqa_decode is not subject to this check and accepts broadcast Q through its LSU load path.

Core Attention

qsa_sparse_attention consumes framework-prepared request-local indices shared by all Q heads. TileLang requires int32 indices, positions, lengths and page tables; -1 pads, and kernels mask invalid or future tokens. Token-only selections may include padding beyond the canonical 2051 slots, and generic paths use the actual runtime width, so the caller owns the selected-token budget; block-hint widths follow the fixed formula below. Supplied metadata is used without dtype conversion. Batch/query dimensions of query, indices, gate and output must form zero-copy row views; otherwise pass framework-prepared flattened tensors.

  • Canonical geometry is 24 Q heads, 2 KV heads and D256. Contiguous, direct-paged and SplitKV kernels also accept positive divisible Hq/Hkv with group size up to 16. Global Q/O use logical heads, and dummy compute-tile rows are zero-filled in shared memory. Structured prefill and warp specialization require canonical geometry; generic paged prefill also supports local 3/1 heads.

  • Paged K/V use [pages, page_size, Hkv, 256]; reference execution also accepts flattened varlen K/V. Native kernels mask missing or out-of-range pages. seq_lens_q counts query rows per request independently of cached causal positions, and padded rows produce zero output.

  • Optional int32 block_indices describes the four-token blocks behind token_indices. Structured paged prefill uses the hint only when the KV page size is divisible by four; otherwise it is ignored and the plain token_indices path runs. With a block hint the widths are tied: token_indices.shape[-1] == block_indices.shape[-1] * 4 + 3, and block_indices.shape[-1] must be in [1, 512]. The widest block-hint selection is therefore 2051 tokens, so wider selections such as 2055 must use the token-only path.

  • num_splits=None selects the prefill/decode schedule, and a positive integer sets the split count explicitly; 1 disables SplitKV. Warp specialization is opt-in. gate applies sigmoid gating, while output projection remains external. return_lse=True exports native softmax statistics; paged execution combines SplitKV statistics without repeating QK or PV.

  • Q/K/V, indices, metadata and optional out accept sliced views and storage offsets, with concrete strides passed as runtime int64. K/V require aligned unit feature strides, and the regular prefill path also requires contiguous K/V. MQA score and prefill enforce the aligned Q/K layout described in MQA score, while MQA decode accepts broadcast Q through its LSU load path. Native paths bind out directly and reject input/output overlap.

The caller prepares valid request IDs, monotonic cumulative lengths ending at the tensor row count, and sequence lengths within KV capacity. Native wrappers check shape, dtype, device and layout without reading device contents. Warmed-up native MQA and logical/physical Core APIs support MUSA Graph capture/replay; replay consumes current metadata, and explicit reference execution is outside this guarantee.

selected_kv_indices= accepts flattened BF16 Q/K/V with physical token indices instead of page translation or causal filtering. Its metadata-only validation supports public MUSA Graph capture/replay after warmup. Optional out and natural-log lse buffers bind directly, and providing lse writes it regardless of return_lse. Empty selections return zero output and -inf LSE. FP8 KV and FP8 compute remain unsupported. softmax_scale aliases scale, and supplying both is an error. Physical indices cannot be combined with logical request/page metadata or gating.

backend='tilelang' (the default) uses native kernels for supported MUSA layouts and head geometry and raises otherwise, with no automatic reference fallback; reference selects the PyTorch oracle, and auto is not supported.

Entrypoints

mate.qsa_indexer_mqa_score(q_idx: torch.Tensor, block_keys: torch.Tensor, *, block_table: torch.Tensor | None = None, page_size: int | None = None, cu_seqlens_q: torch.Tensor | None = None, token_to_req: torch.Tensor | None = None, num_requests: int | None = None, query_positions: torch.Tensor | None = None, seq_lens_q: torch.Tensor | None = None, block_size: int = 4, scale: float | None = None, score_scale: float | None = None, backend: str = 'tilelang') → torch.Tensor

Compute QSA MQA block scores.

Parameters:

  • q_idx – Projected indexer queries, [B, Q, 4, 128]. Varlen [total_q, 4, 128] requires explicit reference execution; use qsa_mqa_prefill for native packed rows.

  • block_keys – Logical compressed keys [B, N_B, 128]. With block_table this instead is a paged cache [pages, page_size, 1, 128] (the head axis may be omitted), supported only by the explicit reference. Native paged scoring uses qsa_mqa_decode.

  • query_positions – Local request positions. If omitted, positions are inferred from the padded/varlen query layout.

  • seq_lens_q – Optional request lengths. Padded rows beyond the length are masked.

  • score_scale – Optional packed-row denominator. scale and score_scale are mutually exclusive; the default is scale=1/sqrt(D).

  • backend – "reference" forces the FP32 PyTorch oracle; "tilelang" requires the supported MUSA ABI (sliced inputs are accepted); TileLang is the default; unsupported native combinations raise an error. This entrypoint does not select Top-K.

Returns:

  • torch.Tensor – FP32 scores with shape [B, Q, N_B] or [total_q, N_B]. A block whose last token is not visible has score -inf.

  • Raises – ValueError: Inconsistent layouts, metadata, scaling, or backend options. TypeError: Unsupported argument types or metadata dtypes.

mate.qsa_mqa(q: torch.Tensor, k: torch.Tensor, row_starts: torch.Tensor, row_ends: torch.Tensor, score_scale: float | None = None, *, scale: float | None = None, backend: str = 'tilelang') → torch.Tensor

Compute packed-prefill MQA logits.

Parameters:

  • q – Queries with shape [rows,heads,dim].

  • k – Shared keys with shape [keys,dim] or [keys,1,dim].

  • row_starts – Inclusive visible-key start per query row.

  • row_ends – Exclusive visible-key end per query row.

  • score_scale – Positive score denominator, exclusive with scale.

  • scale – Positive score multiplier; defaults to 1/sqrt(dim).

  • backend – tilelang (default) or explicit reference.

Returns:

FP32 [rows,keys] scores with -inf outside each visible interval. TileLang supports BF16/FP16 H4/D128 inputs on MUSA.

Return type:

torch.Tensor

Raises:

  • ValueError – Invalid shape, scaling, or unsupported explicit backend.

  • TypeError – Unsupported argument types or metadata dtypes.

row_starts/row_ends identify the visible compressed-key interval for each query row. score_scale uses the denominator convention; scale is an optional multiplier used by the MATE API.

mate.qsa_mqa_decode(q: torch.Tensor, k_cache: torch.Tensor, page_table: torch.Tensor, context_lens: torch.Tensor, max_model_len: int, score_scale: float | None = None, *, scale: float | None = None, backend: str = 'tilelang') → torch.Tensor

Compute paged-decode MQA logits with packed-row metadata.

Parameters:

  • q – Queries, [B,heads,dim] or [B,1,heads,dim].

  • k_cache – Compressed paged keys, [pages,page_size,1,dim]; the shared head axis may be omitted.

  • page_table – Logical-to-physical page IDs with shape [B,max_pages].

  • context_lens – Visible compressed-key counts, one per request.

  • max_model_len – Output score width in compressed-key units.

  • score_scale – Positive score denominator, exclusive with scale.

  • scale – Positive score multiplier; defaults to 1/sqrt(dim).

  • backend – tilelang (default) or explicit reference.

Returns:

FP32 [B,max_model_len] scores; invisible keys have score -inf. TileLang supports BF16/FP16 H4/D128 inputs on MUSA.

Return type:

torch.Tensor

Raises:

  • ValueError – Invalid shape, scaling, or backend options.

  • TypeError – Unsupported argument types or metadata dtypes.

mate.qsa_mqa_prefill(q: torch.Tensor, k: torch.Tensor, row_starts: torch.Tensor, row_ends: torch.Tensor, score_scale: float | None = None, *, scale: float | None = None, backend: str = 'tilelang') → torch.Tensor

Compute packed-prefill MQA logits.

Parameters:

  • q – Queries with shape [rows,heads,dim].

  • k – Shared keys with shape [keys,dim] or [keys,1,dim].

  • row_starts – Inclusive visible-key start per query row.

  • row_ends – Exclusive visible-key end per query row.

  • score_scale – Positive score denominator, exclusive with scale.

  • scale – Positive score multiplier; defaults to 1/sqrt(dim).

  • backend – tilelang (default) or explicit reference.

Returns:

FP32 [rows,keys] scores with -inf outside each visible interval. TileLang supports BF16/FP16 H4/D128 inputs on MUSA.

Return type:

torch.Tensor

Raises:

  • ValueError – Invalid shape, scaling, or unsupported explicit backend.

  • TypeError – Unsupported argument types or metadata dtypes.

row_starts/row_ends identify the visible compressed-key interval for each query row. score_scale uses the denominator convention; scale is an optional multiplier used by the MATE API.

mate.qsa_sparse_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, token_indices: torch.Tensor | None = None, *, block_indices: torch.Tensor | None = None, cu_seqlens_q: torch.Tensor | None = None, token_to_req: torch.Tensor | None = None, num_requests: int | None = None, cu_seqlens_k: torch.Tensor | None = None, seq_lens_q: torch.Tensor | None = None, seq_lens_k: torch.Tensor | None = None, query_positions: torch.Tensor | None = None, block_table: torch.Tensor | None = None, page_size: int | None = None, scale: float | None = None, gate: torch.Tensor | None = None, num_splits: int | None = None, return_lse: bool = False, backend: str = 'tilelang', warp_specialized: bool = False, assume_compact_blocks: bool = False, out: torch.Tensor | None = None, selected_kv_indices: torch.Tensor | None = None, softmax_scale: float | None = None, k_scale: float = 1.0, v_scale: float = 1.0, compute_dtype: torch.dtype = torch.bfloat16, lse: torch.Tensor | None = None)

Run QSA core attention over a per-query compact token set.

Parameters:

  • query – [B,Q,Hq,D] or flattened [total_q,Hq,D] queries.

  • key – Padded or paged KV storage; flattened-varlen requires reference.

  • value – Values with the same layout, dtype, and device as key.

  • token_indices – Framework-prepared int32 request-local token IDs; -1 pads.

  • selected_kv_indices – Alternative int32 [rows, selected_tokens] physical IDs into flattened BF16 K/V. Mutually exclusive with token_indices. The caller handles causal filtering and page translation. Duplicates are retained; negative and out-of-capacity IDs are masked.

  • softmax_scale – Alias for scale; do not supply both.

  • k_scale – Must be 1 for BF16 K/V.

  • v_scale – Must be 1 for BF16 K/V.

  • compute_dtype – BF16 for physical indices; FP8 is not enabled.

  • lse – Optional FP32 [rows, Hq] natural-log LSE output for physical indices. Written even when return_lse=False; empty rows are -inf. The supplied buffer is bound directly to the merge kernel.

  • block_indices – Optional four-token blocks that produced token indices.

  • query_positions – Request-local causal positions; inferred if omitted.

  • seq_lens_k – Actual KV lengths, bounded by the supplied cache capacity.

  • block_table – Optional logical-to-physical page mapping.

  • num_splits – None chooses a schedule; 1 disables SplitKV.

  • gate – Optional per-query/head gate, applied through sigmoid.

  • return_lse – Also return FP32 natural-log log-sum-exp. TileLang exports native softmax statistics; paged inputs use a SplitKV merge.

  • backend – tilelang (default) or explicit reference.

  • out – Optional writable, nonoverlapping positive-stride output view matching the query shape, dtype, and device. Batch/query axes must form a zero-copy row view. Input/output overlap is rejected.

Returns:

Query-shaped output, or (output, lse) when requested. LSE has shape query.shape[:-1]. The output is exactly out when supplied.

Return type:

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

Raises:

  • ValueError – Invalid shape, output layout, scaling, incompatible hints, or unsupported explicit backend.

  • TypeError – Unsupported argument types or metadata dtypes.

TileLang consumes int32 indices, positions, lengths and page tables without dtype conversion. Padded batch/query axes of query, indices, gate and out must collapse to zero-copy row views; otherwise pass framework-prepared flattened tensors. Indices use -1 padding. The same list is shared by all Q heads; keys/values are grouped by Hq // Hkv (12 for the Qwen configuration). Contiguous padded KV, flattened varlen KV, and paged KV are accepted by the reference path. Native nonpaged KV requires padded rank-4 queries and num_splits=1; native flattened queries and SplitKV require paged KV.

out may be a nonoverlapping strided view with the query shape, dtype and device. The returned output aliases that exact view; unused storage around a slice is preserved. Strides are runtime kernel parameters.

num_splits=None uses the measured prefill/decode heuristic. A positive integer selects that split count explicitly; 1 disables SplitKV.

block_indices is an optional optimization hint for the canonical QSA layout. It must be the 4-token block selection that produced token_indices. When both are supplied, official-width paged prefill transfers one encoded descriptor per block and derives the current tail from query metadata. Omitting it preserves arbitrary selected-token semantics and the original dispatch path. assume_compact_blocks is an opt-in producer contract: set it only when the indexer guarantees valid block ids occupy a prefix of each row, followed by -1 padding, and every selected block is causally visible. It selects compact dynamic prefill without scanning or reading back a proof. TileLang dispatch accepts BF16/FP16 Q/K/V with D256, positive divisible Hq/Hkv and GQA groups up to 16. Q, K, V, indices, metadata and out may be sliced views with storage offsets. Nonoverlapping output is required; the returned tensor is exactly out when supplied. Head geometry, dtype and algorithm tiles specialize the kernels; logical lengths and concrete strides are runtime arguments, with int64 stride arithmetic. Supported head shapes use TileLang by default. Unsupported combinations raise an error; the PyTorch reference requires backend="reference".

Vector K/V gathers require unit feature stride and 16-byte alignment; unsupported layouts raise an error. Aligned outer K/V slices retain their independent runtime strides. The regular prefill path requires contiguous K/V; its WS variant supports aligned outer slices. Q/O head padding stays inside the shared-memory compute tile. Native paths bind out directly; dense stores retain TME and sliced stores use LSU. Gate/output overlap and input/output overlap raise errors.

The structured TME prefill and WS specializations retain canonical 24/2/256 geometry; unsupported optional hints use the ordinary path. SplitKV accumulates and merges in FP32; partial output storage can use the input dtype. Its fused merge applies the gate in that merge; direct paths use an in-place TileLang post-op. Native logical and physical paths export LSE without repeating QK or PV, and support public MUSA Graph capture/replay after warmup. warp_specialized is an opt-in MUSA prefill experiment that splits QK and PV consumer groups and uses a two-stage K/V/score ring; it is not enabled by default. For the canonical short-context structured prefill window, metadata may additionally order complete physical descriptors to improve the following sparse gather; that guard is disabled outside the measured window.