Struct NystromState 

pub struct NystromState { /* private fields */ }

Streaming Nyström attention state for ONE GQA group.

State splits along the GQA grain, because the operator does:

• SHARED per KV group (NystromGroup) — the exact window ring, the sink buffer and the key landmarks K̃. Under GQA every Q head of a group reads the SAME k/v rows, so all three are bit-identical across the group; storing them once per group instead of once per Q head is the point of this split (identical arithmetic, ×heads_per_kv less window memory). K̃ = seg_means(ks, t, d, m_eff) is a pure function of the group’s keys and of t (which fixes m_eff), so it is shareable for the same reason the keys are.
• PRIVATE per Q head (NystromHead) — the far accumulators T̂/Ẑ and their per-landmark running maxima, the QUERY landmarks Q̃, and the mixing matrix M = pinv(exp(Q̃K̃ᵀ/√d)). Q̃ is built from that head’s own queries, so M and the far field it drives are per-Q-head and cannot be shared: the far mass a head accumulates is contracted against its own query landmarks.

Lifecycle: new(m, w, sink) → prefill(prompt) once → step() per decode token (single-head façade), or new_group/prefill_group/ step_group for a whole GQA group at once. All buffers are flat Vec<f32>, row-major; the skeleton path performs no allocations inside step().

Implementations

impl NystromState

pub fn new(m: usize, w: usize, sink: usize) -> Self

Single-head state (heads_per_kv == 1, and the shape the kernel unit tests use).

m — landmark budget (≥ 4; see O1_DEFAULT_M), w — exact window width (validated setting is 128), sink — permanent exact sink keys (validated default is 4; 0 reproduces the sink-free kernel bit-for-bit). Rectifier defaults to O1_DEFAULT_RECT; override with with_rect (the golden-parity test pins it explicitly).

pub fn new_group(m: usize, w: usize, sink: usize, q_heads: usize) -> Self

State for one GQA group of q_heads query heads sharing a KV head. The window/sink/K̃ are stored ONCE for the group; each Q head keeps its own far field, Q̃ and M.

pub fn with_rect(self, rect: O1Rect) -> Self

Select the skeleton rectifier for every head of the group (builder; see O1Rect).

pub fn num_q_heads(&self) -> usize

Query heads in this group.

pub fn far_len(&self, head: usize) -> usize

Keys absorbed into head head’s far field. Exposed for the delayed-insertion invariant test: eviction is a GROUP event, but each head must absorb the evicted key EXACTLY once, so this must equal the number of evictions — never a multiple of it.

pub fn prefill( &mut self, qs: &[f32], ks: &[f32], vs: &[f32], t: usize, d: usize, dv: usize, )

Absorb the whole prompt for a single-head state — see prefill_group.

pub fn prefill_group( &mut self, qs: &[&[f32]], ks: &[f32], vs: &[f32], t: usize, d: usize, dv: usize, )

Absorb the whole prompt for a GQA group: freeze each head’s landmarks and M, then replay the prompt through the step() state semantics (window fill + delayed far insertion). qs[h] is that head’s [t][d] query block; ks is [t][d] and vs is [t][dv] — the group’s shared keys/values, row-major.

pub fn step(&mut self, q: &[f32], k: &[f32], v: &[f32], out: &mut [f32])

One decode step for a single-head state — see step_group.

pub fn step_group( &mut self, q_all: &[f32], k: &[f32], v: &[f32], out_all: &mut [f32], )

One decode step for the whole GQA group. Inserts the group’s (k, v) ONCE, evicting the oldest window key into every head’s far accumulators, then writes each head’s attention output. q_all is [q_heads][d], out_all is [q_heads][dv].

pub fn memory_bytes(&self) -> usize

Heap bytes held by this group’s state (shared window + sinks + K̃, plus each head’s skeleton and scratch) — feeds the honest “KV+state” memory line, same discipline as counting linear_state for the linear core.

Trait Implementations

impl Clone for NystromState

fn clone(&self) -> NystromState

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for NystromState

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

Formats the value using the given formatter.