prefill#
Batched differentiable GDN prefill with materialized numerical boundaries.
Functions
Execute all eight GDN prefill sites with autograd through their QDQ operators. |
- matmul_gdn(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, cu_seqlens=None, cu_seqlens_cpu=None, state_v_first=False, chunk_size=64, cp_context=None, prefill_lengths=None)#
Execute all eight GDN prefill sites with autograd through their QDQ operators.
This path batches chunk-local matmuls and carries states across chunks without detaching them. It uses FP32 working arithmetic (FP64 for double inputs) and casts outputs back to Q’s dtype. It materializes the inverse and intermediates; it is a numerical-emulation backend, not the fused FLA implementation.
- Parameters:
sites (LinearAttentionMatmulSites)
policy (LinearAttentionConfig)