Kernels:

---

jburtoft

/

kda-neuron-kernels

like 0

---
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](https://github.com/fla-org/flash-linear-attention).
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`](https://github.com/huggingface/kernels) library on a
Trainium machine:

```python
from kernels import get_kernel
k = get_kernel("jburtoft/kda-neuron-kernels", version=2, trust_remote_code=True)
# k.kda_chunk_step_exact(...), k.kda_chunk_step_exact_bwd(...), k.kda_recurrent_fwd(...), etc.
```

**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`.

## 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

```python
import torch
import torch.nn.functional as F
from kda_neuron_kernels.build.torch_neuron import kda_chunk_step_exact

# Prefill one 128-token chunk for a single (batch, head) slice.
S, Dk = 128, 128
q_raw = torch.randn(S, Dk)
k_raw = torch.randn(S, Dk)
v     = torch.randn(S, Dk)
g     = -torch.rand(S, Dk) * 0.01   # per-channel log-decay (negative)
beta  = torch.rand(S)               # per-token scalar

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

state = torch.zeros(Dk, Dk, dtype=torch.float32).to("neuron")
chunk_out, state = kda_chunk_step_exact(
    q.to("neuron"), k.to("neuron"), v.to("neuron"),
    beta_bc.to("neuron"), g.to("neuron"), state,
)
# 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

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

S, D = 256, 128   # S must be divisible by 128
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-channel 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)

out, final_state = kda_chunked_fused(q, k, v, g, beta, initial_state=None)  # .to("neuron") for hardware
loss = out.sum()
loss.backward()   # gradients flow through the NKI backward
# 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)](https://github.com/fla-org/flash-linear-attention)
  — 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.