aprender-serve 0.70.2

use crate::error::Result;
use crate::gguf::quantized::{OwnedQuantizedTensor, QuantizedTensorRef};
use crate::gguf::{GGUFConfig, GGUFModel, OwnedQuantizedModel};
use std::f32::consts::E;

/// 0.69.1 load-time warning for Qwen3.5-0.8B (#4032, known issue #4030).
#[path = "qwen35_known_issue.rs"]
pub mod qwen35_known_issue;

/// SiLU activation function
pub fn silu(x: f32) -> f32 {
    x / (1.0 + (-x as f32).exp())
}

/// Softplus activation function
pub fn softplus(x: f32) -> f32 {
    if x > 20.0 {
        x
    } else {
        (1.0 + x.exp()).ln()
    }
}

/// L2 Normalization with epsilon over the given slice.
pub fn l2_norm(x: &mut [f32], eps: f32) {
    let mut sq_sum = 0.0;
    for &v in x.iter() {
        sq_sum += v * v;
    }
    let scale = 1.0 / (sq_sum + eps).sqrt();
    for v in x.iter_mut() {
        *v *= scale;
    }
}

/// Scale each `head_dim`-wide head of `x` to unit L2 norm on its own, as llama.cpp's
/// `build_gdn_l2_norm` does for Gated `DeltaNet` q and k (`rms_norm(x, eps / n) / sqrt(n)`
/// over one head is `x_h / sqrt(sum(x_h^2) + eps)`). One norm over all heads scales every
/// head by the other heads' magnitudes.
pub fn l2_norm_per_head(x: &mut [f32], head_dim: usize, eps: f32) {
    for head in x.chunks_exact_mut(head_dim) {
        l2_norm(head, eps);
    }
}

/// Multiply `x` elementwise by `sigmoid(gate)`: the output gate of Qwen3.5's full-attention
/// layers (llama.cpp's `attn_gated` node), applied before the output projection.
pub fn apply_sigmoid_gate(x: &mut [f32], gate: &[f32]) {
    for (o, g) in x.iter_mut().zip(gate) {
        *o *= 1.0 / (1.0 + (-*g).exp());
    }
}

/// Gated RMSNorm: (RMSNorm(input) * silu(gate))
/// Norm is computed independently over chunks of size `head_v_dim`.
pub fn gated_rmsnorm(
    input: &[f32],
    gate: &[f32],
    weight: &[f32],
    eps: f32,
    head_v_dim: usize,
    output: &mut [f32],
) {
    assert_eq!(input.len(), gate.len());
    assert_eq!(input.len(), output.len());
    assert_eq!(weight.len(), head_v_dim);
    assert_eq!(input.len() % head_v_dim, 0);

    for (chunk_in, (chunk_gate, chunk_out)) in input.chunks_exact(head_v_dim).zip(
        gate.chunks_exact(head_v_dim)
            .zip(output.chunks_exact_mut(head_v_dim)),
    ) {
        let mut sq_sum = 0.0;
        for &v in chunk_in {
            sq_sum += v * v;
        }
        let rms_scale = 1.0 / ((sq_sum / head_v_dim as f32) + eps).sqrt();

        for i in 0..head_v_dim {
            let norm_v = chunk_in[i] * rms_scale * weight[i];
            chunk_out[i] = norm_v * silu(chunk_gate[i]);
        }
    }
}

/// Causal Conv1d over seq for a single time step.
/// state is `(kernel_size - 1) * channels` floats.
/// For each channel, we shift the past inputs left and insert the new input.
pub fn causal_conv1d(
    input: &[f32],
    state: &mut [f32],
    weight: &[f32],
    kernel_size: usize,
    channels: usize,
    output: &mut [f32],
) {
    assert_eq!(input.len(), channels);
    assert_eq!(state.len(), (kernel_size - 1) * channels);
    assert_eq!(weight.len(), kernel_size * channels);
    assert_eq!(output.len(), channels);

    for c in 0..channels {
        let mut sum = 0.0;
        let s_offset = c * (kernel_size - 1);
        let w_offset = c * kernel_size;

        for k in 0..(kernel_size - 1) {
            sum += state[s_offset + k] * weight[w_offset + k];
        }
        sum += input[c] * weight[w_offset + kernel_size - 1];

        // shift left (oldest is at 0, newest is at kernel_size - 2)
        for k in 0..(kernel_size - 2) {
            state[s_offset + k] = state[s_offset + k + 1];
        }
        if kernel_size > 1 {
            state[s_offset + kernel_size - 2] = input[c];
        }

        output[c] = sum;
    }
}

/// The gated delta-rule recurrence for a single token exactly as delta-net-base.cpp computes it,
/// for a file whose Gated `DeltaNet` has one key/query head per value head
/// (`num_k_heads == num_v_heads`, `head_k_dim == head_v_dim`) — Qwen3.5-0.8B and -2B.
///
/// This is [`delta_rule_recurrence_gqa`] at ratio 1; it is kept as its own entry point so the
/// 0.8B call sites and the GPU parity tests read unchanged.
pub fn delta_rule_recurrence(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    beta: &[f32],
    gate: &[f32],
    state: &mut [f32],
    output: &mut [f32],
    num_v_heads: usize,
    head_v_dim: usize,
) {
    delta_rule_recurrence_gqa(
        q,
        k,
        v,
        beta,
        gate,
        state,
        output,
        num_v_heads,
        head_v_dim,
        num_v_heads,
        head_v_dim,
    );
}

/// The gated delta-rule recurrence for a single token, with **grouped** key/query heads
/// (PMAT-3477, #3346/#3510).
///
/// Qwen3.5 4B and 9B carry `linear_num_value_heads = 32` against `linear_num_key_heads = 16`,
/// and 27B carries 48 against 16, so `q` and `k` are narrower than `v`. A **GGUF** file shares
/// them TILED:
///
/// > **value head `h` reads key/query head `h % num_k_heads`.** With `num_k_heads = 16` and
/// > `num_v_heads = 32`, value heads 0 and 16 share key head 0, value heads 1 and 17 share key
/// > head 1, and so on.
///
/// This is not the HF checkpoint's own order, and that is the trap. HF
/// (`modeling_qwen3_next.py::GatedDeltaNet`) stores the value heads grouped by key head and
/// expands q/k with `repeat_interleave` (`h / ratio`); the **conversion permutes the value
/// heads out of that order**, and everything downstream is written for the permuted file:
///
/// * `llama.cpp/conversion/qwen.py:455-464` — `_LinearAttentionVReorderBase`, which
///   `Qwen3_5TextModel` (line 639, `MODEL_ARCH.QWEN35`) is built from: *"reorders V heads from
///   grouped to tiled order for ggml broadcast … The HF weights store V heads grouped by K
///   head … ggml binary ops use tiled broadcast … We reorder V heads to tiled order so
///   `ggml_repeat` can replace the expensive interleaved repeat"*. `modify_tensors`
///   (lines 568-613) permutes **every** v-head-indexed tensor together — the v rows of
///   `in_proj_qkv`, `in_proj_z` (the gate), `in_proj_a`/`in_proj_b` (dt and beta),
///   `A_log`/`dt_bias`, the v channels of `conv1d`, and the v columns of `out_proj` — so the
///   loader needs no compensating permutation anywhere: only this index changes.
/// * `llama.cpp/src/models/qwen35.cpp:436-441` — the graph expands q and k with
///   `ggml_repeat_4d`, which tiles.
/// * `llama.cpp/ggml/src/ggml-cpu/ops.cpp:10976-10977` — the fused kernel that skips that
///   repeat reads `iq1 = iv1 % neq1; ik1 = iv1 % nek1;`.
///
/// MEASURED, so that nobody re-opens this from a text sample: with this mapping the 4B CPU
/// forward agrees with `llama-eval-callback` (`-ub 1`, autoregressive graph, same token
/// stream) on **every** `DeltaNet` intermediate of layer 0 at position 0 — `attn_norm`
/// -26.4532 vs -26.4532, `q_conv_predelta` 1.4485 vs 1.4614, `k_conv_predelta` 13.3249 vs
/// 13.3100, `gate` -18.2186 vs -18.2173, `beta_sigmoid` 21.0785 vs 21.0794, `attn_output`
/// 0.9502 vs 0.9535, `new_state` 38.7674 vs 38.7736, `l_out` 0.2917 vs 0.3056 (tensor sums;
/// the residual is the Q4_K/Q5_K activation quantisation, which is ~0.3% on a 8192-wide
/// sum). The *incoherent* 4B/9B text that survived this mapping was never a DeltaNet defect
/// at all — [`crate::gguf::qwen35_load::load_qwen35_layers`] was reading a bare
/// `block_count` key that never matched, so it built 24 layers for a 32-block file.
///
/// Shapes (this is the contract the CUDA `DeltaRuleRecurrenceKernel` / `Qwen35CudaModel`
/// mirror — the kernel's `num_v_heads`/`head_v_dim` pair gains `num_k_heads`/`head_k_dim` and
/// the same `h % num_k_heads` read):
///
/// | argument | length |
/// |---|---|
/// | `q`, `k` | `num_k_heads * head_k_dim` |
/// | `v`, `output` | `num_v_heads * head_v_dim` |
/// | `beta`, `gate` (dt) | `num_v_heads` — every gate is per VALUE head |
/// | `state` | `num_v_heads * head_v_dim * head_k_dim` |
///
/// The recurrent state of value head `h` is `S ∈ R^(head_k_dim × head_v_dim)` laid out so that
/// `S[i][j] = state[h * head_v_dim * head_k_dim + j * head_k_dim + i]` — memory row `j` is
/// column `j` of `S`, `i` runs over the key dim and `j` over the value dim. Every Qwen3.5 size
/// ships `head_k_dim == head_v_dim == 128`, so the block is square in practice; the two dims
/// are kept distinct here so a future file with a rectangular state is a shape, not a rewrite.
///
/// The scale is `1/sqrt(head_k_dim)` (the query/key width, as `fla`'s
/// `chunk_gated_delta_rule` defaults it) — identical to the previous
/// `1/sqrt(head_v_dim)` on every file that exists today.
///
/// # Panics
/// If any slice length disagrees with the table above, or if `num_v_heads` is not a positive
/// multiple of `num_k_heads`.
pub fn delta_rule_recurrence_gqa(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    beta: &[f32],
    gate: &[f32],
    state: &mut [f32],
    output: &mut [f32],
    num_k_heads: usize,
    head_k_dim: usize,
    num_v_heads: usize,
    head_v_dim: usize,
) {
    assert!(
        num_k_heads > 0 && num_v_heads % num_k_heads == 0,
        "Gated DeltaNet: num_v_heads ({num_v_heads}) must be a positive multiple of num_k_heads \
         ({num_k_heads})"
    );
    assert_eq!(q.len(), num_k_heads * head_k_dim);
    assert_eq!(k.len(), num_k_heads * head_k_dim);
    assert_eq!(v.len(), num_v_heads * head_v_dim);
    assert_eq!(beta.len(), num_v_heads);
    assert_eq!(gate.len(), num_v_heads);
    assert_eq!(state.len(), num_v_heads * head_v_dim * head_k_dim);
    assert_eq!(output.len(), num_v_heads * head_v_dim);

    let scale = 1.0 / (head_k_dim as f32).sqrt();

    for h in 0..num_v_heads {
        // ggml_repeat (tiled): value head h reads key/query head h % num_k_heads. The GGUF
        // conversion already permuted the value heads into this order — see the doc comment.
        let kh = h % num_k_heads;
        let q_h = &q[kh * head_k_dim..(kh + 1) * head_k_dim];
        let k_h = &k[kh * head_k_dim..(kh + 1) * head_k_dim];
        let v_h = &v[h * head_v_dim..(h + 1) * head_v_dim];
        let beta_val = beta[h];
        let gate_val = gate[h];

        let state_stride = head_v_dim * head_k_dim;
        let state_offset = h * state_stride;
        delta_rule_head(
            q_h,
            k_h,
            v_h,
            beta_val,
            gate_val,
            &mut state[state_offset..state_offset + state_stride],
            &mut output[h * head_v_dim..(h + 1) * head_v_dim],
            head_k_dim,
            scale,
        );
    }
}

/// One value head of [`delta_rule_recurrence_gqa`] for one token: `s_h` is that head's
/// `head_v_dim × head_k_dim` state and `out_h` its `head_v_dim` outputs. The per-token path and
/// the batched prefill's per-head walk (#4228) both run this body, so their arithmetic is the same.
#[allow(clippy::too_many_arguments)]
#[inline]
pub(crate) fn delta_rule_head(
    q_h: &[f32],
    k_h: &[f32],
    v_h: &[f32],
    beta_val: f32,
    gate_val: f32,
    s_h: &mut [f32],
    out_h: &mut [f32],
    head_k_dim: usize,
    scale: f32,
) {
    let head_v_dim = out_h.len();
    // 1. S_h *= exp(gate_val)
    let exp_gate = gate_val.exp();
    for s in s_h.iter_mut() {
        *s *= exp_gate;
    }

    // 2. delta = (v_h - S_h^T * k_h) * beta_val
    // Note: s_h[j * head_k_dim + i] is S[i][j].
    // So row j of s_h in memory is column j of S.
    // sum = dot(row j of s_h, k_h)
    let mut delta = vec![0.0; head_v_dim];
    for j in 0..head_v_dim {
        let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim];
        let mut sum = 0.0;
        for i in 0..head_k_dim {
            sum += row_j[i] * k_h[i];
        }
        delta[j] = (v_h[j] - sum) * beta_val;
    }

    // 3. S_h += k_h * delta^T
    for j in 0..head_v_dim {
        let row_j = &mut s_h[j * head_k_dim..(j + 1) * head_k_dim];
        let d_j = delta[j];
        for i in 0..head_k_dim {
            row_j[i] += k_h[i] * d_j;
        }
    }

    // 4. out_h = S_h^T * q_h * scale
    for j in 0..head_v_dim {
        let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim];
        let mut sum = 0.0;
        for i in 0..head_k_dim {
            sum += row_j[i] * q_h[i];
        }
        out_h[j] = sum * scale;
    }
}

/// Partial NEOX RoPE, exactly as llama.cpp ggml_rope_multi for text input (every position stream is the
/// token position, so the sections do not change theta): for pair j in 0..n_rot/2 of each head,
/// (x[j], x[j + n_rot/2]) rotates by theta_j = pos * theta_scale^j with theta_scale = freq_base^(-2/n_rot),
/// computed iteratively as ggml does; dims n_rot..head_dim pass through unrotated.
pub fn apply_partial_neox_rope(
    x: &mut [f32],
    num_heads: usize,
    head_dim: usize,
    n_rot: usize,
    pos: usize,
    freq_base: f32,
) {
    let half = n_rot / 2;
    let theta_scale = freq_base.powf(-2.0 / n_rot as f32);
    for h in 0..num_heads {
        let base = h * head_dim;
        let mut theta = pos as f32;
        for j in 0..half {
            let (sin, cos) = theta.sin_cos();
            let (a, b) = (x[base + j], x[base + j + half]);
            x[base + j] = a * cos - b * sin;
            x[base + j + half] = a * sin + b * cos;
            theta *= theta_scale;
        }
    }
}

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

    #[test]
    fn test_silu() {
        assert!((silu(0.0) - 0.0).abs() < 1e-6);
        assert!((silu(1.0) - (1.0 / (1.0 + (-1.0_f32).exp()))).abs() < 1e-6);
    }

    #[test]
    fn test_l2_norm() {
        let mut x = [1.0, 2.0, 2.0];
        l2_norm(&mut x, 0.0);
        assert!((x[0] - 1.0 / 3.0).abs() < 1e-6);
        assert!((x[1] - 2.0 / 3.0).abs() < 1e-6);
        assert!((x[2] - 2.0 / 3.0).abs() < 1e-6);
    }

    #[test]
    fn test_causal_conv1d() {
        let input = [1.0, 2.0];
        let mut state = [0.1, 0.2, 0.3, 0.4]; // 2 past states per channel. Channels = 2.
        let weight = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; // kernel_size = 3. format: [w_0, w_1, w_2] per channel. Wait, weight memory is [kernel_size * channels].
                                                     // Is it channel outer or time outer?
                                                     // causal_conv1d uses `w_offset = c * kernel_size`, so channel outer.
                                                     // weight for c=0: [1.0, 2.0, 3.0]. weight for c=1: [4.0, 5.0, 6.0].
                                                     // state for c=0: [0.1, 0.2]. state for c=1: [0.3, 0.4].

        let mut output = [0.0, 0.0];
        causal_conv1d(&input, &mut state, &weight, 3, 2, &mut output);

        // c = 0: sum = 0.1 * 1.0 + 0.2 * 2.0 + 1.0 * 3.0 = 0.1 + 0.4 + 3.0 = 3.5
        assert!((output[0] - 3.5).abs() < 1e-5);
        // c = 1: sum = 0.3 * 4.0 + 0.4 * 5.0 + 2.0 * 6.0 = 1.2 + 2.0 + 12.0 = 15.2
        assert!((output[1] - 15.2).abs() < 1e-5);

        // State should shift: [0.2, 1.0, 0.4, 2.0]
        assert!((state[0] - 0.2).abs() < 1e-5);
        assert!((state[1] - 1.0).abs() < 1e-5);
        assert!((state[2] - 0.4).abs() < 1e-5);
        assert!((state[3] - 2.0).abs() < 1e-5);
    }

    #[test]
    fn test_delta_rule_recurrence() {
        let num_v_heads = 1;
        let head_v_dim = 2;
        let q = [1.0, 2.0];
        let k = [0.5, 0.5];
        let v = [1.0, -1.0];
        let beta = [0.5];
        let gate = [0.0]; // exp(0) = 1.0

        // State S is 2x2.
        // s_h[j * 2 + i] is S[i][j].
        // Memory: [S[0][0], S[1][0], S[0][1], S[1][1]]
        // Let S be identity.
        let mut state = [
            1.0, 0.0, // column 0
            0.0, 1.0, // column 1
        ];
        let mut output = [0.0, 0.0];

        delta_rule_recurrence(
            &q,
            &k,
            &v,
            &beta,
            &gate,
            &mut state,
            &mut output,
            num_v_heads,
            head_v_dim,
        );

        // 1. S_h *= 1.0 -> same
        // 2. delta[j] = (v[j] - S^T * k) * 0.5
        // S^T * k = [1 0; 0 1] * [0.5, 0.5] = [0.5, 0.5]
        // v - S^T k = [0.5, -1.5]
        // delta = [0.25, -0.75]
        // 3. S += k * delta^T
        // k * delta^T = [0.5; 0.5] * [0.25, -0.75] = [0.125, -0.375; 0.125, -0.375]
        // New S = [1.125, -0.375; 0.125, 0.625]
        // 4. out = S^T * q * scale (scale = 1/sqrt(2) = 0.7071)
        // S^T = [1.125, 0.125; -0.375, 0.625]
        // S^T * [1.0, 2.0] = [1.125 + 0.25, -0.375 + 1.25] = [1.375, 0.875]
        // out = [1.375 / sqrt(2), 0.875 / sqrt(2)]

        let s2 = 2.0_f32.sqrt();
        assert!((output[0] - 1.375 / s2).abs() < 1e-5);
        assert!((output[1] - 0.875 / s2).abs() < 1e-5);

        // Check state
        assert!((state[0] - 1.125).abs() < 1e-5); // S[0][0]
        assert!((state[1] - 0.125).abs() < 1e-5); // S[1][0]
        assert!((state[2] - -0.375).abs() < 1e-5); // S[0][1]
        assert!((state[3] - 0.625).abs() < 1e-5); // S[1][1]
    }
}

/// Per-sequence decode state for [`Qwen35Model`]: the attention layers' KV cache plus each
/// Gated `DeltaNet` layer's causal-conv window and recurrent state. Build one with
/// [`Qwen35Model::new_state`].
pub struct Qwen35State {
    pub(crate) conv_states: Vec<Vec<f32>>,
    pub(crate) ssm_states: Vec<Vec<f32>>,
    pub(crate) kv_cache: crate::gguf::OwnedQuantizedKVCache,
}

impl Qwen35State {
    pub(crate) fn new(
        num_layers: usize,
        max_seq_len: usize,
        head_dim: usize,
        num_kv_heads: usize,
        num_k_heads: usize,
        head_k_dim: usize,
        num_v_heads: usize,
        head_value_dim: usize,
    ) -> Self {
        let convalue_dim = head_k_dim * num_k_heads * 2 + head_value_dim * num_v_heads;
        Self {
            conv_states: vec![vec![0.0; convalue_dim * 3]; num_layers],
            // One [head_k_dim x head_value_dim] recurrent state per VALUE head — the key width
            // is the state's row length, which is only the same as the value width because
            // every Qwen3.5 size ships head_k_dim == head_v_dim (PMAT-3477).
            ssm_states: vec![vec![0.0; num_v_heads * head_value_dim * head_k_dim]; num_layers],
            kv_cache: crate::gguf::OwnedQuantizedKVCache::new(
                num_layers,
                num_kv_heads * head_dim,
                max_seq_len,
            ),
        }
    }

    /// Return the state to position 0 in place (#3595): every conv window and
    /// recurrent state zeroed and the KV cache emptied — what
    /// [`Qwen35Model::new_state`] starts from, without reallocating it.
    pub fn reset(&mut self) {
        for v in self
            .conv_states
            .iter_mut()
            .chain(self.ssm_states.iter_mut())
        {
            v.fill(0.0);
        }
        self.kv_cache.reset();
    }
}

/// A copy of a [`Qwen35State`]'s recurrent half at one position (#4214): the
/// conv windows and recurrent states, which cannot be rolled back, plus how
/// many KV rows it held — the rows themselves stay in the state, where a later
/// position never overwrites them.
#[derive(Debug, Clone)]
pub struct Qwen35Checkpoint {
    conv_states: Vec<Vec<f32>>,
    ssm_states: Vec<Vec<f32>>,
    kv_len: usize,
}

impl Qwen35State {
    /// Copy the state as it holds now.
    #[must_use]
    pub fn checkpoint(&self) -> Qwen35Checkpoint {
        Qwen35Checkpoint {
            conv_states: self.conv_states.clone(),
            ssm_states: self.ssm_states.clone(),
            kv_len: self.kv_cache.len(),
        }
    }

    /// Return to `checkpoint`. `false` (the state unchanged) when the state no
    /// longer holds the KV rows the copy was taken over.
    pub fn restore(&mut self, checkpoint: &Qwen35Checkpoint) -> bool {
        if self.kv_cache.len() < checkpoint.kv_len
            || self.conv_states.len() != checkpoint.conv_states.len()
            || self.ssm_states.len() != checkpoint.ssm_states.len()
        {
            return false;
        }
        for (dst, src) in self
            .conv_states
            .iter_mut()
            .zip(&checkpoint.conv_states)
            .chain(self.ssm_states.iter_mut().zip(&checkpoint.ssm_states))
        {
            dst.copy_from_slice(src);
        }
        self.kv_cache.truncate(checkpoint.kv_len);
        true
    }
}

pub(crate) struct Qwen35OwnedDeltaNetLayer {
    pub(crate) attn_norm: Vec<f32>,
    pub(crate) attn_qkv: OwnedQuantizedTensor,
    pub(crate) attn_gate: OwnedQuantizedTensor,
    pub(crate) ssm_alpha: OwnedQuantizedTensor,
    pub(crate) ssm_beta: OwnedQuantizedTensor,
    pub(crate) ssm_a: Vec<f32>,
    pub(crate) ssm_dt_bias: Vec<f32>,
    pub ssm_conv1d_weight: Vec<f32>,
    pub(crate) ssm_norm_weight: Vec<f32>,
    pub(crate) ssm_out: OwnedQuantizedTensor,
    pub(crate) post_attention_norm: Vec<f32>,
    pub(crate) ffn_gate: OwnedQuantizedTensor,
    pub(crate) ffn_up: OwnedQuantizedTensor,
    pub(crate) ffn_down: OwnedQuantizedTensor,
}

pub(crate) struct Qwen35OwnedAttentionLayer {
    pub(crate) attn_norm: Vec<f32>,
    pub(crate) attn_q: OwnedQuantizedTensor,
    pub(crate) attn_k: OwnedQuantizedTensor,
    pub(crate) attn_v: OwnedQuantizedTensor,
    pub(crate) attn_q_norm: Vec<f32>,
    pub(crate) attn_k_norm: Vec<f32>,
    pub(crate) attn_output: OwnedQuantizedTensor,
    pub(crate) post_attention_norm: Vec<f32>,
    pub(crate) ffn_gate: OwnedQuantizedTensor,
    pub(crate) ffn_up: OwnedQuantizedTensor,
    pub(crate) ffn_down: OwnedQuantizedTensor,
}

pub(crate) enum Qwen35OwnedLayer {
    DeltaNet(Qwen35OwnedDeltaNetLayer),
    Attention(Qwen35OwnedAttentionLayer),
}

/// Qwen3.5 / Qwen3.8 hybrid decoder on the CPU (#3091): Gated `DeltaNet` layers (short causal
/// conv, per-head recurrent state, gated delta rule) interleaved with gated full-attention
/// layers, on top of the dense model's embeddings, final norm and `lm_head`.
pub struct Qwen35Model<'a> {
    pub(crate) base: &'a OwnedQuantizedModel,
    pub(crate) layers: Vec<Qwen35OwnedLayer>,
    pub(crate) head_dim: usize,
    pub(crate) num_kv_heads: usize,
    pub(crate) num_v_heads: usize,
    pub(crate) head_v_dim: usize,
    pub(crate) num_k_heads: usize,
    pub(crate) head_k_dim: usize,
    pub(crate) conv_kernel: usize,
    pub(crate) rope_sections: [usize; 4],
}

fn load_f32_vec(tensor_ref: &QuantizedTensorRef, data: &[u8]) -> Result<Vec<f32>> {
    if tensor_ref.qtype != crate::gguf::types::GGUF_TYPE_F32 {
        return Err(crate::error::RealizarError::FormatError {
            reason: format!(
                "qwen35: expected an F32 tensor, found GGUF type {}",
                tensor_ref.qtype
            ),
        });
    }
    let bytes = data
        .get(tensor_ref.offset..tensor_ref.offset + tensor_ref.byte_size)
        .ok_or_else(|| crate::error::RealizarError::FormatError {
            reason: format!(
                "qwen35: F32 tensor at byte {} (+{}) lies outside the file",
                tensor_ref.offset, tensor_ref.byte_size
            ),
        })?;
    // Decoded, not reinterpreted: an mmap offset carries no f32 alignment guarantee.
    Ok(bytes
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| f32::from_le_bytes(*c))
        .collect())
}

/// Metadata map of a GGUF file, as `GGUFModel::metadata` carries it.
type GGUFMetadata = std::collections::HashMap<String, crate::gguf::types::GGUFValue>;

/// The Gated `DeltaNet` shape Qwen3.5 records in GGUF metadata, under either the
/// `qwen2.*` or the `qwen35.*` key prefix.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Qwen35SsmMeta {
    pub(crate) head_k_dim: usize,
    pub(crate) num_k_heads: usize,
    pub(crate) num_v_heads: usize,
    pub(crate) conv_kernel: usize,
    pub(crate) rope_sections: [usize; 4],
}

/// One unsigned metadata value, accepting either GGUF integer width.
fn metadata_usize(metadata: &GGUFMetadata, key: &str) -> Option<usize> {
    match metadata.get(key) {
        Some(crate::gguf::types::GGUFValue::UInt32(v)) => Some(*v as usize),
        Some(crate::gguf::types::GGUFValue::Int32(v)) => Some(*v as usize),
        _ => None,
    }
}

/// `key`, else `fallback`, else `default` — the `qwen2.*`/`qwen35.*` prefix pair.
fn metadata_usize_or(metadata: &GGUFMetadata, key: &str, fallback: &str, default: usize) -> usize {
    metadata_usize(metadata, key)
        .or_else(|| metadata_usize(metadata, fallback))
        .unwrap_or(default)
}

/// The four M-RoPE section widths. Defaulted per section: a section the array does not
/// reach, or carries with a non-integer type, keeps its default.
fn rope_sections_from_metadata(metadata: &GGUFMetadata) -> [usize; 4] {
    let mut rope_sections = [16, 24, 24, 0];
    if let Some(crate::gguf::types::GGUFValue::Array(arr)) = metadata
        .get("qwen2.rope.dimension_sections")
        .or_else(|| metadata.get("qwen35.rope.dimension_sections"))
    {
        for (i, val) in arr.iter().take(4).enumerate() {
            if let crate::gguf::types::GGUFValue::UInt32(v) = val {
                rope_sections[i] = *v as usize;
            } else if let crate::gguf::types::GGUFValue::Int32(v) = val {
                rope_sections[i] = *v as usize;
            }
        }
    }
    rope_sections
}

/// Read the Gated `DeltaNet` shape, falling back to the Qwen3.5-0.8B defaults key by key.
pub(crate) fn qwen35_ssm_meta(metadata: &GGUFMetadata) -> Qwen35SsmMeta {
    Qwen35SsmMeta {
        head_k_dim: metadata_usize_or(
            metadata,
            "qwen2.ssm.state_size",
            "qwen35.ssm.state_size",
            16,
        ),
        num_k_heads: metadata_usize_or(
            metadata,
            "qwen2.ssm.group_count",
            "qwen35.ssm.group_count",
            1,
        ),
        num_v_heads: metadata_usize_or(
            metadata,
            "qwen2.ssm.time_step_rank",
            "qwen35.ssm.time_step_rank",
            16,
        ),
        conv_kernel: metadata_usize_or(
            metadata,
            "qwen2.ssm.conv_kernel",
            "qwen35.ssm.conv_kernel",
            4,
        ),
        rope_sections: rope_sections_from_metadata(metadata),
    }
}

/// The widths every hybrid layer's tensors are dequantised with.
struct Qwen35LayerDims {
    hidden_dim: usize,
    intermediate_dim: usize,
    num_heads: usize,
    num_kv_heads: usize,
    num_v_heads: usize,
    conv_dim: usize,
    value_dim: usize,
}

/// Own one Gated `DeltaNet` layer's tensors.
fn own_deltanet_layer(
    d: &crate::gguf::qwen35_load::Qwen35DeltaNetLayer,
    data: &[u8],
    dims: &Qwen35LayerDims,
) -> Result<Qwen35OwnedDeltaNetLayer> {
    Ok(Qwen35OwnedDeltaNetLayer {
        attn_norm: load_f32_vec(&d.attn_norm, data)?,
        attn_qkv: OwnedQuantizedTensor::from_ref_with_dims(
            &d.attn_qkv,
            data,
            dims.hidden_dim,
            dims.conv_dim,
        ),
        attn_gate: OwnedQuantizedTensor::from_ref_with_dims(
            &d.attn_gate,
            data,
            dims.hidden_dim,
            dims.value_dim,
        ),
        ssm_alpha: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ssm_alpha,
            data,
            dims.hidden_dim,
            dims.num_v_heads,
        ),
        ssm_beta: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ssm_beta,
            data,
            dims.hidden_dim,
            dims.num_v_heads,
        ),
        ssm_a: load_f32_vec(&d.ssm_a, data)?,
        ssm_dt_bias: load_f32_vec(&d.ssm_dt_bias, data)?,
        ssm_conv1d_weight: load_f32_vec(&d.ssm_conv1d_weight, data)?,
        ssm_norm_weight: load_f32_vec(&d.ssm_norm_weight, data)?,
        ssm_out: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ssm_out,
            data,
            dims.value_dim,
            dims.hidden_dim,
        ),
        post_attention_norm: load_f32_vec(&d.post_attention_norm, data)?,
        ffn_gate: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ffn_gate,
            data,
            dims.hidden_dim,
            dims.intermediate_dim,
        ),
        ffn_up: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ffn_up,
            data,
            dims.hidden_dim,
            dims.intermediate_dim,
        ),
        ffn_down: OwnedQuantizedTensor::from_ref_with_dims(
            &d.ffn_down,
            data,
            dims.intermediate_dim,
            dims.hidden_dim,
        ),
    })
}

/// Own one full-attention layer's tensors. The head width is `attn_q_norm`'s length
/// (256 for Qwen3.5 standard attention), not `hidden_dim / num_heads`, and `attn_q` is
/// gated, so it is twice as wide as the head fan-out.
fn own_attention_layer(
    a: &crate::gguf::qwen35_load::Qwen35AttentionLayer,
    data: &[u8],
    dims: &Qwen35LayerDims,
) -> Result<Qwen35OwnedAttentionLayer> {
    let attn_q_norm = load_f32_vec(&a.attn_q_norm, data)?;
    let true_head_dim = attn_q_norm.len();

    Ok(Qwen35OwnedAttentionLayer {
        attn_norm: load_f32_vec(&a.attn_norm, data)?,
        attn_q: OwnedQuantizedTensor::from_ref_with_dims(
            &a.attn_q,
            data,
            dims.hidden_dim,
            dims.num_heads * true_head_dim * 2,
        ),
        attn_k: OwnedQuantizedTensor::from_ref_with_dims(
            &a.attn_k,
            data,
            dims.hidden_dim,
            dims.num_kv_heads * true_head_dim,
        ),
        attn_v: OwnedQuantizedTensor::from_ref_with_dims(
            &a.attn_v,
            data,
            dims.hidden_dim,
            dims.num_kv_heads * true_head_dim,
        ),
        attn_q_norm,
        attn_k_norm: load_f32_vec(&a.attn_k_norm, data)?,
        attn_output: OwnedQuantizedTensor::from_ref_with_dims(
            &a.attn_output,
            data,
            dims.num_heads * true_head_dim,
            dims.hidden_dim,
        ),
        post_attention_norm: load_f32_vec(&a.post_attention_norm, data)?,
        ffn_gate: OwnedQuantizedTensor::from_ref_with_dims(
            &a.ffn_gate,
            data,
            dims.hidden_dim,
            dims.intermediate_dim,
        ),
        ffn_up: OwnedQuantizedTensor::from_ref_with_dims(
            &a.ffn_up,
            data,
            dims.hidden_dim,
            dims.intermediate_dim,
        ),
        ffn_down: OwnedQuantizedTensor::from_ref_with_dims(
            &a.ffn_down,
            data,
            dims.intermediate_dim,
            dims.hidden_dim,
        ),
    })
}

/// Attention head width is `key_length` (= `attn_q_norm`'s width, 256), NOT
/// `hidden / heads` (128): the KV cache and the attention loop used 128-wide heads over
/// 256-wide q/k/v. A file with no full-attention layer keeps `fallback`.
fn attention_head_dim(layers: &[Qwen35OwnedLayer], fallback: usize) -> usize {
    layers
        .iter()
        .find_map(|l| match l {
            Qwen35OwnedLayer::Attention(a) => Some(a.attn_q_norm.len()),
            Qwen35OwnedLayer::DeltaNet(_) => None,
        })
        .unwrap_or(fallback)
}

impl<'a> Qwen35Model<'a> {
    /// Build the shared base (config, token embeddings, final norm, `lm_head`) without the
    /// dense layer loader, which refuses hybrid files.
    ///
    /// # Errors
    /// Missing or malformed base tensors or metadata.
    pub fn create_base_model(
        model: &crate::gguf::GGUFModel,
        data: &[u8],
    ) -> crate::error::Result<crate::gguf::OwnedQuantizedModel> {
        // #4032: Qwen3.5-0.8B is a known RED in 0.69.1 (#4030): warn, never refuse.
        qwen35_known_issue::warn_if_known_issue(model);
        let config = crate::gguf::config::ValidatedModelConfig::from_gguf(model)?.into_inner();

        let token_embedding = model.get_tensor_f32("token_embd.weight", data)?;
        let output_norm_weight = model.get_tensor_f32("output_norm.weight", data)?;

        let lm_head_ref =
            crate::gguf::QuantizedGGUFTransformer::get_tensor_ref(model, data, "output.weight")
                .or_else(|_| {
                    crate::gguf::QuantizedGGUFTransformer::get_tensor_ref(
                        model,
                        data,
                        "token_embd.weight",
                    )
                })?;
        let lm_head_weight = crate::gguf::OwnedQuantizedTensor::from_ref_with_dims(
            &lm_head_ref,
            data,
            config.hidden_dim,
            config.vocab_size,
        );

        Ok(crate::gguf::OwnedQuantizedModel {
            config,
            token_embedding,
            position_embedding: None,
            layers: vec![],
            encoder_layers: vec![],
            encoder_output_norm_weight: None,
            encoder_output_norm_bias: None,
            output_norm_weight,
            output_norm_bias: None,
            lm_head_weight,
            lm_head_bias: None,
            // A CUDA build carries these on every OwnedQuantizedModel (as in loading.rs); the
            // Qwen3.5 base never uses them (the GPU path has no Gated DeltaNet yet, #3090).
            #[cfg(feature = "cuda")]
            cuda_executor: None,
            #[cfg(feature = "cuda")]
            cuda_kernel_count: std::sync::atomic::AtomicU64::new(0),
            #[cfg(feature = "cuda")]
            cached_weight_names: std::sync::Mutex::new(std::collections::HashSet::new()),
        })
    }

    /// Load every hybrid layer of `model` on top of `base`.
    ///
    /// # Errors
    /// A layer tensor that is missing, has an unexpected type, or lies outside the file.
    pub fn from_model_and_layers(
        base: &'a OwnedQuantizedModel,
        model: &GGUFModel,
        data: &[u8],
    ) -> Result<Self> {
        let refs = crate::gguf::qwen35_load::load_qwen35_layers(model, data)?;
        let meta = qwen35_ssm_meta(&model.metadata);

        let head_v_dim = meta.head_k_dim;
        let num_heads = base.config.num_heads;
        let num_kv_heads = base.config.num_kv_heads;
        let head_dim = base.config.hidden_dim / num_heads;

        let key_dim = meta.head_k_dim * meta.num_k_heads;
        let value_dim = head_v_dim * meta.num_v_heads;
        let dims = Qwen35LayerDims {
            hidden_dim: base.config.hidden_dim,
            intermediate_dim: base.config.intermediate_dim,
            num_heads,
            num_kv_heads,
            num_v_heads: meta.num_v_heads,
            conv_dim: key_dim * 2 + value_dim,
            value_dim,
        };

        let mut owned = Vec::with_capacity(refs.len());
        for layer_ref in refs {
            owned.push(match layer_ref {
                crate::gguf::qwen35_load::Qwen35Layer::DeltaNet(d) => {
                    Qwen35OwnedLayer::DeltaNet(own_deltanet_layer(&d, data, &dims)?)
                },
                crate::gguf::qwen35_load::Qwen35Layer::Attention(a) => {
                    Qwen35OwnedLayer::Attention(own_attention_layer(&a, data, &dims)?)
                },
            });
        }

        let attn_head_dim = attention_head_dim(&owned, head_dim);
        Ok(Self {
            base,
            layers: owned,
            head_dim: attn_head_dim,
            num_kv_heads,
            num_v_heads: meta.num_v_heads,
            head_v_dim,
            num_k_heads: meta.num_k_heads,
            head_k_dim: meta.head_k_dim,
            conv_kernel: meta.conv_kernel,
            rope_sections: meta.rope_sections,
        })
    }

    /// A fresh decode state with room for `max_seq_len` positions.
    #[must_use]
    pub fn new_state(&self, max_seq_len: usize) -> Qwen35State {
        Qwen35State::new(
            self.layers.len(),
            max_seq_len,
            self.head_dim,
            self.num_kv_heads,
            self.num_k_heads,
            self.head_k_dim,
            self.num_v_heads,
            self.head_v_dim,
        )
    }

    /// Run one token at `position` through every layer, updating `cache`, and return the
    /// logits.
    ///
    /// # Errors
    /// A matmul or shape failure in any layer.
    pub fn forward_single_qwen35(
        &self,
        token_id: u32,
        cache: &mut Qwen35State,
        position: usize,
    ) -> Result<Vec<f32>> {
        let mut hidden = self.base.token_embedding()[(token_id as usize)
            * self.base.config.hidden_dim
            ..(token_id as usize + 1) * self.base.config.hidden_dim]
            .to_vec();

        let mut normed = vec![0.0; self.base.config.hidden_dim];
        let mut post_attn_normed = vec![0.0; self.base.config.hidden_dim];

        let fn_softplus = |x: f32| -> f32 {
            if x > 20.0 {
                x
            } else {
                (1.0 + x.exp()).ln()
            }
        };

        for (il, layer) in self.layers.iter().enumerate() {
            match layer {
                Qwen35OwnedLayer::DeltaNet(d) => {
                    self.forward_deltanet(
                        d,
                        &mut hidden,
                        cache,
                        il,
                        position,
                        &mut normed,
                        &mut post_attn_normed,
                    )?;
                },
                Qwen35OwnedLayer::Attention(a) => {
                    self.forward_attention(
                        a,
                        &mut hidden,
                        cache,
                        il,
                        position,
                        &mut normed,
                        &mut post_attn_normed,
                    )?;
                },
            }
        }
        let mut out_normed = vec![0.0; self.base.config.hidden_dim];
        crate::gguf::ops::rms_norm_into(
            &hidden,
            self.base.output_norm_weight(),
            self.base.config.eps,
            &mut out_normed,
        );

        let mut logits = vec![0.0; self.base.config.vocab_size];
        self.base
            .fused_matmul_into(&out_normed, self.base.lm_head_weight(), &mut logits)?;

        cache.kv_cache.advance();

        Ok(logits)
    }

    /// One Gated `DeltaNet` layer, in place on `hidden`.
    ///
    /// `pub(crate)` (PMAT-3477, #3090) so the GPU model's per-layer parity test can
    /// drive the CPU reference layer by layer with teacher forcing. No logic change.
    pub(crate) fn forward_deltanet(
        &self,
        d: &Qwen35OwnedDeltaNetLayer,
        hidden: &mut [f32],
        cache: &mut Qwen35State,
        il: usize,
        position: usize,
        normed: &mut [f32],
        post_attn_normed: &mut [f32],
    ) -> Result<()> {
        crate::gguf::ops::rms_norm_into(hidden, &d.attn_norm, self.base.config.eps, normed);
        let conv_dim = self.head_k_dim * self.num_k_heads * 2 + self.head_v_dim * self.num_v_heads;
        let v_dim = self.head_v_dim * self.num_v_heads;
        let mut conv_in = vec![0.0; conv_dim];
        self.base
            .fused_matmul_into(normed, &d.attn_qkv, &mut conv_in)?;
        let mut dt_raw = vec![0.0; self.num_v_heads];
        self.base
            .fused_matmul_into(normed, &d.ssm_alpha, &mut dt_raw)?;
        let mut beta = vec![0.0; self.num_v_heads];
        self.base
            .fused_matmul_into(normed, &d.ssm_beta, &mut beta)?;
        let mut gate = vec![0.0; v_dim];
        self.base
            .fused_matmul_into(normed, &d.attn_gate, &mut gate)?;

        let ssm_out_in = self.deltanet_mix_token(d, cache, il, &conv_in, &dt_raw, &mut beta, &gate);

        let mut ssm_out = vec![0.0; self.base.config.hidden_dim];
        self.base
            .fused_matmul_into(&ssm_out_in, &d.ssm_out, &mut ssm_out)?;

        for i in 0..self.base.config.hidden_dim {
            hidden[i] += ssm_out[i];
        }

        crate::gguf::ops::rms_norm_into(
            hidden,
            &d.post_attention_norm,
            self.base.config.eps,
            post_attn_normed,
        );
        let mut ffn_gate = vec![0.0; d.ffn_gate.out_dim];
        self.base
            .fused_matmul_into(post_attn_normed, &d.ffn_gate, &mut ffn_gate)?;
        let mut ffn_up = vec![0.0; d.ffn_up.out_dim];
        self.base
            .fused_matmul_into(post_attn_normed, &d.ffn_up, &mut ffn_up)?;

        for i in 0..ffn_gate.len() {
            let x = ffn_gate[i];
            let silu = x / (1.0 + (-x as f32).exp());
            ffn_up[i] *= silu;
        }
        let mut ffn_down = vec![0.0; d.ffn_down.out_dim];
        self.base
            .fused_matmul_into(&ffn_up, &d.ffn_down, &mut ffn_down)?;
        for i in 0..self.base.config.hidden_dim {
            hidden[i] += ffn_down[i];
        }
        Ok(())
    }

    /// One full-attention layer, in place on `hidden`.
    ///
    /// `pub(crate)` (PMAT-3477, #3090) so the GPU parity test can advance the CPU
    /// reference across the interleaved attention layers while it compares the
    /// `DeltaNet` ones. No logic change.
    pub(crate) fn forward_attention(
        &self,
        a: &Qwen35OwnedAttentionLayer,
        hidden: &mut [f32],
        cache: &mut Qwen35State,
        il: usize,
        position: usize,
        normed: &mut [f32],
        post_attn_normed: &mut [f32],
    ) -> Result<()> {
        crate::gguf::ops::rms_norm_into(hidden, &a.attn_norm, self.base.config.eps, normed);

        let mut q_full = vec![0.0; a.attn_q.out_dim];
        let mut k = vec![0.0; a.attn_k.out_dim];
        let mut v = vec![0.0; a.attn_v.out_dim];
        self.base
            .fused_matmul_into(normed, &a.attn_q, &mut q_full)?;
        self.base.fused_matmul_into(normed, &a.attn_k, &mut k)?;
        self.base.fused_matmul_into(normed, &a.attn_v, &mut v)?;

        let attn_out_in = self.attention_mix_token(a, cache, il, position, &q_full, &mut k, &v);
        let mut attn_out = vec![0.0; self.base.config.hidden_dim];
        self.base
            .fused_matmul_into(&attn_out_in, &a.attn_output, &mut attn_out)?;
        for i in 0..self.base.config.hidden_dim {
            hidden[i] += attn_out[i];
        }

        crate::gguf::ops::rms_norm_into(
            hidden,
            &a.post_attention_norm,
            self.base.config.eps,
            post_attn_normed,
        );
        let mut ffn_gate = vec![0.0; a.ffn_gate.out_dim];
        self.base
            .fused_matmul_into(post_attn_normed, &a.ffn_gate, &mut ffn_gate)?;
        let mut ffn_up = vec![0.0; a.ffn_up.out_dim];
        self.base
            .fused_matmul_into(post_attn_normed, &a.ffn_up, &mut ffn_up)?;
        for i in 0..ffn_gate.len() {
            let x = ffn_gate[i];
            let silu = x / (1.0 + (-x as f32).exp());
            ffn_up[i] *= silu;
        }
        let mut ffn_down = vec![0.0; a.ffn_down.out_dim];
        self.base
            .fused_matmul_into(&ffn_up, &a.ffn_down, &mut ffn_down)?;
        for i in 0..self.base.config.hidden_dim {
            hidden[i] += ffn_down[i];
        }
        Ok(())
    }
    /// Prefill `tokens` at positions `start_pos..`, updating `cache`, and return the logits of
    /// the LAST token (#4228).
    ///
    /// Layer-major over chunks of [`QWEN35_PREFILL_CHUNK`] tokens: every projection and FFN
    /// matmul runs once per chunk over all its rows (Q4_K weights stream from DRAM once per
    /// tile, not once per token), while the `DeltaNet` recurrence and causal attention still
    /// run token by token in order. Bit-identical to calling [`Self::forward_single_qwen35`]
    /// once per token, which FALSIFY-4228-004 checks.
    ///
    /// The whole prompt must fit the KV cache: `cache` length plus `tokens.len()` may not
    /// exceed its `max_len`. Past that bound the per-token path stops appending and then
    /// indexes past the cache, so a chunk that straddles it has no per-token behaviour to
    /// match; the overflow is refused before any state is touched (FALSIFY-4228-007).
    ///
    /// # Errors
    /// An empty `tokens`, a prompt that overflows the KV cache, or a matmul or shape
    /// failure in any layer.
    pub fn forward_prefill_qwen35(
        &self,
        tokens: &[u32],
        cache: &mut Qwen35State,
        start_pos: usize,
    ) -> Result<Vec<f32>> {
        if tokens.is_empty() {
            return Err(crate::error::RealizarError::InvalidShape {
                reason: "forward_prefill_qwen35: no tokens".to_string(),
            });
        }
        let (used, cap) = (cache.kv_cache.len(), cache.kv_cache.max_len());
        if used + tokens.len() > cap {
            return Err(crate::error::RealizarError::InvalidShape {
                reason: format!(
                    "forward_prefill_qwen35: {} tokens at cache length {used} overflow max_seq_len {cap}",
                    tokens.len()
                ),
            });
        }
        let mut logits = Vec::new();
        for (ci, chunk) in tokens.chunks(QWEN35_PREFILL_CHUNK).enumerate() {
            logits =
                self.prefill_chunk_qwen35(chunk, cache, start_pos + ci * QWEN35_PREFILL_CHUNK)?;
        }
        Ok(logits)
    }

    fn prefill_chunk_qwen35(
        &self,
        chunk: &[u32],
        cache: &mut Qwen35State,
        pos0: usize,
    ) -> Result<Vec<f32>> {
        let hd = self.base.config.hidden_dim;
        let eps = self.base.config.eps;
        let b = chunk.len();
        let emb = self.base.token_embedding();
        let mut hidden = vec![0.0; b * hd];
        for (t, &tok) in chunk.iter().enumerate() {
            let at = tok as usize * hd;
            hidden[t * hd..(t + 1) * hd].copy_from_slice(&emb[at..at + hd]);
        }
        let mut normed = vec![0.0; b * hd];

        for (il, layer) in self.layers.iter().enumerate() {
            match layer {
                Qwen35OwnedLayer::DeltaNet(d) => {
                    rms_norm_rows(&hidden, &d.attn_norm, eps, hd, &mut normed);
                    let conv_in = self.matmul_rows(&normed, b, &d.attn_qkv)?;
                    let dt_raw = self.matmul_rows(&normed, b, &d.ssm_alpha)?;
                    let mut beta = self.matmul_rows(&normed, b, &d.ssm_beta)?;
                    let gate = self.matmul_rows(&normed, b, &d.attn_gate)?;
                    let mix = self
                        .deltanet_mix_rows(d, cache, il, b, &conv_in, &dt_raw, &mut beta, &gate);
                    let out = self.matmul_rows(&mix, b, &d.ssm_out)?;
                    add_into(&mut hidden, &out);
                    self.ffn_rows(
                        &mut hidden,
                        b,
                        &d.post_attention_norm,
                        [&d.ffn_gate, &d.ffn_up, &d.ffn_down],
                    )?;
                },
                Qwen35OwnedLayer::Attention(a) => {
                    rms_norm_rows(&hidden, &a.attn_norm, eps, hd, &mut normed);
                    let q_full = self.matmul_rows(&normed, b, &a.attn_q)?;
                    let mut k = self.matmul_rows(&normed, b, &a.attn_k)?;
                    let v = self.matmul_rows(&normed, b, &a.attn_v)?;
                    let mix = self.attention_mix_rows(a, cache, il, pos0, b, &q_full, &mut k, &v);
                    let out = self.matmul_rows(&mix, b, &a.attn_output)?;
                    add_into(&mut hidden, &out);
                    self.ffn_rows(
                        &mut hidden,
                        b,
                        &a.post_attention_norm,
                        [&a.ffn_gate, &a.ffn_up, &a.ffn_down],
                    )?;
                },
            }
        }
        cache.kv_cache.advance_by(b);

        let mut out_normed = vec![0.0; hd];
        crate::gguf::ops::rms_norm_into(
            &hidden[(b - 1) * hd..],
            self.base.output_norm_weight(),
            eps,
            &mut out_normed,
        );
        let mut logits = vec![0.0; self.base.config.vocab_size];
        self.base
            .fused_matmul_into(&out_normed, self.base.lm_head_weight(), &mut logits)?;
        Ok(logits)
    }

    /// `b` rows of `weight.in_dim` through `weight`, token-major out. Q4_K/Q5_K/Q6_K with
    /// `b > 1` take the multi-row GEMMs (each bit-identical to its per-row matvec); every other
    /// type loops the same per-row `fused_matmul_into` the per-token forward calls.
    /// A `DeltaNet` layer's mix for `b` token rows, bitwise [`Self::deltanet_mix_token`] per row.
    /// The conv window runs in token order; then each value head walks every token on its own
    /// thread, because a head's recurrence state is touched by that head alone.
    #[allow(clippy::too_many_arguments)]
    fn deltanet_mix_rows(
        &self,
        d: &Qwen35OwnedDeltaNetLayer,
        cache: &mut Qwen35State,
        il: usize,
        b: usize,
        conv_in: &[f32],
        dt_raw: &[f32],
        beta: &mut [f32],
        gate: &[f32],
    ) -> Vec<f32> {
        use rayon::prelude::*;
        let (cd, nv) = (d.attn_qkv.out_dim, d.ssm_alpha.out_dim);
        let (hk, hv, nk) = (self.head_k_dim, self.head_v_dim, self.num_k_heads);
        let vd = hv * nv;
        let preps: Vec<[Vec<f32>; 4]> = (0..b)
            .map(|t| {
                self.deltanet_prep_token(
                    d,
                    cache,
                    il,
                    &conv_in[t * cd..(t + 1) * cd],
                    &dt_raw[t * nv..(t + 1) * nv],
                    &mut beta[t * nv..(t + 1) * nv],
                )
            })
            .collect();

        let scale = 1.0 / (hk as f32).sqrt();
        let beta: &[f32] = beta;
        let mut by_head = vec![0.0f32; nv * b * hv];
        cache.ssm_states[il][..]
            .par_chunks_mut(hv * hk)
            .zip(by_head.par_chunks_mut(b * hv))
            .enumerate()
            .for_each(|(h, (s_h, out))| {
                let kh = h % nk;
                for (t, [q, k, v, dt]) in preps.iter().enumerate() {
                    delta_rule_head(
                        &q[kh * hk..(kh + 1) * hk],
                        &k[kh * hk..(kh + 1) * hk],
                        &v[h * hv..(h + 1) * hv],
                        beta[t * nv + h],
                        dt[h],
                        s_h,
                        &mut out[t * hv..(t + 1) * hv],
                        hk,
                        scale,
                    );
                }
            });

        let mut mix = vec![0.0; b * vd];
        let mut out_h = vec![0.0; vd];
        for t in 0..b {
            for h in 0..nv {
                let at = h * b * hv + t * hv;
                out_h[h * hv..(h + 1) * hv].copy_from_slice(&by_head[at..at + hv]);
            }
            gated_rmsnorm(
                &out_h,
                &gate[t * vd..(t + 1) * vd],
                &d.ssm_norm_weight,
                self.base.config.eps,
                hv,
                &mut mix[t * vd..(t + 1) * vd],
            );
        }
        mix
    }

    /// A full-attention layer's mix for `b` token rows at positions `pos0..pos0 + b`, bitwise
    /// [`Self::attention_mix_token`] per row. Every row's K/V is appended first, in order; row
    /// `t` then reads positions `0..=pos0 + t` only, so the rows attend in parallel.
    #[allow(clippy::too_many_arguments)]
    fn attention_mix_rows(
        &self,
        a: &Qwen35OwnedAttentionLayer,
        cache: &mut Qwen35State,
        il: usize,
        pos0: usize,
        b: usize,
        q_full: &[f32],
        k: &mut [f32],
        v: &[f32],
    ) -> Vec<f32> {
        use rayon::prelude::*;
        let (qd, kd) = (a.attn_q.out_dim, a.attn_k.out_dim);
        let od = a.attn_output.in_dim;
        let qg: Vec<(Vec<f32>, Vec<f32>)> = (0..b)
            .map(|t| {
                self.attention_append_token(
                    a,
                    cache,
                    il,
                    pos0 + t,
                    &q_full[t * qd..(t + 1) * qd],
                    &mut k[t * kd..(t + 1) * kd],
                    &v[t * kd..(t + 1) * kd],
                )
            })
            .collect();
        let cache: &Qwen35State = cache;
        let mut mix = vec![0.0; b * od];
        mix.par_chunks_mut(od)
            .zip(qg.par_iter())
            .enumerate()
            .for_each(|(t, (out, (q, gate)))| {
                out.copy_from_slice(&self.attention_attend_token(cache, il, pos0 + t, q, gate));
            });
        mix
    }

    fn matmul_rows(
        &self,
        input: &[f32],
        b: usize,
        weight: &OwnedQuantizedTensor,
    ) -> Result<Vec<f32>> {
        let (in_dim, out_dim) = (weight.in_dim, weight.out_dim);
        let mut out = vec![0.0; b * out_dim];
        if weight.qtype == crate::gguf::GGUF_TYPE_Q4_K && b > 1 {
            crate::quantize::fused_q4k_multirow_matmul_f32_into(
                &weight.data,
                &input[..b * in_dim],
                b,
                in_dim,
                out_dim,
                &mut out,
            )?;
        } else if weight.qtype == crate::gguf::GGUF_TYPE_Q5_K && b > 1 {
            crate::quantize::fused_q5k_multirow_matmul_into(
                &weight.data,
                &input[..b * in_dim],
                b,
                in_dim,
                out_dim,
                &mut out,
            )?;
        } else if weight.qtype == crate::gguf::GGUF_TYPE_Q6_K && b > 1 {
            crate::quantize::fused_q6k_multirow_matmul_into(
                &weight.data,
                &input[..b * in_dim],
                b,
                in_dim,
                out_dim,
                &mut out,
            )?;
        } else {
            for t in 0..b {
                self.base.fused_matmul_into(
                    &input[t * in_dim..(t + 1) * in_dim],
                    weight,
                    &mut out[t * out_dim..(t + 1) * out_dim],
                )?;
            }
        }
        Ok(out)
    }

    /// The post-attention SwiGLU FFN over `b` rows of `hidden`, residual added in place; the
    /// same arithmetic as the tail of `forward_deltanet`/`forward_attention`.
    fn ffn_rows(
        &self,
        hidden: &mut [f32],
        b: usize,
        post_norm: &[f32],
        [ffn_gate, ffn_up, ffn_down]: [&OwnedQuantizedTensor; 3],
    ) -> Result<()> {
        let hd = self.base.config.hidden_dim;
        let mut normed = vec![0.0; b * hd];
        rms_norm_rows(hidden, post_norm, self.base.config.eps, hd, &mut normed);
        let gate = self.matmul_rows(&normed, b, ffn_gate)?;
        let mut up = self.matmul_rows(&normed, b, ffn_up)?;
        for (u, &x) in up.iter_mut().zip(&gate) {
            let silu = x / (1.0 + (-x as f32).exp());
            *u *= silu;
        }
        let down = self.matmul_rows(&up, b, ffn_down)?;
        add_into(hidden, &down);
        Ok(())
    }

    /// The conv window, SiLU, the per-head q/k L2 norms, `dt` and the in-place beta sigmoid of
    /// [`Self::deltanet_mix_token`]: everything before the recurrence. Returns `[q, k, v, dt]`.
    #[allow(clippy::too_many_arguments)]
    fn deltanet_prep_token(
        &self,
        d: &Qwen35OwnedDeltaNetLayer,
        cache: &mut Qwen35State,
        il: usize,
        conv_in: &[f32],
        dt_raw: &[f32],
        beta: &mut [f32],
    ) -> [Vec<f32>; 4] {
        let fn_softplus = |x: f32| -> f32 {
            if x > 20.0 {
                x
            } else {
                (1.0 + x.exp()).ln()
            }
        };

        let conv_dim = self.head_k_dim * self.num_k_heads * 2 + self.head_v_dim * self.num_v_heads;
        let mut conv_out = vec![0.0; conv_dim];
        causal_conv1d(
            conv_in,
            &mut cache.conv_states[il][..],
            &d.ssm_conv1d_weight,
            self.conv_kernel,
            conv_dim,
            &mut conv_out,
        );
        for x in conv_out.iter_mut() {
            *x = *x / (1.0 + (-*x).exp()); // silu/sigmoid
        }
        // SiLU ONCE, as llama.cpp (ggml_silu on conv_output_raw). A second identical loop applied silu(silu(x)).

        let k_dim = self.head_k_dim * self.num_k_heads;
        let v_dim = self.head_v_dim * self.num_v_heads;
        let mut q = conv_out[0..k_dim].to_vec();
        let mut k = conv_out[k_dim..k_dim * 2].to_vec();
        let v = conv_out[k_dim * 2..conv_dim].to_vec();

        // Per-head L2 normalisation of q and k, as llama.cpp build_gdn_l2_norm:
        // rms_norm(x, eps/n) * 1/sqrt(n) over ne[0] = head_k_dim, i.e. x_h / sqrt(sum(x_h^2) + eps)
        // for EACH head. The earlier global norm over all heads (applied twice) scaled every head wrong.
        l2_norm_per_head(&mut q, self.head_k_dim, self.base.config.eps);
        l2_norm_per_head(&mut k, self.head_k_dim, self.base.config.eps);

        // Gated DeltaNet applies NO RoPE: llama.cpp build_layer_attn_linear has none (position comes from the causal conv).

        let mut dt = vec![0.0; self.num_v_heads];
        for (i, val) in dt_raw.iter().enumerate() {
            dt[i] = fn_softplus(val + d.ssm_dt_bias[i]) * d.ssm_a[i];
        }

        // sigmoid ONCE, as llama.cpp: beta = sigmoid(ssm_beta . x). (It was applied twice.)
        for x in beta.iter_mut() {
            *x = 1.0 / (1.0 + (-*x).exp());
        }

        // q and k are num_k_heads wide, v is num_v_heads wide: on 4B/9B (32 value heads to 16
        // key heads) and 27B (48 to 16) the recurrence shares each key/query head across
        // num_v_heads / num_k_heads value heads (PMAT-3477, #3346/#3510). On 0.8B and 2B the
        // ratio is 1 and this is the pre-GQA arithmetic, unchanged.
        [q, k, v, dt]
    }

    /// The per-token middle of a `DeltaNet` layer: causal conv, q/k L2 norms, the delta-rule
    /// recurrence and the gated RMS norm. Takes the four input projections of one token and
    /// returns the `ssm_out` input. Shared by the per-token forward and the batched prefill
    /// (#4228), so both run the exact same recurrence arithmetic.
    #[allow(clippy::too_many_arguments)]
    fn deltanet_mix_token(
        &self,
        d: &Qwen35OwnedDeltaNetLayer,
        cache: &mut Qwen35State,
        il: usize,
        conv_in: &[f32],
        dt_raw: &[f32],
        beta: &mut [f32],
        gate: &[f32],
    ) -> Vec<f32> {
        let [q, k, v, dt] = self.deltanet_prep_token(d, cache, il, conv_in, dt_raw, beta);
        let v_dim = self.head_v_dim * self.num_v_heads;
        let mut out_h = vec![0.0; v_dim];
        delta_rule_recurrence_gqa(
            &q,
            &k,
            &v,
            beta,
            &dt,
            &mut cache.ssm_states[il][..],
            &mut out_h,
            self.num_k_heads,
            self.head_k_dim,
            self.num_v_heads,
            self.head_v_dim,
        );

        let mut ssm_out_in = vec![0.0; v_dim];
        gated_rmsnorm(
            &out_h,
            gate,
            &d.ssm_norm_weight,
            self.base.config.eps,
            self.head_v_dim,
            &mut ssm_out_in,
        );
        // No normalisation after the gated RMS norm: llama.cpp goes from build_norm_gated straight to ssm_out.
        // (Two extra global L2 norms here forced every DeltaNet output to unit length.)
        ssm_out_in
    }

    /// The per-token middle of a full-attention layer: q/gate split, per-head norms, partial
    /// NEOX RoPE at `position`, the KV append, causal attention and the sigmoid output gate.
    /// Returns the `attn_output` input. Shared with the batched prefill (#4228).
    #[allow(clippy::too_many_arguments)]
    fn attention_mix_token(
        &self,
        a: &Qwen35OwnedAttentionLayer,
        cache: &mut Qwen35State,
        il: usize,
        position: usize,
        q_full: &[f32],
        k: &mut [f32],
        v: &[f32],
    ) -> Vec<f32> {
        let (q, gate) = self.attention_append_token(a, cache, il, position, q_full, k, v);
        self.attention_attend_token(cache, il, position, &q, &gate)
    }

    /// The mutating half of [`Self::attention_mix_token`]: splits q from its gate, applies the
    /// per-head norms and the partial NEOX RoPE at `position`, and appends this token's K/V to
    /// layer `il`. Returns the roped `q` and the gate.
    #[allow(clippy::too_many_arguments)]
    fn attention_append_token(
        &self,
        a: &Qwen35OwnedAttentionLayer,
        cache: &mut Qwen35State,
        il: usize,
        position: usize,
        q_full: &[f32],
        k: &mut [f32],
        v: &[f32],
    ) -> (Vec<f32>, Vec<f32>) {
        let num_heads = self.base.config.num_heads;
        let head_dim = a.attn_q_norm.len(); // 256
        let mut q = vec![0.0; num_heads * head_dim];
        let mut gate = vec![0.0; num_heads * head_dim];

        for h in 0..num_heads {
            let offset_q_full = h * head_dim * 2;
            let offset_q = h * head_dim;

            // Split into Q and gate
            q[offset_q..offset_q + head_dim]
                .copy_from_slice(&q_full[offset_q_full..offset_q_full + head_dim]);
            gate[offset_q..offset_q + head_dim]
                .copy_from_slice(&q_full[offset_q_full + head_dim..offset_q_full + head_dim * 2]);
        }

        crate::gguf::ops::apply_per_head_rms_norm(
            &mut q,
            &a.attn_q_norm,
            num_heads,
            self.base.config.eps,
        );
        crate::gguf::ops::apply_per_head_rms_norm(
            k,
            &a.attn_k_norm,
            self.base.config.num_kv_heads,
            self.base.config.eps,
        );

        // Partial NEOX RoPE over n_rot = 2 * sum(rope.dimension_sections) = 64 of each 256-wide head, as
        // llama.cpp ggml_rope_multi for text input. The base apply_rope assumed head_dim = hidden/heads = 128
        // and a full rotation, both wrong for Qwen3.5 attention.
        let n_rot = 2 * self.rope_sections.iter().sum::<usize>();
        let freq_base = self.base.config.rope_theta;
        apply_partial_neox_rope(&mut q, num_heads, head_dim, n_rot, position, freq_base);
        apply_partial_neox_rope(
            k,
            self.base.config.num_kv_heads,
            head_dim,
            n_rot,
            position,
            freq_base,
        );

        cache.kv_cache.append(il, k, v);
        (q, gate)
    }

    /// The read-only half of [`Self::attention_mix_token`]: causal attention of the roped `q`
    /// over positions `0..=position` of layer `il`'s KV cache, then the sigmoid output gate.
    /// Reads only rows the caller already appended, so the batched prefill (#4228) runs it for
    /// every token of a chunk in parallel once the chunk's K/V is in.
    fn attention_attend_token(
        &self,
        cache: &Qwen35State,
        il: usize,
        position: usize,
        q: &[f32],
        gate: &[f32],
    ) -> Vec<f32> {
        let k_cache = cache.kv_cache.get_k(il);
        let v_cache = cache.kv_cache.get_v(il);

        let mut attn_out_in = vec![0.0; q.len()];
        let num_kv_heads = self.base.config.num_kv_heads;
        let num_heads = self.base.config.num_heads;
        let group_size = num_heads / num_kv_heads;
        let head_dim = self.head_dim;

        for h in 0..num_heads {
            let kv_h = h / group_size;
            let q_h = &q[h * head_dim..(h + 1) * head_dim];

            let mut scores = vec![0.0; position + 1];
            for p in 0..=position {
                let mut dot = 0.0;
                let k_p = &k_cache[p * (num_kv_heads * head_dim) + kv_h * head_dim
                    ..p * (num_kv_heads * head_dim) + (kv_h + 1) * head_dim];
                for i in 0..head_dim {
                    dot += q_h[i] * k_p[i];
                }
                scores[p] = dot / (head_dim as f32).sqrt();
            }
            crate::gguf::ops::softmax(&mut scores);

            let out_h = &mut attn_out_in[h * head_dim..(h + 1) * head_dim];
            for p in 0..=position {
                let w = scores[p];
                let v_p = &v_cache[p * (num_kv_heads * head_dim) + kv_h * head_dim
                    ..p * (num_kv_heads * head_dim) + (kv_h + 1) * head_dim];
                for i in 0..head_dim {
                    out_h[i] += w * v_p[i];
                }
            }
        }
        // Output gate, as llama.cpp: attn_output * sigmoid(gate) before the output projection
        // (the attn_gated node). The gate split off the joint Q projection was never applied.
        apply_sigmoid_gate(&mut attn_out_in, gate);
        attn_out_in
    }
}

/// Prompt tokens per layer-major prefill chunk (#4228): enough rows to amortise each Q4_K
/// weight tile's DRAM pass (the multi-row GEMM measured 8.4x at 32 rows, 8.6x at 64), small
/// enough that a chunk's activations stay cache-resident.
pub const QWEN35_PREFILL_CHUNK: usize = 64;

/// Row-wise RMS norm: each `width`-wide row of `x` into the same row of `out`.
fn rms_norm_rows(x: &[f32], weight: &[f32], eps: f32, width: usize, out: &mut [f32]) {
    for (row, o) in x.chunks_exact(width).zip(out.chunks_exact_mut(width)) {
        crate::gguf::ops::rms_norm_into(row, weight, eps, o);
    }
}

/// `acc += x`, element-wise.
fn add_into(acc: &mut [f32], x: &[f32]) {
    for (a, &v) in acc.iter_mut().zip(x) {
        *a += v;
    }
}

/// The per-token CPU REFERENCE the one engine is tested against: prefill
/// `input_tokens` one position at a time through the Qwen3.5 CPU forward, then
/// decode up to `gen_config.max_tokens` more with the dense path's token choice.
/// Returns the prompt followed by the new tokens.
///
/// Test-only since #4263: every production verb generates through
/// `realizar::session`, and this loop is kept as the independent oracle it is
/// compared with — it shares no loop, no state sizing and no batched prefill
/// with the engine.
///
/// # Errors
/// An empty prompt, a layer the Qwen3.5 loader cannot read, or a forward-pass failure.
#[cfg(test)]
pub(crate) fn qwen35_reference_generate(
    mapped: &crate::gguf::MappedGGUFModel,
    base: &OwnedQuantizedModel,
    input_tokens: &[u32],
    gen_config: &crate::gguf::QuantizedGenerateConfig,
) -> Result<Vec<u32>> {
    use rand::SeedableRng;
    if input_tokens.is_empty() {
        return Err(crate::error::RealizarError::InvalidShape {
            reason: "qwen35_reference_generate: prompt cannot be empty".to_string(),
        });
    }
    qwen35_check_context(input_tokens.len(), base.config.context_length)?;
    let qwen = Qwen35Model::from_model_and_layers(base, &mapped.model, mapped.data())?;
    let max_seq_len = input_tokens.len() + gen_config.max_tokens + 1;
    let mut state = qwen.new_state(max_seq_len);
    let mut rng = rand::rngs::StdRng::seed_from_u64(gen_config.seed);

    let mut logits = Vec::new();
    for (pos, &token) in input_tokens.iter().enumerate() {
        logits = qwen.forward_single_qwen35(token, &mut state, pos)?;
    }
    let mut tokens = input_tokens.to_vec();
    for _ in 0..gen_config.max_tokens {
        let next = if gen_config.temperature == 0.0 || gen_config.top_k == 1 {
            crate::gguf::ops::argmax(&logits)
        } else {
            OwnedQuantizedModel::sample_topk_seeded(
                &logits,
                gen_config.temperature,
                gen_config.top_k,
                gen_config.top_p,
                &mut rng,
            )
        };
        tokens.push(next);
        if gen_config.stop_tokens.contains(&next) || tokens.len() >= max_seq_len {
            break;
        }
        logits = qwen.forward_single_qwen35(next, &mut state, tokens.len() - 1)?;
    }
    Ok(tokens)
}

/// Why a Qwen3.5 run served its tokens from the CPU forward.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Qwen35CpuReason {
    /// The caller asked for the CPU (`--no-gpu` / `force_cpu`).
    Requested,
    /// This binary was built without the `cuda` feature, so there is no GPU
    /// backend to route to — the hybrid forward itself exists on both (#3090).
    NoCudaBackend,
}

/// Which backend serves a Qwen3.5 (Gated `DeltaNet`) GGUF.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Qwen35Route {
    /// `Qwen35CudaModel::forward_single` (#3090).
    Gpu,
    /// `Qwen35Model::forward_single_qwen35` (#3091).
    Cpu(Qwen35CpuReason),
}

/// Pure: which backend a Qwen3.5 GGUF is routed to.
///
/// Split out from the dispatcher so the decision is falsifiable on a host with
/// no GPU and in a build without `cuda`: `cuda_backend` is the caller's
/// `cfg!(feature = "cuda")`, not a device probe (a device that is present but
/// unusable is a *fallback*, which is printed — see
/// [`run_qwen35_generate_dispatch`] — never a silent re-route).
#[must_use]
pub fn qwen35_route(no_gpu: bool, cuda_backend: bool) -> Qwen35Route {
    if no_gpu {
        Qwen35Route::Cpu(Qwen35CpuReason::Requested)
    } else if cuda_backend {
        Qwen35Route::Gpu
    } else {
        Qwen35Route::Cpu(Qwen35CpuReason::NoCudaBackend)
    }
}

/// The bound every other generate path checks (GH-167), which the Qwen3.5 dispatch
/// skipped: #3596's 262k rung prefilled 263,089 positions into a model that declares
/// 262,144, on both the GPU and the CPU route, and nothing refused it.
///
/// # Errors
/// [`crate::error::RealizarError::ContextLimitExceeded`] when the prompt is longer
/// than the model's declared `context_length`.
pub fn qwen35_check_context(prompt_len: usize, context_length: usize) -> Result<()> {
    if prompt_len > context_length {
        return Err(crate::error::RealizarError::ContextLimitExceeded {
            provided: prompt_len,
            maximum: context_length,
        });
    }
    Ok(())
}

/// The one-line notice a route owes the user, or `None` when it owes none.
///
/// #3477: the GPU case used to print `[qwen35: Gated DeltaNet runs on the CPU;
/// the GPU backend does not implement it yet (#3090)]` — which is now false: the
/// hybrid runs on the GPU. What a user still needs told is the case where they
/// asked for the GPU and this *binary* cannot give them one. Asking for the CPU
/// explicitly is not news.
#[must_use]
pub fn qwen35_route_notice(route: Qwen35Route) -> Option<&'static str> {
    match route {
        Qwen35Route::Cpu(Qwen35CpuReason::NoCudaBackend) => Some(
            "[qwen35: this binary has no CUDA backend (built without --features cuda); \
             the Gated DeltaNet forward runs on the CPU (#3091)]",
        ),
        Qwen35Route::Gpu | Qwen35Route::Cpu(Qwen35CpuReason::Requested) => None,
    }
}

/// The prefix of the loud, never-silent CPU fallback for the hybrid GPU path.
pub const QWEN35_GPU_FALLBACK_PREFIX: &str = "warning: GPU (CUDA) qwen35 path rejected";

/// Positions the F2 hybrid guard forwards on both backends before it will let
/// the GPU serve a token. Same cap as the dense guard's `gpu_probe`.
#[cfg(feature = "cuda")]
pub(crate) const QWEN35_F2_PROBE_MAX: usize = 64;

/// The F2 runtime guard for the hybrid: forward the real prompt through BOTH
/// backends and accept the GPU only if every real position agrees.
///
/// The dense `validate_gpu_first_token` cannot serve here — its CPU reference is
/// `forward_single_with_cache` and its GPU probe `forward_gpu_resident`, neither
/// of which exists for this architecture. The decision itself is shared:
/// `f2_multi_position_report` with its floors (0.95 / 0.98 / 0.90), so the
/// hybrid is judged by the same rule as every other GPU path. The probe path is
/// [`F2ProbePath::Batched`] (#3596): `qwen35_gpu_decode` prefills the prompt with the
/// batched prefill and decodes from there, so the probe does exactly that.
///
/// Both states are throwaway: the guard allocates its own, and the generation
/// that follows allocates another.
/// What the F2 guard concluded — and, separately, whether it concluded
/// anything at all.
///
/// The distinction exists for the receipt (#3604): three of this guard's exits
/// return "let the GPU serve" WITHOUT having compared a single position —
/// `SKIP_PARITY_GATE=1`, a probe shorter than two tokens, a CPU reference that
/// would not run. A receipt written on any of those would let a one-token
/// prompt "validate" the (model, apr, device) triple for every prompt after it.
/// So the wrapper writes a receipt on `Accepted` only.
#[cfg(feature = "cuda")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum F2Verdict {
    /// CPU and GPU agreed on `positions_judged` positions.
    Accepted {
        /// Probe positions plus the one greedy decode step.
        positions_judged: usize,
    },
    /// They disagreed, or the GPU probe failed: fail closed.
    Rejected,
    /// Nothing was compared. The GPU may serve, but nothing was proved.
    NotJudged,
}

#[cfg(feature = "cuda")]
impl F2Verdict {
    const fn lets_the_gpu_serve(self) -> bool {
        !matches!(self, Self::Rejected)
    }
}

#[cfg(feature = "cuda")]
fn f2_validate_qwen35(
    gpu: &mut crate::gguf::cuda::Qwen35CudaModel<'_>,
    cpu: &Qwen35Model<'_>,
    probe_context: &[u32],
) -> F2Verdict {
    // Same escape hatch as the dense gate, and the same one `apr parity` uses.
    if std::env::var("SKIP_PARITY_GATE").is_ok_and(|v| v == "1") {
        return F2Verdict::NotJudged;
    }
    let probe = &probe_context[probe_context.len().saturating_sub(QWEN35_F2_PROBE_MAX)..];
    // A one-token probe has no REAL position (≥1) to judge; position 0 is the
    // context-less near-tie the dense gate excludes for the same reason.
    if probe.len() < 2 {
        return F2Verdict::NotJudged;
    }
    let Some(cpu_per_pos) = f2_qwen35_cpu_reference(cpu, probe) else {
        return F2Verdict::NotJudged; // the CPU forward itself failed: nothing to judge against.
    };
    let decode_token = cpu_per_pos
        .get(probe.len().saturating_sub(1))
        .map_or(0, |l| crate::infer::argmax_u32(l));
    let gpu_per_pos = match f2_qwen35_gpu_logits(gpu, probe, decode_token) {
        Ok(v) => v,
        Err(msg) => {
            eprintln!("{msg}");
            return F2Verdict::Rejected; // fail closed.
        },
    };
    let report = crate::infer::f2_multi_position_report(&cpu_per_pos, &gpu_per_pos);
    if report.accepted {
        F2Verdict::Accepted {
            positions_judged: cpu_per_pos.len(),
        }
    } else {
        eprintln!(
            "{}",
            crate::infer::f2_divergence_msg(&report, crate::infer::F2ProbePath::Batched)
        );
        F2Verdict::Rejected
    }
}

/// The outcome of the receipted guard, for the caller and — once #3606's
/// `StageTimings` lands — for `apr run --json`, where `source` is what makes
/// a cached validation distinguishable from a fresh one (#3604 `done_when` 5).
#[cfg(feature = "cuda")]
#[derive(Debug, Clone, PartialEq)]
pub struct F2Outcome {
    /// May the GPU serve this run?
    pub accepted: bool,
    /// `"receipt"` when the forward was skipped on a matching receipt,
    /// `"fresh"` when it ran, `"not-judged"` when it ran but compared nothing.
    pub source: &'static str,
    /// Wall time of the guard itself, EXCLUDING the hash. ~0 on a receipt hit.
    pub validate_ms: f64,
    /// Wall time of hashing the model file — the price of the receipt's key,
    /// paid on every run, reported separately so it cannot hide in either
    /// number above.
    pub sha256_ms: f64,
    /// Where the receipt was read from or written to, if a cache dir exists.
    pub receipt_path: Option<std::path::PathBuf>,
}

/// [`f2_validate_qwen35`] behind its receipt (#3604).
///
/// Validate once per (model sha256, apr version, device); later runs of the
/// same triple read the receipt and skip the forward; `--revalidate` forces a
/// fresh run and rewrites it. Every path that is not a three-key match
/// validates — see `f2_receipt.rs` for the table — and a receipt is written on
/// [`F2Verdict::Accepted`] only, never on a verdict that judged nothing.
///
/// Nothing in here can fail the run except the guard's own rejection: an
/// unwritable cache directory is reported and the next run simply validates
/// again.
#[cfg(feature = "cuda")]
fn f2_validate_qwen35_receipted(
    gpu: &mut crate::gguf::cuda::Qwen35CudaModel<'_>,
    cpu: &Qwen35Model<'_>,
    probe_context: &[u32],
    model_bytes: &[u8],
    device_name: &str,
) -> F2Outcome {
    let hash = Qwen35ModelHash::of(model_bytes);
    f2_validate_qwen35_receipted_hashed(gpu, cpu, probe_context, &hash, device_name)
}

/// The receipt key's model half, hashed once (#3595).
///
/// [`f2_validate_qwen35_receipted`] hashes the whole file on every call — 8 s
/// on the 27B. A caller that keeps the model resident across calls hashes it
/// once, when it loads, and hands this to
/// [`f2_validate_qwen35_receipted_hashed`].
#[cfg(feature = "cuda")]
#[derive(Debug, Clone)]
pub(crate) struct Qwen35ModelHash {
    /// Hex sha256 of the model file's bytes.
    pub(crate) sha256: String,
    /// Wall time the hash took, reported in the guard's line.
    pub(crate) sha256_ms: f64,
}

#[cfg(feature = "cuda")]
impl Qwen35ModelHash {
    pub(crate) fn of(model_bytes: &[u8]) -> Self {
        let start = std::time::Instant::now();
        let sha256 = crate::gguf::f2_receipt::model_sha256(model_bytes);
        Self {
            sha256,
            sha256_ms: start.elapsed().as_secs_f64() * 1000.0,
        }
    }
}

/// [`f2_validate_qwen35_receipted`] with the model already hashed.
#[cfg(feature = "cuda")]
pub(crate) fn f2_validate_qwen35_receipted_hashed(
    gpu: &mut crate::gguf::cuda::Qwen35CudaModel<'_>,
    cpu: &Qwen35Model<'_>,
    probe_context: &[u32],
    hash: &Qwen35ModelHash,
    device_name: &str,
) -> F2Outcome {
    use crate::gguf::f2_receipt::{
        apr_version, decide, read_receipt, receipt_dir, receipt_path, revalidate_requested,
        unix_now, write_receipt, F2Decision, F2Receipt, F2ReceiptKey, F2_RECEIPT_SCHEMA,
    };

    let key = F2ReceiptKey {
        model_sha256: hash.sha256.clone(),
        apr_version: apr_version(),
        device: device_name.to_string(),
    };
    let sha256_ms = hash.sha256_ms;

    let path = receipt_dir().map(|d| receipt_path(&d, &key.model_sha256));
    let found = match path.as_deref() {
        Some(p) => read_receipt(p),
        None => Ok(None),
    };

    match decide(found, &key, revalidate_requested()) {
        F2Decision::Skip { receipt } => {
            let age_s = unix_now().saturating_sub(receipt.validated_at);
            eprintln!(
                "F2 guard: receipt matches (model sha256 {}…, apr {}, {}) — validated {}s ago on {} positions; CPU reference forward skipped [source=receipt, sha256 {:.0} ms]. `apr run --revalidate` forces a fresh run.",
                &key.model_sha256[..12],
                key.apr_version,
                key.device,
                age_s,
                receipt.positions_judged,
                sha256_ms
            );
            return F2Outcome {
                accepted: true,
                source: "receipt",
                validate_ms: 0.0,
                sha256_ms,
                receipt_path: path,
            };
        },
        F2Decision::Validate(reason) => {
            eprintln!("F2 guard: validating on this run ({reason}) [source=fresh]");
        },
    }

    let start = std::time::Instant::now();
    let verdict = f2_validate_qwen35(gpu, cpu, probe_context);
    let validate_ms = start.elapsed().as_secs_f64() * 1000.0;

    let source = match verdict {
        F2Verdict::Accepted { positions_judged } => {
            match path.as_deref() {
                Some(p) => {
                    let receipt = F2Receipt {
                        schema: F2_RECEIPT_SCHEMA,
                        key,
                        validated_at: unix_now(),
                        positions_judged,
                    };
                    match write_receipt(p, &receipt) {
                        Ok(()) => eprintln!(
                            "F2 guard: passed in {validate_ms:.0} ms on {positions_judged} positions; receipt written to {} — the next run of this (model, apr, device) skips it.",
                            p.display()
                        ),
                        Err(e) => eprintln!(
                            "F2 guard: passed in {validate_ms:.0} ms, but the receipt could not be written ({e}); the next run validates again. Set APR_F2_RECEIPT_DIR to a writable directory."
                        ),
                    }
                }
                None => eprintln!(
                    "F2 guard: passed in {validate_ms:.0} ms; no cache directory (no HOME, XDG_CACHE_HOME or APR_F2_RECEIPT_DIR), so no receipt — every run validates."
                ),
            }
            "fresh"
        },
        F2Verdict::NotJudged => {
            eprintln!(
                "F2 guard: nothing was compared on this run (probe too short, SKIP_PARITY_GATE, or the CPU reference did not run); the GPU serves, and NO receipt is written."
            );
            "not-judged"
        },
        F2Verdict::Rejected => "fresh",
    };

    F2Outcome {
        accepted: verdict.lets_the_gpu_serve(),
        source,
        validate_ms,
        sha256_ms,
        receipt_path: path,
    }
}

/// CPU logits at every probe position plus one greedy decode step — the
/// reference half of [`f2_validate_qwen35`].
#[cfg(feature = "cuda")]
fn f2_qwen35_cpu_reference(cpu: &Qwen35Model<'_>, probe: &[u32]) -> Option<Vec<Vec<f32>>> {
    let mut state = cpu.new_state(probe.len() + 2);
    let mut per_pos: Vec<Vec<f32>> = Vec::with_capacity(probe.len() + 1);
    for (pos, &tok) in probe.iter().enumerate() {
        per_pos.push(cpu.forward_single_qwen35(tok, &mut state, pos).ok()?);
    }
    let last = per_pos.last()?;
    let next = crate::infer::argmax_u32(last);
    per_pos.push(
        cpu.forward_single_qwen35(next, &mut state, probe.len())
            .ok()?,
    );
    Some(per_pos)
}

/// GPU logits for the same positions, from a fresh device state.
#[cfg(feature = "cuda")]
fn f2_qwen35_gpu_logits(
    gpu: &mut crate::gguf::cuda::Qwen35CudaModel<'_>,
    probe: &[u32],
    decode_token: u32,
) -> std::result::Result<Vec<Vec<f32>>, String> {
    let steps = probe.len() + 1;
    // #3596: the probe takes the path a real run takes — the batched prefill over the
    // probe, then one decode step — in a state sized for exactly those `steps`
    // positions, not the request's whole KV.
    let mut state = gpu
        .new_state_with_len(steps)
        .map_err(|e| format!("F2 qwen35 probe: the device state would not allocate: {e}"))?;
    let every_position: Vec<usize> = (0..probe.len()).collect();
    let mut per_pos = gpu
        .prefill_logits_at(probe, &mut state, 0, &every_position)
        .map_err(|e| crate::infer::gpu_forward_failure_msg(0, steps, &e))?;
    match gpu.forward_single(decode_token, &mut state, probe.len()) {
        Ok(logits) => per_pos.push(logits),
        Err(e) => {
            return Err(crate::infer::gpu_forward_failure_msg(
                probe.len(),
                steps,
                &e,
            ))
        },
    }
    Ok(per_pos)
}

#[cfg(test)]
mod qwen35_route_tests {
    use super::{
        qwen35_check_context, qwen35_route, qwen35_route_notice, Qwen35CpuReason, Qwen35Route,
    };

    // #3596: the 262k rung's 263,089-position prompt ran on a 262,144-position model.
    // The boundary is inclusive (a prompt of exactly `context_length` is served), as
    // on every other generate path (GH-167).
    #[test]
    fn a_prompt_past_the_declared_context_is_refused_with_both_numbers() {
        const CTX: usize = 262_144;
        for (len, ok) in [
            (1, true),
            (CTX - 1, true),
            (CTX, true),
            (CTX + 1, false),
            (263_089, false),
        ] {
            let got = qwen35_check_context(len, CTX);
            assert_eq!(got.is_ok(), ok, "prompt of {len} against {CTX}: {got:?}");
            if let Err(e) = got {
                let msg = e.to_string();
                assert!(
                    msg.contains(&len.to_string()) && msg.contains(&CTX.to_string()),
                    "{msg}"
                );
            }
        }
    }

    // #3090/#3477: with a CUDA build and no --no-gpu, the hybrid goes to the GPU.
    // This is the whole point of the ticket; if it ever reads Cpu again, `apr run
    // --gpu` is silently serving CPU tokens.
    #[test]
    fn a_cuda_build_that_was_not_told_otherwise_routes_to_the_gpu() {
        assert_eq!(qwen35_route(false, true), Qwen35Route::Gpu);
        assert_eq!(qwen35_route_notice(Qwen35Route::Gpu), None);
    }

    #[test]
    fn no_gpu_wins_over_a_present_cuda_backend() {
        assert_eq!(
            qwen35_route(true, true),
            Qwen35Route::Cpu(Qwen35CpuReason::Requested)
        );
        // The user asked for this: it is not news.
        assert_eq!(
            qwen35_route_notice(Qwen35Route::Cpu(Qwen35CpuReason::Requested)),
            None
        );
    }

    // A binary without the cuda feature still runs the model — but the user who
    // asked for a GPU must be told why they did not get one.
    #[test]
    fn a_build_without_cuda_says_so() {
        let route = qwen35_route(false, false);
        assert_eq!(route, Qwen35Route::Cpu(Qwen35CpuReason::NoCudaBackend));
        let notice = qwen35_route_notice(route).expect("the no-backend case owes a notice");
        assert!(
            notice.contains("no CUDA backend"),
            "the notice must name the missing backend, not the architecture: {notice}"
        );
        assert!(
            !notice.contains("#3090"),
            "#3090 is the GPU forward, which now exists — citing it here is the  withdrawn 'the GPU does not implement it' notice: {notice}"
        );
    }

    // Both CPU reasons are reachable from the routing function, so neither arm
    // of the notice is dead code.
    #[test]
    fn every_cpu_reason_is_produced_by_the_router() {
        for (no_gpu, cuda, want) in [
            (true, true, Qwen35CpuReason::Requested),
            (true, false, Qwen35CpuReason::Requested),
            (false, false, Qwen35CpuReason::NoCudaBackend),
        ] {
            assert_eq!(qwen35_route(no_gpu, cuda), Qwen35Route::Cpu(want));
        }
    }
}

#[cfg(test)]
mod qwen35_ssm_meta_tests {
    use super::{qwen35_ssm_meta, Qwen35SsmMeta};
    use crate::gguf::types::GGUFValue;
    use std::collections::HashMap;

    fn md(pairs: &[(&str, GGUFValue)]) -> HashMap<String, GGUFValue> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_string(), v.clone()))
            .collect()
    }

    #[test]
    fn test_empty_metadata_gives_the_qwen35_0_8b_defaults() {
        assert_eq!(
            qwen35_ssm_meta(&HashMap::new()),
            Qwen35SsmMeta {
                head_k_dim: 16,
                num_k_heads: 1,
                num_v_heads: 16,
                conv_kernel: 4,
                rope_sections: [16, 24, 24, 0],
            }
        );
    }

    #[test]
    fn test_qwen2_keys_win_and_qwen35_keys_are_the_fallback() {
        let meta = qwen35_ssm_meta(&md(&[
            ("qwen2.ssm.state_size", GGUFValue::UInt32(128)),
            ("qwen35.ssm.state_size", GGUFValue::UInt32(999)),
            ("qwen35.ssm.group_count", GGUFValue::UInt32(2)),
            ("qwen35.ssm.time_step_rank", GGUFValue::Int32(32)),
            ("qwen35.ssm.conv_kernel", GGUFValue::UInt32(4)),
        ]));
        assert_eq!(meta.head_k_dim, 128, "the qwen2 key is preferred");
        assert_eq!(meta.num_k_heads, 2, "no qwen2 key: the qwen35 one is read");
        assert_eq!(meta.num_v_heads, 32, "Int32 is accepted like UInt32");
        assert_eq!(meta.conv_kernel, 4);
    }

    #[test]
    fn test_a_non_integer_value_falls_through_to_the_default() {
        // A string where an integer belongs is not a value: it must not be read as 0.
        let meta = qwen35_ssm_meta(&md(&[(
            "qwen2.ssm.conv_kernel",
            GGUFValue::String("four".to_string()),
        )]));
        assert_eq!(meta.conv_kernel, 4);
    }

    #[test]
    fn test_rope_sections_are_defaulted_per_section() {
        // Two entries given: the last two sections keep their defaults, and a
        // non-integer entry keeps its own.
        let meta = qwen35_ssm_meta(&md(&[(
            "qwen2.rope.dimension_sections",
            GGUFValue::Array(vec![
                GGUFValue::UInt32(8),
                GGUFValue::Float32(1.0),
                GGUFValue::Int32(12),
            ]),
        )]));
        assert_eq!(meta.rope_sections, [8, 24, 12, 0]);
    }
}

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

    #[test]
    fn test_l2_norm_per_head_normalises_each_head_on_its_own() {
        // Two heads whose magnitudes differ 10x: each must come out as [0.6, 0.8].
        let mut x = [3.0, 4.0, 30.0, 40.0];
        l2_norm_per_head(&mut x, 2, 1e-12);
        for (got, want) in x.iter().zip([0.6, 0.8, 0.6, 0.8]) {
            assert!((got - want).abs() < 1e-6, "{x:?}");
        }
    }

    #[test]
    fn test_apply_sigmoid_gate_scales_by_sigmoid() {
        let mut x = [2.0, 2.0, 2.0];
        apply_sigmoid_gate(&mut x, &[0.0, 40.0, -40.0]);
        assert!((x[0] - 1.0).abs() < 1e-6, "sigmoid(0) = 0.5: {x:?}");
        assert!((x[1] - 2.0).abs() < 1e-6, "sigmoid(40) ~ 1: {x:?}");
        assert!(x[2].abs() < 1e-6, "sigmoid(-40) ~ 0: {x:?}");
    }

    #[test]
    fn test_partial_neox_rope_is_identity_at_position_zero() {
        let orig: Vec<f32> = (0..8u8).map(|i| f32::from(i) + 1.0).collect();
        let mut x = orig.clone();
        apply_partial_neox_rope(&mut x, 1, 8, 4, 0, 10_000.0);
        assert_eq!(x, orig);
    }

    #[test]
    fn test_partial_neox_rope_rotates_half_pairs_and_leaves_the_tail() {
        // head_dim 8, n_rot 4: the pairs are (0,2) and (1,3) (NEOX halves, not (0,1),(2,3)),
        // and dims 4..8 are not rotated at all.
        let base = 10_000.0f32;
        let mut x = [0.0, 1.0, 0.0, 0.0, 5.0, 6.0, 7.0, 8.0];
        apply_partial_neox_rope(&mut x, 1, 8, 4, 3, base);
        let theta1 = 3.0 * base.powf(-2.0 / 4.0); // pair j = 1
        assert!(
            (x[1] - theta1.cos()).abs() < 1e-6 && (x[3] - theta1.sin()).abs() < 1e-6,
            "{x:?}"
        );
        assert_eq!(
            (x[0], x[2]),
            (0.0, 0.0),
            "the (0,2) pair was all zeros: {x:?}"
        );
        assert_eq!(&x[4..], &[5.0, 6.0, 7.0, 8.0]);
    }
}

/// Proof-obligation bodies for `contracts/qwen35-hybrid-forward-v1.yaml` (QHF-INV-001..004,
/// QHF-INV-006, QHF-CON-007).
#[cfg(test)]
#[path = "forward_qwen35_contract_tests.rs"]
mod qhf_contract_tests;

/// The Gated `DeltaNet` head mapping when `num_v_heads > num_k_heads`
/// (Qwen3.5 4B/9B/27B) — PMAT-3477, #3346/#3510.
#[cfg(test)]
#[path = "forward_qwen35_gqa_tests.rs"]
mod qwen35_gqa_tests;

#[cfg(test)]
#[path = "forward_qwen35_prefill_tests.rs"]
mod qwen35_prefill_tests;