kda

Differentiable KDA prefill with stable per-channel decay interactions.

Functions

matmul_kda

Run KDA chunk prefill with actual-operand QDQ and differentiable state carry.

matmul_kda(q, k, v, g, beta, *, sites, policy, w_quantizer, state_qdq=False, state_format='fp8_e4m3', scale=None, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, use_gate_in_kernel=False, use_beta_sigmoid_in_kernel=False, allow_neg_eigval=False, A_log=None, dt_bias=None, safe_gate=False, lower_bound=None, cu_seqlens=None, cu_seqlens_cpu=None, state_v_first=False, chunk_size=64, cp_context=None, disable_recompute=False, return_intermediate_states=False, prefill_lengths=None)

Run KDA chunk prefill with actual-operand QDQ and differentiable state carry.

Gates are per-key-channel natural log retentions. Causal interactions form exp(prefix[i] - prefix[j]) directly, avoiding an overflowing inverse decay. Raw-gate activation, when requested, follows the selected FLA gate formula.