Struct LogLinearState 

pub struct LogLinearState { /* private fields */ }

Available on crate feature alloc only.

Hierarchical stack of matrix states, one per active Fenwick level.

Storage is fixed at max_levels slots; each slot is a d_k x d_v matrix (zeros when inactive). The active mask records which slots currently hold a real bucket. The size field counts tokens pushed so far — equivalently, after size = t pushes, the bits of t indicate which levels are active.

Paper reference

Han Guo, Songlin Yang, Tarushii Goel, Eric P. Xing, Tri Dao, Yoon Kim. Log-Linear Attention. ICLR 2026. arXiv:2506.04761, §2-§3.

Implementations

impl LogLinearState

pub fn new(max_levels: usize, d_k: usize, d_v: usize) -> Self

Create a new state with all max_levels matrices zero-initialized.

Panics

Panics in debug mode if max_levels == 0, d_k == 0, or d_v == 0.

pub fn max_levels(&self) -> usize

Hierarchy depth cap (max_levels). Storage is always padded to this size.

pub fn d_k(&self) -> usize

Per-head key dimension.

pub fn d_v(&self) -> usize

Per-head value dimension.

pub fn size(&self) -> u64

Number of tokens pushed so far. Equivalent to t in the paper.

pub fn active_level_count(&self) -> usize

Number of currently active levels = popcount(size).

Always ≤ max_levels. After exhausting capacity (size ≥ 2^max_levels), the highest level absorbs further carries (see push_leaf).

pub fn is_active(&self, level: usize) -> bool

Whether level ℓ currently holds a real bucket.

Panics

Panics in debug mode if level >= max_levels.

pub fn level(&self, level: usize) -> &AttentionState

Borrow level ℓ’s matrix state (zero matrix if inactive).

Panics

Panics in debug mode if level >= max_levels.

pub fn push_leaf(&mut self, k: &[f64], v: &[f64])

Push a new leaf bucket holding the outer product k * v^T, then run carry-propagation upward.

Algorithm (paper §2.1):

1. Set level 0 to k * v^T. If level 0 was already active, the new leaf would collide — but classical Fenwick increment means that case happens iff the previous push produced a carry that did NOT consume level 0. By construction the invariant holds: after every prior push, level 0 is active iff bit 0 of size is set (== size is odd). So before push: level0_active iff size_was_odd. We treat this with standard binary-increment: place the new bucket at level 0 pre-emptively, then run the standard carry loop.

In the paper this is the carry-propagation form of the Fenwick scan; in irithyll terms it’s an in-place rewrite of the level stack, no allocation past max_levels.

Capacity overflow

If a carry would propagate above level max_levels - 1, the excess bucket is folded into the topmost level via matrix addition. This preserves the invariant “total information captured by the Fenwick tree” at the cost of resolution at the very deepest scale — equivalent to the paper’s note that max_levels = ⌊log₂(T_max)⌋ + 1 should be chosen so T_max exceeds the expected stream length.

Arguments

• k — key vector, length d_k.
• v — value vector, length d_v.

pub fn reset(&mut self)

Reset all levels to zero and clear size. After reset, state() returns all zeros and active_level_count() == 0.

pub fn flat_state(&self) -> &[f64]

Flat view of the padded state — concatenation of all max_levels levels in row-major order.

Length is always max_levels * d_k * d_v, regardless of active_level_count(). Inactive levels contribute all-zero blocks. This is the constant-shape contract required by AttentionLayer::state() consumers.

pub fn query_mixed(&self, q: &[f64], lambdas: &[f64], out: &mut [f64])

Compute the λ-weighted readout Σ_ℓ λ_ℓ · q^T · S^(ℓ) over all max_levels slots and write into out (length d_v).

Inactive levels contribute zero (their S^(ℓ) is the zero matrix). The caller supplies lambdas of length max_levels (typically a softplus-softmax mix bounding Σ λ ≤ 1).

Arguments

• q — query vector, length d_k.
• lambdas — per-level non-negative mix weights, length max_levels.
• out — output buffer, length d_v. Overwritten.

Panics

Panics in debug mode if q.len() != d_k, lambdas.len() != max_levels, or out.len() != d_v.

Trait Implementations

impl Clone for LogLinearState

fn clone(&self) -> LogLinearState

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for LogLinearState

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.