use super::checkpoint::{read_buffer, read_u32, write_buffer, write_u32};
use super::*;
use crate::engine::qwen::{QwenModel, QwenTensor};
use crate::settings::QWEN_DEFAULT_PREFILL_CHUNK;
use rustc_hash::FxHashMap as HashMap;
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 HYPER_FOLDED_ROWS: u32 = HC_RANK + HC;
const HYPER_QUANT_GROUP: u32 = 64;
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 GDN_FUSED_INPUT: u32 = GDN_QKV + GDN_CONTROLS;
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 QSA_HEADS: u32 = 4;
const QSA_KV_HEADS: u32 = 1;
const QSA_DIM: u32 = 128;
const QSA_WIDTH: u32 = (QSA_HEADS + QSA_KV_HEADS) * QSA_DIM;
const QSA_FUSED_QKV: u32 = ATTN_WIDTH * 2 + ATTN_KV_WIDTH * 2;
const QSA_FUSED_INPUT: u32 = QSA_FUSED_QKV + QSA_WIDTH;
const QSA_RATIO: u32 = 4;
const QSA_TOP_K: u32 = 512;
const MTP_TOKEN_RESERVE: u32 = 3;
const MTP_CAPTURE_ROWS: usize = MTP_TOKEN_RESERVE as usize + 1;
const VERIFY_ROWS: u64 = MTP_CAPTURE_ROWS as u64;
const EXPERTS: u32 = 512;
const EXPERTS_USED: usize = 10;
const EXPERT_WIDTH: u32 = 640;
const VOCAB: u32 = 248_320;
const SAMPLING_TOP_K: u32 = 32;
const DECODE_DENSE_BUDGET: u32 = 2_048;
const PREFILL_DENSE_BUDGET: u32 = 32_768;
const QSA_VERIFY_DENSE_LIMIT: u32 = 16_384;
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 = 5;
const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024;
const QMM_SPLIT_TARGET_TILES: u32 = 512;
const QMM_SPLIT_ELEMENTS: u64 = QMM_SPLIT_TARGET_TILES as u64 * 32 * 32;

fn affine_qmm_splits(in_dim: u32, out_dim: u32, rows: u32, group: u32) -> u32 {
    let alignment = group.max(32);
    let tiles = u64::from(out_dim.div_ceil(32)) * u64::from(rows.div_ceil(32));
    let mut splits = ((u64::from(QMM_SPLIT_TARGET_TILES) / tiles.max(1)) as u32)
        .min(in_dim / alignment)
        .max(1);
    while splits > 1 && !in_dim.is_multiple_of(splits * alignment) {
        splits -= 1;
    }
    splits
}

const fn dense_attention_budget(rows: u32) -> u32 {
    if rows == 1 {
        DECODE_DENSE_BUDGET
    } else if rows > VERIFY_ROWS as u32 {
        PREFILL_DENSE_BUDGET
    } else {
        DECODE_DENSE_BUDGET
    }
}

const fn packed_dense_qsa(rows: u32, tokens: u32) -> bool {
    rows <= VERIFY_ROWS as u32 && tokens < QSA_VERIFY_DENSE_LIMIT
}

type TokenDistribution = Vec<(usize, f64)>;
type TokenDistributions = Vec<TokenDistribution>;
type DraftStep = (Vec<f32>, Option<TokenDistribution>);

#[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,
        qsa_raw: Buffer,
        qsa_pooled: Buffer,
    },
}

struct Scratch {
    rows: u32,
    token_ids: Buffer,
    hidden: Buffer,
    hc: Buffer,
    batch_final_hc: Buffer,
    hc_norm: Buffer,
    hc_mix: Buffer,
    rank: Buffer,
    block: Buffer,
    injection: Buffer,
    hyper_rms: Buffer,
    qkv: Buffer,
    controls: Buffer,
    gdn_fused: Buffer,
    gdn_fused_qkv: Buffer,
    gdn_fused_controls: Buffer,
    gdn_z: Buffer,
    gdn_b: Buffer,
    gdn_a: Buffer,
    gdn_q: Buffer,
    gdn_k: Buffer,
    gdn_out: Buffer,
    router: Buffer,
    qmm_partials: Buffer,
    route_ids: Buffer,
    route_weights: Buffer,
    route_counts: Buffer,
    route_map: Buffer,
    route_work: Buffer,
    gate: Buffer,
    up: Buffer,
    mid: Buffer,
    expert: Buffer,
    moe: Buffer,
    shared: Buffer,
    q_packed: Buffer,
    qsa_fused: Buffer,
    qsa_fused_q: Buffer,
    qsa_fused_k: Buffer,
    qsa_fused_v: Buffer,
    qsa_fused_indexer: Buffer,
    q: Buffer,
    q_gate: Buffer,
    k: Buffer,
    k_rope: Buffer,
    v: Buffer,
    attention: Buffer,
    qsa_qk: Buffer,
    qsa_q: Buffer,
    qsa_scores: Buffer,
    qsa_selected: Buffer,
    causal_mask: Buffer,
    dense_query: Buffer,
    dense_keys: Buffer,
    dense_values: 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,
    sampling_selected: Buffer,
    sampling_probabilities: Buffer,
}

impl Scratch {
    fn new(context: u32, capacity: u32) -> Result<Self, String> {
        let qsa_blocks = qsa_block_capacity(context)?;
        let capacity = capacity.max(VERIFY_ROWS as u32);
        let row_count = u64::from(capacity);
        let route_rows = row_count * EXPERTS_USED as u64;
        let route_work_items = route_rows.div_ceil(64) + u64::from(EXPERTS);
        let route_work_bytes = (8 + u64::from(EXPERTS) * 4 + route_rows * 4 + 7) & !7;
        let rows = |width: u32| u64::from(width) * row_count;
        let gdn_fused = Buffer::floats(rows(GDN_FUSED_INPUT))?;
        let gdn_fused_qkv = gdn_fused.view(0, u64::from(GDN_QKV) * 4)?;
        let gdn_fused_controls =
            gdn_fused.view(u64::from(GDN_QKV) * 4, u64::from(GDN_CONTROLS) * 4)?;
        let qsa_fused = Buffer::floats(QSA_FUSED_INPUT.into())?;
        let qsa_fused_q = qsa_fused.view(0, u64::from(ATTN_WIDTH * 2) * 4)?;
        let qsa_fused_k =
            qsa_fused.view(u64::from(ATTN_WIDTH * 2) * 4, u64::from(ATTN_KV_WIDTH) * 4)?;
        let qsa_fused_v = qsa_fused.view(
            u64::from(ATTN_WIDTH * 2 + ATTN_KV_WIDTH) * 4,
            u64::from(ATTN_KV_WIDTH) * 4,
        )?;
        let qsa_fused_indexer = qsa_fused.view(
            u64::from(ATTN_WIDTH * 2 + ATTN_KV_WIDTH * 2) * 4,
            u64::from(QSA_WIDTH) * 4,
        )?;
        Ok(Self {
            rows: capacity,
            token_ids: Buffer::bytes(row_count * 4)?,
            hidden: Buffer::floats(rows(HIDDEN))?,
            hc: Buffer::floats(rows(HC_WIDTH))?,
            batch_final_hc: Buffer::floats(HC_WIDTH.into())?,
            hc_norm: Buffer::floats(rows(HC_WIDTH))?,
            hc_mix: Buffer::floats(rows(HC_WIDTH))?,
            rank: Buffer::floats(rows(HC_RANK))?,
            block: Buffer::floats(rows(HIDDEN))?,
            injection: Buffer::floats(rows(HC))?,
            hyper_rms: Buffer::floats(rows(HC))?,
            qkv: Buffer::floats(rows(GDN_QKV))?,
            controls: Buffer::floats(rows(GDN_CONTROLS))?,
            gdn_fused,
            gdn_fused_qkv,
            gdn_fused_controls,
            gdn_z: Buffer::floats(rows(GDN_VALUE))?,
            gdn_b: Buffer::floats(rows(GDN_HEADS_V))?,
            gdn_a: Buffer::floats(rows(GDN_HEADS_V))?,
            gdn_q: Buffer::floats(rows(GDN_HEADS_K * HEAD_DIM))?,
            gdn_k: Buffer::floats(rows(GDN_HEADS_K * HEAD_DIM))?,
            gdn_out: Buffer::bytes(rows(GDN_VALUE) * 4 + u64::from(GDN_QKV) * 3 * 2)?,
            router: Buffer::floats(rows(EXPERTS))?,
            qmm_partials: Buffer::floats(QMM_SPLIT_ELEMENTS)?,
            route_ids: Buffer::bytes(row_count * EXPERTS_USED as u64 * 4)?,
            route_weights: Buffer::floats(row_count * EXPERTS_USED as u64)?,
            route_counts: Buffer::bytes(u64::from(EXPERTS) * 4)?,
            route_map: Buffer::bytes(u64::from(EXPERTS) * row_count * 4)?,
            route_work: Buffer::bytes(route_work_bytes + route_work_items * 8)?,
            gate: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?,
            up: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?,
            mid: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?,
            expert: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(HIDDEN))?,
            moe: Buffer::floats(rows(HIDDEN))?,
            shared: Buffer::floats(rows(HIDDEN))?,
            q_packed: Buffer::floats(rows(ATTN_WIDTH * 2))?,
            qsa_fused,
            qsa_fused_q,
            qsa_fused_k,
            qsa_fused_v,
            qsa_fused_indexer,
            q: Buffer::floats(rows(ATTN_WIDTH))?,
            q_gate: Buffer::floats(rows(ATTN_WIDTH))?,
            k: Buffer::floats(rows(ATTN_KV_WIDTH))?,
            k_rope: Buffer::floats(rows(ATTN_KV_WIDTH))?,
            v: Buffer::floats(rows(ATTN_KV_WIDTH))?,
            attention: Buffer::floats(rows(ATTN_WIDTH))?,
            qsa_qk: Buffer::floats(rows(QSA_WIDTH))?,
            qsa_q: Buffer::floats(rows(QSA_HEADS * QSA_DIM))?,
            qsa_scores: Buffer::floats(u64::from(qsa_blocks) * row_count)?,
            qsa_selected: Buffer::bytes(u64::from(QSA_TOP_K) * row_count * 4)?,
            causal_mask: Buffer::bytes(u64::from(PREFILL_DENSE_BUDGET) * row_count)?,
            dense_query: Buffer::bytes(rows(ATTN_WIDTH) * 2)?,
            dense_keys: Buffer::bytes(
                u64::from(PREFILL_DENSE_BUDGET) * u64::from(ATTN_KV_WIDTH) * 2,
            )?,
            dense_values: Buffer::bytes(
                u64::from(PREFILL_DENSE_BUDGET) * u64::from(ATTN_KV_WIDTH) * 2,
            )?,
            ple_packed: Buffer::bytes(row_count * PLE_HEADS as u64 * 80)?,
            ple_scales: Buffer::bytes(row_count * PLE_HEADS as u64 * 10)?,
            ple_biases: Buffer::bytes(row_count * PLE_HEADS as u64 * 10)?,
            ple_embedding: Buffer::floats(rows(HIDDEN))?,
            ple_key: Buffer::floats(rows(HC_WIDTH))?,
            ple_value: Buffer::floats(rows(HIDDEN))?,
            ple_gated: Buffer::floats(rows(HC_WIDTH))?,
            ple_norm: Buffer::floats(rows(HC_WIDTH))?,
            ple_output: Buffer::floats(rows(HC_WIDTH))?,
            logits: Buffer::floats(u64::from(VOCAB) * row_count.min(VERIFY_ROWS))?,
            sampling_selected: Buffer::bytes(row_count * u64::from(SAMPLING_TOP_K) * 4)?,
            sampling_probabilities: Buffer::floats(row_count * u64::from(SAMPLING_TOP_K))?,
        })
    }
}

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

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

#[cfg(test)]
struct RecurrentSnapshot {
    gdn: Vec<Option<(Buffer, Buffer)>>,
    ple_conv: Buffer,
    ple_history: [i32; PLE_HISTORY],
    hidden: Buffer,
    logits: Vec<f32>,
    position: u32,
    tokens: usize,
}

struct VerifyCapture {
    gdn: Vec<Option<GdnVerifyCapture>>,
    ple_conv: Buffer,
    ple_input: Buffer,
    hidden: Buffer,
}

struct GdnVerifyCapture {
    conv: Buffer,
    recurrent: Buffer,
    qkv: Buffer,
    q: Buffer,
    k: Buffer,
    v: Buffer,
    controls: Buffer,
}

struct QwenMtp {
    attention: LayerState,
    #[cfg(test)]
    snapshot: RecurrentSnapshot,
    verify: VerifyCapture,
    timing: bool,
    cycles: u64,
    drafted: u64,
    accepted: u64,
    accepted_depth_total: u64,
    accepted_depth_max: u64,
    verifier_passes: u64,
    verifier_ns: u64,
    repair_ns: u64,
}

struct QwenDraft {
    token: i32,
    #[cfg(test)]
    logits: Vec<f32>,
    probabilities: Vec<(usize, f64)>,
}

struct HyperPack {
    r1: Buffer,
    r2: Buffer,
}

pub(in crate::engine) struct QwenExecutor {
    model: QwenModel,
    weights: HashMap<String, Buffer>,
    weight_views: HashMap<usize, QwenWeightView>,
    states: Vec<LayerState>,
    ple_contract: PleContract,
    ple_state: PleState,
    scratch: Scratch,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    position: u32,
    context: u32,
    prefill_chunk: usize,
    power_percent: u8,
    prefill_layer_average: Vec<f64>,
    decode_average: f64,
    checkpoint_tag: [u8; 32],
    mtp: Option<QwenMtp>,
    pending_sampled: Option<i32>,
    folding_sampled: bool,
    hyper_packs: HashMap<String, HyperPack>,
    packed_moe: bool,
    fused_gdn: bool,
    fused_gdn_input: bool,
    fused_qsa_input: bool,
    _context: Context,
}

pub(in crate::engine) struct QwenResidentState {
    states: Vec<LayerState>,
    ple_state: PleState,
    hidden: Buffer,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    position: u32,
    checkpoint_tag: [u8; 32],
    mtp: Option<QwenMtp>,
    pending_sampled: Option<i32>,
    folding_sampled: bool,
}

impl QwenExecutor {
    #[cfg(test)]
    pub(super) fn open(model: QwenModel, context: u32) -> Result<Self, String> {
        Self::open_configured_with_hyper(
            model,
            context,
            0,
            100,
            EngineSpeculativeSettings::default(),
            true,
            true,
        )
    }

    pub(super) fn open_configured(
        model: QwenModel,
        context: u32,
        prefill_chunk: u32,
        power_percent: u8,
        speculative: EngineSpeculativeSettings,
    ) -> Result<Self, String> {
        Self::open_configured_with_hyper(
            model,
            context,
            prefill_chunk,
            power_percent,
            speculative,
            qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_HYPER"),
            true,
        )
    }

    fn open_configured_with_hyper(
        mut model: QwenModel,
        context: u32,
        prefill_chunk: u32,
        power_percent: u8,
        speculative: EngineSpeculativeSettings,
        fused_hyper: bool,
        materialize: bool,
    ) -> Result<Self, String> {
        let ple_contract = ple_contract(&model)?;
        let native = Context::open_qwen(model.memory().admission)?;
        let weights = if materialize {
            let weights = Self::materialize_weights(&model)?;
            model.release_materialized_mappings();
            weights
        } else {
            HashMap::default()
        };
        let weight_views = model
            .tensors()
            .filter_map(|tensor| {
                weights.get(&tensor.name).map(|buffer| {
                    let bytes = tensor.range.end - tensor.range.start;
                    (
                        tensor as *const QwenTensor as usize,
                        QwenWeightView {
                            tensor: buffer.raw().cast_const(),
                            map: std::ptr::null(),
                            size: bytes,
                            offset: 0,
                            bytes,
                        },
                    )
                })
            })
            .collect();
        let states = allocate_states(context)?;
        let mtp = speculative
            .glm_mtp
            .then(|| allocate_mtp(context, speculative.glm_mtp_timing))
            .transpose()?;
        let mut executor = Self {
            model,
            weights,
            weight_views,
            states,
            ple_contract,
            ple_state: allocate_ple_state()?,
            scratch: Scratch::new(context, VERIFY_ROWS as u32)?,
            logits: vec![0.0; VOCAB as usize],
            tokens: Vec::new(),
            position: 0,
            context,
            prefill_chunk: if prefill_chunk == 0 {
                QWEN_DEFAULT_PREFILL_CHUNK as usize
            } else {
                prefill_chunk as usize
            },
            power_percent: if power_percent == 0 {
                100
            } else {
                power_percent
            },
            prefill_layer_average: vec![0.0; LAYERS],
            decode_average: 0.0,
            checkpoint_tag: [0; 32],
            mtp,
            pending_sampled: None,
            folding_sampled: false,
            hyper_packs: HashMap::default(),
            packed_moe: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GATE_UP"),
            fused_gdn: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GDN"),
            fused_gdn_input: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GDN_INPROJ"),
            fused_qsa_input: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_QSA_QKV"),
            _context: native,
        };
        if fused_hyper {
            executor.prepare_hyper_packs()?;
        }
        Ok(executor)
    }

    fn materialize_weights(model: &QwenModel) -> Result<HashMap<String, Buffer>, String> {
        const COPY_CHUNK: usize = 8 * 1024 * 1024;
        let mut tensors = model
            .tensors()
            .filter(|tensor| !tensor.name.starts_with("ngram."))
            .collect::<Vec<_>>();
        tensors.sort_by(|left, right| left.name.cmp(&right.name));
        let mut weights = HashMap::with_capacity_and_hasher(tensors.len(), Default::default());
        let mut staging = vec![0; COPY_CHUNK];
        for tensor in tensors {
            let bytes = tensor.range.end - tensor.range.start;
            let target = Buffer::untracked(bytes)?;
            for offset in (0..bytes).step_by(COPY_CHUNK) {
                let count = usize::try_from((bytes - offset).min(COPY_CHUNK as u64))
                    .map_err(|_| format!("{} chunk size exceeds this platform", tensor.name))?;
                model.read_tensor_at(tensor, offset, &mut staging[..count])?;
                target.write(offset, &staging[..count])?;
                std::thread::yield_now();
            }
            weights.insert(tensor.name.clone(), target);
        }
        Self::pack_gdn_inputs(model, &mut weights)?;
        Self::pack_qsa_inputs(model, &mut weights)?;
        Ok(weights)
    }

    fn pack_gdn_inputs(
        model: &QwenModel,
        weights: &mut HashMap<String, Buffer>,
    ) -> Result<(), String> {
        for layer in (0..LAYERS).filter(|layer| layer % 4 != 3) {
            let prefix = format!("language_model.model.layers.{layer}.linear_attn");
            for suffix in ["weight", "scales", "biases"] {
                let names = ["in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a"]
                    .map(|projection| format!("{prefix}.{projection}.{suffix}"));
                let sizes = names
                    .each_ref()
                    .map(|name| {
                        model
                            .tensor(name)
                            .map(|tensor| tensor.range.end - tensor.range.start)
                    })
                    .into_iter()
                    .collect::<Result<Vec<_>, _>>()?;
                let total = sizes.iter().sum();
                let fused = Buffer::untracked(total)?;
                let commands = Commands::begin()?;
                let mut offset = 0;
                for (name, bytes) in names.iter().zip(&sizes) {
                    fused.copy_from(
                        offset,
                        weights
                            .get(name)
                            .ok_or_else(|| format!("Qwen materialization missed {name}"))?,
                        0,
                        *bytes,
                        "packing Qwen GDN input projections",
                    )?;
                    offset += bytes;
                }
                commands.finish()?;
                offset = 0;
                for (name, bytes) in names.into_iter().zip(sizes) {
                    weights.insert(name, fused.view(offset, bytes)?);
                    offset += bytes;
                }
                weights.insert(format!("{prefix}.in_proj_fused.{suffix}"), fused);
            }
        }
        Ok(())
    }

    fn pack_qsa_inputs(
        model: &QwenModel,
        weights: &mut HashMap<String, Buffer>,
    ) -> Result<(), String> {
        for layer in (0..LAYERS).filter(|layer| layer % 4 == 3) {
            let prefix = format!("language_model.model.layers.{layer}.self_attn");
            let q = model.tensor(&format!("{prefix}.q_proj.weight"))?;
            let indexer = model.tensor(&format!("{prefix}.indexer.index_qk_proj.weight"))?;
            let include_indexer = q.quant_bits == indexer.quant_bits
                && q.group_size == indexer.group_size
                && q.shape.last() == indexer.shape.last();
            for suffix in ["weight", "scales", "biases"] {
                let mut names = ["q_proj", "k_proj", "v_proj"]
                    .map(|projection| format!("{prefix}.{projection}.{suffix}"))
                    .to_vec();
                if include_indexer {
                    names.push(format!("{prefix}.indexer.index_qk_proj.{suffix}"));
                }
                let sizes = names
                    .iter()
                    .map(|name| {
                        model
                            .tensor(name)
                            .map(|tensor| tensor.range.end - tensor.range.start)
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                let total = sizes.iter().sum();
                let fused = Buffer::untracked(total)?;
                let commands = Commands::begin()?;
                let mut offset = 0;
                for (name, bytes) in names.iter().zip(&sizes) {
                    fused.copy_from(
                        offset,
                        weights
                            .get(name)
                            .ok_or_else(|| format!("Qwen materialization missed {name}"))?,
                        0,
                        *bytes,
                        "packing Qwen QSA input projections",
                    )?;
                    offset += bytes;
                }
                commands.finish()?;
                offset = 0;
                for (name, bytes) in names.into_iter().zip(sizes) {
                    weights.insert(name, fused.view(offset, bytes)?);
                    offset += bytes;
                }
                weights.insert(format!("{prefix}.qkv_fused.{suffix}"), fused);
            }
        }
        Ok(())
    }

    fn prepare_hyper_packs(&mut self) -> Result<(), String> {
        let mut prefixes = Vec::with_capacity(LAYERS * 2 + usize::from(self.mtp.is_some()) * 2);
        for layer in 0..LAYERS {
            for block in ["attn", "mlp"] {
                prefixes.push(format!(
                    "language_model.model.layers.{layer}.{block}_hyper_connection"
                ));
            }
        }
        if self.mtp.is_some() {
            prefixes.push("mtp.layers.0.attn_hyper_connection".to_owned());
            prefixes.push("mtp.layers.0.mlp_hyper_connection".to_owned());
        }

        let commands = Commands::begin()?;
        for prefix in prefixes {
            let pack = self.prepare_hyper_pack(&prefix)?;
            self.hyper_packs.insert(prefix, pack);
        }
        commands.finish()
    }

    fn prepare_hyper_pack(&self, prefix: &str) -> Result<HyperPack, String> {
        let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?;
        let inject = self.weight(&format!("{prefix}.block_inject_weight.weight"))?;
        let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?;
        let r1 = Buffer::bytes(hyper_pack_bytes(HYPER_FOLDED_ROWS, HC_WIDTH))?;
        let r2 = Buffer::bytes(hyper_pack_bytes(HC_WIDTH, HC_RANK))?;
        self.quantize_hyper_rows(down, &r1, HC_WIDTH, HC_RANK, 0, HYPER_FOLDED_ROWS)?;
        self.quantize_hyper_rows(inject, &r1, HC_WIDTH, HC, HC_RANK, HYPER_FOLDED_ROWS)?;
        self.quantize_hyper_rows(up, &r2, HC_RANK, HC_WIDTH, 0, HC_WIDTH)?;
        Ok(HyperPack { r1, r2 })
    }

    #[allow(clippy::too_many_arguments)]
    fn quantize_hyper_rows(
        &self,
        weight: Weight<'_>,
        pack: &Buffer,
        columns: u32,
        rows: u32,
        destination_row: u32,
        destination_rows: u32,
    ) -> Result<(), String> {
        if weight.tensor.dtype != "BF16" || weight.tensor.shape != [rows as u64, columns as u64] {
            return Err(format!(
                "{} has an incompatible BF16 matrix shape",
                weight.tensor.name
            ));
        }
        let mut values = args();
        values.u[0] = columns;
        values.u[1] = rows;
        values.u[2] = destination_row;
        values.u[3] = destination_rows;
        values.u[12] = 32;
        self.dispatch(
            c"kernel_qwen_quantize_hyper_q8",
            pack,
            None,
            None,
            None,
            &[self.view(weight)],
            &values,
            columns / HYPER_QUANT_GROUP,
            rows,
        )
    }

    pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
        if self.position > 0 && self.mtp.is_some() {
            let _ = self.mtp_step(token, self.position - 1, false, None)?;
        }
        self.eval_target(token)
    }

    fn eval_target(&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
            ));
        }
        let started = Instant::now();
        let commands = Commands::begin()?;
        self.begin_token(token)?;
        for layer in 0..LAYERS {
            if layer == 1 {
                self.ple(token)?;
            }
            self.encode_layer(layer)?;
        }
        self.final_output()?;
        commands.finish()?;
        self.scratch.logits.read_f32(&mut self.logits)?;
        self.tokens.push(token);
        self.position += 1;
        throttle(
            &mut self.decode_average,
            started.elapsed(),
            self.power_percent,
        );
        Ok(())
    }

    fn eval_target_batch(
        &mut self,
        tokens: &[i32],
        sampling: Option<(f32, f32, f32, i32)>,
    ) -> Result<(Vec<Vec<f32>>, Option<TokenDistributions>), String> {
        if tokens.is_empty() || tokens.len() > MTP_CAPTURE_ROWS {
            return Err("Qwen verifier width is outside the MTPLX S=2..4 contract".into());
        }
        if self.position + tokens.len() as u32 > self.context {
            return Err(format!(
                "the Qwen executor supports {} tokens per session",
                self.context
            ));
        }
        let rows = tokens.len() as u32;
        let profiled = std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some();
        let total_started = Instant::now();
        let commands = Commands::begin()?;
        self.begin_tokens(tokens)?;
        for layer in 0..LAYERS {
            let prefix = format!("language_model.model.layers.{layer}");
            if layer == 1 {
                if let Some(mtp) = &self.mtp {
                    mtp.verify.ple_input.copy_from(
                        0,
                        &self.scratch.hc,
                        0,
                        u64::from(rows * HC_WIDTH) * 4,
                        "capturing Qwen verifier PLE input rows",
                    )?;
                }
                let state_out = self
                    .mtp
                    .as_ref()
                    .ok_or("Qwen MTP is not configured")?
                    .verify
                    .ple_conv
                    .view(0, u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?;
                self.ple_batch_with_state(tokens, Some(&state_out))?;
            }
            self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?;
            match &self.states[layer] {
                LayerState::Gdn { .. } => self.gdn_rows(&prefix, layer, rows, true)?,
                state @ LayerState::Attention { .. } => {
                    self.attention_rows(&prefix, state, self.position, rows)?
                }
            }
            self.hyper_write_rows(rows)?;
            self.encode_moe_rows(&prefix, rows)?;
        }
        self.final_output_rows(rows)?;
        if let Some(mtp) = &self.mtp {
            mtp.verify.hidden.copy_from(
                0,
                &self.scratch.hc,
                0,
                u64::from(rows) * u64::from(HC_WIDTH) * 4,
                "capturing Qwen verifier hidden rows",
            )?;
        }
        let sparse = if let Some((temperature, _, _, top_k)) = sampling {
            self.encode_sparse_probabilities(&self.scratch.logits, rows, temperature, top_k)?
        } else {
            false
        };
        let encoded = total_started.elapsed();
        commands.finish()?;
        let executed = total_started.elapsed();
        let probabilities = if sparse {
            let (_, top_p, min_p, top_k) = sampling.expect("sampling configuration disappeared");
            self.read_sparse_probabilities(rows, top_p, min_p, top_k)?
        } else {
            None
        };
        let logits = if probabilities.is_some() {
            Vec::new()
        } else {
            let mut flattened = vec![0.0; tokens.len() * VOCAB as usize];
            self.scratch.logits.read_f32(&mut flattened)?;
            flattened
                .chunks_exact(VOCAB as usize)
                .map(<[f32]>::to_vec)
                .collect::<Vec<_>>()
        };
        if profiled {
            eprintln!(
                "ds4: Qwen verifier rows={rows} encode={:.3} ms gpu={:.3} ms read={:.3} ms",
                encoded.as_secs_f64() * 1_000.0,
                (executed - encoded).as_secs_f64() * 1_000.0,
                (total_started.elapsed() - executed).as_secs_f64() * 1_000.0,
            );
        }
        if let Some(last) = logits.last() {
            self.logits.copy_from_slice(last);
        }
        self.tokens.extend_from_slice(tokens);
        self.position += rows;
        throttle(
            &mut self.decode_average,
            total_started.elapsed(),
            self.power_percent,
        );
        Ok((logits, probabilities))
    }

    fn read_logits_row(&mut self, row: usize) -> Result<(), String> {
        self.logits = self.logits_row(row)?;
        Ok(())
    }

    fn logits_row(&self, row: usize) -> Result<Vec<f32>, String> {
        let mut logits = vec![0.0; VOCAB as usize];
        self.scratch
            .logits
            .view(row as u64 * u64::from(VOCAB) * 4, u64::from(VOCAB) * 4)?
            .read_f32(&mut logits)?;
        Ok(logits)
    }

    fn eval_prefill_batch(
        &mut self,
        tokens: &[i32],
        mtp_tokens: Option<&[i32]>,
    ) -> Result<(), String> {
        if tokens.is_empty() || tokens.len() > self.scratch.rows as usize {
            return Err(format!(
                "Qwen prefill batch width must be in 1..={}",
                self.scratch.rows,
            ));
        }
        if mtp_tokens.is_some_and(|next| next.len() != tokens.len()) {
            return Err("Qwen MTP prompt-history rows do not match the target rows".into());
        }
        if self.position + tokens.len() as u32 > self.context {
            return Err(format!(
                "the Qwen executor supports {} tokens per session",
                self.context
            ));
        }
        let rows = tokens.len() as u32;
        let profile = std::env::var_os("DS4SERVER_QWEN38_PREFILL_PROFILE").is_some();
        let mut commands = Commands::begin()?;
        self.begin_tokens(tokens)?;
        for layer in 0..LAYERS {
            let started = Instant::now();
            let prefix = format!("language_model.model.layers.{layer}");
            if layer == 1 {
                self.ple_batch(tokens)?;
            }
            self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?;
            match &self.states[layer] {
                LayerState::Gdn { .. } => self.gdn_rows(&prefix, layer, rows, false)?,
                state @ LayerState::Attention { .. } => {
                    self.attention_rows(&prefix, state, self.position, rows)?
                }
            }
            self.hyper_write_rows(rows)?;
            self.encode_moe_rows(&prefix, rows)?;
            if self.power_percent < 100 || profile {
                commands.finish()?;
                if profile {
                    eprintln!(
                        "DS4SERVER_QWEN38_PREFILL layer={layer} rows={rows} ms={:.3}",
                        started.elapsed().as_secs_f64() * 1_000.0
                    );
                }
                if self.power_percent < 100 {
                    throttle(
                        &mut self.prefill_layer_average[layer],
                        started.elapsed(),
                        self.power_percent,
                    );
                }
                commands = Commands::begin()?;
            }
        }
        self.mix_output_rows("language_model.model.hyper_connection_mixer", rows)?;
        let final_block = self.scratch.block.view(
            u64::from(rows - 1) * u64::from(HIDDEN) * 4,
            u64::from(HIDDEN) * 4,
        )?;
        self.affine_mv_into(
            &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?,
            &final_block,
            &self.scratch.logits,
            HIDDEN,
            VOCAB,
        )?;
        self.scratch.batch_final_hc.copy_from(
            0,
            &self.scratch.hc,
            u64::from(rows - 1) * u64::from(HC_WIDTH) * 4,
            u64::from(HC_WIDTH) * 4,
            "preserving the final Qwen prefill hidden row",
        )?;
        commands.finish()?;
        self.scratch.logits.read_f32(&mut self.logits)?;
        if let Some(next) = mtp_tokens {
            self.mtp_history_batch(next, self.position)?;
        }
        self.tokens.extend_from_slice(tokens);
        self.position += rows;
        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;
        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,
        )?;
        Ok(())
    }

    fn begin_tokens(&self, tokens: &[i32]) -> Result<(), String> {
        if tokens.is_empty() || tokens.len() > self.scratch.rows as usize {
            return Err(format!(
                "Qwen batch width is outside the allocated {}-row capacity",
                self.scratch.rows
            ));
        }
        if let Some(&token) = tokens
            .iter()
            .find(|&&token| token < 0 || token as u32 >= VOCAB)
        {
            return Err(format!("token {token} is outside the Qwen vocabulary"));
        }
        self.scratch.token_ids.write_i32(tokens)?;
        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] = tokens.len() as u32;
        self.dispatch(
            c"kernel_qwen_affine_embedding_batch",
            &self.scratch.hidden,
            Some(&self.scratch.token_ids),
            None,
            None,
            &[
                self.view(embedding.packed),
                self.view(embedding.scales),
                self.view(embedding.biases),
            ],
            &args,
            HIDDEN,
            tokens.len() as u32,
        )?;
        args.u[0] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_repeat4",
            &self.scratch.hc,
            Some(&self.scratch.hidden),
            None,
            None,
            &[],
            &args,
            HIDDEN * HC * tokens.len() as u32,
            1,
        )
    }

    fn ple(&mut self, token: i32) -> Result<(), String> {
        self.ple_batch(&[token])
    }

    fn ple_batch(&mut self, tokens: &[i32]) -> Result<(), String> {
        self.ple_batch_with_state(tokens, None)
    }

    fn ple_batch_with_state(
        &mut self,
        tokens: &[i32],
        state_out: Option<&Buffer>,
    ) -> Result<(), String> {
        let read_started = Instant::now();
        let mut history = self.ple_state.history;
        let mut requested = Vec::with_capacity(tokens.len() * PLE_HEADS);
        for &token in tokens {
            requested.extend(ple_rows(&self.ple_contract, history, token)?);
            history = advance_ple_history(history, token);
        }
        let mut packed_stage = vec![0_u8; tokens.len() * PLE_HEADS * 80];
        let mut scales_stage = vec![0_u8; tokens.len() * PLE_HEADS * 10];
        let mut biases_stage = vec![0_u8; tokens.len() * PLE_HEADS * 10];
        for (index, bytes) in self.model.read_ple_rows(&requested)?.iter().enumerate() {
            packed_stage[index * 80..][..80].copy_from_slice(&bytes[..80]);
            scales_stage[index * 10..][..10].copy_from_slice(&bytes[80..90]);
            biases_stage[index * 10..][..10].copy_from_slice(&bytes[90..]);
        }
        if std::env::var_os("DS4SERVER_QWEN38_PLE_PROFILE").is_some() {
            eprintln!(
                "DS4SERVER_QWEN38_PLE_READ rows={} us={:.3}",
                tokens.len(),
                read_started.elapsed().as_secs_f64() * 1_000_000.0
            );
        }
        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 mut dequant = args();
        dequant.u[0] = PLE_HEAD_DIM;
        dequant.u[1] = PLE_HEADS as u32 * tokens.len() 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 * tokens.len() as u32,
            1,
        )?;
        let prefix = "language_model.model.layers.1.ple";
        self.bf16_batch(
            self.weight(&format!("{prefix}.key_proj.weight"))?,
            &self.scratch.ple_embedding,
            &self.scratch.ple_key,
            HIDDEN,
            HC_WIDTH,
            tokens.len() as u32,
        )?;
        self.bf16_batch(
            self.weight(&format!("{prefix}.value_proj.weight"))?,
            &self.scratch.ple_embedding,
            &self.scratch.ple_value,
            HIDDEN,
            HIDDEN,
            tokens.len() as u32,
        )?;
        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 = rms_args(HC_WIDTH, HIDDEN);
            norm.u[4] = tokens.len() as u32;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                output,
                Some(input),
                None,
                None,
                &[self.view(self.weight(&format!("{prefix}.{name}.weight"))?)],
                &norm,
                HC * tokens.len() as u32,
                1,
            )?;
        }
        let mut gate = args();
        gate.u[0] = HIDDEN;
        gate.u[10] = 32 * 4;
        gate.u[12] = 640;
        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,
            tokens.len() as u32,
        )?;
        let mut norm = rms_args(HC_WIDTH, HIDDEN);
        norm.u[4] = tokens.len() as u32;
        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 * tokens.len() as u32,
            1,
        )?;
        let mut conv = args();
        conv.u[0] = HC_WIDTH;
        conv.u[4] = tokens.len() as u32;
        self.dispatch_with_d(
            c"kernel_qwen_ple_conv",
            &self.scratch.ple_output,
            Some(&self.scratch.ple_gated),
            Some(&self.scratch.ple_norm),
            Some(&self.ple_state.conv),
            Some(state_out.unwrap_or(&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 * tokens.len() as u32;
        self.dispatch(
            c"kernel_qwen_add",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            Some(&self.scratch.ple_output),
            None,
            &[],
            &add,
            HC_WIDTH * tokens.len() as u32,
            1,
        )?;
        std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm);
        self.ple_state.history = history;
        Ok(())
    }

    fn encode_layer(&mut self, layer: usize) -> Result<(), String> {
        let prefix = format!("language_model.model.layers.{layer}");
        self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?;
        match &self.states[layer] {
            LayerState::Gdn { .. } => self.gdn(&prefix, layer)?,
            state @ LayerState::Attention { .. } => {
                self.attention(&prefix, state, self.position)?
            }
        }
        self.hyper_write()?;
        self.encode_moe(&prefix, &layer.to_string())
    }

    fn encode_moe(&mut self, prefix: &str, _layer_name: &str) -> Result<(), String> {
        self.encode_moe_rows(prefix, 1)
    }

    fn encode_moe_rows(&mut self, prefix: &str, rows: u32) -> Result<(), String> {
        self.hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), rows)?;
        self.affine_batch_into(
            &self.affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)?,
            &self.scratch.block,
            &self.scratch.router,
            HIDDEN,
            EXPERTS,
            rows,
            true,
        )?;
        let mut softmax_args = args();
        softmax_args.u[0] = EXPERTS;
        softmax_args.u[12] = EXPERTS / 4;
        self.dispatch(
            c"kernel_qwen_softmax_precise_512",
            &self.scratch.router,
            Some(&self.scratch.router),
            None,
            None,
            &[],
            &softmax_args,
            rows,
            1,
        )?;
        let mut route_args = args();
        route_args.u[0] = EXPERTS;
        route_args.u[12] = EXPERTS;
        self.dispatch(
            c"kernel_qwen_route_top10",
            &self.scratch.route_ids,
            Some(&self.scratch.route_weights),
            Some(&self.scratch.router),
            None,
            &[],
            &route_args,
            rows,
            1,
        )?;
        self.routed_experts_rows(prefix, rows)?;
        self.shared_expert_rows(prefix, rows)?;
        let mut inject_args = args();
        inject_args.u[0] = HIDDEN;
        inject_args.u[4] = rows;
        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 * rows,
            1,
        )?;
        std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm);
        Ok(())
    }

    fn hyper_read(&self, prefix: &str) -> Result<(), String> {
        self.hyper_read_rows(prefix, 1)
    }

    fn hyper_read_rows(&self, prefix: &str, rows: u32) -> 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"))?;
        for (weight, shape) in [
            (norm, &[HC_WIDTH as u64][..]),
            (down, &[HC_RANK as u64, HC_WIDTH as u64][..]),
            (up, &[HC_WIDTH as u64, HC_RANK as u64][..]),
            (inject, &[HC as u64, HC_WIDTH as u64][..]),
        ] {
            if weight.tensor.dtype != "BF16" || weight.tensor.shape != shape {
                return Err(format!(
                    "{} has an incompatible BF16 tensor shape",
                    weight.tensor.name
                ));
            }
        }
        if rows == 1
            && let Some(pack) = self.hyper_packs.get(prefix)
        {
            let mut values = args();
            values.u[12] = 1024;
            self.dispatch_with_d(
                c"kernel_qwen_hyper_v3_r1",
                &self.scratch.rank,
                Some(&self.scratch.hc),
                Some(&pack.r1),
                Some(&self.scratch.injection),
                Some(&self.scratch.hyper_rms),
                &[self.view(norm)],
                &values,
                HYPER_FOLDED_ROWS.div_ceil(32),
                rows,
            )?;
            self.dispatch_with_d(
                c"kernel_qwen_hyper_v3_r2",
                &self.scratch.block,
                Some(&self.scratch.hc),
                Some(&pack.r2),
                Some(&self.scratch.rank),
                Some(&self.scratch.hyper_rms),
                &[self.view(norm)],
                &values,
                HIDDEN.div_ceil(32),
                rows,
            )?;
            return Ok(());
        }
        let mut rms = rms_args(HC_WIDTH, HIDDEN);
        rms.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_zero_rms",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            None,
            None,
            &[self.view(norm)],
            &rms,
            HC * rows,
            1,
        )?;
        self.bf16_batch(
            down,
            &self.scratch.hc_norm,
            &self.scratch.rank,
            HC_WIDTH,
            HC_RANK,
            rows,
        )?;
        let mut unary = args();
        unary.u[0] = HC_RANK * rows;
        self.dispatch(
            c"kernel_qwen_silu_div4",
            &self.scratch.rank,
            Some(&self.scratch.rank),
            None,
            None,
            &[],
            &unary,
            HC_RANK * rows,
            1,
        )?;
        self.bf16_batch(
            up,
            &self.scratch.rank,
            &self.scratch.hc_mix,
            HC_RANK,
            HC_WIDTH,
            rows,
        )?;
        unary.u[0] = HC_WIDTH * rows;
        self.dispatch(
            c"kernel_qwen_sigmoid",
            &self.scratch.hc_mix,
            Some(&self.scratch.hc_mix),
            None,
            None,
            &[],
            &unary,
            HC_WIDTH * rows,
            1,
        )?;
        let mut mix = args();
        mix.u[0] = HIDDEN;
        mix.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_hyper_mix",
            &self.scratch.block,
            Some(&self.scratch.hc_norm),
            Some(&self.scratch.hc_mix),
            None,
            &[],
            &mix,
            HIDDEN * rows,
            1,
        )?;
        self.bf16_batch(
            inject,
            &self.scratch.hc_norm,
            &self.scratch.injection,
            HC_WIDTH,
            HC,
            rows,
        )?;
        unary.u[0] = HC * rows;
        self.dispatch(
            c"kernel_qwen_sigmoid2_div4",
            &self.scratch.injection,
            Some(&self.scratch.injection),
            None,
            None,
            &[],
            &unary,
            HC * rows,
            1,
        )
    }

    fn hyper_write(&mut self) -> Result<(), String> {
        self.hyper_write_rows(1)
    }

    fn hyper_write_rows(&mut self, rows: u32) -> Result<(), String> {
        let mut values = args();
        values.u[0] = HIDDEN;
        values.u[4] = rows;
        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 * rows,
            1,
        )?;
        std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm);
        Ok(())
    }

    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());
        };
        let (qkv, controls) = self.gdn_input(prefix)?;
        if self.fused_gdn && self.position > 0 {
            let mut fused = args();
            fused.u[0] = GDN_VALUE;
            fused.u[1] = GDN_VALUE + GDN_HEADS_V;
            fused.u[10] = (HEAD_DIM * 4 + 8) * 4;
            fused.u[12] = 256;
            self.dispatch_with_d(
                c"kernel_qwen_gdn_step_fused",
                &self.scratch.gdn_out,
                Some(qkv),
                Some(controls),
                Some(recurrent),
                Some(conv),
                &[
                    self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?),
                    self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?),
                    self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?),
                    self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?),
                ],
                &fused,
                1,
                GDN_HEADS_V,
            )?;
            let next_conv = self
                .scratch
                .gdn_out
                .view(u64::from(GDN_VALUE) * 4, u64::from(GDN_QKV) * 3 * 2)?;
            let mut commit = args();
            commit.u[0] = GDN_QKV * 3;
            self.dispatch(
                c"kernel_qwen_copy_bf16",
                conv,
                Some(&next_conv),
                None,
                None,
                &[],
                &commit,
                GDN_QKV * 3,
                1,
            )?;
        } else {
            let mut conv_args = args();
            conv_args.u[0] = GDN_QKV;
            conv_args.u[12] = 1024;
            self.dispatch(
                c"kernel_qwen_gdn_conv_norm",
                qkv,
                Some(qkv),
                Some(conv),
                None,
                &[self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?)],
                &conv_args,
                GDN_QKV / 1024,
                1,
            )?;
            let mut control_args = args();
            control_args.u[0] = GDN_VALUE;
            control_args.u[1] = GDN_VALUE + GDN_HEADS_V;
            self.dispatch(
                c"kernel_qwen_gdn_controls",
                controls,
                None,
                None,
                None,
                &[
                    self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?),
                    self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?),
                ],
                &control_args,
                GDN_HEADS_V,
                1,
            )?;
            let mut delta_args = control_args;
            delta_args.u[12] = 64;
            self.dispatch(
                c"kernel_qwen_gdn_delta_packed",
                &self.scratch.gdn_out,
                Some(qkv),
                Some(controls),
                Some(recurrent),
                &[],
                &delta_args,
                HEAD_DIM / 16,
                GDN_HEADS_V,
            )?;
            let mut gate_args = args();
            gate_args.u[12] = 32;
            self.dispatch(
                c"kernel_qwen_gdn_norm_gate",
                &self.scratch.gdn_out,
                Some(&self.scratch.gdn_out),
                Some(controls),
                None,
                &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)],
                &gate_args,
                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 gdn_input<'a>(&'a self, prefix: &str) -> Result<(&'a Buffer, &'a Buffer), String> {
        let qkv = self.affine(
            &format!("{prefix}.linear_attn.in_proj_qkv"),
            HIDDEN,
            GDN_QKV,
            None,
        )?;
        if self.fused_gdn_input {
            for (name, width) in [
                ("in_proj_z", GDN_VALUE),
                ("in_proj_b", GDN_HEADS_V),
                ("in_proj_a", GDN_HEADS_V),
            ] {
                let projection =
                    self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?;
                if projection.bits != qkv.bits || projection.group != qkv.group {
                    return Err(format!(
                        "{prefix}.linear_attn has incompatible fused input quantization"
                    ));
                }
            }
            let fused = |suffix: &str, bytes: u64| -> Result<QwenWeightView, String> {
                let name = format!("{prefix}.linear_attn.in_proj_fused.{suffix}");
                let buffer = self
                    .weights
                    .get(&name)
                    .ok_or_else(|| format!("Qwen materialization missed {name}"))?;
                Ok(QwenWeightView {
                    tensor: buffer.raw().cast_const(),
                    map: std::ptr::null(),
                    size: bytes,
                    offset: 0,
                    bytes,
                })
            };
            let packed = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN) * u64::from(qkv.bits) / 8;
            let parameters = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN / qkv.group) * 2;
            let mut values = args();
            values.u[0] = HIDDEN;
            values.u[1] = GDN_FUSED_INPUT;
            values.u[2] = qkv.bits;
            values.u[3] = qkv.group;
            values.u[11] = 1;
            values.u[12] = 64;
            self.dispatch(
                self.affine_qmv_kernel(&qkv, false, HIDDEN, GDN_FUSED_INPUT)?,
                &self.scratch.gdn_fused,
                Some(&self.scratch.block),
                None,
                None,
                &[
                    fused("weight", packed)?,
                    fused("scales", parameters)?,
                    fused("biases", parameters)?,
                ],
                &values,
                GDN_FUSED_INPUT.div_ceil(8),
                1,
            )?;
            return Ok((
                &self.scratch.gdn_fused_qkv,
                &self.scratch.gdn_fused_controls,
            ));
        }

        self.affine_mv_into(
            &qkv,
            &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,
            )?;
        }
        Ok((&self.scratch.qkv, &self.scratch.controls))
    }

    fn gdn_rows(
        &self,
        prefix: &str,
        layer: usize,
        rows: u32,
        capture_verify: bool,
    ) -> Result<(), String> {
        let LayerState::Gdn { conv, recurrent } = &self.states[layer] else {
            return Err("Qwen GDN graph received attention state".into());
        };
        let verify_capture = if capture_verify {
            Some(
                self.mtp
                    .as_ref()
                    .ok_or("Qwen MTP is not configured")?
                    .verify
                    .gdn
                    .get(layer)
                    .and_then(Option::as_ref)
                    .ok_or("Qwen verifier GDN capture layout is invalid")?,
            )
        } else {
            None
        };
        let qkv_out = verify_capture.map_or(&self.scratch.qkv, |capture| &capture.qkv);
        let q_out = verify_capture.map_or(&self.scratch.gdn_q, |capture| &capture.q);
        let k_out = verify_capture.map_or(&self.scratch.gdn_k, |capture| &capture.k);
        let conv_out = verify_capture.map_or(conv, |capture| &capture.conv);
        let recurrent_out = verify_capture.map_or(recurrent, |capture| &capture.recurrent);
        if self.fused_gdn_input {
            self.gdn_fused_input_rows(prefix, rows, qkv_out)?;
        } else {
            self.affine_batch_into(
                &self.affine(
                    &format!("{prefix}.linear_attn.in_proj_qkv"),
                    HIDDEN,
                    GDN_QKV,
                    None,
                )?,
                &self.scratch.block,
                qkv_out,
                HIDDEN,
                GDN_QKV,
                rows,
                true,
            )?;
            for (name, target, width) in [
                ("in_proj_z", &self.scratch.gdn_z, GDN_VALUE),
                ("in_proj_b", &self.scratch.gdn_b, GDN_HEADS_V),
                ("in_proj_a", &self.scratch.gdn_a, GDN_HEADS_V),
            ] {
                self.affine_batch_into(
                    &self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?,
                    &self.scratch.block,
                    target,
                    HIDDEN,
                    width,
                    rows,
                    true,
                )?;
            }
            let mut pack = args();
            pack.u[4] = rows;
            self.dispatch_with_d(
                c"kernel_qwen_pack_gdn_controls",
                &self.scratch.controls,
                Some(&self.scratch.gdn_z),
                Some(&self.scratch.gdn_b),
                Some(&self.scratch.gdn_a),
                None,
                &[],
                &pack,
                GDN_CONTROLS * rows,
                1,
            )?;
        }
        encode_gdn_conv_norm_rows(
            qkv_out,
            q_out,
            k_out,
            &self.scratch.gdn_out,
            conv,
            conv_out,
            self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?),
            rows,
        )?;
        let mut controls = args();
        controls.u[0] = GDN_VALUE;
        controls.u[1] = GDN_VALUE + GDN_HEADS_V;
        controls.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_gdn_controls",
            &self.scratch.controls,
            None,
            None,
            None,
            &[
                self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?),
                self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?),
            ],
            &controls,
            GDN_HEADS_V * rows,
            1,
        )?;
        if let Some(capture) = verify_capture {
            for (target, source, width, message) in [
                (
                    &capture.v,
                    &self.scratch.gdn_out,
                    GDN_VALUE,
                    "capturing Qwen verifier GDN value rows",
                ),
                (
                    &capture.controls,
                    &self.scratch.controls,
                    GDN_CONTROLS,
                    "capturing Qwen verifier GDN control rows",
                ),
            ] {
                target.copy_from(0, source, 0, u64::from(rows * width) * 4, message)?;
            }
        }
        encode_gdn_delta_rows(
            &self.scratch.gdn_out,
            q_out,
            k_out,
            &self.scratch.controls,
            recurrent,
            recurrent_out,
            rows,
        )?;
        let mut norm = controls;
        norm.u[12] = 32;
        self.dispatch(
            c"kernel_qwen_gdn_norm_gate",
            &self.scratch.gdn_z,
            Some(&self.scratch.gdn_out),
            Some(&self.scratch.controls),
            None,
            &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)],
            &norm,
            GDN_HEADS_V,
            rows,
        )?;
        self.affine_batch_into(
            &self.affine(
                &format!("{prefix}.linear_attn.out_proj"),
                GDN_VALUE,
                HIDDEN,
                None,
            )?,
            &self.scratch.gdn_z,
            &self.scratch.hidden,
            GDN_VALUE,
            HIDDEN,
            rows,
            true,
        )
    }

    fn gdn_fused_input_rows(
        &self,
        prefix: &str,
        rows: u32,
        qkv_out: &Buffer,
    ) -> Result<(), String> {
        let qkv = self.affine(
            &format!("{prefix}.linear_attn.in_proj_qkv"),
            HIDDEN,
            GDN_QKV,
            None,
        )?;
        for (name, width) in [
            ("in_proj_z", GDN_VALUE),
            ("in_proj_b", GDN_HEADS_V),
            ("in_proj_a", GDN_HEADS_V),
        ] {
            let projection =
                self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?;
            if projection.bits != qkv.bits || projection.group != qkv.group {
                return Err(format!(
                    "{prefix}.linear_attn has incompatible fused input quantization"
                ));
            }
        }
        let fused = |suffix: &str, bytes: u64| -> Result<QwenWeightView, String> {
            let name = format!("{prefix}.linear_attn.in_proj_fused.{suffix}");
            let buffer = self
                .weights
                .get(&name)
                .ok_or_else(|| format!("Qwen materialization missed {name}"))?;
            Ok(QwenWeightView {
                tensor: buffer.raw().cast_const(),
                map: std::ptr::null(),
                size: bytes,
                offset: 0,
                bytes,
            })
        };
        let packed = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN) * u64::from(qkv.bits) / 8;
        let parameters = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN / qkv.group) * 2;
        let mut values = args();
        values.u[0] = HIDDEN;
        values.u[1] = GDN_FUSED_INPUT;
        values.u[2] = qkv.bits;
        values.u[3] = qkv.group;
        values.u[4] = rows;
        values.u[11] = 1;
        values.u[12] = 64;
        let views = [
            fused("weight", packed)?,
            fused("scales", parameters)?,
            fused("biases", parameters)?,
        ];
        if rows > VERIFY_ROWS as u32 {
            if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 {
                return Err("Qwen batched prefill requires the Metal 4 tensor API".to_string());
            }
            values.u[10] = 0;
            values.u[12] = 128;
            self.dispatch(
                self.affine_qmm_kernel(&qkv, rows)?,
                &self.scratch.gdn_fused,
                Some(&self.scratch.block),
                None,
                None,
                &views,
                &values,
                GDN_FUSED_INPUT.div_ceil(64),
                rows.div_ceil(if rows <= 32 { 32 } else { 64 }),
            )?;
        } else {
            self.dispatch(
                self.affine_qmv_batch_kernel(&qkv, HIDDEN, GDN_FUSED_INPUT)?,
                &self.scratch.gdn_fused,
                Some(&self.scratch.block),
                None,
                None,
                &views,
                &values,
                GDN_FUSED_INPUT.div_ceil(8),
                1,
            )?;
        }
        let mut unpack = args();
        unpack.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_unpack_gdn_inputs",
            qkv_out,
            Some(&self.scratch.gdn_fused),
            Some(&self.scratch.controls),
            None,
            &[],
            &unpack,
            GDN_FUSED_INPUT * rows,
            1,
        )
    }

    fn attention(&self, prefix: &str, state: &LayerState, position: u32) -> Result<(), String> {
        let LayerState::Attention {
            kv,
            qsa_raw,
            qsa_pooled,
        } = state
        else {
            return Err("Qwen attention graph received GDN state".into());
        };
        let (q_packed, k, v, qsa_qk) = self.attention_input(prefix)?;
        let mut index = args();
        index.u[0] = QSA_DIM;
        index.u[1] = QSA_HEADS;
        index.u[2] = position;
        self.dispatch(
            c"kernel_qwen_qsa_store_raw",
            qsa_raw,
            Some(qsa_qk),
            Some(&self.scratch.qsa_q),
            None,
            &[],
            &index,
            QSA_WIDTH,
            1,
        )?;
        self.indexer_norm_rope(
            &self.scratch.qsa_q,
            &self.scratch.qsa_qk,
            self.weight(&format!("{prefix}.self_attn.indexer.q_layernorm.weight"))?,
            QSA_HEADS,
            position,
        )?;
        if (position + 1).is_multiple_of(QSA_RATIO) {
            let mut pool = args();
            pool.u[0] = QSA_DIM;
            pool.u[1] = position / QSA_RATIO;
            pool.u[2] = QSA_RATIO;
            pool.u[3] = position + 1 - QSA_RATIO;
            pool.f[0] = 1.0e-6;
            pool.f[1] = 10_000_000.0;
            self.dispatch(
                c"kernel_qwen_qsa_pool_key",
                qsa_pooled,
                Some(qsa_raw),
                None,
                None,
                &[self
                    .view(self.weight(&format!("{prefix}.self_attn.indexer.k_layernorm.weight"))?)],
                &pool,
                QSA_DIM,
                1,
            )?;
        }
        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(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,
            position,
        )?;
        self.head_norm_rope(
            k,
            &self.scratch.k_rope,
            self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?,
            ATTN_KV_HEADS,
            position,
        )?;
        let mut store = args();
        store.u[0] = ATTN_KV_WIDTH;
        store.u[1] = position;
        self.dispatch(
            c"kernel_qwen_store_kv_bf16",
            kv,
            Some(&self.scratch.k_rope),
            Some(v),
            None,
            &[],
            &store,
            ATTN_KV_WIDTH,
            1,
        )?;
        self.attend(kv, qsa_pooled, position)?;
        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 attention_input<'a>(
        &'a self,
        prefix: &str,
    ) -> Result<(&'a Buffer, &'a Buffer, &'a Buffer, &'a Buffer), String> {
        let q = self.affine(
            &format!("{prefix}.self_attn.q_proj"),
            HIDDEN,
            ATTN_WIDTH * 2,
            None,
        )?;
        if self.fused_qsa_input
            && self
                .weights
                .contains_key(&format!("{prefix}.self_attn.qkv_fused.weight"))
        {
            for (name, width) in [("k_proj", ATTN_KV_WIDTH), ("v_proj", ATTN_KV_WIDTH)] {
                let projection =
                    self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?;
                if projection.bits != q.bits || projection.group != q.group {
                    return Err(format!(
                        "{prefix}.self_attn has incompatible fused input quantization"
                    ));
                }
            }
            let indexer = self.affine(
                &format!("{prefix}.self_attn.indexer.index_qk_proj"),
                HIDDEN,
                QSA_WIDTH,
                None,
            )?;
            let include_indexer = indexer.bits == q.bits && indexer.group == q.group;
            let out_dim = if include_indexer {
                QSA_FUSED_INPUT
            } else {
                QSA_FUSED_QKV
            };
            let fused = |suffix: &str, bytes: u64| -> Result<QwenWeightView, String> {
                let name = format!("{prefix}.self_attn.qkv_fused.{suffix}");
                let buffer = self
                    .weights
                    .get(&name)
                    .ok_or_else(|| format!("Qwen materialization missed {name}"))?;
                Ok(QwenWeightView {
                    tensor: buffer.raw().cast_const(),
                    map: std::ptr::null(),
                    size: bytes,
                    offset: 0,
                    bytes,
                })
            };
            let packed = u64::from(out_dim) * u64::from(HIDDEN) * u64::from(q.bits) / 8;
            let parameters = u64::from(out_dim) * u64::from(HIDDEN / q.group) * 2;
            let mut values = args();
            values.u[0] = HIDDEN;
            values.u[1] = out_dim;
            values.u[2] = q.bits;
            values.u[3] = q.group;
            values.u[11] = 1;
            values.u[12] = 64;
            self.dispatch(
                self.affine_qmv_kernel(&q, false, HIDDEN, out_dim)?,
                &self.scratch.qsa_fused,
                Some(&self.scratch.block),
                None,
                None,
                &[
                    fused("weight", packed)?,
                    fused("scales", parameters)?,
                    fused("biases", parameters)?,
                ],
                &values,
                out_dim.div_ceil(8),
                1,
            )?;
            if !include_indexer {
                self.affine_mv_into(
                    &indexer,
                    &self.scratch.block,
                    &self.scratch.qsa_qk,
                    HIDDEN,
                    QSA_WIDTH,
                )?;
            }
            return Ok((
                &self.scratch.qsa_fused_q,
                &self.scratch.qsa_fused_k,
                &self.scratch.qsa_fused_v,
                if include_indexer {
                    &self.scratch.qsa_fused_indexer
                } else {
                    &self.scratch.qsa_qk
                },
            ));
        }

        for (name, output, width) in [
            ("q_proj", &self.scratch.q_packed, ATTN_WIDTH * 2),
            ("k_proj", &self.scratch.k, ATTN_KV_WIDTH),
            ("v_proj", &self.scratch.v, ATTN_KV_WIDTH),
            ("indexer.index_qk_proj", &self.scratch.qsa_qk, QSA_WIDTH),
        ] {
            self.affine_mv_into(
                &self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?,
                &self.scratch.block,
                output,
                HIDDEN,
                width,
            )?;
        }
        Ok((
            &self.scratch.q_packed,
            &self.scratch.k,
            &self.scratch.v,
            &self.scratch.qsa_qk,
        ))
    }

    fn attention_rows(
        &self,
        prefix: &str,
        state: &LayerState,
        position: u32,
        rows: u32,
    ) -> Result<(), String> {
        let LayerState::Attention {
            kv,
            qsa_raw,
            qsa_pooled,
        } = state
        else {
            return Err("Qwen attention graph received GDN state".into());
        };
        self.affine_batch_into(
            &self.affine(
                &format!("{prefix}.self_attn.indexer.index_qk_proj"),
                HIDDEN,
                QSA_WIDTH,
                None,
            )?,
            &self.scratch.block,
            &self.scratch.qsa_qk,
            HIDDEN,
            QSA_WIDTH,
            rows,
            true,
        )?;
        let mut index = args();
        index.u[0] = QSA_DIM;
        index.u[1] = QSA_HEADS;
        index.u[2] = position;
        index.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_qsa_store_raw",
            qsa_raw,
            Some(&self.scratch.qsa_qk),
            Some(&self.scratch.qsa_q),
            None,
            &[],
            &index,
            QSA_WIDTH,
            rows,
        )?;
        self.indexer_norm_rope_rows(
            &self.scratch.qsa_q,
            &self.scratch.qsa_qk,
            self.weight(&format!("{prefix}.self_attn.indexer.q_layernorm.weight"))?,
            QSA_HEADS,
            position,
            rows,
        )?;
        for row in 0..rows {
            let row_position = position + row;
            if (row_position + 1).is_multiple_of(QSA_RATIO) {
                let mut pool = args();
                pool.u[0] = QSA_DIM;
                pool.u[1] = row_position / QSA_RATIO;
                pool.u[2] = QSA_RATIO;
                pool.u[3] = row_position + 1 - QSA_RATIO;
                pool.f[0] = 1.0e-6;
                pool.f[1] = 10_000_000.0;
                self.dispatch(
                    c"kernel_qwen_qsa_pool_key",
                    qsa_pooled,
                    Some(qsa_raw),
                    None,
                    None,
                    &[self.view(
                        self.weight(&format!("{prefix}.self_attn.indexer.k_layernorm.weight"))?,
                    )],
                    &pool,
                    QSA_DIM,
                    1,
                )?;
            }
        }
        for (name, output, width) in [
            ("q_proj", &self.scratch.q_packed, ATTN_WIDTH * 2),
            ("k_proj", &self.scratch.k, ATTN_KV_WIDTH),
            ("v_proj", &self.scratch.v, ATTN_KV_WIDTH),
        ] {
            self.affine_batch_into(
                &self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?,
                &self.scratch.block,
                output,
                HIDDEN,
                width,
                rows,
                true,
            )?;
        }
        let mut split = args();
        split.u[0] = ATTN_HEADS;
        split.u[1] = ATTN_DIM;
        split.u[4] = rows;
        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 * rows,
            1,
        )?;
        self.head_norm_rope_rows(
            &self.scratch.q,
            &self.scratch.attention,
            self.weight(&format!("{prefix}.self_attn.q_norm.weight"))?,
            ATTN_HEADS,
            position,
            rows,
        )?;
        self.head_norm_rope_rows(
            &self.scratch.k,
            &self.scratch.k_rope,
            self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?,
            ATTN_KV_HEADS,
            position,
            rows,
        )?;
        let mut store = args();
        store.u[0] = ATTN_KV_WIDTH;
        store.u[1] = position;
        store.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_store_kv_bf16",
            kv,
            Some(&self.scratch.k_rope),
            Some(&self.scratch.v),
            None,
            &[],
            &store,
            ATTN_KV_WIDTH,
            rows,
        )?;
        if position + rows <= dense_attention_budget(rows) {
            let mut dense = args();
            dense.u[0] = ATTN_HEADS;
            dense.u[1] = ATTN_KV_HEADS;
            dense.u[2] = ATTN_DIM;
            dense.u[3] = position + rows;
            dense.u[4] = rows;
            dense.u[5] = position;
            dense.u[12] = 1024;
            if rows * (ATTN_HEADS / ATTN_KV_HEADS) > 32 {
                let mut prepare = args();
                prepare.u[0] = ATTN_KV_WIDTH;
                prepare.u[1] = ATTN_DIM;
                prepare.u[3] = position + rows;
                self.dispatch(
                    c"kernel_qwen_prepare_dense_kv_f16",
                    &self.scratch.dense_keys,
                    Some(kv),
                    Some(&self.scratch.dense_values),
                    None,
                    &[],
                    &prepare,
                    ATTN_KV_WIDTH,
                    position + rows,
                )?;
                call(
                    unsafe {
                        ds4_gpu_qwen_attention_flash_staged_tensor(
                            self.scratch.q.raw(),
                            self.scratch.attention.raw(),
                            self.scratch.dense_keys.raw(),
                            self.scratch.dense_values.raw(),
                            position,
                            rows,
                            position + rows,
                            ATTN_HEADS,
                            ATTN_KV_HEADS,
                            ATTN_DIM,
                        )
                    },
                    "running Qwen prefill FlashAttention",
                )?;
                return self.finish_attention_rows(prefix, rows);
            }
            let mut mask = args();
            mask.u[0] = position + rows;
            mask.u[4] = rows;
            mask.u[5] = position;
            self.dispatch(
                c"kernel_qwen_causal_mask",
                &self.scratch.causal_mask,
                None,
                None,
                None,
                &[],
                &mask,
                position + rows,
                rows,
            )?;
            let mut prepare = args();
            prepare.u[0] = ATTN_KV_WIDTH;
            prepare.u[1] = ATTN_DIM;
            prepare.u[3] = position + rows;
            self.dispatch(
                c"kernel_qwen_prepare_dense_kv",
                &self.scratch.dense_keys,
                Some(kv),
                Some(&self.scratch.dense_values),
                None,
                &[],
                &prepare,
                ATTN_KV_WIDTH,
                position + rows,
            )?;
            let mut query = args();
            query.u[0] = ATTN_WIDTH * rows;
            self.dispatch(
                c"kernel_qwen_float_to_bf16",
                &self.scratch.dense_query,
                Some(&self.scratch.attention),
                None,
                None,
                &[],
                &query,
                ATTN_WIDTH * rows,
                1,
            )?;
            self.dispatch_with_d(
                c"kernel_qwen_dense_attention_rows",
                &self.scratch.q,
                Some(&self.scratch.dense_query),
                Some(&self.scratch.dense_keys),
                Some(&self.scratch.dense_values),
                Some(&self.scratch.causal_mask),
                &[],
                &dense,
                ATTN_HEADS,
                rows,
            )?;
        } else {
            let complete_blocks = (position + rows) / QSA_RATIO;
            let mut score = args();
            score.u[0] = QSA_DIM;
            score.u[1] = QSA_HEADS;
            score.u[2] = complete_blocks;
            score.u[3] = position + 1;
            score.u[5] = QSA_RATIO;
            score.f[0] = (QSA_DIM as f32).sqrt().recip();
            self.dispatch(
                c"kernel_qwen_qsa_scores",
                &self.scratch.qsa_scores,
                Some(&self.scratch.qsa_qk),
                Some(qsa_pooled),
                None,
                &[],
                &score,
                complete_blocks,
                rows,
            )?;
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(
                        self.scratch.qsa_selected.raw(),
                        self.scratch.qsa_scores.raw(),
                        complete_blocks,
                        rows,
                        QSA_TOP_K,
                    )
                },
                "selecting batched Qwen QSA blocks",
            )?;
            if packed_dense_qsa(rows, position + rows) {
                let mut mask = args();
                mask.u[0] = position + rows;
                mask.u[4] = rows;
                mask.u[5] = QSA_RATIO;
                mask.u[6] = QSA_TOP_K;
                mask.u[8] = position + 1;
                mask.u[12] = 512;
                self.dispatch(
                    c"kernel_qwen_qsa_mask",
                    &self.scratch.causal_mask,
                    Some(&self.scratch.qsa_selected),
                    None,
                    None,
                    &[],
                    &mask,
                    rows,
                    1,
                )?;
                let mut dense = args();
                dense.u[0] = ATTN_HEADS;
                dense.u[1] = ATTN_KV_HEADS;
                dense.u[2] = ATTN_DIM;
                dense.u[3] = position + rows;
                dense.u[4] = rows;
                dense.u[12] = 1024;
                self.dispatch(
                    c"kernel_qwen_dense_attention_masked",
                    &self.scratch.q,
                    Some(&self.scratch.attention),
                    Some(kv),
                    Some(&self.scratch.causal_mask),
                    &[],
                    &dense,
                    ATTN_HEADS,
                    rows,
                )?;
                return self.finish_attention_rows(prefix, rows);
            }
            let mut order = args();
            order.u[0] = QSA_TOP_K;
            self.dispatch(
                c"kernel_qwen_qsa_sort_blocks",
                &self.scratch.qsa_selected,
                None,
                None,
                None,
                &[],
                &order,
                rows,
                1,
            )?;
            let mut sparse = args();
            sparse.u[0] = ATTN_HEADS;
            sparse.u[1] = ATTN_KV_HEADS;
            sparse.u[2] = ATTN_DIM;
            sparse.u[3] = position + rows;
            sparse.u[4] = QSA_TOP_K;
            sparse.u[5] = QSA_RATIO;
            sparse.u[8] = position + 1;
            sparse.u[12] = 256;
            self.dispatch(
                c"kernel_qwen_sparse_attention",
                &self.scratch.q,
                Some(&self.scratch.attention),
                Some(kv),
                Some(&self.scratch.qsa_selected),
                &[],
                &sparse,
                ATTN_HEADS,
                rows,
            )?;
        }
        self.finish_attention_rows(prefix, rows)
    }

    fn finish_attention_rows(&self, prefix: &str, rows: u32) -> Result<(), String> {
        let mut gate = args();
        gate.u[0] = ATTN_WIDTH * rows;
        self.dispatch(
            c"kernel_qwen_gate_attention",
            &self.scratch.attention,
            Some(&self.scratch.q),
            Some(&self.scratch.q_gate),
            None,
            &[],
            &gate,
            ATTN_WIDTH * rows,
            1,
        )?;
        self.affine_batch_into(
            &self.affine(
                &format!("{prefix}.self_attn.o_proj"),
                ATTN_WIDTH,
                HIDDEN,
                None,
            )?,
            &self.scratch.attention,
            &self.scratch.hidden,
            ATTN_WIDTH,
            HIDDEN,
            rows,
            true,
        )
    }

    fn indexer_norm_rope(
        &self,
        input: &Buffer,
        output: &Buffer,
        weight: Weight<'_>,
        heads: u32,
        position: u32,
    ) -> Result<(), String> {
        self.indexer_norm_rope_rows(input, output, weight, heads, position, 1)
    }

    fn indexer_norm_rope_rows(
        &self,
        input: &Buffer,
        output: &Buffer,
        weight: Weight<'_>,
        heads: u32,
        position: u32,
        rows: u32,
    ) -> Result<(), String> {
        let mut values = args();
        values.u[0] = QSA_DIM;
        values.u[1] = 64;
        values.u[2] = heads;
        values.u[3] = position;
        values.u[12] = QSA_DIM.div_ceil(4).div_ceil(32) * 32;
        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,
            heads,
            rows,
        )
    }

    fn attend(&self, kv: &Buffer, qsa_pooled: &Buffer, position: u32) -> Result<(), String> {
        self.attend_row(
            kv,
            qsa_pooled,
            position,
            &self.scratch.attention,
            &self.scratch.q,
            &self.scratch.qsa_qk,
            &self.scratch.qsa_scores,
            &self.scratch.qsa_selected,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn attend_row(
        &self,
        kv: &Buffer,
        qsa_pooled: &Buffer,
        position: u32,
        query: &Buffer,
        output: &Buffer,
        qsa_query: &Buffer,
        qsa_scores: &Buffer,
        qsa_selected: &Buffer,
    ) -> Result<(), String> {
        let tokens = position + 1;
        let mut values = args();
        values.u[0] = ATTN_HEADS;
        values.u[1] = ATTN_KV_HEADS;
        values.u[2] = ATTN_DIM;
        values.u[3] = tokens;
        values.u[4] = 1;
        values.u[5] = position;
        if tokens <= DECODE_DENSE_BUDGET {
            values.u[12] = 1024;
            return self.dispatch(
                c"kernel_qwen_dense_attention",
                output,
                Some(query),
                Some(kv),
                None,
                &[],
                &values,
                ATTN_HEADS,
                1,
            );
        }

        let complete_blocks = tokens / QSA_RATIO;
        let mut score = args();
        score.u[0] = QSA_DIM;
        score.u[1] = QSA_HEADS;
        score.u[2] = complete_blocks;
        score.f[0] = (QSA_DIM as f32).sqrt().recip();
        self.dispatch(
            c"kernel_qwen_qsa_scores",
            qsa_scores,
            Some(qsa_query),
            Some(qsa_pooled),
            None,
            &[],
            &score,
            complete_blocks,
            1,
        )?;
        call(
            unsafe {
                ds4_gpu_indexer_topk_tensor(
                    qsa_selected.raw(),
                    qsa_scores.raw(),
                    complete_blocks,
                    1,
                    QSA_TOP_K,
                )
            },
            "selecting Qwen QSA blocks",
        )?;
        if packed_dense_qsa(1, tokens) {
            let mut mask = args();
            mask.u[0] = tokens;
            mask.u[4] = 1;
            mask.u[5] = QSA_RATIO;
            mask.u[6] = QSA_TOP_K;
            mask.u[8] = position + 1;
            mask.u[12] = 512;
            self.dispatch(
                c"kernel_qwen_qsa_mask",
                &self.scratch.causal_mask,
                Some(qsa_selected),
                None,
                None,
                &[],
                &mask,
                1,
                1,
            )?;
            values.u[4] = 1;
            values.u[12] = 1024;
            return self.dispatch(
                c"kernel_qwen_dense_attention_masked",
                output,
                Some(query),
                Some(kv),
                Some(&self.scratch.causal_mask),
                &[],
                &values,
                ATTN_HEADS,
                1,
            );
        }
        let mut order = args();
        order.u[0] = QSA_TOP_K;
        self.dispatch(
            c"kernel_qwen_qsa_sort_blocks",
            qsa_selected,
            None,
            None,
            None,
            &[],
            &order,
            1,
            1,
        )?;
        values.u[4] = QSA_TOP_K;
        values.u[5] = QSA_RATIO;
        values.u[6] = complete_blocks * QSA_RATIO;
        values.u[7] = tokens - complete_blocks * QSA_RATIO;
        values.u[12] = 256;
        self.dispatch(
            c"kernel_qwen_sparse_attention",
            output,
            Some(query),
            Some(kv),
            Some(qsa_selected),
            &[],
            &values,
            ATTN_HEADS,
            1,
        )
    }

    fn head_norm_rope(
        &self,
        input: &Buffer,
        output: &Buffer,
        weight: Weight<'_>,
        heads: u32,
        position: u32,
    ) -> Result<(), String> {
        self.head_norm_rope_rows(input, output, weight, heads, position, 1)
    }

    fn head_norm_rope_rows(
        &self,
        input: &Buffer,
        output: &Buffer,
        weight: Weight<'_>,
        heads: u32,
        position: u32,
        rows: u32,
    ) -> Result<(), String> {
        let mut values = args();
        values.u[0] = ATTN_DIM;
        values.u[1] = 64;
        values.u[2] = heads;
        values.u[3] = position;
        values.u[12] = ATTN_DIM.div_ceil(4).div_ceil(32) * 32;
        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,
            heads,
            rows,
        )
    }

    #[cfg(test)]
    fn routed_experts(&self, prefix: &str) -> Result<(), String> {
        self.routed_experts_rows(prefix, 1)
    }

    fn routed_experts_rows(&self, prefix: &str, rows: u32) -> Result<(), String> {
        if rows > VERIFY_ROWS as u32 {
            self.prepare_routed_expert_map(rows)?;
            return self.routed_experts_sorted_rows(prefix, rows);
        }
        self.routed_experts_grouped_rows(prefix, rows)
    }

    fn routed_experts_grouped_rows(&self, prefix: &str, rows: u32) -> Result<(), String> {
        let gate = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.gate_proj"),
            HIDDEN,
            EXPERT_WIDTH,
        )?;
        let up = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.up_proj"),
            HIDDEN,
            EXPERT_WIDTH,
        )?;
        if self.packed_moe {
            self.affine_pair_batch_into(
                &gate,
                &up,
                &self.scratch.block,
                &self.scratch.gate,
                &self.scratch.up,
                HIDDEN,
                EXPERT_WIDTH,
                true,
                rows,
                false,
            )?;
        } else {
            for (table, target) in [(&gate, &self.scratch.gate), (&up, &self.scratch.up)] {
                self.affine_gather_batch_into(
                    table,
                    &self.scratch.block,
                    target,
                    HIDDEN,
                    EXPERT_WIDTH,
                    rows,
                    false,
                )?;
            }
        }
        let mut swiglu = args();
        swiglu.u[0] = EXPERT_WIDTH * EXPERTS_USED as u32 * rows;
        self.dispatch(
            c"kernel_qwen_swiglu",
            &self.scratch.mid,
            Some(&self.scratch.gate),
            Some(&self.scratch.up),
            None,
            &[],
            &swiglu,
            EXPERT_WIDTH * EXPERTS_USED as u32 * rows,
            1,
        )?;
        let table = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.down_proj"),
            EXPERT_WIDTH,
            HIDDEN,
        )?;
        self.affine_gather_batch_into(
            &table,
            &self.scratch.mid,
            &self.scratch.expert,
            EXPERT_WIDTH,
            HIDDEN,
            rows,
            true,
        )?;
        let mut sum = args();
        sum.u[0] = HIDDEN * rows;
        sum.u[1] = HIDDEN;
        self.dispatch(
            c"kernel_qwen_weighted_sum10",
            &self.scratch.moe,
            Some(&self.scratch.expert),
            Some(&self.scratch.route_weights),
            None,
            &[],
            &sum,
            HIDDEN * rows,
            1,
        )
    }

    fn routed_experts_sorted_rows(&self, prefix: &str, rows: u32) -> Result<(), String> {
        let mut gather = args();
        gather.u[0] = HIDDEN;
        gather.u[4] = rows;
        gather.u[6] = self.scratch.rows * EXPERTS_USED as u32;
        gather.u[8] = EXPERTS_USED as u32;
        self.dispatch(
            c"kernel_qwen_route_gather_rows",
            &self.scratch.expert,
            Some(&self.scratch.block),
            Some(&self.scratch.route_work),
            None,
            &[],
            &gather,
            HIDDEN,
            rows * EXPERTS_USED as u32,
        )?;

        let gate = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.gate_proj"),
            HIDDEN,
            EXPERT_WIDTH,
        )?;
        let up = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.up_proj"),
            HIDDEN,
            EXPERT_WIDTH,
        )?;
        self.affine_sorted_qmm_into(
            &gate,
            &self.scratch.expert,
            &self.scratch.gate,
            HIDDEN,
            EXPERT_WIDTH,
            rows,
        )?;
        self.affine_sorted_qmm_into(
            &up,
            &self.scratch.expert,
            &self.scratch.up,
            HIDDEN,
            EXPERT_WIDTH,
            rows,
        )?;
        let mut swiglu = args();
        swiglu.u[0] = EXPERT_WIDTH * EXPERTS_USED as u32 * rows;
        self.dispatch(
            c"kernel_qwen_swiglu_bf16",
            &self.scratch.mid,
            Some(&self.scratch.gate),
            Some(&self.scratch.up),
            None,
            &[],
            &swiglu,
            EXPERT_WIDTH * EXPERTS_USED as u32 * rows,
            1,
        )?;
        let down = self.affine_table(
            &format!("{prefix}.mlp.switch_mlp.down_proj"),
            EXPERT_WIDTH,
            HIDDEN,
        )?;
        self.affine_sorted_qmm_into(
            &down,
            &self.scratch.mid,
            &self.scratch.expert,
            EXPERT_WIDTH,
            HIDDEN,
            rows,
        )?;
        let mut sum = args();
        sum.u[0] = HIDDEN * rows;
        sum.u[1] = HIDDEN;
        sum.u[4] = rows;
        sum.u[8] = EXPERTS_USED as u32;
        self.dispatch(
            c"kernel_qwen_weighted_sum10_sorted",
            &self.scratch.moe,
            Some(&self.scratch.expert),
            Some(&self.scratch.route_weights),
            Some(&self.scratch.route_work),
            &[],
            &sum,
            HIDDEN * rows,
            1,
        )
    }

    #[cfg(test)]
    fn shared_expert(&self, prefix: &str) -> Result<(), String> {
        self.shared_expert_rows(prefix, 1)
    }

    fn shared_expert_rows(&self, prefix: &str, rows: u32) -> Result<(), String> {
        let gate = self.affine(
            &format!("{prefix}.mlp.shared_expert.gate_proj"),
            HIDDEN,
            EXPERT_WIDTH,
            None,
        )?;
        let up = self.affine(
            &format!("{prefix}.mlp.shared_expert.up_proj"),
            HIDDEN,
            EXPERT_WIDTH,
            None,
        )?;
        if self.packed_moe && rows <= VERIFY_ROWS as u32 {
            self.affine_pair_batch_into(
                &gate,
                &up,
                &self.scratch.block,
                &self.scratch.gate,
                &self.scratch.up,
                HIDDEN,
                EXPERT_WIDTH,
                false,
                rows,
                false,
            )?;
        } else {
            for (weight, target) in [(&gate, &self.scratch.gate), (&up, &self.scratch.up)] {
                self.affine_batch_into(
                    weight,
                    &self.scratch.block,
                    target,
                    HIDDEN,
                    EXPERT_WIDTH,
                    rows,
                    true,
                )?;
            }
        }
        let mut swiglu = args();
        swiglu.u[0] = EXPERT_WIDTH * rows;
        self.dispatch(
            c"kernel_qwen_swiglu",
            &self.scratch.mid,
            Some(&self.scratch.gate),
            Some(&self.scratch.up),
            None,
            &[],
            &swiglu,
            EXPERT_WIDTH * rows,
            1,
        )?;
        self.affine_batch_into(
            &self.affine(
                &format!("{prefix}.mlp.shared_expert.down_proj"),
                EXPERT_WIDTH,
                HIDDEN,
                None,
            )?,
            &self.scratch.mid,
            &self.scratch.shared,
            EXPERT_WIDTH,
            HIDDEN,
            rows,
            true,
        )?;
        self.affine_batch_into(
            &self.affine(&format!("{prefix}.mlp.shared_expert_gate"), HIDDEN, 1, None)?,
            &self.scratch.block,
            &self.scratch.gate,
            HIDDEN,
            1,
            rows,
            true,
        )?;
        let mut accumulate = args();
        accumulate.u[0] = HIDDEN;
        accumulate.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_accumulate_sigmoid_scalar",
            &self.scratch.moe,
            Some(&self.scratch.shared),
            Some(&self.scratch.gate),
            None,
            &[],
            &accumulate,
            HIDDEN * rows,
            1,
        )
    }

    fn final_output(&mut self) -> Result<(), String> {
        self.final_output_rows(1)
    }

    fn final_output_rows(&mut self, rows: u32) -> Result<(), String> {
        self.mix_output_rows("language_model.model.hyper_connection_mixer", rows)?;
        self.affine_batch_into(
            &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?,
            &self.scratch.block,
            &self.scratch.logits,
            HIDDEN,
            VOCAB,
            rows,
            true,
        )
    }

    fn mix_output(&self, prefix: &str) -> Result<(), String> {
        self.mix_output_rows(prefix, 1)
    }

    fn mix_output_rows(&self, prefix: &str, rows: u32) -> 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 mut rms = rms_args(HC_WIDTH, HIDDEN);
        rms.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_zero_rms",
            &self.scratch.hc_norm,
            Some(&self.scratch.hc),
            None,
            None,
            &[self.view(norm)],
            &rms,
            HC * rows,
            1,
        )?;
        self.bf16_batch(
            down,
            &self.scratch.hc_norm,
            &self.scratch.rank,
            HC_WIDTH,
            HC_RANK,
            rows,
        )?;
        let mut unary = args();
        unary.u[0] = HC_RANK * rows;
        self.dispatch(
            c"kernel_qwen_silu_div4",
            &self.scratch.rank,
            Some(&self.scratch.rank),
            None,
            None,
            &[],
            &unary,
            HC_RANK * rows,
            1,
        )?;
        self.bf16_batch(
            up,
            &self.scratch.rank,
            &self.scratch.hc_mix,
            HC_RANK,
            HC_WIDTH,
            rows,
        )?;
        unary.u[0] = HC_WIDTH * rows;
        self.dispatch(
            c"kernel_qwen_sigmoid",
            &self.scratch.hc_mix,
            Some(&self.scratch.hc_mix),
            None,
            None,
            &[],
            &unary,
            HC_WIDTH * rows,
            1,
        )?;
        let mut mix = args();
        mix.u[0] = HIDDEN;
        mix.u[4] = rows;
        self.dispatch(
            c"kernel_qwen_hyper_mix",
            &self.scratch.block,
            Some(&self.scratch.hc_norm),
            Some(&self.scratch.hc_mix),
            None,
            &[],
            &mix,
            HIDDEN * rows,
            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 affine_table(&self, prefix: &str, in_dim: u32, out_dim: u32) -> Result<Affine<'_>, String> {
        let mut table = self.affine(prefix, in_dim, out_dim, Some(0))?;
        table.packed.expert = None;
        table.scales.expert = None;
        table.biases.expert = None;
        Ok(table)
    }

    fn affine_qmv_kernel(
        &self,
        weight: &Affine<'_>,
        gather: bool,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<&'static CStr, String> {
        let alignment = match weight.bits {
            4 => 512,
            8 => 256,
            _ => 0,
        };
        let fast = alignment != 0
            && out_dim.is_multiple_of(8)
            && in_dim.is_multiple_of(alignment)
            && self.affine_is_direct(weight);
        match (gather, fast, weight.bits, weight.group) {
            (false, true, 4, 32) => Ok(c"kernel_qwen_affine_qmv_fast_b4g32"),
            (false, true, 4, 64) => Ok(c"kernel_qwen_affine_qmv_fast_b4g64"),
            (false, true, 8, 64) => Ok(c"kernel_qwen_affine_qmv_fast_b8g64"),
            (true, true, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b4g32"),
            (true, true, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b4g64"),
            (true, true, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b8g64"),
            (false, false, 4, 32) => Ok(c"kernel_qwen_affine_qmv_safe_b4g32"),
            (false, false, 4, 64) => Ok(c"kernel_qwen_affine_qmv_safe_b4g64"),
            (false, false, 8, 64) => Ok(c"kernel_qwen_affine_qmv_safe_b8g64"),
            (true, false, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b4g32"),
            (true, false, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b4g64"),
            (true, false, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b8g64"),
            _ => Err(format!(
                "Qwen affine QMV does not support {}-bit group-{} weights",
                weight.bits, weight.group
            )),
        }
    }

    fn affine_qmv_batch_kernel(
        &self,
        weight: &Affine<'_>,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<&'static CStr, String> {
        let alignment = match weight.bits {
            4 => 512,
            8 => 256,
            _ => 0,
        };
        let fast = alignment != 0
            && out_dim.is_multiple_of(8)
            && in_dim.is_multiple_of(alignment)
            && self.affine_is_direct(weight);
        match (fast, weight.bits, weight.group) {
            (true, 4, 32) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b4g32"),
            (true, 4, 64) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b4g64"),
            (true, 8, 64) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b8g64"),
            (false, 4, 32) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b4g32"),
            (false, 4, 64) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b4g64"),
            (false, 8, 64) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b8g64"),
            _ => Err(format!(
                "Qwen affine batch QMV does not support {}-bit group-{} weights",
                weight.bits, weight.group
            )),
        }
    }

    fn affine_qmm_kernel(&self, weight: &Affine<'_>, rows: u32) -> Result<&'static CStr, String> {
        match (rows <= 32, weight.bits, weight.group) {
            (true, 4, 32) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g32_bm32"),
            (true, 4, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g64_bm32"),
            (true, 8, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b8g64_bm32"),
            (false, 4, 32) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g32_bm64"),
            (false, 4, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g64_bm64"),
            (false, 8, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b8g64_bm64"),
            _ => Err(format!(
                "Qwen affine QMM does not support {}-bit group-{} weights",
                weight.bits, weight.group
            )),
        }
    }

    fn affine_gather_batch_kernel(
        &self,
        weight: &Affine<'_>,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<&'static CStr, String> {
        let alignment = match weight.bits {
            4 => 512,
            8 => 256,
            _ => 0,
        };
        let fast = alignment != 0
            && out_dim.is_multiple_of(8)
            && in_dim.is_multiple_of(alignment)
            && self.affine_is_direct(weight);
        match (fast, weight.bits, weight.group) {
            (true, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b4g32"),
            (true, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b4g64"),
            (true, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b8g64"),
            (false, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b4g32"),
            (false, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b4g64"),
            (false, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b8g64"),
            _ => Err(format!(
                "Qwen affine gather batch does not support {}-bit group-{} weights",
                weight.bits, weight.group
            )),
        }
    }

    fn affine_pair_qmv_batch_kernel(
        &self,
        gate: &Affine<'_>,
        up: &Affine<'_>,
        gather: bool,
        in_dim: u32,
        out_dim: u32,
    ) -> Result<&'static CStr, String> {
        let alignment = match gate.bits {
            4 => 512,
            8 => 256,
            _ => 0,
        };
        let fast = alignment != 0
            && out_dim.is_multiple_of(8)
            && in_dim.is_multiple_of(alignment)
            && self.affine_is_direct(gate)
            && self.affine_is_direct(up);
        match (gather, fast, gate.bits, gate.group) {
            (false, true, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b4g32"),
            (false, true, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b4g64"),
            (false, true, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b8g64"),
            (true, true, 4, 32) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g32"),
            (true, true, 4, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g64"),
            (true, true, 8, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b8g64"),
            (false, false, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b4g32"),
            (false, false, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b4g64"),
            (false, false, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b8g64"),
            (true, false, 4, 32) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g32"),
            (true, false, 4, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g64"),
            (true, false, 8, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b8g64"),
            _ => Err(format!(
                "Qwen affine batch pair does not support {}-bit group-{} weights",
                gate.bits, gate.group
            )),
        }
    }

    fn affine_pair_qmv_wide_kernel(
        &self,
        gate: &Affine<'_>,
        up: &Affine<'_>,
    ) -> Result<&'static CStr, String> {
        match (gate.bits, gate.group, up.bits, up.group) {
            (4, 32, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b4g32"),
            (4, 64, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b4g64"),
            (8, 64, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b8g64"),
            _ => Err("Qwen gate/up projections use incompatible quantization layouts".into()),
        }
    }

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

    fn affine_is_direct(&self, weight: &Affine<'_>) -> bool {
        [&weight.packed, &weight.scales, &weight.biases]
            .into_iter()
            .all(|part| {
                self.weight_views
                    .contains_key(&(part.tensor as *const QwenTensor as usize))
                    || part.tensor.range.start.is_multiple_of(4)
            })
    }

    fn view(&self, weight: Weight<'_>) -> QwenWeightView {
        if let Some(view) = self
            .weight_views
            .get(&(weight.tensor as *const QwenTensor as usize))
        {
            let experts = weight.expert.map_or(1, |_| EXPERTS as u64);
            let bytes = view.bytes / experts;
            return QwenWeightView {
                tensor: view.tensor,
                map: view.map,
                size: view.size,
                offset: u64::from(weight.expert.unwrap_or(0)) * bytes,
                bytes,
            };
        }
        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 {
            tensor: std::ptr::null(),
            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> {
        self.affine_mv_with_precision(weight, input, output, in_dim, out_dim, true)
    }

    fn affine_mv_with_precision(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        bf16_output: bool,
    ) -> 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;
        values.u[11] = u32::from(bf16_output);
        values.u[12] = 64;
        self.dispatch(
            self.affine_qmv_kernel(weight, false, in_dim, out_dim)?,
            output,
            Some(input),
            None,
            None,
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim.div_ceil(8),
            1,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn affine_batch_into(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        rows: u32,
        bf16_output: bool,
    ) -> Result<(), String> {
        if rows == 1 {
            return self.affine_mv_with_precision(
                weight,
                input,
                output,
                in_dim,
                out_dim,
                bf16_output,
            );
        }
        if rows > VERIFY_ROWS as u32 {
            if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 {
                return Err("Qwen batched prefill requires the Metal 4 tensor API".to_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;
            values.u[4] = rows;
            values.u[10] = 0;
            values.u[11] = u32::from(bf16_output);
            values.u[12] = 128;
            let splits = if bf16_output {
                affine_qmm_splits(in_dim, out_dim, rows, weight.group)
            } else {
                1
            };
            let partial_elements = u64::from(splits) * u64::from(rows) * u64::from(out_dim);
            if splits > 1 && (splits > 256 || partial_elements > QMM_SPLIT_ELEMENTS) {
                return Err(
                    "Qwen affine split-K exceeds the supported reduction/scratch bound".into(),
                );
            }
            values.u[8] = splits;
            self.dispatch(
                self.affine_qmm_kernel(weight, rows)?,
                if splits > 1 {
                    &self.scratch.qmm_partials
                } else {
                    output
                },
                Some(input),
                None,
                None,
                &[
                    self.view(weight.packed),
                    self.view(weight.scales),
                    self.view(weight.biases),
                ],
                &values,
                out_dim.div_ceil(64),
                rows.div_ceil(if rows <= 32 { 32 } else { 64 }) * splits,
            )?;
            if splits > 1 {
                values.u[12] = 64;
                self.dispatch(
                    c"kernel_qwen_affine_splitk_reduce",
                    output,
                    Some(&self.scratch.qmm_partials),
                    None,
                    None,
                    &[],
                    &values,
                    (rows * out_dim).div_ceil(2),
                    1,
                )?;
            }
            return Ok(());
        }
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[2] = weight.bits;
        values.u[3] = weight.group;
        values.u[4] = rows;
        values.u[11] = u32::from(bf16_output);
        values.u[12] = 64;
        self.dispatch(
            self.affine_qmv_batch_kernel(weight, in_dim, out_dim)?,
            output,
            Some(input),
            None,
            None,
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim.div_ceil(8),
            1,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn affine_gather_batch_into(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        rows: u32,
        per_expert_input: bool,
    ) -> Result<(), String> {
        if rows > VERIFY_ROWS as u32 {
            return self.affine_gather_qmm_into(
                weight,
                input,
                output,
                in_dim,
                out_dim,
                rows,
                per_expert_input,
            );
        }
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[2] = weight.bits;
        values.u[3] = weight.group;
        values.u[4] = rows;
        values.u[5] = u32::from(per_expert_input);
        values.u[8] = EXPERTS_USED as u32;
        values.u[11] = 1;
        values.u[12] = 64;
        self.dispatch(
            self.affine_gather_batch_kernel(weight, in_dim, out_dim)?,
            output,
            Some(input),
            Some(&self.scratch.route_ids),
            None,
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim.div_ceil(8),
            rows * EXPERTS_USED as u32,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn affine_pair_batch_into(
        &self,
        gate: &Affine<'_>,
        up: &Affine<'_>,
        input: &Buffer,
        gate_output: &Buffer,
        up_output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        gather: bool,
        rows: u32,
        per_expert_input: bool,
    ) -> Result<(), String> {
        if gate.bits != up.bits || gate.group != up.group {
            return Err("Qwen gate/up projections use different quantization layouts".into());
        }
        if gather && rows > VERIFY_ROWS as u32 {
            self.affine_gather_qmm_into(
                gate,
                input,
                gate_output,
                in_dim,
                out_dim,
                rows,
                per_expert_input,
            )?;
            return self.affine_gather_qmm_into(
                up,
                input,
                up_output,
                in_dim,
                out_dim,
                rows,
                per_expert_input,
            );
        }
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[2] = gate.bits;
        values.u[3] = gate.group;
        values.u[4] = rows;
        values.u[5] = u32::from(per_expert_input);
        values.u[8] = EXPERTS_USED as u32;
        values.u[11] = 1;
        values.u[12] = 64;
        let wide = !gather && rows > 1;
        self.dispatch(
            if wide {
                self.affine_pair_qmv_wide_kernel(gate, up)?
            } else {
                self.affine_pair_qmv_batch_kernel(gate, up, gather, in_dim, out_dim)?
            },
            gate_output,
            Some(input),
            Some(&self.scratch.route_ids),
            Some(up_output),
            &[
                self.view(gate.packed),
                self.view(gate.scales),
                self.view(gate.biases),
                self.view(up.packed),
                self.view(up.scales),
                self.view(up.biases),
            ],
            &values,
            out_dim.div_ceil(8) * 2,
            if wide {
                1
            } else if gather {
                rows * EXPERTS_USED as u32
            } else {
                rows
            },
        )
    }

    fn prepare_routed_expert_map(&self, rows: u32) -> Result<(), String> {
        let mut values = args();
        values.u[4] = rows;
        values.u[8] = EXPERTS_USED as u32;
        values.u[10] = EXPERTS * EXPERTS_USED as u32 * 2;
        values.u[12] = EXPERTS;
        self.dispatch(
            c"kernel_qwen_route_map",
            &self.scratch.route_counts,
            Some(&self.scratch.route_ids),
            Some(&self.scratch.route_map),
            Some(&self.scratch.route_work),
            &[],
            &values,
            1,
            1,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn affine_gather_qmm_into(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        rows: u32,
        per_expert_input: bool,
    ) -> Result<(), String> {
        if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 {
            return Err("Qwen grouped prefill MoE requires the Metal 4 tensor API".into());
        }
        let kernel = match (weight.bits, weight.group) {
            (4, 32) => c"kernel_qwen_affine_gather_qmm_mpp_b4g32",
            (4, 64) => c"kernel_qwen_affine_gather_qmm_mpp_b4g64",
            (8, 64) => c"kernel_qwen_affine_gather_qmm_mpp_b8g64",
            _ => {
                return Err(format!(
                    "Qwen grouped affine QMM does not support {}-bit group-{} weights",
                    weight.bits, weight.group
                ));
            }
        };
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[4] = rows;
        values.u[5] = u32::from(per_expert_input);
        values.u[8] = EXPERTS_USED as u32;
        values.u[10] = 64 * 64 * 4;
        values.u[11] = 1;
        values.u[12] = 128;
        let max_work_items = rows
            .saturating_mul(EXPERTS_USED as u32)
            .div_ceil(64)
            .saturating_add(EXPERTS);
        self.dispatch_with_d(
            kernel,
            output,
            Some(input),
            Some(&self.scratch.route_map),
            Some(&self.scratch.route_work),
            Some(&self.scratch.route_counts),
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim.div_ceil(64),
            max_work_items,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn affine_sorted_qmm_into(
        &self,
        weight: &Affine<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        rows: u32,
    ) -> Result<(), String> {
        let kernel = match (weight.bits, weight.group) {
            (4, 32) => c"kernel_qwen_affine_sorted_qmm_mpp_b4g32",
            (4, 64) => c"kernel_qwen_affine_sorted_qmm_mpp_b4g64",
            (8, 64) => c"kernel_qwen_affine_sorted_qmm_mpp_b8g64",
            _ => {
                return Err(format!(
                    "Qwen sorted affine QMM does not support {}-bit group-{} weights",
                    weight.bits, weight.group
                ));
            }
        };
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        values.u[4] = rows;
        values.u[6] = self.scratch.rows * EXPERTS_USED as u32;
        values.u[8] = EXPERTS_USED as u32;
        values.u[11] = 1;
        values.u[12] = 128;
        let max_work_items = rows
            .saturating_mul(EXPERTS_USED as u32)
            .div_ceil(64)
            .saturating_add(EXPERTS);
        self.dispatch_with_d(
            kernel,
            output,
            Some(input),
            Some(&self.scratch.route_map),
            Some(&self.scratch.route_work),
            Some(&self.scratch.route_counts),
            &[
                self.view(weight.packed),
                self.view(weight.scales),
                self.view(weight.biases),
            ],
            &values,
            out_dim.div_ceil(64),
            max_work_items,
        )
    }

    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;
        let (mode, rows_per_group, threads, threadgroup_bytes) = if in_dim <= 64 {
            (3, 32, 32, 0)
        } else if in_dim >= 16 * out_dim {
            (2, 4, 256, 64 * 4)
        } else if out_dim >= 4096 {
            (1, 32, 256, 0)
        } else {
            (0, 16, 128, 0)
        };
        values.u[4] = mode;
        values.u[10] = threadgroup_bytes;
        values.u[12] = threads;
        self.dispatch(
            c"kernel_qwen_bf16_gemv",
            output,
            Some(input),
            None,
            None,
            &[self.view(weight)],
            &values,
            out_dim.div_ceil(rows_per_group),
            1,
        )
    }

    fn bf16_batch(
        &self,
        weight: Weight<'_>,
        input: &Buffer,
        output: &Buffer,
        in_dim: u32,
        out_dim: u32,
        rows: u32,
    ) -> Result<(), String> {
        if rows == VERIFY_ROWS as u32 {
            return self.bf16_mv4(weight, input, output, in_dim, out_dim);
        }
        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
            ));
        }
        if rows > VERIFY_ROWS as u32 {
            if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 {
                return Err("Qwen batched prefill requires the Metal 4 tensor API".to_string());
            }
            let mut values = args();
            values.u[0] = in_dim;
            values.u[1] = out_dim;
            values.u[4] = rows;
            values.u[12] = 128;
            let tile_m = if rows <= 32 { 32 } else { 64 };
            return self.dispatch(
                if tile_m == 32 {
                    c"kernel_qwen_bf16_qmm_mpp_bm32"
                } else {
                    c"kernel_qwen_bf16_qmm_mpp_bm64"
                },
                output,
                Some(input),
                None,
                None,
                &[self.view(weight)],
                &values,
                out_dim.div_ceil(64),
                rows.div_ceil(tile_m),
            );
        }
        let mut values = args();
        values.u[0] = in_dim;
        values.u[1] = out_dim;
        let (mode, rows_per_group, threads, threadgroup_bytes) = if in_dim <= 64 {
            (3, 32, 32, 0)
        } else if in_dim >= 16 * out_dim {
            (2, 4, 256, 64 * 4)
        } else if out_dim >= 4096 {
            (1, 32, 256, 0)
        } else {
            (0, 16, 128, 0)
        };
        values.u[4] = mode;
        values.u[10] = threadgroup_bytes;
        values.u[12] = threads;
        self.dispatch(
            c"kernel_qwen_bf16_gemv",
            output,
            Some(input),
            None,
            None,
            &[self.view(weight)],
            &values,
            out_dim.div_ceil(rows_per_group),
            rows,
        )
    }

    fn bf16_mv4(
        &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;
        values.u[12] = 128;
        self.dispatch(
            c"kernel_qwen_bf16_gemv4",
            output,
            Some(input),
            None,
            None,
            &[self.view(weight)],
            &values,
            out_dim.div_ceil(4),
            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)
    }

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

    fn encode_sparse_probabilities(
        &self,
        logits: &Buffer,
        rows: u32,
        temperature: f32,
        top_k: i32,
    ) -> Result<bool, String> {
        if temperature <= 0.0 || top_k <= 0 || top_k as u32 > SAMPLING_TOP_K {
            return Ok(false);
        }
        let top_k = top_k as u32;
        call(
            unsafe {
                ds4_gpu_indexer_topk_tensor(
                    self.scratch.sampling_selected.raw(),
                    logits.raw(),
                    VOCAB,
                    rows,
                    top_k,
                )
            },
            "selecting Qwen sampling support",
        )?;
        let mut values = args();
        values.u[0] = VOCAB;
        values.u[1] = top_k;
        values.u[4] = rows;
        values.u[12] = 256;
        values.f[0] = temperature.recip();
        self.dispatch(
            c"kernel_qwen_sparse_probabilities",
            &self.scratch.sampling_probabilities,
            Some(logits),
            Some(&self.scratch.sampling_selected),
            None,
            &[],
            &values,
            rows,
            1,
        )?;
        Ok(true)
    }

    fn read_sparse_probabilities(
        &self,
        rows: u32,
        top_p: f32,
        min_p: f32,
        top_k: i32,
    ) -> Result<Option<TokenDistributions>, String> {
        if top_k <= 0 || top_k as u32 > SAMPLING_TOP_K {
            return Ok(None);
        }
        let top_k = top_k as u32;
        let count = rows as usize * top_k as usize;
        let mut ids = vec![0; count];
        let mut probabilities = vec![0.0; count];
        self.scratch.sampling_selected.read_i32(&mut ids)?;
        self.scratch
            .sampling_probabilities
            .read_f32(&mut probabilities)?;
        let top_p = if top_p <= 0.0 || top_p > 1.0 {
            1.0
        } else {
            f64::from(top_p)
        };
        let minimum = f64::from(min_p.max(0.0));
        let mut distributions = Vec::with_capacity(rows as usize);
        for row in 0..rows as usize {
            let start = row * top_k as usize;
            let mut support = ids[start..start + top_k as usize]
                .iter()
                .copied()
                .zip(probabilities[start..start + top_k as usize].iter().copied())
                .filter(|(token, probability)| {
                    *token >= 0 && *token < VOCAB as i32 && probability.is_finite()
                })
                .map(|(token, probability)| (token as usize, f64::from(probability)))
                .collect::<Vec<_>>();
            support.sort_unstable_by(|left, right| {
                right
                    .1
                    .total_cmp(&left.1)
                    .then_with(|| left.0.cmp(&right.0))
            });
            let mut kept = 0.0;
            let mut end = 0;
            for (_, probability) in &support {
                if end > 0 && *probability < minimum * support[0].1 {
                    break;
                }
                if kept >= top_p {
                    break;
                }
                kept += *probability;
                end += 1;
            }
            support.truncate(end);
            if support.is_empty() || !kept.is_finite() || kept <= 0.0 {
                return Ok(None);
            }
            for (_, probability) in &mut support {
                *probability /= kept;
            }
            support.sort_unstable_by_key(|(token, _)| *token);
            distributions.push(support);
        }
        Ok(Some(distributions))
    }

    #[cfg(test)]
    fn target_probabilities_row(
        &self,
        row: usize,
        temperature: f32,
        top_p: f32,
        min_p: f32,
        top_k: i32,
    ) -> Result<TokenDistribution, String> {
        let logits = self
            .scratch
            .logits
            .view(row as u64 * u64::from(VOCAB) * 4, u64::from(VOCAB) * 4)?;
        let commands = Commands::begin()?;
        let sparse = self.encode_sparse_probabilities(&logits, 1, temperature, top_k)?;
        commands.finish()?;
        if sparse
            && let Some(mut rows) = self.read_sparse_probabilities(1, top_p, min_p, top_k)?
            && let Some(probabilities) = rows.pop()
        {
            return Ok(probabilities);
        }
        Ok(qwen_sampling_probabilities(
            &self.logits_row(row)?,
            temperature,
            top_p,
            min_p,
            top_k,
        ))
    }

    fn mtp_history_batch(&mut self, tokens: &[i32], position: u32) -> Result<(), String> {
        let commands = Commands::begin()?;
        self.encode_mtp_history_batch(tokens, position)?;
        commands.finish()
    }

    fn encode_mtp_history_batch(&mut self, tokens: &[i32], position: u32) -> Result<(), String> {
        if tokens.is_empty() || tokens.len() > self.scratch.rows as usize {
            return Err("Qwen MTP prompt-history width exceeds the prefill batch".into());
        }
        if position + tokens.len() as u32 > self.context + MTP_TOKEN_RESERVE {
            return Err("Qwen MTP prompt history exceeds its bounded cache".into());
        }
        if let Some(&token) = tokens
            .iter()
            .find(|&&token| token < 0 || token as u32 >= VOCAB)
        {
            return Err(format!("token {token} is outside the Qwen vocabulary"));
        }

        let mtp = self.mtp.take().ok_or("Qwen MTP is not configured")?;
        let result = (|| {
            let rows = tokens.len() as u32;

            let mut hidden_norm = rms_args(HC_WIDTH, HC_WIDTH);
            hidden_norm.u[4] = rows;
            hidden_norm.u[11] = 1;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                &self.scratch.hc_norm,
                Some(&self.scratch.hc),
                None,
                None,
                &[self.view(self.weight("mtp.pre_fc_norm_hidden.weight")?)],
                &hidden_norm,
                HC * rows,
                1,
            )?;
            self.bf16_batch(
                self.weight("mtp.fc_hidden.weight")?,
                &self.scratch.hc_norm,
                &self.scratch.hc_mix,
                HIDDEN,
                HIDDEN,
                HC * rows,
            )?;

            self.scratch.token_ids.write_i32(tokens)?;
            let embedding =
                self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?;
            let mut embed = args();
            embed.u[0] = HIDDEN;
            embed.u[2] = embedding.bits;
            embed.u[3] = embedding.group;
            embed.u[4] = rows;
            self.dispatch(
                c"kernel_qwen_affine_embedding_batch",
                &self.scratch.hidden,
                Some(&self.scratch.token_ids),
                None,
                None,
                &[
                    self.view(embedding.packed),
                    self.view(embedding.scales),
                    self.view(embedding.biases),
                ],
                &embed,
                HIDDEN,
                rows,
            )?;
            let mut embedding_norm = rms_args(HIDDEN, HIDDEN);
            embedding_norm.u[4] = rows;
            embedding_norm.u[11] = 1;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                &self.scratch.block,
                Some(&self.scratch.hidden),
                None,
                None,
                &[self.view(self.weight("mtp.pre_fc_norm_embedding.weight")?)],
                &embedding_norm,
                rows,
                1,
            )?;
            self.bf16_batch(
                self.weight("mtp.fc_embedding.weight")?,
                &self.scratch.block,
                &self.scratch.hidden,
                HIDDEN,
                HIDDEN,
                rows,
            )?;
            let mut repeat = args();
            repeat.u[0] = HIDDEN;
            repeat.u[4] = rows;
            self.dispatch(
                c"kernel_qwen_repeat4",
                &self.scratch.hc_norm,
                Some(&self.scratch.hidden),
                None,
                None,
                &[],
                &repeat,
                HC_WIDTH * rows,
                1,
            )?;
            let mut add = args();
            add.u[0] = HC_WIDTH * rows;
            self.dispatch(
                c"kernel_qwen_add",
                &self.scratch.hc,
                Some(&self.scratch.hc_mix),
                Some(&self.scratch.hc_norm),
                None,
                &[],
                &add,
                HC_WIDTH * rows,
                1,
            )?;

            let prefix = "mtp.layers.0";
            self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?;
            self.attention_rows(prefix, &mtp.attention, position, rows)?;
            self.hyper_write_rows(rows)?;
            self.encode_moe_rows(prefix, rows)?;
            Ok(())
        })();
        self.mtp = Some(mtp);
        result
    }

    fn mtp_step(
        &mut self,
        token: i32,
        position: u32,
        draft_requested: bool,
        sampling: Option<(f32, f32, f32, i32)>,
    ) -> Result<Option<DraftStep>, String> {
        let mut mtp = self.mtp.take().ok_or("Qwen MTP is not configured")?;
        let started = Instant::now();
        let result = (|| {
            if token < 0 || token as u32 >= VOCAB {
                return Err(format!("token {token} is outside the Qwen vocabulary"));
            }
            if position >= self.context + MTP_TOKEN_RESERVE {
                return Err("Qwen MTP position exceeds its bounded cache".into());
            }

            let embedding =
                self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?;
            let mut embed = args();
            embed.u[0] = HIDDEN;
            embed.u[2] = embedding.bits;
            embed.u[3] = embedding.group;
            embed.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),
                ],
                &embed,
                HIDDEN,
                1,
            )?;
            let mut norm = rms_args(HIDDEN, HIDDEN);
            norm.u[11] = 1;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                &self.scratch.block,
                Some(&self.scratch.hidden),
                None,
                None,
                &[self.view(self.weight("mtp.pre_fc_norm_embedding.weight")?)],
                &norm,
                1,
                1,
            )?;
            self.bf16_mv(
                self.weight("mtp.fc_embedding.weight")?,
                &self.scratch.block,
                &self.scratch.hidden,
                HIDDEN,
                HIDDEN,
            )?;
            norm.u[0] = HC_WIDTH;
            norm.u[1] = HC_WIDTH;
            self.dispatch(
                c"kernel_qwen_zero_rms",
                &self.scratch.hc_norm,
                Some(&self.scratch.hc),
                None,
                None,
                &[self.view(self.weight("mtp.pre_fc_norm_hidden.weight")?)],
                &norm,
                1,
                1,
            )?;
            let hidden_weight = self.weight("mtp.fc_hidden.weight")?;
            self.bf16_mv4(
                hidden_weight,
                &self.scratch.hc_norm,
                &self.scratch.hc_mix,
                HIDDEN,
                HIDDEN,
            )?;
            let mut repeat = args();
            repeat.u[0] = HIDDEN;
            self.dispatch(
                c"kernel_qwen_repeat4",
                &self.scratch.hc_norm,
                Some(&self.scratch.hidden),
                None,
                None,
                &[],
                &repeat,
                HC_WIDTH,
                1,
            )?;
            let mut add = args();
            add.u[0] = HC_WIDTH;
            self.dispatch(
                c"kernel_qwen_add",
                &self.scratch.hc,
                Some(&self.scratch.hc_mix),
                Some(&self.scratch.hc_norm),
                None,
                &[],
                &add,
                HC_WIDTH,
                1,
            )?;
            let prefix = "mtp.layers.0";
            self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?;
            self.attention(prefix, &mtp.attention, position)?;
            self.hyper_write()?;
            self.encode_moe(prefix, "MTP")?;

            if !draft_requested {
                commands.finish()?;
                return Ok(None);
            }
            self.mix_output("mtp.hyper_connection_mixer")?;
            self.affine_mv_into(
                &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?,
                &self.scratch.block,
                &self.scratch.logits,
                HIDDEN,
                VOCAB,
            )?;
            let sparse = if let Some((temperature, _, _, top_k)) = sampling {
                self.encode_sparse_probabilities(&self.scratch.logits, 1, temperature, top_k)?
            } else {
                false
            };
            commands.finish()?;
            mtp.drafted += 1;
            if mtp.timing {
                eprintln!(
                    "ds4: Qwen MTP proposal at position {position} in {:.1} ms",
                    started.elapsed().as_secs_f64() * 1000.0
                );
            }
            let probabilities = if sparse {
                let (_, top_p, min_p, top_k) =
                    sampling.expect("sampling configuration disappeared");
                self.read_sparse_probabilities(1, top_p, min_p, top_k)?
                    .and_then(|mut rows| rows.pop())
            } else {
                None
            };
            let logits = if probabilities.is_some() {
                Vec::new()
            } else {
                self.scratch.logits.read_f32(&mut self.logits)?;
                self.logits.clone()
            };
            Ok(Some((logits, probabilities)))
        })();
        self.mtp = Some(mtp);
        result
    }

    #[cfg(test)]
    fn capture_recurrent(&mut self) -> Result<(), String> {
        let commands = Commands::begin()?;
        let mtp = self.mtp.as_mut().ok_or("Qwen MTP is not configured")?;
        let snapshot = &mut mtp.snapshot;
        for (state, saved) in self.states.iter().zip(&snapshot.gdn) {
            match (state, saved) {
                (LayerState::Gdn { conv, recurrent }, Some((saved_conv, saved_recurrent))) => {
                    saved_conv.copy_from(
                        0,
                        conv,
                        0,
                        u64::from(GDN_QKV) * 3 * 2,
                        "capturing Qwen MTP GDN convolution state",
                    )?;
                    saved_recurrent.copy_from(
                        0,
                        recurrent,
                        0,
                        u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                        "capturing Qwen MTP GDN recurrent state",
                    )?;
                }
                (LayerState::Attention { .. }, None) => {}
                _ => return Err("Qwen MTP capture layout is invalid".into()),
            }
        }
        snapshot.ple_conv.copy_from(
            0,
            &self.ple_state.conv,
            0,
            u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
            "capturing Qwen MTP PLE state",
        )?;
        snapshot.hidden.copy_from(
            0,
            &self.scratch.hc,
            0,
            u64::from(HC_WIDTH) * 4,
            "capturing Qwen MTP target hidden state",
        )?;
        snapshot.ple_history = self.ple_state.history;
        snapshot.logits.copy_from_slice(&self.logits);
        snapshot.position = self.position;
        snapshot.tokens = self.tokens.len();
        commands.finish()
    }

    #[cfg(test)]
    fn restore_recurrent(&mut self) -> Result<(), String> {
        let commands = Commands::begin()?;
        let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?;
        let snapshot = &mtp.snapshot;
        for (state, saved) in self.states.iter().zip(&snapshot.gdn) {
            match (state, saved) {
                (LayerState::Gdn { conv, recurrent }, Some((saved_conv, saved_recurrent))) => {
                    conv.copy_from(
                        0,
                        saved_conv,
                        0,
                        u64::from(GDN_QKV) * 3 * 2,
                        "restoring Qwen MTP GDN convolution state",
                    )?;
                    recurrent.copy_from(
                        0,
                        saved_recurrent,
                        0,
                        u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                        "restoring Qwen MTP GDN recurrent state",
                    )?;
                }
                (LayerState::Attention { .. }, None) => {}
                _ => return Err("Qwen MTP restore layout is invalid".into()),
            }
        }
        self.ple_state.conv.copy_from(
            0,
            &snapshot.ple_conv,
            0,
            u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
            "restoring Qwen MTP PLE state",
        )?;
        self.scratch.hc.copy_from(
            0,
            &snapshot.hidden,
            0,
            u64::from(HC_WIDTH) * 4,
            "restoring Qwen MTP target hidden state",
        )?;
        self.ple_state.history = snapshot.ple_history;
        self.logits.copy_from_slice(&snapshot.logits);
        self.position = snapshot.position;
        self.tokens.truncate(snapshot.tokens);
        commands.finish()
    }

    fn draft_limit(&self, max_tokens: u32, reasoning: ReasoningMode, token: i32) -> u32 {
        if self.mtp.is_none()
            || self.position == 0
            || max_tokens <= 1
            || self.model.is_stop_token_for_reasoning(token, reasoning)
        {
            return 0;
        }
        let limit = max_tokens
            .saturating_sub(1)
            .min(MTP_TOKEN_RESERVE)
            .min(self.context.saturating_sub(self.position + 1));
        #[cfg(test)]
        if let Ok(depth) = std::env::var("DS4SERVER_QWEN38_DRAFT_DEPTH") {
            return limit.min(depth.parse().expect("invalid Qwen test draft depth"));
        }
        limit
    }

    fn draft_greedy(
        &mut self,
        token: i32,
        count: u32,
        reasoning: ReasoningMode,
    ) -> Result<Vec<QwenDraft>, String> {
        let mut drafts = Vec::with_capacity(count as usize);
        let mut input = token;
        let start = self.position - 1;
        for step in 0..count {
            let logits = self
                .mtp_step(input, start + step, true, None)?
                .ok_or("Qwen MTP did not return draft logits")?
                .0;
            let token = argmax(&logits);
            drafts.push(QwenDraft {
                token,
                #[cfg(test)]
                logits,
                probabilities: Vec::new(),
            });
            input = token;
            if self.model.is_stop_token_for_reasoning(token, reasoning) {
                break;
            }
        }
        Ok(drafts)
    }

    #[allow(clippy::too_many_arguments)]
    fn draft_sampled(
        &mut self,
        token: i32,
        count: u32,
        reasoning: ReasoningMode,
        temperature: f32,
        top_p: f32,
        min_p: f32,
        top_k: i32,
        rng: &mut Rng,
    ) -> Result<Vec<QwenDraft>, String> {
        let mut drafts = Vec::with_capacity(count as usize);
        let mut input = token;
        let start = self.position - 1;
        for step in 0..count {
            let (logits, probabilities) = self
                .mtp_step(
                    input,
                    start + step,
                    true,
                    Some((temperature, top_p, min_p, top_k)),
                )?
                .ok_or("Qwen MTP did not return draft logits")?;
            let probabilities = probabilities.unwrap_or_else(|| {
                qwen_sampling_probabilities(&logits, temperature, top_p, min_p, top_k)
            });
            let token = sample_probabilities_f64(&probabilities, rng);
            drafts.push(QwenDraft {
                token,
                #[cfg(test)]
                logits,
                probabilities,
            });
            input = token;
            if self.model.is_stop_token_for_reasoning(token, reasoning) {
                break;
            }
        }
        Ok(drafts)
    }

    fn canonicalize_batch_row(&self, row: usize) -> Result<(), String> {
        if row == 0 {
            return Ok(());
        }
        let commands = Commands::begin()?;
        self.scratch.hc.copy_from(
            0,
            &self.scratch.hc,
            row as u64 * u64::from(HC_WIDTH) * 4,
            u64::from(HC_WIDTH) * 4,
            "selecting the committed Qwen verify hidden row",
        )?;
        commands.finish()
    }

    fn commit_final_mtp_batch_row(
        &mut self,
        drafts: &[QwenDraft],
        batch_start: u32,
    ) -> Result<(), String> {
        let rows = u32::try_from(drafts.len()).map_err(|_| "Qwen MTP draft width overflow")?;
        if rows == 0 {
            return Ok(());
        }
        self.commit_verifier_recurrent_state()?;
        let commands = Commands::begin()?;
        self.scratch.batch_final_hc.copy_from(
            0,
            &self.scratch.hc,
            u64::from(rows) * u64::from(HC_WIDTH) * 4,
            u64::from(HC_WIDTH) * 4,
            "preserving the final Qwen verify hidden row",
        )?;
        self.scratch.hc.copy_from(
            0,
            &self
                .mtp
                .as_ref()
                .ok_or("Qwen MTP is not configured")?
                .verify
                .hidden,
            0,
            u64::from(rows * HC_WIDTH) * 4,
            "restoring authoritative Qwen MTP history rows",
        )?;
        let tokens = drafts.iter().map(|draft| draft.token).collect::<Vec<_>>();
        self.encode_mtp_history_batch(&tokens, batch_start)?;
        self.scratch.hc.copy_from(
            0,
            &self.scratch.batch_final_hc,
            0,
            u64::from(HC_WIDTH) * 4,
            "restoring the final Qwen verify hidden row",
        )?;
        commands.finish()
    }

    fn commit_verifier_recurrent_state(&mut self) -> Result<(), String> {
        let mtp = self.mtp.as_mut().ok_or("Qwen MTP is not configured")?;
        for (state, capture) in self.states.iter_mut().zip(&mut mtp.verify.gdn) {
            match (state, capture) {
                (
                    LayerState::Gdn { conv, recurrent },
                    Some(GdnVerifyCapture {
                        conv: verified_conv,
                        recurrent: verified_recurrent,
                        ..
                    }),
                ) => {
                    std::mem::swap(conv, verified_conv);
                    std::mem::swap(recurrent, verified_recurrent);
                }
                (LayerState::Attention { .. }, None) => {}
                _ => return Err("Qwen verifier GDN state layout is invalid".into()),
            }
        }
        std::mem::swap(&mut self.ple_state.conv, &mut mtp.verify.ple_conv);
        Ok(())
    }

    fn commit_verified_window_from_current(
        &mut self,
        committed: &[i32],
        committed_logits: &[f32],
    ) -> Result<(), String> {
        let keep_rows = committed.len();
        if keep_rows == 0 || keep_rows > MTP_CAPTURE_ROWS {
            return Err("Qwen verifier commit width is invalid".into());
        }
        if committed_logits.len() != VOCAB as usize {
            return Err("Qwen verifier commit logits are invalid".into());
        }
        let batch_start = self.position;

        let rows = keep_rows as u32;
        let commands = Commands::begin()?;
        {
            let ple = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?;
            self.scratch.hc.copy_from(
                0,
                &ple.verify.ple_input,
                0,
                u64::from(rows * HC_WIDTH) * 4,
                "restoring captured Qwen verifier PLE input rows",
            )?;
        }
        if rows == 1 {
            self.ple(committed[0])?;
            self.encode_layer(1)?;
        } else {
            self.ple_batch(committed)?;
            let ple_layer = "language_model.model.layers.1";
            self.hyper_read_rows(&format!("{ple_layer}.attn_hyper_connection"), rows)?;
            self.gdn_rows(ple_layer, 1, rows, false)?;
            self.hyper_write_rows(rows)?;
            self.encode_moe_rows(ple_layer, rows)?;
        }

        for layer in 0..LAYERS {
            if layer == 1 {
                continue;
            }
            let LayerState::Gdn { conv, recurrent } = &self.states[layer] else {
                continue;
            };
            let capture = self
                .mtp
                .as_ref()
                .and_then(|mtp| mtp.verify.gdn.get(layer))
                .and_then(Option::as_ref)
                .ok_or("Qwen verifier GDN capture layout is invalid")?;
            self.scratch.gdn_out.copy_from(
                0,
                &capture.v,
                0,
                u64::from(rows * GDN_VALUE) * 4,
                "restoring captured Qwen verifier GDN value rows",
            )?;
            encode_gdn_delta_rows(
                &self.scratch.gdn_out,
                &capture.q,
                &capture.k,
                &capture.controls,
                recurrent,
                recurrent,
                rows,
            )?;
            let mut commit = args();
            commit.u[0] = GDN_QKV;
            commit.u[4] = rows;
            self.dispatch(
                c"kernel_qwen_gdn_commit_conv_rows",
                conv,
                Some(&capture.qkv),
                None,
                None,
                &[],
                &commit,
                GDN_QKV,
                1,
            )?;
        }

        if rows > 1 {
            let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?;
            self.scratch.hc.copy_from(
                0,
                &mtp.verify.hidden,
                0,
                u64::from((rows - 1) * HC_WIDTH) * 4,
                "restoring authoritative Qwen MTP history rows",
            )?;
            self.encode_mtp_history_batch(&committed[1..], batch_start)?;
        }
        let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?;
        self.scratch.hc.copy_from(
            0,
            &mtp.verify.hidden,
            (u64::from(rows) - 1) * u64::from(HC_WIDTH) * 4,
            u64::from(HC_WIDTH) * 4,
            "selecting the committed Qwen verifier hidden row",
        )?;
        commands.finish()?;

        self.tokens.extend_from_slice(committed);
        self.position = batch_start + rows;
        self.logits.copy_from_slice(committed_logits);
        Ok(())
    }

    #[cfg(test)]
    fn commit_verified_window(
        &mut self,
        keep_rows: usize,
        tokens: &[i32],
        committed_logits: &[f32],
    ) -> Result<(), String> {
        self.restore_recurrent()?;
        let baseline_tokens = self.tokens.len();
        if tokens.len() != baseline_tokens + keep_rows {
            return Err("Qwen verifier commit token span is invalid".into());
        }
        self.commit_verified_window_from_current(&tokens[baseline_tokens..], committed_logits)
    }

    #[allow(clippy::too_many_arguments)]
    fn restore_or_replay(
        &mut self,
        committed: &[i32],
        committed_logits: &[f32],
        batch_start: u32,
        baseline_tokens: usize,
        ple_history: [i32; PLE_HISTORY],
        captures_complete: bool,
    ) -> Result<bool, String> {
        if captures_complete {
            self.position = batch_start;
            self.tokens.truncate(baseline_tokens);
            self.ple_state.history = ple_history;
            match self.commit_verified_window_from_current(committed, committed_logits) {
                Ok(()) => return Ok(true),
                Err(error) => {
                    eprintln!("ds4: Qwen capture-commit refused ({error}); falling back to replay")
                }
            }
        }
        let mut tokens = self.tokens[..baseline_tokens].to_vec();
        tokens.extend_from_slice(committed);
        self.replay_tokens(&tokens)?;
        Ok(false)
    }

    fn replay_tokens(&mut self, tokens: &[i32]) -> Result<(), String> {
        let tag = self.checkpoint_tag;
        let counters = self.mtp.as_ref().map(|mtp| {
            (
                mtp.cycles,
                mtp.drafted,
                mtp.accepted,
                mtp.accepted_depth_total,
                mtp.accepted_depth_max,
                mtp.verifier_passes,
                mtp.verifier_ns,
                mtp.repair_ns,
            )
        });
        self.reset()?;
        for &token in tokens {
            self.eval(token)?;
        }
        self.checkpoint_tag = tag;
        if let (Some(mtp), Some(counters)) = (&mut self.mtp, counters) {
            (
                mtp.cycles,
                mtp.drafted,
                mtp.accepted,
                mtp.accepted_depth_total,
                mtp.accepted_depth_max,
                mtp.verifier_passes,
                mtp.verifier_ns,
                mtp.repair_ns,
            ) = counters;
        }
        Ok(())
    }

    fn note_mtp_verification(&mut self, accepted: usize, started: Instant) {
        if let Some(mtp) = &mut self.mtp {
            mtp.accepted = mtp.accepted.saturating_add(accepted as u64);
            mtp.accepted_depth_total = mtp.accepted_depth_total.saturating_add(accepted as u64);
            mtp.accepted_depth_max = mtp.accepted_depth_max.max(accepted as u64);
            mtp.verifier_passes = mtp.verifier_passes.saturating_add(1);
            mtp.verifier_ns = mtp
                .verifier_ns
                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
        }
    }

    pub(super) fn eval_speculative_greedy(
        &mut self,
        token: i32,
        max_tokens: u32,
        reasoning: ReasoningMode,
        cancelled: &std::sync::atomic::AtomicBool,
    ) -> Result<Vec<i32>, String> {
        if self.mtp.is_none() {
            self.eval_target(token)?;
            return Ok(vec![token]);
        }
        if let Some(mtp) = &mut self.mtp {
            mtp.cycles = mtp.cycles.saturating_add(1);
        }
        let count = self.draft_limit(max_tokens, reasoning, token);
        if count == 0 || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
            if self.position > 0 && !self.model.is_stop_token_for_reasoning(token, reasoning) {
                let _ = self.mtp_step(token, self.position - 1, false, None)?;
            }
            self.eval_target(token)?;
            return Ok(vec![token]);
        }
        let drafts = self.draft_greedy(token, count, reasoning)?;
        if drafts.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
            self.eval_target(token)?;
            return Ok(vec![token]);
        }
        let started = Instant::now();
        let baseline_tokens = self.tokens.len();
        let batch_start = self.position;
        let baseline_ple_history = self.ple_state.history;
        let mut verify_tokens = Vec::with_capacity(drafts.len() + 1);
        verify_tokens.push(token);
        verify_tokens.extend(drafts.iter().map(|draft| draft.token));
        let (verify_logits, _) = self.eval_target_batch(&verify_tokens, None)?;
        let target_logits = &verify_logits[..drafts.len()];
        let accepted = drafts
            .iter()
            .zip(target_logits)
            .take_while(|(draft, target)| argmax(target) == draft.token)
            .count();
        if accepted != drafts.len() {
            let repair = Instant::now();
            let mut committed = Vec::with_capacity(accepted + 1);
            committed.push(token);
            committed.extend(drafts[..accepted].iter().map(|draft| draft.token));
            let _ = self.restore_or_replay(
                &committed,
                &verify_logits[accepted],
                batch_start,
                baseline_tokens,
                baseline_ple_history,
                true,
            )?;
            if let Some(mtp) = &mut self.mtp {
                mtp.repair_ns = mtp
                    .repair_ns
                    .saturating_add(u64::try_from(repair.elapsed().as_nanos()).unwrap_or(u64::MAX));
            }
        } else if drafts.last().is_some_and(|last| {
            !self
                .model
                .is_stop_token_for_reasoning(last.token, reasoning)
        }) {
            let commit = Instant::now();
            self.commit_final_mtp_batch_row(&drafts, batch_start)?;
            self.read_logits_row(drafts.len())?;
            if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() {
                eprintln!(
                    "ds4: Qwen verifier full commit rows={} ms={:.3}",
                    drafts.len(),
                    commit.elapsed().as_secs_f64() * 1_000.0
                );
            }
        } else {
            self.commit_verifier_recurrent_state()?;
            self.canonicalize_batch_row(drafts.len())?;
            self.read_logits_row(drafts.len())?;
        }
        self.note_mtp_verification(accepted, started);
        let mut emitted = Vec::with_capacity(accepted + 1);
        emitted.push(token);
        emitted.extend(drafts[..accepted].iter().map(|draft| draft.token));
        Ok(emitted)
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn eval_speculative_sampled(
        &mut self,
        token: i32,
        max_tokens: u32,
        reasoning: ReasoningMode,
        temperature: f32,
        top_p: f32,
        min_p: f32,
        top_k: i32,
        rng: &mut Rng,
        cancelled: &std::sync::atomic::AtomicBool,
    ) -> Result<Vec<i32>, String> {
        let folded_primary = self.folding_sampled;
        self.folding_sampled = false;
        if self.mtp.is_none() {
            self.eval_target(token)?;
            return Ok(if folded_primary {
                Vec::new()
            } else {
                vec![token]
            });
        }
        if let Some(mtp) = &mut self.mtp {
            mtp.cycles = mtp.cycles.saturating_add(1);
        }
        let count = self.draft_limit(
            max_tokens.saturating_add(u32::from(folded_primary)),
            reasoning,
            token,
        );
        if count == 0 || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
            if self.position > 0 && !self.model.is_stop_token_for_reasoning(token, reasoning) {
                let _ = self.mtp_step(token, self.position - 1, false, None)?;
            }
            self.eval_target(token)?;
            return Ok(if folded_primary {
                Vec::new()
            } else {
                vec![token]
            });
        }
        let drafts = self.draft_sampled(
            token,
            count,
            reasoning,
            temperature,
            top_p,
            min_p,
            top_k,
            rng,
        )?;
        if drafts.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
            self.eval_target(token)?;
            return Ok(if folded_primary {
                Vec::new()
            } else {
                vec![token]
            });
        }
        let started = Instant::now();
        let baseline_tokens = self.tokens.len();
        let batch_start = self.position;
        let baseline_ple_history = self.ple_state.history;
        let mut verify_tokens = Vec::with_capacity(drafts.len() + 1);
        verify_tokens.push(token);
        verify_tokens.extend(drafts.iter().map(|draft| draft.token));
        let (verify_logits, verify_probabilities) =
            self.eval_target_batch(&verify_tokens, Some((temperature, top_p, min_p, top_k)))?;
        let verify_probabilities = verify_probabilities.unwrap_or_else(|| {
            verify_logits
                .iter()
                .map(|logits| qwen_sampling_probabilities(logits, temperature, top_p, min_p, top_k))
                .collect()
        });
        let mut accepted = 0;
        let mut replacement = None;
        for ((draft, target), index) in drafts.iter().zip(&verify_probabilities).zip(0..) {
            let (sampled, was_draft) = exact_speculative_sample_from_probabilities(
                target,
                &draft.probabilities,
                draft.token,
                rng,
            );
            if !was_draft {
                replacement = Some((index, sampled));
                break;
            }
            accepted += 1;
        }
        if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() {
            eprintln!(
                "ds4: Qwen MTP cycle primary={token} folded={folded_primary} drafts={:?} accepted={accepted} replacement={replacement:?}",
                drafts.iter().map(|draft| draft.token).collect::<Vec<_>>()
            );
        }
        let mut emitted = Vec::with_capacity(accepted + 2);
        emitted.push(token);
        emitted.extend(drafts[..accepted].iter().map(|draft| draft.token));
        if let Some((slot, replacement)) = replacement {
            let repair = Instant::now();
            let mut committed = Vec::with_capacity(accepted + 1);
            committed.push(token);
            committed.extend(drafts[..accepted].iter().map(|draft| draft.token));
            let captured = if verify_logits.is_empty() {
                let committed_logits = self.logits_row(slot)?;
                self.restore_or_replay(
                    &committed,
                    &committed_logits,
                    batch_start,
                    baseline_tokens,
                    baseline_ple_history,
                    true,
                )?
            } else {
                self.restore_or_replay(
                    &committed,
                    &verify_logits[slot],
                    batch_start,
                    baseline_tokens,
                    baseline_ple_history,
                    true,
                )?
            };
            if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() {
                eprintln!(
                    "ds4: Qwen verifier commit rows={} ms={:.3}",
                    slot + 1,
                    repair.elapsed().as_secs_f64() * 1_000.0
                );
            }
            emitted.push(replacement);
            let newly_emitted = emitted.len() - usize::from(folded_primary);
            if captured
                && newly_emitted < max_tokens as usize
                && !self
                    .model
                    .is_stop_token_for_reasoning(replacement, reasoning)
            {
                self.pending_sampled = Some(replacement);
            } else {
                if !self
                    .model
                    .is_stop_token_for_reasoning(replacement, reasoning)
                {
                    let _ = self.mtp_step(replacement, self.position - 1, false, None)?;
                }
                self.eval_target(replacement)?;
            }
            if let Some(mtp) = &mut self.mtp {
                mtp.repair_ns = mtp
                    .repair_ns
                    .saturating_add(u64::try_from(repair.elapsed().as_nanos()).unwrap_or(u64::MAX));
            }
        } else if drafts.last().is_some_and(|last| {
            !self
                .model
                .is_stop_token_for_reasoning(last.token, reasoning)
        }) {
            let commit = Instant::now();
            self.commit_final_mtp_batch_row(&drafts, batch_start)?;
            self.read_logits_row(drafts.len())?;
            if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() {
                eprintln!(
                    "ds4: Qwen verifier full commit rows={} ms={:.3}",
                    drafts.len(),
                    commit.elapsed().as_secs_f64() * 1_000.0
                );
            }
            let newly_emitted = emitted.len() - usize::from(folded_primary);
            if newly_emitted < max_tokens as usize {
                let bonus = sample_probabilities_f64(&verify_probabilities[drafts.len()], rng);
                emitted.push(bonus);
                self.pending_sampled = Some(bonus);
            }
        } else {
            self.commit_verifier_recurrent_state()?;
            self.canonicalize_batch_row(drafts.len())?;
            self.read_logits_row(drafts.len())?;
        }
        self.note_mtp_verification(accepted, started);
        if folded_primary {
            debug_assert_eq!(emitted.first(), Some(&token));
            emitted.remove(0);
        }
        Ok(emitted)
    }

    pub(super) fn prefill(
        &mut self,
        tokens: &[i32],
        mut progress: impl FnMut(u32) -> bool,
    ) -> Result<usize, String> {
        let end = self
            .position
            .checked_add(u32::try_from(tokens.len()).map_err(|_| "Qwen prefill is too large")?)
            .ok_or("Qwen prefill position overflow")?;
        if end > self.context {
            return Err(format!(
                "the Qwen executor supports {} tokens per session",
                self.context
            ));
        }
        if tokens.len() <= MTP_CAPTURE_ROWS {
            let mut completed = 0;
            for &token in tokens {
                if !progress(completed) {
                    break;
                }
                self.eval(token)?;
                completed += 1;
            }
            progress(completed);
            return Ok(completed as usize);
        }

        let has_mtp = self.mtp.is_some();
        if self.position > 0 && has_mtp {
            let _ = self.mtp_step(tokens[0], self.position - 1, false, None)?;
        }
        let body = tokens.len() - 1;
        let capacity = body.clamp(1, self.prefill_chunk) as u32;
        let batch_scratch = Scratch::new(self.context, capacity)?;
        let decode_scratch = std::mem::replace(&mut self.scratch, batch_scratch);
        let mut completed = 0;
        let result = (|| {
            while completed < body {
                if !progress(completed as u32) {
                    break;
                }
                let rows = (body - completed).min(self.prefill_chunk);
                self.eval_prefill_batch(
                    &tokens[completed..completed + rows],
                    has_mtp.then_some(&tokens[completed + 1..completed + rows + 1]),
                )?;
                completed += rows;
            }
            if completed == body && progress(completed as u32) {
                self.eval_target(tokens[body])?;
                completed += 1;
                let commands = Commands::begin()?;
                self.scratch.batch_final_hc.copy_from(
                    0,
                    &self.scratch.hc,
                    0,
                    u64::from(HC_WIDTH) * 4,
                    "preserving the final scalar Qwen prefill hidden row",
                )?;
                commands.finish()?;
            }
            Ok::<(), String>(())
        })();
        let batch_scratch = std::mem::replace(&mut self.scratch, decode_scratch);
        if completed != 0 {
            let commands = Commands::begin()?;
            self.scratch.hc.copy_from(
                0,
                &batch_scratch.batch_final_hc,
                0,
                u64::from(HC_WIDTH) * 4,
                "restoring the final Qwen prefill hidden row",
            )?;
            commands.finish()?;
        }
        progress(completed as u32);
        result.map(|()| completed)
    }

    pub(super) fn logits(&self) -> &[f32] {
        &self.logits
    }
    pub(super) fn take_pending_sampled_token(&mut self) -> Option<i32> {
        let pending = self.pending_sampled.take();
        self.folding_sampled = pending.is_some();
        pending
    }

    pub(super) fn finalize_pending_sampled_token(&mut self) -> Result<(), String> {
        self.folding_sampled = false;
        let Some(token) = self.pending_sampled.take() else {
            return Ok(());
        };
        let result = (|| {
            if self.position > 0 {
                let _ = self.mtp_step(token, self.position - 1, false, None)?;
            }
            self.eval_target(token)
        })();
        if result.is_err() {
            self.pending_sampled = Some(token);
        }
        result
    }

    pub(super) fn execution_stats(&self) -> ExecutionStats {
        let ple = self.model.ple_stats();
        self.mtp.as_ref().map_or_else(
            || ExecutionStats {
                ssd_enabled: true,
                ssd_cache_bytes: ple.cache_bytes,
                ssd_cache_entries: ple.entries,
                ssd_cache_hits: ple.hits,
                ssd_cache_misses: ple.misses,
                ssd_cache_evictions: ple.evictions,
                ssd_pread_bytes: ple.read_bytes,
                ssd_pread_ms: ple.read_ms,
                ssd_selected_requests: ple.hits.saturating_add(ple.misses),
                ssd_requested_bytes: ple
                    .hits
                    .saturating_add(ple.misses)
                    .saturating_mul(PLE_ROW_BYTES as u64),
                ..ExecutionStats::default()
            },
            |mtp| ExecutionStats {
                speculative_mode: 4,
                speculative_cycles: mtp.cycles,
                drafted_tokens: mtp.drafted,
                accepted_draft_tokens: mtp.accepted,
                accepted_depth_total: mtp.accepted_depth_total,
                accepted_depth_max: mtp.accepted_depth_max,
                verifier_passes: mtp.verifier_passes,
                verifier_ms: mtp.verifier_ns / 1_000_000,
                repair_ms: mtp.repair_ns / 1_000_000,
                ssd_enabled: true,
                ssd_cache_bytes: ple.cache_bytes,
                ssd_cache_entries: ple.entries,
                ssd_cache_hits: ple.hits,
                ssd_cache_misses: ple.misses,
                ssd_cache_evictions: ple.evictions,
                ssd_pread_bytes: ple.read_bytes,
                ssd_pread_ms: ple.read_ms,
                ssd_selected_requests: ple.hits.saturating_add(ple.misses),
                ssd_requested_bytes: ple
                    .hits
                    .saturating_add(ple.misses)
                    .saturating_mul(PLE_ROW_BYTES as u64),
                ..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()?;
        if let Some(mtp) = &self.mtp {
            self.mtp = Some(allocate_mtp(self.context, mtp.timing)?);
        }
        self.logits.fill(0.0);
        self.tokens.clear();
        self.position = 0;
        self.pending_sampled = None;
        self.folding_sampled = false;
        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()?,
            hidden: Buffer::floats(VERIFY_ROWS * u64::from(HC_WIDTH))?,
            logits: vec![0.0; VOCAB as usize],
            tokens: Vec::new(),
            position: 0,
            checkpoint_tag: [0; 32],
            mtp: self
                .mtp
                .as_ref()
                .map(|mtp| allocate_mtp(self.context, mtp.timing))
                .transpose()?,
            pending_sampled: None,
            folding_sampled: false,
        })
    }

    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.scratch.hc, &mut incoming.hidden);
        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);
        std::mem::swap(&mut self.mtp, &mut incoming.mtp);
        std::mem::swap(&mut self.pending_sampled, &mut incoming.pending_sampled);
        std::mem::swap(&mut self.folding_sampled, &mut incoming.folding_sampled);
        *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,
            u32::from(self.mtp.is_some()),
        ] {
            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.scratch.hc,
            0,
            u64::from(HC_WIDTH) * 4,
            &mut chunk,
            progress,
        )?;
        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,
                    qsa_raw,
                    qsa_pooled,
                } => {
                    let [kv_bytes, raw_bytes, pooled_bytes] =
                        attention_checkpoint_bytes(self.position);
                    write_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?;
                    write_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?;
                    write_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &mut chunk, progress)?;
                }
            }
        }
        if let Some(mtp) = &self.mtp {
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &mtp.attention
            else {
                return Err("Qwen MTP checkpoint state is invalid".into());
            };
            let [kv_bytes, raw_bytes, pooled_bytes] =
                attention_checkpoint_bytes(self.position.saturating_sub(1));
            write_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?;
            write_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?;
            write_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &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 checkpoint_mtp = read_u32(&mut file)?;
        if checkpoint_mtp > 1 || (checkpoint_mtp != 0) != self.mtp.is_some() {
            return Err("Qwen checkpoint MTP configuration does not match the executor".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.scratch.hc,
            0,
            u64::from(HC_WIDTH) * 4,
            &mut chunk,
            progress,
        )?;
        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,
                    qsa_raw,
                    qsa_pooled,
                } => {
                    let [kv_bytes, raw_bytes, pooled_bytes] = attention_checkpoint_bytes(position);
                    read_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?;
                    read_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?;
                    read_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &mut chunk, progress)?;
                }
            }
        }
        if let Some(mtp) = &self.mtp {
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &mtp.attention
            else {
                return Err("Qwen MTP checkpoint state is invalid".into());
            };
            let [kv_bytes, raw_bytes, pooled_bytes] =
                attention_checkpoint_bytes(position.saturating_sub(1));
            read_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?;
            read_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?;
            read_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &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)
    }
}

#[allow(clippy::too_many_arguments)]
fn encode_gdn_conv_norm_rows(
    input: &Buffer,
    q: &Buffer,
    k: &Buffer,
    v: &Buffer,
    state: &Buffer,
    state_out: &Buffer,
    weight: QwenWeightView,
    rows: u32,
) -> Result<(), String> {
    let state_bytes = u64::from(GDN_QKV) * 3 * 2;
    // Reuse the existing GDN scratch tail. Parallel rows must all see the old
    // convolution state, even when the caller commits back into that state.
    let staged = (rows > 6)
        .then(|| v.view(u64::from(rows * GDN_VALUE) * 4, state_bytes))
        .transpose()?;
    let mut params = args();
    params.u[4] = rows;
    params.u[12] = 1024;
    dispatch_qwen_with_d(
        c"kernel_qwen_gdn_conv_norm_rows",
        q,
        Some(input),
        Some(k),
        Some(v),
        Some(staged.as_ref().unwrap_or(state_out)),
        &[buffer_weight_view(state, state_bytes), weight],
        &params,
        GDN_QKV / 1024,
        if rows > 6 { rows } else { 1 },
    )?;
    if let Some(staged) = staged {
        state_out.copy_from(
            0,
            &staged,
            0,
            state_bytes,
            "committing Qwen prefill convolution state",
        )?;
    }
    Ok(())
}

fn encode_gdn_delta_rows(
    values: &Buffer,
    q: &Buffer,
    k: &Buffer,
    controls: &Buffer,
    state: &Buffer,
    state_out: &Buffer,
    rows: u32,
) -> Result<(), String> {
    let mut params = args();
    params.u[0] = GDN_VALUE;
    params.u[1] = GDN_VALUE + GDN_HEADS_V;
    params.u[4] = rows;
    params.u[12] = 64;
    dispatch_qwen_with_d(
        c"kernel_qwen_gdn_delta_rows",
        values,
        Some(q),
        Some(k),
        Some(controls),
        Some(state_out),
        &[buffer_weight_view(
            state,
            u64::from(GDN_HEADS_V * HEAD_DIM * HEAD_DIM) * 4,
        )],
        &params,
        HEAD_DIM / 16,
        GDN_HEADS_V,
    )
}

fn allocate_states(context: u32) -> Result<Vec<LayerState>, String> {
    let token_capacity = context
        .checked_add(MTP_TOKEN_RESERVE)
        .ok_or_else(|| "Qwen attention capacity overflows".to_owned())?;
    let block_capacity = qsa_block_capacity(context)?;
    (0..LAYERS)
        .map(|layer| {
            if layer % 4 == 3 {
                Ok(LayerState::Attention {
                    kv: Buffer::bytes(u64::from(token_capacity) * ATTN_KV_WIDTH as u64 * 4)?,
                    qsa_raw: Buffer::bytes(u64::from(token_capacity) * QSA_DIM as u64 * 2)?,
                    qsa_pooled: Buffer::bytes(u64::from(block_capacity) * QSA_DIM as u64 * 2)?,
                })
            } 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_attention_state(context: u32) -> Result<LayerState, String> {
    let token_capacity = context
        .checked_add(MTP_TOKEN_RESERVE)
        .ok_or_else(|| "Qwen MTP attention capacity overflows".to_owned())?;
    Ok(LayerState::Attention {
        kv: Buffer::bytes(u64::from(token_capacity) * ATTN_KV_WIDTH as u64 * 4)?,
        qsa_raw: Buffer::bytes(u64::from(token_capacity) * QSA_DIM as u64 * 2)?,
        qsa_pooled: Buffer::bytes(u64::from(qsa_block_capacity(context)?) * QSA_DIM as u64 * 2)?,
    })
}

#[cfg(test)]
fn allocate_recurrent_snapshot() -> Result<RecurrentSnapshot, String> {
    let gdn = (0..LAYERS)
        .map(|layer| {
            if layer % 4 == 3 {
                Ok(None)
            } else {
                Ok(Some((
                    Buffer::bytes(u64::from(GDN_QKV) * 3 * 2)?,
                    Buffer::bytes(u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4)?,
                )))
            }
        })
        .collect::<Result<_, String>>()?;
    Ok(RecurrentSnapshot {
        gdn,
        ple_conv: Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?,
        ple_history: [0; PLE_HISTORY],
        hidden: Buffer::floats(HC_WIDTH.into())?,
        logits: vec![0.0; VOCAB as usize],
        position: 0,
        tokens: 0,
    })
}

fn allocate_verify_capture() -> Result<VerifyCapture, String> {
    let gdn = (0..LAYERS)
        .map(|layer| {
            if layer % 4 == 3 {
                Ok(None)
            } else {
                Ok(Some(GdnVerifyCapture {
                    conv: Buffer::bytes(u64::from(GDN_QKV) * 3 * 2)?,
                    recurrent: Buffer::bytes(
                        u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                    )?,
                    qkv: Buffer::floats(VERIFY_ROWS * u64::from(GDN_QKV))?,
                    q: Buffer::floats(VERIFY_ROWS * u64::from(GDN_HEADS_K * HEAD_DIM))?,
                    k: Buffer::floats(VERIFY_ROWS * u64::from(GDN_HEADS_K * HEAD_DIM))?,
                    v: Buffer::floats(VERIFY_ROWS * u64::from(GDN_VALUE))?,
                    controls: Buffer::floats(VERIFY_ROWS * u64::from(GDN_CONTROLS))?,
                }))
            }
        })
        .collect::<Result<_, String>>()?;
    Ok(VerifyCapture {
        gdn,
        ple_conv: Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?,
        ple_input: Buffer::floats(VERIFY_ROWS * u64::from(HC_WIDTH))?,
        hidden: Buffer::floats(VERIFY_ROWS * u64::from(HC_WIDTH))?,
    })
}

fn allocate_mtp(context: u32, timing: bool) -> Result<QwenMtp, String> {
    Ok(QwenMtp {
        attention: allocate_attention_state(context)?,
        #[cfg(test)]
        snapshot: allocate_recurrent_snapshot()?,
        verify: allocate_verify_capture()?,
        timing,
        cycles: 0,
        drafted: 0,
        accepted: 0,
        accepted_depth_total: 0,
        accepted_depth_max: 0,
        verifier_passes: 0,
        verifier_ns: 0,
        repair_ns: 0,
    })
}

fn qsa_block_capacity(context: u32) -> Result<u32, String> {
    context
        .checked_add(MTP_TOKEN_RESERVE)
        .and_then(|tokens| tokens.checked_add(QSA_RATIO - 1))
        .map(|tokens| tokens / QSA_RATIO)
        .ok_or_else(|| "Qwen QSA block capacity overflows".to_owned())
}

fn attention_checkpoint_bytes(position: u32) -> [u64; 3] {
    [
        u64::from(position) * ATTN_KV_WIDTH as u64 * 4,
        u64::from(position) * QSA_DIM as u64 * 2,
        u64::from(position / QSA_RATIO) * QSA_DIM as u64 * 2,
    ]
}

#[cfg(test)]
fn reference_qsa_selection(scores: &[f32]) -> Vec<u32> {
    let mut blocks = (0..scores.len() as u32).collect::<Vec<_>>();
    blocks.sort_by(|&left, &right| {
        scores[right as usize]
            .total_cmp(&scores[left as usize])
            .then(left.cmp(&right))
    });
    blocks.truncate(QSA_TOP_K.min(blocks.len() as u32) as usize);
    blocks
}

#[cfg(test)]
fn reference_qsa_tokens(selected: &[u32], tokens: u32) -> Vec<u32> {
    let complete = tokens / QSA_RATIO;
    let mut visible = selected
        .iter()
        .flat_map(|&block| block * QSA_RATIO..(block + 1) * QSA_RATIO)
        .collect::<Vec<_>>();
    visible.extend(complete * QSA_RATIO..tokens);
    visible.sort_unstable();
    visible
}

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]
    }
}

#[cfg(test)]
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)
}

#[cfg(test)]
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()
}

fn qwen_feature_enabled(name: &str) -> bool {
    std::env::var(name).map_or(true, |value| {
        matches!(
            value.trim().to_ascii_lowercase().as_str(),
            "1" | "true" | "yes" | "on"
        )
    })
}

fn hyper_pack_bytes(rows: u32, columns: u32) -> u64 {
    let values = u64::from(rows) * u64::from(columns);
    let groups = values / u64::from(HYPER_QUANT_GROUP);
    values + groups * 4
}

fn rms_args(width: u32, group_size: u32) -> QwenKernelArgs {
    let mut values = args();
    values.u[0] = width;
    values.u[1] = group_size;
    values.u[10] = 33 * 4;
    values.u[12] = if group_size <= 4096 {
        group_size.div_ceil(4).div_ceil(32) * 32
    } else {
        1024
    };
    values.f[0] = 1.0e-6;
    values
}

#[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> {
    dispatch_qwen_with_d(kernel, out, a, b, c, None, weights, values, grid_x, grid_y)
}

fn buffer_weight_view(buffer: &Buffer, bytes: u64) -> QwenWeightView {
    QwenWeightView {
        tensor: buffer.raw().cast_const(),
        map: std::ptr::null(),
        size: bytes,
        offset: 0,
        bytes,
    }
}

#[allow(clippy::too_many_arguments)]
fn dispatch_qwen_with_d(
    kernel: &CStr,
    out: &Buffer,
    a: Option<&Buffer>,
    b: Option<&Buffer>,
    c: Option<&Buffer>,
    d: 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()),
                d.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 peak_rss_bytes() -> i64 {
        let mut usage = std::mem::MaybeUninit::<libc::rusage>::uninit();
        // SAFETY: getrusage initializes the supplied process-local output structure.
        if unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) } == 0 {
            // SAFETY: a successful getrusage call initialized the complete structure.
            unsafe { usage.assume_init().ru_maxrss }
        } else {
            -1
        }
    }

    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]
    fn qsa_selection_and_capacity_cover_the_full_native_window() {
        assert_eq!((QSA_HEADS, QSA_KV_HEADS, QSA_DIM), (4, 1, 128));
        assert_eq!(qsa_block_capacity(1).unwrap(), 1);
        assert_eq!(qsa_block_capacity(2_048).unwrap(), 513);
        assert_eq!(qsa_block_capacity(262_144).unwrap(), 65_537);

        let mut scores = (0..600).map(|index| index as f32).collect::<Vec<_>>();
        scores[10] = 10_000.0;
        scores[11] = 10_000.0;
        let selected = reference_qsa_selection(&scores);
        assert_eq!(selected.len(), QSA_TOP_K as usize);
        assert_eq!(&selected[..2], &[10, 11]);
        let visible = reference_qsa_tokens(&selected, 2_403);
        assert_eq!(&visible[visible.len() - 3..], &[2_400, 2_401, 2_402]);
        assert_eq!(visible.len(), QSA_TOP_K as usize * QSA_RATIO as usize + 3);
        assert!(visible.windows(2).all(|pair| pair[0] < pair[1]));

        for depth in [2_048, 16_384, 65_536, 131_072, 262_144] {
            let complete = depth / QSA_RATIO;
            let tail = depth % QSA_RATIO;
            assert_eq!(complete * QSA_RATIO + tail, depth);
            assert!(complete <= qsa_block_capacity(depth).unwrap());
            assert_eq!(
                attention_checkpoint_bytes(depth),
                [
                    u64::from(depth) * ATTN_KV_WIDTH as u64 * 4,
                    u64::from(depth) * QSA_DIM as u64 * 2,
                    u64::from(complete) * QSA_DIM as u64 * 2,
                ]
            );
        }
    }

    #[test]
    fn dense_attention_crossover_is_phase_specific() {
        assert_eq!(dense_attention_budget(1), DECODE_DENSE_BUDGET);
        assert_eq!(
            dense_attention_budget(VERIFY_ROWS as u32),
            DECODE_DENSE_BUDGET
        );
        assert_eq!(
            dense_attention_budget(VERIFY_ROWS as u32 + 1),
            PREFILL_DENSE_BUDGET
        );
        assert!(packed_dense_qsa(VERIFY_ROWS as u32, 8_192));
        assert!(packed_dense_qsa(1, 8_192));
        assert!(!packed_dense_qsa(VERIFY_ROWS as u32 + 1, 8_192));
        assert!(!packed_dense_qsa(
            VERIFY_ROWS as u32,
            QSA_VERIFY_DENSE_LIMIT
        ));
    }

    #[test]
    #[ignore = "requires Apple Metal"]
    fn qwen_gdn_conv_norm_rows_matches_mtplx() {
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let root = PathBuf::from(std::env::var_os("DS4SERVER_GDN_REFERENCE").expect(
            "generate fixtures with tools/qwen-gdn-reference.py and set DS4SERVER_GDN_REFERENCE",
        ));
        let read = |name: &str| {
            fs::read(root.join(format!("{name}.f32")))
                .unwrap()
                .chunks_exact(4)
                .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap()))
                .collect::<Vec<_>>()
        };
        let upload_bf16 = |values: &[f32]| {
            let bytes = values
                .iter()
                .flat_map(|x| bf16(*x).to_le_bytes())
                .collect::<Vec<_>>();
            let buffer = Buffer::bytes(bytes.len() as u64).unwrap();
            buffer.write(0, &bytes).unwrap();
            buffer
        };
        let state = upload_bf16(&read("state"));
        let weights = upload_bf16(&read("weight"));
        let mut failures = Vec::new();
        for rows in [1_u32, 2, 3, 4, 5, 6, 7, 32, 2048] {
            let input = Buffer::floats(u64::from(rows * GDN_QKV)).unwrap();
            input.write_f32(&read(&format!("{rows}-input"))).unwrap();
            let q = Buffer::floats(u64::from(rows * 2048)).unwrap();
            let k = Buffer::floats(u64::from(rows * 2048)).unwrap();
            let v = Buffer::bytes(u64::from(rows * GDN_VALUE) * 4 + u64::from(GDN_QKV) * 3 * 2)
                .unwrap();
            let state_out = Buffer::bytes(u64::from(GDN_QKV) * 3 * 2).unwrap();
            let mut wall_ms = Vec::new();
            for run in 0..11 {
                let started = Instant::now();
                let commands = Commands::begin().unwrap();
                encode_gdn_conv_norm_rows(
                    &input,
                    &q,
                    &k,
                    &v,
                    &state,
                    &state_out,
                    buffer_weight_view(&weights, u64::from(GDN_QKV) * 4 * 2),
                    rows,
                )
                .unwrap();
                commands.finish().unwrap();
                if run > 0 {
                    wall_ms.push(started.elapsed().as_secs_f64() * 1000.0);
                }
            }
            // Production prefills commit into the input state. Exercise that
            // alias as well; updating it inside a parallel kernel would race.
            let commands = Commands::begin().unwrap();
            state_out
                .copy_from(
                    0,
                    &state,
                    0,
                    u64::from(GDN_QKV) * 3 * 2,
                    "resetting test convolution state",
                )
                .unwrap();
            encode_gdn_conv_norm_rows(
                &input,
                &q,
                &k,
                &v,
                &state_out,
                &state_out,
                buffer_weight_view(&weights, u64::from(GDN_QKV) * 4 * 2),
                rows,
            )
            .unwrap();
            commands.finish().unwrap();
            for (name, buffer, width) in [("q", &q, 2048), ("k", &k, 2048), ("v", &v, GDN_VALUE)] {
                let mut actual = vec![0.0; (rows * width) as usize];
                buffer.read_f32(&mut actual).unwrap();
                let expected = read(&format!("{rows}-{name}"));
                assert_eq!(actual.len(), expected.len());
                let mismatches = actual.iter().zip(&expected).filter(|(a, b)| a != b).count();
                let max_error = actual
                    .iter()
                    .zip(&expected)
                    .map(|(a, b)| (a - b).abs())
                    .fold(0.0_f32, f32::max);
                eprintln!(
                    "GDN_REFERENCE rows={rows} tensor={name} mismatches={mismatches} max_abs={max_error}"
                );
                if mismatches > 0 {
                    failures.push((rows, name, mismatches, max_error));
                }
            }
            let mut bytes = vec![0; GDN_QKV as usize * 3 * 2];
            state_out.read(0, &mut bytes).unwrap();
            let expected = read(&format!("{rows}-state"))
                .iter()
                .flat_map(|x| bf16(*x).to_le_bytes())
                .collect::<Vec<_>>();
            assert_eq!(bytes, expected, "state rows={rows}");
            eprintln!("GDN_TIMING rows={rows} wall_ms={wall_ms:?}");
        }
        assert!(failures.is_empty(), "GDN differs from MTPLX: {failures:?}");
    }

    #[test]
    #[ignore = "requires Apple Metal"]
    fn qwen_gdn_delta_rows_matches_mtplx() {
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let root = PathBuf::from(
            std::env::var_os("DS4SERVER_GDN_REFERENCE")
                .expect("generate tools/qwen-gdn-reference.py --delta fixtures first"),
        );
        let read = |name: &str| {
            let bytes = fs::read(root.join(format!("delta-{name}.f32"))).unwrap();
            assert_eq!(bytes.len() % 4, 0);
            bytes
                .chunks_exact(4)
                .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap()))
                .collect::<Vec<_>>()
        };
        let upload = |values: &[f32]| {
            let buffer = Buffer::floats(values.len() as u64).unwrap();
            buffer.write_f32(values).unwrap();
            buffer
        };
        let initial = read("initial");
        let state = upload(&initial);
        let state_out = Buffer::floats(initial.len() as u64).unwrap();
        let mut failures = Vec::new();
        for rows in [1_u32, 2, 3, 4, 5, 6, 7, 32, 2048] {
            let q = upload(&read(&format!("{rows}-q")));
            let k = upload(&read(&format!("{rows}-k")));
            let v = read(&format!("{rows}-v"));
            let values = upload(&v);
            let g = read(&format!("{rows}-g"));
            let beta = read(&format!("{rows}-beta"));
            let mut controls = vec![0.0; (rows * GDN_CONTROLS) as usize];
            for row in 0..rows as usize {
                let offset = row * GDN_CONTROLS as usize + GDN_VALUE as usize;
                let source = row * GDN_HEADS_V as usize;
                controls[offset..offset + GDN_HEADS_V as usize]
                    .copy_from_slice(&beta[source..source + GDN_HEADS_V as usize]);
                controls[offset + GDN_HEADS_V as usize..offset + 2 * GDN_HEADS_V as usize]
                    .copy_from_slice(&g[source..source + GDN_HEADS_V as usize]);
            }
            let controls = upload(&controls);
            let mut wall_ms = Vec::new();
            for run in 0..11 {
                values.write_f32(&v).unwrap();
                let started = Instant::now();
                let commands = Commands::begin().unwrap();
                encode_gdn_delta_rows(&values, &q, &k, &controls, &state, &state_out, rows)
                    .unwrap();
                commands.finish().unwrap();
                if run > 0 {
                    wall_ms.push(started.elapsed().as_secs_f64() * 1000.0);
                }
            }
            let mut check = |stage: &str| {
                for (name, buffer, count) in [
                    ("out", &values, v.len()),
                    ("state", &state_out, initial.len()),
                ] {
                    let mut actual = vec![0.0; count];
                    buffer.read_f32(&mut actual).unwrap();
                    let expected = read(&format!("{rows}-{name}"));
                    assert_eq!(actual.len(), expected.len());
                    let mismatches = actual.iter().zip(&expected).filter(|(a, b)| a != b).count();
                    let max_error = actual
                        .iter()
                        .zip(&expected)
                        .map(|(a, b)| (a - b).abs())
                        .fold(0.0_f32, f32::max);
                    eprintln!(
                        "GDN_DELTA rows={rows} stage={stage} tensor={name} mismatches={mismatches} max_abs={max_error}"
                    );
                    if mismatches > 0 {
                        failures.push((rows, name, mismatches, max_error));
                    }
                }
            };
            check("whole");
            // Check continuation and production's in-place state commit too.
            // Splitting the token sequence must preserve the same recurrence.
            values.write_f32(&v).unwrap();
            state_out.write_f32(&initial).unwrap();
            let commands = Commands::begin().unwrap();
            let split = rows.div_ceil(2);
            for (start, count) in [(0, split), (split, rows - split)] {
                if count == 0 {
                    continue;
                }
                let view = |buffer: &Buffer, width: u32| {
                    buffer
                        .view(u64::from(start * width) * 4, u64::from(count * width) * 4)
                        .unwrap()
                };
                encode_gdn_delta_rows(
                    &view(&values, GDN_VALUE),
                    &view(&q, 2048),
                    &view(&k, 2048),
                    &view(&controls, GDN_CONTROLS),
                    &state_out,
                    &state_out,
                    count,
                )
                .unwrap();
            }
            commands.finish().unwrap();
            check("split_in_place");
            eprintln!("GDN_DELTA_TIMING rows={rows} wall_ms={wall_ms:?}");
        }
        assert!(
            failures.is_empty(),
            "GDN delta differs from MTPLX: {failures:?}"
        );
    }

    #[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, 1.0, 1.0, 1.0, 1.0,
        ] {
            bytes.extend_from_slice(&bf16(value).to_le_bytes());
        }
        let qsa_weight_offset = bytes.len() as u64;
        bytes.extend_from_slice(&bf16(1.0).to_le_bytes().repeat(QSA_DIM as usize));
        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 {
            tensor: std::ptr::null(),
            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],
            f32::from_bits(u32::from(bf16(8.0 / (1.0 + (-8.0_f32).exp()))) << 16),
        );
        assert_eq!(
            {
                let mut state = [0; 6];
                conv_state.read(0, &mut state).unwrap();
                state
            },
            [0, 0, 0, 0, 0, 64]
        );

        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 mut softmax = args();
        softmax.u[0] = EXPERTS;
        softmax.u[12] = EXPERTS / 4;
        dispatch_qwen(
            c"kernel_qwen_softmax_precise_512",
            &router,
            Some(&router),
            None,
            None,
            &[],
            &softmax,
            1,
            1,
        )
        .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;
        route.u[12] = 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, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        let mut probabilities = [0.0; EXPERTS_USED];
        weights.read_f32(&mut probabilities).unwrap();
        assert_eq!(
            probabilities,
            [
                7.8201294e-5,
                0.0002117157,
                0.0005760193,
                0.0015640259,
                0.004272461,
                0.01159668,
                0.03149414,
                0.08544922,
                0.23242188,
                0.6328125,
            ]
        );

        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 zero_norm = rms_args(4, 2);
        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, f32::from_bits(u32::from(bf16(expected)) << 16));
        }

        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();
        assert_eq!(moe_values, [1.0, 3.3125]);

        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 mut actual_rope = [0.0; 4];
        rope_output.read_f32(&mut actual_rope).unwrap();
        assert_eq!(actual_rope, [-0.72265625, 0.71484375, 0.8984375, 1.46875]);

        let qsa_projected = Buffer::floats(QSA_WIDTH.into()).unwrap();
        let qsa_query = Buffer::floats((QSA_HEADS * QSA_DIM).into()).unwrap();
        let qsa_raw = Buffer::bytes(u64::from(QSA_RATIO * QSA_DIM) * 2).unwrap();
        let qsa_pooled = Buffer::bytes(u64::from(QSA_DIM) * 2).unwrap();
        let mut projected = vec![0.0; QSA_WIDTH as usize];
        for head in 0..QSA_HEADS {
            projected[(head * QSA_DIM) as usize] = 1.0;
        }
        let mut qsa_store = args();
        qsa_store.u[0] = QSA_DIM;
        qsa_store.u[1] = QSA_HEADS;
        for position in 0..QSA_RATIO {
            projected[(QSA_HEADS * QSA_DIM) as usize] = position as f32 + 1.0;
            qsa_projected.write_f32(&projected).unwrap();
            qsa_store.u[2] = position;
            dispatch_qwen(
                c"kernel_qwen_qsa_store_raw",
                &qsa_raw,
                Some(&qsa_projected),
                Some(&qsa_query),
                None,
                &[],
                &qsa_store,
                QSA_WIDTH,
                1,
            )
            .unwrap();
        }
        let mut qsa_pool = args();
        qsa_pool.u[0] = QSA_DIM;
        qsa_pool.u[1] = 0;
        qsa_pool.u[2] = QSA_RATIO;
        qsa_pool.u[3] = 0;
        qsa_pool.f[0] = 1.0e-6;
        qsa_pool.f[1] = 10_000_000.0;
        dispatch_qwen(
            c"kernel_qwen_qsa_pool_key",
            &qsa_pooled,
            Some(&qsa_raw),
            None,
            None,
            &[view(qsa_weight_offset, u64::from(QSA_DIM) * 2)],
            &qsa_pool,
            QSA_DIM,
            1,
        )
        .unwrap();
        let qsa_scores = Buffer::floats(1).unwrap();
        let mut qsa_score = args();
        qsa_score.u[0] = QSA_DIM;
        qsa_score.u[1] = QSA_HEADS;
        qsa_score.u[2] = 1;
        qsa_score.f[0] = (QSA_DIM as f32).sqrt().recip();
        dispatch_qwen(
            c"kernel_qwen_qsa_scores",
            &qsa_scores,
            Some(&qsa_query),
            Some(&qsa_pooled),
            None,
            &[],
            &qsa_score,
            1,
            1,
        )
        .unwrap();
        qsa_scores.read_f32(&mut scalar).unwrap();
        close(scalar[0], 4.0);

        let tie_scores = Buffer::floats(4).unwrap();
        tie_scores.write_f32(&[5.0, 5.0, 4.0, 3.0]).unwrap();
        let tie_selected = Buffer::bytes(8).unwrap();
        call(
            unsafe { ds4_gpu_indexer_topk_tensor(tie_selected.raw(), tie_scores.raw(), 4, 1, 2) },
            "testing deterministic Qwen QSA ties",
        )
        .unwrap();
        let mut tie_ids = [0; 2];
        tie_selected.read_i32(&mut tie_ids).unwrap();
        assert_eq!(tie_ids, [0, 1]);

        let selected_blocks = Buffer::bytes(12).unwrap();
        selected_blocks.write_i32(&[3, 1, 2]).unwrap();
        let mut sort = args();
        sort.u[0] = 3;
        dispatch_qwen(
            c"kernel_qwen_qsa_sort_blocks",
            &selected_blocks,
            None,
            None,
            None,
            &[],
            &sort,
            1,
            1,
        )
        .unwrap();
        let mut sorted = [0; 3];
        selected_blocks.read_i32(&mut sorted).unwrap();
        assert_eq!(sorted, [1, 2, 3]);

        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_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();
        assert_eq!(scalar[0], 8.5);
        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 Apple Metal and allocates the full QSA score window"]
    fn qwen_qsa_selection_and_attention_remain_bounded_at_native_depths() {
        use std::time::Instant;

        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let max_blocks = 262_144 / QSA_RATIO;
        let query = Buffer::floats((QSA_HEADS * QSA_DIM).into()).unwrap();
        query
            .write_f32(&vec![0.0; (QSA_HEADS * QSA_DIM) as usize])
            .unwrap();
        let pooled = Buffer::bytes(u64::from(max_blocks * QSA_DIM) * 2).unwrap();
        pooled
            .write(0, &vec![0; (max_blocks * QSA_DIM * 2) as usize])
            .unwrap();
        let scores = Buffer::floats(max_blocks.into()).unwrap();
        let selected = Buffer::bytes(u64::from(QSA_TOP_K) * 4).unwrap();

        for depth in [2_048, 5_228, 16_384, 65_536, 131_072, 262_144] {
            let blocks = depth / QSA_RATIO;
            let started = Instant::now();
            let mut score = args();
            score.u[0] = QSA_DIM;
            score.u[1] = QSA_HEADS;
            score.u[2] = blocks;
            score.f[0] = (QSA_DIM as f32).sqrt().recip();
            dispatch_qwen(
                c"kernel_qwen_qsa_scores",
                &scores,
                Some(&query),
                Some(&pooled),
                None,
                &[],
                &score,
                blocks,
                1,
            )
            .unwrap();
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(selected.raw(), scores.raw(), blocks, 1, QSA_TOP_K)
                },
                "testing full-window Qwen QSA selection",
            )
            .unwrap();
            let mut order = args();
            order.u[0] = QSA_TOP_K;
            dispatch_qwen(
                c"kernel_qwen_qsa_sort_blocks",
                &selected,
                None,
                None,
                None,
                &[],
                &order,
                1,
                1,
            )
            .unwrap();
            let mut ids = vec![0; QSA_TOP_K as usize];
            selected.read_i32(&mut ids).unwrap();
            assert_eq!(ids, (0..QSA_TOP_K as i32).collect::<Vec<_>>());
            eprintln!(
                "Qwen QSA selection depth {depth}: {:.3} ms",
                started.elapsed().as_secs_f64() * 1_000.0
            );
        }

        let blocks = 5_228 / QSA_RATIO;
        let iterations = 100;
        query
            .write_f32(&vec![1.0; (QSA_HEADS * QSA_DIM) as usize])
            .unwrap();
        pooled
            .write(
                0,
                &(0..u64::from(blocks * QSA_DIM))
                    .flat_map(|index| bf16((index % 127) as f32 / 127.0).to_le_bytes())
                    .collect::<Vec<_>>(),
            )
            .unwrap();
        let mut score = args();
        score.u[0] = QSA_DIM;
        score.u[1] = QSA_HEADS;
        score.u[2] = blocks;
        score.f[0] = (QSA_DIM as f32).sqrt().recip();
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            dispatch_qwen(
                c"kernel_qwen_qsa_scores",
                &scores,
                Some(&query),
                Some(&pooled),
                None,
                &[],
                &score,
                blocks,
                1,
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen QSA scores depth 5228: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(selected.raw(), scores.raw(), blocks, 1, QSA_TOP_K)
                },
                "benchmarking Qwen QSA top-k",
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen QSA top-k depth 5228: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            dispatch_qwen(
                c"kernel_qwen_qsa_scores",
                &scores,
                Some(&query),
                Some(&pooled),
                None,
                &[],
                &score,
                blocks,
                1,
            )
            .unwrap();
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(selected.raw(), scores.raw(), blocks, 1, QSA_TOP_K)
                },
                "benchmarking Qwen QSA score-to-top-k",
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen QSA score-to-top-k depth 5228: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );

        let depth = 5_229_u32;
        let cache = Buffer::bytes(u64::from(depth) * ATTN_KV_WIDTH as u64 * 4).unwrap();
        cache
            .write(
                0,
                &(0..u64::from(depth) * ATTN_KV_WIDTH as u64 * 2)
                    .flat_map(|_| bf16(1.0).to_le_bytes())
                    .collect::<Vec<_>>(),
            )
            .unwrap();
        selected
            .write_i32(&(0..QSA_TOP_K as i32).collect::<Vec<_>>())
            .unwrap();
        let main_query = Buffer::floats(ATTN_WIDTH.into()).unwrap();
        main_query
            .write_f32(&vec![1.0; ATTN_WIDTH as usize])
            .unwrap();
        let output = Buffer::floats(ATTN_WIDTH.into()).unwrap();
        let mask = Buffer::bytes(depth.into()).unwrap();
        let mut mask_args = args();
        mask_args.u[0] = depth;
        mask_args.u[4] = 1;
        mask_args.u[5] = QSA_RATIO;
        mask_args.u[6] = QSA_TOP_K;
        mask_args.u[8] = depth;
        mask_args.u[12] = 512;
        dispatch_qwen(
            c"kernel_qwen_qsa_mask",
            &mask,
            Some(&selected),
            None,
            None,
            &[],
            &mask_args,
            1,
            1,
        )
        .unwrap();
        let mut dense = args();
        dense.u[0] = ATTN_HEADS;
        dense.u[1] = ATTN_KV_HEADS;
        dense.u[2] = ATTN_DIM;
        dense.u[3] = depth;
        dense.u[4] = 1;
        dense.u[12] = 1024;
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            dispatch_qwen(
                c"kernel_qwen_dense_attention_masked",
                &output,
                Some(&main_query),
                Some(&cache),
                Some(&mask),
                &[],
                &dense,
                ATTN_HEADS,
                1,
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen masked dense attention depth {depth}: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            dispatch_qwen(
                c"kernel_qwen_qsa_scores",
                &scores,
                Some(&query),
                Some(&pooled),
                None,
                &[],
                &score,
                blocks,
                1,
            )
            .unwrap();
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(selected.raw(), scores.raw(), blocks, 1, QSA_TOP_K)
                },
                "benchmarking the complete Qwen QSA pipeline",
            )
            .unwrap();
            dispatch_qwen(
                c"kernel_qwen_qsa_mask",
                &mask,
                Some(&selected),
                None,
                None,
                &[],
                &mask_args,
                1,
                1,
            )
            .unwrap();
            dispatch_qwen(
                c"kernel_qwen_dense_attention_masked",
                &output,
                Some(&main_query),
                Some(&cache),
                Some(&mask),
                &[],
                &dense,
                ATTN_HEADS,
                1,
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen complete QSA pipeline depth {depth}: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );
        dispatch_qwen(
            c"kernel_qwen_qsa_mask",
            &mask,
            Some(&selected),
            None,
            None,
            &[],
            &mask_args,
            1,
            1,
        )
        .unwrap();
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..iterations {
            dispatch_qwen(
                c"kernel_qwen_dense_attention_masked",
                &output,
                Some(&main_query),
                Some(&cache),
                Some(&mask),
                &[],
                &dense,
                ATTN_HEADS,
                1,
            )
            .unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "Qwen scattered masked attention depth {depth}: {:.3} ms/call",
            started.elapsed().as_secs_f64() * 1_000.0 / f64::from(iterations),
        );

        let depth = 262_143_u32;
        let tail_start = depth / QSA_RATIO * QSA_RATIO;
        selected
            .write_i32(&(0..QSA_TOP_K as i32).collect::<Vec<_>>())
            .unwrap();
        let cache = Buffer::bytes(u64::from(depth) * ATTN_KV_WIDTH as u64 * 4).unwrap();
        cache
            .write(
                0,
                &vec![0; (QSA_TOP_K * QSA_RATIO * ATTN_KV_WIDTH * 4) as usize],
            )
            .unwrap();
        cache
            .write(
                u64::from(tail_start) * ATTN_KV_WIDTH as u64 * 4,
                &vec![0; ((depth - tail_start) * ATTN_KV_WIDTH * 4) as usize],
            )
            .unwrap();
        let main_query = Buffer::floats(ATTN_WIDTH.into()).unwrap();
        main_query
            .write_f32(&vec![0.0; ATTN_WIDTH as usize])
            .unwrap();
        let output = Buffer::floats(ATTN_WIDTH.into()).unwrap();
        let mut sparse = args();
        sparse.u[0] = ATTN_HEADS;
        sparse.u[1] = ATTN_KV_HEADS;
        sparse.u[2] = ATTN_DIM;
        sparse.u[3] = depth;
        sparse.u[4] = QSA_TOP_K;
        sparse.u[5] = QSA_RATIO;
        sparse.u[6] = tail_start;
        sparse.u[7] = depth - tail_start;
        sparse.u[12] = 256;
        let started = Instant::now();
        dispatch_qwen(
            c"kernel_qwen_sparse_attention",
            &output,
            Some(&main_query),
            Some(&cache),
            Some(&selected),
            &[],
            &sparse,
            ATTN_HEADS,
            1,
        )
        .unwrap();
        let mut values = vec![1.0; ATTN_WIDTH as usize];
        output.read_f32(&mut values).unwrap();
        assert!(values.into_iter().all(|value| value == 0.0));
        eprintln!(
            "Qwen sparse attention depth {depth}, {} visible tokens: {:.3} ms",
            QSA_TOP_K * QSA_RATIO + depth - tail_start,
            started.elapsed().as_secs_f64() * 1_000.0,
        );
    }

    #[test]
    #[ignore = "requires Apple Metal"]
    fn qwen_qsa_dense_mask_accepts_unsorted_blocks() {
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let tokens = 5_229_u32;
        let selected_ids = (0..QSA_TOP_K)
            .map(|index| ((index * 2) % (tokens / QSA_RATIO)) as i32)
            .collect::<Vec<_>>();
        let selected = Buffer::bytes(u64::from(QSA_TOP_K) * 4).unwrap();
        selected.write_i32(&selected_ids).unwrap();
        let mask = Buffer::bytes(tokens.into()).unwrap();
        let mut values = args();
        values.u[0] = tokens;
        values.u[4] = 1;
        values.u[5] = QSA_RATIO;
        values.u[6] = QSA_TOP_K;
        values.u[8] = tokens;
        values.u[12] = 512;
        dispatch_qwen(
            c"kernel_qwen_qsa_mask",
            &mask,
            Some(&selected),
            None,
            None,
            &[],
            &values,
            1,
            1,
        )
        .unwrap();

        let mut actual = vec![0_u8; tokens as usize];
        mask.read(0, &mut actual).unwrap();
        let mut expected = vec![0_u8; tokens as usize];
        for block in selected_ids {
            let token = block as usize * QSA_RATIO as usize;
            expected[token..token + QSA_RATIO as usize].fill(1);
        }
        expected[tokens as usize - 1] = 1;
        assert_eq!(actual, expected);
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_qsa_full_context_preallocation_succeeds() {
        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, 262_144).unwrap();
        assert_eq!(model.memory().kv_and_recurrent, 7_564_812_288);
        let executor = QwenExecutor::open_configured_with_hyper(
            model,
            262_144,
            0,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        assert_eq!(executor.context(), 262_144);
        assert_eq!(qsa_block_capacity(executor.context()).unwrap(), 65_537);
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_qsa_real_layer_executes_the_first_sparse_token() {
        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, DECODE_DENSE_BUDGET + 1).unwrap();
        let mut executor = QwenExecutor::open(model, DECODE_DENSE_BUDGET + 1).unwrap();
        executor.position = DECODE_DENSE_BUDGET;
        executor
            .scratch
            .block
            .write_f32(&vec![0.0; HIDDEN as usize])
            .unwrap();
        let LayerState::Attention { kv, qsa_pooled, .. } = &executor.states[3] else {
            unreachable!()
        };
        kv.write(
            0,
            &vec![0; (DECODE_DENSE_BUDGET * ATTN_KV_WIDTH * 4) as usize],
        )
        .unwrap();
        qsa_pooled
            .write(0, &vec![0; (QSA_TOP_K * QSA_DIM * 2) as usize])
            .unwrap();
        executor
            .attention("language_model.model.layers.3", &executor.states[3], 3)
            .unwrap();
        let mut hidden = vec![0.0; HIDDEN as usize];
        executor.scratch.hidden.read_f32(&mut hidden).unwrap();
        assert!(hidden.iter().all(|value| value.is_finite()));
        let mut selected = vec![0; QSA_TOP_K as usize];
        executor
            .scratch
            .qsa_selected
            .read_i32(&mut selected)
            .unwrap();
        assert_eq!(selected, (0..QSA_TOP_K as i32).collect::<Vec<_>>());
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_qsa_logical_depth_oracles_cover_prefill_decode_and_state() {
        use std::time::Instant;

        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, 262_144).unwrap();
        let mut executor = QwenExecutor::open(model, 262_144).unwrap();
        let oracles = [
            (
                2_048,
                "b6d28407fc87ebef94c56cf8ece0e0be472d45d8b7a8ba6f08d17420412d7411",
                5,
            ),
            (
                16_384,
                "79d71772c9cffd81e84ea21f9f5a0f3d7e63575ef3196fdea8512fec2a286e8d",
                5,
            ),
            (
                65_536,
                "0a1131fe9164b01265781b54a6dcca7da625e9b384a00015a779337a9a3ed35d",
                5,
            ),
            (
                131_072,
                "fb678d10cca29f42322253cde1ed2d4bc9dcfa47b7b6383e5de0f509964677d7",
                5,
            ),
            (
                262_144,
                "01ae755c3366805aa54ccaf11d226c1f90b216b8bd05e910807bfb552a810bc3",
                5,
            ),
        ];
        let mut baseline = None;
        for (depth, expected_digest, expected_token) in oracles {
            executor.reset().unwrap();
            let start = depth - 4;
            let prefix_bytes = u64::from(start.min(DECODE_DENSE_BUDGET)) * ATTN_KV_WIDTH as u64 * 4;
            let pooled_bytes = u64::from(start / QSA_RATIO) * QSA_DIM as u64 * 2;
            let zero_kv = vec![0; prefix_bytes as usize];
            let zero_pooled = vec![0; pooled_bytes as usize];
            for state in &executor.states {
                if let LayerState::Attention { kv, qsa_pooled, .. } = state {
                    kv.write(0, &zero_kv).unwrap();
                    qsa_pooled.write(0, &zero_pooled).unwrap();
                }
            }
            executor.position = start;
            executor.tokens = vec![EOS_TOKEN; start as usize];

            let prefill_started = Instant::now();
            assert_eq!(executor.prefill(&[1, 2, 3], |_| true).unwrap(), 3);
            let prefill_seconds = prefill_started.elapsed().as_secs_f64();
            let decode_started = Instant::now();
            executor.eval(4).unwrap();
            let decode_seconds = decode_started.elapsed().as_secs_f64();
            assert_eq!(executor.position, depth);
            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();
            let digest = digest
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect::<String>();
            let token = executor
                .logits
                .iter()
                .enumerate()
                .max_by(|left, right| left.1.total_cmp(right.1))
                .unwrap()
                .0;
            assert_eq!(digest, expected_digest);
            assert_eq!(token, expected_token);
            let live_attention_bytes =
                attention_checkpoint_bytes(depth).into_iter().sum::<u64>() * 12;
            let prefill_tps = 3.0 / prefill_seconds;
            let decode_tps = 1.0 / decode_seconds;
            eprintln!(
                "Qwen observed depth {depth}: prefill {prefill_tps:.3} tok/s, AR {decode_tps:.3} tok/s"
            );
            if let Some((baseline_prefill, baseline_decode)) = baseline {
                assert!(
                    prefill_tps >= baseline_prefill * 0.9,
                    "Qwen prefill collapsed at depth {depth}: {prefill_tps:.3} versus {baseline_prefill:.3} tok/s",
                );
                assert!(
                    decode_tps >= baseline_decode * 0.9,
                    "Qwen AR collapsed at depth {depth}: {decode_tps:.3} versus {baseline_decode:.3} tok/s",
                );
            } else {
                baseline = Some((prefill_tps, decode_tps));
            }
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &executor.states[3]
            else {
                unreachable!()
            };
            let mut frontier = vec![0; (ATTN_KV_WIDTH * 4 + QSA_DIM * 4) as usize];
            let (kv_frontier, rest) = frontier.split_at_mut((ATTN_KV_WIDTH * 4) as usize);
            let (raw_frontier, pool_frontier) = rest.split_at_mut((QSA_DIM * 2) as usize);
            kv.read(u64::from(depth - 1) * ATTN_KV_WIDTH as u64 * 4, kv_frontier)
                .unwrap();
            qsa_raw
                .read(u64::from(depth - 1) * QSA_DIM as u64 * 2, raw_frontier)
                .unwrap();
            qsa_pooled
                .read(
                    u64::from(depth / QSA_RATIO - 1) * QSA_DIM as u64 * 2,
                    pool_frontier,
                )
                .unwrap();
            let mut resident = None;
            executor.swap_resident_state(&mut resident).unwrap();
            assert_eq!(executor.position, 0);
            executor.swap_resident_state(&mut resident).unwrap();
            assert_eq!(executor.position, depth);
            assert_eq!(executor.tokens.len(), depth as usize);
            let checkpoint = std::env::temp_dir()
                .join(format!("ds4-qwen97-depth-{depth}-{}", std::process::id()));
            let tag = [depth.trailing_zeros() as u8; 32];
            executor
                .save_checkpoint(&checkpoint, tag, &mut |_| {})
                .unwrap();
            assert_eq!(
                fs::metadata(&checkpoint).unwrap().len(),
                116_676_712 + u64::from(depth) * 4 + live_attention_bytes
            );
            assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
            assert_eq!(executor.position, depth);
            assert_eq!(executor.checkpoint_tag, tag);
            fs::remove_file(checkpoint).unwrap();
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &executor.states[3]
            else {
                unreachable!()
            };
            let mut restored = vec![0; frontier.len()];
            let (kv_restored, rest) = restored.split_at_mut((ATTN_KV_WIDTH * 4) as usize);
            let (raw_restored, pool_restored) = rest.split_at_mut((QSA_DIM * 2) as usize);
            kv.read(u64::from(depth - 1) * ATTN_KV_WIDTH as u64 * 4, kv_restored)
                .unwrap();
            qsa_raw
                .read(u64::from(depth - 1) * QSA_DIM as u64 * 2, raw_restored)
                .unwrap();
            qsa_pooled
                .read(
                    u64::from(depth / QSA_RATIO - 1) * QSA_DIM as u64 * 2,
                    pool_restored,
                )
                .unwrap();
            assert_eq!(restored, frontier);
            eprintln!(
                "Qwen depth {depth}: logits {digest}, token {token}, prefill {:.3} tok/s, AR {:.3} tok/s, live attention {} bytes, peak RSS {} bytes",
                prefill_tps,
                decode_tps,
                live_attention_bytes,
                peak_rss_bytes(),
            );
        }
    }

    #[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 observed_next = executor
            .logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;
        eprintln!("Qwen token-1 argmax after projection dispatch: {observed_next}");
        eprintln!(
            "Qwen token-1 logit samples: {:?}",
            [
                executor.logits[0],
                executor.logits[1],
                executor.logits[1000],
                executor.logits[VOCAB as usize - 1],
            ]
        );
        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;
        let LayerState::Attention { qsa_raw, .. } = &executor.states[3] else {
            unreachable!()
        };
        let mut saved_qsa_raw = vec![0; QSA_DIM as usize * 2];
        qsa_raw.read(0, &mut saved_qsa_raw).unwrap();
        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]);
        let LayerState::Attention { qsa_raw, .. } = &executor.states[3] else {
            unreachable!()
        };
        let mut restored_qsa_raw = vec![0; saved_qsa_raw.len()];
        qsa_raw.read(0, &mut restored_qsa_raw).unwrap();
        assert_eq!(restored_qsa_raw, saved_qsa_raw);
        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]);

        let LayerState::Attention {
            qsa_raw,
            qsa_pooled,
            ..
        } = &executor.states[3]
        else {
            unreachable!()
        };
        let mut raw_key = vec![0; QSA_DIM as usize * 2];
        qsa_raw.read(0, &mut raw_key).unwrap();
        assert!(raw_key.iter().any(|&byte| byte != 0));
        assert!(!qsa_pooled.raw().is_null());
        let pooled_pattern = vec![0x5a; QSA_DIM as usize * 2];
        qsa_pooled.write(0, &pooled_pattern).unwrap();
        executor.position = 4;
        executor.tokens = vec![EOS_TOKEN; 4];
        executor
            .save_checkpoint(&checkpoint, [8; 32], &mut |_| {})
            .unwrap();
        assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        let LayerState::Attention { qsa_pooled, .. } = &executor.states[3] else {
            unreachable!()
        };
        let mut restored_pool = vec![0; pooled_pattern.len()];
        qsa_pooled.read(0, &mut restored_pool).unwrap();
        assert_eq!(restored_pool, pooled_pattern);
        assert_eq!(executor.position, 4);
        assert_eq!(executor.checkpoint_tag, [8; 32]);
        fs::remove_file(checkpoint).unwrap();
    }

    #[test]
    #[ignore = "diagnostic trace against the pinned MLX-LM reference"]
    fn qwen_reference_stateful_trace() {
        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 mut executor = QwenExecutor::open(QwenModel::open(&root, 4).unwrap(), 4).unwrap();
        let save = |name: &str, buffer: &Buffer, width: usize| {
            let mut values = vec![0.0; width];
            buffer.read_f32(&mut values).unwrap();
            let bytes = values
                .iter()
                .flat_map(|value| value.to_le_bytes())
                .collect::<Vec<_>>();
            std::fs::write(format!("/private/tmp/ds4_t2_{name}.f32"), bytes).unwrap();
            let mut hash = Sha256::new();
            for value in values {
                hash.update(value.to_le_bytes());
            }
            eprintln!("{name}: {:02x?}", hash.finalize());
        };

        executor.eval_target(1).unwrap();
        executor.eval_target(2).unwrap();
        executor.begin_token(2972).unwrap();
        for layer in 0..LAYERS {
            if layer == 1 {
                let commands = Commands::begin().unwrap();
                executor
                    .dispatch(
                        c"kernel_qwen_zero_rms",
                        &executor.scratch.hc_mix,
                        Some(&executor.scratch.hc),
                        None,
                        None,
                        &[executor.view(
                            executor
                                .weight("language_model.model.layers.1.ple.norm_query.weight")
                                .unwrap(),
                        )],
                        &rms_args(HC_WIDTH, HIDDEN),
                        HC,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                save(
                    "ple_query_before",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                executor.ple(2972).unwrap();
                save("after_ple", &executor.scratch.hc, HC_WIDTH as usize);
                save(
                    "ple_output",
                    &executor.scratch.ple_output,
                    HC_WIDTH as usize,
                );
                save(
                    "ple_embedding",
                    &executor.scratch.ple_embedding,
                    HIDDEN as usize,
                );
                save("ple_key", &executor.scratch.ple_key, HC_WIDTH as usize);
                save("ple_value", &executor.scratch.ple_value, HIDDEN as usize);
                save("ple_query", &executor.scratch.hc_norm, HC_WIDTH as usize);
                save("ple_gated", &executor.scratch.ple_gated, HC_WIDTH as usize);
                save("ple_norm", &executor.scratch.ple_norm, HC_WIDTH as usize);
            }
            if layer == 3 {
                let prefix = "language_model.model.layers.3";
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                save(
                    "layer3_hc_norm",
                    &executor.scratch.hc_norm,
                    HC_WIDTH as usize,
                );
                save("layer3_hc_rank", &executor.scratch.rank, HC_RANK as usize);
                save("layer3_hc_mix", &executor.scratch.hc_mix, HC_WIDTH as usize);
                save("layer3_mixed", &executor.scratch.block, HIDDEN as usize);
                save("layer3_inject", &executor.scratch.injection, HC as usize);
                let commands = Commands::begin().unwrap();
                executor
                    .attention(prefix, &executor.states[3], executor.position)
                    .unwrap();
                commands.finish().unwrap();
                save(
                    "layer3_q_packed",
                    &executor.scratch.q_packed,
                    (ATTN_WIDTH * 2) as usize,
                );
                save(
                    "layer3_q_gate",
                    &executor.scratch.q_gate,
                    ATTN_WIDTH as usize,
                );
                save("layer3_k", &executor.scratch.k, ATTN_KV_WIDTH as usize);
                save("layer3_v", &executor.scratch.v, ATTN_KV_WIDTH as usize);
                save(
                    "layer3_k_rope",
                    &executor.scratch.k_rope,
                    ATTN_KV_WIDTH as usize,
                );
                save(
                    "layer3_raw_attention",
                    &executor.scratch.q,
                    ATTN_WIDTH as usize,
                );
                save(
                    "layer3_gated_attention",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize,
                );
                save("layer3_output", &executor.scratch.hidden, HIDDEN as usize);
                let commands = Commands::begin().unwrap();
                executor
                    .head_norm_rope(
                        &executor.scratch.k,
                        &executor.scratch.qsa_qk,
                        executor
                            .weight("language_model.model.layers.3.self_attn.k_norm.weight")
                            .unwrap(),
                        ATTN_KV_HEADS,
                        0,
                    )
                    .unwrap();
                commands.finish().unwrap();
                save(
                    "layer3_k_norm",
                    &executor.scratch.qsa_qk,
                    ATTN_KV_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                let mut norm_only = args();
                norm_only.u[0] = ATTN_DIM;
                norm_only.u[1] = 64;
                norm_only.u[2] = ATTN_KV_HEADS;
                norm_only.u[11] = 1;
                norm_only.u[12] = 32;
                norm_only.f[0] = 1.0e-6;
                norm_only.f[1] = 10_000_000.0;
                executor
                    .dispatch(
                        c"kernel_qwen_head_norm_rope",
                        &executor.scratch.qsa_qk,
                        Some(&executor.scratch.k),
                        None,
                        None,
                        &[executor.view(
                            executor
                                .weight("language_model.model.layers.3.self_attn.k_norm.weight")
                                .unwrap(),
                        )],
                        &norm_only,
                        1,
                        ATTN_KV_HEADS,
                    )
                    .unwrap();
                commands.finish().unwrap();
                save(
                    "layer3_k_unweighted",
                    &executor.scratch.qsa_qk,
                    ATTN_KV_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                save(
                    "layer3_after_attention",
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
                executor.encode_moe(prefix, "3").unwrap();
            } else {
                let commands = Commands::begin().unwrap();
                executor.encode_layer(layer).unwrap();
                commands.finish().unwrap();
            }
            save(
                &format!("layer{layer}"),
                &executor.scratch.hc,
                HC_WIDTH as usize,
            );
        }
        save("final_widened", &executor.scratch.hc, HC_WIDTH as usize);
        executor.final_output().unwrap();
        save("final_mixed", &executor.scratch.block, HIDDEN as usize);
        save("logits", &executor.scratch.logits, VOCAB as usize);
    }

    #[test]
    #[ignore = "diagnostic trace against the pinned MTPLX verifier"]
    fn qwen_reference_verify_trace() {
        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_configured(&root, 8, true).unwrap();
        let settings = EngineSpeculativeSettings {
            glm_mtp: true,
            ..EngineSpeculativeSettings::default()
        };
        let mut executor = QwenExecutor::open_configured(model, 8, 0, 100, settings).unwrap();
        let save = |name: &str, buffer: &Buffer, width: usize| {
            let mut values = vec![0.0; width];
            buffer.read_f32(&mut values).unwrap();
            let bytes = values
                .iter()
                .flat_map(|value| value.to_le_bytes())
                .collect::<Vec<_>>();
            std::fs::write(format!("/private/tmp/ds4_verify_{name}.f32"), bytes).unwrap();
        };
        let save_i32 = |name: &str, buffer: &Buffer, width: usize| {
            let mut values = vec![0_i32; width];
            buffer.read_i32(&mut values).unwrap();
            let bytes = values
                .iter()
                .flat_map(|value| value.to_le_bytes())
                .collect::<Vec<_>>();
            std::fs::write(format!("/private/tmp/ds4_verify_{name}.i32"), bytes).unwrap();
        };
        executor.eval_target(1).unwrap();
        executor.eval_target(2).unwrap();
        let tokens = [2972, 44934, 28986, 31524];
        executor.begin_tokens(&tokens).unwrap();
        save("embedding", &executor.scratch.hc, HC_WIDTH as usize * 4);
        for layer in 0..LAYERS {
            let prefix = format!("language_model.model.layers.{layer}");
            if layer == 1 {
                executor.ple_batch(&tokens).unwrap();
            }
            executor
                .hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), 4)
                .unwrap();
            if layer == 0 {
                save(
                    "layer0_attn_hc_norm",
                    &executor.scratch.hc_norm,
                    HC_WIDTH as usize * 4,
                );
                save(
                    "layer0_attn_hc_rank",
                    &executor.scratch.rank,
                    HC_RANK as usize * 4,
                );
                save(
                    "layer0_attn_hc_mix",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize * 4,
                );
                save(
                    "layer0_attn_mixed",
                    &executor.scratch.block,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer0_attn_inject",
                    &executor.scratch.injection,
                    HC as usize * 4,
                );
            }
            match &executor.states[layer] {
                LayerState::Gdn { .. } => executor.gdn_rows(&prefix, layer, 4, true).unwrap(),
                state @ LayerState::Attention { .. } => executor
                    .attention_rows(&prefix, state, executor.position, 4)
                    .unwrap(),
            }
            if layer == 11 {
                save(
                    "layer11_attn_mixed",
                    &executor.scratch.block,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer11_q_packed",
                    &executor.scratch.q_packed,
                    (ATTN_WIDTH * 2) as usize * 4,
                );
                save("layer11_k", &executor.scratch.k, ATTN_KV_WIDTH as usize * 4);
                save("layer11_v", &executor.scratch.v, ATTN_KV_WIDTH as usize * 4);
                save(
                    "layer11_index_qk",
                    &executor.scratch.qsa_qk,
                    QSA_WIDTH as usize * 4,
                );
                save(
                    "layer11_attention",
                    &executor.scratch.q,
                    ATTN_WIDTH as usize * 4,
                );
                save(
                    "layer11_gated",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize * 4,
                );
                save(
                    "layer11_attn_output",
                    &executor.scratch.hidden,
                    HIDDEN as usize * 4,
                );
            }
            if layer == 3 {
                save(
                    "layer3_attn_mixed",
                    &executor.scratch.block,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer3_q_packed",
                    &executor.scratch.q_packed,
                    (ATTN_WIDTH * 2) as usize * 4,
                );
                save("layer3_k", &executor.scratch.k, ATTN_KV_WIDTH as usize * 4);
                save("layer3_v", &executor.scratch.v, ATTN_KV_WIDTH as usize * 4);
                save(
                    "layer3_index_qk",
                    &executor.scratch.qsa_qk,
                    QSA_WIDTH as usize * 4,
                );
                save(
                    "layer3_attn_output",
                    &executor.scratch.hidden,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer3_attention_batch",
                    &executor.scratch.q,
                    ATTN_WIDTH as usize * 4,
                );
                save(
                    "layer3_gated_batch",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize * 4,
                );
                let LayerState::Attention { kv, .. } = &executor.states[layer] else {
                    unreachable!()
                };
                let mut packed = vec![0_u8; 6 * ATTN_KV_WIDTH as usize * 4];
                kv.read(0, &mut packed).unwrap();
                for (name, value_offset) in [
                    ("layer3_cache_keys", 0),
                    ("layer3_cache_values", ATTN_KV_WIDTH as usize),
                ] {
                    let mut values = Vec::with_capacity(6 * ATTN_KV_WIDTH as usize);
                    for head in 0..ATTN_KV_HEADS as usize {
                        for token in 0..6_usize {
                            for column in 0..ATTN_DIM as usize {
                                let index = token * ATTN_KV_WIDTH as usize * 2
                                    + value_offset
                                    + head * ATTN_DIM as usize
                                    + column;
                                let byte = index * 2;
                                values.push(f32::from_bits(
                                    u32::from(u16::from_le_bytes([packed[byte], packed[byte + 1]]))
                                        << 16,
                                ));
                            }
                        }
                    }
                    let bytes = values
                        .iter()
                        .flat_map(|value| value.to_le_bytes())
                        .collect::<Vec<_>>();
                    std::fs::write(format!("/private/tmp/ds4_verify_{name}.f32"), bytes).unwrap();
                }
                let mut split = args();
                split.u[0] = ATTN_HEADS;
                split.u[1] = ATTN_DIM;
                split.u[4] = 4;
                executor
                    .dispatch(
                        c"kernel_qwen_split_q_gate",
                        &executor.scratch.q,
                        Some(&executor.scratch.q_packed),
                        Some(&executor.scratch.q_gate),
                        None,
                        &[],
                        &split,
                        ATTN_WIDTH * 4,
                        1,
                    )
                    .unwrap();
                save("layer3_q_raw", &executor.scratch.q, ATTN_WIDTH as usize * 4);
                save(
                    "layer3_q_gate",
                    &executor.scratch.q_gate,
                    ATTN_WIDTH as usize * 4,
                );
                save(
                    "layer3_k_rope",
                    &executor.scratch.k_rope,
                    ATTN_KV_WIDTH as usize * 4,
                );
                executor
                    .head_norm_rope_rows(
                        &executor.scratch.q,
                        &executor.scratch.attention,
                        executor
                            .weight(&format!("{prefix}.self_attn.q_norm.weight"))
                            .unwrap(),
                        ATTN_HEADS,
                        executor.position,
                        4,
                    )
                    .unwrap();
                save(
                    "layer3_q_rope",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize * 4,
                );
                let LayerState::Attention { kv, qsa_pooled, .. } = &executor.states[layer] else {
                    unreachable!()
                };
                for row in 0..4_u32 {
                    let view = |buffer: &Buffer, width: u32| {
                        buffer
                            .view(u64::from(row) * u64::from(width) * 4, u64::from(width) * 4)
                            .unwrap()
                    };
                    executor
                        .attend_row(
                            kv,
                            qsa_pooled,
                            executor.position + row,
                            &view(&executor.scratch.attention, ATTN_WIDTH),
                            &view(&executor.scratch.q, ATTN_WIDTH),
                            &view(&executor.scratch.qsa_qk, QSA_HEADS * QSA_DIM),
                            &view(
                                &executor.scratch.qsa_scores,
                                qsa_block_capacity(executor.context).unwrap(),
                            ),
                            &executor
                                .scratch
                                .qsa_selected
                                .view(
                                    u64::from(row) * u64::from(QSA_TOP_K) * 4,
                                    u64::from(QSA_TOP_K) * 4,
                                )
                                .unwrap(),
                        )
                        .unwrap();
                }
                save(
                    "layer3_attention",
                    &executor.scratch.q,
                    ATTN_WIDTH as usize * 4,
                );
                let mut gate = args();
                gate.u[0] = ATTN_WIDTH * 4;
                executor
                    .dispatch(
                        c"kernel_qwen_gate_attention",
                        &executor.scratch.attention,
                        Some(&executor.scratch.q),
                        Some(&executor.scratch.q_gate),
                        None,
                        &[],
                        &gate,
                        ATTN_WIDTH * 4,
                        1,
                    )
                    .unwrap();
                save(
                    "layer3_gated",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize * 4,
                );
            }
            if layer == 0 {
                save(
                    "layer0_gdn_qkv",
                    &executor.scratch.qkv,
                    GDN_QKV as usize * 4,
                );
                save(
                    "layer0_gdn_controls",
                    &executor.scratch.controls,
                    GDN_CONTROLS as usize * 4,
                );
                save(
                    "layer0_gdn_q",
                    &executor.scratch.gdn_q,
                    (GDN_HEADS_K * HEAD_DIM) as usize * 4,
                );
                save(
                    "layer0_gdn_k",
                    &executor.scratch.gdn_k,
                    (GDN_HEADS_K * HEAD_DIM) as usize * 4,
                );
                save(
                    "layer0_gdn_delta",
                    &executor.scratch.gdn_out,
                    GDN_VALUE as usize * 4,
                );
                save(
                    "layer0_gdn_gated",
                    &executor.scratch.gdn_z,
                    GDN_VALUE as usize * 4,
                );
                save(
                    "layer0_attn_output",
                    &executor.scratch.hidden,
                    HIDDEN as usize * 4,
                );
            }
            executor.hyper_write_rows(4).unwrap();
            if layer == 11 {
                save(
                    "layer11_after_attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize * 4,
                );
            }
            if layer == 3 {
                save(
                    "layer3_after_attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize * 4,
                );
            }
            if layer == 0 {
                save(
                    "layer0_after_attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize * 4,
                );
                executor
                    .hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), 4)
                    .unwrap();
                save(
                    "layer0_mlp_mixed",
                    &executor.scratch.block,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer0_mlp_inject",
                    &executor.scratch.injection,
                    HC as usize * 4,
                );
            }
            if layer == 0 {
                executor
                    .affine_batch_into(
                        &executor
                            .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)
                            .unwrap(),
                        &executor.scratch.block,
                        &executor.scratch.router,
                        HIDDEN,
                        EXPERTS,
                        4,
                        true,
                    )
                    .unwrap();
                save(
                    "layer0_router_logits",
                    &executor.scratch.router,
                    EXPERTS as usize * 4,
                );
                let mut softmax = args();
                softmax.u[0] = EXPERTS;
                softmax.u[12] = EXPERTS / 4;
                executor
                    .dispatch(
                        c"kernel_qwen_softmax_precise_512",
                        &executor.scratch.router,
                        Some(&executor.scratch.router),
                        None,
                        None,
                        &[],
                        &softmax,
                        4,
                        1,
                    )
                    .unwrap();
                let mut route = args();
                route.u[0] = EXPERTS;
                route.u[12] = EXPERTS;
                executor
                    .dispatch(
                        c"kernel_qwen_route_top10",
                        &executor.scratch.route_ids,
                        Some(&executor.scratch.route_weights),
                        Some(&executor.scratch.router),
                        None,
                        &[],
                        &route,
                        4,
                        1,
                    )
                    .unwrap();
                save(
                    "layer0_router_probabilities",
                    &executor.scratch.router,
                    EXPERTS as usize * 4,
                );
                save_i32(
                    "layer0_route_ids",
                    &executor.scratch.route_ids,
                    EXPERTS_USED * 4,
                );
                save(
                    "layer0_route_weights",
                    &executor.scratch.route_weights,
                    EXPERTS_USED * 4,
                );
                executor.routed_experts_rows(&prefix, 4).unwrap();
                save(
                    "layer0_expert_gate",
                    &executor.scratch.gate,
                    EXPERT_WIDTH as usize * EXPERTS_USED * 4,
                );
                save(
                    "layer0_expert_up",
                    &executor.scratch.up,
                    EXPERT_WIDTH as usize * EXPERTS_USED * 4,
                );
                save(
                    "layer0_expert_mid",
                    &executor.scratch.mid,
                    EXPERT_WIDTH as usize * EXPERTS_USED * 4,
                );
                save(
                    "layer0_expert_down",
                    &executor.scratch.expert,
                    HIDDEN as usize * EXPERTS_USED * 4,
                );
                save("layer0_routed", &executor.scratch.moe, HIDDEN as usize * 4);
                executor.shared_expert_rows(&prefix, 4).unwrap();
                save("layer0_shared_gate", &executor.scratch.gate, 4);
                save(
                    "layer0_shared_up",
                    &executor.scratch.up,
                    EXPERT_WIDTH as usize * 4,
                );
                save(
                    "layer0_shared_mid",
                    &executor.scratch.mid,
                    EXPERT_WIDTH as usize * 4,
                );
                save(
                    "layer0_shared_down",
                    &executor.scratch.shared,
                    HIDDEN as usize * 4,
                );
                save(
                    "layer0_mlp_output",
                    &executor.scratch.moe,
                    HIDDEN as usize * 4,
                );
                let mut inject = args();
                inject.u[0] = HIDDEN;
                inject.u[4] = 4;
                executor
                    .dispatch(
                        c"kernel_qwen_hyper_inject",
                        &executor.scratch.hc_norm,
                        Some(&executor.scratch.hc),
                        Some(&executor.scratch.moe),
                        Some(&executor.scratch.injection),
                        &[],
                        &inject,
                        HC_WIDTH * 4,
                        1,
                    )
                    .unwrap();
                std::mem::swap(&mut executor.scratch.hc, &mut executor.scratch.hc_norm);
            } else {
                if layer == 3 {
                    executor
                        .hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), 4)
                        .unwrap();
                    save(
                        "layer3_mlp_mixed",
                        &executor.scratch.block,
                        HIDDEN as usize * 4,
                    );
                }
                executor.encode_moe_rows(&prefix, 4).unwrap();
                if layer == 3 {
                    save(
                        "layer3_mlp_output",
                        &executor.scratch.moe,
                        HIDDEN as usize * 4,
                    );
                }
                if layer == 11 {
                    save(
                        "layer11_mlp_mixed",
                        &executor.scratch.block,
                        HIDDEN as usize * 4,
                    );
                    save(
                        "layer11_router_probabilities",
                        &executor.scratch.router,
                        EXPERTS as usize * 4,
                    );
                    save_i32(
                        "layer11_route_ids",
                        &executor.scratch.route_ids,
                        EXPERTS_USED * 4,
                    );
                    save(
                        "layer11_route_weights",
                        &executor.scratch.route_weights,
                        EXPERTS_USED * 4,
                    );
                    save(
                        "layer11_expert_down",
                        &executor.scratch.expert,
                        HIDDEN as usize * EXPERTS_USED * 4,
                    );
                    save("layer11_routed", &executor.scratch.moe, HIDDEN as usize * 4);
                    save(
                        "layer11_shared_down",
                        &executor.scratch.shared,
                        HIDDEN as usize * 4,
                    );
                    save(
                        "layer11_mlp_output",
                        &executor.scratch.moe,
                        HIDDEN as usize * 4,
                    );
                }
            }
            save(
                &format!("layer{layer}"),
                &executor.scratch.hc,
                HC_WIDTH as usize * 4,
            );
        }
        executor.final_output_rows(4).unwrap();
        let mut logits = vec![0.0; VOCAB as usize * 4];
        executor.scratch.logits.read_f32(&mut logits).unwrap();
        let bytes = logits
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect::<Vec<_>>();
        std::fs::write("/private/tmp/ds4_verify_logits.f32", bytes).unwrap();
    }

    #[test]
    #[ignore = "diagnostic comparison with captured MTPLX SDPA inputs"]
    fn qwen_reference_sdpa_trace() {
        configure_sources().unwrap();
        let read_f32 = |path: &str| {
            fs::read(path)
                .unwrap()
                .chunks_exact(4)
                .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap()))
                .collect::<Vec<_>>()
        };
        let upload_bf16 = |values: &[f32]| {
            let bytes = values
                .iter()
                .flat_map(|value| bf16(*value).to_le_bytes())
                .collect::<Vec<_>>();
            let buffer = Buffer::bytes(bytes.len() as u64).unwrap();
            buffer.write(0, &bytes).unwrap();
            buffer
        };
        let query = upload_bf16(&read_f32("/private/tmp/mtplx_verify_layer3_q_rope.f32"));
        let keys = upload_bf16(&read_f32("/private/tmp/mtplx_verify_layer3_cache_keys.f32"));
        let values = upload_bf16(&read_f32(
            "/private/tmp/mtplx_verify_layer3_cache_values.f32",
        ));
        let mask_values = (0..4_u32)
            .flat_map(|row| (0..6_u32).map(move |token| u8::from(token <= row + 2)))
            .collect::<Vec<_>>();
        let mask = Buffer::bytes(mask_values.len() as u64).unwrap();
        mask.write(0, &mask_values).unwrap();
        let out = Buffer::floats(u64::from(ATTN_WIDTH) * 4).unwrap();
        let mut values_args = args();
        values_args.u[0] = ATTN_HEADS;
        values_args.u[1] = ATTN_KV_HEADS;
        values_args.u[2] = ATTN_DIM;
        values_args.u[3] = 6;
        values_args.u[4] = 4;
        values_args.u[5] = 2;
        values_args.u[12] = 1024;
        let commands = Commands::begin().unwrap();
        dispatch_qwen_with_d(
            c"kernel_qwen_dense_attention_rows",
            &out,
            Some(&query),
            Some(&keys),
            Some(&values),
            Some(&mask),
            &[],
            &values_args,
            ATTN_HEADS,
            4,
        )
        .unwrap();
        commands.finish().unwrap();
        let mut actual = vec![0.0; ATTN_WIDTH as usize * 4];
        out.read_f32(&mut actual).unwrap();
        let expected = read_f32("/private/tmp/mtplx_verify_layer3_attention.f32");
        let hash = |input: &[f32]| {
            let mut digest = Sha256::new();
            for value in input {
                digest.update(value.to_le_bytes());
            }
            digest
                .finalize()
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect::<String>()
        };
        eprintln!("expected={} actual={}", hash(&expected), hash(&actual));
        assert_eq!(actual, expected);
    }

    #[test]
    #[ignore = "diagnostic trace against the pinned MLX-LM reference"]
    fn qwen_reference_layer_trace() {
        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 mut executor = QwenExecutor::open(model, 4).unwrap();
        let show = |label: &str, buffer: &Buffer, width: usize| {
            let mut values = vec![0.0; width];
            buffer.read_f32(&mut values).unwrap();
            let mut digest = Sha256::new();
            for value in &values {
                digest.update(value.to_bits().to_le_bytes());
            }
            let digest = digest
                .finalize()
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect::<String>();
            let at = |index| values[index];
            if matches!(
                label,
                "layer-0-attn-hc-mix"
                    | "ple-embedding"
                    | "ple-key"
                    | "ple-value"
                    | "ple-query"
                    | "ple-gated"
                    | "ple-norm"
                    | "ple-output"
                    | "ple"
                    | "layer-0-attn-input"
                    | "layer-0-gdn-qkv"
                    | "layer-0-gdn-conv"
                    | "layer-0-gdn-controls"
                    | "layer-0-gdn-gated"
                    | "layer-0-attn-output"
                    | "layer-0-after-attn"
                    | "layer-0-mlp-hc-norm"
                    | "layer-0-mlp-hc-down"
                    | "layer-0-mlp-hc-rank"
                    | "layer-0-mlp-hc-up"
                    | "layer-0-mlp-hc-mix"
                    | "layer-0-mlp-hc-inject"
                    | "layer-0-mlp-mixed"
                    | "layer-0-router-logits"
                    | "layer-0-router-probabilities"
                    | "layer-0-expert-gate"
                    | "layer-0-expert-up"
                    | "layer-0-expert-mid"
                    | "layer-0-expert-down"
                    | "layer-0-routed"
                    | "layer-0-shared-gate-raw"
                    | "layer-0-shared-gate-proj"
                    | "layer-0-shared-up"
                    | "layer-0-shared-mid"
                    | "layer-0-shared-down"
                    | "layer-0-total"
                    | "layer-0-output"
                    | "layer-1"
                    | "layer-1-attn-mixed"
                    | "layer-1-attn-inject"
                    | "layer-1-gdn-qkv"
                    | "layer-1-gdn-z"
                    | "layer-1-gdn-b"
                    | "layer-1-gdn-a"
                    | "layer-1-attn-out"
                    | "layer-1-after-attn"
                    | "layer-1-mlp-hc-norm"
                    | "layer-1-mlp-hc-rank"
                    | "layer-1-mlp-hc-up"
                    | "layer-1-mlp-hc-mix"
                    | "layer-1-mlp-inject"
                    | "layer-1-mlp-mixed"
                    | "layer-1-mlp-out"
                    | "layer-2-attn-mixed"
                    | "layer-2-attn-inject"
                    | "layer-2-attn-out"
                    | "layer-2-after-attn"
                    | "layer-2-mlp-mixed"
                    | "layer-2-router-probabilities"
                    | "layer-2-route-weights"
                    | "layer-2-expert-gate"
                    | "layer-2-expert-up"
                    | "layer-2-expert-mid"
                    | "layer-2-expert-down"
                    | "layer-2-routed"
                    | "layer-2-shared-gate-raw"
                    | "layer-2-shared-up"
                    | "layer-2-shared-mid"
                    | "layer-2-shared-down"
                    | "layer-2-total"
                    | "layer-2-mlp-out"
                    | "layer-3-attn-mixed"
                    | "layer-3-attn-inject"
                    | "layer-3-attn-out"
                    | "layer-3-after-attn"
                    | "layer-3-mlp-mixed"
                    | "layer-3-mlp-out"
                    | "layer-3-q"
                    | "layer-3-gate"
                    | "layer-3-v"
                    | "layer-3-attention"
                    | "layer-3-hidden"
                    | "layer-2"
                    | "layer-3"
                    | "layer-47"
                    | "final-widened"
                    | "final-mixed"
                    | "logits"
            ) {
                let bytes = values
                    .iter()
                    .flat_map(|value| value.to_le_bytes())
                    .collect::<Vec<_>>();
                let path = match label {
                    "ple-embedding" => "/private/tmp/ds4_ple_embedding.f32",
                    "ple-key" => "/private/tmp/ds4_ple_key.f32",
                    "ple-value" => "/private/tmp/ds4_ple_value.f32",
                    "ple-query" => "/private/tmp/ds4_ple_query.f32",
                    "ple-gated" => "/private/tmp/ds4_ple_gated.f32",
                    "ple-norm" => "/private/tmp/ds4_ple_norm.f32",
                    "ple-output" => "/private/tmp/ds4_ple_output.f32",
                    "ple" => "/private/tmp/ds4_after_ple.f32",
                    "layer-0-attn-hc-mix" => "/private/tmp/ds4_hc_mix.f32",
                    "layer-0-attn-input" => "/private/tmp/ds4_attn_mixed.f32",
                    "layer-0-gdn-qkv" => "/private/tmp/ds4_gdn_qkv.f32",
                    "layer-0-gdn-conv" => "/private/tmp/ds4_gdn_conv.f32",
                    "layer-0-gdn-controls" => "/private/tmp/ds4_gdn_controls.f32",
                    "layer-0-gdn-gated" => "/private/tmp/ds4_gdn_gated.f32",
                    "layer-0-after-attn" => "/private/tmp/ds4_after_attn.f32",
                    "layer-0-mlp-hc-norm" => "/private/tmp/ds4_mlp_hc_norm.f32",
                    "layer-0-mlp-hc-down" => "/private/tmp/ds4_mlp_hc_down.f32",
                    "layer-0-mlp-hc-rank" => "/private/tmp/ds4_mlp_hc_rank.f32",
                    "layer-0-mlp-hc-up" => "/private/tmp/ds4_mlp_hc_up.f32",
                    "layer-0-mlp-hc-mix" => "/private/tmp/ds4_mlp_hc_mix.f32",
                    "layer-0-mlp-hc-inject" => "/private/tmp/ds4_mlp_hc_inject.f32",
                    "layer-0-mlp-mixed" => "/private/tmp/ds4_mlp_mixed.f32",
                    "layer-0-router-logits" => "/private/tmp/ds4_router_logits.f32",
                    "layer-0-router-probabilities" => "/private/tmp/ds4_router_probabilities.f32",
                    "layer-0-expert-gate" => "/private/tmp/ds4_moe_expert_gate.f32",
                    "layer-0-expert-up" => "/private/tmp/ds4_moe_expert_up.f32",
                    "layer-0-expert-mid" => "/private/tmp/ds4_moe_expert_mid.f32",
                    "layer-0-expert-down" => "/private/tmp/ds4_moe_expert_down.f32",
                    "layer-0-routed" => "/private/tmp/ds4_moe_routed.f32",
                    "layer-0-shared-gate-raw" => "/private/tmp/ds4_moe_shared_gate_raw.f32",
                    "layer-0-shared-gate-proj" => "/private/tmp/ds4_moe_shared_gate_proj.f32",
                    "layer-0-shared-up" => "/private/tmp/ds4_moe_shared_up.f32",
                    "layer-0-shared-mid" => "/private/tmp/ds4_moe_shared_mid.f32",
                    "layer-0-shared-down" => "/private/tmp/ds4_moe_shared_down.f32",
                    "layer-0-total" => "/private/tmp/ds4_moe_total.f32",
                    "layer-0-output" => "/private/tmp/ds4_layer0_output.f32",
                    "layer-1" => "/private/tmp/ds4_layer1_output.f32",
                    "layer-1-attn-mixed" => "/private/tmp/ds4_layer1_attn_mixed.f32",
                    "layer-1-attn-inject" => "/private/tmp/ds4_layer1_attn_inject.f32",
                    "layer-1-gdn-qkv" => "/private/tmp/ds4_layer1_gdn_qkv.f32",
                    "layer-1-gdn-z" => "/private/tmp/ds4_layer1_gdn_z.f32",
                    "layer-1-gdn-b" => "/private/tmp/ds4_layer1_gdn_b.f32",
                    "layer-1-gdn-a" => "/private/tmp/ds4_layer1_gdn_a.f32",
                    "layer-1-attn-out" => "/private/tmp/ds4_layer1_attn_out.f32",
                    "layer-1-after-attn" => "/private/tmp/ds4_layer1_after_attn.f32",
                    "layer-1-mlp-hc-norm" => "/private/tmp/ds4_layer1_mlp_hc_norm.f32",
                    "layer-1-mlp-hc-rank" => "/private/tmp/ds4_layer1_mlp_hc_rank.f32",
                    "layer-1-mlp-hc-up" => "/private/tmp/ds4_layer1_mlp_hc_up.f32",
                    "layer-1-mlp-hc-mix" => "/private/tmp/ds4_layer1_mlp_hc_mix.f32",
                    "layer-1-mlp-inject" => "/private/tmp/ds4_layer1_mlp_hc_inject.f32",
                    "layer-1-mlp-mixed" => "/private/tmp/ds4_layer1_mlp_mixed.f32",
                    "layer-1-mlp-out" => "/private/tmp/ds4_layer1_mlp_out.f32",
                    "layer-2-attn-mixed" => "/private/tmp/ds4_layer2_attn_mixed.f32",
                    "layer-2-attn-inject" => "/private/tmp/ds4_layer2_attn_inject.f32",
                    "layer-2-attn-out" => "/private/tmp/ds4_layer2_attn_out.f32",
                    "layer-2-after-attn" => "/private/tmp/ds4_layer2_after_attn.f32",
                    "layer-2-mlp-mixed" => "/private/tmp/ds4_layer2_mlp_mixed.f32",
                    "layer-2-router-probabilities" => {
                        "/private/tmp/ds4_layer2_router_probabilities.f32"
                    }
                    "layer-2-route-weights" => "/private/tmp/ds4_layer2_route_weights.f32",
                    "layer-2-expert-gate" => "/private/tmp/ds4_layer2_moe_expert_gate.f32",
                    "layer-2-expert-up" => "/private/tmp/ds4_layer2_moe_expert_up.f32",
                    "layer-2-expert-mid" => "/private/tmp/ds4_layer2_moe_expert_mid.f32",
                    "layer-2-expert-down" => "/private/tmp/ds4_layer2_moe_expert_down.f32",
                    "layer-2-routed" => "/private/tmp/ds4_layer2_moe_routed.f32",
                    "layer-2-shared-gate-raw" => "/private/tmp/ds4_layer2_moe_shared_gate_raw.f32",
                    "layer-2-shared-up" => "/private/tmp/ds4_layer2_moe_shared_up.f32",
                    "layer-2-shared-mid" => "/private/tmp/ds4_layer2_moe_shared_mid.f32",
                    "layer-2-shared-down" => "/private/tmp/ds4_layer2_moe_shared_down.f32",
                    "layer-2-total" => "/private/tmp/ds4_layer2_moe_total.f32",
                    "layer-2-mlp-out" => "/private/tmp/ds4_layer2_mlp_out.f32",
                    "layer-2" => "/private/tmp/ds4_layer2_output.f32",
                    "layer-3-attn-mixed" => "/private/tmp/ds4_layer3_attn_mixed.f32",
                    "layer-3-attn-inject" => "/private/tmp/ds4_layer3_attn_inject.f32",
                    "layer-3-attn-out" => "/private/tmp/ds4_layer3_attn_out.f32",
                    "layer-3-after-attn" => "/private/tmp/ds4_layer3_after_attn.f32",
                    "layer-3-mlp-mixed" => "/private/tmp/ds4_layer3_mlp_mixed.f32",
                    "layer-3-mlp-out" => "/private/tmp/ds4_layer3_mlp_out.f32",
                    "layer-3-q" => "/private/tmp/ds4_layer3_raw_attention.f32",
                    "layer-3-gate" => "/private/tmp/ds4_layer3_gate.f32",
                    "layer-3-v" => "/private/tmp/ds4_layer3_value.f32",
                    "layer-3-attention" => "/private/tmp/ds4_layer3_gated_attention.f32",
                    "layer-3-hidden" => "/private/tmp/ds4_layer3_manual_out.f32",
                    "layer-3" => "/private/tmp/ds4_layer3_output.f32",
                    "layer-47" => "/private/tmp/ds4_layer47_output.f32",
                    "final-widened" => "/private/tmp/ds4_final_widened.f32",
                    "final-mixed" => "/private/tmp/ds4_final_mixed.f32",
                    "logits" => "/private/tmp/ds4_logits.f32",
                    _ => "/private/tmp/ds4_attn_out.f32",
                };
                std::fs::write(path, bytes).unwrap();
            }
            eprintln!(
                "{label}: [{:.9}, {:.9}, {:.9}, {:.9}] sha256={}",
                at(0),
                at(1.min(width - 1)),
                at(1000.min(width - 1)),
                at((HIDDEN as usize - 1).min(width - 1)),
                digest,
            );
        };
        let sigmoid_inputs = (0..=u16::MAX)
            .filter_map(|bits| {
                let value = f32::from_bits(u32::from(bits) << 16);
                (value.is_finite() && value.abs() <= 32.0).then_some(value)
            })
            .collect::<Vec<_>>();
        let mut sigmoid_outputs = Vec::with_capacity(sigmoid_inputs.len());
        for input in sigmoid_inputs.chunks(HC_WIDTH as usize) {
            executor.scratch.hc_mix.write_f32(input).unwrap();
            let commands = Commands::begin().unwrap();
            let mut unary = args();
            unary.u[0] = input.len() as u32;
            executor
                .dispatch(
                    c"kernel_qwen_sigmoid",
                    &executor.scratch.hc_norm,
                    Some(&executor.scratch.hc_mix),
                    None,
                    None,
                    &[],
                    &unary,
                    input.len() as u32,
                    1,
                )
                .unwrap();
            commands.finish().unwrap();
            let mut output = vec![0.0; input.len()];
            executor.scratch.hc_norm.read_f32(&mut output).unwrap();
            sigmoid_outputs.extend(output);
        }
        for (path, values) in [
            ("/private/tmp/ds4_sigmoid_inputs.f32", &sigmoid_inputs),
            ("/private/tmp/ds4_sigmoid_outputs.f32", &sigmoid_outputs),
        ] {
            let bytes = values
                .iter()
                .flat_map(|value| value.to_bits().to_le_bytes())
                .collect::<Vec<_>>();
            std::fs::write(path, bytes).unwrap();
        }
        executor.begin_token(1).unwrap();
        eprintln!(
            "ple-rows: {:?}",
            ple_rows(&executor.ple_contract, executor.ple_state.history, 1).unwrap()
        );
        show("embed", &executor.scratch.hc, HC_WIDTH as usize);
        for layer in 0..LAYERS {
            let layer_started = Instant::now();
            if layer == 1 {
                executor.ple(1).unwrap();
                show(
                    "ple-embedding",
                    &executor.scratch.ple_embedding,
                    HIDDEN as usize,
                );
                show("ple-key", &executor.scratch.ple_key, HC_WIDTH as usize);
                show("ple-value", &executor.scratch.ple_value, HIDDEN as usize);
                show("ple-query", &executor.scratch.hc_norm, HC_WIDTH as usize);
                show("ple-gated", &executor.scratch.ple_gated, HC_WIDTH as usize);
                show("ple-norm", &executor.scratch.ple_norm, HC_WIDTH as usize);
                show(
                    "ple-output",
                    &executor.scratch.ple_output,
                    HC_WIDTH as usize,
                );
                show("ple", &executor.scratch.hc, HC_WIDTH as usize);
            }
            if layer == 0 {
                let prefix = "language_model.model.layers.0";
                let commands = Commands::begin().unwrap();
                let norm = executor
                    .weight(&format!("{prefix}.attn_hyper_connection.hc_norm.weight"))
                    .unwrap();
                let down = executor
                    .weight(&format!(
                        "{prefix}.attn_hyper_connection.input_mix_weight_down.weight"
                    ))
                    .unwrap();
                executor
                    .dispatch(
                        c"kernel_qwen_zero_rms",
                        &executor.scratch.hc_norm,
                        Some(&executor.scratch.hc),
                        None,
                        None,
                        &[executor.view(norm)],
                        &rms_args(HC_WIDTH, HIDDEN),
                        HC,
                        1,
                    )
                    .unwrap();
                executor
                    .bf16_mv(
                        down,
                        &executor.scratch.hc_norm,
                        &executor.scratch.rank,
                        HC_WIDTH,
                        HC_RANK,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-attn-hc-down-raw",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-attn-hc-norm",
                    &executor.scratch.hc_norm,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-0-attn-hc-rank",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                let commands = Commands::begin().unwrap();
                let up = executor
                    .weight(&format!(
                        "{prefix}.attn_hyper_connection.input_mix_weight_up.weight"
                    ))
                    .unwrap();
                executor
                    .bf16_mv(
                        up,
                        &executor.scratch.rank,
                        &executor.scratch.hc_mix,
                        HC_RANK,
                        HC_WIDTH,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-attn-hc-up-raw",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                let mut sigmoid = args();
                sigmoid.u[0] = HC_WIDTH;
                executor
                    .dispatch(
                        c"kernel_qwen_sigmoid",
                        &executor.scratch.hc_mix,
                        Some(&executor.scratch.hc_mix),
                        None,
                        None,
                        &[],
                        &sigmoid,
                        HC_WIDTH,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-attn-hc-mix",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-0-attn-hc-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                show(
                    "layer-0-attn-input",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .affine_mv_into(
                        &executor
                            .affine(
                                &format!("{prefix}.linear_attn.in_proj_qkv"),
                                HIDDEN,
                                GDN_QKV,
                                None,
                            )
                            .unwrap(),
                        &executor.scratch.block,
                        &executor.scratch.qkv,
                        HIDDEN,
                        GDN_QKV,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show("layer-0-gdn-qkv", &executor.scratch.qkv, GDN_QKV as usize);
                let commands = Commands::begin().unwrap();
                executor.gdn(prefix, layer).unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-gdn-gated",
                    &executor.scratch.gdn_out,
                    GDN_VALUE as usize,
                );
                show("layer-0-gdn-conv", &executor.scratch.qkv, GDN_QKV as usize);
                show(
                    "layer-0-gdn-controls",
                    &executor.scratch.controls,
                    GDN_CONTROLS as usize,
                );
                show(
                    "layer-0-attn-output",
                    &executor.scratch.hidden,
                    HIDDEN as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-after-attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.mlp_hyper_connection"))
                    .unwrap();
                executor
                    .affine_mv_into(
                        &executor
                            .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)
                            .unwrap(),
                        &executor.scratch.block,
                        &executor.scratch.router,
                        HIDDEN,
                        EXPERTS,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-mlp-hc-norm",
                    &executor.scratch.hc_norm,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-0-mlp-hc-rank",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                show(
                    "layer-0-mlp-hc-mix",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-0-mlp-hc-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                let down = executor
                    .weight(&format!(
                        "{prefix}.mlp_hyper_connection.input_mix_weight_down.weight"
                    ))
                    .unwrap();
                let up = executor
                    .weight(&format!(
                        "{prefix}.mlp_hyper_connection.input_mix_weight_up.weight"
                    ))
                    .unwrap();
                let commands = Commands::begin().unwrap();
                executor
                    .bf16_mv(
                        down,
                        &executor.scratch.hc_norm,
                        &executor.scratch.rank,
                        HC_WIDTH,
                        HC_RANK,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-mlp-hc-down",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                let commands = Commands::begin().unwrap();
                let mut unary = args();
                unary.u[0] = HC_RANK;
                executor
                    .dispatch(
                        c"kernel_qwen_silu_div4",
                        &executor.scratch.rank,
                        Some(&executor.scratch.rank),
                        None,
                        None,
                        &[],
                        &unary,
                        HC_RANK,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-mlp-hc-rank",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .bf16_mv(
                        up,
                        &executor.scratch.rank,
                        &executor.scratch.hc_mix,
                        HC_RANK,
                        HC_WIDTH,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-mlp-hc-up",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                unary.u[0] = HC_WIDTH;
                executor
                    .dispatch(
                        c"kernel_qwen_sigmoid",
                        &executor.scratch.hc_mix,
                        Some(&executor.scratch.hc_mix),
                        None,
                        None,
                        &[],
                        &unary,
                        HC_WIDTH,
                        1,
                    )
                    .unwrap();
                let mut mix = args();
                mix.u[0] = HIDDEN;
                executor
                    .dispatch(
                        c"kernel_qwen_hyper_mix",
                        &executor.scratch.block,
                        Some(&executor.scratch.hc_norm),
                        Some(&executor.scratch.hc_mix),
                        None,
                        &[],
                        &mix,
                        HIDDEN,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-mlp-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                show(
                    "layer-0-router-logits",
                    &executor.scratch.router,
                    EXPERTS as usize,
                );
                executor.encode_moe(prefix, "0").unwrap();
                show(
                    "layer-0-router-probabilities",
                    &executor.scratch.router,
                    EXPERTS as usize,
                );
                let mut route_ids = [0_i32; EXPERTS_USED];
                let mut route_weights = [0.0_f32; EXPERTS_USED];
                executor.scratch.route_ids.read_i32(&mut route_ids).unwrap();
                executor
                    .scratch
                    .route_weights
                    .read_f32(&mut route_weights)
                    .unwrap();
                eprintln!("layer-0-routes: {route_ids:?} {route_weights:?}");
                let commands = Commands::begin().unwrap();
                executor.routed_experts(prefix).unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-expert-gate",
                    &executor.scratch.gate,
                    (EXPERT_WIDTH as usize) * EXPERTS_USED,
                );
                show(
                    "layer-0-expert-up",
                    &executor.scratch.up,
                    (EXPERT_WIDTH as usize) * EXPERTS_USED,
                );
                show(
                    "layer-0-expert-mid",
                    &executor.scratch.mid,
                    (EXPERT_WIDTH as usize) * EXPERTS_USED,
                );
                show(
                    "layer-0-expert-down",
                    &executor.scratch.expert,
                    (HIDDEN as usize) * EXPERTS_USED,
                );
                show("layer-0-routed", &executor.scratch.moe, HIDDEN as usize);
                let commands = Commands::begin().unwrap();
                executor.shared_expert(prefix).unwrap();
                commands.finish().unwrap();
                show("layer-0-shared-gate-raw", &executor.scratch.gate, 1);
                show(
                    "layer-0-shared-up",
                    &executor.scratch.up,
                    EXPERT_WIDTH as usize,
                );
                show(
                    "layer-0-shared-mid",
                    &executor.scratch.mid,
                    EXPERT_WIDTH as usize,
                );
                show(
                    "layer-0-shared-down",
                    &executor.scratch.shared,
                    HIDDEN as usize,
                );
                show("layer-0-total", &executor.scratch.moe, HIDDEN as usize);
                let commands = Commands::begin().unwrap();
                executor
                    .affine_mv_into(
                        &executor
                            .affine(
                                &format!("{prefix}.mlp.shared_expert.gate_proj"),
                                HIDDEN,
                                EXPERT_WIDTH,
                                None,
                            )
                            .unwrap(),
                        &executor.scratch.block,
                        &executor.scratch.gate,
                        HIDDEN,
                        EXPERT_WIDTH,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-0-shared-gate-proj",
                    &executor.scratch.gate,
                    EXPERT_WIDTH as usize,
                );
                show("layer-0-output", &executor.scratch.hc, HC_WIDTH as usize);
            } else if layer == 1 {
                let prefix = "language_model.model.layers.1";
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-1-attn-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                show(
                    "layer-1-attn-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.gdn(prefix, layer).unwrap();
                commands.finish().unwrap();
                show("layer-1-gdn-qkv", &executor.scratch.qkv, GDN_QKV as usize);
                show(
                    "layer-1-gdn-z",
                    &executor.scratch.controls,
                    GDN_VALUE as usize,
                );
                let b = executor
                    .scratch
                    .controls
                    .view(u64::from(GDN_VALUE) * 4, u64::from(GDN_HEADS_V) * 4)
                    .unwrap();
                let a = executor
                    .scratch
                    .controls
                    .view(
                        u64::from(GDN_VALUE + GDN_HEADS_V) * 4,
                        u64::from(GDN_HEADS_V) * 4,
                    )
                    .unwrap();
                show("layer-1-gdn-b", &b, GDN_HEADS_V as usize);
                show("layer-1-gdn-a", &a, GDN_HEADS_V as usize);
                show(
                    "layer-1-attn-out",
                    &executor.scratch.hidden,
                    HIDDEN as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                show(
                    "layer-1-after-attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.mlp_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-1-mlp-hc-norm",
                    &executor.scratch.hc_norm,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-1-mlp-hc-rank",
                    &executor.scratch.rank,
                    HC_RANK as usize,
                );
                let commands = Commands::begin().unwrap();
                let up = executor
                    .weight(&format!(
                        "{prefix}.mlp_hyper_connection.input_mix_weight_up.weight"
                    ))
                    .unwrap();
                executor
                    .bf16_mv(
                        up,
                        &executor.scratch.rank,
                        &executor.scratch.qkv,
                        HC_RANK,
                        HC_WIDTH,
                    )
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-1-mlp-hc-up",
                    &executor.scratch.qkv,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-1-mlp-hc-mix",
                    &executor.scratch.hc_mix,
                    HC_WIDTH as usize,
                );
                show(
                    "layer-1-mlp-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                executor.encode_moe(prefix, "1").unwrap();
                show(
                    "layer-1-mlp-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                show("layer-1-mlp-out", &executor.scratch.moe, HIDDEN as usize);
            } else if layer == 2 {
                let prefix = "language_model.model.layers.2";
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-2-attn-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                show(
                    "layer-2-attn-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.gdn(prefix, layer).unwrap();
                commands.finish().unwrap();
                show(
                    "layer-2-attn-out",
                    &executor.scratch.hidden,
                    HIDDEN as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                show(
                    "layer-2-after-attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.mlp_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-2-mlp-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                executor.encode_moe(prefix, "2").unwrap();
                show(
                    "layer-2-router-probabilities",
                    &executor.scratch.router,
                    EXPERTS as usize,
                );
                show(
                    "layer-2-route-weights",
                    &executor.scratch.route_weights,
                    EXPERTS_USED,
                );
                let mut route_ids = [0_i32; EXPERTS_USED];
                executor.scratch.route_ids.read_i32(&mut route_ids).unwrap();
                let route_bytes = route_ids
                    .iter()
                    .flat_map(|value| value.to_le_bytes())
                    .collect::<Vec<_>>();
                std::fs::write("/private/tmp/ds4_layer2_route_ids.i32", route_bytes).unwrap();
                eprintln!("layer-2-routes: {route_ids:?}");
                let up_table = executor
                    .affine_table(
                        "language_model.model.layers.2.mlp.switch_mlp.up_proj",
                        HIDDEN,
                        EXPERT_WIDTH,
                    )
                    .unwrap();
                let mut probe_args = args();
                probe_args.u[0] = HIDDEN;
                probe_args.u[1] = EXPERT_WIDTH;
                probe_args.u[2] = up_table.bits;
                probe_args.u[3] = up_table.group;
                probe_args.u[4] = EXPERTS_USED as u32;
                let commands = Commands::begin().unwrap();
                executor
                    .dispatch(
                        c"kernel_qwen_probe_affine_table",
                        &executor.scratch.hc_norm,
                        None,
                        Some(&executor.scratch.route_ids),
                        None,
                        &[
                            executor.view(up_table.packed),
                            executor.view(up_table.scales),
                            executor.view(up_table.biases),
                        ],
                        &probe_args,
                        EXPERTS_USED as u32,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                let mut gpu_probe = [0_i32; EXPERTS_USED * 3];
                executor.scratch.hc_norm.read_i32(&mut gpu_probe).unwrap();
                let tensor_word = |weight: Weight<'_>, index: usize, bytes: usize| {
                    let data = executor.model.tensor_bytes(weight.tensor).unwrap();
                    let start = index * bytes;
                    match bytes {
                        2 => u16::from_le_bytes(data[start..start + 2].try_into().unwrap()) as u32,
                        4 => u32::from_le_bytes(data[start..start + 4].try_into().unwrap()),
                        _ => unreachable!(),
                    }
                };
                let packed_columns = HIDDEN as usize / (32 / up_table.bits as usize);
                let groups_per_row = HIDDEN as usize / up_table.group as usize;
                let mut cpu_probe = [0_u32; EXPERTS_USED * 3];
                for (slot, expert) in route_ids.iter().copied().enumerate() {
                    let row = expert as usize * EXPERT_WIDTH as usize;
                    cpu_probe[slot * 3] = tensor_word(up_table.packed, row * packed_columns, 4);
                    cpu_probe[slot * 3 + 1] = tensor_word(up_table.scales, row * groups_per_row, 2);
                    cpu_probe[slot * 3 + 2] = tensor_word(up_table.biases, row * groups_per_row, 2);
                }
                eprintln!("layer-2-up-probe-cpu: {cpu_probe:?}");
                eprintln!("layer-2-up-probe-gpu: {gpu_probe:?}");
                let commands = Commands::begin().unwrap();
                executor.routed_experts(prefix).unwrap();
                commands.finish().unwrap();
                show(
                    "layer-2-expert-gate",
                    &executor.scratch.gate,
                    EXPERT_WIDTH as usize * EXPERTS_USED,
                );
                show(
                    "layer-2-expert-up",
                    &executor.scratch.up,
                    EXPERT_WIDTH as usize * EXPERTS_USED,
                );
                show(
                    "layer-2-expert-mid",
                    &executor.scratch.mid,
                    EXPERT_WIDTH as usize * EXPERTS_USED,
                );
                show(
                    "layer-2-expert-down",
                    &executor.scratch.expert,
                    HIDDEN as usize * EXPERTS_USED,
                );
                show("layer-2-routed", &executor.scratch.moe, HIDDEN as usize);
                let commands = Commands::begin().unwrap();
                executor.shared_expert(prefix).unwrap();
                commands.finish().unwrap();
                show("layer-2-shared-gate-raw", &executor.scratch.gate, 1);
                show(
                    "layer-2-shared-up",
                    &executor.scratch.up,
                    EXPERT_WIDTH as usize,
                );
                show(
                    "layer-2-shared-mid",
                    &executor.scratch.mid,
                    EXPERT_WIDTH as usize,
                );
                show(
                    "layer-2-shared-down",
                    &executor.scratch.shared,
                    HIDDEN as usize,
                );
                show("layer-2-total", &executor.scratch.moe, HIDDEN as usize);
                show("layer-2-mlp-out", &executor.scratch.moe, HIDDEN as usize);
            } else if layer == 3 {
                let prefix = "language_model.model.layers.3";
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                executor
                    .attention(prefix, &executor.states[3], executor.position)
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-3-attn-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                show(
                    "layer-3-attn-inject",
                    &executor.scratch.injection,
                    HC as usize,
                );
                show("layer-3-q", &executor.scratch.q, ATTN_WIDTH as usize);
                show(
                    "layer-3-gate",
                    &executor.scratch.q_gate,
                    ATTN_WIDTH as usize,
                );
                show(
                    "layer-3-k",
                    &executor.scratch.k_rope,
                    ATTN_KV_WIDTH as usize,
                );
                show("layer-3-v", &executor.scratch.v, ATTN_KV_WIDTH as usize);
                show(
                    "layer-3-attention",
                    &executor.scratch.attention,
                    ATTN_WIDTH as usize,
                );
                show("layer-3-hidden", &executor.scratch.hidden, HIDDEN as usize);
                show(
                    "layer-3-attn-out",
                    &executor.scratch.hidden,
                    HIDDEN as usize,
                );
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                show(
                    "layer-3-after-attn",
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.mlp_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                show(
                    "layer-3-mlp-mixed",
                    &executor.scratch.block,
                    HIDDEN as usize,
                );
                executor.encode_moe(prefix, "3").unwrap();
                show("layer-3-mlp-out", &executor.scratch.moe, HIDDEN as usize);
            } else {
                executor.encode_layer(layer).unwrap();
            }
            if matches!(layer, 0 | 1 | 2 | 3 | 47) {
                show(
                    &format!("layer-{layer}"),
                    &executor.scratch.hc,
                    HC_WIDTH as usize,
                );
            }
            eprintln!(
                "layer-{layer}-ms: {:.3}",
                layer_started.elapsed().as_secs_f64() * 1_000.0
            );
        }
        show("final-widened", &executor.scratch.hc, HC_WIDTH as usize);
        executor.final_output().unwrap();
        show("final-mixed", &executor.scratch.block, HIDDEN as usize);
        show("logits", &executor.scratch.logits, VOCAB as usize);
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_affine_qmv_profile() {
        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 mut executor = QwenExecutor::open(model, 4).unwrap();
        executor.scratch.block.fill(0.0, HIDDEN.into()).unwrap();
        let projection = "language_model.model.layers.3.mlp.shared_expert.gate_proj";
        let weight = executor
            .affine(projection, HIDDEN, EXPERT_WIDTH, None)
            .unwrap();
        executor
            .affine_mv_into(
                &weight,
                &executor.scratch.block,
                &executor.scratch.gate,
                HIDDEN,
                EXPERT_WIDTH,
            )
            .unwrap();
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..1_000 {
            executor
                .affine_mv_into(
                    &weight,
                    &executor.scratch.block,
                    &executor.scratch.gate,
                    HIDDEN,
                    EXPERT_WIDTH,
                )
                .unwrap();
        }
        let encoded = started.elapsed();
        commands.finish().unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_QMV encode_us={:.3} total_us={:.3}",
            encoded.as_secs_f64() * 1_000.0,
            started.elapsed().as_secs_f64() * 1_000.0,
        );
        let prefix = "language_model.model.layers.3.attn_hyper_connection";
        executor.scratch.hc.fill(0.25, HC_WIDTH.into()).unwrap();
        executor.hyper_read(prefix).unwrap();
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..100 {
            executor.hyper_read(prefix).unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_HYPER_EXACT us={:.3}",
            started.elapsed().as_secs_f64() * 10_000.0
        );
        let pack = executor.prepare_hyper_pack(prefix).unwrap();
        executor.hyper_packs.insert(prefix.to_owned(), pack);
        executor.hyper_read(prefix).unwrap();
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        for _ in 0..100 {
            executor.hyper_read(prefix).unwrap();
        }
        commands.finish().unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_HYPER_V3 us={:.3}",
            started.elapsed().as_secs_f64() * 10_000.0
        );
    }

    #[test]
    #[ignore = "requires pinned Qwen embeddings and Apple Metal; stage profiler opt-in"]
    fn qwen_stage_counters_preserve_async_batches() {
        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 executor = QwenExecutor::open_configured_with_hyper(
            QwenModel::open(&root, 4).unwrap(),
            4,
            0,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        let count = u64::from(HC_WIDTH) * 4;
        let copied = Buffer::floats(count).unwrap();
        unsafe { ds4_gpu_busy_stats_reset() };
        let mut commands = Commands::begin().unwrap();
        executor.begin_tokens(&[1, 2, 3, 4]).unwrap();
        copied
            .copy_from(0, &executor.scratch.hc, 0, count * 4, "stage test copy")
            .unwrap();
        let mut values = args();
        values.u[0] = HIDDEN;
        executor
            .dispatch(
                c"kernel_qwen_repeat4",
                &executor.scratch.hc,
                Some(&executor.scratch.hidden),
                None,
                None,
                &[],
                &values,
                HC_WIDTH * 4,
                1,
            )
            .unwrap();
        commands.flush().unwrap();
        copied
            .copy_from(
                0,
                &executor.scratch.hc,
                0,
                count * 4,
                "stage test async copy",
            )
            .unwrap();
        commands.finish().unwrap();
        let mut stats = GpuBusyStats::default();
        unsafe { ds4_gpu_busy_stats_get(&mut stats) };
        assert_eq!(stats.command_buffers, 2);
        let mut hidden = vec![0.0; HIDDEN as usize * 4];
        let mut actual = vec![0.0; count as usize];
        executor.scratch.hidden.read_f32(&mut hidden).unwrap();
        copied.read_f32(&mut actual).unwrap();
        for (row, result) in actual.chunks_exact(HC_WIDTH as usize).enumerate() {
            for repeated in result.chunks_exact(HIDDEN as usize) {
                assert_eq!(
                    repeated,
                    &hidden[row * HIDDEN as usize..(row + 1) * HIDDEN as usize]
                );
                assert!(repeated.iter().all(|x| x.is_finite()));
            }
        }
    }

    #[test]
    #[ignore = "requires the pinned Qwen metadata and Apple Metal; no model inference"]
    fn qwen_lazy_target_probabilities_match_batched_sampling() {
        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 executor = QwenExecutor::open_configured_with_hyper(
            QwenModel::open(&root, 4).unwrap(),
            4,
            0,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        let mut state = 0x1234_5678_u32;
        let mut logits = (0..VOCAB * VERIFY_ROWS as u32)
            .map(|_| {
                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                (state >> 8) as f32 / 1_048_576.0 - 8.0
            })
            .collect::<Vec<_>>();
        // Broad mass in row 0; concentrated rows exercise actual nucleus and
        // min-p truncation instead of always keeping the whole top-k support.
        for row in 1..VERIFY_ROWS as usize {
            for (token, value) in [(7, 24.0), (17, 23.0), (27, 22.0)] {
                logits[row * VOCAB as usize + token] = value + row as f32;
            }
        }
        executor.scratch.logits.write_f32(&logits).unwrap();
        for (temperature, top_p, min_p, top_k) in [
            (1.0, 0.95, 0.0, 20),
            (0.5, 0.5, 0.2, 8),
            (1.0, 1.0, 0.0, 32),
            (0.0, 0.95, 0.0, 20),
            (1.0, 0.95, 0.0, 33),
        ] {
            let commands = Commands::begin().unwrap();
            let sparse = executor
                .encode_sparse_probabilities(
                    &executor.scratch.logits,
                    VERIFY_ROWS as u32,
                    temperature,
                    top_k,
                )
                .unwrap();
            commands.finish().unwrap();
            let batched = if sparse {
                executor
                    .read_sparse_probabilities(VERIFY_ROWS as u32, top_p, min_p, top_k)
                    .unwrap()
                    .unwrap()
            } else {
                logits
                    .chunks_exact(VOCAB as usize)
                    .map(|row| qwen_sampling_probabilities(row, temperature, top_p, min_p, top_k))
                    .collect()
            };
            for (row, expected) in batched.iter().enumerate() {
                assert_eq!(
                    executor
                        .target_probabilities_row(row, temperature, top_p, min_p, top_k)
                        .unwrap(),
                    *expected,
                    "row={row} temperature={temperature} top_k={top_k}"
                );
            }
            // Rejection at each draft depth and full acceptance plus bonus:
            // laziness must consume the same draws and preserve their ordering.
            for reject_at in 0..=3 {
                let mut eager_rng = Rng::new_qwen(12345);
                let mut lazy_rng = Rng::new_qwen(12345);
                let mut materialized = 0;
                for (row, eager) in batched.iter().take(3).enumerate() {
                    let lazy = executor
                        .target_probabilities_row(row, temperature, top_p, min_p, top_k)
                        .unwrap();
                    materialized += 1;
                    let (draft, proposal) = if row == reject_at {
                        let missing = (0..VOCAB as usize)
                            .find(|token| !eager.iter().any(|(id, _)| id == token))
                            .unwrap();
                        (missing as i32, vec![(missing, 1.0)])
                    } else {
                        (eager[0].0 as i32, eager.clone())
                    };
                    let expected = exact_speculative_sample_from_probabilities(
                        eager,
                        &proposal,
                        draft,
                        &mut eager_rng,
                    );
                    let actual = exact_speculative_sample_from_probabilities(
                        &lazy,
                        &proposal,
                        draft,
                        &mut lazy_rng,
                    );
                    assert_eq!(actual, expected);
                    assert_eq!(actual.1, row != reject_at);
                    if !actual.1 {
                        break;
                    }
                }
                if reject_at == 3 {
                    let bonus = executor
                        .target_probabilities_row(3, temperature, top_p, min_p, top_k)
                        .unwrap();
                    materialized += 1;
                    assert_eq!(
                        sample_probabilities_f64(&bonus, &mut lazy_rng),
                        sample_probabilities_f64(&batched[3], &mut eager_rng)
                    );
                }
                assert_eq!(materialized, reject_at + 1);
                assert_eq!(eager_rng.unit_f64(), lazy_rng.unit_f64());
            }
        }
        let mut preserved = vec![0.0; logits.len()];
        executor.scratch.logits.read_f32(&mut preserved).unwrap();
        assert_eq!(preserved, logits);

        // Diagnostic only: alternate eager/lazy order on identical resident
        // logits, including dispatch, wait and readback. Not a model benchmark.
        for needed in 1..=VERIFY_ROWS as usize {
            let mut eager_times = Vec::new();
            let mut lazy_times = Vec::new();
            for trial in 0..33 {
                for lazy in [trial % 2 == 0, trial % 2 != 0] {
                    let started = Instant::now();
                    if lazy {
                        for row in 0..needed {
                            executor
                                .target_probabilities_row(row, 1.0, 0.95, 0.0, 20)
                                .unwrap();
                        }
                    } else {
                        let commands = Commands::begin().unwrap();
                        assert!(
                            executor
                                .encode_sparse_probabilities(
                                    &executor.scratch.logits,
                                    VERIFY_ROWS as u32,
                                    1.0,
                                    20,
                                )
                                .unwrap()
                        );
                        commands.finish().unwrap();
                        executor
                            .read_sparse_probabilities(VERIFY_ROWS as u32, 0.95, 0.0, 20)
                            .unwrap()
                            .unwrap();
                    }
                    let elapsed = started.elapsed().as_secs_f64() * 1_000.0;
                    if trial > 0 {
                        if lazy {
                            &mut lazy_times
                        } else {
                            &mut eager_times
                        }
                        .push(elapsed);
                    }
                }
            }
            eager_times.sort_by(f64::total_cmp);
            lazy_times.sort_by(f64::total_cmp);
            eprintln!(
                "QWEN_DISTRIBUTION_TIMING needed={needed} trials=32 eager_median_ms={:.6} lazy_median_ms={:.6}",
                (eager_times[15] + eager_times[16]) / 2.0,
                (lazy_times[15] + lazy_times[16]) / 2.0,
            );
        }
    }

    #[test]
    fn affine_split_k_matches_reference_shapes_and_scratch_bound() {
        for (out_dim, rows, expected) in [
            (GDN_QKV, 32, 1),
            (GDN_QKV, 64, 1),
            (EXPERTS, 32, 20),
            (EXPERTS, 64, 10),
            (EXPERTS, 2048, 1),
            (GDN_HEADS_V, 32, 40),
            (GDN_HEADS_V, 64, 40),
            (GDN_HEADS_V, 2048, 4),
            (1, 32, 40),
            (1, 2048, 8),
        ] {
            for rows in [rows - 1, rows] {
                let splits = affine_qmm_splits(HIDDEN, out_dim, rows, 64);
                assert_eq!(splits, expected, "N={out_dim} M={rows}");
                assert!(HIDDEN.is_multiple_of(splits * 64));
                if splits > 1 {
                    assert!(
                        u64::from(splits) * u64::from(rows) * u64::from(out_dim)
                            <= QMM_SPLIT_ELEMENTS
                    );
                }
            }
        }
    }

    #[test]
    #[ignore = "requires the pinned Qwen artifact set and Apple Metal"]
    fn qwen_fused_gdn_input_rows_match_separate_projections() {
        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, 32).unwrap();
        let mut executor = QwenExecutor::open_configured_with_hyper(
            model,
            32,
            2048,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        let prefix = "language_model.model.layers.0";
        let projections = [
            ("in_proj_qkv", GDN_QKV),
            ("in_proj_z", GDN_VALUE),
            ("in_proj_b", GDN_HEADS_V),
            ("in_proj_a", GDN_HEADS_V),
        ];
        // Load only these four projections, not the full resident model.
        for suffix in ["weight", "scales", "biases"] {
            let mut packed = Vec::new();
            for (name, _) in projections {
                let tensor = executor
                    .model
                    .tensor(&format!("{prefix}.linear_attn.{name}.{suffix}"))
                    .unwrap();
                let offset = packed.len();
                packed.resize(offset + (tensor.range.end - tensor.range.start) as usize, 0);
                executor
                    .model
                    .read_tensor_at(tensor, 0, &mut packed[offset..])
                    .unwrap();
            }
            let buffer = Buffer::bytes(packed.len() as u64).unwrap();
            buffer.write(0, &packed).unwrap();
            executor.weights.insert(
                format!("{prefix}.linear_attn.in_proj_fused.{suffix}"),
                buffer,
            );
        }
        for rows in [4_u32, 5, 32, 65, 2048] {
            executor.scratch.block = Buffer::floats(u64::from(rows * HIDDEN)).unwrap();
            executor.scratch.gdn_fused = Buffer::floats(u64::from(rows * GDN_FUSED_INPUT)).unwrap();
            executor.scratch.controls = Buffer::floats(u64::from(rows * GDN_CONTROLS)).unwrap();
            let qkv = Buffer::floats(u64::from(rows * GDN_QKV)).unwrap();
            let mut state = 0x1234_5678_u32;
            let input = (0..rows * HIDDEN)
                .map(|_| {
                    state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                    ((state >> 24) as i32 - 128) as f32 / 32.0
                })
                .collect::<Vec<_>>();
            executor.scratch.block.write_f32(&input).unwrap();
            let commands = Commands::begin().unwrap();
            executor.gdn_fused_input_rows(prefix, rows, &qkv).unwrap();
            commands.finish().unwrap();
            let mut fused = vec![0.0; (rows * GDN_FUSED_INPUT) as usize];
            executor.scratch.gdn_fused.read_f32(&mut fused).unwrap();
            let mut offset = 0;
            for (name, width) in projections {
                let weight = executor
                    .affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)
                    .unwrap();
                let output = Buffer::floats(u64::from(rows * width)).unwrap();
                let commands = Commands::begin().unwrap();
                // Unsplitted reference accumulation: fusion changes N, so the
                // independent projections must not select different split-K trees.
                executor
                    .affine_batch_into(
                        &weight,
                        &executor.scratch.block,
                        &output,
                        HIDDEN,
                        width,
                        rows,
                        false,
                    )
                    .unwrap();
                commands.finish().unwrap();
                let mut separate = vec![0.0; (rows * width) as usize];
                output.read_f32(&mut separate).unwrap();
                for row in 0..rows as usize {
                    for column in 0..width as usize {
                        assert_eq!(
                            fused[row * GDN_FUSED_INPUT as usize + offset + column],
                            f32::from_bits(
                                u32::from(bf16(separate[row * width as usize + column])) << 16
                            ),
                            "{name} rows={rows} row={row} column={column}"
                        );
                    }
                }
                offset += width as usize;
            }
            let mut actual_qkv = vec![0.0; (rows * GDN_QKV) as usize];
            let mut controls = vec![0.0; (rows * GDN_CONTROLS) as usize];
            qkv.read_f32(&mut actual_qkv).unwrap();
            executor.scratch.controls.read_f32(&mut controls).unwrap();
            for row in 0..rows as usize {
                let start = row * GDN_FUSED_INPUT as usize;
                assert_eq!(
                    &actual_qkv[row * GDN_QKV as usize..(row + 1) * GDN_QKV as usize],
                    &fused[start..start + GDN_QKV as usize]
                );
                assert_eq!(
                    &controls[row * GDN_CONTROLS as usize..(row + 1) * GDN_CONTROLS as usize],
                    &fused[start + GDN_QKV as usize..start + GDN_FUSED_INPUT as usize]
                );
            }
            eprintln!("GDN fused input/unpack rows={rows}: exact match");
        }
    }

    #[test]
    #[ignore = "requires the pinned Qwen artifact set and Apple Metal"]
    fn qwen_affine_qmv_and_qmm_match_mlx_reference() {
        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, 32).unwrap();
        let executor = QwenExecutor::open_configured_with_hyper(
            model,
            32,
            0,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        let projection = "language_model.model.layers.0.linear_attn.in_proj_qkv";
        let input = Buffer::floats(u64::from(HIDDEN)).unwrap();
        let mut state = 0x1234_5678_u32;
        let values = (0..HIDDEN)
            .map(|_| {
                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                ((state >> 24) as i32 - 128) as f32 / 32.0
            })
            .collect::<Vec<_>>();
        input.write_f32(&values).unwrap();
        let batch = Buffer::floats(u64::from(GDN_QKV)).unwrap();
        {
            let weight = executor.affine(projection, HIDDEN, GDN_QKV, None).unwrap();
            let commands = Commands::begin().unwrap();
            executor
                .affine_batch_into(&weight, &input, &batch, HIDDEN, GDN_QKV, 1, true)
                .unwrap();
            commands.finish().unwrap();
        }
        let mut actual = vec![0.0; GDN_QKV as usize];
        batch.read_f32(&mut actual).unwrap();
        let digest = |values: &[f32]| {
            let mut hash = Sha256::new();
            for value in values {
                hash.update(value.to_le_bytes());
            }
            hash.finalize()
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect::<String>()
        };
        assert_eq!(
            digest(&actual),
            "8952fadfdb1ef4450fd23b0060d62ea9b118367733014985a94c2d815a2be203"
        );

        // Component diagnostics generated through MTPLX's Qwen projection
        // modules by tools/qwen-qmm-reference.py, not a full-model oracle.
        let router = "language_model.model.layers.0.mlp.gate";
        let control = "language_model.model.layers.0.linear_attn.in_proj_b";
        let shared_gate = "language_model.model.layers.0.mlp.shared_expert_gate";
        let mut mismatches = Vec::new();
        for (projection, out_dim, rows, expected) in [
            (
                projection,
                GDN_QKV,
                2,
                "0c7b939af0828bfeaf3d4ca68f099ab4fb086ba8e85a541f6908181d459965a1",
            ),
            (
                projection,
                GDN_QKV,
                3,
                "e08f1156aada307117fb8ee2239f6b1a426ce3a4dd46708c97f7cfee459c977a",
            ),
            (
                projection,
                GDN_QKV,
                4,
                "3e7208003c3862f5c7885687570b63ad34ad12db64b000c00d006764f959b804",
            ),
            (
                router,
                EXPERTS,
                2,
                "fdb3b452855432df5f24f9b6a8edbaf617c3b2f8602c1268ebcf8bda7f2cd405",
            ),
            (
                router,
                EXPERTS,
                3,
                "0401c505701fbb89264ed700f2ea645403334a4fd8851531110ef1052f1d7634",
            ),
            (
                router,
                EXPERTS,
                4,
                "efc592819cff2ab804dc0b6a96ad4cfab2288c19b91d2f4d22a61c8043096b65",
            ),
            (
                control,
                GDN_HEADS_V,
                2,
                "98d78a9497397ce1d5530f187ed3bd1578421873442c142a5a416622b9462c62",
            ),
            (
                control,
                GDN_HEADS_V,
                3,
                "6663bc7695a8a25ec6c84ab1bccd18a722313eceaecd741f7c38d2ed0de59f67",
            ),
            (
                control,
                GDN_HEADS_V,
                4,
                "ff72d3959cf71e35a6819fb8988d1e2a02667b0f5d6aa988e1990daf916d2fdf",
            ),
            (
                shared_gate,
                1,
                2,
                "21039f5cc6e4a9daf8fb42a1536cfecdbf34fa52664a7a5ee569accf67e27d84",
            ),
            (
                shared_gate,
                1,
                3,
                "60d311720d130f368bc0bc1bc1ad639c34e9a81f799dab4b9559184f80859b43",
            ),
            (
                shared_gate,
                1,
                4,
                "08666f58ac460ffb8321ad9b488b1d87fad378e9a9eb768350b6104256ecb198",
            ),
            (
                projection,
                GDN_QKV,
                32_u32,
                "1d6076d65844f0c918d40085fea3af9873e86bf818f03ea4f705ac0211134fdf",
            ),
            (
                projection,
                GDN_QKV,
                64_u32,
                "b595beefdc06cc3b8a8408770e0d2b9fad609a6f0701bb94a511afebe60963a8",
            ),
            (
                projection,
                GDN_QKV,
                2048,
                "a379c62625a3a30e24c04d4769e4397046f08eb766f6cc5c97f31e952a6d74d0",
            ),
            (
                router,
                EXPERTS,
                32,
                "a90d762711f4d1790cb51da622a51f8b006b35765d135023802d9bc265335354",
            ),
            (
                router,
                EXPERTS,
                64,
                "ca0f5c0e66423e243ce42efcfd0fdc66239a175098b5f0ff348ae6a31c569dda",
            ),
            (
                router,
                EXPERTS,
                2048,
                "01df7fe984eba144d7b0d8bd9cb141523a064187c555c4e7e1540d5f848c7574",
            ),
            (
                control,
                GDN_HEADS_V,
                32,
                "ce617fb6b4f99d78657babf59f117fb20ffde6394a1061f7951738213d3f5d88",
            ),
            (
                control,
                GDN_HEADS_V,
                64,
                "2eb144aea6e13618f1185505c27262244a403ae631e43296955dffd94da9c8d1",
            ),
            (
                control,
                GDN_HEADS_V,
                2048,
                "53ddb40a6126337ef9b6953891c4d64a50fa62751c499bcb34b41f9875367bc3",
            ),
            (
                shared_gate,
                1,
                32,
                "449e04b448011791f2fce0d66ccb456b47d279e0c6d2f3c649ae7c7202f675cc",
            ),
            (
                shared_gate,
                1,
                64,
                "79070bd2b9599660473250728980bbff6771638bcd516e7caecdf674fb7f80f9",
            ),
            (
                shared_gate,
                1,
                2048,
                "a172b606e52564e89c7f47ece757c4f29c8b55e966658dd5221cce1debeae328",
            ),
        ] {
            let qmm_values = values
                .iter()
                .map(|&value| f32::from_bits(u32::from(bf16(value)) << 16))
                .cycle()
                .take(rows as usize * HIDDEN as usize)
                .collect::<Vec<_>>();
            let input = Buffer::floats(u64::from(rows) * u64::from(HIDDEN)).unwrap();
            input.write_f32(&qmm_values).unwrap();
            let batch = Buffer::floats(u64::from(rows) * u64::from(out_dim)).unwrap();
            {
                let weight = executor.affine(projection, HIDDEN, out_dim, None).unwrap();
                let commands = Commands::begin().unwrap();
                executor
                    .affine_batch_into(&weight, &input, &batch, HIDDEN, out_dim, rows, true)
                    .unwrap();
                commands.finish().unwrap();
            }
            let mut actual = vec![0.0; rows as usize * out_dim as usize];
            batch.read_f32(&mut actual).unwrap();
            let actual_digest = digest(&actual);
            eprintln!("QMM {projection} rows={rows}: {actual_digest}");
            if actual_digest != expected {
                mismatches.push(format!(
                    "{projection} rows={rows}: {actual_digest}, expected {expected}"
                ));
            }

            {
                // Isolated operator timing, not serving throughput. One warmup
                // then 32 samples, including dispatch and completion wait.
                let weight = executor.affine(projection, HIDDEN, out_dim, None).unwrap();
                let mut times = Vec::new();
                let mut gpu_times = Vec::new();
                for trial in 0..33 {
                    unsafe { ds4_gpu_busy_stats_reset() };
                    let started = Instant::now();
                    let commands = Commands::begin().unwrap();
                    executor
                        .affine_batch_into(&weight, &input, &batch, HIDDEN, out_dim, rows, true)
                        .unwrap();
                    commands.finish().unwrap();
                    let elapsed = started.elapsed().as_secs_f64() * 1e6;
                    let mut stats = GpuBusyStats::default();
                    unsafe { ds4_gpu_busy_stats_get(&mut stats) };
                    assert!(stats.command_buffers > 0 && stats.busy_seconds > 0.0);
                    if trial > 0 {
                        times.push(elapsed);
                        gpu_times.push(stats.busy_seconds * 1e6);
                    }
                }
                times.sort_by(f64::total_cmp);
                gpu_times.sort_by(f64::total_cmp);
                eprintln!(
                    "AFFINE {projection} rows={rows} median_us={:.3} median_gpu_us={:.3}",
                    (times[15] + times[16]) / 2.0,
                    (gpu_times[15] + gpu_times[16]) / 2.0
                );
            }
            if rows <= 4 {
                continue;
            }

            // A partial tile must preserve the caller's unused output rows.
            batch
                .fill(-12345.0, u64::from(rows) * u64::from(out_dim))
                .unwrap();
            let weight = executor.affine(projection, HIDDEN, out_dim, None).unwrap();
            let commands = Commands::begin().unwrap();
            executor
                .affine_batch_into(&weight, &input, &batch, HIDDEN, out_dim, rows - 1, true)
                .unwrap();
            commands.finish().unwrap();
            let mut partial = vec![0.0; actual.len()];
            batch.read_f32(&mut partial).unwrap();
            let written = (rows - 1) as usize * out_dim as usize;
            assert_eq!(&partial[..written], &actual[..written]);
            assert!(partial[written..].iter().all(|&value| value == -12345.0));
        }

        assert!(mismatches.is_empty(), "{}", mismatches.join("\n"));

        let weight = executor
            .weight("language_model.model.layers.1.ple.key_proj.weight")
            .unwrap();
        for (rows, expected) in [
            (
                32_u32,
                "cd327b392235feb9837d0a18e3d4524f788d9a3b3cffaff03436e9280627a29e",
            ),
            (
                64_u32,
                "73f56c67ffb2edc6ee82c0491346223b6d8c0cd1b9fd59fa58447ba8bdcb03aa",
            ),
        ] {
            let qmm_values = values
                .iter()
                .map(|&value| f32::from_bits(u32::from(bf16(value)) << 16))
                .cycle()
                .take(rows as usize * HIDDEN as usize)
                .collect::<Vec<_>>();
            let input = Buffer::floats(u64::from(rows) * u64::from(HIDDEN)).unwrap();
            input.write_f32(&qmm_values).unwrap();
            let batch = Buffer::floats(u64::from(rows) * u64::from(GDN_QKV)).unwrap();
            {
                let commands = Commands::begin().unwrap();
                executor
                    .bf16_batch(weight, &input, &batch, HIDDEN, GDN_QKV, rows)
                    .unwrap();
                commands.finish().unwrap();
            }
            let mut actual = vec![0.0; rows as usize * GDN_QKV as usize];
            batch.read_f32(&mut actual).unwrap();
            assert_eq!(digest(&actual), expected);
        }
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_sorted_routed_prefill_matches_grouped_output() {
        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 capacity = 2048;
        let model = QwenModel::open(&root, capacity).unwrap();
        let mut executor = QwenExecutor::open_configured_with_hyper(
            model,
            capacity,
            0,
            100,
            EngineSpeculativeSettings::default(),
            false,
            false,
        )
        .unwrap();
        executor.scratch = Scratch::new(capacity, capacity).unwrap();
        for rows in [5, 6, 7, 32, 64, 65, 137, 2048] {
            let inputs = (0..rows * HIDDEN)
                .map(|index| {
                    let value = ((index.wrapping_mul(17) % 257) as f32 - 128.0) / 64.0;
                    f32::from_bits(u32::from(bf16(value)) << 16)
                })
                .collect::<Vec<_>>();
            let routes = (0..rows as usize)
                .flat_map(|row| {
                    (0..EXPERTS_USED)
                        .map(move |slot| ((row / 100 * 13 + slot * 37) % EXPERTS as usize) as i32)
                })
                .collect::<Vec<_>>();
            let weights = (0..rows)
                .flat_map(|_| (0..EXPERTS_USED).map(|slot| (slot + 1) as f32 / 55.0))
                .collect::<Vec<_>>();
            executor.scratch.block.write_f32(&inputs).unwrap();
            executor.scratch.route_ids.write_i32(&routes).unwrap();
            executor.scratch.route_weights.write_f32(&weights).unwrap();

            let commands = Commands::begin().unwrap();
            executor.prepare_routed_expert_map(rows).unwrap();
            executor
                .routed_experts_grouped_rows("language_model.model.layers.0", rows)
                .unwrap();
            commands.finish().unwrap();
            let mut grouped = vec![0.0; rows as usize * HIDDEN as usize];
            executor.scratch.moe.read_f32(&mut grouped).unwrap();

            let commands = Commands::begin().unwrap();
            executor
                .routed_experts_rows("language_model.model.layers.0", rows)
                .unwrap();
            commands.finish().unwrap();
            let mut direct = vec![0.0; grouped.len()];
            executor.scratch.moe.read_f32(&mut direct).unwrap();

            let mut dot = 0.0_f64;
            let mut left = 0.0_f64;
            let mut right = 0.0_f64;
            let mut mean = 0.0_f64;
            let mut maximum = 0.0_f64;
            for (&expected, &actual) in grouped.iter().zip(&direct) {
                dot += f64::from(expected) * f64::from(actual);
                left += f64::from(expected).powi(2);
                right += f64::from(actual).powi(2);
                let error = f64::from((actual - expected).abs());
                mean += error;
                maximum = maximum.max(error);
            }
            let cosine = dot / (left.sqrt() * right.sqrt());
            mean /= grouped.len() as f64;
            eprintln!(
                "DS4SERVER_QWEN38_SORTED_MOE rows={rows} cosine={cosine:.9} mae={mean:.9} max={maximum:.9}"
            );
            assert!(cosine >= 0.995 && mean <= 0.001 && maximum <= 0.06);
            assert_eq!(grouped, direct);
        }
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_layer_profile() {
        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 depth = std::env::var("DS4SERVER_QWEN38_PROFILE_DEPTH")
            .ok()
            .and_then(|value| value.parse().ok())
            .unwrap_or(1_u32);
        let context = std::env::var("DS4SERVER_QWEN38_PROFILE_CONTEXT")
            .ok()
            .and_then(|value| value.parse().ok())
            .unwrap_or(depth + 1);
        let model = QwenModel::open(&root, context).unwrap();
        let mut executor = QwenExecutor::open(model, context).unwrap();
        executor.eval(1).unwrap();
        executor.position = depth;
        let whole_started = Instant::now();
        executor.eval_target(2).unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_WHOLE depth={depth} ms={:.3}",
            whole_started.elapsed().as_secs_f64() * 1_000.0
        );
        executor.position = depth;
        let commands = Commands::begin().unwrap();
        executor.begin_token(2).unwrap();
        commands.finish().unwrap();
        for layer in 0..LAYERS {
            if layer == 1 {
                let started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor.ple(2).unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_PLE ms={:.3}",
                    started.elapsed().as_secs_f64() * 1_000.0
                );
            }
            let started = Instant::now();
            if layer == 19 {
                let prefix = "language_model.model.layers.19";
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=attn_hyper ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor
                    .attention(prefix, &executor.states[layer], executor.position)
                    .unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=attention ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor.hyper_write().unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=attn_write ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor
                    .hyper_read(&format!("{prefix}.mlp_hyper_connection"))
                    .unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=mlp_hyper ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor
                    .affine_mv_into(
                        &executor
                            .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)
                            .unwrap(),
                        &executor.scratch.block,
                        &executor.scratch.router,
                        HIDDEN,
                        EXPERTS,
                    )
                    .unwrap();
                let mut softmax_args = args();
                softmax_args.u[0] = EXPERTS;
                softmax_args.u[12] = EXPERTS / 4;
                executor
                    .dispatch(
                        c"kernel_qwen_softmax_precise_512",
                        &executor.scratch.router,
                        Some(&executor.scratch.router),
                        None,
                        None,
                        &[],
                        &softmax_args,
                        1,
                        1,
                    )
                    .unwrap();
                let mut route_args = args();
                route_args.u[0] = EXPERTS;
                route_args.u[12] = EXPERTS;
                executor
                    .dispatch(
                        c"kernel_qwen_route_top10",
                        &executor.scratch.route_ids,
                        Some(&executor.scratch.route_weights),
                        Some(&executor.scratch.router),
                        None,
                        &[],
                        &route_args,
                        1,
                        1,
                    )
                    .unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=router ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor.routed_experts(prefix).unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=routed ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                executor.shared_expert(prefix).unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=shared ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
                let stage_started = Instant::now();
                let commands = Commands::begin().unwrap();
                let mut inject_args = args();
                inject_args.u[0] = HIDDEN;
                executor
                    .dispatch(
                        c"kernel_qwen_hyper_inject",
                        &executor.scratch.hc_norm,
                        Some(&executor.scratch.hc),
                        Some(&executor.scratch.moe),
                        Some(&executor.scratch.injection),
                        &[],
                        &inject_args,
                        HC_WIDTH,
                        1,
                    )
                    .unwrap();
                executor
                    .scratch
                    .hc
                    .copy_from(
                        0,
                        &executor.scratch.hc_norm,
                        0,
                        u64::from(HC_WIDTH) * 4,
                        "committing profiled Qwen MoE hyper streams",
                    )
                    .unwrap();
                commands.finish().unwrap();
                eprintln!(
                    "DS4SERVER_QWEN38_STAGE layer={layer} stage=mlp_write ms={:.3}",
                    stage_started.elapsed().as_secs_f64() * 1_000.0
                );
            } else {
                let commands = Commands::begin().unwrap();
                executor.encode_layer(layer).unwrap();
                commands.finish().unwrap();
            }
            eprintln!(
                "DS4SERVER_QWEN38_LAYER layer={layer} ms={:.3}",
                started.elapsed().as_secs_f64() * 1_000.0
            );
        }
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        executor.final_output().unwrap();
        commands.finish().unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_HEAD ms={:.3}",
            started.elapsed().as_secs_f64() * 1_000.0
        );
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_verify_layer_profile() {
        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_configured(&root, 8, true).unwrap();
        let settings = EngineSpeculativeSettings {
            glm_mtp: true,
            ..EngineSpeculativeSettings::default()
        };
        let mut executor = QwenExecutor::open_configured(model, 8, 0, 100, settings).unwrap();
        executor.eval_target(1).unwrap();
        executor.eval_target(2).unwrap();
        let tokens = [2, 3, 4, 5];
        let commands = Commands::begin().unwrap();
        executor.begin_tokens(&tokens).unwrap();
        commands.finish().unwrap();
        for layer in 0..LAYERS {
            let prefix = format!("language_model.model.layers.{layer}");
            let started = Instant::now();
            let commands = Commands::begin().unwrap();
            if layer == 1 {
                executor.ple_batch(&tokens).unwrap();
            }
            executor
                .hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), 4)
                .unwrap();
            match &executor.states[layer] {
                LayerState::Gdn { .. } => executor.gdn_rows(&prefix, layer, 4, true).unwrap(),
                state @ LayerState::Attention { .. } => executor
                    .attention_rows(&prefix, state, executor.position, 4)
                    .unwrap(),
            }
            executor.hyper_write_rows(4).unwrap();
            commands.finish().unwrap();
            let attention_elapsed = started.elapsed();
            let commands = Commands::begin().unwrap();
            executor.encode_moe_rows(&prefix, 4).unwrap();
            commands.finish().unwrap();
            if layer == 20 {
                eprintln!(
                    "DS4SERVER_QWEN38_VERIFY_STAGE layer={layer} attention_ms={:.3} moe_ms={:.3}",
                    attention_elapsed.as_secs_f64() * 1_000.0,
                    (started.elapsed() - attention_elapsed).as_secs_f64() * 1_000.0
                );
            }
            eprintln!(
                "DS4SERVER_QWEN38_VERIFY_LAYER layer={layer} ms={:.3}",
                started.elapsed().as_secs_f64() * 1_000.0
            );
        }
        let started = Instant::now();
        let commands = Commands::begin().unwrap();
        executor.final_output_rows(4).unwrap();
        commands.finish().unwrap();
        eprintln!(
            "DS4SERVER_QWEN38_VERIFY_HEAD ms={:.3}",
            started.elapsed().as_secs_f64() * 1_000.0
        );
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_target_submission_profile() {
        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 mut executor = QwenExecutor::open(model, 4).unwrap();
        for token in [1, 2] {
            let total = Instant::now();
            let commands = Commands::begin().unwrap();
            let encoding = Instant::now();
            executor.begin_token(token).unwrap();
            for layer in 0..LAYERS {
                if layer == 1 {
                    executor.ple(token).unwrap();
                }
                executor.encode_layer(layer).unwrap();
            }
            executor.final_output().unwrap();
            let encoding_ms = encoding.elapsed().as_secs_f64() * 1_000.0;
            let gpu = Instant::now();
            commands.finish().unwrap();
            let gpu_ms = gpu.elapsed().as_secs_f64() * 1_000.0;
            let read = Instant::now();
            executor
                .scratch
                .logits
                .read_f32(&mut executor.logits)
                .unwrap();
            eprintln!(
                "DS4SERVER_QWEN38_SUBMISSION token={token} encode_ms={encoding_ms:.3} gpu_ms={gpu_ms:.3} read_ms={:.3} total_ms={:.3}",
                read.elapsed().as_secs_f64() * 1_000.0,
                total.elapsed().as_secs_f64() * 1_000.0,
            );
        }
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_m5_performance_gate() {
        use crate::engine::ChatTurn;
        use std::sync::atomic::AtomicBool;

        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 context = std::env::var("DS4SERVER_QWEN38_CONTEXT")
            .ok()
            .map_or(4_096, |value| value.parse().unwrap());
        let measured = std::env::var("DS4SERVER_QWEN38_MEASURED")
            .ok()
            .map_or(128, |value| value.parse::<u32>().unwrap());
        let runs = std::env::var("DS4SERVER_QWEN38_RUNS")
            .ok()
            .map_or(3, |value| value.parse::<usize>().unwrap());
        assert!(runs > 0);
        let mtp = std::env::var_os("DS4SERVER_QWEN38_MTP").is_some();
        let model = QwenModel::open_configured(&root, context, mtp).unwrap();
        let mut content = std::env::var_os("DS4SERVER_QWEN38_PROMPT_FILE").map_or_else(
            || {
                Ok::<_, std::io::Error>(
                    "Write a concise Rust implementation of binary search and explain its invariants."
                        .to_owned(),
                )
            },
            |path| {
                std::fs::read_to_string(path).map(|input| {
                    format!(
                        "Give a concise summary of the following project README. Return the summary directly without asking questions.\n\n{input}"
                    )
                })
            },
        )
        .unwrap();
        if let Ok(repeat) = std::env::var("DS4SERVER_QWEN38_PROMPT_REPEAT") {
            content = content.repeat(repeat.parse().unwrap());
        }
        let prompt = model.render_conversation(
            "",
            &[ChatTurn {
                user: true,
                tool: false,
                system: false,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content,
            }],
            ReasoningMode::Direct,
        );
        let settings = EngineSpeculativeSettings {
            glm_mtp: mtp,
            glm_mtp_timing: std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some(),
            ..EngineSpeculativeSettings::default()
        };
        let mut executor = QwenExecutor::open_configured(model, context, 0, 100, settings).unwrap();
        let cancelled = AtomicBool::new(false);
        let mut prefill_rates = Vec::with_capacity(runs);
        let mut decode_rates = Vec::with_capacity(runs);
        let mut expected_hash = None;
        for run in 1..=runs {
            executor.reset().unwrap();
            let stats_before = executor.execution_stats();
            let prefill = Instant::now();
            assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len());
            let prefill_seconds = prefill.elapsed().as_secs_f64();
            let mut rng = Rng::new_qwen(12_345);
            let decode = Instant::now();
            let mut generated = Vec::new();
            let mut cycle_ms = Vec::new();
            let mut cycle_tokens = Vec::new();
            while generated.len() < measured as usize {
                let token = executor.take_pending_sampled_token().unwrap_or_else(|| {
                    sample_from_logits(executor.logits(), 1.0, 0.95, 0.0, 20, &mut rng)
                });
                let cycle_started = Instant::now();
                let cycle = executor
                    .eval_speculative_sampled(
                        token,
                        measured - generated.len() as u32,
                        ReasoningMode::Direct,
                        1.0,
                        0.95,
                        0.0,
                        20,
                        &mut rng,
                        &cancelled,
                    )
                    .unwrap();
                cycle_ms.push(cycle_started.elapsed().as_secs_f64() * 1_000.0);
                cycle_tokens.push(cycle.len());
                generated.extend(cycle);
            }
            let decode_seconds = decode.elapsed().as_secs_f64();
            let hash = Sha256::digest(
                generated
                    .iter()
                    .flat_map(|token| token.to_le_bytes())
                    .collect::<Vec<_>>(),
            )
            .iter()
            .map(|byte| format!("{byte:02x}"))
            .collect::<String>();
            if let Some(expected) = &expected_hash {
                assert_eq!(&hash, expected, "same-process token hash drift");
            } else {
                expected_hash = Some(hash.clone());
            }
            let prefill_tps = prompt.len() as f64 / prefill_seconds;
            let decode_tps = generated.len() as f64 / decode_seconds;
            prefill_rates.push(prefill_tps);
            decode_rates.push(decode_tps);
            let stats = executor.execution_stats();
            let first_ms = cycle_ms[0];
            let first_tokens = cycle_tokens[0];
            let steady_seconds = (decode_seconds - first_ms / 1_000.0).max(f64::EPSILON);
            let steady_tps = (generated.len() - first_tokens) as f64 / steady_seconds;
            cycle_ms.sort_by(f64::total_cmp);
            let p50_ms = cycle_ms[cycle_ms.len() / 2];
            let p95_ms = cycle_ms[cycle_ms.len().saturating_mul(95).div_ceil(100) - 1];
            eprintln!("DS4SERVER_QWEN38_TOKENS run={run} {generated:?}");
            eprintln!(
                "DS4SERVER_QWEN38_PERF run={run} mtp={mtp} context={context} prompt={} prefill_tps={prefill_tps:.6} measured={} decode_tps={decode_tps:.6} first_ms={first_ms:.6} steady_tps={steady_tps:.6} p50_ms={p50_ms:.6} p95_ms={p95_ms:.6} peak_rss_bytes={} token_sha256={hash} cycles={} drafted={} accepted={} verifier_passes={} verifier_ms={} repair_ms={} ple_entries={} ple_hits={} ple_misses={} ple_evictions={} ple_pread_bytes={} ple_pread_ms={}",
                prompt.len(),
                generated.len(),
                peak_rss_bytes(),
                stats.speculative_cycles,
                stats.drafted_tokens,
                stats.accepted_draft_tokens,
                stats.verifier_passes,
                stats.verifier_ms,
                stats.repair_ms,
                stats.ssd_cache_entries,
                stats
                    .ssd_cache_hits
                    .saturating_sub(stats_before.ssd_cache_hits),
                stats
                    .ssd_cache_misses
                    .saturating_sub(stats_before.ssd_cache_misses),
                stats
                    .ssd_cache_evictions
                    .saturating_sub(stats_before.ssd_cache_evictions),
                stats
                    .ssd_pread_bytes
                    .saturating_sub(stats_before.ssd_pread_bytes),
                stats.ssd_pread_ms.saturating_sub(stats_before.ssd_pread_ms),
            );
            if std::env::var_os("DS4SERVER_QWEN38_PROFILE_AFTER_PREFILL").is_some() {
                let commands = Commands::begin().unwrap();
                executor.begin_token(1).unwrap();
                commands.finish().unwrap();
                for layer in 0..LAYERS {
                    let started = Instant::now();
                    let commands = Commands::begin().unwrap();
                    if layer == 1 {
                        executor.ple(1).unwrap();
                    }
                    if layer == 19 {
                        let prefix = format!("language_model.model.layers.{layer}");
                        executor
                            .hyper_read(&format!("{prefix}.attn_hyper_connection"))
                            .unwrap();
                        executor
                            .attention(&prefix, &executor.states[layer], executor.position)
                            .unwrap();
                        executor.hyper_write().unwrap();
                        commands.finish().unwrap();
                        let attention_elapsed = started.elapsed();
                        let commands = Commands::begin().unwrap();
                        executor.encode_moe(&prefix, &layer.to_string()).unwrap();
                        commands.finish().unwrap();
                        eprintln!(
                            "DS4SERVER_QWEN38_POST_PREFILL_STAGE layer={layer} attention_ms={:.3} moe_ms={:.3}",
                            attention_elapsed.as_secs_f64() * 1_000.0,
                            (started.elapsed() - attention_elapsed).as_secs_f64() * 1_000.0,
                        );
                    } else {
                        executor.encode_layer(layer).unwrap();
                        commands.finish().unwrap();
                    }
                    eprintln!(
                        "DS4SERVER_QWEN38_POST_PREFILL layer={layer} ms={:.3}",
                        started.elapsed().as_secs_f64() * 1_000.0
                    );
                }
            }
        }
        prefill_rates.sort_by(f64::total_cmp);
        decode_rates.sort_by(f64::total_cmp);
        eprintln!(
            "DS4SERVER_QWEN38_MEDIAN mtp={mtp} runs={runs} context={context} prefill_tps={:.6} decode_tps={:.6} token_sha256={}",
            prefill_rates[runs / 2],
            decode_rates[runs / 2],
            expected_hash.unwrap(),
        );
    }

    #[test]
    #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"]
    fn qwen_mtp_matches_target_and_restores_its_cache() {
        fn digest(logits: &[f32]) -> [u8; 32] {
            let mut hash = Sha256::new();
            for value in logits {
                hash.update(value.to_bits().to_le_bytes());
            }
            hash.finalize().into()
        }

        fn digest_hex(values: &[f32]) -> String {
            digest(values)
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect()
        }

        fn target_state_digest(executor: &QwenExecutor) -> [u8; 32] {
            fn buffer(hash: &mut Sha256, buffer: &Buffer, bytes: u64) {
                let mut values = vec![0; usize::try_from(bytes).unwrap()];
                buffer.read(0, &mut values).unwrap();
                hash.update(values);
            }
            let mut hash = Sha256::new();
            hash.update(executor.position.to_le_bytes());
            for token in &executor.tokens {
                hash.update(token.to_le_bytes());
            }
            for value in &executor.logits {
                hash.update(value.to_bits().to_le_bytes());
            }
            for value in executor.ple_state.history {
                hash.update(value.to_le_bytes());
            }
            buffer(
                &mut hash,
                &executor.ple_state.conv,
                u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
            );
            buffer(&mut hash, &executor.scratch.hc, u64::from(HC_WIDTH) * 4);
            for state in &executor.states {
                match state {
                    LayerState::Gdn { conv, recurrent } => {
                        buffer(&mut hash, conv, u64::from(GDN_QKV) * 3 * 2);
                        buffer(
                            &mut hash,
                            recurrent,
                            u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4,
                        );
                    }
                    LayerState::Attention {
                        kv,
                        qsa_raw,
                        qsa_pooled,
                    } => {
                        let [kv_bytes, raw_bytes, pooled_bytes] =
                            attention_checkpoint_bytes(executor.position);
                        buffer(&mut hash, kv, kv_bytes);
                        buffer(&mut hash, qsa_raw, raw_bytes);
                        buffer(&mut hash, qsa_pooled, pooled_bytes);
                    }
                }
            }
            hash.finalize().into()
        }

        fn mtp_attention_digest(executor: &QwenExecutor) -> [u8; 32] {
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &executor.mtp.as_ref().unwrap().attention
            else {
                unreachable!()
            };
            let capacity = executor.context + MTP_TOKEN_RESERVE;
            let mut hash = Sha256::new();
            for (buffer, bytes) in [
                (kv, u64::from(capacity) * ATTN_KV_WIDTH as u64 * 4),
                (qsa_raw, u64::from(capacity) * QSA_DIM as u64 * 2),
                (
                    qsa_pooled,
                    u64::from(qsa_block_capacity(executor.context).unwrap()) * QSA_DIM as u64 * 2,
                ),
            ] {
                let mut values = vec![0; usize::try_from(bytes).unwrap()];
                buffer.read(0, &mut values).unwrap();
                hash.update(values);
            }
            hash.finalize().into()
        }

        fn committed_mtp_attention_digests(executor: &QwenExecutor) -> [[u8; 32]; 3] {
            let LayerState::Attention {
                kv,
                qsa_raw,
                qsa_pooled,
            } = &executor.mtp.as_ref().unwrap().attention
            else {
                unreachable!()
            };
            let mut digests = [[0; 32]; 3];
            for ((buffer, bytes), digest) in [kv, qsa_raw, qsa_pooled]
                .into_iter()
                .zip(attention_checkpoint_bytes(
                    executor.position.saturating_sub(1),
                ))
                .zip(&mut digests)
            {
                let mut values = vec![0; usize::try_from(bytes).unwrap()];
                buffer.read(0, &mut values).unwrap();
                *digest = Sha256::digest(values).into();
            }
            digests
        }

        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 settings = EngineSpeculativeSettings {
            glm_mtp: true,
            ..EngineSpeculativeSettings::default()
        };
        let model = QwenModel::open_configured(&root, 8, true).unwrap();
        let mut executor = QwenExecutor::open_configured(model, 8, 0, 100, settings).unwrap();
        executor.eval(1).unwrap();
        executor.eval(2).unwrap();
        let seeded = digest(executor.logits());
        assert_eq!(
            seeded,
            [
                0xc2, 0x00, 0xc4, 0x87, 0x20, 0x30, 0xe5, 0xbe, 0x98, 0xbf, 0x22, 0x1d, 0xa3, 0x52,
                0xf9, 0xc1, 0x9e, 0x9b, 0x15, 0xd5, 0x73, 0xb3, 0x12, 0x24, 0x00, 0xf1, 0xfe, 0x01,
                0xea, 0xba, 0xf8, 0x8e,
            ],
            "two-token target logits must match the direct MTPLX Family-Serve oracle",
        );
        let mut token = argmax(executor.logits());
        let mut target_tokens = Vec::new();
        let mut target_state_digests = Vec::new();
        for _ in 0..4 {
            target_tokens.push(token);
            executor.eval_target(token).unwrap();
            target_state_digests.push(target_state_digest(&executor));
            token = argmax(executor.logits());
        }
        executor.reset().unwrap();
        executor.eval(1).unwrap();
        executor.eval(2).unwrap();
        assert_eq!(digest(executor.logits()), seeded);
        let LayerState::Attention { kv, .. } = &executor.mtp.as_ref().unwrap().attention else {
            unreachable!()
        };
        let mut mtp_frontier = vec![0; ATTN_KV_WIDTH as usize * 4];
        kv.read(0, &mut mtp_frontier).unwrap();
        assert!(mtp_frontier.iter().any(|byte| *byte != 0));

        let pinned = executor
            .draft_greedy(target_tokens[0], 3, ReasoningMode::Direct)
            .unwrap();
        assert_eq!(
            pinned
                .iter()
                .map(|draft| (draft.token, digest(&draft.logits)))
                .collect::<Vec<_>>(),
            vec![
                (
                    59,
                    [
                        0x97, 0xc3, 0xcd, 0x28, 0x9b, 0x4f, 0xc6, 0xe4, 0x27, 0x3b, 0x4f, 0x9e,
                        0xb7, 0x0c, 0xee, 0x84, 0x08, 0xb7, 0xd9, 0x25, 0xd3, 0x4f, 0x6f, 0xa6,
                        0xf0, 0xcf, 0xe1, 0xfe, 0x5c, 0x90, 0x94, 0x25,
                    ],
                ),
                (
                    84,
                    [
                        0xe1, 0x2f, 0x01, 0xed, 0x8e, 0x93, 0x97, 0x9e, 0xe9, 0x37, 0x2c, 0x29,
                        0x86, 0x4f, 0x75, 0x62, 0x45, 0xe6, 0xa6, 0xc1, 0xd3, 0xdf, 0xf3, 0x94,
                        0xd3, 0x39, 0x8b, 0xe7, 0x5b, 0xf7, 0x0f, 0x6d,
                    ],
                ),
                (
                    19,
                    [
                        0x9a, 0x9c, 0xec, 0x9e, 0xdb, 0xad, 0xc6, 0xde, 0xf7, 0xdc, 0xfe, 0x46,
                        0x97, 0x65, 0x84, 0xbd, 0xaa, 0xfa, 0xf2, 0x9c, 0x06, 0x79, 0x60, 0x15,
                        0x72, 0xd0, 0x3b, 0x6a, 0x87, 0xca, 0xc7, 0x8c,
                    ],
                ),
            ]
        );
        let mut mtp_mixed = vec![0.0; HIDDEN as usize];
        executor.scratch.block.read_f32(&mut mtp_mixed).unwrap();
        assert_eq!(
            digest(&mtp_mixed),
            [
                0x71, 0x7c, 0x9a, 0x14, 0xc5, 0x14, 0xc0, 0x17, 0x13, 0xb6, 0xe0, 0x5a, 0xc8, 0x4e,
                0x9b, 0x64, 0xf0, 0x39, 0xd4, 0xc8, 0xac, 0x07, 0xfc, 0xd9, 0x40, 0xcb, 0x87, 0x3b,
                0x73, 0x0d, 0x72, 0x7c,
            ]
        );

        executor.capture_recurrent().unwrap();
        let (target_logits, _) = executor.eval_target_batch(&target_tokens, None).unwrap();
        assert_eq!(
            target_logits
                .iter()
                .map(|logits| digest_hex(logits))
                .collect::<Vec<_>>(),
            [
                "16c3853061e2e7bec6a526d1b78b55f000c1c20ac41b46923ca35067fe96370f",
                "4f6790bc259a5572e53098293e7f9d7d8fa3ed02945f0e708d12384026b1175d",
                "b33d082c969dc86a4b1cf4fbf59f37d595e5ad1c19e41aad1ca6dfe868ce148d",
                "f2f67462a96225d305594eb4d6158a0fe122daa41cdb450c59bd67a6459b30d5",
            ],
            "batched target logits must match the MTPLX verifier oracle",
        );
        let expected_hidden = [
            "a7028c4ff1b3f198e9ffdca11fe2cd9d7818347de9eb6023cd74f54e628da625",
            "e86befa3fccd4f09da7693f299d2435a102fcedfa7248375743cab50dff14f84",
            "ae7e8b65a4586b28ebc651362513823adf91df3273c88cc8398f29197e9ebd16",
            "f8a0138552c7160007f0f174a0be00739f98dca509556f405ee219cd9373e14c",
        ];
        let expected_next = [
            "fed33afcd846e6150285b146f0a28bacad5b2dc2e6c588151398c24a08af5bc3",
            "ad69cf9e54d0510b3d578f8fbceac555c69a4d96d8558181a0e272a48962fef0",
            "bc6db53c482bd54bf0f2b12dee8dbb275371478d866c7dcec0694082d1186db0",
            "52179bf6753eb79aa9e272d8c196aa64c7e9baa7b0d1aa5057e7d460d7b81ed6",
        ];
        for slot in (0..target_tokens.len()).rev() {
            let mut committed = vec![1, 2];
            committed.extend_from_slice(&target_tokens[..=slot]);
            executor
                .commit_verified_window(slot + 1, &committed, &target_logits[slot])
                .unwrap();
            assert_eq!(executor.tokens, committed);
            assert_eq!(executor.position, 3 + slot as u32);
            assert_eq!(digest(executor.logits()), digest(&target_logits[slot]));
            let mut hidden = vec![0.0; HC_WIDTH as usize];
            executor.scratch.hc.read_f32(&mut hidden).unwrap();
            assert_eq!(digest_hex(&hidden), expected_hidden[slot]);
            executor.eval_target(42).unwrap();
            assert_eq!(
                digest_hex(executor.logits()),
                expected_next[slot],
                "committed verifier slot {slot}"
            );
        }
        let fallback_tokens = [target_tokens[0], target_tokens[1]];
        executor
            .restore_or_replay(
                &fallback_tokens,
                &target_logits[1],
                2,
                2,
                [EOS_TOKEN; PLE_HISTORY],
                false,
            )
            .unwrap();
        assert_eq!(
            target_state_digest(&executor),
            target_state_digests[1],
            "ordinary-forward fallback state",
        );
        executor.reset().unwrap();
        executor.eval(1).unwrap();
        executor.eval(2).unwrap();

        let checkpoint =
            std::env::temp_dir().join(format!("ds4-qwen98-mtp-checkpoint-{}", std::process::id()));
        executor
            .save_checkpoint(&checkpoint, [98; 32], &mut |_| {})
            .unwrap();
        executor.reset().unwrap();
        assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        let LayerState::Attention { kv, .. } = &executor.mtp.as_ref().unwrap().attention else {
            unreachable!()
        };
        let mut restored = vec![0; mtp_frontier.len()];
        kv.read(0, &mut restored).unwrap();
        assert_eq!(restored, mtp_frontier);
        assert_eq!(executor.checkpoint_tag(), [98; 32]);

        let first = target_tokens[0];
        let checkpoint_mtp = committed_mtp_attention_digests(&executor);
        executor.eval(first).unwrap();
        let finalized_target = target_state_digest(&executor);
        let finalized_mtp = committed_mtp_attention_digests(&executor);
        assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        assert_eq!(committed_mtp_attention_digests(&executor), checkpoint_mtp);
        executor.pending_sampled = Some(first);
        executor.finalize_pending_sampled_token().unwrap();
        assert!(executor.pending_sampled.is_none());
        assert_eq!(target_state_digest(&executor), finalized_target);
        assert_eq!(committed_mtp_attention_digests(&executor), finalized_mtp);
        assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());

        let emitted = executor
            .eval_speculative_greedy(
                first,
                4,
                ReasoningMode::Direct,
                &std::sync::atomic::AtomicBool::new(false),
            )
            .unwrap();
        assert_eq!(emitted, target_tokens[..emitted.len()]);
        assert_eq!(
            digest(executor.logits()),
            digest(&target_logits[emitted.len() - 1])
        );
        let stats = executor.execution_stats();
        assert_eq!(stats.speculative_cycles, 1);
        assert_eq!(stats.drafted_tokens, 3);
        assert_eq!(stats.accepted_draft_tokens, (emitted.len() - 1) as u64);
        assert_eq!(stats.verifier_passes, 1);
        let mtp_before_eos = mtp_attention_digest(&executor);
        assert_eq!(
            executor
                .eval_speculative_greedy(
                    EOS_TOKEN,
                    4,
                    ReasoningMode::Direct,
                    &std::sync::atomic::AtomicBool::new(false),
                )
                .unwrap(),
            [EOS_TOKEN],
        );
        assert_eq!(mtp_attention_digest(&executor), mtp_before_eos);
        fs::remove_file(checkpoint).unwrap();
    }
}