Yesterday
-
README and docstrings: correct KDA name, update install guidance, fix cross-reference
-
v1.0: KDA NKI kernels for HF Transformers on AWS Trainium
kda-kernel project
Model-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 >= …
