reference#

Small differentiable oracles, independent of FLA and Triton.

These references prioritize explicit arithmetic and gradients over speed. Gates are natural logarithms; inputs are already normalized/activated. Packed sequence boundaries are read on the CPU. Accumulation follows the input dtype (use float64 for algebra tests).

Functions

chunk_gdn_reference

Exact GDN chunk algebra with optional state/W fake quantization.

recurrent_delta_rule_reference

GDN (scalar decay) or KDA (per-key decay) exact token recurrence.

state_fp8_qdq_reference

Dynamic E4M3 QDQ with identity STE and one scale per [Dk, block_v] tile.
chunk_gdn_reference(q, k, v, g, beta, *, chunk_size=64, scale=None, initial_state=None, cu_seqlens=None, state_v_first=False, state_qdq=False, state_qdq_block_v=64, w_quantizer=None)#

Exact GDN chunk algebra with optional state/W fake quantization.

The solve is a unit-lower triangular solve. w_quantizer sees the complete materialized [B,T,Hv,Dk] WY operand once, with its own autograd semantics. State QDQ occurs on the initial state and each chunk’s final state, after readout.

Parameters:

  • q (Tensor)

  • k (Tensor)

  • v (Tensor)

  • g (Tensor)

  • beta (Tensor)

  • chunk_size (int)

  • scale (float | None)

  • initial_state (Tensor | None)

  • cu_seqlens (Tensor | None)

  • state_v_first (bool)

  • state_qdq (bool)

  • state_qdq_block_v (int)

  • w_quantizer (Callable[[Tensor], Tensor] | None)

Return type:

tuple[Tensor, Tensor]

recurrent_delta_rule_reference(q, k, v, g, beta, *, scale=None, initial_state=None, cu_seqlens=None, state_v_first=False)#

GDN (scalar decay) or KDA (per-key decay) exact token recurrence.

Inputs have shapes q,k:[B,T,Hk,Dk], v:[B,T,Hv,Dv], beta:[B,T,Hv], and g:[B,T,Hv] (GDN) or [B,T,Hv,Dk] (KDA). State is [N,Hv,Dk,Dv], or [N,Hv,Dv,Dk] with state_v_first. Returns output and final state, preserving gradients through the initial state.

Parameters:

  • q (Tensor)

  • k (Tensor)

  • v (Tensor)

  • g (Tensor)

  • beta (Tensor)

  • scale (float | None)

  • initial_state (Tensor | None)

  • cu_seqlens (Tensor | None)

  • state_v_first (bool)

Return type:

tuple[Tensor, Tensor]

state_fp8_qdq_reference(state, block_v=64)#

Dynamic E4M3 QDQ with identity STE and one scale per [Dk, block_v] tile.

state is in key-first layout [..., Dk, Dv]. Scales and rounding do not contribute derivatives. The zero tile uses scale one; partial value tiles are valid.

Parameters:

  • state (Tensor)

  • block_v (int)

Return type:

Tensor