Commit History
-
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