Linear Attention

Chunked linear-attention and state-space layers with a recurrent matrix state S per (sequence, head), served by the open-source FROST kernels (CuTe DSL, nvidia-cutlass-dsl) through the Python frontend. There is no C++ API for these operations. Every variant has a forward, a full backward, and a pair of summary operations for context parallelism.

The ops live in cudnn.linear_attention.ops and are PyTorch custom operators with autograd formulas; the summary ops are <op>_summary and <op>_summary_bwd. The same node types are available on the Python graph API.

Support Matrix

Requires an SM100 / SM103 / SM107 GPU (Blackwell data center) and the CuTe DSL runtime the package depends on (nvidia-cutlass-dsl, a required dependency except on Windows ARM64; these kernels need 4.7.0 or newer and the engines decline with an error that names the installed version otherwise). On Windows ARM64, the CuTe DSL and TVM-FFI runtime dependencies are excluded from the base installation, so the base install does not enable these kernels there. One linear-attention node per graph.

T is the total token count of the packed batch, N the number of sequences. Dtypes are kernel-native and strict: callers convert, and every gradient comes back in its input’s dtype. cu_seqlens is int32 or int64.

Layout

All operands use the packed THD layout: tokens of every sequence concatenated along the first axis, with cu_seqlens ([N + 1]) giving the sequence boundaries. A dense batch of N equal-length sequences is cu_seqlens = [0, T, 2T, ...] over the flattened tokens; no copy is needed.

Operands may be strided views: the engines build their TMA descriptors from the buffer strides, and the only requirement is a contiguous innermost (head-dim) axis. A column slice of a wider per-token buffer (for example the q, k, v pieces of one fused projection) is accepted directly, forward and backward, without .contiguous(). Tensors are rank 3 ([T, heads, dim]); a hidden-fused [T, heads * dim] slice needs a .view to rank 3, which is free when heads * dim is contiguous.

GDP runs on two timelines: q and g have one row per real token, while k, v and beta have n rows per token. They come either expanded ([T * n, ...], update-major: the n updates of token t are rows t * n to t * n + n - 1) or token-major (k [T, n, HK, K], v [T, n, HV, V], beta [T, n, HO], with the caller’s token stride), so the per-token column slices of a joint [V | K | Q] projection ([T, n * H * D] viewed as [T, n, H, D]) are consumed in place; dk, dv and dbeta come back in the input’s shape. num_householder == 1 is exactly GDN.

Mamba-2 uses the SSD vocabulary: x plays v, B plays k, C plays q, dt is the gate. The output is y = C^T S + D x, gated by z sigmoid(z) when z is given.

Semantics worth knowing

• States and checkpoints. output_final_state=True returns the per-sequence state after the last token. checkpoint_every_n_tokens > 0 also returns the state series, one entry every n tokens strictly before each sequence end, as a non-differentiable dump that the backward reuses instead of recomputing. The rows are io dtype by default; every variant also takes checkpoint_dtype=torch.float32 for fp32 rows (the backward reuses io rows at the chunk cadence and regenerates the dense series from any other rows). state_indices turns initial_state into a pool [N_pool, ...] addressed per sequence, written back in place; forward only and exclusive with checkpoints.
• Summaries. <op>_summary returns the end state of each sequence from a given start state together with the span’s state-transition operator transition: for the delta-rule variants a [N, HO, K, K] fp32 matrix (the product of the span’s decayed erase steps, in the backward’s stored orientation), for Mamba-2 a [N, H] fp32 scalar per head (the log2 of the decay product), so that across spans S_final = S_init * transition + S_H. <op>_summary_bwd is the mirror for the state gradient. Together they implement context parallelism over sequence spans. Both are non-differentiable.
• batch_invariant. Each sequence’s results are bitwise independent of the batch composition (whole-sequence scheduling; split-K load balancing is disabled).
• Clamps and guards. Mamba-2’s dt clamp zeroes the gate gradient where active. GDN-2’s beta_guard is recomputed, not differentiated, in the backward.
• Engines. The FROST engines are the only engines for these nodes and are pure pass-through: check_support verifies the kernel-native dtypes and head dims listed above and declines everything else. A plan can be pinned by name through plan_name (for example gdn_frost, mamba2_frost, mamba2_summary_frost).

Example

import torch
from cudnn.linear_attention.ops import gated_delta_net, mamba2
T, H, K, V = 4096, 16, 128, 128
cu_seqlens = torch.tensor([0, 1000, 2500, T], dtype=torch.int32, device="cuda")
q = torch.randn(T, H, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)
k = torch.randn(T, H, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)
v = torch.randn(T, H, V, dtype=torch.bfloat16, device="cuda", requires_grad=True)
g = -torch.nn.functional.softplus(torch.randn(T, H, device="cuda"))
beta = torch.sigmoid(torch.randn(T, H, device="cuda"))
o, final_state = gated_delta_net(q, k, v, g, beta, cu_seqlens, output_final_state=True, use_qk_l2norm_in_kernel=True)
o.sum().backward()
P, N, G = 64, 128, 2
x = torch.randn(T, H, P, dtype=torch.bfloat16, device="cuda")
dt = torch.randn(T, H, device="cuda")
A = -torch.exp(torch.randn(H, device="cuda"))
B = torch.randn(T, G, N, dtype=torch.bfloat16, device="cuda")
C = torch.randn(T, G, N, dtype=torch.bfloat16, device="cuda")
y, final_state = mamba2(x, dt, A, B, C, cu_seqlens, D=torch.ones(H, device="cuda"), output_final_state=True)

Benchmarks against the reference implementations are in benchmark/linear_attention/.