Split Rust code into domain modules

src/engine/metal/checkpoint.rs

@@ -0,0 +1,367 @@
+use super::*;
+
+impl Executor {
+    pub(in crate::engine) fn save_checkpoint(
+        &mut self,
+        path: &Path,
+        tag: [u8; 32],
+        progress: &mut impl FnMut(u64),
+    ) -> Result<(), String> {
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent).map_err(|error| error.to_string())?;
+        }
+        let temporary = path.with_extension("tmp");
+        let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
+        let mut reported = 0_u64;
+        let mut pending = 0_u64;
+        self.write_checkpoint(&mut file, tag, &mut |bytes| {
+            reported += bytes;
+            pending += bytes;
+            if pending >= CHECKPOINT_IO_CHUNK as u64 {
+                progress(pending);
+                pending = 0;
+            }
+        })?;
+        progress(pending);
+        let total = file.metadata().map_err(|error| error.to_string())?.len();
+        progress(total.saturating_sub(reported));
+        file.sync_all().map_err(|error| error.to_string())?;
+        fs::rename(&temporary, path).map_err(|error| error.to_string())?;
+        self.checkpoint_tag = tag;
+        Ok(())
+    }
+
+    pub(in crate::engine) fn load_checkpoint(
+        &mut self,
+        path: &Path,
+        progress: &mut impl FnMut(u64),
+    ) -> Result<bool, String> {
+        let mut file = match File::open(path) {
+            Ok(file) => file,
+            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
+            Err(error) => return Err(error.to_string()),
+        };
+        let mut reported = 0_u64;
+        let mut pending = 0_u64;
+        self.read_checkpoint(&mut file, &mut |bytes| {
+            reported += bytes;
+            pending += bytes;
+            if pending >= CHECKPOINT_IO_CHUNK as u64 {
+                progress(pending);
+                pending = 0;
+            }
+        })?;
+        progress(pending);
+        let total = file.metadata().map_err(|error| error.to_string())?.len();
+        progress(total.saturating_sub(reported));
+        Ok(true)
+    }
+
+    fn write_checkpoint(
+        &self,
+        file: &mut File,
+        tag: [u8; 32],
+        progress: &mut impl FnMut(u64),
+    ) -> Result<(), String> {
+        let shape = self.model.shape;
+        file.write_all(CHECKPOINT_MAGIC)
+            .map_err(|error| error.to_string())?;
+        for value in [
+            CHECKPOINT_VERSION,
+            self.session.context,
+            self.session.raw_cap,
+            shape.layers,
+            shape.head_dim as u32,
+            shape.indexer_head_dim as u32,
+            shape.vocab as u32,
+            u32::from(self.quality),
+        ] {
+            write_u32(file, value)?;
+        }
+        write_u64(file, self.model.main.len())?;
+        write_u64(file, self.model_modified.0)?;
+        write_u32(file, self.model_modified.1)?;
+        for weight in [self.weights.token_embedding, self.weights.output] {
+            write_u64(file, weight.offset)?;
+            write_u64(file, weight.bytes)?;
+            write_u32(file, weight.kind)?;
+        }
+        file.write_all(&tag).map_err(|error| error.to_string())?;
+        let token_count = u32::try_from(self.tokens.len())
+            .map_err(|_| "KV checkpoint has too many tokens".to_owned())?;
+        let raw_live = token_count.min(self.session.raw_cap);
+        write_u32(file, token_count)?;
+        write_u32(file, raw_live)?;
+        for &token in &self.tokens {
+            write_u32(file, token as u32)?;
+        }
+        for &logit in &self.logits {
+            write_u32(file, logit.to_bits())?;
+        }
+        for layer in &self.session.layers {
+            write_u32(
+                file,
+                layer.compression.as_ref().map_or(0, |state| state.rows),
+            )?;
+        }
+        for layer in &self.session.layers {
+            write_u32(file, layer.indexer.as_ref().map_or(0, |state| state.rows))?;
+        }
+
+        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
+        let raw_first = token_count - raw_live;
+        for layer in &self.session.layers {
+            for position in raw_first..token_count {
+                let physical = position % self.session.raw_cap;
+                write_buffer(
+                    file,
+                    &layer.raw_cache,
+                    u64::from(physical) * shape.head_dim * 4,
+                    shape.head_dim * 4,
+                    &mut chunk,
+                    progress,
+                )?;
+            }
+            if let Some(state) = &layer.compression {
+                write_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.head_dim * 2,
+                    &mut chunk,
+                    progress,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
+                write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
+                write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
+            }
+            if let Some(state) = &layer.indexer {
+                write_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.indexer_head_dim * 4,
+                    &mut chunk,
+                    progress,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
+                write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
+                write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
+            }
+        }
+        Ok(())
+    }
+
+    fn read_checkpoint(
+        &mut self,
+        file: &mut File,
+        progress: &mut impl FnMut(u64),
+    ) -> Result<(), String> {
+        let mut magic = [0; 8];
+        file.read_exact(&mut magic)
+            .map_err(|error| error.to_string())?;
+        if &magic != CHECKPOINT_MAGIC {
+            return Err("KV checkpoint has an invalid signature".into());
+        }
+        let shape = self.model.shape;
+        let header = [
+            CHECKPOINT_VERSION,
+            self.session.context,
+            self.session.raw_cap,
+            shape.layers,
+            shape.head_dim as u32,
+            shape.indexer_head_dim as u32,
+            shape.vocab as u32,
+            u32::from(self.quality),
+        ];
+        for expected in header {
+            if read_u32(file)? != expected {
+                return Err("KV checkpoint does not match the current executor".into());
+            }
+        }
+        if read_u64(file)? != self.model.main.len() {
+            return Err("KV checkpoint was written for a different model".into());
+        }
+        if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 {
+            return Err("KV checkpoint model file has changed".into());
+        }
+        for weight in [self.weights.token_embedding, self.weights.output] {
+            if read_u64(file)? != weight.offset
+                || read_u64(file)? != weight.bytes
+                || read_u32(file)? != weight.kind
+            {
+                return Err("KV checkpoint was written for a different model layout".into());
+            }
+        }
+        let mut checkpoint_tag = [0; 32];
+        file.read_exact(&mut checkpoint_tag)
+            .map_err(|error| error.to_string())?;
+        let token_count = read_u32(file)?;
+        let raw_live = read_u32(file)?;
+        if token_count > self.session.context || raw_live != token_count.min(self.session.raw_cap) {
+            return Err("KV checkpoint token count is invalid".into());
+        }
+        let mut tokens = Vec::with_capacity(token_count as usize);
+        for _ in 0..token_count {
+            let token = read_u32(file)?;
+            if u64::from(token) >= shape.vocab {
+                return Err("KV checkpoint contains an invalid token".into());
+            }
+            tokens.push(token as i32);
+        }
+        let mut logits = Vec::with_capacity(shape.vocab as usize);
+        for _ in 0..shape.vocab {
+            logits.push(f32::from_bits(read_u32(file)?));
+        }
+        let mut compressed_rows = Vec::with_capacity(shape.layers as usize);
+        let mut indexer_rows = Vec::with_capacity(shape.layers as usize);
+        for layer in 0..shape.layers {
+            let rows = read_u32(file)?;
+            let ratio = compression_ratio(layer);
+            if rows != token_count.checked_div(ratio).unwrap_or(0) {
+                return Err("KV checkpoint compressed row count is invalid".into());
+            }
+            compressed_rows.push(rows);
+        }
+        for layer in 0..shape.layers {
+            let rows = read_u32(file)?;
+            let expected = if compression_ratio(layer) == 4 {
+                token_count / 4
+            } else {
+                0
+            };
+            if rows != expected {
+                return Err("KV checkpoint indexer row count is invalid".into());
+            }
+            indexer_rows.push(rows);
+        }
+
+        self.reset()?;
+        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
+        let raw_first = token_count - raw_live;
+        for (index, layer) in self.session.layers.iter_mut().enumerate() {
+            for position in raw_first..token_count {
+                let physical = position % self.session.raw_cap;
+                read_buffer(
+                    file,
+                    &layer.raw_cache,
+                    u64::from(physical) * shape.head_dim * 4,
+                    shape.head_dim * 4,
+                    &mut chunk,
+                    progress,
+                )?;
+            }
+            if let Some(state) = &mut layer.compression {
+                state.rows = compressed_rows[index];
+                read_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.head_dim * 2,
+                    &mut chunk,
+                    progress,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
+                read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
+                read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
+            }
+            if let Some(state) = &mut layer.indexer {
+                state.rows = indexer_rows[index];
+                read_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.indexer_head_dim * 4,
+                    &mut chunk,
+                    progress,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
+                read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
+                read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
+            }
+        }
+        let mut trailing = [0];
+        if file
+            .read(&mut trailing)
+            .map_err(|error| error.to_string())?
+            != 0
+        {
+            self.reset()?;
+            return Err("KV checkpoint has trailing data".into());
+        }
+        self.session.position = token_count;
+        self.tokens = tokens;
+        self.logits = logits;
+        self.checkpoint_tag = checkpoint_tag;
+        Ok(())
+    }
+}
+
+fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 {
+    let coefficient = if ratio == 4 { 2 } else { 1 };
+    coefficient * head_dim * coefficient * u64::from(ratio) * 4
+}
+
+fn write_buffer(
+    file: &mut File,
+    buffer: &Buffer,
+    mut offset: u64,
+    mut bytes: u64,
+    chunk: &mut [u8],
+    progress: &mut impl FnMut(u64),
+) -> Result<(), String> {
+    while bytes != 0 {
+        let length = bytes.min(chunk.len() as u64) as usize;
+        buffer.read(offset, &mut chunk[..length])?;
+        file.write_all(&chunk[..length])
+            .map_err(|error| error.to_string())?;
+        progress(length as u64);
+        offset += length as u64;
+        bytes -= length as u64;
+    }
+    Ok(())
+}
+
+fn read_buffer(
+    file: &mut File,
+    buffer: &Buffer,
+    mut offset: u64,
+    mut bytes: u64,
+    chunk: &mut [u8],
+    progress: &mut impl FnMut(u64),
+) -> Result<(), String> {
+    while bytes != 0 {
+        let length = bytes.min(chunk.len() as u64) as usize;
+        file.read_exact(&mut chunk[..length])
+            .map_err(|error| error.to_string())?;
+        buffer.write(offset, &chunk[..length])?;
+        progress(length as u64);
+        offset += length as u64;
+        bytes -= length as u64;
+    }
+    Ok(())
+}
+
+fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn read_u32(file: &mut File) -> Result<u32, String> {
+    let mut bytes = [0; 4];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u32::from_le_bytes(bytes))
+}
+
+fn read_u64(file: &mut File) -> Result<u64, String> {
+    let mut bytes = [0; 8];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u64::from_le_bytes(bytes))
+}

src/engine/metal/gpu.rs

@@ -0,0 +1,926 @@
+use super::*;
+
+#[repr(C)]
+pub(super) struct GpuTensor {
+    _private: [u8; 0],
+}
+
+unsafe extern "C" {
+    pub(super) fn ds4_gpu_init() -> i32;
+    pub(super) fn ds4_gpu_cleanup();
+    pub(super) fn ds4_gpu_set_model_map_range(
+        model_map: *const c_void,
+        model_size: u64,
+        map_offset: u64,
+        map_size: u64,
+        max_tensor_bytes: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_set_quality(quality: bool);
+    pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
+    pub(super) fn ds4_gpu_tensor_view(
+        base: *const GpuTensor,
+        offset: u64,
+        bytes: u64,
+    ) -> *mut GpuTensor;
+    pub(super) fn ds4_gpu_tensor_free(tensor: *mut GpuTensor);
+    pub(super) fn ds4_gpu_tensor_fill_f32(tensor: *mut GpuTensor, value: f32, count: u64) -> i32;
+    pub(super) fn ds4_gpu_tensor_read(
+        tensor: *const GpuTensor,
+        offset: u64,
+        data: *mut c_void,
+        bytes: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_tensor_write(
+        tensor: *mut GpuTensor,
+        offset: u64,
+        data: *const c_void,
+        bytes: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_tensor_copy(
+        dst: *mut GpuTensor,
+        dst_offset: u64,
+        src: *const GpuTensor,
+        src_offset: u64,
+        bytes: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_tensor_copy_f32_to_f16(
+        dst: *mut GpuTensor,
+        dst_offset: u64,
+        src: *const GpuTensor,
+        src_offset: u64,
+        count: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_begin_commands() -> i32;
+    pub(super) fn ds4_gpu_end_commands() -> i32;
+
+    pub(super) fn ds4_gpu_embed_tokens_hc_tensor(
+        out: *mut GpuTensor,
+        tokens: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        vocab: u32,
+        rows: u32,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+
+    pub(super) fn ds4_gpu_embed_token_hc_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        vocab: u32,
+        token: u32,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_rms_norm_plain_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        n: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_rms_norm_weight_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        n: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor(
+        out: *mut GpuTensor,
+        scale: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u32,
+        output: u32,
+        x: *const GpuTensor,
+        rows: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_f16_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_q8_0_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_q8_0_pair_tensor(
+        out_a: *mut GpuTensor,
+        out_b: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight_a: u64,
+        weight_b: u64,
+        input: u64,
+        output_a: u64,
+        output_b: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_f16_pair_tensor(
+        out_a: *mut GpuTensor,
+        out_b: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight_a: u64,
+        weight_b: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_f16_pair_compressor_store_tensor(
+        out_kv: *mut GpuTensor,
+        out_score: *mut GpuTensor,
+        state_kv: *mut GpuTensor,
+        state_score: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight_kv: u64,
+        weight_score: u64,
+        ape: u64,
+        ape_type: u32,
+        input: u64,
+        width: u32,
+        x: *const GpuTensor,
+        ratio: u32,
+        pos: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_split_weighted_sum_norm_tensor(
+        out: *mut GpuTensor,
+        norm: *mut GpuTensor,
+        split: *mut GpuTensor,
+        mix: *const GpuTensor,
+        residual: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        scale: u64,
+        base: u64,
+        norm_weight: u64,
+        embd: u32,
+        hc: u32,
+        iterations: u32,
+        eps: f32,
+        norm_eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
+        q_out: *mut GpuTensor,
+        q: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        q_weight: u64,
+        q_n: u32,
+        kv_out: *mut GpuTensor,
+        kv: *const GpuTensor,
+        kv_weight: u64,
+        kv_n: u32,
+        rows: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
+        out: *mut GpuTensor,
+        half: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        rot: u32,
+        pos: u32,
+        original_context: u32,
+        inverse: bool,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_head_rms_norm_tensor(
+        x: *mut GpuTensor,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_head_rms_norm_rope_tail_tensor(
+        x: *mut GpuTensor,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        rot: u32,
+        pos: u32,
+        original_context: u32,
+        inverse: bool,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_rope_tail_tensor(
+        x: *mut GpuTensor,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        rot: u32,
+        pos: u32,
+        original_context: u32,
+        inverse: bool,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_decode_heads_tensor(
+        heads_out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw_kv: *const GpuTensor,
+        n_raw: u32,
+        raw_cap: u32,
+        raw_start: u32,
+        comp_kv: *const GpuTensor,
+        comp_f16: u32,
+        n_comp: u32,
+        comp_mask: *const GpuTensor,
+        use_mask: u32,
+        heads: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_decode_raw_batch_heads_tensor(
+        heads: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw_kv: *const GpuTensor,
+        tokens: u32,
+        pos: u32,
+        n_raw: u32,
+        raw_cap: u32,
+        raw_start: u32,
+        window: u32,
+        heads_count: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_decode_mixed_batch_heads_tensor(
+        heads: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw_kv: *const GpuTensor,
+        compressed: *const GpuTensor,
+        compressed_f16: u32,
+        compressed_mask: *const GpuTensor,
+        use_mask: u32,
+        tokens: u32,
+        pos: u32,
+        n_raw: u32,
+        raw_cap: u32,
+        raw_start: u32,
+        n_comp: u32,
+        window: u32,
+        ratio: u32,
+        heads_count: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
+        heads: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw_kv: *const GpuTensor,
+        comp_kv: *const GpuTensor,
+        comp_f16: u32,
+        topk: *const GpuTensor,
+        tokens: u32,
+        pos: u32,
+        n_raw: u32,
+        raw_cap: u32,
+        raw_start: u32,
+        n_comp: u32,
+        top_k: u32,
+        window: u32,
+        ratio: u32,
+        heads: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_indexer_score_one_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        index_comp: *const GpuTensor,
+        n_comp: u32,
+        heads: u32,
+        head_dim: u32,
+        scale: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_indexer_scores_decode_batch_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        index_comp: *const GpuTensor,
+        n_comp: u32,
+        tokens: u32,
+        pos: u32,
+        heads: u32,
+        head_dim: u32,
+        ratio: u32,
+        scale: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_indexer_topk_tensor(
+        selected: *mut GpuTensor,
+        scores: *const GpuTensor,
+        n_comp: u32,
+        tokens: u32,
+        top_k: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_dsv4_indexer_qat_tensor(
+        x: *mut GpuTensor,
+        rows: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_compressor_update_tensor(
+        kv: *const GpuTensor,
+        score: *const GpuTensor,
+        state_kv: *mut GpuTensor,
+        state_score: *mut GpuTensor,
+        cache: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        ape: u64,
+        ape_type: u32,
+        norm: u64,
+        norm_type: u32,
+        head_dim: u32,
+        ratio: u32,
+        pos: u32,
+        row: u32,
+        rot: u32,
+        original_context: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        rms_eps: f32,
+        state_already_stored: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_compressor_prefill_state_ratio4_tensor(
+        state_kv: *mut GpuTensor,
+        state_score: *mut GpuTensor,
+        kv_tail: *const GpuTensor,
+        score_tail: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        ape: u64,
+        ape_type: u32,
+        head_dim: u32,
+        pos: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_dsv4_fp8_kv_quantize_tensor(
+        x: *mut GpuTensor,
+        rows: u32,
+        head_dim: u32,
+        rot: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_kv_fp8_store_raw_tensor(
+        kv: *mut GpuTensor,
+        raw_cache: *mut GpuTensor,
+        raw_cap: u32,
+        row: u32,
+        head_dim: u32,
+        rot: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_store_raw_kv_batch_tensor(
+        raw_cache: *mut GpuTensor,
+        kv: *const GpuTensor,
+        raw_cap: u32,
+        pos: u32,
+        rows: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_compressor_prefill_tensor(
+        cache: *mut GpuTensor,
+        state_kv: *mut GpuTensor,
+        state_score: *mut GpuTensor,
+        kv: *const GpuTensor,
+        score: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        ape: u64,
+        ape_type: u32,
+        norm: u64,
+        norm_type: u32,
+        head_dim: u32,
+        ratio: u32,
+        pos: u32,
+        rows: u32,
+        rot: u32,
+        original_context: u32,
+        quantize_fp8: bool,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        rms_eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_compressor_prefill_ratio4_replay_tensor(
+        cache: *mut GpuTensor,
+        state_kv: *mut GpuTensor,
+        state_score: *mut GpuTensor,
+        kv: *const GpuTensor,
+        score: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        ape: u64,
+        ape_type: u32,
+        norm: u64,
+        norm_type: u32,
+        head_dim: u32,
+        pos: u32,
+        rows: u32,
+        rot: u32,
+        original_context: u32,
+        quantize_fp8: bool,
+        freq_base: f32,
+        freq_scale: f32,
+        ext_factor: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        rms_eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_prefill_raw_heads_tensor(
+        heads: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw: *const GpuTensor,
+        rows: u32,
+        window: u32,
+        heads_count: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_prefill_static_mixed_heads_tensor(
+        heads: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw: *const GpuTensor,
+        compressed: *const GpuTensor,
+        compressed_f16: u32,
+        rows: u32,
+        compressed_rows: u32,
+        window: u32,
+        ratio: u32,
+        heads_count: u32,
+        head_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_indexer_scores_prefill_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        compressed: *const GpuTensor,
+        compressed_rows: u32,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        ratio: u32,
+        scale: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_output_q8_batch_f16_tensor(
+        out_half: *mut GpuTensor,
+        low: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight_a: u64,
+        weight_b: u64,
+        group_dim: u64,
+        rank: u64,
+        groups: u32,
+        output: u64,
+        heads: *const GpuTensor,
+        rows: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_output_q8_batch_tensor(
+        out: *mut GpuTensor,
+        low: *mut GpuTensor,
+        group_scratch: *mut GpuTensor,
+        low_scratch: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight_a: u64,
+        weight_b: u64,
+        group_dim: u64,
+        rank: u64,
+        groups: u32,
+        output: u64,
+        heads: *const GpuTensor,
+        rows: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_router_select_batch_tensor(
+        selected: *mut GpuTensor,
+        weights: *mut GpuTensor,
+        probs: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        bias: u64,
+        hash: u64,
+        hash_rows: u32,
+        expert_groups: u32,
+        groups_used: u32,
+        has_bias: bool,
+        hash_mode: bool,
+        logits: *const GpuTensor,
+        tokens: *const GpuTensor,
+        experts: u32,
+        experts_used: u32,
+        scale: f32,
+        rows: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_routed_moe_batch_tensor(
+        out: *mut GpuTensor,
+        gate: *mut GpuTensor,
+        up: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        experts_out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        down_weight: u64,
+        gate_type: u32,
+        down_type: u32,
+        gate_expert_bytes: u64,
+        gate_row_bytes: u64,
+        down_expert_bytes: u64,
+        down_row_bytes: u64,
+        input: u32,
+        middle: u32,
+        output: u32,
+        selected: *const GpuTensor,
+        weights: *const GpuTensor,
+        total_experts: u32,
+        used_experts: u32,
+        clamp: f32,
+        x: *const GpuTensor,
+        layer: u32,
+        rows: u32,
+        mid_f16: *mut bool,
+        force_resident: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_swiglu_tensor(
+        out: *mut GpuTensor,
+        gate: *const GpuTensor,
+        up: *const GpuTensor,
+        count: u32,
+        clamp: f32,
+        scale: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_expand_split_half_tensor(
+        out: *mut GpuTensor,
+        block_half: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_expand_split_tensor(
+        out: *mut GpuTensor,
+        block: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_expand_add_split_tensor(
+        out: *mut GpuTensor,
+        block: *const GpuTensor,
+        add: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_output_low_q8_tensor(
+        low: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        group_dim: u64,
+        rank: u64,
+        groups: u32,
+        heads: *const GpuTensor,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_q8_0_hc_expand_tensor(
+        out_hc: *mut GpuTensor,
+        block_out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_router_select_tensor(
+        selected: *mut GpuTensor,
+        weights: *mut GpuTensor,
+        probs: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        bias: u64,
+        hash: u64,
+        hash_rows: u32,
+        token: u32,
+        experts: u32,
+        used: u32,
+        scale: f32,
+        expert_groups: u32,
+        groups_used: u32,
+        has_bias: bool,
+        hash_mode: bool,
+        logits: *const GpuTensor,
+    ) -> i32;
+    pub(super) fn ds4_gpu_routed_moe_one_tensor(
+        out: *mut GpuTensor,
+        gate: *mut GpuTensor,
+        up: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        experts_out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        down_weight: u64,
+        gate_type: u32,
+        down_type: u32,
+        gate_expert_bytes: u64,
+        gate_row_bytes: u64,
+        down_expert_bytes: u64,
+        down_row_bytes: u64,
+        input: u32,
+        middle: u32,
+        output: u32,
+        selected: *const GpuTensor,
+        weights: *const GpuTensor,
+        total_experts: u32,
+        used_experts: u32,
+        clamp: f32,
+        x: *const GpuTensor,
+        add: *const GpuTensor,
+        layer: u32,
+        force_resident: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
+        gate: *mut GpuTensor,
+        up: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        clamp: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_shared_down_hc_expand_q8_0_tensor(
+        out_hc: *mut GpuTensor,
+        shared_out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        middle: *const GpuTensor,
+        routed: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_output_hc_weights_tensor(
+        out: *mut GpuTensor,
+        pre: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        scale: u64,
+        base: u64,
+        hc: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_weighted_sum_norm_tensor(
+        out: *mut GpuTensor,
+        norm: *mut GpuTensor,
+        residual: *const GpuTensor,
+        weights: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        norm_weight: u64,
+        embd: u32,
+        hc: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_hc_weighted_sum_tensor(
+        out: *mut GpuTensor,
+        residual: *const GpuTensor,
+        weights: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+}
+
+pub(super) struct Context;
+
+impl Context {
+    pub(super) fn open(model: &Model, quality: bool) -> Result<Self, String> {
+        check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
+        let data_offset = model.main.data_offset();
+        if let Err(error) = check(
+            unsafe {
+                ds4_gpu_set_model_map_range(
+                    model.main.map_ptr().cast(),
+                    model.main.len(),
+                    data_offset,
+                    model.main.len() - data_offset,
+                    model.main.max_tensor_bytes(),
+                )
+            },
+            "model mapping",
+        ) {
+            unsafe { ds4_gpu_cleanup() };
+            return Err(error);
+        }
+        unsafe { ds4_gpu_set_quality(quality) };
+        Ok(Self)
+    }
+}
+
+impl Drop for Context {
+    fn drop(&mut self) {
+        unsafe { ds4_gpu_cleanup() };
+    }
+}
+
+pub(super) struct Commands(bool);
+
+impl Commands {
+    pub(super) fn begin() -> Result<Self, String> {
+        check(
+            unsafe { ds4_gpu_begin_commands() },
+            "beginning Metal commands",
+        )?;
+        Ok(Self(true))
+    }
+
+    pub(super) fn finish(mut self) -> Result<(), String> {
+        self.0 = false;
+        check(
+            unsafe { ds4_gpu_end_commands() },
+            "executing Metal commands",
+        )
+    }
+}
+
+impl Drop for Commands {
+    fn drop(&mut self) {
+        if self.0 {
+            unsafe { ds4_gpu_end_commands() };
+        }
+    }
+}
+
+pub(super) struct Buffer(NonNull<GpuTensor>);
+
+impl Buffer {
+    pub(super) fn floats(count: u64) -> Result<Self, String> {
+        Self::bytes(count * 4)
+    }
+
+    pub(super) fn bytes(bytes: u64) -> Result<Self, String> {
+        NonNull::new(unsafe { ds4_gpu_tensor_alloc(bytes) })
+            .map(Self)
+            .ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
+    }
+
+    pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
+        NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) })
+            .map(Self)
+            .ok_or_else(|| "Metal could not create a tensor view".to_owned())
+    }
+
+    pub(super) fn read_f32(&self, values: &mut [f32]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_read(
+                    self.raw(),
+                    0,
+                    values.as_mut_ptr().cast(),
+                    std::mem::size_of_val(values) as u64,
+                )
+            },
+            "reading Metal output",
+        )
+    }
+
+    pub(super) fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_read(
+                    self.raw(),
+                    offset,
+                    values.as_mut_ptr().cast(),
+                    values.len() as u64,
+                )
+            },
+            "reading a Metal buffer",
+        )
+    }
+
+    pub(super) fn write(&self, offset: u64, values: &[u8]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_write(
+                    self.raw(),
+                    offset,
+                    values.as_ptr().cast(),
+                    values.len() as u64,
+                )
+            },
+            "restoring a Metal buffer",
+        )
+    }
+
+    pub(super) fn write_i32(&self, values: &[i32]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_write(
+                    self.raw(),
+                    0,
+                    values.as_ptr().cast(),
+                    std::mem::size_of_val(values) as u64,
+                )
+            },
+            "uploading tokens",
+        )
+    }
+
+    pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> {
+        call(
+            unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },
+            "initializing a Metal buffer",
+        )
+    }
+
+    pub(super) fn raw(&self) -> *mut GpuTensor {
+        self.0.as_ptr()
+    }
+}
+
+impl Drop for Buffer {
+    fn drop(&mut self) {
+        unsafe { ds4_gpu_tensor_free(self.raw()) };
+    }
+}