Struct ExpertStack 

pub struct ExpertStack<B: Backend> {
    pub gate_up: Vec<Box<dyn Linear<B>>>,
    pub down: Vec<Box<dyn Linear<B>>>,
    pub gate_stacked: Option<B::QuantStore>,
    pub up_stacked: Option<B::QuantStore>,
    pub down_stacked: Option<B::QuantStore>,
}

Per-layer expert weights, materialised as [num_experts]-long vectors of Box<dyn Linear<B>>. Each entry runs the corresponding expert’s fused [gate; up] projection or its down projection.

B::Buffer is hidden behind Linear<B> so this struct is generic over backend, but Phase 2’s only consumer (moe_forward_cpu) is CPU- only — generic moe_forward<B> is deferred until the trait gains scaled-accumulate + cheap buffer slicing.

Fields

gate_up: Vec<Box<dyn Linear<B>>>

Fused [gate; up] projection per expert. Output shape per token: [2 * expert_intermediate] — the lower half is gate, upper is up.

down: Vec<Box<dyn Linear<B>>>

down projection per expert. Output shape per token: [hidden_size].

gate_stacked: Option<B::QuantStore>

Stacked-experts representation for backends that have a batched MoE indirect-dispatch kernel (Metal gemv_q4kw_moe_id_f32 / gemv_q6kw_moe_id_f32). Holds all experts for one matmul role in a single B::QuantStore with byte stride between expert slabs, so a single dispatch can cover all selected (token, expert) pairs at decode m=1.

None on backends without the kernel (CPU, CUDA-without-MoE-kernel) and on quant flavours that don’t have a stacked path yet — callers fall back to the per-expert gate_up / down Linears in those cases.

up_stacked: Option<B::QuantStore>

down_stacked: Option<B::QuantStore>

Implementations

impl<B: Backend> ExpertStack<B>

pub fn from_dense_stacks( gate_stack: &[f32], up_stack: &[f32], down_stack: &[f32], num_experts: usize, hidden_size: usize, expert_intermediate: usize, ) -> Result<Self>

Build from raw fp32 stacked tensors (test helper). Caller has already dequantised and laid out the data: gate_stack: [num_experts * expert_inter * hidden] up_stack: [num_experts * expert_inter * hidden] down_stack: [num_experts * hidden * expert_inter] Each per-expert slice is row-major in the natural Linear shape.

pub fn load_from_gguf( gguf: &GgufFile, layer_idx: usize, num_experts: usize, hidden_size: usize, expert_intermediate: usize, ) -> Result<Self>

Load all experts for one MoE layer from a GGUF file. Names follow the GGUF convention: blk.{layer_idx}.ffn_{gate,up,down}_exps.weight.

The loader picks between two strategies based on the on-disk dtype of the expert tensors:

• Quantised path (Q4_K / Q6_K only): each expert’s gate || up becomes a single QuantLinear<B> (Fused QuantStore — gate + up share n_cols = hidden), and down is a plain QuantLinear<B>. Block bytes stay compressed in backend memory; per-call dequant happens inside gemm_quant.
• Dense fallback (everything else, e.g. F32 / F16 / Q5_K until a kernel ships): eager-dequant to fp32 and wrap DenseLinear<B>. Memory inflates ~7× vs Q4_K_M but the algorithm is correctness-equivalent and this is the path the synthetic-MoE test fixtures need.

The runtime dispatcher (moe_forward<B>) doesn’t see which path was taken — it just calls Linear::forward per (token, expert).

pub fn open_and_load( path: impl AsRef<Path>, layer_idx: usize, num_experts: usize, hidden_size: usize, expert_intermediate: usize, ) -> Result<Self>

Convenience: open a GGUF and load layer layer_idx. The GGUF stays open inside this call only — for multi-layer loads use Self::load_from_gguf with a shared GgufFile.

pub fn num_experts(&self) -> usize

num_experts for the layer (consistency check helper).

Returns the per-expert Vec length, OR — when the stacked-only path is in effect (Metal MoE fast path with empty per-expert Vecs) — falls back to a stored count via the stacked variants. In the stacked-only case there’s no Vec to count, so this method is mostly used by tests on the per-expert path.