fix(neuron): trim the host heap on the load path too (#318)

The previous commit trimmed on drop, which misses the case that has
been happening all day: a load that fails part-way drops its partial
model inside `load_dense_inner`, not through `Job::DropArch`, so the
trim on unload never runs for it. Until an hour ago every qwen4_exp
load on this hardware failed part-way.

A load that succeeds also leaves transients worth returning — the f32
round-trips the quantiser makes to give each tensor ownership of its
storage, and the cache read buffers — so the trim goes after the load
either way, exactly as `trim_device_pool` already did on the line
above. The two now sit together and mean the same thing on either side
of the bus.

Refs #318, #325.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015v1FQqShiRp9rrJEEcgNER

.gitea/workflows/deploy.yml

@@ -317,21 +317,82 @@ jobs:
           # restart decision below, exactly like the drop-in.
           neuron_toml_changed="${{ steps.neuronconf.outputs.changed }}"
 
-          # ── Desired per-model systemd drop-in ─────────────────────────
-          # model.conf carries NEURON_MAX_PROMPT_TOKENS so the context cap
-          # is deterministic per host and rolled out (with a restart) by
+          # ── Desired per-host systemd drop-in ──────────────────────────
+          # model.conf carries the environment neuron is started with, so
+          # it is deterministic per host and rolled out (with a restart) by
           # this workflow, not hand-edited. It sorts after local.conf, so a
           # deploy-managed value wins over any manual local override of the
-          # same variable. See doc/context-limits.md.
+          # same variable.
+          #
+          # RUST_LOG is fleet-uniform and unconditional. The code default
+          # is `info`, which is too coarse to diagnose a load or a routing
+          # decision after the fact, and before this it was `debug` on
+          # exactly one host because someone set it there by hand. A
+          # verbosity you have to remember to set is a verbosity you find
+          # out you did not have at the moment you needed it. Bumping a
+          # module (`debug,neuron::harness::candle=trace`) is a one-line
+          # change here, applied everywhere, with an audit trail.
+          #
+          # NEURON_MAX_PROMPT_TOKENS is per-host and stays conditional;
+          # see doc/context-limits.md.
           conf=/etc/systemd/system/neuron.service.d/model.conf
           config_changed=${neuron_toml_changed:-0}
+          rust_log="debug"
+          desired=$(printf '%s\n%s\n%s' \
+            "# Managed by .gitea/workflows/deploy.yml - do not edit by hand." \
+            "[Service]" \
+            "Environment=RUST_LOG=${rust_log}")
           if [ -n "${max_prompt_tokens}" ]; then
-            desired=$(printf '%s\n%s\n%s\n%s' \
-              "# Managed by .gitea/workflows/deploy.yml - do not edit by hand." \
-              "# Per-model context cap; see doc/context-limits.md." \
-              "[Service]" \
+            desired=$(printf '%s\n%s' "${desired}" \
               "Environment=NEURON_MAX_PROMPT_TOKENS=${max_prompt_tokens}")
-            [ "${desired}" = "$(cat "${conf}" 2>/dev/null || true)" ] || config_changed=1
+          fi
+          [ "${desired}" = "$(cat "${conf}" 2>/dev/null || true)" ] || config_changed=1
+
+          # ── Desired registry drop-in ──────────────────────────────────
+          # Every model fetch goes through rustingface rather than
+          # huggingface.co: it holds the upstream credential, and it
+          # caches to the operator's store, so the second host to want a
+          # checkpoint does not re-cross the internet for it.
+          #
+          # This makes rustingface a hard dependency for every cold load
+          # on the fleet. That is the point, not a side effect. An outage
+          # on our own host is one we can end; an outage at huggingface.co
+          # is one we can only wait out, and a serving stack that stops
+          # when someone else's CDN does is not sovereign in any sense
+          # worth the word.
+          #
+          # So do NOT add a fallback to huggingface.co here. It would read
+          # as a robustness improvement and would delete the property this
+          # exists for: the fleet would silently resume depending on a
+          # third party at exactly the moments that dependency matters.
+          #
+          # registry.conf sorts after local.conf, so the deploy-managed
+          # values win over a hand-set HF_TOKEN — which is the point:
+          # this is the mechanism that stops the registry being a
+          # hand-synced per-host decision (#287).
+          #
+          # Mode 0600, and compared by hash under sudo, because it
+          # carries a token. Written ONLY when the secret is non-empty:
+          # pointing HF_ENDPOINT at a registry we cannot authenticate to
+          # turns every cold load into a 401, which is strictly worse
+          # than fetching from upstream.
+          rf_conf=/etc/systemd/system/neuron.service.d/registry.conf
+          rf_token='${{ secrets.RUSTINGFACE_TOKEN }}'
+          rf_endpoint='https://rf.internal'
+          if [ -n "${rf_token}" ]; then
+            rf_desired=$(printf '%s\n%s\n%s\n%s\n%s\n%s' \
+              "# Managed by .gitea/workflows/deploy.yml - do not edit by hand." \
+              "# All model fetches go through rustingface, which holds the" \
+              "# upstream Hugging Face credential and caches to our store." \
+              "[Service]" \
+              "Environment=HF_ENDPOINT=${rf_endpoint}" \
+              "Environment=HF_TOKEN=${rf_token}")
+            rf_have=$(sudo /usr/bin/sha256sum "${rf_conf}" 2>/dev/null | cut -d' ' -f1 || true)
+            rf_want=$(printf '%s\n' "${rf_desired}" | sha256sum | cut -d' ' -f1)
+            [ "${rf_have}" = "${rf_want}" ] || config_changed=1
+          else
+            echo "WARNING: RUSTINGFACE_TOKEN is empty - leaving ${rf_conf} as-is."
+            echo "         Fetches use whatever HF_ENDPOINT/HF_TOKEN the host already has."
           fi
 
           # ── Package version gate (manifest rationale: see deploy-cortex) ──
@@ -364,15 +425,24 @@ jobs:
           echo "installed=${installed} published=${latest:-unknown} pkg_changed=${pkg_changed} config_changed=${config_changed} — deploying"
 
           # Write the drop-in (staged in gitea_ci's dir, installed root-owned).
-          # Guarded on max_prompt_tokens as well as config_changed: since
-          # a neuron.toml change also sets config_changed, an entry with
-          # no max_prompt_tokens would otherwise reach this with `desired`
-          # unset.
-          if [ "${config_changed}" -eq 1 ] && [ -n "${max_prompt_tokens}" ]; then
+          # `desired` is always set now that RUST_LOG is unconditional, so
+          # config_changed alone is the guard.
+          if [ "${config_changed}" -eq 1 ]; then
             printf '%s\n' "${desired}" > /var/lib/gitea_ci/model.conf
             sudo /usr/bin/install -o root -g root -m 0644 -D /var/lib/gitea_ci/model.conf "${conf}"
             rm -f /var/lib/gitea_ci/model.conf
-            echo "applied ${conf}: NEURON_MAX_PROMPT_TOKENS=${max_prompt_tokens}"
+            echo "applied ${conf}: RUST_LOG=${rust_log} NEURON_MAX_PROMPT_TOKENS=${max_prompt_tokens:-<unset>}"
+          fi
+
+          # Staged 0600 in gitea_ci's own dir before install, so the token
+          # is never briefly world-readable on the way in.
+          if [ -n "${rf_token}" ] && [ "${rf_have:-}" != "${rf_want:-}" ]; then
+            umask 077
+            printf '%s\n' "${rf_desired}" > /var/lib/gitea_ci/registry.conf
+            sudo /usr/bin/install -o root -g root -m 0600 -D /var/lib/gitea_ci/registry.conf "${rf_conf}"
+            rm -f /var/lib/gitea_ci/registry.conf
+            umask 022
+            echo "applied ${rf_conf}: HF_ENDPOINT=${rf_endpoint} (token not logged)"
           fi
 
           if systemctl is-active --quiet neuron.service; then

CLAUDE.md

@@ -1061,3 +1061,85 @@ ampere flavour.
 - Deferred by operator decision: Z-Image base variant, operator
   fine-tunes, Z-Image-Edit (`images/edits`). Qwen-Image-2512 is the
   parked "frontier image" tier (own epic if taken).
+
+## 2026-09-02 addendum: qwen4_exp (Qwen3.8-Flash-Next) text path
+
+`Qwen/Qwen3.8-Flash-Next` is `model_type: qwen4_exp` — a new
+architecture generation, not a Qwen3 variant. Epic #307; the port spec
+is `doc/qwen4_exp-port-spec.md` and remains the reference for every
+dimension. The text path is implemented in `crates/neuron/src/harness/
+arch/qwen4_exp/` and routed in `harness/candle.rs`; it loads a
+checkpoint and returns logits.
+
+**What is genuinely new**, and where each lives:
+
+- `hyper.rs` — **four residual streams**, not one. The inter-layer
+  residual is `hidden_size * hc_count` (10240) wide, initialised as four
+  copies of the token embedding. There is **no `input_layernorm`, no
+  `post_attention_layernorm` and no final `model.norm`** in the
+  checkpoint: the `hc_norm` inside each hyper-connection does all three
+  jobs. Do not go looking for the missing tensor.
+- `ple.rs` — a hashed n-gram embedding on **zero-indexed layer 1**
+  (`ple_layer_ids: [2]` is one-indexed upstream) over a 320,001,536-row
+  table, 28.44% of the model's parameters. Three pieces: addressing
+  (`NGramHasher`), the gather (`NGramTable` trait — residency is #310),
+  and consumption (`PleBlock`).
+- `qsa.rs` — sparse attention. Blocks of 4 past positions are scored by
+  a 4-head side channel and only the best 512 attended, which is what
+  makes a 262k window affordable. Keeps a **second KV cache** (12
+  layers × 1 head × 128 dims, ~3 KiB/token) that every `kv_budget_mb`
+  arithmetic must account for.
+- `decoder.rs`, `model.rs` — the composition.
+
+**Reused from `arch/qwen3_5/` unchanged**: `rope.rs` (identical
+interleaved M-RoPE config), the GatedDeltaNet for 36 of 48 layers, and
+the MoE routing. Those reuses needed three small seams in `qwen3_5`,
+all behaviour-preserving: `apply_partial_rotary` / `cos_sin_at` on the
+rotary, `GatedDeltaNetParams` on the linear attention, and `from_parts`
+constructors on the MoE and MLP.
+
+### Sharp edges
+
+- **Read `model.safetensors.index.json` before writing a loader.** It is
+  170 KB, needs no auth, and `curl`-able from HF at the pinned revision.
+  It caught a real bug: the QSA indexer is at
+  `...self_attn.indexer.index_qk_proj.weight`, a **submodule of the
+  attention block**, where the spec's prose implied a sibling. That
+  failure would otherwise have surfaced after a 360 GB download.
+- **The three PLE buffers are I64 in a bf16 checkpoint.** Read through
+  the model's VarBuilder they round multipliers of ~1e13 and prime vocab
+  sizes of ~2e7 into nonsense; the table is then sized wrong and the
+  addressing still returns rows. Force the dtype.
+- **A QSA mask must never reach the flash-attn path.**
+  `attention_context` used to take `Option<&Tensor>` and read
+  `is_some()` as "causal", never looking at the tensor — so an arbitrary
+  mask was silently replaced by a dense causal one. It is now
+  `AttnMask::{None, Causal, Additive}`, and `Additive` forces the eager
+  path. Whether QSA can use flash-attn at all (it needs a *gather*, not
+  a mask) is open and now on the critical path.
+- **PLE's short conv is dilated** by `ngram_size`, so it reaches 9
+  positions back and its cached context is 9 wide, not `kernel_size`. It
+  cannot reuse `run_causal_conv1d`, which assumes dilation 1.
+- **`output_gate_type: sigmoid` and `mamba_ssm_dtype: float32`** are the
+  checkpoint's explicit choices, against our SiLU fallback and #284's
+  bf16 round-trip default. A port must not answer an upstream choice
+  with our own.
+
+### What the text path does not yet do
+
+- **No batched decode.** `supports_kv_snapshot` is false, so it serves
+  one request at a time whatever `max_in_flight` says. Extending
+  `snapshot.rs` to carry PLE's two state slots and the indexer cache is
+  its own piece of work.
+- **No tensor parallelism** — absent from `TP_SUPPORTED_MODEL_TYPES`, so
+  a `tensor_parallel` request is refused rather than silently served
+  from one device.
+- **The n-gram table's residency is undecided** (#310). The shipped
+  `ShardedNGramTable` holds the shards wherever the VarBuilder put them,
+  which for 27 GB is the problem that issue exists to solve; it is the
+  control, not the answer.
+- **Gate 6 (full-model logits parity against a reference trace) is
+  open.** Gates 4 and 5 — QSA below budget is exactly dense, and above
+  budget matches a naive oracle — are closed.
+- **MTP (#313), vision (#314), quantisation and rollout (#315)** are
+  untouched.

Cargo.lock

@@ -3147,6 +3147,8 @@ dependencies = [
  "half",
  "hf-hub",
  "image",
+ "libc",
+ "memmap2",
  "minijinja",
  "minijinja-contrib",
  "rayon",

asset/sudoers.d/neuron-host.conf

@@ -18,6 +18,14 @@ gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/neuron/neuron.toml
 # NEURON_MAX_PROMPT_TOKENS: gitea_ci stages it in its own dir, then
 # installs it root-owned. Exact source/dest paths; see doc/context-limits.md.
 gitea_ci ALL=(root) NOPASSWD: /usr/bin/install -o root -g root -m 0644 -D /var/lib/gitea_ci/model.conf /etc/systemd/system/neuron.service.d/model.conf
+# deploy.yml also writes the registry drop-in, which points every model
+# fetch at rustingface and carries the client token for it. Mode 0600
+# rather than model.conf's 0644: this one holds a credential, and a
+# world-readable systemd drop-in is a world-readable token.
+gitea_ci ALL=(root) NOPASSWD: /usr/bin/install -o root -g root -m 0600 -D /var/lib/gitea_ci/registry.conf /etc/systemd/system/neuron.service.d/registry.conf
+# Reading it back to decide whether it changed needs root too, since
+# 0600 excludes gitea_ci. Restricted to the one file.
+gitea_ci ALL=(root) NOPASSWD: /usr/bin/sha256sum /etc/systemd/system/neuron.service.d/registry.conf
 gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl start neuron.service
 gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl stop neuron.service
 gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl enable --now neuron.service

crates/cortex-core/src/harness.rs

@@ -172,6 +172,60 @@ pub struct ModelSpec {
     /// a belief. A request's own `chat_template_kwargs` still wins.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub preserve_thinking: Option<bool>,
+
+    /// Where the routed experts of a sparse MoE live (#318).
+    ///
+    /// `None` — the default, and every model on the fleet today — keeps
+    /// them on the device with everything else.
+    ///
+    /// `Some(ExpertResidency::Host)` quantises them into host memory
+    /// instead and moves the activations across for the expert matmul.
+    /// That is the only way `qwen4_exp` loads at all: its routed
+    /// experts are 73.0 GB at q4k against 63.7 GB of VRAM across both
+    /// of beast's cards, so no quantisation or tensor-parallel split
+    /// makes them fit. It is slow — a token's experts are ~1.35 GB of
+    /// host reads, so decode is bounded near 37 tok/s by DDR5
+    /// bandwidth before compute — and it is the difference between a
+    /// model that serves and one that does not.
+    ///
+    /// Explicitly opted into rather than inferred from a failed device
+    /// load, because #318's whole point is that placement is measured
+    /// per machine rather than decided once: an automatic fallback
+    /// would quietly pick for every host and hide the measurement it
+    /// exists to inform.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub expert_residency: Option<ExpertResidency>,
+
+    /// Concurrent running requests for *this* model, overriding the
+    /// host's `[harness.candle.admission] max_in_flight`.
+    ///
+    /// Capacity is a property of the model as much as the host. A
+    /// device-resident dense model amortises its weight reads across a
+    /// batch almost perfectly, which is why beast runs 8. A sparse MoE
+    /// does not: *n* concurrent tokens select up to `10n` distinct
+    /// experts, so expert traffic scales close to linearly with the
+    /// batch — #309 measured eight concurrent streams returning only
+    /// **1.73×** the total decode throughput of one, with the experts
+    /// still on the device. With them in host memory behind DDR5 the
+    /// same effect runs against a bus an order of magnitude slower.
+    ///
+    /// So a host that serves both wants different answers for each,
+    /// and the load spec is where per-model operator decisions already
+    /// travel (`sampling`, `preserve_thinking`).
+    ///
+    /// `None` takes the host's setting, which is every model today.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_in_flight: Option<usize>,
+}
+
+/// Where a sparse MoE's routed experts are held (#318).
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum ExpertResidency {
+    /// On the device, with the rest of the model. The default.
+    Device,
+    /// In host memory, with activations crossing for each expert.
+    Host,
 }
 
 /// Per-model token budget advertised by the catalogue or neuron.

crates/cortex-gateway/src/router.rs

@@ -545,6 +545,12 @@ async fn profile_to_spec(
         // this host was configured to hold resident.
         sampling: profile.sampling.clone(),
         preserve_thinking: profile.preserve_thinking,
+        // Expert residency (#318) is a property of the machine, not of
+        // the catalogue: whether a host must hold the experts in RAM
+        // depends on its VRAM, and cortex routes to several. The neuron
+        // decides from its own config.
+        expert_residency: None,
+        max_in_flight: None,
     }
 }
 

crates/helexa-bench/src/client.rs

@@ -218,6 +218,11 @@ impl TargetClient {
             // ever used to load a model the host has no config for.
             sampling: None,
             preserve_thinking: None,
+            // Same limitation, same reason: `GET /models` does not
+            // report where the experts were placed (#318), so a
+            // reconstructed spec cannot carry it either.
+            expert_residency: None,
+            max_in_flight: None,
         })
     }
 

crates/neuron/Cargo.toml

@@ -104,6 +104,18 @@ minijinja-contrib = { version = "2", features = ["pycompat"] }
 # tp `fused_load` module to read per-rank slices of fused QKV tensors
 # without materialising the full tensor on device.
 safetensors = "0.8"
+# The qwen4_exp PLE n-gram table is 102.4 GB of BF16 across 128 shards
+# (#310). It is a pure gather of ~5 KB per token, so it is read straight
+# out of the checkpoint by mmap rather than made resident anywhere —
+# the page cache decides what stays warm, and almost all of the table is
+# cold at any instant.
+memmap2 = "0.9"
+# `malloc_trim` after a model drop (#318). Host-resident experts are
+# ~1.4 MB per allocation and there are 73,728 of them; glibc raises its
+# mmap threshold as it observes frees, so past a point they come from
+# the arena and are never returned to the OS. Measured: 64.6 GB still
+# resident ten minutes after unload, on a 123 GB host.
+libc = "0.2"
 # Vision capability for Qwen3.6 (Stage A of the vision plan in
 # doc/vision-qwen3_6-spec.md). `image` decodes PNG/JPEG/etc from
 # the bytes embedded in `data:image/...;base64,...` content parts;

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,341 @@
 //! 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, Device, 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 QuantizedExpert {
+    /// Where this expert's weights are, when that is knowable.
+    ///
+    /// `None` for a `QMatMul` that dequantised on construction — it
+    /// holds a plain tensor whose device is the load device by
+    /// construction, so the caller's fallback is right.
+    fn device(&self) -> Option<Device> {
+        match &self.gate {
+            QMatMul::QTensor(t) => Some(t.device()),
+            _ => None,
+        }
+    }
+}
+
+/// The quantisation type actually usable for rows of `row` elements.
+///
+/// ggml quantises along the last dimension, and a row must be a whole
+/// number of blocks. The k-quants (`Q2K`..`Q8K`) block on 256; the
+/// older types block on 32. `qwen4_exp`'s `down_proj` reduces over
+/// `moe_intermediate_size = 640`, which is 2.5 blocks of 256 — so a
+/// `quant = "q4k"` load fails on the very first layer with
+/// `quantized tensor must have their last dim divisible by block size
+/// [2560, 640] 256`, having read 131 shards to get there.
+///
+/// **This table is llama.cpp's**, from `tensor_type_fallback` in
+/// `src/llama-quant.cpp`, and the principle behind it is the part
+/// worth stating: **every fallback spends bits rather than saving
+/// them.** A k-quant carries a two-level scale hierarchy that the
+/// legacy 32-block types do not, so it is better than a legacy type at
+/// equal width — dropping `Q4K` (4.5 bpw) to `Q4_0` (also 4.5) would
+/// hold the memory constant and take the quality hit silently.
+/// Upstream declines that trade and pays a bit instead:
+///
+/// | requested | fallback | bpw |
+/// |---|---|---|
+/// | `Q2K`, `Q3K` | `Q4_0` | 2.63 / 3.44 -> 4.5 |
+/// | `Q4K` | `Q5_0` | 4.5 -> 5.5 |
+/// | `Q5K` | `Q5_1` | 5.5 -> 6.0 |
+/// | `Q6K`, `Q8K` | `Q8_0` | 6.56 -> 8.5 |
+///
+/// For `qwen4_exp` that is not free: `down_proj` is exactly a third of
+/// the 120.8 B routed parameters, so `Q5_0` rather than `Q4_0` costs
+/// **+5.0 GB** across the expert bank (67.9 -> 73.0 GB). We take it
+/// anyway. The cheaper table was an unmeasured divergence in the one
+/// direction that flatters the constraint we are straining against,
+/// and this matrix is the one a mixed k-quant scheme normally treats
+/// as quality-critical. If the 5 GB turns out to buy nothing, #315 is
+/// where that gets measured and reversed on evidence.
+///
+/// A row that divides 32 either — vanishingly rare, and impossible for
+/// this checkpoint since 640 and 2560 are both multiples of 32 — falls
+/// back to F16 as upstream does, loudly. It is a 3.5x blowup on that
+/// tensor, so it must never be quiet.
+pub(crate) fn quant_for_row(requested: GgmlDType, row: usize) -> candle_core::Result<GgmlDType> {
+    if row.is_multiple_of(requested.block_size()) {
+        return Ok(requested);
+    }
+    let fallback = match requested {
+        GgmlDType::Q2K | GgmlDType::Q3K => GgmlDType::Q4_0,
+        GgmlDType::Q4K => GgmlDType::Q5_0,
+        GgmlDType::Q5K => GgmlDType::Q5_1,
+        GgmlDType::Q6K | GgmlDType::Q8K => GgmlDType::Q8_0,
+        // Already a 32-block type (or not a quant at all): there is no
+        // smaller block to demote to, so the F16 check below decides.
+        other => other,
+    };
+    if row.is_multiple_of(fallback.block_size()) {
+        return Ok(fallback);
+    }
+    tracing::warn!(
+        ?requested,
+        ?fallback,
+        row,
+        "row divides neither the requested block size nor its 32-wide \
+         fallback; storing this tensor as F16 — 3.5x the bytes"
+    );
+    Ok(GgmlDType::F16)
+}
+
+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.
+    ///
+    /// The two halves may not end up the same dtype: see
+    /// [`quant_for_row`]. `qwen4_exp`'s `down_proj` reduces over 640,
+    /// which no k-quant can block.
+    /// `onto` is where the quantised experts come to rest, which need
+    /// not be where the fused source tensor was read (#318). Passing
+    /// `Device::Cpu` for a CUDA load leaves the experts in host memory
+    /// and is what makes a model whose experts exceed VRAM loadable at
+    /// all. The source is dropped either way, so the peak is one
+    /// layer's fused pair regardless.
+    /// Quantise where the source already is. Only the tests want this
+    /// now — the loaders name a destination explicitly, because for
+    /// `qwen4_exp` the answer is not the source's device.
+    #[cfg(test)]
+    pub(crate) fn quantize_banked(
+        gate_up: &Tensor,
+        down: &Tensor,
+        intermediate: usize,
+        dtype: GgmlDType,
+    ) -> candle_core::Result<Self> {
+        Self::quantize_banked_onto(gate_up, down, intermediate, dtype, gate_up.device())
+    }
+
+    pub(crate) fn quantize_banked_onto(
+        gate_up: &Tensor,
+        down: &Tensor,
+        intermediate: usize,
+        dtype: GgmlDType,
+        onto: &Device,
+    ) -> candle_core::Result<Self> {
+        let n = gate_up.dims()[0];
+        // Decided once, from the shapes, rather than per expert —
+        // every expert in a layer has the same two row widths, and
+        // 512 identical log lines per layer would bury the fact.
+        let gate_up_dtype = quant_for_row(dtype, gate_up.dims()[2])?;
+        let down_dtype = quant_for_row(dtype, down.dims()[2])?;
+        if gate_up_dtype != dtype || down_dtype != dtype {
+            tracing::info!(
+                requested = ?dtype,
+                gate_up_row = gate_up.dims()[2],
+                gate_up_dtype = ?gate_up_dtype,
+                down_row = down.dims()[2],
+                down_dtype = ?down_dtype,
+                "expert quantisation: a row width does not divide the requested \
+                 block size; substituting a 32-block type for that half"
+            );
+        }
+        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, dtype: GgmlDType| -> candle_core::Result<QMatMul> {
+                let dims = t.dims().to_vec();
+                // This copy to host was always here — it is how the
+                // source comes to own its storage. Building the owned
+                // tensor on `onto` rather than on `t.device()` is the
+                // whole of host residency: the bytes are already in
+                // RAM at this point, and the only question is whether
+                // they go back.
+                let vals = t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?;
+                let owned = Tensor::from_vec(vals, dims, onto)?;
+                QMatMul::from_qtensor(QTensor::quantize(&owned, dtype)?)
+            };
+            out.push(QuantizedExpert {
+                gate: q(gu.narrow(0, 0, intermediate)?, gate_up_dtype)?,
+                up: q(gu.narrow(0, intermediate, intermediate)?, gate_up_dtype)?,
+                down: q(down.i(e)?, down_dtype)?,
+            });
+        }
+        Ok(Experts::Quantized(out))
+    }
+
+    /// The quantised tensors, named for the on-disk cache (#322).
+    ///
+    /// `None` for any layout that is not quantised — there is nothing
+    /// to cache about weights that were never transformed.
+    pub(crate) fn cache_entries(&self) -> Option<Vec<(String, std::sync::Arc<QTensor>)>> {
+        let Experts::Quantized(v) = self else {
+            return None;
+        };
+        let of = |m: &QMatMul| match m {
+            QMatMul::QTensor(t) => Some(t.clone()),
+            // A dequantised-on-load QMatMul has no quantised tensor to
+            // write. `CANDLE_DEQUANTIZE_ALL` produces these; caching
+            // half a layer would be worse than caching none of it.
+            _ => None,
+        };
+        let mut out = Vec::with_capacity(v.len() * 3);
+        for (i, q) in v.iter().enumerate() {
+            out.push((format!("e{i}.gate"), of(&q.gate)?));
+            out.push((format!("e{i}.up"), of(&q.up)?));
+            out.push((format!("e{i}.down"), of(&q.down)?));
+        }
+        Some(out)
+    }
+
+    /// Rebuild from what [`Self::cache_entries`] wrote.
+    ///
+    /// Errors rather than filling gaps: a layer missing an expert is a
+    /// truncated artifact, and serving it would mean one expert of 512
+    /// quietly doing nothing.
+    pub(crate) fn from_cache_entries(
+        mut entries: std::collections::HashMap<String, QTensor>,
+        num_experts: usize,
+    ) -> candle_core::Result<Self> {
+        let mut v = Vec::with_capacity(num_experts);
+        for i in 0..num_experts {
+            let mut take = |part: &str| -> candle_core::Result<QMatMul> {
+                let name = format!("e{i}.{part}");
+                match entries.remove(&name) {
+                    Some(t) => QMatMul::from_qtensor(t),
+                    None => candle_core::bail!("cached expert layer is missing '{name}'"),
+                }
+            };
+            v.push(QuantizedExpert {
+                gate: take("gate")?,
+                up: take("up")?,
+                down: take("down")?,
+            });
+        }
+        Ok(Experts::Quantized(v))
+    }
+
+    /// 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, and on the caller's device, so the
+                // routing arithmetic upstream does not have to know
+                // which layout or residency it got.
+                let dtype = xs.dtype();
+                let home = q.device().unwrap_or_else(|| xs.device().clone());
+                // Host-resident experts (#318): the activations go to
+                // the weights rather than the weights to the
+                // activations, because the weights are ~2.4 GB a layer
+                // and the activations are the rows routed to one
+                // expert. `to_device` is a clone when the devices
+                // already match, so the device-resident path pays
+                // nothing for this.
+                let xs32 = xs.to_dtype(DType::F32)?.to_device(&home)?;
+                let lhs = candle_nn::ops::silu(&q.gate.forward(&xs32)?)?;
+                let rhs = q.up.forward(&xs32)?;
+                q.down
+                    .forward(&(lhs * rhs)?)?
+                    .to_device(xs.device())?
+                    .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 +385,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 +452,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 +524,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 +575,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 +606,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 +660,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 +692,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 +708,384 @@ 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.
+    ///
+    /// Experts restored from cache land on the device they are asked
+    /// for, not the one the VarBuilder happens to use (#318, #322).
+    ///
+    /// This is the defect that made a host-resident load fill a 32 GB
+    /// card in 37 seconds: the fresh-quantisation path honoured the
+    /// residency target and the *cache* path did not, so 22 cached
+    /// layers at 1.42 GB each went to VRAM at about a gigabyte a
+    /// second while the code reported host residency in its logs. The
+    /// two paths produce the same experts and must place them the same
+    /// way; nothing but a test says so.
+    #[test]
+    fn cached_experts_land_where_they_are_asked_for() {
+        let dev = Device::Cpu;
+        let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let fresh =
+            Experts::quantize_banked_onto(&gate_up, &down, inter, GgmlDType::Q4K, &dev).unwrap();
+
+        let dir = tempfile::tempdir().unwrap();
+        let cache =
+            crate::harness::isq_cache::IsqCache::new(dir.path(), "test/model", Some("rev0"), "q4k")
+                .unwrap();
+        let entries = fresh.cache_entries().unwrap();
+        let refs: Vec<(&str, &QTensor)> = entries
+            .iter()
+            .map(|(n, t)| (n.as_str(), t.as_ref()))
+            .collect();
+        cache.try_store("layer-0", &refs).unwrap();
+
+        let restored =
+            Experts::from_cache_entries(cache.load("layer-0", &dev).unwrap(), n_experts).unwrap();
+        let Experts::Quantized(v) = &restored else {
+            panic!("a cache load must produce quantised experts");
+        };
+        for (e, q) in v.iter().enumerate() {
+            assert!(
+                q.device()
+                    .expect("a quantised expert knows its device")
+                    .same_device(&dev),
+                "expert {e} came back on the wrong device"
+            );
+        }
+    }
+
+    /// Host-resident experts compute the same function as device-
+    /// resident ones (#318).
+    ///
+    /// This is the contract that makes the whole placement decision
+    /// safe to take: *where* the weights live may change what a load
+    /// costs, and must not change what it computes. The activations
+    /// cross to the weights and the result crosses back, and nothing
+    /// else about the expert is different.
+    ///
+    /// On a CPU-only test host both "devices" are the same one, so this
+    /// pins the plumbing — the dtype round trip, the `to_device` on the
+    /// way out, the caller getting its own device and dtype back —
+    /// rather than the transfer itself. The transfer is exercised on
+    /// hardware, where the two devices genuinely differ; what a unit
+    /// test can hold is that residency is not silently part of the
+    /// arithmetic.
+    #[test]
+    fn host_resident_experts_compute_the_same_function() {
+        let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+
+        let on_device =
+            Experts::quantize_banked_onto(&gate_up, &down, inter, GgmlDType::Q4K, &Device::Cpu)
+                .unwrap();
+        let on_host =
+            Experts::quantize_banked_onto(&gate_up, &down, inter, GgmlDType::Q4K, &Device::Cpu)
+                .unwrap();
+
+        let xs = randn(&[3, hidden]);
+        for e in 0..n_experts {
+            let a: Vec<f32> = on_device
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            let b = on_host.forward_one(e, &xs).unwrap();
+            // The caller's device and dtype come back unchanged,
+            // whatever the weights did in between.
+            assert_eq!(b.dtype(), xs.dtype());
+            assert!(b.device().same_device(xs.device()));
+            let b: Vec<f32> = b.flatten_all().unwrap().to_vec1().unwrap();
+            assert_eq!(a, b, "expert {e} computed differently off-device");
+        }
+    }
+
+    /// Experts restored from the on-disk cache compute the same
+    /// function as the ones that filled it (#322).
+    ///
+    /// The cache's whole risk is that it returns something plausible:
+    /// the model loads, it serves, and its weights are quietly not the
+    /// ones it was asked for. So this asserts the restored experts are
+    /// *bit-identical* in output, not close — there is no arithmetic
+    /// between writing and reading that could legitimately differ.
+    ///
+    /// At the checkpoint's real widths, because they are what decide
+    /// the types involved: 2560 blocks under Q4K, 640 does not and
+    /// falls back to Q5_0, so this exercises a layer whose halves are
+    /// *different* quantisation types — exactly the case a container
+    /// storing one type per file could not represent, and the reason
+    /// the artifact is GGUF.
+    #[test]
+    fn cached_experts_compute_the_same_function() {
+        let dev = Device::Cpu;
+        let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let fresh = Experts::quantize_banked(&gate_up, &down, inter, GgmlDType::Q4K).unwrap();
+
+        let entries = fresh
+            .cache_entries()
+            .expect("quantised experts must be cacheable");
+        // The two halves really did take different types; if they had
+        // not, this test would not be covering what it claims to.
+        let dtypes: std::collections::HashSet<GgmlDType> =
+            entries.iter().map(|(_, t)| t.dtype()).collect();
+        assert!(
+            dtypes.len() > 1,
+            "expected a mix of quant types across gate/up and down, got {dtypes:?}"
+        );
+
+        let dir = tempfile::tempdir().unwrap();
+        let cache =
+            crate::harness::isq_cache::IsqCache::new(dir.path(), "test/model", Some("rev0"), "q4k")
+                .unwrap();
+        let refs: Vec<(&str, &QTensor)> = entries
+            .iter()
+            .map(|(n, t)| (n.as_str(), t.as_ref()))
+            .collect();
+        cache
+            .try_store("layer-0", &refs)
+            .expect("the write must succeed, and say why if it does not");
+
+        let restored = Experts::from_cache_entries(
+            cache
+                .load("layer-0", &dev)
+                .expect("the artifact just written"),
+            n_experts,
+        )
+        .unwrap();
+
+        let xs = randn(&[3, hidden]);
+        for e in 0..n_experts {
+            let a: Vec<f32> = fresh
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            let b: Vec<f32> = restored
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            assert_eq!(a, b, "expert {e} changed value on the cache round trip");
+        }
+    }
+
+    /// `qwen4_exp`'s real expert widths, which no k-quant can block.
+    ///
+    /// `moe_intermediate_size` is 640, so `down_proj` is `[2560, 640]`
+    /// and its 640-element rows are 2.5 blocks of 256. The load died
+    /// on beast at decoder layer 0 with `quantized tensor must have
+    /// their last dim divisible by block size [2560, 640] 256`, after
+    /// resolving all 131 shards.
+    ///
+    /// The fixture that missed it was 256 wide in both dimensions —
+    /// the smallest width a k-quant accepts — so every row divided
+    /// and the orientation could not matter. This one uses the
+    /// checkpoint's own numbers, scaled down only in expert count.
+    #[test]
+    fn the_real_expert_widths_quantise_despite_a_640_wide_reduction() {
+        assert_eq!(
+            quant_for_row(GgmlDType::Q4K, 2560).unwrap(),
+            GgmlDType::Q4K,
+            "gate_up reduces over hidden_size, which is 10 whole blocks"
+        );
+        assert_eq!(
+            quant_for_row(GgmlDType::Q4K, 640).unwrap(),
+            GgmlDType::Q5_0,
+            "down reduces over moe_intermediate_size, which is not"
+        );
+        // llama.cpp's `tensor_type_fallback` table, which spends bits
+        // rather than saving them: a legacy 32-block type is worse than
+        // a k-quant at equal width, so upstream pays one to avoid the
+        // quality cliff. Q4_0 above would be 4.5 bpw like Q4K and would
+        // save 5.0 GB across the expert bank — that is the trade we
+        // declined, and this is where it is pinned.
+        assert_eq!(quant_for_row(GgmlDType::Q5K, 640).unwrap(), GgmlDType::Q5_1);
+        assert_eq!(quant_for_row(GgmlDType::Q2K, 640).unwrap(), GgmlDType::Q4_0);
+        assert_eq!(quant_for_row(GgmlDType::Q3K, 640).unwrap(), GgmlDType::Q4_0);
+        assert_eq!(quant_for_row(GgmlDType::Q6K, 640).unwrap(), GgmlDType::Q8_0);
+        assert_eq!(
+            quant_for_row(GgmlDType::Q8_0, 640).unwrap(),
+            GgmlDType::Q8_0
+        );
+        // 100 divides neither 256 nor 32: F16, as upstream does.
+        assert_eq!(quant_for_row(GgmlDType::Q4K, 100).unwrap(), GgmlDType::F16);
+
+        // And the whole path runs at those widths.
+        let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let quantised = Experts::quantize_banked(&gate_up, &down, inter, GgmlDType::Q4K)
+            .expect("the real widths must quantise");
+        let xs = randn(&[2, hidden]);
+        let got = quantised.forward_one(0, &xs).expect("forward");
+        assert_eq!(got.dims(), &[2, hidden]);
+        assert!(
+            got.flatten_all()
+                .unwrap()
+                .to_vec1::<f32>()
+                .unwrap()
+                .iter()
+                .all(|v| v.is_finite()),
+            "a substituted dtype must still produce finite output"
+        );
+    }
+
+    /// 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/config.rs

@@ -197,10 +197,18 @@ impl TextConfig {
             .map(|d| d.eq_ignore_ascii_case("float32") || d.eq_ignore_ascii_case("f32"))
     }
 
+    /// Both spellings upstream uses. The shipped checkpoint's
+    /// `layer_types` says `full_attention` (36 linear, 12 full), but
+    /// `Qwen4ExpTextConfig` *derives* the array with
+    /// `qwen_sparse_attention` when the checkpoint omits it. Matching
+    /// only the first is a silent misclassification: every layer would
+    /// read as linear attention, and a model that is 3:1 linear would
+    /// look almost right until its full-attention tensors were not
+    /// found.
     pub fn is_full_attention(&self, layer: usize) -> bool {
         self.layer_types
             .get(layer)
-            .is_some_and(|t| t == "full_attention")
+            .is_some_and(|t| t == "full_attention" || t == "qwen_sparse_attention")
     }
 
     fn fill_layer_types(&mut self) -> Result<()> {
@@ -474,6 +482,35 @@ mod tests {
     }
 
     /// A one-indexed field that reads 0 is a config written against the
+    /// Upstream has two names for the same layer, and we must read both.
+    ///
+    /// The shipped `config.json` spells the attention layers
+    /// `full_attention`. `Qwen4ExpTextConfig` derives them as
+    /// `qwen_sparse_attention` when the checkpoint omits the array —
+    /// observed directly from `transformers` `main`, which reports
+    /// `['linear_attention', ..., 'qwen_sparse_attention']` for a config
+    /// carrying only `full_attention_interval`.
+    ///
+    /// Reading only the shipped spelling would classify every layer as
+    /// linear on such a config. That is not a load error at the point it
+    /// happens — it is a model built with the wrong 12 layers, which
+    /// fails later and elsewhere.
+    #[test]
+    fn both_upstream_spellings_of_a_full_attention_layer_are_read() {
+        let mut cfg = Config::from_config_json(SHIPPED).unwrap().text_config;
+
+        cfg.layer_types = vec!["linear_attention".into(), "full_attention".into()];
+        assert!(!cfg.is_full_attention(0));
+        assert!(cfg.is_full_attention(1), "the shipped spelling");
+
+        cfg.layer_types = vec!["linear_attention".into(), "qwen_sparse_attention".into()];
+        assert!(!cfg.is_full_attention(0));
+        assert!(
+            cfg.is_full_attention(1),
+            "the derived spelling must classify as full attention too"
+        );
+    }
+
     /// wrong convention, not a layer.
     #[test]
     fn rejects_a_zero_ple_layer_id() {

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

@@ -46,6 +46,7 @@ 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.
@@ -72,6 +73,7 @@ impl DecoderLayer {
         rotary: Arc<RotaryEmbedding>,
         layer_idx: usize,
         vb: &ShardedVarBuilder,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let layer_type = cfg
             .layer_types
@@ -83,6 +85,30 @@ impl DecoderLayer {
                     cfg.layer_types.len()
                 )
             })?;
+        Self::load_typed(
+            cfg,
+            rotary,
+            layer_type,
+            cfg.ple_layers().contains(&layer_idx),
+            vb,
+            opts,
+        )
+    }
+
+    /// 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,
+        vb: &ShardedVarBuilder,
+        opts: &super::LoadOptions<'_>,
+    ) -> Result<Self> {
         let mixer = match layer_type {
             "full_attention" => {
                 Mixer::Full(Box::new(Attention::load(cfg, rotary, &vb.pp("self_attn"))?))
@@ -91,8 +117,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 +135,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 +153,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, &vb.pp("mlp"), opts)?,
             ple,
         })
     }
@@ -190,6 +216,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,50 @@ pub mod hyper;
 pub mod linear_attn;
 pub mod model;
 pub mod moe;
+pub mod mtp;
 pub mod ple;
 pub mod qsa;
+
+/// What a load needs to know beyond the config and the weights.
+///
+/// Threading these one at a time got the loader signatures to eight
+/// arguments, at which point clippy is making a readability point
+/// rather than a stylistic one: `load(cfg, dtype, quant, device, vb,
+/// paths, cache, onto)` has two devices and two ways to name the
+/// checkpoint in it, and nothing at the call site says which is which.
+///
+/// They also belong together. Every one of them is a decision about
+/// *this load* rather than about the architecture — what precision, on
+/// what hardware, reading from where, reusing what.
+pub struct LoadOptions<'a> {
+    /// In-situ quantisation type (#315). `None` keeps the checkpoint's
+    /// precision, which for this architecture fits nowhere.
+    pub quant: Option<candle_core::quantized::GgmlDType>,
+    /// The checkpoint's own files, for mmapping the PLE n-gram table
+    /// where it already lies (#310). Empty falls back to loading it
+    /// through the VarBuilder, which puts 102.4 GB wherever that
+    /// points.
+    pub safetensors_paths: &'a [std::path::PathBuf],
+    /// Reuse of a previous load's quantisation (#322).
+    pub isq_cache: Option<&'a crate::harness::isq_cache::IsqCache>,
+    /// Where the routed experts come to rest (#318) — the load device
+    /// for an ordinary model, host memory for one whose experts exceed
+    /// VRAM.
+    pub experts_onto: &'a candle_core::Device,
+    /// A VarBuilder rooted at the checkpoint but bound to
+    /// [`Self::experts_onto`], used to read the fused expert tensors
+    /// when that is not the load device.
+    ///
+    /// Without it, host-resident experts still arrive *through* VRAM:
+    /// `get` puts the layer's 5.03 GB fused pair on the load device,
+    /// the quantiser copies it to the host, and the source is dropped
+    /// — 48 times, 241 GB of PCIe traffic, against a pool that has to
+    /// absorb two multi-gigabyte blocks and 1536 small ones per layer.
+    /// Measured: the load reached layer 22 of 48 and ran out with 31 GB
+    /// nominally free, which is fragmentation rather than occupancy —
+    /// the same failure the image path trims after every job to avoid.
+    ///
+    /// Reading them straight into host memory removes the transit
+    /// rather than managing it.
+    pub experts_vb: Option<&'a candle_nn::var_builder::ShardedVarBuilder<'a>>,
+}

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

@@ -44,7 +44,8 @@ use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
 use super::config::{Config, TextConfig};
 use super::decoder::DecoderLayer;
 use super::hyper::HyperConnection;
-use super::ple::{NGramHasher, NGramTable, ShardedNGramTable};
+use super::ple::{MmapNGramTable, 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 +90,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,11 +124,21 @@ 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.
+    /// `safetensors_paths` is the checkpoint's own file set, used only
+    /// to mmap the PLE n-gram table where it already lies (#310). Empty
+    /// falls back to [`ShardedNGramTable`], which puts 102.4 GB
+    /// wherever `vb` points — the oracle, not a residency to serve
+    /// from.
     pub fn load(
         cfg: &Config,
         dtype: DType,
         device: &Device,
         vb: &ShardedVarBuilder,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let text = &cfg.text_config;
         // `RotaryEmbedding` reads a qwen3_5 config; the rope block is the
@@ -131,10 +154,29 @@ impl Qwen4ExpForCausalLM {
         let embed_tokens = Embedding::new(embed_weight, text.hidden_size);
 
         let layers_vb = root.pp("layers");
+
+        // Before the layers, not after them. The table is mapped rather
+        // than made resident (#310), so its position in the load order
+        // costs nothing — but everything it can get wrong is a name or
+        // an extent, and those are worth finding in the second before
+        // the expert loop rather than the half-hour after it. Loading
+        // it last meant a checkpoint whose shard names had moved would
+        // report that only once 48 layers of experts had been
+        // quantised, or — since they do not currently fit — never.
+        let ngram = match text.ple_layers().first() {
+            Some(layer) => Some(load_ngram(
+                text,
+                &layers_vb.pp(*layer).pp("ple"),
+                opts.safetensors_paths,
+                device,
+            )?),
+            None => None,
+        };
+
         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, &layers_vb.pp(i), opts)
                     .with_context(|| format!("load decoder layer {i}"))?,
             );
         }
@@ -156,11 +198,6 @@ impl Qwen4ExpForCausalLM {
             .get((text.vocab_size, text.hidden_size), "weight")
             .context("load 'lm_head.weight'")?;
 
-        let ngram = match text.ple_layers().first() {
-            Some(layer) => Some(load_ngram(text, &layers_vb.pp(*layer).pp("ple"))?),
-            None => None,
-        };
-
         Ok(Self {
             embed_tokens,
             layers,
@@ -176,7 +213,44 @@ impl Qwen4ExpForCausalLM {
 
     /// `input_ids` is `(B, T)`; `offset` is the sequence position the
     /// first of them sits at. Returns logits `(B, T, vocab)`.
+    /// Logits for **every** position, `(B, T, vocab)`.
+    ///
+    /// Not the serving path — [`Self::forward`] is, and it returns only
+    /// the last position because that is all sampling needs and all the
+    /// harness accepts. This exists because the upstream reference
+    /// returns every position, so full-model parity (#323) has to
+    /// compare every position or compare nothing.
+    ///
+    /// The two are pinned together by a test rather than by
+    /// inspection: `forward` must equal the last row of this. A
+    /// second path that is only exercised by tests is worth exactly
+    /// what proves it agrees with the first.
+    pub fn logits_all_positions(&mut self, input_ids: &Tensor, offset: usize) -> Result<Tensor> {
+        let (_, h) = self.forward_inner(input_ids, offset)?;
+        let x = self.mixer.collapse(&h)?;
+        Ok(self.lm_head.forward(&x)?)
+    }
+
+    /// Logits for the **last** position, `(B, 1, 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 +283,101 @@ 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)?)
+        // Only the last position, `(B, 1, vocab)` — the same contract
+        // `qwen3_5` states and the harness relies on: `squeeze_to_vocab`
+        // strips leading singleton dims and refuses anything else, so a
+        // full `(B, T, vocab)` reaches it as `(T, vocab)` and is
+        // rejected. Returning every position is what a *reference*
+        // implementation does, which is why the parity fixture holds
+        // them all and [`Self::logits_all_positions`] exists to be
+        // compared against it.
+        let l = x.dim(1)?;
+        Ok((self.lm_head.forward(&x.i((.., l - 1.., ..))?)?, 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<()> {
@@ -234,7 +402,12 @@ impl Qwen4ExpForCausalLM {
     }
 }
 
-fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbedding> {
+fn load_ngram(
+    cfg: &TextConfig,
+    ple_vb: &ShardedVarBuilder,
+    safetensors_paths: &[std::path::PathBuf],
+    device: &Device,
+) -> Result<NGramEmbedding> {
     let vb = ple_vb.pp("ple_embedding");
     let heads = cfg.ngram_heads();
     let head_dim = cfg.ngram_head_dim();
@@ -260,12 +433,28 @@ fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbed
         cfg.split_ngram_parts
     );
 
-    let table = ShardedNGramTable::load(
-        &vb.pp("ngram_embedding"),
-        cfg.split_ngram_parts,
-        rows,
-        head_dim,
-    )?;
+    // The shard names are whatever this VarBuilder would have asked
+    // for, so the prefix is taken from it rather than rebuilt from the
+    // layer index — the two cannot disagree that way, and getting it
+    // wrong is a 102.4 GB read of the wrong thing.
+    let ngram_vb = vb.pp("ngram_embedding");
+    let table: Box<dyn NGramTable> = if safetensors_paths.is_empty() {
+        Box::new(ShardedNGramTable::load(
+            &ngram_vb,
+            cfg.split_ngram_parts,
+            rows,
+            head_dim,
+        )?)
+    } else {
+        Box::new(MmapNGramTable::open(
+            safetensors_paths,
+            &ngram_vb.prefix(),
+            cfg.split_ngram_parts,
+            rows,
+            head_dim,
+            device,
+        )?)
+    };
 
     Ok(NGramEmbedding {
         hasher: NGramHasher::new(
@@ -276,7 +465,7 @@ fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbed
             offsets,
             cfg.eos_token_id,
         )?,
-        table: Box::new(table),
+        table,
         heads,
         head_dim,
         carried: Vec::new(),
@@ -320,6 +509,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 +725,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 +824,23 @@ 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();
+        // Passing the path exercises the mmap table (#310), which is
+        // the residency production uses — the whole suite then runs
+        // against the gather that actually serves, not the oracle.
+        let model = Qwen4ExpForCausalLM::load(
+            &cfg,
+            DType::F32,
+            &Device::Cpu,
+            &vb,
+            &super::super::LoadOptions {
+                quant: None,
+                safetensors_paths: std::slice::from_ref(&path),
+                isq_cache: None,
+                experts_onto: &Device::Cpu,
+                experts_vb: None,
+            },
+        )
+        .unwrap();
         (dir, model)
     }
 
@@ -558,13 +853,24 @@ mod tests {
         let (_dir, mut model) = build();
         let ids = Tensor::from_vec(vec![5u32, 9, 3, 12, 7], (1, 5), &Device::Cpu).unwrap();
 
+        // The serving path returns the last position only, `(B, 1, V)`
+        // — the contract `qwen3_5` states and `squeeze_to_vocab`
+        // enforces.
         let logits = model.forward(&ids, 0).unwrap();
-        assert_eq!(logits.dims(), &[1, 5, 32]);
+        assert_eq!(logits.dims(), &[1, 1, 32]);
         let values: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
         assert!(
             values.iter().all(|v| v.is_finite()),
             "a wrong stream width or a dead mask row shows up as NaN"
         );
+
+        // The parity path returns every position, and both must be
+        // finite: a dead mask row shows up in one and not the other.
+        model.clear_kv_cache().unwrap();
+        let all = model.logits_all_positions(&ids, 0).unwrap();
+        assert_eq!(all.dims(), &[1, 5, 32]);
+        let values: Vec<f32> = all.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(values.iter().all(|v| v.is_finite()));
     }
 
     /// A decode step continues the prefill: the caches, the rotary
@@ -603,4 +909,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_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,
+    vb: &ShardedVarBuilder,
+    opts: &super::LoadOptions<'_>,
+) -> Result<Qwen3_5MoeBlock> {
+    let (quant, isq_cache, experts_onto) = (opts.quant, opts.isq_cache, opts.experts_onto);
     let (h, inter) = (cfg.hidden_size, cfg.moe_intermediate_size);
     anyhow::ensure!(
         cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && inter > 0,
@@ -58,14 +77,78 @@ pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Qwen3_5MoeBlock>
         None,
     );
 
-    let experts_vb = vb.pp("experts");
-    let gate_up = experts_vb
-        .get((cfg.num_experts, inter * 2, h), "gate_up_proj")
-        .with_context(|| format!("load '{}/gate_up_proj'", experts_vb.prefix()))?;
-    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)?;
+    // Read the fused tensors wherever the experts are going to live.
+    // When that is the host, this is a VarBuilder bound to it, walked
+    // to the same prefix — so the 5.03 GB pair never touches the
+    // device on its way to being quantised into RAM.
+    let experts_vb = match opts.experts_vb {
+        Some(host_root) => {
+            let mut at = host_root.clone();
+            for segment in vb.prefix().split('.').filter(|s| !s.is_empty()) {
+                at = at.pp(segment);
+            }
+            at.pp("experts")
+        }
+        None => vb.pp("experts"),
+    };
+    // The cache key is the layer's own tensor prefix, so it cannot
+    // collide across layers and needs no separate index to be passed
+    // down alongside it.
+    let cache_key = experts_vb.prefix().replace(['.', '/'], "-");
+
+    // A hit skips reading the fused tensors at all, which is the point:
+    // the cost being avoided is the 5.03 GB read *and* the ~32 s of
+    // quantisation, not just the latter.
+    let cached = match (quant, isq_cache) {
+        (Some(_), Some(cache)) => cache.load(&cache_key, experts_onto).and_then(|entries| {
+            match Experts::from_cache_entries(entries, cfg.num_experts) {
+                Ok(e) => Some(e),
+                Err(e) => {
+                    tracing::warn!(key = %cache_key, error = %e,
+                            "isq cache: layer artifact rejected, re-quantising");
+                    None
+                }
+            }
+        }),
+        _ => None,
+    };
+
+    let experts = match cached {
+        Some(experts) => experts,
+        None => {
+            let gate_up = experts_vb
+                .get((cfg.num_experts, inter * 2, h), "gate_up_proj")
+                .with_context(|| format!("load '{}/gate_up_proj'", experts_vb.prefix()))?;
+            let down = experts_vb
+                .get((cfg.num_experts, h, inter), "down_proj")
+                .with_context(|| format!("load '{}/down_proj'", experts_vb.prefix()))?;
+            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.
+            match quant {
+                Some(dtype) => {
+                    let experts =
+                        Experts::quantize_banked_onto(&gate_up, &down, inter, dtype, experts_onto)
+                            .with_context(|| format!("quantize experts to {dtype:?}"))?;
+                    if let (Some(cache), Some(entries)) = (isq_cache, experts.cache_entries()) {
+                        let refs: Vec<(&str, &candle_core::quantized::QTensor)> = entries
+                            .iter()
+                            .map(|(n, t)| (n.as_str(), t.as_ref()))
+                            .collect();
+                        cache.store(&cache_key, &refs);
+                    }
+                    experts
+                }
+                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 +176,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 +197,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,273 @@
+//! 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,
+                &mtp.pp("layers").pp(0),
+                // The MTP head is not on the serving path (#313), so it
+                // has nothing to gain from the cache and nothing to
+                // stale it — and its experts go wherever the rest of it
+                // does.
+                &super::LoadOptions {
+                    quant,
+                    safetensors_paths: &[],
+                    isq_cache: None,
+                    experts_onto: vb.device(),
+                    experts_vb: None,
+                },
+            )
+            .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

@@ -66,7 +66,7 @@
 //! reimplement. See `doc/qwen4_exp-port-spec.md` §2.
 
 use anyhow::{Context, Result, ensure};
-use candle_core::{D, Module, Tensor};
+use candle_core::{D, DType, Device, Module, Tensor};
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
@@ -337,6 +337,216 @@ impl NGramTable for ShardedNGramTable {
     }
 }
 
+/// The table read straight out of the checkpoint by mmap (#310).
+///
+/// The shipping residency. Nothing is made resident anywhere: the 128
+/// shard tensors are mapped where they already lie in the checkpoint's
+/// safetensors files, and a gather copies the requested rows out of
+/// the mapping. The page cache decides what stays warm, which for this
+/// access pattern is the right decision-maker — a token needs sixteen
+/// rows of 160 values, about 5 KB, out of 102.4 GB, so almost the whole
+/// table is cold at any instant and nothing is served by making it
+/// otherwise.
+///
+/// This is what makes the model loadable at all. [`ShardedNGramTable`]
+/// puts those 102.4 GB wherever the `VarBuilder` points, which for a
+/// CUDA load is a 32 GB card; it stays as the oracle this is diffed
+/// against, not as a residency anyone would choose.
+///
+/// Mapping the checkpoint's own tensors in place — rather than building
+/// a separate artifact — is FreeToken's `source_from_safetensors`
+/// trick, and the reason it works here is that the shards are already
+/// exactly what we want: contiguous, equal-stride, row-major, one
+/// extent each. There is nothing to convert and nothing to keep in
+/// sync.
+///
+/// Unlike the sharded gather this needs no sort-and-permute: rows are
+/// copied in the order asked for, so the inverse permutation that has
+/// to be right there cannot be wrong here.
+pub struct MmapNGramTable {
+    /// One mapping per checkpoint file that holds shards. Kept for the
+    /// table's lifetime because every `gather` reads through them.
+    maps: Vec<memmap2::Mmap>,
+    /// Byte position of each shard's row 0: which mapping, and where in
+    /// it. Resolved once at load so a gather is arithmetic.
+    shards: Vec<ShardExtent>,
+    rows_per_shard: usize,
+    head_dim: usize,
+    dtype: DType,
+    row_bytes: usize,
+    device: Device,
+}
+
+/// Where one shard's bytes live.
+struct ShardExtent {
+    map: usize,
+    offset: usize,
+}
+
+impl MmapNGramTable {
+    /// `prefix` is the full tensor-name prefix of the shards, e.g.
+    /// `model.language_model.layers.1.ple.ple_embedding.ngram_embedding`
+    /// — the shards are then `{prefix}.shard_{s}.weight`.
+    ///
+    /// `files` is the checkpoint's safetensors set; shards may be
+    /// spread across any of them (in the real checkpoint, 33 of the
+    /// 131 files carry them), so each is opened and asked what it has.
+    pub fn open(
+        files: &[std::path::PathBuf],
+        prefix: &str,
+        parts: usize,
+        rows: usize,
+        head_dim: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        ensure!(parts > 0, "split_ngram_parts must be > 0");
+        ensure!(
+            rows.is_multiple_of(parts),
+            "{rows} rows do not divide evenly into {parts} shards"
+        );
+        let rows_per_shard = rows / parts;
+
+        let mut maps: Vec<memmap2::Mmap> = Vec::new();
+        let mut found: Vec<Option<ShardExtent>> = (0..parts).map(|_| None).collect();
+        let mut dtype: Option<DType> = None;
+
+        for path in files {
+            let file = std::fs::File::open(path)
+                .with_context(|| format!("open {} for the n-gram table", path.display()))?;
+            // SAFETY: the same contract candle's own safetensors
+            // loading takes — the checkpoint is not modified while we
+            // serve from it. A truncated or rewritten file underneath a
+            // live mapping is a corrupt-model problem, not one this can
+            // defend against.
+            let map = unsafe { memmap2::Mmap::map(&file) }
+                .with_context(|| format!("mmap {}", path.display()))?;
+            let (header_len, meta) =
+                safetensors::SafeTensors::read_metadata(&map).map_err(|e| {
+                    anyhow::anyhow!("read safetensors metadata from {}: {e}", path.display())
+                })?;
+            // `data_offsets` are relative to the start of the data
+            // section, which follows the 8-byte header length and the
+            // header itself.
+            let data_start = 8 + header_len;
+
+            let mut used = false;
+            for (s, slot) in found.iter_mut().enumerate() {
+                if slot.is_some() {
+                    continue;
+                }
+                let name = format!("{prefix}.shard_{s}.weight");
+                let Some(info) = meta.info(&name) else {
+                    continue;
+                };
+                ensure!(
+                    info.shape == [rows_per_shard, head_dim],
+                    "'{name}' is {:?}, expected [{rows_per_shard}, {head_dim}]",
+                    info.shape
+                );
+                let dt = st_dtype(info.dtype)
+                    .with_context(|| format!("'{name}' has a dtype the gather cannot read"))?;
+                match dtype {
+                    None => dtype = Some(dt),
+                    Some(seen) => ensure!(
+                        seen == dt,
+                        "shard dtypes disagree: {seen:?} then {dt:?} at '{name}'"
+                    ),
+                }
+                *slot = Some(ShardExtent {
+                    map: maps.len(),
+                    offset: data_start + info.data_offsets.0,
+                });
+                used = true;
+            }
+            if used {
+                maps.push(map);
+            }
+        }
+
+        let dtype = dtype.context("no n-gram shard was found in any checkpoint file")?;
+        let shards = found
+            .into_iter()
+            .enumerate()
+            .map(|(s, e)| e.with_context(|| format!("'{prefix}.shard_{s}.weight' is in no file")))
+            .collect::<Result<Vec<_>>>()?;
+
+        let row_bytes = head_dim * dtype.size_in_bytes();
+        // Every extent must actually contain its rows: a header that
+        // disagrees with the file length would otherwise surface as a
+        // panic inside a gather, mid-request.
+        for (s, extent) in shards.iter().enumerate() {
+            let end = extent.offset + rows_per_shard * row_bytes;
+            ensure!(
+                end <= maps[extent.map].len(),
+                "shard {s} claims bytes {}..{end} of a {}-byte file",
+                extent.offset,
+                maps[extent.map].len()
+            );
+        }
+
+        tracing::info!(
+            parts,
+            rows,
+            head_dim,
+            ?dtype,
+            files = maps.len(),
+            gib = (rows * row_bytes) as f64 / (1024.0 * 1024.0 * 1024.0),
+            "ple: n-gram table mapped from the checkpoint, not made resident (#310)"
+        );
+
+        Ok(Self {
+            maps,
+            shards,
+            rows_per_shard,
+            head_dim,
+            dtype,
+            row_bytes,
+            device: device.clone(),
+        })
+    }
+}
+
+impl NGramTable for MmapNGramTable {
+    fn gather(&self, rows: &[i64]) -> Result<Tensor> {
+        let total = (self.shards.len() * self.rows_per_shard) as i64;
+        if let Some(bad) = rows.iter().find(|r| **r < 0 || **r >= total) {
+            anyhow::bail!("row {bad} is outside the table's 0..{total}");
+        }
+        let mut out = vec![0u8; rows.len() * self.row_bytes];
+        for (slot, &row) in rows.iter().enumerate() {
+            let shard = row as usize / self.rows_per_shard;
+            let within = row as usize % self.rows_per_shard;
+            let extent = &self.shards[shard];
+            let from = extent.offset + within * self.row_bytes;
+            out[slot * self.row_bytes..(slot + 1) * self.row_bytes]
+                .copy_from_slice(&self.maps[extent.map][from..from + self.row_bytes]);
+        }
+        Ok(Tensor::from_raw_buffer(
+            &out,
+            self.dtype,
+            &[rows.len(), self.head_dim],
+            &self.device,
+        )?)
+    }
+
+    fn head_dim(&self) -> usize {
+        self.head_dim
+    }
+}
+
+/// safetensors' dtype vocabulary is wider than candle's. Only the
+/// float widths a checkpoint would store an embedding table in are
+/// accepted; anything else is a checkpoint we do not understand rather
+/// than one to guess at.
+fn st_dtype(dtype: safetensors::Dtype) -> Result<DType> {
+    Ok(match dtype {
+        safetensors::Dtype::BF16 => DType::BF16,
+        safetensors::Dtype::F16 => DType::F16,
+        safetensors::Dtype::F32 => DType::F32,
+        other => anyhow::bail!("unsupported n-gram table dtype {other:?}"),
+    })
+}
+
 /// The consumption half of PLE: `Qwen4ExpTextPLE` minus the table.
 ///
 /// Given the gathered n-gram embedding `e` and the four-stream residual
@@ -491,6 +701,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.
@@ -883,6 +1113,88 @@ mod tests {
 
     // ---- the table ----
 
+    /// The mmap table against the device-resident one, row for row.
+    ///
+    /// This is the differential #310 asks for: [`ShardedNGramTable`] is
+    /// the oracle — obviously correct, and unusable at 102.4 GB — and
+    /// [`MmapNGramTable`] is what serves. They must agree on every row
+    /// of every shard, in whatever order they are asked, or the model
+    /// quietly reads somebody else's n-gram.
+    ///
+    /// The row list is chosen to break the things that can break: shard
+    /// boundaries in both directions, the first and last row of the
+    /// table, a repeat, and an order that is not sorted (the oracle
+    /// sorts internally and inverts the permutation; the mmap path does
+    /// not, so a disagreement here is exactly that inverse being wrong
+    /// on one side).
+    #[test]
+    fn the_mapped_table_agrees_with_the_resident_one() {
+        let dev = Device::Cpu;
+        let (parts, rows_per_shard, head_dim) = (4usize, 5usize, 3usize);
+        let prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding";
+
+        // Row r holds [r, r + 0.5, r + 0.25] — distinct per row, and
+        // not a pattern a stride mistake would reproduce.
+        let mut tensors = std::collections::HashMap::new();
+        for s in 0..parts {
+            let mut v = Vec::with_capacity(rows_per_shard * head_dim);
+            for r in 0..rows_per_shard {
+                let global = (s * rows_per_shard + r) as f32;
+                v.extend_from_slice(&[global, global + 0.5, global + 0.25]);
+            }
+            tensors.insert(
+                format!("{prefix}.shard_{s}.weight"),
+                Tensor::from_vec(v, (rows_per_shard, head_dim), &dev).unwrap(),
+            );
+        }
+
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&tensors, &path).unwrap();
+
+        let vb = unsafe {
+            candle_nn::var_builder::ShardedSafeTensors::var_builder(&[&path], DType::F32, &dev)
+                .unwrap()
+        };
+        let total = parts * rows_per_shard;
+        let oracle = ShardedNGramTable::load(&vb.pp(prefix), parts, total, head_dim).unwrap();
+        let mapped = MmapNGramTable::open(
+            std::slice::from_ref(&path),
+            prefix,
+            parts,
+            total,
+            head_dim,
+            &dev,
+        )
+        .unwrap();
+
+        assert_eq!(mapped.head_dim(), oracle.head_dim());
+
+        let asked: Vec<i64> = vec![19, 0, 5, 4, 9, 10, 19, 1, 15, 14, 7, 7];
+        let flat = |t: Tensor| -> Vec<f32> { t.flatten_all().unwrap().to_vec1().unwrap() };
+        let from_oracle = flat(oracle.gather(&asked).unwrap());
+        let from_mapped = flat(mapped.gather(&asked).unwrap());
+        assert_eq!(
+            from_mapped, from_oracle,
+            "the mapped gather disagrees with the resident one"
+        );
+        // And that both are right, rather than wrong together.
+        assert_eq!(&from_mapped[..3], &[19.0, 19.5, 19.25]);
+        assert_eq!(&from_mapped[3..6], &[0.0, 0.5, 0.25]);
+
+        // Every row of the table, in order, so no shard is untouched.
+        let all: Vec<i64> = (0..total as i64).collect();
+        assert_eq!(
+            flat(mapped.gather(&all).unwrap()),
+            flat(oracle.gather(&all).unwrap())
+        );
+
+        // The empty gather and the out-of-range guard behave alike.
+        assert_eq!(mapped.gather(&[]).unwrap().dims(), &[0, head_dim]);
+        assert!(mapped.gather(&[total as i64]).is_err());
+        assert!(mapped.gather(&[-1]).is_err());
+    }
+
     /// Rows come back in the order asked for, not the order they were
     /// fetched in. Two shards of three rows, requested out of order and
     /// across the boundary: the shard-major gather returns them sorted,

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

@@ -134,7 +134,13 @@ impl BlockSelector {
     /// Returns `[n_blocks]`.
     pub fn scores(&self, q: &Tensor, pooled: &Tensor) -> candle_core::Result<Tensor> {
         debug_assert_eq!(q.dims2()?, (self.n_heads, self.head_dim));
-        let per_head = q.matmul(&pooled.t()?)?.relu()?;
+        // f32 for the same reason `pool` is: this is a *selection*
+        // function. See the note there and on the batched path in
+        // `Indexer::positions_from`.
+        let per_head = q
+            .to_dtype(DType::F32)?
+            .matmul(&pooled.to_dtype(DType::F32)?.t()?)?
+            .relu()?;
         per_head.sum(0)? / (self.head_dim as f64).sqrt()
     }
 
@@ -293,6 +299,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) {
@@ -490,15 +511,30 @@ impl Indexer {
         // in one batched matmul: every head and query shares the same
         // block keys, so fold them into the row axis rather than
         // broadcasting the keys across heads.
-        let q_rows = q.reshape((batch, self.n_heads * seq_len, self.head_dim))?;
-        let keys = blocks.squeeze(1)?.transpose(1, 2)?.contiguous()?;
+        //
+        // In f32, matching upstream, which casts both operands before
+        // the matmul (`torch.matmul(q.float(), block_key_states.float()
+        // .transpose(-1, -2))`). It is not a rounding nicety: these
+        // scores feed a top-k over as many as 65,536 blocks, and the
+        // output is the *set of positions a query may attend*. Two
+        // blocks whose f32 scores differ by less than a bf16 tick —
+        // 0.03125 near 4.0 — compare equal in bf16 and the tie-break
+        // decides instead, so bf16 does not perturb the answer, it
+        // changes which past the model sees. Same argument as `pool`,
+        // which we already got right.
+        let q_rows = q
+            .reshape((batch, self.n_heads * seq_len, self.head_dim))?
+            .to_dtype(DType::F32)?;
+        let keys = blocks
+            .squeeze(1)?
+            .transpose(1, 2)?
+            .contiguous()?
+            .to_dtype(DType::F32)?;
         let per_head = q_rows
             .matmul(&keys)?
             .reshape((batch, self.n_heads, seq_len, n_blocks))?
             .relu()?;
-        let scores = (per_head.sum(1)? / (self.head_dim as f64).sqrt())?
-            .to_dtype(DType::F32)?
-            .to_vec3::<f32>()?;
+        let scores = (per_head.sum(1)? / (self.head_dim as f64).sqrt())?.to_vec3::<f32>()?;
 
         let mut out = Vec::with_capacity(batch);
         for row in scores.iter().take(batch) {
@@ -634,6 +670,63 @@ mod tests {
     /// mean is exactly 0.25 — and so is an f32 mean cast back to bf16,
     /// because the true 0.2507 sits below the first bf16 tick above
     /// 0.25. Both mistakes land on the same wrong number, which is why
+    /// Scoring is f32 even when the cache is bf16, because it decides
+    /// *which past a query sees*, not merely by how much.
+    ///
+    /// Upstream casts both operands before the matmul
+    /// (`torch.matmul(q.float(), block_key_states.float().T)`); we did
+    /// not, and ran the matmul, the relu and the sum over heads in the
+    /// model dtype.
+    ///
+    /// The construction is the whole argument. Two blocks, four heads,
+    /// per-head contributions of 1.0 and — for one head of one block —
+    /// 1.0078125, which is 1 + 2^-7 and therefore exact in bf16. Their
+    /// sums are 4.0 and 4.0078125, and bf16's tick near 4.0 is 2^-5 =
+    /// 0.03125, so the second rounds back onto the first. The inputs
+    /// survive the dtype; the *comparison* does not. In bf16 the two
+    /// blocks tie and `select`'s index tie-break picks the wrong one,
+    /// silently, for as long as the model is served.
+    #[test]
+    fn scoring_is_f32_even_when_the_cache_is_bf16() {
+        let dev = Device::Cpu;
+        // Budget 4 over 4-wide blocks keeps exactly one block, so the
+        // comparison decides the answer rather than being absorbed by a
+        // budget large enough to take both.
+        let qsa = BlockSelector::new(4, 2, 4, 4).unwrap();
+
+        // Block 0 reads dimension 0 of each head, block 1 reads
+        // dimension 1, so the two scores are independent sums.
+        let pooled = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (2, 2), &dev)
+            .unwrap()
+            .to_dtype(DType::BF16)
+            .unwrap();
+        let q = Tensor::from_vec(
+            vec![
+                1.0f32, 1.0, // head 0: block 0 <- 1.0, block 1 <- 1.0
+                1.0, 1.0, // head 1
+                1.0, 1.0, // head 2
+                1.0, 1.0078125, // head 3: block 1 gets the extra tick
+            ],
+            (4, 2),
+            &dev,
+        )
+        .unwrap()
+        .to_dtype(DType::BF16)
+        .unwrap();
+
+        let scores: Vec<f32> = qsa.scores(&q, &pooled).unwrap().to_vec1().unwrap();
+        assert!(
+            scores[1] > scores[0],
+            "bf16 scoring lost the distinction between the blocks: {scores:?}"
+        );
+        // And the selection that rests on it picks the right block.
+        assert_eq!(
+            qsa.select(&scores, 8).unwrap(),
+            vec![4, 5, 6, 7],
+            "the higher-scoring block should be the one kept"
+        );
+    }
+
     /// this asserts the value rather than the dtype.
     #[test]
     fn pooling_accumulates_in_f32() {
@@ -952,4 +1045,202 @@ mod tests {
             "with a one-block budget exactly one of the two blocks wins"
         );
     }
+
+    // ---- gate 5: the batched path against a naive one ----
+
+    /// A deliberately naive QSA, written the way the reference is: one
+    /// query at a time, pooling and scoring and selecting in plain f32
+    /// rather than in one batched matmul.
+    ///
+    /// It shares only the projection and the rotation with the
+    /// implementation — the weights, which are not the subject. The
+    /// pooling, the per-query visibility, the scoring, the top-k and
+    /// the tail assembly are all done independently here, because those
+    /// are what gate 5 is asking about.
+    fn oracle(
+        ix: &Indexer,
+        hidden: &Tensor,
+        rope: &RotaryEmbedding,
+        past_len: usize,
+    ) -> Vec<Vec<Vec<usize>>> {
+        let (batch, seq_len, _) = hidden.dims3().unwrap();
+        let (q, keys) = ix.project(hidden).unwrap();
+        let total = past_len + seq_len;
+        let block = ix.selector.block_size;
+        let n_blocks = total / block;
+
+        // Raw cached keys, as plain numbers.
+        let raw: Vec<Vec<Vec<f32>>> = (0..batch)
+            .map(|b| {
+                (0..total)
+                    .map(|t| keys.i(b).unwrap().i(t).unwrap().to_vec1::<f32>().unwrap())
+                    .collect()
+            })
+            .collect();
+
+        let mut out = Vec::new();
+        for (b, rows) in raw.iter().enumerate() {
+            // Pool each block by hand: the mean of four rows.
+            let pooled: Vec<Vec<f32>> = rows
+                .chunks_exact(block)
+                .take(n_blocks)
+                .map(|chunk| {
+                    let mut acc = vec![0.0f64; ix.head_dim];
+                    for row in chunk {
+                        for (a, v) in acc.iter_mut().zip(row) {
+                            *a += *v as f64;
+                        }
+                    }
+                    acc.iter().map(|a| (a / block as f64) as f32).collect()
+                })
+                .collect();
+
+            // Norm and rotate the pooled keys with the layer's own
+            // weights, one block at a time and at its own start.
+            let mut block_keys = Vec::with_capacity(n_blocks);
+            for (blk, row) in pooled.iter().enumerate() {
+                let t =
+                    Tensor::from_vec(row.clone(), (1, 1, 1, ix.head_dim), keys.device()).unwrap();
+                let normed = ix.k_layernorm.forward(&t).unwrap();
+                let (cos, sin) = rope.cos_sin_at(&[blk * block]).unwrap();
+                block_keys.push(
+                    apply_partial_rotary(&normed, &cos, &sin, rope.rotary_dim())
+                        .unwrap()
+                        .flatten_all()
+                        .unwrap()
+                        .to_vec1::<f32>()
+                        .unwrap(),
+                );
+            }
+
+            let mut per_query = Vec::with_capacity(seq_len);
+            for i in 0..seq_len {
+                let pos = past_len + i;
+                // This query's heads, rotated at its own position.
+                let qi = q
+                    .i(b)
+                    .unwrap()
+                    .narrow(1, i, 1)
+                    .unwrap()
+                    .unsqueeze(0)
+                    .unwrap();
+                let (cos, sin) = rope.cos_sin_at(&[pos]).unwrap();
+                let qi = apply_partial_rotary(&qi, &cos, &sin, rope.rotary_dim()).unwrap();
+                let heads: Vec<Vec<f32>> = (0..ix.n_heads)
+                    .map(|h| {
+                        qi.i(0)
+                            .unwrap()
+                            .i(h)
+                            .unwrap()
+                            .i(0)
+                            .unwrap()
+                            .to_vec1::<f32>()
+                            .unwrap()
+                    })
+                    .collect();
+
+                let visible = pos + 1;
+                let visible_blocks = visible / block;
+                let mut scored: Vec<(usize, f32)> = (0..visible_blocks)
+                    .map(|blk| {
+                        let score: f32 = heads
+                            .iter()
+                            .map(|h| {
+                                let dot: f32 =
+                                    h.iter().zip(&block_keys[blk]).map(|(a, b)| a * b).sum();
+                                dot.max(0.0)
+                            })
+                            .sum::<f32>()
+                            / (ix.head_dim as f32).sqrt();
+                        (blk, score)
+                    })
+                    .collect();
+
+                // Top-k, ties to the earlier block.
+                scored.sort_by(|a, b| {
+                    b.1.partial_cmp(&a.1)
+                        .unwrap_or(std::cmp::Ordering::Equal)
+                        .then(a.0.cmp(&b.0))
+                });
+                scored.truncate(ix.selector.block_topk().min(visible_blocks));
+                let mut chosen: Vec<usize> = scored.iter().map(|(blk, _)| *blk).collect();
+                chosen.sort_unstable();
+
+                let mut positions = Vec::new();
+                for blk in chosen {
+                    positions.extend(blk * block..(blk + 1) * block);
+                }
+                positions.extend(visible_blocks * block..visible);
+                per_query.push(positions);
+            }
+            out.push(per_query);
+        }
+        out
+    }
+
+    /// **Gate 5.** Above the budget, where the selection actually binds,
+    /// the batched implementation must choose what the naive one does.
+    ///
+    /// Gate 4 proves the layer is dense attention below the budget,
+    /// which is a strong statement about the easy case and says nothing
+    /// about the hard one. This is the hard one: 23 positions, a budget
+    /// of two blocks out of five, so three blocks are discarded per
+    /// query and the ordering of the scores is what decides which.
+    #[test]
+    fn above_budget_selection_matches_the_naive_reference() {
+        let (hidden_size, head_dim) = (8, 8);
+        let rope = rope_for(head_dim);
+        let h = hidden(2, 23, hidden_size, 8.8);
+
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 8);
+        assert!(!ix.selector.is_dense(23), "the budget must bind");
+        let got = ix.visible_positions(&h, &rope, 0).unwrap();
+
+        let reference = indexer(hidden_size, 4, head_dim, 4, 8);
+        let want = oracle(&reference, &h, &rope, 0);
+
+        assert_eq!(got.len(), want.len());
+        for (b, (rows_got, rows_want)) in got.iter().zip(&want).enumerate() {
+            for (i, (g, w)) in rows_got.iter().zip(rows_want).enumerate() {
+                assert_eq!(g, w, "batch {b}, query {i}");
+            }
+        }
+
+        // Guard against agreeing on nothing: at least one query must
+        // actually have discarded a block, or this proves only that two
+        // implementations of dense attention agree.
+        let discarded = got[0]
+            .iter()
+            .enumerate()
+            .filter(|(i, positions)| positions.len() < i + 1)
+            .count();
+        assert!(
+            discarded > 0,
+            "no query discarded a block; the comparison is vacuous"
+        );
+    }
+
+    /// The same, continued across a decode boundary: the reference sees
+    /// one flat sequence, the implementation sees a prefill and then a
+    /// step reading its own cache.
+    #[test]
+    fn a_bound_decode_step_matches_the_naive_reference() {
+        let (hidden_size, head_dim) = (8, 8);
+        let rope = rope_for(head_dim);
+        let h = hidden(1, 23, hidden_size, 9.9);
+
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 8);
+        let prefill = h.narrow(1, 0, 22).unwrap().contiguous().unwrap();
+        ix.visible_positions(&prefill, &rope, 0).unwrap();
+        let step = h.narrow(1, 22, 1).unwrap().contiguous().unwrap();
+        let got = ix.visible_positions(&step, &rope, 22).unwrap();
+
+        let reference = indexer(hidden_size, 4, head_dim, 4, 8);
+        let want = oracle(&reference, &h, &rope, 0);
+
+        assert_eq!(
+            got[0][0], want[0][22],
+            "the last query, after a decode step"
+        );
+    }
 }

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),
 }
@@ -1050,6 +1050,16 @@ pub enum ModelArch {
     // are real; forward bails "not implemented yet". See
     // `arch/qwen3_5.rs` for the open architecture work.
     Qwen3_5Dense(super::arch::qwen3_5::Qwen3_5ForCausalLM),
+
+    // Qwen3.8-Flash-Next (model_type "qwen4_exp") — a new architecture
+    // generation, not a Qwen3 variant: four residual streams joined by
+    // hyper-connections, QSA sparse attention, and a hashed n-gram
+    // table on one layer. Boxed for the same reason as LlamaDense.
+    //
+    // No snapshot or batched-decode support yet, so it falls through to
+    // the "no batched-decode support" arms and serves one request at a
+    // time. See `arch/qwen4_exp/` and the epic on #307.
+    Qwen4Exp(Box<super::arch::qwen4_exp::model::Qwen4ExpForCausalLM>),
 }
 
 impl ModelArch {
@@ -1066,6 +1076,7 @@ impl ModelArch {
             ModelArch::LlamaQuantized(m) => m.forward(input, offset)?,
             ModelArch::LlamaDense(m) => m.forward(input, offset)?,
             ModelArch::Qwen3_5Dense(m) => m.forward(input, offset)?,
+            ModelArch::Qwen4Exp(m) => m.forward(input, offset)?,
         };
         squeeze_to_vocab(&raw)
     }
@@ -1093,6 +1104,10 @@ impl ModelArch {
                 m.clear_kv_cache();
                 Ok(())
             }
+            // Clears the KV caches, the QSA indexer's separate cache and
+            // PLE's carried ids together — a survivor of any one of them
+            // reads the previous request's text.
+            ModelArch::Qwen4Exp(m) => m.clear_kv_cache(),
         }
     }
 
@@ -1101,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"),
         }
     }
@@ -1118,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"),
         }
     }
@@ -1266,6 +1283,29 @@ impl ModelArch {
     }
 }
 
+/// Apply a model's `max_in_flight` override to the host's admission
+/// settings (#318).
+///
+/// A free function so the policy can be tested without constructing a
+/// harness — and so it is one place rather than the three that read
+/// the limit: the controller, and the two batch-engine gates. Those
+/// gates reading `self.admission_cfg` directly is what would have made
+/// an override that looked applied do nothing: admission would bound
+/// the model at 1 while the engine still started with 8 slots.
+fn admission_with_override(
+    host: &crate::config::AdmissionConfig,
+    model: Option<usize>,
+) -> crate::config::AdmissionConfig {
+    let mut cfg = host.clone();
+    // Zero is not "unlimited" and not "none"; it is a config that would
+    // admit nothing. Ignore it rather than serve a model that can never
+    // run a request.
+    if let Some(n) = model.filter(|n| *n > 0) {
+        cfg.max_in_flight = n;
+    }
+    cfg
+}
+
 /// Squeeze any leading singleton dims off the logits tensor so the
 /// caller gets a rank-1 `[vocab_size]` slice ready for sampling. Bails
 /// on a non-singleton leading dim (would mean a batched forward, which
@@ -1332,8 +1372,14 @@ impl LlamaDense {
 /// value. New entries land alongside a new `ModelArch` variant + a
 /// dispatch branch in `load_arch_dense` (plus, for TP, a parallel
 /// pattern in `tp_qwen3.rs`).
-const DENSE_SUPPORTED_MODEL_TYPES: &[&str] =
-    &["llama", "qwen3", "qwen3_5", "qwen3_moe", "qwen3_next"];
+pub(crate) const DENSE_SUPPORTED_MODEL_TYPES: &[&str] = &[
+    "llama",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_moe",
+    "qwen3_next",
+    "qwen4_exp",
+];
 
 /// Pre-flight check the operator's `config.json` against the set of
 /// architectures the dense path actually knows how to build. Surfaces
@@ -1388,6 +1434,11 @@ pub(crate) fn check_dense_config_supported(config_json: &str, model_id: &str) ->
 /// families than the TP path because each TP-aware module is a real
 /// chunk of work (`tp_qwen3.rs` is the only one shipped today).
 #[cfg(feature = "cuda")]
+// `qwen4_exp` is deliberately absent: the single-GPU path loads it, but
+// there is no `tp_qwen4_exp.rs`, so a `tensor_parallel` request for it
+// is refused here rather than silently served from one device. This
+// list is `cfg(cuda)`, so no CPU-build test can cover it — the gate is
+// the list itself.
 const TP_SUPPORTED_MODEL_TYPES: &[&str] = &["qwen3", "qwen3_5", "qwen3_next"];
 
 /// TP-side counterpart to `check_dense_config_supported`. Gates the
@@ -2707,6 +2758,15 @@ 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 isq_cache = self.isq_cache_for(
+            &spec.model_id,
+            &source_id.scheme,
+            spec.quant.as_deref(),
+            &safetensors_paths,
+        );
 
         let arch = tokio::task::spawn_blocking(move || -> Result<ModelArch> {
             let cfg_text = std::fs::read_to_string(&config_path).context("read config.json")?;
@@ -2802,6 +2862,49 @@ impl CandleHarness {
                         .context("build Qwen3-Next dense model")?;
                     Ok(ModelArch::Qwen3_5Dense(model))
                 }
+                "qwen4_exp" => {
+                    // Qwen3.8-Flash-Next. Like qwen3_5 it loads through
+                    // the sharded backend; unlike it, the config is a
+                    // separate type — the field vocabulary overlaps but
+                    // the model-level plumbing does not.
+                    let cfg = super::arch::qwen4_exp::config::Config::from_config_json(&cfg_text)
+                        .context("parse qwen4_exp config.json")?;
+                    let sharded_vb = unsafe {
+                        candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                            &safetensors_paths,
+                            dtype,
+                            &device_for_load,
+                        )
+                        .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,
+                        &device_for_load,
+                        &sharded_vb,
+                        &super::arch::qwen4_exp::LoadOptions {
+                            quant: isq,
+                            safetensors_paths: &safetensors_paths,
+                            isq_cache: isq_cache.as_ref(),
+                            // The CPU path already has everything in
+                            // host memory, so residency is not a
+                            // choice here.
+                            experts_onto: &device_for_load,
+                            experts_vb: None,
+                        },
+                    )
+                    .context("build qwen4_exp model")?;
+                    Ok(ModelArch::Qwen4Exp(Box::new(model)))
+                }
                 other => {
                     // Defensive: `check_dense_config_supported` already
                     // gated on the supported set, so this branch is
@@ -2820,6 +2923,120 @@ impl CandleHarness {
         Ok((tokenizer_path, arch))
     }
 
+    /// Whether this spec's `quant` names a GGUF the repo actually has.
+    ///
+    /// `quant` is overloaded (#320): the TP path reads it as in-situ
+    /// quantisation of safetensors, while the single-GPU paths read it
+    /// as "select the GGUF whose filename contains this string". A repo
+    /// with no GGUF at all cannot mean the second, so rather than fail
+    /// with an error about a file that was never going to exist, the
+    /// same string is read the way the TP path reads it.
+    ///
+    /// This is deliberately narrow. A repo that *does* carry GGUFs keeps
+    /// the old behaviour exactly, including the old error when none of
+    /// them match the requested quant — that case is a typo, not an ISQ
+    /// request, and turning it into a silent 360 GB dense load would be
+    /// a worse answer than the error.
+    /// The host's admission settings with this model's override
+    /// applied (#318).
+    ///
+    /// Only `max_in_flight` is overridable: the queue depth, the wait
+    /// deadline and the per-principal share are policies about how a
+    /// host treats its callers, and they should not vary by which
+    /// model a caller happened to ask for. How many requests a model
+    /// can actually run at once is a property of the model.
+    fn admission_for(&self, spec: &ModelSpec) -> crate::config::AdmissionConfig {
+        let cfg = admission_with_override(&self.admission_cfg, spec.max_in_flight);
+        if cfg.max_in_flight != self.admission_cfg.max_in_flight {
+            tracing::info!(
+                model = %spec.model_id,
+                host_default = self.admission_cfg.max_in_flight,
+                model_override = cfg.max_in_flight,
+                "admission: this model overrides the host's max_in_flight (#318)"
+            );
+        }
+        cfg
+    }
+
+    /// The quantisation cache for this load, or `None` when the quant
+    /// is unset or the checkpoint's revision cannot be identified
+    /// (#322).
+    ///
+    /// It lives **inside the source's own cache root**, beside the
+    /// weights it is derived from, rather than behind a config key
+    /// naming a path. That is the fleet's existing design: an operator
+    /// who wants this on faster storage symlinks it there, exactly as
+    /// the hot model weights on beast are symlinked from the spinning
+    /// `/archive3` cache onto NVMe. Adding a second, differently-shaped
+    /// way to say the same thing would leave two places to get it
+    /// wrong.
+    ///
+    /// A revision we cannot read means we cannot prove two checkpoints
+    /// are the same bytes, and a cache key that might collide is worse
+    /// than no cache — so an unrecognised layout disables it rather
+    /// than guessing.
+    fn isq_cache_for(
+        &self,
+        model_id: &str,
+        scheme: &str,
+        quant: Option<&str>,
+        safetensors_paths: &[PathBuf],
+    ) -> Option<super::isq_cache::IsqCache> {
+        let quant = quant.filter(|q| !q.is_empty())?;
+        let root = self.hf_cache_for(scheme).ok()?.path().join("neuron-isq");
+        let revision = super::isq_cache::revision_from_paths(safetensors_paths);
+        let cache = super::isq_cache::IsqCache::new(&root, model_id, revision, quant);
+        match &cache {
+            Some(c) => tracing::info!(model = model_id, quant, dir = %c.dir().display(),
+                "isq cache: enabled for this load (#322)"),
+            None => tracing::info!(
+                model = model_id,
+                quant,
+                "isq cache: disabled — the checkpoint revision could not be identified"
+            ),
+        }
+        cache
+    }
+
+    async fn quant_selects_a_gguf(
+        &self,
+        spec: &ModelSpec,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<bool> {
+        let Some(quant) = spec.quant.as_deref().filter(|q| !q.is_empty()) else {
+            return Ok(false);
+        };
+        let api = self.hf_api_for(&source_id.scheme)?;
+        let info = api
+            .model(source_id.repo_path())
+            .info()
+            .await
+            .with_context(|| format!("fetch HF repo info for {source_id}"))?;
+        let names: Vec<String> = info
+            .siblings
+            .iter()
+            .map(|s| s.rfilename.to_lowercase())
+            .collect();
+        let any_gguf = names.iter().any(|n| n.ends_with(".gguf"));
+        if any_gguf {
+            return Ok(true);
+        }
+        let any_safetensors = names.iter().any(|n| n.ends_with(".safetensors"));
+        if any_safetensors {
+            tracing::info!(
+                model = %source_id,
+                quant = %quant,
+                "no GGUF in this repo; reading `quant` as in-situ quantization of \
+                 the safetensors (#320)"
+            );
+            return Ok(false);
+        }
+        anyhow::bail!(
+            "repo {source_id} has neither a GGUF nor safetensors; `quant = {quant:?}` \
+             cannot be satisfied"
+        )
+    }
+
     /// Resolve a model spec to local GGUF and tokenizer file paths via
     /// hf-hub. Downloads on first use; subsequent calls are cached.
     async fn resolve_files(
@@ -3804,7 +4021,7 @@ impl Harness for CandleHarness {
             context_profile,
         ) = if let Some(w) = &worker {
             // CUDA path: resolve, then load in the worker.
-            if spec.quant.is_some() {
+            if self.quant_selects_a_gguf(spec, &source_id).await? {
                 let (gguf_path, tokenizer_path) = self.resolve_files(spec, &source_id).await?;
                 let handle = w
                     .load_gguf(gguf_path, spec.model_id.clone())
@@ -3828,36 +4045,65 @@ impl Harness for CandleHarness {
                 let (config_path, tokenizer_path, safetensors_paths) =
                     self.resolve_dense_files(spec, &source_id).await?;
                 let meta = VisionMeta::from_config_path(&config_path);
-                // Prefix snapshots (#11) exist only for the in-tree
-                // qwen3_5 arch — which serves BOTH its model_types
-                // (qwen3_5 and the flat-config qwen3_next MoE family);
-                // the worker holds the ModelArch so the async side
-                // decides from config.json instead.
-                let snapshot_capable = matches!(
-                    config_model_type(&config_path).as_deref(),
-                    Some(super::arch::qwen3_5::MODEL_TYPE)
-                        | Some(super::arch::qwen3_5::MODEL_TYPE_NEXT)
-                );
                 // Context-limit physics (#67): single-GPU → world_size 1.
-                // `None` for non-qwen3_5 dense archs.
+                // `None` for dense archs with no profile builder — they
+                // fall back to the static prompt cap.
                 let context_profile =
-                    super::context_limit::profile_from_qwen3_5_config(&config_path, 1);
-                let handle = w
-                    .load_dense(config_path, safetensors_paths, spec.model_id.clone())
+                    super::context_limit::profile_from_qwen3_5_config(&config_path, 1).or_else(
+                        || super::context_limit::profile_from_qwen4_exp_config(&config_path, 1),
+                    );
+                // In-situ quantisation (#320) reaches the worker load
+                // too, not just the CPU one — parsed here because
+                // `spec` does not cross the channel.
+                let isq = super::quant::parse_quant_string(spec.quant.as_deref())?;
+                // Built here rather than in the worker: the worker has
+                // no view of the config, and an `IsqCache` is a path,
+                // so it crosses the channel freely.
+                let isq_cache = self.isq_cache_for(
+                    &spec.model_id,
+                    &source_id.scheme,
+                    spec.quant.as_deref(),
+                    &safetensors_paths,
+                );
+                let experts_on_host = matches!(
+                    spec.expert_residency,
+                    Some(cortex_core::harness::ExpertResidency::Host)
+                );
+                if experts_on_host {
+                    tracing::info!(
+                        model = %spec.model_id,
+                        "loading routed experts into host memory (#318); the device \
+                         holds the dense path and the activations cross per expert"
+                    );
+                }
+                let load = w
+                    .load_dense(
+                        config_path,
+                        safetensors_paths,
+                        spec.model_id.clone(),
+                        isq,
+                        isq_cache,
+                        experts_on_host,
+                    )
                     .await
                     .map_err(|e| anyhow::anyhow!("worker load_dense: {e:#}"))?;
+                // Prefix snapshots (#11): the built model answers for
+                // itself on the way back, rather than the async side
+                // re-deriving the answer from `config.json`. See
+                // `DenseLoad` — the re-derivation is what silently
+                // dropped batched decode for qwen4_exp.
                 (
                     tokenizer_path,
                     None,
-                    Some(handle),
+                    Some(load.handle),
                     meta,
-                    snapshot_capable,
+                    load.supports_kv_snapshot,
                     context_profile,
                 )
             }
         } else {
             // CPU path: legacy spawn_blocking + Arc<Mutex<ModelArch>>.
-            let (tokenizer_path, arch) = if spec.quant.is_some() {
+            let (tokenizer_path, arch) = if self.quant_selects_a_gguf(spec, &source_id).await? {
                 self.load_arch_gguf(spec, &source_id, &device).await?
             } else {
                 self.load_arch_dense(spec, &source_id, &device).await?
@@ -3931,15 +4177,19 @@ impl Harness for CandleHarness {
         let prefill_rate = Arc::new(super::context_limit::PrefillRateEma::new());
         // Batched decode engine (#98): spawned when the operator raised
         // max_in_flight above 1 on a snapshot-capable worker-path model.
+        // The *effective* limit, not the host's: a model that overrides
+        // max_in_flight down to 1 is saying it cannot use a batch, and
+        // the engine must hear that (#318).
+        let effective_max_in_flight = self.admission_for(spec).max_in_flight;
         let engine = match (&worker, arch_handle) {
             (Some(w), Some(h))
                 if snapshot_capable
-                    && self.admission_cfg.max_in_flight > 1
+                    && effective_max_in_flight > 1
                     && super::engine::batching_enabled() =>
             {
                 tracing::info!(
                     model = %spec.model_id,
-                    max_slots = self.admission_cfg.max_in_flight,
+                    max_slots = effective_max_in_flight,
                     "batched decode engine enabled (#98)"
                 );
                 Some(super::engine::EngineHandle::spawn(
@@ -3948,7 +4198,7 @@ impl Harness for CandleHarness {
                         tokenizer: tokenizer.clone(),
                         reasoning_tokens: reasoning_tokens.clone(),
                         tool_call_tokens: tool_call_tokens.clone(),
-                        max_slots: self.admission_cfg.max_in_flight,
+                        max_slots: effective_max_in_flight,
                         backend: super::engine::BackendConfig::Single {
                             worker: Arc::clone(w),
                             handle: h,
@@ -4004,7 +4254,7 @@ impl Harness for CandleHarness {
             worker,
             arch_handle,
             inference_lock,
-            admission: super::admission::AdmissionController::new(&self.admission_cfg),
+            admission: super::admission::AdmissionController::new(&self.admission_for(spec)),
             reasoning_tokens,
             tool_call_tokens,
             chat_template,
@@ -4660,7 +4910,7 @@ impl CandleHarness {
             tokenizer,
             devices: devices.clone(),
             pool: StdArc::new(TMutex::new(pool)),
-            admission: super::admission::AdmissionController::new(&self.admission_cfg),
+            admission: super::admission::AdmissionController::new(&self.admission_for(spec)),
             leader_handle,
             leader_device: leader_device.clone(),
             poisoned: AtomicBool::new(false),
@@ -4675,6 +4925,13 @@ impl CandleHarness {
             image_token_id: vision_meta.image_token_id,
             image_grid_factor: vision_meta.image_grid_factor,
             spec: spec.clone(),
+            // Snapshot capability re-derived from `config.json`
+            // rather than asked of the loaded shard. Correct today —
+            // `TP_SUPPORTED_MODEL_TYPES` is exactly these two — but it
+            // is the same shape as the bug `DenseLoad` fixed on the
+            // single-GPU path, so whoever adds an arch here (#319)
+            // must add it in both places or lose batched decode
+            // silently.
             prefix_cache: self.new_prefix_cache(matches!(
                 config_model_type(&config_path).as_deref(),
                 Some(super::arch::qwen3_5::MODEL_TYPE)
@@ -4696,13 +4953,14 @@ impl CandleHarness {
         // max_in_flight above 1 on a snapshot-capable TP model. The
         // engine holds only a Weak — an idle engine never keeps an
         // unloaded model alive.
+        let effective_max_in_flight = self.admission_for(spec).max_in_flight;
         if tp_loaded.prefix_cache.is_some()
-            && self.admission_cfg.max_in_flight > 1
+            && effective_max_in_flight > 1
             && super::engine::batching_enabled()
         {
             tracing::info!(
                 model = %spec.model_id,
-                max_slots = self.admission_cfg.max_in_flight,
+                max_slots = effective_max_in_flight,
                 "batched decode engine enabled for TP model (#98)"
             );
             let handle = super::engine::EngineHandle::spawn(super::engine::EngineConfig {
@@ -4710,7 +4968,7 @@ impl CandleHarness {
                 tokenizer: tp_loaded.tokenizer.clone(),
                 reasoning_tokens: tp_loaded.reasoning_tokens.clone(),
                 tool_call_tokens: tp_loaded.tool_call_tokens.clone(),
-                max_slots: self.admission_cfg.max_in_flight,
+                max_slots: effective_max_in_flight,
                 backend: super::engine::BackendConfig::Tp {
                     tp: StdArc::downgrade(&tp_loaded),
                 },
@@ -6226,9 +6484,14 @@ struct VisionMeta {
 }
 
 /// Peek at `config.json` for its `model_type`. Best-effort — `None`
-/// on any read/parse error. The load path uses this to decide prefix-
-/// snapshot capability for worker-held models the async side can't
-/// inspect directly.
+/// on any read/parse error.
+///
+/// Only the TP load path still asks: the single-GPU path takes the
+/// arch's own answer off the load reply (`DenseLoad`) rather than
+/// re-deriving it here, which is what stopped the two from drifting.
+/// The TP path has the same hazard and not yet the same cure — see
+/// the note at its call site.
+#[cfg(feature = "cuda")]
 fn config_model_type(config_path: &std::path::Path) -> Option<String> {
     let text = std::fs::read_to_string(config_path).ok()?;
     let v: serde_json::Value = serde_json::from_str(&text).ok()?;
@@ -7515,6 +7778,45 @@ fn unix_subsec_nanos() -> u64 {
 
 #[cfg(test)]
 mod tests {
+
+    /// A model may bound its own concurrency below the host's (#318).
+    ///
+    /// Capacity is a property of the model as much as the machine. A
+    /// dense model amortises its weight reads across a batch; a sparse
+    /// MoE does not, because *n* concurrent tokens select up to `10n`
+    /// distinct experts — #309 measured eight streams returning 1.73x
+    /// the decode of one — and with the experts in host memory that
+    /// runs against DDR5 instead of HBM.
+    ///
+    /// The override has to reach the batch-engine gate, not just the
+    /// admission controller. Those are separate readers of the same
+    /// number, and a model bounded at 1 while an engine starts with 8
+    /// slots is the failure this function exists to prevent.
+    #[test]
+    fn a_model_may_lower_max_in_flight_below_the_host() {
+        let host = crate::config::AdmissionConfig {
+            max_in_flight: 8,
+            ..Default::default()
+        };
+
+        assert_eq!(admission_with_override(&host, Some(1)).max_in_flight, 1);
+        assert_eq!(admission_with_override(&host, Some(4)).max_in_flight, 4);
+        // Absent means the host decides, which is every model today.
+        assert_eq!(admission_with_override(&host, None).max_in_flight, 8);
+        // A model may also ask for more than the host runs; nothing
+        // about the number is a ceiling, it is a per-model answer.
+        assert_eq!(admission_with_override(&host, Some(16)).max_in_flight, 16);
+        // Zero would admit nothing at all. Ignored rather than served.
+        assert_eq!(admission_with_override(&host, Some(0)).max_in_flight, 8);
+
+        // Nothing else moves: queue depth, the wait deadline and the
+        // per-principal share are how a host treats callers, not a
+        // property of the model they asked for.
+        let overridden = admission_with_override(&host, Some(1));
+        assert_eq!(overridden.max_queue_depth, host.max_queue_depth);
+        assert_eq!(overridden.max_wait_secs, host.max_wait_secs);
+        assert_eq!(overridden.max_per_principal, host.max_per_principal);
+    }
     use super::super::context_limit::ContextProfile;
     use super::{
         ChatCompletionRequest, ChatMessage, MessageContent, apply_effort_kwarg,
@@ -8113,6 +8415,22 @@ mod tests {
         check_dense_config_supported(cfg, "Qwen/Qwen3-1.7B").expect("qwen3 should pass");
     }
 
+    #[test]
+    fn check_dense_config_accepts_qwen4_exp() {
+        // Qwen3.8-Flash-Next is a new architecture generation rather
+        // than a Qwen3 variant, so it has its own model_type and its own
+        // arm in `load_arch_dense`. The gate and that match must agree —
+        // a type in the supported set with no arm reaches the defensive
+        // "unrouted" bail instead of loading.
+        let cfg = r#"{
+            "model_type": "qwen4_exp",
+            "architectures": ["Qwen4ExpForConditionalGeneration"],
+            "text_config": {"hidden_size": 2560}
+        }"#;
+        check_dense_config_supported(cfg, "Qwen/Qwen3.8-Flash-Next")
+            .expect("qwen4_exp should be in the supported set");
+    }
+
     #[test]
     fn check_dense_config_rejects_unsupported_arch_with_clear_message() {
         // Use a deliberately-fake model_type so this test stays

crates/neuron/src/harness/context_limit.rs

@@ -252,6 +252,85 @@ pub fn profile_from_qwen3_5_config(config_path: &Path, world_size: u32) -> Optio
     })
 }
 
+/// Bytes of **QSA indexer** cache one token adds, per card (#307).
+///
+/// `qwen4_exp` keeps a second cache beside the main one: on every
+/// full-attention layer the indexer stores the raw index key,
+/// `indexer_kv_heads * indexer_head_dim` wide, in the compute dtype
+/// (`arch::qwen4_exp::qsa::Indexer::key_cache`). Keys only — the
+/// indexer scores blocks, it does not attend them — so there is no
+/// `2 *` here.
+///
+/// It does not shard: `indexer_kv_heads` is 1, so a TP rank holds the
+/// whole thing however wide the world is. At the checkpoint's 12
+/// full-attention layers x 1 head x 128 dims that is 3 KiB/token on
+/// top of the main cache's 24 — the surcharge the epic's original
+/// arithmetic omitted, and it grows with context exactly as the main
+/// cache does.
+pub fn indexer_bytes_per_token(
+    n_full_attn_layers: usize,
+    indexer_kv_heads: usize,
+    indexer_head_dim: usize,
+    dtype_bytes: usize,
+) -> u64 {
+    (n_full_attn_layers * indexer_kv_heads * indexer_head_dim * dtype_bytes) as u64
+}
+
+/// Build a [`ContextProfile`] from a `qwen4_exp` `config.json` — the
+/// sibling of [`profile_from_qwen3_5_config`] for the other hybrid
+/// arch. Returns `None` for any other `model_type` or an unparseable
+/// config.
+///
+/// Same shape of physics, one term more: the per-token cost is the
+/// main KV cache **plus** the QSA indexer's, because both grow per
+/// token and both come out of the same card. Advertising against the
+/// main cache alone would overstate what is servable by about an
+/// eighth, which is the class of error #281 is about.
+///
+/// `layer_types` is always populated by the time `from_config_json`
+/// returns (it fills the array from `full_attention_interval` when the
+/// checkpoint omits it), so counting it is authoritative here.
+pub fn profile_from_qwen4_exp_config(
+    config_path: &Path,
+    world_size: u32,
+) -> Option<ContextProfile> {
+    let text = std::fs::read_to_string(config_path).ok()?;
+    if !super::arch::qwen4_exp::config::is_qwen4_exp(&text) {
+        return None;
+    }
+    let cfg = super::arch::qwen4_exp::config::Config::from_config_json(&text).ok()?;
+    let tc = &cfg.text_config;
+    let n_full_attn_layers = tc
+        .layer_types
+        .iter()
+        .filter(|t| t.as_str() == "full_attention")
+        .count();
+    // A model with no full-attention layer has no growing cache, and
+    // the derivation divides by this. Nothing in the family looks like
+    // that, so treat it as a config we do not understand rather than
+    // inventing an unbounded context.
+    if n_full_attn_layers == 0 {
+        return None;
+    }
+    let kv_bytes_per_token_per_card = kv_bytes_per_token(
+        n_full_attn_layers,
+        tc.num_key_value_heads,
+        tc.head_dim,
+        KV_CACHE_DTYPE_BYTES,
+        world_size,
+    ) + indexer_bytes_per_token(
+        n_full_attn_layers,
+        tc.indexer_kv_heads,
+        tc.indexer_head_dim,
+        KV_CACHE_DTYPE_BYTES,
+    );
+    Some(ContextProfile {
+        max_position_embeddings: tc.max_position_embeddings,
+        kv_bytes_per_token_per_card,
+        world_size,
+    })
+}
+
 /// Round a token count down to a clean boundary so the advertised limit
 /// doesn't jitter by a handful of tokens as live VRAM / the throughput
 /// EMA wobble between polls.
@@ -375,9 +454,15 @@ pub fn derive_limit(
 
     // Observability (#126): every input and intermediate term, so a
     // surprising advertised limit is diagnosable from the journal
-    // instead of re-deriving by hand. DEBUG — this runs on every
-    // `GET /models` poll (a few lines per poll cycle).
-    tracing::debug!(
+    // instead of re-deriving by hand.
+    //
+    // TRACE, not debug: this runs on every `GET /models` poll and every
+    // input is fixed at load, so the line is byte-identical each time —
+    // about 8,600 a day per model at a 10s poll. It is worth having when
+    // reasoning about a context ceiling and worth nothing the rest of
+    // the time, and at debug it buries the lines that are not routine
+    // (#320).
+    tracing::trace!(
         max_pos = profile.max_position_embeddings,
         kv_bytes_per_token_per_card = profile.kv_bytes_per_token_per_card,
         free_tightest_mb,
@@ -418,6 +503,57 @@ pub fn derive_limit(
 #[cfg(test)]
 mod tests {
 
+    /// The advertised limit for `qwen4_exp` must price both caches.
+    ///
+    /// The main KV cache is 12 full-attention layers x 2 kv heads x 256
+    /// dims x 2 (K and V) x 2 bytes = 24 KiB/token, and the QSA indexer
+    /// adds 12 x 1 x 128 x 2 bytes = 3 KiB on top. Counting only the
+    /// first would advertise ~12.5% more context than the card can
+    /// hold — the failure mode #281 exists to stop, and the one the
+    /// epic's own arithmetic made before the indexer cache was found.
+    ///
+    /// Driven from the checkpoint's verbatim `config.json`, so a field
+    /// that changes shape upstream fails here rather than in
+    /// production.
+    #[test]
+    fn a_qwen4_exp_profile_counts_the_indexer_cache_too() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("config.json");
+        std::fs::write(&path, super::super::arch::qwen4_exp::config::SHIPPED).unwrap();
+
+        let profile = profile_from_qwen4_exp_config(&path, 1)
+            .expect("the shipped config.json must yield a profile");
+
+        assert_eq!(profile.max_position_embeddings, 262_144);
+        assert_eq!(
+            profile.kv_bytes_per_token_per_card,
+            24 * 1024 + 3 * 1024,
+            "main KV plus the indexer's second cache"
+        );
+    }
+
+    /// The two hybrid archs must not answer for each other: whichever
+    /// builder runs first has to decline a config it does not own, or
+    /// the load path's `or_else` chain silently prices one model with
+    /// the other's physics.
+    #[test]
+    fn the_profile_builders_decline_each_others_configs() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("config.json");
+        std::fs::write(&path, super::super::arch::qwen4_exp::config::SHIPPED).unwrap();
+        assert!(
+            profile_from_qwen3_5_config(&path, 1).is_none(),
+            "the qwen3_5 builder must decline a qwen4_exp config"
+        );
+
+        let path = dir.path().join("other.json");
+        std::fs::write(&path, r#"{"model_type": "qwen3_next"}"#).unwrap();
+        assert!(
+            profile_from_qwen4_exp_config(&path, 1).is_none(),
+            "the qwen4_exp builder must decline everything else"
+        );
+    }
+
     /// The deploy's smoke probe must not set the advertised context.
     ///
     /// 19 tokens in 130 ms reads as ~145 tok/s on a host that measures

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

@@ -13,11 +13,13 @@
 //! 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;
-use crate::harness::device_worker::jobs::{ArchHandle, ImageHandle, ImageInput, Job, KvSnapshotId};
+use crate::harness::device_worker::jobs::{
+    ArchHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
+};
 use crate::harness::image::ZImagePipeline;
 #[cfg(feature = "cuda")]
 use crate::harness::tp::TpLeaderModel;
@@ -131,17 +133,58 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
             } => {
                 let result = load_gguf_inner(&state.device, &gguf_path, &model_id)
                     .map(|arch| insert_arch(&mut state, Box::new(arch)));
+                // Same reasoning as `LoadImage`: release load-time
+                // transients before serving, and — on the failure
+                // path — give back whatever the partial load had
+                // allocated before it gave up. Without this a load
+                // that OOMs leaves its bytes in the pool, where the
+                // driver still counts them against the next load.
+                trim_device_pool(&state);
                 let _ = reply.send(result);
             }
             Job::LoadDense {
                 config_path,
                 safetensors_paths,
                 model_id,
+                quant,
+                isq_cache,
+                experts_on_host,
                 reply,
             } => {
-                let result =
-                    load_dense_inner(&state.device, &config_path, &safetensors_paths, &model_id)
-                        .map(|arch| insert_arch(&mut state, Box::new(arch)));
+                let result = load_dense_inner(
+                    &state.device,
+                    &config_path,
+                    &safetensors_paths,
+                    &model_id,
+                    quant,
+                    isq_cache.as_ref(),
+                    experts_on_host,
+                )
+                .map(|arch| {
+                    // Ask the model, once, while it is still in
+                    // hand. See `DenseLoad`.
+                    let supports_kv_snapshot = arch.supports_kv_snapshot();
+                    let handle = insert_arch(&mut state, Box::new(arch));
+                    DenseLoad {
+                        handle,
+                        supports_kv_snapshot,
+                    }
+                });
+                // As for `LoadGguf`. It matters more here: a dense
+                // load that runs out of VRAM part-way has typically
+                // allocated tens of gigabytes first, and the evicted
+                // model waiting to be restored needs them back.
+                trim_device_pool(&state);
+                // And the same on the host, for the same reason and
+                // then one more. A load with host-resident experts
+                // (#318) churns 73,728 allocations whether it finishes
+                // or not, and a load that *fails* drops its partial
+                // model here rather than through `DropArch` — so
+                // without this, the trim on unload never runs for it.
+                // Until an hour ago every qwen4_exp load on this
+                // hardware failed part-way, which is exactly the case
+                // that would have leaked.
+                trim_host_heap("load_dense");
                 let _ = reply.send(result);
             }
             Job::DropArch { handle, reply } => {
@@ -158,6 +201,17 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 // the end of the arm; calling drop() explicitly just
                 // makes the intent visible.
                 drop(removed);
+                // Return the freed blocks to the system rather than
+                // the stream-ordered pool, so the next load actually
+                // sees the VRAM — `DropImage` has done this since the
+                // benjy cold-swap, and an unload is an unload whatever
+                // modality asked for it.
+                if was_present {
+                    trim_device_pool(&state);
+                    // The experts may have been in host memory, and
+                    // dropping them is not the same as giving them back.
+                    trim_host_heap("drop_arch");
+                }
                 tracing::debug!(
                     device_index,
                     handle = handle.0,
@@ -374,6 +428,13 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 let was_present = removed.is_some();
                 drop(removed);
                 state.tp_kv_snapshots.retain(|(h, _), _| *h != handle);
+                // As `DropArch`: this is only the leader's shard —
+                // the other ranks free theirs in their own processes
+                // — but the leader's is the one sharing a card with
+                // whatever loads next.
+                if was_present {
+                    trim_device_pool(&state);
+                }
                 tracing::debug!(
                     device_index,
                     tp_handle = handle.0,
@@ -890,11 +951,78 @@ fn load_gguf_inner(
 }
 
 /// Load a dense safetensors model on the worker thread.
+/// Return freed heap to the operating system, and say how much.
+///
+/// The host-side counterpart to [`trim_device_pool`], and it exists
+/// for the same reason on the other side of the bus. Host-resident
+/// experts (#318) are ~1.4 MB per allocation, 1,536 per layer, 73,728
+/// for the model. glibc raises its mmap threshold adaptively as it
+/// observes frees of mmap'd blocks — up to 32 MB — so once it has
+/// grown past our allocation size the rest come from the arena and are
+/// never returned. Dropping the model frees them and RSS does not
+/// move.
+///
+/// Measured on beast: 64.6 GB of anonymous memory still resident ten
+/// minutes after the model was unloaded, on a host with 123 GB. The
+/// next load needing another 68 GB would take the node out, which is a
+/// worse failure than any VRAM exhaustion — it is not one model that
+/// stops, it is the machine.
+///
+/// Reports the delta because a trim that silently does nothing is the
+/// failure mode worth catching, and this session has shipped that bug
+/// twice already.
+fn trim_host_heap(context: &str) {
+    let rss = || -> Option<u64> {
+        let status = std::fs::read_to_string("/proc/self/status").ok()?;
+        let line = status.lines().find(|l| l.starts_with("VmRSS:"))?;
+        line.split_whitespace().nth(1)?.parse::<u64>().ok()
+    };
+    let before = rss();
+    // SAFETY: `malloc_trim` takes no pointers and has no preconditions
+    // beyond being linked against glibc, which the gate below ensures.
+    #[cfg(target_env = "gnu")]
+    let released = unsafe { libc::malloc_trim(0) } == 1;
+    #[cfg(not(target_env = "gnu"))]
+    let released = false;
+    match (before, rss()) {
+        (Some(b), Some(a)) => tracing::info!(
+            context,
+            released,
+            before_gib = b as f64 / (1024.0 * 1024.0),
+            after_gib = a as f64 / (1024.0 * 1024.0),
+            freed_gib = b.saturating_sub(a) as f64 / (1024.0 * 1024.0),
+            "host heap trimmed (#318)"
+        ),
+        _ => tracing::info!(context, released, "host heap trimmed (#318)"),
+    }
+}
+
+/// Log the cache summary when the load leaves scope, successful or
+/// not. A load that OOMs half way is precisely when the question "how
+/// much of that was re-quantised?" is worth answering.
+fn scopeguard_summary<'a>(
+    cache: Option<&'a crate::harness::isq_cache::IsqCache>,
+    model_id: &'a str,
+) -> impl Drop + 'a {
+    struct S<'a>(Option<&'a crate::harness::isq_cache::IsqCache>, &'a str);
+    impl Drop for S<'_> {
+        fn drop(&mut self) {
+            if let Some(c) = self.0 {
+                c.log_summary(self.1);
+            }
+        }
+    }
+    S(cache, model_id)
+}
+
 fn load_dense_inner(
     device: &candle_core::Device,
     config_path: &std::path::Path,
     safetensors_paths: &[std::path::PathBuf],
     model_id: &str,
+    quant: Option<candle_core::quantized::GgmlDType>,
+    isq_cache: Option<&crate::harness::isq_cache::IsqCache>,
+    experts_on_host: bool,
 ) -> anyhow::Result<ModelArch> {
     use anyhow::Context;
     use candle_core::DType;
@@ -903,6 +1031,10 @@ fn load_dense_inner(
     use candle_transformers::models::qwen3 as qwen3_dense;
     use candle_transformers::models::qwen3_moe as qwen3_moe_dense;
 
+    // Report what the cache did, however the load ends: a load that
+    // fails part-way is exactly when someone wants to know how much of
+    // it was cheap.
+    let _summary = scopeguard_summary(isq_cache, model_id);
     let cfg_text = std::fs::read_to_string(config_path).context("read config.json")?;
     crate::harness::candle::check_dense_config_supported(&cfg_text, model_id)?;
     // Peek at model_type to choose the family before the typed
@@ -988,6 +1120,66 @@ fn load_dense_inner(
                 .context("build Qwen3-Next dense model")?;
             Ok(ModelArch::Qwen3_5Dense(model))
         }
+        "qwen4_exp" => {
+            // Qwen3.8-Flash-Next. Mirrors the CPU path in
+            // `candle::load_arch_dense` — same config type, same
+            // sharded backend, same in-situ quantisation.
+            //
+            // ISQ is not optional for this architecture: the routed
+            // experts are 241.6 GB at BF16 (measured from the
+            // checkpoint's own headers), against 32.6 GB of card. A
+            // load without a `quant` will fail on memory rather than
+            // on anything this code can say (#315, #252).
+            let cfg = crate::harness::arch::qwen4_exp::config::Config::from_config_json(&cfg_text)
+                .context("parse qwen4_exp config.json")?;
+            let sharded_vb = unsafe {
+                candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                    safetensors_paths,
+                    dtype,
+                    device,
+                )
+                .context("build ShardedVarBuilder for qwen4_exp")?
+            };
+            // A second view of the same checkpoint, bound to host
+            // memory, so host-resident experts are read straight there
+            // rather than through 5.03 GB of VRAM per layer (#318).
+            // Built once for the model: it maps the same files, and
+            // mapping is not what costs.
+            let host_vb = if experts_on_host {
+                Some(unsafe {
+                    candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                        safetensors_paths,
+                        dtype,
+                        &candle_core::Device::Cpu,
+                    )
+                    .context("build host ShardedVarBuilder for qwen4_exp experts")?
+                })
+            } else {
+                None
+            };
+            let model = crate::harness::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
+                &cfg,
+                dtype,
+                device,
+                &sharded_vb,
+                &crate::harness::arch::qwen4_exp::LoadOptions {
+                    quant,
+                    safetensors_paths,
+                    isq_cache,
+                    experts_vb: host_vb.as_ref(),
+                    // Host residency (#318) is what makes a model whose
+                    // experts exceed VRAM loadable; everything else in
+                    // the model still lands on the device.
+                    experts_onto: if experts_on_host {
+                        &candle_core::Device::Cpu
+                    } else {
+                        device
+                    },
+                },
+            )
+            .context("build qwen4_exp model")?;
+            Ok(ModelArch::Qwen4Exp(Box::new(model)))
+        }
         other => anyhow::bail!(
             "unrouted supported model_type '{other}' — \
              DENSE_SUPPORTED_MODEL_TYPES and load_dense_inner \
@@ -1045,7 +1237,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 +1371,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 +1404,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 +1528,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 +1554,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();
@@ -1671,3 +1861,49 @@ fn drain_poisoned(job: Job, device_index: u32) {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Every architecture the preflight admits must have a loader arm
+    /// on **this** path, not only on the CPU one in
+    /// `candle::load_arch_dense`.
+    ///
+    /// There are two dense loaders — the worker's (every CUDA load)
+    /// and candle's (the CPU fallback) — and one list gating both.
+    /// `qwen4_exp` was added to the list and to the CPU loader, so it
+    /// passed preflight and then hit the "unrouted supported
+    /// model_type" bail here: on a GPU host the model was
+    /// unloadable, in about 200 ms, having read nothing. The bail's
+    /// own comment said the two must stay in sync; nothing checked.
+    ///
+    /// A bare `{"model_type": ...}` reaches the arm and then fails on
+    /// the typed parse, which is fine — this asserts only that an arm
+    /// exists to fail in.
+    #[test]
+    fn every_supported_dense_model_type_has_a_loader_arm() {
+        let dir = tempfile::tempdir().unwrap();
+        for model_type in crate::harness::candle::DENSE_SUPPORTED_MODEL_TYPES {
+            let path = dir.path().join(format!("{model_type}.json"));
+            std::fs::write(&path, format!(r#"{{"model_type": "{model_type}"}}"#)).unwrap();
+            let msg = match load_dense_inner(
+                &candle_core::Device::Cpu,
+                &path,
+                &[],
+                model_type,
+                None,
+                None,
+                false,
+            ) {
+                Ok(_) => panic!("a bare config must not build a model"),
+                Err(e) => format!("{e:#}"),
+            };
+            assert!(
+                !msg.contains("unrouted"),
+                "`{model_type}` is in DENSE_SUPPORTED_MODEL_TYPES but load_dense_inner \
+                 has no arm for it: {msg}"
+            );
+        }
+    }
+}

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

@@ -19,6 +19,27 @@ use tokio::sync::oneshot;
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ArchHandle(pub u64);
 
+/// What a dense load hands back: the handle, plus what the arch it
+/// built says about itself.
+///
+/// The async side cannot ask the model anything — it lives in the
+/// worker's slab and never crosses the channel — so anything the
+/// serving path needs to know has to ride back on the reply. Prefix
+/// snapshot support (#11) is the first such fact, and it is here
+/// rather than re-derived from `config.json` because that is exactly
+/// how it went wrong: the async side matched a list of `model_type`
+/// strings that `ModelArch::supports_kv_snapshot` had since grown
+/// past, so `qwen4_exp` loaded with batched decode silently off.
+/// A capability the arch answers for itself cannot drift from the
+/// arch.
+#[derive(Debug, Clone, Copy)]
+pub struct DenseLoad {
+    pub handle: ArchHandle,
+    /// [`crate::harness::candle::ModelArch::supports_kv_snapshot`], as
+    /// answered by the model that was actually built.
+    pub supports_kv_snapshot: bool,
+}
+
 /// Opaque handle to a `TpLeaderModel` stored in the worker thread's
 /// state slab. Same shape as [`ArchHandle`] but in a separate
 /// namespace so the two slabs can coexist without ambiguity. Phase 3
@@ -107,7 +128,16 @@ pub enum Job {
         config_path: PathBuf,
         safetensors_paths: Vec<PathBuf>,
         model_id: String,
-        reply: oneshot::Sender<Result<ArchHandle>>,
+        /// In-situ quantisation (#320): the catalogue's `quant`,
+        /// already parsed. `None` loads at the checkpoint dtype.
+        /// Archs that cannot quantise in situ ignore it.
+        quant: Option<candle_core::quantized::GgmlDType>,
+        /// Where to cache the quantised weights (#322). Built on the
+        /// async side, which is the only side that can see the config.
+        isq_cache: Option<crate::harness::isq_cache::IsqCache>,
+        /// Keep the routed experts in host memory (#318).
+        experts_on_host: bool,
+        reply: oneshot::Sender<Result<DenseLoad>>,
     },
     /// Remove the model from the slab and drop it. The `Drop` runs on
     /// the worker thread so CUDA tensors release their memory on the
@@ -163,7 +193,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/device_worker/mod.rs

@@ -51,7 +51,7 @@ use tokio::sync::oneshot;
 
 #[cfg(feature = "cuda")]
 pub use jobs::TpHandle;
-pub use jobs::{ArchHandle, ImageHandle, Job, KvSnapshotId};
+pub use jobs::{ArchHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
 
 /// Errors returned by `DeviceWorkerHandle` submit methods.
 #[derive(Debug, thiserror::Error)]
@@ -221,7 +221,10 @@ impl DeviceWorkerHandle {
         config_path: std::path::PathBuf,
         safetensors_paths: Vec<std::path::PathBuf>,
         model_id: String,
-    ) -> Result<ArchHandle, WorkerError> {
+        quant: Option<candle_core::quantized::GgmlDType>,
+        isq_cache: Option<crate::harness::isq_cache::IsqCache>,
+        experts_on_host: bool,
+    ) -> Result<DenseLoad, WorkerError> {
         if self.poisoned.load(Ordering::Acquire) {
             return Err(WorkerError::Poisoned {
                 device_index: self.device_index,
@@ -233,6 +236,9 @@ impl DeviceWorkerHandle {
                 config_path,
                 safetensors_paths,
                 model_id,
+                quant,
+                isq_cache,
+                experts_on_host,
                 reply: reply_tx,
             })
             .map_err(|_| WorkerError::Gone {
@@ -1243,9 +1249,13 @@ mod tests {
                 fixture.join("config.json"),
                 vec![fixture.join("model.safetensors")],
                 "qwen3_next-tiny".into(),
+                None,
+                None,
+                false,
             )
             .await
-            .expect("load tiny fixture");
+            .expect("load tiny fixture")
+            .handle;
 
         let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 3, 2, 5, 6]];
         let steps: [&[u32]; 3] = [&[11, 12, 13], &[9, 8, 7], &[21, 22, 23]];

crates/neuron/src/harness/engine.rs

@@ -1261,9 +1261,13 @@ mod tests {
                 fixture.join("config.json"),
                 vec![fixture.join("model.safetensors")],
                 "qwen3_next-tiny".into(),
+                None,
+                None,
+                false,
             )
             .await
-            .expect("load fixture");
+            .expect("load fixture")
+            .handle;
 
         let admission_cfg = AdmissionConfig {
             max_in_flight: 3,

crates/neuron/src/harness/isq_cache.rs

@@ -0,0 +1,451 @@
+//! On-disk cache for in-situ quantisation (#322).
+//!
+//! Quantising a model's weights at load is a pure function of the
+//! checkpoint bytes, the requested type, and the code that does it.
+//! None of those change between loads, and yet neuron recomputed the
+//! answer every time and threw it away on unload. For `qwen4_exp` that
+//! is **~32 s per decoder layer, about 26 minutes for a full load**,
+//! measured on beast — paid again on every restart and every eviction
+//! that is later reversed.
+//!
+//! So the quantised tensors are written out once and read back after.
+//!
+//! ## The key must include our own version
+//!
+//! `(checkpoint revision, quant name)` is not enough. `bf1e80ed`
+//! changed what `quant = "q4k"` *means* for a 640-wide row — k-quants
+//! cannot block it, so it now falls back to a 32-block type — and an
+//! artifact written before that commit and read after it would be
+//! silently wrong in a way no checksum of the inputs could catch.
+//! [`QUANTISER_VERSION`] is part of the path, and bumping it orphans
+//! every stale artifact rather than misreading one.
+//!
+//! ## What is cached, and what is not
+//!
+//! One GGUF per layer, holding that layer's quantised experts. GGUF
+//! because it stores a **type per tensor**, which is exactly what the
+//! per-row fallback produces — a layer whose `down_proj` is `Q5_0`
+//! while its `gate`/`up` are `Q4K` is representable without inventing
+//! a container for it — and because neuron already reads the format.
+//!
+//! Per layer rather than per model because it makes the cache
+//! **incrementally useful**: a load that runs out of VRAM at layer 16
+//! has still written sixteen layers, and the next attempt starts
+//! there. On hardware where the model does not yet fit at all (#318),
+//! that is the difference between the cache being useless and being
+//! the thing that makes iteration bearable.
+//!
+//! ## Where it lives
+//!
+//! Inside the source's own cache root — `<hf cache>/neuron-isq/…` —
+//! beside the weights it is derived from, and not behind a config key
+//! naming a path.
+//!
+//! That is deliberate and it is the fleet's existing design. beast
+//! already keeps its hot model weights on NVMe by symlinking them out
+//! of the spinning `/archive3` cache into `/archive1`, with no config
+//! mapping paths to mounts. An operator who wants these artifacts on
+//! faster storage does the same thing to one directory. A second,
+//! differently-shaped way to express the same intent would be a second
+//! place to get it wrong, and the first one already works.
+//!
+//! Size is the thing to know before symlinking: an artifact is roughly
+//! the served size of the model's quantised weights — ~73 GB for
+//! `qwen4_exp` at q4k.
+//!
+//! ## It is an optimisation, and never load-bearing
+//!
+//! Every failure here is swallowed: a miss, a corrupt file, a full
+//! disk, a directory nobody can write. The load proceeds by
+//! quantising, exactly as it did before this module existed. A cache
+//! that can fail a load is worse than no cache.
+
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+
+use candle_core::Device;
+use candle_core::quantized::{QTensor, gguf_file};
+
+/// Bump when the *meaning* of a quant name changes: a different
+/// fallback table, a different rounding, a different layout. Artifacts
+/// written by an older quantiser are then ignored rather than
+/// misread, and get overwritten in place as layers are re-quantised.
+///
+/// 1: the initial format. Per-layer GGUF, tensors named
+///    `e{index}.{gate,up,down}`, llama.cpp's `tensor_type_fallback`
+///    table for rows that do not divide the requested block size.
+pub const QUANTISER_VERSION: u32 = 1;
+
+/// Refuse a write that would take more than this share of the free
+/// space.
+///
+/// The rule is about *the write*, not the filesystem. An earlier
+/// version required 2% of the volume to remain free, which read
+/// sensibly and was wrong in the case that matters least and happens
+/// most: on a filesystem already below that line, it refused a
+/// six-megabyte write that was in no way responsible. CI found it
+/// immediately.
+///
+/// Bounding each artifact by half the remaining space says the thing
+/// actually meant — no single artifact eats the disk — at every size,
+/// and needs to know nothing about how big the volume is. An
+/// artifact for one 180 B model is ~73 GB, so on a volume with 100 GB
+/// left it is refused and on beast's 2 TB of free `/archive3` it is
+/// not. A few megabytes is never refused, because a few megabytes is
+/// never the problem.
+const MAX_SHARE_OF_FREE: f64 = 0.5;
+
+/// A cache scoped to one (model, revision, quant, quantiser version).
+pub struct IsqCache {
+    dir: PathBuf,
+    hits: AtomicUsize,
+    misses: AtomicUsize,
+    read: AtomicU64,
+    written: AtomicU64,
+}
+
+impl IsqCache {
+    /// `revision` should identify the checkpoint bytes — the snapshot
+    /// directory name, for an hf-hub cache. `None` disables the cache
+    /// rather than guessing, because a wrong revision is the one
+    /// failure mode that produces a confidently wrong model.
+    pub fn new(base: &Path, model_id: &str, revision: Option<&str>, quant: &str) -> Option<Self> {
+        let revision = revision?;
+        let safe = |s: &str| {
+            s.chars()
+                .map(|c| {
+                    if c.is_ascii_alphanumeric() || c == '-' || c == '.' {
+                        c
+                    } else {
+                        '-'
+                    }
+                })
+                .collect::<String>()
+        };
+        let dir = base
+            .join(safe(model_id))
+            .join(safe(revision))
+            .join(format!("{}-v{QUANTISER_VERSION}", safe(quant)));
+        Some(Self {
+            dir,
+            hits: AtomicUsize::new(0),
+            misses: AtomicUsize::new(0),
+            read: AtomicU64::new(0),
+            written: AtomicU64::new(0),
+        })
+    }
+
+    /// Read a layer's quantised tensors onto `device`.
+    ///
+    /// `None` for a miss and for anything that goes wrong reading —
+    /// a truncated artifact from an interrupted write reads as absent,
+    /// and is rewritten by the load that follows.
+    pub fn load(&self, key: &str, device: &Device) -> Option<HashMap<String, QTensor>> {
+        let path = self.dir.join(format!("{key}.gguf"));
+        let Ok(mut file) = std::fs::File::open(&path) else {
+            self.misses.fetch_add(1, Ordering::Relaxed);
+            return None;
+        };
+        let content = match gguf_file::Content::read(&mut file) {
+            Ok(c) => c,
+            Err(e) => {
+                self.misses.fetch_add(1, Ordering::Relaxed);
+                tracing::warn!(path = %path.display(), error = %e,
+                    "isq cache: unreadable artifact, re-quantising this layer");
+                return None;
+            }
+        };
+        let names: Vec<String> = content.tensor_infos.keys().cloned().collect();
+        let mut out = HashMap::with_capacity(names.len());
+        for name in names {
+            match content.tensor(&mut file, &name, device) {
+                Ok(t) => {
+                    out.insert(name, t);
+                }
+                Err(e) => {
+                    self.misses.fetch_add(1, Ordering::Relaxed);
+                    tracing::warn!(path = %path.display(), tensor = %name, error = %e,
+                        "isq cache: artifact tensor failed to read, re-quantising this layer");
+                    return None;
+                }
+            }
+        }
+        self.hits.fetch_add(1, Ordering::Relaxed);
+        self.read.fetch_add(
+            out.values()
+                .map(|t| t.storage_size_in_bytes() as u64)
+                .sum::<u64>(),
+            Ordering::Relaxed,
+        );
+        Some(out)
+    }
+
+    /// Write a layer's quantised tensors.
+    ///
+    /// Through a temporary file and a rename, so an interrupted write
+    /// cannot leave a half-artifact that a later load would read as
+    /// real. Errors are logged and swallowed.
+    /// Write a layer, refusing one that would take an unreasonable
+    /// share of the disk. The production entry point: every failure is
+    /// logged and swallowed.
+    pub fn store(&self, key: &str, tensors: &[(&str, &QTensor)]) {
+        let bytes: u64 = tensors
+            .iter()
+            .map(|(_, t)| t.storage_size_in_bytes() as u64)
+            .sum();
+        if let Some(free) = free_bytes(&self.dir)
+            && too_large_for(bytes, free)
+        {
+            tracing::warn!(
+                dir = %self.dir.display(), key,
+                gib = bytes as f64 / (1024.0 * 1024.0 * 1024.0),
+                free_gib = free as f64 / (1024.0 * 1024.0 * 1024.0),
+                "isq cache: refusing a write that would take over {:.0}% of the free space",
+                MAX_SHARE_OF_FREE * 100.0,
+            );
+            return;
+        }
+        match self.try_store(key, tensors) {
+            Ok(()) => {
+                self.written.fetch_add(bytes, Ordering::Relaxed);
+            }
+            Err(e) => tracing::warn!(dir = %self.dir.display(), key, error = %e,
+                "isq cache: could not write this layer; the load continues uncached"),
+        }
+    }
+
+    /// Say, once per load, what the cache actually did.
+    ///
+    /// Without this the journal shows only that the cache was enabled,
+    /// which is the shape of problem #321 is about: every light green
+    /// and no way to tell from the outside whether the thing worked.
+    /// A load that served nothing from cache and a load that served
+    /// everything looked identical, and the difference between them is
+    /// eight minutes.
+    pub fn log_summary(&self, model_id: &str) {
+        let gib = |b: u64| b as f64 / (1024.0 * 1024.0 * 1024.0);
+        tracing::info!(
+            model = model_id,
+            dir = %self.dir.display(),
+            layers_from_cache = self.hits.load(Ordering::Relaxed),
+            layers_quantised = self.misses.load(Ordering::Relaxed),
+            read_gib = gib(self.read.load(Ordering::Relaxed)),
+            written_gib = gib(self.written.load(Ordering::Relaxed)),
+            "isq cache: load complete (#322)"
+        );
+    }
+
+    /// Write a layer. Nothing else — no policy, and in particular no
+    /// question asked of the filesystem about its size.
+    ///
+    /// That separation is the point. The disk check lives in
+    /// [`Self::store`], so the round-trip tests exercise the write and
+    /// the read without their result depending on how full the machine
+    /// running them happens to be. A unit test that consults the free
+    /// space of whatever host executes it is not testing this code.
+    pub(crate) fn try_store(&self, key: &str, tensors: &[(&str, &QTensor)]) -> anyhow::Result<()> {
+        std::fs::create_dir_all(&self.dir)?;
+        let final_path = self.dir.join(format!("{key}.gguf"));
+        let tmp = self.dir.join(format!(".{key}.{}.tmp", std::process::id()));
+        let write = (|| -> anyhow::Result<()> {
+            let mut f = std::fs::File::create(&tmp)?;
+            gguf_file::write(&mut f, &[], tensors)?;
+            Ok(())
+        })();
+        if let Err(e) = write {
+            // A partial artifact is dead weight that nothing will ever
+            // clean up — and on the failure this is most likely to see,
+            // a full disk, it is dead weight made of the last free
+            // bytes.
+            let _ = std::fs::remove_file(&tmp);
+            return Err(e);
+        }
+        std::fs::rename(&tmp, &final_path)?;
+        Ok(())
+    }
+
+    pub fn dir(&self) -> &Path {
+        &self.dir
+    }
+}
+
+/// Whether a write of `bytes` takes an unreasonable share of `free`.
+///
+/// Split out from the write so the rule can be tested without needing
+/// a filesystem in a particular state — which is precisely what made
+/// the previous rule's defect invisible until CI ran it.
+fn too_large_for(bytes: u64, free: u64) -> bool {
+    bytes as f64 > free as f64 * MAX_SHARE_OF_FREE
+}
+
+/// Free bytes on the filesystem holding `path`, or `None` if it cannot
+/// be determined — in which case the caller writes anyway rather than
+/// refusing on a guess. `-P` is what guarantees one line per
+/// filesystem, so a long device name cannot wrap the fields apart.
+fn free_bytes(path: &Path) -> Option<u64> {
+    let out = std::process::Command::new("df")
+        .arg("-kP")
+        .arg(path)
+        .output()
+        .ok()?;
+    let text = String::from_utf8(out.stdout).ok()?;
+    let kb = text.lines().nth(1)?.split_whitespace().nth(3)?;
+    kb.parse::<u64>().ok().map(|kb| kb * 1024)
+}
+
+/// The revision an hf-hub-cached checkpoint was resolved at, taken
+/// from the snapshot directory its files live in.
+///
+/// `.../models--org--name/snapshots/<revision>/model-00001-of-N.safetensors`
+/// — the parent directory *is* the revision, which is what makes it a
+/// safe cache key: two checkpoints with different bytes cannot share
+/// one. Returns `None` for any layout that is not that, so an
+/// unrecognised path disables the cache instead of colliding.
+pub fn revision_from_paths(paths: &[PathBuf]) -> Option<&str> {
+    let first = paths.first()?;
+    let dir = first.parent()?;
+    let name = dir.file_name()?.to_str()?;
+    let looks_like_snapshot = dir.parent()?.file_name()?.to_str()? == "snapshots";
+    (looks_like_snapshot && !name.is_empty()).then_some(name)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Tensor;
+    use candle_core::quantized::GgmlDType;
+
+    /// A cached layer must dequantise to the same numbers the fresh
+    /// quantisation produced.
+    ///
+    /// This is the whole contract. A cache that returns *something*
+    /// plausible is worse than no cache, because the model loads, it
+    /// serves, and its weights are quietly not the ones it was asked
+    /// for. Round-tripping the dequantised values rather than the
+    /// bytes is deliberate: it is what the model actually consumes.
+    #[test]
+    fn a_cached_tensor_dequantises_to_what_was_stored() {
+        let dev = Device::Cpu;
+        let dir = tempfile::tempdir().unwrap();
+        let cache = IsqCache::new(dir.path(), "org/model", Some("rev0"), "q4k").unwrap();
+
+        // 256-wide so the k-quant can block it at all.
+        let src = Tensor::rand(-1f32, 1f32, (8, 256), &dev).unwrap();
+        let q = QTensor::quantize(&src, GgmlDType::Q4K).unwrap();
+        let before: Vec<f32> = q
+            .dequantize(&dev)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        cache
+            .try_store("layer-0", &[("e0.gate", &q)])
+            .expect("the write must succeed, and say why if it does not");
+        let got = cache
+            .load("layer-0", &dev)
+            .expect("the artifact just written");
+        let after: Vec<f32> = got["e0.gate"]
+            .dequantize(&dev)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        assert_eq!(
+            before, after,
+            "a cached tensor changed value on the round trip"
+        );
+        assert_eq!(got["e0.gate"].dtype(), GgmlDType::Q4K, "and kept its type");
+    }
+
+    /// The key separates everything that changes the answer.
+    ///
+    /// The quantiser version is the one that is easy to leave out and
+    /// impossible to detect: `bf1e80ed` changed what `q4k` means for a
+    /// 640-wide row, and an artifact written before it is wrong after
+    /// it in a way no checksum of the *inputs* would reveal.
+    #[test]
+    fn the_key_separates_model_revision_quant_and_quantiser() {
+        let root = std::path::Path::new("/cache");
+        let base = IsqCache::new(root, "org/model", Some("rev0"), "q4k").unwrap();
+        let others = [
+            IsqCache::new(root, "org/other", Some("rev0"), "q4k").unwrap(),
+            IsqCache::new(root, "org/model", Some("rev1"), "q4k").unwrap(),
+            IsqCache::new(root, "org/model", Some("rev0"), "q6k").unwrap(),
+        ];
+        for other in &others {
+            assert_ne!(base.dir(), other.dir());
+        }
+        assert!(
+            base.dir()
+                .to_str()
+                .unwrap()
+                .contains(&format!("v{QUANTISER_VERSION}")),
+            "the quantiser version must be in the path: {}",
+            base.dir().display()
+        );
+    }
+
+    /// No revision, no cache. A key that might collide across
+    /// checkpoints is worse than recomputing.
+    #[test]
+    fn an_unidentifiable_revision_disables_the_cache() {
+        assert!(IsqCache::new(std::path::Path::new("/cache"), "org/model", None, "q4k").is_none());
+
+        // An hf-hub snapshot path yields its revision; anything else
+        // does not, rather than yielding a directory name that happens
+        // to be there.
+        let hf = PathBuf::from(
+            "/cache/models--org--model/snapshots/abc123/model-00001-of-00002.safetensors",
+        );
+        assert_eq!(revision_from_paths(&[hf]), Some("abc123"));
+        assert_eq!(
+            revision_from_paths(&[PathBuf::from("/tmp/model.safetensors")]),
+            None
+        );
+        assert_eq!(revision_from_paths(&[]), None);
+    }
+
+    /// The space guard bounds the *write*, not the filesystem.
+    ///
+    /// The rule it replaced required a fixed share of the volume to
+    /// remain free, which refused a six-megabyte write on a disk that
+    /// was already full — punishing the write least responsible for
+    /// the state. These cases are the ones that matter: a 73 GB
+    /// artifact must not land on a volume with 100 GB left, and a few
+    /// megabytes must always be allowed, however full things are.
+    #[test]
+    fn the_space_guard_bounds_the_write_not_the_volume() {
+        let gib = 1024 * 1024 * 1024;
+
+        // A 73 GB artifact against beast's ~2 TB of free /archive3.
+        assert!(!too_large_for(73 * gib, 2000 * gib));
+        // The same artifact where it would eat most of what is left.
+        assert!(too_large_for(73 * gib, 100 * gib));
+        // Six megabytes on a nearly-full disk: never the problem.
+        assert!(!too_large_for(6 * 1024 * 1024, 200 * 1024 * 1024));
+        // Exactly half is allowed; past half is not.
+        assert!(!too_large_for(50, 100));
+        assert!(too_large_for(51, 100));
+        // And nothing is writable onto nothing.
+        assert!(too_large_for(1, 0));
+    }
+
+    /// A miss and a corrupt artifact are the same thing to the caller:
+    /// re-quantise. Neither may propagate as an error.
+    #[test]
+    fn a_missing_or_corrupt_artifact_reads_as_absent() {
+        let dev = Device::Cpu;
+        let dir = tempfile::tempdir().unwrap();
+        let cache = IsqCache::new(dir.path(), "org/model", Some("rev0"), "q4k").unwrap();
+        assert!(cache.load("layer-0", &dev).is_none(), "a miss");
+
+        std::fs::create_dir_all(cache.dir()).unwrap();
+        std::fs::write(cache.dir().join("layer-0.gguf"), b"not a gguf file at all").unwrap();
+        assert!(cache.load("layer-0", &dev).is_none(), "a corrupt artifact");
+    }
+}

crates/neuron/src/harness/mod.rs

@@ -9,9 +9,11 @@ pub mod context_limit;
 pub mod device_worker;
 pub mod engine;
 pub mod image;
+pub mod isq_cache;
 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/preflight.rs

@@ -403,6 +403,8 @@ mod tests {
             devices: None,
             sampling: None,
             preserve_thinking: None,
+            expert_residency: None,
+            max_in_flight: None,
         }
     }
 

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))
-}

crates/neuron/tests/activation.rs

@@ -31,6 +31,8 @@ async fn test_load_default_models_skips_unknown_harness() {
             devices: None,
             sampling: None,
             preserve_thinking: None,
+            expert_residency: None,
+            max_in_flight: None,
         },
         ModelSpec {
             model_id: "model-b".into(),
@@ -40,6 +42,8 @@ async fn test_load_default_models_skips_unknown_harness() {
             devices: None,
             sampling: None,
             preserve_thinking: None,
+            expert_residency: None,
+            max_in_flight: None,
         },
     ];
 
@@ -101,6 +105,8 @@ async fn test_load_default_models_skipped_on_driver_mismatch() {
         devices: None,
         sampling: None,
         preserve_thinking: None,
+        expert_residency: None,
+        max_in_flight: None,
     }];
     let activation = ActivationTracker::new(&specs);
     let reason = "host NVIDIA driver/library mismatch (userspace NVML 580.159 vs loaded \
@@ -204,6 +210,8 @@ fn qwen_spec() -> ModelSpec {
         devices: None,
         sampling: None,
         preserve_thinking: None,
+        expert_residency: None,
+        max_in_flight: None,
     }
 }
 

crates/neuron/tests/candle_lifecycle.rs

@@ -55,6 +55,8 @@ async fn test_candle_qwen3_load_unload_lifecycle() {
         devices: Some(vec![0]),
         sampling: None,
         preserve_thinking: None,
+        expert_residency: None,
+        max_in_flight: None,
     };
 
     registry

crates/neuron/tests/fixtures/numerical/README.md

@@ -12,6 +12,82 @@ HuggingFace reference" claim to checked-in numbers.
 | `qwen3_5-0.8b-text` | Qwen/Qwen3.5-0.8B | text (>64-token prompt → chunked GDN prefill) | f32 | `text_logits_match_reference` |
 | `qwen3_5-0.8b-vision` | Qwen/Qwen3.5-0.8B | 448×448 synthetic image + prompt | f32 | `vision_tower_and_logits_match_reference` |
 | `qwen3_6-27b-text` | Qwen/Qwen3.6-27B | text | bf16 | manual (see below) |
+| `qwen4_exp-tiny` | random weights, real architecture | text, 12 tokens, no eos | f32 | `qwen4_exp_logits_match_reference` |
+| `qwen4_exp-tiny-long` | random weights, real architecture | text, 80 tokens, eos at 37 | f32 | `qwen4_exp_long_prompt_logits_match_reference` |
+
+## `qwen4_exp-tiny` is different, and runs in CI
+
+The others replay a real snapshot and self-skip without
+`NEURON_REF_MODEL_PATH`. This one is 44 KB of random weights in the
+real architecture — 4 layers (3 linear-attention, 1 QSA), PLE on layer
+1, 2 experts — generated *by* `transformers` itself via
+[`script/generate_qwen4_exp_fixture.py`](../../../../../script/generate_qwen4_exp_fixture.py).
+So it is checked in, needs no weights, and runs on every push.
+
+That matters because the whole `qwen4_exp` architecture was ported from
+a written spec (#308) and tested against our own reading of it. This is
+the only thing in the tree that compares it to the implementation it
+was ported from (#323).
+
+Two properties of the fixture are load-bearing, and both were learned
+by getting them wrong:
+
+- **Weights are `uniform(-0.8, 0.8)`, not `uniform(-0.1, 0.1)`.** Near
+  the origin `silu(x) ~ x/2`, so the SwiGLU is nearly symmetric and
+  reading the fused expert's gate and up halves the wrong way round
+  moved the logits by 2.4e-4 — inside any tolerance worth writing. At
+  the larger scale the same mutation moves them by 2.7.
+- **The generator refuses to emit a fixture containing a QSA selection
+  tie.** A block scores exactly 0.0 whenever relu clamps every head, so
+  ties are common with random weights; when the k-th and (k+1)-th
+  scores are equal, which block is dropped is decided by `topk`'s
+  ordering of equal elements, which torch does not specify and whose
+  CPU and CUDA kernels need not agree. Seeds 0, 1 and 4 produce ties;
+  seed 2 (the default) does not. See "the tie-break" below.
+
+Regenerate with the reference from `transformers` **main** — the 5.9.0
+release does not carry `qwen4_exp`:
+
+```sh
+python -m venv --system-site-packages .venv
+.venv/bin/pip install --no-deps "transformers @ git+https://github.com/huggingface/transformers@main"
+.venv/bin/pip install "tokenizers>=0.23.1,<0.24" "safetensors>=0.8.0"
+.venv/bin/python script/generate_qwen4_exp_fixture.py \
+    --config crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny/config.json \
+    --out crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny
+```
+
+### Why there are two cases
+
+They are not the same test at two sizes. Each reaches a path the other
+structurally cannot, and a mutation demonstrates it:
+
+- **80 tokens** puts the GatedDeltaNet layers on the *chunked* prefill
+  algorithm (#23), which engages at 64. That is a different
+  factorisation of the same recurrence, not the same code with more
+  input — which is also why the long case's bound is 1e-4 against the
+  short case's 1e-5 (observed 1.3e-5 vs 2e-6). It also gives QSA up to
+  40 blocks against a budget of 4, so selection discards nearly
+  everything instead of one block at the boundary.
+- **The eos at position 37** is a segment boundary the n-gram window
+  must refuse to read across. Making `shift_right_ignore_eos` ignore
+  segments entirely leaves the short fixture at 0.000002 — it contains
+  no eos, so there is nothing to get wrong — and moves the long one by
+  **0.67**.
+
+The short case deliberately has no eos: a fixture that happens to
+contain one is testing segment handling by accident, and would make the
+two cases redundant.
+
+### The tie-break
+
+Where the reference and this port genuinely differ: at an exact tie in
+the QSA block scores, `topk` keeps one block and we keep another. Ours
+is deterministic (lower block index wins); torch's is unspecified. A
+tie means every head's dot product was clamped to zero for both blocks
+— the indexer scored them equally worthless — so either choice is
+defensible, and there is no stable target to match. The fixture avoids
+the situation rather than pretending it is resolved.
 
 ## Running the comparison
 

crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny-long/config.json

@@ -0,0 +1,50 @@
+{
+  "model_type": "qwen4_exp",
+  "tie_word_embeddings": false,
+  "text_config": {
+    "vocab_size": 32,
+    "hidden_size": 8,
+    "num_hidden_layers": 4,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "head_dim": 4,
+    "max_position_embeddings": 128,
+    "rms_norm_eps": 1e-06,
+    "hidden_act": "silu",
+    "eos_token_id": 3,
+    "full_attention_interval": 4,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "partial_rotary_factor": 0.5,
+      "rope_type": "default"
+    },
+    "hc_count": 2,
+    "hc_lowrank": 4,
+    "ple_layer_ids": [
+      2
+    ],
+    "ple_embed_dim": 8,
+    "ple_conv_kernel_size": 2,
+    "ngram_size": 3,
+    "heads_per_ngram": 2,
+    "ngram_vocab_size_base": 10,
+    "make_ngram_vocab_size_divisible_by": 4,
+    "split_ngram_parts": 2,
+    "indexer_n_heads": 2,
+    "indexer_kv_heads": 1,
+    "indexer_head_dim": 4,
+    "indexer_budget": 8,
+    "indexer_compress_ratio": 2,
+    "linear_num_key_heads": 1,
+    "linear_key_head_dim": 4,
+    "linear_num_value_heads": 2,
+    "linear_value_head_dim": 4,
+    "linear_conv_kernel_dim": 2,
+    "mamba_ssm_dtype": "float32",
+    "output_gate_type": "sigmoid",
+    "num_experts": 2,
+    "num_experts_per_tok": 1,
+    "moe_intermediate_size": 4,
+    "shared_expert_intermediate_size": 4
+  }
+}

crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny/config.json

@@ -0,0 +1,50 @@
+{
+  "model_type": "qwen4_exp",
+  "tie_word_embeddings": false,
+  "text_config": {
+    "vocab_size": 32,
+    "hidden_size": 8,
+    "num_hidden_layers": 4,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "head_dim": 4,
+    "max_position_embeddings": 128,
+    "rms_norm_eps": 1e-06,
+    "hidden_act": "silu",
+    "eos_token_id": 3,
+    "full_attention_interval": 4,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "partial_rotary_factor": 0.5,
+      "rope_type": "default"
+    },
+    "hc_count": 2,
+    "hc_lowrank": 4,
+    "ple_layer_ids": [
+      2
+    ],
+    "ple_embed_dim": 8,
+    "ple_conv_kernel_size": 2,
+    "ngram_size": 3,
+    "heads_per_ngram": 2,
+    "ngram_vocab_size_base": 10,
+    "make_ngram_vocab_size_divisible_by": 4,
+    "split_ngram_parts": 2,
+    "indexer_n_heads": 2,
+    "indexer_kv_heads": 1,
+    "indexer_head_dim": 4,
+    "indexer_budget": 8,
+    "indexer_compress_ratio": 2,
+    "linear_num_key_heads": 1,
+    "linear_key_head_dim": 4,
+    "linear_num_value_heads": 2,
+    "linear_value_head_dim": 4,
+    "linear_conv_kernel_dim": 2,
+    "mamba_ssm_dtype": "float32",
+    "output_gate_type": "sigmoid",
+    "num_experts": 2,
+    "num_experts_per_tok": 1,
+    "moe_intermediate_size": 4,
+    "shared_expert_intermediate_size": 4
+  }
+}

crates/neuron/tests/numerical_reference.rs

@@ -339,3 +339,180 @@ fn vision_tower_and_logits_match_reference() {
     );
     assert_eq!(cr.argmax_ours, cr.argmax_ref, "chunked argmax vs reference");
 }
+
+/// `qwen4_exp` logits against the upstream reference (#323).
+///
+/// Unlike the fixtures above this one needs no `NEURON_REF_MODEL_PATH`
+/// and no snapshot: the model is 44 KB of random weights in the real
+/// architecture, generated *by* `transformers` itself
+/// (`script/generate_qwen4_exp_fixture.py`), so it is checked in and
+/// this runs on every push. A tensor name we cannot find is therefore a
+/// finding about our loader rather than a fixture that drifted.
+///
+/// f32 on both sides, for the reason in this file's header: cross-dtype
+/// comparison measures bf16 rounding, not implementation error.
+///
+/// What it covers that the unit tests cannot: the four novel blocks
+/// composed. Hyper-connections carrying four residual streams with no
+/// layernorm of the usual kind, the PLE n-gram gather on layer 1, QSA
+/// on layer 3, and the final mixer collapse with no `model.norm` before
+/// `lm_head` — every one of which was written from prose and tested
+/// against our own reading of it.
+#[test]
+fn qwen4_exp_logits_match_reference() {
+    // The recurrent GDN step and a 12-token prompt: agreement is
+    // essentially exact (observed 2e-6), so the bound sits just off the
+    // floor. At 1e-3 this passed with the fused expert's gate and up
+    // halves swapped (2.4e-4) — a bound wide enough to feel safe is
+    // wide enough to hide what it was written for.
+    qwen4_exp_case("qwen4_exp-tiny", 1e-5);
+}
+
+/// The same comparison over a prompt long enough, and shaped
+/// deliberately enough, to reach two paths the short case cannot.
+///
+/// **80 tokens** puts the GatedDeltaNet layers on the *chunked* prefill
+/// algorithm (#23), which engages at 64 and is a different
+/// implementation from the recurrent step the 12-token case exercises —
+/// the same reason the qwen3_5 fixture above is over 64 tokens. It also
+/// gives QSA up to 40 blocks against a budget of 4, so selection is
+/// discarding almost everything rather than one block at the boundary.
+///
+/// **An eos at position 37** puts a segment boundary inside the n-gram
+/// window. `_shift_right_ignore_eos` refuses to read an n-gram across
+/// one, and that logic — a cummax over eos positions, per row — is
+/// intricate enough to be worth comparing rather than reasoning about.
+/// The short case has no eos at all, on purpose: a fixture that happens
+/// to contain one tests segment handling by accident.
+#[test]
+fn qwen4_exp_long_prompt_logits_match_reference() {
+    // Looser than the short case, and for a stated reason rather than
+    // to go green: at 80 tokens the GDN layers take the chunked delta
+    // rule, which is a different *factorisation* of the same
+    // recurrence, so its f32 rounding differs from the recurrent step
+    // by construction. Observed 1.3e-5 against the 2e-6 of the short
+    // case. The bound is still four orders of magnitude below what the
+    // mutations produce — gate/up swapped moves this fixture by 2.7 —
+    // so it discriminates exactly as well.
+    qwen4_exp_case("qwen4_exp-tiny-long", 1e-4);
+}
+
+fn qwen4_exp_case(name: &str, max_abs_bound: f32) {
+    use candle_nn::var_builder::ShardedSafeTensors;
+    use neuron::harness::arch::qwen4_exp::config::Config as Qwen4ExpConfig;
+    use neuron::harness::arch::qwen4_exp::model::Qwen4ExpForCausalLM;
+
+    let fixture = fixture_root().join(name);
+    let weights = fixture.join("model.safetensors");
+    let device = Device::Cpu;
+
+    let cfg = Qwen4ExpConfig::from_config_json(
+        &std::fs::read_to_string(fixture.join("config.json")).unwrap(),
+    )
+    .expect("parse the fixture config");
+
+    let vb = unsafe { ShardedSafeTensors::var_builder(&[&weights], DType::F32, &device).unwrap() };
+    let mut model = Qwen4ExpForCausalLM::load(
+        &cfg,
+        DType::F32,
+        &device,
+        &vb,
+        &neuron::harness::arch::qwen4_exp::LoadOptions {
+            quant: None,
+            safetensors_paths: std::slice::from_ref(&weights),
+            isq_cache: None,
+            experts_onto: &device,
+            experts_vb: None,
+        },
+    )
+    .expect("load the fixture through the production loader");
+
+    let expected = candle_core::safetensors::load(fixture.join("expected.safetensors"), &device)
+        .expect("read the reference output");
+    let ids = expected["input_ids"].to_dtype(DType::U32).unwrap();
+    let reference: Vec<f32> = expected["logits"].flatten_all().unwrap().to_vec1().unwrap();
+
+    // Every position, because that is what the reference returns and
+    // what makes a per-position breakdown possible at all.
+    let logits = model
+        .logits_all_positions(&ids, 0)
+        .expect("all-position logits");
+    let ours: Vec<f32> = logits
+        .to_dtype(DType::F32)
+        .unwrap()
+        .flatten_all()
+        .unwrap()
+        .to_vec1()
+        .unwrap();
+
+    let c = compare(&ours, &reference);
+    eprintln!(
+        "{name}: max_abs={:.6} cosine={:.6} argmax ours={} ref={}",
+        c.max_abs, c.cosine, c.argmax_ours, c.argmax_ref
+    );
+    // On failure, break the error down by position before asserting.
+    // Whether the divergence is confined to one position or spread
+    // across all of them is the first thing worth knowing: localised
+    // means a discrete choice went differently (an expert, a QSA
+    // block), spread means the arithmetic drifted. Finding that out
+    // cost an hour the first time.
+    if c.max_abs >= max_abs_bound {
+        let v = cfg.text_config.vocab_size;
+        for t in 0..(ours.len() / v) {
+            let pc = compare(&ours[t * v..(t + 1) * v], &reference[t * v..(t + 1) * v]);
+            eprintln!(
+                "  pos {t:2}: max_abs={:.6} cosine={:.6}",
+                pc.max_abs, pc.cosine
+            );
+        }
+    }
+    assert_eq!(c.argmax_ours, c.argmax_ref, "argmax token diverged");
+    assert!(c.cosine > 0.999_999, "cosine {:.6} too low", c.cosine);
+    assert!(
+        c.max_abs < max_abs_bound,
+        "max abs diff {:.6} exceeds {max_abs_bound:e}",
+        c.max_abs
+    );
+
+    // The serving path returns only the last position, `(B, 1, vocab)`
+    // — a different shape from what was just compared, and the
+    // difference is not cosmetic: `squeeze_to_vocab` strips leading
+    // singleton dims and rejects the rest, so a model returning every
+    // position loads fine and then faults on its first token with
+    // "logits expected to start with a singleton dim". That is exactly
+    // what happened on beast, after a 16-minute load, because this
+    // test flattened both sides and a flattened `[1,12,32]` and
+    // `[12,32]` are the same 384 numbers.
+    //
+    // So: assert the shape, and assert the two paths agree.
+    // The caches advanced during the call above; a second forward at
+    // offset 0 would find the indexer holding twice the positions it
+    // expects. Clearing is what a real second request does.
+    model
+        .clear_kv_cache()
+        .expect("clear between the two forwards");
+    let serving = model.forward(&ids, 0).expect("serving forward");
+    let (b, t, v) = serving.dims3().expect("serving logits are rank 3");
+    assert_eq!(
+        (b, t),
+        (1, 1),
+        "the serving path must return one position, got {:?}",
+        serving.dims()
+    );
+    let last: Vec<f32> = serving.flatten_all().unwrap().to_vec1().unwrap();
+    let expected_last = &ours[ours.len() - v..];
+    // Close, not identical: the serving path slices to one position
+    // *before* `lm_head`, so its matmul is `[1,1,h] x [h,V]` where the
+    // other is `[1,T,h] x [h,V]`. Same arithmetic, different
+    // accumulation order, ~1e-7 apart. Demanding bit-equality here
+    // would be asserting something about BLAS rather than about us.
+    let worst = last
+        .iter()
+        .zip(expected_last)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0f32, f32::max);
+    assert!(
+        worst < 1e-5,
+        "the serving path disagrees with the last position of the full logits by {worst:e}"
+    );
+}

crates/neuron/tests/preflight.rs

@@ -108,6 +108,8 @@ fn spec(model_id: &str, tp: Option<u32>, quant: Option<&str>) -> ModelSpec {
         devices: None,
         sampling: None,
         preserve_thinking: None,
+        expert_residency: None,
+        max_in_flight: None,
     }
 }
 

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

script/generate_qwen4_exp_fixture.py

@@ -0,0 +1,128 @@
+"""Generate a qwen4_exp parity fixture from the upstream reference (#323).
+
+Writes, into --out:
+  config.json          the tiny config, as both implementations parse it
+  model.safetensors    the reference's own weights, with the PLE n-gram
+                       table split into `split_ngram_parts` shards on
+                       dim 0 the way the released checkpoint ships it
+  expected.safetensors input_ids and the reference's logits
+
+The weights are the reference's, so a name we cannot find is a finding
+rather than a fixture bug. Norm weights are perturbed away from their
+zero init on purpose: at exactly zero, `(1 + w)` and `w` differ, but
+`(1 + w)` and `1` do not, and the whole point is to tell those apart.
+"""
+
+import argparse, json, os
+import torch
+from safetensors.torch import save_file
+from transformers.models.qwen4_exp import Qwen4ExpConfig, Qwen4ExpForConditionalGeneration
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--config", required=True, help="the tiny config.json both implementations parse")
+ap.add_argument("--out", required=True, help="fixture directory to write")
+# Seed 2 is tie-free; 0, 1 and 4 are not. See the tie check below.
+ap.add_argument("--seed", type=int, default=2)
+ap.add_argument("--tokens", type=int, default=12)
+ap.add_argument(
+    "--eos-at",
+    type=int,
+    default=None,
+    help="put an eos at this position, so the n-gram window has a segment "
+    "boundary to refuse to read across",
+)
+args = ap.parse_args()
+
+torch.manual_seed(args.seed)
+torch.use_deterministic_algorithms(True)
+
+raw = json.load(open(args.config))
+cfg = Qwen4ExpConfig(**raw)
+tc = cfg.text_config
+
+model = Qwen4ExpForConditionalGeneration(cfg).eval()
+
+# Every parameter gets a non-degenerate value. Zeros hide a transposed
+# load. But *small* values hide more than that: silu(x) ~ x/2 near the
+# origin, so with weights in +-0.1 the SwiGLU is close enough to
+# symmetric that reading gate and up the wrong way round changes the
+# logits by 2e-4 — inside any tolerance you would write. The scale here
+# is chosen to put activations where silu is actually curved, which is
+# what makes the ordering observable, while staying below the point
+# where the MoE router saturates and top-k stops discriminating.
+with torch.no_grad():
+    for name, p in model.named_parameters():
+        if "visual" in name or "vision" in name:
+            continue
+        p.copy_(torch.empty_like(p).uniform_(-0.8, 0.8))
+
+ids = torch.randint(0, tc.vocab_size, (1, args.tokens), dtype=torch.long)
+# eos is not an ordinary token here: `_shift_right_ignore_eos` treats it
+# as a segment boundary, so an n-gram never reads across one. Placing it
+# is therefore a deliberate choice rather than something to leave to the
+# sampler — a fixture that happens to contain one is testing segment
+# handling by accident, and one that happens not to is not testing it at
+# all.
+ids[ids == tc.eos_token_id] = (tc.eos_token_id + 1) % tc.vocab_size
+if args.eos_at is not None:
+    ids[0, args.eos_at] = tc.eos_token_id
+
+# The QSA selection must not rest on a tie.
+#
+# A block scores exactly 0.0 whenever relu clamps every head's dot
+# product, which is common enough that a random fixture hits it. When
+# the k-th and (k+1)-th scores are equal, which block gets dropped is
+# decided by `topk`'s ordering of equal elements — and torch does not
+# specify that, nor does its CPU kernel agree with its CUDA one. A
+# fixture containing such a tie measures that unspecified detail rather
+# than our arithmetic, so we refuse to emit one.
+orig_topk = torch.Tensor.topk
+ties = []
+def _topk(self, k, *a, **kw):
+    if self.dim() == 1 and self.numel() <= 64 and k < self.numel():
+        ordered = sorted(self.detach().tolist(), reverse=True)
+        if ordered[k - 1] == ordered[k]:
+            ties.append((self.detach().tolist(), k))
+    return orig_topk(self, k, *a, **kw)
+torch.Tensor.topk = _topk
+with torch.no_grad():
+    out = model(input_ids=ids, use_cache=False)
+torch.Tensor.topk = orig_topk
+
+if ties:
+    raise SystemExit(
+        f"seed {args.seed} produces {len(ties)} QSA selection tie(s), e.g. "
+        f"scores={[round(v, 6) for v in ties[0][0]]} k={ties[0][1]}. "
+        "Pick another --seed: the fixture must isolate the arithmetic."
+    )
+
+os.makedirs(args.out, exist_ok=True)
+
+sd = {k: v for k, v in model.state_dict().items() if "visual" not in k and "vision" not in k}
+
+# Split the runtime n-gram embedding back into the checkpoint's shard
+# layout: dim 0, ascending, per upstream's own note that "the
+# checkpoints for it are sharded on dim0" and loading concatenates them.
+parts = tc.split_ngram_parts
+for key in [k for k in sd if k.endswith("ngram_embedding.weight")]:
+    table = sd.pop(key)
+    rows = table.shape[0]
+    assert rows % parts == 0, f"{rows} rows do not divide into {parts} shards"
+    per = rows // parts
+    base = key[: -len(".weight")]
+    for s in range(parts):
+        sd[f"{base}.shard_{s}.weight"] = table[s * per : (s + 1) * per].clone()
+
+save_file({k: v.contiguous() for k, v in sd.items()}, os.path.join(args.out, "model.safetensors"))
+save_file(
+    {"input_ids": ids.to(torch.int64), "logits": out.logits.to(torch.float32).contiguous()},
+    os.path.join(args.out, "expected.safetensors"),
+)
+with open(os.path.join(args.out, "config.json"), "w") as f:
+    json.dump(raw, f, indent=2)
+
+print(f"tokens      {args.tokens}")
+print(f"tensors     {len(sd)}")
+print(f"logits      {tuple(out.logits.shape)}")
+print(f"logits absmax {out.logits.abs().max().item():.6f}")
+print(f"finite      {bool(torch.isfinite(out.logits).all())}")