Commit History
-
v1.7: add batched recurrent backward (kda_recurrent_bwd_batched); metadata version=3
Jim BurtoftCompletes 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.1bb2101
-
docs: document version=2 (pinned) and revision=main (always-latest) load options; note no zero-arg default in kernels lib
Jim Burtoft15b4561
-
v1.6: kernel-repo distribution fix
Jim Burtoft- 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.eacd70d
-
v1.5: exact chunked backward (kda_chunk_step_exact_bwd)
Jim BurtoftGradient 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).6709b3f -
v1.4: forward optimizations (decode key-fold, batched decode, exact prefill, head-interleaved prefill)
Jim Burtoft- 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).5532501
-
README: document kernel-type repo + get_kernel load path
kda-kernel project24e7c33 -
Add kernels-library metadata.json (neuron backend) for kernel-type repo
kda-kernel projectf1f52b0 -
v1.3: fused multi-chunk backward
kda-kernel projectAdd 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.ee0b350 -
v1.2: training / backward support
kda-kernel projectAdd 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.07e0355
-
v1.1: T-KDA-02 state-decay axis fix + fully model-agnostic
kda-kernel project- 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.d8005f6 -
v1.0: KDA NKI kernels for HF Transformers on AWS Trainium
kda-kernel projectModel-agnostic NKI implementation of KDA (Kernel-based Decomposed Attention) linear attention. Primary reference model: upstage/Solar-Open2-250B. v1.0 exposes three raw kernel entry points: - kda_recurrent_fwd (TKG/decode) - kda_recurrent_fwd_state (TKG with final-state emit) - kda_chunk_step (CTE/prefill, C=128) Parity vs fla-core PyTorch reference (trn2.3xlarge LNC=2, 2026-08-05): Recurrent (S=128 random): cos_sim = 0.99977 (bf16 accumulation floor) Chunked (C=128 random): cos_sim = 0.99988 (scalar-mean approx ceiling) Performance (single (b,h), trn2.3xlarge LNC=2): Chunked wall: 87 μs/chunk, MFU 1.19% BF16 / 4.71% FP32 Recurrent wall: 838 μs at S=128 (6.55 μs/token amortized) Includes: - T-KDA-01 accuracy fix (cos_sim 0.78 -> 0.99988 on Solar Open 2 g-scale) - Task 008 -6.7% wall-clock optimization (nc_transpose substitution) Full-layer NeuronKDA(nn.Module) wrapper is planned for v1.1. Requires: - PyTorch Native Beta 3+ (torch-neuronx 2.11.3+) - NKI >= …f145768
-
initial commit
Verified signature9c58c3f