kda-neuron-kernels
Neuron NKI kernels for KDA (Kernel-based Decomposed Attention) linear attention.
Model-agnostic implementation of the KDA algorithm described in the
flash-linear-attention (fla-core) library.
Compatible with any HuggingFace Transformers model whose attention layer follows
the KDA algorithm. Runs on AWS Trainium (trn2) under PyTorch Native.
This is a
kernel-type repository (build variant
torch-neuron, backend
neuron).
Load it with the
kernels library on a
Trainium machine:
1from kernels import get_kernel
2# Recommended: pin the stable major version (currently the latest).
3k = get_kernel("jburtoft/kda-neuron-kernels", version=3, trust_remote_code=True)
4# Or always track the latest published kernels:
5# k = get_kernel("jburtoft/kda-neuron-kernels", revision="main", trust_remote_code=True)
6# k.kda_chunk_step_exact(...), k.kda_chunk_step_exact_bwd(...), k.kda_recurrent_fwd(...), etc.
get_kernel requires either version= or revision= — there is no zero-arg
default (this is a kernels-library requirement, not a repo setting). version=2
is the current latest major version; revision="main" always tracks the newest
published tree. Both resolve to the same content today.
Runtime requirement (important): load this from a PyTorch Native
(torch-neuronx) environment where torch is a CPU/Neuron build and torch.neuron
is registered (e.g. the DLAMI venvs aws_neuronx_venv_pytorch_2_9_nxd_inference
or a PyTorch-Native Beta venv). The kernels library selects the build variant
from the active torch backend; in a CUDA torch build (some DLAMI base venvs ship
torch ...+cuXXX) it will detect backend cuda and refuse the neuron variant
with "backend (neuron) does not match system backend (cuda)". If you hit that,
switch to a Neuron/PyTorch-Native venv (verify with
python -c "from kernels.backends import _backend; print(_backend().name)" → should
print neuron).
What this package provides
Inference (forward)
kda_recurrent_fwd(q, k, v, g, beta) — decode / token-generation per-token
recurrence. One (batch, head) invocation processes S tokens sequentially.kda_recurrent_fwd_state(q, k, v, g, beta) — same, and also returns the final
recurrent state for prefill→decode hand-off.kda_chunk_step(q, k, v, beta, g_cumsum, g_last, state_in) — prefill per-chunk
step. Processes one 128-token chunk given the state from the previous chunk. Uses a
scalar-mean decay approximation for the intra-chunk term — see the warning below.kda_chunk_step_exact(q, k, v, beta, g, state_in) — numerically exact
per-channel prefill (sub-chunk + WY reformulation). Use this when the model's gate
decay is non-trivial (see warning below). ~1.05× the latency of kda_chunk_step.kda_chunk_step_exact_multihead(q, k, v, beta, g, state_in) — head-interleaved
exact prefill over NV heads ([NV, C, dk] shapes).kda_decode_batch(q, k, v, g, beta, state_in) — batched multi-(request, head)
decode; advances all B*nv items one token in a single call. Shapes [B, nv, dk].kda_chunk_step_exact_bwd(q, k, v, beta, g, state_in, d_output, dS_final) —
exact chunked BACKWARD (gradient of kda_chunk_step_exact). Returns
dq, dk, dv, dg, dbeta, dstate_in, matching fla-core autograd at cos_sim 1.0
for all five gradients in every gate regime (g = 0.01…2.0). Pair this with
kda_chunk_step_exact for training. The older approximate chunked backward has
a broken dg (cos_sim ≈ 0.09 in ALL regimes) and NaNs at large decay.
Training (differentiable, loss.backward()-ready)
kda_recurrent(q, k, v, g, beta, initial_state=None) → (output, final_state).
Differentiable; routes through the NKI recurrent backward. Requires zero
initial_state.kda_chunked(q, k, v, g, beta, initial_state=None) → (output, final_state).
Differentiable; loops chunks in Python around the single-chunk kernels. Supports
state carry-over across chunks.kda_chunked_fused(q, k, v, g, beta, initial_state=None) → (output, final_state).
Same numerics as kda_chunked, but processes all chunks in one NKI launch per
direction. Faster on multi-chunk sequences. Recommended for training.- Raw kernels also exported:
kda_recurrent_fwd_v2, kda_chunk_step_v2,
kda_recurrent_bwd, kda_chunk_bwd, kda_fused_chunked_fwd, kda_fused_chunked_bwd.
kda_recurrent_bwd_batched(q, k, v, g, beta, state_stack, d_output, dS_final)
— batched recurrent backward over all B*nv (request, head) items in one launch
(flat work list + ping-pong state). ~1.1–1.24× over a per-(b,h) loop; all six
gradients bit-match the single-item kda_recurrent_bwd.
Requirements
- Hardware: AWS Trainium (tested on trn2.3xlarge).
- SDK / runtime: PyTorch Native (
device="neuron"), torch-neuronx 2.11+, PyTorch 2.11+.
- NKI ≥ 0.4.0.
kernels ≥ 0.15.2 (to load via get_kernel).transformers with KernelConfig support, if using the KernelConfig path.
Usage
Inference — direct kernel calls
1import torch
2import torch.nn.functional as F
3from kda_neuron_kernels.build.torch_neuron import kda_chunk_step_exact
4
5# Prefill one 128-token chunk for a single (batch, head) slice.
6S, Dk = 128, 128
7q_raw = torch.randn(S, Dk)
8k_raw = torch.randn(S, Dk)
9v = torch.randn(S, Dk)
10g = -torch.rand(S, Dk) * 0.01 # per-channel log-decay (negative)
11beta = torch.rand(S) # per-token scalar
12
13# Caller preprocessing (fla-core convention): L2-norm q, k and scale q by 1/sqrt(Dk).
14q = F.normalize(q_raw, p=2, dim=-1) * (Dk ** -0.5)
15k = F.normalize(k_raw, p=2, dim=-1)
16beta_bc = beta.unsqueeze(-1).expand(S, Dk).contiguous()
17
18state = torch.zeros(Dk, Dk, dtype=torch.float32).to("neuron")
19chunk_out, state = kda_chunk_step_exact(
20 q.to("neuron"), k.to("neuron"), v.to("neuron"),
21 beta_bc.to("neuron"), g.to("neuron"), state,
22)
23# chunk_out: (128, 128) per-token output; state: (128, 128) carries to the next chunk.
See tests/example_usage.py for a fully-worked example.
Training
1import torch, torch.nn.functional as F
2from kda_neuron_kernels.build.torch_neuron import kda_chunked_fused
3
4S, D = 256, 128 # S must be divisible by 128
5q = F.normalize(torch.randn(S, D), p=2, dim=-1) * (D ** -0.5)
6k = F.normalize(torch.randn(S, D), p=2, dim=-1)
7v = torch.randn(S, D) * 0.3
8g = -torch.rand(S, D) * 0.01 # per-channel log-decay
9beta = (torch.rand(S, 1) - 0.5 + 1.0).expand(S, D).contiguous() # per-token scalar bcast
10for t in (q, k, v, g, beta):
11 t.requires_grad_(True)
12
13out, final_state = kda_chunked_fused(q, k, v, g, beta, initial_state=None) # .to("neuron") for hardware
14loss = out.sum()
15loss.backward() # gradients flow through the NKI backward
16# q.grad, k.grad, v.grad, g.grad, beta.grad now populated
Kernels operate per (batch, head); loop B*H in the caller.
Input contract
Callers pass raw q, k already L2-normed, with q additionally scaled by
1/sqrt(dk) (fla-core convention). The kernels compute all decay-related scaling
internally from g.
For kda_chunk_step (and _v2):
q, k: L2-normed q (scaled by 1/sqrt(dk)), L2-normed k — shape (128, 128)v: value tensor — (128, 128)beta: per-token scalar, broadcast to (128, 128)g_cumsum: per-channel cumsum(g) within the chunk — (128, 128)g_last: g_cumsum[-1:, :] broadcast to (128, 128)state_in: recurrent state from the previous chunk — (128, 128)- Returns
(chunk_out, state_out), each (128, 128).
For kda_recurrent_fwd:
q, k: (S, 128) L2-normed (q scaled)v, g, beta: (S, 128) (beta per-token scalar, broadcast across the dim)- Returns
output (S, 128).
Constraints
head_k_dim == head_v_dim == 128 (matches the NeuronCore SBUF partition width).
Other head dims are not supported.chunk_size == 128 for the chunked kernels; S must be divisible by 128.- float32 inputs.
kda_recurrent (training wrapper) requires zero initial_state; use kda_chunked
for state carry-over across sequence packs.
⚠️ Functional warning — chunked gate-decay approximation
kda_chunk_step (and its training wrappers kda_chunked / kda_chunked_fused) use a
scalar-mean approximation for the intra-chunk attention decay
(exp(mean_c(gc)_i - mean_c(gc)_j) instead of the exact per-channel
exp(gc_i - gc_j)). This is a compute/accuracy tradeoff.
The approximation is only accurate for small gate decay. Measured single-chunk
cosine similarity vs the fla-core reference:
gate scale g | cos_sim (kda_chunk_step) |
|---|
| ~0.01 (small) | ~0.99 |
| ~0.3 | ~0.49 |
| ~2.0 | ~0.22 |
If your model has non-trivial gate decay, use kda_chunk_step_exact (or
kda_chunk_step_exact_multihead), which is numerically exact (cos_sim ≥ 0.9999999
across all gate regimes) at ~1.05× the latency. The recurrent kernels
(kda_recurrent_fwd, kda_recurrent) are exact in all regimes.
Because kda_chunked / kda_chunked_fused differentiate the approximate forward,
their dg gradient is the exact gradient of the approximate forward — self-consistent
for training with these kernels, but not equal to the exact-per-channel dg unless the
approximation is accurate (i.e. small gate decay).
Parity
Against the fla-core naive_recurrent_kda / naive_chunk_kda PyTorch references
(random inputs, g_scale=0.01, seq_len=128, single (batch, head)):
| Kernel | cos_sim vs fla | max_abs_diff |
|---|
kda_recurrent_fwd (S=128) | 1.00000 | 3.4e-8 |
kda_chunk_step (C=128, small gate) | 0.99988 | 1.2e-3 |
kda_chunk_step_exact (C=128, all gate regimes) | ≥ 0.9999999 | ~1e-6 |
Backward gradients vs fla-core autograd, across gate regimes (g = 0.01 / 0.3 / 1.0 / 2.0):
| Backward kernel | dq | dk | dv | dg | dbeta |
|---|
kda_chunk_step_exact_bwd (all regimes) | 1.0000 | 1.0000 | 1.0000 | 1.0000 | 1.0000 |
approximate kda_chunk_bwd @ g=0.01 | 0.9999 | 0.9999 | 0.9999 | 0.085 | 0.9999 |
approximate kda_chunk_bwd @ g=0.3 | 0.990 | 0.990 | 0.991 | 0.121 | 0.993 |
approximate kda_chunk_bwd @ g=2.0 | NaN | NaN | NaN | NaN | NaN |
The exact backward fixes the approximate kernel's uncorrelated dg (which is wrong
in every regime, not just at high decay) and its NaN at large gate decay.
Training backward gradients (kda_recurrent, kda_chunked) verified end-to-end
through loss.backward() against fla-core autograd: recurrent all five gradients
cos_sim ≥ 0.9998; chunked dq/dk/dv/dbeta ≥ 0.9998 (with the dg caveat above —
use kda_chunk_step_exact_bwd to fix it).
Performance
Measured on trn2.3xlarge, LNC=2, single logical core, single (batch, head) invocation.
Prefill (chunked)
| Metric | Value |
|---|
| Wall-clock per chunk (C=128) | 87 μs |
| Per-token effective | 0.68 μs |
| Achieved TFLOPS | 1.89 |
| MFU (BF16 peak 158 TFLOPS/LNC=2) | 1.19% |
| MBU (empirical 218 GB/s/LNC=2) | 3.11% |
Decode (recurrent)
| Metric | Value |
|---|
| Wall-clock per call (S=128) | 527 μs |
| Per-token (amortized) | 4.1 μs |
The recurrent kernel is overhead-dominated at small sequence lengths. For real decode
throughput, batch tokens (or requests via kda_decode_batch): per-token wall-clock
drops from ~70 μs at S=1 to ~6 μs at S=128, and kda_decode_batch amortizes launch
overhead across a whole serving batch.
Training (fused vs Python chunk loop, fwd+bwd)
| S | Chunks | kda_chunked (loop) | kda_chunked_fused |
|---|
| 256 | 2 | 695 μs | 336 μs |
| 512 | 4 | 1320 μs | 338 μs |
| 1024 | 8 | 2588 μs | 529 μs |
The fused path pays per-launch overhead once instead of per chunk, so its advantage
grows with sequence length. Prefer kda_chunked_fused for training.
MFU/MBU denominators are per LNC=2 core on trn2 (NeuronCore-v3), from the AWS
Trainium2 architecture guide and empirical measurements.
Not provided
- A full
nn.Module drop-in replacement for a HuggingFace Kda layer (planned).
References
- Algorithm: flash-linear-attention (fla-core)
— KDA is defined in
fla/ops/kda/.
License
Apache-2.0. This is an inference/training runtime kernel package, not a model. The
fla-core algorithm reference is MIT-licensed and compatible.
