Trait QuantWeights 

pub trait QuantWeights: Send + Sync {
    // Required methods
    fn out_features(&self) -> u32;
    fn in_features(&self) -> u32;
    fn gemv(
        &self,
        gpu: &dyn GpuBackend,
        x: DevicePtr,
        y: DevicePtr,
        stream: u64,
    ) -> Result<()>;

    // Provided methods
    fn gemv_gate_up_with(
        &self,
        other: &Self,
        gpu: &dyn GpuBackend,
        x: DevicePtr,
        gate_y: DevicePtr,
        up_y: DevicePtr,
        stream: u64,
    ) -> Result<()>
       where Self: Sized { ... }
    fn gemv_silu_gate(
        &self,
        _gpu: &dyn GpuBackend,
        _gate: DevicePtr,
        _up: DevicePtr,
        _y: DevicePtr,
        _stream: u64,
    ) -> Result<()> { ... }
    fn gemv_silu_gate_resid(
        &self,
        _gpu: &dyn GpuBackend,
        _gate: DevicePtr,
        _up: DevicePtr,
        _x_resid: DevicePtr,
        _y: DevicePtr,
        _stream: u64,
    ) -> Result<()> { ... }
}

A quantised weight tensor that can drive matvec / matmul ops on a GpuBackend. Implementations live with each backend’s weight loader (Metal: MlxInt8Weight; future CUDA: Nvfp4Weight, Fp8DenseWeight, …).

Required Methods

fn out_features(&self) -> u32

Output dimension N of the underlying [N, K] weight.

fn in_features(&self) -> u32

Input dimension K of the underlying [N, K] weight.

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

Decode-path matvec: y = self @ x.

x is a BF16 buffer of length in_features(); y must hold at least out_features() BF16 slots.

Provided Methods

fn gemv_gate_up_with( &self, other: &Self, gpu: &dyn GpuBackend, x: DevicePtr, gate_y: DevicePtr, up_y: DevicePtr, stream: u64, ) -> Result<()>

where Self: Sized,

Dual-output GEMV with shared input: gate_y = self @ x, up_y = other @ x. The default impl is two serial gemv calls — correct on any backend, just slower than the fused kernel some backends ship (e.g. Metal’s mlx_int8_gemv_gate_up). Backends with a fused dual-output path override this to halve x-side memory bandwidth and remove a launch.

where Self: Sized keeps this method off the dyn-trait surface (it’s only callable through generic-parameter dispatch, which is what the forward modules use anyway).

fn gemv_silu_gate( &self, _gpu: &dyn GpuBackend, _gate: DevicePtr, _up: DevicePtr, _y: DevicePtr, _stream: u64, ) -> Result<()>

Fused FFN tail: y = self @ (silu(gate) ⊙ up).

Default impl errors — backends that ship the fused kernel override; backends that don’t can either error out (forcing the caller to do the unfused dance) or override with their own composition.

fn gemv_silu_gate_resid( &self, _gpu: &dyn GpuBackend, _gate: DevicePtr, _up: DevicePtr, _x_resid: DevicePtr, _y: DevicePtr, _stream: u64, ) -> Result<()>

Same as Self::gemv_silu_gate but additionally folds the layer-output residual addition into the same kernel: y[n] = x_resid[n] + sum_k self[n, k] * (silu(gate[k]) ⊙ up[k]).

Default impl errors. Backends that ship a _resid variant of the fused kernel override.

Implementations on Foreign Types

impl QuantWeights for MlxInt8Weight

fn out_features(&self) -> u32

fn in_features(&self) -> u32

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

fn gemv_gate_up_with( &self, other: &Self, gpu: &dyn GpuBackend, x: DevicePtr, gate_y: DevicePtr, up_y: DevicePtr, stream: u64, ) -> Result<()>

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

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

Implementors