fix(qwen4_exp): point the CUDA leader load at the moved parse_quant_string

Moving the function out of tp/worker.rs left dispatch.rs calling it
through worker's `use` import, which is private — so the call compiled
nowhere except in my head. The CUDA type-check caught it; the CPU
build cannot, because that call site is cfg(cuda) and every local check
I ran was blind to it.

That is what the job is for, and it is the second time this port has
relied on it. Worth remembering that "fmt, clippy and test are green"
means green on roughly half of this crate.

crates/neuron/src/harness/arch/mod.rs

@@ -22,3 +22,4 @@
 
 pub mod qwen3_5;
 pub mod qwen4_exp;
+pub mod snapshot;

crates/neuron/src/harness/arch/qwen3_5/decoder.rs

@@ -18,6 +18,7 @@ use candle_core::{Module, Tensor};
 use candle_nn::var_builder::ShardedVarBuilder;
 use std::sync::Arc;
 
+use super::super::snapshot::LayerKvSnapshot;
 use super::TextConfig;
 use super::full_attn::Qwen3_5Attention;
 use super::linear_attn::GatedDeltaNet;
@@ -25,7 +26,6 @@ use super::mlp::Qwen3_5MLP;
 use super::moe::Qwen3_5MoeBlock;
 use super::rmsnorm::Qwen3_5RmsNorm;
 use super::rope::RotaryEmbedding;
-use super::snapshot::LayerKvSnapshot;
 
 /// One of the two attention flavours sitting in a decoder layer's
 /// attention slot. Full-attention layers need the rotary table and

crates/neuron/src/harness/arch/qwen3_5/mlp.rs

@@ -19,9 +19,12 @@ pub struct Qwen3_5MLP {
 impl Qwen3_5MLP {
     /// Construct directly from pre-built projections.
     ///
-    /// `qwen4_exp` ships its routed experts as fused 3D tensors rather
-    /// than per-expert modules, so its loader slices them and builds the
-    /// experts here — same SwiGLU, different storage.
+    /// Test-only again: `qwen4_exp` briefly built its 512 experts this
+    /// way before the fused tensors stayed banked (`Experts::Banked`),
+    /// which removed the last caller outside tests. The differential
+    /// test that proves banking neutral still needs it, to express the
+    /// same weights in the layout it is being compared against.
+    #[cfg(test)]
     pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
         Self {
             gate_proj,

crates/neuron/src/harness/arch/qwen3_5/mod.rs

@@ -79,7 +79,6 @@ pub mod mlp;
 pub mod moe;
 pub mod rmsnorm;
 pub mod rope;
-pub mod snapshot;
 pub mod vision;
 
 use decoder::Qwen3_5DecoderLayer;
@@ -534,15 +533,18 @@ impl Qwen3_5Model {
     /// counter as one consistent prefix snapshot (#11). Only valid at
     /// a token boundary — i.e. between forward calls, which is the
     /// only time the caller can reach this anyway.
-    pub fn snapshot_kv_cache(&self) -> candle_core::Result<snapshot::KvCacheSnapshot> {
+    pub fn snapshot_kv_cache(
+        &self,
+    ) -> candle_core::Result<crate::harness::arch::snapshot::KvCacheSnapshot> {
         let layers = self
             .layers
             .iter()
             .map(|l| l.snapshot_kv())
             .collect::<candle_core::Result<Vec<_>>>()?;
-        Ok(snapshot::KvCacheSnapshot {
+        Ok(crate::harness::arch::snapshot::KvCacheSnapshot {
             layers,
             rope_delta: self.rope_delta,
+            ple: None,
         })
     }
 
@@ -550,7 +552,7 @@ impl Qwen3_5Model {
     /// snapshot. The snapshot stays valid for further restores.
     pub fn restore_kv_cache(
         &mut self,
-        snap: &snapshot::KvCacheSnapshot,
+        snap: &crate::harness::arch::snapshot::KvCacheSnapshot,
     ) -> candle_core::Result<()> {
         if snap.layers.len() != self.layers.len() {
             candle_core::bail!(
@@ -582,7 +584,7 @@ impl Qwen3_5Model {
     /// new token per batch row — with each row at its own sequence
     /// position `positions[i]` (typically `prefix_lens[i] + step`).
     /// The cache must hold batched state (see
-    /// `snapshot::assemble_batch`); `attn_mask` is the padding mask
+    /// `crate::harness::arch::snapshot::assemble_batch`); `attn_mask` is the padding mask
     /// from [`Self::batch_decode_mask`] (or `None` when no row is
     /// padded). Text-only: `rope_delta` is ignored — positions are
     /// explicit and vision requests never enter the batch path.
@@ -1084,14 +1086,16 @@ impl Qwen3_5ForCausalLM {
     }
 
     /// See [`Qwen3_5Model::snapshot_kv_cache`].
-    pub fn snapshot_kv_cache(&self) -> candle_core::Result<snapshot::KvCacheSnapshot> {
+    pub fn snapshot_kv_cache(
+        &self,
+    ) -> candle_core::Result<crate::harness::arch::snapshot::KvCacheSnapshot> {
         self.base.snapshot_kv_cache()
     }
 
     /// See [`Qwen3_5Model::restore_kv_cache`].
     pub fn restore_kv_cache(
         &mut self,
-        snap: &snapshot::KvCacheSnapshot,
+        snap: &crate::harness::arch::snapshot::KvCacheSnapshot,
     ) -> candle_core::Result<()> {
         self.base.restore_kv_cache(snap)
     }

crates/neuron/src/harness/arch/qwen3_5/moe.rs

@@ -29,18 +29,153 @@
 //! replaces the loop behind the same `forward` signature.
 
 use anyhow::{Context, Result};
-use candle_core::{DType, Module, Tensor};
+use candle_core::quantized::{GgmlDType, QMatMul, QTensor};
+use candle_core::{DType, IndexOp, Module, Tensor};
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
 use super::TextConfig;
 use super::mlp::Qwen3_5MLP;
 
+/// How a layer's routed experts are stored.
+///
+/// Two shapes, because the right one depends on whether the model fits.
+///
+/// `PerExpert` is the Qwen3-Next layout: one module per expert, sliced
+/// out at load. Fine when the weights are device-resident anyway.
+///
+/// `Banked` keeps the checkpoint's own fused tensors whole and takes a
+/// view per routed expert. `qwen4_exp` needs it: slicing its 512
+/// experts into modules copies **241.6 GB** at BF16 — 5.03 GB per
+/// layer — which fits in neither 64 GB of VRAM nor 123 GB of host RAM,
+/// at any quantisation. It is also the layout a host-resident bank with
+/// a device-side gather wants (#318), so the fused form is where that
+/// work starts rather than something it would have to undo.
+///
+/// Banking costs nothing per call: a slice along dim 0 of a contiguous
+/// tensor is contiguous, and `x.matmul(w.t())` is what `Linear` does
+/// anyway, so the views feed the same GEMMs without a copy.
+pub(crate) enum Experts {
+    PerExpert(Vec<Qwen3_5MLP>),
+    /// Quantised in situ at load (#315). One `(gate, up, down)` triple
+    /// per expert.
+    ///
+    /// This is what makes `qwen4_exp` loadable at all: 120.8 B routed
+    /// parameters are 241.6 GB at BF16 and 67.9 GB at q4k, against
+    /// beast's 64 GB of VRAM and ~110 GB of usable host RAM. Nothing
+    /// about the arithmetic works until they shrink.
+    ///
+    /// Quantising *from* the banked tensor is what keeps the peak
+    /// bounded: one layer's fused BF16 tensor is 5.03 GB, and it is
+    /// dropped once its 512 experts are quantised, so the full 241.6 GB
+    /// never exists at once.
+    Quantized(Vec<QuantizedExpert>),
+    Banked {
+        /// `(num_experts, 2 * intermediate, hidden)` — gate rows first,
+        /// then up. Reversing them computes `silu(up) * gate`, which is
+        /// a different function of the same weights (#312).
+        gate_up: Tensor,
+        /// `(num_experts, hidden, intermediate)`
+        down: Tensor,
+        intermediate: usize,
+    },
+}
+
+/// One expert's three projections, quantised.
+pub(crate) struct QuantizedExpert {
+    gate: QMatMul,
+    up: QMatMul,
+    down: QMatMul,
+}
+
+impl Experts {
+    pub(crate) fn len(&self) -> usize {
+        match self {
+            Experts::PerExpert(v) => v.len(),
+            Experts::Banked { gate_up, .. } => gate_up.dims()[0],
+            Experts::Quantized(v) => v.len(),
+        }
+    }
+
+    /// Quantise a banked pair into per-expert triples.
+    ///
+    /// Takes the fused tensors by reference and drops nothing itself —
+    /// the caller owns the peak, and should hold one layer at a time.
+    pub(crate) fn quantize_banked(
+        gate_up: &Tensor,
+        down: &Tensor,
+        intermediate: usize,
+        dtype: GgmlDType,
+    ) -> candle_core::Result<Self> {
+        let n = gate_up.dims()[0];
+        let mut out = Vec::with_capacity(n);
+        for e in 0..n {
+            let gu = gate_up.i(e)?;
+            // The source must own its storage outright.
+            // `QTensor::quantize` reads `src.storage()` — the whole
+            // backing buffer, ignoring shape and offset — so a view into
+            // the fused tensor hands it every expert's bytes while the
+            // shape claims one, and candle panics on a block-count
+            // mismatch inside `from_float`. Neither `contiguous()` nor
+            // `copy()` saves you here: a slice along dim 0 is already
+            // contiguous, so both are free to hand back something still
+            // sharing the parent's buffer. Round-tripping through a Vec
+            // is the part that actually allocates, and it is also where
+            // the f32 the quantiser wants comes from.
+            let q = |t: Tensor| -> candle_core::Result<QMatMul> {
+                let dims = t.dims().to_vec();
+                let vals = t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?;
+                let owned = Tensor::from_vec(vals, dims, t.device())?;
+                QMatMul::from_qtensor(QTensor::quantize(&owned, dtype)?)
+            };
+            out.push(QuantizedExpert {
+                gate: q(gu.narrow(0, 0, intermediate)?)?,
+                up: q(gu.narrow(0, intermediate, intermediate)?)?,
+                down: q(down.i(e)?)?,
+            });
+        }
+        Ok(Experts::Quantized(out))
+    }
+
+    /// Expert `e` applied to the rows routed to it.
+    pub(crate) fn forward_one(&self, e: usize, xs: &Tensor) -> candle_core::Result<Tensor> {
+        match self {
+            Experts::PerExpert(v) => v[e].forward(xs),
+            Experts::Banked {
+                gate_up,
+                down,
+                intermediate,
+            } => {
+                let gu = gate_up.i(e)?;
+                let gate_w = gu.narrow(0, 0, *intermediate)?;
+                let up_w = gu.narrow(0, *intermediate, *intermediate)?;
+                let lhs = candle_nn::ops::silu(&xs.matmul(&gate_w.t()?)?)?;
+                let rhs = xs.matmul(&up_w.t()?)?;
+                (lhs * rhs)?.matmul(&down.i(e)?.t()?)
+            }
+            Experts::Quantized(v) => {
+                let q = &v[e];
+                // The quantised kernels accumulate in f32 whatever the
+                // activation dtype, so the cast is explicit here rather
+                // than implied — and the result comes back in the
+                // caller's dtype so the routing arithmetic upstream
+                // does not have to know which layout it got.
+                let dtype = xs.dtype();
+                let xs32 = xs.to_dtype(DType::F32)?;
+                let lhs = candle_nn::ops::silu(&q.gate.forward(&xs32)?)?;
+                let rhs = q.up.forward(&xs32)?;
+                q.down.forward(&(lhs * rhs)?)?.to_dtype(dtype)
+            }
+        }
+    }
+}
+
 pub struct Qwen3_5MoeBlock {
     /// Router: `(num_experts, hidden)`, checkpoint name `mlp.gate`.
     gate: Linear,
-    /// Routed experts, checkpoint names `mlp.experts.{i}.{gate,up,down}_proj`.
-    experts: Vec<Qwen3_5MLP>,
+    /// Routed experts. Per-expert modules for `qwen3_5`; the
+    /// checkpoint's fused tensors for `qwen4_exp` — see [`Experts`].
+    experts: Experts,
     /// Always-on expert, `mlp.shared_expert.*`. `None` when the config
     /// declares no shared expert (Qwen3-30B-A3B style).
     shared_expert: Option<Qwen3_5MLP>,
@@ -62,7 +197,7 @@ impl Qwen3_5MoeBlock {
     /// `qwen3_5` config.
     pub(crate) fn from_parts(
         gate: Linear,
-        experts: Vec<Qwen3_5MLP>,
+        experts: Experts,
         shared_expert: Option<Qwen3_5MLP>,
         shared_expert_gate: Option<Linear>,
         num_experts_per_tok: usize,
@@ -129,7 +264,7 @@ impl Qwen3_5MoeBlock {
 
         Ok(Self {
             gate,
-            experts,
+            experts: Experts::PerExpert(experts),
             shared_expert,
             shared_expert_gate,
             num_experts_per_tok: cfg.num_experts_per_tok,
@@ -201,13 +336,13 @@ impl Module for Qwen3_5MoeBlock {
         )?;
 
         let mut ys = xs_flat.zeros_like()?;
-        for (e, expert) in self.experts.iter().enumerate() {
+        for e in 0..self.experts.len() {
             if tokens_for[e].is_empty() {
                 continue;
             }
             let rows = Tensor::new(tokens_for[e].as_slice(), xs.device())?;
             let picked = xs_flat.index_select(&rows, 0)?;
-            let out = expert.forward(&picked)?;
+            let out = self.experts.forward_one(e, &picked)?;
             let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
                 .to_dtype(out.dtype())?
                 .reshape(((), 1))?;
@@ -252,7 +387,7 @@ mod tests {
 
         let block = Qwen3_5MoeBlock {
             gate: Linear::new(randn(&[n_exp, hidden]), None),
-            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            experts: Experts::PerExpert((0..n_exp).map(|_| rand_mlp(hidden, inter)).collect()),
             shared_expert: Some(rand_mlp(hidden, inter)),
             shared_expert_gate: Some(Linear::new(randn(&[1, hidden]), None)),
             num_experts_per_tok: top_k,
@@ -283,8 +418,9 @@ mod tests {
             let mut expect = vec![0f32; hidden];
             for &e in selected {
                 let w = probs[e] / denom;
-                let out: Vec<f32> = block.experts[e]
-                    .forward(&row)
+                let out: Vec<f32> = block
+                    .experts
+                    .forward_one(e, &row)
                     .unwrap()
                     .flatten_all()
                     .unwrap()
@@ -336,7 +472,7 @@ mod tests {
         let (hidden, inter, n_exp) = (4, 2, 3);
         let block = Qwen3_5MoeBlock {
             gate: Linear::new(randn(&[n_exp, hidden]), None),
-            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            experts: Experts::PerExpert((0..n_exp).map(|_| rand_mlp(hidden, inter)).collect()),
             shared_expert: None,
             shared_expert_gate: None,
             num_experts_per_tok: 1,
@@ -368,8 +504,9 @@ mod tests {
             .max_by(|&a, &b| exps[a].partial_cmp(&exps[b]).unwrap())
             .unwrap();
         let w = exps[best] / sum;
-        let out: Vec<f32> = block.experts[best]
-            .forward(&xs.reshape(((), hidden)).unwrap())
+        let out: Vec<f32> = block
+            .experts
+            .forward_one(best, &xs.reshape(((), hidden)).unwrap())
             .unwrap()
             .flatten_all()
             .unwrap()
@@ -383,4 +520,155 @@ mod tests {
             );
         }
     }
+
+    /// Banking is a storage change, so the two layouts must compute the
+    /// same function. Same weights, expressed both ways, asserted
+    /// identical — otherwise "we only changed where the bytes live" is
+    /// a claim rather than a fact.
+    #[test]
+    fn banked_experts_equal_per_expert_modules() {
+        let (n_experts, hidden, inter) = (3usize, 4usize, 2usize);
+
+        // One source of truth for the weights, laid out both ways.
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+
+        let per_expert: Vec<Qwen3_5MLP> = (0..n_experts)
+            .map(|e| {
+                let gu = gate_up.i(e).unwrap();
+                Qwen3_5MLP::from_weights(
+                    Linear::new(gu.narrow(0, 0, inter).unwrap().contiguous().unwrap(), None),
+                    Linear::new(
+                        gu.narrow(0, inter, inter).unwrap().contiguous().unwrap(),
+                        None,
+                    ),
+                    Linear::new(down.i(e).unwrap().contiguous().unwrap(), None),
+                )
+            })
+            .collect();
+
+        let sliced = Experts::PerExpert(per_expert);
+        let banked = Experts::Banked {
+            gate_up,
+            down,
+            intermediate: inter,
+        };
+        assert_eq!(sliced.len(), banked.len());
+
+        let xs = randn(&[5, hidden]);
+        for e in 0..n_experts {
+            let a: Vec<f32> = sliced
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            let b: Vec<f32> = banked
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            assert_eq!(a.len(), b.len());
+            for (x, y) in a.iter().zip(b.iter()) {
+                assert!(
+                    (x - y).abs() < 1e-5,
+                    "expert {e} differs between layouts: {a:?} vs {b:?}"
+                );
+            }
+        }
+    }
+
+    /// And the halves are not interchangeable: reading `gate_up` the
+    /// other way round computes silu(up) * gate, which the banked path
+    /// would do just as silently as the sliced one.
+    #[test]
+    fn banked_reads_the_gate_before_the_up() {
+        let dev = Device::Cpu;
+        // gate row picks channel 0, up row picks channel 1.
+        let gate_up = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (1, 2, 2), &dev).unwrap();
+        let down = Tensor::from_vec(vec![1.0f32, 0.0], (1, 2, 1), &dev).unwrap();
+        let banked = Experts::Banked {
+            gate_up,
+            down,
+            intermediate: 1,
+        };
+        let xs = Tensor::from_vec(vec![1.0f32, 2.0], (1, 2), &dev).unwrap();
+        let got: Vec<f32> = banked
+            .forward_one(0, &xs)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // silu(gate.x) * (up.x) = silu(1) * 2
+        let want = (1.0f32 / (1.0 + (-1.0f32).exp())) * 2.0;
+        assert!((got[0] - want).abs() < 1e-5, "got {got:?} want {want}");
+        // the swapped reading would be silu(2) * 1
+        let swapped = (2.0f32 / (1.0 + (-2.0f32).exp())) * 1.0;
+        assert!(
+            (got[0] - swapped).abs() > 1e-3,
+            "this is silu(up) * gate: {got:?}"
+        );
+    }
+
+    /// In-situ quantisation has to approximate the dense expert, not
+    /// merely produce numbers of the right shape.
+    ///
+    /// Dimensions are 256 because the k-quants block on 256 elements —
+    /// a tiny fixture silently cannot be quantised at all, which is the
+    /// first thing that bites when this is tried on a toy model.
+    ///
+    /// The assertion is relative error against the dense output, with a
+    /// noise baseline: a q8_0 expert should be close, a q4k expert
+    /// looser, and both far nearer the truth than an unrelated expert
+    /// is. Without that last comparison a "quantised" path that
+    /// returned anything smoothly wrong would pass.
+    #[test]
+    fn quantised_experts_approximate_the_dense_ones() {
+        let (n_experts, hidden, inter) = (2usize, 256usize, 256usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let banked = Experts::Banked {
+            gate_up: gate_up.clone(),
+            down: down.clone(),
+            intermediate: inter,
+        };
+        let xs = randn(&[4, hidden]);
+
+        let rel_err = |a: &[f32], b: &[f32]| -> f32 {
+            let num: f32 = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum();
+            let den: f32 = a.iter().map(|x| x * x).sum::<f32>().max(1e-12);
+            (num / den).sqrt()
+        };
+        let out = |e: &Experts, i: usize| -> Vec<f32> {
+            e.forward_one(i, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap()
+        };
+
+        let dense0 = out(&banked, 0);
+        // Baseline: a different expert on the same input. Any quantised
+        // reading must beat this by a wide margin.
+        let noise = rel_err(&dense0, &out(&banked, 1));
+
+        for (dtype, tol) in [(GgmlDType::Q8_0, 0.05f32), (GgmlDType::Q4K, 0.20f32)] {
+            let q = Experts::quantize_banked(&gate_up, &down, inter, dtype).unwrap();
+            assert_eq!(q.len(), n_experts);
+            let err = rel_err(&dense0, &out(&q, 0));
+            assert!(err < tol, "{dtype:?} relative error {err:.4} exceeds {tol}");
+            assert!(
+                err < noise / 2.0,
+                "{dtype:?} error {err:.4} is not clearly better than an unrelated \
+                 expert ({noise:.4}) — the quantised path may not be reading \
+                 this expert at all"
+            );
+        }
+    }
 }

crates/neuron/src/harness/arch/qwen4_exp/decoder.rs

@@ -41,11 +41,14 @@ use crate::harness::arch::qwen3_5::linear_attn::GatedDeltaNet;
 use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
 use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
 
+use candle_core::quantized::GgmlDType;
+
 use super::config::TextConfig;
 use super::full_attn::Attention;
 use super::hyper::HyperConnection;
 use super::ple::PleBlock;
 use super::{linear_attn, moe};
+use crate::harness::arch::snapshot::LayerKvSnapshot;
 
 /// The attention slot: one layer in four is full attention with the QSA
 /// indexer, the rest are the recurrent delta rule.
@@ -71,6 +74,7 @@ impl DecoderLayer {
         cfg: &TextConfig,
         rotary: Arc<RotaryEmbedding>,
         layer_idx: usize,
+        quant: Option<GgmlDType>,
         vb: &ShardedVarBuilder,
     ) -> Result<Self> {
         let layer_type = cfg
@@ -83,6 +87,30 @@ impl DecoderLayer {
                     cfg.layer_types.len()
                 )
             })?;
+        Self::load_typed(
+            cfg,
+            rotary,
+            layer_type,
+            cfg.ple_layers().contains(&layer_idx),
+            quant,
+            vb,
+        )
+    }
+
+    /// Load a layer whose type the caller knows, rather than one looked
+    /// up in `layer_types`.
+    ///
+    /// The MTP head's layer (#313) needs this: it lives under
+    /// `mtp.layers.0` with its own single-entry type list, and it
+    /// carries no PLE wherever the main model puts one.
+    pub fn load_typed(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        layer_type: &str,
+        with_ple: bool,
+        quant: Option<GgmlDType>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
         let mixer = match layer_type {
             "full_attention" => {
                 Mixer::Full(Box::new(Attention::load(cfg, rotary, &vb.pp("self_attn"))?))
@@ -91,8 +119,8 @@ impl DecoderLayer {
                 Mixer::Linear(Box::new(linear_attn::load(cfg, &vb.pp("linear_attn"))?))
             }
             other => anyhow::bail!(
-                "unknown layer_type '{other}' for layer {layer_idx} (expected \
-                 'full_attention' or 'linear_attention')"
+                "unknown layer_type '{other}' (expected 'full_attention' or \
+                 'linear_attention')"
             ),
         };
 
@@ -109,7 +137,7 @@ impl DecoderLayer {
             )
         };
 
-        let ple = if cfg.ple_layers().contains(&layer_idx) {
+        let ple = if with_ple {
             Some(PleBlock::load(
                 &vb.pp("ple"),
                 cfg.hidden_size,
@@ -127,7 +155,7 @@ impl DecoderLayer {
             attn_hc: hc("attn_hyper_connection")?,
             mlp_hc: hc("mlp_hyper_connection")?,
             mixer,
-            mlp: moe::load(cfg, &vb.pp("mlp"))?,
+            mlp: moe::load(cfg, quant, &vb.pp("mlp"))?,
             ple,
         })
     }
@@ -190,6 +218,65 @@ impl DecoderLayer {
         Ok(())
     }
 
+    /// Capture this layer's attention-side state.
+    pub fn snapshot_kv(&self) -> candle_core::Result<LayerKvSnapshot> {
+        Ok(match &self.mixer {
+            Mixer::Full(attn) => {
+                let (kv, indexer_keys) = attn.snapshot_kv();
+                LayerKvSnapshot::FullSparse { kv, indexer_keys }
+            }
+            Mixer::Linear(net) => {
+                let (conv_state, recurrent_state) = net.snapshot_state()?;
+                LayerKvSnapshot::Linear {
+                    conv_state,
+                    recurrent_state,
+                }
+            }
+        })
+    }
+
+    /// Replace this layer's attention-side state. A variant mismatch
+    /// means the snapshot came from a different model, or from the same
+    /// model with a different `layer_types` — either way, restoring it
+    /// would put a recurrent state into an attention layer.
+    pub fn restore_kv(&mut self, snap: &LayerKvSnapshot) -> candle_core::Result<()> {
+        match (&mut self.mixer, snap) {
+            (Mixer::Full(attn), LayerKvSnapshot::FullSparse { kv, indexer_keys }) => {
+                attn.restore_kv(kv.as_ref(), indexer_keys.as_ref())
+            }
+            (
+                Mixer::Linear(net),
+                LayerKvSnapshot::Linear {
+                    conv_state,
+                    recurrent_state,
+                },
+            ) => net.restore_state(conv_state.as_ref(), recurrent_state.as_ref()),
+            _ => candle_core::bail!(
+                "restore_kv: snapshot layer kind does not match this layer's mixer kind"
+            ),
+        }
+    }
+
+    /// The PLE block's rolling conv context, if this layer carries one.
+    pub fn snapshot_ple(&self) -> candle_core::Result<Option<Tensor>> {
+        match &self.ple {
+            Some(block) => block.snapshot_state(),
+            None => Ok(None),
+        }
+    }
+
+    /// Restore the PLE conv context. Refuses a snapshot that disagrees
+    /// with this layer about whether PLE lives here at all.
+    pub fn restore_ple(&mut self, conv_state: Option<&Tensor>) -> candle_core::Result<()> {
+        match &mut self.ple {
+            Some(block) => block.restore_state(conv_state),
+            None if conv_state.is_none() => Ok(()),
+            None => candle_core::bail!(
+                "restore_ple: snapshot carries PLE state but this layer has no PLE block"
+            ),
+        }
+    }
+
     pub fn has_ple(&self) -> bool {
         self.ple.is_some()
     }

crates/neuron/src/harness/arch/qwen4_exp/full_attn.rs

@@ -170,6 +170,34 @@ impl Attention {
         Ok(self.o_proj.forward(&gated)?)
     }
 
+    /// Capture both caches at one token boundary.
+    ///
+    /// The same reason `clear_kv_cache` drops them together: a
+    /// snapshot holding only the attention K/V restores a layer whose
+    /// indexer still holds the previous request's keys, and the
+    /// selection it computes from them is wrong without being invalid.
+    pub fn snapshot_kv(&self) -> (Option<(Tensor, Tensor)>, Option<Tensor>) {
+        let kv = match (self.kv_cache.k(), self.kv_cache.v()) {
+            (Some(k), Some(v)) => Some((k.clone(), v.clone())),
+            _ => None,
+        };
+        (kv, self.indexer.snapshot_keys())
+    }
+
+    /// Replace both caches from a snapshot.
+    pub fn restore_kv(
+        &mut self,
+        kv: Option<&(Tensor, Tensor)>,
+        indexer_keys: Option<&Tensor>,
+    ) -> candle_core::Result<()> {
+        self.kv_cache.reset();
+        if let Some((k, v)) = kv {
+            self.kv_cache.append(k, v)?;
+        }
+        self.indexer.restore_keys(indexer_keys);
+        Ok(())
+    }
+
     /// Both caches, together. Dropping only one leaves the indexer
     /// selecting blocks of a previous request's text.
     pub fn clear_kv_cache(&mut self) {

crates/neuron/src/harness/arch/qwen4_exp/mod.rs

@@ -25,6 +25,8 @@
 //!   from `qwen3_5` with two settings corrected.
 //! - [`model`] — embeddings in, logits out; four streams, one
 //!   n-gram lookup, and no final norm.
+//! - [`mtp`] — the shipped draft head: one more decoder layer, fed
+//!   the target's pre-mixer stream and the next token's embedding.
 //! - [`moe`] — the FFN every layer has, loaded from fused expert
 //!   tensors instead of 512 modules.
 //! - [`ple`] — the hashed n-gram table on layer 1: how its rows
@@ -39,5 +41,6 @@ pub mod hyper;
 pub mod linear_attn;
 pub mod model;
 pub mod moe;
+pub mod mtp;
 pub mod ple;
 pub mod qsa;

crates/neuron/src/harness/arch/qwen4_exp/model.rs

@@ -34,6 +34,7 @@
 //! See `doc/qwen4_exp-port-spec.md`.
 
 use anyhow::{Context, Result, ensure};
+use candle_core::quantized::GgmlDType;
 use candle_core::{DType, Device, IndexOp, Module, Tensor};
 use candle_nn::var_builder::ShardedVarBuilder;
 use candle_nn::{Embedding, Linear};
@@ -45,6 +46,7 @@ use super::config::{Config, TextConfig};
 use super::decoder::DecoderLayer;
 use super::hyper::HyperConnection;
 use super::ple::{NGramHasher, NGramTable, ShardedNGramTable};
+use crate::harness::arch::snapshot::{KvCacheSnapshot, PleSnapshot};
 
 /// The hashed n-gram lookup, hoisted to the model because it is a
 /// function of the input ids alone.
@@ -89,6 +91,18 @@ impl NGramEmbedding {
         Ok(stacked)
     }
 
+    /// The carried ids, for a prefix snapshot. Two integers per batch
+    /// row — small, but they decide which n-grams the next token
+    /// hashes, so a restore that forgets them silently addresses the
+    /// wrong rows of a 320-million-row table.
+    fn snapshot(&self) -> Vec<Vec<i64>> {
+        self.carried.clone()
+    }
+
+    fn restore(&mut self, carried: &[Vec<i64>]) {
+        self.carried = carried.to_vec();
+    }
+
     fn reset(&mut self) {
         for row in &mut self.carried {
             row.clear();
@@ -111,9 +125,14 @@ pub struct Qwen4ExpForCausalLM {
 }
 
 impl Qwen4ExpForCausalLM {
+    /// `quant` quantises the routed experts in situ (#315). `None`
+    /// keeps the checkpoint's precision, which for this architecture is
+    /// useful only on a model small enough to hold: 241.6 GB of BF16
+    /// experts fit neither 64 GB of VRAM nor 123 GB of host RAM.
     pub fn load(
         cfg: &Config,
         dtype: DType,
+        quant: Option<GgmlDType>,
         device: &Device,
         vb: &ShardedVarBuilder,
     ) -> Result<Self> {
@@ -134,7 +153,7 @@ impl Qwen4ExpForCausalLM {
         let mut layers = Vec::with_capacity(text.num_hidden_layers);
         for i in 0..text.num_hidden_layers {
             layers.push(
-                DecoderLayer::load(text, rotary.clone(), i, &layers_vb.pp(i))
+                DecoderLayer::load(text, rotary.clone(), i, quant, &layers_vb.pp(i))
                     .with_context(|| format!("load decoder layer {i}"))?,
             );
         }
@@ -177,6 +196,24 @@ impl Qwen4ExpForCausalLM {
     /// `input_ids` is `(B, T)`; `offset` is the sequence position the
     /// first of them sits at. Returns logits `(B, T, vocab)`.
     pub fn forward(&mut self, input_ids: &Tensor, offset: usize) -> Result<Tensor> {
+        Ok(self.forward_inner(input_ids, offset)?.0)
+    }
+
+    /// Logits, plus the four-stream residual as it stood *before* the
+    /// mixer collapsed it.
+    ///
+    /// The MTP draft head consumes that pre-mixer stream, not the 2560
+    /// the LM head sees (#313), so it has to leave the model here rather
+    /// than be reconstructed from the logits — which it could not be.
+    pub fn forward_with_hidden(
+        &mut self,
+        input_ids: &Tensor,
+        offset: usize,
+    ) -> Result<(Tensor, Tensor)> {
+        self.forward_inner(input_ids, offset)
+    }
+
+    fn forward_inner(&mut self, input_ids: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
         let (batch, seq_len) = input_ids.dims2()?;
 
         // Four identical streams — the residual this architecture
@@ -209,7 +246,92 @@ impl Qwen4ExpForCausalLM {
 
         // The mixer's hc_norm is the only normalisation before the head.
         let x = self.mixer.collapse(&h)?;
-        Ok(self.lm_head.forward(&x)?)
+        Ok((self.lm_head.forward(&x)?, h))
+    }
+
+    /// Capture every piece of per-request state at one token boundary:
+    /// the attention K/V and the QSA indexer cache per full-attention
+    /// layer, the GatedDeltaNet conv and recurrent states per linear
+    /// layer, and PLE's conv context and carried ids.
+    ///
+    /// There is no `rope_delta` here — this architecture's text path
+    /// takes positions from the caller's offset rather than tracking a
+    /// vision-induced skew, so the field is recorded as zero. When the
+    /// vision tower lands (#314) that stops being true.
+    pub fn snapshot_kv_cache(&self) -> candle_core::Result<KvCacheSnapshot> {
+        let layers = self
+            .layers
+            .iter()
+            .map(|l| l.snapshot_kv())
+            .collect::<candle_core::Result<Vec<_>>>()?;
+        let conv_state = self
+            .layers
+            .iter()
+            .map(|l| l.snapshot_ple())
+            .collect::<candle_core::Result<Vec<_>>>()?
+            .into_iter()
+            .flatten()
+            .next();
+        Ok(KvCacheSnapshot {
+            layers,
+            rope_delta: 0,
+            ple: self.ngram.as_ref().map(|n| PleSnapshot {
+                conv_state,
+                carried_ids: n.snapshot(),
+            }),
+        })
+    }
+
+    /// Replace the live state from a snapshot. The snapshot stays valid
+    /// for further restores.
+    pub fn restore_kv_cache(&mut self, snap: &KvCacheSnapshot) -> candle_core::Result<()> {
+        if snap.layers.len() != self.layers.len() {
+            candle_core::bail!(
+                "restore_kv_cache: snapshot has {} layers, model has {}",
+                snap.layers.len(),
+                self.layers.len()
+            );
+        }
+        // A snapshot from a model whose PLE layer sits elsewhere would
+        // otherwise restore the conv context onto the wrong layer and
+        // leave this one holding the previous request's.
+        if snap.ple.is_some() != self.ngram.is_some() {
+            candle_core::bail!(
+                "restore_kv_cache: snapshot {} PLE state, model {} a PLE layer",
+                if snap.ple.is_some() {
+                    "carries"
+                } else {
+                    "has no"
+                },
+                if self.ngram.is_some() {
+                    "has"
+                } else {
+                    "has none"
+                }
+            );
+        }
+        for (layer, layer_snap) in self.layers.iter_mut().zip(snap.layers.iter()) {
+            layer.restore_kv(layer_snap)?;
+        }
+        let conv = snap.ple.as_ref().and_then(|p| p.conv_state.as_ref());
+        for layer in self.layers.iter_mut() {
+            if layer.has_ple() {
+                layer.restore_ple(conv)?;
+            }
+        }
+        if let (Some(ngram), Some(ple)) = (&mut self.ngram, snap.ple.as_ref()) {
+            ngram.restore(&ple.carried_ids);
+        }
+        Ok(())
+    }
+
+    /// `embed_tokens` applied to `input_ids`.
+    ///
+    /// The MTP head reuses this embedding
+    /// (`mtp_use_dedicated_embeddings: false`) but does not own it, so
+    /// the speculative loop borrows it from the target model.
+    pub fn embed(&self, input_ids: &Tensor) -> Result<Tensor> {
+        Ok(self.embed_tokens.forward(input_ids)?)
     }
 
     pub fn clear_kv_cache(&mut self) -> Result<()> {
@@ -320,6 +442,7 @@ fn rotary_for(cfg: &TextConfig, dtype: DType, device: &Device) -> Result<RotaryE
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::harness::arch::snapshot::LayerKvSnapshot;
     use candle_nn::var_builder::ShardedSafeTensors;
     use std::collections::HashMap;
 
@@ -535,6 +658,95 @@ mod tests {
                 }
             }
         }
+
+        // ---- the MTP draft head ----
+        // A second full-attention layer plus the two fc_* projections
+        // and their pre-norms. Emitted unconditionally: the head is what
+        // #313 builds on and the fixture is the only place its loader
+        // gets exercised.
+        {
+            let m = "mtp";
+            t.insert(format!("{m}.pre_fc_norm_embedding.weight"), rand(vec![h]));
+            t.insert(format!("{m}.pre_fc_norm_hidden.weight"), rand(vec![wide]));
+            t.insert(format!("{m}.fc_embedding.weight"), rand(vec![h, h]));
+            t.insert(format!("{m}.fc_hidden.weight"), rand(vec![h, h]));
+            t.insert(
+                format!("{m}.hyper_connection_mixer.hc_norm.weight"),
+                rand(vec![wide]),
+            );
+            t.insert(
+                format!("{m}.hyper_connection_mixer.input_mix_weight_down.weight"),
+                rand(vec![cfg.hc_lowrank, wide]),
+            );
+            t.insert(
+                format!("{m}.hyper_connection_mixer.input_mix_weight_up.weight"),
+                rand(vec![wide, cfg.hc_lowrank]),
+            );
+            let l = format!("{m}.layers.0");
+            for which in ["attn_hyper_connection", "mlp_hyper_connection"] {
+                t.insert(format!("{l}.{which}.hc_norm.weight"), rand(vec![wide]));
+                t.insert(
+                    format!("{l}.{which}.input_mix_weight_down.weight"),
+                    rand(vec![cfg.hc_lowrank, wide]),
+                );
+                t.insert(
+                    format!("{l}.{which}.input_mix_weight_up.weight"),
+                    rand(vec![wide, cfg.hc_lowrank]),
+                );
+                t.insert(
+                    format!("{l}.{which}.block_inject_weight.weight"),
+                    rand(vec![hc, wide]),
+                );
+            }
+            let a = format!("{l}.self_attn");
+            let heads = cfg.num_attention_heads * cfg.head_dim;
+            t.insert(format!("{a}.q_proj.weight"), rand(vec![heads * 2, h]));
+            t.insert(
+                format!("{a}.k_proj.weight"),
+                rand(vec![cfg.num_key_value_heads * cfg.head_dim, h]),
+            );
+            t.insert(
+                format!("{a}.v_proj.weight"),
+                rand(vec![cfg.num_key_value_heads * cfg.head_dim, h]),
+            );
+            t.insert(format!("{a}.o_proj.weight"), rand(vec![h, heads]));
+            t.insert(format!("{a}.q_norm.weight"), rand(vec![cfg.head_dim]));
+            t.insert(format!("{a}.k_norm.weight"), rand(vec![cfg.head_dim]));
+            let fused = (cfg.indexer_n_heads + cfg.indexer_kv_heads) * cfg.indexer_head_dim;
+            t.insert(
+                format!("{a}.indexer.index_qk_proj.weight"),
+                rand(vec![fused, h]),
+            );
+            t.insert(
+                format!("{a}.indexer.q_layernorm.weight"),
+                rand(vec![cfg.indexer_head_dim]),
+            );
+            t.insert(
+                format!("{a}.indexer.k_layernorm.weight"),
+                rand(vec![cfg.indexer_head_dim]),
+            );
+            let mm = format!("{l}.mlp");
+            t.insert(format!("{mm}.gate.weight"), rand(vec![cfg.num_experts, h]));
+            t.insert(
+                format!("{mm}.experts.gate_up_proj"),
+                rand(vec![cfg.num_experts, cfg.moe_intermediate_size * 2, h]),
+            );
+            t.insert(
+                format!("{mm}.experts.down_proj"),
+                rand(vec![cfg.num_experts, h, cfg.moe_intermediate_size]),
+            );
+            for pj in ["gate_proj", "up_proj"] {
+                t.insert(
+                    format!("{mm}.shared_expert.{pj}.weight"),
+                    rand(vec![cfg.shared_expert_intermediate_size, h]),
+                );
+            }
+            t.insert(
+                format!("{mm}.shared_expert.down_proj.weight"),
+                rand(vec![h, cfg.shared_expert_intermediate_size]),
+            );
+            t.insert(format!("{mm}.shared_expert_gate.weight"), rand(vec![1, h]));
+        }
         t
     }
 
@@ -545,7 +757,7 @@ mod tests {
         candle_core::safetensors::save(&tensors(&cfg.text_config), &path).unwrap();
         let vb =
             unsafe { ShardedSafeTensors::var_builder(&[&path], DType::F32, &Device::Cpu).unwrap() };
-        let model = Qwen4ExpForCausalLM::load(&cfg, DType::F32, &Device::Cpu, &vb).unwrap();
+        let model = Qwen4ExpForCausalLM::load(&cfg, DType::F32, None, &Device::Cpu, &vb).unwrap();
         (dir, model)
     }
 
@@ -603,4 +815,279 @@ mod tests {
         assert!(ngram.table.gather(&[(TABLE_ROWS - 1) as i64]).is_ok());
         assert!(ngram.table.gather(&[TABLE_ROWS as i64]).is_err());
     }
+
+    fn tensor_vec(t: &Tensor) -> Vec<f32> {
+        t.flatten_all()
+            .unwrap()
+            .to_dtype(DType::F32)
+            .unwrap()
+            .to_vec1()
+            .unwrap()
+    }
+
+    fn logits_vec(t: &Tensor) -> Vec<f32> {
+        t.flatten_all().unwrap().to_vec1().unwrap()
+    }
+
+    /// The property prefix caching rests on: a snapshot must *rewind*
+    /// the live state, not merely coexist with it.
+    ///
+    /// So the state is deliberately moved past the snapshot before the
+    /// restore — two further decode steps — and the restored step is
+    /// then required to reproduce the original exactly. An
+    /// implementation that captured nothing would not reproduce it; one
+    /// that restored only the attention K/V and left the indexer cache
+    /// at seven positions would trip the divergence guard rather than
+    /// return wrong numbers.
+    ///
+    /// The control is that a *cleared* model cannot take this step at
+    /// all: with an empty cache and a query at position 4 the guard
+    /// fires. That is what distinguishes restore from clear here, and
+    /// it is asserted rather than assumed.
+    #[test]
+    fn a_restored_snapshot_rewinds_rather_than_merely_continuing() {
+        let (_dir, mut model) = build();
+        let prompt = Tensor::from_vec(vec![5u32, 9, 3, 12], (1, 4), &Device::Cpu).unwrap();
+        let next = Tensor::from_vec(vec![7u32], (1, 1), &Device::Cpu).unwrap();
+
+        model.forward(&prompt, 0).unwrap();
+        let snap = model.snapshot_kv_cache().unwrap();
+        let want = logits_vec(&model.forward(&next, 4).unwrap());
+
+        // Move the live state well past the snapshot.
+        model.forward(&next, 5).unwrap();
+        model.forward(&next, 6).unwrap();
+
+        model.restore_kv_cache(&snap).unwrap();
+        let got = logits_vec(&model.forward(&next, 4).unwrap());
+        assert_eq!(got, want, "restore must rewind to the snapshot's boundary");
+
+        // Control: a cleared model cannot even take the step, so the
+        // restore above did something a clear does not.
+        model.clear_kv_cache().unwrap();
+        // `{:#}` to walk the context chain: the forward wraps the
+        // failure in "layer 1", and the guard is what it wraps.
+        let err = format!("{:#}", model.forward(&next, 4).unwrap_err());
+        assert!(
+            err.contains("diverged"),
+            "expected the cache-divergence guard, got: {err}"
+        );
+    }
+
+    /// The snapshot carries PLE's state, and PLE is the piece with two
+    /// distinct kinds of it — a rolling conv context and the carried
+    /// n-gram ids. Restoring one without the other reads the wrong rows
+    /// of the table, so the snapshot has to hold both.
+    #[test]
+    fn the_snapshot_carries_ple_state() {
+        let (_dir, mut model) = build();
+        let prompt = Tensor::from_vec(vec![5u32, 9, 3, 12], (1, 4), &Device::Cpu).unwrap();
+        model.forward(&prompt, 0).unwrap();
+
+        let snap = model.snapshot_kv_cache().unwrap();
+        let ple = snap.ple.as_ref().expect("layer 1 carries PLE");
+        assert!(ple.conv_state.is_some(), "the dilated conv context");
+        assert_eq!(ple.carried_ids.len(), 1, "one batch row");
+        assert_eq!(
+            ple.carried_ids[0].len(),
+            2,
+            "ngram_size - 1 ids carried across the step"
+        );
+        assert!(snap.size_bytes() > 0);
+    }
+
+    /// Round-trip the snapshot itself, not the logits it leads to.
+    ///
+    /// Taken from llama.cpp's `test-save-load-state` for this
+    /// architecture (their #27941, which fixed the indexer cache
+    /// dropping `ext.x`/`ext.y` on restore): comparing *generated
+    /// output* cannot see a field dropped on the way back in, because a
+    /// dropped field need not change the next token. Comparing the
+    /// captured state can. Their fix moved 198 bytes of 335,692 — a
+    /// logits comparison would never have noticed.
+    ///
+    /// So: snapshot, restore, snapshot again, and require the two to
+    /// agree field for field. Anything `restore_kv_cache` silently
+    /// fails to put back shows up here as a shape or value difference.
+    #[test]
+    fn a_snapshot_survives_its_own_round_trip() {
+        let (_dir, mut model) = build();
+        let prompt = Tensor::from_vec(vec![5u32, 9, 3, 12, 7], (1, 5), &Device::Cpu).unwrap();
+        model.forward(&prompt, 0).unwrap();
+
+        let first = model.snapshot_kv_cache().unwrap();
+        model.restore_kv_cache(&first).unwrap();
+        let second = model.snapshot_kv_cache().unwrap();
+
+        assert_eq!(first.layer_count(), second.layer_count());
+        assert_eq!(
+            first.size_bytes(),
+            second.size_bytes(),
+            "total captured bytes"
+        );
+
+        for (i, (a, b)) in first.layers.iter().zip(second.layers.iter()).enumerate() {
+            match (a, b) {
+                (
+                    LayerKvSnapshot::FullSparse {
+                        kv: ka,
+                        indexer_keys: ia,
+                    },
+                    LayerKvSnapshot::FullSparse {
+                        kv: kb,
+                        indexer_keys: ib,
+                    },
+                ) => {
+                    assert_eq!(ka.is_some(), kb.is_some(), "layer {i} kv presence");
+                    if let (Some((k1, v1)), Some((k2, v2))) = (ka, kb) {
+                        assert_eq!(k1.dims(), k2.dims(), "layer {i} k");
+                        assert_eq!(v1.dims(), v2.dims(), "layer {i} v");
+                        assert_eq!(tensor_vec(k1), tensor_vec(k2), "layer {i} k values");
+                    }
+                    // The field their fix was about: dropped silently,
+                    // and invisible to anything that looks at outputs.
+                    assert_eq!(
+                        ia.is_some(),
+                        ib.is_some(),
+                        "layer {i} indexer cache presence"
+                    );
+                    if let (Some(x), Some(y)) = (ia, ib) {
+                        assert_eq!(x.dims(), y.dims(), "layer {i} indexer keys");
+                        assert_eq!(tensor_vec(x), tensor_vec(y), "layer {i} indexer values");
+                    }
+                }
+                (
+                    LayerKvSnapshot::Linear {
+                        conv_state: ca,
+                        recurrent_state: ra,
+                    },
+                    LayerKvSnapshot::Linear {
+                        conv_state: cb,
+                        recurrent_state: rb,
+                    },
+                ) => {
+                    assert_eq!(ca.is_some(), cb.is_some(), "layer {i} conv presence");
+                    assert_eq!(ra.is_some(), rb.is_some(), "layer {i} recurrent presence");
+                    if let (Some(x), Some(y)) = (ca, cb) {
+                        assert_eq!(tensor_vec(x), tensor_vec(y), "layer {i} conv values");
+                    }
+                    if let (Some(x), Some(y)) = (ra, rb) {
+                        assert_eq!(tensor_vec(x), tensor_vec(y), "layer {i} recurrent values");
+                    }
+                }
+                _ => panic!("layer {i}: snapshot variant changed across a round trip"),
+            }
+        }
+
+        let (pa, pb) = (first.ple.as_ref().unwrap(), second.ple.as_ref().unwrap());
+        assert_eq!(pa.carried_ids, pb.carried_ids, "PLE carried ids");
+        assert_eq!(pa.conv_state.is_some(), pb.conv_state.is_some());
+        if let (Some(x), Some(y)) = (&pa.conv_state, &pb.conv_state) {
+            assert_eq!(tensor_vec(x), tensor_vec(y), "PLE conv context");
+        }
+    }
+
+    /// A snapshot from a differently-shaped model must be refused
+    /// rather than restored onto whatever lines up.
+    #[test]
+    fn a_mismatched_snapshot_is_refused() {
+        let (_dir, mut model) = build();
+        let prompt = Tensor::from_vec(vec![5u32, 9], (1, 2), &Device::Cpu).unwrap();
+        model.forward(&prompt, 0).unwrap();
+        let mut snap = model.snapshot_kv_cache().unwrap();
+
+        snap.layers.pop();
+        let err = model.restore_kv_cache(&snap).unwrap_err().to_string();
+        assert!(err.contains("layers"), "got: {err}");
+
+        let mut snap = model.snapshot_kv_cache().unwrap();
+        snap.ple = None;
+        let err = model.restore_kv_cache(&snap).unwrap_err().to_string();
+        assert!(err.contains("PLE"), "got: {err}");
+    }
+
+    /// The draft head, end to end on the fixture: the target's
+    /// pre-mixer stream and the next token's embedding go in, and both
+    /// of the head's outputs come out — the collapse for `lm_head` and
+    /// the multi stream that a second draft step would consume.
+    #[test]
+    fn the_mtp_head_drafts_from_the_targets_pre_mixer_stream() {
+        use crate::harness::arch::qwen4_exp::mtp::MtpHead;
+        let (_dir, mut model) = build();
+        let cfg = Config::from_config_json(TINY).unwrap();
+        let dir2 = tempfile::tempdir().unwrap();
+        let path = dir2.path().join("model.safetensors");
+        candle_core::safetensors::save(&tensors(&cfg.text_config), &path).unwrap();
+        let vb =
+            unsafe { ShardedSafeTensors::var_builder(&[&path], DType::F32, &Device::Cpu).unwrap() };
+        let rotary = Arc::new(rotary_for(&cfg.text_config, DType::F32, &Device::Cpu).unwrap());
+        let mut head = MtpHead::load(&cfg.text_config, rotary.clone(), None, &vb).unwrap();
+
+        let prompt = Tensor::from_vec(vec![5u32, 9, 3, 12], (1, 4), &Device::Cpu).unwrap();
+        let (_logits, hidden) = model.forward_with_hidden(&prompt, 0).unwrap();
+        let wide = cfg.text_config.stream_width();
+        assert_eq!(hidden.dims(), &[1, 4, wide], "the pre-mixer stream");
+
+        // The head owns its own KV and indexer caches, so it has to walk
+        // the prompt alongside the target before it can draft from
+        // position 4 — a draft head that only ever sees the newest token
+        // is decoding against an empty past.
+        let ids: Vec<u32> = prompt.flatten_all().unwrap().to_vec1().unwrap();
+        for (i, id) in ids.iter().enumerate() {
+            let tok = Tensor::from_vec(vec![*id], (1, 1), &Device::Cpu).unwrap();
+            let e = model.embed(&tok).unwrap();
+            let h_i = hidden.narrow(1, i, 1).unwrap().contiguous().unwrap();
+            let (cos, sin) = rotary.plain_cos_sin(i, 1).unwrap();
+            head.forward(&e, &h_i, None, &cos, &sin, i).unwrap();
+        }
+
+        let next = Tensor::from_vec(vec![7u32], (1, 1), &Device::Cpu).unwrap();
+        let e = model.embed(&next).unwrap();
+        let last = hidden.narrow(1, 3, 1).unwrap().contiguous().unwrap();
+        let (cos, sin) = rotary.plain_cos_sin(4, 1).unwrap();
+
+        let step = head.forward(&e, &last, None, &cos, &sin, 4).unwrap();
+        assert_eq!(
+            step.sample_hidden.dims(),
+            &[1, 1, cfg.text_config.hidden_size]
+        );
+        assert_eq!(step.multi_stream.dims(), &[1, 1, wide]);
+        let v: Vec<f32> = step.sample_hidden.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(v.iter().all(|x| x.is_finite()));
+
+        // Chained: the second step reads the head's own multi stream,
+        // not the target's and not the collapse.
+        let (cos, sin) = rotary.plain_cos_sin(5, 1).unwrap();
+        let step2 = head
+            .forward(&e, &step.multi_stream, None, &cos, &sin, 5)
+            .unwrap();
+        assert_eq!(step2.multi_stream.dims(), &[1, 1, wide]);
+    }
+
+    /// Handing the head the collapsed hidden state instead of the
+    /// pre-mixer stream is the mistake the shapes invite — both are
+    /// "the hidden state" and only one is right. It is refused by width
+    /// rather than silently broadcast.
+    #[test]
+    fn the_head_refuses_the_collapsed_hidden_state() {
+        use crate::harness::arch::qwen4_exp::mtp::MtpHead;
+        let cfg = Config::from_config_json(TINY).unwrap();
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&tensors(&cfg.text_config), &path).unwrap();
+        let vb =
+            unsafe { ShardedSafeTensors::var_builder(&[&path], DType::F32, &Device::Cpu).unwrap() };
+        let rotary = Arc::new(rotary_for(&cfg.text_config, DType::F32, &Device::Cpu).unwrap());
+        let mut head = MtpHead::load(&cfg.text_config, rotary.clone(), None, &vb).unwrap();
+
+        let h = cfg.text_config.hidden_size;
+        let collapsed = Tensor::zeros((1, 1, h), DType::F32, &Device::Cpu).unwrap();
+        let e = Tensor::zeros((1, 1, h), DType::F32, &Device::Cpu).unwrap();
+        let (cos, sin) = rotary.plain_cos_sin(0, 1).unwrap();
+        let err = head
+            .forward(&e, &collapsed, None, &cos, &sin, 0)
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains("pre-final-mixer"), "got: {err}");
+    }
 }

crates/neuron/src/harness/arch/qwen4_exp/moe.rs

@@ -22,26 +22,45 @@
 //! of error that produces a model that talks and reasons slightly
 //! worse. It is asserted below rather than commented.
 //!
-//! The fused layout is the better one for a grouped GEMM, and keeping
-//! it would be the faster path; that is an algorithm change, and the
-//! spec is explicit that this stage is not it. Sparsity is extreme —
-//! 10 of 512 with an intermediate of 640 — so the routing, not the
-//! GEMM shape, is what dominates today.
+//! ## Why the fused tensors stay fused
+//!
+//! Slicing them into 512 per-expert modules copies **241.6 GB** at
+//! BF16, 5.03 GB per layer. That fits in neither beast's 64 GB of VRAM
+//! nor its 123 GB of host RAM, and no quantisation rescues it: #309
+//! measured 59.4 GB of 63.7 used at 4.25 bpw with only *four* layers'
+//! experts offloaded. A model that cannot be held cannot be loaded.
+//!
+//! So the block takes `Experts::Banked` and views one expert per
+//! routed token instead. A slice along dim 0 of a contiguous tensor is
+//! contiguous, and `x.matmul(w.t())` is what `Linear` does anyway, so
+//! the views feed the same GEMMs with no copy. It is also the layout a
+//! host-resident bank with a device-side gather wants (#318).
 //!
 //! See `doc/qwen4_exp-port-spec.md` §6.
 
 use anyhow::{Context, Result};
-use candle_core::{IndexOp, Tensor};
+use candle_core::Tensor;
+use candle_core::quantized::GgmlDType;
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
 use crate::harness::arch::qwen3_5::mlp::Qwen3_5MLP;
-use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
+use crate::harness::arch::qwen3_5::moe::{Experts, Qwen3_5MoeBlock};
 
 use super::config::TextConfig;
 
 /// `vb` should be `.pp(...)`-ed to the layer's `mlp` prefix.
-pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Qwen3_5MoeBlock> {
+///
+/// `quant` quantises the routed experts in situ (#315). Without it the
+/// experts stay banked at the checkpoint's precision, which is only
+/// useful for a model small enough to hold — 120.8 B routed parameters
+/// are 241.6 GB at BF16 and fit nowhere on this hardware, so a real
+/// load of `Qwen3.8-Flash-Next` always passes one.
+pub fn load(
+    cfg: &TextConfig,
+    quant: Option<GgmlDType>,
+    vb: &ShardedVarBuilder,
+) -> Result<Qwen3_5MoeBlock> {
     let (h, inter) = (cfg.hidden_size, cfg.moe_intermediate_size);
     anyhow::ensure!(
         cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && inter > 0,
@@ -65,7 +84,19 @@ pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Qwen3_5MoeBlock>
     let down = experts_vb
         .get((cfg.num_experts, h, inter), "down_proj")
         .with_context(|| format!("load '{}/down_proj'", experts_vb.prefix()))?;
-    let experts = split_fused_experts(&gate_up, &down, inter)?;
+    check_fused_experts(&gate_up, &down, inter)?;
+    // The fused pair is the peak: 5.03 GB per layer at BF16. Quantising
+    // here and letting it drop at the end of this function is what keeps
+    // the whole 241.6 GB from ever existing.
+    let experts = match quant {
+        Some(dtype) => Experts::quantize_banked(&gate_up, &down, inter, dtype)
+            .with_context(|| format!("quantize experts to {dtype:?}"))?,
+        None => Experts::Banked {
+            gate_up,
+            down,
+            intermediate: inter,
+        },
+    };
 
     let (shared_expert, shared_expert_gate) = if cfg.shared_expert_intermediate_size > 0 {
         let shared = Qwen3_5MLP::load_with_dims(
@@ -93,14 +124,17 @@ pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Qwen3_5MoeBlock>
     ))
 }
 
-/// Slice `(E, 2I, H)` and `(E, H, I)` into `E` SwiGLU experts.
+/// Validate the fused pair before it is banked.
 ///
-/// Gate first, up second — see the module note.
-pub(crate) fn split_fused_experts(
+/// The block will take views of these per routed expert, so a
+/// disagreement here surfaces as a shape error mid-forward on whichever
+/// token first routes to the offending expert — which is a worse place
+/// to find it than at load.
+pub(crate) fn check_fused_experts(
     gate_up: &Tensor,
     down: &Tensor,
     intermediate: usize,
-) -> Result<Vec<Qwen3_5MLP>> {
+) -> Result<()> {
     let (experts, fused, _) = gate_up.dims3()?;
     anyhow::ensure!(
         fused == intermediate * 2,
@@ -111,78 +145,27 @@ pub(crate) fn split_fused_experts(
         "down_proj holds {} experts, gate_up_proj holds {experts}",
         down.dims3()?.0
     );
-
-    let mut out = Vec::with_capacity(experts);
-    for e in 0..experts {
-        let fused_e = gate_up.i(e)?;
-        out.push(Qwen3_5MLP::from_weights(
-            Linear::new(fused_e.narrow(0, 0, intermediate)?.contiguous()?, None),
-            Linear::new(
-                fused_e
-                    .narrow(0, intermediate, intermediate)?
-                    .contiguous()?,
-                None,
-            ),
-            Linear::new(down.i(e)?.contiguous()?, None),
-        ));
-    }
-    Ok(out)
+    Ok(())
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use candle_core::{DType, Device, Module};
+    use candle_core::{DType, Device};
 
-    /// One expert, two hidden dims, one intermediate. The gate row picks
-    /// the first input channel and the up row picks the second, so
-    /// swapping them swaps which channel is squashed by the SiLU — and
-    /// the two answers differ by more than any tolerance.
+    /// The loader validates the fused pair up front. Without this the
+    /// first token to route to a bad expert fails mid-forward, which is
+    /// a worse place to learn the checkpoint disagrees with the config.
     #[test]
-    fn the_gate_is_the_first_half_of_the_fused_tensor() {
-        let dev = Device::Cpu;
-        // gate_up: [[gate row], [up row]] = [[1, 0], [0, 1]]
-        let gate_up = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (1, 2, 2), &dev).unwrap();
-        // down: (E=1, H=2, I=1)
-        let down = Tensor::from_vec(vec![2.0f32, 3.0], (1, 2, 1), &dev).unwrap();
-
-        let experts = split_fused_experts(&gate_up, &down, 1).unwrap();
-        assert_eq!(experts.len(), 1);
-
-        let x = Tensor::from_vec(vec![1.0f32, 2.0], (1, 2), &dev).unwrap();
-        let got: Vec<f32> = experts[0]
-            .forward(&x)
-            .unwrap()
-            .flatten_all()
-            .unwrap()
-            .to_vec1()
-            .unwrap();
-
-        // gate(x) = 1, up(x) = 2  ->  silu(1) * 2  ->  down
-        let silu1 = 1.0f32 / (1.0 + (-1.0f32).exp());
-        let want = [2.0 * silu1 * 2.0, 3.0 * silu1 * 2.0];
-        for (g, w) in got.iter().zip(want.iter()) {
-            assert!((g - w).abs() < 1e-5, "got {got:?} want {want:?}");
-        }
-
-        // The swapped reading: silu(2) * 1, which is a different model.
-        let silu2 = 2.0f32 / (1.0 + (-2.0f32).exp());
-        assert!(
-            (got[0] - 2.0 * silu2).abs() > 1e-3,
-            "this is silu(up) * gate, not silu(gate) * up: {got:?}"
-        );
-    }
-
-    #[test]
-    fn a_fused_tensor_of_the_wrong_width_is_rejected() {
+    fn a_fused_pair_that_disagrees_with_the_config_is_rejected() {
         let dev = Device::Cpu;
         let gate_up = Tensor::zeros((2, 6, 4), DType::F32, &dev).unwrap();
         let down = Tensor::zeros((2, 4, 3), DType::F32, &dev).unwrap();
-        assert!(split_fused_experts(&gate_up, &down, 3).is_ok());
+        assert!(check_fused_experts(&gate_up, &down, 3).is_ok());
         // 2 x 4 != 6
-        assert!(split_fused_experts(&gate_up, &down, 4).is_err());
+        assert!(check_fused_experts(&gate_up, &down, 4).is_err());
         // expert counts must agree
         let down_short = Tensor::zeros((1, 4, 3), DType::F32, &dev).unwrap();
-        assert!(split_fused_experts(&gate_up, &down_short, 3).is_err());
+        assert!(check_fused_experts(&gate_up, &down_short, 3).is_err());
     }
 }

crates/neuron/src/harness/arch/qwen4_exp/mtp.rs

@@ -0,0 +1,263 @@
+//! The multi-token-prediction head — `qwen4_exp`'s shipped draft model.
+//!
+//! The checkpoint carries a complete second decoder layer under `mtp.*`:
+//! full attention with its own QSA indexer, its own 512-expert MoE, both
+//! hyper-connections, and a final mixer. It reuses `embed_tokens` and
+//! `lm_head` (`mtp_use_dedicated_embeddings: false`), so what is new
+//! here is only how the target model's state and the next token's
+//! embedding are fused into an input for it.
+//!
+//! ```text
+//! e = fc_embedding(pre_fc_norm_embedding(embed(t)))   [T, 2560]
+//! h = pre_fc_norm_hidden(h.flatten(-2))               ONE norm over 10240
+//! h = fc_hidden(h.view(T, hc, 2560))                  per stream, shared matrix
+//! h = e.unsqueeze(-2) + h                             embedding into every stream
+//! h = h.flatten(-2)                                   [T, 10240] -> the layer
+//! ```
+//!
+//! vLLM calls this `residual_linear_shared`, as against the
+//! `Linear(2H, H)` + repeat other MTP variants use. Three parts of it
+//! cannot be recovered from the tensor shapes, and all three were read
+//! from `vllm/models/qwen4_exp/nvidia/mtp.py` rather than inferred:
+//!
+//! 1. **`pre_fc_norm_hidden` is ungrouped** — one normalisation over the
+//!    whole 10240, unlike every other 10240-wide norm in this
+//!    architecture. vLLM has a `GroupedGemmaRMSNorm` and deliberately
+//!    does not use it here. Grouping it rescales the draft input in four
+//!    independent pieces, which costs acceptance rate and reports
+//!    nothing.
+//! 2. **The hidden state is the *pre-final-mixer* multi-stream** — `h`
+//!    after the target's last decoder layer, before its mixer collapses
+//!    it. Not the 2560 the LM head sees.
+//! 3. **The head emits both** — a collapsed stream for `lm_head` and the
+//!    pre-final-mixer stream for its own next step. A head that returned
+//!    only the first cannot be chained past one token.
+//!
+//! PLE is off in this layer while the stream count is kept, which the
+//! absence of `mtp.layers.0.ple.*` in the checkpoint agrees with.
+//!
+//! **The head keeps its own KV and indexer caches, and they are not the
+//! target's.** It therefore has to walk the prompt alongside the target
+//! before it can draft: a head fed only the newest token is attending
+//! over an empty past, and this architecture says so rather than
+//! guessing, because the indexer's divergence check fires. That makes
+//! prefill part of the speculative loop's cost, not just decode.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §9 and #313.
+
+use anyhow::{Context, Result};
+use candle_core::quantized::GgmlDType;
+use candle_core::{D, Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+use std::sync::Arc;
+
+use crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm;
+use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
+
+use super::config::TextConfig;
+use super::decoder::DecoderLayer;
+use super::hyper::HyperConnection;
+
+/// One draft step's two outputs.
+#[derive(Debug)]
+pub struct MtpStep {
+    /// `[B, L, hidden]` — the mixer's collapse, for `lm_head`.
+    pub sample_hidden: Tensor,
+    /// `[B, L, hidden * hc_count]` — pre-final-mixer, and the input to
+    /// the *next* draft step. Chaining reads this, not the collapse.
+    pub multi_stream: Tensor,
+}
+
+pub struct MtpHead {
+    pre_fc_norm_embedding: Qwen3_5RmsNorm,
+    /// Ungrouped over `hidden * hc_count` — see the module note.
+    pre_fc_norm_hidden: Qwen3_5RmsNorm,
+    fc_embedding: Linear,
+    /// Shared across the streams: one `hidden -> hidden` matrix applied
+    /// to each branch, not a `wide -> wide` one.
+    fc_hidden: Linear,
+    layer: DecoderLayer,
+    mixer: HyperConnection,
+    hidden_size: usize,
+    hc_count: usize,
+}
+
+impl MtpHead {
+    /// `vb` is the checkpoint root; the head hangs off `mtp.`.
+    pub fn load(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        quant: Option<GgmlDType>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let mtp = vb.pp("mtp");
+        let (h, hc) = (cfg.hidden_size, cfg.hc_count);
+        Ok(Self {
+            pre_fc_norm_embedding: Qwen3_5RmsNorm::load(
+                &mtp.pp("pre_fc_norm_embedding"),
+                h,
+                cfg.rms_norm_eps,
+            )?,
+            // `load`, not `load_grouped`: the one 10240-wide norm in
+            // this architecture that is not per-stream.
+            pre_fc_norm_hidden: Qwen3_5RmsNorm::load(
+                &mtp.pp("pre_fc_norm_hidden"),
+                h * hc,
+                cfg.rms_norm_eps,
+            )?,
+            fc_embedding: linear(&mtp, "fc_embedding", h, h)?,
+            fc_hidden: linear(&mtp, "fc_hidden", h, h)?,
+            layer: DecoderLayer::load_typed(
+                cfg,
+                rotary,
+                "full_attention",
+                false,
+                quant,
+                &mtp.pp("layers").pp(0),
+            )
+            .context("load mtp.layers.0")?,
+            mixer: HyperConnection::load(
+                &mtp.pp("hyper_connection_mixer"),
+                h,
+                hc,
+                cfg.hc_lowrank,
+                cfg.rms_norm_eps,
+                false,
+            )
+            .context("load mtp.hyper_connection_mixer")?,
+            hidden_size: h,
+            hc_count: hc,
+        })
+    }
+
+    /// One draft step.
+    ///
+    /// `token_embed` is `embed_tokens(t)` for the token being predicted
+    /// *from*, `[B, L, hidden]`. `hidden` is the pre-final-mixer stream,
+    /// `[B, L, hidden * hc_count]` — from the target model on the first
+    /// step, from this head's own previous [`MtpStep::multi_stream`]
+    /// afterwards.
+    pub fn forward(
+        &mut self,
+        token_embed: &Tensor,
+        hidden: &Tensor,
+        causal_mask: Option<&Tensor>,
+        cos: &Tensor,
+        sin: &Tensor,
+        past_len: usize,
+    ) -> Result<MtpStep> {
+        let (b, l, wide) = hidden.dims3()?;
+        anyhow::ensure!(
+            wide == self.hidden_size * self.hc_count,
+            "mtp: expected the pre-final-mixer stream ({} wide), got {wide} — \
+             passing the collapsed hidden state here is the likely mistake",
+            self.hidden_size * self.hc_count
+        );
+
+        let e = self
+            .fc_embedding
+            .forward(&self.pre_fc_norm_embedding.forward(token_embed)?)?;
+
+        // Normalise the flattened stream, *then* split it: one
+        // normalisation over the whole width.
+        let h = self.pre_fc_norm_hidden.forward(hidden)?.reshape((
+            b,
+            l,
+            self.hc_count,
+            self.hidden_size,
+        ))?;
+        let h = self.fc_hidden.forward(&h)?;
+
+        // The embedding enters every stream.
+        let h = h
+            .broadcast_add(&e.unsqueeze(D::Minus2)?)?
+            .flatten_from(D::Minus2)?;
+
+        let h = self
+            .layer
+            .forward(&h, None, causal_mask, cos, sin, past_len)
+            .context("mtp draft layer")?;
+
+        Ok(MtpStep {
+            sample_hidden: self.mixer.collapse(&h)?,
+            multi_stream: h,
+        })
+    }
+
+    /// Drop the draft layer's caches. The head has its own KV and its
+    /// own indexer cache, independent of the target model's.
+    pub fn clear_kv_cache(&mut self) -> Result<()> {
+        self.layer.clear_kv_cache()
+    }
+}
+
+fn linear(vb: &ShardedVarBuilder, name: &str, in_dim: usize, out_dim: usize) -> Result<Linear> {
+    let weight = vb
+        .pp(name)
+        .get((out_dim, in_dim), "weight")
+        .with_context(|| format!("load '{}/{name}/weight'", vb.prefix()))?;
+    Ok(Linear::new(weight, None))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+
+    /// `pre_fc_norm_hidden` normalises the whole 10240 at once, where
+    /// every other norm this wide in the architecture normalises four
+    /// streams of 2560 independently. The port spec asserted the
+    /// opposite until vLLM settled it, and the difference is invisible
+    /// downstream — it rescales the draft input and costs acceptance
+    /// rate.
+    ///
+    /// So it is pinned by value, on an input the two readings cannot
+    /// agree on: two streams whose magnitudes differ threefold. Grouped,
+    /// each is normalised to its own scale and both come out all-ones.
+    /// Ungrouped, the larger stream stays larger.
+    #[test]
+    fn the_hidden_pre_norm_is_ungrouped() {
+        let dev = Device::Cpu;
+        let (hidden, hc) = (4usize, 2usize);
+        let wide = hidden * hc;
+        // Zero weight: (1 + w) = 1, so this is a pure RMS normalisation.
+        let norm =
+            Qwen3_5RmsNorm::from_weight(Tensor::zeros(wide, DType::F32, &dev).unwrap(), 1e-6, None);
+        let h = Tensor::from_vec(
+            vec![1.0f32, 1.0, 1.0, 1.0, 3.0, 3.0, 3.0, 3.0],
+            (1, 1, wide),
+            &dev,
+        )
+        .unwrap();
+
+        let got: Vec<f32> = norm
+            .forward(&h)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // rms over all eight = sqrt((4*1 + 4*9)/8) = sqrt(5)
+        let rms = 5.0f32.sqrt();
+        let want = [
+            1.0 / rms,
+            1.0 / rms,
+            1.0 / rms,
+            1.0 / rms,
+            3.0 / rms,
+            3.0 / rms,
+            3.0 / rms,
+            3.0 / rms,
+        ];
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-5, "got {got:?} want {want:?}");
+        }
+        // The grouped reading would flatten both streams to all-ones.
+        assert!(
+            (got[4] - 1.0).abs() > 1e-3,
+            "this is the grouped reading, which the spec had and vLLM does not: {got:?}"
+        );
+    }
+}

crates/neuron/src/harness/arch/qwen4_exp/ple.rs

@@ -491,6 +491,26 @@ impl PleBlock {
         (self.kernel_size - 1) * self.dilation
     }
 
+    /// Capture the rolling conv context for a prefix snapshot.
+    ///
+    /// Deep-copied rather than shared. The buffer is replaced by
+    /// assignment today, so a shallow clone would be sound — but
+    /// `narrow(..).contiguous()` on an already-contiguous slice is a
+    /// view, so whether the snapshot shares storage depends on a
+    /// property of the *data*, not of the code. That is exactly the
+    /// kind of conditional aliasing the GatedDeltaNet states are
+    /// deep-copied to avoid, and nine positions of a 10240-wide signal
+    /// is not worth being clever about.
+    pub fn snapshot_state(&self) -> candle_core::Result<Option<Tensor>> {
+        self.conv_state.as_ref().map(Tensor::copy).transpose()
+    }
+
+    /// Replace the rolling conv context from a snapshot.
+    pub fn restore_state(&mut self, conv_state: Option<&Tensor>) -> candle_core::Result<()> {
+        self.conv_state = conv_state.map(Tensor::copy).transpose()?;
+        Ok(())
+    }
+
     /// Drop the rolling conv context. Call alongside a KV-cache clear;
     /// the carried n-gram ids ([`NGramHasher::context_len`]) reset with
     /// it.

crates/neuron/src/harness/arch/qwen4_exp/qsa.rs

@@ -293,6 +293,21 @@ impl Indexer {
         &self.selector
     }
 
+    /// Capture the raw key cache for a prefix snapshot.
+    ///
+    /// A shallow clone, like the attention K/V it travels with: the
+    /// cache is only ever *replaced* by a `cat` into a fresh
+    /// allocation, never written in place, so the snapshot stays valid
+    /// after the live cache moves on.
+    pub fn snapshot_keys(&self) -> Option<Tensor> {
+        self.key_cache.clone()
+    }
+
+    /// Replace the key cache from a snapshot.
+    pub fn restore_keys(&mut self, keys: Option<&Tensor>) {
+        self.key_cache = keys.cloned();
+    }
+
     /// Drop the indexer's cache. Must happen with the main KV clear —
     /// a stale indexer cache selects blocks of another request's text.
     pub fn clear_cache(&mut self) {

crates/neuron/src/harness/arch/snapshot.rs

@@ -1,5 +1,12 @@
 //! Cache-state snapshots for prefix KV caching (#11).
 //!
+//! Shared across architectures rather than owned by one: `qwen3_5`,
+//! its TP mirror `tp_qwen3_5`, and `qwen4_exp` all capture state
+//! through these types, and `ModelArch::snapshot_kv_cache` hands one
+//! back whatever the architecture. It lived under `arch/qwen3_5/`
+//! while that was the only arch with a snapshot, which stopped being
+//! true.
+//!
 //! A snapshot captures everything `clear_kv_cache` would destroy, at
 //! one consistent token boundary:
 //!
@@ -40,6 +47,17 @@ pub enum LayerKvSnapshot {
         conv_state: Option<Tensor>,
         recurrent_state: Option<Tensor>,
     },
+    /// A full-attention layer that also keeps a QSA indexer cache
+    /// (`qwen4_exp`). The indexer's raw keys are a second, separately
+    /// shaped cache — one head of 128 dims per position — and a
+    /// snapshot that captured only `kv` would restore a layer whose
+    /// indexer still holds the *previous* request's keys, selecting
+    /// blocks of somebody else's text. It fails fluently, so the two
+    /// travel together or not at all.
+    FullSparse {
+        kv: Option<(Tensor, Tensor)>,
+        indexer_keys: Option<Tensor>,
+    },
 }
 
 /// One consistent cache snapshot of a `Qwen3_5Model` (or its TP
@@ -50,6 +68,9 @@ pub enum LayerKvSnapshot {
 pub struct KvCacheSnapshot {
     pub(crate) layers: Vec<LayerKvSnapshot>,
     pub(crate) rope_delta: i64,
+    /// `qwen4_exp`'s PLE state. `None` for architectures without a
+    /// hashed n-gram layer, which is every other one today.
+    pub(crate) ple: Option<PleSnapshot>,
 }
 
 impl KvCacheSnapshot {
@@ -79,8 +100,44 @@ impl KvCacheSnapshot {
                     conv_state.as_ref().map(t_bytes).unwrap_or(0)
                         + recurrent_state.as_ref().map(t_bytes).unwrap_or(0)
                 }
+                LayerKvSnapshot::FullSparse { kv, indexer_keys } => {
+                    kv.as_ref()
+                        .map(|(k, v)| t_bytes(k) + t_bytes(v))
+                        .unwrap_or(0)
+                        + indexer_keys.as_ref().map(t_bytes).unwrap_or(0)
+                }
             })
-            .sum()
+            .sum::<u64>()
+            + self.ple.as_ref().map(PleSnapshot::size_bytes).unwrap_or(0)
+    }
+}
+
+/// PLE's per-request state (`qwen4_exp`), captured at the same token
+/// boundary as the layers.
+///
+/// Two pieces with different natures. `conv_state` is the dilated
+/// short conv's left context — nine positions of a 10240-wide signal,
+/// deep-copied like the GatedDeltaNet states because the conv writes
+/// its rolling buffer in place. `carried_ids` is the previous
+/// `ngram_size - 1` **token ids** per batch row, which upstream keeps
+/// as a third "conv state" on the same layer; they are ids, not
+/// activations, and cost two integers per sequence.
+pub struct PleSnapshot {
+    pub(crate) conv_state: Option<Tensor>,
+    pub(crate) carried_ids: Vec<Vec<i64>>,
+}
+
+impl PleSnapshot {
+    fn size_bytes(&self) -> u64 {
+        self.conv_state
+            .as_ref()
+            .map(|t| (t.elem_count() * t.dtype().size_in_bytes()) as u64)
+            .unwrap_or(0)
+            + self
+                .carried_ids
+                .iter()
+                .map(|r| (r.len() * std::mem::size_of::<i64>()) as u64)
+                .sum::<u64>()
     }
 }
 
@@ -145,6 +202,7 @@ pub fn assemble_batch(seqs: &[(&KvCacheSnapshot, usize)]) -> candle_core::Result
     }
     Ok(BatchedKvState {
         snapshot: KvCacheSnapshot {
+            ple: None,
             layers,
             rope_delta: 0,
         },
@@ -159,6 +217,14 @@ fn assemble_layer(
     padded_len: usize,
 ) -> candle_core::Result<LayerKvSnapshot> {
     match &seqs[0].0.layers[li] {
+        // Batched decode for the sparse-attention layers is not built
+        // yet. Assembling them as if they were plain full-attention
+        // would drop the indexer cache and hand every row a selection
+        // computed from another request's keys, which produces text.
+        LayerKvSnapshot::FullSparse { .. } => candle_core::bail!(
+            "assemble_batch: layer {li} carries a QSA indexer cache; batched decode \
+             for this architecture is not implemented"
+        ),
         LayerKvSnapshot::Full(_) => {
             let mut ks = Vec::with_capacity(seqs.len());
             let mut vs = Vec::with_capacity(seqs.len());
@@ -242,6 +308,10 @@ pub fn extract_row(
     let mut layers = Vec::with_capacity(snap.layers.len());
     for (li, layer) in snap.layers.iter().enumerate() {
         layers.push(match layer {
+            LayerKvSnapshot::FullSparse { .. } => candle_core::bail!(
+                "extract_row: layer {li} carries a QSA indexer cache; batched decode \
+                 for this architecture is not implemented"
+            ),
             LayerKvSnapshot::Full(Some((k, v))) => {
                 let (b, _h, s, _d) = k.dims4()?;
                 if row >= b {
@@ -274,6 +344,7 @@ pub fn extract_row(
         });
     }
     Ok(KvCacheSnapshot {
+        ple: None,
         layers,
         rope_delta: 0,
     })
@@ -312,7 +383,7 @@ fn pad_seq(t: &Tensor, padded_len: usize) -> candle_core::Result<Tensor> {
 
 #[cfg(test)]
 mod tests {
-    use super::super::{Qwen3_5Model, RopeParameters, TextConfig};
+    use crate::harness::arch::qwen3_5::{Qwen3_5Model, RopeParameters, TextConfig};
     use candle_core::{DType, Device, Tensor};
     use std::collections::HashMap;
 

crates/neuron/src/harness/candle.rs

@@ -555,7 +555,7 @@ impl LoadedHandle {
 #[derive(Clone)]
 pub enum KvSnapshotRef {
     Worker(super::device_worker::KvSnapshotId),
-    Local(Arc<super::arch::qwen3_5::snapshot::KvCacheSnapshot>),
+    Local(Arc<super::arch::snapshot::KvCacheSnapshot>),
     #[cfg(feature = "cuda")]
     Tp(u64),
 }
@@ -1116,15 +1116,16 @@ impl ModelArch {
     /// candle-transformers archs keep theirs private, so they stay on
     /// the clear-every-request path.
     pub fn supports_kv_snapshot(&self) -> bool {
-        matches!(self, ModelArch::Qwen3_5Dense(_))
+        matches!(self, ModelArch::Qwen3_5Dense(_) | ModelArch::Qwen4Exp(_))
     }
 
     /// Capture the live cache state as a prefix snapshot. See
     /// `arch/qwen3_5/snapshot.rs` for what a snapshot contains and the
     /// copy-semantics constraints.
-    pub fn snapshot_kv_cache(&self) -> Result<super::arch::qwen3_5::snapshot::KvCacheSnapshot> {
+    pub fn snapshot_kv_cache(&self) -> Result<super::arch::snapshot::KvCacheSnapshot> {
         match self {
             ModelArch::Qwen3_5Dense(m) => Ok(m.snapshot_kv_cache()?),
+            ModelArch::Qwen4Exp(m) => Ok(m.snapshot_kv_cache()?),
             _ => anyhow::bail!("snapshot_kv_cache: architecture has no snapshot support"),
         }
     }
@@ -1133,10 +1134,11 @@ impl ModelArch {
     /// restore-instead-of-clear half of prefix caching.
     pub fn restore_kv_cache(
         &mut self,
-        snap: &super::arch::qwen3_5::snapshot::KvCacheSnapshot,
+        snap: &super::arch::snapshot::KvCacheSnapshot,
     ) -> Result<()> {
         match self {
             ModelArch::Qwen3_5Dense(m) => Ok(m.restore_kv_cache(snap)?),
+            ModelArch::Qwen4Exp(m) => Ok(m.restore_kv_cache(snap)?),
             _ => anyhow::bail!("restore_kv_cache: architecture has no snapshot support"),
         }
     }
@@ -2733,6 +2735,9 @@ impl CandleHarness {
             self.resolve_dense_files(spec, source_id).await?;
         let device_for_load = device.clone();
         let model_id_for_log = spec.model_id.clone();
+        // Resolved out here: `spec` is borrowed and does not outlive the
+        // blocking task.
+        let isq = super::quant::parse_quant_string(spec.quant.as_deref())?;
 
         let arch = tokio::task::spawn_blocking(move || -> Result<ModelArch> {
             let cfg_text = std::fs::read_to_string(&config_path).context("read config.json")?;
@@ -2843,9 +2848,19 @@ impl CandleHarness {
                         )
                         .context("build ShardedVarBuilder for qwen4_exp")?
                     };
+                    // ISQ is not optional for this architecture in
+                    // practice: without it the routed experts are
+                    // 241.6 GB. `quant` comes from the catalogue, and a
+                    // load without one will fail on memory rather than
+                    // on anything this code can say (#315).
+                    // ISQ is not really optional here: without it the
+                    // routed experts are 241.6 GB and the load fails on
+                    // memory rather than on anything this code can say.
+                    // The catalogue's `quant` is the knob (#315, #252).
                     let model = super::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
                         &cfg,
                         dtype,
+                        isq,
                         &device_for_load,
                         &sharded_vb,
                     )

crates/neuron/src/harness/device_worker/dispatch.rs

@@ -13,7 +13,7 @@
 //! ARCH model state in this state slab will gain a companion
 //! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
 
-use crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot;
+use crate::harness::arch::snapshot::KvCacheSnapshot;
 use crate::harness::candle::ModelArch;
 #[cfg(feature = "cuda")]
 use crate::harness::device_worker::jobs::TpHandle;
@@ -1045,7 +1045,7 @@ fn tp_load_shard_inner(
         "qwen3_5" | "qwen3_next" => {
             let cfg = crate::harness::tp::tp_qwen3_5::Config::from_config_json(config_json)
                 .context("parse Qwen3-Next Config JSON for leader load")?;
-            let quant_dtype = crate::harness::tp::worker::parse_quant_string(quant)?;
+            let quant_dtype = crate::harness::quant::parse_quant_string(quant)?;
             TpLeaderModel::Qwen3_5(crate::harness::tp::tp_qwen3_5::TpQwen3_5ForCausalLM::load(
                 cfg,
                 &vb,
@@ -1179,7 +1179,7 @@ fn tp_assemble_kv_batch(
         })?;
         pairs.push((snap, *len));
     }
-    let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
+    let batch = crate::harness::arch::snapshot::assemble_batch(&pairs)?;
     let model = tp_models
         .get_mut(&handle)
         .ok_or_else(|| anyhow::anyhow!("TpAssembleKvBatch: no model for handle {}", handle.0))?;
@@ -1212,9 +1212,8 @@ fn tp_extract_kv_rows(
         .snapshot_kv_cache()?;
     let mut total = 0u64;
     for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
-        let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
-            &live, row, prefix_len, padded_len, steps,
-        )?;
+        let snap =
+            crate::harness::arch::snapshot::extract_row(&live, row, prefix_len, padded_len, steps)?;
         total += snap.size_bytes();
         state.tp_kv_snapshots.insert((handle, id), snap);
     }
@@ -1337,7 +1336,7 @@ fn assemble_kv_batch(
         })?;
         pairs.push((snap, *len));
     }
-    let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
+    let batch = crate::harness::arch::snapshot::assemble_batch(&pairs)?;
     let arch = models
         .get_mut(&handle)
         .ok_or_else(|| anyhow::anyhow!("AssembleKvBatch: no model for handle {}", handle.0))?;
@@ -1363,9 +1362,8 @@ fn extract_kv_rows(
         .snapshot_kv_cache()?;
     let mut out = Vec::with_capacity(rows.len());
     for &(row, prefix_len) in rows {
-        let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
-            &live, row, prefix_len, padded_len, steps,
-        )?;
+        let snap =
+            crate::harness::arch::snapshot::extract_row(&live, row, prefix_len, padded_len, steps)?;
         let id = KvSnapshotId(state.next_kv_snapshot_id);
         state.next_kv_snapshot_id = state.next_kv_snapshot_id.wrapping_add(1);
         let bytes = snap.size_bytes();

crates/neuron/src/harness/device_worker/jobs.rs

@@ -163,7 +163,7 @@ pub enum Job {
     /// state and install it as the model's live state (#98). `seqs`
     /// pairs each snapshot id with its true token length; attention
     /// K/V is right-padded to the batch max and `cat`ed on dim 0 (see
-    /// `arch::qwen3_5::snapshot::assemble_batch`). Replies with the
+    /// `arch::snapshot::assemble_batch`). Replies with the
     /// padded uniform KV length. The source snapshots remain stored —
     /// the caller drops them via `DropKvSnapshot` when the sequences
     /// leave the batch.

crates/neuron/src/harness/mod.rs

@@ -12,6 +12,7 @@ pub mod image;
 pub mod prefix_cache;
 pub mod preflight;
 pub mod preprocess;
+pub(crate) mod quant;
 pub mod reasoning_budget;
 pub mod reasoning_effort;
 pub mod sampling;

crates/neuron/src/harness/quant.rs

@@ -0,0 +1,49 @@
+//! Catalogue quant names to ggml dtypes.
+//!
+//! Shared rather than owned by the TP path: `quant = "q4k"` in
+//! `models.toml` means the same thing whichever loader reads it. It sat
+//! behind `cfg(feature = "cuda")` while only the TP worker used it,
+//! which put it out of reach of the single-GPU path — where, for
+//! `qwen4_exp`, in-situ quantisation is the difference between a model
+//! that loads and 241.6 GB that does not (#315).
+
+/// Parse a `ModelSpec.quant` string into a `GgmlDType`. Accepts the
+/// common ggml format names (case-insensitive). `None` and `Some("")`
+/// both map to "no quantization".
+///
+/// Supported: `q4_0`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q8_1`,
+/// `q2k`/`q2_k`, `q3k`/`q3_k`, `q4k`/`q4_k`, `q5k`/`q5_k`,
+/// `q6k`/`q6_k`, `q8k`/`q8_k`, `f16`, `bf16`, `f32`. The underscore
+/// is optional and the prefix is case-insensitive.
+pub(crate) fn parse_quant_string(
+    s: Option<&str>,
+) -> anyhow::Result<Option<candle_core::quantized::GgmlDType>> {
+    use candle_core::quantized::GgmlDType;
+    let s = match s {
+        Some(s) if !s.is_empty() => s,
+        _ => return Ok(None),
+    };
+    let normalised = s.to_ascii_lowercase().replace('_', "");
+    let dtype = match normalised.as_str() {
+        "q40" => GgmlDType::Q4_0,
+        "q41" => GgmlDType::Q4_1,
+        "q50" => GgmlDType::Q5_0,
+        "q51" => GgmlDType::Q5_1,
+        "q80" => GgmlDType::Q8_0,
+        "q81" => GgmlDType::Q8_1,
+        "q2k" => GgmlDType::Q2K,
+        "q3k" => GgmlDType::Q3K,
+        "q4k" | "q4km" => GgmlDType::Q4K,
+        "q5k" | "q5km" => GgmlDType::Q5K,
+        "q6k" => GgmlDType::Q6K,
+        "q8k" => GgmlDType::Q8K,
+        "f16" => GgmlDType::F16,
+        "bf16" => GgmlDType::BF16,
+        "f32" => GgmlDType::F32,
+        other => anyhow::bail!(
+            "unknown quant '{other}' (expected one of: q4_0, q4_1, q5_0, q5_1, q8_0, \
+             q8_1, q2k, q3k, q4k, q5k, q6k, q8k, f16, bf16, f32)"
+        ),
+    };
+    Ok(Some(dtype))
+}

crates/neuron/src/harness/tp/mod.rs

@@ -138,7 +138,7 @@ impl TpLeaderModel {
     /// Capture rank 0's cache state for a prefix snapshot (#11).
     pub fn snapshot_kv_cache(
         &self,
-    ) -> candle_core::Result<crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot> {
+    ) -> candle_core::Result<crate::harness::arch::snapshot::KvCacheSnapshot> {
         match self {
             TpLeaderModel::Qwen3(_) => {
                 candle_core::bail!("snapshot_kv_cache: qwen3 (dense) has no snapshot support")
@@ -150,7 +150,7 @@ impl TpLeaderModel {
     /// Replace rank 0's live cache state with a stored snapshot.
     pub fn restore_kv_cache(
         &mut self,
-        snap: &crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot,
+        snap: &crate::harness::arch::snapshot::KvCacheSnapshot,
     ) -> candle_core::Result<()> {
         match self {
             TpLeaderModel::Qwen3(_) => {

crates/neuron/src/harness/tp/tp_qwen3_5.rs

@@ -45,10 +45,10 @@ use super::tp_linear::{ColumnParallelLinear, RowParallelLinear};
 use crate::harness::arch::qwen3_5::linear_attn::repeat_interleave;
 use crate::harness::arch::qwen3_5::rmsnorm::{Qwen3_5RmsNorm, Qwen3_5RmsNormGated, l2norm};
 use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
-use crate::harness::arch::qwen3_5::snapshot::{KvCacheSnapshot, LayerKvSnapshot};
 use crate::harness::arch::qwen3_5::splice_runs;
 use crate::harness::arch::qwen3_5::vision::VisionTower;
 pub use crate::harness::arch::qwen3_5::{Config, TextConfig};
+use crate::harness::arch::snapshot::{KvCacheSnapshot, LayerKvSnapshot};
 
 // ─── linear-attention (Gated DeltaNet) ──────────────────────────────
 
@@ -1700,6 +1700,7 @@ impl TpQwen3_5Model {
         Ok(KvCacheSnapshot {
             layers,
             rope_delta: self.rope_delta,
+            ple: None,
         })
     }
 

crates/neuron/src/harness/tp/worker.rs

@@ -16,6 +16,11 @@ use anyhow::Result;
 use std::collections::HashMap;
 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
 
+// Only the CUDA shard-load path calls this; the CPU build of this
+// module has no caller.
+#[cfg(feature = "cuda")]
+use crate::harness::quant::parse_quant_string;
+
 use super::nccl_state::NcclState;
 use super::rpc::{WorkerRequest, WorkerResponse};
 
@@ -119,7 +124,7 @@ impl WorkerModel {
     /// `TpLeaderModel` support first.
     fn snapshot_kv_cache(
         &self,
-    ) -> candle_core::Result<crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot> {
+    ) -> candle_core::Result<crate::harness::arch::snapshot::KvCacheSnapshot> {
         match self {
             WorkerModel::Qwen3(_) => {
                 candle_core::bail!("snapshot_kv_cache: qwen3 (dense) has no snapshot support")
@@ -130,7 +135,7 @@ impl WorkerModel {
 
     fn restore_kv_cache(
         &mut self,
-        snap: &crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot,
+        snap: &crate::harness::arch::snapshot::KvCacheSnapshot,
     ) -> candle_core::Result<()> {
         match self {
             WorkerModel::Qwen3(_) => {
@@ -227,7 +232,7 @@ struct WorkerState {
     /// pool and shared across all ranks. Dropped with the shard on
     /// `UnloadModel`.
     #[cfg(feature = "cuda")]
-    kv_snapshots: HashMap<(String, u64), crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot>,
+    kv_snapshots: HashMap<(String, u64), crate::harness::arch::snapshot::KvCacheSnapshot>,
     /// Placeholder mirroring `models` on the non-cuda build.
     #[cfg(not(feature = "cuda"))]
     #[allow(dead_code)]
@@ -689,11 +694,10 @@ impl WorkerState {
             };
             pairs.push((snap, *len));
         }
-        let result =
-            crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs).and_then(|batch| {
-                model.restore_kv_cache(&batch.snapshot)?;
-                Ok(batch.padded_len)
-            });
+        let result = crate::harness::arch::snapshot::assemble_batch(&pairs).and_then(|batch| {
+            model.restore_kv_cache(&batch.snapshot)?;
+            Ok(batch.padded_len)
+        });
         match result {
             Ok(padded_len) => WorkerResponse::KvBatchAssembled {
                 padded_len: padded_len as u64,
@@ -756,7 +760,7 @@ impl WorkerState {
         };
         let mut total_bytes = 0u64;
         for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
-            match crate::harness::arch::qwen3_5::snapshot::extract_row(
+            match crate::harness::arch::snapshot::extract_row(
                 &live, row, prefix_len, padded_len, steps,
             ) {
                 Ok(snap) => {
@@ -1036,45 +1040,3 @@ impl WorkerState {
         }
     }
 }
-
-/// Parse a `ModelSpec.quant` string into a `GgmlDType`. Accepts the
-/// common ggml format names (case-insensitive). `None` and `Some("")`
-/// both map to "no quantization".
-///
-/// Supported: `q4_0`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q8_1`,
-/// `q2k`/`q2_k`, `q3k`/`q3_k`, `q4k`/`q4_k`, `q5k`/`q5_k`,
-/// `q6k`/`q6_k`, `q8k`/`q8_k`, `f16`, `bf16`, `f32`. The underscore
-/// is optional and the prefix is case-insensitive.
-#[cfg(feature = "cuda")]
-pub(crate) fn parse_quant_string(
-    s: Option<&str>,
-) -> anyhow::Result<Option<candle_core::quantized::GgmlDType>> {
-    use candle_core::quantized::GgmlDType;
-    let s = match s {
-        Some(s) if !s.is_empty() => s,
-        _ => return Ok(None),
-    };
-    let normalised = s.to_ascii_lowercase().replace('_', "");
-    let dtype = match normalised.as_str() {
-        "q40" => GgmlDType::Q4_0,
-        "q41" => GgmlDType::Q4_1,
-        "q50" => GgmlDType::Q5_0,
-        "q51" => GgmlDType::Q5_1,
-        "q80" => GgmlDType::Q8_0,
-        "q81" => GgmlDType::Q8_1,
-        "q2k" => GgmlDType::Q2K,
-        "q3k" => GgmlDType::Q3K,
-        "q4k" | "q4km" => GgmlDType::Q4K,
-        "q5k" | "q5km" => GgmlDType::Q5K,
-        "q6k" => GgmlDType::Q6K,
-        "q8k" => GgmlDType::Q8K,
-        "f16" => GgmlDType::F16,
-        "bf16" => GgmlDType::BF16,
-        "f32" => GgmlDType::F32,
-        other => anyhow::bail!(
-            "unknown quant '{other}' (expected one of: q4_0, q4_1, q5_0, q5_1, q8_0, \
-             q8_1, q2k, q3k, q4k, q5k, q6k, q8k, f16, bf16, f32)"
-        ),
-    };
-    Ok(Some(dtype))
-}

doc/qwen4_exp-port-spec.md

@@ -470,11 +470,21 @@ conventions — this is a solved problem in our tree.
 
 **What is new is `group_size`.** `Qwen4ExpTextRMSNorm(dim, group_size)`
 reshapes the last axis to `(-1, group_size)`, normalises each group
-independently, and flattens back. Every 10240-wide norm in this model
-(`hc_norm`, `norm_key`, `norm_query`, `norm_conv`,
-`mtp.pre_fc_norm_hidden`) uses `group_size = 2560`, i.e. four
-independent RMS normalisations, not one over 10240. A single norm over
-the full 10240 vector compiles, runs, and is wrong.
+independently, and flattens back. The 10240-wide norms in the decoder
+(`hc_norm`, `norm_key`, `norm_query`, `norm_conv`) use
+`group_size = 2560`, i.e. four independent RMS normalisations, not one
+over 10240. A single norm over the full 10240 vector compiles, runs,
+and is wrong.
+
+**`mtp.pre_fc_norm_hidden` is the exception, and this document had it
+wrong.** It is 10240 wide and **ungrouped** — one normalisation over
+the whole vector. vLLM distinguishes the two cases explicitly: its
+hyper-connection uses `GroupedGemmaRMSNorm` with
+`group_size = hidden_size`, and the MTP head uses a plain
+`GemmaRMSNorm(hidden_size * hc_count)`. Both classes exist in that
+tree, so the choice is deliberate rather than an oversight. Grouping it
+would rescale the draft head's input in four independent pieces and
+show up as a poor acceptance rate, not as an error. See §9.
 
 Variance is computed in f32 and the `(1 + w)` shift is applied in f32
 before the multiply — same reasoning as the note already in
@@ -501,14 +511,41 @@ mtp.layers.0.*                                  # a complete full_attention
 
 **Transformers does not implement the MTP head.**
 `Qwen4ExpPreTrainedModel._keys_to_ignore_on_load_unexpected = [r"^mtp.*"]`
-— the weights are dropped on load. So for #313 the upstream reference is
-*not* usable, and the authorities are instead llama.cpp PR #27836
-(`--spec-type draft-mtp`) and vLLM's `speculative-config
-{"method": "mtp"}` path. The tensor names above and the two `fc_*`
-projections make the intended wiring clear (project the next token's
-embedding and the accepted hidden state, sum, run one decoder layer,
-reuse `lm_head`), but the accept rule and the proposal depth must come
-from one of those two implementations, not from this document.
+— the weights are dropped on load.
+
+**The wiring, settled against vLLM** (`vllm/models/qwen4_exp/nvidia/
+mtp.py`, read 2026-09-02). vLLM calls this shape
+`residual_linear_shared`, as against the `Linear(2H, H)` + repeat that
+other MTP variants use:
+
+```
+e = fc_embedding(pre_fc_norm_embedding(embed(t)))         # [T, 2560]
+h = pre_fc_norm_hidden(h.flatten(-2))                      # ONE norm over 10240
+h = fc_hidden(h.view(T, hc, 2560))                         # per stream, shared matrix
+h = e.unsqueeze(-2) + h                                    # embedding broadcast to all four
+h = h.flatten(-2)                                          # [T, 10240] -> the decoder layer
+```
+
+Three things here are not guessable from the tensor shapes:
+
+1. **`pre_fc_norm_hidden` is applied to the flattened 10240 and is
+   ungrouped** — see §8, which this document previously got wrong.
+2. **The hidden state is the *pre-final-mixer* multi-stream**, i.e. `h`
+   after the last decoder layer and *before* `hyper_connection_mixer`
+   collapses it. Not the 2560 the LM head sees. vLLM calls this
+   "scheme A"; on later draft steps the head consumes its own previous
+   multi-stream rather than the target's.
+3. **The head emits two things per step**: the collapsed `[T, 2560]`
+   for `lm_head`, and the pre-final-mixer `[T, 10240]` for the next
+   draft step.
+
+PLE is forced off in the draft layer while the HC stream count is kept
+(`ngram_context=None`), which is consistent with `mtp.layers.0` shipping
+no `ple.*` tensors.
+
+The accept rule is standard (greedy prefix match, or the Leviathan
+rejection scheme) and is not architecture-specific; this document
+previously conflated it with the wiring above.
 
 `mtp.layers.0` carries its own full 512-expert MoE — 2.61 B params,
 5.2 GB at BF16. It is not free to keep resident and its placement is