Making Attention Residuals kernel-efficient.

A residual stream is already a depth router.

A standard PreNorm residual stack updates its state as h_l = h_l−1 + f_l−1(h_l−1). Unroll the recurrence and each layer receives the embedding plus every earlier layer output with a coefficient of one. This is useful for optimization, but rigid as information routing.

That fixed aggregation can let hidden-state magnitude grow with depth and can dilute any one layer's relative contribution. The sequence axis escaped a similar fixed recurrence by adopting attention. The Kimi Team's Attention Residuals report asks the natural question: why not give the depth axis learned, content-dependent routing too?

Attention Residuals make depth selective.

AttnRes replaces uniform addition with a softmax-weighted sum over previous representations. Each layer owns a learned pseudo-query w_l. Earlier layer or block outputs serve as both values and normalized routing keys.

Full AttnRes attends over every preceding layer output. Block AttnRes compresses layers into a modest number of completed block summaries. Inside a block, each layer also sees the partial sum produced so far. This bounds persistent history while retaining content-dependent cross-depth access.

The equation is small. The eager graph is not.

A readable PyTorch implementation stacks values, materializes RMS-normalized keys, computes depth logits, writes probabilities, and launches separate reductions and elementwise operations. Backpropagation retains or reconstructs many of those intermediates. Depth is much smaller than sequence length—but at large token counts, moving an extra tensor across HBM is still real work.

The optimization target is therefore not a new approximation. It is a better schedule: batch what is shared, defer what is sequential, and keep stable-softmax state minimal.

Two phases, aligned with data lifetime.

Hydra-2P treats completed history and the partial block as two softmax partitions. The first can be computed early and in parallel for all layer queries. The second arrives sequentially and can be merged exactly.

Phase 1: sweep shared history once

For each `(layer, token)` row, the Triton kernel loads the pseudo-query and walks the completed blocks. Each value is normalized in registers to form a key, dotted with the query, and immediately folded into an online-softmax accumulator.

The phase emits three sufficient statistics per row: numerator o, maximum m, and denominator ℓ. It never writes the normalized key or source probability.

Phase 2: merge the live partial exactly

When a layer's partial-block value becomes available, Phase 2 normalizes and scores that singleton source. Let `(o₁, m₁, ℓ₁)` describe completed history and `(o₂, m₂, ℓ₂)` describe the partial. A common maximum makes their merge stable and exact:

The `has_partial` mask handles the first layer in a block without a sentinel contaminating the denominator. Disabled partials contribute exactly zero mass.

Fusion is useful only if the math survives it.

The implementation keeps routing and reduction state in FP32 for FP16 and BF16 inputs. RMS variance, dot-product logits, maxima, denominators, numerators, outputs, and gradient accumulation all use FP32 internally before results are cast to the requested dtype.

This contract is backed by a materializing reference implementation and tests comparing both forward values and manually derived gradients.

Save statistics. Recompute probabilities.

Saving a depth-attention probability tensor would undermine much of the forward fusion. Hydra-2P instead saves the output and final row `(max, denominator)`. During backward, each source score is recomputed and its weight reconstructed.

From that scalar, the kernel computes the direct value gradient, the normalized-key gradient, the RMSNorm Jacobian, and the query contribution. Completed-block gradients and query gradients reduce through FP32 atomics; the partial gradient has single-writer ownership.

The current repository exposes four Triton kernels: two forward phases and two backward ownership paths. This is deliberately transparent research code: the readable PyTorch implementation carries the same manual gradient, while the Triton version makes memory ownership explicit.

The fastest tile is an empirical fact.

Kernel performance depends on hidden dimension, token count, depth, dtype, GPU architecture, register pressure, and occupancy. A single compile-time configuration is rarely ideal across the shape space.

The plan tuner evaluates valid hidden-dimension tiles and warp counts using the metric users actually care about: combined forward + backward latency. The winner is cached under a key containing GPU name, dtype, query shape, block shape, and partial shape.

COPYpython -m attnres_kernel.benchmark \
  --layers 8 \
  --blocks 8 \
  --tokens 4096 \
  --dim 128 \
  --tune

This makes the system “self-evolving” in a practical, bounded sense: it learns execution policy from measured hardware behavior while leaving the mathematical operator fixed.

Correctness is verified. GPU speed is a measurement task.

The `v0.2.0` suite covers Hydra-2P, recurrent KDA, KDA write traces, Cross-Axis Delta Routing, manual backward parity, exact partition merges, integrated-layer gradients, and CUDA backend hooks.

The development host did not expose an NVIDIA runtime, so the project does not manufacture a GPU speedup number. The benchmark is synchronized, measures median end-to-end forward + backward time, and records the selected plan. Credible results should report GPU model, PyTorch and Triton versions, dtype, full `S/N/M/D` shape, absolute timings, and speedup.

KDA adds a second, complementary memory axis.

Hydra-2P routes across model depth. Kimi Delta Attention (KDA) routes across sequence time using a fixed-size recurrent matrix state. The repository now integrates both without pretending they are the same operation.

The new `KimiDeltaAttention` module follows the Kimi Linear parameterization: short causal convolutions and SiLU for Q/K/V, L2-normalized Q/K, low-rank channel decay, a head-wise β gate, head RMSNorm, and a gated output projection. Its portable recurrent state uses FP32 and is exactly streamable across sequence segments.

`KDAAttnResLayer` first uses Hydra-2P to construct the current layer input, normalizes it, applies KDA across tokens, and accumulates the KDA update into the live AttnRes block. The functional KDA API can dispatch to FLA—and therefore FlashKDA on compatible installations—while retaining the recurrent implementation as a correctness oracle.

CADR routes what changed—and whether KDA found it novel.

Cross-Axis Delta Routing replaces cumulative depth values with residual deltas and adds a second normalized key derived from KDA's gated write error, write energy, and memory retention. The descriptor changes source selection but never enters the hidden value stream.

The CADR implementation extends Hydra-2P's completed/live split to both key spaces. Two fused Triton forward kernels stream delta/descriptor pairs and merge the live source exactly. Two backward kernels reconstruct probabilities and differentiate the delta value path, both RMSNorm key paths, and both pseudo-query families.

The new `CADRKDAAttnResLayer` closes the loop: CADR constructs the layer input, KDA mixes across tokens, and KDA emits the descriptor cached with the next residual delta. The architecture document contains the complete derivation, tensor contract, kernel ownership model, diagnostics, ablations, and novelty boundary.

The repository is built to be read, tested, and changed.

The public API validates shapes and dispatches supported CUDA inputs to Triton, with a custom PyTorch fallback everywhere else. The full implementation is MIT licensed at github.com/manishklach/attnres-kernel-lab.

COPYimport torch
from attnres_kernel import attnres

S, N, B, T, D = 8, 8, 2, 2048, 128
q = torch.randn(S, D, device="cuda", dtype=torch.float16,
                requires_grad=True)
blocks = torch.randn(N, B, T, D, device="cuda",
                     dtype=torch.float16, requires_grad=True)
partials = torch.randn(S, B, T, D, device="cuda",
                       dtype=torch.float16, requires_grad=True)
has_partial = torch.tensor([False] + [True] * 7,
                           device="cuda")

output = attnres(q, blocks, partials, has_partial)
output.square().mean().backward()

For transformer integration, `prepare_inter_block` computes Phase 1 ahead of the sequential layer group, and `merge_partial` consumes each live state. `AttnResPlan` makes backend, tile, warp count, and epsilon explicit when experiments require reproducible dispatch.