Module kda

Module kda 

Source
Expand description

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.