use super::checkpoint::{read_buffer, read_u32, write_buffer, write_u32};
use super::*;
use crate::engine::qwen::{QwenModel, QwenTensor};
use std::ffi::CStr;

const HIDDEN: u32 = 2_560;
const HC: u32 = 4;
const HC_WIDTH: u32 = HIDDEN * HC;
const HC_RANK: u32 = 320;
const LAYERS: usize = 48;
const GDN_HEADS_K: u32 = 16;
const GDN_HEADS_V: u32 = 48;
const HEAD_DIM: u32 = 128;
const GDN_QKV: u32 = (GDN_HEADS_K * 2 + GDN_HEADS_V) * HEAD_DIM;
const GDN_VALUE: u32 = GDN_HEADS_V * HEAD_DIM;
const GDN_CONTROLS: u32 = GDN_VALUE + GDN_HEADS_V * 2;
const ATTN_HEADS: u32 = 24;
const ATTN_KV_HEADS: u32 = 2;
const ATTN_DIM: u32 = 256;
const ATTN_WIDTH: u32 = ATTN_HEADS * ATTN_DIM;
const ATTN_KV_WIDTH: u32 = ATTN_KV_HEADS * ATTN_DIM;
const EXPERTS: u32 = 512;
const EXPERTS_USED: usize = 10;
const EXPERT_WIDTH: u32 = 640;
const VOCAB: u32 = 248_320;
const DENSE_BUDGET: u32 = 2_048;
const PLE_HEADS: usize = 16;
const PLE_HEAD_DIM: u32 = 160;
const PLE_HISTORY: usize = 2;
const PLE_CONV_STATE: u32 = 9;
const EOS_TOKEN: i32 = 248_044;
const PLE_ROW_BYTES: usize = 100;
const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01";
const CHECKPOINT_VERSION: u32 = 2;
const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024;

#[derive(Clone, Copy)]
struct Weight<'a> {
    tensor: &'a QwenTensor,
    expert: Option<u32>,
}

struct Affine<'a> {
    packed: Weight<'a>,
    scales: Weight<'a>,
    biases: Weight<'a>,
    bits: u32,
    group: u32,
}

enum LayerState {
    Gdn { conv: Buffer, recurrent: Buffer },
    Attention { kv: Buffer },
}

struct Scratch {
    hidden: Buffer,
    hc: Buffer,
    hc_norm: Buffer,
    hc_mix: Buffer,
    rank: Buffer,
    block: Buffer,
    injection: Buffer,
    qkv: Buffer,
    controls: Buffer,
    gdn_raw: Buffer,
    gdn_out: Buffer,
    router: Buffer,
    route_ids: Buffer,
    route_weights: Buffer,
    gate: Buffer,
    up: Buffer,
    mid: Buffer,
    expert: Buffer,
    moe: Buffer,
    shared: Buffer,
    q_packed: Buffer,
    q: Buffer,
    q_gate: Buffer,
    k: Buffer,
    k_rope: Buffer,
    v: Buffer,
    attention: Buffer,
    ple_packed: Buffer,
    ple_scales: Buffer,
    ple_biases: Buffer,
    ple_embedding: Buffer,
    ple_key: Buffer,
    ple_value: Buffer,
    ple_gated: Buffer,
    ple_norm: Buffer,
    ple_output: Buffer,
    logits: Buffer,
}

impl Scratch {
    fn new() -> Result<Self, String> {
        Ok(Self {
            hidden: Buffer::floats(HIDDEN.into())?,
            hc: Buffer::floats(HC_WIDTH.into())?,
            hc_norm: Buffer::floats(HC_WIDTH.into())?,
            hc_mix: Buffer::floats(HC_WIDTH.into())?,
            rank: Buffer::floats(HC_RANK.into())?,
            block: Buffer::floats(HIDDEN.into())?,
            injection: Buffer::floats(HC.into())?,
            qkv: Buffer::floats(GDN_QKV.into())?,
            controls: Buffer::floats(GDN_CONTROLS.into())?,
            gdn_raw: Buffer::floats(GDN_VALUE.into())?,
            gdn_out: Buffer::floats(GDN_VALUE.into())?,
            router: Buffer::floats(EXPERTS.into())?,
            route_ids: Buffer::bytes((EXPERTS_USED * 4) as u64)?,
            route_weights: Buffer::floats(EXPERTS_USED as u64)?,
            gate: Buffer::floats(EXPERT_WIDTH.into())?,
            up: Buffer::floats(EXPERT_WIDTH.into())?,
            mid: Buffer::floats(EXPERT_WIDTH.into())?,
            expert: Buffer::floats(HIDDEN.into())?,
            moe: Buffer::floats(HIDDEN.into())?,
            shared: Buffer::floats(HIDDEN.into())?,
            q_packed: Buffer::floats((ATTN_WIDTH * 2).into())?,
            q: Buffer::floats(ATTN_WIDTH.into())?,
            q_gate: Buffer::floats(ATTN_WIDTH.into())?,
            k: Buffer::floats(ATTN_KV_WIDTH.into())?,
            k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?,
            v: Buffer::floats(ATTN_KV_WIDTH.into())?,
            attention: Buffer::floats(ATTN_WIDTH.into())?,
            ple_packed: Buffer::bytes((PLE_HEADS as u64) * 80)?,
            ple_scales: Buffer::bytes((PLE_HEADS as u64) * 10)?,
            ple_biases: Buffer::bytes((PLE_HEADS as u64) * 10)?,
            ple_embedding: Buffer::floats(HIDDEN.into())?,
            ple_key: Buffer::floats(HC_WIDTH.into())?,
            ple_value: Buffer::floats(HIDDEN.into())?,
            ple_gated: Buffer::floats(HC_WIDTH.into())?,
            ple_norm: Buffer::floats(HC_WIDTH.into())?,
            ple_output: Buffer::floats(HC_WIDTH.into())?,
            logits: Buffer::floats(VOCAB.into())?,
        })
    }
}

struct PleContract {
    multipliers: [i64; 3],
    sizes: [i64; PLE_HEADS],
    offsets: [i64; PLE_HEADS],
}

struct PleState {
    history: [i32; PLE_HISTORY],
    conv: Buffer,
}

pub(in crate::engine) struct QwenExecutor {
    model: QwenModel,
    states: Vec<LayerState>,
    ple_contract: PleContract,
    ple_state: PleState,
    scratch: Scratch,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    position: u32,
    context: u32,
    checkpoint_tag: [u8; 32],
    _context: Context,
}

pub(in crate::engine) struct QwenResidentState {
    states: Vec<LayerState>,
    ple_state: PleState,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    position: u32,
    checkpoint_tag: [u8; 32],
}

impl QwenExecutor {
    pub(super) fn open(model: QwenModel, context: u32) -> Result<Self, String> {
        let ple_contract = ple_contract(&model)?;
        let native = Context::open_qwen(model.memory().admission)?;
        let states = allocate_states(context)?;
        Ok(Self {
            model,
            states,
            ple_contract,
            ple_state: allocate_ple_state()?,
            scratch: Scratch::new()?,
            logits: vec![0.0; VOCAB as usize],
            tokens: Vec::new(),
            position: 0,
            context,
            checkpoint_tag: [0; 32],
            _context: native,
        })
    }

    pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
        if token < 0 || token as u32 >= VOCAB {
            return Err(format!("token {token} is outside the Qwen vocabulary"));
        }
        if self.position >= self.context {
            return Err(format!(
                "the Qwen executor supports {} tokens per session",
                self.context
            ));
        }
        if self.position + 1 > DENSE_BUDGET {
            return Err(
                "Qwen sparse QSA selection is required beyond 2048 tokens; native QSA belongs to issue #97"
                    .into(),
            );
        }

        self.begin_token(token)?;

        for layer in 0..LAYERS {
            if layer == 1 {
                self.ple(token)?;
            }
            self.encode_layer(layer)?;
        }
        self.final_output()?;
        self.tokens.push(token);
        self.position += 1;
        Ok(())
    }

    fn begin_token(&self, token: i32) -> Result<(), String> {
        let embedding = self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?;
        let mut args = args();
        args.u[0] = HIDDEN;
        args.u[2] = embedding.bits;
        args.u[3] = embedding.group;
        args.u[4] = token as u32;
        let commands = Commands::begin()?;
        self.dispatch(
            c"kernel_qwen_affine_embedding",
            &self.scratch.hidden,
            None,
            None,
            None,
            &[
                self.view(embedding.packed),
                self.view(embedding.scales),
                self.view(embedding.biases),
            ],
            &args,
            HIDDEN,
            1,
        )?;
        args.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_repeat4",
            &self.scratch.hc,
            Some(&self.scratch.hidden),
            None,
            None,
            &[],
            &args,
            HC_WIDTH,
            1,
        )?;
        commands.finish()
    }

    fn ple(&mut self, token: i32) -> Result<(), String> {
        let rows = ple_rows(&self.ple_contract, self.ple_state.history, token)?;
        let packed = self.model.tensor("ngram.weight")?;
        let scales = self.model.tensor("ngram.scales")?;
        let biases = self.model.tensor("ngram.biases")?;
        let mut packed_stage = [0_u8; PLE_HEADS * 80];
        let mut scales_stage = [0_u8; PLE_HEADS * 10];
        let mut biases_stage = [0_u8; PLE_HEADS * 10];
        let packed_bytes = self.model.tensor_bytes(packed)?;
        let scales_bytes = self.model.tensor_bytes(scales)?;
        let biases_bytes = self.model.tensor_bytes(biases)?;
        for (head, &row) in rows.iter().enumerate() {
            let bytes = gather_ple_row(row, packed_bytes, scales_bytes, biases_bytes)?;
            packed_stage[head * 80..][..80].copy_from_slice(&bytes[..80]);
            scales_stage[head * 10..][..10].copy_from_slice(&bytes[80..90]);
            biases_stage[head * 10..][..10].copy_from_slice(&bytes[90..]);
        }
        self.scratch.ple_packed.write(0, &packed_stage)?;
        self.scratch.ple_scales.write(0, &scales_stage)?;
        self.scratch.ple_biases.write(0, &biases_stage)?;

        let commands = Commands::begin()?;
        let mut dequant = args();
        dequant.u[0] = PLE_HEAD_DIM;
        dequant.u[1] = PLE_HEADS as u32;
        dequant.u[2] = 4;
        dequant.u[3] = 32;
        self.dispatch(
            c"kernel_qwen_ple_dequant",
            &self.scratch.ple_embedding,
            Some(&self.scratch.ple_packed),
            Some(&self.scratch.ple_scales),
            Some(&self.scratch.ple_biases),
            &[],
            &dequant,
            HIDDEN,
            1,
        )?;
        let prefix = "language_model.model.layers.1.ple";
        self.bf16_mv(
            self.weight(&format!("{prefix}.key_proj.weight"))?,
            &self.scratch.ple_embedding,
            &self.scratch.ple_key,
            HIDDEN,
            HC_WIDTH,
        )?;
        self.bf16_mv(
            self.weight(&format!("{prefix}.value_proj.weight"))?,
            &self.scratch.ple_embedding,
            &self.scratch.ple_value,
            HIDDEN,
            HIDDEN,
        )?;
        for (input, output, name) in [
            (&self.scratch.ple_key, &self.scratch.ple_key, "norm_key"),
            (&self.scratch.hc, &self.scratch.hc_norm, "norm_query"),
        ] {
            let mut norm = args();
            norm.u[0] = HC_WIDTH;
            norm.u[1] = HIDDEN;
            norm.f[0] = 1.0e-6;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                output,
                Some(input),
                None,
                None,
                &[self.view(self.weight(&format!("{prefix}.{name}.weight"))?)],
                &norm,
                HC,
                1,
            )?;
        }
        let mut gate = args();
        gate.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_ple_gate",
            &self.scratch.ple_gated,
            Some(&self.scratch.ple_key),
            Some(&self.scratch.hc_norm),
            Some(&self.scratch.ple_value),
            &[],
            &gate,
            HC,
            1,
        )?;
        let mut norm = args();
        norm.u[0] = HC_WIDTH;
        norm.u[1] = HIDDEN;
        norm.f[0] = 1.0e-6;
        self.dispatch(
            c"kernel_qwen_zero_rms",
            &self.scratch.ple_norm,
            Some(&self.scratch.ple_gated),
            None,
            None,
            &[self.view(self.weight(&format!("{prefix}.norm_conv.weight"))?)],
            &norm,
            HC,
            1,
        )?;
        let mut conv = args();
        conv.u[0] = HC_WIDTH;
        self.dispatch(
            c"kernel_qwen_ple_conv",
            &self.scratch.ple_output,
            Some(&self.scratch.ple_gated),
            Some(&self.scratch.ple_norm),
            Some(&self.ple_state.conv),
            &[self.view(self.weight(&format!("{prefix}.conv_weight"))?)],
            &conv,
            HC_WIDTH,
            1,
        )?;
        let mut add = args();
        add.u[0] = HC_WIDTH;
        self.dispatch(
            c"kernel_qwen_add",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            Some(&self.scratch.ple_output),
            None,
            &[],
            &add,
            HC_WIDTH,
            1,
        )?;
        self.scratch.hc.copy_from(
            0,
            &self.scratch.hc_norm,
            0,
            u64::from(HC_WIDTH) * 4,
            "committing Qwen PLE injection",
        )?;
        commands.finish()?;
        self.ple_state.history = advance_ple_history(self.ple_state.history, token);
        Ok(())
    }

    fn encode_layer(&mut self, layer: usize) -> Result<(), String> {
        let prefix = format!("language_model.model.layers.{layer}");
        let commands = Commands::begin()?;
        self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?;
        match &self.states[layer] {
            LayerState::Gdn { .. } => self.gdn(&prefix, layer)?,
            LayerState::Attention { .. } => self.attention(&prefix, layer)?,
        }
        self.hyper_write()?;
        self.hyper_read(&format!("{prefix}.mlp_hyper_connection"))?;
        self.affine_mv_into(
            &self.affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)?,
            &self.scratch.block,
            &self.scratch.router,
            HIDDEN,
            EXPERTS,
        )?;
        let mut route_args = args();
        route_args.u[0] = EXPERTS;
        self.dispatch(
            c"kernel_qwen_route_top10",
            &self.scratch.route_ids,
            Some(&self.scratch.route_weights),
            Some(&self.scratch.router),
            None,
            &[],
            &route_args,
            1,
            1,
        )?;
        commands.finish()?;

        let mut ids = [0_i32; EXPERTS_USED];
        let mut weights = [0.0_f32; EXPERTS_USED];
        self.scratch.route_ids.read_i32(&mut ids)?;
        self.scratch.route_weights.read_f32(&mut weights)?;
        self.scratch.moe.fill(0.0, HIDDEN.into())?;
        let commands = Commands::begin()?;
        for (&expert, &weight) in ids.iter().zip(&weights) {
            if !(0..EXPERTS as i32).contains(&expert) || !weight.is_finite() || weight < 0.0 {
                let mut router = vec![0.0; EXPERTS as usize];
                self.scratch.router.read_f32(&mut router)?;
                let non_finite = router.iter().filter(|value| !value.is_finite()).count();
                let mut block = vec![0.0; HIDDEN as usize];
                self.scratch.block.read_f32(&mut block)?;
                let block_non_finite = block.iter().filter(|value| !value.is_finite()).count();
                return Err(format!(
                    "Qwen layer {layer} router produced invalid expert {expert} with weight {weight} ({non_finite} non-finite logits, {block_non_finite} non-finite inputs)"
                ));
            }
            self.expert(&prefix, expert as u32, weight)?;
        }
        self.shared_expert(&prefix)?;
        let mut inject_args = args();
        inject_args.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_hyper_inject",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            Some(&self.scratch.moe),
            Some(&self.scratch.injection),
            &[],
            &inject_args,
            HC_WIDTH,
            1,
        )?;
        self.scratch.hc.copy_from(
            0,
            &self.scratch.hc_norm,
            0,
            u64::from(HC_WIDTH) * 4,
            "committing Qwen MoE hyper streams",
        )?;
        commands.finish()
    }

    fn hyper_read(&self, prefix: &str) -> Result<(), String> {
        let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?;
        let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?;
        let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?;
        let inject = self.weight(&format!("{prefix}.block_inject_weight.weight"))?;
        let mut rms = args();
        rms.u[0] = HC_WIDTH;
        rms.u[1] = HIDDEN;
        rms.f[0] = 1.0e-6;
        self.dispatch(
            c"kernel_qwen_zero_rms",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            None,
            None,
            &[self.view(norm)],
            &rms,
            HC,
            1,
        )?;
        self.bf16_mv(
            down,
            &self.scratch.hc_norm,
            &self.scratch.rank,
            HC_WIDTH,
            HC_RANK,
        )?;
        let mut unary = args();
        unary.u[0] = HC_RANK;
        self.dispatch(
            c"kernel_qwen_silu_div4",
            &self.scratch.rank,
            Some(&self.scratch.rank),
            None,
            None,
            &[],
            &unary,
            HC_RANK,
            1,
        )?;
        self.bf16_mv(
            up,
            &self.scratch.rank,
            &self.scratch.hc_mix,
            HC_RANK,
            HC_WIDTH,
        )?;
        unary.u[0] = HC_WIDTH;
        self.dispatch(
            c"kernel_qwen_sigmoid",
            &self.scratch.hc_mix,
            Some(&self.scratch.hc_mix),
            None,
            None,
            &[],
            &unary,
            HC_WIDTH,
            1,
        )?;
        let mut mix = args();
        mix.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_hyper_mix",
            &self.scratch.block,
            Some(&self.scratch.hc_norm),
            Some(&self.scratch.hc_mix),
            None,
            &[],
            &mix,
            HIDDEN,
            1,
        )?;
        self.bf16_mv(
            inject,
            &self.scratch.hc_norm,
            &self.scratch.injection,
            HC_WIDTH,
            HC,
        )?;
        unary.u[0] = HC;
        self.dispatch(
            c"kernel_qwen_sigmoid2_div4",
            &self.scratch.injection,
            Some(&self.scratch.injection),
            None,
            None,
            &[],
            &unary,
            HC,
            1,
        )
    }

    fn hyper_write(&self) -> Result<(), String> {
        let mut values = args();
        values.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_hyper_inject",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            Some(&self.scratch.hidden),
            Some(&self.scratch.injection),
            &[],
            &values,
            HC_WIDTH,
            1,
        )?;
        self.scratch.hc.copy_from(
            0,
            &self.scratch.hc_norm,
            0,
            u64::from(HC_WIDTH) * 4,
            "committing Qwen attention hyper streams",
        )
    }

    fn gdn(&self, prefix: &str, layer: usize) -> Result<(), String> {
        let LayerState::Gdn { conv, recurrent } = &self.states[layer] else {
            return Err("Qwen GDN graph received attention state".into());
        };
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.linear_attn.in_proj_qkv"),
                HIDDEN,
                GDN_QKV,
                None,
            )?,
            &self.scratch.block,
            &self.scratch.qkv,
            HIDDEN,
            GDN_QKV,
        )?;
        for (name, offset, width) in [
            ("in_proj_z", 0_u64, GDN_VALUE),
            ("in_proj_b", u64::from(GDN_VALUE), GDN_HEADS_V),
            ("in_proj_a", u64::from(GDN_VALUE + GDN_HEADS_V), GDN_HEADS_V),
        ] {
            let target = self
                .scratch
                .controls
                .view(offset * 4, u64::from(width) * 4)?;
            self.affine_mv_into(
                &self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?,
                &self.scratch.block,
                &target,
                HIDDEN,
                width,
            )?;
        }
        let mut conv_args = args();
        conv_args.u[0] = GDN_QKV;
        self.dispatch(
            c"kernel_qwen_conv_silu",
            &self.scratch.qkv,
            Some(&self.scratch.qkv),
            Some(conv),
            None,
            &[self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?)],
            &conv_args,
            GDN_QKV,
            1,
        )?;
        let mut step = args();
        step.u[0] = HEAD_DIM;
        step.u[1] = GDN_HEADS_K;
        step.u[2] = GDN_HEADS_V;
        step.u[3] = GDN_VALUE;
        step.u[4] = GDN_VALUE + GDN_HEADS_V;
        step.f[0] = 1.0e-6;
        self.dispatch(
            c"kernel_qwen_gdn_step",
            &self.scratch.gdn_raw,
            Some(&self.scratch.gdn_out),
            Some(&self.scratch.controls),
            Some(recurrent),
            &[
                self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?),
                self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?),
            ],
            &step,
            HEAD_DIM,
            GDN_HEADS_V,
        )?;
        let mut gate = args();
        gate.u[0] = HEAD_DIM;
        gate.u[1] = GDN_HEADS_V;
        gate.f[0] = 1.0e-6;
        self.dispatch(
            c"kernel_qwen_gdn_norm_gate",
            &self.scratch.gdn_out,
            Some(&self.scratch.gdn_raw),
            Some(&self.scratch.controls),
            None,
            &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)],
            &gate,
            GDN_HEADS_V,
            1,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.linear_attn.out_proj"),
                GDN_VALUE,
                HIDDEN,
                None,
            )?,
            &self.scratch.gdn_out,
            &self.scratch.hidden,
            GDN_VALUE,
            HIDDEN,
        )
    }

    fn attention(&self, prefix: &str, layer: usize) -> Result<(), String> {
        if self.position + 1 > DENSE_BUDGET {
            return Err(
                "Qwen sparse QSA selection is required beyond 2048 tokens; native QSA belongs to issue #97"
                    .into(),
            );
        }
        let LayerState::Attention { kv } = &self.states[layer] else {
            return Err("Qwen attention graph received GDN state".into());
        };
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.self_attn.q_proj"),
                HIDDEN,
                ATTN_WIDTH * 2,
                None,
            )?,
            &self.scratch.block,
            &self.scratch.q_packed,
            HIDDEN,
            ATTN_WIDTH * 2,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.self_attn.k_proj"),
                HIDDEN,
                ATTN_KV_WIDTH,
                None,
            )?,
            &self.scratch.block,
            &self.scratch.k,
            HIDDEN,
            ATTN_KV_WIDTH,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.self_attn.v_proj"),
                HIDDEN,
                ATTN_KV_WIDTH,
                None,
            )?,
            &self.scratch.block,
            &self.scratch.v,
            HIDDEN,
            ATTN_KV_WIDTH,
        )?;
        let mut split = args();
        split.u[0] = ATTN_HEADS;
        split.u[1] = ATTN_DIM;
        self.dispatch(
            c"kernel_qwen_split_q_gate",
            &self.scratch.q,
            Some(&self.scratch.q_packed),
            Some(&self.scratch.q_gate),
            None,
            &[],
            &split,
            ATTN_WIDTH,
            1,
        )?;
        self.head_norm_rope(
            &self.scratch.q,
            &self.scratch.attention,
            self.weight(&format!("{prefix}.self_attn.q_norm.weight"))?,
            ATTN_HEADS,
        )?;
        self.head_norm_rope(
            &self.scratch.k,
            &self.scratch.k_rope,
            self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?,
            ATTN_KV_HEADS,
        )?;
        let mut store = args();
        store.u[0] = ATTN_KV_WIDTH;
        store.u[1] = self.position;
        self.dispatch(
            c"kernel_qwen_store_kv_bf16",
            kv,
            Some(&self.scratch.k_rope),
            Some(&self.scratch.v),
            None,
            &[],
            &store,
            ATTN_KV_WIDTH,
            1,
        )?;
        let mut dense = args();
        dense.u[0] = ATTN_HEADS;
        dense.u[1] = ATTN_KV_HEADS;
        dense.u[2] = ATTN_DIM;
        dense.u[3] = self.position + 1;
        self.dispatch(
            c"kernel_qwen_dense_attention",
            &self.scratch.q,
            Some(&self.scratch.attention),
            Some(kv),
            None,
            &[],
            &dense,
            ATTN_HEADS,
            1,
        )?;
        let mut gate = args();
        gate.u[0] = ATTN_WIDTH;
        self.dispatch(
            c"kernel_qwen_gate_attention",
            &self.scratch.attention,
            Some(&self.scratch.q),
            Some(&self.scratch.q_gate),
            None,
            &[],
            &gate,
            ATTN_WIDTH,
            1,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.self_attn.o_proj"),
                ATTN_WIDTH,
                HIDDEN,
                None,
            )?,
            &self.scratch.attention,
            &self.scratch.hidden,
            ATTN_WIDTH,
            HIDDEN,
        )
    }

    fn head_norm_rope(
        &self,
        input: &Buffer,
        output: &Buffer,
        weight: Weight<'_>,
        heads: u32,
    ) -> Result<(), String> {
        let mut values = args();
        values.u[0] = ATTN_DIM;
        values.u[1] = 64;
        values.u[2] = heads;
        values.u[3] = self.position;
        values.f[0] = 1.0e-6;
        values.f[1] = 10_000_000.0;
        self.dispatch(
            c"kernel_qwen_head_norm_rope",
            output,
            Some(input),
            None,
            None,
            &[self.view(weight)],
            &values,
            ATTN_DIM,
            heads,
        )
    }

    fn expert(&self, prefix: &str, expert: u32, weight: f32) -> Result<(), String> {
        for (name, target) in [
            ("gate_proj", &self.scratch.gate),
            ("up_proj", &self.scratch.up),
        ] {
            self.affine_mv_into(
                &self.affine(
                    &format!("{prefix}.mlp.switch_mlp.{name}"),
                    HIDDEN,
                    EXPERT_WIDTH,
                    Some(expert),
                )?,
                &self.scratch.block,
                target,
                HIDDEN,
                EXPERT_WIDTH,
            )?;
        }
        let mut swiglu = args();
        swiglu.u[0] = EXPERT_WIDTH;
        self.dispatch(
            c"kernel_qwen_swiglu",
            &self.scratch.mid,
            Some(&self.scratch.gate),
            Some(&self.scratch.up),
            None,
            &[],
            &swiglu,
            EXPERT_WIDTH,
            1,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.mlp.switch_mlp.down_proj"),
                EXPERT_WIDTH,
                HIDDEN,
                Some(expert),
            )?,
            &self.scratch.mid,
            &self.scratch.expert,
            EXPERT_WIDTH,
            HIDDEN,
        )?;
        let mut accumulate = args();
        accumulate.u[0] = HIDDEN;
        accumulate.f[0] = weight;
        self.dispatch(
            c"kernel_qwen_accumulate",
            &self.scratch.moe,
            Some(&self.scratch.expert),
            None,
            None,
            &[],
            &accumulate,
            HIDDEN,
            1,
        )
    }

    fn shared_expert(&self, prefix: &str) -> Result<(), String> {
        for (name, target) in [
            ("gate_proj", &self.scratch.gate),
            ("up_proj", &self.scratch.up),
        ] {
            self.affine_mv_into(
                &self.affine(
                    &format!("{prefix}.mlp.shared_expert.{name}"),
                    HIDDEN,
                    EXPERT_WIDTH,
                    None,
                )?,
                &self.scratch.block,
                target,
                HIDDEN,
                EXPERT_WIDTH,
            )?;
        }
        let mut swiglu = args();
        swiglu.u[0] = EXPERT_WIDTH;
        self.dispatch(
            c"kernel_qwen_swiglu",
            &self.scratch.mid,
            Some(&self.scratch.gate),
            Some(&self.scratch.up),
            None,
            &[],
            &swiglu,
            EXPERT_WIDTH,
            1,
        )?;
        self.affine_mv_into(
            &self.affine(
                &format!("{prefix}.mlp.shared_expert.down_proj"),
                EXPERT_WIDTH,
                HIDDEN,
                None,
            )?,
            &self.scratch.mid,
            &self.scratch.shared,
            EXPERT_WIDTH,
            HIDDEN,
        )?;
        self.affine_mv_into(
            &self.affine(&format!("{prefix}.mlp.shared_expert_gate"), HIDDEN, 1, None)?,
            &self.scratch.block,
            &self.scratch.gate,
            HIDDEN,
            1,
        )?;
        let mut accumulate = args();
        accumulate.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_accumulate_sigmoid_scalar",
            &self.scratch.moe,
            Some(&self.scratch.shared),
            Some(&self.scratch.gate),
            None,
            &[],
            &accumulate,
            HIDDEN,
            1,
        )
    }

    fn final_output(&mut self) -> Result<(), String> {
        let commands = Commands::begin()?;
        self.final_mix()?;
        self.affine_mv_into(
            &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?,
            &self.scratch.block,
            &self.scratch.logits,
            HIDDEN,
            VOCAB,
        )?;
        commands.finish()?;
        self.scratch.logits.read_f32(&mut self.logits)
    }

    fn final_mix(&self) -> Result<(), String> {
        let prefix = "language_model.model.hyper_connection_mixer";
        let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?;
        let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?;
        let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?;
        let mut rms = args();
        rms.u[0] = HC_WIDTH;
        rms.u[1] = HIDDEN;
        rms.f[0] = 1.0e-6;
        self.dispatch(
            c"kernel_qwen_zero_rms",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            None,
            None,
            &[self.view(norm)],
            &rms,
            HC,
            1,
        )?;
        self.bf16_mv(
            down,
            &self.scratch.hc_norm,
            &self.scratch.rank,
            HC_WIDTH,
            HC_RANK,
        )?;
        let mut unary = args();
        unary.u[0] = HC_RANK;
        self.dispatch(
            c"kernel_qwen_silu_div4",
            &self.scratch.rank,
            Some(&self.scratch.rank),
            None,
            None,
            &[],
            &unary,
            HC_RANK,
            1,
        )?;
        self.bf16_mv(
            up,
            &self.scratch.rank,
            &self.scratch.hc_mix,
            HC_RANK,
            HC_WIDTH,
        )?;
        unary.u[0] = HC_WIDTH;
        self.dispatch(
            c"kernel_qwen_sigmoid",
            &self.scratch.hc_mix,
            Some(&self.scratch.hc_mix),
            None,
            None,
            &[],
            &unary,
            HC_WIDTH,
            1,
        )?;
        let mut mix = args();
        mix.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_hyper_mix",
            &self.scratch.block,
            Some(&self.scratch.hc_norm),
            Some(&self.scratch.hc_mix),
            None,
            &[],
            &mix,
            HIDDEN,
            1,
        )
    }

    fn affine(
        &self,
        prefix: &str,
        in_dim: u32,
        out_dim: u32,
        expert: Option<u32>,
    ) -> Result<Affine<'_>, String> {
        let packed = self.weight(&format!("{prefix}.weight"))?;
        let scales = self.weight(&format!("{prefix}.scales"))?;
        let biases = self.weight(&format!("{prefix}.biases"))?;
        let bits = packed
            .tensor
            .quant_bits
            .ok_or_else(|| format!("{prefix} is not affine quantized"))?;
        let group = packed
            .tensor
            .group_size
            .ok_or_else(|| format!("{prefix} has no affine group size"))?
            as u32;
        if !matches!(bits, 2 | 4 | 8) || !in_dim.is_multiple_of(group) {
            return Err(format!("{prefix} has an incompatible affine layout"));
        }
        let expected = if expert.is_some() {
            vec![
                EXPERTS as u64,
                out_dim as u64,
                (in_dim / (32 / bits)) as u64,
            ]
        } else {
            vec![out_dim as u64, (in_dim / (32 / bits)) as u64]
        };
        if packed.tensor.shape != expected {
            return Err(format!("{prefix} has an incompatible executor shape"));
        }
        let expected_parameters = if expert.is_some() {
            vec![EXPERTS as u64, out_dim as u64, (in_dim / group) as u64]
        } else {
            vec![out_dim as u64, (in_dim / group) as u64]
        };
        if packed.tensor.dtype != "U32"
            || scales.tensor.dtype != "BF16"
            || biases.tensor.dtype != "BF16"
            || scales.tensor.shape != expected_parameters
            || biases.tensor.shape != expected_parameters
        {
            return Err(format!("{prefix} has incompatible affine parameters"));
        }
        Ok(Affine {
            packed: Weight {
                tensor: packed.tensor,
                expert,
            },
            scales: Weight {
                tensor: scales.tensor,
                expert,
            },
            biases: Weight {
                tensor: biases.tensor,
                expert,
            },
            bits,
            group,
        })
    }

    fn weight(&self, name: &str) -> Result<Weight<'_>, String> {
        self.model.tensor(name).map(|tensor| Weight {
            tensor,
            expert: None,
        })
    }

    fn view(&self, weight: Weight<'_>) -> QwenWeightView {
        let (map, _) = self.model.map(weight.tensor.map);
        let experts = weight.expert.map_or(1, |_| EXPERTS as u64);
        let bytes = (weight.tensor.range.end - weight.tensor.range.start) / experts;
        let offset = weight.tensor.range.start + u64::from(weight.expert.unwrap_or(0)) * bytes;
        QwenWeightView {
            map: map.as_ptr().cast(),
            size: map.len() as u64,
            offset,
            bytes,
        }
    }

    fn affine_mv_into(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<(), String> {
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[2] = weight.bits;
        values.u[3] = weight.group;
        self.dispatch(
            c"kernel_qwen_affine_mv",
            output,
            Some(input),
            None,
            None,
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim,
            1,
        )
    }

    fn bf16_mv(
        &self,
        weight: Weight<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<(), String> {
        if weight.tensor.dtype != "BF16" || weight.tensor.shape != [out_dim as u64, in_dim as u64] {
            return Err(format!(
                "{} has an incompatible BF16 matrix shape",
                weight.tensor.name
            ));
        }
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        self.dispatch(
            c"kernel_qwen_bf16_mv",
            output,
            Some(input),
            None,
            None,
            &[self.view(weight)],
            &values,
            out_dim,
            1,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn dispatch(
        &self,
        kernel: &CStr,
        out: &Buffer,
        a: Option<&Buffer>,
        b: Option<&Buffer>,
        c: Option<&Buffer>,
        weights: &[QwenWeightView],
        values: &QwenKernelArgs,
        grid_x: u32,
        grid_y: u32,
    ) -> Result<(), String> {
        dispatch_qwen(kernel, out, a, b, c, weights, values, grid_x, grid_y)
    }

    pub(super) fn prefill(
        &mut self,
        tokens: &[i32],
        mut progress: impl FnMut(u32) -> bool,
    ) -> Result<usize, String> {
        let mut completed = 0;
        for &token in tokens {
            if !progress(completed) {
                break;
            }
            self.eval(token)?;
            completed += 1;
        }
        progress(completed);
        Ok(completed as usize)
    }

    pub(super) fn logits(&self) -> &[f32] {
        &self.logits
    }
    pub(super) fn execution_stats(&self) -> ExecutionStats {
        ExecutionStats::default()
    }
    pub(super) fn model(&self) -> &QwenModel {
        &self.model
    }
    pub(super) fn context(&self) -> u32 {
        self.context
    }
    pub(super) fn position(&self) -> u32 {
        self.position
    }
    pub(super) fn tokens(&self) -> &[i32] {
        &self.tokens
    }
    pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
        self.checkpoint_tag
    }
    pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) {
        self.checkpoint_tag = tag;
    }

    pub(super) fn reset(&mut self) -> Result<(), String> {
        self.states = allocate_states(self.context)?;
        self.ple_state = allocate_ple_state()?;
        self.logits.fill(0.0);
        self.tokens.clear();
        self.position = 0;
        self.checkpoint_tag = [0; 32];
        Ok(())
    }

    pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
        if !tokens.starts_with(&self.tokens) {
            self.reset()?;
        }
        Ok(self.tokens.len())
    }

    fn blank_resident(&self) -> Result<QwenResidentState, String> {
        Ok(QwenResidentState {
            states: allocate_states(self.context)?,
            ple_state: allocate_ple_state()?,
            logits: vec![0.0; VOCAB as usize],
            tokens: Vec::new(),
            position: 0,
            checkpoint_tag: [0; 32],
        })
    }

    pub(super) fn swap_resident_state(
        &mut self,
        state: &mut Option<QwenResidentState>,
    ) -> Result<(), String> {
        let mut incoming = state.take().map_or_else(|| self.blank_resident(), Ok)?;
        std::mem::swap(&mut self.states, &mut incoming.states);
        std::mem::swap(&mut self.ple_state, &mut incoming.ple_state);
        std::mem::swap(&mut self.logits, &mut incoming.logits);
        std::mem::swap(&mut self.tokens, &mut incoming.tokens);
        std::mem::swap(&mut self.position, &mut incoming.position);
        std::mem::swap(&mut self.checkpoint_tag, &mut incoming.checkpoint_tag);
        *state = Some(incoming);
        Ok(())
    }

    pub(in crate::engine) fn save_checkpoint(
        &mut self,
        path: &Path,
        tag: [u8; 32],
        progress: &mut impl FnMut(u64),
    ) -> Result<(), String> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(|error| error.to_string())?;
        }
        let temporary = path.with_extension("tmp");
        let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
        file.write_all(CHECKPOINT_MAGIC)
            .map_err(|error| error.to_string())?;
        for value in [
            CHECKPOINT_VERSION,
            self.context,
            self.position,
            VOCAB,
            LAYERS as u32,
        ] {
            write_u32(&mut file, value)?;
        }
        file.write_all(&self.model.checkpoint_identity())
            .map_err(|error| error.to_string())?;
        file.write_all(&tag).map_err(|error| error.to_string())?;
        for token in self.ple_state.history {
            write_u32(&mut file, token as u32)?;
        }
        for &token in &self.tokens {
            write_u32(&mut file, token as u32)?;
        }
        for &logit in &self.logits {
            write_u32(&mut file, logit.to_bits())?;
        }
        let mut chunk = vec![0; CHECKPOINT_CHUNK];
        write_buffer(
            &mut file,
            &self.ple_state.conv,
            0,
            u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
            &mut chunk,
            progress,
        )?;
        for state in &self.states {
            match state {
                LayerState::Gdn { conv, recurrent } => {
                    write_buffer(
                        &mut file,
                        conv,
                        0,
                        u64::from(GDN_QKV) * 3 * 2,
                        &mut chunk,
                        progress,
                    )?;
                    write_buffer(
                        &mut file,
                        recurrent,
                        0,
                        u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                        &mut chunk,
                        progress,
                    )?;
                }
                LayerState::Attention { kv } => {
                    write_buffer(
                        &mut file,
                        kv,
                        0,
                        u64::from(self.position) * ATTN_KV_WIDTH as u64 * 4,
                        &mut chunk,
                        progress,
                    )?;
                }
            }
        }
        file.sync_all().map_err(|error| error.to_string())?;
        fs::rename(temporary, path).map_err(|error| error.to_string())?;
        self.checkpoint_tag = tag;
        Ok(())
    }

    pub(in crate::engine) fn load_checkpoint(
        &mut self,
        path: &Path,
        progress: &mut impl FnMut(u64),
    ) -> Result<bool, String> {
        let mut file = match File::open(path) {
            Ok(file) => file,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
            Err(error) => return Err(error.to_string()),
        };
        let mut magic = [0; 8];
        file.read_exact(&mut magic)
            .map_err(|error| error.to_string())?;
        if &magic != CHECKPOINT_MAGIC {
            return Err("Qwen checkpoint has an invalid signature".into());
        }
        for expected in [CHECKPOINT_VERSION, self.context] {
            if read_u32(&mut file)? != expected {
                return Err("Qwen checkpoint does not match the current executor".into());
            }
        }
        let position = read_u32(&mut file)?;
        if position > self.context
            || read_u32(&mut file)? != VOCAB
            || read_u32(&mut file)? != LAYERS as u32
        {
            return Err("Qwen checkpoint shape is invalid".into());
        }
        let mut identity = [0; 32];
        file.read_exact(&mut identity)
            .map_err(|error| error.to_string())?;
        if identity != self.model.checkpoint_identity() {
            return Err("Qwen checkpoint model identity changed".into());
        }
        let mut tag = [0; 32];
        file.read_exact(&mut tag)
            .map_err(|error| error.to_string())?;
        let mut ple_history = [0; PLE_HISTORY];
        for token in &mut ple_history {
            let value = read_u32(&mut file)?;
            if value >= VOCAB {
                return Err("Qwen checkpoint PLE history is invalid".into());
            }
            *token = value as i32;
        }
        let mut tokens = Vec::with_capacity(position as usize);
        for _ in 0..position {
            let token = read_u32(&mut file)?;
            if token >= VOCAB {
                return Err("Qwen checkpoint token is invalid".into());
            }
            tokens.push(token as i32);
        }
        let mut logits = Vec::with_capacity(VOCAB as usize);
        for _ in 0..VOCAB {
            logits.push(f32::from_bits(read_u32(&mut file)?));
        }
        self.reset()?;
        let mut chunk = vec![0; CHECKPOINT_CHUNK];
        read_buffer(
            &mut file,
            &self.ple_state.conv,
            0,
            u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
            &mut chunk,
            progress,
        )?;
        for state in &self.states {
            match state {
                LayerState::Gdn { conv, recurrent } => {
                    read_buffer(
                        &mut file,
                        conv,
                        0,
                        u64::from(GDN_QKV) * 3 * 2,
                        &mut chunk,
                        progress,
                    )?;
                    read_buffer(
                        &mut file,
                        recurrent,
                        0,
                        u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                        &mut chunk,
                        progress,
                    )?;
                }
                LayerState::Attention { kv } => read_buffer(
                    &mut file,
                    kv,
                    0,
                    u64::from(position) * ATTN_KV_WIDTH as u64 * 4,
                    &mut chunk,
                    progress,
                )?,
            }
        }
        let mut trailing = [0];
        if file
            .read(&mut trailing)
            .map_err(|error| error.to_string())?
            != 0
        {
            self.reset()?;
            return Err("Qwen checkpoint has trailing data".into());
        }
        self.position = position;
        self.tokens = tokens;
        self.logits = logits;
        self.ple_state.history = ple_history;
        self.checkpoint_tag = tag;
        Ok(true)
    }
}

fn allocate_states(context: u32) -> Result<Vec<LayerState>, String> {
    (0..LAYERS)
        .map(|layer| {
            if layer % 4 == 3 {
                Ok(LayerState::Attention {
                    kv: Buffer::bytes(u64::from(context) * ATTN_KV_WIDTH as u64 * 4)?,
                })
            } else {
                let conv = Buffer::bytes(u64::from(GDN_QKV) * 3 * 2)?;
                let recurrent =
                    Buffer::floats(u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64)?;
                conv.fill(0.0, u64::from(GDN_QKV) * 3 / 2)?;
                recurrent.fill(
                    0.0,
                    u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64,
                )?;
                Ok(LayerState::Gdn { conv, recurrent })
            }
        })
        .collect()
}

fn allocate_ple_state() -> Result<PleState, String> {
    let conv = Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?;
    conv.fill(0.0, u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 / 2)?;
    Ok(PleState {
        history: [EOS_TOKEN; PLE_HISTORY],
        conv,
    })
}

fn ple_contract(model: &QwenModel) -> Result<PleContract, String> {
    let multipliers = read_i64_array::<3>(
        model,
        "language_model.model.layers.1.ple.ple_embedding.layer_multipliers",
    )?;
    let sizes = read_i64_array::<PLE_HEADS>(
        model,
        "language_model.model.layers.1.ple.ple_embedding.ngram_heads_vocab_sizes",
    )?;
    let offsets = read_i64_array::<PLE_HEADS>(
        model,
        "language_model.model.layers.1.ple.ple_embedding.ngram_heads_offsets",
    )?;
    let expected_multipliers = official_ple_multipliers();
    let mut expected_sizes = [0; PLE_HEADS];
    let mut expected_offsets = [0; PLE_HEADS];
    let mut total = 0_i64;
    let mut prime = 19_999_999_i64;
    for head in 0..PLE_HEADS {
        prime = next_prime(prime);
        expected_sizes[head] = prime;
        expected_offsets[head] = total;
        total += prime;
    }
    if multipliers != expected_multipliers || sizes != expected_sizes || offsets != expected_offsets
    {
        return Err("Qwen PLE hash parameters do not match the official contract".into());
    }
    for (name, shape, bits, group) in [
        ("ngram.weight", [320_001_536, 20], Some(4), Some(32)),
        ("ngram.scales", [320_001_536, 5], Some(4), Some(32)),
        ("ngram.biases", [320_001_536, 5], Some(4), Some(32)),
    ] {
        let tensor = model.tensor(name)?;
        if tensor.shape != shape || tensor.quant_bits != bits || tensor.group_size != group {
            return Err(format!("{name} does not match the Qwen PLE row layout"));
        }
    }
    Ok(PleContract {
        multipliers,
        sizes,
        offsets,
    })
}

fn read_i64_array<const N: usize>(model: &QwenModel, name: &str) -> Result<[i64; N], String> {
    let tensor = model.tensor(name)?;
    if tensor.dtype != "I64" || tensor.shape != [N as u64] {
        return Err(format!("{name} does not match the Qwen PLE integer layout"));
    }
    let bytes = model.tensor_bytes(tensor)?;
    if bytes.len() != N * 8 {
        return Err(format!("{name} has an invalid byte length"));
    }
    Ok(std::array::from_fn(|index| {
        i64::from_le_bytes(bytes[index * 8..index * 8 + 8].try_into().unwrap())
    }))
}

fn official_ple_multipliers() -> [i64; 3] {
    const GAMMA: u64 = 0x9e37_79b9_7f4a_7c15;
    const M1: u64 = 0xbf58_476d_1ce4_e5b9;
    const M2: u64 = 0x94d0_49bb_1331_11eb;
    let bound = (i64::MAX / VOCAB as i64 / 2) as u64;
    std::array::from_fn(|index| {
        let mut value = 1234_u64.wrapping_add(GAMMA.wrapping_mul(index as u64 + 1));
        value = value.wrapping_add(GAMMA);
        value = (value ^ (value >> 30)).wrapping_mul(M1);
        value = (value ^ (value >> 27)).wrapping_mul(M2);
        value ^= value >> 31;
        (2 * (value % bound) + 1) as i64
    })
}

fn next_prime(mut value: i64) -> i64 {
    loop {
        value += 1;
        if value % 2 != 0
            && (3..=((value as f64).sqrt() as i64))
                .step_by(2)
                .all(|divisor| value % divisor != 0)
        {
            return value;
        }
    }
}

fn ple_rows(
    contract: &PleContract,
    history: [i32; PLE_HISTORY],
    token: i32,
) -> Result<[u64; PLE_HEADS], String> {
    if token < 0 || token as u32 >= VOCAB {
        return Err(format!("token {token} is outside the Qwen vocabulary"));
    }
    let shifted = [token as i64, history[1] as i64, history[0] as i64];
    let two = shifted[0].wrapping_mul(contract.multipliers[0])
        ^ shifted[1].wrapping_mul(contract.multipliers[1]);
    let three = two ^ shifted[2].wrapping_mul(contract.multipliers[2]);
    Ok(std::array::from_fn(|head| {
        let mixed = if head < 8 { two } else { three };
        (contract.offsets[head] + mixed.rem_euclid(contract.sizes[head])) as u64
    }))
}

fn advance_ple_history(history: [i32; PLE_HISTORY], token: i32) -> [i32; PLE_HISTORY] {
    if token == EOS_TOKEN {
        [EOS_TOKEN; PLE_HISTORY]
    } else {
        [history[1], token]
    }
}

fn gather_ple_row(
    row: u64,
    packed: &[u8],
    scales: &[u8],
    biases: &[u8],
) -> Result<[u8; PLE_ROW_BYTES], String> {
    let row = usize::try_from(row).map_err(|_| "Qwen PLE row exceeds this platform")?;
    let mut value = [0; PLE_ROW_BYTES];
    value[..80].copy_from_slice(ple_row_slice(packed, row, 80)?);
    value[80..90].copy_from_slice(ple_row_slice(scales, row, 10)?);
    value[90..].copy_from_slice(ple_row_slice(biases, row, 10)?);
    Ok(value)
}

fn ple_row_slice(data: &[u8], row: usize, width: usize) -> Result<&[u8], String> {
    let start = row
        .checked_mul(width)
        .ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?;
    data.get(start..start + width)
        .ok_or_else(|| format!("Qwen PLE row {row} is truncated"))
}

fn args() -> QwenKernelArgs {
    QwenKernelArgs::default()
}

#[allow(clippy::too_many_arguments)]
fn dispatch_qwen(
    kernel: &CStr,
    out: &Buffer,
    a: Option<&Buffer>,
    b: Option<&Buffer>,
    c: Option<&Buffer>,
    weights: &[QwenWeightView],
    values: &QwenKernelArgs,
    grid_x: u32,
    grid_y: u32,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_qwen_dispatch(
                kernel.as_ptr(),
                out.raw(),
                a.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()),
                b.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()),
                c.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()),
                weights.as_ptr(),
                weights.len() as u32,
                values,
                grid_x,
                grid_y,
            )
        },
        kernel.to_str().unwrap_or("running a Qwen Metal kernel"),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use memmap2::MmapOptions;
    use sha2::{Digest, Sha256};
    use std::fs;
    use std::path::PathBuf;

    fn bf16(value: f32) -> u16 {
        let bits = value.to_bits();
        ((bits + 0x7fff + ((bits >> 16) & 1)) >> 16) as u16
    }

    fn close(actual: f32, expected: f32) {
        assert!(
            (actual - expected).abs() <= 2.0e-4 * expected.abs().max(1.0),
            "{actual} != {expected}"
        );
    }

    #[test]
    fn qwen_ple_hash_contract_matches_golden_boundaries() {
        assert_eq!(
            official_ple_multipliers(),
            [23_703_573_157_769, 20_109_073_645_365, 8_052_911_324_071]
        );
        let mut contract = PleContract {
            multipliers: official_ple_multipliers(),
            sizes: [0; PLE_HEADS],
            offsets: [0; PLE_HEADS],
        };
        let mut prime = 19_999_999;
        let mut offset = 0;
        for head in 0..PLE_HEADS {
            prime = next_prime(prime);
            contract.sizes[head] = prime;
            contract.offsets[head] = offset;
            offset += prime;
        }
        let initial = ple_rows(&contract, [EOS_TOKEN; 2], 1).unwrap();
        let repeated = ple_rows(&contract, [1, 1], 1).unwrap();
        let boundary = ple_rows(&contract, [EOS_TOKEN, 42], 43).unwrap();
        assert_eq!(
            initial,
            [
                16_121_432,
                28_938_500,
                59_087_997,
                73_487_090,
                81_148_277,
                104_500_129,
                120_276_032,
                149_373_875,
                176_283_436,
                184_305_849,
                216_528_839,
                231_080_079,
                257_961_536,
                266_068_568,
                289_043_455,
                305_959_965,
            ]
        );
        assert_eq!(
            repeated,
            [
                6_868_091,
                38_325_817,
                54_054_700,
                68_075_137,
                82_949_816,
                101_241_419,
                138_678_867,
                155_262_032,
                176_541_251,
                196_154_476,
                215_703_237,
                234_413_824,
                254_220_543,
                274_027_268,
                293_962_951,
                313_640_732,
            ]
        );
        assert_eq!(
            boundary,
            [
                18_529_343,
                23_547_650,
                56_056_978,
                73_570_159,
                88_581_601,
                113_585_506,
                121_091_299,
                151_099_148,
                175_585_266,
                184_439_538,
                216_587_431,
                222_082_137,
                250_284_866,
                278_847_169,
                281_781_121,
                317_050_322,
            ]
        );
        assert_eq!(contract.offsets[0], 0);
        assert_eq!(contract.offsets[8], 160_000_374);
        assert_eq!(contract.offsets[15] + contract.sizes[15], 320_001_446);
        for (head, &row) in initial.iter().enumerate() {
            assert!(
                (contract.offsets[head]..contract.offsets[head] + contract.sizes[head])
                    .contains(&(row as i64))
            );
        }
        assert_eq!(320_001_536 % 128, 0);
        let history = advance_ple_history([EOS_TOKEN; 2], 1);
        assert_eq!(advance_ple_history(history, 2), [1, 2]);
        assert_eq!(advance_ple_history([1, 2], EOS_TOKEN), [EOS_TOKEN; 2]);

        let hash_chunks = |chunk: usize| {
            let mut history = [EOS_TOKEN; PLE_HISTORY];
            [1, 2, EOS_TOKEN, 3, 4]
                .chunks(chunk)
                .flat_map(|tokens| {
                    tokens
                        .iter()
                        .map(|&token| {
                            let rows = ple_rows(&contract, history, token).unwrap();
                            history = advance_ple_history(history, token);
                            rows
                        })
                        .collect::<Vec<_>>()
                })
                .collect::<Vec<_>>()
        };
        assert_eq!(hash_chunks(1), hash_chunks(2));
        assert_eq!(hash_chunks(1), hash_chunks(5));
    }

    #[test]
    fn qwen_ple_gather_copies_only_the_requested_row() {
        let packed = (0..240).map(|value| value as u8).collect::<Vec<_>>();
        let scales = (0..30).map(|value| (value + 17) as u8).collect::<Vec<_>>();
        let biases = (0..30).map(|value| (value + 47) as u8).collect::<Vec<_>>();
        let direct = gather_ple_row(1, &packed, &scales, &biases).unwrap();
        assert_eq!(&direct[..80], &packed[80..160]);
        assert_eq!(&direct[80..90], &scales[10..20]);
        assert_eq!(&direct[90..], &biases[10..20]);
        assert_eq!(
            gather_ple_row(1, &packed, &scales, &biases).unwrap(),
            direct
        );
        assert!(
            gather_ple_row(3, &packed, &scales, &biases)
                .unwrap_err()
                .contains("truncated")
        );
    }

    #[test]
    #[ignore = "requires Apple Metal"]
    fn qwen_metal_primitives_match_reference_vectors() {
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();

        let input = Buffer::floats(64).unwrap();
        input.write_f32(&vec![1.0; 64]).unwrap();
        let output = Buffer::floats(64).unwrap();
        let path = std::env::temp_dir().join(format!(
            "ds4-qwen96-weights-{}-{}",
            std::process::id(),
            std::thread::current().name().unwrap_or("test")
        ));
        // Safetensors headers are not guaranteed to align the data region.
        let mut bytes = vec![0];
        for _ in 0..8 {
            bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes());
        }
        for value in [
            0.5, -1.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0,
        ] {
            bytes.extend_from_slice(&bf16(value).to_le_bytes());
        }
        fs::write(&path, bytes).unwrap();
        let file = File::open(&path).unwrap();
        // SAFETY: this test owns the read-only file for the lifetime of the mapping.
        let map = unsafe { MmapOptions::new().map(&file).unwrap() };
        let view = |offset, size| QwenWeightView {
            map: map.as_ptr().cast(),
            size: map.len() as u64,
            offset,
            bytes: size,
        };

        let mut affine = args();
        affine.u[0] = 64;
        affine.u[1] = 1;
        affine.u[2] = 4;
        affine.u[3] = 64;
        dispatch_qwen(
            c"kernel_qwen_affine_mv",
            &output,
            Some(&input),
            None,
            None,
            &[view(1, 32), view(33, 2), view(35, 2)],
            &affine,
            1,
            1,
        )
        .unwrap();
        let mut scalar = [0.0];
        output.read_f32(&mut scalar).unwrap();
        close(scalar[0], 32.0);

        let conv_input = Buffer::floats(1).unwrap();
        conv_input.write_f32(&[2.0]).unwrap();
        let conv_state = Buffer::bytes(6).unwrap();
        conv_state.write(0, &[0; 6]).unwrap();
        let conv_output = Buffer::floats(1).unwrap();
        let mut conv = args();
        conv.u[0] = 1;
        dispatch_qwen(
            c"kernel_qwen_conv_silu",
            &conv_output,
            Some(&conv_input),
            Some(&conv_state),
            None,
            &[view(41, 8)],
            &conv,
            1,
            1,
        )
        .unwrap();
        conv_output.read_f32(&mut scalar).unwrap();
        close(scalar[0], 8.0 / (1.0 + (-8.0_f32).exp()));
        assert_eq!(
            {
                let mut state = [0; 6];
                conv_state.read(0, &mut state).unwrap();
                state
            },
            [0, 0, 0, 0, 0, 64]
        );

        let qkv = Buffer::floats(6).unwrap();
        qkv.write_f32(&[3.0, 4.0, 0.0, 2.0, 5.0, 7.0]).unwrap();
        let controls = Buffer::floats(4).unwrap();
        controls.write_f32(&[0.0; 4]).unwrap();
        let recurrent = Buffer::floats(4).unwrap();
        recurrent.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap();
        let raw = Buffer::floats(2).unwrap();
        let gated = Buffer::floats(2).unwrap();
        let mut step = args();
        step.u[0] = 2;
        step.u[1] = 1;
        step.u[2] = 1;
        step.u[3] = 2;
        step.u[4] = 3;
        step.f[0] = 1.0e-6;
        dispatch_qwen(
            c"kernel_qwen_gdn_step",
            &raw,
            Some(&qkv),
            Some(&controls),
            Some(&recurrent),
            &[view(37, 2), view(39, 2)],
            &step,
            2,
            1,
        )
        .unwrap();
        let mut raw_values = [0.0; 2];
        raw.read_f32(&mut raw_values).unwrap();
        close(raw_values[0], 2.0506096);
        close(raw_values[1], 2.9698484);
        let mut norm = args();
        norm.u[0] = 2;
        norm.u[1] = 1;
        norm.f[0] = 1.0e-6;
        dispatch_qwen(
            c"kernel_qwen_gdn_norm_gate",
            &gated,
            Some(&raw),
            Some(&controls),
            None,
            &[view(49, 4)],
            &norm,
            1,
            1,
        )
        .unwrap();
        let mut gated_values = [0.0; 2];
        gated.read_f32(&mut gated_values).unwrap();
        let rms = ((raw_values[0].powi(2) + raw_values[1].powi(2)) * 0.5 + 1.0e-6).sqrt();
        close(gated_values[0], raw_values[0] / rms * 0.5);
        close(gated_values[1], raw_values[1] / rms * 0.5);

        let router = Buffer::floats(EXPERTS.into()).unwrap();
        let mut logits = vec![-100.0; EXPERTS as usize];
        for (index, value) in logits.iter_mut().take(EXPERTS_USED).enumerate() {
            *value = index as f32;
        }
        router.write_f32(&logits).unwrap();
        let ids = Buffer::bytes((EXPERTS_USED * 4) as u64).unwrap();
        let weights = Buffer::floats(EXPERTS_USED as u64).unwrap();
        let mut route = args();
        route.u[0] = EXPERTS;
        dispatch_qwen(
            c"kernel_qwen_route_top10",
            &ids,
            Some(&weights),
            Some(&router),
            None,
            &[],
            &route,
            1,
            1,
        )
        .unwrap();
        let mut selected = [0; EXPERTS_USED];
        ids.read_i32(&mut selected).unwrap();
        assert_eq!(selected, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);
        let mut probabilities = [0.0; EXPERTS_USED];
        weights.read_f32(&mut probabilities).unwrap();
        close(probabilities.iter().sum(), 1.0);

        let norm_input = Buffer::floats(4).unwrap();
        norm_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap();
        let norm_output = Buffer::floats(4).unwrap();
        let mut zero_norm = args();
        zero_norm.u[0] = 4;
        zero_norm.u[1] = 2;
        zero_norm.f[0] = 1.0e-6;
        dispatch_qwen(
            c"kernel_qwen_zero_rms",
            &norm_output,
            Some(&norm_input),
            None,
            None,
            &[view(53, 8)],
            &zero_norm,
            2,
            1,
        )
        .unwrap();
        let mut normalized = [0.0; 4];
        norm_output.read_f32(&mut normalized).unwrap();
        for (actual, expected) in normalized.into_iter().zip([
            1.0 / (2.5_f32 + 1.0e-6).sqrt(),
            2.0 / (2.5_f32 + 1.0e-6).sqrt(),
            3.0 / (12.5_f32 + 1.0e-6).sqrt(),
            4.0 / (12.5_f32 + 1.0e-6).sqrt(),
        ]) {
            close(actual, expected);
        }

        let hyper_input = Buffer::floats(8).unwrap();
        hyper_input
            .write_f32(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
            .unwrap();
        let hyper_weights = Buffer::floats(8).unwrap();
        hyper_weights.write_f32(&[1.0; 8]).unwrap();
        let mixed = Buffer::floats(2).unwrap();
        let mut hyper = args();
        hyper.u[0] = 2;
        dispatch_qwen(
            c"kernel_qwen_hyper_mix",
            &mixed,
            Some(&hyper_input),
            Some(&hyper_weights),
            None,
            &[],
            &hyper,
            2,
            1,
        )
        .unwrap();
        let mut mixed_values = [0.0; 2];
        mixed.read_f32(&mut mixed_values).unwrap();
        assert_eq!(mixed_values, [4.0, 5.0]);
        let injection = Buffer::floats(4).unwrap();
        injection.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap();
        let combined = Buffer::floats(8).unwrap();
        dispatch_qwen(
            c"kernel_qwen_hyper_inject",
            &combined,
            Some(&hyper_input),
            Some(&mixed),
            Some(&injection),
            &[],
            &hyper,
            8,
            1,
        )
        .unwrap();
        let mut combined_values = [0.0; 8];
        combined.read_f32(&mut combined_values).unwrap();
        assert_eq!(
            combined_values,
            [5.0, 7.0, 11.0, 14.0, 17.0, 21.0, 23.0, 28.0]
        );

        let expert_gate = Buffer::floats(2).unwrap();
        expert_gate.write_f32(&[0.0, 2.0]).unwrap();
        let expert_up = Buffer::floats(2).unwrap();
        expert_up.write_f32(&[4.0, 3.0]).unwrap();
        let expert = Buffer::floats(2).unwrap();
        let mut moe = args();
        moe.u[0] = 2;
        dispatch_qwen(
            c"kernel_qwen_swiglu",
            &expert,
            Some(&expert_gate),
            Some(&expert_up),
            None,
            &[],
            &moe,
            2,
            1,
        )
        .unwrap();
        let accumulated = Buffer::floats(2).unwrap();
        accumulated.fill(0.0, 2).unwrap();
        moe.f[0] = 0.25;
        dispatch_qwen(
            c"kernel_qwen_accumulate",
            &accumulated,
            Some(&expert),
            None,
            None,
            &[],
            &moe,
            2,
            1,
        )
        .unwrap();
        let shared = Buffer::floats(2).unwrap();
        shared.write_f32(&[2.0, 4.0]).unwrap();
        let shared_gate = Buffer::floats(1).unwrap();
        shared_gate.write_f32(&[0.0]).unwrap();
        dispatch_qwen(
            c"kernel_qwen_accumulate_sigmoid_scalar",
            &accumulated,
            Some(&shared),
            Some(&shared_gate),
            None,
            &[],
            &moe,
            2,
            1,
        )
        .unwrap();
        let mut moe_values = [0.0; 2];
        accumulated.read_f32(&mut moe_values).unwrap();
        close(moe_values[0], 1.0);
        close(
            moe_values[1],
            0.25 * (2.0 / (1.0 + (-2.0_f32).exp()) * 3.0) + 2.0,
        );

        let rope_input = Buffer::floats(4).unwrap();
        rope_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap();
        let rope_output = Buffer::floats(4).unwrap();
        let mut rope = args();
        rope.u[0] = 4;
        rope.u[1] = 4;
        rope.u[2] = 1;
        rope.u[3] = 1;
        rope.f[0] = 1.0e-6;
        rope.f[1] = 10_000.0;
        dispatch_qwen(
            c"kernel_qwen_head_norm_rope",
            &rope_output,
            Some(&rope_input),
            None,
            None,
            &[view(53, 8)],
            &rope,
            4,
            1,
        )
        .unwrap();
        let rms = (7.5_f32 + 1.0e-6).sqrt();
        let theta = [1.0_f32, 0.01];
        let normalized = [1.0 / rms, 2.0 / rms, 3.0 / rms, 4.0 / rms];
        let expected_rope = [
            normalized[0] * theta[0].cos() - normalized[2] * theta[0].sin(),
            normalized[1] * theta[1].cos() - normalized[3] * theta[1].sin(),
            normalized[2] * theta[0].cos() + normalized[0] * theta[0].sin(),
            normalized[3] * theta[1].cos() + normalized[1] * theta[1].sin(),
        ];
        let mut actual_rope = [0.0; 4];
        rope_output.read_f32(&mut actual_rope).unwrap();
        for (actual, expected) in actual_rope.into_iter().zip(expected_rope) {
            close(actual, expected);
        }

        let cache = Buffer::bytes(16).unwrap();
        let key = Buffer::floats(2).unwrap();
        let value = Buffer::floats(2).unwrap();
        let mut store = args();
        store.u[0] = 2;
        for (position, (keys, values)) in [([1.0, 0.0], [2.0, 3.0]), ([0.0, 1.0], [5.0, 7.0])]
            .into_iter()
            .enumerate()
        {
            key.write_f32(&keys).unwrap();
            value.write_f32(&values).unwrap();
            store.u[1] = position as u32;
            dispatch_qwen(
                c"kernel_qwen_store_kv_bf16",
                &cache,
                Some(&key),
                Some(&value),
                None,
                &[],
                &store,
                2,
                1,
            )
            .unwrap();
        }
        let query = Buffer::floats(2).unwrap();
        query.write_f32(&[1.0, 0.0]).unwrap();
        let attention = Buffer::floats(2).unwrap();
        let mut dense = args();
        dense.u[0] = 1;
        dense.u[1] = 1;
        dense.u[2] = 2;
        dense.u[3] = 2;
        dispatch_qwen(
            c"kernel_qwen_dense_attention",
            &attention,
            Some(&query),
            Some(&cache),
            None,
            &[],
            &dense,
            1,
            1,
        )
        .unwrap();
        let first = (1.0_f32 / 2.0_f32.sqrt()).exp();
        let probability = first / (first + 1.0);
        let mut actual_attention = [0.0; 2];
        attention.read_f32(&mut actual_attention).unwrap();
        close(
            actual_attention[0],
            probability * 2.0 + (1.0 - probability) * 5.0,
        );
        close(
            actual_attention[1],
            probability * 3.0 + (1.0 - probability) * 7.0,
        );

        let ple_packed = Buffer::bytes(80).unwrap();
        ple_packed.write(0, &[0x33; 80]).unwrap();
        let ple_scales = Buffer::bytes(10).unwrap();
        ple_scales
            .write(0, &bf16(0.5).to_le_bytes().repeat(5))
            .unwrap();
        let ple_biases = Buffer::bytes(10).unwrap();
        ple_biases
            .write(0, &bf16(-1.0).to_le_bytes().repeat(5))
            .unwrap();
        let ple_embedding = Buffer::floats(160).unwrap();
        let mut dequant = args();
        dequant.u[0] = 160;
        dequant.u[1] = 1;
        dequant.u[2] = 4;
        dequant.u[3] = 32;
        dispatch_qwen(
            c"kernel_qwen_ple_dequant",
            &ple_embedding,
            Some(&ple_packed),
            Some(&ple_scales),
            Some(&ple_biases),
            &[],
            &dequant,
            160,
            1,
        )
        .unwrap();
        let mut embedding = [0.0; 160];
        ple_embedding.read_f32(&mut embedding).unwrap();
        assert!(embedding.into_iter().all(|value| value == 0.5));

        let ple_key = Buffer::floats(8).unwrap();
        ple_key.write_f32(&[1.0; 8]).unwrap();
        let ple_query = Buffer::floats(8).unwrap();
        ple_query.write_f32(&[1.0; 8]).unwrap();
        let ple_value = Buffer::floats(2).unwrap();
        ple_value.write_f32(&[2.0, 3.0]).unwrap();
        let ple_gated = Buffer::floats(8).unwrap();
        let mut gate = args();
        gate.u[0] = 2;
        dispatch_qwen(
            c"kernel_qwen_ple_gate",
            &ple_gated,
            Some(&ple_key),
            Some(&ple_query),
            Some(&ple_value),
            &[],
            &gate,
            4,
            1,
        )
        .unwrap();
        let expected_gate = 1.0 / (1.0 + (-2.0_f32.sqrt().sqrt()).exp());
        let mut gated = [0.0; 8];
        ple_gated.read_f32(&mut gated).unwrap();
        for stream in 0..4 {
            close(gated[stream * 2], 2.0 * expected_gate);
            close(gated[stream * 2 + 1], 3.0 * expected_gate);
        }

        let ple_normalized = Buffer::floats(1).unwrap();
        ple_normalized.write_f32(&[2.0]).unwrap();
        let ple_gate_value = Buffer::floats(1).unwrap();
        ple_gate_value.write_f32(&[0.5]).unwrap();
        let ple_state = Buffer::bytes(18).unwrap();
        ple_state.write(0, &[0; 18]).unwrap();
        let ple_output = Buffer::floats(1).unwrap();
        let mut conv = args();
        conv.u[0] = 1;
        dispatch_qwen(
            c"kernel_qwen_ple_conv",
            &ple_output,
            Some(&ple_gate_value),
            Some(&ple_normalized),
            Some(&ple_state),
            &[view(41, 8)],
            &conv,
            1,
            1,
        )
        .unwrap();
        ple_output.read_f32(&mut scalar).unwrap();
        close(scalar[0], 0.5 + 8.0 / (1.0 + (-8.0_f32).exp()));
        let mut ple_history = [0; 18];
        ple_state.read(0, &mut ple_history).unwrap();
        assert_eq!(&ple_history[16..], &bf16(2.0).to_le_bytes());

        drop(map);
        drop(file);
        fs::remove_file(path).unwrap();
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_core_boundary_and_checkpoint_are_stable() {
        configure_sources().unwrap();
        let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE")
            .map(PathBuf::from)
            .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory");
        let model = QwenModel::open(&root, 4).unwrap();
        let residency_before = model.mapped_residency().unwrap();
        let mut executor = QwenExecutor::open(model, 4).unwrap();
        executor.eval(1).unwrap();
        let residency_after = executor.model().mapped_residency().unwrap();
        assert!(residency_after.0 <= executor.model().memory().resident_core);
        assert!(residency_after.1 <= executor.model().memory().mapped_ple);
        eprintln!(
            "Qwen mapped residency core/PLE before {:?}, after {:?}",
            residency_before, residency_after
        );
        assert!(executor.logits.iter().all(|value| value.is_finite()));
        let mut digest = Sha256::new();
        for value in &executor.logits {
            digest.update(value.to_bits().to_le_bytes());
        }
        let digest: [u8; 32] = digest.finalize().into();
        assert_eq!(
            digest,
            [
                137, 244, 133, 253, 201, 214, 196, 144, 249, 130, 28, 63, 124, 75, 32, 40, 16, 148,
                148, 123, 5, 50, 90, 165, 101, 44, 223, 164, 62, 54, 164, 86,
            ]
        );
        let reference = [
            executor.logits[0],
            executor.logits[1],
            executor.logits[1000],
            executor.logits[VOCAB as usize - 1],
        ];
        let next = executor
            .logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0 as i32;
        for (actual, expected) in
            reference
                .into_iter()
                .zip([6.406_557, 2.082_818_3, -3.058_045_6, -0.121_010_3])
        {
            close(actual, expected);
        }
        assert_eq!(next, 89_648);

        let checkpoint = std::env::temp_dir().join(format!(
            "ds4-qwen96-checkpoint-{}-{}",
            std::process::id(),
            std::thread::current().name().unwrap_or("test")
        ));
        executor
            .save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
            .unwrap();

        executor.eval(next).unwrap();
        let continued = [
            executor.logits[0],
            executor.logits[1],
            executor.logits[1000],
            executor.logits[VOCAB as usize - 1],
        ];
        let LayerState::Gdn { conv, recurrent } = &executor.states[0] else {
            unreachable!()
        };
        let mut continued_conv = vec![0; GDN_QKV as usize * 3 * 2];
        conv.read(0, &mut continued_conv).unwrap();
        let mut continued_recurrent = vec![0; 4 * 1024];
        recurrent.read(0, &mut continued_recurrent).unwrap();
        let mut continued_ple = vec![0; HC_WIDTH as usize * PLE_CONV_STATE as usize * 2];
        executor.ple_state.conv.read(0, &mut continued_ple).unwrap();

        assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        assert_eq!(executor.position, 1);
        assert_eq!(executor.tokens, [1]);
        assert_eq!(executor.ple_state.history, [EOS_TOKEN, 1]);
        assert_eq!(executor.checkpoint_tag, [7; 32]);
        assert_eq!(
            [
                executor.logits[0],
                executor.logits[1],
                executor.logits[1000],
                executor.logits[VOCAB as usize - 1],
            ],
            reference
        );
        executor.eval(next).unwrap();
        for (actual, expected) in [
            executor.logits[0],
            executor.logits[1],
            executor.logits[1000],
            executor.logits[VOCAB as usize - 1],
        ]
        .into_iter()
        .zip(continued)
        {
            close(actual, expected);
        }
        let LayerState::Gdn { conv, recurrent } = &executor.states[0] else {
            unreachable!()
        };
        let mut resumed_conv = vec![0; continued_conv.len()];
        conv.read(0, &mut resumed_conv).unwrap();
        assert_eq!(resumed_conv, continued_conv);
        let mut resumed_recurrent = vec![0; continued_recurrent.len()];
        recurrent.read(0, &mut resumed_recurrent).unwrap();
        assert_eq!(resumed_recurrent, continued_recurrent);
        let mut resumed_ple = vec![0; continued_ple.len()];
        executor.ple_state.conv.read(0, &mut resumed_ple).unwrap();
        assert_eq!(resumed_ple, continued_ple);

        executor.reset().unwrap();
        assert_eq!(executor.position, 0);
        assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);
        let mut reset_ple = vec![1; continued_ple.len()];
        executor.ple_state.conv.read(0, &mut reset_ple).unwrap();
        assert!(reset_ple.into_iter().all(|byte| byte == 0));
        executor.eval(1).unwrap();
        for (actual, expected) in [
            executor.logits[0],
            executor.logits[1],
            executor.logits[1000],
            executor.logits[VOCAB as usize - 1],
        ]
        .into_iter()
        .zip(reference)
        {
            close(actual, expected);
        }

        executor.reset().unwrap();
        executor.eval(EOS_TOKEN).unwrap();
        assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);

        executor.position = DENSE_BUDGET;
        executor.context = DENSE_BUDGET + 1;
        let error = executor.eval(1).unwrap_err();
        assert!(error.contains("issue #97"));
        fs::remove_file(checkpoint).unwrap();
    }
}