Enum WeightStorage 

pub enum WeightStorage {
    F32(Arc<[f32]>),
    BF16(Arc<[bf16]>),
    Q4_0 {
        words: Arc<[u32]>,
        bytes_len: usize,
        in_features: usize,
        out_features: usize,
    },
    WithLoRA {
        base: Box<WeightStorage>,
        lora_a: Arc<[f32]>,
        lora_b: Arc<[f32]>,
        rank: usize,
        alpha: f32,
        in_features: usize,
        out_features: usize,
    },
}

Weight tensor storage that preserves source precision.

Projection weights (Q/K/V/O for attention, gate/up/down for FFN, and the output lm_head matrix) stay in whatever dtype the source checkpoint used — f32 when that’s how it was saved, bf16 for modern HF checkpoints that ship bf16 to halve weight memory. Activations are usually f32, but LlamaModel’s cached decode path (Phase D increment A, “BF16-throughout decode”) threads the KV cache’s dtype through the whole activation stream — when that’s BF16, weights must ALSO be BF16 (matmul’s dtype gate only allows same-dtype homogeneous or (lhs=F32, rhs=BF16); a BF16-activation × F32-weight matmul is rejected, the opposite direction from the classic f32-activation-over-bf16-weight quantized-serving case). F32 activations still route the mixed (A:F32, B:BF16) → F32 precision via VulkanBackend::matmul unchanged.

Norm gains and biases are NOT covered by this enum — they’re small and precision-sensitive, so they stay Arc<[f32]>.

Cloning is cheap (Arc bump) for both variants. Use WeightStorage::const_like to emit a Tensor Const node with the right dtype.

Variants

F32(Arc<[f32]>)

BF16(Arc<[bf16]>)

Q4_0

GGML Q4_0 blocks (raw byte stream), laid out as [out_features, in_features / 32] blocks (18 bytes each, llama.cpp convention). Stored as Arc<[u32]> — the byte stream reinterpreted as u32 words so subsequent forward passes just Arc-clone (cheap) rather than recopying the bytes. The graph sees this directly as a U32 tensor; matmul dispatch goes through Op::QMatMul.

bytes_len is the original byte count (u32_len * 4) so the const_like shape computation doesn’t accidentally round up.

Fields

words: Arc<[u32]>

bytes_len: usize

in_features: usize

out_features: usize

WithLoRA

Base weight wrapped with a trainable LoRA (Low-Rank Adaptation) update: effective weight W_eff = base + (alpha / rank) · A · B where A has shape [in_features, rank] and B has shape [rank, out_features] (both stored in the same layout convention as F32 weights — [in, out]).

Used for PEFT-style inference with frozen base weights (which can be F32, BF16, or Q4_0) plus small trainable adapter matrices. The adapter is cheap to apply — for a 2560×2560 projection at rank 8 the LoRA path is ~0.5% of the base matmul cost.

Fields

base: Box<WeightStorage>

lora_a: Arc<[f32]>

[in_features, rank] adapter A (HF’s lora_A transposed).

lora_b: Arc<[f32]>

[rank, out_features] adapter B (HF’s lora_B transposed).

rank: usize

alpha: f32

LoRA scaling factor; effective scale is alpha / rank.

in_features: usize

out_features: usize

Implementations

impl WeightStorage

pub fn elem_count(&self) -> usize

pub fn dtype(&self) -> DType

pub fn const_like(&self, anchor: &Tensor, shape: Shape) -> Result<Tensor, Error>

Emit a Const node on anchor’s graph matching this storage’s dtype. Used everywhere the forward pass wraps a weight into a Tensor.

For Q4_0, the emitted tensor is a 1-D U32 const of length bytes.len() / 4 holding the raw block byte stream. Callers must pair this with Tensor::qmatmul rather than matmul.

Returns Err for WithLoRA — the base + LoRA update can only be applied via apply_linear so the right graph structure is built.

pub fn apply_linear_with_bias( &self, x: &Tensor, in_features: usize, out_features: usize, bias: Arc<[f32]>, ) -> Result<Tensor, Error>

Produce X @ W + bias for this weight storage. Bias is a length-out_features Arc<f32> materialized fresh on the receiver’s graph and broadcast across the leading dims of the output.

Equivalent to the per-port apply_linear_with_bias helpers that several ports inlined — promoted here so call sites stop drifting.

pub fn apply_linear( &self, x: &Tensor, in_features: usize, out_features: usize, ) -> Result<Tensor, Error>

Produce X @ W (with optional bias) for this weight storage. Dispatches to matmul for F32/BF16 weights and to qmatmul for Q4_0. x’s dtype must be compatible with this weight’s dtype per matmul’s gate: same-dtype homogeneous (e.g. BF16 activations × BF16 weights — LlamaModel’s BF16-throughout decode path, Phase D increment A) or (x=F32, weight=BF16) (f32-activation quantized-weight serving). (x=BF16, weight=F32) is NOT supported — cast the weight or the activation to match before calling.

Errors

Returns a typed error rather than panicking on a dimension or dtype mismatch. This used to .unwrap() the inner matmul/qmatmul and assert_eq! the stored-vs-requested dimensions, which made a production path panic — a mismatched checkpoint (a GGUF whose stored shape disagrees with the model config) is a data condition, not a programming invariant, so it must surface as Err.

pub fn apply_linear_dyn_m( &self, x: &Tensor, in_features: usize, out_features: usize, row_count: DynScalar, ) -> Result<Tensor, Error>

Data-determined-M variant of Self::apply_linear: computes only row_count rows of the capacity-buffer projection x @ W, the rest left zeroed — the per-expert sparse-MoE FFN path. x is a [capacity, in_features] buffer whose first row_count rows are the routed tokens; the returned [capacity, out_features] buffer’s tail (rows row_count..capacity) stays exactly zero (see Tensor::matmul_dyn_m).

F32-only (mirrors matmul_dyn_m); other weight encodings surface a typed build-time error. No bias is applied — a bias would fill the un-computed tail, which the sparse-MoE caller relies on being zero for a correct index_add scatter-back.

pub fn with_lora( self, lora_a: Arc<[f32]>, lora_b: Arc<[f32]>, rank: usize, alpha: f32, in_features: usize, out_features: usize, ) -> Self

Wrap this weight storage with a LoRA adapter. Asserts that the adapter shapes match in_features/out_features. Panics if the base is already a WithLoRA (nested adapters aren’t supported; merge them explicitly if needed).

Trait Implementations

impl Clone for WeightStorage

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for WeightStorage

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl From<Arc<[bf16]>> for WeightStorage

fn from(a: Arc<[bf16]>) -> Self

Converts to this type from the input type.

impl From<Arc<[f32]>> for WeightStorage

fn from(a: Arc<[f32]>) -> Self

Converts to this type from the input type.

impl From<Vec<bf16>> for WeightStorage

fn from(v: Vec<bf16>) -> Self

Converts to this type from the input type.

impl From<Vec<f32>> for WeightStorage

fn from(v: Vec<f32>) -> Self

Converts to this type from the input type.