//! QSAIndexer shared projection/preparation/cache prefix from qwen4_exp.py.
//! Original compiled/fused/eager routing shares one persistent cache and core.
use super::array::{Array, Dtype};
use super::binary::{Binary, binary};
use super::cache::QsaCache;
use super::mlp::Linear;
use super::qsa_core::{Config, Core, Source};
use super::stream::Stream;
use super::{indexing, ops, qsa_prepare, reduce, rope};

pub(super) struct Indexer<'a> {
    pub(super) projection: Option<Linear<'a>>,
    pub(super) heads: i32,
    pub(super) kv_heads: i32,
    pub(super) head_dim: i32,
    pub(super) block_topk: i32,
    pub(super) ratio: i32,
    pub(super) q_norm: &'a Array,
    pub(super) k_norm: &'a Array,
    pub(super) inv_freq: &'a Array,
    pub(super) mrope_axes: Option<&'a Array>,
    pub(super) eps: f64,
    pub(super) scaling: f64,
    pub(super) scratch: u64,
    pub(super) compiled: Option<(Vec<usize>, Core<'a>)>,
}

pub(super) struct PreparedRows {
    pub(super) query: Array,
    pub(super) pooled: Array,
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) total: i32,
}

#[cfg(test)]
#[path = "issue107_indexer.rs"]
mod issue107;

impl<'a> Indexer<'a> {
    #[expect(
        clippy::too_many_arguments,
        reason = "original static compiled-route predicate"
    )]
    fn compiled_supported(
        &self,
        source: &Array,
        cache: &QsaCache,
        pos: i32,
        supplied: bool,
        decode: bool,
        mode: super::QsaCompiledMode,
        options: &super::QsaIndexerOptions,
        stream: Stream,
    ) -> bool {
        use super::{QsaCompiledMode, stream::Device};
        let s = source.layout();
        if !options.fused
            || !options.compiled
            || stream.device() != Device::Gpu
            || s.shape().len() != 3
            || s.shape()[0] != 1
            || s.shape()[1] <= 0
            || self.kv_heads != 1
        {
            return false;
        }
        let rows = s.shape()[1];
        if !decode
            && ((mode == QsaCompiledMode::UpdateOnly && options.phase_prefill)
                || (rows as u32 >= options.prefill_min_rows
                    && (!matches!(
                        mode,
                        QsaCompiledMode::PrefillBlocks | QsaCompiledMode::UpdateOnly
                    ) || rows as u32 != options.prefill_compile_rows)))
        {
            return false;
        }
        if cache.ratio != self.ratio
            || pos != cache.kv.offset
            || !matches!(s.dtype(), Dtype::BF16 | Dtype::F16 | Dtype::F32)
            || !(1..=128).contains(&self.head_dim)
            || !(1..=512).contains(&self.block_topk)
            || self.ratio <= 0
        {
            return false;
        }
        let freq = self.inv_freq.layout();
        if freq.shape().len() != 1
            || freq.dtype() != Dtype::F32
            || freq.shape()[0] <= 0
            || i64::from(freq.shape()[0]) * 2 > i64::from(self.head_dim)
            || [self.q_norm, self.k_norm]
                .iter()
                .any(|a| a.layout().shape() != [self.head_dim] || a.layout().dtype() != s.dtype())
        {
            return false;
        }
        let width = (i64::from(self.heads) + i64::from(self.kv_heads)) * i64::from(self.head_dim);
        if supplied {
            if i64::from(s.shape()[2]) != width {
                return false;
            }
        } else {
            let Some(projection) = self.projection else {
                return false;
            };
            let weight = projection.weight().layout();
            if weight.shape().len() != 2 || i64::from(weight.shape()[0]) != width {
                return false;
            }
            let same_dtype = match projection {
                Linear::Dense(_) => weight.dtype() == s.dtype(),
                Linear::Quantized(a) => {
                    a.scales.layout().dtype() == s.dtype() && a.biases.layout().dtype() == s.dtype()
                }
            };
            if !same_dtype {
                return false;
            }
        }
        for backing in [&cache.raw_keys, &cache.pooled].into_iter().flatten() {
            let a = backing.borrow();
            let l = a.layout();
            if l.shape().len() != 3
                || l.shape()[0] != 1
                || l.shape()[2] != self.head_dim
                || l.dtype() != s.dtype()
            {
                return false;
            }
        }
        if pos > 0
            && cache
                .raw_keys
                .as_ref()
                .is_none_or(|a| a.borrow().layout().shape()[1] < pos)
        {
            return false;
        }
        if let Some(pool) = &cache.pooled {
            if !(0..=pool.borrow().layout().shape()[1]).contains(&cache.pooled_len) {
                return false;
            }
        } else if cache.pooled_len != 0 {
            return false;
        }
        let logical = (i64::from(pos) + i64::from(rows)) / i64::from(self.ratio);
        if logical - i64::from(cache.pooled_len).min(logical)
            > (i64::from(rows) + i64::from(self.ratio) - 1) / i64::from(self.ratio)
        {
            return false;
        }
        if decode
            && !matches!(mode, QsaCompiledMode::UpdateOnly | QsaCompiledMode::Blocks)
            && options.gather_decode
        {
            return false;
        }
        matches!(
            mode,
            QsaCompiledMode::UpdateOnly | QsaCompiledMode::PrefillBlocks
        ) || s.dtype() != Dtype::F32
            || super::sparse_attention::nax_available()
    }

    fn selector_supported(&self, q: &Array, p: &Array, prefill: bool, stream: Stream) -> bool {
        let q = q.layout();
        let p = p.layout();
        stream.device() == super::stream::Device::Gpu
            && [q.dtype(), p.dtype()]
                .iter()
                .all(|d| matches!(d, Dtype::BF16 | Dtype::F16 | Dtype::F32))
            && q.shape().len() == 4
            && p.shape().len() == 3
            && q.shape()[0] == 1
            && p.shape()[0] == 1
            && q.shape()[1] > i32::from(prefill)
            && q.shape()[2] > 0
            && q.shape()[3] > 0
            && p.shape()[1] > 0
            && p.shape()[2] == q.shape()[3]
            && (1..=512).contains(&self.block_topk)
            && self.ratio > 0
            && (prefill
                || (q.dtype() != Dtype::F32 && p.dtype() != Dtype::F32)
                || super::sparse_attention::nax_available())
    }

    /// QSAIndexer.__call__/_call_rows. Eligibility alone chooses a lane;
    /// errors after selection propagate, never rerun through an eager fallback.
    #[expect(
        clippy::too_many_arguments,
        reason = "original indexer inputs and evaluator"
    )]
    pub(super) fn call(
        &mut self,
        hidden: &Array,
        pos: i32,
        cache: &mut QsaCache,
        qk_rows: Option<&Array>,
        options: &super::QsaIndexerOptions,
        streams: &super::stream::Streams,
        vision: Option<(Option<&Array>, i32)>,
        stream: Stream,
    ) -> Result<(Option<super::attention::EagerSelection>, &'static str), String> {
        use super::{
            QsaSelectMode,
            attention::{self, EagerSelection},
            qsa_prefill,
        };
        let shape = hidden.layout().shape().to_vec();
        if shape.len() != 3
            || shape[0] != 1
            || shape[1] <= 0
            || pos < 0
            || self.ratio <= 0
            || self.block_topk <= 0
        {
            return Err("QSA serves a positive single sequence with valid geometry".into());
        }
        let rows = shape[1];
        let decode = rows == 1;
        let total = pos.checked_add(rows).ok_or("QSA position overflow")?;
        let mode = options.compiled_mode(
            rows as u32,
            total as u32,
            self.ratio as u32,
            self.block_topk as u32,
        );
        let vision = vision.filter(|_| {
            self.mrope_axes.is_some() && crate::engine::qwen_vision::cache_policy() != 0
        });
        if vision.is_none()
            && self.compiled_supported(
                qk_rows.unwrap_or(hidden),
                cache,
                pos,
                qk_rows.is_some(),
                decode,
                mode,
                options,
                stream,
            )
        {
            return self
                .call_compiled(
                    hidden,
                    pos,
                    cache,
                    qk_rows,
                    mode,
                    options.prefill_workspace,
                    stream,
                )
                .map(|a| (a, "compiled"));
        }
        let Some(prepared) =
            self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, vision, stream)?
        else {
            return Ok((None, "update_only"));
        };
        let q = prepared.query;
        let pooled = prepared.pooled;
        let logical = pooled.layout().shape()[1];
        let large_prefill = !decode && options.large_prefill(rows as u32, pos as u32);
        if large_prefill && self.selector_supported(&q, &pooled, true, stream) {
            let mut values = qsa_prefill::select(
                &q,
                &pooled,
                [pos, total, logical].map(qsa_prepare::Position::Host),
                self.block_topk,
                self.ratio,
                None,
                QsaSelectMode::Blocks,
                options.prefill_workspace,
                stream,
            )?
            .into_iter();
            attention::prefill_count("metal_selector")?;
            return Ok((
                Some(EagerSelection::PrefillBlocks(
                    values.next().ok_or("QSA missing prefill ids")?,
                    values.next().ok_or("QSA missing prefill validity")?,
                )),
                "metal_prefill",
            ));
        }
        let backing = cache.pooled.as_ref().map(|a| a.borrow().clone());
        let fused = options.fused
            && (decode || (rows as u32) < options.prefill_min_rows)
            && !(decode && options.gather_decode)
            && backing
                .as_ref()
                .is_some_and(|p| self.selector_supported(&q, p, false, stream));
        if !fused {
            return attention::select_eager(
                &q,
                pos,
                cache,
                total,
                self.block_topk,
                i32::try_from(options.tile_rows).map_err(|_| "QSA tile overflow")?,
                &super::QsaEagerOptions {
                    large_prefill,
                    flash: options.flash,
                    gather_decode: options.gather_decode,
                    gather: options.gather,
                    gather_max_rows: options.gather_max_rows,
                    gather_min_context: options.gather_min_context,
                },
                streams,
                stream,
            )
            .map(|a| (Some(a), "eager"));
        }
        let mode = if decode && options.flash {
            QsaSelectMode::Blocks
        } else if !decode
            && options.gather
            && rows as u32 <= options.gather_max_rows
            && total as u32 >= options.gather_min_context
        {
            QsaSelectMode::RowTokens
        } else {
            QsaSelectMode::DenseMask
        };
        let mut values = self
            .select_fused(
                &q,
                pos,
                backing.as_ref().unwrap(),
                total,
                logical,
                mode,
                stream,
            )?
            .into_iter();
        let first = values.next().ok_or("QSA missing fused selection")?;
        let out = match mode {
            QsaSelectMode::Blocks => {
                let width = first.layout().shape()[1];
                EagerSelection::FlashBlocks(
                    ops::squeeze(
                        &ops::slice(&first, &[0, 0], &[1, width], &[1, 1], stream)?,
                        Some(&[0]),
                        stream,
                    )?,
                    ((pos + 1) / self.ratio) * self.ratio,
                )
            }
            QsaSelectMode::RowTokens => {
                EagerSelection::RowTokens(first, values.next().ok_or("QSA missing fused validity")?)
            }
            QsaSelectMode::DenseMask => EagerSelection::Dense(first),
        };
        Ok((Some(out), "legacy_fused"))
    }

    /// Identity seal is of captured parameter descriptors, not their bytes,
    /// call shapes or environment. Replacing a leaf resets the whole bank.
    fn compiled_core(&mut self, prefill_workspace: u64) -> Result<&mut Core<'a>, String> {
        let mut signature = vec![self.q_norm.id(), self.k_norm.id(), self.inv_freq.id()];
        if let Some(projection) = self.projection {
            signature.push(projection.weight().id());
        }
        if let Some(Linear::Quantized(layer)) = self.projection {
            signature.extend([layer.scales.id(), layer.biases.id()]);
        }
        if self
            .compiled
            .as_ref()
            .is_none_or(|(old, _)| *old != signature)
        {
            let core = Core::new(
                Config {
                    heads: self.heads,
                    kv_heads: self.kv_heads,
                    dim: self.head_dim,
                    topk: self.block_topk,
                    ratio: self.ratio,
                    eps: self.eps,
                    scale: self.scaling,
                    scratch: self.scratch,
                    prefill_workspace,
                    ..Config::default()
                },
                [
                    self.q_norm.clone(),
                    self.k_norm.clone(),
                    self.inv_freq.clone(),
                ],
                self.projection,
            )?;
            self.compiled = Some((signature, core));
        }
        Ok(&mut self.compiled.as_mut().unwrap().1)
    }

    /// Original _call_rows_compiled: reserve first, commit explicit outputs
    /// without evaluating frontiers. Attention alone advances kv.offset.
    #[expect(
        clippy::too_many_arguments,
        reason = "original compiled indexer inputs and resolved workspace"
    )]
    pub(super) fn call_compiled(
        &mut self,
        hidden: &Array,
        pos: i32,
        cache: &mut QsaCache,
        qk_rows: Option<&Array>,
        mode: super::QsaCompiledMode,
        prefill_workspace: u64,
        stream: Stream,
    ) -> Result<Option<super::attention::EagerSelection>, String> {
        use super::{
            QsaCompiledMode,
            attention::{self, EagerSelection},
            cache::object,
        };
        use qsa_prepare::Position;
        let rows = hidden
            .layout()
            .shape()
            .get(1)
            .copied()
            .ok_or("QSA hidden rank")?;
        let total = pos.checked_add(rows).ok_or("QSA total overflow")?;
        if self.ratio <= 0 {
            return Err("QSA ratio must be positive".into());
        }
        let logical = total / self.ratio;
        let value = qk_rows.unwrap_or(hidden);
        let (raw, pool) = cache.ensure_compiled_backings(
            value.layout().dtype(),
            pos,
            rows,
            self.head_dim,
            stream,
        )?;
        let out = self.compiled_core(prefill_workspace)?.select(
            if qk_rows.is_some() {
                Source::QkRows
            } else {
                Source::Hidden
            },
            value,
            &raw.borrow(),
            &pool.borrow(),
            [pos, total, logical, cache.pooled_len.min(logical)].map(Position::Host),
            mode,
            None,
            stream,
        )?;
        cache.raw_keys = Some(object(out.raw_keys));
        cache.pooled = Some(object(out.pooled));
        cache.pooled_len = logical;
        cache.pooled_f32_t = None;
        let mut selected = out.selection.into_iter();
        Ok(Some(match mode {
            QsaCompiledMode::UpdateOnly => return Ok(None),
            QsaCompiledMode::Blocks => {
                let ids = selected.next().ok_or("QSA missing block ids")?;
                let mut end = ids.layout().shape().to_vec();
                end[0] = 1;
                let first =
                    ops::slice(&ids, &vec![0; end.len()], &end, &vec![1; end.len()], stream)?;
                EagerSelection::FlashBlocks(
                    ops::squeeze(&first, Some(&[0]), stream)?,
                    ((pos + 1) / self.ratio) * self.ratio,
                )
            }
            QsaCompiledMode::PrefillBlocks => {
                attention::prefill_count("compiled_selector")?;
                EagerSelection::PrefillBlocks(
                    selected.next().ok_or("QSA missing prefill ids")?,
                    selected.next().ok_or("QSA missing prefill validity")?,
                )
            }
            QsaCompiledMode::RowTokens => EagerSelection::RowTokens(
                selected.next().ok_or("QSA missing row ids")?,
                selected.next().ok_or("QSA missing row validity")?,
            ),
            QsaCompiledMode::DenseMask => {
                let mask = selected.next().ok_or("QSA missing dense mask")?;
                EagerSelection::Dense(ops::slice(
                    &mask,
                    &[0; 4],
                    &[1, 1, rows, total],
                    &[1; 4],
                    stream,
                )?)
            }
        }))
    }

    /// Original _select_fused: chunk query rows (32MiB score scratch), never
    /// logical history; keep dense specialization at backing capacity+tail.
    #[expect(
        clippy::too_many_arguments,
        reason = "original selector mode and frontier arguments"
    )]
    pub(super) fn select_fused(
        &self,
        q: &Array,
        pos: i32,
        pooled: &Array,
        total: i32,
        logical: i32,
        mode: super::QsaSelectMode,
        stream: Stream,
    ) -> Result<Vec<Array>, String> {
        use super::{QsaSelectMode, qsa_select};
        let rows = q.layout().shape().get(1).copied().ok_or("QSA query rank")?;
        let blocks = pooled
            .layout()
            .shape()
            .get(1)
            .copied()
            .ok_or("QSA pooled rank")?;
        if rows <= 0 {
            return Err("QSA fused rows must be positive".into());
        }
        let chunk_rows = u64::try_from(rows)
            .unwrap()
            .min((self.scratch / (blocks.max(1) as u64 * 4)).max(1))
            as i32;
        let capacity = if mode == QsaSelectMode::DenseMask {
            Some(
                blocks
                    .checked_add(1)
                    .and_then(|n| n.checked_mul(self.ratio))
                    .ok_or("QSA dense capacity overflow")?,
            )
        } else {
            None
        };
        let mut chunks = Vec::new();
        for start in (0..rows).step_by(chunk_rows as usize) {
            let mut end = q.layout().shape().to_vec();
            end[1] = start.saturating_add(chunk_rows).min(rows);
            let q = ops::slice(q, &[0, start, 0, 0], &end, &[1; 4], stream)?;
            chunks.push(qsa_select::select(
                &q,
                pooled,
                [
                    qsa_prepare::Position::Host(
                        pos.checked_add(start).ok_or("QSA position overflow")?,
                    ),
                    qsa_prepare::Position::Host(total),
                    qsa_prepare::Position::Host(logical),
                ],
                self.block_topk,
                self.ratio,
                capacity,
                mode,
                stream,
            )?);
        }
        let mut outputs = if chunks.len() == 1 {
            chunks.pop().unwrap()
        } else {
            (0..chunks[0].len())
                .map(|leaf| {
                    ops::concatenate(
                        &chunks.iter().map(|c| c[leaf].clone()).collect::<Vec<_>>(),
                        if mode == QsaSelectMode::DenseMask {
                            2
                        } else {
                            0
                        },
                        stream,
                    )
                })
                .collect::<Result<Vec<_>, _>>()?
        };
        if mode == QsaSelectMode::DenseMask {
            outputs[0] = ops::slice(&outputs[0], &[0; 4], &[1, 1, rows, total], &[1; 4], stream)?;
        }
        Ok(outputs)
    }

    fn prepare_queries(
        &self,
        q: &Array,
        pos: i32,
        fused: bool,
        vision: Option<(Option<&Array>, i32)>,
        stream: Stream,
    ) -> Result<Array, String> {
        if vision.is_none()
            && fused
            && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream)
        {
            return qsa_prepare::queries(
                q,
                self.q_norm,
                self.inv_freq,
                qsa_prepare::Position::Host(pos),
                self.eps,
                self.scaling,
                stream,
            );
        }
        let q = ops::rms_norm(q, Some(self.q_norm), self.eps as f32, stream)?;
        let positions = indexing::arange_range_i32(pos, pos + q.layout().shape()[1], stream)?;
        let (cos, sin) = match (vision, self.mrope_axes) {
            (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
                &positions,
                table,
                delta,
                self.inv_freq,
                axes,
                self.scaling,
                stream,
            )?,
            _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
        };
        rope::apply_partial(&q, &cos, &sin, stream)
    }

    fn pool_keys(
        &self,
        fresh: &Array,
        old: i32,
        total: i32,
        fused: bool,
        vision: Option<(Option<&Array>, i32)>,
        stream: Stream,
    ) -> Result<Array, String> {
        if vision.is_none()
            && fused
            && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream)
        {
            return qsa_prepare::pool(
                fresh,
                self.k_norm,
                self.inv_freq,
                qsa_prepare::Position::Host(old),
                self.ratio,
                self.eps,
                self.scaling,
                stream,
            );
        }
        let fresh = ops::reshape(fresh, &[1, total - old, self.ratio, self.head_dim], stream)?;
        let pooled = reduce::mean_f32(
            &ops::astype(&fresh, Dtype::F32, false, stream)?,
            &[2],
            false,
            stream,
        )?;
        let pooled = ops::astype(&pooled, fresh.layout().dtype(), false, stream)?;
        let pooled = ops::rms_norm(&pooled, Some(self.k_norm), self.eps as f32, stream)?;
        let ratio = Array::new(
            &[],
            Dtype::I32,
            super::scalar_buffer(&self.ratio.to_le_bytes())?,
        )?;
        let positions = binary(
            &indexing::arange_range_i32(old, total, stream)?,
            &ratio,
            Binary::Multiply,
            stream,
        )?;
        let (cos, sin) = match (vision, self.mrope_axes) {
            (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
                &positions,
                table,
                delta,
                self.inv_freq,
                axes,
                self.scaling,
                stream,
            )?,
            _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
        };
        ops::squeeze(
            &rope::apply_partial(
                &ops::expand_dims(&pooled, &[2], stream)?,
                &cos,
                &sin,
                stream,
            )?,
            Some(&[2]),
            stream,
        )
    }

    pub(super) fn extend_pooled(
        &self,
        cache: &mut QsaCache,
        total: i32,
        fused: bool,
        vision: Option<(Option<&Array>, i32)>,
        stream: Stream,
    ) -> Result<Option<Array>, String> {
        let nb = total / self.ratio;
        let old = cache.pooled_len.min(nb);
        if nb > old {
            let raw = cache
                .raw_keys
                .as_ref()
                .ok_or("QSA pooling requires raw cache")?
                .borrow();
            let fresh = ops::slice(
                &raw,
                &[0, old * self.ratio, 0],
                &[1, nb * self.ratio, self.head_dim],
                &[1; 3],
                stream,
            )?;
            drop(raw);
            cache.write_pooled(
                &self.pool_keys(&fresh, old, nb, fused, vision, stream)?,
                old,
                nb,
                stream,
            )?;
        }
        if nb == 0 {
            return Ok(None);
        }
        let pooled = cache
            .pooled
            .as_ref()
            .ok_or("QSA pooled cache missing")?
            .borrow();
        ops::slice(&pooled, &[0; 3], &[1, nb, self.head_dim], &[1; 3], stream).map(Some)
    }

    /// Common noncompiled _call_rows prefix. None means the complete-block
    /// budget fits: cache updates still happen, but lazy query work is unused.
    /// Cache.kv.offset is advanced by Attention, never by this indexer.
    #[expect(
        clippy::too_many_arguments,
        reason = "original preparation inputs and request positions"
    )]
    pub(super) fn prepare_rows(
        &self,
        hidden: &Array,
        pos: i32,
        cache: &mut QsaCache,
        qk_rows: Option<&Array>,
        fused: bool,
        vision: Option<(Option<&Array>, i32)>,
        stream: Stream,
    ) -> Result<Option<PreparedRows>, String> {
        let shape = hidden.layout().shape().to_vec();
        if shape.len() != 3
            || shape[0] != 1
            || shape[1] <= 0
            || pos < 0
            || cache.kv.offset != pos
            || self.heads <= 0
            || self.head_dim <= 0
            || self.ratio <= 0
            || self.block_topk <= 0
            || cache.ratio != self.ratio
        {
            return Err("invalid QSA indexer dimensions/frontier".into());
        }
        let total = pos.checked_add(shape[1]).ok_or("QSA position overflow")?;
        let width = self
            .heads
            .checked_mul(self.head_dim)
            .ok_or("QSA query width overflow")?;
        let full_width = width
            .checked_add(self.head_dim)
            .ok_or("QSA projection width overflow")?;
        let qk = match qk_rows {
            Some(rows) => rows.clone(),
            None => self
                .projection
                .ok_or("QSA hidden source requires a projection")?
                .apply(hidden, stream)?,
        };
        if qk.layout().shape() != [1, shape[1], full_width] {
            return Err("QSA projection must have (heads+1)*head_dim columns".into());
        }
        let parts = ops::split(&qk, &[width], -1, stream)?;
        let query = ops::reshape(&parts[0], &[1, shape[1], self.heads, self.head_dim], stream)?;
        let keys = ops::reshape(&parts[1], &[1, shape[1], self.head_dim], stream)?;
        let query = self.prepare_queries(&query, pos, fused, vision, stream)?;
        cache.write_raw(&keys, stream)?;
        let pooled = self.extend_pooled(cache, total, fused, vision, stream)?;
        if total / self.ratio <= self.block_topk {
            return Ok(None);
        }
        Ok(Some(PreparedRows {
            query,
            pooled: pooled.ok_or("QSA selection requires pooled keys")?,
            total,
        }))
    }
}

#[test]
#[ignore = "requires Metal and original ongoing indexer routing receipts; no weights"]
fn mtplx_qsa_canonical_indexer_routes_match_reference() {
    use super::super::{
        gpu::Context,
        qwen_mtplx_tests::{capture_dispatches, pattern, qsa_indexer_options},
    };
    use super::{
        QsaCompiledMode, allocator,
        attention::EagerSelection,
        cache::object,
        mlp,
        stream::{Device, Streams},
    };
    use sha2::{Digest, Sha256};
    super::configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(Device::Gpu).unwrap();
    let cases = include_str!("../../../../tests/fixtures/mtplx-custom-kernels.jsonl")
        .lines()
        .chain(include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-connected.jsonl").lines())
        .map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
        .filter(|c| c["kernel"].as_str().unwrap().starts_with("qsa_indexer_"))
        .collect::<Vec<_>>();
    assert_eq!(cases.len(), 200);
    let weight = mlp::fixture_linear(640, 2560, 4, 64, 75);
    let q_norm = Array::new(&[128], Dtype::BF16, pattern(128, 72)).unwrap();
    let k_norm = Array::new(&[128], Dtype::BF16, pattern(128, 73)).unwrap();
    let freqs = cases[0]["frequencies"]
        .as_array()
        .unwrap()
        .iter()
        .flat_map(|v| (v.as_f64().unwrap() as f32).to_le_bytes())
        .collect::<Vec<_>>();
    let inv_freq = Array::new(&[32], Dtype::F32, super::scalar_buffer(&freqs).unwrap()).unwrap();
    let make = || Indexer {
        projection: Some(Linear::Quantized(&weight)),
        heads: 4,
        kv_heads: 1,
        head_dim: 128,
        block_topk: 512,
        ratio: 4,
        q_norm: &q_norm,
        k_norm: &k_norm,
        inv_freq: &inv_freq,
        mrope_axes: None,
        eps: 1e-6,
        scaling: 1.,
        scratch: 32 * 1024 * 1024,
        compiled: None,
    };
    let mut indexer = make();
    let mut cache = QsaCache::new(4);
    let source = |shape: &[i32], salt| {
        Array::new(
            shape,
            Dtype::BF16,
            pattern(shape.iter().product::<i32>() as u32, salt),
        )
        .unwrap()
    };
    let digest = |a: &Array| {
        let a =
            ops::contiguous(&ops::astype(a, Dtype::F32, false, gpu).unwrap(), false, gpu).unwrap();
        ops::evaluate(&streams, std::slice::from_ref(&a), gpu, false).unwrap();
        let mut bytes = vec![0; a.layout().nbytes()];
        a.buffer().read(a.offset(), &mut bytes).unwrap();
        Sha256::digest(bytes)
            .iter()
            .map(|b| format!("{b:02x}"))
            .collect::<String>()
    };
    for c in &cases {
        super::eval::execution_tests::observed(|| {
            let n = |s: &str| c[s].as_i64().unwrap() as i32;
            let step = n("step");
            let rows = n("rows");
            let pos = n("pos");
            let options = qsa_indexer_options(&c["options"]);
            if step == 0 {
                indexer = make();
                cache = QsaCache::new(4);
                let prefix = n("prefix");
                if prefix > 0 {
                    cache
                        .write_raw(&source(&[1, prefix, 128], 88), gpu)
                        .unwrap();
                    indexer
                        .extend_pooled(&mut cache, prefix, options.fused, None, gpu)
                        .unwrap();
                    let kv = ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu)
                        .unwrap();
                    cache.kv.update_and_fetch(&kv, &kv, gpu).unwrap();
                    let leaves = [
                        &cache.raw_keys,
                        &cache.pooled,
                        &cache.pooled_f32_t,
                        &cache.kv.keys,
                        &cache.kv.values,
                    ]
                    .into_iter()
                    .flatten()
                    .map(|a| a.borrow().clone())
                    .collect::<Vec<_>>();
                    ops::evaluate(&streams, &leaves, gpu, false).unwrap();
                }
            }
            assert_eq!(cache.kv.offset, pos);
            let hidden = source(&[1, rows, 2560], 77 + step as u32);
            let qk = c["supplied"].as_bool().unwrap().then(|| {
                ops::slice(
                    &source(&[1, rows, n("stride")], 121 + step as u32),
                    &[0; 3],
                    &[1, rows, 640],
                    &[1; 3],
                    gpu,
                )
                .unwrap()
            });
            let ((selected, lane), dispatches) = capture_dispatches(|| {
                indexer
                    .call(
                        &hidden,
                        pos,
                        &mut cache,
                        qk.as_ref(),
                        &options,
                        &streams,
                        None,
                        gpu,
                    )
                    .unwrap()
            });
            assert_eq!(lane, c["lane"].as_str().unwrap(), "{}", c["kernel"]);
            if matches!(lane, "compiled" | "metal_prefill" | "legacy_fused") {
                assert!(dispatches.is_empty(), "{} must stay lazy", c["kernel"]);
            }
            if let Some(counts) = c.get("graph_counts") {
                let actual =
                    indexer
                        .compiled
                        .as_ref()
                        .map_or(serde_json::json!([0, 0, 0]), |(_, core)| {
                            let r = core.report();
                            serde_json::json!([r["compiled_calls"], r["traces"], r["entry_count"]])
                        });
                assert_eq!(actual, *counts, "{}", c["kernel"]);
            }
            let (kind, outputs, tail) = match selected {
                None => ("none", vec![], None),
                Some(EagerSelection::Dense(a)) => ("dense", vec![a], None),
                Some(EagerSelection::PrefillBlocks(a, b)) => ("flash_prefill", vec![a, b], None),
                Some(EagerSelection::FlashBlocks(a, t)) => ("flash", vec![a], Some(t)),
                Some(EagerSelection::RowTokens(a, b)) => ("gather_rows", vec![a, b], None),
                Some(EagerSelection::DecodeTokens(a)) => ("decode_tokens", vec![a], None),
            };
            assert_eq!(kind, c["kind"].as_str().unwrap(), "{}", c["kernel"]);
            assert_eq!(tail, c["tail"].as_i64().map(|v| v as i32));
            let mut leaves = outputs.clone();
            leaves.extend(
                [&cache.raw_keys, &cache.pooled, &cache.pooled_f32_t]
                    .into_iter()
                    .flatten()
                    .map(|a| a.borrow().clone()),
            );
            ops::evaluate(&streams, &leaves, gpu, false).unwrap();
            assert_eq!(
                serde_json::json!(outputs.iter().map(&digest).collect::<Vec<_>>()),
                c["sha256"],
                "{} output",
                c["kernel"]
            );
            assert_eq!(cache.kv.offset, pos);
            assert_eq!(cache.pooled_len, n("pooled_len"));
            for (key, leaf, axis) in [
                ("raw_cap", &cache.raw_keys, 1),
                ("pooled_cap", &cache.pooled, 1),
                ("mirror_cap", &cache.pooled_f32_t, 3),
            ] {
                assert_eq!(
                    leaf.as_ref()
                        .map_or(0, |a| a.borrow().layout().shape()[axis]),
                    n(key),
                    "{} {key}",
                    c["kernel"]
                );
            }
            for (key, a) in [
                ("raw", &cache.raw_keys),
                ("pooled", &cache.pooled),
                ("mirror", &cache.pooled_f32_t),
            ] {
                let hash = a.as_ref().map(|a| digest(&a.borrow()));
                assert_eq!(
                    serde_json::json!(hash),
                    c["state_sha256"][key],
                    "{} {key}",
                    c["kernel"]
                );
            }
            let kv = ops::transpose(
                &source(&[1, rows, 2, 256], 140 + step as u32),
                &[0, 2, 1, 3],
                gpu,
            )
            .unwrap();
            let (k, v) = cache.kv.update_and_fetch(&kv, &kv, gpu).unwrap();
            ops::evaluate(
                &streams,
                &[k.borrow().clone(), v.borrow().clone()],
                gpu,
                false,
            )
            .unwrap();
        });
    }
    // Reuse the original 720 pre-reservation eligibility receipts, including
    // inconsistent frontiers and present-but-nonbucketed or empty backings.
    let routing: serde_json::Value = serde_json::from_str(include_str!(
        "../../../../tests/fixtures/mtplx-qsa-compiled-routes.json"
    ))
    .unwrap();
    let indexer = make();
    let mut checked = 0;
    for c in routing["cases"].as_array().unwrap() {
        let options = qsa_indexer_options(&routing["options"][c[0].as_u64().unwrap() as usize]);
        let state = &routing["states"][c[1].as_u64().unwrap() as usize];
        let n = |i: usize| state[i].as_i64().unwrap_or(0) as i32;
        let mut cache = QsaCache::new(4);
        cache.kv.offset = n(2);
        cache.pooled_len = n(5);
        cache.raw_keys = state[3]
            .as_i64()
            .map(|cap| object(Array::unallocated(&[1, cap as i32, 128], Dtype::BF16).unwrap()));
        cache.pooled = state[4]
            .as_i64()
            .map(|cap| object(Array::unallocated(&[1, cap as i32, 128], Dtype::BF16).unwrap()));
        let mode = match c[2].as_str().unwrap() {
            "update_only" => QsaCompiledMode::UpdateOnly,
            "blocks" => QsaCompiledMode::Blocks,
            "prefill_blocks" => QsaCompiledMode::PrefillBlocks,
            "row_tokens" => QsaCompiledMode::RowTokens,
            "dense_mask" => QsaCompiledMode::DenseMask,
            _ => panic!("unknown mode"),
        };
        assert_eq!(
            options.compiled_mode(n(1) as u32, (n(0) + n(1)) as u32, 4, 512),
            mode
        );
        for supplied in [false, true] {
            let source =
                Array::unallocated(&[1, n(1), if supplied { 640 } else { 2560 }], Dtype::BF16)
                    .unwrap();
            let (accepted, dispatches) = capture_dispatches(|| {
                indexer.compiled_supported(
                    &source,
                    &cache,
                    n(0),
                    supplied,
                    n(1) == 1,
                    mode,
                    &options,
                    gpu,
                )
            });
            assert!(dispatches.is_empty());
            assert_eq!(
                accepted,
                c[if supplied { 4 } else { 3 }].as_bool().unwrap(),
                "{c} {state} supplied{supplied}"
            );
            checked += 1;
        }
    }
    assert_eq!(checked, 720);
    assert!(indexer.compiled.is_none());
}

#[test]
#[ignore = "requires Metal and original indexer lifecycle receipts; no model weights"]
fn mtplx_qsa_indexer_compiled_commit_matches_reference() {
    use super::super::{
        gpu::Context,
        qwen_mtplx_tests::{capture_dispatches, pattern},
    };
    use super::{
        QsaCompiledMode, allocator,
        attention::{self, EagerSelection},
        mlp,
        stream::{Device, Streams},
    };
    use sha2::{Digest, Sha256};
    super::configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(Device::Gpu).unwrap();
    let cases = include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-commit.jsonl")
        .lines()
        .map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
        .collect::<Vec<_>>();
    assert_eq!(cases.len(), 26);
    let renewed = |a: &Array| {
        Array::new(
            a.layout().shape(),
            a.layout().dtype(),
            a.buffer()
                .view(a.offset(), a.layout().nbytes() as u64)
                .unwrap(),
        )
        .unwrap()
    };
    let norm = |salt| Array::new(&[128], Dtype::BF16, pattern(128, salt)).unwrap();
    let replacement = norm(72);
    let q_norm = [norm(72), replacement.clone(), replacement];
    let k_norm = [norm(73), norm(73)];
    let invalid = Array::unallocated(&[127], Dtype::BF16).unwrap();
    let bytes = cases[0]["frequencies"]
        .as_array()
        .unwrap()
        .iter()
        .flat_map(|v| (v.as_f64().unwrap() as f32).to_le_bytes())
        .collect::<Vec<_>>();
    let freq = Array::new(&[32], Dtype::F32, super::scalar_buffer(&bytes).unwrap()).unwrap();
    let freq = [renewed(&freq), freq];
    let weight = Array::new(&[640, 128], Dtype::BF16, pattern(640 * 128, 231)).unwrap();
    let weights = [renewed(&weight), weight];
    let mut layers = vec![mlp::fixture_linear(640, 128, 4, 64, 75)];
    for leaf in 0..3 {
        let a = layers.last().unwrap();
        layers.push(mlp::QuantizedLinear {
            weight: if leaf == 0 {
                renewed(&a.weight)
            } else {
                a.weight.clone()
            },
            scales: if leaf == 1 {
                renewed(&a.scales)
            } else {
                a.scales.clone()
            },
            biases: if leaf == 2 {
                renewed(&a.biases)
            } else {
                a.biases.clone()
            },
            bits: a.bits,
            group: a.group,
        });
    }
    for retain in [true, false] {
        for quantized in [false, true] {
            let mut indexer = Indexer {
                projection: Some(if quantized {
                    Linear::Quantized(&layers[0])
                } else {
                    Linear::Dense(&weights[0])
                }),
                heads: 4,
                kv_heads: 1,
                head_dim: 128,
                block_topk: 8,
                ratio: 4,
                q_norm: &q_norm[0],
                k_norm: &k_norm[0],
                inv_freq: &freq[0],
                mrope_axes: None,
                eps: 1e-6,
                scaling: 1.,
                scratch: 32 * 1024 * 1024,
                compiled: None,
            };
            let mut cache = QsaCache::new(4);
            for c in cases.iter().filter(|c| c["quantized"] == quantized) {
                super::eval::execution_tests::observed(|| {
                    let n = |key: &str| c[key].as_i64().unwrap() as i32;
                    let step = n("step");
                    let rows = n("rows");
                    if step >= 2 {
                        indexer.scratch = 1024 * 1024;
                    }
                    if step == 3 {
                        indexer.q_norm = &q_norm[1];
                    }
                    if step == 4 {
                        indexer.k_norm = &k_norm[1];
                    }
                    if step == 5 {
                        indexer.inv_freq = &freq[1];
                    }
                    if step == 6 {
                        indexer.projection = Some(if quantized {
                            Linear::Quantized(&layers[1])
                        } else {
                            Linear::Dense(&weights[1])
                        });
                    }
                    if quantized && step == 7 {
                        indexer.projection = Some(Linear::Quantized(&layers[2]));
                    }
                    if quantized && step == 8 {
                        indexer.projection = Some(Linear::Quantized(&layers[3]));
                    }
                    if step == 9 {
                        cache.reserve_indexer_capacity(2048, 512, gpu).unwrap();
                    }
                    if step == 10 {
                        assert_eq!(cache.trim(3), 3);
                    }
                    let pos = cache.kv.offset;
                    assert_eq!(pos, n("pos"));
                    let workspace = if step >= 2 {
                        2 * 1024 * 1024
                    } else {
                        128 * 1024 * 1024
                    };
                    if step == 11 {
                        let report = indexer.compiled.as_ref().unwrap().1.report();
                        let signature = indexer.compiled.as_ref().unwrap().0.clone();
                        indexer.q_norm = &invalid;
                        assert!(indexer.compiled_core(workspace).is_err());
                        assert_eq!(indexer.compiled.as_ref().unwrap().0, signature);
                        assert_eq!(indexer.compiled.as_ref().unwrap().1.report(), report);
                        indexer.q_norm = &q_norm[1];
                    }
                    if step == 12 {
                        // A new Rust handle of the same descriptor is not a new leaf.
                        indexer.q_norm = &q_norm[2];
                    }
                    let hidden = Array::new(
                        &[1, rows, 128],
                        Dtype::BF16,
                        pattern(rows as u32 * 128, 300 + step as u32),
                    )
                    .unwrap();
                    let supplied = c["supplied"].as_bool().unwrap();
                    let qk = supplied.then(|| {
                        let a = Array::new(
                            &[1, rows, 704],
                            Dtype::BF16,
                            pattern(rows as u32 * 704, 320 + step as u32),
                        )
                        .unwrap();
                        ops::slice(&a, &[0; 3], &[1, rows, 640], &[1; 3], gpu).unwrap()
                    });
                    if cache.pooled.is_some() {
                        cache.pooled_f32_view(cache.pooled_len, gpu).unwrap();
                    }
                    let old = [cache.raw_keys.clone(), cache.pooled.clone()];
                    let before = old
                        .iter()
                        .flatten()
                        .map(|a| a.borrow().clone())
                        .collect::<Vec<_>>();
                    ops::evaluate(&streams, &before, gpu, false).unwrap();
                    drop(before);
                    let old = if retain {
                        Some(old)
                    } else {
                        drop(old);
                        None
                    };
                    let mode = match c["mode"].as_str().unwrap() {
                        "update_only" => QsaCompiledMode::UpdateOnly,
                        "blocks" => QsaCompiledMode::Blocks,
                        "prefill_blocks" => QsaCompiledMode::PrefillBlocks,
                        "row_tokens" => QsaCompiledMode::RowTokens,
                        "dense_mask" => QsaCompiledMode::DenseMask,
                        _ => panic!("mode"),
                    };
                    let engagement = attention::qsa_prefill_engagement()
                        .unwrap()
                        .get("compiled_selector")
                        .copied()
                        .unwrap_or(0);
                    let (result, dispatches) = capture_dispatches(|| {
                        indexer
                            .call_compiled(
                                &hidden,
                                pos,
                                &mut cache,
                                qk.as_ref(),
                                mode,
                                workspace,
                                gpu,
                            )
                            .unwrap()
                    });
                    assert!(dispatches.is_empty(), "cache commit must stay lazy");
                    assert_eq!(cache.kv.offset, pos);
                    assert_eq!(cache.pooled_len, n("pooled_len"));
                    assert!(cache.pooled_f32_t.is_none());
                    assert_eq!(
                        indexer.compiled.as_ref().unwrap().1.report(),
                        c["report"],
                        "{step}/{quantized}"
                    );
                    assert_eq!(
                        attention::qsa_prefill_engagement()
                            .unwrap()
                            .get("compiled_selector")
                            .copied()
                            .unwrap_or(0),
                        engagement + u64::from(mode == QsaCompiledMode::PrefillBlocks)
                    );
                    let (kind, mut outputs, tail) = match result {
                        None => ("none", vec![], None),
                        Some(EagerSelection::FlashBlocks(a, tail)) => {
                            ("flash", vec![a], Some(tail))
                        }
                        Some(EagerSelection::PrefillBlocks(a, b)) => {
                            ("flash_prefill", vec![a, b], None)
                        }
                        Some(EagerSelection::RowTokens(a, b)) => ("gather_rows", vec![a, b], None),
                        Some(EagerSelection::Dense(a)) => ("dense", vec![a], None),
                        _ => panic!("unexpected return contract"),
                    };
                    assert_eq!(kind, c["kind"].as_str().unwrap());
                    assert_eq!(tail, c["tail"].as_i64().map(|v| v as i32));
                    outputs.extend([
                        cache.raw_keys.as_ref().unwrap().borrow().clone(),
                        cache.pooled.as_ref().unwrap().borrow().clone(),
                    ]);
                    if let Some(old) = old {
                        for (old, new) in old.iter().zip([&cache.raw_keys, &cache.pooled]) {
                            if let Some(old) = old {
                                assert!(!std::rc::Rc::ptr_eq(old, new.as_ref().unwrap()));
                                outputs.push(old.borrow().clone());
                            }
                        }
                    }
                    for (a, shape) in outputs.iter().zip(c["shapes"].as_array().unwrap()) {
                        assert_eq!(serde_json::json!(a.layout().shape()), *shape);
                    }
                    let outputs = outputs
                        .iter()
                        .map(|a| {
                            ops::contiguous(
                                &ops::astype(a, Dtype::F32, false, gpu).unwrap(),
                                false,
                                gpu,
                            )
                            .unwrap()
                        })
                        .collect::<Vec<_>>();
                    ops::evaluate(&streams, &outputs, gpu, false).unwrap();
                    for (leaf, a) in outputs.iter().enumerate() {
                        let mut bytes = vec![0; a.layout().nbytes()];
                        a.buffer().read(a.offset(), &mut bytes).unwrap();
                        let hash = Sha256::digest(bytes)
                            .iter()
                            .map(|b| format!("{b:02x}"))
                            .collect::<String>();
                        assert_eq!(
                            hash,
                            c["sha256"][leaf].as_str().unwrap(),
                            "step{step} quant{quantized} retain{retain} leaf{leaf}"
                        );
                    }
                    // A real attention invocation advances this after indexer commit.
                    cache.kv.offset = pos + rows;
                });
            }
        }
    }
}

#[test]
#[ignore = "requires Apple Metal, original stateful QSAIndexer receipts, no model weights"]
fn mtplx_qsa_indexer_prefix_matches_reference() {
    use super::super::{
        gpu::Context,
        qwen_mtplx_tests::{capture_dispatches, pattern, verify_typed},
    };
    use super::{
        QsaEagerOptions, allocator,
        attention::{self, EagerSelection},
        configure_sources, mlp,
        stream::{Device, Streams},
    };
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(Device::Gpu).unwrap();
    let tf32 = std::env::var("MLX_ENABLE_TF32").ok().as_deref() != Some("0");
    let cases = include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-prefix.jsonl")
        .lines()
        .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
        .filter(|case| case["tf32"].as_bool() == Some(tf32))
        .collect::<Vec<_>>();
    assert_eq!(cases.len(), 153);
    for prepared in [false, true] {
        for tag in ["bfloat16", "float16", "float32"] {
            let dtype = match tag {
                "bfloat16" => Dtype::BF16,
                "float16" => Dtype::F16,
                _ => Dtype::F32,
            };
            for profile in 0..if dtype == Dtype::BF16 { 7 } else { 5 } {
                let sequence = cases
                    .iter()
                    .filter(|c| c["dtype"] == tag && c["profile"] == profile)
                    .collect::<Vec<_>>();
                assert_eq!(sequence.len(), 9);
                let first = sequence[0];
                let dim = first["dim"].as_i64().unwrap() as i32;
                let rot = first["rot"].as_i64().unwrap() as i32;
                let fused = first["fused"].as_bool().unwrap();
                let kind = first["kind"].as_str().unwrap();
                let source = |shape: &[i32], salt: u32, dtype: Dtype| {
                    let count = shape.iter().product::<i32>();
                    ops::astype(
                        &Array::new(shape, Dtype::BF16, pattern(count as u32, salt)).unwrap(),
                        dtype,
                        false,
                        gpu,
                    )
                    .unwrap()
                };
                let weight = source(&[5 * dim, 64], 187, dtype);
                let quant = match kind {
                    "q4" => Some(mlp::fixture_linear(5 * dim as u32, 64, 4, 64, 187)),
                    "q8" => Some(mlp::fixture_linear(5 * dim as u32, 64, 8, 64, 187)),
                    _ => None,
                };
                let norm_dtype = if kind == "promoted" {
                    Dtype::F32
                } else {
                    dtype
                };
                let q_norm = source(&[dim], 185, norm_dtype);
                let k_norm = source(&[dim], 186, norm_dtype);
                let frequencies = (0..rot / 2)
                    .flat_map(|i| ((1. / f64::from(i + 1)) as f32).to_le_bytes())
                    .collect::<Vec<_>>();
                let inv_freq = Array::new(
                    &[rot / 2],
                    Dtype::F32,
                    super::scalar_buffer(&frequencies).unwrap(),
                )
                .unwrap();
                let indexer = Indexer {
                    projection: Some(
                        quant
                            .as_ref()
                            .map_or(Linear::Dense(&weight), Linear::Quantized),
                    ),
                    heads: 4,
                    kv_heads: 1,
                    head_dim: dim,
                    block_topk: 512,
                    ratio: 4,
                    q_norm: &q_norm,
                    k_norm: &k_norm,
                    inv_freq: &inv_freq,
                    mrope_axes: None,
                    eps: 1e-6,
                    scaling: 1.125,
                    scratch: 32 * 1024 * 1024,
                    compiled: None,
                };
                let options = QsaEagerOptions {
                    large_prefill: false,
                    flash: false,
                    gather_decode: true,
                    gather: true,
                    gather_max_rows: 8,
                    gather_min_context: 0,
                };
                let mut cache = QsaCache::new(4);
                cache.reserve_indexer_capacity(256, 256, gpu).unwrap();
                for case in sequence {
                    super::eval::execution_tests::observed(|| {
                        let num = |s: &str| case[s].as_i64().unwrap() as i32;
                        let rows = num("rows");
                        let step = num("step");
                        if num("trim") > 0 {
                            assert_eq!(cache.trim(num("trim")), num("trim"));
                        }
                        let pos = cache.kv.offset;
                        assert_eq!(pos, num("pos"));
                        let old = cache.pooled_len;
                        let hidden = source(&[1, rows, 64], 190 + step as u32 * 7, dtype);
                        let qk = if kind == "supplied" {
                            Some(
                                ops::slice(
                                    &source(&[1, rows, 5 * dim * 2], 191 + step as u32 * 7, dtype),
                                    &[0; 3],
                                    &[1, rows, 5 * dim * 2],
                                    &[1, 1, 2],
                                    gpu,
                                )
                                .unwrap(),
                            )
                        } else {
                            None
                        };
                        if prepared {
                            let mut inputs = vec![
                                hidden.clone(),
                                q_norm.clone(),
                                k_norm.clone(),
                                weight.clone(),
                            ];
                            if let Some(qk) = &qk {
                                inputs.push(qk.clone());
                            }
                            ops::evaluate(&streams, &inputs, gpu, false).unwrap();
                        }
                        let (prefix, calls) = capture_dispatches(|| {
                            indexer
                                .prepare_rows(
                                    &hidden,
                                    pos,
                                    &mut cache,
                                    qk.as_ref(),
                                    fused,
                                    None,
                                    gpu,
                                )
                                .unwrap()
                        });
                        assert!(calls.is_empty(), "QSA prefix must remain lazy");
                        assert_eq!(cache.kv.offset, num("offset"));
                        assert_eq!(cache.pooled_len, num("pooled_len"));
                        let mut outputs = vec![cache.raw_keys.as_ref().unwrap().borrow().clone()];
                        if let Some(pooled) = &cache.pooled {
                            outputs.extend([
                                pooled.borrow().clone(),
                                cache.pooled_f32_t.as_ref().unwrap().borrow().clone(),
                            ]);
                        }
                        let lane = if let Some(prefix) = prefix {
                            assert_eq!(prefix.total, pos + rows);
                            assert_eq!(prefix.pooled.layout().shape(), [1, cache.pooled_len, dim]);
                            let selection = attention::select_eager(
                                &prefix.query,
                                pos,
                                &mut cache,
                                prefix.total,
                                512,
                                0,
                                &options,
                                &streams,
                                gpu,
                            )
                            .unwrap();
                            match selection {
                                EagerSelection::RowTokens(ids, valid) => {
                                    outputs.extend([ids, valid]);
                                    "gather_rows"
                                }
                                EagerSelection::DecodeTokens(ids) => {
                                    outputs.push(ids);
                                    "decode_tokens"
                                }
                                EagerSelection::Dense(mask) => {
                                    outputs.push(mask);
                                    "dense"
                                }
                                _ => panic!("unexpected eager route"),
                            }
                        } else {
                            "none"
                        };
                        assert_eq!(lane, case["lane"].as_str().unwrap());
                        for (i, a) in outputs.iter().enumerate() {
                            let shape = case["shapes"][i]
                                .as_array()
                                .unwrap()
                                .iter()
                                .map(|d| d.as_i64().unwrap() as i32)
                                .collect::<Vec<_>>();
                            assert_eq!(a.layout().shape(), shape, "{} output{i}", case["kernel"]);
                        }
                        drop((hidden, qk));
                        let outputs = outputs
                            .iter()
                            // NumPy's reference tobytes is logical row order;
                            // the cache's F32 mirror is a transposed view.
                            .map(|a| {
                                ops::contiguous(
                                    &ops::astype(a, Dtype::F32, false, gpu).unwrap(),
                                    false,
                                    gpu,
                                )
                                .unwrap()
                            })
                            .collect::<Vec<_>>();
                        let (_, calls) = capture_dispatches(|| {
                            ops::evaluate(&streams, &outputs, gpu, false).unwrap()
                        });
                        let custom = fused && dim <= 128 && norm_dtype == dtype;
                        assert_eq!(
                            calls
                                .iter()
                                .filter(|c| c.0.contains("qsa_prepare_q"))
                                .count(),
                            usize::from(custom && lane != "none")
                        );
                        assert_eq!(
                            calls.iter().filter(|c| c.0.contains("qsa_pool_k")).count(),
                            usize::from(custom && cache.pooled_len > old)
                        );
                        let buffers = outputs.iter().map(Array::buffer).collect::<Vec<_>>();
                        let checks = buffers
                            .iter()
                            .zip(&outputs)
                            .map(|(b, a)| (&**b, a.layout().size() as u32, 4))
                            .collect::<Vec<_>>();
                        verify_typed(case["kernel"].as_str().unwrap(), &checks);
                        cache.kv.offset = pos + rows; // Attention's frontier advance.
                    });
                }
            }
        }
    }
}