Completes the optimization set -- Task 021's batched recurrent backward (1.1-1.24x,
all 6 gradients bit-match the per-(b,h) kernel) was validated but not previously
packaged. Now exported and documented. Bump metadata version 2 -> 3.
docs: document version=2 (pinned) and revision=main (always-latest) load options; note no zero-arg default in kernels lib
Jim Burtoft
15b4561
v1.6: kernel-repo distribution fix
Jim Burtoft
eacd70d
- Bump metadata.json version 1 -> 2 so kernels' get_kernel(version=2) resolves
the current tree.
- README: document the PyTorch-Native runtime requirement for get_kernel (kernels
backend detection needs torch.neuron / a non-CUDA torch build).
Also (out-of-band): syncing this tree to the kernel-type repo endpoint
(huggingface.co/kernels/jburtoft/kda-neuron-kernels), which get_kernel reads and
which was frozen at an old commit -- the model-repo pushes (v1.1..v1.5) never
reached it.
Gradient of the exact 16-sub-chunk + WY forward. Fixes the approximate chunked
backward whose dg gradient is uncorrelated with the true gradient (cos_sim ~0.09)
in ALL gate regimes and which NaNs at large decay.
kda_chunk_step_exact_bwd matches fla-core autograd at cos_sim 1.0 for
dq/dk/dv/dg/dbeta in every regime (g=0.01..2.0), no overflow. Validated via a
NumPy analytic backward against the torch autograd oracle, then the NKI kernel
against fla-core autograd directly.
Also folds in the confidentiality-sanitization docstring cleanups deferred from
v1.4. No kernel logic changed in existing files (verified statement-level).
- kda_recurrent_fwd: decode key-fold, -37% wall-clock at S=128 (838->527us),
parity improved to 1.0000, Sync engine 31.4%->0.6%.
- kda_decode_batch: batched multi-(request,head) decode, up to 15.3x faster per
decode step for a serving batch.
- kda_chunk_step_exact: numerically exact per-channel prefill (16-sub-chunk + WY),
cos_sim >= 0.9999999 in all gate regimes at ~1.05x latency. The scalar-mean
kda_chunk_step drops to 0.22-0.49 cos_sim once gate decay is non-trivial.
- kda_chunk_step_exact_multihead: head-interleaved exact prefill, ~1.38x at NV=4.
All forward-only, parity re-verified on trn2 (Beta 4 + Beta 5).
Add kernels-library metadata.json (neuron backend) for kernel-type repo
kda-kernel project
f1f52b0
v1.3: fused multi-chunk backward
kda-kernel project
ee0b350
Add kda_chunked_fused + raw fused kernels (kda_fused_chunked_fwd/_bwd): all
chunks in one NKI launch per direction, state/dSn carried in SBUF. Same numerics
as kda_chunked (cos_sim ~= 1.0 end-to-end through loss.backward()).
Perf (fwd+bwd, trn2, single b,h): 2.07x @ 2 chunks, 3.90x @ 4, 4.89x @ 8 vs the
Python chunk-loop; speedup grows with sequence length (fused fwd+bwd nearly flat
because per-chunk launch overhead is paid once).
New exports: kda_chunked_fused, kda_fused_chunked_fwd, kda_fused_chunked_bwd.
v1.2: training / backward support
kda-kernel project
07e0355
Add NKI backward kernels + torch.autograd.Function wrappers (kda_recurrent,
kda_chunked) making KDA trainable on Trainium.
New exports:
kda_recurrent, kda_chunked -- differentiable, loss.backward()-ready
kda_recurrent_fwd_v2, kda_chunk_step_v2 -- fwd + saved backward intermediates
kda_recurrent_bwd, kda_chunk_bwd -- raw backward kernels
Parity vs fla-core naive autograd (sim trn2-pinned + trn2 hardware):
recurrent: all 5 grads cos_sim >= 0.9998 (exact)
chunked: dq/dk/dv/dbeta >= 0.9998; dg limited by scalar-mean intra-chunk approx
(self-consistent gradient of our forward, cos_sim=1.0)
Perf (trn2, single b,h): recurrent bwd S=128 2031us; chunked bwd C=128 174us.
NKI fwd+bwd is 9x (recurrent) / 86.6x (chunked) faster than torch.compile,
and compiles in seconds vs 200-585s.
Deferred: fused multi-chunk backward (single NKI launch) for long-seq throughput.
- Fix T-KDA-02: both kernels decayed state per-V column instead of per-K row.
Recurrent parity vs fla-core naive_recurrent_kda 0.99977 -> 1.00000
(max_abs_diff 2.65e-3 -> 3.4e-8). Recurrent kernel also simplified (per-K decay
is a single tensor_scalar vs the old transpose-scale-transpose). Chunked per-K
state decay fixed for cross-chunk carry; forward-output parity unchanged at
0.99988 (scalar-mean intra-chunk approximation is the remaining limit).
- All 6 recurrent backward gradients now match fla autograd at cos_sim=1.0.
- Removed all references to specific consumer models; package is now fully
model-agnostic (attributes only the fla-core KDA algorithm).
Verified on simulator + trn2.3xlarge hardware.
v1.0: KDA NKI kernels for HF Transformers on AWS Trainium
