mod checkpoint;
mod gpu;

use gpu::*;

use super::gguf::{F16, Gguf, Q8_0, Tensor as GgufTensor};
use super::{Model, ModelFamily};
use std::env;
use std::ffi::c_void;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::Path;
use std::ptr::NonNull;
use std::time::UNIX_EPOCH;

const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01";
const CHECKPOINT_VERSION: u32 = 1;
const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024;
const DEFAULT_PREFILL_CHUNK: u32 = 4096;

const SOURCES: [(&str, &str); 19] = [
    ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"),
    ("DS4_METAL_DENSE_SOURCE", "dense.metal"),
    ("DS4_METAL_MOE_SOURCE", "moe.metal"),
    ("DS4_METAL_DSV4_HC_SOURCE", "dsv4_hc.metal"),
    ("DS4_METAL_UNARY_SOURCE", "unary.metal"),
    ("DS4_METAL_DSV4_KV_SOURCE", "dsv4_kv.metal"),
    ("DS4_METAL_DSV4_ROPE_SOURCE", "dsv4_rope.metal"),
    ("DS4_METAL_DSV4_MISC_SOURCE", "dsv4_misc.metal"),
    ("DS4_METAL_ARGSORT_SOURCE", "argsort.metal"),
    ("DS4_METAL_CPY_SOURCE", "cpy.metal"),
    ("DS4_METAL_CONCAT_SOURCE", "concat.metal"),
    ("DS4_METAL_GET_ROWS_SOURCE", "get_rows.metal"),
    ("DS4_METAL_SUM_ROWS_SOURCE", "sum_rows.metal"),
    ("DS4_METAL_SOFTMAX_SOURCE", "softmax.metal"),
    ("DS4_METAL_REPEAT_SOURCE", "repeat.metal"),
    ("DS4_METAL_GLU_SOURCE", "glu.metal"),
    ("DS4_METAL_NORM_SOURCE", "norm.metal"),
    ("DS4_METAL_BIN_SOURCE", "bin.metal"),
    ("DS4_METAL_SET_ROWS_SOURCE", "set_rows.metal"),
];

// The Metal boundary uses this only to decide whether diagnostic logs get ANSI
// color. App output is never a terminal, so Rust owns the one non-Metal symbol
// that ds4_metal.m expects without pulling in ds4.c.
#[unsafe(no_mangle)]
extern "C" fn ds4_log_is_tty(_stream: *mut c_void) -> bool {
    false
}

pub(crate) fn configure_sources() -> Result<(), String> {
    let bundled = env::current_exe().ok().and_then(|path| {
        path.parent()?
            .parent()
            .map(|path| path.join("Resources/metal"))
    });
    let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("metal");
    let directory = bundled
        .filter(|path| path.join("dense.metal").is_file())
        .unwrap_or(checkout);
    for (name, file) in SOURCES {
        let path = directory.join(file);
        if !path.is_file() {
            return Err(format!(
                "required Metal kernel is missing: {}",
                path.display()
            ));
        }
        // SAFETY: main calls this before Iced starts any application threads.
        unsafe { env::set_var(name, path) };
    }
    Ok(())
}

#[derive(Clone, Copy)]
struct Weight {
    offset: u64,
    kind: u32,
    bytes: u64,
    dims: [u64; 3],
}

impl Weight {
    fn bind(model: &Gguf, name: &str) -> Result<Self, String> {
        let tensor = model.tensor(name)?;
        Ok(Self::from_tensor(tensor))
    }

    fn optional(model: &Gguf, name: &str) -> Option<Self> {
        model.tensors.get(name).map(Self::from_tensor)
    }

    fn from_tensor(tensor: &GgufTensor) -> Self {
        let mut dims = [1; 3];
        for (to, from) in dims.iter_mut().zip(&tensor.dims) {
            *to = *from;
        }
        Self {
            offset: tensor.offset,
            kind: tensor.kind,
            bytes: tensor.bytes,
            dims,
        }
    }
}

struct Layer {
    hc_attn_fn: Weight,
    hc_attn_scale: Weight,
    hc_attn_base: Weight,
    attn_norm: Weight,
    attn_q_a: Weight,
    attn_q_a_norm: Weight,
    attn_q_b: Weight,
    attn_kv: Weight,
    attn_kv_norm: Weight,
    attn_sinks: Weight,
    attn_output_a: Weight,
    attn_output_b: Weight,
    attn_compressor: Option<CompressorWeights>,
    indexer: Option<IndexerWeights>,
    hc_ffn_fn: Weight,
    hc_ffn_scale: Weight,
    hc_ffn_base: Weight,
    ffn_norm: Weight,
    router: Weight,
    router_bias: Option<Weight>,
    router_hash: Option<Weight>,
    expert_gate: Weight,
    expert_up: Weight,
    expert_down: Weight,
    shared_gate: Weight,
    shared_up: Weight,
    shared_down: Weight,
}

#[derive(Clone, Copy)]
struct CompressorWeights {
    ape: Weight,
    kv: Weight,
    gate: Weight,
    norm: Weight,
}

#[derive(Clone, Copy)]
struct IndexerWeights {
    q: Weight,
    proj: Weight,
    compressor: CompressorWeights,
}

impl Layer {
    fn bind(model: &Gguf, index: u32) -> Result<Self, String> {
        let required = |suffix: &str| Weight::bind(model, &format!("blk.{index}.{suffix}"));
        let optional = |suffix: &str| Weight::optional(model, &format!("blk.{index}.{suffix}"));
        let attn_compressor = (index >= 2)
            .then(|| {
                Ok::<_, String>(CompressorWeights {
                    ape: required("attn_compressor_ape.weight")?,
                    kv: required("attn_compressor_kv.weight")?,
                    gate: required("attn_compressor_gate.weight")?,
                    norm: required("attn_compressor_norm.weight")?,
                })
            })
            .transpose()?;
        let indexer = (compression_ratio(index) == 4)
            .then(|| {
                Ok::<_, String>(IndexerWeights {
                    q: required("indexer.attn_q_b.weight")?,
                    proj: required("indexer.proj.weight")?,
                    compressor: CompressorWeights {
                        ape: required("indexer_compressor_ape.weight")?,
                        kv: required("indexer_compressor_kv.weight")?,
                        gate: required("indexer_compressor_gate.weight")?,
                        norm: required("indexer_compressor_norm.weight")?,
                    },
                })
            })
            .transpose()?;
        let layer = Self {
            hc_attn_fn: required("hc_attn_fn.weight")?,
            hc_attn_scale: required("hc_attn_scale.weight")?,
            hc_attn_base: required("hc_attn_base.weight")?,
            attn_norm: required("attn_norm.weight")?,
            attn_q_a: required("attn_q_a.weight")?,
            attn_q_a_norm: required("attn_q_a_norm.weight")?,
            attn_q_b: required("attn_q_b.weight")?,
            attn_kv: required("attn_kv.weight")?,
            attn_kv_norm: required("attn_kv_a_norm.weight")?,
            attn_sinks: required("attn_sinks.weight")?,
            attn_output_a: required("attn_output_a.weight")?,
            attn_output_b: required("attn_output_b.weight")?,
            attn_compressor,
            indexer,
            hc_ffn_fn: required("hc_ffn_fn.weight")?,
            hc_ffn_scale: required("hc_ffn_scale.weight")?,
            hc_ffn_base: required("hc_ffn_base.weight")?,
            ffn_norm: required("ffn_norm.weight")?,
            router: required("ffn_gate_inp.weight")?,
            router_bias: optional("exp_probs_b.bias"),
            router_hash: optional("ffn_gate_tid2eid.weight"),
            expert_gate: required("ffn_gate_exps.weight")?,
            expert_up: required("ffn_up_exps.weight")?,
            expert_down: required("ffn_down_exps.weight")?,
            shared_gate: required("ffn_gate_shexp.weight")?,
            shared_up: required("ffn_up_shexp.weight")?,
            shared_down: required("ffn_down_shexp.weight")?,
        };
        for (name, weight) in [
            ("attention q", layer.attn_q_b),
            ("attention output A", layer.attn_output_a),
            ("attention output B", layer.attn_output_b),
            ("shared gate", layer.shared_gate),
            ("shared up", layer.shared_up),
            ("shared down", layer.shared_down),
        ] {
            if weight.kind != Q8_0 {
                return Err(format!("Flash Metal path requires Q8_0 {name} weights"));
            }
        }
        if let Some(indexer) = layer.indexer
            && (!matches!(indexer.q.kind, F16 | Q8_0) || indexer.proj.kind != F16)
        {
            return Err("Flash Metal path requires F16/Q8 indexer weights".into());
        }
        Ok(layer)
    }
}

struct Weights {
    token_embedding: Weight,
    output_hc_base: Weight,
    output_hc_fn: Weight,
    output_hc_scale: Weight,
    output_norm: Weight,
    output: Weight,
    layers: Vec<Layer>,
}

impl Weights {
    fn bind(model: &Model) -> Result<Self, String> {
        if model.shape.family != ModelFamily::DeepSeek {
            return Err("the Rust Metal executor currently supports DeepSeek Flash only".into());
        }
        let main = &model.main;
        let output = Weight::bind(main, "output.weight")?;
        if output.kind != Q8_0 {
            return Err("Flash Metal path requires a Q8_0 output weight".into());
        }
        let layers = (0..model.shape.layers)
            .map(|index| Layer::bind(main, index))
            .collect::<Result<_, _>>()?;
        Ok(Self {
            token_embedding: Weight::bind(main, "token_embd.weight")?,
            output_hc_base: Weight::bind(main, "output_hc_base.weight")?,
            output_hc_fn: Weight::bind(main, "output_hc_fn.weight")?,
            output_hc_scale: Weight::bind(main, "output_hc_scale.weight")?,
            output_norm: Weight::bind(main, "output_norm.weight")?,
            output,
            layers,
        })
    }
}

struct Scratch {
    current_hc: Buffer,
    next_hc: Buffer,
    flat_hc: Buffer,
    hc_mix: Buffer,
    hc_split: Buffer,
    current: Buffer,
    norm: Buffer,
    q_rank: Buffer,
    q_rank_norm: Buffer,
    q: Buffer,
    kv_raw: Buffer,
    kv: Buffer,
    compressed_kv: Buffer,
    compressed_score: Buffer,
    compressed_stage: Buffer,
    indexer_q: Buffer,
    indexer_weights: Buffer,
    indexer_scores: Buffer,
    indexer_selected: Buffer,
    heads: Buffer,
    attention_low: Buffer,
    attention_out: Buffer,
    router_logits: Buffer,
    router_probs: Buffer,
    router_selected: Buffer,
    router_weights: Buffer,
    routed_gate: Buffer,
    routed_up: Buffer,
    routed_mid: Buffer,
    routed_experts: Buffer,
    routed_out: Buffer,
    shared_gate: Buffer,
    shared_up: Buffer,
    shared_mid: Buffer,
    shared_out: Buffer,
    output_pre: Buffer,
    output_weights: Buffer,
    output_embedding: Buffer,
    output_norm: Buffer,
    logits: Buffer,
}

struct BatchScratch {
    tokens: Buffer,
    current_hc: Buffer,
    next_hc: Buffer,
    after_attention_hc: Buffer,
    flat_hc: Buffer,
    hc_mix: Buffer,
    hc_split: Buffer,
    current: Buffer,
    norm: Buffer,
    q_rank: Buffer,
    q_rank_norm: Buffer,
    q: Buffer,
    q_half: Buffer,
    kv_raw: Buffer,
    kv: Buffer,
    compressed_kv: Buffer,
    compressed_score: Buffer,
    compressed_stage: Buffer,
    indexer_q: Buffer,
    indexer_weights: Buffer,
    indexer_scores: Buffer,
    indexer_selected: Buffer,
    heads: Buffer,
    attention_low: Buffer,
    attention_group_tmp: Buffer,
    attention_low_tmp: Buffer,
    attention_out: Buffer,
    router_logits: Buffer,
    router_probs: Buffer,
    router_selected: Buffer,
    router_weights: Buffer,
    routed_gate: Buffer,
    routed_up: Buffer,
    routed_mid: Buffer,
    routed_experts: Buffer,
    routed_out: Buffer,
    shared_gate: Buffer,
    shared_up: Buffer,
    shared_mid: Buffer,
    shared_out: Buffer,
}

impl BatchScratch {
    fn allocate(model: &Model, pos: u32, rows: u32) -> Result<Self, String> {
        let shape = model.shape;
        let rows = u64::from(rows);
        let hc_dim = shape.hc * shape.embd;
        let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
        let q_dim = shape.heads * shape.head_dim;
        let group_dim = shape.head_dim * (shape.heads / shape.out_groups);
        let low_dim = shape.out_groups * shape.lora_o;
        let routed = shape.experts_used * shape.ff_expert;
        Ok(Self {
            tokens: Buffer::bytes(rows * 4)?,
            current_hc: Buffer::floats(rows * hc_dim)?,
            next_hc: Buffer::floats(rows * hc_dim)?,
            after_attention_hc: Buffer::floats(rows * hc_dim)?,
            flat_hc: Buffer::floats(rows * hc_dim)?,
            hc_mix: Buffer::floats(rows * mix_hc)?,
            hc_split: Buffer::floats(rows * mix_hc)?,
            current: Buffer::floats(rows * shape.embd)?,
            norm: Buffer::floats(rows * shape.embd)?,
            q_rank: Buffer::floats(rows * shape.lora_q)?,
            q_rank_norm: Buffer::floats(rows * shape.lora_q)?,
            q: Buffer::floats(rows * q_dim)?,
            q_half: Buffer::bytes(rows * q_dim.max(shape.embd) * 2)?,
            kv_raw: Buffer::floats(rows * shape.head_dim)?,
            kv: Buffer::floats(rows * shape.head_dim)?,
            compressed_kv: Buffer::floats(rows * 2 * shape.head_dim)?,
            compressed_score: Buffer::floats(rows * 2 * shape.head_dim)?,
            compressed_stage: Buffer::floats(rows * shape.head_dim)?,
            indexer_q: Buffer::floats(rows * shape.indexer_heads * shape.indexer_head_dim)?,
            indexer_weights: Buffer::floats(rows * shape.indexer_heads)?,
            indexer_scores: Buffer::floats(
                rows * (u64::from(pos) + rows).div_ceil(4).saturating_add(2),
            )?,
            indexer_selected: Buffer::bytes(rows * shape.indexer_top_k * 4)?,
            heads: Buffer::floats(rows * q_dim)?,
            attention_low: Buffer::floats(rows * low_dim)?,
            attention_group_tmp: Buffer::floats(rows * group_dim)?,
            attention_low_tmp: Buffer::floats(rows * shape.lora_o)?,
            attention_out: Buffer::floats(rows * shape.embd)?,
            router_logits: Buffer::floats(rows * shape.experts)?,
            router_probs: Buffer::floats(rows * shape.experts)?,
            router_selected: Buffer::bytes(rows * shape.experts_used * 4)?,
            router_weights: Buffer::floats(rows * shape.experts_used)?,
            routed_gate: Buffer::floats(rows * routed)?,
            routed_up: Buffer::floats(rows * routed)?,
            routed_mid: Buffer::floats(rows * routed)?,
            routed_experts: Buffer::floats(rows * shape.experts_used * shape.embd)?,
            routed_out: Buffer::floats(rows * shape.embd)?,
            shared_gate: Buffer::floats(rows * shape.ff_expert)?,
            shared_up: Buffer::floats(rows * shape.ff_expert)?,
            shared_mid: Buffer::floats(rows * shape.ff_expert)?,
            shared_out: Buffer::floats(rows * shape.embd)?,
        })
    }
}

impl Scratch {
    fn allocate(model: &Model, context: u32) -> Result<Self, String> {
        let shape = model.shape;
        let hc_dim = shape.hc * shape.embd;
        let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
        let q_dim = shape.heads * shape.head_dim;
        let low_dim = shape.out_groups * shape.lora_o;
        let routed = shape.experts_used * shape.ff_expert;
        Ok(Self {
            current_hc: Buffer::floats(hc_dim)?,
            next_hc: Buffer::floats(hc_dim)?,
            flat_hc: Buffer::floats(hc_dim)?,
            hc_mix: Buffer::floats(mix_hc)?,
            hc_split: Buffer::floats(mix_hc)?,
            current: Buffer::floats(shape.embd)?,
            norm: Buffer::floats(shape.embd)?,
            q_rank: Buffer::floats(shape.lora_q)?,
            q_rank_norm: Buffer::floats(shape.lora_q)?,
            q: Buffer::floats(q_dim)?,
            kv_raw: Buffer::floats(shape.head_dim)?,
            kv: Buffer::floats(shape.head_dim)?,
            compressed_kv: Buffer::floats(2 * shape.head_dim)?,
            compressed_score: Buffer::floats(2 * shape.head_dim)?,
            compressed_stage: Buffer::floats(shape.head_dim)?,
            indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?,
            indexer_weights: Buffer::floats(shape.indexer_heads)?,
            indexer_scores: Buffer::floats(u64::from(context / 4 + 2))?,
            indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?,
            heads: Buffer::floats(q_dim)?,
            attention_low: Buffer::floats(low_dim)?,
            attention_out: Buffer::floats(shape.embd)?,
            router_logits: Buffer::floats(shape.experts)?,
            router_probs: Buffer::floats(shape.experts)?,
            router_selected: Buffer::bytes(shape.experts_used * 4)?,
            router_weights: Buffer::floats(shape.experts_used)?,
            routed_gate: Buffer::floats(routed)?,
            routed_up: Buffer::floats(routed)?,
            routed_mid: Buffer::floats(routed)?,
            routed_experts: Buffer::floats(shape.experts_used * shape.embd)?,
            routed_out: Buffer::floats(shape.embd)?,
            shared_gate: Buffer::floats(shape.ff_expert)?,
            shared_up: Buffer::floats(shape.ff_expert)?,
            shared_mid: Buffer::floats(shape.ff_expert)?,
            shared_out: Buffer::floats(shape.embd)?,
            output_pre: Buffer::floats(shape.hc)?,
            output_weights: Buffer::floats(shape.hc)?,
            output_embedding: Buffer::floats(shape.embd)?,
            output_norm: Buffer::floats(shape.embd)?,
            logits: Buffer::floats(shape.vocab)?,
        })
    }
}

/// Rust-owned DeepSeek Flash session. The ObjC boundary only schedules the
/// unchanged Metal kernels; token/session/graph state lives here.
pub(super) struct Session {
    scratch: Scratch,
    layers: Vec<LayerState>,
    context: u32,
    prefill_cap: u32,
    raw_cap: u32,
    position: u32,
}

struct LayerState {
    raw_cache: Buffer,
    compression: Option<CompressionState>,
    indexer: Option<CompressionState>,
}

struct CompressionState {
    ratio: u32,
    cache: Buffer,
    state_kv: Buffer,
    state_score: Buffer,
    rows: u32,
}

impl LayerState {
    fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result<Self, String> {
        let shape = model.shape;
        let ratio = compression_ratio(index);
        let compression = if ratio == 0 {
            None
        } else {
            let coefficient = if ratio == 4 { 2 } else { 1 };
            let width = coefficient * shape.head_dim;
            let state_rows = coefficient * ratio as u64;
            let state_kv = Buffer::floats(width * state_rows)?;
            let state_score = Buffer::floats(width * state_rows)?;
            state_kv.fill(0.0, width * state_rows)?;
            state_score.fill(f32::NEG_INFINITY, width * state_rows)?;
            Some(CompressionState {
                ratio,
                cache: Buffer::bytes((context / ratio + 2) as u64 * shape.head_dim * 2)?,
                state_kv,
                state_score,
                rows: 0,
            })
        };
        let indexer = if ratio == 4 {
            let width = 2 * shape.indexer_head_dim;
            let state_rows = 2 * u64::from(ratio);
            let state_kv = Buffer::floats(width * state_rows)?;
            let state_score = Buffer::floats(width * state_rows)?;
            state_kv.fill(0.0, width * state_rows)?;
            state_score.fill(f32::NEG_INFINITY, width * state_rows)?;
            Some(CompressionState {
                ratio,
                cache: Buffer::floats(u64::from(context / ratio + 2) * shape.indexer_head_dim)?,
                state_kv,
                state_score,
                rows: 0,
            })
        } else {
            None
        };
        Ok(Self {
            raw_cache: Buffer::floats(raw_cap as u64 * shape.head_dim)?,
            compression,
            indexer,
        })
    }
}

impl Session {
    fn new(model: &Model, context: u32, prefill_chunk: u32) -> Result<Self, String> {
        if context == 0 {
            return Err("context must contain at least one token".into());
        }
        let prefill_cap = prefill_chunk.max(1).min(context);
        let raw_cap = prefill_cap
            .saturating_add(model.shape.sliding_window as u32)
            .min(context)
            .div_ceil(256)
            .saturating_mul(256)
            .min(8192)
            .min(context)
            .max(1);
        let layers = (0..model.shape.layers)
            .map(|index| LayerState::allocate(model, index, context, raw_cap))
            .collect::<Result<_, _>>()?;
        Ok(Self {
            scratch: Scratch::allocate(model, context)?,
            layers,
            context,
            prefill_cap,
            raw_cap,
            position: 0,
        })
    }
}

// SAFETY: declaration order is required because Rust drops fields in order.
// `session` must release every Buffer before `_context` calls ds4_gpu_cleanup(),
// and `_context` must drop before `model` unmaps memory wrapped without copying
// by native/metal/ds4_metal.m:10329. This intentionally differs from
// ../ds4/ds4.c:56287-56288; do not reorder these fields to match it.
pub(super) struct Executor {
    weights: Weights,
    session: Session,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    quality: bool,
    checkpoint_tag: [u8; 32],
    model_modified: (u64, u32),
    _context: Context,
    model: Model,
}

impl Executor {
    pub(super) fn open(
        model: Model,
        context: u32,
        quality: bool,
        prefill_chunk: u32,
    ) -> Result<Self, String> {
        let weights = Weights::bind(&model)?;
        let context_handle = Context::open(&model, quality)?;
        let session = Session::new(
            &model,
            context,
            if prefill_chunk == 0 {
                DEFAULT_PREFILL_CHUNK
            } else {
                prefill_chunk
            },
        )?;
        let model_modified = fs::metadata(model.main.path())
            .and_then(|metadata| metadata.modified())
            .ok()
            .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
            .map(|duration| (duration.as_secs(), duration.subsec_nanos()))
            .unwrap_or_default();
        Ok(Self {
            weights,
            session,
            logits: vec![0.0; model.shape.vocab as usize],
            tokens: Vec::new(),
            quality,
            checkpoint_tag: [0; 32],
            model_modified,
            _context: context_handle,
            model,
        })
    }

    pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
        if token < 0 || token as u64 >= self.model.shape.vocab {
            return Err(format!("token {token} is outside the vocabulary"));
        }
        if self.session.position >= self.session.context {
            return Err(format!(
                "the Rust Metal executor currently supports {} tokens per session",
                self.session.context
            ));
        }
        let commands = Commands::begin()?;
        self.encode_token(token as u32)?;
        commands.finish()?;
        self.session.scratch.logits.read_f32(&mut self.logits)?;
        self.session.position += 1;
        self.tokens.push(token);
        Ok(())
    }

    pub(super) fn prefill(
        &mut self,
        tokens: &[i32],
        mut progress: impl FnMut(u32) -> bool,
    ) -> Result<usize, String> {
        if tokens.len() < 4 && self.session.position != 0 {
            for (index, &token) in tokens.iter().enumerate() {
                self.eval(token)?;
                if !progress(self.session.position) {
                    return Ok(index + 1);
                }
            }
            return Ok(tokens.len());
        }

        let end = self
            .session
            .position
            .checked_add(u32::try_from(tokens.len()).map_err(|_| "prefill is too large")?)
            .ok_or("prefill position overflow")?;
        if end > self.session.context {
            return Err(format!(
                "the Rust Metal executor currently supports {} tokens per session",
                self.session.context
            ));
        }
        let mut consumed = 0_usize;
        while consumed < tokens.len() {
            if !progress(self.session.position) {
                return Ok(consumed);
            }
            let pos = self.session.position;
            let mut cap = self.session.prefill_cap;
            if pos != 0 {
                cap = cap.min(self.session.raw_cap);
                let offset = pos % self.session.prefill_cap;
                if offset != 0 {
                    cap = cap.min(self.session.prefill_cap - offset);
                }
            }
            let rows = (tokens.len() - consumed).min(cap as usize);
            self.eval_batch(&tokens[consumed..consumed + rows])?;
            consumed += rows;
            if !progress(self.session.position) {
                return Ok(consumed);
            }
        }
        Ok(tokens.len())
    }

    fn eval_batch(&mut self, tokens: &[i32]) -> Result<(), String> {
        let rows = u32::try_from(tokens.len()).map_err(|_| "prefill batch is too large")?;
        if rows == 0 || rows > self.session.prefill_cap {
            return Err("prefill batch exceeds the configured prefill workspace".into());
        }
        if tokens
            .iter()
            .any(|token| *token < 0 || *token as u64 >= self.model.shape.vocab)
        {
            return Err("prefill contains a token outside the vocabulary".into());
        }
        let mut batch = BatchScratch::allocate(&self.model, self.session.position, rows)?;
        batch.tokens.write_i32(tokens)?;
        let map = self.model.main.map_ptr().cast();
        let size = self.model.main.len();
        let shape = self.model.shape;
        let pos = self.session.position;

        let commands = Commands::begin()?;
        call(
            unsafe {
                ds4_gpu_embed_tokens_hc_tensor(
                    batch.current_hc.raw(),
                    batch.tokens.raw(),
                    map,
                    size,
                    self.weights.token_embedding.offset,
                    shape.vocab as u32,
                    rows,
                    shape.embd as u32,
                    shape.hc as u32,
                )
            },
            "batch token embedding",
        )?;
        commands.finish()?;

        for (index, (weights, state)) in self
            .weights
            .layers
            .iter()
            .zip(&mut self.session.layers)
            .enumerate()
        {
            let commands = Commands::begin()?;
            encode_batch_layer(
                &batch,
                state,
                weights,
                shape,
                map,
                size,
                index as u32,
                pos,
                rows,
                self.session.raw_cap,
            )?;
            commands.finish()?;
            std::mem::swap(&mut batch.current_hc, &mut batch.next_hc);
        }

        let commands = Commands::begin()?;
        call(
            unsafe {
                ds4_gpu_tensor_copy(
                    self.session.scratch.current_hc.raw(),
                    0,
                    batch.current_hc.raw(),
                    u64::from(rows - 1) * shape.hc * shape.embd * 4,
                    shape.hc * shape.embd * 4,
                )
            },
            "selecting the final prefill row",
        )?;
        encode_output(&self.session.scratch, &self.weights, shape, map, size)?;
        commands.finish()?;
        self.session.scratch.logits.read_f32(&mut self.logits)?;
        self.session.position += rows;
        self.tokens.extend_from_slice(tokens);
        Ok(())
    }

    pub(super) fn logits(&self) -> &[f32] {
        &self.logits
    }

    pub(super) fn model(&self) -> &Model {
        &self.model
    }

    pub(super) fn context(&self) -> u32 {
        self.session.context
    }

    pub(super) fn position(&self) -> u32 {
        self.session.position
    }

    pub(super) fn reset(&mut self) -> Result<(), String> {
        self.session = Session::new(&self.model, self.session.context, self.session.prefill_cap)?;
        self.tokens.clear();
        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())
    }

    pub(super) fn tokens(&self) -> &[i32] {
        &self.tokens
    }

    pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
        self.checkpoint_tag
    }

    fn encode_token(&mut self, token: u32) -> Result<(), String> {
        let shape = self.model.shape;
        let map = self.model.main.map_ptr().cast();
        let size = self.model.main.len();
        let scratch = &mut self.session.scratch;
        call(
            unsafe {
                ds4_gpu_embed_token_hc_tensor(
                    scratch.current_hc.raw(),
                    map,
                    size,
                    self.weights.token_embedding.offset,
                    shape.vocab as u32,
                    token,
                    shape.embd as u32,
                    shape.hc as u32,
                )
            },
            "token embedding",
        )?;
        for (index, layer) in self.weights.layers.iter().enumerate() {
            let state = &mut self.session.layers[index];
            encode_layer(
                scratch,
                state,
                layer,
                shape,
                map,
                size,
                index as u32,
                self.session.position,
                token,
                self.session.raw_cap,
            )?;
        }
        encode_output(scratch, &self.weights, shape, map, size)
    }
}

#[allow(clippy::too_many_arguments)]
fn refresh_ratio4_compressor_state(
    s: &BatchScratch,
    state: &CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    rows: u32,
    head_dim: u32,
) -> Result<(), String> {
    if rows < 4 {
        return Ok(());
    }
    let width = 2 * u64::from(head_dim);
    let tail = s
        .norm
        .view(u64::from(rows - 4) * shape.embd * 4, 4 * shape.embd * 4)?;
    f16_rows(
        &s.compressed_kv,
        weights.kv,
        shape.embd,
        width,
        &tail,
        4,
        map,
        size,
    )?;
    f16_rows(
        &s.compressed_score,
        weights.gate,
        shape.embd,
        width,
        &tail,
        4,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_compressor_prefill_state_ratio4_tensor(
                state.state_kv.raw(),
                state.state_score.raw(),
                s.compressed_kv.raw(),
                s.compressed_score.raw(),
                map,
                size,
                weights.ape.offset,
                weights.ape.kind,
                head_dim,
                pos + rows - 4,
            )
        },
        "refreshing ratio-4 compressor state",
    )
}

#[allow(clippy::too_many_arguments)]
fn compress_attention_batch(
    s: &BatchScratch,
    state: &mut CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    rows: u32,
    original: u32,
    freq_base: f32,
    freq_scale: f32,
    ext: f32,
    attn_factor: f32,
) -> Result<u32, String> {
    let ratio = state.ratio;
    let chunk = rows / ratio;
    if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
        let before = if pos == 0 { 0 } else { state.rows };
        let target = s
            .compressed_stage
            .view(0, u64::from(chunk) * shape.head_dim * 4)?;
        let result = if pos == 0 || ratio != 4 {
            unsafe {
                ds4_gpu_compressor_prefill_tensor(
                    target.raw(),
                    state.state_kv.raw(),
                    state.state_score.raw(),
                    s.compressed_kv.raw(),
                    s.compressed_score.raw(),
                    map,
                    size,
                    weights.ape.offset,
                    weights.ape.kind,
                    weights.norm.offset,
                    weights.norm.kind,
                    shape.head_dim as u32,
                    ratio,
                    pos,
                    rows,
                    shape.rot as u32,
                    original,
                    true,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                    shape.rms_epsilon,
                )
            }
        } else {
            unsafe {
                ds4_gpu_compressor_prefill_ratio4_replay_tensor(
                    target.raw(),
                    state.state_kv.raw(),
                    state.state_score.raw(),
                    s.compressed_kv.raw(),
                    s.compressed_score.raw(),
                    map,
                    size,
                    weights.ape.offset,
                    weights.ape.kind,
                    weights.norm.offset,
                    weights.norm.kind,
                    shape.head_dim as u32,
                    pos,
                    rows,
                    shape.rot as u32,
                    original,
                    true,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                    shape.rms_epsilon,
                )
            }
        };
        call(result, "batch attention compression")?;
        if chunk != 0 {
            call(
                unsafe {
                    ds4_gpu_tensor_copy_f32_to_f16(
                        state.cache.raw(),
                        u64::from(before) * shape.head_dim * 2,
                        target.raw(),
                        0,
                        u64::from(chunk) * shape.head_dim,
                    )
                },
                "committing compressed attention KV",
            )?;
        }
        state.rows = before + chunk;
        if ratio == 4 {
            refresh_ratio4_compressor_state(
                s,
                state,
                weights,
                shape,
                map,
                size,
                pos,
                rows,
                shape.head_dim as u32,
            )?;
        }
    } else {
        let width = if ratio == 4 { 2 } else { 1 } * shape.head_dim;
        let target = s.compressed_stage.view(0, shape.head_dim * 4)?;
        for row in 0..rows {
            let absolute = pos + row;
            let emit = (absolute + 1).is_multiple_of(ratio);
            let kv = s
                .compressed_kv
                .view(u64::from(row) * width * 4, width * 4)?;
            let score = s
                .compressed_score
                .view(u64::from(row) * width * 4, width * 4)?;
            call(
                unsafe {
                    ds4_gpu_compressor_update_tensor(
                        kv.raw(),
                        score.raw(),
                        state.state_kv.raw(),
                        state.state_score.raw(),
                        target.raw(),
                        map,
                        size,
                        weights.ape.offset,
                        weights.ape.kind,
                        weights.norm.offset,
                        weights.norm.kind,
                        shape.head_dim as u32,
                        ratio,
                        absolute,
                        0,
                        shape.rot as u32,
                        original,
                        freq_base,
                        freq_scale,
                        ext,
                        attn_factor,
                        shape.rope_beta_fast,
                        shape.rope_beta_slow,
                        shape.rms_epsilon,
                        false,
                    )
                },
                "updating batched attention compression",
            )?;
            if emit {
                call(
                    unsafe {
                        ds4_gpu_dsv4_fp8_kv_quantize_tensor(
                            target.raw(),
                            1,
                            shape.head_dim as u32,
                            shape.rot as u32,
                        )
                    },
                    "rounding compressed attention KV",
                )?;
                call(
                    unsafe {
                        ds4_gpu_tensor_copy_f32_to_f16(
                            state.cache.raw(),
                            u64::from(state.rows) * shape.head_dim * 2,
                            target.raw(),
                            0,
                            shape.head_dim,
                        )
                    },
                    "committing compressed attention KV",
                )?;
                state.rows += 1;
            }
        }
    }
    Ok(state.rows)
}

#[allow(clippy::too_many_arguments)]
fn compress_index_batch(
    s: &BatchScratch,
    state: &mut CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    rows: u32,
    original: u32,
    freq_base: f32,
    freq_scale: f32,
    ext: f32,
    attn_factor: f32,
) -> Result<(), String> {
    let ratio = state.ratio;
    if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
        let before = if pos == 0 { 0 } else { state.rows };
        let chunk = rows / ratio;
        let target = state.cache.view(
            u64::from(before) * shape.indexer_head_dim * 4,
            u64::from(chunk) * shape.indexer_head_dim * 4,
        )?;
        let result = if pos == 0 {
            unsafe {
                ds4_gpu_compressor_prefill_tensor(
                    target.raw(),
                    state.state_kv.raw(),
                    state.state_score.raw(),
                    s.compressed_kv.raw(),
                    s.compressed_score.raw(),
                    map,
                    size,
                    weights.ape.offset,
                    weights.ape.kind,
                    weights.norm.offset,
                    weights.norm.kind,
                    shape.indexer_head_dim as u32,
                    ratio,
                    pos,
                    rows,
                    shape.rot as u32,
                    original,
                    false,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                    shape.rms_epsilon,
                )
            }
        } else {
            unsafe {
                ds4_gpu_compressor_prefill_ratio4_replay_tensor(
                    target.raw(),
                    state.state_kv.raw(),
                    state.state_score.raw(),
                    s.compressed_kv.raw(),
                    s.compressed_score.raw(),
                    map,
                    size,
                    weights.ape.offset,
                    weights.ape.kind,
                    weights.norm.offset,
                    weights.norm.kind,
                    shape.indexer_head_dim as u32,
                    pos,
                    rows,
                    shape.rot as u32,
                    original,
                    false,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                    shape.rms_epsilon,
                )
            }
        };
        call(result, "batch indexer compression")?;
        if chunk != 0 {
            call(
                unsafe {
                    ds4_gpu_dsv4_indexer_qat_tensor(
                        target.raw(),
                        chunk,
                        shape.indexer_head_dim as u32,
                    )
                },
                "batch compressed index rounding",
            )?;
        }
        state.rows = before + chunk;
        refresh_ratio4_compressor_state(
            s,
            state,
            weights,
            shape,
            map,
            size,
            pos,
            rows,
            shape.indexer_head_dim as u32,
        )?;
    } else {
        let width = 2 * shape.indexer_head_dim;
        for row in 0..rows {
            let absolute = pos + row;
            let emit = (absolute + 1).is_multiple_of(ratio);
            let kv = s
                .compressed_kv
                .view(u64::from(row) * width * 4, width * 4)?;
            let score = s
                .compressed_score
                .view(u64::from(row) * width * 4, width * 4)?;
            let target = state.cache.view(
                u64::from(state.rows) * shape.indexer_head_dim * 4,
                shape.indexer_head_dim * 4,
            )?;
            call(
                unsafe {
                    ds4_gpu_compressor_update_tensor(
                        kv.raw(),
                        score.raw(),
                        state.state_kv.raw(),
                        state.state_score.raw(),
                        target.raw(),
                        map,
                        size,
                        weights.ape.offset,
                        weights.ape.kind,
                        weights.norm.offset,
                        weights.norm.kind,
                        shape.indexer_head_dim as u32,
                        ratio,
                        absolute,
                        0,
                        shape.rot as u32,
                        original,
                        freq_base,
                        freq_scale,
                        ext,
                        attn_factor,
                        shape.rope_beta_fast,
                        shape.rope_beta_slow,
                        shape.rms_epsilon,
                        false,
                    )
                },
                "updating batched indexer compression",
            )?;
            if emit {
                call(
                    unsafe {
                        ds4_gpu_dsv4_indexer_qat_tensor(
                            target.raw(),
                            1,
                            shape.indexer_head_dim as u32,
                        )
                    },
                    "rounding compressed index",
                )?;
                state.rows += 1;
            }
        }
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn encode_batch_layer(
    s: &BatchScratch,
    state: &mut LayerState,
    w: &Layer,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    layer: u32,
    pos: u32,
    rows: u32,
    raw_cap: u32,
) -> Result<(), String> {
    let hc_dim = shape.hc * shape.embd;
    let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
    let q_dim = shape.heads * shape.head_dim;
    let compressed = layer >= 2;
    let ratio = compression_ratio(layer);
    let freq_base = if compressed {
        shape.compress_rope_base
    } else {
        shape.rope_base
    };
    let freq_scale = if compressed {
        1.0 / shape.rope_scale
    } else {
        1.0
    };
    let ext = if compressed && shape.rope_scale > 1.0 {
        1.0
    } else {
        0.0
    };
    let attn_factor = if ext != 0.0 {
        1.0 / (1.0 + 0.1 * (1.0 / freq_scale).ln())
    } else {
        1.0
    };
    let original = if compressed {
        shape.original_context as u32
    } else {
        0
    };

    call(
        unsafe {
            ds4_gpu_hc_rms_scale_project_f16_tensor(
                s.hc_mix.raw(),
                s.flat_hc.raw(),
                map,
                size,
                w.hc_attn_fn.offset,
                hc_dim as u32,
                mix_hc as u32,
                s.current_hc.raw(),
                rows,
                shape.rms_epsilon,
            )
        },
        "batch attention HC projection",
    )?;
    call(
        unsafe {
            ds4_gpu_hc_split_weighted_sum_norm_tensor(
                s.current.raw(),
                s.norm.raw(),
                s.hc_split.raw(),
                s.hc_mix.raw(),
                s.current_hc.raw(),
                map,
                size,
                w.hc_attn_scale.offset,
                w.hc_attn_base.offset,
                w.attn_norm.offset,
                shape.embd as u32,
                shape.hc as u32,
                shape.hc_sinkhorn as u32,
                shape.hc_epsilon,
                shape.rms_epsilon,
            )
        },
        "batch attention HC mix",
    )?;
    q8_rows(
        &s.q_rank,
        w.attn_q_a,
        shape.embd,
        shape.lora_q,
        &s.norm,
        rows,
        map,
        size,
    )?;
    q8_rows(
        &s.kv_raw,
        w.attn_kv,
        shape.embd,
        shape.head_dim,
        &s.norm,
        rows,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
                s.q_rank_norm.raw(),
                s.q_rank.raw(),
                map,
                size,
                w.attn_q_a_norm.offset,
                shape.lora_q as u32,
                s.kv.raw(),
                s.kv_raw.raw(),
                w.attn_kv_norm.offset,
                shape.head_dim as u32,
                rows,
                shape.rms_epsilon,
            )
        },
        "batch Q/KV norm",
    )?;
    let fused_q = unsafe {
        ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
            s.q.raw(),
            s.q_half.raw(),
            map,
            size,
            w.attn_q_b.offset,
            shape.lora_q,
            q_dim,
            s.q_rank_norm.raw(),
            rows,
            shape.heads as u32,
            shape.head_dim as u32,
            shape.rot as u32,
            pos,
            original,
            false,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
            shape.rope_beta_fast,
            shape.rope_beta_slow,
            shape.rms_epsilon,
        )
    } != 0;
    if !fused_q {
        q8_rows(
            &s.q,
            w.attn_q_b,
            shape.lora_q,
            q_dim,
            &s.q_rank_norm,
            rows,
            map,
            size,
        )?;
        call(
            unsafe {
                ds4_gpu_head_rms_norm_tensor(
                    s.q.raw(),
                    rows,
                    shape.heads as u32,
                    shape.head_dim as u32,
                    shape.rms_epsilon,
                )
            },
            "batch Q norm",
        )?;
        call(
            unsafe {
                ds4_gpu_rope_tail_tensor(
                    s.q.raw(),
                    rows,
                    shape.heads as u32,
                    shape.head_dim as u32,
                    shape.rot as u32,
                    pos,
                    original,
                    false,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                )
            },
            "batch Q RoPE",
        )?;
    }
    call(
        unsafe {
            ds4_gpu_rope_tail_tensor(
                s.kv.raw(),
                rows,
                shape.head_kv as u32,
                shape.head_dim as u32,
                shape.rot as u32,
                pos,
                original,
                false,
                freq_base,
                freq_scale,
                ext,
                attn_factor,
                shape.rope_beta_fast,
                shape.rope_beta_slow,
            )
        },
        "batch KV RoPE",
    )?;
    call(
        unsafe {
            ds4_gpu_dsv4_fp8_kv_quantize_tensor(
                s.kv.raw(),
                rows,
                shape.head_dim as u32,
                shape.rot as u32,
            )
        },
        "batch KV rounding",
    )?;
    call(
        unsafe {
            ds4_gpu_store_raw_kv_batch_tensor(
                state.raw_cache.raw(),
                s.kv.raw(),
                raw_cap,
                pos,
                rows,
                shape.head_dim as u32,
            )
        },
        "batch raw KV store",
    )?;

    let mut compressed_rows = 0;
    if let (Some(weights), Some(compression)) =
        (w.attn_compressor.as_ref(), state.compression.as_mut())
    {
        let width = if ratio == 4 { 2 } else { 1 } * shape.head_dim;
        f16_rows(
            &s.compressed_kv,
            weights.kv,
            shape.embd,
            width,
            &s.norm,
            rows,
            map,
            size,
        )?;
        f16_rows(
            &s.compressed_score,
            weights.gate,
            shape.embd,
            width,
            &s.norm,
            rows,
            map,
            size,
        )?;
        compressed_rows = compress_attention_batch(
            s,
            compression,
            *weights,
            shape,
            map,
            size,
            pos,
            rows,
            original,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
        )?;
    }

    if ratio == 4 {
        let weights = w
            .indexer
            .ok_or("ratio-4 layer is missing indexer weights")?;
        let indexer = state
            .indexer
            .as_mut()
            .ok_or("ratio-4 layer is missing indexer state")?;
        let width = 2 * shape.indexer_head_dim;
        f16_rows(
            &s.compressed_kv,
            weights.compressor.kv,
            shape.embd,
            width,
            &s.norm,
            rows,
            map,
            size,
        )?;
        f16_rows(
            &s.compressed_score,
            weights.compressor.gate,
            shape.embd,
            width,
            &s.norm,
            rows,
            map,
            size,
        )?;
        compress_index_batch(
            s,
            indexer,
            weights.compressor,
            shape,
            map,
            size,
            pos,
            rows,
            original,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
        )?;
        matmul_rows(
            &s.indexer_q,
            weights.q,
            shape.lora_q,
            shape.indexer_heads * shape.indexer_head_dim,
            &s.q_rank_norm,
            rows,
            map,
            size,
        )?;
        call(
            unsafe {
                ds4_gpu_rope_tail_tensor(
                    s.indexer_q.raw(),
                    rows,
                    shape.indexer_heads as u32,
                    shape.indexer_head_dim as u32,
                    shape.rot as u32,
                    pos,
                    original,
                    false,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                )
            },
            "batch indexer Q RoPE",
        )?;
        call(
            unsafe {
                ds4_gpu_dsv4_indexer_qat_tensor(
                    s.indexer_q.raw(),
                    rows * shape.indexer_heads as u32,
                    shape.indexer_head_dim as u32,
                )
            },
            "batch indexer Q rounding",
        )?;
        f16_rows(
            &s.indexer_weights,
            weights.proj,
            shape.embd,
            shape.indexer_heads,
            &s.norm,
            rows,
            map,
            size,
        )?;
    }

    if pos == 0 && compressed_rows == 0 {
        call(
            unsafe {
                ds4_gpu_attention_prefill_raw_heads_tensor(
                    s.heads.raw(),
                    map,
                    size,
                    w.attn_sinks.offset,
                    s.q.raw(),
                    s.kv.raw(),
                    rows,
                    shape.sliding_window as u32,
                    shape.heads as u32,
                    shape.head_dim as u32,
                )
            },
            "batch raw prefill attention",
        )?;
    } else if pos == 0 && ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
        let indexer = state
            .indexer
            .as_ref()
            .ok_or("ratio-4 layer is missing indexer state")?;
        call(
            unsafe {
                ds4_gpu_indexer_scores_prefill_tensor(
                    s.indexer_scores.raw(),
                    s.indexer_q.raw(),
                    s.indexer_weights.raw(),
                    indexer.cache.raw(),
                    compressed_rows,
                    rows,
                    shape.indexer_heads as u32,
                    shape.indexer_head_dim as u32,
                    ratio,
                    ((shape.indexer_heads * shape.indexer_head_dim) as f32)
                        .sqrt()
                        .recip(),
                )
            },
            "batch indexer scoring",
        )?;
        call(
            unsafe {
                ds4_gpu_indexer_topk_tensor(
                    s.indexer_selected.raw(),
                    s.indexer_scores.raw(),
                    compressed_rows,
                    rows,
                    shape.indexer_top_k as u32,
                )
            },
            "batch indexer top-k",
        )?;
        let compression = state
            .compression
            .as_ref()
            .ok_or("compressed layer is missing attention state")?;
        call(
            unsafe {
                ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
                    s.heads.raw(),
                    map,
                    size,
                    w.attn_sinks.offset,
                    s.q.raw(),
                    state.raw_cache.raw(),
                    compression.cache.raw(),
                    1,
                    s.indexer_selected.raw(),
                    rows,
                    pos,
                    rows,
                    raw_cap,
                    0,
                    compressed_rows,
                    shape.indexer_top_k as u32,
                    shape.sliding_window as u32,
                    ratio,
                    shape.heads as u32,
                    shape.head_dim as u32,
                )
            },
            "batch indexed prefill attention",
        )?;
    } else if pos == 0 {
        let compression = state
            .compression
            .as_ref()
            .ok_or("compressed layer is missing attention state")?;
        call(
            unsafe {
                ds4_gpu_attention_prefill_static_mixed_heads_tensor(
                    s.heads.raw(),
                    map,
                    size,
                    w.attn_sinks.offset,
                    s.q.raw(),
                    s.kv.raw(),
                    compression.cache.raw(),
                    1,
                    rows,
                    compressed_rows,
                    shape.sliding_window as u32,
                    ratio,
                    shape.heads as u32,
                    shape.head_dim as u32,
                )
            },
            "batch mixed prefill attention",
        )?;
    } else {
        let (n_raw, raw_start) = raw_batch_span(pos, rows, raw_cap, shape.sliding_window as u32);
        if ratio == 0 {
            call(
                unsafe {
                    ds4_gpu_attention_decode_raw_batch_heads_tensor(
                        s.heads.raw(),
                        map,
                        size,
                        w.attn_sinks.offset,
                        s.q.raw(),
                        state.raw_cache.raw(),
                        rows,
                        pos,
                        n_raw,
                        raw_cap,
                        raw_start,
                        shape.sliding_window as u32,
                        shape.heads as u32,
                        shape.head_dim as u32,
                    )
                },
                "batch resumed raw attention",
            )?;
        } else if ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
            let indexer = state
                .indexer
                .as_ref()
                .ok_or("ratio-4 layer is missing indexer state")?;
            call(
                unsafe {
                    ds4_gpu_indexer_scores_decode_batch_tensor(
                        s.indexer_scores.raw(),
                        s.indexer_q.raw(),
                        s.indexer_weights.raw(),
                        indexer.cache.raw(),
                        compressed_rows,
                        rows,
                        pos,
                        shape.indexer_heads as u32,
                        shape.indexer_head_dim as u32,
                        ratio,
                        ((shape.indexer_heads * shape.indexer_head_dim) as f32)
                            .sqrt()
                            .recip(),
                    )
                },
                "batch resumed indexer scoring",
            )?;
            call(
                unsafe {
                    ds4_gpu_indexer_topk_tensor(
                        s.indexer_selected.raw(),
                        s.indexer_scores.raw(),
                        compressed_rows,
                        rows,
                        shape.indexer_top_k as u32,
                    )
                },
                "batch resumed indexer top-k",
            )?;
            let compression = state
                .compression
                .as_ref()
                .ok_or("compressed layer is missing attention state")?;
            call(
                unsafe {
                    ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
                        s.heads.raw(),
                        map,
                        size,
                        w.attn_sinks.offset,
                        s.q.raw(),
                        state.raw_cache.raw(),
                        compression.cache.raw(),
                        1,
                        s.indexer_selected.raw(),
                        rows,
                        pos,
                        n_raw,
                        raw_cap,
                        raw_start,
                        compressed_rows,
                        shape.indexer_top_k as u32,
                        shape.sliding_window as u32,
                        ratio,
                        shape.heads as u32,
                        shape.head_dim as u32,
                    )
                },
                "batch resumed indexed attention",
            )?;
        } else {
            let compression = state
                .compression
                .as_ref()
                .ok_or("compressed layer is missing attention state")?;
            call(
                unsafe {
                    ds4_gpu_attention_decode_mixed_batch_heads_tensor(
                        s.heads.raw(),
                        map,
                        size,
                        w.attn_sinks.offset,
                        s.q.raw(),
                        state.raw_cache.raw(),
                        compression.cache.raw(),
                        1,
                        std::ptr::null(),
                        0,
                        rows,
                        pos,
                        n_raw,
                        raw_cap,
                        raw_start,
                        compressed_rows,
                        shape.sliding_window as u32,
                        ratio,
                        shape.heads as u32,
                        shape.head_dim as u32,
                    )
                },
                "batch resumed mixed attention",
            )?;
        }
    }
    call(
        unsafe {
            ds4_gpu_rope_tail_tensor(
                s.heads.raw(),
                rows,
                shape.heads as u32,
                shape.head_dim as u32,
                shape.rot as u32,
                pos,
                original,
                true,
                freq_base,
                freq_scale,
                ext,
                attn_factor,
                shape.rope_beta_fast,
                shape.rope_beta_slow,
            )
        },
        "batch inverse attention RoPE",
    )?;
    let group_dim = shape.head_dim * (shape.heads / shape.out_groups);
    let half_output = unsafe {
        ds4_gpu_attention_output_q8_batch_f16_tensor(
            s.q_half.raw(),
            s.attention_low.raw(),
            map,
            size,
            w.attn_output_a.offset,
            w.attn_output_b.offset,
            group_dim,
            shape.lora_o,
            shape.out_groups as u32,
            shape.embd,
            s.heads.raw(),
            rows,
        )
    } != 0;
    if half_output {
        call(
            unsafe {
                ds4_gpu_hc_expand_split_half_tensor(
                    s.after_attention_hc.raw(),
                    s.q_half.raw(),
                    s.current_hc.raw(),
                    s.hc_split.raw(),
                    shape.embd as u32,
                    shape.hc as u32,
                )
            },
            "batch attention HC expansion",
        )?;
    } else {
        call(
            unsafe {
                ds4_gpu_attention_output_q8_batch_tensor(
                    s.attention_out.raw(),
                    s.attention_low.raw(),
                    s.attention_group_tmp.raw(),
                    s.attention_low_tmp.raw(),
                    map,
                    size,
                    w.attn_output_a.offset,
                    w.attn_output_b.offset,
                    group_dim,
                    shape.lora_o,
                    shape.out_groups as u32,
                    shape.embd,
                    s.heads.raw(),
                    rows,
                )
            },
            "batch attention output",
        )?;
        call(
            unsafe {
                ds4_gpu_hc_expand_split_tensor(
                    s.after_attention_hc.raw(),
                    s.attention_out.raw(),
                    s.current_hc.raw(),
                    s.hc_split.raw(),
                    shape.embd as u32,
                    shape.hc as u32,
                )
            },
            "batch attention HC expansion",
        )?;
    }

    call(
        unsafe {
            ds4_gpu_hc_rms_scale_project_f16_tensor(
                s.hc_mix.raw(),
                s.flat_hc.raw(),
                map,
                size,
                w.hc_ffn_fn.offset,
                hc_dim as u32,
                mix_hc as u32,
                s.after_attention_hc.raw(),
                rows,
                shape.rms_epsilon,
            )
        },
        "batch FFN HC projection",
    )?;
    call(
        unsafe {
            ds4_gpu_hc_split_weighted_sum_norm_tensor(
                s.current.raw(),
                s.norm.raw(),
                s.hc_split.raw(),
                s.hc_mix.raw(),
                s.after_attention_hc.raw(),
                map,
                size,
                w.hc_ffn_scale.offset,
                w.hc_ffn_base.offset,
                w.ffn_norm.offset,
                shape.embd as u32,
                shape.hc as u32,
                shape.hc_sinkhorn as u32,
                shape.hc_epsilon,
                shape.rms_epsilon,
            )
        },
        "batch FFN HC mix",
    )?;
    f16_rows(
        &s.router_logits,
        w.router,
        shape.embd,
        shape.experts,
        &s.norm,
        rows,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_router_select_batch_tensor(
                s.router_selected.raw(),
                s.router_weights.raw(),
                s.router_probs.raw(),
                map,
                size,
                w.router_bias.map_or(0, |value| value.offset),
                w.router_hash.map_or(0, |value| value.offset),
                w.router_hash.map_or(0, |value| value.dims[1] as u32),
                0,
                0,
                w.router_bias.is_some(),
                w.router_hash.is_some(),
                s.router_logits.raw(),
                s.tokens.raw(),
                shape.experts as u32,
                shape.experts_used as u32,
                shape.expert_weight_scale,
                rows,
            )
        },
        "batch expert routing",
    )?;
    let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]);
    let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]);
    let mut mid_f16 = false;
    call(
        unsafe {
            ds4_gpu_routed_moe_batch_tensor(
                s.routed_out.raw(),
                s.routed_gate.raw(),
                s.routed_up.raw(),
                s.routed_mid.raw(),
                s.routed_experts.raw(),
                map,
                size,
                w.expert_gate.offset,
                w.expert_up.offset,
                w.expert_down.offset,
                w.expert_gate.kind,
                w.expert_down.kind,
                w.expert_gate.dims[1] * gate_row,
                gate_row,
                w.expert_down.dims[1] * down_row,
                down_row,
                w.expert_gate.dims[0] as u32,
                w.expert_down.dims[0] as u32,
                w.expert_down.dims[1] as u32,
                s.router_selected.raw(),
                s.router_weights.raw(),
                shape.experts as u32,
                shape.experts_used as u32,
                shape.swiglu_clamp,
                s.norm.raw(),
                layer,
                rows,
                &mut mid_f16,
                false,
            )
        },
        "batch routed experts",
    )?;
    q8_rows(
        &s.shared_gate,
        w.shared_gate,
        shape.embd,
        shape.ff_expert,
        &s.norm,
        rows,
        map,
        size,
    )?;
    q8_rows(
        &s.shared_up,
        w.shared_up,
        shape.embd,
        shape.ff_expert,
        &s.norm,
        rows,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_swiglu_tensor(
                s.shared_mid.raw(),
                s.shared_gate.raw(),
                s.shared_up.raw(),
                rows * shape.ff_expert as u32,
                shape.swiglu_clamp,
                1.0,
            )
        },
        "batch shared expert activation",
    )?;
    q8_rows(
        &s.shared_out,
        w.shared_down,
        shape.ff_expert,
        shape.embd,
        &s.shared_mid,
        rows,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_hc_expand_add_split_tensor(
                s.next_hc.raw(),
                s.routed_out.raw(),
                s.shared_out.raw(),
                s.after_attention_hc.raw(),
                s.hc_split.raw(),
                shape.embd as u32,
                shape.hc as u32,
            )
        },
        "batch FFN HC expansion",
    )
}

// Mirrors one fixed DS4 tape invocation; grouping these scalar dimensions would
// only hide the Metal kernel contract behind another one-use type.
#[allow(clippy::too_many_arguments)]
fn encode_layer(
    s: &Scratch,
    state: &mut LayerState,
    w: &Layer,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    layer: u32,
    pos: u32,
    token: u32,
    raw_cap: u32,
) -> Result<(), String> {
    let hc_dim = shape.hc * shape.embd;
    let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
    let compressed = layer >= 2;
    let freq_base = if compressed {
        shape.compress_rope_base
    } else {
        shape.rope_base
    };
    let freq_scale = if compressed {
        1.0 / shape.rope_scale
    } else {
        1.0
    };
    let ext = if compressed && shape.rope_scale > 1.0 {
        1.0
    } else {
        0.0
    };
    let attn_factor = if ext != 0.0 {
        1.0 / (1.0 + 0.1 * (1.0 / freq_scale).ln())
    } else {
        1.0
    };
    let original = if compressed {
        shape.original_context as u32
    } else {
        0
    };
    let raw_cache = state.raw_cache.raw();
    call(
        unsafe {
            ds4_gpu_rms_norm_plain_tensor(
                s.flat_hc.raw(),
                s.current_hc.raw(),
                hc_dim as u32,
                shape.rms_epsilon,
            )
        },
        "attention HC norm",
    )?;
    f16(
        &s.hc_mix,
        w.hc_attn_fn,
        hc_dim,
        mix_hc,
        &s.flat_hc,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_hc_split_weighted_sum_norm_tensor(
                s.current.raw(),
                s.norm.raw(),
                s.hc_split.raw(),
                s.hc_mix.raw(),
                s.current_hc.raw(),
                map,
                size,
                w.hc_attn_scale.offset,
                w.hc_attn_base.offset,
                w.attn_norm.offset,
                shape.embd as u32,
                shape.hc as u32,
                shape.hc_sinkhorn as u32,
                shape.hc_epsilon,
                shape.rms_epsilon,
            )
        },
        "attention HC mix",
    )?;
    call(
        unsafe {
            ds4_gpu_matmul_q8_0_pair_tensor(
                s.q_rank.raw(),
                s.kv_raw.raw(),
                map,
                size,
                w.attn_q_a.offset,
                w.attn_kv.offset,
                shape.embd,
                shape.lora_q,
                shape.head_dim,
                s.norm.raw(),
                1,
            )
        },
        "Q/KV projection",
    )?;
    call(
        unsafe {
            ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
                s.q_rank_norm.raw(),
                s.q_rank.raw(),
                map,
                size,
                w.attn_q_a_norm.offset,
                shape.lora_q as u32,
                s.kv.raw(),
                s.kv_raw.raw(),
                w.attn_kv_norm.offset,
                shape.head_dim as u32,
                1,
                shape.rms_epsilon,
            )
        },
        "Q/KV norm",
    )?;
    q8(
        &s.q,
        w.attn_q_b,
        shape.lora_q,
        shape.heads * shape.head_dim,
        &s.q_rank_norm,
        map,
        size,
    )?;
    let fused_q_rope = unsafe {
        ds4_gpu_head_rms_norm_rope_tail_tensor(
            s.q.raw(),
            1,
            shape.heads as u32,
            shape.head_dim as u32,
            shape.rot as u32,
            pos,
            original,
            false,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
            shape.rope_beta_fast,
            shape.rope_beta_slow,
            shape.rms_epsilon,
        )
    } != 0;
    if !fused_q_rope {
        call(
            unsafe {
                ds4_gpu_head_rms_norm_tensor(
                    s.q.raw(),
                    1,
                    shape.heads as u32,
                    shape.head_dim as u32,
                    shape.rms_epsilon,
                )
            },
            "Q norm",
        )?;
        call(
            unsafe {
                ds4_gpu_rope_tail_tensor(
                    s.q.raw(),
                    1,
                    shape.heads as u32,
                    shape.head_dim as u32,
                    shape.rot as u32,
                    pos,
                    original,
                    false,
                    freq_base,
                    freq_scale,
                    ext,
                    attn_factor,
                    shape.rope_beta_fast,
                    shape.rope_beta_slow,
                )
            },
            "Q RoPE",
        )?;
    }
    call(
        unsafe {
            ds4_gpu_rope_tail_tensor(
                s.kv.raw(),
                1,
                shape.head_kv as u32,
                shape.head_dim as u32,
                shape.rot as u32,
                pos,
                original,
                false,
                freq_base,
                freq_scale,
                ext,
                attn_factor,
                shape.rope_beta_fast,
                shape.rope_beta_slow,
            )
        },
        "KV RoPE",
    )?;
    let raw_row = pos % raw_cap;
    call(
        unsafe {
            ds4_gpu_kv_fp8_store_raw_tensor(
                s.kv.raw(),
                raw_cache,
                raw_cap,
                raw_row,
                shape.head_dim as u32,
                shape.rot as u32,
            )
        },
        "KV cache write",
    )?;
    let (n_raw, raw_start) = raw_decode_span(pos, raw_cap, shape.sliding_window as u32);
    if let (Some(weights), Some(compression)) =
        (w.attn_compressor.as_ref(), state.compression.as_mut())
    {
        update_compression(
            s,
            compression,
            *weights,
            shape,
            map,
            size,
            pos,
            original,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
        )?;
    }
    if let (Some(weights), Some(indexer)) = (w.indexer, state.indexer.as_mut()) {
        update_indexer_compression(
            s,
            indexer,
            weights.compressor,
            shape,
            map,
            size,
            pos,
            original,
            freq_base,
            freq_scale,
            ext,
            attn_factor,
        )?;
    }
    let (comp_cache, n_comp) = state
        .compression
        .as_ref()
        .map_or((std::ptr::null(), 0), |compression| {
            (compression.cache.raw().cast_const(), compression.rows)
        });
    let indexed = if n_comp > 1_024 {
        match (w.indexer, state.indexer.as_ref()) {
            (Some(weights), Some(indexer)) if indexer.rows > shape.indexer_top_k as u32 => {
                matmul(
                    &s.indexer_q,
                    weights.q,
                    shape.lora_q,
                    shape.indexer_heads * shape.indexer_head_dim,
                    &s.q_rank_norm,
                    map,
                    size,
                )?;
                call(
                    unsafe {
                        ds4_gpu_rope_tail_tensor(
                            s.indexer_q.raw(),
                            1,
                            shape.indexer_heads as u32,
                            shape.indexer_head_dim as u32,
                            shape.rot as u32,
                            pos,
                            original,
                            false,
                            freq_base,
                            freq_scale,
                            ext,
                            attn_factor,
                            shape.rope_beta_fast,
                            shape.rope_beta_slow,
                        )
                    },
                    "indexer Q RoPE",
                )?;
                call(
                    unsafe {
                        ds4_gpu_dsv4_indexer_qat_tensor(
                            s.indexer_q.raw(),
                            shape.indexer_heads as u32,
                            shape.indexer_head_dim as u32,
                        )
                    },
                    "indexer Q quantization",
                )?;
                f16(
                    &s.indexer_weights,
                    weights.proj,
                    shape.embd,
                    shape.indexer_heads,
                    &s.norm,
                    map,
                    size,
                )?;
                call(
                    unsafe {
                        ds4_gpu_indexer_score_one_tensor(
                            s.indexer_scores.raw(),
                            s.indexer_q.raw(),
                            s.indexer_weights.raw(),
                            indexer.cache.raw(),
                            indexer.rows,
                            shape.indexer_heads as u32,
                            shape.indexer_head_dim as u32,
                            ((shape.indexer_heads * shape.indexer_head_dim) as f32)
                                .sqrt()
                                .recip(),
                        )
                    },
                    "indexer scoring",
                )?;
                call(
                    unsafe {
                        ds4_gpu_indexer_topk_tensor(
                            s.indexer_selected.raw(),
                            s.indexer_scores.raw(),
                            indexer.rows,
                            1,
                            shape.indexer_top_k as u32,
                        )
                    },
                    "indexer top-k selection",
                )?;
                Some(indexer)
            }
            _ => None,
        }
    } else {
        None
    };
    if indexed.is_some() {
        call(
            unsafe {
                ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
                    s.heads.raw(),
                    map,
                    size,
                    w.attn_sinks.offset,
                    s.q.raw(),
                    raw_cache,
                    comp_cache,
                    1,
                    s.indexer_selected.raw(),
                    1,
                    pos,
                    n_raw,
                    raw_cap,
                    raw_start,
                    n_comp,
                    shape.indexer_top_k as u32,
                    shape.sliding_window as u32,
                    compression_ratio(layer),
                    shape.heads as u32,
                    shape.head_dim as u32,
                )
            },
            "indexed attention",
        )?;
    } else {
        call(
            unsafe {
                ds4_gpu_attention_decode_heads_tensor(
                    s.heads.raw(),
                    map,
                    size,
                    w.attn_sinks.offset,
                    s.q.raw(),
                    raw_cache,
                    n_raw,
                    raw_cap,
                    raw_start,
                    comp_cache,
                    1,
                    n_comp,
                    std::ptr::null(),
                    0,
                    shape.heads as u32,
                    shape.head_dim as u32,
                )
            },
            "attention",
        )?;
    }
    call(
        unsafe {
            ds4_gpu_rope_tail_tensor(
                s.heads.raw(),
                1,
                shape.heads as u32,
                shape.head_dim as u32,
                shape.rot as u32,
                pos,
                original,
                true,
                freq_base,
                freq_scale,
                ext,
                attn_factor,
                shape.rope_beta_fast,
                shape.rope_beta_slow,
            )
        },
        "inverse attention RoPE",
    )?;
    let group_dim = shape.head_dim * (shape.heads / shape.out_groups);
    call(
        unsafe {
            ds4_gpu_attention_output_low_q8_tensor(
                s.attention_low.raw(),
                map,
                size,
                w.attn_output_a.offset,
                group_dim,
                shape.lora_o,
                shape.out_groups as u32,
                s.heads.raw(),
            )
        },
        "attention low projection",
    )?;
    call(
        unsafe {
            ds4_gpu_matmul_q8_0_hc_expand_tensor(
                s.next_hc.raw(),
                s.attention_out.raw(),
                map,
                size,
                w.attn_output_b.offset,
                shape.out_groups * shape.lora_o,
                shape.embd,
                s.attention_low.raw(),
                s.current_hc.raw(),
                s.hc_split.raw(),
                shape.embd as u32,
                shape.hc as u32,
            )
        },
        "attention output",
    )?;
    call(
        unsafe {
            ds4_gpu_rms_norm_plain_tensor(
                s.flat_hc.raw(),
                s.next_hc.raw(),
                hc_dim as u32,
                shape.rms_epsilon,
            )
        },
        "FFN HC norm",
    )?;
    f16(
        &s.hc_mix,
        w.hc_ffn_fn,
        hc_dim,
        mix_hc,
        &s.flat_hc,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_hc_split_weighted_sum_norm_tensor(
                s.current.raw(),
                s.norm.raw(),
                s.hc_split.raw(),
                s.hc_mix.raw(),
                s.next_hc.raw(),
                map,
                size,
                w.hc_ffn_scale.offset,
                w.hc_ffn_base.offset,
                w.ffn_norm.offset,
                shape.embd as u32,
                shape.hc as u32,
                shape.hc_sinkhorn as u32,
                shape.hc_epsilon,
                shape.rms_epsilon,
            )
        },
        "FFN HC mix",
    )?;
    f16(
        &s.router_logits,
        w.router,
        shape.embd,
        shape.experts,
        &s.norm,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_router_select_tensor(
                s.router_selected.raw(),
                s.router_weights.raw(),
                s.router_probs.raw(),
                map,
                size,
                w.router_bias.map_or(0, |v| v.offset),
                w.router_hash.map_or(0, |v| v.offset),
                w.router_hash.map_or(0, |v| v.dims[1] as u32),
                token,
                shape.experts as u32,
                shape.experts_used as u32,
                shape.expert_weight_scale,
                0,
                0,
                w.router_bias.is_some(),
                w.router_hash.is_some(),
                s.router_logits.raw(),
            )
        },
        "expert routing",
    )?;
    let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]);
    let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]);
    call(
        unsafe {
            ds4_gpu_routed_moe_one_tensor(
                s.routed_out.raw(),
                s.routed_gate.raw(),
                s.routed_up.raw(),
                s.routed_mid.raw(),
                s.routed_experts.raw(),
                map,
                size,
                w.expert_gate.offset,
                w.expert_up.offset,
                w.expert_down.offset,
                w.expert_gate.kind,
                w.expert_down.kind,
                w.expert_gate.dims[1] * gate_row,
                gate_row,
                w.expert_down.dims[1] * down_row,
                down_row,
                w.expert_gate.dims[0] as u32,
                w.expert_down.dims[0] as u32,
                w.expert_down.dims[1] as u32,
                s.router_selected.raw(),
                s.router_weights.raw(),
                shape.experts as u32,
                shape.experts_used as u32,
                shape.swiglu_clamp,
                s.norm.raw(),
                std::ptr::null(),
                layer,
                false,
            )
        },
        "routed experts",
    )?;
    call(
        unsafe {
            ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
                s.shared_gate.raw(),
                s.shared_up.raw(),
                s.shared_mid.raw(),
                map,
                size,
                w.shared_gate.offset,
                w.shared_up.offset,
                shape.embd,
                shape.ff_expert,
                s.norm.raw(),
                shape.swiglu_clamp,
            )
        },
        "shared expert gate/up",
    )?;
    call(
        unsafe {
            ds4_gpu_shared_down_hc_expand_q8_0_tensor(
                s.current_hc.raw(),
                s.shared_out.raw(),
                map,
                size,
                w.shared_down.offset,
                shape.ff_expert,
                shape.embd,
                s.shared_mid.raw(),
                s.routed_out.raw(),
                s.next_hc.raw(),
                s.hc_split.raw(),
                shape.embd as u32,
                shape.hc as u32,
            )
        },
        "shared expert output",
    )
}

#[allow(clippy::too_many_arguments)]
fn update_compression(
    s: &Scratch,
    state: &mut CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    original: u32,
    freq_base: f32,
    freq_scale: f32,
    ext: f32,
    attn_factor: f32,
) -> Result<(), String> {
    let emit = update_compressor_stage(
        s,
        state,
        weights,
        shape,
        map,
        size,
        pos,
        original,
        freq_base,
        freq_scale,
        ext,
        attn_factor,
        shape.head_dim as u32,
    )?;
    if emit {
        call(
            unsafe {
                ds4_gpu_dsv4_fp8_kv_quantize_tensor(
                    s.compressed_stage.raw(),
                    1,
                    shape.head_dim as u32,
                    shape.rot as u32,
                )
            },
            "compressed KV quantization",
        )?;
        call(
            unsafe {
                ds4_gpu_tensor_copy_f32_to_f16(
                    state.cache.raw(),
                    state.rows as u64 * shape.head_dim * 2,
                    s.compressed_stage.raw(),
                    0,
                    shape.head_dim,
                )
            },
            "compressed KV cache write",
        )?;
        state.rows += 1;
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn update_indexer_compression(
    s: &Scratch,
    state: &mut CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    original: u32,
    freq_base: f32,
    freq_scale: f32,
    ext: f32,
    attn_factor: f32,
) -> Result<(), String> {
    let emit = update_compressor_stage(
        s,
        state,
        weights,
        shape,
        map,
        size,
        pos,
        original,
        freq_base,
        freq_scale,
        ext,
        attn_factor,
        shape.indexer_head_dim as u32,
    )?;
    if emit {
        call(
            unsafe {
                ds4_gpu_dsv4_indexer_qat_tensor(
                    s.compressed_stage.raw(),
                    1,
                    shape.indexer_head_dim as u32,
                )
            },
            "compressed index quantization",
        )?;
        call(
            unsafe {
                ds4_gpu_tensor_copy(
                    state.cache.raw(),
                    state.rows as u64 * shape.indexer_head_dim * 4,
                    s.compressed_stage.raw(),
                    0,
                    shape.indexer_head_dim * 4,
                )
            },
            "compressed index cache write",
        )?;
        state.rows += 1;
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn update_compressor_stage(
    s: &Scratch,
    state: &CompressionState,
    weights: CompressorWeights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
    pos: u32,
    original: u32,
    freq_base: f32,
    freq_scale: f32,
    ext: f32,
    attn_factor: f32,
    head_dim: u32,
) -> Result<bool, String> {
    let width = if state.ratio == 4 { 2 } else { 1 } * head_dim;
    let fused = unsafe {
        ds4_gpu_matmul_f16_pair_compressor_store_tensor(
            s.compressed_kv.raw(),
            s.compressed_score.raw(),
            state.state_kv.raw(),
            state.state_score.raw(),
            map,
            size,
            weights.kv.offset,
            weights.gate.offset,
            weights.ape.offset,
            weights.ape.kind,
            shape.embd,
            width,
            s.norm.raw(),
            state.ratio,
            pos,
        )
    };
    if fused < 0 {
        return Err("Metal failed while storing compressor projections".into());
    }
    if fused == 0 {
        call(
            unsafe {
                ds4_gpu_matmul_f16_pair_tensor(
                    s.compressed_kv.raw(),
                    s.compressed_score.raw(),
                    map,
                    size,
                    weights.kv.offset,
                    weights.gate.offset,
                    shape.embd,
                    width as u64,
                    s.norm.raw(),
                    1,
                )
            },
            "compressor projection",
        )?;
    }
    call(
        unsafe {
            ds4_gpu_compressor_update_tensor(
                s.compressed_kv.raw(),
                s.compressed_score.raw(),
                state.state_kv.raw(),
                state.state_score.raw(),
                s.compressed_stage.raw(),
                map,
                size,
                weights.ape.offset,
                weights.ape.kind,
                weights.norm.offset,
                weights.norm.kind,
                head_dim,
                state.ratio,
                pos,
                0,
                shape.rot as u32,
                original,
                freq_base,
                freq_scale,
                ext,
                attn_factor,
                shape.rope_beta_fast,
                shape.rope_beta_slow,
                shape.rms_epsilon,
                fused > 0,
            )
        },
        "compressor update",
    )?;
    Ok((pos + 1).is_multiple_of(state.ratio))
}

fn compression_ratio(layer: u32) -> u32 {
    if layer < 2 {
        0
    } else if layer.is_multiple_of(2) {
        4
    } else {
        128
    }
}

fn raw_decode_span(pos: u32, raw_cap: u32, window: u32) -> (u32, u32) {
    let n_raw = (pos + 1)
        .min(raw_cap)
        .min(if window == 0 { raw_cap } else { window });
    (n_raw, (pos + 1 - n_raw) % raw_cap)
}

fn raw_batch_span(pos: u32, rows: u32, raw_cap: u32, window: u32) -> (u32, u32) {
    if rows == 0 || raw_cap == 0 {
        return (0, 0);
    }
    let last = pos + rows - 1;
    let needed = rows
        .saturating_add(if rows == 1 {
            window.saturating_sub(1)
        } else {
            window
        })
        .min(last + 1)
        .min(raw_cap);
    (needed, (last + 1 - needed) % raw_cap)
}

fn encode_output(
    s: &Scratch,
    w: &Weights,
    shape: super::Shape,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    let hc_dim = shape.hc * shape.embd;
    call(
        unsafe {
            ds4_gpu_rms_norm_plain_tensor(
                s.flat_hc.raw(),
                s.current_hc.raw(),
                hc_dim as u32,
                shape.rms_epsilon,
            )
        },
        "output HC norm",
    )?;
    f16(
        &s.output_pre,
        w.output_hc_fn,
        hc_dim,
        shape.hc,
        &s.flat_hc,
        map,
        size,
    )?;
    call(
        unsafe {
            ds4_gpu_output_hc_weights_tensor(
                s.output_weights.raw(),
                s.output_pre.raw(),
                map,
                size,
                w.output_hc_scale.offset,
                w.output_hc_base.offset,
                shape.hc as u32,
                shape.hc_epsilon,
            )
        },
        "output HC weights",
    )?;
    let fused_sum_norm = unsafe {
        ds4_gpu_hc_weighted_sum_norm_tensor(
            s.output_embedding.raw(),
            s.output_norm.raw(),
            s.current_hc.raw(),
            s.output_weights.raw(),
            map,
            size,
            w.output_norm.offset,
            shape.embd as u32,
            shape.hc as u32,
            shape.rms_epsilon,
        )
    } != 0;
    if !fused_sum_norm {
        call(
            unsafe {
                ds4_gpu_hc_weighted_sum_tensor(
                    s.output_embedding.raw(),
                    s.current_hc.raw(),
                    s.output_weights.raw(),
                    shape.embd as u32,
                    shape.hc as u32,
                )
            },
            "output HC collapse",
        )?;
        call(
            unsafe {
                ds4_gpu_rms_norm_weight_tensor(
                    s.output_norm.raw(),
                    s.output_embedding.raw(),
                    map,
                    size,
                    w.output_norm.offset,
                    shape.embd as u32,
                    shape.rms_epsilon,
                )
            },
            "output norm",
        )?;
    }
    q8(
        &s.logits,
        w.output,
        shape.embd,
        shape.vocab,
        &s.output_norm,
        map,
        size,
    )
}

fn f16(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_matmul_f16_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                input,
                output,
                x.raw(),
                1,
            )
        },
        "F16 projection",
    )
}

#[allow(clippy::too_many_arguments)]
fn f16_rows(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    rows: u32,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_matmul_f16_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                input,
                output,
                x.raw(),
                u64::from(rows),
            )
        },
        "batch F16 projection",
    )
}

#[allow(clippy::too_many_arguments)]
fn matmul_rows(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    rows: u32,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    match weight.kind {
        F16 => f16_rows(out, weight, input, output, x, rows, map, size),
        Q8_0 => q8_rows(out, weight, input, output, x, rows, map, size),
        kind => Err(format!("unsupported Metal batch projection type {kind}")),
    }
}

fn matmul(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    match weight.kind {
        F16 => f16(out, weight, input, output, x, map, size),
        Q8_0 => q8(out, weight, input, output, x, map, size),
        kind => Err(format!("unsupported Metal projection type {kind}")),
    }
}

fn q8(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_matmul_q8_0_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                input,
                output,
                x.raw(),
                1,
            )
        },
        "Q8 projection",
    )
}

#[allow(clippy::too_many_arguments)]
fn q8_rows(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    rows: u32,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_matmul_q8_0_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                input,
                output,
                x.raw(),
                u64::from(rows),
            )
        },
        "batch Q8 projection",
    )
}

fn call(result: i32, operation: &str) -> Result<(), String> {
    if result == 0 {
        Err(format!("Metal failed while {operation}"))
    } else {
        Ok(())
    }
}

fn check(result: i32, operation: &str) -> Result<(), String> {
    call(result, operation)
}

#[cfg(test)]
mod tests {
    use super::{raw_batch_span, raw_decode_span};

    #[test]
    fn decode_raw_span_tracks_the_logical_sliding_window() {
        assert_eq!(raw_decode_span(0, 4_352, 128), (1, 0));
        assert_eq!(raw_decode_span(127, 4_352, 128), (128, 0));
        assert_eq!(raw_decode_span(128, 4_352, 128), (128, 1));
        assert_eq!(raw_decode_span(362, 4_352, 128), (128, 235));
        assert_eq!(raw_decode_span(4_500, 4_352, 128), (128, 4_373 % 4_352));
    }

    #[test]
    fn batch_raw_span_includes_the_chunk_and_previous_window() {
        assert_eq!(raw_batch_span(0, 4_096, 4_352, 128), (4_096, 0));
        assert_eq!(raw_batch_span(4_096, 1_925, 4_352, 128), (2_053, 3_968));
        assert_eq!(raw_batch_span(4_500, 1, 4_352, 128), (128, 21));
    }
}
