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.

chunk_kda_reference

KDA chunk oracle with causal per-channel decays and optional operand QDQ.

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.

state_qdq_reference

Dynamic state-tile QDQ with detached scales and identity STE.
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, state_format='fp8_e4m3', w_quantizer=None, matmul=None, arithmetic=None, inverse_fn=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. Optional callbacks expose the eight named matmuls and elementwise boundaries. Matmul callbacks receive conventional lhs @ rhs shapes, without padding.

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)

  • state_format (str)

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

  • matmul (Callable[[str, Tensor, Tensor], Tensor] | None)

  • arithmetic (Callable[[str, Tensor], Tensor] | None)

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

Return type:

tuple[Tensor, Tensor]

chunk_kda_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, state_format='fp8_e4m3', w_quantizer=None, matmul=None, inverse_fn=None)#

KDA chunk oracle with causal per-channel decays and optional operand QDQ.

The eight-site callback receives conventional three-dimensional matmul operands. Interaction sites are invoked one query row at a time with decay already applied to the right-hand keys. This avoids inverse-prefix factors.

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

state_qdq_reference(state, block_v=64, state_format='fp8_e4m3')#

Dynamic state-tile QDQ with detached scales and identity STE.

Parameters:

  • state (Tensor)

  • block_v (int)

  • state_format (str)