Enum WeightMatrix 

pub enum WeightMatrix {
    F32(Tensor),
    Quantized {
        data: WeightBytes,
        rows: usize,
        cols: usize,
        kind: QuantKind,
    },
    Mxfp4 {
        packed: WeightBytes,
        scale: WeightBytes,
        rows: usize,
        cols: usize,
    },
}

Variants

F32(Tensor)

Quantized

Fields

data: WeightBytes

rows: usize

cols: usize

kind: QuantKind

Mxfp4

MXFP4 (OCP Microscaling 4-bit float, Kimi K3’s real routed-expert format): unlike every Quantized kind above, which store one interleaved block buffer per row, Kimi K3’s real checkpoint stores the packed 4-bit codes and per-group E8M0 scales as two separate tensors (confirmed against a real shard header, see ferrox_quant’s MXFP4 module docs) – so this variant holds two independently zero-copy-mappable buffers instead of Quantized’s single data buffer. apply/apply_batch dispatch to ferrox_quant::dot_mxfp4_row_f32, which reads directly from these buffers without ever materializing a dequantized f32 copy of the whole matrix – the same zero-copy-mmap-plus-fused-dot discipline as every Quantized kind, letting a real MXFP4 checkpoint’s resident memory stay close to its on-disk size instead of the ~8x larger eager-f32-dequant footprint.

Fields

packed: WeightBytes

scale: WeightBytes

rows: usize

cols: usize

Implementations

impl WeightMatrix

pub fn rows(&self) -> usize

pub fn quant_kind(&self) -> Option<QuantKind>

The block format, or None for the two non-block storages (F32, safetensors-pair Mxfp4). This is the key every kernel-availability table is indexed by.

pub fn cols(&self) -> usize

pub fn dequant_row(&self, r: usize) -> Vec<f32>

Dequantizes exactly one row to f32, without touching any other row’s bytes. This is what makes a quantized embedding table usable directly: token lookup reads row_bytes bytes and dequantizes cols values, instead of the whole vocabulary tensor ever being widened to f32 (which for a large-vocab model is a multi-GB allocation that exists only to be indexed one row at a time).

pub fn apply(&self, x: &[f32]) -> Vec<f32>

Computes W @ x for a single activation vector x of length self.cols(), returning a vector of length self.rows(). Parallelized over output rows with rayon, same decomposition as matmul_f32.

With --features metal / --features cuda, when the matching dense GPU env selects a device (see [metal_dense_enabled] / [cuda_dense_enabled]), quantized kinds that have a GPU kernel go through [Self::apply_gpu] first so dense Llama-class decode uses the GPU instead of only MoE expert placement.

pub fn apply_three( a: &Self, b: &Self, c: &Self, x: &[f32], ) -> (Vec<f32>, Vec<f32>, Vec<f32>)

CPU-only matvec (NEON/AVX/scalar via ferrox-quant). Used by Self::apply after Metal miss/disable, and by GPU parity tests that must not recurse into [Self::apply_gpu]. Applies three independent matrices to the same activation, overlapping their parallel regions instead of running them one after another.

Decode opens one rayon fork-join per weight matrix – roughly seven per layer – and the measured CPU decode deficit is scheduling, not kernels (ferrox scales 1.40x/2.93x from 1 to 6 threads where llama.cpp scales 1.99x/4.39x, while beating llama at one thread). q/k/v share an input and are independent, so their regions can coexist and let rayon’s work-stealing fill threads that would otherwise idle at the tail of each one.

CPU only. On a GPU backend each apply submits and waits on its own command buffer, and Metal decode is already at or ahead of parity – there is nothing to win and a live path to disturb.

pub fn apply_cpu(&self, x: &[f32]) -> Vec<f32>

pub fn apply_cpu_q8(&self, act: &Q8Activations) -> Option<Vec<f32>>

INT_DOT matvec against a pre-quantized Q8_0 activation (shared gate/up).

pub fn dot_pair_cpu_q8( &self, row: usize, act: &Q8Activations, ) -> Option<(f32, f32)>

Two contiguous rows × one Q8 act (shared act loads). Q4_0 uses ferrox_quant::dot_q4_0_q8_2row; Q8_0 falls back to two singles.

pub fn dot_row_cpu_q8(&self, row: usize, act: &Q8Activations) -> Option<f32>

Single-row INT_DOT against pre-quantized Q8_0 acts (llama mul_mat_id inner loop). Returns None if this matrix is not Q4_0/Q8_0 INT_DOT.

pub fn apply_batch(&self, x_batch: &[f32], batch_size: usize) -> Vec<f32>

Computes W @ X for a batch of activation vectors at once: x_batch is batch_size rows of self.cols() elements each, flattened row-major; returns batch_size rows of self.rows() elements each, flattened row-major ([batch, rows], matching the layout Tensor/Decoder expect for chaining into further matmuls).

This is not just a convenience wrapper: for a quantized matrix, each weight row’s bytes are read from memory once and dotted against every activation in the batch, instead of once per apply call. For a memory-bandwidth-bound quantized matmul – which fused Q8_0/Q4_0 dot products are, since the whole point of keeping weights quantized is that reading them is the bottleneck, not the arithmetic – processing batch_size positions this way costs roughly the same memory traffic as processing one position, not batch_size times as much. This is the same reason speculative-decoding verification and batched prefill are faster per-token than sequential single-token decode on real hardware: it turns batch_size separate reads of the same weights into one.

With Metal dense enabled, dispatches a single batched Metal command buffer — Q4_K/Q6_K use [ferrox_metal::gpu::launch_q4_k_matmul_batch] / [ferrox_metal::gpu::launch_q6_k_matmul_batch] when batch_size >= 2; other kinds use [ferrox_metal::gpu::launch_matvec_batch]. Falls back to per-row Self::apply if the batch launch fails.

pub fn quantize_batch_acts( &self, x_batch: &[f32], batch_size: usize, ) -> Option<BatchActs>

Quantize x_batch once, in the activation format this matrix’s INT_DOT batch path consumes, for sharing across every projection that reads the same input (q/k/v on one normed batch; gate/up on another). Returns None when Self::apply_batch would not use quantized activations for this matrix — GPU dispatch, INT_DOT off, unsupported kind or width — so callers can pass the result straight to Self::apply_batch_with_acts unconditionally.

pub fn apply_batch_with_acts( &self, x_batch: &[f32], batch_size: usize, shared: Option<&BatchActs>, ) -> Vec<f32>

Self::apply_batch, optionally reusing a shared pre-quantized activation batch from Self::quantize_batch_acts. A shared value whose format or length does not match this matrix is simply ignored (the activations are re-quantized locally), so mixed-kind projection groups stay correct.

pub fn resident_bytes(&self) -> usize

Bytes actually resident in memory for this matrix – the number that matters for “can this model’s weights fit in RAM/VRAM at all,” as opposed to the always-4x-larger f32-expanded size.

pub fn probe_kernels(&self, role: &'static str)

Eagerly resolve, and record, every kernel lookup this matrix’s dispatch paths will make later, without dispatching anything.

Call once per weight while the model is being built, with role naming the tensor ("attn_q", "ffn_down", …). The predicates consulted here are the same functions the hot path consults, so the recorded prediction cannot drift from the decision. See crate::kernel_registry for why this exists and crate::kernel_registry::seal for what is done with it.

Observation only: nothing here influences a later dispatch.

pub fn probe_kernels_into( &self, reg: &Registry, role: &'static str, loc: &'static Location<'static>, )

Self::probe_kernels against an explicit registry and call site, so tests can probe into an instance of their own instead of the process-wide one.

pub fn probe_kernels_for( &self, reg: &Registry, backend: Backend, role: &'static str, loc: &'static Location<'static>, )

Self::probe_kernels_into against an explicit backend rather than active_backend. Lets a test on a CPU-only build ask what a Metal or CUDA run would resolve – which is the only way the kernel-coverage tests can run under plain cargo test --workspace, where every GPU feature is off.