gdn

GDN adapter for serving-aligned recurrent-state QAT.

Functions

gdn_state_qat

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

gdn_state_qat(q, k, v, g, beta, *, policy, state_quantizer=None, read_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, cu_seqlens=None, cu_seqlens_cpu=None, state_v_first=False, chunk_size=64, cp_context=None, prefill_lengths=None, sequence_lengths=None, gate_inputs=None)

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

Validate prepared scalar log gates and promote inputs to FP32 working values. 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. precision="vllm" requires gate_inputs=(raw_g, raw_beta, A_log, bias) for native gate preparation; the prepared g and beta supply the adjoint. replayssm_two_precision uses read_quantizer for FP8 block128 reads, requires an unquantized recurrence anchor, and does not return partial-window state. Its adjoint reconstructs one replay window at a time during backward.

Parameters:

policy (LinearAttentionConfig)