Fast Attention Residuals

Fast Attention Residuals (Fast-AttnRes) makes Attention Residuals a single PyTorch operation: pass ordered full-width residual sources and one learned query, get one full-width residual back. The same attnres(values, query) call handles standard and sliced low-rank AttnRes in Full and Block schedules, with packed tensors or ordered source lists.

Start with the standard quickstart, choose a Full or Block schedule, and try sliced LR-AttnRes for a smaller routing query.

Why use Fast-AttnRes

• One PyTorch call: full-width output and ordinary first-order autograd, including gradients through routing keys and shared residual sources.
• One schedule primitive: Full and Block use the same operator; the caller controls source order, block sums, and learned queries.
• BF16 training: CUDA BF16 values and queries, FP32 internal reductions, torch.compile, and CUDA Graph replay on H100 and B200.

Training performance

On the 24-layer Full workload, Fast-AttnRes reduces step latency by 5.95% on H100 SXM and 22.74% on B200 against native FLA Triton checkpoint 1 (median of three paired seed estimates).

The figures measure complete BF16 CUDA Graph training steps on H100 SXM and B200, including forward, backward, and the optimizer update. Full and Block use the same per-read Fast-AttnRes backend.

The 24-layer headline uses three seeds with 120 paired rounds each. The 8-layer comparisons use one seed with 40 paired rounds per configuration. See the benchmark protocol and reproducible results for workloads, source versions, confidence intervals, and reproduction commands.

Compiled BF16 training steps

Comparators are native FLA Triton checkpoint 1, Liger 0.8.2, and Catswe phase 1. Unsupported and failed arms remain labelled. Quarter-rank comparisons against standard FLA compare different routing equations. The D2048 results retain disclosed normalization-rounding deviations. On H100, standard D2048 is 0.28% slower than FLA, within the declared 1% parity band.

Quarter-rank routing

These figures compare our R=D/4 kernel with our R=D kernel on each workload. Values and outputs retain width D; only the routing rank changes. Quarter rank reduces step latency in all ten measured device/workload pairs.

Install

Install the release with the pinned CUDA runtime:

python -m pip install --index-url https://download.pytorch.org/whl/cu126 torch==2.13.0
python -m pip install "fast-attnres[cuda]==2.1.0"

For development, clone the repository and run python -m pip install -e ".[cuda]".

Version 2 requires CUDA BF16 tensors; the CPU/FP32 execution and exported reference from version 1 are no longer part of the public API.

The minimum PyTorch version is 2.10. The pinned H100 runtime is Python 3.11, PyTorch 2.13.0 with CUDA 12.6, and Triton 3.7.1. On B200, use cu130 instead of cu126; multi-query qualification here covers H100.

Quickstart: standard AttnRes

import torch
from attnres import attnres

values = torch.randn(8, 2, 1024, device="cuda", dtype=torch.bfloat16, requires_grad=True)
source_list = tuple(values.unbind(0))
query = torch.randn(1024, device="cuda", dtype=torch.bfloat16, requires_grad=True)
read = torch.compile(attnres, fullgraph=True, dynamic=False)
output = read(source_list, query)
output.square().mean().backward()

The public signature is attnres(values, query, *, eps=2**-23, scale=1.0). values is either packed [S, ..., D] or an ordered list/tuple of [ ..., D ] tensors. The query is [R], with 1 <= R <= D; the output retains width D. Values, query, output, and first-order operator gradients are CUDA BF16. Internal FP32 accumulators may be used for normalization, logits, softmax, and reductions.

Multi-query AttnRes

attnres_multi reads the same sources with multiple learned queries. Each query has its own source weights and full-width output:

from attnres import attnres_multi

queries = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
read_multi = torch.compile(attnres_multi, fullgraph=True, dynamic=False)
outputs = read_multi(source_list, queries)  # [2, ..., 1024]
outputs.float().square().mean().backward()

Queries have shape [M, R], with 1 <= M <= 16; R=D gives standard routing and R<D uses the final R source coordinates. Packed, list and tuple sources use the same API. Set output_layout="token_major" for [..., M, D] output. For per-token query selection, attnres_selected(sources, query_bank, ids) takes an [E, R] bank and integer [..., K] IDs, returning [..., K, D].

Complete training-step times

The four-layer QK/V fixture measures compiled forward, loss, backward and AdamW on synthetic tokens. One query shares a read between Q/K and V; two queries route those roles separately. The two independent reads provide the matched baseline for fusion. These are fixture timings, not the 24-layer results above. Bars are medians of 40 randomized rounds; the reduction uses paired time ratios. See the measurements for the workload, runtime and raw samples. Performance depends on shape.

Available since 2.1.0. The graph uses PyTorch 2.10.0 / Triton 3.6.0 on H100. Small packed and selected-query reads can be slower than separate compiled reads; benchmark your shapes before adopting fusion.

Equation

For source value v_s, take its final R coordinates as the implicit key:

t_s       = v_s[..., D-R:D]
r_s       = sqrt(mean(t_s ** 2) + eps)
k_s       = t_s / r_s
score_s   = scale * dot(k_s, query)
p_s       = softmax(score, axis=source)_s
output    = sum_s p_s * v_s

Normalization and softmax run independently at each carried batch or token position. The output is not normalized, source-count weighted, or source averaged. See docs/equation.md for the complete contract.

Full and Block schedules

Full supplies the embedding and every preceding writer output. Block supplies the embedding, completed block sums, and an optional current partial sum. The caller owns block boundaries and sums a partial block before the read; it is passed as one ordinary source and is not averaged.

full_output = attnres((embedding, *writers), query)

completed = (embedding, first_block, second_block)
block_sources = completed if partial is None else completed + (partial,)
block_output = attnres(block_sources, query)

Every read evaluates its routing weights and mixture from the supplied sources. Both schedules use the same query, eps, scale, and operator contract.

Sliced LR-AttnRes

Sliced LR-AttnRes keeps full-width values and output while using a shorter implicit key and query:

rank = 64
query = torch.randn(rank, device="cuda", dtype=torch.bfloat16)
output = attnres(source_list, query)  # [..., D], BF16

For a trainable static query, use LearnedQuery(rank):

from attnres import LearnedQuery
learned_query = LearnedQuery(rank).to(device="cuda", dtype=torch.bfloat16)
output = attnres(source_list, learned_query())

Standard AttnRes is the R == D case; sliced routing uses R < D and the final R value coordinates. Projected keys, routing priors, and architectural changes are outside this package contract.

Optimization targets common ranks: powers of two from 16 upward, plus widths such as 384 and 640. Other ranks retain the same mathematical support.

Validation scope

Correctness checks compare BF16 outputs and first-order gradients with an independent BF16 PyTorch reference with FP32 internal accumulation at rtol=0.05 and atol=0.05. Coverage includes packed/list sources, repeated reads, partial Blocks, changed inputs, non-contiguous layouts, shared sources, and compiled replay. See the validation protocol.

Timing reports identify the exact source, device, runtime, workload, and measurement boundary. Failed, incomplete, and inconclusive comparisons remain visible in the results.

License

Fast-AttnRes is released under the MIT License. FLA-derived source list attribution remains in NOTICE.