Struct MlxInt8Weight 

pub struct MlxInt8Weight {
    pub packed: DevicePtr,
    pub scales: DevicePtr,
    pub biases: DevicePtr,
    pub out_features: u32,
    pub in_features: u32,
    pub group_size: u32,
}

One MLX-int8 quantized linear weight resident on the GPU.

The fields are public because the consumer (transformer layer implementation) usually owns the pointers and frees them in batch at model teardown — there’s no per-weight Drop here.

Fields

packed: DevicePtr

[out_features, in_features / 4] packed bytes (uint32 words).

scales: DevicePtr

[out_features, in_features / group_size] per-group BF16 scales.

biases: DevicePtr

[out_features, in_features / group_size] per-group BF16 biases.

out_features: u32

in_features: u32

group_size: u32

Implementations

impl MlxInt8Weight

pub fn load( gpu: &dyn GpuBackend, st: &SafeTensors<'_>, base: &str, group_size: u32, ) -> Result<Self>

Load a (.weight, .scales, .biases) triplet from a parsed safetensors blob and upload to the GPU. base is the tensor name minus the suffix (e.g. "language_model.model.embed_tokens").

pub fn dequantize_to( &self, gpu: &dyn GpuBackend, out: DevicePtr, stream: u64, ) -> Result<()>

Materialize the full dequantized weight as BF16 into out, which must be a DevicePtr to a buffer of at least out_features * in_features * 2 bytes. Runs the mlx_int8_dequant Metal kernel under the hood.

pub fn gemv( &self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, stream: u64, ) -> Result<()>

Decode-path matvec: y = self_dequant @ x. x must be BF16 [in_features]; y must be a BF16 buffer with at least out_features slots. Runs the fused mlx_int8_gemv kernel.

pub fn gemv_silu_gate_resid( &self, gpu: &dyn GpuBackend, gate: DevicePtr, up: DevicePtr, x_resid: DevicePtr, y: DevicePtr, stream: u64, ) -> Result<()>

Like gemv_silu_gate, but additionally folds the residual stream addition into the same kernel: y[n] = x_resid[n] + sum_k self[n, k] * (silu(gate[k]) ⊙ up[k]) Eliminates the trailing bf16_add and the FFN-out staging buffer on the decoder layer’s exit.

pub fn gemv_silu_gate( &self, gpu: &dyn GpuBackend, gate: DevicePtr, up: DevicePtr, y: DevicePtr, stream: u64, ) -> Result<()>

Decode-path FFN-residual fusion: y = self @ (silu(gate) ⊙ up) Runs the fused mlx_int8_gemv_silu_gate kernel — replaces the silu_gate → gemv(down_proj) pair with a single launch and no INTERMEDIATE-sized staging buffer.

pub fn gemm( &self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, m: u32, stream: u64, ) -> Result<()>

Prefill-path GEMM: Y = X @ self_dequant^T. X is BF16 [m, in_features]; Y is BF16 [m, out_features]. Runs the fused mlx_int8_gemm kernel — straightforward correctness reference; tile-optimised replacement is a follow-on PR.

pub fn release(&self, gpu: &dyn GpuBackend) -> Result<()>

Free the three GPU buffers backing this weight. Idempotent if the pointers are null. Call this at model teardown — there’s no Drop because MlxInt8Weight is intentionally Copy-friendly (the DevicePtrs are u64 handles, not owners).