kda

KDA adapter for serving-aligned recurrent-state QAT.

Functions

kda_state_qat

Adapt Megatron's FLA-style KDA call to serving-aligned state QAT.

kda_state_qat(q, k, v, g, beta, *, policy, state_quantizer=None, 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, sequence_lengths=None, gate_inputs=None)

Adapt Megatron’s FLA-style KDA call to serving-aligned state QAT.

Prepare per-key-channel gates and optional beta activation using FLA’s formula. The shared training forward runs the chunked prefix and recurrent suffix with native forward values and configured state QDQ. FLA supplies the prefix backward; a Torch recurrence supplies the suffix adjoint. state_quantizer owns state QDQ; passing None or a disabled quantizer retains the unquantized native path.