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: &SharedPagedKv, ) -> 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 the paged attention kernel 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 and, per attention arm, by every_paged_attention_arm_is_bit_identical_to_its_contiguous_twin.

This used to refuse gpt-oss outright, because the paged kernel had no attention-sink term and no sliding-window arm and would have answered differently from the contiguous path without saying so. It now mirrors all three arms of that dispatch, so the refusal is gone rather than merely relaxed.

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_with_hidden( &self, tokens: &[usize], start_pos: usize, kv_caches: &mut [KvCache], ) -> (Vec<Vec<f32>>, Vec<Vec<f32>>)

Self::forward_batch that also hands back the final-layer hidden state for every position instead of dropping it.

forward_batch computes these and throws them away; a hidden-state-conditioned drafter (EAGLE, MTP, dFlash) needs exactly the vector for the last verified position, so recomputing it would mean running the target model twice for something the first pass already had in hand. The extra cost here is one copy of [batch x hidden], which is why forward_batch keeps its move-only path for the prefill case that does not want them.

Returns (logits_per_position, hidden_per_position), both indexed by position in tokens.

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_batch_last_paged( &self, tokens: &[usize], start_pos: usize, kv_caches: &mut [PagedKvCache], stores: &SharedPagedKv, ) -> Result<Vec<f32>, PagedStoreExhausted>

Self::forward_batch_last over paged KV: the prefill twin of Self::forward_token_paged.

Why this gathers instead of paging the kernel

forward_hidden_batch’s fast arm hands cache.k / cache.v to causal_gqa_attention_prefill_shared_kv_windowed, which is Rayon over [query-block x head] against one flat KV buffer. That blocking is why CPU prefill is not the per-query path, and a block table cannot be handed to it as a slice.

The alternative was a second blocked kernel that reads through the table. This file has just finished paying for what a second copy of a rule costs: the paged decode path silently lost the window arm, the sink term, the attention softcap, the embedding scale and the final logit softcap, one at a time, because it was a copy. A prefill kernel is a much larger surface to keep in step than any of those. So the pages are materialised, the ONE prefill implementation every other path uses runs against them, and the new rows go back.

Bit-identity is therefore by construction rather than by agreement between two kernels: this calls the same function with the same values. What the tests pin is that the gather and the scatter are faithful, not that two implementations of attention happen to match.

The cost is one KV-sized copy per layer per call, against the matmuls that dominate prefill. Decode is untouched: it still reads through the block table and copies nothing, which is where page sharing pays.

Failure is checked before anything is written

Every layer’s blocks are reserved up front, so a store too small for the batch refuses with PagedStoreExhausted having mutated no layer. A partial append would leave some layers longer than others, and no caller can recover from that.

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>