decode#

Explicit token-state and encoded-update replay references for QAT.

Classes

EncodedLinearAttentionTensor

Fake decoded values plus detached scale metadata; no compressed storage claim.

LinearAttentionCarry

Explicit anchor/update state, position, and codec contract across calls.

ReplayEntry

An already computed, encoded rank-one update and its log retention.

Functions

recurrent_decode

Run the selected training implementation with explicit differentiable carry.

recurrent_decode_reference

Run one preactivated sequence [T,H,D] with explicit token/replay write events.

class EncodedLinearAttentionTensor#

Bases: object

Fake decoded values plus detached scale metadata; no compressed storage claim.

__init__(values, scales, format, block_v)#
Parameters:
  • values (Tensor)

  • scales (Tensor | None)

  • format (str)

  • block_v (int | None)

Return type:

None

block_v: int | None#
format: str#
scales: Tensor | None#
values: Tensor#
class LinearAttentionCarry#

Bases: object

Explicit anchor/update state, position, and codec contract across calls.

__init__(anchor, entries, position, started, signature, value_basis='identity')#
Parameters:
Return type:

None

anchor: EncodedLinearAttentionTensor#
property cursor#

Number of encoded updates since the last anchor refresh.

entries: tuple[ReplayEntry, ...]#
position: int#
reconstruct(*, original_basis=True)#

Replay entries, returning the original value basis unless explicitly disabled.

signature: str#
started: bool#
value_basis: str = 'identity'#
class ReplayEntry#

Bases: object

An already computed, encoded rank-one update and its log retention.

__init__(key, update, log_retention)#
Parameters:
Return type:

None

key: EncodedLinearAttentionTensor#
log_retention: Tensor#
update: EncodedLinearAttentionTensor#
recurrent_decode(q, k, v, g, beta, *, config, state_qdq=False, state_format='fp8_e4m3', block_v=64, initial_state=None, carry=None, position=0, scale=None, checkpoint_interval=8)#

Run the selected training implementation with explicit differentiable carry.

The fused encode-once path carries the same reconstructed state incrementally and exposes the anchor and encoded updates needed for continuation gradients.

Parameters:

config (LinearAttentionDecodeConfig)

recurrent_decode_reference(q, k, v, g, beta, *, config, state_qdq=False, state_format='fp8_e4m3', block_v=64, initial_state=None, carry=None, position=0, scale=None)#

Run one preactivated sequence [T,H,D] with explicit token/replay write events.

Keys and value heads must already be aligned. Scalar GDN or per-key-channel KDA log gates are accepted. Outputs and all returned carry values retain their graphs. An empty call performs no write or initial-state quantization.

Parameters:

config (LinearAttentionDecodeConfig)