Struct Decoder 

pub struct Decoder {
    pub config: ModelConfig,
    pub embedding: WeightMatrix,
    pub layers: Vec<LayerWeights>,
    pub final_norm: Vec<f32>,
    pub output_head: WeightMatrix,
    pub gpu_vram_budget_bytes: Option<u64>,
    pub gpt_oss: Option<GptOssWeights>,
    pub execution_plan: ExecutionPlan,
    pub plan_cache: Mutex<HashMap<PlanGeometry, FusedOpCaps>>,
}

Fields

config: ModelConfig

embedding: WeightMatrix

[vocab_size, hidden_dim]. A WeightMatrix rather than an eagerly-widened f32 Tensor, so a quantized token_embd.weight stays quantized on disk/mmap and token lookup dequantizes one row at a time (WeightMatrix::dequant_row) – a large-vocab model’s embedding table is multi-GB in f32 and only ever read row-wise.

layers: Vec<LayerWeights>

final_norm: Vec<f32>

output_head: WeightMatrix

gpu_vram_budget_bytes: Option<u64>

Real VRAM budget for GPU-resident routed experts. None (both constructors below set it) means every expert always runs on CPU – the exact behavior this field’s absence had before it existed. Some(bytes) makes each forward call build ONE global ResidencyPlan (Decoder::residency_plan) across every layer’s actual resident expert sizes and observed activation counts against this single budget – the budget is never re-spent per layer – dispatching device-placed routed experts through ferrox_moe::run_expert_placed (a real CUDA kernel when the cuda feature is compiled in and the expert’s quant kind has one; a correct CPU fallback otherwise, so setting this on a non-cuda build is harmless, just never GPU-accelerated). Shared experts and a dense layer’s sole expert always run on CPU regardless – every token activates them, so there’s no routing decision to offload the way routed-expert placement is. Rebuilding the plan on every forward call is real but not yet performance-tuned; a real, disclosed limit, not a correctness gap.

gpt_oss: Option<GptOssWeights>

Some only for the gpt-oss family. See GptOssWeights. When set, every layer runs the gpt-oss CPU graph (attention sinks, alternating SWA, biased router + experts, swiglu_oai), GPU offload is refused at load time, and the paged-KV decode path is refused at call time — neither implements sinks, and answering with a different distribution is the failure this replaces.

execution_plan: ExecutionPlan

Load-time execution plan (family, fused-op caps, SWA/RoPE policy). Built once; hot path must not re-resolve architecture strings. See crate::execution_plan.

plan_cache: Mutex<HashMap<PlanGeometry, FusedOpCaps>>

Cache key hit → fused caps last used for that geometry (enables decode/prefill plan reuse without rebuilding residency).

Implementations

impl Decoder

pub fn probe_kernels(&self)

Eagerly resolve every kernel lookup this model’s dispatch paths will make, and record it in ferrox_core::kernel_registry before anything runs.

Call once, at the end of loading, immediately before ferrox_core::kernel_registry::seal. Nothing here dispatches or decides anything: it asks the same predicates the hot path asks and writes the answers down, so a kernel that is missing becomes a startup line instead of an unexplained benchmark row.

Routed experts held in an ExpertBacking::Stored layer are not probed – they exist only as byte ranges until a token routes to them, and materialising every expert here would defeat the bounded expert store. Their kinds are the same as the resident case, and a dispatch-site miss still trips the sealed registry.

pub fn new_random_small( config: ModelConfig, n_layers: usize, vocab_size: usize, ) -> Self

Builds a decoder with correctly-shaped, randomly initialized weights for config, but overrides n_layers and vocab_size with small test-scale numbers so it can actually be allocated and run inside a CI sandbox. Use this to validate the forward-pass plumbing only, never to draw conclusions about real model quality.

pub fn forward_token( &self, token_id: usize, pos: usize, kv_caches: &mut [KvCache], ) -> Vec<f32>

Runs one decode step for token_id at position pos, updating kv_caches (one per layer) in place, and returns the logits over the (test-scale) vocabulary.

pub fn forward_token_paged( &self, token_id: usize, pos: usize, kv_caches: &mut [PagedKvCache], stores: &mut [PagedKvStore], ) -> Result<Vec<f32>, PagedStoreExhausted>

Same computation as forward_token, but each layer’s K/V cache is a PagedKvCache (block-table-indexed into a per-layer PagedKvStore) instead of a KvCache’s contiguous buffer – exercises causal_gqa_attention_paged in a real decode loop instead of only in isolation. kv_caches/stores are parallel per-layer arrays, mirroring forward_token’s kv_caches: &mut [KvCache]. Must produce bit-identical output to forward_token given stores sized so no layer ever exhausts its blocks – pinned by forward_token_paged_matches_forward_token_bit_identical.

pub fn expert_store_stats(&self) -> Option<ExpertStoreStats>

The shared expert store’s live counters, when this model runs with store-backed (streamed) routed experts – None for fully resident models. Every store-backed layer shares one store, so the first one found speaks for the whole model.

pub fn forward_batch( &self, tokens: &[usize], start_pos: usize, kv_caches: &mut [KvCache], ) -> Vec<Vec<f32>>

Processes multiple new positions in one call instead of calling forward_token once per position. tokens[i] is the token at absolute position start_pos + i; all positions attend causally (position i sees positions 0..=i of this batch plus everything already in kv_caches, nothing later).

The attention block’s Q/K/V/O projections and the MoE router are computed as batched matmuls (WeightMatrix::apply_batch), which for quantized weights means each weight row is read from memory once and dotted against every position in the batch, not once per position – see apply_batch’s doc comment for why that’s a real memory-bandwidth saving, not just fewer function calls. The expert FFN stage is not batched: which expert(s) a position routes to is data-dependent per position, so positions routed to different experts can’t share a single matmul the way the shared Q/K/V/router projections can. RoPE and attention itself (causal masking, softmax) are also per-position, since they’re cheap relative to the matmuls and batching them would add complexity for little benefit.

This is what makes prompt-lookup speculative decoding (speculative module) actually save work rather than just reshuffle it: verifying k draft tokens costs one batched call here, not k calls to forward_token.

Thin wrapper over Self::forward_hidden_batch + output_head.

pub fn forward_batch_last( &self, tokens: &[usize], start_pos: usize, kv_caches: &mut [KvCache], ) -> Vec<f32>

Self::forward_batch for the common case where only the final position’s logits are wanted: prefill a prompt, then sample the next token. Runs output_head on one row instead of all batch_size of them.

The KV cache and every hidden state are identical either way — only the vocabulary projection is skipped, and only for rows whose logits the caller was going to drop. That projection is not a rounding error: it is [batch x hidden] x [hidden x vocab], which for a large-vocabulary model with a small body is a large share of prefill. V*H / (V*H + L*P_layer) comes to 30% on Gemma-3-1B, 21% on Llama-3.2-1B and SmolLM2, 23% on Gemma-2-2B. llama.cpp does not do this work at all during pp512 — llama_batch_get_one leaves logits unset, so inp_out_ids selects a single row.

Self::forward_batch stays for the callers that genuinely need every row: speculative verification checks each draft position, and /v1/embeddings pools over all of them.

pub fn forward_hidden_batch( &self, tokens: &[usize], start_pos: usize, kv_caches: &mut [KvCache], ) -> Vec<Vec<f32>>

Like Self::forward_batch, but returns final RMS-normed hidden states (pre-output_head) — one hidden_dim vector per input token. Used by /v1/embeddings pooling (mean / last).

pub fn forward_multi_seq( &self, tokens: &[usize], positions: &[usize], kv_caches: &mut [Vec<KvCache>], ) -> Vec<Vec<f32>>

Continuous-batching primitive: one decode step across N independent sequences, each contributing exactly one new token at its own current position, sharing every layer’s projection/router matmuls the same way forward_batch shares them across positions of a single sequence – but each sequence keeps its own KvCache, independent seq_len, and independent position, so sequences admitted/evicted at different times can still share one batched matmul per step (this is what “continuous” batching means: the batch membership can change every step, unlike forward_batch’s fixed-size prompt-processing batch). kv_caches[s][l] is sequence s’s layer-l cache; tokens[s]/positions[s] is that sequence’s next token and its position within its own history. Returns one logits vector per sequence, same order as tokens.

Must produce bit-identical output to calling forward_token once per sequence with that sequence’s own cache/position – batching independent sequences together is a scheduling detail, not a math change (no sequence’s attention ever reads another sequence’s cache).

impl Decoder

pub fn from_gguf( path: impl AsRef<Path>, config: ModelConfig, ) -> Result<Self, LoadError>

Loads real weights from path for the given config. config supplies the architecture shape (layer count, head counts, MoE topology); tensor names are resolved against it using the llama.cpp naming convention described in the module docs.

A config.moe.n_experts <= 1 model is treated as dense: expert weights are read from the plain blk.N.ffn_{gate,up,down}.weight tensor names rather than the packed 3D _exps variant.

pub fn from_gguf_with_expert_cache( path: impl AsRef<Path>, config: ModelConfig, expert_cache_bytes: Option<u64>, ) -> Result<Self, LoadError>

Like from_gguf, but with expert_cache_bytes: Some(budget) routed experts are NOT loaded resident: each layer holds only byte-range layouts, and expert bytes are read on demand through one bounded, lease-protected ExpertStore shared by every layer (a single global byte budget; see ferrox_core::expert_store). Dense layers, shared experts, attention, embeddings, and the output head stay resident/mapped exactly as before – only routed experts stream. Layers whose expert tensors are F32/BF16 fall back to resident loading (the store exists for the quantized case). Output is bit-identical to the resident path – same bytes, same kernels – pinned by the roundtrip suite’s equivalence test.

Trait Implementations

impl Engine for Decoder

type State = Vec<KvCache>

fn new_state(&self) -> Vec<KvCache>

Builds fresh (empty) per-layer state for a new request.

fn vocab_size(&self) -> usize

fn forward_token( &self, token_id: usize, pos: usize, state: &mut Self::State, ) -> Vec<f32>