Module kda 

Kimi K3’s KDA (Kimi Delta Attention): a gated delta-rule linear attention mechanism used on the majority of Kimi K3’s layers (69 of 93, per AttentionKind::KimiHybrid), interleaved with Gated MLA (ferrox_models::mla) on the remainder.

Transcribed directly from real reference source fetched live (not guessed, not derived by analogy to other gated-linear-attention designs):

• moonshotai/Kimi-K3’s modeling_kimi_linear.py, KimiDeltaAttention (q/k/v projections, short causal convolutions, decay-gate and beta projections, output gate, RMSNormGated).
• fla-org/flash-linear-attention’s fla/ops/kda/naive.py (naive_recurrent_kda) for the core state recurrence: per position, decay the state by exp(g), add a rank-1 correction beta * k ⊗ (v - kᵀS), read the output as qᵀS.
• That same project’s fla/ops/kda/fused_recurrent.py Triton kernel source (the kernel actually invoked in decode, read directly since naive_recurrent_kda takes its inputs pre-transformed) for the exact preprocessing: L2-normalize q/k (eps=1e-6) before the recurrence, then scale q by head_dim^-0.5.
• fla’s fla/modules/conv/short_conv.py (ShortConvolution: a depthwise causal Conv1d, padding=kernel_size-1, no bias) and fla/modules/fused_norm_gate.py (FusedRMSNormGated) for the short causal conv and output-gate formulas.

Two real, non-obvious facts confirmed by reading source rather than assuming standard conventions:

1. The decay gate g is per-(head, key-dim), not one scalar per head: g = gate_lower_bound * sigmoid(exp(A_log) * (raw_g + dt_bias)), with A_log/gate_lower_bound per-head but raw_g/dt_bias per-(head, dim). State decay S *= exp(g) is applied per key dimension, broadcast across the value dimension (S is [head_dim, head_dim] per head here, since KDA’s K and V head dims are both head_dim — unlike Gated MLA, where they differ).
2. FusedRMSNormGated’s output-gate activation is sigmoid, not silu/swish (the more common choice in gated-linear-attention literature) — confirmed directly by KimiDeltaAttention.__init__ passing activation='sigmoid' explicitly.

KdaConfig::use_full_rank_gate is true for Kimi K3’s real configuration, so only that output-gate path (g_proj projecting hidden_size -> num_heads*head_dim directly) is implemented; the real reference’s low-rank alternative (g_a_proj/g_b_proj) is unused by Kimi K3 and intentionally not implemented here.

Not yet wired into Decoder’s forward pass. Tested here against synthetic weights, cross-validated against an independent Python transcription of the same real algorithm run one position at a time (matching this module’s incremental decode API) over a 5-position sequence — long enough to exercise the short causal conv’s full window (short_conv_kernel_size = 4) past its initial zero-padded steps.

Structs

KdaAttnWeights

KdaState

Per-layer decode-time state: the short causal convs’ recent-input history (up to short_conv_kernel_size - 1 raw projected vectors each) and the recurrent state S ([n_heads, head_dim, head_dim], flattened, zero-initialized) – fundamentally different from ferrox_core::cache::KvCache’s growing K/V history, since KDA’s per-layer state is fixed-size regardless of sequence length.

Functions

kda_forward_token

One decode step.