Module kda 

Kimi K3 KDA CPU reference — a new backend, not GDN / Mamba-2.

Production geometry: head_dim=128, short_conv_kernel_size=4, use_full_rank_gate=true, gate_lower_bound=Some(-5). Decay stays low-rank f_a/f_b; the output gate is full-rank g_proj (unlike GLM-5.3’s g_a/g_b). The 0.40B twin omits gate_lower_bound; HF then runs FLA’s unbounded -exp(A_log)*softplus path (None here).

Recurrence (decode, prenorm q/k):

S <- S * diag(exp(g_t))     // decay on KEY axis, per channel
delta <- (v_t - S^T k_t) * sigmoid(beta_t)
S <- S + k_t ⊗ delta
o_t <- S^T q_t / sqrt(d)

Conv state is [channels, kernel] (FLA ShortConvolution cache W=kernel). Slot 0 is shifted out. beta is a raw logit; the step applies sigmoid.

Structs

KdaConfig

KDA geometry. Tiny dims are legal for CPU tests; production is 128/4.

KdaState

Per-sequence KDA state. Both buffers are FP32, read-modify-write.

Constants

KDA_L2_EPS

FLA use_qk_l2norm_in_kernel eps. CUDA k3_kda_recurrent_step_f32 uses the same.

Functions

bounded_gate

KDA forget-gate in log space (FLA use_gate_in_kernel).

conv_update

Causal depthwise conv + SiLU. Shifts Atlas-width state left, writes x into the last slot, then y[c] = silu(dot(w[c], state[c])).

cuda_kda_enabled

LinearAttention BoundLayer uses CUDA kda_decode unless K3_CUDA_KDA=0. Projections, AttnRes, and MLP stay on the host either way.

full_rank_output_gate

Full-rank output gate: sigmoid(g) ⊙ RMSNorm(core) per head.

kda_decode_token

One decode token: conv update then recurrent step.

kda_from

Map parsed ModelConfig onto KDA geometry.

kda_recurrent_step

One-token KDA core. qkv is post-conv [3 * qkv_dim] (q|k|v). Updates state.recurrent in place. q/k are L2-normalised here.