Hugging Face's logo

jburtoft
/
kda-neuron-kernels

Kernels
kernel
neuron
trainium
kda
linear-attention
fla-core
training
backward
Model card Files
xet
 Community
kda-neuron-kernels / README.md
jburtoft's picture
jburtoft
v1.3: fused multi-chunk backward (kda_chunked_fused, 2-5x faster training)
6a4016a verified
preview code
|
Raw
History Blame Contribute Delete
19.8 kB
metadata 
license: apache-2.0
library_name: kernels
tags:
  - kernel
  - neuron
  - trainium
  - kda
  - linear-attention
  - fla-core
  - training
  - backward

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.

Compatible with AWS Trainium (trn2 tested; trn1 / trn3 not yet verified). Runs under PyTorch Native (Beta 3+ / Beta 4).

What this package provides

Inference forward kernels (v1.0/v1.1) — three raw NKI entry points:

  • kda_recurrent_fwd(q, k, v, g, beta) — TKG (decode) per-token recurrence. One (batch, head) invocation processes S tokens sequentially.
  • kda_recurrent_fwd_state(q, k, v, g, beta) — same as above, also returns the final recurrent state for CTE→TKG hand-off.
  • kda_chunk_step(q, k, v, beta_broadcast, g_cumsum, g_last, state_in) — CTE (prefill) per-chunk step. Processes one 128-token chunk given the state from the previous chunk.

Training / backward (v1.2/v1.3) — differentiable torch.autograd.Function wrappers plus the underlying backward kernels:

  • kda_recurrent(q, k, v, g, beta, initial_state=None) → (output, final_state), differentiable. loss.backward() routes through the NKI recurrent backward.
  • kda_chunked(q, k, v, g, beta, initial_state=None) → (output, final_state), differentiable. Loops chunks in Python around the single-chunk kernels.
  • kda_chunked_fused(q, k, v, g, beta, initial_state=None) (v1.3) → same numerics as kda_chunked, but processes all chunks in one NKI launch per direction (state carried in SBUF). 2–5× faster on hardware for multi-chunk sequences (speedup grows with sequence length). Recommended for training.
  • Raw kernels also exported: kda_recurrent_fwd_v2, kda_chunk_step_v2, kda_recurrent_bwd, kda_chunk_bwd, and the fused kda_fused_chunked_fwd, kda_fused_chunked_bwd.

Backward gradients verified against fla-core autograd on both the NKI simulator (trn2-pinned) and real trn2 hardware — see the Parity and Performance sections.

Not yet provided: a full NeuronKDA(nn.Module) drop-in replacement for HF Transformers' Kda layer (gated on the upstream transformers KDA integration finalizing).

For a full-layer example, see the sibling kernel package: jburtoft/qwen35-deltanet-neuron-kernels which does the equivalent for Qwen3.5's Gated DeltaNet.

Installation & Environment

Requires:

  • PyTorch Native Beta 3+ (torch-neuronx 2.11.3+, PyTorch 2.11+)
  • NKI ≥ 0.4.0 (tested on NKI 0.5.0 / SDK 2.31 / Beta 4)
  • transformers with KernelConfig support (main branch commit ≥ 4b0a02931b, i.e. 5.10.0.dev0 or later)
  • kernels==0.15.2

Consult Neuron documentation for the current Beta setup.

Usage

Direct kernel calls (v1.0 supported path) 

import torch
import torch.nn.functional as F
from kda_neuron_kernels.build.torch_neuron import (
    kda_recurrent_fwd, kda_recurrent_fwd_state, kda_chunk_step,
)

# Example: prefill one chunk of 128 tokens for a single (batch, head) slice
B, S, H, Dk = 1, 128, 64, 128
q_raw = torch.randn(B, S, H, Dk)
k_raw = torch.randn(B, S, H, Dk)
v     = torch.randn(B, S, H, Dk)  # v_dim == k_dim == 128
g     = -torch.rand(B, S, H, Dk) * 0.01  # per-dim log-decay, negative
beta  = torch.rand(B, S, H)  # per-head scalar

# Wrapper preprocessing:
# - L2-norm q, k and scale q by 1/sqrt(Dk) (fla-core convention)
q = F.normalize(q_raw, p=2, dim=-1) * (Dk ** -0.5)
k = F.normalize(k_raw, p=2, dim=-1)

# Dispatch chunked kernel per (b, h)
state = torch.zeros(Dk, Dk, dtype=torch.float32).to("neuron")
for b in range(B):
    for h in range(H):
        q_c = q[b, :, h].contiguous().to("neuron")   # (128, 128)
        k_c = k[b, :, h].contiguous().to("neuron")
        v_c = v[b, :, h].contiguous().to("neuron")
        g_c = g[b, :, h].contiguous()
        gc = torch.cumsum(g_c, dim=0).to("neuron")
        gl = gc[-1:, :].expand(128, Dk).contiguous()
        beta_c = beta[b, :, h]
        beta_bc = beta_c.unsqueeze(-1).expand(128, Dk).contiguous().to("neuron")

        chunk_out, state = kda_chunk_step(q_c, k_c, v_c, beta_bc, gc, gl, state)
        # chunk_out: (128, 128) float32 per-token output
        # state: (128, 128) carries to next chunk

For a fully-worked example (with all wrapper preprocessing done in PyTorch), see tests/example_usage.py in this repo.

Via KernelConfig (v1.1 planned)

v1.1 will support the following API (currently blocked on NeuronKDA full-layer wrapper implementation, planned once the upstream transformers KDA integration finalizes):

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig({
    "Kda": "jburtoft/kda-neuron-kernels:NeuronKDA",
})

model = AutoModelForCausalLM.from_pretrained(
    "<your-kda-model>",   # any HF model whose attention layer is KDA
    dtype=torch.bfloat16,
    kernel_config=kernel_config,
    device_map="neuron",
    trust_remote_code=True,
)

Training (v1.2 — differentiable) 

import torch, torch.nn.functional as F
from kda_neuron_kernels.build.torch_neuron import kda_chunked_fused, kda_recurrent

S, D = 256, 128   # S must be divisible by 128 for the chunked paths
q = F.normalize(torch.randn(S, D), p=2, dim=-1) * (D ** -0.5)
k = F.normalize(torch.randn(S, D), p=2, dim=-1)
v = torch.randn(S, D) * 0.3
g = -torch.rand(S, D) * 0.01                                   # per-K raw log-decay
beta = (torch.rand(S, 1) - 0.5 + 1.0).expand(S, D).contiguous()  # per-token scalar bcast
for t in (q, k, v, g, beta):
    t.requires_grad_(True)

# CPU -> simulator; or .to("neuron") for trn2 hardware.
# kda_chunked_fused: single NKI launch for all chunks (fastest). kda_chunked: Python chunk loop.
out, final_state = kda_chunked_fused(q, k, v, g, beta, initial_state=None)
loss = out.sum()
loss.backward()          # gradients flow through the fused NKI backward
# q.grad, k.grad, v.grad, g.grad, beta.grad now populated

# Decode-style recurrent training is also available (zero initial_state only):
# out, fs = kda_recurrent(q, k, v, g, beta)

Per (batch, head): loop B*H in the caller. kda_recurrent currently requires zero initial_state; use kda_chunked for state carry-over across sequence packs.

Wrapper Contract (Read This)

All three kernels take raw q, k already L2-normed by the caller, with q additionally scaled by 1/sqrt(dk) (fla-core convention). The kernels compute all decay-related pre-scaling (exp(gc_mean), exp(-gc_mean), exp(gc), exp(g_last-gc), exp(g_last)) internally from gc = cumsum(g).

For kda_chunk_step:

  • q, k: RAW L2-normed q, k (q scaled by 1/sqrt(dk)) — shape (128, 128)
  • v: raw value tensor — shape (128, 128)
  • beta_broadcast: per-token scalar beta broadcast to (128, 128)
  • g_cumsum: per-dim cumsum(g) within the chunk — shape (128, 128)
  • g_last: g_cumsum[-1:, :] broadcast to (128, 128)
  • state_in: recurrent state from previous chunk — shape (128, 128)
  • Returns: (chunk_out (128, 128), state_out (128, 128))

For kda_recurrent_fwd:

  • query, key: (S, 128) RAW L2-normed
  • value, g_in, beta_in: (S, 128) per-dim / broadcast form (see nki_kda.py docstring)
  • Returns: output (S, 128) per-token output

Hard Constraints

  • head_k_dim == head_v_dim == 128 — matches NeuronCore SBUF partition width. Not currently portable to other head dims.
  • chunk_size == 128 for kda_chunk_step.
  • float32 inputs (kernel internally handles precision).
  • trn2 tested. trn1 / trn3 not verified in v1.0.

Parity

Measured against fla-core naive_recurrent_kda and naive_chunk_kda PyTorch references on random inputs at a typical KDA operating regime (g_scale=0.01, seq_len=128, single (b, h)):

Kernel cos_sim vs fla reference max_abs_diff
kda_recurrent_fwd (S=128) 1.00000 3.4e-8
kda_chunk_step (C=128) 0.99988 1.22e-3

Precision floor sources:

  • Recurrent: after the v1.1 K-vs-V axis fix (see Fix history), the recurrent kernel matches fla to max_abs_diff ~= 3e-8 (fp32 rounding floor) across S = 16, 128, 256, 512. The v1.0 number (0.99977) was dominated by the state-decay axis bug, not by BF16 accumulation.
  • Chunked: scalar-mean approximation in the intra-chunk attention (exp(gc_mean_i - gc_mean_j) instead of exact per-dim exp(gc_i - gc_j)). A deliberate O(BT^2) vs O(BT^2 * K) tradeoff; the ceiling (~0.99988) is inherent to the algorithm choice and unaffected by the v1.1 fix.

Note: an earlier version of the chunked kernel (predating v1.0) had a latent accuracy bug (cos_sim ≈ 0.78 on random inputs at typical KDA g-scale). This bug was fixed in v1.0. A second axis-convention bug (state decayed per-V instead of per-K) affecting both kernels was fixed in v1.1. See "Fix history" below.

Performance

Measured on trn2.3xlarge, LNC=2, single logical core, PyTorch Native Beta 4, SDK 2.31, single (batch, head) invocation:

Chunked (prefill)

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%
MFU (FP32 peak 40 TFLOPS/LNC=2) 4.71%
Achieved HBM GB/s 6.77
MBU (empirical peak 218 GB/s/LNC=2) 3.11%

Recurrent (decode)

Metric Value
Wall-clock per call (S=128) 838 μs
Per-token (amortized) 6.55 μs
Achieved TFLOPS 0.023
MFU (BF16 peak 158 TFLOPS/LNC=2) 0.01%
MFU (FP32 peak 40 TFLOPS/LNC=2) 0.06%
Achieved HBM GB/s 0.47
MBU (empirical peak 218 GB/s/LNC=2) 0.22%

The recurrent kernel is overhead-dominated at S=128 (per-token wall ≫ per-token useful work). For real decode throughput, batch multiple tokens per invocation:

S per invocation Per-token wall-clock
1 70.4 μs (all overhead)
8 9.7 μs
128 6.55 μs
512 6.38 μs

Increasing S from 1 to 8 gives a 7.3× per-token improvement with no kernel changes. This is the dominant lever for decode throughput.

Comparison to torch.compile(backend="neuron") on the same reference

The direct "is NKI worth it?" question, measured on the same trn2.3xlarge with the same fla-core naive_* PyTorch reference compiled through the Neuron XLA backend:

Recurrent (decode) -- NKI vs torch.compile(naive_recurrent_kda)

S NKI (μs) torch.compile (μs) NKI vs torch.compile
1 70 66 0.94× (torch marginally faster in the overhead-dominated regime)
8 77 97 1.26×
32 231 272 1.17×
128 848 939 1.11×

At S ≥ 8, NKI is 11-26% faster per invocation. Both converge to ~6.5-7.3 μs per token as S grows. The NKI advantage is largest in the S=8-32 range where per-call overhead is amortized but the sequential recurrence stays short.

Chunked (prefill) -- NKI vs torch.compile(naive_chunk_kda)

C NKI (μs) torch.compile (μs) NKI vs torch.compile
128 87 1660 19.08×

On prefill, NKI is 19× faster than the same algorithm compiled through the Neuron XLA backend. The NKI kernel packs the 24 128×128 matmuls of the Neumann series plus all elementwise ops into a single NEFF with all intermediates staying in SBUF; the XLA-compiled path does many HBM round-trips.

Compilation time (first-run)

Path S=1 S=8 S=32 S=128 recurrent C=128 chunked
NKI ~8s ~8s ~9s ~9s ~2s
torch.compile 1.9s 3.4s 10.6s 45.3s 284s (4.7 min)

torch.compile is competitive on tiny recurrent-workload first-runs (S ≤ 8) but takes 5x longer at S=128 and 142× longer on chunked because the entire Neumann-series graph has to be traced through Python loops and lowered by XLA. NKI's @nki.jit compiles the whole algorithm as one function.

Bottom line: NKI is a modest win for short recurrent workloads (1.1-1.3× per-invocation) and a step-change win for chunked prefill (19×) on both wall-clock and compile time. Full report at torch_compile_comparison.md in the kda-kernel project working tree.

Peak references

All MFU / MBU denominators cited above are per LNC=2 core on trn2 (NeuronCore-v3), from AWS documentation and empirical measurements:

  • PE peak: 158 BF16 TFLOPS (spec) / 40 FP32 TFLOPS (spec). Trainium2 architecture guide.
  • HBM peak: 750 GB/s theoretical (3 TB/s device / 8 physical cores × 2 per LNC=2), 218 GB/s empirical achievable (per project measurements on similar workloads).

Peak numbers are per LNC=2 core on trn2 (NeuronCore-v3), from the AWS Trainium2 architecture guide and empirical measurements.

Fix history

v1.3 (2026-08-06) — Fused multi-chunk backward

Added kda_chunked_fused + the raw fused kernels (kda_fused_chunked_fwd, kda_fused_chunked_bwd): all chunks processed in one NKI launch per direction, recurrent state / dSn carried in SBUF across the internal chunk loop. Same numerics as kda_chunked (verified cos_sim ≈ 1.0 end-to-end through loss.backward()), but eliminates the per-chunk launch overhead of the Python chunk loop.

Performance (fwd+bwd wall-clock, trn2, single (b,h)):

S Chunks Python-loop Fused Speedup
256 2 695 μs 336 μs 2.07×
512 4 1320 μs 338 μs 3.90×
1024 8 2588 μs 529 μs 4.89×

The fused fwd+bwd is nearly flat from S=256→512 (launch overhead paid once, not NC times); the speedup grows with sequence length. Use kda_chunked_fused for training.

v1.2 (2026-08-06) — Training / backward support

Added NKI backward kernels + torch.autograd.Function wrappers (kda_recurrent, kda_chunked), making KDA trainable on Trainium. Backward math verified against fla-core autograd on the simulator (trn2-pinned) and real trn2 hardware, end-to-end through loss.backward().

Parity (backward, vs fla-core naive autograd):

  • Recurrent: all 5 gradients (dq, dk, dv, dg, dbeta) cos_sim ≥ 0.9998 — exact algorithm.
  • Chunked: dq/dk/dv/dbeta cos_sim ≥ 0.9998; dg differs from fla (chunked forward uses the scalar-mean intra-chunk-attention approximation, so our dg is the exact gradient of our forward — self-consistent cos_sim=1.0 — but differs from fla's exact per-dim dg).

Performance (backward, trn2, single (b,h), from perf/backward_perf.md in the source project):

  • recurrent bwd S=128: 2031 μs (2.4× the forward); chunked bwd C=128: 174 μs (2.0× the forward).
  • vs torch.compile(backend="neuron") fwd+bwd: 9.0× faster (recurrent) and 86.6× faster (chunked); NKI compiles in seconds vs 200-585 s for torch.compile's autograd graph.

Two NKI techniques used in the backward: reverse-cumsum via a single matmul (U @ dgc with upper-tri ones), and per-free-dim reductions via tensor_reduce(axis=(1,)).

Deferred: a fused multi-chunk backward (one NKI launch for all chunks) to remove the Python chunk-loop's per-chunk launch overhead at long sequences.

v1.1 (2026-08-05) — State-decay axis fix (T-KDA-02)

Accuracy fix (T-KDA-02): Both kernels decayed the recurrent state per-V column (state[k, v] *= exp(g[v])) instead of per-K row (state[k, v] *= exp(g[k])) as fla-core's canonical KDA does. The bug was invisible with uniform g (per-dim identical values) but diverged on per-dim varying g at typical KDA scale.

  • Recurrent (kda_recurrent_fwd): parity vs fla went from cos_sim 0.99977 → 1.00000 (max_abs_diff 2.65e-3 → 3.4e-8). The fix also simplifies the kernel — the per-V decay required a transpose-scale-transpose dance (2 nc_transpose + 1 tensor_scalar per token); per-K decay is a single tensor_scalar on the partition axis, so v1.1 is also marginally faster.
  • Chunked (kda_chunk_step): the per-K state decay is fixed (matters for state carry-over across chunks). Forward-output parity is unchanged at 0.99988 because that number is dominated by the separate scalar-mean intra-chunk-attention approximation.

All 6 recurrent backward gradients (dq, dk, dv, dg, dbeta, dinitial_state) now match fla naive_recurrent_kda autograd at cos_sim = 1.0. (Backward kernels themselves ship separately; see the kda-backward work.)

v1.0 (2026-08-05) — Initial public release

Accuracy fix (T-KDA-01): The pre-release version of the chunked kernel had a wrong-sign / extra-factor pattern in FOUR internal steps (k_beta for QK/A construction, k_beta * exp_gc for k_cumdecay, q_c * exp_gc for attn_inter, k_c * exp_gl_minus_gc for k_state_decay) that produced cos_sim ≈ 0.78 vs fla-core reference at typical KDA g-scale. This is fixed in v1.0: the wrapper contract now passes raw q, k (previously required wrapper to pre-multiply by exp(±gc_mean)), and the kernel computes all decay flavors internally, correctly distinguishing between the four different scaled versions of k needed.

Prefill performance optimization: 9× nc_matmul(stationary=X, moving=eye) transpose-via-matmul calls replaced with nisa.nc_transpose(dst, data=X). This gives the compiler an explicit transpose hint and delivers a −6.7% wall-clock improvement (93.3 → 87.1 μs per chunk) with zero parity risk.

Not addressed in v1.1 (deferred to v2.0)

  • Chunked dg parity vs fla: because the chunked forward uses a scalar-mean intra-chunk-attention approximation, the g-gradient differs from fla's exact per-dim form (dq/dk/dv/dbeta match at 0.99986). Fixing requires exact per-dim intra-chunk attention (O(BT^2·K) instead of O(BT^2)).
  • Deeper R2 fix: operand-order refactoring to compute QK.T directly and skip the transpose pairs (estimated additional ~5-10% wall-clock).
  • Split-Neumann across LNC=2: dividing the Neumann series across two physical cores (up to 2× wall-clock, requires cross-core state management).
  • Full-layer wrapper: NeuronKDA(nn.Module) drop-in replacement for HF Transformers' Kda layer (once upstream integration finalizes).
  • Backward / training kernels: NKI backward kernels for training (chunked + recurrent) are in development; the math is verified but the NKI ports are not yet in this package.

References

License

Apache-2.0. This kernel package is an inference-runtime component, not a fine-tuned model. The fla-core algorithm reference is MIT-licensed and compatible.

Contributing

This package's development happens in the internal kda-kernel project. External contributions welcome via PR to this HuggingFace Hub repo. For issues affecting the underlying kernels, please file on the Neuron team's internal ticketing system (not on public GitHub) so we can route them correctly.