linear-attention (Helion)

Seven chunked linear-attention variants written in Helion and packaged for the kernels library: linear attention, Simple GLA, Retention, GLA, Delta Rule, Gated Delta Rule, and KDA. The kernel ships pre-tuned configs (via helion.aot_kernel), so the first call compiles a shipped configuration instead of running an autotuning search. They were tuned over the six shapes in flash-linear-attention's benchmarks/ops/registry.py; other shapes may run slower.

How to use

# make sure `kernels` is installed: `pip install -U kernels`
import torch
from kernels import get_kernel

kernel = get_kernel("HelionDSL/linear-attention", version=0)

b, t, h, d = 2, 1024, 8, 64
q = torch.randn(b, t, h, d, device="cuda", dtype=torch.bfloat16)
k = torch.nn.functional.normalize(
    torch.randn_like(q, dtype=torch.float32), dim=-1
).to(q)
v = torch.randn_like(q)
g = -torch.rand_like(q, dtype=torch.float32) * 0.1
beta = torch.rand(b, t, h, device="cuda", dtype=torch.bfloat16)

out, _ = kernel.chunk_kda(q, k, v, g, beta)

Inputs use the FLA-style (batch, time, heads, dim) layout. Functions return (output, final_state), with final_state=None unless requested.

Use version=0, or revision="<commit>" with a commit from the v0 branch.

HelionDSL is a trusted kernel publisher on the Hub, so trust_remote_code=True is not required.

Available functions

• chunk_linear_attn(q, k, v, ...) — causal linear attention.
• chunk_simple_gla(q, k, v, g, ...) — scalar-decay GLA.
• chunk_retention(q, k, v, ...) — RetNet fixed per-head decay.
• chunk_gla(q, k, v, g, ...) — per-channel GLA.
• chunk_delta_rule(q, k, v, beta, ...) — Delta Rule.
• chunk_gated_delta_rule(q, k, v, g, beta, ...) — Gated Delta Rule.
• chunk_kda(q, k, v, g, beta, ...) — Kimi Delta Attention, including variable-length forward.

Available layers

• chunk_kimi_delta_attention — KDA mapping for kernelize().
• chunk_gated_delta_rule — Gated Delta Rule mapping for kernelize().

Notes

• Backend: Helion → Triton. Correctness is verified against PyTorch recurrent references.
• Pre-tuned for H100 (sm90) and B200 (sm100). Other GPUs run untuned default configs; retune with HELION_AOT_MODE=collect.
• Gradients: all seven functions support backward. Untuned Delta Rule calls with head dimension 128 or larger are rejected; KDA variable-length and fused-input paths are forward only.
• chunk_linear_attn requires normalize=False. Simple GLA and Retention do not support initial_state.