prism-inference 0.1.20

//! Explicit layer-major CPU prompt tiles with unchanged token arithmetic.

use std::fmt::Write as _;

use super::{
    Arithmetic, Attention, Evaluation, FullAttention, Geometry, KvCache, LayerState, Qwen35,
    Transforms, add, checked_product, finite, resize_into, rope_frequencies, row_into, sigmoid,
    small_float,
};
use crate::kernels::{
    MatvecBackend, RowParallelism,
    prefill::{ProjectionTile, ROW_TASK_SCRATCH_BYTES},
    round_f16_slice,
};
use crate::model::Tensor;
use crate::model::lora::{RankOneAdapter, Scratch as AdapterScratch};
use crate::{MAX_PREFILL_BLOCK_TOKENS, Result, invalid};

#[path = "prefill_gdn.rs"]
mod gdn;

#[derive(Clone, Copy, Default)]
pub(super) struct TileReuse<'a> {
    basis: Option<&'a Tensor>,
    prepared: Option<&'a Tensor>,
}

const fn reborrow_executor<'a, 'b: 'a>(
    executor: &'a mut Option<&'b mut dyn crate::kernels::Pq2ProjectionExecutor>,
) -> Option<&'a mut (dyn crate::kernels::Pq2ProjectionExecutor + 'a)> {
    match executor {
        Some(exec) => Some(*exec),
        None => None,
    }
}

/// Lifetime counters for explicitly selected CPU tiled prompt ingestion.
///
/// Counts include provisional work later cancelled or rolled back. They are not
/// performance measurements, committed-token counters or worker/ISA attribution.
#[derive(Clone, Copy, Default, Debug, Eq, PartialEq)]
pub struct TiledPrefillStatistics {
    /// Successfully completed causal token tiles.
    pub completed_tiles: u64,
    /// Tokens in successfully completed tiles, including later request rollback.
    pub completed_tokens: u64,
    /// Packed matrix projections dispatched across multiple tile tokens.
    pub packed_projections: u64,
    /// Packed weight blocks decoded once and shared across tile tokens.
    pub decoded_weight_blocks: u64,
    /// Recurrent heads evaluated by the separately selected FP32 chunk algebra.
    pub chunkwise_gdn_heads: u64,
}

pub(super) struct TiledPrefill {
    tokens: usize,
    scratch_bytes: usize,
    hidden: Vec<Vec<f32>>,
    normalized: Vec<Vec<f32>>,
    up: Vec<Vec<f32>>,
    gate: Vec<Vec<f32>>,
    transformed: Vec<Vec<f32>>,
    mixed: Vec<Vec<f32>>,
    recurrent_gate: Vec<Vec<f32>>,
    beta: Vec<Vec<f32>>,
    alpha: Vec<Vec<f32>>,
    attended: Vec<Vec<f32>>,
    attention: Vec<Vec<f32>>,
    frequencies: Vec<Vec<(f32, f32)>>,
    gdn_conv_outputs: Vec<f32>,
    gdn_head_outputs: Vec<f32>,
    gdn_key_dots: Vec<f32>,
    chunkwise_gdn: bool,
    projector: ProjectionTile,
    statistics: TiledPrefillStatistics,
}

impl TiledPrefill {
    #[cfg_attr(
        not(feature = "parallel-rows"),
        allow(
            clippy::missing_const_for_fn,
            reason = "The delegated setter drops the replaced pool owner in parallel builds; retain one signature across feature sets."
        )
    )]
    pub(super) fn set_row_parallelism(&mut self, parallelism: RowParallelism) {
        self.projector.set_row_parallelism(parallelism);
    }

    fn new(
        tokens: usize,
        width: usize,
        scratch_bytes: usize,
        backend: MatvecBackend,
        parallelism: RowParallelism,
    ) -> Result<Self> {
        let create = || -> Result<Vec<Vec<f32>>> {
            let mut rows = Vec::new();
            rows.try_reserve_exact(tokens)?;
            for _ in 0..tokens {
                let mut row = Vec::new();
                row.try_reserve_exact(width)?;
                rows.push(row);
            }
            Ok(rows)
        };
        let mut frequencies = Vec::new();
        frequencies.try_reserve_exact(tokens)?;
        for _ in 0..tokens {
            frequencies.push(Vec::new());
        }
        let mut projector = ProjectionTile::with_backend(backend);
        projector.set_row_parallelism(parallelism);
        Ok(Self {
            tokens,
            scratch_bytes,
            hidden: create()?,
            normalized: create()?,
            up: create()?,
            gate: create()?,
            transformed: create()?,
            mixed: create()?,
            recurrent_gate: create()?,
            beta: create()?,
            alpha: create()?,
            attended: create()?,
            attention: create()?,
            frequencies,
            gdn_conv_outputs: Vec::new(),
            gdn_head_outputs: Vec::new(),
            gdn_key_dots: Vec::new(),
            chunkwise_gdn: false,
            projector,
            statistics: TiledPrefillStatistics::default(),
        })
    }
}

impl Qwen35 {
    // Rebuild execution identity from the selected typed policy whenever an
    // adapter reset replaces its trailing scale contract. Never infer a live
    // numerical policy by parsing an old identity string.
    pub(super) fn append_tiled_prefill_identity(&self, identity: &mut String) -> Result<()> {
        if let Some(tile) = &self.tiled_prefill {
            identity.try_reserve_exact(64)?;
            write!(identity, ":cpu-tiled-ffn-v1:{}", tile.tokens)?;
            if tile.chunkwise_gdn {
                identity.push_str(":gdn-wy-f32-v1");
            }
        }
        Ok(())
    }

    /// Conservative extra payload/metadata reservation for a CPU prompt tile.
    ///
    /// This is additional to live state, ordinary graph scratch, checkpoints,
    /// worker stacks and allocator overhead. Both packed activation contracts
    /// and floating fallback high-water storage are charged simultaneously.
    ///
    /// # Errors
    /// Rejects tile sizes other than 8, 16 or 32 and size overflow.
    pub fn tiled_prefill_scratch_bytes(&self, tokens: usize) -> Result<usize> {
        if !matches!(tokens, 8 | 16 | 32) {
            return Err(invalid("CPU prefill tiles must contain 8, 16 or 32 tokens"));
        }
        let width = self
            .geometry
            .embedding
            .max(self.geometry.feed_forward)
            .max(self.geometry.convolution_width())
            .max(self.geometry.value_width())
            .max(self.geometry.query_gate_width())
            .max(self.geometry.kv_width());
        checked_product(&[tokens, width, 80])?
            .checked_add(checked_product(&[tokens, 2048])?)
            // The four fixed 32x32 coefficient/product tables and their small
            // temporary vectors exist even for an eight-token, narrow graph.
            .and_then(|bytes| bytes.checked_add(64 * 1024))
            // Fixed task slice/error descriptors stay bounded at 16 workers
            // and 32 token outputs, including later worker-policy changes.
            .and_then(|bytes| bytes.checked_add(ROW_TASK_SCRATCH_BYTES))
            .ok_or_else(|| invalid("CPU prefill scratch reservation overflow"))
    }

    /// Select experimental CPU tiled state-only prefill before the first token.
    ///
    /// This consumes no new weights and leaves ordinary one-token decode and
    /// default block ingestion unchanged. FFN packed blocks are decoded once
    /// across a tile; attention/GDN/convolution remain strictly causal in their
    /// existing token arithmetic order. This is not a chunkwise WY GDN kernel.
    /// Separate identity prevents mixing runtime prefix/cache plans. No private
    /// service or application default selects this unqualified source path.
    ///
    /// The explicit reservation covers additional tile storage, and must be
    /// charged by a model owner without reducing existing reserves. Cancellation
    /// checks occur between tokens, layers and packed matrix projections. An
    /// error may leave provisional layer state; as for [`Decoder::ingest_block`],
    /// the owner must restore the entire request checkpoint or reset before reuse.
    ///
    /// # Errors
    /// Rejects nonempty/already-configured graphs, unsupported sizes, insufficient
    /// reservation and failed allocation before changing selection or identity.
    pub fn with_tiled_prefill(mut self, tokens: usize, scratch_bytes: usize) -> Result<Self> {
        let required = self.tiled_prefill_scratch_bytes(tokens)?;
        if self.position != 0 || self.tiled_prefill.is_some() || scratch_bytes < required {
            return Err(invalid(
                "CPU tiled prefill requires an empty graph and sufficient scratch",
            ));
        }
        let tile = TiledPrefill::new(
            tokens,
            self.geometry
                .embedding
                .max(self.geometry.feed_forward)
                .max(self.geometry.convolution_width())
                .max(self.geometry.value_width())
                .max(self.geometry.query_gate_width())
                .max(self.geometry.kv_width()),
            required,
            self.requested_backend(),
            self.workspace.projection.matvec.clone_row_parallelism(),
        )?;
        let suffix = format!(":cpu-tiled-ffn-v1:{tokens}");
        self.identity.try_reserve_exact(suffix.len())?;
        self.identity.push_str(&suffix);
        self.provenance = std::sync::Arc::new(());
        self.tiled_prefill = Some(tile);
        Ok(self)
    }

    /// Select a separately identified FP32 triangular GDN chunk algorithm.
    ///
    /// Requires explicit tiled prefill selection on an empty graph. This changes
    /// floating evaluation order; it does not promise bit-identical recurrence.
    /// The initial qualification tolerance is `1e-3 + 1e-3 * abs(reference)` for
    /// complete logits and state, including long-prefix continuation and restore.
    /// This source API establishes no actual-model acceptance or speed benefit.
    /// Decay interval products avoid division and handle exactly zero decay.
    ///
    /// # Errors
    /// Rejects nonempty, untiled or already selected graphs and failed identity allocation.
    pub fn with_chunkwise_gdn_prefill(mut self) -> Result<Self> {
        let Some(tile) = self.tiled_prefill.as_mut() else {
            return Err(invalid("chunkwise GDN requires explicit tiled prefill"));
        };
        if self.position != 0 || tile.chunkwise_gdn {
            return Err(invalid("chunkwise GDN requires an empty, unselected graph"));
        }
        let suffix = ":gdn-wy-f32-v1";
        self.identity.try_reserve_exact(suffix.len())?;
        self.identity.push_str(suffix);
        self.provenance = std::sync::Arc::new(());
        tile.chunkwise_gdn = true;
        Ok(self)
    }

    /// Extra reserved tile payload/metadata, zero for ordinary ingestion.
    #[must_use]
    pub fn tiled_prefill_reserved_bytes(&self) -> usize {
        self.tiled_prefill
            .as_ref()
            .map_or(0, |tile| tile.scratch_bytes)
    }

    /// Provisional tile work, separate from one-vector backend diagnostics.
    #[must_use]
    pub fn tiled_prefill_statistics(&self) -> TiledPrefillStatistics {
        self.tiled_prefill
            .as_ref()
            .map_or_else(TiledPrefillStatistics::default, |tile| {
                let mut statistics = tile.statistics;
                statistics.packed_projections = tile.projector.packed_projections;
                statistics.decoded_weight_blocks = tile.projector.decoded_blocks;
                statistics
            })
    }

    pub(super) fn ingest_prefill_block(
        &mut self,
        tokens: &[u32],
        check: &dyn Fn() -> Result<()>,
    ) -> Result<()> {
        const fn coerce_executor<'a>(
            executor: &'a mut Box<dyn crate::kernels::Pq2ProjectionExecutor + Send>,
        ) -> &'a mut (dyn crate::kernels::Pq2ProjectionExecutor + 'a) {
            &mut **executor
        }
        let mut executor = self.projection_executor.take();
        let result = self.ingest_prefill_block_with_executor(
            tokens,
            check,
            executor.as_mut().map(coerce_executor),
        );
        self.projection_executor = executor;
        result
    }

    pub(super) fn ingest_prefill_block_with_executor(
        &mut self,
        tokens: &[u32],
        check: &dyn Fn() -> Result<()>,
        mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    ) -> Result<()> {
        if tokens.is_empty()
            || tokens.len() > MAX_PREFILL_BLOCK_TOKENS
            || tokens
                .iter()
                .any(|&token| usize::try_from(token).map_or(true, |value| value >= self.vocabulary))
        {
            return Err(invalid("invalid native prefill block bounds or tokens"));
        }
        if self.tiled_prefill.is_some()
            && self
                .position
                .checked_add(tokens.len())
                .is_none_or(|end| end > self.context_limit)
        {
            return Err(invalid("CPU prefill block exceeds remaining context"));
        }
        let Some(mut tile) = self.tiled_prefill.take() else {
            for &token in tokens {
                check()?;
                self.evaluate_token_with_executor(token, None, reborrow_executor(&mut executor))?;
                check()?;
            }
            return Ok(());
        };
        let result = (|| {
            for tokens in tokens.chunks(tile.tokens) {
                check()?;
                self.ingest_cpu_tile(tokens, check, &mut tile, reborrow_executor(&mut executor))?;
                tile.statistics.completed_tiles = tile.statistics.completed_tiles.saturating_add(1);
                tile.statistics.completed_tokens = tile
                    .statistics
                    .completed_tokens
                    .saturating_add(u64::try_from(tokens.len())?);
                check()?;
            }
            Ok(())
        })();
        self.tiled_prefill = Some(tile);
        result
    }

    #[allow(
        clippy::too_many_lines,
        reason = "Keep causal attention commits, FFN tile consumption and cancellation visibly ordered."
    )]
    fn ingest_cpu_tile(
        &mut self,
        tokens: &[u32],
        check: &dyn Fn() -> Result<()>,
        tile: &mut TiledPrefill,
        mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    ) -> Result<()> {
        let count = tokens.len();
        let base = self.position;
        for (token, hidden) in tokens.iter().zip(&mut tile.hidden) {
            check()?;
            row_into(&self.embedding, usize::try_from(*token)?, hidden)?;
            if self.transforms.inverse_embedding {
                self.transforms
                    .rotate_with_parallelism(hidden, tile.projector.row_parallelism())?;
                self.transforms.sign(hidden);
            }
            if let Some(adapter) = &self.adapter {
                adapter.embedding(
                    usize::try_from(*token)?,
                    self.adapter_scale,
                    hidden,
                    &mut self.workspace.adapter,
                )?;
            }
        }
        // A partial layer-major attempt is not a committed token prefix. Mark
        // position nonempty before any state mutation so explicit reset always
        // clears provisional state, including failure during the first tile.
        self.position = base + count;
        let mut frequencies_ready = false;
        for ((layer, state), delta) in self
            .layers
            .iter()
            .zip(&mut self.state)
            .zip(&mut self.workspace.deltas)
        {
            if count > 1 && tile.projector.row_parallelism().workers() > 1 {
                let hidden_ref = &tile.hidden[..count];
                let arithmetic = self.arithmetic;
                let norm = &layer.norm;
                let epsilon = self.geometry.epsilon;
                tile.projector.row_parallelism().update_chunks(
                    &mut tile.normalized[..count],
                    1,
                    |base_tok, chunk| {
                        for (offset, normalized) in chunk.iter_mut().enumerate() {
                            arithmetic.normalize(
                                &hidden_ref[base_tok + offset],
                                norm,
                                epsilon,
                                normalized,
                            )?;
                        }
                        Ok(())
                    },
                )?;
            } else {
                for (hidden, normalized) in tile.hidden[..count]
                    .iter()
                    .zip(&mut tile.normalized[..count])
                {
                    self.arithmetic.normalize(
                        hidden,
                        &layer.norm,
                        self.geometry.epsilon,
                        normalized,
                    )?;
                }
            }
            match (&layer.attention, &mut *state) {
                (Attention::Recurrent(weights), _) => {
                    gdn::recurrent_tile(
                        weights,
                        state,
                        delta,
                        &self.geometry,
                        self.arithmetic,
                        &self.transforms,
                        self.adapter.as_ref(),
                        self.adapter_scale,
                        &mut self.workspace.adapter,
                        &mut self.workspace.projection.matvec,
                        count,
                        check,
                        tile,
                        reborrow_executor(&mut executor),
                    )?;
                }
                (Attention::Full(weights), LayerState::Full { keys, values }) => {
                    check()?;
                    let _ = project_tile_triple(
                        &self.transforms,
                        &self.geometry,
                        self.arithmetic,
                        self.adapter.as_ref(),
                        self.adapter_scale,
                        &mut self.workspace.adapter,
                        &mut self.workspace.projection.matvec,
                        [&weights.query, &weights.key, &weights.value],
                        &tile.normalized[..count],
                        &mut tile.mixed[..count],
                        &mut tile.recurrent_gate[..count],
                        &mut tile.beta[..count],
                        &mut tile.transformed[..count],
                        &mut tile.projector,
                        reborrow_executor(&mut executor),
                        check,
                    )?;
                    full_attention_tile(
                        weights,
                        &self.geometry,
                        self.arithmetic,
                        &self.transforms,
                        base,
                        count,
                        &mut frequencies_ready,
                        keys,
                        values,
                        &mut self.workspace.projection,
                        self.adapter.as_ref(),
                        self.adapter_scale,
                        &mut self.workspace.adapter,
                        &mut self.workspace.attention_scratch,
                        tile,
                        check,
                    )?;
                    check()?;
                    let _ = project_tile(
                        &self.transforms,
                        &self.geometry,
                        self.arithmetic,
                        self.adapter.as_ref(),
                        self.adapter_scale,
                        &mut self.workspace.adapter,
                        &mut self.workspace.projection.matvec,
                        TileReuse::default(),
                        &weights.output,
                        &tile.attended[..count],
                        &mut tile.attention[..count],
                        &mut tile.transformed[..count],
                        &mut tile.projector,
                        reborrow_executor(&mut executor),
                        check,
                    )?;
                }
                _ => return Err(invalid("CPU tiled prefill layer/cache mismatch")),
            }
            for index in 0..count {
                check()?;
                add(&mut tile.hidden[index], &tile.attention[index]);
            }
            if count > 1 && tile.projector.row_parallelism().workers() > 1 {
                let hidden_ref = &tile.hidden[..count];
                let arithmetic = self.arithmetic;
                let post_norm = &layer.post_norm;
                let epsilon = self.geometry.epsilon;
                tile.projector.row_parallelism().update_chunks(
                    &mut tile.normalized[..count],
                    1,
                    |base_tok, chunk| {
                        for (offset, normalized) in chunk.iter_mut().enumerate() {
                            arithmetic.normalize(
                                &hidden_ref[base_tok + offset],
                                post_norm,
                                epsilon,
                                normalized,
                            )?;
                        }
                        Ok(())
                    },
                )?;
            } else {
                for (hidden, normalized) in tile.hidden[..count]
                    .iter()
                    .zip(&mut tile.normalized[..count])
                {
                    self.arithmetic.normalize(
                        hidden,
                        &layer.post_norm,
                        self.geometry.epsilon,
                        normalized,
                    )?;
                }
            }
            check()?;
            project_tile_ffn_swiglu(
                &self.transforms,
                &self.geometry,
                self.arithmetic,
                self.adapter.as_ref(),
                self.adapter_scale,
                &mut self.workspace.adapter,
                &mut self.workspace.projection.matvec,
                &layer.gate,
                &layer.up,
                &tile.normalized[..count],
                &mut tile.gate[..count],
                &mut tile.up[..count],
                &mut tile.transformed[..count],
                &mut tile.projector,
                reborrow_executor(&mut executor),
                check,
            )?;
            check()?;
            let _ = project_tile(
                &self.transforms,
                &self.geometry,
                self.arithmetic,
                self.adapter.as_ref(),
                self.adapter_scale,
                &mut self.workspace.adapter,
                &mut self.workspace.projection.matvec,
                TileReuse::default(),
                &layer.down,
                &tile.up[..count],
                &mut tile.normalized[..count],
                &mut tile.transformed[..count],
                &mut tile.projector,
                reborrow_executor(&mut executor),
                check,
            )?;
            for (hidden, output) in tile.hidden[..count]
                .iter_mut()
                .zip(&tile.normalized[..count])
            {
                add(hidden, output);
                finite(hidden, "CPU tiled decoder layer output")?;
            }
            check()?;
        }
        Ok(())
    }
}

#[allow(
    clippy::too_many_arguments,
    clippy::too_many_lines,
    clippy::suboptimal_flops,
    reason = "Keep causal KV append ordering, exact FP16 score/value rounding and tile scratch explicit."
)]
fn full_attention_tile(
    weights: &FullAttention,
    g: &Geometry,
    arithmetic: Arithmetic,
    transforms: &Transforms,
    base: usize,
    count: usize,
    frequencies_ready: &mut bool,
    keys: &mut KvCache,
    values: &mut KvCache,
    projection: &mut super::ProjectionScratch,
    adapter: Option<&RankOneAdapter>,
    adapter_scale: f32,
    adapter_scratch: &mut AdapterScratch,
    attention_scratch: &mut super::AttentionScratch,
    tile: &mut TiledPrefill,
    check: &dyn Fn() -> Result<()>,
) -> Result<()> {
    if count <= 1 || tile.projector.row_parallelism().workers() <= 1 {
        for index in 0..count {
            check()?;
            if !*frequencies_ready {
                rope_frequencies(g, base + index, &mut tile.frequencies[index])?;
            }
            std::mem::swap(
                &mut attention_scratch.frequencies,
                &mut tile.frequencies[index],
            );
            let mut evaluation = Evaluation {
                geometry: g,
                arithmetic,
                transforms,
                position: base + index,
                frequencies_ready: true,
                projection,
                adapter,
                adapter_scale,
                adapter_scratch,
                executor: None,
            };
            let step_result = evaluation.full_attention_state(
                weights,
                &tile.mixed[index],
                &mut tile.recurrent_gate[index],
                &mut tile.beta[index],
                keys,
                values,
                &mut attention_scratch.query,
                &mut attention_scratch.gate,
                &mut attention_scratch.scores,
                &mut attention_scratch.frequencies,
                &mut tile.attended[index],
            );
            std::mem::swap(
                &mut attention_scratch.frequencies,
                &mut tile.frequencies[index],
            );
            step_result?;
            keys.reserve_append()?;
            values.reserve_append()?;
            keys.append(&tile.recurrent_gate[index]);
            values.append(&tile.beta[index]);
            check()?;
        }
        *frequencies_ready = true;
        return Ok(());
    }

    check()?;
    if !*frequencies_ready {
        for (index, frequencies) in tile.frequencies.iter_mut().enumerate().take(count) {
            rope_frequencies(g, base + index, frequencies)?;
        }
        *frequencies_ready = true;
    }
    let query_gate_len = g.head * g.heads;
    for index in 0..count {
        resize_into(query_gate_len, &mut tile.up[index])?;
        resize_into(query_gate_len, &mut tile.gate[index])?;
        resize_into(query_gate_len, &mut tile.attended[index])?;
    }
    let mixed_ref: &[Vec<f32>] = &tile.mixed[..count];
    let frequencies_ref: &[Vec<(f32, f32)>] = &tile.frequencies[..count];
    let query_norm: &[f32] = &weights.query_norm;
    tile.projector.row_parallelism().update_paired_chunks(
        &mut tile.up[..count],
        1,
        &mut tile.gate[..count],
        1,
        |base_tok, query_chunk, gate_chunk| {
            for (offset, (query_tok, gate_tok)) in query_chunk
                .iter_mut()
                .zip(gate_chunk.iter_mut())
                .enumerate()
            {
                let tok = base_tok + offset;
                let projected = &mixed_ref[tok];
                let frequencies = &frequencies_ref[tok];
                for (head_idx, (query_head, gate_head)) in query_tok
                    .chunks_exact_mut(g.head)
                    .zip(gate_tok.chunks_exact_mut(g.head))
                    .enumerate()
                {
                    let src = &projected[head_idx * 2 * g.head..(head_idx + 1) * 2 * g.head];
                    query_head.copy_from_slice(&src[..g.head]);
                    gate_head.copy_from_slice(&src[g.head..]);
                    arithmetic.normalize_in_place(query_head, query_norm, g.epsilon)?;
                    arithmetic.attention_rope(query_head, g, frequencies)?;
                    round_f16_slice(query_head);
                }
                finite(query_tok, "binary16 attention query")?;
            }
            Ok(())
        },
    )?;
    let key_norm: &[f32] = &weights.key_norm;
    tile.projector.row_parallelism().update_paired_chunks(
        &mut tile.recurrent_gate[..count],
        1,
        &mut tile.beta[..count],
        1,
        |base_tok, key_chunk, value_chunk| {
            for (offset, (key_tok, value_tok)) in
                key_chunk.iter_mut().zip(value_chunk.iter_mut()).enumerate()
            {
                let tok = base_tok + offset;
                let frequencies = &frequencies_ref[tok];
                for (key_head, value_head) in key_tok
                    .chunks_exact_mut(g.head)
                    .zip(value_tok.chunks_exact_mut(g.head))
                {
                    arithmetic.normalize_in_place(key_head, key_norm, g.epsilon)?;
                    arithmetic.attention_rope(key_head, g, frequencies)?;
                    round_f16_slice(key_head);
                    round_f16_slice(value_head);
                }
                finite(key_tok, "binary16 key cache")?;
                finite(value_tok, "binary16 value cache")?;
            }
            Ok(())
        },
    )?;
    for index in 0..count {
        if index > 0 {
            check()?;
        }
        keys.reserve_append()?;
        values.reserve_append()?;
        keys.append(&tile.recurrent_gate[index]);
        values.append(&tile.beta[index]);
        check()?;
    }
    let scale = 1.0 / small_float(g.head)?.sqrt();
    let max_score_len = base + count;
    let token_score_stride = 4 * max_score_len;
    resize_into(count * token_score_stride, &mut attention_scratch.scores)?;
    let query_ref: &[Vec<f32>] = &tile.up[..count];
    let gate_ref: &[Vec<f32>] = &tile.gate[..count];
    let current_keys: &[Vec<f32>] = &tile.recurrent_gate[..count];
    let current_values: &[Vec<f32>] = &tile.beta[..count];
    let key_cache: &KvCache = keys;
    let value_cache: &KvCache = values;
    let heads_per_kv = g.heads / g.kv_heads;
    tile.projector.row_parallelism().update_paired_chunks(
        &mut tile.attended[..count],
        1,
        &mut attention_scratch.scores,
        token_score_stride,
        |base_tok, attended_chunk, scores_chunk| {
            let chunk_tokens = attended_chunk.len();
            let mut head = 0;
            while head < g.heads {
                let kv_head = head / heads_per_kv;
                let offset = kv_head * g.head;
                if head + 4 <= g.heads && kv_head == (head + 3) / heads_per_kv {
                    for (pos, past) in key_cache
                        .tokens()
                        .take(base + base_tok + chunk_tokens.saturating_sub(1))
                        .enumerate()
                    {
                        let past_key = &past[offset..offset + g.head];
                        let first_t = (pos + 1).saturating_sub(base + base_tok);
                        for t in first_t..chunk_tokens {
                            let tok = base_tok + t;
                            let queries = std::array::from_fn(|i| {
                                let h = head + i;
                                &query_ref[tok][h * g.head..(h + 1) * g.head]
                            });
                            let dots = arithmetic.attention_dot_four(queries, past_key)?;
                            let tok_scores = &mut scores_chunk
                                [t * token_score_stride..(t + 1) * token_score_stride];
                            tok_scores[pos] = dots[0] * scale;
                            tok_scores[max_score_len + pos] = dots[1] * scale;
                            tok_scores[2 * max_score_len + pos] = dots[2] * scale;
                            tok_scores[3 * max_score_len + pos] = dots[3] * scale;
                        }
                    }
                    for (t, attended_tok) in attended_chunk.iter_mut().enumerate() {
                        let tok = base_tok + t;
                        let score_len = base + tok + 1;
                        let queries = std::array::from_fn(|i| {
                            let h = head + i;
                            &query_ref[tok][h * g.head..(h + 1) * g.head]
                        });
                        let current_dots = arithmetic.attention_current_dot_four(
                            queries,
                            &current_keys[tok][offset..offset + g.head],
                        )?;
                        let tok_scores =
                            &mut scores_chunk[t * token_score_stride..(t + 1) * token_score_stride];
                        let (s01, s23) = tok_scores.split_at_mut(2 * max_score_len);
                        let (s0, s1) = s01.split_at_mut(max_score_len);
                        let (s2, s3) = s23.split_at_mut(max_score_len);
                        let s0 = &mut s0[..score_len];
                        let s1 = &mut s1[..score_len];
                        let s2 = &mut s2[..score_len];
                        let s3 = &mut s3[..score_len];
                        s0[score_len - 1] = current_dots[0] * scale;
                        s1[score_len - 1] = current_dots[1] * scale;
                        s2[score_len - 1] = current_dots[2] * scale;
                        s3[score_len - 1] = current_dots[3] * scale;
                        for scores in [&mut *s0, &mut *s1, &mut *s2, &mut *s3] {
                            arithmetic.attention_softmax(scores)?;
                            round_f16_slice(scores);
                        }
                        let probs: [&[f32]; 4] = [s0, s1, s2, s3];
                        let (r01, r23) = attended_tok[head * g.head..(head + 4) * g.head]
                            .split_at_mut(2 * g.head);
                        let (r0, r1) = r01.split_at_mut(g.head);
                        let (r2, r3) = r23.split_at_mut(g.head);
                        for (tile_idx, (((out0, out1), out2), out3)) in r0
                            .chunks_mut(crate::kernels::pinned_attention::VALUE_TILE)
                            .zip(r1.chunks_mut(crate::kernels::pinned_attention::VALUE_TILE))
                            .zip(r2.chunks_mut(crate::kernels::pinned_attention::VALUE_TILE))
                            .zip(r3.chunks_mut(crate::kernels::pinned_attention::VALUE_TILE))
                            .enumerate()
                        {
                            let dimension = tile_idx * crate::kernels::pinned_attention::VALUE_TILE;
                            arithmetic.attention_value_tile_four(
                                value_cache,
                                &current_values[tok]
                                    [offset + dimension..offset + dimension + out0.len()],
                                offset + dimension,
                                probs,
                                [out0, out1, out2, out3],
                            )?;
                        }
                        for (i, result) in [r0, r1, r2, r3].into_iter().enumerate() {
                            let h = head + i;
                            for (out, &gate) in result
                                .iter_mut()
                                .zip(&gate_ref[tok][h * g.head..(h + 1) * g.head])
                            {
                                *out *= sigmoid(gate);
                            }
                        }
                    }
                    head += 4;
                    continue;
                }
                for (pos, past) in key_cache
                    .tokens()
                    .take(base + base_tok + chunk_tokens.saturating_sub(1))
                    .enumerate()
                {
                    let past_key = &past[offset..offset + g.head];
                    let first_t = (pos + 1).saturating_sub(base + base_tok);
                    for t in first_t..chunk_tokens {
                        let tok = base_tok + t;
                        let q = &query_ref[tok][head * g.head..(head + 1) * g.head];
                        scores_chunk[t * token_score_stride + pos] =
                            arithmetic.attention_dot(q, past_key)? * scale;
                    }
                }
                for (t, attended_tok) in attended_chunk.iter_mut().enumerate() {
                    let tok = base_tok + t;
                    let score_len = base + tok + 1;
                    let q = &query_ref[tok][head * g.head..(head + 1) * g.head];
                    let scores = &mut scores_chunk
                        [t * token_score_stride..t * token_score_stride + score_len];
                    scores[score_len - 1] = arithmetic
                        .attention_current_dot(q, &current_keys[tok][offset..offset + g.head])?
                        * scale;
                    arithmetic.attention_softmax(scores)?;
                    round_f16_slice(scores);
                    let result = &mut attended_tok[head * g.head..(head + 1) * g.head];
                    for (tile_idx, out) in result
                        .chunks_mut(crate::kernels::pinned_attention::VALUE_TILE)
                        .enumerate()
                    {
                        let dimension = tile_idx * crate::kernels::pinned_attention::VALUE_TILE;
                        arithmetic.attention_value_tile(
                            value_cache,
                            &current_values[tok]
                                [offset + dimension..offset + dimension + out.len()],
                            offset + dimension,
                            scores,
                            out,
                        )?;
                    }
                    for (out, &gate) in result
                        .iter_mut()
                        .zip(&gate_ref[tok][head * g.head..(head + 1) * g.head])
                    {
                        *out *= sigmoid(gate);
                    }
                }
                head += 1;
            }
            Ok(())
        },
    )
}

fn transform_tile_inputs<'a>(
    transforms: &Transforms,
    geometry: &Geometry,
    tensor: &Tensor,
    reuse_basis: bool,
    inputs: &'a [Vec<f32>],
    transformed: &'a mut [Vec<f32>],
    parallelism: &RowParallelism,
) -> Result<&'a [Vec<f32>]> {
    let basis = transforms.basis(tensor);
    if basis.0 {
        if !reuse_basis {
            if inputs.len() > 1 && parallelism.workers() > 1 {
                parallelism.update_chunks(transformed, 1, |base, chunk| {
                    for (offset, target) in chunk.iter_mut().enumerate() {
                        transforms.transform_into(
                            tensor,
                            &inputs[base + offset],
                            geometry,
                            target,
                        )?;
                    }
                    Ok(())
                })?;
            } else {
                for (input, transformed) in inputs.iter().zip(transformed.iter_mut()) {
                    transforms.transform_into_with_parallelism(
                        tensor,
                        input,
                        geometry,
                        parallelism,
                        transformed,
                    )?;
                }
            }
        }
        Ok(&*transformed)
    } else {
        Ok(inputs)
    }
}

#[allow(
    clippy::too_many_arguments,
    reason = "Keep exact transforms, original-basis adapter and bounded tile storage explicit."
)]
fn project_tile_pair<'a>(
    transforms: &Transforms,
    geometry: &Geometry,
    arithmetic: Arithmetic,
    adapter: Option<&RankOneAdapter>,
    adapter_scale: f32,
    adapter_scratch: &mut AdapterScratch,
    projection_scratch: &mut crate::kernels::MatvecScratch,
    tensors: [&'a Tensor; 2],
    inputs: &[Vec<f32>],
    outputs_a: &mut [Vec<f32>],
    outputs_b: &mut [Vec<f32>],
    transformed: &mut [Vec<f32>],
    projector: &mut ProjectionTile,
    mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    check: &dyn Fn() -> Result<()>,
) -> Result<TileReuse<'a>> {
    if let Some(exec) = reborrow_executor(&mut executor)
        && tensors[0].kind == 142
        && tensors[1].kind == 142
        && transforms.basis(tensors[0]) == transforms.basis(tensors[1])
        && exec.supports_f32_group(&tensors)
    {
        let activation = transform_tile_inputs(
            transforms,
            geometry,
            tensors[0],
            false,
            inputs,
            transformed,
            projector.row_parallelism(),
        )?;
        check()?;
        exec.project_f32_pair_tokens(
            tensors,
            activation,
            matches!(arithmetic, Arithmetic::PinnedSwift),
            outputs_a,
            outputs_b,
        )?;
        let rows_a = tensors[0]
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        let rows_b = tensors[1]
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        for (out_a, out_b) in outputs_a.iter().zip(outputs_b.iter()) {
            if out_a.len() != rows_a || out_b.len() != rows_b {
                return Err(invalid("projection executor returned incomplete output"));
            }
            finite(out_a, "projection executor output")?;
            finite(out_b, "projection executor output")?;
        }
        if let Some(adapter) = adapter {
            for (input, output) in inputs.iter().zip(outputs_a) {
                adapter.linear(
                    &tensors[0].name,
                    input,
                    adapter_scale,
                    output,
                    adapter_scratch,
                )?;
            }
            for (input, output) in inputs.iter().zip(outputs_b) {
                adapter.linear(
                    &tensors[1].name,
                    input,
                    adapter_scale,
                    output,
                    adapter_scratch,
                )?;
            }
        }
        return Ok(TileReuse {
            basis: Some(tensors[1]),
            prepared: None,
        });
    }
    let reuse = project_tile(
        transforms,
        geometry,
        arithmetic,
        adapter,
        adapter_scale,
        adapter_scratch,
        projection_scratch,
        TileReuse::default(),
        tensors[0],
        inputs,
        outputs_a,
        transformed,
        projector,
        reborrow_executor(&mut executor),
        check,
    )?;
    check()?;
    project_tile(
        transforms,
        geometry,
        arithmetic,
        adapter,
        adapter_scale,
        adapter_scratch,
        projection_scratch,
        reuse,
        tensors[1],
        inputs,
        outputs_b,
        transformed,
        projector,
        reborrow_executor(&mut executor),
        check,
    )
}

#[allow(
    clippy::too_many_arguments,
    reason = "Keep exact transforms, original-basis adapter and bounded tile storage explicit."
)]
fn project_tile_ffn_swiglu(
    transforms: &Transforms,
    geometry: &Geometry,
    arithmetic: Arithmetic,
    adapter: Option<&RankOneAdapter>,
    adapter_scale: f32,
    adapter_scratch: &mut AdapterScratch,
    projection_scratch: &mut crate::kernels::MatvecScratch,
    gate: &Tensor,
    up: &Tensor,
    inputs: &[Vec<f32>],
    gate_outputs: &mut [Vec<f32>],
    up_outputs: &mut [Vec<f32>],
    transformed: &mut [Vec<f32>],
    projector: &mut ProjectionTile,
    mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    check: &dyn Fn() -> Result<()>,
) -> Result<()> {
    if adapter.is_none()
        && let Some(exec) = reborrow_executor(&mut executor)
        && gate.kind == 142
        && up.kind == 142
        && transforms.basis(gate) == transforms.basis(up)
        && exec.supports_f32_ffn_swiglu(gate, up)
    {
        let activation = transform_tile_inputs(
            transforms,
            geometry,
            gate,
            false,
            inputs,
            transformed,
            projector.row_parallelism(),
        )?;
        check()?;
        exec.project_f32_ffn_swiglu_tokens(
            gate,
            up,
            activation,
            matches!(arithmetic, Arithmetic::PinnedSwift),
            up_outputs,
        )?;
        let rows = up
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        for out in up_outputs.iter() {
            if out.len() != rows {
                return Err(invalid("projection executor returned incomplete output"));
            }
            finite(out, "projection executor output")?;
        }
        return Ok(());
    }
    let _ = project_tile_pair(
        transforms,
        geometry,
        arithmetic,
        adapter,
        adapter_scale,
        adapter_scratch,
        projection_scratch,
        [up, gate],
        inputs,
        up_outputs,
        gate_outputs,
        transformed,
        projector,
        executor,
        check,
    )?;
    let count = inputs.len();
    if count > 1 && projector.row_parallelism().workers() > 1 {
        let gate_ref = &gate_outputs[..count];
        projector
            .row_parallelism()
            .update_chunks(&mut up_outputs[..count], 1, |base_tok, chunk| {
                for (offset, up_out) in chunk.iter_mut().enumerate() {
                    arithmetic.swiglu_in_place(&gate_ref[base_tok + offset], up_out)?;
                }
                Ok(())
            })
    } else {
        for (up_out, gate_out) in up_outputs[..count].iter_mut().zip(&gate_outputs[..count]) {
            arithmetic.swiglu_in_place_with_parallelism(
                gate_out,
                up_out,
                projector.row_parallelism(),
            )?;
        }
        Ok(())
    }
}

#[allow(
    clippy::too_many_arguments,
    clippy::too_many_lines,
    reason = "Keep exact transforms, original-basis adapter and bounded tile storage explicit."
)]
fn project_tile_triple<'a>(
    transforms: &Transforms,
    geometry: &Geometry,
    arithmetic: Arithmetic,
    adapter: Option<&RankOneAdapter>,
    adapter_scale: f32,
    adapter_scratch: &mut AdapterScratch,
    projection_scratch: &mut crate::kernels::MatvecScratch,
    tensors: [&'a Tensor; 3],
    inputs: &[Vec<f32>],
    outputs_a: &mut [Vec<f32>],
    outputs_b: &mut [Vec<f32>],
    outputs_c: &mut [Vec<f32>],
    transformed: &mut [Vec<f32>],
    projector: &mut ProjectionTile,
    mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    check: &dyn Fn() -> Result<()>,
) -> Result<TileReuse<'a>> {
    if let Some(exec) = reborrow_executor(&mut executor)
        && tensors[0].kind == 142
        && tensors[1].kind == 142
        && tensors[2].kind == 142
        && transforms.basis(tensors[0]) == transforms.basis(tensors[1])
        && transforms.basis(tensors[0]) == transforms.basis(tensors[2])
        && exec.supports_f32_group(&tensors)
    {
        let activation = transform_tile_inputs(
            transforms,
            geometry,
            tensors[0],
            false,
            inputs,
            transformed,
            projector.row_parallelism(),
        )?;
        check()?;
        exec.project_f32_triple_tokens(
            tensors,
            activation,
            matches!(arithmetic, Arithmetic::PinnedSwift),
            outputs_a,
            outputs_b,
            outputs_c,
        )?;
        let rows_a = tensors[0]
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        let rows_b = tensors[1]
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        let rows_c = tensors[2]
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        for ((out_a, out_b), out_c) in outputs_a.iter().zip(outputs_b.iter()).zip(outputs_c.iter())
        {
            if out_a.len() != rows_a || out_b.len() != rows_b || out_c.len() != rows_c {
                return Err(invalid("projection executor returned incomplete output"));
            }
            finite(out_a, "projection executor output")?;
            finite(out_b, "projection executor output")?;
            finite(out_c, "projection executor output")?;
        }
        if let Some(adapter) = adapter {
            for (tensor, outputs) in [
                (tensors[0], &mut *outputs_a),
                (tensors[1], &mut *outputs_b),
                (tensors[2], &mut *outputs_c),
            ] {
                for (input, output) in inputs.iter().zip(outputs) {
                    adapter.linear(&tensor.name, input, adapter_scale, output, adapter_scratch)?;
                }
            }
        }
        return Ok(TileReuse {
            basis: Some(tensors[2]),
            prepared: None,
        });
    }
    let reuse = project_tile_pair(
        transforms,
        geometry,
        arithmetic,
        adapter,
        adapter_scale,
        adapter_scratch,
        projection_scratch,
        [tensors[0], tensors[1]],
        inputs,
        outputs_a,
        outputs_b,
        transformed,
        projector,
        reborrow_executor(&mut executor),
        check,
    )?;
    check()?;
    project_tile(
        transforms,
        geometry,
        arithmetic,
        adapter,
        adapter_scale,
        adapter_scratch,
        projection_scratch,
        reuse,
        tensors[2],
        inputs,
        outputs_c,
        transformed,
        projector,
        reborrow_executor(&mut executor),
        check,
    )
}

#[allow(
    clippy::too_many_arguments,
    reason = "Keep exact transforms, original-basis adapter and bounded tile storage explicit."
)]
fn project_tile<'a>(
    transforms: &Transforms,
    geometry: &Geometry,
    arithmetic: Arithmetic,
    adapter: Option<&RankOneAdapter>,
    adapter_scale: f32,
    adapter_scratch: &mut AdapterScratch,
    projection_scratch: &mut crate::kernels::MatvecScratch,
    previous: TileReuse<'a>,
    tensor: &'a Tensor,
    inputs: &[Vec<f32>],
    outputs: &mut [Vec<f32>],
    transformed: &mut [Vec<f32>],
    projector: &mut ProjectionTile,
    mut executor: Option<&mut dyn crate::kernels::Pq2ProjectionExecutor>,
    check: &dyn Fn() -> Result<()>,
) -> Result<TileReuse<'a>> {
    let basis = transforms.basis(tensor);
    let reuse_basis = previous
        .basis
        .is_some_and(|prev| transforms.basis(prev) == basis);
    let reuse_prepared = reuse_basis
        && previous
            .prepared
            .is_some_and(|prev| transforms.basis(prev) == basis);
    let activation = transform_tile_inputs(
        transforms,
        geometry,
        tensor,
        reuse_basis,
        inputs,
        transformed,
        projector.row_parallelism(),
    )?;
    let prepared = if let Some(exec) = reborrow_executor(&mut executor) {
        let cols = tensor.dimensions.first().copied().unwrap_or(0);
        let rows = tensor
            .dimensions
            .get(1..)
            .map_or(0, |dims| dims.iter().product::<usize>());
        if tensor.kind == 142 && cols.is_multiple_of(128) && exec.supports_f32(tensor) {
            check()?;
            exec.project_f32_tokens(
                tensor,
                activation,
                matches!(arithmetic, Arithmetic::PinnedSwift),
                outputs,
            )?;
            for out in outputs.iter() {
                if out.len() != rows {
                    return Err(invalid("projection executor returned incomplete output"));
                }
                finite(out, "projection executor output")?;
            }
            None
        } else if (tensor.kind == 142 && cols.is_multiple_of(256) && exec.supports(tensor))
            || ((tensor.kind == 143 || (tensor.kind == 142 && !cols.is_multiple_of(256)))
                && exec.supports_q8(tensor))
        {
            for (input, output) in activation.iter().zip(outputs.iter_mut()) {
                check()?;
                projection_scratch
                    .prepare(tensor, input)?
                    .matvec_with_executor(tensor, output, exec)?;
            }
            None
        } else {
            projector.project_with_previous(
                previous.prepared.filter(|_| reuse_prepared),
                tensor,
                activation,
                outputs,
                matches!(arithmetic, Arithmetic::PinnedSwift),
                check,
            )?;
            Some(tensor)
        }
    } else {
        projector.project_with_previous(
            previous.prepared.filter(|_| reuse_prepared),
            tensor,
            activation,
            outputs,
            matches!(arithmetic, Arithmetic::PinnedSwift),
            check,
        )?;
        Some(tensor)
    };
    if let Some(adapter) = adapter {
        for (input, output) in inputs.iter().zip(outputs) {
            adapter.linear(&tensor.name, input, adapter_scale, output, adapter_scratch)?;
        }
    }
    Ok(TileReuse {
        basis: Some(tensor),
        prepared,
    })
}