Sparse attention

Lightning indexer

lightning_indexer selects a shared pool of key positions using multi-head weighted ReLU scores. It returns indices, not attention outputs or differentiable scores:

from attn_gym.sparse import lightning_indexer

indices = lightning_indexer(q, k, weights, topk=128, causal=True)
# q: [B, T, H, D], k: [B, T, D], weights: [B, T, H]
# indices: [B, T, 128], int32

# DeepSeek compressed sparse attention: one candidate per 4 consecutive tokens.
indices = lightning_indexer(q, k_compressed, weights, topk=128, causal=True, compress_ratio=4)
# k_compressed: [B, T // 4, D]; query t may select candidates s < (t + 1) // 4

The score for a query/candidate pair is sum_h(weights[h] * relu(dot(q[h], k))) / sqrt(H * D). Weights may be negative. Without packed offsets, k holds S = T // compress_ratio candidates, each summarizing compress_ratio consecutive tokens (a trailing partial window forms no candidate). For packed documents, pass paired cu_seqlens and cu_seqlens_k: compression and causal positions restart per document, and output indices are document-local. With causal selection, query t considers only candidates s < (t + 1) // compress_ratio, those whose tokens all lie at positions <= t (candidates 0..t with the default ratio of 1); rows with fewer than topk candidates contain -1 padding. compress_ratio != 1 requires causal=True. Mask padding before gathering: a raw PyTorch index of -1 selects the last position rather than an invalid position. topk=0 returns an empty last dimension. Output order and tie-breaking are unspecified, including between repeated calls.

Implementations and device dispatch

• impl="reference" evaluates PyTorch scoring and Top-K on CPU or CUDA. FP16/BF16/FP32 inputs accumulate in FP32; FP64 inputs retain FP64. It materializes the score intermediates and is intended for correctness checks and small inputs.
• impl="fused" (the default) selects CuTe on SM100/SM103, or Triton on other NVIDIA GPUs with compute capability 9.0 or newer, including Hopper. Both optimized implementations keep their selection state on chip. CuTe separates score generation from radix Top-K and reuses a per-call FP32 score slab capped at 32 MiB and 1024 query rows, independent of batch size. Large inputs are processed in slabs rather than an unbounded quadratic score allocation. This scratch is additional to the returned indices; Triton needs no global score scratch.
• kernel_options={"backend": "cute"} or {"backend": "triton"} overrides that choice. Omit options for automatic selection. Options are rejected for impl="reference". Unsupported shapes, missing dependencies, and launch errors propagate; there is no retry with another backend.

Restriction	CuTe	Triton
GPU	SM100 or SM103	SM90 or newer
Input dtype	FP16 or BF16, shared by all inputs	FP16 or BF16, shared by all inputs
Heads H	Positive and even	1..256
Head dimension D	Positive, divisible by 16	8..256, divisible by 8
Sequence length T	1..2**20	1..2**20
topk	Non-negative; missing candidates pad with -1	Non-negative; missing candidates pad with -1; per-tile cost grows with topk
Layout	Q/K last stride 1, bases and non-singleton outer strides 16-byte aligned; weights may be strided	Q/K last stride 1, bases and outer strides 16-byte aligned; weights may be strided

For nonzero topk, CuTe limits candidate capacity S to 2**22 so at least one query pair's FP32 scores fit its 32 MiB workspace. This includes inactive packed capacity.

CuTe accepts independently permuted or padded outer dimensions and broadcast inputs without materializing contiguous copies. Weights need only element alignment, not TMA alignment.

CuTe additionally requires the optional linear dependencies. It supports SM100 and SM103 (including GB300), not every Blackwell variant; other supported devices use Triton by default.

Compilation and training scope

Both fused backends support the public API under torch.compile(fullgraph=True), including changing batch/sequence sizes with dynamic=True. Backend selection and TMA setup happen inside a private operator on real CUDA tensors, rather than during tracing. Static-shape CUDA Graph capture/replay is supported after warming up compilation:

import torch

compiled_indexer = torch.compile(lightning_indexer, fullgraph=True, dynamic=True)
indices = compiled_indexer(q, k, weights, 128, causal=True)

Selection is nondifferentiable. Inputs may require gradients, but the integer result has no gradient function. Gather attention can train its own Q/K/V computation with these indices held fixed. It does not propagate gradients through the selection step into the indexer's queries, keys, or scoring weights. This API supplies no surrogate gradient or indexer-training loss.

Selection with NaN/Inf scores is unspecified and may differ across backends.

attn_gym.sparse.indexer.lightning_indexer(q, k, weights, topk, *, causal=False, compress_ratio=1, cu_seqlens=None, cu_seqlens_k=None, impl=Impl.FUSED, kernel_options=None)

Return the Top-K candidate indices for every (batch, query) row.

Computes a multi-head weighted ReLU score for each query-candidate pair and selects the topk highest-scoring candidates per query::

dots[b, t, h, s] = q[b, t, h, :] · k[b, s, :]
score[b, t, s]   = sum_h(w[b, t, h] * relu(dots[b, t, h, s]))
                   / sqrt(H * D)
output[b, t, :]  = topk(score[b, t, :]).indices

Parameters:

Name	Type	Description	Default
q	Tensor	Query tensor, [B, T, H, D].	required
k	Tensor	Key candidate pool shared across heads, [B, S, D]. Without packed offsets, S = T // compress_ratio. With offsets, candidates from each document are concatenated in document order; S may include inactive capacity.	required
weights	Tensor	Per-head weights, [B, T, H]. May be negative.	required
topk	int	Non-negative output width. Rows with fewer valid candidates are padded with -1, including when S is zero or topk exceeds S.	required
causal	bool	If True, query t can only select candidates s < (t + 1) // compress_ratio, i.e. those whose covered tokens all lie at positions <= t (candidates 0..t when compress_ratio=1).	False
compress_ratio	int	Number of consecutive tokens summarized by each candidate, as in DeepSeek compressed sparse attention. A trailing partial window forms no candidate. In packed mode grouping and the causal query position restart at each document boundary: document d has (cu_seqlens[d + 1] - cu_seqlens[d]) // compress_ratio candidates. Values other than 1 require causal=True.	1
cu_seqlens	Tensor | None	Packed offsets shaped [N + 1] for batch-one inputs, as contiguous int32 on q.device; they start at zero, never decrease, may repeat for empty sequences, and may end before T. Query rows beyond the endpoint return only -1. This API processes whole documents, not query chunks with cached history. Must be supplied together with cu_seqlens_k. Offset values are caller invariants, not host-validated.	None
cu_seqlens_k	Tensor | None	Packed candidate offsets shaped [N + 1], contiguous int32 on q.device. They start at zero, never decrease, and end at or before S. Each document's span must equal its query length divided by compress_ratio, rounded down. Returned indices are zero-based within that document's candidate pool, suitable for gather_attn with these offsets. Without packed offsets, indices are positions in the batch element's pool.	None
impl	Impl | str	Impl.REFERENCE (or "reference") uses eager PyTorch on CPU or CUDA; Impl.FUSED (default, or "fused") uses optimized CUDA kernels.	FUSED
kernel_options	dict[str, str] | None	Fused backend override: {"backend": "cute"} or {"backend": "triton"}. Omit options to select CuTe on SM100/SM103 and Triton on other Hopper-or-newer NVIDIA GPUs. There is no fallback when the selected backend rejects a shape or layout.	None

Returns:

Type	Description
Tensor	Contiguous [B, T, topk] INT32 indices. Order and tie-breaking are unspecified, and
Tensor	CuTe's may differ across repeated calls. Under
Tensor	torch.use_deterministic_algorithms(True) each fused backend is repeatable for
Tensor	identical inputs (results may differ between backends), and CuTe returns valid
Tensor	indices ascending, resolving equal ordered-FP32 scores toward lower indices.
Tensor	Rows with fewer than topk valid candidates contain -1 padding after the valid
Tensor	indices; topk=0 returns an empty last dimension. Packed selection excludes other
Tensor	documents before top-k, not by filtering its result.

Fused backends support torch.compile(fullgraph=True) and CUDA Graph replay. Both require FP16/BF16 inputs and T <= 220. CuTe requires SM100/SM103, even H, D divisible by 16, and Q/K with unit last strides and 16-byte-aligned bases and non-singleton outer strides. Other Q/K strides may vary independently; weights may have arbitrary strides and need only element alignment. Triton requires SM90 or newer, H <= 256, D <= 256 divisible by 8, and Q/K with unit last strides and 16-byte-aligned bases and outer strides; weights may be strided. Its register-resident selection makes per-tile cost grow with topk. CuTe reuses a per-call FP32 score workspace capped at 32 MiB and 1024 query rows. For nonzero topk, candidate capacity S must be at most 222 so one query pair fits that budget. Large inputs use slabs rather than an unbounded quadratic score allocation. This workspace is additional to the returned indices and is not shared across calls. Packed CuTe calls also prepare an INT32 [T, 2] bounds buffer, reused by scoring and top-k to restrict work to each document. Triton computes those bounds in its kernel prologue without a separate buffer.

Selection with NaN/Inf scores is unspecified and may differ across backends. Indices are nondifferentiable even when inputs require gradients. Training attention over the selected positions does not propagate gradients through selection into q, k, or weights.