fix(neuron): pass the model id into the worker streaming path

The metrics call added in the previous commit referenced `model_id` in
`stream_inference_via_worker`, which has no such binding: that function
gets the model name from its tracing span, and a span field is not a
value. Caught by the CUDA type-check, which is the only job that
compiles either streaming path — both are `#[cfg(feature = "cuda")]`
and neither builds locally.

Threaded through as an owned `String` (the function runs in a spawned
task) following the `_inner` clone pattern the call site already uses
for the other captured state.

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

Cargo.lock

@@ -3147,6 +3147,10 @@ dependencies = [
  "half",
  "hf-hub",
  "image",
+ "libc",
+ "memmap2",
+ "metrics",
+ "metrics-exporter-prometheus",
  "minijinja",
  "minijinja-contrib",
  "rayon",

asset/monitoring/README.md

@@ -6,10 +6,18 @@ bench UI shows.
 
 ## Topology
 
-- **cortex** runs on `hanzalova.internal` (10.6.0.46) and is the **only**
-  Prometheus target — it exposes every `cortex_*` metric on `:31314`, already
-  labelled by `{node,model}` / `{node,device}` from its neuron poller. neuron
-  has no `/metrics` endpoint.
+- **cortex** runs on `hanzalova.internal` (10.6.0.46) and exposes every
+  `cortex_*` metric on `:31314`, already labelled by `{node,model}` /
+  `{node,device}` from its neuron poller. This is the fleet-shaped view:
+  who is loaded where, how busy, how much VRAM.
+- **neuron** additionally serves `/metrics` on its own API port
+  (`:13131`) on each GPU host. This is the request-shaped view — where a
+  request's time actually went, as histograms rather than polled gauges.
+  It exists because cortex's numbers are a ~10s poll of `/health`
+  republished as gauges, which cannot see inside a one-second forward and
+  carries no distribution. Scrape job: `prometheus-neuron.scrape.yml`.
+  **No firewalld rule needed** — 13131 is already open for the control
+  plane (`data/neuron-firewalld.xml`).
 - **Prometheus + Grafana** run on `golgafrinchans.kosherinata.internal` as
   podman quadlets. Host ports (from the registered band — 9090 is Cockpit,
   3000 was avoided): **Prometheus `:26559`** (container config bind-mounted at

asset/monitoring/prometheus-cortex.scrape.yml

@@ -4,9 +4,12 @@
 # Cockpit on that box):  curl -X POST http://localhost:26559/-/reload
 #
 # cortex runs on hanzalova.internal (10.6.0.46) and exposes ALL fleet metrics
-# on :31314 — it is the only Prometheus target in the stack (neuron has no
-# /metrics; cortex's poller already relabels everything by {node,model} /
-# {node,device}). Reaching it cross-DC (kosherinata -> hanzalova) requires the
+# on :31314; cortex's poller already relabels everything by {node,model} /
+# {node,device}. It is no longer the only target — neuron now serves its own
+# /metrics on :13131 with the per-request phase histograms cortex's ~10s
+# gauge poll cannot express (see prometheus-neuron.scrape.yml); the two are
+# complementary, not redundant.
+# Reaching cortex cross-DC (kosherinata -> hanzalova) requires the
 # firewalld rule in asset/monitoring/firewalld-cortex-metrics.sh, which opens
 # :31314 to this host's mesh IP (10.3.101.4) only.
 #

asset/monitoring/prometheus-neuron.scrape.yml

@@ -0,0 +1,40 @@
+# helexa neuron scrape job — append into the `scrape_configs:` list of
+# /etc/prometheus/prometheus.yml on golgafrinchans.kosherinata.internal,
+# then reload:  curl -X POST http://localhost:26559/-/reload
+#
+# WHY A SECOND TARGET, when cortex already publishes fleet metrics.
+#
+# cortex's numbers come from polling each neuron's /health every ~10s and
+# republishing it as gauges. That is the right shape for routing decisions
+# and the wrong shape for asking where a request's time went: an EMA
+# sampled at 10s cannot see inside a one-second forward, and gauges carry
+# no distribution. The finest split cortex can express is prefill vs
+# decode.
+#
+# neuron serves /metrics on its own API port (13131) — the port is already
+# open in data/neuron-firewalld.xml for the cortex control plane, so unlike
+# the cortex scrape this needs NO new firewalld rule.
+#
+# What is here that cortex cannot give you:
+#   neuron_decode_phase_seconds{model,phase=forward|sample|emit|residual}
+#     Decode time by phase, summed per request, as a real histogram.
+#     The `_sum` series is the point — ratio a phase against the whole to
+#     see where decode goes:
+#       sum(rate(neuron_decode_phase_seconds_sum{phase="forward"}[5m]))
+#         / sum(rate(neuron_decode_phase_seconds_sum[5m]))
+#     `residual` is decode_ms minus the three measured phases, emitted as
+#     its own series so unbracketed time is visible rather than divided
+#     silently among the phases that were measured.
+#   neuron_prefill_seconds / neuron_decode_seconds  (histograms)
+#   neuron_completion_tokens_total  (the per-token denominator)
+#
+# A model whose serving path does not measure phases emits NO phase series
+# at all rather than zeros — absent means "nobody looked", not "free".
+- job_name: neuron
+  static_configs:
+    - targets:
+        - 'beast.hanzalova.internal:13131'
+        - 'benjy.hanzalova.internal:13131'
+        - 'quadbrat.hanzalova.internal:13131'
+      labels:
+        fleet: hanzalova

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-core/src/openai.rs

@@ -210,6 +210,19 @@ pub struct HelexaTiming {
     pub prefill_ms: u64,
     pub decode_ms: u64,
     pub prefill_tokens: u64,
+    /// Decode-loop phase split, when the serving path measured it.
+    /// Absent (rather than zero) on paths that do not, so a consumer
+    /// can tell "not measured" from "measured as free". Together these
+    /// should roughly account for `decode_ms`; the residual is time in
+    /// the loop that no bracket covers.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub forward_ms: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub sample_ms: Option<u64>,
+    /// Detokenise + hand to the consumer. Includes bounded-channel
+    /// backpressure, so a slow reader is charged here.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub emit_ms: Option<u64>,
 }
 
 /// Sub-counts of `Usage::completion_tokens`.

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

@@ -47,6 +47,8 @@ cuda-integration = ["cuda"]
 cortex-core.workspace = true
 tokio.workspace = true
 axum.workspace = true
+metrics.workspace = true
+metrics-exporter-prometheus.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 reqwest.workspace = true
@@ -104,6 +106,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/api.rs

@@ -10,6 +10,7 @@ use crate::wire::{openai_chat, openai_responses};
 use axum::Router;
 use axum::extract::{Path, State};
 use axum::http::StatusCode;
+use axum::http::header;
 use axum::response::sse::{Event, KeepAlive, Sse};
 use axum::response::{IntoResponse, Json};
 use axum::routing::{get, post};
@@ -43,6 +44,7 @@ pub struct NeuronState {
 /// Build the neuron API router.
 pub fn neuron_routes() -> Router<Arc<NeuronState>> {
     Router::new()
+        .route("/metrics", get(metrics_handler))
         .route("/version", get(version_handler))
         .route("/discovery", get(discovery_handler))
         .route("/health", get(health_handler))
@@ -64,6 +66,29 @@ async fn version_handler() -> Json<cortex_core::build_info::BuildInfo> {
     Json(crate::version::build_info())
 }
 
+/// `GET /metrics` — Prometheus text format.
+///
+/// Served from the API port rather than a second listener, so a scrape
+/// target needs no extra firewall rule. A missing recorder is a 503
+/// rather than an empty 200: an empty scrape looks exactly like a
+/// healthy process with nothing to say, and would leave the fleet
+/// silently unmonitored.
+async fn metrics_handler() -> impl IntoResponse {
+    match crate::metrics::render() {
+        Some(body) => (
+            StatusCode::OK,
+            [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
+            body,
+        )
+            .into_response(),
+        None => (
+            StatusCode::SERVICE_UNAVAILABLE,
+            "metrics recorder not installed",
+        )
+            .into_response(),
+    }
+}
+
 async fn discovery_handler(State(state): State<Arc<NeuronState>>) -> Json<DiscoveryResponse> {
     Json(state.discovery.clone())
 }

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

@@ -20,6 +20,7 @@
 //! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
 //! the pattern set by `tp_qwen3.rs`.
 
+pub mod phase_probe;
 pub mod qwen3_5;
 pub mod qwen4_exp;
 pub mod snapshot;

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

@@ -0,0 +1,167 @@
+//! Opt-in phase timing inside a model's forward pass.
+//!
+//! The request-level split (`PhaseTiming` on the wire) says whether a
+//! decode step's time went into the forward, the sampler or the emit.
+//! When the answer is "the forward", this says *where* in the forward —
+//! embedding, a gather, the layer stack, the head.
+//!
+//! # Why this needs a device synchronise, and why it is off by default
+//!
+//! candle's CUDA ops are asynchronous on the stream. An `Instant` taken
+//! around GPU work therefore measures how fast kernels were *queued*,
+//! not how long they ran, and whichever phase happens to contain the
+//! next real synchronisation point silently absorbs everyone else's
+//! cost. The numbers look plausible and attribute the time to the wrong
+//! place — a worse failure than no numbers at all.
+//!
+//! So each mark synchronises the device first. That is only affordable
+//! when the phases being measured are much larger than a sync: at a few
+//! tens of microseconds per sync, a handful of marks is noise against a
+//! step measured in hundreds of milliseconds, and material against one
+//! measured in single milliseconds. Since that ratio is a property of
+//! how slow the model currently is rather than something safe to assume,
+//! the probe stays off unless `NEURON_FORWARD_PHASES` is set.
+
+use candle_core::Device;
+use std::sync::OnceLock;
+use std::time::Instant;
+
+fn enabled() -> bool {
+    static ON: OnceLock<bool> = OnceLock::new();
+    *ON.get_or_init(|| std::env::var("NEURON_FORWARD_PHASES").is_ok())
+}
+
+/// Accumulates named phases within one forward pass.
+///
+/// Construct at the top of the forward, call [`mark`] at each boundary,
+/// and [`emit`] at the end. Disabled, every method is a predictable
+/// branch and nothing is allocated.
+///
+/// [`mark`]: PhaseProbe::mark
+/// [`emit`]: PhaseProbe::emit
+pub struct PhaseProbe {
+    on: bool,
+    device: Device,
+    last: Instant,
+    marks: Vec<(&'static str, f64)>,
+}
+
+impl PhaseProbe {
+    pub fn new(device: &Device) -> Self {
+        Self::with_enabled(device, enabled())
+    }
+
+    /// Constructor with the gate supplied explicitly.
+    ///
+    /// [`new`] reads the environment through a `OnceLock`, which a test
+    /// cannot toggle — so without this the enabled path would be
+    /// unreachable from a test, and a probe that only ever runs
+    /// disabled is not evidence that it records anything.
+    ///
+    /// [`new`]: PhaseProbe::new
+    pub fn with_enabled(device: &Device, on: bool) -> Self {
+        // Start from a synchronised device, or the first phase inherits
+        // whatever was still in flight when the forward was entered.
+        if on {
+            let _ = device.synchronize();
+        }
+        Self {
+            on,
+            device: device.clone(),
+            last: Instant::now(),
+            marks: Vec::new(),
+        }
+    }
+
+    /// Close the phase that ended here and open the next one.
+    pub fn mark(&mut self, name: &'static str) {
+        if !self.on {
+            return;
+        }
+        // A failed synchronise means the elapsed time below is not the
+        // phase's own. Record it as such rather than reporting a number
+        // that looks like a measurement.
+        let synced = self.device.synchronize().is_ok();
+        let ms = self.last.elapsed().as_secs_f64() * 1e3;
+        self.marks
+            .push((if synced { name } else { "UNSYNCED" }, ms));
+        self.last = Instant::now();
+    }
+
+    /// The phases recorded so far, in order.
+    #[cfg(test)]
+    pub fn marks(&self) -> &[(&'static str, f64)] {
+        &self.marks
+    }
+
+    /// Emit one line carrying every phase of this forward.
+    pub fn emit(self, rows: usize) {
+        if !self.on || self.marks.is_empty() {
+            return;
+        }
+        let total: f64 = self.marks.iter().map(|(_, ms)| ms).sum();
+        let detail = self
+            .marks
+            .iter()
+            .map(|(n, ms)| format!("{n}={ms:.3}"))
+            .collect::<Vec<_>>()
+            .join(" ");
+        tracing::info!(rows, total_ms = total, %detail, "forward phases");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A disabled probe must record nothing at all — not zero-valued
+    /// phases. Zeros in a log read as "this phase was free", which is
+    /// the same failure the wire-side `Option` exists to prevent.
+    #[test]
+    fn a_disabled_probe_records_nothing() {
+        let mut p = PhaseProbe::with_enabled(&Device::Cpu, false);
+        p.mark("embed");
+        p.mark("layers");
+        assert!(
+            p.marks().is_empty(),
+            "disabled probe recorded {:?}",
+            p.marks()
+        );
+    }
+
+    /// An enabled probe records one phase per mark, in call order, and
+    /// each phase measures only the span since the previous mark rather
+    /// than since the start.
+    #[test]
+    fn phases_are_recorded_in_order_and_do_not_accumulate() {
+        let mut p = PhaseProbe::with_enabled(&Device::Cpu, true);
+        p.mark("embed");
+        std::thread::sleep(std::time::Duration::from_millis(20));
+        p.mark("layers");
+        p.mark("head");
+
+        let names: Vec<_> = p.marks().iter().map(|(n, _)| *n).collect();
+        assert_eq!(names, ["embed", "layers", "head"]);
+
+        let ms = |i: usize| p.marks()[i].1;
+        // "layers" contains the sleep; the marks either side must not.
+        // A probe that timed from the start instead of from the previous
+        // mark would make "head" >= "layers", which is the mistake this
+        // pins.
+        assert!(
+            ms(1) >= 15.0,
+            "layers should carry the sleep, got {:?}",
+            p.marks()
+        );
+        assert!(
+            ms(2) < 10.0,
+            "head must not inherit it, got {:?}",
+            p.marks()
+        );
+        assert!(
+            ms(0) < 10.0,
+            "embed must not inherit it, got {:?}",
+            p.marks()
+        );
+    }
+}

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

@@ -30,7 +30,7 @@
 
 use anyhow::{Context, Result};
 use candle_core::quantized::{GgmlDType, QMatMul, QTensor};
-use candle_core::{DType, IndexOp, Module, Tensor};
+use candle_core::{DType, Device, IndexOp, Module, Tensor};
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
@@ -88,6 +88,20 @@ pub(crate) struct QuantizedExpert {
     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
@@ -98,39 +112,59 @@ pub(crate) struct QuantizedExpert {
 /// `quantized tensor must have their last dim divisible by block size
 /// [2560, 640] 256`, having read 131 shards to get there.
 ///
-/// Substituting a 32-block type for that half is what llama.cpp does
-/// with the same constraint, and it costs almost nothing here: `Q4_0`
-/// and `Q4K` are both 4.5 bpw, so the memory arithmetic the placement
-/// decisions rest on is unchanged. The quality is a little worse per
-/// bit — `Q4_0` has one scale per 32 values where `Q4K` has a
-/// two-level hierarchy — which is a real trade, but it applies to one
-/// of the three expert matrices and the alternative is not loading.
+/// **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:
 ///
-/// Errors rather than falls back further if the row will not even
-/// divide 32: that is a config we do not understand, and silently
-/// leaving those weights dense would be a memory surprise, not a
-/// mercy.
+/// | 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::Q4K => GgmlDType::Q4_0,
-        GgmlDType::Q5K => GgmlDType::Q5_0,
+        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: there is nothing coarser to fall
-        // back to, so this is the error case below.
+        // 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);
     }
-    candle_core::bail!(
-        "cannot quantise a row of {row} elements as {requested:?} (block \
-         {}) or {fallback:?} (block {}): the row divides neither",
-        requested.block_size(),
-        fallback.block_size(),
-    )
+    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 {
@@ -150,11 +184,31 @@ impl Experts {
     /// 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 —
@@ -189,8 +243,14 @@ impl Experts {
             // 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, t.device())?;
+                let owned = Tensor::from_vec(vals, dims, onto)?;
                 QMatMul::from_qtensor(QTensor::quantize(&owned, dtype)?)
             };
             out.push(QuantizedExpert {
@@ -202,6 +262,65 @@ impl Experts {
         Ok(Experts::Quantized(out))
     }
 
+    /// Where the routed experts live, when that is knowable (#318).
+    pub(crate) fn device(&self) -> Option<Device> {
+        match self {
+            Experts::Quantized(v) => v.first().and_then(QuantizedExpert::device),
+            _ => None,
+        }
+    }
+
+    /// 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 {
@@ -223,13 +342,25 @@ impl Experts {
                 // The quantised kernels accumulate in f32 whatever the
                 // activation dtype, so the cast is explicit here rather
                 // than implied — and the result comes back in the
-                // caller's dtype so the routing arithmetic upstream
-                // does not have to know which layout it got.
+                // 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 xs32 = xs.to_dtype(DType::F32)?;
+                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_dtype(dtype)
+                q.down
+                    .forward(&(lhs * rhs)?)?
+                    .to_device(xs.device())?
+                    .to_dtype(dtype)
             }
         }
     }
@@ -392,6 +523,9 @@ impl Module for Qwen3_5MoeBlock {
         let (b, l, hidden) = xs.dims3()?;
         let xs_flat = xs.reshape(((), hidden))?;
 
+        // Routing stays where the activations are: it is one small
+        // matmul against `gate` and a host-side top-k, and the router
+        // is device-resident even when the experts are not.
         let (tokens_for, weights_for) = route_scatter(
             &self.gate,
             &xs_flat,
@@ -400,19 +534,51 @@ impl Module for Qwen3_5MoeBlock {
             self.norm_topk_prob,
         )?;
 
-        let mut ys = xs_flat.zeros_like()?;
+        // Cross to the experts once for the whole layer, not once per
+        // expert (#318). With host-resident experts the old shape was
+        // 10 round trips per layer and 480 per token; this is one in
+        // and one out. `to_device` is a clone when they already match,
+        // so a device-resident model pays nothing.
+        //
+        // The measurement this was written against says the transfers
+        // were never the dominant cost — see `expert_ms` below — but
+        // one crossing per layer is the right shape regardless of what
+        // dominates.
+        let home = self.experts.device();
+        let off_device = home.is_some();
+        let routed_dev = home.unwrap_or_else(|| xs.device().clone());
+        let xs_routed = xs_flat.to_device(&routed_dev)?;
+
+        let started = std::time::Instant::now();
+        let mut ys = xs_routed.zeros_like()?;
         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 rows = Tensor::new(tokens_for[e].as_slice(), &routed_dev)?;
+            let picked = xs_routed.index_select(&rows, 0)?;
             let out = self.experts.forward_one(e, &picked)?;
-            let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
+            let w = Tensor::new(weights_for[e].as_slice(), &routed_dev)?
                 .to_dtype(out.dtype())?
                 .reshape(((), 1))?;
             ys = ys.index_add(&rows, &out.broadcast_mul(&w)?, 0)?;
         }
+        // Where the time actually goes, per layer, when the experts are
+        // not on the device. Logged at trace so a normal load is
+        // unaffected; `NEURON_EXPERT_TIMING=1` is what a measurement
+        // run sets. Without this the only figure available is
+        // end-to-end tok/s, which cannot distinguish "the transfers are
+        // slow" from "the CPU matmul is slow" — and those have opposite
+        // fixes.
+        if off_device && std::env::var("NEURON_EXPERT_TIMING").is_ok() {
+            tracing::info!(
+                expert_ms = started.elapsed().as_secs_f64() * 1e3,
+                rows = xs_routed.dims()[0],
+                active = tokens_for.iter().filter(|t| !t.is_empty()).count(),
+                "moe: routed experts (#318)"
+            );
+        }
+        let mut ys = ys.to_device(xs.device())?;
 
         if let (Some(shared), Some(gate)) = (&self.shared_expert, &self.shared_expert_gate) {
             let mix = candle_nn::ops::sigmoid(&gate.forward(&xs_flat)?)?;
@@ -687,6 +853,174 @@ mod tests {
     /// 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]`
@@ -708,16 +1042,25 @@ mod tests {
         );
         assert_eq!(
             quant_for_row(GgmlDType::Q4K, 640).unwrap(),
-            GgmlDType::Q4_0,
+            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.
-        assert!(quant_for_row(GgmlDType::Q4K, 100).is_err());
+        // 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);

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

@@ -41,8 +41,6 @@ use crate::harness::arch::qwen3_5::linear_attn::GatedDeltaNet;
 use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
 use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
 
-use candle_core::quantized::GgmlDType;
-
 use super::config::TextConfig;
 use super::full_attn::Attention;
 use super::hyper::HyperConnection;
@@ -74,8 +72,8 @@ impl DecoderLayer {
         cfg: &TextConfig,
         rotary: Arc<RotaryEmbedding>,
         layer_idx: usize,
-        quant: Option<GgmlDType>,
         vb: &ShardedVarBuilder,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let layer_type = cfg
             .layer_types
@@ -92,8 +90,8 @@ impl DecoderLayer {
             rotary,
             layer_type,
             cfg.ple_layers().contains(&layer_idx),
-            quant,
             vb,
+            opts,
         )
     }
 
@@ -108,8 +106,8 @@ impl DecoderLayer {
         rotary: Arc<RotaryEmbedding>,
         layer_type: &str,
         with_ple: bool,
-        quant: Option<GgmlDType>,
         vb: &ShardedVarBuilder,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let mixer = match layer_type {
             "full_attention" => {
@@ -155,7 +153,7 @@ impl DecoderLayer {
             attn_hc: hc("attn_hyper_connection")?,
             mlp_hc: hc("mlp_hyper_connection")?,
             mixer,
-            mlp: moe::load(cfg, quant, &vb.pp("mlp"))?,
+            mlp: moe::load(cfg, &vb.pp("mlp"), opts)?,
             ple,
         })
     }

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

@@ -44,3 +44,47 @@ 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

@@ -34,7 +34,6 @@
 //! See `doc/qwen4_exp-port-spec.md`.
 
 use anyhow::{Context, Result, ensure};
-use candle_core::quantized::GgmlDType;
 use candle_core::{DType, Device, IndexOp, Module, Tensor};
 use candle_nn::var_builder::ShardedVarBuilder;
 use candle_nn::{Embedding, Linear};
@@ -45,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::phase_probe::PhaseProbe;
 use crate::harness::arch::snapshot::{KvCacheSnapshot, PleSnapshot};
 
 /// The hashed n-gram lookup, hoisted to the model because it is a
@@ -129,12 +129,17 @@ impl Qwen4ExpForCausalLM {
     /// 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,
-        quant: Option<GgmlDType>,
         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
@@ -150,10 +155,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, quant, &layers_vb.pp(i))
+                DecoderLayer::load(text, rotary.clone(), i, &layers_vb.pp(i), opts)
                     .with_context(|| format!("load decoder layer {i}"))?,
             );
         }
@@ -175,11 +199,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,
@@ -195,6 +214,25 @@ 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)
     }
@@ -218,6 +256,7 @@ impl Qwen4ExpForCausalLM {
 
         // Four identical streams — the residual this architecture
         // carries between layers, not one.
+        let mut probe = PhaseProbe::new(&self.device);
         let x = self.embed_tokens.forward(input_ids)?;
         let mut h = x.repeat((1, 1, self.hc_count))?;
 
@@ -227,6 +266,7 @@ impl Qwen4ExpForCausalLM {
         } else {
             None
         };
+        probe.mark("embed");
 
         // Computed once: it is a function of the ids, and one layer
         // consumes it.
@@ -237,16 +277,33 @@ impl Qwen4ExpForCausalLM {
             }
             None => None,
         };
+        // The n-gram table is mapped from the checkpoint rather than
+        // resident (#310), so this phase is a host page-fault path, not
+        // device work: it is where a cold or evicted mapping shows up.
+        probe.mark("ngram");
 
         for (i, layer) in self.layers.iter_mut().enumerate() {
             h = layer
                 .forward(&h, ngram.as_ref(), mask.as_ref(), &cos, &sin, offset)
                 .with_context(|| format!("layer {i}"))?;
         }
+        probe.mark("layers");
 
         // The mixer's hc_norm is the only normalisation before the head.
         let x = self.mixer.collapse(&h)?;
-        Ok((self.lm_head.forward(&x)?, h))
+        // 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)?;
+        let logits = self.lm_head.forward(&x.i((.., l - 1.., ..))?)?;
+        probe.mark("head");
+        probe.emit(seq_len);
+        Ok((logits, h))
     }
 
     /// Capture every piece of per-request state at one token boundary:
@@ -356,7 +413,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();
@@ -382,12 +444,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(
@@ -398,7 +476,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(),
@@ -757,7 +835,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, None, &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)
     }
 
@@ -770,13 +864,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

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

@@ -40,7 +40,6 @@
 
 use anyhow::{Context, Result};
 use candle_core::Tensor;
-use candle_core::quantized::GgmlDType;
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
@@ -58,9 +57,10 @@ use super::config::TextConfig;
 /// load of `Qwen3.8-Flash-Next` always passes one.
 pub fn load(
     cfg: &TextConfig,
-    quant: Option<GgmlDType>,
     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,
@@ -77,25 +77,77 @@ pub fn load(
         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()))?;
-    check_fused_experts(&gate_up, &down, inter)?;
-    // The fused pair is the peak: 5.03 GB per layer at BF16. Quantising
-    // here and letting it drop at the end of this function is what keeps
-    // the whole 241.6 GB from ever existing.
-    let experts = match quant {
-        Some(dtype) => Experts::quantize_banked(&gate_up, &down, inter, dtype)
-            .with_context(|| format!("quantize experts to {dtype:?}"))?,
-        None => Experts::Banked {
-            gate_up,
-            down,
-            intermediate: inter,
-        },
+    // 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 {

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

@@ -113,8 +113,18 @@ impl MtpHead {
                 rotary,
                 "full_attention",
                 false,
-                quant,
                 &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(

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
@@ -475,10 +685,35 @@ impl PleBlock {
             Tensor::cat(&[&pad, &prepended], 2)?
         });
 
-        // Pad both sides by the full reach, then keep the left-aligned
-        // window: output t sees inputs t - j*dilation only.
-        let out = prepended.conv1d(&self.conv1d_weight, reach, 1, self.dilation, wide)?;
-        let out = candle_nn::ops::silu(&out.narrow(2, 0, prep_len)?)?;
+        // Written as `kernel_size` shifted broadcast multiplies rather
+        // than a `conv1d` with `groups = wide`. candle decomposes a
+        // grouped convolution into one single-group call per group plus
+        // a `groups`-way concat, so the depthwise form costs O(wide)
+        // kernel launches — ~40k of them for the 10240-wide stream here,
+        // which dominates the whole forward pass. This form is O(kernel
+        // _size) launches and computes the same sum:
+        //
+        //     out[t] = sum_j w[.., j] * x[t - (kernel_size - 1 - j) * dilation]
+        //
+        // with positions left of the sequence reading as zero. Tap
+        // `kernel_size - 1` is the current position, so it always
+        // contributes and seeds the accumulator.
+        let tap = |j: usize| -> candle_core::Result<Tensor> {
+            self.conv1d_weight.narrow(2, j, 1)?.reshape((1, wide, 1))
+        };
+        let mut acc = prepended.broadcast_mul(&tap(self.kernel_size - 1)?)?;
+        for j in 0..self.kernel_size - 1 {
+            let lag = (self.kernel_size - 1 - j) * self.dilation;
+            // Every position this tap reads is left of the sequence, so
+            // it contributes only the zeros the padding would have held.
+            if lag >= prep_len {
+                continue;
+            }
+            let pad = Tensor::zeros((b, wide, lag), prepended.dtype(), prepended.device())?;
+            let shifted = Tensor::cat(&[&pad, &prepended.narrow(2, 0, prep_len - lag)?], 2)?;
+            acc = (acc + shifted.broadcast_mul(&tap(j)?)?)?;
+        }
+        let out = candle_nn::ops::silu(&acc)?;
         out.narrow(2, prep_len - l, l)?
             .transpose(1, 2)?
             .contiguous()
@@ -842,6 +1077,59 @@ mod tests {
         assert_eq!(ple.left_context(), dilation);
     }
 
+    /// Which tap reaches how far back, at the shipped `kernel_size`.
+    ///
+    /// Tap `j` must read position `t - (kernel_size - 1 - j) * dilation`,
+    /// so the weights run far-to-near and the *last* one is the current
+    /// position. Every other conv test here uses `kernel_size = 2`, where
+    /// that mapping and its reverse agree on the only non-zero lag — so
+    /// none of them can see a transposed tap order. The shipped config is
+    /// `kernel_size = 4`, where the four lags are 9, 6, 3 and 0: a single
+    /// impulse with four distinct weights lands each tap in its own
+    /// position, and reversing the order swaps the 9 and 3 responses.
+    #[test]
+    fn conv_taps_run_far_to_near_across_the_full_kernel() {
+        let dev = Device::Cpu;
+        let (wide, dilation, kernel_size) = (1, 3, 4);
+        // Far to near: w[0] reaches 9 back, w[3] is the current position.
+        let conv_w =
+            Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], (wide, 1, kernel_size), &dev).unwrap();
+        let mut ple = block(1, 1, kernel_size, dilation, eye(1), eye(1), conv_w);
+
+        let mut impulse = vec![0.0f32; 10];
+        impulse[0] = 1.0;
+        let x = Tensor::from_vec(impulse, (1, 10, wide), &dev).unwrap();
+        let got: Vec<f32> = ple
+            .short_conv(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let silu = |v: f32| v / (1.0 + (-v).exp());
+        // Weight 4 at lag 0, 3 at lag 3, 2 at lag 6, 1 at lag 9.
+        let want = [
+            silu(4.0),
+            0.0,
+            0.0,
+            silu(3.0),
+            0.0,
+            0.0,
+            silu(2.0),
+            0.0,
+            0.0,
+            silu(1.0),
+        ];
+        for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
+            assert!(
+                (g - w).abs() < 1e-5,
+                "position {i}: got {got:?} want {want:?}"
+            );
+        }
+        assert_eq!(ple.left_context(), (kernel_size - 1) * dilation);
+    }
+
     /// Decoding one token at a time must give the same answer as running
     /// the whole sequence at once — the property the rolling state
     /// exists for, and the one that breaks silently when the cached
@@ -903,6 +1191,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()
     }
 
@@ -505,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) {
@@ -649,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() {

crates/neuron/src/harness/candle.rs

@@ -25,6 +25,11 @@ use cortex_core::openai::{
     ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, MessageContent,
 };
 
+// `PhaseTiming` is constructed only on the two cuda-gated streaming
+// paths (device-worker and TP), so importing it unconditionally warns
+// on a CPU-only build and clippy runs with `-D warnings`.
+#[cfg(feature = "cuda")]
+use crate::wire::PhaseTiming;
 use crate::wire::{
     FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
     detect_reasoning_token_pair, detect_tool_call_token_pair, openai_chat as wire_chat,
@@ -1283,6 +1288,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
@@ -2738,6 +2766,12 @@ impl CandleHarness {
         // 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")?;
@@ -2860,9 +2894,18 @@ impl CandleHarness {
                     let model = super::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
                         &cfg,
                         dtype,
-                        isq,
                         &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)))
@@ -2899,6 +2942,67 @@ impl CandleHarness {
     /// 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,
@@ -3406,11 +3510,13 @@ impl CandleHarness {
                 let reasoning_tokens_inner = loaded.reasoning_tokens.clone();
                 let tool_call_tokens_inner = loaded.tool_call_tokens.clone();
                 let tool_schemas_inner = tool_schemas.clone();
+                let model_id_inner = model_id.clone();
                 tokio::spawn(
                     async move {
                         let _admit = admit;
                         let _inference_guard = loaded_for_task.inference_lock.lock().await;
                         match stream_inference_via_worker(
+                            model_id_inner,
                             worker,
                             handle,
                             tokenizer,
@@ -3957,8 +4063,35 @@ impl Harness for CandleHarness {
                 // 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)
+                    .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
@@ -4051,15 +4184,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(
@@ -4068,7 +4205,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,
@@ -4124,7 +4261,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,
@@ -4780,7 +4917,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),
@@ -4823,13 +4960,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 {
@@ -4837,7 +4975,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),
                 },
@@ -5191,6 +5329,9 @@ impl CandleHarness {
                 // read them; populated at the prefill→decode boundary inside.
                 let mut prefill_ms_measured: u32 = 0;
                 let mut decode_start: Option<std::time::Instant> = None;
+                let mut phase_forward = std::time::Duration::ZERO;
+                let mut phase_sample = std::time::Duration::ZERO;
+                let mut phase_emit = std::time::Duration::ZERO;
 
                 // Hoisted out of `'work` so the completion path below
                 // can report it (#269) — the reuse count is decided
@@ -5351,6 +5492,12 @@ impl CandleHarness {
                         };
                     // Decode-phase timer for the Finish prefill/decode split (#85).
                     decode_start = Some(std::time::Instant::now());
+                    // Phase split (see `PhaseTiming`). Instrumented here
+                    // as well as on the single-GPU path so a healthy
+                    // model is a control for a slow one: both archs share
+                    // this loop, this sampler and a 248,320 vocabulary,
+                    // so a phase that is cheap here and expensive there
+                    // is a property of the arch, not of the loop.
 
                     if Some(next_token) == eos_id {
                         finish_reason = FinishReason::Stop;
@@ -5437,6 +5584,7 @@ impl CandleHarness {
                         }
 
                         for index in 0..max_new.saturating_sub(1) {
+                            let t_forward = std::time::Instant::now();
                             let logits_vec = match pool
                                 .generate_step(
                                     &model_id,
@@ -5452,6 +5600,8 @@ impl CandleHarness {
                                     break 'work;
                                 }
                             };
+                            phase_forward += t_forward.elapsed();
+                            let t_sample = std::time::Instant::now();
                             let logits = match Tensor::new(logits_vec.as_slice(), &Device::Cpu) {
                                 Ok(t) => t,
                                 Err(e) => {
@@ -5480,6 +5630,11 @@ impl CandleHarness {
                                     break 'work;
                                 }
                             };
+                            phase_sample += t_sample.elapsed();
+                            // NB: this per-step VRAM query is a second
+                            // channel round-trip per token. It is charged
+                            // to neither phase, so it lands in the
+                            // residual between the three and `decode_ms`.
                             // Always await the query (even when the
                             // trace! is filtered out by RUST_LOG): the
                             // channel hop is ~tens of µs, comparable to
@@ -5499,6 +5654,7 @@ impl CandleHarness {
                                 finish_reason = FinishReason::Stop;
                                 break;
                             }
+                            let t_emit = std::time::Instant::now();
                             all_tokens.push(next_token);
                             match handle_tool_call_marker(
                                 next_token,
@@ -5580,6 +5736,7 @@ impl CandleHarness {
                                     "TP stream: decode_stream step failed"
                                 ),
                             }
+                            phase_emit += t_emit.elapsed();
                         }
                     }
                 }
@@ -5644,8 +5801,24 @@ impl CandleHarness {
                         decode_ms = decode_start
                             .map(|d| d.elapsed().as_millis() as u32)
                             .unwrap_or(0),
+                        forward_ms = phase_forward.as_millis(),
+                        sample_ms = phase_sample.as_millis(),
+                        emit_ms = phase_emit.as_millis(),
                         "chat_completion (stream): done"
                     );
+                    crate::metrics::record_finish(
+                        &model_id,
+                        prefill_ms_measured,
+                        decode_start
+                            .map(|d| d.elapsed().as_millis() as u32)
+                            .unwrap_or(0),
+                        all_tokens.len() as u32,
+                        Some((
+                            phase_forward.as_millis() as u32,
+                            phase_sample.as_millis() as u32,
+                            phase_emit.as_millis() as u32,
+                        )),
+                    );
                     let _ = tx
                         .send(InferenceEvent::Finish {
                             reason: finish_reason,
@@ -5659,6 +5832,11 @@ impl CandleHarness {
                                     .map(|d| d.elapsed().as_millis() as u32)
                                     .unwrap_or(0),
                                 prefill_tokens: prompt_len as u32,
+                                phases: Some(PhaseTiming {
+                                    forward_ms: phase_forward.as_millis() as u32,
+                                    sample_ms: phase_sample.as_millis() as u32,
+                                    emit_ms: phase_emit.as_millis() as u32,
+                                }),
                             }),
                         })
                         .await;
@@ -7105,6 +7283,9 @@ async fn run_inference_via_worker(
 #[cfg(feature = "cuda")]
 #[allow(clippy::too_many_arguments)]
 async fn stream_inference_via_worker(
+    // Owned: this runs in a spawned task, and the only other source of
+    // the model name here is the tracing span, which is not a value.
+    model_id: String,
     worker: Arc<super::device_worker::DeviceWorkerHandle>,
     handle: super::device_worker::ArchHandle,
     tokenizer: Tokenizer,
@@ -7238,6 +7419,13 @@ async fn stream_inference_via_worker(
     };
     // Decode-phase timer for the Finish prefill/decode split (#85).
     let decode_start = std::time::Instant::now();
+    // Phase split of the decode loop. Each boundary below is already a
+    // synchronisation point -- the worker hands back CPU-side logits --
+    // so this adds no device sync and measures elapsed work rather than
+    // kernel-launch time.
+    let mut phase_forward = std::time::Duration::ZERO;
+    let mut phase_sample = std::time::Duration::ZERO;
+    let mut phase_emit = std::time::Duration::ZERO;
 
     // Per-token routing. `tokenizers::DecodeStream` carries five
     // generic parameters (`M, N, PT, PP, D`) which makes naming
@@ -7344,10 +7532,14 @@ async fn stream_inference_via_worker(
     }
 
     for index in 0..max_new.saturating_sub(1) {
+        let t_forward = std::time::Instant::now();
         let logits_vec = worker
             .forward_logits(handle, vec![next_token], prompt_len + index)
             .await
             .map_err(|e| anyhow::anyhow!("decode step {index}: {e:#}"))?;
+        phase_forward += t_forward.elapsed();
+
+        let t_sample = std::time::Instant::now();
         let logits = Tensor::new(logits_vec.as_slice(), &Device::Cpu)?;
         next_token = match sample_with_penalty(
             &logits,
@@ -7367,11 +7559,15 @@ async fn stream_inference_via_worker(
                 return Err(e);
             }
         };
+        phase_sample += t_sample.elapsed();
         if Some(next_token) == eos_id {
             finish_reason = FinishReason::Stop;
             break;
         }
-        if !route_token!(next_token) {
+        let t_emit = std::time::Instant::now();
+        let routed = route_token!(next_token);
+        phase_emit += t_emit.elapsed();
+        if !routed {
             return Ok(finish_reason.as_openai_str().to_string());
         }
     }
@@ -7393,8 +7589,22 @@ async fn stream_inference_via_worker(
         finish_reason = ?finish_reason,
         prefill_ms = prefill_elapsed.as_millis(),
         decode_ms = decode_start.elapsed().as_millis(),
+        forward_ms = phase_forward.as_millis(),
+        sample_ms = phase_sample.as_millis(),
+        emit_ms = phase_emit.as_millis(),
         "chat_completion (stream): done"
     );
+    crate::metrics::record_finish(
+        &model_id,
+        prefill_elapsed.as_millis() as u32,
+        decode_start.elapsed().as_millis() as u32,
+        all_tokens.len() as u32,
+        Some((
+            phase_forward.as_millis() as u32,
+            phase_sample.as_millis() as u32,
+            phase_emit.as_millis() as u32,
+        )),
+    );
     let _ = tx
         .send(InferenceEvent::Finish {
             reason: finish_reason,
@@ -7406,6 +7616,11 @@ async fn stream_inference_via_worker(
                 prefill_ms: prefill_elapsed.as_millis() as u32,
                 decode_ms: decode_start.elapsed().as_millis() as u32,
                 prefill_tokens: prefill_prompt_len as u32,
+                phases: Some(PhaseTiming {
+                    forward_ms: phase_forward.as_millis() as u32,
+                    sample_ms: phase_sample.as_millis() as u32,
+                    emit_ms: phase_emit.as_millis() as u32,
+                }),
             }),
         })
         .await;
@@ -7626,6 +7841,7 @@ fn run_inference_streaming(
             prefill_ms: prefill_elapsed.as_millis() as u32,
             decode_ms: decode_start.elapsed().as_millis() as u32,
             prefill_tokens: prompt_tokens.len() as u32,
+            phases: None,
         }),
     });
     Ok(())
@@ -7647,6 +7863,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,

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

@@ -147,6 +147,8 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 safetensors_paths,
                 model_id,
                 quant,
+                isq_cache,
+                experts_on_host,
                 reply,
             } => {
                 let result = load_dense_inner(
@@ -155,6 +157,8 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                     &safetensors_paths,
                     &model_id,
                     quant,
+                    isq_cache.as_ref(),
+                    experts_on_host,
                 )
                 .map(|arch| {
                     // Ask the model, once, while it is still in
@@ -171,6 +175,16 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 // 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 } => {
@@ -194,6 +208,9 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 // 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,
@@ -934,12 +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;
@@ -948,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
@@ -1053,12 +1140,42 @@ fn load_dense_inner(
                 )
                 .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,
-                quant,
                 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)))
@@ -1770,11 +1887,18 @@ mod tests {
         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) {
-                    Ok(_) => panic!("a bare config must not build a model"),
-                    Err(e) => format!("{e:#}"),
-                };
+            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 \

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

@@ -132,6 +132,11 @@ pub enum Job {
         /// 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

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

@@ -222,6 +222,8 @@ impl DeviceWorkerHandle {
         safetensors_paths: Vec<std::path::PathBuf>,
         model_id: String,
         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 {
@@ -235,6 +237,8 @@ impl DeviceWorkerHandle {
                 safetensors_paths,
                 model_id,
                 quant,
+                isq_cache,
+                experts_on_host,
                 reply: reply_tx,
             })
             .map_err(|_| WorkerError::Gone {
@@ -1246,6 +1250,8 @@ mod tests {
                 vec![fixture.join("model.safetensors")],
                 "qwen3_next-tiny".into(),
                 None,
+                None,
+                false,
             )
             .await
             .expect("load tiny fixture")

crates/neuron/src/harness/engine.rs

@@ -663,6 +663,7 @@ async fn finish_slot(slot: &mut Slot, reason: FinishReason, rate: &PrefillRateEm
                 prefill_ms: slot.prefill_ms,
                 decode_ms: slot.decode_start.elapsed().as_millis() as u32,
                 prefill_tokens: slot.prefill_tokens,
+                phases: None,
             },
         })
         .await;
@@ -1151,6 +1152,7 @@ mod tests {
                     prefill_ms: 6000,
                     decode_ms: 900_000,
                     prefill_tokens: 7413,
+                    phases: None,
                 },
             })
             .await
@@ -1262,6 +1264,8 @@ mod tests {
                 vec![fixture.join("model.safetensors")],
                 "qwen3_next-tiny".into(),
                 None,
+                None,
+                false,
             )
             .await
             .expect("load fixture")

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,6 +9,7 @@ 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;

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/lib.rs

@@ -5,6 +5,7 @@ pub mod cuda;
 pub mod discovery;
 pub mod harness;
 pub mod health;
+pub mod metrics;
 pub mod startup;
 pub mod version;
 pub mod wire;

crates/neuron/src/main.rs

@@ -220,6 +220,12 @@ async fn daemon(args: Args) -> Result<()> {
     // host look down to anything probing `/health` during pre-warm.
     // The pre-warm task runs in the background instead — `/health`
     // surfaces its progress via the activation field.
+    // Before the listener binds, so the very first scrape after a
+    // restart finds a recorder rather than a 503.
+    if let Err(e) = neuron::metrics::install() {
+        tracing::warn!(error = %e, "metrics recorder not installed; /metrics will 503");
+    }
+
     let app = api::neuron_routes().with_state(Arc::clone(&state));
     let addr: std::net::SocketAddr = format!("0.0.0.0:{port}").parse()?;
     let listener = tokio::net::TcpListener::bind(addr).await?;

crates/neuron/src/metrics.rs

@@ -0,0 +1,173 @@
+//! Prometheus metrics for neuron.
+//!
+//! neuron had no `/metrics` surface at all: every signal it produced
+//! reached Prometheus only by cortex polling `/health` every ~10 s and
+//! republishing it as gauges. That is enough for routing decisions and
+//! useless for asking where a request's time went — an EMA sampled at
+//! 10 s cannot see inside a one-second forward, and the finest split
+//! that existed anywhere was prefill versus decode.
+//!
+//! Served from the main API port rather than a second listener: neuron
+//! already runs an axum server, and a separate port is one more thing
+//! to open in the firewall for no benefit.
+
+use anyhow::Result;
+use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
+use std::sync::OnceLock;
+
+static HANDLE: OnceLock<PrometheusHandle> = OnceLock::new();
+
+/// Install the process-wide recorder. Idempotent in effect: a second
+/// call leaves the first handle in place.
+pub fn install() -> Result<()> {
+    let handle = with_buckets(PrometheusBuilder::new())?
+        .install_recorder()
+        .map_err(|e| anyhow::anyhow!("failed to install Prometheus recorder: {e}"))?;
+    describe();
+    let _ = HANDLE.set(handle);
+    Ok(())
+}
+
+/// Render the current metrics, or `None` if the recorder was never
+/// installed. A `None` here must surface as an error rather than an
+/// empty body: an empty scrape is indistinguishable from a healthy
+/// process with nothing to report.
+pub fn render() -> Option<String> {
+    HANDLE.get().map(|h| h.render())
+}
+
+/// Explicit buckets, or `metrics-exporter-prometheus` renders every
+/// histogram as a **summary** — `{quantile=...}` series, no `_bucket`,
+/// no aggregation across hosts, and every quantile decaying to `0`
+/// when idle so "no traffic" looks like "instant". The same trap
+/// cortex's exporter documents; the fix has to be repeated per metric.
+fn with_buckets(builder: PrometheusBuilder) -> Result<PrometheusBuilder> {
+    // Per-request phase sums: microseconds for an emit on a fast model,
+    // minutes for the forward on a long turn.
+    let phase_seconds = &[
+        0.001, 0.005, 0.025, 0.1, 0.25, 1.0, 2.5, 5.0, 15.0, 30.0, 60.0, 300.0,
+    ];
+    let seconds_short = &[
+        0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0, 120.0,
+    ];
+    let seconds_long = &[
+        0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1200.0,
+    ];
+    builder
+        .set_buckets_for_metric(
+            Matcher::Full("neuron_decode_phase_seconds".into()),
+            phase_seconds,
+        )?
+        .set_buckets_for_metric(
+            Matcher::Full("neuron_prefill_seconds".into()),
+            seconds_short,
+        )?
+        .set_buckets_for_metric(Matcher::Full("neuron_decode_seconds".into()), seconds_long)
+        .map_err(|e| anyhow::anyhow!("failed to configure histogram buckets: {e}"))
+}
+
+fn describe() {
+    metrics::describe_histogram!(
+        "neuron_decode_phase_seconds",
+        "Decode-loop time by phase, summed per request. The `_sum` \
+         series is the point: ratio one phase against the total to see \
+         where decode goes, without needing per-step resolution."
+    );
+    metrics::describe_histogram!("neuron_prefill_seconds", "Prefill wall-clock per request");
+    metrics::describe_histogram!("neuron_decode_seconds", "Decode wall-clock per request");
+    metrics::describe_counter!(
+        "neuron_completion_tokens_total",
+        "Tokens generated, the denominator for per-token phase cost"
+    );
+}
+
+/// Record one finished inference.
+///
+/// `phases` is optional because not every serving path measures them,
+/// and a path that does not must emit **nothing** rather than zeros —
+/// a zero sample pulls the histogram's mean down and reads as "this
+/// phase is free here", which is the opposite of "nobody looked".
+///
+/// The residual is emitted as its own phase. Whatever the brackets
+/// failed to cover is then visible as a labelled series instead of
+/// being silently divided among the phases that were measured.
+pub fn record_finish(
+    model: &str,
+    prefill_ms: u32,
+    decode_ms: u32,
+    completion_tokens: u32,
+    phases: Option<(u32, u32, u32)>,
+) {
+    let m = model.to_string();
+    metrics::histogram!("neuron_prefill_seconds", "model" => m.clone())
+        .record(prefill_ms as f64 / 1e3);
+    metrics::histogram!("neuron_decode_seconds", "model" => m.clone())
+        .record(decode_ms as f64 / 1e3);
+    metrics::counter!("neuron_completion_tokens_total", "model" => m.clone())
+        .increment(completion_tokens as u64);
+
+    let Some((forward_ms, sample_ms, emit_ms)) = phases else {
+        return;
+    };
+    for (phase, ms) in [
+        ("forward", forward_ms),
+        ("sample", sample_ms),
+        ("emit", emit_ms),
+    ] {
+        metrics::histogram!("neuron_decode_phase_seconds", "model" => m.clone(), "phase" => phase)
+            .record(ms as f64 / 1e3);
+    }
+    // Saturating: the phases are measured inside the decode window, so
+    // they cannot legitimately exceed it, but clock granularity can
+    // round them a millisecond over and an underflow here would report
+    // an enormous residual.
+    let residual =
+        decode_ms.saturating_sub(forward_ms.saturating_add(sample_ms).saturating_add(emit_ms));
+    metrics::histogram!("neuron_decode_phase_seconds", "model" => m, "phase" => "residual")
+        .record(residual as f64 / 1e3);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// An unmeasured path must contribute no phase samples at all.
+    /// Zeros would drag the histogram's mean toward zero and read as
+    /// "this phase costs nothing on this model".
+    #[test]
+    fn phases_absent_means_no_phase_series() {
+        let handle = with_buckets(PrometheusBuilder::new())
+            .expect("buckets")
+            .install_recorder()
+            .expect("recorder");
+        describe();
+
+        record_finish("measured", 100, 1000, 64, Some((900, 50, 10)));
+        record_finish("unmeasured", 100, 1000, 64, None);
+        let out = handle.render();
+
+        assert!(
+            out.contains("neuron_decode_phase_seconds") && out.contains(r#"model="measured""#),
+            "measured model must emit phases:\n{out}"
+        );
+        for line in out
+            .lines()
+            .filter(|l| l.contains("neuron_decode_phase_seconds"))
+        {
+            assert!(
+                !line.contains(r#"model="unmeasured""#),
+                "unmeasured model must emit no phase series, got: {line}"
+            );
+        }
+        // Buckets, not a summary: `histogram_quantile()` needs `_bucket`.
+        assert!(
+            out.contains("neuron_decode_phase_seconds_bucket"),
+            "must export as a histogram, not a summary:\n{out}"
+        );
+        // 1000 - (900 + 50 + 10) = 40 ms of residual, as its own series.
+        assert!(
+            out.contains(r#"phase="residual""#),
+            "residual must be its own series:\n{out}"
+        );
+    }
+}

crates/neuron/src/wire/event.rs

@@ -121,6 +121,41 @@ pub struct FinishTiming {
     /// the elapsed `prefill_ms` drops while this stays the full prompt
     /// length, so a high implied rate is itself the cache-hit signal.
     pub prefill_tokens: u32,
+    /// How the decode loop's wall-clock divides, when the serving path
+    /// measured it. `None` means "this path does not report phases",
+    /// which is deliberately distinct from "measured, and they were
+    /// zero" — a phase split that cannot express its own absence
+    /// invites reading an unmeasured path as an idle one.
+    pub phases: Option<PhaseTiming>,
+}
+
+/// Where a decode loop's wall-clock actually went, summed over every
+/// step of the request.
+///
+/// These are sums across steps, not a breakdown of one step, and each
+/// boundary is a point the loop already synchronises at: the forward
+/// hands back CPU-side logits, so no device synchronisation is added
+/// and the numbers are elapsed work rather than kernel-launch time.
+///
+/// The three should roughly account for `decode_ms`; a large residual
+/// is itself a finding (it is time inside the loop that none of the
+/// brackets cover).
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct PhaseTiming {
+    /// Submitting the step to the device worker and getting logits
+    /// back: the channel hop, the model forward, and the device-to-host
+    /// copy of the logits row. For a large vocabulary that copy is not
+    /// negligible -- 248,320 f32 is ~993 KB per step.
+    pub forward_ms: u32,
+    /// Wrapping the logits in a CPU tensor and sampling: repeat/presence
+    /// penalties plus the `LogitsProcessor`. Scales with vocabulary,
+    /// not with model size, and top-p sorts the distribution.
+    pub sample_ms: u32,
+    /// Detokenising the sampled token and handing the delta to the
+    /// consumer. **Includes backpressure**: the send awaits a bounded
+    /// channel, so a slow reader is charged here. This is emit cost,
+    /// not detokenisation cost, and the two are not separated.
+    pub emit_ms: u32,
 }
 
 /// Why a stream stopped. Stays small on purpose — anything that

crates/neuron/src/wire/mod.rs

@@ -22,6 +22,6 @@ pub mod openai_chat;
 pub mod openai_responses;
 
 pub use event::{
-    FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
+    FinishReason, FinishTiming, InferenceEvent, PhaseTiming, ReasoningTokenPair, ToolCallTokenPair,
     detect_reasoning_token_pair, detect_tool_call_token_pair,
 };

crates/neuron/src/wire/openai_chat.rs

@@ -444,6 +444,9 @@ fn usage_chunk(
                 prefill_ms: t.prefill_ms as u64,
                 decode_ms: t.decode_ms as u64,
                 prefill_tokens: t.prefill_tokens as u64,
+                forward_ms: t.phases.map(|p| p.forward_ms as u64),
+                sample_ms: t.phases.map(|p| p.sample_ms as u64),
+                emit_ms: t.phases.map(|p| p.emit_ms as u64),
             }),
         }),
         extra: serde_json::Value::Object(Default::default()),
@@ -596,6 +599,9 @@ pub async fn collect_chat_completion(
                 prefill_ms: t.prefill_ms as u64,
                 decode_ms: t.decode_ms as u64,
                 prefill_tokens: t.prefill_tokens as u64,
+                forward_ms: t.phases.map(|p| p.forward_ms as u64),
+                sample_ms: t.phases.map(|p| p.sample_ms as u64),
+                emit_ms: t.phases.map(|p| p.emit_ms as u64),
             }),
         }),
         extra: serde_json::Value::Object(Default::default()),
@@ -685,6 +691,7 @@ mod tests {
         );
     }
     use super::*;
+    use crate::wire::event::PhaseTiming;
 
     // ── Non-streaming collector (#285) ────────────────────────────
 
@@ -799,6 +806,7 @@ mod tests {
                 prefill_ms: 12,
                 decode_ms: 34,
                 prefill_tokens: 100,
+                phases: None,
             }),
         })
         .await
@@ -813,6 +821,65 @@ mod tests {
         assert_eq!((t.prefill_ms, t.decode_ms, t.prefill_tokens), (12, 34, 100));
     }
 
+    /// The phase split must survive to the wire, and an unmeasured path
+    /// must serialise as **absent** rather than as zeros.
+    ///
+    /// That distinction is the reason the field is an `Option` at all:
+    /// a consumer reading `forward_ms: 0` off a path that never
+    /// measured would conclude the forward is free. The serialised
+    /// form is asserted, not just the struct, because
+    /// `skip_serializing_if` is what actually enforces it and a plain
+    /// field comparison cannot see it.
+    #[tokio::test]
+    async fn phase_timing_reaches_the_wire_and_absence_is_not_zero() {
+        async fn timing_json(phases: Option<PhaseTiming>) -> serde_json::Value {
+            let (tx, rx) = mpsc::channel(8);
+            tx.send(InferenceEvent::Finish {
+                reason: FinishReason::Stop,
+                prompt_tokens: 10,
+                completion_tokens: 5,
+                reasoning_tokens: 0,
+                cached_tokens: 0,
+                timing: Some(FinishTiming {
+                    prefill_ms: 12,
+                    decode_ms: 340,
+                    prefill_tokens: 10,
+                    phases,
+                }),
+            })
+            .await
+            .expect("send");
+            drop(tx);
+            let r = collect_chat_completion(rx, "id".into(), 1, "m".into())
+                .await
+                .expect("collected");
+            serde_json::to_value(r.usage.expect("usage").helexa_timing.expect("timing"))
+                .expect("serialise")
+        }
+
+        let measured = timing_json(Some(PhaseTiming {
+            forward_ms: 300,
+            sample_ms: 25,
+            emit_ms: 4,
+        }))
+        .await;
+        assert_eq!(measured["forward_ms"], 300);
+        assert_eq!(measured["sample_ms"], 25);
+        assert_eq!(measured["emit_ms"], 4);
+
+        let unmeasured = timing_json(None).await;
+        for f in ["forward_ms", "sample_ms", "emit_ms"] {
+            assert!(
+                unmeasured.get(f).is_none(),
+                "unmeasured {f} must be absent, not {:?} -- a zero here reads as \
+                 'this phase is free' rather than 'nobody measured it'",
+                unmeasured.get(f)
+            );
+        }
+        // The prefill/decode split is unconditional and must be unaffected.
+        assert_eq!(unmeasured["decode_ms"], 340);
+    }
+
     /// A dropped engine slot ends the channel with no Finish. Returning
     /// a 200 with partial text would hide a fatal worker fault.
     #[tokio::test]
@@ -898,6 +965,7 @@ mod tests {
                 prefill_ms: 200,
                 decode_ms: 1500,
                 prefill_tokens: 128,
+                phases: None,
             }),
         })
         .await

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/api.rs

@@ -99,6 +99,49 @@ async fn test_version_endpoint() {
     assert!(!body.git_sha.is_empty());
 }
 
+#[tokio::test]
+async fn test_metrics_endpoint_serves_prometheus_text() {
+    // The recorder is process-wide and other tests in this binary may
+    // have installed it already; either way the route must answer, and
+    // must never answer 200-with-empty-body (indistinguishable from a
+    // healthy process reporting nothing).
+    let _ = neuron::metrics::install();
+    neuron::metrics::record_finish("probe-model", 100, 1000, 64, Some((900, 50, 10)));
+    let url = spawn_neuron(fake_discovery()).await;
+
+    let resp = reqwest::Client::new()
+        .get(format!("{url}/metrics"))
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(resp.status(), 200);
+    let ct = resp
+        .headers()
+        .get(reqwest::header::CONTENT_TYPE)
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or_default()
+        .to_string();
+    assert!(
+        ct.starts_with("text/plain"),
+        "Prometheus needs a text/plain scrape, got {ct:?}"
+    );
+
+    let body = resp.text().await.expect("body");
+    // The exporter renders only metrics that carry samples -- a
+    // `describe_*` on its own produces nothing -- so the scrape has to
+    // be checked after a recorded request, which is also the path that
+    // actually matters: record here, readable there.
+    assert!(
+        body.contains("neuron_decode_phase_seconds"),
+        "a recorded request must appear in the scrape:\n{body}"
+    );
+    assert!(
+        body.contains(r#"phase="forward""#) && body.contains(r#"phase="residual""#),
+        "phase labels must survive to the scrape:\n{body}"
+    );
+}
+
 #[tokio::test]
 async fn test_health_endpoint() {
     let url = spawn_neuron(fake_discovery()).await;

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

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