metrale_gpu_runtime/
mlx_int8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launches for 8-bit affine weights in the MLX layout: load a
4//! `(.weight, .scales, .biases)` triplet and run the dequant, GEMV, fused GEMV and
5//! GEMM kernels of `kernels/metal/common/mlx_int8_*.metal`.
6//!
7//! Layout, as [`MlxInt8Weight::load`] checks it:
8//! - `{base}.weight`: `U32` `[out_features, in_features / 4]`; each word packs four
9//!   unsigned bytes, the low byte holding the lowest column.
10//! - `{base}.scales`, `{base}.biases`: `BF16` `[out_features, in_features / G]`,
11//!   where `G` is `group_size`.
12//!
13//! The kernels dequantise `w[r, c] = byte * scales[r, c / G] + biases[r, c / G]`,
14//! where `byte` is byte `c % 4` of `packed[r, c / 4]`.
15//!
16//! Owner: gpu-runtime (metal).
17//! Invariants: none beyond the types.
18
19use anyhow::{Context, Result, bail};
20use safetensors::SafeTensors;
21use serde_json::Value as JsonValue;
22
23use crate::gpu::{DevicePtr, GpuBackend, KernelArg};
24
25/// 2026-09-25: Quantization metadata from the model's `config.json`.
26#[derive(Debug, Clone, Copy)]
27pub struct MlxQuantConfig {
28    pub bits: u32,
29    pub group_size: u32,
30}
31
32impl MlxQuantConfig {
33    /// 2026-09-25: Read `bits` and `group_size` from the top-level `quantization`
34    /// block, or from `quantization_config` when `quantization` is absent. `None`
35    /// when the block or either key is missing or not an unsigned integer.
36    pub fn from_config(config: &JsonValue) -> Option<Self> {
37        let q = config
38            .get("quantization")
39            .or_else(|| config.get("quantization_config"))?;
40        let bits = q.get("bits")?.as_u64()? as u32;
41        let group_size = q.get("group_size")?.as_u64()? as u32;
42        Some(Self { bits, group_size })
43    }
44}
45
46/// 2026-09-25: One 8-bit affine linear weight resident on the GPU. There is no
47/// `Drop`: the owner frees the three buffers with [`Self::release`].
48pub struct MlxInt8Weight {
49    /// 2026-09-25: `[out_features, in_features / 4]` packed bytes (uint32 words).
50    pub packed: DevicePtr,
51    /// 2026-09-25: `[out_features, in_features / group_size]` per-group BF16 scales.
52    pub scales: DevicePtr,
53    /// 2026-09-25: `[out_features, in_features / group_size]` per-group BF16 biases.
54    pub biases: DevicePtr,
55    pub out_features: u32,
56    pub in_features: u32,
57    pub group_size: u32,
58}
59
60impl MlxInt8Weight {
61    /// 2026-09-25: Load the `{base}.weight`, `{base}.scales` and `{base}.biases`
62    /// tensors from a parsed safetensors blob and upload them. Fails on a missing
63    /// tensor, a weight that is not 2-D `U32`, scales or biases that are not `BF16`
64    /// of shape `[out_features, in_features / group_size]`, or an `in_features`
65    /// that `group_size` does not divide.
66    pub fn load(
67        gpu: &dyn GpuBackend,
68        st: &SafeTensors,
69        base: &str,
70        group_size: u32,
71    ) -> Result<Self> {
72        let weight_name = format!("{base}.weight");
73        let scales_name = format!("{base}.scales");
74        let biases_name = format!("{base}.biases");
75
76        let weight = st
77            .tensor(&weight_name)
78            .with_context(|| format!("missing tensor {weight_name}"))?;
79        let scales = st
80            .tensor(&scales_name)
81            .with_context(|| format!("missing tensor {scales_name}"))?;
82        let biases = st
83            .tensor(&biases_name)
84            .with_context(|| format!("missing tensor {biases_name}"))?;
85
86        if weight.dtype() != safetensors::Dtype::U32 {
87            bail!(
88                "{weight_name}: expected U32 (MLX 8-bit packed), got {:?}",
89                weight.dtype()
90            );
91        }
92        if scales.dtype() != safetensors::Dtype::BF16 || biases.dtype() != safetensors::Dtype::BF16
93        {
94            bail!(
95                "{base}.scales/biases: expected BF16, got scales={:?}, biases={:?}",
96                scales.dtype(),
97                biases.dtype()
98            );
99        }
100
101        let weight_shape = weight.shape();
102        if weight_shape.len() != 2 {
103            bail!(
104                "{weight_name}: expected 2-D weight tensor, got rank {}",
105                weight_shape.len()
106            );
107        }
108        let out_features = weight_shape[0] as u32;
109        let packed_cols = weight_shape[1] as u32;
110        let in_features = packed_cols * 4;
111
112        let groups_per_row = in_features
113            .checked_div(group_size)
114            .filter(|&g| g * group_size == in_features)
115            .ok_or_else(|| {
116                anyhow::anyhow!(
117                    "{base}: in_features {in_features} not divisible by group_size {group_size}"
118                )
119            })?;
120
121        let expected = [out_features as usize, groups_per_row as usize];
122        if scales.shape() != expected {
123            bail!(
124                "{}.scales: expected shape {:?}, got {:?}",
125                base,
126                expected,
127                scales.shape()
128            );
129        }
130        if biases.shape() != expected {
131            bail!(
132                "{}.biases: expected shape {:?}, got {:?}",
133                base,
134                expected,
135                biases.shape()
136            );
137        }
138
139        let packed_ptr = gpu.alloc(weight.data().len())?;
140        gpu.copy_h2d(weight.data(), packed_ptr)?;
141        let scales_ptr = gpu.alloc(scales.data().len())?;
142        gpu.copy_h2d(scales.data(), scales_ptr)?;
143        let biases_ptr = gpu.alloc(biases.data().len())?;
144        gpu.copy_h2d(biases.data(), biases_ptr)?;
145
146        Ok(Self {
147            packed: packed_ptr,
148            scales: scales_ptr,
149            biases: biases_ptr,
150            out_features,
151            in_features,
152            group_size,
153        })
154    }
155
156    /// 2026-09-25: Write the full dequantised weight as BF16 to `out`, which must
157    /// hold at least `out_features * in_features * 2` bytes. Runs `mlx_int8_dequant`.
158    pub fn dequantize_to(&self, gpu: &dyn GpuBackend, out: DevicePtr, stream: u64) -> Result<()> {
159        let kernel = gpu.kernel("mlx_int8_dequant", "mlx_int8_dequant")?;
160        // 2026-09-25: One thread per element; the kernel bounds-checks the ragged
161        // last column tile.
162        let block_x: u32 = 16;
163        let block_y: u32 = 1;
164        let grid_x = self.in_features.div_ceil(block_x);
165        let grid_y = self.out_features;
166        gpu.launch_typed(
167            kernel,
168            [grid_x, grid_y, 1],
169            [block_x, block_y, 1],
170            0,
171            stream,
172            &[
173                KernelArg::Bytes(&self.out_features.to_le_bytes()),
174                KernelArg::Bytes(&self.in_features.to_le_bytes()),
175                KernelArg::Bytes(&self.group_size.to_le_bytes()),
176                KernelArg::Buffer(self.packed),
177                KernelArg::Buffer(self.scales),
178                KernelArg::Buffer(self.biases),
179                KernelArg::Buffer(out),
180            ],
181        )
182    }
183
184    /// 2026-09-25: `y = W x`, with `x` BF16 `[in_features]` and `y` a BF16 buffer of
185    /// at least `out_features` elements. Runs `mlx_int8_gemv`.
186    pub fn gemv(
187        &self,
188        gpu: &dyn GpuBackend,
189        x: DevicePtr,
190        y: DevicePtr,
191        stream: u64,
192    ) -> Result<()> {
193        let kernel = gpu.kernel("mlx_int8_gemv", "mlx_int8_gemv")?;
194        // 2026-09-25: Four rows per threadgroup, one 32-lane simdgroup per row,
195        // reduced with `simd_sum`, so no reduction crosses simdgroups.
196        const ROWS_PER_TG: u32 = 4;
197        const SIMDGROUP_SIZE: u32 = 32;
198        let threads_per_tg: u32 = ROWS_PER_TG * SIMDGROUP_SIZE;
199        let row_groups = self.out_features.div_ceil(ROWS_PER_TG);
200        gpu.launch_typed(
201            kernel,
202            [row_groups, 1, 1],
203            [threads_per_tg, 1, 1],
204            0,
205            stream,
206            &[
207                KernelArg::Bytes(&self.out_features.to_le_bytes()),
208                KernelArg::Bytes(&self.in_features.to_le_bytes()),
209                KernelArg::Bytes(&self.group_size.to_le_bytes()),
210                KernelArg::Buffer(self.packed),
211                KernelArg::Buffer(self.scales),
212                KernelArg::Buffer(self.biases),
213                KernelArg::Buffer(x),
214                KernelArg::Buffer(y),
215            ],
216        )
217    }
218
219    /// 2026-09-25: [`Self::gemv_silu_gate`] plus the residual, in one launch:
220    ///   `y[n] = x_resid[n] + sum_k W[n, k] * (silu(gate[k]) * up[k])`.
221    pub fn gemv_silu_gate_resid(
222        &self,
223        gpu: &dyn GpuBackend,
224        gate: DevicePtr,
225        up: DevicePtr,
226        x_resid: DevicePtr,
227        y: DevicePtr,
228        stream: u64,
229    ) -> Result<()> {
230        let kernel = gpu.kernel("mlx_int8_gemv_silu_gate", "mlx_int8_gemv_silu_gate_resid")?;
231        const ROWS_PER_TG: u32 = 4;
232        const SIMDGROUP_SIZE: u32 = 32;
233        let threads_per_tg: u32 = ROWS_PER_TG * SIMDGROUP_SIZE;
234        let row_groups = self.out_features.div_ceil(ROWS_PER_TG);
235        gpu.launch_typed(
236            kernel,
237            [row_groups, 1, 1],
238            [threads_per_tg, 1, 1],
239            0,
240            stream,
241            &[
242                KernelArg::Bytes(&self.out_features.to_le_bytes()),
243                KernelArg::Bytes(&self.in_features.to_le_bytes()),
244                KernelArg::Bytes(&self.group_size.to_le_bytes()),
245                KernelArg::Buffer(self.packed),
246                KernelArg::Buffer(self.scales),
247                KernelArg::Buffer(self.biases),
248                KernelArg::Buffer(gate),
249                KernelArg::Buffer(up),
250                KernelArg::Buffer(x_resid),
251                KernelArg::Buffer(y),
252            ],
253        )
254    }
255
256    /// 2026-09-25: `y = W (silu(gate) * up)` in one launch of
257    /// `mlx_int8_gemv_silu_gate`; `silu(gate) * up` is never written to memory.
258    pub fn gemv_silu_gate(
259        &self,
260        gpu: &dyn GpuBackend,
261        gate: DevicePtr,
262        up: DevicePtr,
263        y: DevicePtr,
264        stream: u64,
265    ) -> Result<()> {
266        let kernel = gpu.kernel("mlx_int8_gemv_silu_gate", "mlx_int8_gemv_silu_gate")?;
267        const ROWS_PER_TG: u32 = 4;
268        const SIMDGROUP_SIZE: u32 = 32;
269        let threads_per_tg: u32 = ROWS_PER_TG * SIMDGROUP_SIZE;
270        let row_groups = self.out_features.div_ceil(ROWS_PER_TG);
271        gpu.launch_typed(
272            kernel,
273            [row_groups, 1, 1],
274            [threads_per_tg, 1, 1],
275            0,
276            stream,
277            &[
278                KernelArg::Bytes(&self.out_features.to_le_bytes()),
279                KernelArg::Bytes(&self.in_features.to_le_bytes()),
280                KernelArg::Bytes(&self.group_size.to_le_bytes()),
281                KernelArg::Buffer(self.packed),
282                KernelArg::Buffer(self.scales),
283                KernelArg::Buffer(self.biases),
284                KernelArg::Buffer(gate),
285                KernelArg::Buffer(up),
286                KernelArg::Buffer(y),
287            ],
288        )
289    }
290
291    /// 2026-09-25: `Y = X W^T`, with `X` BF16 `[m, in_features]` and `Y` BF16
292    /// `[m, out_features]`. Runs `mlx_int8_gemm`, one thread per output element.
293    pub fn gemm(
294        &self,
295        gpu: &dyn GpuBackend,
296        x: DevicePtr,
297        y: DevicePtr,
298        m: u32,
299        stream: u64,
300    ) -> Result<()> {
301        let kernel = gpu.kernel("mlx_int8_gemm", "mlx_int8_gemm")?;
302        let block_x: u32 = 16;
303        let block_y: u32 = 16;
304        let grid_x = self.out_features.div_ceil(block_x);
305        let grid_y = m.div_ceil(block_y);
306        gpu.launch_typed(
307            kernel,
308            [grid_x, grid_y, 1],
309            [block_x, block_y, 1],
310            0,
311            stream,
312            &[
313                KernelArg::Bytes(&m.to_le_bytes()),
314                KernelArg::Bytes(&self.out_features.to_le_bytes()),
315                KernelArg::Bytes(&self.in_features.to_le_bytes()),
316                KernelArg::Bytes(&self.group_size.to_le_bytes()),
317                KernelArg::Buffer(x),
318                KernelArg::Buffer(self.packed),
319                KernelArg::Buffer(self.scales),
320                KernelArg::Buffer(self.biases),
321                KernelArg::Buffer(y),
322            ],
323        )
324    }
325
326    /// 2026-09-25: Free the three GPU buffers. `DevicePtr` is a plain handle, so
327    /// nothing frees them automatically; call this once.
328    pub fn release(&self, gpu: &dyn GpuBackend) -> Result<()> {
329        gpu.free(self.packed)?;
330        gpu.free(self.scales)?;
331        gpu.free(self.biases)?;
332        Ok(())
333    }
334}
335
336/// 2026-09-25: `gate_y = gate x` and `up_y = up x` in one launch of
337/// `mlx_int8_gemv_gate_up`, which reads `x` once for both. Both weights must share
338/// `(out_features, in_features, group_size)`; only `debug_assert_eq!` checks it.
339pub fn gemv_gate_up(
340    gpu: &dyn GpuBackend,
341    gate: &MlxInt8Weight,
342    up: &MlxInt8Weight,
343    x: DevicePtr,
344    gate_y: DevicePtr,
345    up_y: DevicePtr,
346    stream: u64,
347) -> Result<()> {
348    debug_assert_eq!(gate.out_features, up.out_features);
349    debug_assert_eq!(gate.in_features, up.in_features);
350    debug_assert_eq!(gate.group_size, up.group_size);
351    let kernel = gpu.kernel("mlx_int8_gemv_gate_up", "mlx_int8_gemv_gate_up")?;
352    const ROWS_PER_TG: u32 = 4;
353    const SIMDGROUP_SIZE: u32 = 32;
354    let threads_per_tg: u32 = ROWS_PER_TG * SIMDGROUP_SIZE;
355    let row_groups = gate.out_features.div_ceil(ROWS_PER_TG);
356    gpu.launch_typed(
357        kernel,
358        [row_groups, 1, 1],
359        [threads_per_tg, 1, 1],
360        0,
361        stream,
362        &[
363            KernelArg::Bytes(&gate.out_features.to_le_bytes()),
364            KernelArg::Bytes(&gate.in_features.to_le_bytes()),
365            KernelArg::Bytes(&gate.group_size.to_le_bytes()),
366            KernelArg::Buffer(gate.packed),
367            KernelArg::Buffer(gate.scales),
368            KernelArg::Buffer(gate.biases),
369            KernelArg::Buffer(up.packed),
370            KernelArg::Buffer(up.scales),
371            KernelArg::Buffer(up.biases),
372            KernelArg::Buffer(x),
373            KernelArg::Buffer(gate_y),
374            KernelArg::Buffer(up_y),
375        ],
376    )
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382
383    #[test]
384    fn parse_quant_config_from_mlx_layout() {
385        let cfg: JsonValue = serde_json::json!({
386            "quantization": { "bits": 8, "group_size": 64, "mode": "affine" },
387            "model_type": "qwen3_5",
388        });
389        let q = MlxQuantConfig::from_config(&cfg).expect("expected quant block");
390        assert_eq!(q.bits, 8);
391        assert_eq!(q.group_size, 64);
392    }
393
394    #[test]
395    fn parse_quant_config_falls_back_to_quantization_config() {
396        let cfg: JsonValue = serde_json::json!({
397            "quantization_config": { "bits": 8, "group_size": 64 },
398        });
399        let q = MlxQuantConfig::from_config(&cfg).expect("expected quant_config block");
400        assert_eq!(q.bits, 8);
401        assert_eq!(q.group_size, 64);
402    }
403
404    #[test]
405    fn parse_quant_config_returns_none_when_absent() {
406        let cfg: JsonValue = serde_json::json!({ "model_type": "qwen3_5" });
407        assert!(MlxQuantConfig::from_config(&cfg).is_none());
408    }
409}
metrale_gpu_runtime/mlx_int8.rs

metrale_gpu_runtime/
mlx_int8.rs
