lattice-inference 0.7.2

//! Decode-time single-token attention kernel.
//!
//! Single-token (seq_len=1) attention against a full cached K/V sequence.
//! NEON SIMD path accelerates QK dot products, softmax, and V accumulation.

use crate::attention::gqa::GqaConfig;
#[cfg(target_arch = "aarch64")]
use crate::forward::cpu::simd_config;

/// Release-active overflow guard for the decode shape products.
///
/// Every product below feeds either a length assertion or an unsafe NEON
/// pointer offset in the decode kernels. A wrapping `usize` product (e.g. an
/// absurd `num_heads` near `usize::MAX / head_dim`) would make the length
/// checks accept impossible buffers and let the SIMD path read out of bounds.
/// Reject overflow before any unsafe dispatch, matching the FlatKVCache
/// checked-arithmetic precedent.
#[inline]
fn assert_decode_no_overflow(cfg: GqaConfig, kv_seq_len: usize, score_stride: usize) {
    assert!(
        cfg.num_heads.checked_mul(cfg.head_dim).is_some(),
        "decode shape overflow: num_heads * head_dim"
    );
    assert!(
        cfg.num_kv_heads.checked_mul(cfg.head_dim).is_some(),
        "decode shape overflow: num_kv_heads * head_dim"
    );
    assert!(
        cfg.num_heads.checked_mul(score_stride).is_some(),
        "decode shape overflow: num_heads * score_stride"
    );
    assert!(
        kv_seq_len.checked_mul(cfg.kv_dim()).is_some(),
        "decode shape overflow: kv_seq_len * kv_dim"
    );
}

/// **Unstable**: compute decode-time attention scores for one query token.
///
/// Layouts:
/// - `q_buf`: `[num_heads * head_dim]`
/// - `k_buf`: `[kv_seq_len, num_kv_heads * head_dim]`
/// - `scores`: at least `num_heads * score_stride`
///
/// Writes scaled QK scores to `scores[h * score_stride..][..kv_seq_len]`.
/// Panics if shapes are inconsistent or `num_heads` is not divisible by
/// `num_kv_heads`.
pub fn decode_attention_scores(
    q_buf: &[f32],
    k_buf: &[f32],
    scores: &mut [f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    assert!(cfg.num_kv_heads > 0, "num_kv_heads must be > 0");
    assert_eq!(
        cfg.num_heads % cfg.num_kv_heads,
        0,
        "num_heads must be divisible by num_kv_heads"
    );
    assert_decode_no_overflow(cfg, kv_seq_len, score_stride);
    assert_eq!(q_buf.len(), cfg.q_dim(), "q_buf length mismatch");
    assert_eq!(
        k_buf.len(),
        kv_seq_len * cfg.kv_dim(),
        "k_buf length mismatch"
    );
    assert!(
        scores.len() >= cfg.num_heads * score_stride,
        "scores buffer too small"
    );
    assert!(score_stride >= kv_seq_len, "score_stride < kv_seq_len");

    #[cfg(target_arch = "aarch64")]
    {
        if simd_config().neon_enabled {
            unsafe {
                decode_scores_neon(q_buf, k_buf, scores, kv_seq_len, cfg, score_stride);
            }
            return;
        }
    }

    decode_scores_scalar(q_buf, k_buf, scores, kv_seq_len, cfg, score_stride);
}

/// **Unstable**: compute full decode-time attention for one query token.
///
/// Layouts:
/// - `q_buf`: `[num_heads * head_dim]`
/// - `k_buf`: `[kv_seq_len, num_kv_heads * head_dim]`
/// - `v_buf`: `[kv_seq_len, num_kv_heads * head_dim]`
/// - `attn_out`: `[num_heads * head_dim]`
/// - `scores`: at least `num_heads * score_stride`
///
/// Computes scores, applies row softmax in-place to `scores`, then writes
/// `attn_out[h, d] = sum_i softmax(scores[h, i]) * V[i, kv_h, d]`.
/// Panics if shapes are inconsistent or `num_heads` is not divisible by
/// `num_kv_heads`.
pub fn decode_attention(
    q_buf: &[f32],
    k_buf: &[f32],
    v_buf: &[f32],
    attn_out: &mut [f32],
    scores: &mut [f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    assert!(cfg.num_kv_heads > 0, "num_kv_heads must be > 0");
    assert_eq!(
        cfg.num_heads % cfg.num_kv_heads,
        0,
        "num_heads must be divisible by num_kv_heads"
    );
    assert_decode_no_overflow(cfg, kv_seq_len, score_stride);
    assert_eq!(q_buf.len(), cfg.q_dim(), "q_buf length mismatch");
    assert_eq!(attn_out.len(), cfg.q_dim(), "attn_out length mismatch");
    assert!(
        scores.len() >= cfg.num_heads * score_stride,
        "scores buffer too small"
    );

    if kv_seq_len == 0 {
        attn_out.fill(0.0);
        return;
    }

    assert_eq!(
        k_buf.len(),
        kv_seq_len * cfg.kv_dim(),
        "k_buf length mismatch"
    );
    assert_eq!(
        v_buf.len(),
        kv_seq_len * cfg.kv_dim(),
        "v_buf length mismatch"
    );
    assert!(score_stride >= kv_seq_len, "score_stride < kv_seq_len");

    #[cfg(target_arch = "aarch64")]
    {
        if simd_config().neon_enabled {
            unsafe {
                decode_scores_neon(q_buf, k_buf, scores, kv_seq_len, cfg, score_stride);
                softmax_decode_neon(scores, cfg.num_heads, kv_seq_len, score_stride);
                attn_out.fill(0.0);
                accumulate_decode_v_neon(attn_out, scores, v_buf, kv_seq_len, cfg, score_stride);
            }
            return;
        }
    }

    decode_scores_scalar(q_buf, k_buf, scores, kv_seq_len, cfg, score_stride);
    softmax_decode_scores(scores, cfg.num_heads, kv_seq_len, score_stride);
    attn_out.fill(0.0);
    accumulate_decode_v_scalar(attn_out, scores, v_buf, kv_seq_len, cfg, score_stride);
}

fn decode_scores_scalar(
    q_buf: &[f32],
    k_buf: &[f32],
    scores: &mut [f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    let groups = cfg.groups();
    let head_dim = cfg.head_dim;
    let kv_dim = cfg.kv_dim();
    let scale = 1.0f32 / (head_dim as f32).sqrt();

    for kv_h in 0..cfg.num_kv_heads {
        let group_start = kv_h * groups;
        for ki in 0..kv_seq_len {
            let k_off = ki * kv_dim + kv_h * head_dim;
            for gi in 0..groups {
                let h = group_start + gi;
                let q_off = h * head_dim;
                let mut dot = 0.0f32;
                for d in 0..head_dim {
                    dot += q_buf[q_off + d] * k_buf[k_off + d];
                }
                scores[h * score_stride + ki] = dot * scale;
            }
        }
    }
}

fn softmax_decode_scores(
    scores: &mut [f32],
    num_heads: usize,
    kv_seq_len: usize,
    score_stride: usize,
) {
    for h in 0..num_heads {
        let row = &mut scores[h * score_stride..h * score_stride + kv_seq_len];
        let mut m = f32::NEG_INFINITY;
        let mut l = 0.0f32;
        for &s in row.iter() {
            let m_new = m.max(s);
            #[allow(clippy::float_cmp)]
            let alpha = if m == f32::NEG_INFINITY {
                0.0
            } else {
                (m - m_new).exp()
            };
            l = l * alpha + (s - m_new).exp();
            m = m_new;
        }
        // ADR-080 C1: route the fail-closed final decision through the
        // shared row-finalizer (#780). Behavior-preserving: the online
        // (Welford-style) running normalizer `l` here is arithmetically
        // independent of this split — `l` and `m` are already fully
        // resolved from the read-only pass above. Splitting the fused
        // `exp(*s - m) * inv_l` into a separate exp pass then
        // `finalize_row`'s multiply-by-`inv` is the same floating-point
        // operations in the same order (no added/removed rounding step).
        // The added `l.is_finite()` guard (vs. the prior bare `l > 0.0`)
        // does not change reachable behavior: in every non-finite-score
        // scenario reachable here, `l` becomes NaN (not a finite-but-huge
        // or `+inf` value) before this point, so `l > 0.0` and
        // `l.is_finite() && l > 0.0` already agree on every reachable input.
        for s in row.iter_mut() {
            *s = (*s - m).exp();
        }
        crate::attention::softmax_row::finalize_row(row, l);
    }
}

fn accumulate_decode_v_scalar(
    attn_out: &mut [f32],
    scores: &[f32],
    v_buf: &[f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    let groups = cfg.groups();
    let head_dim = cfg.head_dim;
    let kv_dim = cfg.kv_dim();

    for h in 0..cfg.num_heads {
        let kv_h = h / groups;
        let out_off = h * head_dim;
        let score_off = h * score_stride;
        for ki in 0..kv_seq_len {
            let w = scores[score_off + ki];
            let v_off = ki * kv_dim + kv_h * head_dim;
            for d in 0..head_dim {
                attn_out[out_off + d] += w * v_buf[v_off + d];
            }
        }
    }
}

#[cfg(target_arch = "aarch64")]
#[inline]
fn fast_exp_decode(x: f32) -> f32 {
    let x = x.clamp(-87.0, 88.0);
    let val = (12_102_203.0f32 * x + 1_065_353_216.0f32) as i32;
    f32::from_bits(val as u32)
}

#[cfg(target_arch = "aarch64")]
#[inline]
#[target_feature(enable = "neon")]
unsafe fn fast_exp_neon_decode(
    x: std::arch::aarch64::float32x4_t,
) -> std::arch::aarch64::float32x4_t {
    use std::arch::aarch64::*;
    let x = vmaxq_f32(x, vdupq_n_f32(-87.0));
    let x = vminq_f32(x, vdupq_n_f32(88.0));
    let t = vfmaq_f32(vdupq_n_f32(1_065_353_216.0), x, vdupq_n_f32(12_102_203.0));
    vreinterpretq_f32_s32(vcvtq_s32_f32(t))
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn decode_scores_neon(
    q_buf: &[f32],
    k_buf: &[f32],
    scores: &mut [f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    use std::arch::aarch64::*;

    let groups = cfg.groups();
    let head_dim = cfg.head_dim;
    let kv_dim = cfg.kv_dim();
    let scale = 1.0f32 / (head_dim as f32).sqrt();

    let hd_chunks = head_dim / 16;

    for kv_h in 0..cfg.num_kv_heads {
        let group_start = kv_h * groups;
        for gi in 0..groups {
            let h = group_start + gi;
            let q_ptr = q_buf.as_ptr().add(h * head_dim);

            for ki in 0..kv_seq_len {
                let k_ptr = k_buf.as_ptr().add(ki * kv_dim + kv_h * head_dim);

                let mut acc0 = vdupq_n_f32(0.0);
                let mut acc1 = vdupq_n_f32(0.0);
                let mut acc2 = vdupq_n_f32(0.0);
                let mut acc3 = vdupq_n_f32(0.0);

                for c in 0..hd_chunks {
                    let off = c * 16;
                    acc0 = vfmaq_f32(acc0, vld1q_f32(q_ptr.add(off)), vld1q_f32(k_ptr.add(off)));
                    acc1 = vfmaq_f32(
                        acc1,
                        vld1q_f32(q_ptr.add(off + 4)),
                        vld1q_f32(k_ptr.add(off + 4)),
                    );
                    acc2 = vfmaq_f32(
                        acc2,
                        vld1q_f32(q_ptr.add(off + 8)),
                        vld1q_f32(k_ptr.add(off + 8)),
                    );
                    acc3 = vfmaq_f32(
                        acc3,
                        vld1q_f32(q_ptr.add(off + 12)),
                        vld1q_f32(k_ptr.add(off + 12)),
                    );
                }
                let sum = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
                let mut dot = vaddvq_f32(sum);

                for d in (hd_chunks * 16)..head_dim {
                    dot += *q_ptr.add(d) * *k_ptr.add(d);
                }

                scores[h * score_stride + ki] = dot * scale;
            }
        }
    }
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn softmax_decode_neon(
    scores: &mut [f32],
    num_heads: usize,
    kv_seq_len: usize,
    score_stride: usize,
) {
    use std::arch::aarch64::*;

    let chunks = kv_seq_len / 16;

    for h in 0..num_heads {
        let ptr = scores.as_mut_ptr().add(h * score_stride);

        // Pass 1: find max with 4x unroll
        let mut m0 = vdupq_n_f32(f32::NEG_INFINITY);
        let mut m1 = m0;
        let mut m2 = m0;
        let mut m3 = m0;
        for c in 0..chunks {
            let base = c * 16;
            m0 = vmaxq_f32(m0, vld1q_f32(ptr.add(base) as *const f32));
            m1 = vmaxq_f32(m1, vld1q_f32(ptr.add(base + 4) as *const f32));
            m2 = vmaxq_f32(m2, vld1q_f32(ptr.add(base + 8) as *const f32));
            m3 = vmaxq_f32(m3, vld1q_f32(ptr.add(base + 12) as *const f32));
        }
        let vmax = vmaxq_f32(vmaxq_f32(m0, m1), vmaxq_f32(m2, m3));
        let mut max_val = vmaxvq_f32(vmax);
        for i in (chunks * 16)..kv_seq_len {
            let s = *ptr.add(i);
            // `f32::max` returns the finite operand when the other is NaN. That
            // would both drop a tail NaN AND erase an already-NaN chunk max (from
            // the NEON `FMAX` reduction) behind a later finite tail lane. Force
            // the row to fail closed whenever either side is NaN, so NaN
            // propagates across the chunk/tail boundary like the scalar
            // `softmax_decode_scores` path. `+inf`/`-inf` still flow through
            // `max` below (they are not NaN) and are handled by the branch below.
            if max_val.is_nan() || s.is_nan() {
                max_val = f32::NAN;
                break;
            }
            max_val = max_val.max(s);
        }

        // Fail closed on a non-finite row max (e.g. a `+inf` score). The fast-exp
        // path turns `finite - inf = -inf` into a tiny positive value and
        // `inf - inf = NaN` into 0, so `sum` stays small-but-finite-positive and
        // the normalizer would hand all probability mass to finite losers. Emit a
        // zero row instead, matching the scalar `softmax_decode_scores` behavior
        // (`l > 0.0` is false for a non-finite normalizer → `row.fill(0.0)`).
        if !max_val.is_finite() {
            for i in 0..kv_seq_len {
                *ptr.add(i) = 0.0;
            }
            continue;
        }

        // Pass 2: exp(x - max) + accumulate sum
        let vmax_bcast = vdupq_n_f32(max_val);
        let mut s0 = vdupq_n_f32(0.0);
        let mut s1 = s0;
        let mut s2 = s0;
        let mut s3 = s0;
        for c in 0..chunks {
            let base = c * 16;
            let e0 = fast_exp_neon_decode(vsubq_f32(
                vld1q_f32(ptr.add(base) as *const f32),
                vmax_bcast,
            ));
            let e1 = fast_exp_neon_decode(vsubq_f32(
                vld1q_f32(ptr.add(base + 4) as *const f32),
                vmax_bcast,
            ));
            let e2 = fast_exp_neon_decode(vsubq_f32(
                vld1q_f32(ptr.add(base + 8) as *const f32),
                vmax_bcast,
            ));
            let e3 = fast_exp_neon_decode(vsubq_f32(
                vld1q_f32(ptr.add(base + 12) as *const f32),
                vmax_bcast,
            ));
            vst1q_f32(ptr.add(base), e0);
            vst1q_f32(ptr.add(base + 4), e1);
            vst1q_f32(ptr.add(base + 8), e2);
            vst1q_f32(ptr.add(base + 12), e3);
            s0 = vaddq_f32(s0, e0);
            s1 = vaddq_f32(s1, e1);
            s2 = vaddq_f32(s2, e2);
            s3 = vaddq_f32(s3, e3);
        }
        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)));
        for i in (chunks * 16)..kv_seq_len {
            let e = fast_exp_decode(*ptr.add(i) - max_val);
            *ptr.add(i) = e;
            sum += e;
        }

        // Pass 3: normalize. Require a finite positive sum so a finite row max
        // with a `NaN` score (sum becomes `NaN`) also fails closed, matching the
        // scalar path rather than normalizing against `NaN`.
        if sum.is_finite() && sum > 0.0 {
            let vinv = vdupq_n_f32(1.0 / sum);
            for c in 0..chunks {
                let base = c * 16;
                vst1q_f32(
                    ptr.add(base),
                    vmulq_f32(vld1q_f32(ptr.add(base) as *const f32), vinv),
                );
                vst1q_f32(
                    ptr.add(base + 4),
                    vmulq_f32(vld1q_f32(ptr.add(base + 4) as *const f32), vinv),
                );
                vst1q_f32(
                    ptr.add(base + 8),
                    vmulq_f32(vld1q_f32(ptr.add(base + 8) as *const f32), vinv),
                );
                vst1q_f32(
                    ptr.add(base + 12),
                    vmulq_f32(vld1q_f32(ptr.add(base + 12) as *const f32), vinv),
                );
            }
            let inv = 1.0 / sum;
            for i in (chunks * 16)..kv_seq_len {
                *ptr.add(i) *= inv;
            }
        } else {
            for i in 0..kv_seq_len {
                *ptr.add(i) = 0.0;
            }
        }
    }
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn accumulate_decode_v_neon(
    attn_out: &mut [f32],
    scores: &[f32],
    v_buf: &[f32],
    kv_seq_len: usize,
    cfg: GqaConfig,
    score_stride: usize,
) {
    use std::arch::aarch64::*;

    let groups = cfg.groups();
    let head_dim = cfg.head_dim;
    let kv_dim = cfg.kv_dim();
    let hd_chunks = head_dim / 16;

    for h in 0..cfg.num_heads {
        let kv_h = h / groups;
        let out_ptr = attn_out.as_mut_ptr().add(h * head_dim);
        let score_off = h * score_stride;

        for ki in 0..kv_seq_len {
            let w = scores[score_off + ki];
            if w == 0.0 {
                continue;
            }
            let vw = vdupq_n_f32(w);
            let v_ptr = v_buf.as_ptr().add(ki * kv_dim + kv_h * head_dim);

            for c in 0..hd_chunks {
                let off = c * 16;
                vst1q_f32(
                    out_ptr.add(off),
                    vfmaq_f32(
                        vld1q_f32(out_ptr.add(off) as *const f32),
                        vw,
                        vld1q_f32(v_ptr.add(off)),
                    ),
                );
                vst1q_f32(
                    out_ptr.add(off + 4),
                    vfmaq_f32(
                        vld1q_f32(out_ptr.add(off + 4) as *const f32),
                        vw,
                        vld1q_f32(v_ptr.add(off + 4)),
                    ),
                );
                vst1q_f32(
                    out_ptr.add(off + 8),
                    vfmaq_f32(
                        vld1q_f32(out_ptr.add(off + 8) as *const f32),
                        vw,
                        vld1q_f32(v_ptr.add(off + 8)),
                    ),
                );
                vst1q_f32(
                    out_ptr.add(off + 12),
                    vfmaq_f32(
                        vld1q_f32(out_ptr.add(off + 12) as *const f32),
                        vw,
                        vld1q_f32(v_ptr.add(off + 12)),
                    ),
                );
            }
            for d in (hd_chunks * 16)..head_dim {
                *out_ptr.add(d) += w * *v_ptr.add(d);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn small_data(n: usize, seed: u32) -> Vec<f32> {
        (0..n)
            .map(|i| {
                let mut x = (i as u32)
                    .wrapping_mul(seed)
                    .wrapping_add(1_013_904_223 ^ seed.wrapping_mul(0x9E37_79B9));
                x ^= x >> 16;
                x = x.wrapping_mul(0x7FEB_352D);
                x ^= x >> 16;
                ((x & 0x00FF_FFFF) as f32 / 16_777_216.0 - 0.5) * 0.125
            })
            .collect()
    }

    #[test]
    fn test_decode_attention_zero_kvlen() {
        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 4,
            head_dim: 8,
        };
        let q = small_data(cfg.q_dim(), 1);
        let mut out = vec![99.0f32; cfg.q_dim()];
        let mut scores = vec![0.0f32; cfg.num_heads];
        decode_attention(&q, &[], &[], &mut out, &mut scores, 0, cfg, 1);
        assert!(out.iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_decode_attention_mha() {
        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 4,
            head_dim: 8,
        };
        let kv_seq_len = 5;
        let q = small_data(cfg.q_dim(), 1);
        let k = small_data(kv_seq_len * cfg.kv_dim(), 2);
        let v = small_data(kv_seq_len * cfg.kv_dim(), 3);
        let mut out = vec![0.0f32; cfg.q_dim()];
        let mut scores = vec![0.0f32; cfg.num_heads * kv_seq_len];
        decode_attention(
            &q,
            &k,
            &v,
            &mut out,
            &mut scores,
            kv_seq_len,
            cfg,
            kv_seq_len,
        );
        assert!(out.iter().all(|x| x.is_finite()));
    }

    #[test]
    fn test_decode_attention_gqa() {
        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 2,
            head_dim: 8,
        };
        let kv_seq_len = 5;
        let q = small_data(cfg.q_dim(), 1);
        let k = small_data(kv_seq_len * cfg.kv_dim(), 2);
        let v = small_data(kv_seq_len * cfg.kv_dim(), 3);
        let mut out = vec![0.0f32; cfg.q_dim()];
        let mut scores = vec![0.0f32; cfg.num_heads * kv_seq_len];
        decode_attention(
            &q,
            &k,
            &v,
            &mut out,
            &mut scores,
            kv_seq_len,
            cfg,
            kv_seq_len,
        );
        assert!(out.iter().all(|x| x.is_finite()));
    }

    #[test]
    fn test_decode_scores_only() {
        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 2,
            head_dim: 8,
        };
        let kv_seq_len = 3;
        let q = small_data(cfg.q_dim(), 1);
        let k = small_data(kv_seq_len * cfg.kv_dim(), 2);
        let mut scores = vec![0.0f32; cfg.num_heads * kv_seq_len];
        decode_attention_scores(&q, &k, &mut scores, kv_seq_len, cfg, kv_seq_len);
        assert!(scores.iter().all(|x| x.is_finite()));
    }

    #[test]
    #[should_panic(expected = "num_heads must be divisible by num_kv_heads")]
    fn test_non_divisible_heads_panics() {
        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 3,
            head_dim: 8,
        };
        let q = vec![0.0f32; cfg.q_dim()];
        let mut out = vec![0.0f32; cfg.q_dim()];
        let mut scores = vec![0.0f32; 4 * 4];
        decode_attention(&q, &[], &[], &mut out, &mut scores, 0, cfg, 1);
    }

    /// FP-NEON-SOFTMAX: NEON decode path (Schraudolph fast_exp) vs scalar path parity.
    ///
    /// The NEON softmax uses the Schraudolph bit-trick exp (~5% per-element error).
    /// Bias cancels approximately in normalization, so the final attention output
    /// should agree with the scalar path within 5% relative tolerance.
    ///
    /// Tested with adversarial score patterns: uniform, monotonic ramp, one-hot spike.
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_decode_softmax_neon_vs_scalar_parity() {
        // LCG data generator matching the pattern used elsewhere in this file.
        fn lcg_data(n: usize, seed: u32) -> Vec<f32> {
            (0..n)
                .map(|i| {
                    let mut x = (i as u32)
                        .wrapping_mul(seed)
                        .wrapping_add(1_013_904_223 ^ seed.wrapping_mul(0x9E37_79B9));
                    x ^= x >> 16;
                    x = x.wrapping_mul(0x7FEB_352D);
                    x ^= x >> 16;
                    ((x & 0x00FF_FFFF) as f32 / 16_777_216.0 - 0.5) * 2.0
                })
                .collect()
        }

        let cfg = GqaConfig {
            num_heads: 4,
            num_kv_heads: 2,
            head_dim: 32,
        };

        // Adversarial score patterns: uniform, monotonic ramp, one-hot spike.
        let kv_seq_lens: &[(usize, &str)] = &[(16, "uniform"), (32, "ramp"), (8, "one-hot")];

        for &(kv_seq_len, pattern_name) in kv_seq_lens {
            let q = lcg_data(cfg.q_dim(), 42);
            let v = lcg_data(kv_seq_len * cfg.kv_dim(), 44);

            // Construct adversarial Q/K pairs for each score pattern.
            // We run both paths end-to-end from Q/K/V to exercise the full pipeline.
            let (q_patched, k_patched) = match pattern_name {
                "uniform" => {
                    // All-equal scores: uniform Q so all QK dots are equal.
                    let q_eq = vec![1.0f32 / (cfg.head_dim as f32).sqrt(); cfg.q_dim()];
                    let k_eq = vec![1.0f32; kv_seq_len * cfg.kv_dim()];
                    (q_eq, k_eq)
                }
                "ramp" => {
                    // Monotonic ramp: K[i] increases with i so scores are monotone.
                    let q_r = q.clone();
                    let k_r: Vec<f32> = (0..kv_seq_len * cfg.kv_dim())
                        .map(|i| (i / cfg.kv_dim()) as f32 * 0.1)
                        .collect();
                    (q_r, k_r)
                }
                _ => {
                    // One-hot spike: only position 0 has a large K, rest are zero.
                    let q_s = q.clone();
                    let mut k_s = vec![0.0f32; kv_seq_len * cfg.kv_dim()];
                    for d in 0..cfg.kv_dim() {
                        k_s[d] = 10.0;
                    }
                    (q_s, k_s)
                }
            };

            // Run NEON path (dispatches to NEON on aarch64 when neon_enabled).
            let mut out_neon = vec![0.0f32; cfg.q_dim()];
            let mut scores_neon = vec![0.0f32; cfg.num_heads * kv_seq_len];
            decode_attention(
                &q_patched,
                &k_patched,
                &v,
                &mut out_neon,
                &mut scores_neon,
                kv_seq_len,
                cfg,
                kv_seq_len,
            );

            // Run scalar path directly.
            let mut out_scalar = vec![0.0f32; cfg.q_dim()];
            let mut scores_scalar = vec![0.0f32; cfg.num_heads * kv_seq_len];
            decode_scores_scalar(
                &q_patched,
                &k_patched,
                &mut scores_scalar,
                kv_seq_len,
                cfg,
                kv_seq_len,
            );
            softmax_decode_scores(&mut scores_scalar, cfg.num_heads, kv_seq_len, kv_seq_len);
            out_scalar.fill(0.0);
            accumulate_decode_v_scalar(
                &mut out_scalar,
                &scores_scalar,
                &v,
                kv_seq_len,
                cfg,
                kv_seq_len,
            );

            // Assert outputs match within a mixed tolerance budget:
            //   atol=5e-3 (absolute floor; Schraudolph fast_exp ~5% error on [-1, 1] weights
            //              contributes an absolute error floor ≈ max_weight * 5%)
            //   rtol=0.05 (5% relative, accounting for Schraudolph fast_exp error)
            //
            // We use |a-b| <= atol + rtol * max(|a|, |b|) which avoids division-by-zero
            // instability when both values are near zero (e.g. tail positions after softmax).
            let atol = 5e-3f32;
            let rtol = 0.05f32;
            for (i, (&n, &s)) in out_neon.iter().zip(out_scalar.iter()).enumerate() {
                assert!(
                    n.is_finite() && s.is_finite(),
                    "decode_attention pattern={pattern_name} idx={i}: non-finite output neon={n} scalar={s}"
                );
                let abs_err = (n - s).abs();
                let tol = atol + rtol * n.abs().max(s.abs());
                assert!(
                    abs_err <= tol,
                    "decode_attention NEON vs scalar parity failure: pattern={pattern_name} \
                     idx={i} neon={n:.6} scalar={s:.6} abs_err={abs_err:.6} tol={tol:.6}"
                );
            }
        }
    }

    /// A `+inf` score must fail closed (zero row), matching the scalar path,
    /// rather than letting finite-loser fast-exp mass win the softmax. Uses
    /// `kv_seq_len=16` so the NEON path exercises its chunked max-reduction.
    #[test]
    fn test_decode_softmax_pos_inf_score_fails_closed() {
        let cfg = GqaConfig {
            num_heads: 1,
            num_kv_heads: 1,
            head_dim: 1,
        };
        let kv_seq_len = 16;
        let q = vec![1.0e20f32];
        let mut k = vec![0.0f32; kv_seq_len * cfg.kv_dim()];
        k[5] = 1.0e20; // q·k overflows f32 -> +inf score at position 5
        let v: Vec<f32> = (0..kv_seq_len).map(|i| (i as f32 + 1.0) * 10.0).collect();
        let mut out = vec![0.0f32; cfg.q_dim()];
        let mut scores = vec![0.0f32; cfg.num_heads * kv_seq_len];
        decode_attention(
            &q,
            &k,
            &v,
            &mut out,
            &mut scores,
            kv_seq_len,
            cfg,
            kv_seq_len,
        );
        assert!(
            out.iter().all(|x| x.is_finite()),
            "non-finite output: {out:?}"
        );
        assert!(
            out.iter().all(|&x| x == 0.0),
            "expected fail-closed zero row, got {out:?}"
        );
    }

    /// ADR-080 C1: `softmax_decode_scores` (the scalar/portable path, called
    /// directly here rather than through `decode_attention` which always
    /// dispatches to NEON on aarch64) must fail closed on a lone NaN score,
    /// not leave it corrupting the row. A second, unpoisoned head in the
    /// same batched buffer must stay unaffected.
    #[test]
    fn test_decode_softmax_scalar_nan_score_fails_closed() {
        let kv_seq_len = 4;
        let num_heads = 2;
        let score_stride = kv_seq_len;
        let mut scores = vec![0.0f32; num_heads * score_stride];
        // head 0: poisoned with a NaN score.
        scores[0..kv_seq_len].copy_from_slice(&[1.0, f32::NAN, 0.5, -0.2]);
        // head 1: well-formed, must be unaffected.
        scores[score_stride..score_stride + kv_seq_len].copy_from_slice(&[0.1, 0.2, 0.3, 0.05]);

        softmax_decode_scores(&mut scores, num_heads, kv_seq_len, score_stride);

        let head0 = &scores[0..kv_seq_len];
        assert!(
            head0.iter().all(|&v| v == 0.0),
            "NaN-poisoned head must zero out: {head0:?}"
        );
        let head1 = &scores[score_stride..score_stride + kv_seq_len];
        assert!(
            head1.iter().all(|v| v.is_finite()),
            "sibling head must be unaffected: {head1:?}"
        );
        let sum: f32 = head1.iter().sum();
        assert!((sum - 1.0).abs() < 1e-4, "head1 must sum to ~1: {head1:?}");
    }

    /// A `NaN` score in the scalar tail (`kv_seq_len % 16 != 0`) must fail closed
    /// like the scalar path. Rust `f32::max` ignores NaN, so the NEON tail max
    /// scan must detect it explicitly; otherwise the finite chunk lanes normalize
    /// against a finite max and diverge from `softmax_decode_scores`. Driven at
    /// the softmax level (the NaN row is not reachable from a single q·k product).
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_decode_softmax_neon_nan_tail_fails_closed() {
        let kv_seq_len = 17usize; // one 16-wide chunk + a 1-lane scalar tail
        let stride = kv_seq_len;
        let mut row_neon = vec![0.5f32; kv_seq_len];
        row_neon[16] = f32::NAN; // NaN in the scalar tail, finite chunk
        let mut row_scalar = row_neon.clone();
        unsafe {
            softmax_decode_neon(&mut row_neon, 1, kv_seq_len, stride);
        }
        softmax_decode_scores(&mut row_scalar, 1, kv_seq_len, stride);
        assert!(
            row_neon.iter().all(|x| x.is_finite()),
            "neon row non-finite: {row_neon:?}"
        );
        assert!(
            row_neon.iter().all(|&x| x == 0.0),
            "neon tail-NaN row should fail closed, got {row_neon:?}"
        );
        assert!(
            row_scalar.iter().all(|&x| x == 0.0),
            "scalar tail-NaN row should fail closed, got {row_scalar:?}"
        );
    }

    /// A `NaN` inside the 16-wide chunk must not be erased by a later finite
    /// scalar-tail lane (Rust `f32::max` drops NaN). The chunk reduction yields a
    /// NaN `max_val`, which the tail scan must preserve so the row fails closed,
    /// matching `softmax_decode_scores`.
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_decode_softmax_neon_chunk_nan_finite_tail_fails_closed() {
        let kv_seq_len = 17usize; // one 16-wide chunk + a finite 1-lane tail
        let stride = kv_seq_len;
        let mut row_neon = vec![0.5f32; kv_seq_len];
        row_neon[5] = f32::NAN; // NaN in the chunk; index 16 (tail) stays finite
        let mut row_scalar = row_neon.clone();
        unsafe {
            softmax_decode_neon(&mut row_neon, 1, kv_seq_len, stride);
        }
        softmax_decode_scores(&mut row_scalar, 1, kv_seq_len, stride);
        assert!(
            row_neon.iter().all(|x| x.is_finite()),
            "neon row non-finite: {row_neon:?}"
        );
        assert!(
            row_neon.iter().all(|&x| x == 0.0),
            "neon chunk-NaN row should fail closed, got {row_neon:?}"
        );
        assert!(
            row_scalar.iter().all(|&x| x == 0.0),
            "scalar chunk-NaN row should fail closed, got {row_scalar:?}"
        );
    }

    /// A `NaN` inside a 16-wide chunk when `kv_seq_len % 16 == 0` (NO scalar
    /// tail) must still fail closed. This is the subtlest fail-closed case: with
    /// no tail lane, the explicit `is_nan` scan in the tail loop never runs, so
    /// correctness rests entirely on the chunk reduction (`vmaxq_f32` /
    /// `vmaxvq_f32`, ARM `FMAX`) *propagating* the NaN into `max_val`, which the
    /// `!max_val.is_finite()` branch then rejects. A future swap to `FMAXNM`
    /// (`vmaxnmq_f32`, numeric-max that suppresses NaN) would silently normalize
    /// the NaN row against a finite max instead — this test is the regression
    /// guard for that. Must match `softmax_decode_scores` (all-zero row).
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_decode_softmax_neon_chunk_nan_no_tail_fails_closed() {
        let kv_seq_len = 16usize; // exactly one 16-wide chunk, zero scalar tail
        let stride = kv_seq_len;
        let mut row_neon = vec![0.5f32; kv_seq_len];
        row_neon[5] = f32::NAN; // NaN in the chunk; no tail lane exists to catch it
        let mut row_scalar = row_neon.clone();
        unsafe {
            softmax_decode_neon(&mut row_neon, 1, kv_seq_len, stride);
        }
        softmax_decode_scores(&mut row_scalar, 1, kv_seq_len, stride);
        assert!(
            row_neon.iter().all(|x| x.is_finite()),
            "neon row non-finite: {row_neon:?}"
        );
        assert!(
            row_neon.iter().all(|&x| x == 0.0),
            "neon no-tail chunk-NaN row should fail closed, got {row_neon:?}"
        );
        assert!(
            row_scalar.iter().all(|&x| x == 0.0),
            "scalar no-tail chunk-NaN row should fail closed, got {row_scalar:?}"
        );
    }

    /// An absurd `num_heads` makes `num_heads * head_dim` wrap `usize`; the
    /// release-active overflow guard must reject it before any length check or
    /// unsafe NEON dispatch can act on the wrapped (too-small) product.
    #[test]
    #[should_panic(expected = "decode shape overflow")]
    fn test_decode_overflow_guard_panics() {
        let cfg = GqaConfig {
            num_heads: usize::MAX / 8 + 1,
            num_kv_heads: 1,
            head_dim: 8,
        };
        let mut scores: Vec<f32> = Vec::new();
        decode_attention_scores(&[], &[], &mut scores, 1, cfg, 8);
    }
}