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 device_view(&self) -> O1DeviceView<'_>

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).

Examples found in repository

examples/o1_recovery.rs (line 94)

fn trial(m: usize, w: usize, sink: usize, amp: f32, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    let t = 8 * m + w + depth;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let qhat = unit(&mut s);
    let q: Vec<f32> = qhat.iter().map(|x| x * rd).collect();

    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }

    let mut ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let mut vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();
    for j in 0..t {
        vs[j * dv] = 0.0;
    }
    let p = t - depth;
    for c in 0..d {
        ks[p * d + c] += amp * qhat[c];
    }
    vs[p * dv] = 1.0;

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);

    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let mut v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    v_new[0] = 0.0;
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        let dot: f32 = (0..d).map(|c| q[c] * ks[j * d + c]).sum();
        logits.push(dot / rd);
    }
    let dot: f32 = (0..d).map(|c| q[c] * k_new[c]).sum();
    logits.push(dot / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let (mut den, mut num0) = (0f64, 0f64);
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        num0 += e * if j == t { v_new[0] } else { vs[j * dv] } as f64;
    }
    let want0 = (num0 / den) as f32;
    if want0.abs() < 1e-4 {
        return f32::NAN;
    }
    got[0] / want0
}


/// The background channel on its own: no needle, no signal, just the
/// skeleton's error in estimating a far field of ordinary keys.
///
/// The recovery statistic divides by the exact answer, so it mixes two
/// error sources — the needle's own estimate and the denominator built
/// from every background key. Removing the needle isolates the second,
/// and it is the one that should follow the central-limit law: a sum
/// over ~n far keys approximated from m landmarks.
///
/// Returned: ‖o_o1 − o_exact‖ / ‖o_exact‖ over the whole output vector.
/// An AMPLITUDE, deliberately — its square is the energy, and the two
/// differ by exactly a factor of two in any log-log slope, which is why
/// only one of them needs measuring.
fn background_error(m: usize, w: usize, sink: usize, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    // Sequence length is FIXED, not 8*m + …. Tying it to m — which the
    // recovery sweep above does, to guarantee m_eff reaches m — means a
    // larger budget also gets a longer far field, and the sweep varies
    // two things at once. Here only m moves. `8 * 64` keeps m_eff = m up
    // to the largest budget measured.
    let t = 8 * 64 + w + depth;
    let _ = m;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let q: Vec<f32> = unit(&mut s).iter().map(|x| x * rd).collect();
    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }
    let ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);
    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        logits.push((0..d).map(|c| q[c] * ks[j * d + c]).sum::<f32>() / rd);
    }
    logits.push((0..d).map(|c| q[c] * k_new[c]).sum::<f32>() / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut den = 0f64;
    let mut want = vec![0f64; dv];
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        let src = if j == t { &v_new[..] } else { &vs[j * dv..(j + 1) * dv] };
        for (c, wc) in want.iter_mut().enumerate() {
            *wc += e * src[c] as f64;
        }
    }
    let want: Vec<f32> = want.iter().map(|x| (x / den) as f32).collect();
    let num: f32 = got.iter().zip(&want).map(|(g, x)| (g - x) * (g - x)).sum::<f32>().sqrt();
    let den2: f32 = want.iter().map(|x| x * x).sum::<f32>().sqrt();
    num / den2.max(1e-9)
}

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.

Examples found in repository

examples/o1_recovery.rs (line 95)

fn trial(m: usize, w: usize, sink: usize, amp: f32, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    let t = 8 * m + w + depth;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let qhat = unit(&mut s);
    let q: Vec<f32> = qhat.iter().map(|x| x * rd).collect();

    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }

    let mut ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let mut vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();
    for j in 0..t {
        vs[j * dv] = 0.0;
    }
    let p = t - depth;
    for c in 0..d {
        ks[p * d + c] += amp * qhat[c];
    }
    vs[p * dv] = 1.0;

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);

    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let mut v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    v_new[0] = 0.0;
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        let dot: f32 = (0..d).map(|c| q[c] * ks[j * d + c]).sum();
        logits.push(dot / rd);
    }
    let dot: f32 = (0..d).map(|c| q[c] * k_new[c]).sum();
    logits.push(dot / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let (mut den, mut num0) = (0f64, 0f64);
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        num0 += e * if j == t { v_new[0] } else { vs[j * dv] } as f64;
    }
    let want0 = (num0 / den) as f32;
    if want0.abs() < 1e-4 {
        return f32::NAN;
    }
    got[0] / want0
}


/// The background channel on its own: no needle, no signal, just the
/// skeleton's error in estimating a far field of ordinary keys.
///
/// The recovery statistic divides by the exact answer, so it mixes two
/// error sources — the needle's own estimate and the denominator built
/// from every background key. Removing the needle isolates the second,
/// and it is the one that should follow the central-limit law: a sum
/// over ~n far keys approximated from m landmarks.
///
/// Returned: ‖o_o1 − o_exact‖ / ‖o_exact‖ over the whole output vector.
/// An AMPLITUDE, deliberately — its square is the energy, and the two
/// differ by exactly a factor of two in any log-log slope, which is why
/// only one of them needs measuring.
fn background_error(m: usize, w: usize, sink: usize, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    // Sequence length is FIXED, not 8*m + …. Tying it to m — which the
    // recovery sweep above does, to guarantee m_eff reaches m — means a
    // larger budget also gets a longer far field, and the sweep varies
    // two things at once. Here only m moves. `8 * 64` keeps m_eff = m up
    // to the largest budget measured.
    let t = 8 * 64 + w + depth;
    let _ = m;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let q: Vec<f32> = unit(&mut s).iter().map(|x| x * rd).collect();
    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }
    let ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);
    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        logits.push((0..d).map(|c| q[c] * ks[j * d + c]).sum::<f32>() / rd);
    }
    logits.push((0..d).map(|c| q[c] * k_new[c]).sum::<f32>() / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut den = 0f64;
    let mut want = vec![0f64; dv];
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        let src = if j == t { &v_new[..] } else { &vs[j * dv..(j + 1) * dv] };
        for (c, wc) in want.iter_mut().enumerate() {
            *wc += e * src[c] as f64;
        }
    }
    let want: Vec<f32> = want.iter().map(|x| (x / den) as f32).collect();
    let num: f32 = got.iter().zip(&want).map(|(g, x)| (g - x) * (g - x)).sum::<f32>().sqrt();
    let den2: f32 = want.iter().map(|x| x * x).sum::<f32>().sqrt();
    num / den2.max(1e-9)
}

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.

Examples found in repository

examples/o1_recovery.rs (line 101)

fn trial(m: usize, w: usize, sink: usize, amp: f32, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    let t = 8 * m + w + depth;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let qhat = unit(&mut s);
    let q: Vec<f32> = qhat.iter().map(|x| x * rd).collect();

    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }

    let mut ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let mut vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();
    for j in 0..t {
        vs[j * dv] = 0.0;
    }
    let p = t - depth;
    for c in 0..d {
        ks[p * d + c] += amp * qhat[c];
    }
    vs[p * dv] = 1.0;

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);

    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let mut v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    v_new[0] = 0.0;
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        let dot: f32 = (0..d).map(|c| q[c] * ks[j * d + c]).sum();
        logits.push(dot / rd);
    }
    let dot: f32 = (0..d).map(|c| q[c] * k_new[c]).sum();
    logits.push(dot / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let (mut den, mut num0) = (0f64, 0f64);
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        num0 += e * if j == t { v_new[0] } else { vs[j * dv] } as f64;
    }
    let want0 = (num0 / den) as f32;
    if want0.abs() < 1e-4 {
        return f32::NAN;
    }
    got[0] / want0
}


/// The background channel on its own: no needle, no signal, just the
/// skeleton's error in estimating a far field of ordinary keys.
///
/// The recovery statistic divides by the exact answer, so it mixes two
/// error sources — the needle's own estimate and the denominator built
/// from every background key. Removing the needle isolates the second,
/// and it is the one that should follow the central-limit law: a sum
/// over ~n far keys approximated from m landmarks.
///
/// Returned: ‖o_o1 − o_exact‖ / ‖o_exact‖ over the whole output vector.
/// An AMPLITUDE, deliberately — its square is the energy, and the two
/// differ by exactly a factor of two in any log-log slope, which is why
/// only one of them needs measuring.
fn background_error(m: usize, w: usize, sink: usize, depth: usize, seed: u64) -> f32 {
    let (d, dv) = (64usize, 8usize);
    // Sequence length is FIXED, not 8*m + …. Tying it to m — which the
    // recovery sweep above does, to guarantee m_eff reaches m — means a
    // larger budget also gets a longer far field, and the sweep varies
    // two things at once. Here only m moves. `8 * 64` keeps m_eff = m up
    // to the largest budget measured.
    let t = 8 * 64 + w + depth;
    let _ = m;
    let mut s = seed;
    let rd = (d as f32).sqrt();

    let mut unit = |s: &mut u64| -> Vec<f32> {
        let v: Vec<f32> = (0..d).map(|_| unif(s)).collect();
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter().map(|x| x / n).collect()
    };
    let q: Vec<f32> = unit(&mut s).iter().map(|x| x * rd).collect();
    let mut qs = Vec::with_capacity(t * d);
    for _ in 0..t {
        let u = unit(&mut s);
        qs.extend(u.iter().map(|x| x * rd));
    }
    let ks: Vec<f32> = (0..t * d).map(|_| unif(&mut s) * 1.732).collect();
    let vs: Vec<f32> = (0..t * dv).map(|_| unif(&mut s)).collect();

    let mut st = NystromState::new(m, w, sink);
    st.prefill(&qs, &ks, &vs, t, d, dv);
    let k_new: Vec<f32> = (0..d).map(|_| unif(&mut s) * 1.732).collect();
    let v_new: Vec<f32> = (0..dv).map(|_| unif(&mut s)).collect();
    let mut got = vec![0f32; dv];
    st.step(&q, &k_new, &v_new, &mut got);

    let mut logits = Vec::with_capacity(t + 1);
    for j in 0..t {
        logits.push((0..d).map(|c| q[c] * ks[j * d + c]).sum::<f32>() / rd);
    }
    logits.push((0..d).map(|c| q[c] * k_new[c]).sum::<f32>() / rd);
    let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut den = 0f64;
    let mut want = vec![0f64; dv];
    for (j, &l) in logits.iter().enumerate() {
        let e = ((l - mx) as f64).exp();
        den += e;
        let src = if j == t { &v_new[..] } else { &vs[j * dv..(j + 1) * dv] };
        for (c, wc) in want.iter_mut().enumerate() {
            *wc += e * src[c] as f64;
        }
    }
    let want: Vec<f32> = want.iter().map(|x| (x / den) as f32).collect();
    let num: f32 = got.iter().zip(&want).map(|(g, x)| (g - x) * (g - x)).sum::<f32>().sqrt();
    let den2: f32 = want.iter().map(|x| x * x).sum::<f32>().sqrt();
    num / den2.max(1e-9)
}

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.