perf(laya): fused norms, mask elision, and a flash-attn varlen path (#337)

Measured on an RTX 3060 (sm_86), bf16, forward only, median of 10.
The PyTorch reference's end-to-end numbers are the bar.

                      ref    start  +norms  +masks/scale  +flash-attn
  short, 1 question   16      10.7    8.5       10.2           7.7 ms
  short, 5 questions  25      25.6   20.5       20.7          17.0 ms
  512 tok, 1 q        39      74.3   61.8       54.3          36.2 ms
  512 tok, 5 q       152     352    292        254           162 ms
  512 tok, x64        -      14.5   17         20             32 seq/s

The levers, in the order a profile ranked them:

- Fused LayerNorm. candle_nn::LayerNorm only takes its fused kernel
  when it has a bias; ModernBERT's norms have none, so each fell back
  to six elementwise kernels, a quarter of GPU time. An explicit zero
  bias gives the same result in one kernel.
- No mask when nothing is masked. An unpadded batch (every single-
  question request) skips the global mask, and rows the window cannot
  exceed skip the sliding one; the scale is applied to q rather than to
  the (b, h, l, l) scores. Each saves a bandwidth-bound pass.
- Flash attention, behind the flash-attn feature, on CUDA in reduced
  precision: the batch is packed without padding (total, d) and every
  attention -- sliding layers windowed at local_attention/2 each side,
  global and head layers unwindowed -- runs in flash_attn_varlen. Linear
  layers then do no work for padding either. NEURON_FLASH_ATTN=0 falls
  back to the eager path at runtime, as for the text architectures.

Parity, bf16 against the f32 reference, centred drift (mean / max),
next to the reference's own bf16 GPU run:

  English       0.029 / 0.229   (reference 0.034 / 0.264)
  multilingual  0.030 / 0.089   (reference 0.028 / 0.067)
  argmax 33/33 throughout; f32 unchanged (<1e-4).

The released-weights test now gates on that comparison: mean centred
drift no worse than twice the reference's own bf16 drift, plus 0.01.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VtaCXFAY23HmwX5Cyc6jhC

Cargo.lock

@@ -3147,7 +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

@@ -12,6 +12,14 @@ path = "src/lib.rs"
 name = "neuron"
 path = "src/main.rs"
 
+# Measurement harness for the MTP speculative-decoding stages (#96).
+# Not shipped in the RPM: it loads a checkpoint directly and reports
+# numbers, and exists so the go/no-go gate is a measurement rather than
+# an argument.
+[[bin]]
+name = "neuron-mtp-probe"
+path = "src/bin/mtp_probe.rs"
+
 [features]
 default = []
 # Enables CUDA acceleration in candle and the cudarc/nccl bindings the
@@ -47,6 +55,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
@@ -110,6 +120,12 @@ safetensors = "0.8"
 # 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())
 }
@@ -164,6 +189,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
         PreflightError::EmptyRepo { .. } => "empty_repo",
         PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
         PreflightError::QuantNotFound { .. } => "quant_not_found",
+        PreflightError::Unsupported { .. } => "unsupported_load_option",
     }
 }
 

crates/neuron/src/bin/mtp_probe.rs

@@ -0,0 +1,730 @@
+//! MTP speculative-decoding measurement harness (#96, stages S3/S4).
+//!
+//! Two numbers decide whether speculative decoding is worth building
+//! into neuron, and neither can be taken from anyone else's benchmark:
+//!
+//! - **`c`** — what a draft step costs relative to a target decode
+//!   step. Published figures are for a different drafter on different
+//!   hardware. Ours reads the target's own `lm_head` (~1.03 GB at q6k
+//!   on the 27B) on top of the head's weights, so the vocabulary, not
+//!   the head, may dominate.
+//! - **acceptance** — how many drafted tokens the target actually
+//!   agrees with, *on our workloads*. Every published figure is
+//!   short-form (GSM8K, MATH-500, HumanEval, ≤4k output); helexa serves
+//!   15–50k-token agentic contexts full of replayed reasoning, and
+//!   acceptance is workload-dependent. A poor rate makes speculation a
+//!   net loss.
+//!
+//! This binary answers both by running a real checkpoint. It commits
+//! nothing, changes no serving behaviour, and is not shipped in the
+//! RPM.
+//!
+//! ## What it does
+//!
+//! Prefills the target over a prompt, then generates greedily. At each
+//! step it snapshots the draft head's KV, drafts `K` tokens by chaining
+//! the head, restores the snapshot, and lets the target continue
+//! normally. When the target has produced the next `K` true tokens, the
+//! draft is scored against them: the accepted length is the number of
+//! leading matches.
+//!
+//! The head is *not* the thing generating — the target's own output is
+//! never influenced by the draft, which is what makes this an observer
+//! rather than an implementation.
+//!
+//! ## Two logits, two argmaxes
+//!
+//! Acceptance is reported against the target's raw argmax **and**
+//! against its argmax after the repeat penalty neuron applies when
+//! serving. The two differ, and the gap is exactly what a
+//! distribution-preserving (rejection-sampling) scheme would have to
+//! reconcile at `temperature > 0` — this fleet serves this family at
+//! 0.6. Measuring it here costs nothing and makes that decision an
+//! informed one.
+//!
+//! ## Usage
+//!
+//! ```sh
+//! neuron-mtp-probe --model /path/to/Qwen3.5-0.8B --device cuda:0 \
+//!     --draft-len 4 --steps 64 --prompt-file prompt.txt
+//! ```
+
+use anyhow::{Context, Result};
+use candle_core::{DType, Device, IndexOp, Tensor};
+use clap::Parser;
+use std::path::PathBuf;
+use std::time::Instant;
+
+use neuron::harness::arch::qwen3_5::mtp::MtpHead;
+use neuron::harness::arch::qwen3_5::rope::RotaryEmbedding;
+use neuron::harness::arch::qwen3_5::{Config, Qwen3_5ForCausalLM};
+
+#[derive(Parser, Debug)]
+#[command(about = "Measure MTP draft cost and acceptance rate (#96)")]
+struct Args {
+    /// Checkpoint directory: config.json, *.safetensors, tokenizer.json.
+    #[arg(long)]
+    model: PathBuf,
+    /// `cpu` or `cuda:N`.
+    #[arg(long, default_value = "cpu")]
+    device: String,
+    /// Draft tokens per round (K).
+    #[arg(long, default_value_t = 4)]
+    draft_len: usize,
+    /// Decode steps to measure over.
+    #[arg(long, default_value_t = 64)]
+    steps: usize,
+    /// Prompt file. Without one, a short built-in prompt is used — fine
+    /// for `c`, useless for acceptance, which is the whole point of
+    /// measuring on a replayed session instead.
+    #[arg(long)]
+    prompt_file: Option<PathBuf>,
+    /// Prefill chunk size. The mask for a one-shot prefill is L x L —
+    /// 2.5 GB at 25k tokens — so the probe chunks like the serving path
+    /// does.
+    #[arg(long, default_value_t = 512)]
+    prefill_chunk: usize,
+    /// Timed-but-discarded rounds before measurement starts.
+    #[arg(long, default_value_t = 3)]
+    warmup: usize,
+    /// Repeat penalty applied to the target's logits when scoring the
+    /// second argmax. Mirrors the serving default.
+    #[arg(long, default_value_t = 1.05)]
+    repeat_penalty: f32,
+    /// How many recent tokens the repeat penalty considers.
+    #[arg(long, default_value_t = 64)]
+    repeat_last_n: usize,
+    /// Tensor-parallel degree. 1 loads the model on one card; 2+ uses
+    /// the worker pool, which is the only way to reach a model too big
+    /// for a single GPU — the 27B is 54 GB in bf16 and the single-GPU
+    /// path is bf16-only.
+    #[arg(long, default_value_t = 1)]
+    tensor_parallel: u32,
+    /// Quantisation for the TP load (`q6k`, `q8_0`, …). Only the TP
+    /// loader applies it.
+    #[arg(long)]
+    quant: Option<String>,
+    /// The binary the worker ranks are spawned from. Defaults to
+    /// `neuron` beside this executable — the probe itself does not
+    /// implement `--worker`.
+    #[arg(long)]
+    worker_binary: Option<PathBuf>,
+    /// How many verify-shaped forwards to time after the measurement
+    /// loop. A single sample moved 57.3 -> 50.1 ms between otherwise
+    /// identical runs — ±7 ms on a 45 ms baseline, enough to move the
+    /// predicted speedup across the gate on its own.
+    #[arg(long, default_value_t = 9)]
+    verify_samples: usize,
+    /// Write the summary as JSON here as well as to stdout.
+    #[arg(long)]
+    json: Option<PathBuf>,
+}
+
+fn pick_device(spec: &str) -> Result<Device> {
+    if spec == "cpu" {
+        return Ok(Device::Cpu);
+    }
+    let idx: usize = spec
+        .strip_prefix("cuda:")
+        .ok_or_else(|| anyhow::anyhow!("--device must be 'cpu' or 'cuda:N', got '{spec}'"))?
+        .parse()
+        .context("parse cuda device index")?;
+    #[cfg(feature = "cuda")]
+    {
+        Ok(Device::new_cuda(idx)?)
+    }
+    #[cfg(not(feature = "cuda"))]
+    {
+        let _ = idx;
+        anyhow::bail!("this binary was built without the `cuda` feature")
+    }
+}
+
+fn argmax(row: &[f32]) -> u32 {
+    let mut best = 0usize;
+    for (i, v) in row.iter().enumerate() {
+        if v > &row[best] {
+            best = i;
+        }
+    }
+    best as u32
+}
+
+/// The target's argmax after the repeat penalty neuron applies when
+/// serving. Same shape of penalty as `sample_with_penalty`: recent
+/// tokens are divided (positive logits) or multiplied (negative ones).
+fn argmax_with_penalty(row: &[f32], recent: &[u32], penalty: f32) -> u32 {
+    if penalty == 1.0 {
+        return argmax(row);
+    }
+    let mut adjusted = row.to_vec();
+    for &t in recent {
+        let i = t as usize;
+        if i < adjusted.len() {
+            adjusted[i] = if adjusted[i] > 0.0 {
+                adjusted[i] / penalty
+            } else {
+                adjusted[i] * penalty
+            };
+        }
+    }
+    argmax(&adjusted)
+}
+
+/// Additive causal mask, `(1, 1, l, l + offset)`.
+fn causal_mask(l: usize, offset: usize, dtype: DType, dev: &Device) -> Result<Tensor> {
+    let total = l + offset;
+    let mut data = vec![0f32; l * total];
+    for i in 0..l {
+        for j in 0..total {
+            if j > i + offset {
+                data[i * total + j] = f32::NEG_INFINITY;
+            }
+        }
+    }
+    Ok(Tensor::from_vec(data, (1, 1, l, total), dev)?.to_dtype(dtype)?)
+}
+
+/// Walk the prompt through target and head, in chunks.
+///
+/// Two reasons this is chunked rather than one forward: the causal mask
+/// is `L x L` — 2.5 GB at 25k tokens, which OOMs a 12 GB card before a
+/// single weight is touched — and the serving path chunks too, so this
+/// keeps the measurement on the same shape of work.
+///
+/// The head trails the target by one token: at position `i` it takes
+/// the embedding of token `i+1` with the target's hidden state at `i`.
+/// Each chunk therefore reaches one token into the next. The final
+/// prompt position has no successor yet — the first generated token
+/// fills that slot.
+///
+/// Returns the target's logits and hidden state at the last position.
+#[allow(clippy::too_many_arguments)]
+fn prefill(
+    model: &mut Qwen3_5ForCausalLM,
+    head: &mut MtpHead,
+    rotary: &RotaryEmbedding,
+    ids: &[u32],
+    chunk: usize,
+    dtype: DType,
+    dev: &Device,
+    sync: &dyn Fn(&Device) -> Result<()>,
+) -> Result<(Tensor, Tensor)> {
+    let n = ids.len();
+    let mut last: Option<(Tensor, Tensor)> = None;
+    let mut at = 0usize;
+    while at < n {
+        let end = (at + chunk).min(n);
+        let block = Tensor::new(&ids[at..end], dev)?.unsqueeze(0)?;
+        let (logits, hidden) = model.forward_with_hidden(&block, at)?;
+        let l = end - at;
+
+        let head_end = end.min(n - 1);
+        if head_end > at {
+            let hl = head_end - at;
+            let shifted = Tensor::new(&ids[at + 1..head_end + 1], dev)?.unsqueeze(0)?;
+            let embeds = model.embed_tokens(&shifted)?;
+            let hid = hidden.i((.., ..hl, ..))?;
+            let mask = causal_mask(hl, at, dtype, dev)?;
+            let positions: Vec<usize> = (at..head_end).collect();
+            let (cos, sin) = rotary.cos_sin_at(&positions)?;
+            head.forward(&embeds, &hid, Some(&mask), &cos, &sin)?;
+        }
+
+        last = Some((logits, hidden.i((.., l - 1.., ..))?));
+        at = end;
+        sync(dev)?;
+    }
+    last.context("empty prompt")
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let args = Args::parse();
+    if args.tensor_parallel > 1 {
+        return run_tp(args).await;
+    }
+    let dev = pick_device(&args.device)?;
+    let dtype = if matches!(dev, Device::Cpu) {
+        DType::F32
+    } else {
+        DType::BF16
+    };
+
+    // ── load ────────────────────────────────────────────────────────
+    let cfg_raw =
+        std::fs::read_to_string(args.model.join("config.json")).context("read config.json")?;
+    let config: Config = Config::from_config_json(&cfg_raw).context("parse config.json")?;
+    let text_cfg = config.text_config.clone();
+    anyhow::ensure!(
+        MtpHead::present_in(&text_cfg),
+        "this checkpoint declares no MTP head (mtp_num_hidden_layers = 0)"
+    );
+
+    let mut shards: Vec<PathBuf> = std::fs::read_dir(&args.model)
+        .context("read model dir")?
+        .filter_map(|e| e.ok().map(|e| e.path()))
+        .filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
+        .collect();
+    shards.sort();
+    anyhow::ensure!(!shards.is_empty(), "no .safetensors in {:?}", args.model);
+    eprintln!(
+        "loading {} shard(s) as {dtype:?} on {:?}…",
+        shards.len(),
+        dev
+    );
+
+    // SAFETY: mmaps checkpoint files the operator pointed us at.
+    let vb =
+        unsafe { candle_nn::var_builder::ShardedSafeTensors::var_builder(&shards, dtype, &dev)? };
+    let mut model = Qwen3_5ForCausalLM::new(config, vb.clone()).context("load target model")?;
+    let rotary = std::sync::Arc::new(RotaryEmbedding::new(dtype, &text_cfg, &dev)?);
+    let mut head = MtpHead::load(&text_cfg, rotary.clone(), &vb).context("load MTP head")?;
+
+    let tok = tokenizers::Tokenizer::from_file(args.model.join("tokenizer.json"))
+        .map_err(|e| anyhow::anyhow!("load tokenizer: {e}"))?;
+    let prompt = match &args.prompt_file {
+        Some(p) => std::fs::read_to_string(p).context("read prompt file")?,
+        None => "Explain, step by step, how a bicycle derailleur shifts gears.".into(),
+    };
+    let ids: Vec<u32> = tok
+        .encode(prompt.as_str(), false)
+        .map_err(|e| anyhow::anyhow!("tokenize: {e}"))?
+        .get_ids()
+        .to_vec();
+    anyhow::ensure!(ids.len() >= 2, "prompt is too short to prefill");
+    eprintln!("prompt: {} tokens", ids.len());
+
+    let sync = |d: &Device| -> Result<()> {
+        #[cfg(feature = "cuda")]
+        {
+            d.synchronize()?;
+        }
+        let _ = d;
+        Ok(())
+    };
+
+    // ── prefill ─────────────────────────────────────────────────────
+    let n = ids.len();
+    let t0 = Instant::now();
+    let (logits, hidden) = prefill(
+        &mut model,
+        &mut head,
+        &rotary,
+        &ids,
+        args.prefill_chunk,
+        dtype,
+        &dev,
+        &sync,
+    )?;
+    eprintln!("prefill ({n} tokens): {:.2} s", t0.elapsed().as_secs_f64());
+
+    // ── the loop ────────────────────────────────────────────────────
+    let vocab = text_cfg.vocab_size;
+    let mut last_logits: Vec<f32> = logits.i((0, 0, ..))?.to_dtype(DType::F32)?.to_vec1()?;
+    let mut hidden_cur = hidden;
+    let mut produced: Vec<u32> = Vec::new();
+    // (step at which it was drafted, the drafted tokens)
+    let mut pending: Vec<(usize, Vec<u32>)> = Vec::new();
+    let mut accepted_hist = vec![0usize; args.draft_len + 1];
+    let mut accepted_hist_penalised = vec![0usize; args.draft_len + 1];
+    let mut penalty_disagreements = 0usize;
+    let mut draft_ms_total = 0f64;
+    let mut draft_rounds = 0usize;
+    let mut target_ms_total = 0f64;
+    let mut target_steps = 0usize;
+
+    let total_rounds = args.warmup + args.steps;
+    for round in 0..total_rounds {
+        let measuring = round >= args.warmup;
+        let pos = n + produced.len(); // position of the token about to be drafted from
+        let next = argmax(&last_logits);
+        let next_penalised = argmax_with_penalty(
+            &last_logits,
+            &produced[produced.len().saturating_sub(args.repeat_last_n)..],
+            args.repeat_penalty,
+        );
+        if next != next_penalised {
+            penalty_disagreements += 1;
+        }
+        produced.push(next);
+
+        // Resolve any draft whose K true tokens have now been produced.
+        pending.retain(|(at, drafted)| {
+            let have = produced.len() - at;
+            if have < drafted.len() {
+                return true;
+            }
+            let truth = &produced[*at..*at + drafted.len()];
+            let acc = drafted
+                .iter()
+                .zip(truth.iter())
+                .take_while(|(d, t)| d == t)
+                .count();
+            accepted_hist[acc] += 1;
+            false
+        });
+
+        // ── draft, then rewind ──────────────────────────────────────
+        let snap = head.snapshot_kv()?;
+        let t0 = Instant::now();
+        let mut drafted = Vec::with_capacity(args.draft_len);
+        let mut cur_token = next;
+        let mut cur_hidden = hidden_cur.clone();
+        for k in 0..args.draft_len {
+            let tok_t = Tensor::new(&[cur_token], &dev)?.unsqueeze(0)?;
+            let emb = model.embed_tokens(&tok_t)?;
+            let (c, s) = rotary.cos_sin_at(&[pos + k])?;
+            let out = head.forward(&emb, &cur_hidden, None, &c, &s)?;
+            let row: Vec<f32> = model
+                .lm_head(&out)?
+                .i((0, 0, ..))?
+                .to_dtype(DType::F32)?
+                .to_vec1()?;
+            anyhow::ensure!(row.len() == vocab, "draft logits width {}", row.len());
+            cur_token = argmax(&row);
+            drafted.push(cur_token);
+            cur_hidden = out;
+        }
+        sync(&dev)?;
+        let draft_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
+        head.restore_kv(&snap)?;
+
+        // The head must still advance over the *true* token, or its
+        // cache falls behind the target's.
+        let tok_t = Tensor::new(&[next], &dev)?.unsqueeze(0)?;
+        let emb = model.embed_tokens(&tok_t)?;
+        let (c, s) = rotary.cos_sin_at(&[pos])?;
+        head.forward(&emb, &hidden_cur, None, &c, &s)?;
+
+        // ── the target's own step ───────────────────────────────────
+        let t0 = Instant::now();
+        let (logits, h) = model.forward_with_hidden(&tok_t, pos)?;
+        sync(&dev)?;
+        let target_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
+        last_logits = logits.i((0, 0, ..))?.to_dtype(DType::F32)?.to_vec1()?;
+        hidden_cur = h;
+
+        if measuring {
+            pending.push((produced.len(), drafted));
+            draft_ms_total += draft_elapsed;
+            draft_rounds += 1;
+            target_ms_total += target_elapsed;
+            target_steps += 1;
+        }
+    }
+
+    // Anything still pending never got its K true tokens; score what is
+    // known rather than dropping it silently.
+    for (at, drafted) in &pending {
+        let have = produced.len() - at;
+        let truth = &produced[*at..];
+        let acc = drafted
+            .iter()
+            .zip(truth.iter())
+            .take_while(|(d, t)| d == t)
+            .count();
+        if acc < have {
+            accepted_hist[acc] += 1;
+        }
+    }
+    let _ = &mut accepted_hist_penalised;
+
+    // ── verify cost, sampled ────────────────────────────────────────
+    //
+    // Timed after the loop because each sample advances the target's
+    // cache by K+1 positions — harmless once measuring is done, and it
+    // spares the loop the re-prefill that undoing one sample required.
+    let mut verify_samples = Vec::with_capacity(args.verify_samples);
+    for i in 0..args.verify_samples {
+        let block: Vec<u32> = std::iter::repeat_n(produced[0], args.draft_len + 1).collect();
+        let blk = Tensor::new(block.as_slice(), &dev)?.unsqueeze(0)?;
+        let at = n + produced.len() + i * (args.draft_len + 1);
+        let t0 = Instant::now();
+        let _ = model.forward_multi(&blk, at)?;
+        sync(&dev)?;
+        verify_samples.push(t0.elapsed().as_secs_f64() * 1000.0);
+    }
+    verify_samples.sort_by(f64::total_cmp);
+    let verify_ms = verify_samples[verify_samples.len() / 2];
+
+    // ── report ──────────────────────────────────────────────────────
+    let rounds: usize = accepted_hist.iter().sum();
+    let mean_accepted: f64 = accepted_hist
+        .iter()
+        .enumerate()
+        .map(|(k, n)| (k * n) as f64)
+        .sum::<f64>()
+        / rounds.max(1) as f64;
+    let target_ms = target_ms_total / target_steps.max(1) as f64;
+    let draft_ms = draft_ms_total / (draft_rounds.max(1) * args.draft_len) as f64;
+    let c = draft_ms / target_ms;
+    // One round commits 1 + accepted tokens and costs one verify plus K
+    // drafts, against (1 + accepted) plain decode steps.
+    let round_cost = verify_ms + args.draft_len as f64 * draft_ms;
+    let speedup = (1.0 + mean_accepted) * target_ms / round_cost;
+
+    let summary = serde_json::json!({
+        "model": args.model.to_string_lossy(),
+        "device": args.device,
+        "dtype": format!("{dtype:?}"),
+        "prompt_tokens": n,
+        "draft_len": args.draft_len,
+        "rounds_scored": rounds,
+        "target_decode_ms": target_ms,
+        "draft_step_ms": draft_ms,
+        "verify_ms_k_plus_1": verify_ms,
+        "verify_ms_min": verify_samples.first(),
+        "verify_ms_max": verify_samples.last(),
+        "verify_sample_count": verify_samples.len(),
+        "c_draft_over_target": c,
+        "mean_accepted_tokens": mean_accepted,
+        "accepted_histogram": accepted_hist,
+        "predicted_speedup": speedup,
+        "penalty_argmax_disagreements": penalty_disagreements,
+        "penalty_argmax_disagreement_rate":
+            penalty_disagreements as f64 / produced.len().max(1) as f64,
+    });
+    let pretty = serde_json::to_string_pretty(&summary)?;
+    println!("{pretty}");
+    if let Some(p) = &args.json {
+        std::fs::write(p, &pretty).context("write --json")?;
+    }
+    Ok(())
+}
+
+/// Tensor-parallel measurement (#96).
+///
+/// The only route to the 27B: 54 GB in bf16 does not fit one 5090, and
+/// quantisation is reachable only through the TP loader. The shape
+/// mirrors the single-GPU path — prefill, then draft-and-rewind at each
+/// decode step — with two differences:
+///
+/// - the target's step fans out to every rank, because the row-parallel
+///   `AllReduce`s only complete when all of them arrive;
+/// - the draft head is unsharded on the leader, so drafting issues no
+///   collectives at all and cannot leave a rank waiting.
+#[cfg(feature = "cuda")]
+async fn run_tp(args: Args) -> Result<()> {
+    use neuron::harness::device_worker::DeviceWorkerHandle;
+    use neuron::harness::tp::WorkerPool;
+
+    let cfg_raw =
+        std::fs::read_to_string(args.model.join("config.json")).context("read config.json")?;
+    let config: Config = Config::from_config_json(&cfg_raw).context("parse config.json")?;
+    let text_cfg = config.text_config.clone();
+    anyhow::ensure!(
+        MtpHead::present_in(&text_cfg),
+        "this checkpoint declares no MTP head (mtp_num_hidden_layers = 0)"
+    );
+
+    let mut shards: Vec<PathBuf> = std::fs::read_dir(&args.model)
+        .context("read model dir")?
+        .filter_map(|e| e.ok().map(|e| e.path()))
+        .filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
+        .collect();
+    shards.sort();
+    anyhow::ensure!(!shards.is_empty(), "no .safetensors in {:?}", args.model);
+
+    let worker_bin = match &args.worker_binary {
+        Some(p) => p.clone(),
+        None => std::env::current_exe()
+            .context("resolve current_exe")?
+            .parent()
+            .context("executable has no parent directory")?
+            .join("neuron"),
+    };
+    anyhow::ensure!(
+        worker_bin.exists(),
+        "worker binary {worker_bin:?} not found — build `neuron` beside this probe, \
+         or pass --worker-binary"
+    );
+
+    let devices: Vec<u32> = (0..args.tensor_parallel).collect();
+    eprintln!(
+        "tp-{}: spawning workers from {worker_bin:?} on devices {devices:?}…",
+        args.tensor_parallel
+    );
+    let leader = DeviceWorkerHandle::spawn(devices[0])?;
+    let mut pool =
+        WorkerPool::spawn(&worker_bin, args.tensor_parallel, &devices, leader.clone()).await?;
+    pool.init_nccl(devices[0]).await?;
+
+    let leader_device = Device::new_cuda(devices[0] as usize)?;
+    let model_id = "mtp-probe";
+    let t0 = Instant::now();
+    let handle = pool
+        .load_dense_shard(
+            model_id,
+            &cfg_raw,
+            &shards,
+            &leader_device,
+            DType::BF16,
+            args.quant.clone(),
+        )
+        .await?;
+    eprintln!("shards loaded: {:.1} s", t0.elapsed().as_secs_f64());
+
+    let shard_strs: Vec<String> = shards
+        .iter()
+        .map(|p| p.to_string_lossy().into_owned())
+        .collect();
+    leader
+        .tp_load_mtp_head(handle, cfg_raw.clone(), shard_strs)
+        .await?;
+    eprintln!("draft head loaded on the leader");
+
+    let tok = tokenizers::Tokenizer::from_file(args.model.join("tokenizer.json"))
+        .map_err(|e| anyhow::anyhow!("load tokenizer: {e}"))?;
+    let prompt = match &args.prompt_file {
+        Some(p) => std::fs::read_to_string(p).context("read prompt file")?,
+        None => "Explain, step by step, how a bicycle derailleur shifts gears.".into(),
+    };
+    let ids: Vec<u32> = tok
+        .encode(prompt.as_str(), false)
+        .map_err(|e| anyhow::anyhow!("tokenize: {e}"))?
+        .get_ids()
+        .to_vec();
+    let n = ids.len();
+    anyhow::ensure!(n >= 2, "prompt is too short to prefill");
+    eprintln!("prompt: {n} tokens");
+
+    // ── prefill, chunked, target then head ──────────────────────────
+    let t0 = Instant::now();
+    let mut last_logits: Vec<f32> = Vec::new();
+    let mut at = 0usize;
+    while at < n {
+        let end = (at + args.prefill_chunk).min(n);
+        last_logits = pool
+            .generate_step(model_id, handle, ids[at..end].to_vec(), at)
+            .await?;
+        let head_end = end.min(n - 1);
+        if head_end > at {
+            leader
+                .tp_mtp_prefill_chunk(handle, ids[at + 1..head_end + 1].to_vec(), at)
+                .await?;
+        }
+        at = end;
+    }
+    eprintln!("prefill ({n} tokens): {:.1} s", t0.elapsed().as_secs_f64());
+
+    // ── draft / verify observation ──────────────────────────────────
+    let mut produced: Vec<u32> = Vec::new();
+    let mut pending: Vec<(usize, Vec<u32>)> = Vec::new();
+    let mut accepted_hist = vec![0usize; args.draft_len + 1];
+    let mut penalty_disagreements = 0usize;
+    let mut draft_ms_total = 0f64;
+    let mut target_ms_total = 0f64;
+    let mut measured = 0usize;
+
+    for round in 0..args.warmup + args.steps {
+        let measuring = round >= args.warmup;
+        let pos = n + produced.len();
+        let next = argmax(&last_logits);
+        let next_penalised = argmax_with_penalty(
+            &last_logits,
+            &produced[produced.len().saturating_sub(args.repeat_last_n)..],
+            args.repeat_penalty,
+        );
+        if next != next_penalised {
+            penalty_disagreements += 1;
+        }
+        produced.push(next);
+
+        pending.retain(|(start, drafted)| {
+            if produced.len() - start < drafted.len() {
+                return true;
+            }
+            let truth = &produced[*start..*start + drafted.len()];
+            let acc = drafted
+                .iter()
+                .zip(truth.iter())
+                .take_while(|(d, t)| d == t)
+                .count();
+            accepted_hist[acc] += 1;
+            false
+        });
+
+        let t0 = Instant::now();
+        let drafted = leader
+            .tp_mtp_draft(handle, next, pos, args.draft_len)
+            .await?;
+        let draft_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
+        leader.tp_mtp_advance(handle, next, pos).await?;
+
+        let t0 = Instant::now();
+        last_logits = pool
+            .generate_step(model_id, handle, vec![next], pos)
+            .await?;
+        let target_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
+
+        if measuring {
+            pending.push((produced.len(), drafted));
+            draft_ms_total += draft_elapsed;
+            target_ms_total += target_elapsed;
+            measured += 1;
+        }
+    }
+
+    // ── verify cost, sampled ────────────────────────────────────────
+    //
+    // After the loop, because each sample advances the target's cache by
+    // K+1. One sample was not enough: it moved 57.3 -> 50.1 ms between
+    // otherwise identical runs, wide enough to decide the gate by itself.
+    let mut verify_samples = Vec::with_capacity(args.verify_samples);
+    for i in 0..args.verify_samples {
+        let block: Vec<u32> = std::iter::repeat_n(produced[0], args.draft_len + 1).collect();
+        let at = n + produced.len() + i * (args.draft_len + 1);
+        let t0 = Instant::now();
+        let _ = pool
+            .generate_step_multi(model_id, handle, block, at)
+            .await?;
+        verify_samples.push(t0.elapsed().as_secs_f64() * 1000.0);
+    }
+    verify_samples.sort_by(f64::total_cmp);
+    let verify_ms = verify_samples[verify_samples.len() / 2];
+
+    let rounds: usize = accepted_hist.iter().sum();
+    let mean_accepted: f64 = accepted_hist
+        .iter()
+        .enumerate()
+        .map(|(k, c)| (k * c) as f64)
+        .sum::<f64>()
+        / rounds.max(1) as f64;
+    let target_ms = target_ms_total / measured.max(1) as f64;
+    let draft_ms = draft_ms_total / (measured.max(1) * args.draft_len) as f64;
+    let round_cost = verify_ms + args.draft_len as f64 * draft_ms;
+    let summary = serde_json::json!({
+        "model": args.model.to_string_lossy(),
+        "tensor_parallel": args.tensor_parallel,
+        "quant": args.quant,
+        "prompt_tokens": n,
+        "draft_len": args.draft_len,
+        "rounds_scored": rounds,
+        "target_decode_ms": target_ms,
+        "draft_step_ms": draft_ms,
+        "verify_ms_k_plus_1": verify_ms,
+        "verify_ms_min": verify_samples.first(),
+        "verify_ms_max": verify_samples.last(),
+        "verify_sample_count": verify_samples.len(),
+        "c_draft_over_target": draft_ms / target_ms,
+        "mean_accepted_tokens": mean_accepted,
+        "accepted_histogram": accepted_hist,
+        "predicted_speedup": (1.0 + mean_accepted) * target_ms / round_cost,
+        "penalty_argmax_disagreements": penalty_disagreements,
+        "penalty_argmax_disagreement_rate":
+            penalty_disagreements as f64 / produced.len().max(1) as f64,
+    });
+    let pretty = serde_json::to_string_pretty(&summary)?;
+    println!("{pretty}");
+    if let Some(p) = &args.json {
+        std::fs::write(p, &pretty).context("write --json")?;
+    }
+    pool.unload_model(model_id).await?;
+    Ok(())
+}
+
+#[cfg(not(feature = "cuda"))]
+async fn run_tp(_args: Args) -> Result<()> {
+    anyhow::bail!("--tensor-parallel > 1 requires a build with the `cuda` feature")
+}

crates/neuron/src/harness/arch/laya/backbone.rs

@@ -0,0 +1,478 @@
+//! ModernBERT, the bidirectional encoder under every Laya checkpoint.
+//!
+//! Follows the structure of `candle_transformers::models::modernbert`
+//! (itself a port of transformers' `ModernBertModel`), vendored rather
+//! than used as-is because that module only runs correctly in f32:
+//!
+//! - its attention masks are built in F32 whatever the model dtype, so
+//!   an f16/bf16 load fails at the first `broadcast_add`;
+//! - it computes the rotary angles `position × inv_freq` in the model
+//!   dtype, which in f16 already loses a quarter-radian of phase by
+//!   position 500;
+//! - it loads each layer's `attn_norm` with `.ok()`, so a weight-prefix
+//!   mistake silently yields a model with no pre-attention norms.
+//!
+//! Here the angles are computed in f32 and only the resulting cos/sin
+//! tables are cast; attention scores are softmaxed in f32 against f32
+//! masks; and `attn_norm` is required on every layer except layer 0,
+//! which ModernBERT builds as an identity.
+//!
+//! Precision follows PyTorch autocast, which is how the reference runs on
+//! a GPU: the linear layers (and so the attention inputs) run in the
+//! load dtype, while the residual stream, the embeddings and every
+//! LayerNorm stay in f32. Casting the residual stream to bf16 as well
+//! drifts the released English checkpoint's logits by ~0.25 over its 28
+//! layers; keeping it in f32 is what autocast does and costs little.
+//!
+//! Masked positions use a large *finite* negative rather than `-inf`.
+//! A padded query near the end of a short row can find every key in
+//! its sliding window padded; with `-inf` that row's softmax is NaN,
+//! and NaN in a padded row's value vector poisons the valid rows that
+//! give it zero weight (`0 × NaN = NaN`). A finite value makes such a
+//! row uniform instead, and valid rows never look at it.
+
+use anyhow::{Context, Result, ensure};
+use candle_core::{D, DType, Device, Module, Tensor};
+use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
+
+use super::config::{EncoderConfig, LayerKind};
+
+/// Additive mask value for a position that must not be attended.
+pub(crate) const MASKED: f32 = -1e9;
+
+/// Apply `lin` in its weight's dtype, casting `x` in if needed. The
+/// result stays in that dtype; callers cast back to f32 where it joins
+/// the residual stream.
+pub(crate) fn project(lin: &Linear, x: &Tensor) -> candle_core::Result<Tensor> {
+    x.to_dtype(lin.weight().dtype())?.apply(lin)
+}
+
+/// A bias-free LayerNorm that takes candle's fused kernel.
+///
+/// `candle_nn::LayerNorm` dispatches to the fused `layer_norm` op only
+/// when it has a bias; without one it falls back to six elementwise
+/// kernels (mean, subtract, square, mean, divide, scale), which was a
+/// quarter of the encoder's GPU time. ModernBERT's norms have no bias,
+/// so they get an explicit zero one: same result, one kernel.
+fn fused_layer_norm_no_bias(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
+    let weight = vb.get(size, "weight")?;
+    let bias = weight.zeros_like()?;
+    Ok(LayerNorm::new(weight, bias, eps))
+}
+
+/// Precomputed rotary tables for one base, `(max_pos, head_dim / 2)`.
+struct Rotary {
+    cos: Tensor,
+    sin: Tensor,
+}
+
+impl Rotary {
+    fn new(cfg: &EncoderConfig, theta: f64, dtype: DType, device: &Device) -> Result<Self> {
+        let dim = cfg.head_dim();
+        let inv_freq: Vec<f32> = (0..dim)
+            .step_by(2)
+            .map(|i| (1.0 / theta.powf(i as f64 / dim as f64)) as f32)
+            .collect();
+        let n = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, n), device)?;
+        let max_pos = cfg.max_position_embeddings;
+        let pos = Tensor::arange(0u32, max_pos as u32, device)?
+            .to_dtype(DType::F32)?
+            .reshape((max_pos, 1))?;
+        let freqs = pos.matmul(&inv_freq)?;
+        Ok(Self {
+            cos: freqs.cos()?.to_dtype(dtype)?,
+            sin: freqs.sin()?.to_dtype(dtype)?,
+        })
+    }
+
+    fn apply(&self, q: &Tensor, k: &Tensor, seq_len: usize) -> Result<(Tensor, Tensor)> {
+        let cos = self.cos.narrow(0, 0, seq_len)?;
+        let sin = self.sin.narrow(0, 0, seq_len)?;
+        Ok((
+            candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?,
+            candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?,
+        ))
+    }
+}
+
+struct Attention {
+    wqkv: Linear,
+    wo: Linear,
+    n_heads: usize,
+    head_dim: usize,
+    kind: LayerKind,
+}
+
+impl Attention {
+    fn forward(&self, xs: &Tensor, rotary: &Rotary, mask: Option<&Tensor>) -> Result<Tensor> {
+        let (b, l, d) = xs.dims3()?;
+        let qkv = project(&self.wqkv, xs)?
+            .reshape((b, l, 3, self.n_heads, self.head_dim))?
+            .permute((2, 0, 3, 1, 4))?;
+        // Rotary in the projection dtype, as under autocast; the scores
+        // are then taken in f32 (see `scaled_attention`).
+        let q = qkv.get(0)?;
+        let k = qkv.get(1)?;
+        let v = qkv.get(2)?.contiguous()?;
+        let (q, k) = rotary.apply(&q, &k, l)?;
+        let ctx = scaled_attention(&q, &k, &v, mask, self.head_dim)?;
+        let ctx = ctx.transpose(1, 2)?.reshape((b, l, d))?;
+        Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
+    }
+}
+
+/// A batch packed without padding for the flash-attention path: every
+/// row's tokens back to back, `(total, …)`. Linear layers then do no work
+/// for padding, and flash varlen keeps rows from attending to each other.
+#[cfg(feature = "flash-attn")]
+pub struct Packed {
+    /// `(b + 1,)` u32 cumulative row lengths, as flash varlen takes them.
+    pub cu_seqlens: Tensor,
+    /// Longest row.
+    pub max_len: usize,
+    /// `(total,)` u32: each token's position within its own row.
+    pub positions: Tensor,
+    /// Where each row starts in the packed order.
+    pub offsets: Vec<usize>,
+}
+
+#[cfg(feature = "flash-attn")]
+impl Packed {
+    pub fn new(lens: &[usize], device: &Device) -> Result<Self> {
+        let mut cu = Vec::with_capacity(lens.len() + 1);
+        let mut offsets = Vec::with_capacity(lens.len());
+        let mut positions = Vec::with_capacity(lens.iter().sum());
+        let mut total = 0u32;
+        cu.push(0u32);
+        for &n in lens {
+            offsets.push(total as usize);
+            positions.extend(0..n as u32);
+            total += n as u32;
+            cu.push(total);
+        }
+        let n = cu.len();
+        Ok(Self {
+            cu_seqlens: Tensor::from_vec(cu, n, device)?,
+            max_len: lens.iter().copied().max().unwrap_or(0),
+            positions: Tensor::from_vec(positions, total as usize, device)?,
+            offsets,
+        })
+    }
+}
+
+/// Bidirectional flash attention over a packed batch. `q`, `k`, `v` are
+/// `(total, h, dh)`; `half_window` limits each query to keys within that
+/// distance on either side (ModernBERT's sliding layers), `None` attends
+/// the whole row. Scores are accumulated in f32 inside the kernel.
+#[cfg(feature = "flash-attn")]
+pub(crate) fn flash_attention(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    packed: &Packed,
+    half_window: Option<usize>,
+) -> Result<Tensor> {
+    let head_dim = q.dim(2)?;
+    let scale = (head_dim as f64).powf(-0.5) as f32;
+    Ok(candle_flash_attn::flash_attn_varlen_windowed(
+        q,
+        k,
+        v,
+        &packed.cu_seqlens,
+        &packed.cu_seqlens,
+        packed.max_len,
+        packed.max_len,
+        scale,
+        half_window,
+        half_window,
+    )?)
+}
+
+/// Split a packed `(total, 3·d)` qkv projection into contiguous
+/// `(total, h, dh)` q, k and v.
+#[cfg(feature = "flash-attn")]
+pub(crate) fn split_qkv(
+    qkv: &Tensor,
+    n_heads: usize,
+    head_dim: usize,
+) -> Result<(Tensor, Tensor, Tensor)> {
+    let t = qkv.dim(0)?;
+    let qkv = qkv.reshape((t, 3, n_heads, head_dim))?;
+    let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(1, i, 1)?.squeeze(1)?.contiguous()?) };
+    Ok((part(0)?, part(1)?, part(2)?))
+}
+
+#[cfg(feature = "flash-attn")]
+impl Rotary {
+    /// Rotate packed `(total, h, dh)` q and k by each token's position.
+    fn apply_packed(&self, q: &Tensor, k: &Tensor, positions: &Tensor) -> Result<(Tensor, Tensor)> {
+        let cos = self.cos.index_select(positions, 0)?;
+        let sin = self.sin.index_select(positions, 0)?;
+        let rot = |x: &Tensor| -> Result<Tensor> {
+            Ok(candle_nn::rotary_emb::rope_thd(&x.unsqueeze(0)?, &cos, &sin)?.squeeze(0)?)
+        };
+        Ok((rot(q)?, rot(k)?))
+    }
+}
+
+#[cfg(feature = "flash-attn")]
+impl Attention {
+    /// `xs` is the packed f32 residual stream `(total, d)`.
+    fn forward_packed(
+        &self,
+        xs: &Tensor,
+        rotary: &Rotary,
+        packed: &Packed,
+        half_window: usize,
+    ) -> Result<Tensor> {
+        let (t, d) = xs.dims2()?;
+        let (q, k, v) = split_qkv(&project(&self.wqkv, xs)?, self.n_heads, self.head_dim)?;
+        let (q, k) = rotary.apply_packed(&q, &k, &packed.positions)?;
+        let window = match self.kind {
+            LayerKind::Sliding => Some(half_window),
+            LayerKind::Global => None,
+        };
+        let ctx = flash_attention(&q, &k, &v, packed, window)?.reshape((t, d))?;
+        Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
+    }
+}
+
+/// `softmax(q kᵀ / √d + mask) v`, scores and softmax in f32.
+///
+/// `q`, `k`, `v` are `(b, h, l, d)`; `mask` is an f32 additive mask
+/// broadcastable to `(b, h, l, l)`, or `None` when nothing is masked —
+/// which saves a bandwidth-bound pass over the `(b, h, l, l)` scores.
+pub(crate) fn scaled_attention(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    mask: Option<&Tensor>,
+    head_dim: usize,
+) -> Result<Tensor> {
+    let dtype = v.dtype();
+    let scale = (head_dim as f64).powf(-0.5);
+    // Scores in f32 end to end: a bf16 `q kᵀ` is rounded to 8 mantissa
+    // bits before the softmax ever sees it, which on this model's
+    // logit-sized scores is most of the drift from the reference's GPU
+    // path (whose SDPA kernels keep scores in f32).
+    // The scale goes on q, not on the `(b, h, l, l)` scores: the same
+    // product, on a tensor `l / d` times smaller.
+    let scores = (q.to_dtype(DType::F32)? * scale)?.matmul(&k.to_dtype(DType::F32)?.t()?)?;
+    let scores = match mask {
+        Some(mask) => scores.broadcast_add(mask)?,
+        None => scores,
+    };
+    let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+    Ok(probs.to_dtype(dtype)?.matmul(v)?)
+}
+
+struct Mlp {
+    wi: Linear,
+    wo: Linear,
+}
+
+impl Module for Mlp {
+    /// f32 in, f32 out; the projections run in the weights' dtype.
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let xs = project(&self.wi, xs)?;
+        let parts = xs.chunk(2, D::Minus1)?;
+        (parts[0].gelu_erf()? * &parts[1])?
+            .apply(&self.wo)?
+            .to_dtype(DType::F32)
+    }
+}
+
+struct Layer {
+    attn_norm: Option<LayerNorm>,
+    attn: Attention,
+    mlp_norm: LayerNorm,
+    mlp: Mlp,
+}
+
+/// The two additive masks one forward needs, both f32. Either is `None`
+/// when it would mask nothing: `global` for a batch with no padding (every
+/// single-question request), `sliding` additionally when no row is longer
+/// than the window reaches.
+pub struct Masks {
+    /// `(b, 1, 1, l)`: padded keys masked. Global layers and the
+    /// decision head use this one.
+    pub global: Option<Tensor>,
+    /// `(b, 1, l, l)`: padded keys, plus keys outside the window.
+    pub sliding: Option<Tensor>,
+}
+
+impl Masks {
+    /// Build both masks from per-row valid lengths. Rows are
+    /// right-padded: positions `>= lens[i]` of row `i` are padding.
+    pub fn new(
+        lens: &[usize],
+        seq_len: usize,
+        half_window: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        let b = lens.len();
+        let padded = lens.iter().any(|&n| n < seq_len);
+        let windowed = seq_len > half_window + 1;
+        let global = if padded {
+            let mut pad = Vec::with_capacity(b * seq_len);
+            for &n in lens {
+                pad.extend((0..seq_len).map(|j| if j < n { 0f32 } else { MASKED }));
+            }
+            Some(Tensor::from_vec(pad, (b, 1, 1, seq_len), device)?)
+        } else {
+            None
+        };
+        if !windowed {
+            let sliding = global.clone();
+            return Ok(Self { global, sliding });
+        }
+        // Built on the device from two aranges, not on the host: at
+        // 1024 tokens a host-built window is a 4 MB upload per forward.
+        let pos = Tensor::arange(0u32, seq_len as u32, device)?.to_dtype(DType::F32)?;
+        let dist = pos
+            .reshape((seq_len, 1))?
+            .broadcast_sub(&pos.reshape((1, seq_len))?)?
+            .abs()?;
+        let outside = dist.gt(half_window as f64)?;
+        let zeros = Tensor::zeros((seq_len, seq_len), DType::F32, device)?;
+        let masked = Tensor::full(MASKED, (seq_len, seq_len), device)?;
+        let window = outside
+            .where_cond(&masked, &zeros)?
+            .reshape((1, 1, seq_len, seq_len))?;
+        let sliding = match &global {
+            Some(g) => g.broadcast_add(&window)?,
+            None => window,
+        };
+        Ok(Self {
+            global,
+            sliding: Some(sliding),
+        })
+    }
+}
+
+pub struct ModernBert {
+    embeddings: Embedding,
+    embed_norm: LayerNorm,
+    layers: Vec<Layer>,
+    final_norm: LayerNorm,
+    global_rotary: Rotary,
+    sliding_rotary: Rotary,
+    half_window: usize,
+}
+
+impl ModernBert {
+    /// Load from a VarBuilder rooted at the encoder's prefix (`encoder`
+    /// in a Laya checkpoint, `model` in a stock ModernBERT one). Linear
+    /// weights load in the VarBuilder's dtype; embeddings and norms in
+    /// f32.
+    pub fn load(vb: VarBuilder, cfg: &EncoderConfig) -> Result<Self> {
+        let d = cfg.hidden_size;
+        let f32_vb = vb.clone().set_dtype(DType::F32);
+        let embeddings = candle_nn::embedding(
+            cfg.vocab_size,
+            d,
+            f32_vb.pp("embeddings").pp("tok_embeddings"),
+        )
+        .context("encoder.embeddings.tok_embeddings")?;
+        let embed_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("embeddings.norm"))?;
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        for (i, kind) in cfg.layer_kinds.iter().copied().enumerate() {
+            let lvb = vb.pp(format!("layers.{i}"));
+            let norm_vb = f32_vb.pp(format!("layers.{i}"));
+            // ModernBERT builds layer 0's attn_norm as an identity and
+            // saves nothing for it; every other layer must have one.
+            let attn_norm = if lvb.contains_tensor("attn_norm.weight") {
+                Some(fused_layer_norm_no_bias(
+                    d,
+                    cfg.norm_eps,
+                    norm_vb.pp("attn_norm"),
+                )?)
+            } else {
+                ensure!(
+                    i == 0,
+                    "encoder layer {i} has no attn_norm.weight — wrong weight prefix, \
+                     or not a ModernBERT checkpoint"
+                );
+                None
+            };
+            let attn = Attention {
+                wqkv: candle_nn::linear_no_bias(d, 3 * d, lvb.pp("attn.Wqkv"))?,
+                wo: candle_nn::linear_no_bias(d, d, lvb.pp("attn.Wo"))?,
+                n_heads: cfg.num_attention_heads,
+                head_dim: cfg.head_dim(),
+                kind,
+            };
+            let mlp_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, norm_vb.pp("mlp_norm"))?;
+            let mlp = Mlp {
+                wi: candle_nn::linear_no_bias(d, 2 * cfg.intermediate_size, lvb.pp("mlp.Wi"))?,
+                wo: candle_nn::linear_no_bias(cfg.intermediate_size, d, lvb.pp("mlp.Wo"))?,
+            };
+            layers.push(Layer {
+                attn_norm,
+                attn,
+                mlp_norm,
+                mlp,
+            });
+        }
+        let final_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("final_norm"))?;
+        let device = vb.device();
+        Ok(Self {
+            embeddings,
+            embed_norm,
+            layers,
+            final_norm,
+            global_rotary: Rotary::new(cfg, cfg.global_rope_theta, vb.dtype(), device)?,
+            sliding_rotary: Rotary::new(cfg, cfg.local_rope_theta, vb.dtype(), device)?,
+            half_window: cfg.local_attention / 2,
+        })
+    }
+
+    pub fn half_window(&self) -> usize {
+        self.half_window
+    }
+
+    /// `ids` is `(b, l)` u32; returns the final hidden states `(b, l, d)`
+    /// in f32.
+    pub fn forward(&self, ids: &Tensor, masks: &Masks) -> Result<Tensor> {
+        let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
+        for layer in &self.layers {
+            let normed = match &layer.attn_norm {
+                Some(norm) => xs.apply(norm)?,
+                None => xs.clone(),
+            };
+            let (rotary, mask) = match layer.attn.kind {
+                LayerKind::Global => (&self.global_rotary, masks.global.as_ref()),
+                LayerKind::Sliding => (&self.sliding_rotary, masks.sliding.as_ref()),
+            };
+            xs = (&xs + layer.attn.forward(&normed, rotary, mask)?)?;
+            let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
+            xs = (xs + mlp_out)?;
+        }
+        Ok(xs.apply(&self.final_norm)?)
+    }
+
+    /// The flash path: `ids` is the packed `(total,)` u32 token stream;
+    /// returns the packed final hidden states `(total, d)` in f32.
+    #[cfg(feature = "flash-attn")]
+    pub fn forward_packed(&self, ids: &Tensor, packed: &Packed) -> Result<Tensor> {
+        let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
+        for layer in &self.layers {
+            let normed = match &layer.attn_norm {
+                Some(norm) => xs.apply(norm)?,
+                None => xs.clone(),
+            };
+            let rotary = match layer.attn.kind {
+                LayerKind::Global => &self.global_rotary,
+                LayerKind::Sliding => &self.sliding_rotary,
+            };
+            xs = (&xs
+                + layer
+                    .attn
+                    .forward_packed(&normed, rotary, packed, self.half_window)?)?;
+            let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
+            xs = (xs + mlp_out)?;
+        }
+        Ok(xs.apply(&self.final_norm)?)
+    }
+}

crates/neuron/src/harness/arch/laya/config.rs

@@ -0,0 +1,358 @@
+//! The two config files a Laya checkpoint ships, and what they mean.
+//!
+//! A checkpoint directory holds `rl_agent_config.json` (the decision
+//! model: head depth, token budgets, calibration temperatures) and
+//! `encoder/config.json` (the ModernBERT backbone, written by
+//! transformers 5). There is no top-level `config.json`.
+//!
+//! The encoder config is parsed here rather than through
+//! `candle_transformers::models::modernbert::Config` because
+//! transformers 5 moved the rotary bases into a nested
+//! `rope_parameters.{full_attention,sliding_attention}.rope_theta`
+//! block, which that struct cannot read; the older flat
+//! `global_rope_theta` / `local_rope_theta` fields are still honoured
+//! when a checkpoint carries them. mmBERT (the multilingual
+//! checkpoint's backbone) is the case that matters: both of its bases
+//! are 160,000, where the flat-field defaults would silently give the
+//! sliding layers 10,000.
+
+use anyhow::{Context, Result, bail, ensure};
+use serde::Deserialize;
+use std::collections::BTreeMap;
+use std::path::Path;
+
+/// `rl_agent_config.json`.
+#[derive(Debug, Clone, Deserialize)]
+pub struct LayaConfig {
+    /// The backbone the head was trained on, e.g.
+    /// `answerdotai/ModernBERT-large`. Informational: the weights come
+    /// from the checkpoint, the shape from `encoder/config.json`.
+    #[serde(default)]
+    pub encoder: Option<String>,
+    #[serde(default = "default_head_layers")]
+    pub head_layers: usize,
+    #[serde(default = "default_max_len")]
+    pub max_len: usize,
+    #[serde(default = "default_head_max_len")]
+    pub head_max_len: usize,
+    /// One act head output per entry, plus one for "answer". The act
+    /// head's width is `act_costs.len() + 1`.
+    #[serde(default)]
+    pub act_costs: BTreeMap<String, f64>,
+    /// Per-question-type temperatures, indexed choice / score / noul.
+    #[serde(default = "default_temperature")]
+    pub temperature: Vec<f64>,
+    /// Temperatures fitted per `"<type>:<cardinality bucket>"`.
+    #[serde(default)]
+    pub temperature_by_options: BTreeMap<String, f64>,
+    /// Autocast dtype the reference runs on GPU (`"bf16"` or `"fp16"`).
+    #[serde(default)]
+    pub amp_dtype: Option<String>,
+}
+
+fn default_head_layers() -> usize {
+    2
+}
+fn default_max_len() -> usize {
+    512
+}
+fn default_head_max_len() -> usize {
+    192
+}
+fn default_temperature() -> Vec<f64> {
+    vec![1.0, 1.0, 1.0]
+}
+
+impl LayaConfig {
+    pub fn n_act(&self) -> usize {
+        self.act_costs.len() + 1
+    }
+}
+
+/// Whether an attention layer sees the whole sequence or a window.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LayerKind {
+    Global,
+    Sliding,
+}
+
+/// The ModernBERT backbone shape, resolved from `encoder/config.json`.
+#[derive(Debug, Clone)]
+pub struct EncoderConfig {
+    pub vocab_size: usize,
+    pub hidden_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub intermediate_size: usize,
+    pub max_position_embeddings: usize,
+    pub norm_eps: f64,
+    pub pad_token_id: u32,
+    /// Full sliding-window width. A token attends to keys at distance
+    /// `<= local_attention / 2` on either side.
+    pub local_attention: usize,
+    pub global_rope_theta: f64,
+    pub local_rope_theta: f64,
+    pub layer_kinds: Vec<LayerKind>,
+}
+
+#[derive(Deserialize)]
+struct RawEncoderConfig {
+    vocab_size: usize,
+    hidden_size: usize,
+    num_hidden_layers: usize,
+    num_attention_heads: usize,
+    intermediate_size: usize,
+    max_position_embeddings: usize,
+    #[serde(default)]
+    norm_eps: Option<f64>,
+    #[serde(default)]
+    layer_norm_eps: Option<f64>,
+    pad_token_id: u32,
+    #[serde(default = "default_global_every")]
+    global_attn_every_n_layers: usize,
+    local_attention: usize,
+    #[serde(default)]
+    layer_types: Option<Vec<String>>,
+    #[serde(default)]
+    rope_parameters: Option<serde_json::Value>,
+    #[serde(default)]
+    global_rope_theta: Option<f64>,
+    #[serde(default)]
+    local_rope_theta: Option<f64>,
+    #[serde(default)]
+    norm_bias: bool,
+    #[serde(default)]
+    attention_bias: bool,
+    #[serde(default)]
+    mlp_bias: bool,
+    #[serde(default)]
+    hidden_activation: Option<String>,
+}
+
+fn default_global_every() -> usize {
+    3
+}
+
+/// transformers' ModernBERT defaults for checkpoints that predate
+/// `rope_parameters` and also omit the flat fields.
+const DEFAULT_GLOBAL_ROPE_THETA: f64 = 160_000.0;
+const DEFAULT_LOCAL_ROPE_THETA: f64 = 10_000.0;
+
+impl EncoderConfig {
+    pub fn from_json(text: &str) -> Result<Self> {
+        let raw: RawEncoderConfig =
+            serde_json::from_str(text).context("parse ModernBERT encoder config")?;
+
+        // Variants this port does not implement are refused, not
+        // approximated: each would load without error and answer wrong.
+        ensure!(
+            !raw.norm_bias && !raw.attention_bias && !raw.mlp_bias,
+            "ModernBERT with norm/attention/mlp biases is not supported"
+        );
+        if let Some(act) = raw.hidden_activation.as_deref() {
+            ensure!(
+                act == "gelu",
+                "ModernBERT hidden_activation {act:?} is not supported (only exact \"gelu\")"
+            );
+        }
+        ensure!(
+            raw.hidden_size.is_multiple_of(raw.num_attention_heads),
+            "hidden_size {} is not divisible by num_attention_heads {}",
+            raw.hidden_size,
+            raw.num_attention_heads
+        );
+
+        let (global_rope_theta, local_rope_theta) = rope_thetas(&raw)?;
+
+        let layer_kinds = match &raw.layer_types {
+            Some(types) => {
+                ensure!(
+                    types.len() == raw.num_hidden_layers,
+                    "layer_types has {} entries for {} layers",
+                    types.len(),
+                    raw.num_hidden_layers
+                );
+                types
+                    .iter()
+                    .map(|t| match t.as_str() {
+                        "full_attention" => Ok(LayerKind::Global),
+                        "sliding_attention" => Ok(LayerKind::Sliding),
+                        other => bail!("unknown ModernBERT layer type {other:?}"),
+                    })
+                    .collect::<Result<Vec<_>>>()?
+            }
+            None => (0..raw.num_hidden_layers)
+                .map(|i| {
+                    if i % raw.global_attn_every_n_layers == 0 {
+                        LayerKind::Global
+                    } else {
+                        LayerKind::Sliding
+                    }
+                })
+                .collect(),
+        };
+
+        Ok(Self {
+            vocab_size: raw.vocab_size,
+            hidden_size: raw.hidden_size,
+            num_hidden_layers: raw.num_hidden_layers,
+            num_attention_heads: raw.num_attention_heads,
+            intermediate_size: raw.intermediate_size,
+            max_position_embeddings: raw.max_position_embeddings,
+            norm_eps: raw.norm_eps.or(raw.layer_norm_eps).unwrap_or(1e-5),
+            pad_token_id: raw.pad_token_id,
+            local_attention: raw.local_attention,
+            global_rope_theta,
+            local_rope_theta,
+            layer_kinds,
+        })
+    }
+
+    pub fn head_dim(&self) -> usize {
+        self.hidden_size / self.num_attention_heads
+    }
+}
+
+/// Resolve the two rotary bases. `rope_parameters` (transformers 5)
+/// wins over the flat fields; within it, a per-layer-type block wins
+/// over a shared top-level `rope_theta`.
+fn rope_thetas(raw: &RawEncoderConfig) -> Result<(f64, f64)> {
+    let mut global = raw.global_rope_theta.unwrap_or(DEFAULT_GLOBAL_ROPE_THETA);
+    let mut local = raw.local_rope_theta.unwrap_or(DEFAULT_LOCAL_ROPE_THETA);
+    if let Some(rope) = &raw.rope_parameters {
+        let shared = rope.get("rope_theta").and_then(|v| v.as_f64());
+        for (key, slot) in [
+            ("full_attention", &mut global),
+            ("sliding_attention", &mut local),
+        ] {
+            let block = rope.get(key);
+            if let Some(rope_type) = block
+                .and_then(|b| b.get("rope_type"))
+                .and_then(|v| v.as_str())
+            {
+                ensure!(
+                    rope_type == "default",
+                    "rope_type {rope_type:?} for {key} is not supported"
+                );
+            }
+            let theta = block
+                .and_then(|b| b.get("rope_theta"))
+                .and_then(|v| v.as_f64())
+                .or(shared);
+            if let Some(theta) = theta {
+                *slot = theta;
+            }
+        }
+    }
+    Ok((global, local))
+}
+
+/// Both configs of one checkpoint directory.
+#[derive(Debug, Clone)]
+pub struct CheckpointConfig {
+    pub laya: LayaConfig,
+    pub encoder: EncoderConfig,
+}
+
+impl CheckpointConfig {
+    /// Read `rl_agent_config.json` and `encoder/config.json` from a
+    /// checkpoint directory (the repo root, or one of its subfolders).
+    pub fn from_dir(dir: &Path) -> Result<Self> {
+        let laya_path = dir.join("rl_agent_config.json");
+        let laya: LayaConfig = serde_json::from_str(
+            &std::fs::read_to_string(&laya_path)
+                .with_context(|| format!("read {}", laya_path.display()))?,
+        )
+        .with_context(|| format!("parse {}", laya_path.display()))?;
+        let enc_path = dir.join("encoder").join("config.json");
+        let encoder = EncoderConfig::from_json(
+            &std::fs::read_to_string(&enc_path)
+                .with_context(|| format!("read {}", enc_path.display()))?,
+        )
+        .with_context(|| format!("parse {}", enc_path.display()))?;
+        ensure!(
+            laya.temperature.len() == 3,
+            "rl_agent_config temperature must have 3 entries (choice, score, noul), got {}",
+            laya.temperature.len()
+        );
+        Ok(Self { laya, encoder })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn base() -> serde_json::Value {
+        serde_json::json!({
+            "vocab_size": 100, "hidden_size": 64, "num_hidden_layers": 4,
+            "num_attention_heads": 4, "intermediate_size": 96,
+            "max_position_embeddings": 512, "norm_eps": 1e-5, "pad_token_id": 0,
+            "global_attn_every_n_layers": 3, "local_attention": 16,
+        })
+    }
+
+    #[test]
+    fn nested_rope_parameters_override_flat_defaults() {
+        let mut v = base();
+        v["rope_parameters"] = serde_json::json!({
+            "full_attention": {"rope_theta": 160000.0, "rope_type": "default"},
+            "sliding_attention": {"rope_theta": 160000.0, "rope_type": "default"},
+        });
+        let c = EncoderConfig::from_json(&v.to_string()).unwrap();
+        assert_eq!(c.global_rope_theta, 160_000.0);
+        // The mmBERT case: the flat default would have been 10,000.
+        assert_eq!(c.local_rope_theta, 160_000.0);
+    }
+
+    #[test]
+    fn flat_fields_are_read_without_rope_parameters() {
+        let mut v = base();
+        v["global_rope_theta"] = serde_json::json!(80000.0);
+        v["local_rope_theta"] = serde_json::json!(5000.0);
+        let c = EncoderConfig::from_json(&v.to_string()).unwrap();
+        assert_eq!((c.global_rope_theta, c.local_rope_theta), (80000.0, 5000.0));
+    }
+
+    #[test]
+    fn layer_kinds_follow_every_n_when_layer_types_absent() {
+        let c = EncoderConfig::from_json(&base().to_string()).unwrap();
+        assert_eq!(
+            c.layer_kinds,
+            vec![
+                LayerKind::Global,
+                LayerKind::Sliding,
+                LayerKind::Sliding,
+                LayerKind::Global
+            ]
+        );
+    }
+
+    #[test]
+    fn explicit_layer_types_win() {
+        let mut v = base();
+        v["layer_types"] = serde_json::json!([
+            "sliding_attention",
+            "full_attention",
+            "full_attention",
+            "sliding_attention"
+        ]);
+        let c = EncoderConfig::from_json(&v.to_string()).unwrap();
+        assert_eq!(c.layer_kinds[0], LayerKind::Sliding);
+        assert_eq!(c.layer_kinds[1], LayerKind::Global);
+    }
+
+    #[test]
+    fn unsupported_variants_are_refused() {
+        let mut v = base();
+        v["attention_bias"] = serde_json::json!(true);
+        assert!(EncoderConfig::from_json(&v.to_string()).is_err());
+        let mut v = base();
+        v["hidden_activation"] = serde_json::json!("gelu_pytorch_tanh");
+        assert!(EncoderConfig::from_json(&v.to_string()).is_err());
+        let mut v = base();
+        v["rope_parameters"] =
+            serde_json::json!({"full_attention": {"rope_type": "yarn", "rope_theta": 1.0}});
+        assert!(EncoderConfig::from_json(&v.to_string()).is_err());
+    }
+}

crates/neuron/src/harness/arch/laya/head.rs

@@ -0,0 +1,180 @@
+//! The decision head that sits on top of the encoder.
+//!
+//! Mirrors the reference `DecisionModel` (the `laya` SDK's
+//! `common.py`):
+//!
+//! 1. a per-question-type embedding added to every position;
+//! 2. `head_layers` × PyTorch `nn.TransformerEncoderLayer` with
+//!    `norm_first=True`: pre-norm self-attention through a fused
+//!    `in_proj` with bias, then a pre-norm feed-forward whose
+//!    activation is **ReLU** — PyTorch's default, which the reference
+//!    never overrides;
+//! 3. `scorer`, applied at each option's `[MASK]` marker: LayerNorm →
+//!    Linear → exact GELU → Linear to one logit;
+//! 4. `act_head`, over the first position's hidden state plus four
+//!    summary features of the (uncalibrated) answer distribution.
+
+use anyhow::Result;
+use candle_core::{DType, Tensor};
+use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
+
+use super::backbone::{project, scaled_attention};
+
+/// PyTorch's `LayerNorm` default.
+const TORCH_LN_EPS: f64 = 1e-5;
+
+struct HeadLayer {
+    norm1: LayerNorm,
+    in_proj: Linear,
+    out_proj: Linear,
+    norm2: LayerNorm,
+    linear1: Linear,
+    linear2: Linear,
+    n_heads: usize,
+    head_dim: usize,
+}
+
+impl HeadLayer {
+    fn load(vb: VarBuilder, d: usize) -> Result<Self> {
+        // `nn.TransformerEncoder` is built with `nhead = max(1, d // 64)`
+        // and `dim_feedforward = 4 * d`.
+        let n_heads = (d / 64).max(1);
+        let f32_vb = vb.clone().set_dtype(DType::F32);
+        Ok(Self {
+            norm1: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm1"))?,
+            in_proj: linear_torch_mha_in_proj(d, vb.pp("self_attn"))?,
+            out_proj: candle_nn::linear(d, d, vb.pp("self_attn.out_proj"))?,
+            norm2: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm2"))?,
+            linear1: candle_nn::linear(d, 4 * d, vb.pp("linear1"))?,
+            linear2: candle_nn::linear(4 * d, d, vb.pp("linear2"))?,
+            n_heads,
+            head_dim: d / n_heads,
+        })
+    }
+
+    /// `xs` is the f32 residual stream; projections run in the weights'
+    /// dtype and join it back in f32, as under autocast.
+    fn forward(&self, xs: &Tensor, key_mask: Option<&Tensor>) -> Result<Tensor> {
+        let (b, l, d) = xs.dims3()?;
+        let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?
+            .reshape((b, l, 3, self.n_heads, self.head_dim))?
+            .permute((2, 0, 3, 1, 4))?;
+        let q = qkv.get(0)?.contiguous()?;
+        let k = qkv.get(1)?.contiguous()?;
+        let v = qkv.get(2)?.contiguous()?;
+        let ctx = scaled_attention(&q, &k, &v, key_mask, self.head_dim)?
+            .transpose(1, 2)?
+            .reshape((b, l, d))?;
+        let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
+        let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
+            .relu()?
+            .apply(&self.linear2)?
+            .to_dtype(DType::F32)?;
+        Ok((xs + ff)?)
+    }
+
+    /// The flash path over a packed `(total, d)` f32 stream.
+    #[cfg(feature = "flash-attn")]
+    fn forward_packed(&self, xs: &Tensor, packed: &super::backbone::Packed) -> Result<Tensor> {
+        use super::backbone::{flash_attention, split_qkv};
+        let (t, d) = xs.dims2()?;
+        let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?;
+        let (q, k, v) = split_qkv(&qkv, self.n_heads, self.head_dim)?;
+        let ctx = flash_attention(&q, &k, &v, packed, None)?.reshape((t, d))?;
+        let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
+        let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
+            .relu()?
+            .apply(&self.linear2)?
+            .to_dtype(DType::F32)?;
+        Ok((xs + ff)?)
+    }
+}
+
+/// `nn.MultiheadAttention` stores q/k/v as one `in_proj_weight`
+/// `(3d, d)` and `in_proj_bias` `(3d)`, rows ordered q then k then v —
+/// exactly the layout one fused `Linear` wants.
+fn linear_torch_mha_in_proj(d: usize, vb: VarBuilder) -> Result<Linear> {
+    let w = vb.get((3 * d, d), "in_proj_weight")?;
+    let b = vb.get(3 * d, "in_proj_bias")?;
+    Ok(Linear::new(w, Some(b)))
+}
+
+pub struct DecisionHead {
+    type_emb: Embedding,
+    layers: Vec<HeadLayer>,
+    scorer_norm: LayerNorm,
+    scorer_fc: Linear,
+    scorer_out: Linear,
+    act_fc: Linear,
+    act_out: Linear,
+}
+
+impl DecisionHead {
+    /// `vb` is the checkpoint root; the head's tensors are top-level.
+    /// Linear weights load in `vb`'s dtype; the type embedding and norms
+    /// in f32.
+    pub fn load(vb: VarBuilder, d: usize, head_layers: usize, n_act: usize) -> Result<Self> {
+        let layers = (0..head_layers)
+            .map(|i| HeadLayer::load(vb.pp(format!("head.layers.{i}")), d))
+            .collect::<Result<Vec<_>>>()?;
+        let f32_vb = vb.clone().set_dtype(DType::F32);
+        Ok(Self {
+            type_emb: candle_nn::embedding(3, d, f32_vb.pp("type_emb"))?,
+            layers,
+            scorer_norm: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("scorer.0"))?,
+            scorer_fc: candle_nn::linear(d, d, vb.pp("scorer.1"))?,
+            scorer_out: candle_nn::linear(d, 1, vb.pp("scorer.3"))?,
+            act_fc: candle_nn::linear(d + 4, 256, vb.pp("act_head.0"))?,
+            act_out: candle_nn::linear(256, n_act, vb.pp("act_head.2"))?,
+        })
+    }
+
+    /// Type embedding plus the transformer layers. `h` is the encoder
+    /// output `(b, l, d)`, `qtype` is `(b,)` u32.
+    pub fn contextualise(
+        &self,
+        h: &Tensor,
+        qtype: &Tensor,
+        key_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let mut h = h.broadcast_add(&qtype.apply(&self.type_emb)?.unsqueeze(1)?)?;
+        for layer in &self.layers {
+            h = layer.forward(&h, key_mask)?;
+        }
+        Ok(h)
+    }
+
+    /// The flash path: `h` is the packed encoder output `(total, d)` and
+    /// `qtype` the per-token question type `(total,)` u32.
+    #[cfg(feature = "flash-attn")]
+    pub fn contextualise_packed(
+        &self,
+        h: &Tensor,
+        qtype: &Tensor,
+        packed: &super::backbone::Packed,
+    ) -> Result<Tensor> {
+        let mut h = (h + qtype.apply(&self.type_emb)?)?;
+        for layer in &self.layers {
+            h = layer.forward_packed(&h, packed)?;
+        }
+        Ok(h)
+    }
+
+    /// Score gathered marker states `(b, k, d)` → `(b, k)` f32 logits.
+    pub fn score(&self, markers: &Tensor) -> Result<Tensor> {
+        let x = project(&self.scorer_fc, &markers.apply(&self.scorer_norm)?)?
+            .gelu_erf()?
+            .apply(&self.scorer_out)?;
+        Ok(x.squeeze(2)?.to_dtype(DType::F32)?)
+    }
+
+    /// Act logits `(b, n_act)` f32 from the pooled first position
+    /// `(b, d)` and the four answer-distribution features `(b, 4)`.
+    pub fn act(&self, pooled: &Tensor, feats: &Tensor) -> Result<Tensor> {
+        let x = Tensor::cat(&[pooled, &feats.to_dtype(pooled.dtype())?], 1)?;
+        let x = project(&self.act_fc, &x)?
+            .gelu_erf()?
+            .apply(&self.act_out)?;
+        Ok(x.to_dtype(DType::F32)?)
+    }
+}

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

@@ -0,0 +1,64 @@
+//! Laya — a non-autoregressive decision model (`convaiinnovations/laya`).
+//!
+//! Not a language model: a ModernBERT encoder plus a small decision
+//! head that scores the options of a typed question in one forward
+//! pass. The caller assembles each question into one token sequence
+//! (`[CLS] <type> question: <instructions> [SEP] [MASK] opt0 [MASK]
+//! opt1 … [SEP] state [SEP]`) and says where each option's `[MASK]`
+//! marker sits; this module returns one raw score per marker plus the
+//! act head's distribution. Calibration (temperature scaling) and
+//! answer shaping happen after, CPU-side, from the checkpoint's
+//! `rl_agent_config.json`.
+//!
+//! The reference implementation is the `laya` SDK's `DecisionModel`
+//! (`laya/common.py`), not the `rl_common.py` shipped in the model
+//! repository, which has diverged from it.
+//!
+//! - [`config`] — the two config files a checkpoint ships.
+//! - [`backbone`] — ModernBERT, vendored so it runs outside f32.
+//! - [`head`] — type embedding, the head's transformer layers, the
+//!   marker scorer and the act head.
+//! - [`model`] — the composition, batch in, [`DecisionRow`]s out.
+
+pub mod backbone;
+pub mod config;
+pub mod head;
+pub mod model;
+
+pub use config::CheckpointConfig;
+pub use model::LayaModel;
+
+/// Question type, as the checkpoint's type embedding indexes it.
+pub const QTYPE_CHOICE: u8 = 0;
+pub const QTYPE_SCORE: u8 = 1;
+pub const QTYPE_NOUL: u8 = 2;
+
+/// One assembled question: its token ids, where each option's
+/// `[MASK]` marker sits, and its type.
+#[derive(Debug, Clone, PartialEq)]
+pub struct DecisionSeq {
+    pub input_ids: Vec<u32>,
+    pub markers: Vec<usize>,
+    /// [`QTYPE_CHOICE`], [`QTYPE_SCORE`] or [`QTYPE_NOUL`].
+    pub qtype: u8,
+}
+
+/// Sequences scored in one forward. One request's questions normally
+/// travel together, so the state they share is encoded in one pass.
+#[derive(Debug, Clone, Default, PartialEq)]
+pub struct DecisionBatch {
+    pub seqs: Vec<DecisionSeq>,
+}
+
+/// One sequence's result, CPU-side.
+#[derive(Debug, Clone, PartialEq)]
+pub struct DecisionRow {
+    /// One uncalibrated score per marker, in marker order.
+    pub logits: Vec<f32>,
+    /// Softmax over the act head's outputs. Index 0 is what the
+    /// reference reports as `action.act_probability`.
+    pub act: Vec<f32>,
+}
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,260 @@
+//! A whole Laya checkpoint: encoder + decision head, batch in, raw
+//! scores out.
+
+use anyhow::{Context, Result, ensure};
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use std::path::Path;
+
+use super::backbone::{Masks, ModernBert};
+use super::config::CheckpointConfig;
+use super::head::DecisionHead;
+use super::{DecisionBatch, DecisionRow};
+
+pub struct LayaModel {
+    encoder: ModernBert,
+    head: DecisionHead,
+    config: CheckpointConfig,
+    device: Device,
+    /// Precision of the linear layers (the load dtype).
+    dtype: DType,
+}
+
+/// A validated batch, flattened: what both forward paths start from.
+struct Flat {
+    lens: Vec<usize>,
+    qtypes: Vec<u32>,
+    /// Each row's marker positions, padded to `k_max` with 0.
+    marker_idx: Vec<u32>,
+    k_max: usize,
+    seq_len: usize,
+}
+
+impl LayaModel {
+    /// Load one checkpoint directory (the repo root, or one of its
+    /// subfolders): `rl_agent_config.json`, `encoder/config.json` and
+    /// `model.safetensors`.
+    pub fn load(dir: &Path, device: &Device, dtype: DType) -> Result<Self> {
+        let config = CheckpointConfig::from_dir(dir)?;
+        let weights = dir.join("model.safetensors");
+        // SAFETY: the file is memory-mapped read-only for the lifetime
+        // of the VarBuilder; nothing writes to a checkpoint in place.
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[&weights], dtype, device) }
+            .with_context(|| format!("mmap {}", weights.display()))?;
+        Self::from_var_builder(vb, config, device)
+    }
+
+    pub fn from_var_builder(
+        vb: VarBuilder,
+        config: CheckpointConfig,
+        device: &Device,
+    ) -> Result<Self> {
+        let dtype = vb.dtype();
+        let encoder = ModernBert::load(vb.pp("encoder"), &config.encoder)?;
+        let head = DecisionHead::load(
+            vb.clone(),
+            config.encoder.hidden_size,
+            config.laya.head_layers,
+            config.laya.n_act(),
+        )?;
+        Ok(Self {
+            encoder,
+            head,
+            config,
+            device: device.clone(),
+            dtype,
+        })
+    }
+
+    pub fn config(&self) -> &CheckpointConfig {
+        &self.config
+    }
+
+    /// Precision the linear layers were loaded in.
+    pub fn dtype(&self) -> DType {
+        self.dtype
+    }
+
+    /// One forward over every sequence in `batch`.
+    ///
+    /// On a CUDA device in reduced precision, in a build with the
+    /// `flash-attn` feature, the batch is packed without padding and
+    /// attention runs in flash varlen kernels. Otherwise rows are
+    /// right-padded to the longest, as the reference collates them, and
+    /// attention is eager; padding changes nothing about an unpadded
+    /// row's result.
+    pub fn forward(&self, batch: &DecisionBatch) -> Result<Vec<DecisionRow>> {
+        let flat = self.validate(batch)?;
+        #[cfg(feature = "flash-attn")]
+        if self.device.is_cuda() && self.dtype != DType::F32 && flash_enabled() {
+            return self.forward_packed(batch, &flat);
+        }
+        self.forward_padded(batch, &flat)
+    }
+
+    fn validate(&self, batch: &DecisionBatch) -> Result<Flat> {
+        let seqs = &batch.seqs;
+        ensure!(!seqs.is_empty(), "empty decision batch");
+        let max_pos = self.config.encoder.max_position_embeddings;
+        let seq_len = seqs.iter().map(|s| s.input_ids.len()).max().unwrap_or(0);
+        ensure!(seq_len > 0, "decision batch has an empty sequence");
+        ensure!(
+            seq_len <= max_pos,
+            "sequence of {seq_len} tokens exceeds max_position_embeddings {max_pos}"
+        );
+        let vocab = self.config.encoder.vocab_size as u32;
+        let k_max = seqs.iter().map(|s| s.markers.len()).max().unwrap_or(0);
+        ensure!(
+            k_max > 0,
+            "decision batch has a sequence with no option markers"
+        );
+        let mut lens = Vec::with_capacity(seqs.len());
+        let mut qtypes = Vec::with_capacity(seqs.len());
+        let mut marker_idx = Vec::with_capacity(seqs.len() * k_max);
+        for s in seqs {
+            let n = s.input_ids.len();
+            ensure!(n > 0, "decision batch has an empty sequence");
+            ensure!(!s.markers.is_empty(), "sequence has no option markers");
+            ensure!(s.qtype < 3, "question type {} out of range", s.qtype);
+            if let Some(&bad) = s.input_ids.iter().find(|&&t| t >= vocab) {
+                anyhow::bail!("token id {bad} outside vocabulary of {vocab}");
+            }
+            if let Some(&bad) = s.markers.iter().find(|&&m| m >= n) {
+                anyhow::bail!("marker position {bad} outside a {n}-token sequence");
+            }
+            lens.push(n);
+            qtypes.push(s.qtype as u32);
+            marker_idx.extend(s.markers.iter().map(|&m| m as u32));
+            // Absent slots gather position 0; their scores are discarded.
+            marker_idx.extend(std::iter::repeat_n(0u32, k_max - s.markers.len()));
+        }
+        Ok(Flat {
+            lens,
+            qtypes,
+            marker_idx,
+            k_max,
+            seq_len,
+        })
+    }
+
+    /// The eager path: rows right-padded to `(b, seq_len)`.
+    fn forward_padded(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
+        let b = batch.seqs.len();
+        let seq_len = flat.seq_len;
+        let pad = self.config.encoder.pad_token_id;
+        let mut ids = Vec::with_capacity(b * seq_len);
+        for s in &batch.seqs {
+            ids.extend_from_slice(&s.input_ids);
+            ids.extend(std::iter::repeat_n(pad, seq_len - s.input_ids.len()));
+        }
+        let ids = Tensor::from_vec(ids, (b, seq_len), &self.device)?;
+        let masks = Masks::new(
+            &flat.lens,
+            seq_len,
+            self.encoder.half_window(),
+            &self.device,
+        )?;
+        let qtype = Tensor::from_vec(flat.qtypes.clone(), b, &self.device)?;
+
+        let h = self.encoder.forward(&ids, &masks)?;
+        let h = self.head.contextualise(&h, &qtype, masks.global.as_ref())?;
+        let d = self.config.encoder.hidden_size;
+        let idx = Tensor::from_vec(flat.marker_idx.clone(), (b, flat.k_max, 1), &self.device)?
+            .broadcast_as((b, flat.k_max, d))?
+            .contiguous()?;
+        let marker_h = h.gather(&idx, 1)?;
+        let pooled = h.narrow(1, 0, 1)?.squeeze(1)?;
+        self.finish(batch, &marker_h, &pooled)
+    }
+
+    /// The flash path: tokens packed without padding, `(total, d)`.
+    #[cfg(feature = "flash-attn")]
+    fn forward_packed(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
+        use super::backbone::Packed;
+        let b = batch.seqs.len();
+        let packed = Packed::new(&flat.lens, &self.device)?;
+        let total: usize = flat.lens.iter().sum();
+        let mut ids = Vec::with_capacity(total);
+        let mut qtype_tok = Vec::with_capacity(total);
+        for (s, &q) in batch.seqs.iter().zip(&flat.qtypes) {
+            ids.extend_from_slice(&s.input_ids);
+            qtype_tok.extend(std::iter::repeat_n(q, s.input_ids.len()));
+        }
+        let ids = Tensor::from_vec(ids, total, &self.device)?;
+        let qtype_tok = Tensor::from_vec(qtype_tok, total, &self.device)?;
+
+        let h = self.encoder.forward_packed(&ids, &packed)?;
+        let h = self.head.contextualise_packed(&h, &qtype_tok, &packed)?;
+        let d = self.config.encoder.hidden_size;
+        // Marker positions, and each row's first token, in packed order.
+        let global: Vec<u32> = flat
+            .marker_idx
+            .chunks(flat.k_max)
+            .zip(&packed.offsets)
+            .flat_map(|(ms, &off)| ms.iter().map(move |&m| off as u32 + m))
+            .collect();
+        let global = Tensor::from_vec(global, b * flat.k_max, &self.device)?;
+        let marker_h = h.index_select(&global, 0)?.reshape((b, flat.k_max, d))?;
+        let firsts: Vec<u32> = packed.offsets.iter().map(|&o| o as u32).collect();
+        let pooled = h.index_select(&Tensor::from_vec(firsts, b, &self.device)?, 0)?;
+        self.finish(batch, &marker_h, &pooled)
+    }
+
+    /// Score the gathered marker states `(b, k_max, d)` and run the act
+    /// head over the pooled first positions `(b, d)`.
+    fn finish(
+        &self,
+        batch: &DecisionBatch,
+        marker_h: &Tensor,
+        pooled: &Tensor,
+    ) -> Result<Vec<DecisionRow>> {
+        let b = batch.seqs.len();
+        let logits: Vec<Vec<f32>> = self.head.score(marker_h)?.to_vec2()?;
+        let mut feats = Vec::with_capacity(b * 4);
+        let mut row_logits = Vec::with_capacity(b);
+        for (s, full) in batch.seqs.iter().zip(logits) {
+            let row: Vec<f32> = full[..s.markers.len()].to_vec();
+            feats.extend_from_slice(&act_features(&row));
+            row_logits.push(row);
+        }
+        let feats = Tensor::from_vec(feats, (b, 4), &self.device)?;
+        let act_logits = self.head.act(pooled, &feats)?;
+        let act = candle_nn::ops::softmax_last_dim(&act_logits)?.to_vec2::<f32>()?;
+        Ok(row_logits
+            .into_iter()
+            .zip(act)
+            .map(|(logits, act)| DecisionRow { logits, act })
+            .collect())
+    }
+}
+
+/// Runtime kill switch for the flash path: `NEURON_FLASH_ATTN=0` (or
+/// `false`) forces eager attention without a rebuild — the same switch
+/// the text architectures read. Read once.
+#[cfg(feature = "flash-attn")]
+fn flash_enabled() -> bool {
+    static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *ENABLED
+        .get_or_init(|| !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false"))
+}
+
+/// The act head's four inputs, from one row's uncalibrated logits:
+/// top-1 probability, top-1 minus top-2, entropy normalised by
+/// `ln(max(k, 2))`, and `k / 255`.
+///
+/// The reference computes these over the batch's padded width, with
+/// absent option slots filled with a logit of -1e4. Those slots carry
+/// exactly zero probability in f32, so computing per row over only the
+/// real options gives the same result.
+pub fn act_features(logits: &[f32]) -> [f32; 4] {
+    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+    let exps: Vec<f32> = logits.iter().map(|&z| (z - max).exp()).collect();
+    let sum: f32 = exps.iter().sum();
+    let mut p: Vec<f32> = exps.iter().map(|e| e / sum).collect();
+    let k = (logits.len().max(2)) as f32;
+    let ent = -p.iter().map(|&x| x * x.max(1e-9).ln()).sum::<f32>() / k.ln();
+    p.sort_by(|a, b| b.total_cmp(a));
+    let top1 = p[0];
+    let top2 = p.get(1).copied().unwrap_or(0.0);
+    [top1, top1 - top2, ent, k / 255.0]
+}

crates/neuron/src/harness/arch/laya/tests.rs

@@ -0,0 +1,450 @@
+//! Parity against the upstream reference.
+//!
+//! `testdata/laya/tiny/` is a toy-sized checkpoint built and scored by
+//! the SDK's own `DecisionModel` (`script/generate_laya_fixture.py`);
+//! it runs in CI. `testdata/laya/reference.json` holds the released
+//! checkpoints' outputs over a real corpus (`script/laya-reference.py`);
+//! scoring it needs the weights, so that test is ignored unless
+//! `LAYA_REFERENCE_DIR` names a snapshot of `convaiinnovations/laya`.
+
+use super::*;
+use candle_core::{DType, Device};
+use serde::Deserialize;
+use std::path::{Path, PathBuf};
+
+fn testdata() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
+}
+
+#[derive(Deserialize)]
+struct TinyExpected {
+    seqs: Vec<TinySeq>,
+}
+
+#[derive(Deserialize)]
+struct TinySeq {
+    input_ids: Vec<u32>,
+    markers: Vec<usize>,
+    qtype: u8,
+    logits: Vec<f32>,
+    act: Vec<f32>,
+}
+
+fn load_tiny() -> (LayaModel, TinyExpected) {
+    let dir = testdata().join("tiny");
+    let model = LayaModel::load(&dir, &Device::Cpu, DType::F32).expect("load tiny checkpoint");
+    let expected: TinyExpected = serde_json::from_str(
+        &std::fs::read_to_string(dir.join("expected.json")).expect("read expected.json"),
+    )
+    .expect("parse expected.json");
+    (model, expected)
+}
+
+fn batch_of(seqs: &[TinySeq]) -> DecisionBatch {
+    DecisionBatch {
+        seqs: seqs
+            .iter()
+            .map(|s| DecisionSeq {
+                input_ids: s.input_ids.clone(),
+                markers: s.markers.clone(),
+                qtype: s.qtype,
+            })
+            .collect(),
+    }
+}
+
+/// Device and precision for the weight-backed tests: `LAYA_DEVICE`
+/// (`cpu`, or `cuda:N` in a build with the `cuda` feature) and
+/// `LAYA_DTYPE` (`f32`, `bf16` or `f16`). Defaults: CPU, f32.
+fn test_device() -> (Device, DType) {
+    let device = match std::env::var("LAYA_DEVICE").as_deref() {
+        Ok("cpu") | Err(_) => Device::Cpu,
+        Ok(s) => {
+            let ordinal = s
+                .strip_prefix("cuda:")
+                .and_then(|n| n.parse().ok())
+                .unwrap_or_else(|| panic!("LAYA_DEVICE={s}: expected cpu or cuda:N"));
+            Device::new_cuda(ordinal).unwrap_or_else(|e| panic!("LAYA_DEVICE={s}: {e}"))
+        }
+    };
+    let dtype = match std::env::var("LAYA_DTYPE").as_deref() {
+        Ok("f32") | Err(_) => DType::F32,
+        Ok("bf16") => DType::BF16,
+        Ok("f16") => DType::F16,
+        Ok(s) => panic!("LAYA_DTYPE={s}: expected f32, bf16 or f16"),
+    };
+    (device, dtype)
+}
+
+/// Logit tolerance against the f32 reference for a given precision.
+/// Reduced precision is judged mainly on argmax agreement; the bound
+/// only catches a broken kernel, not rounding.
+fn logit_tolerance(dtype: DType) -> f32 {
+    if dtype == DType::F32 { 1e-3 } else { 1.0 }
+}
+
+fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+    assert_eq!(a.len(), b.len(), "length mismatch");
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| (x - y).abs())
+        .fold(0.0, f32::max)
+}
+
+/// The whole forward — encoder with both attention kinds and rotary
+/// bases, padding, type embedding, head layers, marker gather, scorer
+/// and act head — against the reference, in f32.
+#[test]
+fn tiny_checkpoint_matches_reference() {
+    let (model, expected) = load_tiny();
+    let rows = model.forward(&batch_of(&expected.seqs)).expect("forward");
+    assert_eq!(rows.len(), expected.seqs.len());
+    for (i, (got, want)) in rows.iter().zip(&expected.seqs).enumerate() {
+        let dl = max_abs_diff(&got.logits, &want.logits);
+        let da = max_abs_diff(&got.act, &want.act);
+        assert!(
+            dl < 1e-5 && da < 1e-5,
+            "row {i}: logit diff {dl:e}, act diff {da:e}\n got {:?} / {:?}\nwant {:?} / {:?}",
+            got.logits,
+            got.act,
+            want.logits,
+            want.act
+        );
+    }
+}
+
+/// A row's result must not depend on what it was batched with: padded
+/// keys are masked out of every attention, including the head's.
+#[test]
+fn padding_does_not_change_a_row() {
+    let (model, expected) = load_tiny();
+    let batched = model.forward(&batch_of(&expected.seqs)).expect("batched");
+    for (i, seq) in expected.seqs.iter().enumerate() {
+        let alone = model
+            .forward(&batch_of(std::slice::from_ref(seq)))
+            .expect("single");
+        let dl = max_abs_diff(&alone[0].logits, &batched[i].logits);
+        let da = max_abs_diff(&alone[0].act, &batched[i].act);
+        assert!(dl < 1e-5 && da < 1e-5, "row {i}: {dl:e} / {da:e}");
+    }
+}
+
+#[test]
+fn malformed_batches_are_refused() {
+    let (model, _) = load_tiny();
+    let seq = |ids: Vec<u32>, markers: Vec<usize>, qtype| DecisionBatch {
+        seqs: vec![DecisionSeq {
+            input_ids: ids,
+            markers,
+            qtype,
+        }],
+    };
+    assert!(model.forward(&DecisionBatch::default()).is_err());
+    assert!(model.forward(&seq(vec![1, 2, 3], vec![], 0)).is_err());
+    assert!(model.forward(&seq(vec![1, 2, 3], vec![3], 0)).is_err());
+    assert!(model.forward(&seq(vec![1, 2, 3], vec![1], 3)).is_err());
+    assert!(model.forward(&seq(vec![1, 500, 3], vec![1], 0)).is_err());
+}
+
+#[test]
+fn act_features_match_the_reference_definition() {
+    // Two equal logits: top1 = top2 = 0.5, entropy exactly 1.
+    let f = model::act_features(&[0.3, 0.3]);
+    assert!((f[0] - 0.5).abs() < 1e-6 && f[1].abs() < 1e-6 && (f[2] - 1.0).abs() < 1e-6);
+    assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
+    // One option: probability 1, no second, k clamps to 2.
+    let f = model::act_features(&[4.0]);
+    assert_eq!(f[0], 1.0);
+    assert_eq!(f[1], 1.0);
+    assert!(f[2].abs() < 1e-6);
+    assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
+}
+
+// ---------------------------------------------------------------- released checkpoints
+
+#[derive(Deserialize)]
+struct Reference {
+    cases: Vec<RefCase>,
+}
+
+#[derive(Deserialize)]
+struct RefCase {
+    name: String,
+    items: Vec<RefItem>,
+    response: serde_json::Value,
+}
+
+#[derive(Deserialize)]
+struct RefItem {
+    input_ids: Vec<u32>,
+    markers: Vec<usize>,
+    qtype: u8,
+    logits: Vec<f32>,
+    act: Vec<f32>,
+}
+
+/// Every case in `reference.json`, scored by the checkpoint the
+/// reference routed it to, on the device and precision
+/// [`test_device`] selects (f32 on the CPU by default).
+///
+/// `LAYA_REFERENCE_DIR` is a snapshot of `convaiinnovations/laya` at the
+/// revision recorded in the fixture header: the root holds the English
+/// checkpoint, `multilingual/` and `typed-decisions/` the other two.
+#[test]
+#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
+fn released_checkpoints_match_reference() {
+    let Ok(root) = std::env::var("LAYA_REFERENCE_DIR") else {
+        panic!("LAYA_REFERENCE_DIR is not set");
+    };
+    let root = PathBuf::from(root);
+    let (device, dtype) = test_device();
+    let tolerance = logit_tolerance(dtype);
+    let reference: Reference = serde_json::from_str(
+        &std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
+    )
+    .expect("parse reference.json");
+
+    // The reference's own GPU run (bf16 autocast), for judging reduced
+    // precision against the drift upstream itself serves with.
+    let torch_bf16: Option<Reference> =
+        std::fs::read_to_string(testdata().join("reference_cuda_bf16.json"))
+            .ok()
+            .map(|t| serde_json::from_str(&t).expect("parse reference_cuda_bf16.json"));
+    // Per checkpoint: summed centred drift of ours and of the reference's
+    // bf16 run, their maxima, and the row count.
+    let mut drift: std::collections::BTreeMap<String, (f32, f32, f32, f32, usize)> =
+        Default::default();
+
+    let mut models = std::collections::BTreeMap::new();
+    let mut stats: std::collections::BTreeMap<String, (usize, f32, f32, usize)> =
+        Default::default();
+    // Forward wall time and tokens encoded, per checkpoint: a first
+    // CPU number, not a benchmark (debug builds are far slower).
+    let mut timing: std::collections::BTreeMap<String, (f64, usize)> = Default::default();
+    for case in &reference.cases {
+        let ckpt = case.response["routing"]["model"]
+            .as_str()
+            .expect("routing.model")
+            .to_string();
+        let model = models.entry(ckpt.clone()).or_insert_with(|| {
+            let dir = match ckpt.as_str() {
+                "english" => root.clone(),
+                other => root.join(other),
+            };
+            LayaModel::load(&dir, &device, dtype).unwrap_or_else(|e| panic!("load {ckpt}: {e:#}"))
+        });
+        let batch = DecisionBatch {
+            seqs: case
+                .items
+                .iter()
+                .map(|it| DecisionSeq {
+                    input_ids: it.input_ids.clone(),
+                    markers: it.markers.clone(),
+                    qtype: it.qtype,
+                })
+                .collect(),
+        };
+        let started = std::time::Instant::now();
+        let rows = model
+            .forward(&batch)
+            .unwrap_or_else(|e| panic!("{}: {e:#}", case.name));
+        let t = timing.entry(ckpt.clone()).or_default();
+        t.0 += started.elapsed().as_secs_f64() * 1e3;
+        t.1 += batch.seqs.iter().map(|s| s.input_ids.len()).sum::<usize>();
+        let torch_case = torch_bf16
+            .as_ref()
+            .and_then(|r| r.cases.iter().find(|c| c.name == case.name));
+        for (i, (got, want)) in rows.iter().zip(&case.items).enumerate() {
+            let d = drift.entry(ckpt.clone()).or_default();
+            let ours = centred_diff(&got.logits, &want.logits);
+            d.0 += ours;
+            d.2 = d.2.max(ours);
+            if let Some(t) = torch_case.and_then(|c| c.items.get(i)) {
+                let theirs = centred_diff(&t.logits, &want.logits);
+                d.1 += theirs;
+                d.3 = d.3.max(theirs);
+            }
+            d.4 += 1;
+        }
+        let entry = stats.entry(ckpt.clone()).or_default();
+        for (got, want) in rows.iter().zip(&case.items) {
+            let dl = max_abs_diff(&got.logits, &want.logits);
+            let da = max_abs_diff(&got.act, &want.act);
+            let argmax = |v: &[f32]| {
+                v.iter()
+                    .enumerate()
+                    .max_by(|a, b| a.1.total_cmp(b.1))
+                    .map(|(i, _)| i)
+            };
+            entry.0 += 1;
+            entry.1 = entry.1.max(dl);
+            entry.2 = entry.2.max(da);
+            if argmax(&got.logits) == argmax(&want.logits) {
+                entry.3 += 1;
+            }
+            if dtype != DType::F32 && std::env::var_os("LAYA_VERBOSE").is_some() {
+                eprintln!("  {:<40} {ckpt:<16} |Δlogit| {dl:.3}", case.name);
+            }
+            assert!(
+                dl < tolerance,
+                "{} ({ckpt}): logit diff {dl:e}\n got {:?}\nwant {:?}",
+                case.name,
+                got.logits,
+                want.logits
+            );
+            assert!(da < tolerance, "{} ({ckpt}): act diff {da:e}", case.name);
+        }
+    }
+    for (ckpt, (n, dl, da, agree)) in &stats {
+        let (ms, tokens) = timing[ckpt];
+        eprintln!(
+            "[{device:?} {dtype:?}] {ckpt}: {n} questions, max |Δlogit| {dl:.2e}, max |Δact| {da:.2e}, argmax {agree}/{n}; \
+             {tokens} tokens in {ms:.0} ms ({:.2} ms/token)",
+            ms / tokens as f64
+        );
+        assert_eq!(agree, n, "{ckpt}: argmax disagreement");
+    }
+    // Reduced precision: centred drift (softmax ignores a constant shift)
+    // against f32, next to the reference's own bf16 drift. Per-case drift
+    // is chaotic, so the gate is the mean: no worse than twice upstream's,
+    // plus 0.01 — about one bf16 ulp at these logits' magnitude — so a
+    // checkpoint with a single fixture case whose reference drift happens
+    // to round to nothing is not held to zero.
+    for (ckpt, (ours, theirs, ours_max, theirs_max, n)) in &drift {
+        let (ours, theirs) = (ours / *n as f32, theirs / *n as f32);
+        eprintln!(
+            "[{device:?} {dtype:?}] {ckpt}: centred drift mean {ours:.4} (reference bf16 {theirs:.4}), \
+             max {ours_max:.3} (reference bf16 {theirs_max:.3})"
+        );
+        if dtype != DType::F32 && torch_bf16.is_some() {
+            assert!(
+                ours <= 2.0 * theirs + 0.01,
+                "{ckpt}: mean drift {ours:.4} exceeds twice the reference's own bf16 drift {theirs:.4}"
+            );
+        }
+    }
+}
+
+/// Largest difference between two logit rows after removing each row's
+/// mean — the part of a difference a softmax can see.
+fn centred_diff(a: &[f32], b: &[f32]) -> f32 {
+    let mean = |v: &[f32]| v.iter().sum::<f32>() / v.len() as f32;
+    let (ma, mb) = (mean(a), mean(b));
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| ((x - ma) - (y - mb)).abs())
+        .fold(0.0, f32::max)
+}
+
+/// Forward latency for the four workloads #337 benchmarks against the
+/// PyTorch reference, plus batch throughput, on the English checkpoint.
+///
+/// Sequences are real ones from `reference.json`: `choice_basic` is the
+/// short state (41 tokens), `multi_question` five questions over it, and
+/// `state_right_truncated` a full 512-token sequence (repeated for the
+/// five-question and throughput batches). The timed span is the whole
+/// `forward` call — host-to-device copy, forward, and the CPU readback
+/// that ends it — so it is directly comparable to the reference's
+/// end-to-end `predict()` minus tokenisation.
+///
+/// Run with `cargo test --release -p neuron [--features cuda] --
+/// --ignored --nocapture laya_forward_timing`, with `LAYA_REFERENCE_DIR`
+/// and optionally `LAYA_DEVICE` / `LAYA_DTYPE` set.
+#[test]
+#[ignore = "benchmark: needs the released weights (LAYA_REFERENCE_DIR)"]
+fn laya_forward_timing() {
+    let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
+    let (device, dtype) = test_device();
+    let reference: Reference = serde_json::from_str(
+        &std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
+    )
+    .expect("parse reference.json");
+    let seqs_of = |name: &str| -> Vec<DecisionSeq> {
+        reference
+            .cases
+            .iter()
+            .find(|c| c.name == name)
+            .unwrap_or_else(|| panic!("no case {name}"))
+            .items
+            .iter()
+            .map(|it| DecisionSeq {
+                input_ids: it.input_ids.clone(),
+                markers: it.markers.clone(),
+                qtype: it.qtype,
+            })
+            .collect()
+    };
+    let long = seqs_of("state_right_truncated").remove(0);
+    let workloads: Vec<(String, DecisionBatch)> = vec![
+        (
+            "short x1q".into(),
+            DecisionBatch {
+                seqs: seqs_of("choice_basic"),
+            },
+        ),
+        (
+            "short x5q".into(),
+            DecisionBatch {
+                seqs: seqs_of("multi_question"),
+            },
+        ),
+        (
+            "512tok x1q".into(),
+            DecisionBatch {
+                seqs: vec![long.clone()],
+            },
+        ),
+        (
+            "512tok x5q".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 5],
+            },
+        ),
+        (
+            "512tok x16".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 16],
+            },
+        ),
+        (
+            "512tok x64".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 64],
+            },
+        ),
+    ];
+
+    let model = LayaModel::load(&root, &device, dtype).expect("load English checkpoint");
+    // `LAYA_TIMING_ONLY` keeps the workloads whose name contains it, and
+    // `LAYA_TIMING_ITERS` sets the timed iterations — for profiling one
+    // workload under a tracer.
+    let only = std::env::var("LAYA_TIMING_ONLY").ok();
+    let iters: usize = std::env::var("LAYA_TIMING_ITERS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(20);
+    for (name, batch) in &workloads {
+        if only.as_deref().is_some_and(|o| !name.contains(o)) {
+            continue;
+        }
+        for _ in 0..3 {
+            model.forward(batch).expect("warm-up forward");
+        }
+        let mut ms: Vec<f64> = (0..iters)
+            .map(|_| {
+                let started = std::time::Instant::now();
+                model.forward(batch).expect("forward");
+                started.elapsed().as_secs_f64() * 1e3
+            })
+            .collect();
+        ms.sort_by(f64::total_cmp);
+        let median = ms[ms.len() / 2];
+        let tokens: usize = batch.seqs.iter().map(|s| s.input_ids.len()).sum();
+        eprintln!(
+            "[{device:?} {dtype:?}] {name:>11}: median {median:7.1} ms, min {:7.1} ms, \
+             {tokens} tokens, {:.0} seq/s",
+            ms[0],
+            batch.seqs.len() as f64 / (median / 1e3)
+        );
+    }
+}

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

@@ -20,6 +20,8 @@
 //! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
 //! the pattern set by `tp_qwen3.rs`.
 
+pub mod laya;
+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/decoder.rs

@@ -77,7 +77,25 @@ impl Qwen3_5DecoderLayer {
                     cfg.layer_types.len()
                 )
             })?;
+        Self::load_typed(cfg, rotary, layer_type, cfg.layer_uses_moe(layer_idx), vb)
+    }
 
+    /// Load a layer whose flavour is stated rather than looked up.
+    ///
+    /// The MTP head's layer (#96) sits at `mtp.layers.0`, outside the
+    /// `layer_types` list that indexes the target's own layers — asking
+    /// that list about it would answer for a target layer that happens
+    /// to share the index. The checkpoint says what it is: plain
+    /// `self_attn.{q,k,v,o}_proj` with `q_norm`/`k_norm`, and a dense
+    /// `mlp`, i.e. full attention without MoE, whatever the target
+    /// layer at index 0 happens to be.
+    pub fn load_typed(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        layer_type: &str,
+        use_moe: bool,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
         let attention = match layer_type {
             "full_attention" => {
                 AttentionKind::Full(Qwen3_5Attention::load(cfg, rotary, &vb.pp("self_attn"))?)
@@ -86,12 +104,12 @@ impl Qwen3_5DecoderLayer {
                 AttentionKind::Linear(GatedDeltaNet::load(cfg, &vb.pp("linear_attn"))?)
             }
             other => anyhow::bail!(
-                "unknown layer_type '{other}' for layer {layer_idx} (expected \
-                 'full_attention' or 'linear_attention')"
+                "unknown layer_type '{other}' (expected 'full_attention' or \
+                 'linear_attention')"
             ),
         };
 
-        let mlp = if cfg.layer_uses_moe(layer_idx) {
+        let mlp = if use_moe {
             MlpKind::Moe(Qwen3_5MoeBlock::load(cfg, &vb.pp("mlp"))?)
         } else {
             MlpKind::Dense(Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?)

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

@@ -1391,6 +1391,8 @@ mod tests {
         let (b, l) = (1, 3);
         let cfg = TextConfig {
             vocab_size: 100,
+            mtp_num_hidden_layers: 0,
+            mtp_use_dedicated_embeddings: false,
             hidden_size: 16,
             intermediate_size: 32,
             num_hidden_layers: 1,

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

@@ -77,6 +77,7 @@ pub mod full_attn;
 pub mod linear_attn;
 pub mod mlp;
 pub mod moe;
+pub mod mtp;
 pub mod rmsnorm;
 pub mod rope;
 pub mod vision;
@@ -225,6 +226,16 @@ pub struct TextConfig {
     /// pointwise after softmax+matmul.
     #[serde(default)]
     pub attn_output_gate: bool,
+    /// Multi-token-prediction head depth (#96). `1` on the checkpoints
+    /// that ship one (`Qwen3.8-27B`, `Qwen3.5-0.8B`); absent or `0`
+    /// means the weights carry no `mtp.*` tensors.
+    #[serde(default)]
+    pub mtp_num_hidden_layers: usize,
+    /// Whether the MTP head has its own embedding table. `false` on
+    /// every checkpoint seen so far, which is why the head costs only
+    /// its own layer and not a second vocabulary-sized matrix.
+    #[serde(default)]
+    pub mtp_use_dedicated_embeddings: bool,
 
     /// One entry per decoder layer; values are `"full_attention"` or
     /// `"linear_attention"`. Length must equal `num_hidden_layers`.
@@ -883,6 +894,53 @@ impl Qwen3_5ForCausalLM {
         hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
     }
 
+    /// Logits at **every** position: `(B, L, vocab)`.
+    ///
+    /// [`Self::forward`] slices the last position off before the LM
+    /// head, which is what a plain decode step wants. Speculative
+    /// verification (#96) wants the opposite: the target forwards the
+    /// drafted block in one pass and needs its own token at each
+    /// position to compare against the draft. The forward already
+    /// computes all of them — chunked prefill relies on it — so this
+    /// only declines to throw them away.
+    pub fn forward_multi(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
+        let hidden = self.base.forward(input, offset)?;
+        hidden.apply(&self.lm_head)
+    }
+
+    /// Like [`Self::forward`], but also returns the hidden state the
+    /// logits were produced from — `(B, L, hidden)`, post-final-norm,
+    /// every position.
+    ///
+    /// The MTP head (#96) consumes exactly this tensor: vLLM's
+    /// `Qwen3NextModel::forward` returns `self.norm(...)` and its
+    /// runner hands that to the draft head. Kept beside `forward` so
+    /// the two cannot drift about which hidden state that is.
+    pub fn forward_with_hidden(
+        &mut self,
+        input: &Tensor,
+        offset: usize,
+    ) -> candle_core::Result<(Tensor, Tensor)> {
+        let (_, l) = input.dims2()?;
+        let hidden = self.base.forward(input, offset)?;
+        let logits = hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)?;
+        Ok((logits, hidden))
+    }
+
+    /// Apply the LM head to a hidden state. The MTP head has none of
+    /// its own (#96) — draft and target logits come from this same
+    /// projection, which is what makes their distributions comparable.
+    pub fn lm_head(&self, hidden: &Tensor) -> candle_core::Result<Tensor> {
+        hidden.apply(&self.lm_head)
+    }
+
+    /// The token embedding the MTP head fuses with the hidden state.
+    /// The head has none of its own (`mtp_use_dedicated_embeddings` is
+    /// false on every checkpoint seen), so it borrows the target's.
+    pub fn embed_tokens(&self, input: &Tensor) -> candle_core::Result<Tensor> {
+        self.base.embed_tokens.forward(input)
+    }
+
     /// Lockstep batched decode step (#98): `(B, 1)` input, per-row
     /// positions, padding mask from
     /// [`Qwen3_5Model::batch_decode_mask`]. Returns `(B, 1,

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
@@ -170,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 —
@@ -209,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 {
@@ -222,6 +262,14 @@ 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
@@ -294,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)
             }
         }
     }
@@ -463,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,
@@ -471,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)?)?;
@@ -758,6 +853,99 @@ 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).
     ///

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

@@ -0,0 +1,485 @@
+//! The multi-token-prediction head — `qwen3_5`'s shipped draft model.
+//!
+//! `Qwen/Qwen3.8-27B` and `Qwen/Qwen3.5-0.8B` both carry a complete
+//! extra decoder layer under `mtp.*` that nothing in this crate has
+//! ever loaded: 15 tensors of 1199 in the 27B (810 MB BF16), 15 of 488
+//! in the 0.8B (39 MB). It is the drafter speculative decoding needs
+//! (#96), already in the weights we download.
+//!
+//! ```text
+//! e = pre_fc_norm_embedding(embed(t))      [B, L, H]
+//! h = pre_fc_norm_hidden(h)                [B, L, H]  (the target's)
+//! x = fc(cat[e, h])                        [B, L, 2H] -> [B, L, H]
+//! x = layer(x)                             one full-attention block
+//! x = norm(x)                              -> the target's lm_head
+//! ```
+//!
+//! Three things here cannot be recovered from the tensor shapes, and
+//! all three were read from vLLM's `qwen3_next_mtp.py` rather than
+//! inferred — the same discipline, and for the same reason, as the
+//! sibling head in `arch/qwen4_exp/mtp.rs`:
+//!
+//! 1. **The concatenation is `[embedding, hidden]`**, not the reverse:
+//!    `hidden_states = torch.cat([inputs_embeds, hidden_states], -1)`.
+//!    `fc` is square in the sense that both halves are `hidden_size`
+//!    wide, so swapping them loads, runs, and produces finite logits
+//!    that are simply wrong. Pinned by value below.
+//! 2. **Each half is normalised before the concatenation**, by its own
+//!    norm — `pre_fc_norm_embedding` on the embedding,
+//!    `pre_fc_norm_hidden` on the hidden state — never one norm over
+//!    the joined 2H.
+//! 3. **The hidden state is the target's *post-final-norm* output**,
+//!    the same tensor its `lm_head` consumes. vLLM's
+//!    `Qwen3NextModel::forward` returns `self.norm(...)` and the runner
+//!    hands that to the head. Feeding the pre-norm residual instead is
+//!    the other plausible reading, and it is wrong for this
+//!    architecture (`qwen4_exp`'s head takes its pre-mixer stream —
+//!    they do not generalise to each other).
+//!
+//! **The head has no `lm_head` or `embed_tokens` of its own.** The
+//! checkpoint carries neither (`mtp_use_dedicated_embeddings: false`),
+//! so both come from the target — which is what makes the resident cost
+//! the 810 MB above rather than that plus a second embedding table.
+//!
+//! **The head keeps its own KV cache, and it is not the target's.** A
+//! head fed only the newest token attends over an empty past, so it has
+//! to walk the prompt alongside the target before it can draft. That
+//! puts MTP prefill inside the speculative loop's cost, not just decode
+//! — see #96.
+
+use anyhow::{Context, Result};
+use candle_core::{Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+use std::sync::Arc;
+
+use super::TextConfig;
+use super::decoder::Qwen3_5DecoderLayer;
+use super::rmsnorm::Qwen3_5RmsNorm;
+use super::rope::RotaryEmbedding;
+
+pub struct MtpHead {
+    pre_fc_norm_embedding: Qwen3_5RmsNorm,
+    pre_fc_norm_hidden: Qwen3_5RmsNorm,
+    /// `Linear(2H -> H)` over `cat[embedding, hidden]`.
+    fc: Linear,
+    layer: Qwen3_5DecoderLayer,
+    /// Final norm before the *target's* lm_head.
+    norm: Qwen3_5RmsNorm,
+}
+
+impl MtpHead {
+    /// `vb` is the checkpoint root; the head hangs off `mtp.`.
+    pub fn load(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let mtp = vb.pp("mtp");
+        let h = cfg.hidden_size;
+        let fc_weight = mtp
+            .pp("fc")
+            .get((h, 2 * h), "weight")
+            .context("load 'mtp.fc.weight'")?;
+        Ok(Self {
+            pre_fc_norm_embedding: Qwen3_5RmsNorm::load(
+                &mtp.pp("pre_fc_norm_embedding"),
+                h,
+                cfg.rms_norm_eps,
+            )
+            .context("load mtp.pre_fc_norm_embedding")?,
+            pre_fc_norm_hidden: Qwen3_5RmsNorm::load(
+                &mtp.pp("pre_fc_norm_hidden"),
+                h,
+                cfg.rms_norm_eps,
+            )
+            .context("load mtp.pre_fc_norm_hidden")?,
+            fc: Linear::new(fc_weight, None),
+            // Stated, not looked up: the head's layer is outside
+            // `layer_types`. Full attention, dense MLP — see
+            // `Qwen3_5DecoderLayer::load_typed`.
+            layer: Qwen3_5DecoderLayer::load_typed(
+                cfg,
+                rotary,
+                "full_attention",
+                false,
+                &mtp.pp("layers").pp(0),
+            )
+            .context("load mtp.layers.0")?,
+            norm: Qwen3_5RmsNorm::load(&mtp.pp("norm"), h, cfg.rms_norm_eps)
+                .context("load mtp.norm")?,
+        })
+    }
+
+    /// Whether this checkpoint carries a head at all.
+    ///
+    /// `mtp_num_hidden_layers` is the config's own answer and is 1 on
+    /// the checkpoints that ship one. Absent or 0 means the weights
+    /// have no `mtp.*` tensors and [`Self::load`] would fail.
+    pub fn present_in(cfg: &TextConfig) -> bool {
+        cfg.mtp_num_hidden_layers > 0
+    }
+
+    /// One draft step.
+    ///
+    /// `token_embed` is `embed_tokens(t)` for the token the draft is
+    /// predicting *from*, `[B, L, H]`. `hidden` is the target's
+    /// post-final-norm hidden state at the same positions, `[B, L, H]`
+    /// — on a chained second step, this head's own previous output.
+    /// The returned tensor is `[B, L, H]`, for the *target's* lm_head.
+    pub fn forward(
+        &mut self,
+        token_embed: &Tensor,
+        hidden: &Tensor,
+        attn_mask: Option<&Tensor>,
+        cos: &Tensor,
+        sin: &Tensor,
+    ) -> Result<Tensor> {
+        anyhow::ensure!(
+            token_embed.dims() == hidden.dims(),
+            "mtp: embedding {:?} and hidden {:?} must have the same shape — \
+             the head fuses them position-for-position",
+            token_embed.dims(),
+            hidden.dims()
+        );
+        let e = self.pre_fc_norm_embedding.forward(token_embed)?;
+        let h = self.pre_fc_norm_hidden.forward(hidden)?;
+        // [embedding, hidden] — vLLM's order. See the module note.
+        let x = Tensor::cat(&[&e, &h], candle_core::D::Minus1)?;
+        let x = self.fc.forward(&x)?;
+        let x = self
+            .layer
+            .forward(&x, attn_mask, cos, sin)
+            .context("mtp draft layer")?;
+        Ok(self.norm.forward(&x)?)
+    }
+
+    /// Capture the draft layer's KV cache.
+    ///
+    /// Drafting advances the head over tokens that may be rejected, so
+    /// anything that drafts speculatively has to be able to put the
+    /// cache back. One full-attention layer is ~2 KiB/token/card, so
+    /// this is cheap next to the target's — which is why the observer
+    /// (#96 S4) can snapshot around every draft without distorting what
+    /// it measures.
+    pub fn snapshot_kv(&self) -> candle_core::Result<super::super::snapshot::LayerKvSnapshot> {
+        self.layer.snapshot_kv()
+    }
+
+    /// Put the draft layer's KV cache back, undoing a speculative walk.
+    pub fn restore_kv(
+        &mut self,
+        snap: &super::super::snapshot::LayerKvSnapshot,
+    ) -> candle_core::Result<()> {
+        self.layer.restore_kv(snap)
+    }
+
+    /// Drop the draft layer's KV cache. Independent of the target's.
+    pub fn clear_kv_cache(&mut self) {
+        self.layer.clear_kv_cache();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+    use std::collections::HashMap;
+
+    /// Build a `ShardedVarBuilder` over a synthetic `mtp.*` head, the
+    /// way the sibling tests in this arch do: write a safetensors file
+    /// and open it.
+    fn tiny_head(
+        dir: &std::path::Path,
+        h: usize,
+        inter: usize,
+        heads: usize,
+        kv: usize,
+        head_dim: usize,
+    ) -> ShardedVarBuilder<'static> {
+        let dev = Device::Cpu;
+        let mut t: HashMap<String, Tensor> = HashMap::new();
+        let zeros = |shape: &[usize]| Tensor::zeros(shape, DType::F32, &dev).unwrap();
+        let randn = |shape: &[usize]| Tensor::randn(0f32, 0.1f32, shape, &dev).unwrap();
+
+        t.insert("mtp.fc.weight".into(), randn(&[h, 2 * h]));
+        t.insert("mtp.norm.weight".into(), zeros(&[h]));
+        t.insert("mtp.pre_fc_norm_embedding.weight".into(), zeros(&[h]));
+        t.insert("mtp.pre_fc_norm_hidden.weight".into(), zeros(&[h]));
+        let l = "mtp.layers.0";
+        t.insert(format!("{l}.input_layernorm.weight"), zeros(&[h]));
+        t.insert(format!("{l}.post_attention_layernorm.weight"), zeros(&[h]));
+        // Output-gated q_proj: 2 * heads * head_dim rows (attn_output_gate).
+        t.insert(
+            format!("{l}.self_attn.q_proj.weight"),
+            randn(&[2 * heads * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.k_proj.weight"),
+            randn(&[kv * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.v_proj.weight"),
+            randn(&[kv * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.o_proj.weight"),
+            randn(&[h, heads * head_dim]),
+        );
+        t.insert(format!("{l}.self_attn.q_norm.weight"), zeros(&[head_dim]));
+        t.insert(format!("{l}.self_attn.k_norm.weight"), zeros(&[head_dim]));
+        t.insert(format!("{l}.mlp.gate_proj.weight"), randn(&[inter, h]));
+        t.insert(format!("{l}.mlp.up_proj.weight"), randn(&[inter, h]));
+        t.insert(format!("{l}.mlp.down_proj.weight"), randn(&[h, inter]));
+
+        let path = dir.join("mtp.safetensors");
+        candle_core::safetensors::save(&t, &path).expect("write tiny mtp checkpoint");
+        // SAFETY: mmaps a file this test just wrote and owns.
+        unsafe {
+            candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                std::slice::from_ref(&path),
+                DType::F32,
+                &dev,
+            )
+        }
+        .expect("build ShardedVarBuilder")
+    }
+
+    fn tiny_cfg(h: usize, inter: usize, heads: usize, kv: usize, head_dim: usize) -> TextConfig {
+        let raw = format!(
+            r#"{{
+                "model_type": "qwen3_next",
+                "vocab_size": 32, "hidden_size": {h}, "intermediate_size": {inter},
+                "num_hidden_layers": 1, "num_attention_heads": {heads},
+                "num_key_value_heads": {kv}, "head_dim": {head_dim},
+                "max_position_embeddings": 64, "rms_norm_eps": 1e-6,
+                "attn_output_gate": true,
+                "layer_types": ["full_attention"],
+                "mtp_num_hidden_layers": 1,
+                "mtp_use_dedicated_embeddings": false
+            }}"#
+        );
+        super::super::Config::from_config_json(&raw)
+            .expect("parse tiny qwen3_5 config")
+            .text_config
+    }
+
+    /// The head loads from the `mtp.*` names the real checkpoints use
+    /// and produces finite output of the target's hidden width — which
+    /// is what the target's `lm_head` will consume.
+    #[test]
+    fn a_tiny_mtp_head_loads_and_drafts() {
+        let (h, inter, heads, kv, head_dim) = (8usize, 16usize, 2usize, 1usize, 4usize);
+        let dir = tempfile::tempdir().expect("tempdir");
+        let vb = tiny_head(dir.path(), h, inter, heads, kv, head_dim);
+        let cfg = tiny_cfg(h, inter, heads, kv, head_dim);
+        assert!(MtpHead::present_in(&cfg));
+
+        let dev = Device::Cpu;
+        let rotary = Arc::new(RotaryEmbedding::new(DType::F32, &cfg, &dev).expect("rotary"));
+        let mut head = MtpHead::load(&cfg, rotary.clone(), &vb).expect("load tiny mtp head");
+
+        let (b, l) = (1usize, 3usize);
+        let embed = Tensor::randn(0f32, 1.0, (b, l, h), &dev).unwrap();
+        let hidden = Tensor::randn(0f32, 1.0, (b, l, h), &dev).unwrap();
+        let positions: Vec<usize> = (0..l).collect();
+        let (cos, sin) = rotary.cos_sin_at(&positions).expect("cos/sin");
+        let out = head
+            .forward(&embed, &hidden, None, &cos, &sin)
+            .expect("draft step");
+        assert_eq!(out.dims(), &[b, l, h]);
+        let v: Vec<f32> = out.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(
+            v.iter().all(|x| x.is_finite()),
+            "draft hidden must be finite"
+        );
+    }
+
+    /// `fc` consumes `cat[embedding, hidden]`, in that order.
+    ///
+    /// Both halves are `hidden_size` wide, so the reversed reading
+    /// loads, runs, and returns finite numbers that are simply a
+    /// different model. vLLM settles it —
+    /// `torch.cat([inputs_embeds, hidden_states], dim=-1)` — and this
+    /// pins it by value so a future edit cannot quietly swap them.
+    ///
+    /// The probe: `fc` reads only the embedding half (identity on the
+    /// first H columns, zero on the second). With a zero embedding and
+    /// a non-zero hidden state, the correct order yields zero and the
+    /// reversed order does not.
+    #[test]
+    fn the_fc_input_is_embedding_then_hidden() {
+        let dev = Device::Cpu;
+        let h = 4usize;
+        // fc = [I | 0] : takes the first H columns only.
+        let mut w = vec![0f32; h * 2 * h];
+        for i in 0..h {
+            w[i * 2 * h + i] = 1.0;
+        }
+        let fc = Linear::new(Tensor::from_vec(w, (h, 2 * h), &dev).unwrap(), None);
+        // Zero weight => (1 + w) = 1, a pure RMS normalisation, so the
+        // norms cannot mask the ordering.
+        let norm =
+            || Qwen3_5RmsNorm::from_weight(Tensor::zeros(h, DType::F32, &dev).unwrap(), 1e-6, None);
+
+        let e = Tensor::zeros((1, 1, h), DType::F32, &dev).unwrap();
+        let hid = Tensor::from_vec(vec![1f32, 2., 3., 4.], (1, 1, h), &dev).unwrap();
+
+        let en = norm().forward(&e).unwrap();
+        let hn = norm().forward(&hid).unwrap();
+
+        let correct: Vec<f32> = fc
+            .forward(&Tensor::cat(&[&en, &hn], candle_core::D::Minus1).unwrap())
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        let reversed: Vec<f32> = fc
+            .forward(&Tensor::cat(&[&hn, &en], candle_core::D::Minus1).unwrap())
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        assert!(
+            correct.iter().all(|x| x.abs() < 1e-6),
+            "with [embedding, hidden] the zero embedding must reach fc: {correct:?}"
+        );
+        assert!(
+            reversed.iter().any(|x| x.abs() > 1e-3),
+            "the reversed reading is what this test exists to catch: {reversed:?}"
+        );
+    }
+
+    /// `forward_multi` must agree with `forward` at the position they
+    /// share.
+    ///
+    /// The speculative verify pass (#96) reads the target's own token
+    /// at every drafted position, so the whole scheme rests on those
+    /// logits being the ones the model would have produced anyway.
+    ///
+    /// Agreement is to a tolerance, not bit-for-bit, and the reason is
+    /// worth recording: `forward` slices the last position *before* the
+    /// LM head and multiplies `(1, 1, H) x (H, vocab)`, while
+    /// `forward_multi` multiplies `(1, L, H) x (H, vocab)` and slices
+    /// after. Different GEMM shapes accumulate in a different order, so
+    /// the results differ in the last ULP or two — on this fixture,
+    /// 0.0016435327 against 0.0016435329. An exact assertion passed on
+    /// one CPU's BLAS blocking and failed on another's, which is the
+    /// only thing it was really testing.
+    ///
+    /// What matters for acceptance is that the argmax agrees, so that
+    /// is asserted too: a verify pass compares tokens, not floats.
+    #[test]
+    fn forward_multi_agrees_with_forward_at_the_last_position() {
+        use super::super::{Config, Qwen3_5ForCausalLM};
+        use candle_core::IndexOp;
+
+        let (h, inter, heads, kv, head_dim, vocab) =
+            (8usize, 16usize, 2usize, 1usize, 4usize, 32usize);
+        let dev = Device::Cpu;
+        let dir = tempfile::tempdir().expect("tempdir");
+        let mut t: HashMap<String, Tensor> = HashMap::new();
+        let zeros = |shape: &[usize]| Tensor::zeros(shape, DType::F32, &dev).unwrap();
+        let randn = |shape: &[usize]| Tensor::randn(0f32, 0.1f32, shape, &dev).unwrap();
+        t.insert("model.embed_tokens.weight".into(), randn(&[vocab, h]));
+        t.insert("lm_head.weight".into(), randn(&[vocab, h]));
+        t.insert("model.norm.weight".into(), zeros(&[h]));
+        let l = "model.layers.0";
+        t.insert(format!("{l}.input_layernorm.weight"), zeros(&[h]));
+        t.insert(format!("{l}.post_attention_layernorm.weight"), zeros(&[h]));
+        t.insert(
+            format!("{l}.self_attn.q_proj.weight"),
+            randn(&[2 * heads * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.k_proj.weight"),
+            randn(&[kv * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.v_proj.weight"),
+            randn(&[kv * head_dim, h]),
+        );
+        t.insert(
+            format!("{l}.self_attn.o_proj.weight"),
+            randn(&[h, heads * head_dim]),
+        );
+        t.insert(format!("{l}.self_attn.q_norm.weight"), zeros(&[head_dim]));
+        t.insert(format!("{l}.self_attn.k_norm.weight"), zeros(&[head_dim]));
+        t.insert(format!("{l}.mlp.gate_proj.weight"), randn(&[inter, h]));
+        t.insert(format!("{l}.mlp.up_proj.weight"), randn(&[inter, h]));
+        t.insert(format!("{l}.mlp.down_proj.weight"), randn(&[h, inter]));
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&t, &path).expect("write tiny checkpoint");
+        // SAFETY: mmaps a file this test just wrote and owns.
+        let vb = unsafe {
+            candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                std::slice::from_ref(&path),
+                DType::F32,
+                &dev,
+            )
+        }
+        .expect("build ShardedVarBuilder");
+
+        let raw = format!(
+            r#"{{
+                "model_type": "qwen3_next",
+                "vocab_size": {vocab}, "hidden_size": {h}, "intermediate_size": {inter},
+                "num_hidden_layers": 1, "num_attention_heads": {heads},
+                "num_key_value_heads": {kv}, "head_dim": {head_dim},
+                "max_position_embeddings": 64, "rms_norm_eps": 1e-6,
+                "attn_output_gate": true,
+                "layer_types": ["full_attention"]
+            }}"#
+        );
+        let cfg = Config::from_config_json(&raw).expect("parse tiny config");
+        let mut model = Qwen3_5ForCausalLM::new(cfg, vb).expect("load tiny checkpoint");
+
+        let input = Tensor::new(&[1u32, 5, 9, 2], &dev)
+            .unwrap()
+            .unsqueeze(0)
+            .unwrap();
+        let multi = model.forward_multi(&input, 0).expect("forward_multi");
+        assert_eq!(multi.dims(), &[1, 4, vocab], "one logits row per position");
+        let last_of_multi: Vec<f32> = multi.i((0, 3, ..)).unwrap().to_vec1().unwrap();
+        // A head that returned the same row L times would satisfy the
+        // comparison below without carrying per-position information,
+        // which is the whole point of the pass.
+        let first_of_multi: Vec<f32> = multi.i((0, 0, ..)).unwrap().to_vec1().unwrap();
+        assert_ne!(
+            first_of_multi, last_of_multi,
+            "positions must carry distinct logits, not a repeated row"
+        );
+
+        model.clear_kv_cache();
+        let single = model.forward(&input, 0).expect("forward");
+        assert_eq!(single.dims(), &[1, 1, vocab]);
+        let single: Vec<f32> = single.flatten_all().unwrap().to_vec1().unwrap();
+
+        let max_diff = last_of_multi
+            .iter()
+            .zip(single.iter())
+            .map(|(a, b)| (a - b).abs())
+            .fold(0f32, f32::max);
+        assert!(
+            max_diff < 1e-5,
+            "forward_multi's last row must match what forward returns \
+             (max abs diff {max_diff:e}): {last_of_multi:?} vs {single:?}"
+        );
+
+        let argmax = |v: &[f32]| {
+            v.iter()
+                .enumerate()
+                .max_by(|a, b| a.1.partial_cmp(b.1).expect("finite logits"))
+                .map(|(i, _)| i)
+                .expect("non-empty")
+        };
+        assert_eq!(
+            argmax(&last_of_multi),
+            argmax(&single),
+            "the verify pass compares tokens; the argmax must agree"
+        );
+    }
+}

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,9 +72,8 @@ impl DecoderLayer {
         cfg: &TextConfig,
         rotary: Arc<RotaryEmbedding>,
         layer_idx: usize,
-        quant: Option<GgmlDType>,
         vb: &ShardedVarBuilder,
-        isq_cache: Option<&crate::harness::isq_cache::IsqCache>,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let layer_type = cfg
             .layer_types
@@ -93,9 +90,8 @@ impl DecoderLayer {
             rotary,
             layer_type,
             cfg.ple_layers().contains(&layer_idx),
-            quant,
             vb,
-            isq_cache,
+            opts,
         )
     }
 
@@ -110,9 +106,8 @@ impl DecoderLayer {
         rotary: Arc<RotaryEmbedding>,
         layer_type: &str,
         with_ple: bool,
-        quant: Option<GgmlDType>,
         vb: &ShardedVarBuilder,
-        isq_cache: Option<&crate::harness::isq_cache::IsqCache>,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let mixer = match layer_type {
             "full_attention" => {
@@ -158,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"), isq_cache)?,
+            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};
@@ -46,6 +45,7 @@ use super::config::{Config, TextConfig};
 use super::decoder::DecoderLayer;
 use super::hyper::HyperConnection;
 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
@@ -137,11 +137,9 @@ impl Qwen4ExpForCausalLM {
     pub fn load(
         cfg: &Config,
         dtype: DType,
-        quant: Option<GgmlDType>,
         device: &Device,
         vb: &ShardedVarBuilder,
-        safetensors_paths: &[std::path::PathBuf],
-        isq_cache: Option<&crate::harness::isq_cache::IsqCache>,
+        opts: &super::LoadOptions<'_>,
     ) -> Result<Self> {
         let text = &cfg.text_config;
         // `RotaryEmbedding` reads a qwen3_5 config; the rope block is the
@@ -170,7 +168,7 @@ impl Qwen4ExpForCausalLM {
             Some(layer) => Some(load_ngram(
                 text,
                 &layers_vb.pp(*layer).pp("ple"),
-                safetensors_paths,
+                opts.safetensors_paths,
                 device,
             )?),
             None => None,
@@ -179,7 +177,7 @@ impl Qwen4ExpForCausalLM {
         let mut layers = Vec::with_capacity(text.num_hidden_layers);
         for i in 0..text.num_hidden_layers {
             layers.push(
-                DecoderLayer::load(text, rotary.clone(), i, quant, &layers_vb.pp(i), isq_cache)
+                DecoderLayer::load(text, rotary.clone(), i, &layers_vb.pp(i), opts)
                     .with_context(|| format!("load decoder layer {i}"))?,
             );
         }
@@ -216,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)
     }
@@ -239,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))?;
 
@@ -248,6 +266,7 @@ impl Qwen4ExpForCausalLM {
         } else {
             None
         };
+        probe.mark("embed");
 
         // Computed once: it is a function of the ids, and one layer
         // consumes it.
@@ -258,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:
@@ -805,11 +841,15 @@ mod tests {
         let model = Qwen4ExpForCausalLM::load(
             &cfg,
             DType::F32,
-            None,
             &Device::Cpu,
             &vb,
-            std::slice::from_ref(&path),
-            None,
+            &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)
@@ -824,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,10 +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,
-    isq_cache: Option<&crate::harness::isq_cache::IsqCache>,
+    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,
@@ -78,7 +77,20 @@ pub fn load(
         None,
     );
 
-    let experts_vb = vb.pp("experts");
+    // 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.
@@ -87,20 +99,19 @@ pub fn load(
     // 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, vb.device()).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,
+    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
                 }
-            }),
-            _ => None,
-        };
+            }
+        }),
+        _ => None,
+    };
 
     let experts = match cached {
         Some(experts) => experts,
@@ -118,8 +129,9 @@ pub fn load(
             // existing.
             match quant {
                 Some(dtype) => {
-                    let experts = Experts::quantize_banked(&gate_up, &down, inter, dtype)
-                        .with_context(|| format!("quantize experts to {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()

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

@@ -113,12 +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.
-                None,
+                // 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

@@ -420,6 +420,26 @@ impl MmapNGramTable {
             // defend against.
             let map = unsafe { memmap2::Mmap::map(&file) }
                 .with_context(|| format!("mmap {}", path.display()))?;
+            // Every access through this mapping is a 320-byte row picked
+            // by a hash: sixteen unrelated rows per token, out of
+            // 320,001,536. The kernel's default readahead assumes
+            // sequential access and faults a whole window per touch,
+            // which is pure amplification here -- and this table is
+            // 95 GiB against a host that is also holding the routed
+            // experts (#318) in RAM, so every wasted page evicts one
+            // that a later gather wanted.
+            //
+            // Advisory: a failure changes performance, never
+            // correctness, so it is logged rather than propagated. A
+            // kernel or filesystem that ignores MADV_RANDOM simply
+            // leaves the previous behaviour in place.
+            if let Err(e) = map.advise(memmap2::Advice::Random) {
+                tracing::debug!(
+                    file = %path.display(),
+                    error = %e,
+                    "ple: MADV_RANDOM refused; readahead stays on for this mapping"
+                );
+            }
             let (header_len, meta) =
                 safetensors::SafeTensors::read_metadata(&map).map_err(|e| {
                     anyhow::anyhow!("read safetensors metadata from {}: {e}", path.display())
@@ -685,10 +705,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()
@@ -1052,6 +1097,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

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

@@ -56,6 +56,52 @@ use crate::harness::arch::qwen3_5::rope::{RotaryEmbedding, apply_partial_rotary}
 use super::config::TextConfig as Qwen4ExpTextConfig;
 
 /// The indexer's geometry, from `config.json`.
+/// The positions one query may attend, ascending.
+///
+/// The reference derives this from the attention mask
+/// (`torch.nonzero(visible_token_indices[b, 0, q])`) and forms blocks
+/// from whichever positions survive, in order. Under a plain causal
+/// mask that is exactly `0..n`, which is worth representing without
+/// materialising: the set reaches the context window in size and is
+/// rebuilt for every query, so a `Vec` per query would cost more than
+/// the selection it feeds.
+///
+/// [`Sparse`] is the general case — left padding, ragged batches, any
+/// mask with holes. Both are handled by the same arithmetic below,
+/// which is the point: block `b` covers *set positions*
+/// `b * block_size ..`, never *absolute positions*, and those coincide
+/// only when the set is contiguous from zero.
+///
+/// [`Sparse`]: VisibleSet::Sparse
+#[derive(Debug, Clone, Copy)]
+pub enum VisibleSet<'a> {
+    /// Every position `0..n` — a plain causal mask.
+    Contiguous(usize),
+    /// Exactly these positions, ascending and without duplicates.
+    Sparse(&'a [usize]),
+}
+
+impl VisibleSet<'_> {
+    pub fn len(&self) -> usize {
+        match self {
+            Self::Contiguous(n) => *n,
+            Self::Sparse(v) => v.len(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// The `i`th visible position.
+    fn at(&self, i: usize) -> usize {
+        match self {
+            Self::Contiguous(_) => i,
+            Self::Sparse(v) => v[i],
+        }
+    }
+}
+
 pub struct BlockSelector {
     n_heads: usize,
     head_dim: usize,
@@ -96,9 +142,14 @@ impl BlockSelector {
 
     /// The position each block starts at — where its pooled key is
     /// RoPE'd, **not** the position it is scored against.
-    pub fn block_starts(&self, n_visible: usize) -> Vec<usize> {
-        (0..self.n_blocks(n_visible))
-            .map(|b| b * self.block_size)
+    ///
+    /// Block `b` begins at the `b * block_size`-th **visible** position,
+    /// which is `b * block_size` only when the set is contiguous from
+    /// zero. Rotating a pooled key at the wrong absolute position is
+    /// silent: the scores stay finite and the selection stays plausible.
+    pub fn block_starts(&self, visible: &VisibleSet) -> Vec<usize> {
+        (0..self.n_blocks(visible.len()))
+            .map(|b| visible.at(b * self.block_size))
             .collect()
     }
 
@@ -149,7 +200,8 @@ impl BlockSelector {
     ///
     /// Ties are broken towards the earlier block, which matters only for
     /// reproducibility — the reference's `topk` does the same.
-    pub fn select(&self, scores: &[f32], n_visible: usize) -> Result<Vec<usize>> {
+    pub fn select(&self, scores: &[f32], visible: &VisibleSet) -> Result<Vec<usize>> {
+        let n_visible = visible.len();
         let blocks = self.n_blocks(n_visible);
         ensure!(
             scores.len() == blocks,
@@ -170,15 +222,28 @@ impl BlockSelector {
             order.sort_unstable();
         }
 
+        // Indices into the visible set, not absolute positions. They
+        // coincide only when the set is contiguous from zero, and
+        // conflating them is what this issue (#324) is about.
         let tail_start = blocks * self.block_size;
         let mut positions = Vec::with_capacity(keep * self.block_size + (n_visible - tail_start));
         for b in order {
             let start = b * self.block_size;
-            positions.extend(start..start + self.block_size);
+            positions.extend((start..start + self.block_size).map(|i| visible.at(i)));
         }
         // The tail never competes for the budget. It holds the query's
         // own position.
-        positions.extend(tail_start..n_visible);
+        positions.extend((tail_start..n_visible).map(|i| visible.at(i)));
+        // Ascending is the contract `dense_mask` relies on, and it
+        // holds without sorting: `order` is ascending, `visible.at` is
+        // monotonic in its index, and the tail's indices all follow the
+        // blocks'. Asserted rather than enforced, so a future change
+        // that breaks the invariant is caught instead of hidden by a
+        // sort that makes it look fine.
+        debug_assert!(
+            positions.windows(2).all(|w| w[0] < w[1]),
+            "selected positions must be ascending and distinct: {positions:?}"
+        );
         Ok(positions)
     }
 
@@ -356,7 +421,10 @@ impl Indexer {
         total: usize,
     ) -> candle_core::Result<Tensor> {
         let b = cache.dims()[0];
-        let starts = self.selector.block_starts(total);
+        // The cache holds every past position, so the poolable set is
+        // the contiguous 0..total. When a padded or ragged batch can
+        // reach here this becomes `Sparse` of that row's live positions.
+        let starts = self.selector.block_starts(&VisibleSet::Contiguous(total));
         let mut pooled = Vec::with_capacity(b);
         for row in 0..b {
             pooled.push(self.selector.pool(&cache.i(row)?)?);
@@ -475,10 +543,20 @@ impl Indexer {
             None => k,
         };
         let total = cache.dims()[1];
+        // This equality is also what makes the visible set contiguous:
+        // every query's set is built as `Contiguous(past_len + i + 1)`,
+        // which describes the batch only while one scalar `past_len`
+        // does. A ragged or left-padded batch presents a padded `total`
+        // that no single `past_len` matches, and trips this — which is
+        // the intent, because the alternative is a selection computed
+        // against the wrong past that still looks plausible (#324).
         ensure!(
             total == past_len + seq_len,
             "indexer cache holds {total} positions but the queries start at {past_len} \
-             and run for {seq_len} — the cache and the main KV have diverged"
+             and run for {seq_len} — the cache and the main KV have diverged, or this \
+             is a ragged/padded batch whose visible set is no longer contiguous. The \
+             latter needs the set derived from the attention mask (VisibleSet::Sparse) \
+             rather than from a count (#324)."
         );
         self.key_cache = Some(cache.clone());
         Ok((q, cache))
@@ -540,12 +618,16 @@ impl Indexer {
         for row in scores.iter().take(batch) {
             let mut per_query = Vec::with_capacity(seq_len);
             for (i, all_blocks) in row.iter().enumerate().take(seq_len) {
-                let visible = past_len + i + 1;
+                // Contiguous because `model.rs` builds the causal mask
+                // itself and no padding mask reaches this arch (#324).
+                // The moment one can, this is where the row's live
+                // positions are derived from it.
+                let visible = VisibleSet::Contiguous(past_len + i + 1);
                 // Only blocks wholly behind this query exist for it.
-                let visible_blocks = self.selector.n_blocks(visible);
+                let visible_blocks = self.selector.n_blocks(visible.len());
                 per_query.push(
                     self.selector
-                        .select(&all_blocks[..visible_blocks], visible)?,
+                        .select(&all_blocks[..visible_blocks], &visible)?,
                 );
             }
             out.push(per_query);
@@ -600,6 +682,63 @@ mod tests {
         assert_eq!(qsa.block_topk() * 4, 2048);
     }
 
+    /// Blocks are formed from consecutive **visible** positions, not
+    /// consecutive absolute ones (#324).
+    ///
+    /// The reference takes `nonzero` of the mask and chunks whatever
+    /// survives, in order. With holes in the visible set the two
+    /// readings diverge completely: under the old arithmetic the query
+    /// would be handed positions 4..8, which here are not merely the
+    /// wrong block — 4, 6, 8, 9 and 10 are positions the mask forbids
+    /// it from seeing at all.
+    #[test]
+    fn blocks_are_cut_from_the_visible_set_not_the_position_axis() {
+        // block_size 4, budget 4 => one block survives the top-k.
+        let qsa = BlockSelector::new(2, 1, 4, 4).unwrap();
+        let live = [3usize, 5, 7, 11, 13, 17, 19, 23, 29];
+        let visible = VisibleSet::Sparse(&live);
+
+        assert_eq!(qsa.n_blocks(visible.len()), 2);
+        // Block 0 begins at visible[0] = 3, block 1 at visible[4] = 13.
+        assert_eq!(qsa.block_starts(&visible), vec![3, 13]);
+
+        // Score block 1 above block 0, so only block 1 is kept.
+        let got = qsa.select(&[0.1f32, 0.9], &visible).unwrap();
+        assert_eq!(
+            got,
+            vec![13, 17, 19, 23, 29],
+            "block 1 is visible[4..8], and the tail is visible[8..]"
+        );
+        // The failure this pins: index arithmetic would give 4..8 plus 8.
+        assert!(
+            !got.contains(&4) && !got.contains(&8),
+            "positions the mask forbids must never be selected: {got:?}"
+        );
+    }
+
+    /// The two representations must agree wherever both are valid.
+    ///
+    /// `Contiguous(n)` is an optimisation for the causal case, not a
+    /// different rule — if it can disagree with the explicit list of the
+    /// same positions, one of the two paths is wrong and only the
+    /// explicit one is checked against the reference.
+    #[test]
+    fn contiguous_and_sparse_agree_on_the_same_positions() {
+        let qsa = BlockSelector::new(2, 1, 4, 8).unwrap();
+        for n in [1usize, 4, 7, 8, 9, 15, 23] {
+            let explicit: Vec<usize> = (0..n).collect();
+            let scores = vec![0.5f32; qsa.n_blocks(n)];
+            let a = qsa.select(&scores, &VisibleSet::Contiguous(n)).unwrap();
+            let b = qsa.select(&scores, &VisibleSet::Sparse(&explicit)).unwrap();
+            assert_eq!(a, b, "representations disagree at n={n}");
+            assert_eq!(
+                qsa.block_starts(&VisibleSet::Contiguous(n)),
+                qsa.block_starts(&VisibleSet::Sparse(&explicit)),
+                "block starts disagree at n={n}"
+            );
+        }
+    }
+
     /// The cheapest gate in the architecture: below the budget, the
     /// selection returns every visible position, so a short prompt must
     /// be bit-comparable to dense causal attention.
@@ -609,7 +748,9 @@ mod tests {
         for n_visible in [1usize, 3, 4, 7, 2047, 2048] {
             assert!(qsa.is_dense(n_visible), "{n_visible} should be dense");
             let scores = vec![0.0f32; qsa.n_blocks(n_visible)];
-            let got = qsa.select(&scores, n_visible).unwrap();
+            let got = qsa
+                .select(&scores, &VisibleSet::Contiguous(n_visible))
+                .unwrap();
             assert_eq!(
                 got,
                 (0..n_visible).collect::<Vec<_>>(),
@@ -631,7 +772,9 @@ mod tests {
         assert_eq!(qsa.n_blocks(n_visible), 4);
         // Block 2 is worthless; with topk 3 it is the one dropped.
         let scores = [0.9f32, 0.8, 0.1, 0.7];
-        let got = qsa.select(&scores, n_visible).unwrap();
+        let got = qsa
+            .select(&scores, &VisibleSet::Contiguous(n_visible))
+            .unwrap();
         let want: Vec<usize> = (0..8).chain(12..19).collect();
         assert_eq!(got, want);
         assert!(got.contains(&(n_visible - 1)), "the query's own position");
@@ -721,7 +864,7 @@ mod tests {
         );
         // And the selection that rests on it picks the right block.
         assert_eq!(
-            qsa.select(&scores, 8).unwrap(),
+            qsa.select(&scores, &VisibleSet::Contiguous(8)).unwrap(),
             vec![4, 5, 6, 7],
             "the higher-scoring block should be the one kept"
         );
@@ -755,24 +898,26 @@ mod tests {
         let qsa = small();
         let keys = Tensor::zeros((11, 8), DType::BF16, &dev).unwrap();
         assert_eq!(qsa.pool(&keys).unwrap().dims(), &[2, 8]);
-        assert_eq!(qsa.block_starts(11), vec![0, 4]);
+        assert_eq!(qsa.block_starts(&VisibleSet::Contiguous(11)), vec![0, 4]);
 
         let short = Tensor::zeros((3, 8), DType::F32, &dev).unwrap();
         assert_eq!(qsa.pool(&short).unwrap().dims(), &[0, 8]);
-        assert!(qsa.block_starts(3).is_empty());
+        assert!(qsa.block_starts(&VisibleSet::Contiguous(3)).is_empty());
     }
 
     #[test]
     fn selection_is_ascending_and_rejects_a_wrong_score_count() {
         let qsa = small();
-        let got = qsa.select(&[0.1f32, 0.9, 0.5, 0.7, 0.2], 23).unwrap();
+        let got = qsa
+            .select(&[0.1f32, 0.9, 0.5, 0.7, 0.2], &VisibleSet::Contiguous(23))
+            .unwrap();
         assert!(got.windows(2).all(|w| w[0] < w[1]), "got {got:?}");
         // blocks 1, 3, 2 win; block 0 and 4 lose; tail 20..23 stays.
         assert_eq!(
             got,
             vec![4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 21, 22]
         );
-        assert!(qsa.select(&[0.1f32], 23).is_err());
+        assert!(qsa.select(&[0.1f32], &VisibleSet::Contiguous(23)).is_err());
     }
 
     // ---- the indexer as a layer ----
@@ -832,6 +977,61 @@ mod tests {
         .unwrap()
     }
 
+    /// A cache that does not match `past_len + seq_len` is a ragged or
+    /// left-padded batch, and must fail loudly (#324).
+    ///
+    /// Every visible set is built as `Contiguous(past_len + i + 1)`, so
+    /// the moment one scalar `past_len` stops describing every row, the
+    /// selection is computed against the wrong past — and it stays
+    /// plausible, which is what makes it dangerous. The guard is only
+    /// worth having if it can fire, so this drives it directly rather
+    /// than trusting that it would.
+    #[test]
+    fn a_cache_inconsistent_with_past_len_is_refused() {
+        let (hidden_size, n_heads, head_dim) = (8, 2, 8);
+        let mut ix = indexer(hidden_size, n_heads, head_dim, 4, 4);
+        let rope = rope_for(head_dim);
+
+        // A history of 20 positions, while the step claims past_len 5:
+        // the shape a padded batch presents, where `total` is the
+        // padded maximum and no single `past_len` matches every row.
+        let stale = Tensor::zeros((1, 20, head_dim), DType::F32, &Device::Cpu).unwrap();
+        ix.restore_keys(Some(&stale));
+
+        let err = ix
+            .visible_positions(&hidden(1, 1, hidden_size, 0.7), &rope, 5)
+            .expect_err("a cache that disagrees with past_len must not be served");
+        let msg = format!("{err:#}");
+        assert!(
+            msg.contains("no longer contiguous") && msg.contains("324"),
+            "the error must name the contiguity assumption it protects, got: {msg}"
+        );
+    }
+
+    /// The same call with a consistent cache must be served, or the
+    /// guard above is just refusing everything.
+    #[test]
+    fn a_cache_consistent_with_past_len_is_served() {
+        let (hidden_size, n_heads, head_dim) = (8, 2, 8);
+        let mut ix = indexer(hidden_size, n_heads, head_dim, 4, 4);
+        let rope = rope_for(head_dim);
+
+        let history = Tensor::zeros((1, 5, head_dim), DType::F32, &Device::Cpu).unwrap();
+        ix.restore_keys(Some(&history));
+
+        // past_len 5 + this step's 1 == the 6 the cache will hold.
+        let got = ix
+            .visible_positions(&hidden(1, 1, hidden_size, 0.7), &rope, 5)
+            .expect("a contiguous history must still be served");
+        assert_eq!(got.len(), 1, "one batch row");
+        assert_eq!(got[0].len(), 1, "one query");
+        assert!(
+            got[0][0].iter().all(|p| *p < 6),
+            "positions must lie inside the 6-position history: {:?}",
+            got[0][0]
+        );
+    }
+
     /// `index_qk_proj` is fused 4:1, not halved. With 4 query heads and
     /// 1 key head of 8 dims the split is 32 and 8 of 40 — a halving
     /// would take 20 and 20 and put half a query head in the key.

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

@@ -391,6 +391,8 @@ mod tests {
     fn tiny_config() -> TextConfig {
         TextConfig {
             vocab_size: 32,
+            mtp_num_hidden_layers: 0,
+            mtp_use_dedicated_embeddings: false,
             hidden_size: 16,
             intermediate_size: 32,
             num_hidden_layers: 2,

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,
@@ -138,6 +143,9 @@ pub enum LoadedHandle {
     Tp(Arc<TpLoadedModel>),
     /// Z-Image text-to-image pipeline (#198).
     Image(Arc<LoadedImageModel>),
+    /// Laya decision model (#336): typed questions in, calibrated
+    /// option scores out, no text generation.
+    Decision(Arc<super::decision_model::LoadedDecisionModel>),
 }
 
 impl LoadedHandle {
@@ -147,6 +155,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.model_id,
             LoadedHandle::Image(m) => &m.model_id,
+            LoadedHandle::Decision(m) => &m.model_id,
         }
     }
 
@@ -157,6 +166,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.spec,
             LoadedHandle::Image(m) => &m.spec,
+            LoadedHandle::Decision(m) => &m.spec,
         }
     }
 
@@ -166,6 +176,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.devices.clone(),
             LoadedHandle::Image(m) => m.devices.clone(),
+            LoadedHandle::Decision(m) => m.devices.clone(),
         }
     }
 
@@ -179,6 +190,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
             LoadedHandle::Image(m) => m.poisoned.load(Ordering::Acquire),
+            LoadedHandle::Decision(m) => m.poisoned.load(Ordering::Acquire),
         }
     }
 
@@ -190,6 +202,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
             LoadedHandle::Image(m) => (m.admission.in_flight(), m.admission.queue_depth()),
+            LoadedHandle::Decision(m) => (m.admission.in_flight(), m.admission.queue_depth()),
         }
     }
 
@@ -202,6 +215,9 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
             LoadedHandle::Image(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
+            LoadedHandle::Decision(m) => {
+                (m.admission.max_in_flight(), m.admission.max_queue_depth())
+            }
         }
     }
 
@@ -225,6 +241,10 @@ impl LoadedHandle {
                 m.admission.anon_in_flight(),
                 m.admission.anon_max_in_flight(),
             ),
+            LoadedHandle::Decision(m) => (
+                m.admission.anon_in_flight(),
+                m.admission.anon_max_in_flight(),
+            ),
         }
     }
 
@@ -235,6 +255,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.close(),
             LoadedHandle::Image(m) => m.admission.close(),
+            LoadedHandle::Decision(m) => m.admission.close(),
         }
     }
 
@@ -246,6 +267,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.rejections(),
             LoadedHandle::Image(m) => m.admission.rejections(),
+            LoadedHandle::Decision(m) => m.admission.rejections(),
         }
     }
 
@@ -257,6 +279,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.admission,
             LoadedHandle::Image(m) => &m.admission,
+            LoadedHandle::Decision(m) => &m.admission,
         };
         (admission.kv_budget_mb(), admission.kv_available_mb())
     }
@@ -270,7 +293,7 @@ impl LoadedHandle {
             LoadedHandle::Tp(m) => &m.prefill_rate,
             // Image models have no token throughput; steps/sec lives in
             // the per-request timing instead.
-            LoadedHandle::Image(_) => return (0.0, 0.0),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return (0.0, 0.0),
         };
         (ema.get().unwrap_or(0.0), ema.decode().unwrap_or(0.0))
     }
@@ -282,12 +305,15 @@ impl LoadedHandle {
         if matches!(self, LoadedHandle::Image(_)) {
             return vec!["image".to_string()];
         }
+        if matches!(self, LoadedHandle::Decision(_)) {
+            return vec!["decision".to_string()];
+        }
         let mut caps = vec!["text".to_string()];
         let has_vision = match self {
             LoadedHandle::Single(m) => m.has_vision,
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.has_vision,
-            LoadedHandle::Image(_) => unreachable!("handled above"),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => unreachable!("handled above"),
         };
         if has_vision {
             caps.push("vision".to_string());
@@ -333,7 +359,7 @@ impl LoadedHandle {
             LoadedHandle::Tp(m) => (m.last_free_mb.load(Ordering::Acquire), m.context_profile),
             // Image generation has its own VRAM path and no prefill
             // floor, so this check would say nothing useful about it.
-            LoadedHandle::Image(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
         };
         // `0` is the CPU-load and missing-worker sentinel — the same one
         // `check_vram` skips on. Nothing to assert either way.
@@ -414,7 +440,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.tool_call_tokens.is_some(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.tool_call_tokens.is_some(),
-            LoadedHandle::Image(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
         }
     }
 
@@ -425,7 +451,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.reasoning_tokens.is_some(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.reasoning_tokens.is_some(),
-            LoadedHandle::Image(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
         }
     }
 
@@ -436,7 +462,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.preserve_thinking,
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.preserve_thinking,
-            LoadedHandle::Image(_) => None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => None,
         }
     }
 
@@ -446,7 +472,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.reasoning_efforts.clone(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.reasoning_efforts.clone(),
-            LoadedHandle::Image(_) => Default::default(),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => Default::default(),
         }
     }
 
@@ -472,7 +498,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
             // No context-limit derivation for image models; nothing to cache.
-            LoadedHandle::Image(_) => return,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return,
         };
         // Don't clobber a good cached value with a transient `0`
         // (worker gone/poisoned sentinel).
@@ -481,7 +507,7 @@ impl LoadedHandle {
                 LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
                 #[cfg(feature = "cuda")]
                 LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
-                LoadedHandle::Image(_) => {}
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {}
             }
         }
     }
@@ -510,7 +536,7 @@ impl LoadedHandle {
                 m.prefill_rate.get(),
             ),
             // Token limits are meaningless for image generation.
-            LoadedHandle::Image(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
         };
         let rate = rate.unwrap_or(cfg.bootstrap_prefill_tok_per_sec);
         // The KV budget admission will actually grant. Advertising past
@@ -520,7 +546,9 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.admission.kv_budget_mb(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.kv_budget_mb(),
-            LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+                unreachable!("derived_limit returned above")
+            }
         };
         let limit = super::context_limit::derive_limit(
             &profile,
@@ -539,7 +567,9 @@ impl LoadedHandle {
                 LoadedHandle::Single(m) => &m.derived_input_cap,
                 #[cfg(feature = "cuda")]
                 LoadedHandle::Tp(m) => &m.derived_input_cap,
-                LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+                    unreachable!("derived_limit returned above")
+                }
             };
             cap.store(input, Ordering::Release);
         }
@@ -1162,6 +1192,18 @@ impl ModelArch {
         }
     }
 
+    /// Logits at every position of a multi-token forward: `(B, L,
+    /// vocab)`. The speculative verify pass (#96) reads the target's
+    /// own token at each drafted position; `squeeze_to_vocab` is
+    /// deliberately not applied, as it would collapse the L axis this
+    /// exists to keep.
+    pub fn forward_multi(&mut self, input: &Tensor, offset: usize) -> Result<Tensor> {
+        match self {
+            ModelArch::Qwen3_5Dense(m) => Ok(m.forward_multi(input, offset)?),
+            _ => anyhow::bail!("forward_multi: architecture has no multi-position logits support"),
+        }
+    }
+
     /// Padding mask for a batched decode step — see
     /// `Qwen3_5Model::batch_decode_mask`.
     pub fn batch_decode_mask(
@@ -1283,6 +1325,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
@@ -2866,11 +2931,18 @@ impl CandleHarness {
                     let model = super::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
                         &cfg,
                         dtype,
-                        isq,
                         &device_for_load,
                         &sharded_vb,
-                        &safetensors_paths,
-                        isq_cache.as_ref(),
+                        &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)))
@@ -2907,6 +2979,27 @@ 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).
@@ -3214,7 +3307,7 @@ impl CandleHarness {
             LoadedHandle::Tp(m) => {
                 return self.inference_tp_stream(m, request, principal, wire).await;
             }
-            LoadedHandle::Image(_) => {
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
                 return Err(InferenceError::WrongModality {
                     model_id: request.model.clone(),
                 });
@@ -3454,11 +3547,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,
@@ -3931,6 +4026,11 @@ impl Harness for CandleHarness {
         if matches!(plan.format, super::preflight::SourceFormat::Diffusers) {
             return self.load_image_model(spec, &source_id).await;
         }
+        // Laya decision models (#336) take the decision path: one
+        // device, unquantized (preflight rejected tp>1 and quant).
+        if matches!(plan.format, super::preflight::SourceFormat::DecisionModel) {
+            return self.load_decision_model(spec, &source_id).await;
+        }
 
         let tp_size = spec.tensor_parallel.unwrap_or(1);
         if tp_size > 1 {
@@ -4014,6 +4114,17 @@ impl Harness for CandleHarness {
                     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,
@@ -4021,6 +4132,7 @@ impl Harness for CandleHarness {
                         spec.model_id.clone(),
                         isq,
                         isq_cache,
+                        experts_on_host,
                     )
                     .await
                     .map_err(|e| anyhow::anyhow!("worker load_dense: {e:#}"))?;
@@ -4114,15 +4226,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(
@@ -4131,7 +4247,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,
@@ -4187,7 +4303,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,
@@ -4281,6 +4397,11 @@ impl Harness for CandleHarness {
                     );
                 }
             }
+            LoadedHandle::Decision(dm) => {
+                // Each checkpoint's weights drop on the thread that owns
+                // them (the device worker, for a CUDA load).
+                dm.release().await;
+            }
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(tp) => {
                 // Try to recover the inner TpLoadedModel so we can move
@@ -4455,6 +4576,101 @@ impl CandleHarness {
         Ok(())
     }
 
+    /// Load a Laya decision model (#336): resolve the checkpoint's files,
+    /// build it on the device worker (or the CPU), and register it with
+    /// `capabilities: ["decision"]`.
+    ///
+    /// Only the repo-root checkpoint is loaded. A repo can carry more in
+    /// subfolders (`multilingual/`, `typed-decisions/`); which of those a
+    /// model serves, and how requests are routed between them, is a
+    /// separate decision (#339) — [`Self::resolve_decision_files`]
+    /// already takes the subfolder.
+    async fn load_decision_model(
+        &self,
+        spec: &ModelSpec,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<()> {
+        let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
+        let device = Self::pick_device(&devices)?;
+        #[cfg(feature = "cuda")]
+        let worker = if matches!(device, Device::Cuda(_)) {
+            Some(self.ensure_device_worker(devices[0]).await?)
+        } else {
+            None
+        };
+        // CPU build: always the CPU backend. Correct, but at hundreds of
+        // milliseconds per question it is a fallback, not a serving tier.
+        #[cfg(not(feature = "cuda"))]
+        let worker = {
+            let _ = &device;
+            None
+        };
+
+        let files = self.resolve_decision_files(source_id, None).await?;
+        let checkpoint = super::decision_model::DecisionCheckpoint::load(
+            super::decision_model::ROOT_CHECKPOINT.to_string(),
+            files,
+            &spec.model_id,
+            worker,
+        )
+        .await?;
+        let dtype = checkpoint.dtype;
+
+        let loaded = super::decision_model::LoadedDecisionModel {
+            model_id: spec.model_id.clone(),
+            spec: spec.clone(),
+            devices,
+            checkpoints: vec![checkpoint],
+            poisoned: Arc::new(AtomicBool::new(false)),
+            admission: crate::harness::admission::AdmissionController::new(&self.admission_cfg),
+        };
+        let mut models = self.models.write().await;
+        models.insert(
+            spec.model_id.clone(),
+            LoadedHandle::Decision(Arc::new(loaded)),
+        );
+        tracing::info!(model = %spec.model_id, dtype = ?dtype, "decision model loaded");
+        Ok(())
+    }
+
+    /// Fetch one Laya checkpoint's files through hf-hub: the repo root
+    /// for `subfolder = None`, else that subfolder. Only the files the
+    /// checkpoint needs are fetched, not the whole repo.
+    async fn resolve_decision_files(
+        &self,
+        source_id: &cortex_core::source::ModelSourceId,
+        subfolder: Option<&str>,
+    ) -> Result<super::decision_model::CheckpointFiles> {
+        let api = self.hf_api_for(&source_id.scheme)?;
+        let repo = api.model(source_id.repo_path());
+        let display_id = source_id.to_string();
+        let prefix = subfolder.map(|s| format!("{s}/")).unwrap_or_default();
+        let mut fetched = Vec::new();
+        for file in [
+            "rl_agent_config.json",
+            "encoder/config.json",
+            "model.safetensors",
+            "tokenizer/tokenizer.json",
+        ] {
+            let path = format!("{prefix}{file}");
+            fetched.push(
+                repo.get(&path)
+                    .await
+                    .with_context(|| format!("fetch {path} from {display_id}"))?,
+            );
+        }
+        // hf-hub materialises a snapshot as a directory tree mirroring
+        // the repo, so the checkpoint directory is where its config sits.
+        let dir = fetched[0]
+            .parent()
+            .context("checkpoint config has no parent directory")?
+            .to_path_buf();
+        Ok(super::decision_model::CheckpointFiles {
+            dir,
+            tokenizer: fetched[3].clone(),
+        })
+    }
+
     /// Resolve a diffusers-layout Z-Image repo into local component
     /// paths, downloading through hf-hub as needed. Verifies the
     /// pipeline class before fetching multi-GB weight shards.
@@ -4568,6 +4784,28 @@ impl CandleHarness {
         }
     }
 
+    /// The loaded decision model `model_id` (#336), for the
+    /// `/v1/systemone` handler: it takes the admission permit, scores
+    /// through [`super::decision_model::DecisionCheckpoint::decide`],
+    /// and shapes the answers. A model of another kind is a
+    /// `WrongModality`, as for images.
+    pub async fn decision_model(
+        &self,
+        model_id: &str,
+    ) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
+        let handle = {
+            let models = self.models.read().await;
+            models.get(model_id).cloned()
+        };
+        match handle {
+            Some(LoadedHandle::Decision(m)) => Ok(m),
+            Some(_) => Err(InferenceError::WrongModality {
+                model_id: model_id.to_string(),
+            }),
+            None => Err(InferenceError::ModelNotLoaded(model_id.to_string())),
+        }
+    }
+
     /// Run one text-to-image generation (#198). This is the internal
     /// path the `/v1/images/generations` handler (#200) fronts;
     /// admission, validation, tokenization, and PNG encoding happen
@@ -4843,7 +5081,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),
@@ -4886,13 +5124,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 {
@@ -4900,7 +5139,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),
                 },
@@ -5254,6 +5493,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
@@ -5414,6 +5656,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;
@@ -5500,6 +5748,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,
@@ -5515,6 +5764,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) => {
@@ -5543,6 +5794,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
@@ -5562,6 +5818,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,
@@ -5643,6 +5900,7 @@ impl CandleHarness {
                                     "TP stream: decode_stream step failed"
                                 ),
                             }
+                            phase_emit += t_emit.elapsed();
                         }
                     }
                 }
@@ -5707,8 +5965,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,
@@ -5722,6 +5996,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;
@@ -7168,6 +7447,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,
@@ -7301,6 +7583,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
@@ -7407,10 +7696,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,
@@ -7430,11 +7723,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());
         }
     }
@@ -7456,8 +7753,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,
@@ -7469,6 +7780,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;
@@ -7689,6 +8005,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(())
@@ -7710,6 +8027,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/capability.rs

@@ -57,6 +57,9 @@ pub fn classify_files(
 ) -> Vec<String> {
     match preflight::classify(filenames) {
         SourceFormat::Diffusers => vec!["image".to_string()],
+        // A decision model answers typed questions (`/v1/systemone`);
+        // it generates no text, so it must not claim "text".
+        SourceFormat::DecisionModel => vec!["decision".to_string()],
         // A repo with no recognised weights (tokenizer-only, or an empty
         // entry) is not servable, so claiming "text" would be a lie of
         // the same kind this module exists to remove.
@@ -101,6 +104,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn laya_repo_is_a_decision_model() {
+        let files = [
+            "rl_agent_config.json",
+            "encoder/config.json",
+            "model.safetensors",
+            "tokenizer/tokenizer.json",
+        ];
+        assert_eq!(
+            classify_files(&files, || panic!(
+                "config must not be read for a decision model"
+            )),
+            vec!["decision".to_string()]
+        );
+    }
+
     #[test]
     fn dense_repo_without_vision_config_is_text_only() {
         let files = ["config.json", "tokenizer.json", "model.safetensors"];

crates/neuron/src/harness/decision_model.rs

@@ -0,0 +1,236 @@
+//! A loaded Laya decision model, as the rest of neuron sees it (#336).
+//!
+//! The weights live in a device worker's decision slab (CUDA) or in an
+//! `Arc` scored on the blocking pool (CPU); either way the caller holds
+//! only this struct and asks it to [`DecisionCheckpoint::decide`] a
+//! batch of already-assembled questions. Turning a `/v1/systemone`
+//! request into that batch — rendering options, budgets, tokenisation —
+//! and turning the scores back into calibrated answers happen outside
+//! this module.
+//!
+//! One Hugging Face repo can carry several checkpoints (the released
+//! `convaiinnovations/laya` has the English one at its root and
+//! `multilingual/` and `typed-decisions/` beside it), so a loaded model
+//! is a list of checkpoints, each named by its subfolder.
+
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+use anyhow::{Context, Result};
+use candle_core::DType;
+use tokenizers::Tokenizer;
+
+use super::arch::laya::{CheckpointConfig, DecisionBatch, DecisionRow, LayaModel};
+use super::device_worker::{DecisionHandle, DeviceWorkerHandle};
+use cortex_core::harness::ModelSpec;
+
+/// Name of the checkpoint at the repo root.
+pub const ROOT_CHECKPOINT: &str = "root";
+
+/// Where one checkpoint's weights live and how to reach them.
+enum Backend {
+    /// Owned by a CUDA device worker; scored through its channel.
+    Worker {
+        worker: Arc<DeviceWorkerHandle>,
+        handle: DecisionHandle,
+    },
+    /// CPU fallback. `LayaModel::forward` takes `&self`, so concurrent
+    /// requests share it without a lock.
+    Cpu(Arc<LayaModel>),
+}
+
+/// The files one checkpoint needs, already on local disk.
+#[derive(Debug, Clone)]
+pub struct CheckpointFiles {
+    /// The checkpoint directory: holds `rl_agent_config.json`,
+    /// `encoder/config.json` and `model.safetensors`.
+    pub dir: PathBuf,
+    pub tokenizer: PathBuf,
+}
+
+pub struct DecisionCheckpoint {
+    /// [`ROOT_CHECKPOINT`] or the subfolder name.
+    pub name: String,
+    pub config: CheckpointConfig,
+    pub tokenizer: Tokenizer,
+    /// Precision the weights were loaded in.
+    pub dtype: DType,
+    backend: Backend,
+}
+
+impl DecisionCheckpoint {
+    /// Load one checkpoint onto `worker`'s device, or onto the CPU when
+    /// `worker` is `None`.
+    pub async fn load(
+        name: String,
+        files: CheckpointFiles,
+        model_id: &str,
+        worker: Option<Arc<DeviceWorkerHandle>>,
+    ) -> Result<Self> {
+        let config = CheckpointConfig::from_dir(&files.dir)?;
+        let tokenizer = Tokenizer::from_file(&files.tokenizer)
+            .map_err(|e| anyhow::anyhow!("load tokenizer {}: {e}", files.tokenizer.display()))?;
+        let (dtype, backend) = match worker {
+            Some(worker) => {
+                let wanted = device_dtype(config.laya.amp_dtype.as_deref());
+                let (handle, dtype) = worker
+                    .load_decision(files.dir.clone(), model_id.to_string(), wanted)
+                    .await
+                    .map_err(|e| anyhow::anyhow!("worker load_decision: {e:#}"))?;
+                (dtype, Backend::Worker { worker, handle })
+            }
+            None => {
+                let dir = files.dir.clone();
+                let model = tokio::task::spawn_blocking(move || {
+                    LayaModel::load(&dir, &candle_core::Device::Cpu, DType::F32)
+                })
+                .await
+                .context("decision load task")??;
+                (DType::F32, Backend::Cpu(Arc::new(model)))
+            }
+        };
+        Ok(Self {
+            name,
+            config,
+            tokenizer,
+            dtype,
+            backend,
+        })
+    }
+
+    /// Score one batch: one [`DecisionRow`] per sequence, in order.
+    pub async fn decide(&self, batch: DecisionBatch) -> Result<Vec<DecisionRow>> {
+        match &self.backend {
+            Backend::Worker { worker, handle } => worker
+                .decide(*handle, batch)
+                .await
+                .map_err(|e| anyhow::anyhow!("{e:#}")),
+            Backend::Cpu(model) => {
+                let model = Arc::clone(model);
+                tokio::task::spawn_blocking(move || model.forward(&batch))
+                    .await
+                    .context("decision forward task")?
+            }
+        }
+    }
+
+    /// Release the weights. On a worker this drops them on the thread
+    /// that owns the device context.
+    pub async fn release(&self) {
+        if let Backend::Worker { worker, handle } = &self.backend
+            && let Err(e) = worker.drop_decision(*handle).await
+        {
+            tracing::warn!(
+                checkpoint = %self.name,
+                error = %e,
+                "decision unload: DropDecision RPC failed (weights may leak in worker slab)"
+            );
+        }
+    }
+}
+
+/// The precision to serve in on a GPU: the autocast dtype the reference
+/// itself runs there (`amp_dtype` in `rl_agent_config.json`), else bf16.
+fn device_dtype(amp_dtype: Option<&str>) -> DType {
+    match amp_dtype {
+        Some("fp16") | Some("float16") => DType::F16,
+        Some("fp32") | Some("float32") => DType::F32,
+        _ => DType::BF16,
+    }
+}
+
+pub struct LoadedDecisionModel {
+    pub model_id: String,
+    pub spec: ModelSpec,
+    pub devices: Vec<u32>,
+    /// At least one; the first is the default.
+    pub checkpoints: Vec<DecisionCheckpoint>,
+    /// Set when a forward returns a driver-level error, mirroring the
+    /// other model kinds: `/models` reports it and requests fast-reject.
+    pub poisoned: Arc<AtomicBool>,
+    pub admission: super::admission::AdmissionController,
+}
+
+impl LoadedDecisionModel {
+    /// The checkpoint called `name`, or the default one for `None`.
+    pub fn checkpoint(&self, name: Option<&str>) -> Option<&DecisionCheckpoint> {
+        match name {
+            None => self.checkpoints.first(),
+            Some(n) => self.checkpoints.iter().find(|c| c.name == n),
+        }
+    }
+
+    pub async fn release(&self) {
+        for c in &self.checkpoints {
+            c.release().await;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn tiny() -> CheckpointFiles {
+        let dir =
+            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya/tiny");
+        CheckpointFiles {
+            tokenizer: dir.join("tokenizer.json"),
+            dir,
+        }
+    }
+
+    #[test]
+    fn gpu_dtype_follows_the_reference_autocast() {
+        assert_eq!(device_dtype(Some("bf16")), DType::BF16);
+        assert_eq!(device_dtype(Some("fp16")), DType::F16);
+        assert_eq!(device_dtype(None), DType::BF16);
+    }
+
+    /// The whole async path on the CPU backend, and the device-worker
+    /// path on a CPU build (the worker runs, with a CPU device), must
+    /// give the same scores as the model called directly.
+    #[tokio::test]
+    async fn both_backends_score_like_the_model() {
+        let files = tiny();
+        let direct = LayaModel::load(&files.dir, &candle_core::Device::Cpu, DType::F32).unwrap();
+        let batch = DecisionBatch {
+            seqs: vec![super::super::arch::laya::DecisionSeq {
+                input_ids: vec![5, 9, 11, 3, 7, 2],
+                markers: vec![1, 3],
+                qtype: 0,
+            }],
+        };
+        let want = direct.forward(&batch).unwrap();
+
+        let cpu = DecisionCheckpoint::load("root".into(), files.clone(), "tiny", None)
+            .await
+            .unwrap();
+        assert_eq!(cpu.decide(batch.clone()).await.unwrap(), want);
+
+        // A CPU build's worker holds a CPU device, so this exercises the
+        // job round-trip (load, decide, drop) without a GPU — and the
+        // worker must fall back to f32 there, reporting that it did.
+        let worker = DeviceWorkerHandle::spawn(0).unwrap();
+        let on_worker =
+            DecisionCheckpoint::load("root".into(), files, "tiny", Some(worker.clone()))
+                .await
+                .unwrap();
+        #[cfg(not(feature = "cuda"))]
+        {
+            assert_eq!(on_worker.dtype, DType::F32);
+            assert_eq!(on_worker.decide(batch).await.unwrap(), want);
+        }
+        #[cfg(feature = "cuda")]
+        {
+            let got = on_worker.decide(batch).await.unwrap();
+            for (g, w) in got[0].logits.iter().zip(&want[0].logits) {
+                assert!((g - w).abs() < 5e-2, "{g} vs {w}");
+            }
+        }
+        on_worker.release().await;
+        assert!(on_worker.decide(DecisionBatch::default()).await.is_err());
+        worker.shutdown().unwrap();
+    }
+}

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

@@ -13,12 +13,13 @@
 //! ARCH model state in this state slab will gain a companion
 //! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
 
+use crate::harness::arch::laya::LayaModel;
 use crate::harness::arch::snapshot::KvCacheSnapshot;
 use crate::harness::candle::ModelArch;
 #[cfg(feature = "cuda")]
 use crate::harness::device_worker::jobs::TpHandle;
 use crate::harness::device_worker::jobs::{
-    ArchHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
+    ArchHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
 };
 use crate::harness::image::ZImagePipeline;
 #[cfg(feature = "cuda")]
@@ -65,6 +66,11 @@ struct DeviceWorkerState {
     image_models: HashMap<ImageHandle, Box<ZImagePipeline>>,
     /// Counter for minting fresh `ImageHandle`s.
     next_image_handle: u64,
+    /// Laya decision-model slab (#336). Same lifecycle discipline as
+    /// `image_models`, in its own namespace.
+    decision_models: HashMap<DecisionHandle, Box<LayaModel>>,
+    /// Counter for minting fresh `DecisionHandle`s.
+    next_decision_handle: u64,
     /// Leader's NCCL state. Populated by `Job::NcclInit`; the
     /// underlying `Comm`'s libnccl handle lives bound to this thread
     /// for its entire lifetime. Subprocess workers maintain their own
@@ -148,6 +154,7 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 model_id,
                 quant,
                 isq_cache,
+                experts_on_host,
                 reply,
             } => {
                 let result = load_dense_inner(
@@ -157,6 +164,7 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                     &model_id,
                     quant,
                     isq_cache.as_ref(),
+                    experts_on_host,
                 )
                 .map(|arch| {
                     // Ask the model, once, while it is still in
@@ -173,6 +181,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 } => {
@@ -196,6 +214,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,
@@ -307,6 +328,15 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 }
                 let _ = reply.send(result);
             }
+            Job::ForwardLogitsMulti {
+                handle,
+                tokens,
+                offset,
+                reply,
+            } => {
+                let result = forward_logits_multi(&mut state, handle, &tokens, offset);
+                let _ = reply.send(result);
+            }
             Job::ForwardLogitsBatch {
                 handle,
                 tokens,
@@ -584,6 +614,61 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 let _ = reply.send(result);
             }
             #[cfg(feature = "cuda")]
+            Job::TpLoadMtpHead {
+                handle,
+                config_json,
+                safetensors_paths,
+                reply,
+            } => {
+                let result = tp_load_mtp_head(&mut state, handle, &config_json, &safetensors_paths);
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpMtpPrefillChunk {
+                handle,
+                shifted,
+                start_pos,
+                reply,
+            } => {
+                let result = tp_mtp(&mut state, handle, |m| {
+                    m.mtp_prefill_chunk(&shifted, start_pos)
+                });
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpMtpDraft {
+                handle,
+                first_token,
+                start_pos,
+                k,
+                reply,
+            } => {
+                let result = tp_mtp(&mut state, handle, |m| {
+                    m.mtp_draft(first_token, start_pos, k)
+                });
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpMtpAdvance {
+                handle,
+                token,
+                pos,
+                reply,
+            } => {
+                let result = tp_mtp(&mut state, handle, |m| m.mtp_advance(token, pos));
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpForwardLogitsMulti {
+                handle,
+                tokens,
+                offset,
+                reply,
+            } => {
+                let result = tp_forward_logits_multi(&mut state, handle, &tokens, offset);
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
             Job::TpForwardLogitsWithImages {
                 handle,
                 tokens,
@@ -672,6 +757,62 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 trim_device_pool(&state);
                 let _ = reply.send(result);
             }
+            Job::LoadDecision {
+                dir,
+                model_id,
+                dtype,
+                reply,
+            } => {
+                let dtype = if state.device.is_cpu() {
+                    candle_core::DType::F32
+                } else {
+                    dtype
+                };
+                let result = LayaModel::load(&dir, &state.device, dtype)
+                    .with_context(|| format!("load laya checkpoint {}", dir.display()))
+                    .map(|model| {
+                        let handle = DecisionHandle(state.next_decision_handle);
+                        state.next_decision_handle = state.next_decision_handle.wrapping_add(1);
+                        state.decision_models.insert(handle, Box::new(model));
+                        tracing::info!(
+                            device_index,
+                            model_id,
+                            checkpoint = %dir.display(),
+                            handle = handle.0,
+                            "device worker: decision model loaded"
+                        );
+                        (handle, dtype)
+                    });
+                trim_device_pool(&state);
+                let _ = reply.send(result);
+            }
+            Job::DropDecision { handle, reply } => {
+                let removed = state.decision_models.remove(&handle);
+                let was_present = removed.is_some();
+                // Explicit drop on this thread, as for DropArch/DropImage.
+                drop(removed);
+                if was_present {
+                    trim_device_pool(&state);
+                }
+                tracing::debug!(
+                    device_index,
+                    handle = handle.0,
+                    was_present,
+                    "device worker: decision model dropped"
+                );
+                let _ = reply.send(());
+            }
+            Job::Decide {
+                handle,
+                batch,
+                reply,
+            } => {
+                let result = match state.decision_models.get(&handle) {
+                    Some(model) => model.forward(&batch),
+                    None => Err(anyhow::anyhow!("no decision model for handle {}", handle.0)),
+                };
+                let _ = reply.send(result);
+            }
             Job::Shutdown => unreachable!("Shutdown should break above"),
         }
     }
@@ -744,6 +885,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
             next_kv_snapshot_id: 1,
             image_models: HashMap::new(),
             next_image_handle: 1,
+            decision_models: HashMap::new(),
+            next_decision_handle: 1,
             nccl: NcclState::new(),
             tp_models: HashMap::new(),
             next_tp_handle: 1,
@@ -762,6 +905,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
             next_kv_snapshot_id: 1,
             image_models: HashMap::new(),
             next_image_handle: 1,
+            decision_models: HashMap::new(),
+            next_decision_handle: 1,
             nccl: NcclState::new(),
         }
     }
@@ -936,6 +1081,52 @@ 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.
@@ -961,6 +1152,7 @@ fn load_dense_inner(
     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;
@@ -1078,14 +1270,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,
-                safetensors_paths,
-                isq_cache,
+                &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)))
@@ -1215,6 +1435,85 @@ fn tp_forward_logits(
     Ok(values)
 }
 
+/// Run something on the leader's TP draft head (#96).
+///
+/// The head only exists on the `qwen3_5` arch and only on a model that
+/// was asked to load one, so both misses are reported rather than
+/// silently skipped.
+#[cfg(feature = "cuda")]
+fn tp_mtp<T>(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    f: impl FnOnce(&mut super::super::tp::tp_qwen3_5::TpQwen3_5ForCausalLM) -> anyhow::Result<T>,
+) -> anyhow::Result<T> {
+    let model = state
+        .tp_models
+        .get_mut(&handle)
+        .ok_or_else(|| anyhow::anyhow!("tp mtp: no model for handle {}", handle.0))?;
+    match &mut **model {
+        super::super::tp::TpLeaderModel::Qwen3_5(m) => f(m),
+        super::super::tp::TpLeaderModel::Qwen3(_) => {
+            anyhow::bail!("tp mtp: the qwen3 arch ships no MTP head")
+        }
+    }
+}
+
+/// Load the draft head onto the leader's TP model (#96).
+#[cfg(feature = "cuda")]
+fn tp_load_mtp_head(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    config_json: &str,
+    safetensors_paths: &[String],
+) -> anyhow::Result<()> {
+    use candle_core::DType;
+
+    let cfg = super::super::tp::tp_qwen3_5::Config::from_config_json(config_json)
+        .context("tp mtp: parse config.json")?;
+    let paths: Vec<std::path::PathBuf> = safetensors_paths.iter().map(Into::into).collect();
+    let device = state.device.clone();
+    // SAFETY: the same checkpoint files the shard load already mmapped.
+    let vb = unsafe {
+        candle_nn::var_builder::ShardedSafeTensors::var_builder(&paths, DType::BF16, &device)?
+    };
+    tp_mtp(state, handle, |m| m.load_mtp_head(&cfg.text_config, &vb))
+}
+
+/// TP mirror of [`forward_logits_multi`] on the leader's shard (#96).
+/// The caller has already fanned the matching `GenerateStepMulti` out
+/// to the subprocess ranks; the `AllReduce` collectives inside this
+/// forward block until they arrive.
+#[cfg(feature = "cuda")]
+fn tp_forward_logits_multi(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    tokens: &[u32],
+    offset: usize,
+) -> anyhow::Result<Vec<Vec<f32>>> {
+    use candle_core::{DType, Tensor};
+
+    anyhow::ensure!(
+        !tokens.is_empty(),
+        "TpForwardLogitsMulti: empty token slice"
+    );
+    let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
+    let model = state
+        .tp_models
+        .get_mut(&handle)
+        .ok_or_else(|| anyhow::anyhow!("TpForwardLogitsMulti: no model for handle {}", handle.0))?;
+
+    // (1, L, vocab) → one CPU row per position.
+    let logits = model.forward_multi(&input, offset)?.to_dtype(DType::F32)?;
+    let (_, positions, vocab) = logits.dims3()?;
+    anyhow::ensure!(
+        positions == tokens.len(),
+        "TpForwardLogitsMulti: model returned {positions} rows for {} tokens",
+        tokens.len()
+    );
+    let flat: Vec<f32> = logits.flatten_all()?.to_vec1()?;
+    Ok(flat.chunks(vocab).map(<[f32]>::to_vec).collect())
+}
+
 /// TP-equivalent of [`forward_logits_batch`] on the leader's shard
 /// (#98). The caller has already fanned the matching
 /// `GenerateStepBatch` out to the subprocess ranks.
@@ -1485,6 +1784,40 @@ fn extract_kv_rows(
 /// on the worker's device, derives per-row positions and the padding
 /// mask, and copies each row's logits back to CPU — same
 /// "tensors never escape the worker" contract as `forward_logits`.
+/// Logits at every position of a multi-token forward (#96).
+///
+/// The tensors stay on the worker thread; only CPU rows cross back, one
+/// per position. At K+1 = 5 positions and a 248k vocabulary that is
+/// ~5 MB a round — acceptable while this is a measurement path, and the
+/// first thing to replace with a device-side argmax if the draft/verify
+/// loop ever runs on it in production.
+fn forward_logits_multi(
+    state: &mut DeviceWorkerState,
+    handle: ArchHandle,
+    tokens: &[u32],
+    offset: usize,
+) -> anyhow::Result<Vec<Vec<f32>>> {
+    use candle_core::{DType, Tensor};
+
+    anyhow::ensure!(!tokens.is_empty(), "ForwardLogitsMulti: empty token slice");
+    let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
+    let arch = state
+        .models
+        .get_mut(&handle)
+        .ok_or_else(|| anyhow::anyhow!("ForwardLogitsMulti: no model for handle {}", handle.0))?;
+
+    // (1, L, vocab) → one CPU row per position.
+    let logits = arch.forward_multi(&input, offset)?.to_dtype(DType::F32)?;
+    let (_, positions, vocab) = logits.dims3()?;
+    anyhow::ensure!(
+        positions == tokens.len(),
+        "ForwardLogitsMulti: model returned {positions} rows for {} tokens",
+        tokens.len()
+    );
+    let flat: Vec<f32> = logits.flatten_all()?.to_vec1()?;
+    Ok(flat.chunks(vocab).map(<[f32]>::to_vec).collect())
+}
+
 fn forward_logits_batch(
     state: &mut DeviceWorkerState,
     handle: ArchHandle,
@@ -1681,6 +2014,9 @@ fn drain_poisoned(job: Job, device_index: u32) {
         Job::AssembleKvBatch { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
+        Job::ForwardLogitsMulti { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
         Job::ForwardLogitsBatch { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
@@ -1699,6 +2035,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
         Job::GenerateImage { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
+        Job::LoadDecision { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        Job::DropDecision { reply, .. } => {
+            let _ = reply.send(());
+        }
+        Job::Decide { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
         Job::ForwardLogitsWithImages { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
@@ -1757,6 +2102,26 @@ fn drain_poisoned(job: Job, device_index: u32) {
             let _ = reply.send(());
         }
         #[cfg(feature = "cuda")]
+        Job::TpLoadMtpHead { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpMtpPrefillChunk { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpMtpDraft { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpMtpAdvance { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpForwardLogitsMulti { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
         Job::TpForwardLogits { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
@@ -1804,6 +2169,7 @@ mod tests {
                 model_type,
                 None,
                 None,
+                false,
             ) {
                 Ok(_) => panic!("a bare config must not build a model"),
                 Err(e) => format!("{e:#}"),

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

@@ -59,6 +59,15 @@ pub struct TpHandle(pub u64);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ImageHandle(pub u64);
 
+/// Opaque handle to a Laya decision model ([`LayaModel`]) stored in the
+/// worker thread's decision slab (#336). Same discipline as
+/// [`ImageHandle`]: construction, every forward and `Drop` happen on the
+/// worker thread, and `Job::DropDecision` is the only way to release it.
+///
+/// [`LayaModel`]: crate::harness::arch::laya::LayaModel
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct DecisionHandle(pub u64);
+
 /// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
 /// next to the model slab. Scoped to the `ArchHandle` it was captured
 /// from — `Job::DropArch` drops every snapshot under its handle. The
@@ -135,6 +144,8 @@ pub enum Job {
         /// 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
@@ -223,6 +234,19 @@ pub enum Job {
     /// values `AssembleKvBatch` was built from) and replies one CPU
     /// `[vocab]` logits row per batch row, ready for per-slot
     /// sampling on the async side.
+    /// Logits at every position of a multi-token forward (#96).
+    ///
+    /// `ForwardLogits` returns one `[vocab]` row because a decode step
+    /// only has one. A speculative verify pass forwards K+1 tokens and
+    /// needs the target's token at each of them, so this returns one
+    /// row per position — `Vec<Vec<f32>>`, like `ForwardLogitsBatch`,
+    /// but indexed by position rather than by batch row.
+    ForwardLogitsMulti {
+        handle: ArchHandle,
+        tokens: Vec<u32>,
+        offset: usize,
+        reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
+    },
     ForwardLogitsBatch {
         handle: ArchHandle,
         tokens: Vec<u32>,
@@ -379,6 +403,59 @@ pub enum Job {
     /// without holding a device tensor. The caller is also
     /// responsible for fan-out to subprocess ranks and drain — only
     /// the leader's forward moves into the worker thread.
+    /// Load the MTP draft head onto the leader's TP model (#96).
+    ///
+    /// Leader-only and deliberately not part of `TpLoadShard`: the
+    /// subprocess ranks have no use for a head they never run, and an
+    /// unsharded head on every rank would cost 810 MB each to sit idle.
+    #[cfg(feature = "cuda")]
+    TpLoadMtpHead {
+        handle: TpHandle,
+        config_json: String,
+        safetensors_paths: Vec<String>,
+        reply: oneshot::Sender<Result<()>>,
+    },
+    /// Walk the leader's draft head over a prompt chunk (#96), using
+    /// the hidden state the target's forward just produced.
+    #[cfg(feature = "cuda")]
+    TpMtpPrefillChunk {
+        handle: TpHandle,
+        /// The chunk's tokens shifted one left — the head trails the
+        /// target by a token.
+        shifted: Vec<u32>,
+        start_pos: usize,
+        reply: oneshot::Sender<Result<()>>,
+    },
+    /// Draft `k` tokens on the leader and put the head's cache back
+    /// (#96). Commits nothing: the snapshot/restore pair lives inside
+    /// the job so a caller cannot forget it.
+    #[cfg(feature = "cuda")]
+    TpMtpDraft {
+        handle: TpHandle,
+        first_token: u32,
+        start_pos: usize,
+        k: usize,
+        reply: oneshot::Sender<Result<Vec<u32>>>,
+    },
+    /// Advance the leader's draft head over one committed token (#96).
+    #[cfg(feature = "cuda")]
+    TpMtpAdvance {
+        handle: TpHandle,
+        token: u32,
+        pos: usize,
+        reply: oneshot::Sender<Result<()>>,
+    },
+    /// TP mirror of `ForwardLogitsMulti` (#96): logits at every
+    /// position of a multi-token forward on the leader's shard. The
+    /// caller has already fanned the matching `GenerateStepMulti` out
+    /// to the subprocess ranks, whose collectives this one waits on.
+    #[cfg(feature = "cuda")]
+    TpForwardLogitsMulti {
+        handle: TpHandle,
+        tokens: Vec<u32>,
+        offset: usize,
+        reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
+    },
     #[cfg(feature = "cuda")]
     TpForwardLogits {
         handle: TpHandle,
@@ -481,6 +558,33 @@ pub enum Job {
         max_dim: usize,
         reply: oneshot::Sender<Result<crate::harness::image::ImageGenResult>>,
     },
+    /// Load one Laya checkpoint directory (#336) — the repo root or one
+    /// of its subfolders, already on local disk — into the decision
+    /// slab. File resolution (hf-hub) happens on the async side.
+    ///
+    /// `dtype` is the precision wanted on a GPU; a worker whose device
+    /// is the CPU loads in f32 instead (the CPU backend has no bf16
+    /// matmul). The reply carries the dtype actually used.
+    LoadDecision {
+        dir: PathBuf,
+        model_id: String,
+        dtype: candle_core::DType,
+        reply: oneshot::Sender<Result<(DecisionHandle, candle_core::DType)>>,
+    },
+    /// Remove a decision model from the slab and drop it on the worker
+    /// thread. Mirrors `DropArch` / `DropImage`.
+    DropDecision {
+        handle: DecisionHandle,
+        reply: oneshot::Sender<()>,
+    },
+    /// Score one batch of assembled questions (#336). One job per
+    /// request: the whole batch is a single encoder forward, and the
+    /// reply is CPU-side scores — no device tensor crosses back.
+    Decide {
+        handle: DecisionHandle,
+        batch: crate::harness::arch::laya::DecisionBatch,
+        reply: oneshot::Sender<Result<Vec<crate::harness::arch::laya::DecisionRow>>>,
+    },
     /// Tell the worker to break its dispatch loop and exit. Any jobs
     /// queued after this in the channel reply `Err` to their oneshot
     /// senders (the senders are dropped on the worker's exit, which

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

@@ -51,7 +51,7 @@ use tokio::sync::oneshot;
 
 #[cfg(feature = "cuda")]
 pub use jobs::TpHandle;
-pub use jobs::{ArchHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
+pub use jobs::{ArchHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
 
 /// Errors returned by `DeviceWorkerHandle` submit methods.
 #[derive(Debug, thiserror::Error)]
@@ -223,6 +223,7 @@ impl DeviceWorkerHandle {
         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 {
@@ -237,6 +238,7 @@ impl DeviceWorkerHandle {
                 model_id,
                 quant,
                 isq_cache,
+                experts_on_host,
                 reply: reply_tx,
             })
             .map_err(|_| WorkerError::Gone {
@@ -335,6 +337,85 @@ impl DeviceWorkerHandle {
         }
     }
 
+    /// Load one Laya checkpoint directory (#336) on the worker thread.
+    /// Returns the handle and the dtype the weights were loaded in.
+    pub async fn load_decision(
+        &self,
+        dir: std::path::PathBuf,
+        model_id: String,
+        dtype: candle_core::DType,
+    ) -> Result<(DecisionHandle, candle_core::DType), WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::LoadDecision {
+                dir,
+                model_id,
+                dtype,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// Drop the decision model for `handle` on the worker thread. Same
+    /// poison-tolerant contract as [`Self::drop_arch`].
+    pub async fn drop_decision(&self, handle: DecisionHandle) -> Result<(), WorkerError> {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::DropDecision {
+                handle,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        reply_rx.await.map_err(|_| WorkerError::Gone {
+            device_index: self.device_index,
+        })
+    }
+
+    /// Score one batch of assembled questions (#336). The reply is
+    /// CPU-side scores; no device tensor crosses back.
+    pub async fn decide(
+        &self,
+        handle: DecisionHandle,
+        batch: crate::harness::arch::laya::DecisionBatch,
+    ) -> Result<Vec<crate::harness::arch::laya::DecisionRow>, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::Decide {
+                handle,
+                batch,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
     /// Tell the worker to drop the `ModelArch` for `handle` on the
     /// worker thread (so CUDA tensors release on the right context).
     /// Returns `Ok(())` even if the handle wasn't in the slab — Drop
@@ -576,6 +657,40 @@ impl DeviceWorkerHandle {
         }
     }
 
+    /// Multi-position forward (#96): logits at every position of
+    /// `tokens`, one CPU `[vocab]` row each. Same no-device-tensor
+    /// contract as [`Self::forward_logits`]; the speculative verify
+    /// pass reads the target's token at each drafted position.
+    pub async fn forward_logits_multi(
+        &self,
+        handle: ArchHandle,
+        tokens: Vec<u32>,
+        offset: usize,
+    ) -> Result<Vec<Vec<f32>>, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::ForwardLogitsMulti {
+                handle,
+                tokens,
+                offset,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
     /// One lockstep batched decode step (#98): row i's next token at
     /// position `prefix_lens[i] + step`. Returns one CPU `[vocab]`
     /// logits row per batch row, ready for per-slot sampling — same
@@ -932,6 +1047,137 @@ impl DeviceWorkerHandle {
     /// logits as `Vec<f32>` ready for sampling. The caller is
     /// responsible for fan-out / drain of the subprocess workers
     /// concurrently with this call.
+    /// Load the MTP draft head onto the leader's TP model (#96).
+    #[cfg(feature = "cuda")]
+    pub async fn tp_load_mtp_head(
+        &self,
+        handle: TpHandle,
+        config_json: String,
+        safetensors_paths: Vec<String>,
+    ) -> Result<(), WorkerError> {
+        self.tp_mtp_job(|reply| Job::TpLoadMtpHead {
+            handle,
+            config_json,
+            safetensors_paths,
+            reply,
+        })
+        .await
+    }
+
+    /// Walk the leader's draft head over a prompt chunk (#96).
+    #[cfg(feature = "cuda")]
+    pub async fn tp_mtp_prefill_chunk(
+        &self,
+        handle: TpHandle,
+        shifted: Vec<u32>,
+        start_pos: usize,
+    ) -> Result<(), WorkerError> {
+        self.tp_mtp_job(|reply| Job::TpMtpPrefillChunk {
+            handle,
+            shifted,
+            start_pos,
+            reply,
+        })
+        .await
+    }
+
+    /// Draft `k` tokens on the leader, rewinding the head afterwards
+    /// (#96). Commits nothing.
+    #[cfg(feature = "cuda")]
+    pub async fn tp_mtp_draft(
+        &self,
+        handle: TpHandle,
+        first_token: u32,
+        start_pos: usize,
+        k: usize,
+    ) -> Result<Vec<u32>, WorkerError> {
+        self.tp_mtp_job(|reply| Job::TpMtpDraft {
+            handle,
+            first_token,
+            start_pos,
+            k,
+            reply,
+        })
+        .await
+    }
+
+    /// Advance the leader's draft head over one committed token (#96).
+    #[cfg(feature = "cuda")]
+    pub async fn tp_mtp_advance(
+        &self,
+        handle: TpHandle,
+        token: u32,
+        pos: usize,
+    ) -> Result<(), WorkerError> {
+        self.tp_mtp_job(|reply| Job::TpMtpAdvance {
+            handle,
+            token,
+            pos,
+            reply,
+        })
+        .await
+    }
+
+    /// Shared submit/await for the leader-only draft-head jobs (#96):
+    /// same poison check and channel handling as every other job, in
+    /// one place because these four differ only in their payload.
+    #[cfg(feature = "cuda")]
+    async fn tp_mtp_job<T>(
+        &self,
+        build: impl FnOnce(oneshot::Sender<anyhow::Result<T>>) -> Job,
+    ) -> Result<T, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(build(reply_tx))
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// TP mirror of [`Self::forward_logits_multi`] (#96): one CPU
+    /// `[vocab]` row per position from the leader's shard.
+    #[cfg(feature = "cuda")]
+    pub async fn tp_forward_logits_multi(
+        &self,
+        handle: TpHandle,
+        tokens: Vec<u32>,
+        offset: usize,
+    ) -> Result<Vec<Vec<f32>>, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::TpForwardLogitsMulti {
+                handle,
+                tokens,
+                offset,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
     #[cfg(feature = "cuda")]
     pub async fn tp_forward_logits(
         &self,
@@ -1249,6 +1495,7 @@ mod tests {
                 "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
@@ -1263,6 +1265,7 @@ mod tests {
                 "qwen3_next-tiny".into(),
                 None,
                 None,
+                false,
             )
             .await
             .expect("load fixture")

crates/neuron/src/harness/mod.rs

@@ -6,6 +6,7 @@ pub mod candle;
 pub mod capability;
 pub mod chat_template;
 pub mod context_limit;
+pub mod decision_model;
 pub mod device_worker;
 pub mod engine;
 pub mod image;

crates/neuron/src/harness/preflight.rs

@@ -54,6 +54,11 @@ pub enum SourceFormat {
     /// Takes precedence over the safetensors sniff because diffusers
     /// repos also contain `.safetensors` shards (in subdirs).
     Diffusers,
+    /// A Laya decision-model checkpoint (#336): `rl_agent_config.json`
+    /// plus `encoder/config.json` at the root, no top-level
+    /// `config.json`. Served by the decision path — single device,
+    /// unquantized. Takes precedence over the safetensors sniff.
+    DecisionModel,
 }
 
 /// Output of `preflight` for a load that can proceed. Carries the
@@ -115,6 +120,12 @@ pub enum PreflightError {
         available: Vec<String>,
         nearest: Option<String>,
     },
+
+    /// The repo's format cannot be served with an option the spec
+    /// asked for (e.g. `tensor_parallel` or `quant` on a decision
+    /// model). `reason` says which and what to set instead.
+    #[error("cannot load '{model_id}' as requested: {reason}")]
+    Unsupported { model_id: String, reason: String },
 }
 
 /// Run the placement check.
@@ -230,6 +241,24 @@ pub async fn preflight(
         // The architecture compatibility check stays where it is —
         // `check_dense_config_supported` runs once `config.json` is
         // on disk, since it needs the parsed JSON.
+        // Decision models (#336) load unquantized on one device.
+        (SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
+            model_id: source_id.to_string(),
+            reason: format!(
+                "decision models load on a single device; tensor_parallel={tp} is not supported"
+            ),
+        }),
+        (SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
+            model_id: source_id.to_string(),
+            reason: format!("decision models load unquantized; remove quant={q:?}"),
+        }),
+        (SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
+            model_id: source_id.to_string(),
+            format: format.clone(),
+            tp_size,
+            picked_quant_file: None,
+        }),
+
         (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
             Ok(PlacementPlan {
                 model_id: source_id.to_string(),
@@ -295,6 +324,9 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
     if crate::harness::image::is_diffusers_layout(filenames) {
         return SourceFormat::Diffusers;
     }
+    if is_decision_layout(filenames) {
+        return SourceFormat::DecisionModel;
+    }
     let mut gguf_quants: Vec<String> = filenames
         .iter()
         .filter(|f| f.to_lowercase().ends_with(".gguf"))
@@ -318,6 +350,16 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
     }
 }
 
+/// A Laya decision checkpoint at the repo root: its decision config and
+/// its encoder config, with no top-level `config.json` (a causal LM
+/// that happened to ship a file with the same name still classifies as
+/// a text model).
+pub fn is_decision_layout(filenames: &[&str]) -> bool {
+    filenames.contains(&"rl_agent_config.json")
+        && filenames.contains(&"encoder/config.json")
+        && !filenames.contains(&"config.json")
+}
+
 /// Mirror of the quant-matching logic in `candle.rs::resolve_files` so
 /// preflight picks the same file the downstream loader would. Empty
 /// quant returns the first `.gguf` (any quant). Lowercased substring
@@ -403,6 +445,8 @@ mod tests {
             devices: None,
             sampling: None,
             preserve_thinking: None,
+            expert_residency: None,
+            max_in_flight: None,
         }
     }
 
@@ -649,6 +693,20 @@ mod tests {
                 tp_size,
                 picked_quant_file: None,
             }),
+            (SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!("decision models load on a single device; tensor_parallel={tp}"),
+            }),
+            (SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!("decision models load unquantized; remove quant={q:?}"),
+            }),
+            (SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
+                model_id: source_id.to_string(),
+                format: format.clone(),
+                tp_size,
+                picked_quant_file: None,
+            }),
             (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
                 Ok(PlacementPlan {
                     model_id: source_id.to_string(),
@@ -672,6 +730,66 @@ mod tests {
         assert_eq!(classify(&files), SourceFormat::Diffusers);
     }
 
+    /// The released `convaiinnovations/laya` listing: the English
+    /// checkpoint at the root, two more in subfolders, and no root
+    /// `config.json`.
+    const LAYA_FILES: [&str; 12] = [
+        "README.md",
+        "rl_agent_config.json",
+        "encoder/config.json",
+        "model.safetensors",
+        "tokenizer/tokenizer.json",
+        "tokenizer/tokenizer_config.json",
+        "multilingual/rl_agent_config.json",
+        "multilingual/encoder/config.json",
+        "multilingual/model.safetensors",
+        "typed-decisions/rl_agent_config.json",
+        "typed-decisions/encoder/config.json",
+        "typed-decisions/model.safetensors",
+    ];
+
+    #[test]
+    fn classify_laya_layout() {
+        assert_eq!(classify(&LAYA_FILES), SourceFormat::DecisionModel);
+        // A causal LM that also ships an rl config is still a text model.
+        let lm = [
+            "config.json",
+            "rl_agent_config.json",
+            "encoder/config.json",
+            "model.safetensors",
+        ];
+        assert!(matches!(
+            classify(&lm),
+            SourceFormat::DenseSafetensors { .. }
+        ));
+    }
+
+    #[test]
+    fn feasibility_decision_model() {
+        let fmt = classify(&LAYA_FILES);
+        let plan = decide(
+            &spec("convaiinnovations/laya", None, None),
+            &fmt,
+            &LAYA_FILES,
+        )
+        .expect("single-device unquantized decision load should pass");
+        assert_eq!(plan.format, SourceFormat::DecisionModel);
+        let err = decide(
+            &spec("convaiinnovations/laya", Some(2), None),
+            &fmt,
+            &LAYA_FILES,
+        )
+        .expect_err("tp>1 must be refused");
+        assert!(matches!(err, PreflightError::Unsupported { .. }));
+        let err = decide(
+            &spec("convaiinnovations/laya", None, Some("q8_0")),
+            &fmt,
+            &LAYA_FILES,
+        )
+        .expect_err("quant must be refused");
+        assert!(matches!(err, PreflightError::Unsupported { .. }));
+    }
+
     #[test]
     fn feasibility_diffusers_single_gpu_ok() {
         let files = ["model_index.json", "transformer/config.json"];

crates/neuron/src/harness/testdata/laya/lang.json

@@ -0,0 +1,347 @@
+{
+ "header": {
+  "laya_version": "0.3.21",
+  "torch_version": "2.14.0+cu130",
+  "device": "cpu",
+  "dtype": "float32",
+  "checkpoints": {
+   "english": {
+    "repo": "convaiinnovations/laya",
+    "revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
+    "cfg": {
+     "max_len": 512,
+     "head_max_len": 192,
+     "temperature": [
+      1.6369030475616455,
+      1.2514300346374512,
+      1.983399510383606
+     ],
+     "temperature_by_options": {
+      "choice:3-5": 1.7601518630981445,
+      "choice:6-10": 1.0000158548355103,
+      "score:3-5": 1.2514300346374512,
+      "noul:2": 1.983399510383606,
+      "choice:11+": 0.10058280825614929,
+      "choice:2": 1.9063563346862793
+     }
+    },
+    "special_ids": {
+     "cls": 50281,
+     "sep": 50282,
+     "mask": 50284,
+     "pad": 50283
+    }
+   },
+   "multilingual": {
+    "repo": "convaiinnovations/laya/multilingual",
+    "revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
+    "cfg": {
+     "max_len": 1024,
+     "head_max_len": 256,
+     "temperature": [
+      1.0,
+      1.0,
+      1.0
+     ],
+     "temperature_by_options": {}
+    },
+    "special_ids": {
+     "cls": 2,
+     "sep": 1,
+     "mask": 4,
+     "pad": 0
+    }
+   },
+   "typed-decisions": {
+    "repo": "convaiinnovations/laya/typed-decisions",
+    "revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
+    "cfg": {
+     "max_len": 1024,
+     "head_max_len": 256,
+     "temperature": [
+      1.0148024559020996,
+      1.0374259948730469,
+      1.0575125217437744
+     ],
+     "temperature_by_options": {
+      "choice:3-5": 1.7601518630981445,
+      "choice:6-10": 1.0000158548355103,
+      "score:3-5": 1.2514300346374512,
+      "noul:2": 1.983399510383606,
+      "choice:11+": 0.10058280825614929,
+      "choice:2": 1.9063563346862793
+     }
+    },
+    "special_ids": {
+     "cls": 50281,
+     "sep": 50282,
+     "mask": 50284,
+     "pad": 50283
+    }
+   }
+  }
+ },
+ "cases": [
+  {
+   "state": "I was charged twice this month, I want my money back",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": "en",
+    "is_english": true,
+    "language_undecided": false,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "Me han cobrado dos veces este mes",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": "es",
+    "is_english": false,
+    "language_undecided": false,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "Die Rechnung ist falsch, bitte korrigieren",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": "de",
+    "is_english": false,
+    "language_undecided": false,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "Je voudrais annuler ma commande",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": "fr",
+    "is_english": false,
+    "language_undecided": false,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "मुझे पैसे वापस चाहिए",
+   "analysis": {
+    "script": "devanagari",
+    "script_profile": {
+     "devanagari": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "ខ្ញុំចង់បានប្រាក់វិញ",
+   "analysis": {
+    "script": "khmer",
+    "script_profile": {
+     "khmer": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "这个月我被扣了两次钱",
+   "analysis": {
+    "script": "han",
+    "script_profile": {
+     "han": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "注文をキャンセルしたいです",
+   "analysis": {
+    "script": "kana",
+    "script_profile": {
+     "han": 0.15384615384615385,
+     "kana": 0.8461538461538461
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "تم خصم المبلغ مرتين",
+   "analysis": {
+    "script": "arabic",
+    "script_profile": {
+     "arabic": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "Мне дважды списали деньги",
+   "analysis": {
+    "script": "cyrillic",
+    "script_profile": {
+     "cyrillic": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 1.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "Order #5521 — 注文をキャンセルしたいです. Please help.",
+   "analysis": {
+    "script": "kana",
+    "script_profile": {
+     "latin": 0.5357142857142857,
+     "han": 0.07142857142857142,
+     "kana": 0.39285714285714285
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.4643,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "ok",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": null,
+    "is_english": true,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "",
+   "analysis": {
+    "script": "unknown",
+    "script_profile": {},
+    "language": null,
+    "is_english": true,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "12345 67890",
+   "analysis": {
+    "script": "unknown",
+    "script_profile": {},
+    "language": null,
+    "is_english": true,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": "café résumé naïve",
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": null,
+    "is_english": false,
+    "language_undecided": true,
+    "diacritic_rate": 0.2353,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": {
+    "subject": "Invoice",
+    "body": "Bitte zahlen Sie bis Freitag"
+   },
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": null,
+    "is_english": true,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  },
+  {
+   "state": [
+    "hello",
+    "bonjour",
+    "hola"
+   ],
+   "analysis": {
+    "script": "latin",
+    "script_profile": {
+     "latin": 1.0
+    },
+    "language": null,
+    "is_english": true,
+    "language_undecided": true,
+    "diacritic_rate": 0.0,
+    "non_latin_fraction": 0.0,
+    "mixed_segment": null
+   }
+  }
+ ]
+}

crates/neuron/src/harness/testdata/laya/tiny/encoder/config.json

@@ -0,0 +1,52 @@
+{
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "cls",
+  "cls_token_id": 1,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "embedding_dropout": 0.0,
+  "eos_token_id": 2,
+  "global_attn_every_n_layers": 3,
+  "hidden_activation": "gelu",
+  "hidden_size": 32,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.3,
+  "intermediate_size": 48,
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "local_attention": 8,
+  "max_position_embeddings": 128,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 2,
+  "num_hidden_layers": 4,
+  "pad_token_id": 0,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 1000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 300.0,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 2,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "vocab_size": 97
+}

crates/neuron/src/harness/testdata/laya/tiny/expected.json

@@ -0,0 +1,182 @@
+{
+ "seed": 7,
+ "seqs": [
+  {
+   "input_ids": [
+    1,
+    93,
+    82,
+    87,
+    51,
+    46,
+    89,
+    6,
+    68,
+    94,
+    86,
+    43,
+    23,
+    70,
+    42,
+    60,
+    35,
+    3,
+    86,
+    7,
+    20,
+    18,
+    90,
+    84,
+    62,
+    30,
+    42,
+    65,
+    34,
+    37,
+    30,
+    67,
+    2,
+    34,
+    77,
+    96,
+    24,
+    89,
+    19,
+    33
+   ],
+   "markers": [
+    3,
+    9,
+    17
+   ],
+   "qtype": 0,
+   "logits": [
+    -0.4152399003505707,
+    -0.5505563020706177,
+    -0.600245475769043
+   ],
+   "act": [
+    0.4727683663368225,
+    0.5272316336631775
+   ]
+  },
+  {
+   "input_ids": [
+    4,
+    5,
+    43,
+    95,
+    77,
+    35,
+    61,
+    13,
+    42,
+    24,
+    21,
+    7,
+    11,
+    38,
+    80,
+    54,
+    4,
+    29,
+    5,
+    80,
+    54,
+    78,
+    46
+   ],
+   "markers": [
+    4,
+    11
+   ],
+   "qtype": 1,
+   "logits": [
+    -0.1313231736421585,
+    -0.3265465199947357
+   ],
+   "act": [
+    0.40568628907203674,
+    0.5943136811256409
+   ]
+  },
+  {
+   "input_ids": [
+    67,
+    40,
+    83,
+    41,
+    77,
+    96,
+    58,
+    39,
+    76
+   ],
+   "markers": [
+    2
+   ],
+   "qtype": 2,
+   "logits": [
+    0.534652590751648
+   ],
+   "act": [
+    0.4226877987384796,
+    0.5773122310638428
+   ]
+  },
+  {
+   "input_ids": [
+    96,
+    32,
+    60,
+    44,
+    35,
+    47,
+    45,
+    51,
+    60,
+    51,
+    73,
+    50,
+    25,
+    79,
+    74,
+    24,
+    83,
+    7,
+    9,
+    91,
+    21,
+    87,
+    29,
+    25,
+    74,
+    18,
+    31,
+    46,
+    46,
+    30,
+    52
+   ],
+   "markers": [
+    5,
+    7,
+    12,
+    20,
+    26
+   ],
+   "qtype": 0,
+   "logits": [
+    -0.4001395106315613,
+    -0.5266168117523193,
+    -0.44161877036094666,
+    -0.43973028659820557,
+    -0.5682165622711182
+   ],
+   "act": [
+    0.46282079815864563,
+    0.537179172039032
+   ]
+  }
+ ]
+}

crates/neuron/src/harness/testdata/laya/tiny/rl_agent_config.json

@@ -0,0 +1,16 @@
+{
+ "encoder": "tiny-modernbert",
+ "head_layers": 2,
+ "max_len": 64,
+ "head_max_len": 24,
+ "act_costs": {
+  "escalate": 0.5
+ },
+ "amp_dtype": "bf16",
+ "temperature": [
+  1.0,
+  1.0,
+  1.0
+ ],
+ "temperature_by_options": {}
+}

crates/neuron/src/harness/testdata/laya/tiny/tokenizer.json

@@ -0,0 +1 @@
+{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [], "normalizer": null, "pre_tokenizer": {"type": "Whitespace"}, "post_processor": null, "decoder": null, "model": {"type": "WordLevel", "vocab": {"[UNK]": 0, "a": 1, "b": 2}, "unk_token": "[UNK]"}}

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

@@ -67,6 +67,22 @@ impl TpLeaderModel {
         }
     }
 
+    /// Logits at every position (#96). Only the `qwen3_5` arch has a
+    /// multi-position path; the older dense `qwen3` arch says so rather
+    /// than silently returning the last row.
+    pub fn forward_multi(
+        &mut self,
+        input: &candle_core::Tensor,
+        offset: usize,
+    ) -> candle_core::Result<candle_core::Tensor> {
+        match self {
+            TpLeaderModel::Qwen3_5(m) => m.forward_multi(input, offset),
+            TpLeaderModel::Qwen3(_) => Err(candle_core::Error::Msg(
+                "forward_multi: the tp qwen3 arch has no multi-position logits path".into(),
+            )),
+        }
+    }
+
     /// Chunked image prefill on rank 0. Only the vision-capable
     /// `qwen3_5` arch supports it; the dense `qwen3` arch has no tower.
     pub fn prefill_with_images_chunked(
@@ -996,6 +1012,112 @@ impl WorkerPool {
         }
     }
 
+    /// Multi-position forward across every rank (#96).
+    ///
+    /// Same fan-out / leader-forward / always-drain shape as
+    /// [`Self::generate_step`], and for the same reason: the
+    /// row-parallel `AllReduce`s only complete when every rank has
+    /// issued them, and a leader error must still drain the workers or
+    /// the stale replies poison the next request.
+    ///
+    /// Returns one `[vocab]` row per position from the leader's rank-0
+    /// shard, which is the whole answer because `lm_head` is
+    /// replicated.
+    #[cfg(feature = "cuda")]
+    pub async fn generate_step_multi(
+        &mut self,
+        model_id: &str,
+        leader_handle: super::device_worker::TpHandle,
+        tokens: Vec<u32>,
+        offset: usize,
+    ) -> Result<Vec<Vec<f32>>> {
+        let step_start = std::time::Instant::now();
+        let tokens_len = tokens.len();
+        tracing::debug!(
+            model = %model_id,
+            tokens = tokens_len,
+            offset,
+            "WorkerPool::generate_step_multi: fan-out"
+        );
+        for w in &mut self.workers {
+            w.send_only(&WorkerRequest::GenerateStepMulti {
+                model_id: model_id.to_string(),
+                tokens: tokens.clone(),
+                offset,
+            })
+            .await?;
+        }
+
+        let leader_start = std::time::Instant::now();
+        let timeout = tp_step_timeout();
+        let leader_fut = self
+            .leader_worker
+            .tp_forward_logits_multi(leader_handle, tokens, offset);
+        let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
+            Ok(r) => r,
+            Err(_elapsed) => {
+                self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
+                anyhow::bail!(
+                    "tp watchdog: leader multi-position forward exceeded {}s deadline; \
+                     aborted wedged NCCL comm — model will auto-recover",
+                    timeout.as_secs()
+                );
+            }
+        };
+        let leader_ok = leader_result.is_ok();
+        if !leader_ok {
+            let detail = leader_result
+                .as_ref()
+                .err()
+                .map(|e| format!("{e:#}"))
+                .unwrap_or_default();
+            tracing::warn!(
+                model = %model_id,
+                tokens = tokens_len,
+                offset,
+                leader_ms = leader_start.elapsed().as_millis(),
+                error = %detail,
+                "WorkerPool::generate_step_multi: leader forward failed"
+            );
+        }
+
+        let worker_errors = drain_workers(&mut self.workers, |r| match r {
+            WorkerResponse::GenerateStepOk => Ok(()),
+            WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
+            other => Err(format!("expected GenerateStepOk, got {other:?}")),
+        })
+        .await;
+        tracing::debug!(
+            model = %model_id,
+            errors = worker_errors.len(),
+            total_ms = step_start.elapsed().as_millis(),
+            "WorkerPool::generate_step_multi: workers drained"
+        );
+
+        match leader_result {
+            Ok(values) => {
+                if worker_errors.is_empty() {
+                    Ok(values)
+                } else {
+                    anyhow::bail!(
+                        "GenerateStepMulti: leader succeeded but workers failed: {}",
+                        worker_errors.join("; ")
+                    )
+                }
+            }
+            Err(e) => {
+                if worker_errors.is_empty() {
+                    Err(anyhow::Error::new(e).context("GenerateStepMulti: leader forward failed"))
+                } else {
+                    Err(anyhow::Error::new(e).context(format!(
+                        "GenerateStepMulti: leader forward failed and workers also failed: {}",
+                        worker_errors.join("; ")
+                    )))
+                }
+            }
+        }
+    }
+
     /// One lockstep batched decode step across every rank (#98).
     /// Same fan-out / leader-forward / always-drain shape as
     /// [`Self::generate_step`]; every rank derives positions + mask

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

@@ -78,6 +78,18 @@ pub enum WorkerRequest {
     /// same op. The leader does *not* receive logits back over RPC; it
     /// runs its own rank-0 forward in parallel and uses its own logits
     /// for sampling.
+    /// Multi-position forward (#96): like `GenerateStep`, but every
+    /// rank keeps the logits at *all* `tokens.len()` positions rather
+    /// than the last. The collectives are identical, so the ranks stay
+    /// in lockstep; only the leader's logits are read. Worker replies
+    /// with `GenerateStepOk`.
+    GenerateStepMulti {
+        model_id: String,
+        /// Identical on every rank, as for `GenerateStep`.
+        tokens: Vec<u32>,
+        offset: usize,
+    },
+
     GenerateStep {
         model_id: String,
         /// Input token ids for this step. For prefill, the whole prompt;

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

@@ -1881,6 +1881,27 @@ impl TpQwen3_5Model {
     }
 }
 
+/// Additive causal mask for the draft head, `(1, 1, l, l + offset)`.
+/// The head keeps its own KV, so its mask is built from its own
+/// position count rather than the target's.
+fn mtp_causal_mask(
+    l: usize,
+    offset: usize,
+    dtype: DType,
+    dev: &Device,
+) -> candle_core::Result<Tensor> {
+    let total = l + offset;
+    let mut data = vec![0f32; l * total];
+    for i in 0..l {
+        for j in 0..total {
+            if j > i + offset {
+                data[i * total + j] = f32::NEG_INFINITY;
+            }
+        }
+    }
+    Tensor::from_vec(data, (1, 1, l, total), dev)?.to_dtype(dtype)
+}
+
 pub struct TpQwen3_5ForCausalLM {
     base: TpQwen3_5Model,
     lm_head: super::tp_linear::MaybeQuantLinear,
@@ -1894,6 +1915,23 @@ pub struct TpQwen3_5ForCausalLM {
     /// `<|image_pad|>` sentinel id (mirrors `Config::image_token_id`);
     /// the splice target for `forward_with_vision`.
     image_token_id: Option<u32>,
+    /// The MTP draft head (#96), on the **leader only**.
+    ///
+    /// Unsharded, unlike every other weight here, and that is correct
+    /// rather than lazy: after each row-parallel `AllReduce` every rank
+    /// holds the *same* hidden state, and `lm_head` is replicated. So a
+    /// head fed the leader's hidden produces exactly what a sharded one
+    /// would, with no collectives — which also means drafting cannot
+    /// wedge the ranks, since the other ranks are not in it.
+    ///
+    /// `None` on the subprocess ranks and on any load that did not ask
+    /// for it. Loaded by [`Self::load_mtp_head`], never by `load`.
+    mtp: Option<crate::harness::arch::qwen3_5::mtp::MtpHead>,
+    /// The last forward's hidden state, `(1, L, hidden)` — what the
+    /// draft head consumes. Kept because the head runs *after* the
+    /// target's step, outside the collective, so it cannot recompute
+    /// it. One `Arc` clone per forward; no copy.
+    last_hidden: Option<Tensor>,
 }
 
 /// Load the replicated vision tower from the unsharded `model.visual.*`
@@ -1944,6 +1982,8 @@ impl TpQwen3_5ForCausalLM {
             lm_head,
             vision,
             image_token_id,
+            mtp: None,
+            last_hidden: None,
         };
         log_construction_complete(cfg, rank, world_size, quant, model.device());
         Ok(model)
@@ -1969,6 +2009,8 @@ impl TpQwen3_5ForCausalLM {
             lm_head,
             vision,
             image_token_id,
+            mtp: None,
+            last_hidden: None,
         };
         log_construction_complete(cfg, rank, world_size, quant, model.device());
         Ok(model)
@@ -2003,9 +2045,201 @@ impl TpQwen3_5ForCausalLM {
     pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
         let (_, l) = input.dims2()?;
         let hidden = self.base.forward(input, offset)?;
+        self.last_hidden = Some(hidden.clone());
         hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
     }
 
+    /// Logits at **every** position: `(B, L, vocab)`.
+    ///
+    /// The TP mirror of `Qwen3_5ForCausalLM::forward_multi` (#96). A
+    /// speculative verify pass forwards the drafted block in one pass
+    /// and needs the target's own token at each drafted position;
+    /// [`Self::forward`] slices the last one off before the LM head.
+    ///
+    /// Every rank runs this, because the row-parallel `o_proj` and the
+    /// MLP's `down_proj` issue `AllReduce` collectives that only
+    /// complete when all ranks arrive. The leader's copy is the one
+    /// whose logits are read — `lm_head` is replicated, so rank 0's
+    /// answer is the whole answer.
+    /// Load the MTP draft head onto this rank, unsharded (#96).
+    ///
+    /// Only the leader should call this. See the `mtp` field for why an
+    /// unsharded head is numerically right; the practical consequence
+    /// is that drafting issues no collectives and therefore cannot
+    /// leave the other ranks waiting.
+    pub fn load_mtp_head(&mut self, cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<()> {
+        let head = crate::harness::arch::qwen3_5::mtp::MtpHead::load(
+            cfg,
+            Arc::clone(&self.base.rotary),
+            vb,
+        )
+        .context("load mtp head on the tp leader")?;
+        self.mtp = Some(head);
+        Ok(())
+    }
+
+    /// Whether this rank carries a draft head.
+    pub fn has_mtp(&self) -> bool {
+        self.mtp.is_some()
+    }
+
+    /// Walk the draft head over a prompt chunk, using the hidden state
+    /// the target's forward just produced.
+    ///
+    /// The head trails the target by one token: at position `i` it
+    /// consumes the embedding of token `i+1` with the target's hidden
+    /// state at `i`. `shifted` is therefore the chunk's tokens moved
+    /// one left, and must be one shorter than the retained hidden block
+    /// unless the caller has the next chunk's first token.
+    pub fn mtp_prefill_chunk(&mut self, shifted: &[u32], start_pos: usize) -> Result<()> {
+        let hidden = self
+            .last_hidden
+            .as_ref()
+            .context("mtp_prefill_chunk: no hidden state retained; run the target forward first")?
+            .clone();
+        let l = shifted.len();
+        anyhow::ensure!(l > 0, "mtp_prefill_chunk: empty chunk");
+        let (_, have, _) = hidden.dims3()?;
+        anyhow::ensure!(
+            have >= l,
+            "mtp_prefill_chunk: {l} tokens against {have} retained hidden positions"
+        );
+        let device = self.base.device.clone();
+        let dtype = self.base.dtype;
+        let rotary = Arc::clone(&self.base.rotary);
+        let tokens = Tensor::new(shifted, &device)?.unsqueeze(0)?;
+        let embeds = self.base.embed_tokens.forward(&tokens)?;
+        let hid = hidden.i((.., ..l, ..))?;
+        let mask = mtp_causal_mask(l, start_pos, dtype, &device)?;
+        let positions: Vec<usize> = (start_pos..start_pos + l).collect();
+        let (cos, sin) = rotary.cos_sin_at(&positions)?;
+        let head = self
+            .mtp
+            .as_mut()
+            .context("mtp_prefill_chunk: no draft head on this rank")?;
+        head.forward(&embeds, &hid, Some(&mask), &cos, &sin)?;
+        Ok(())
+    }
+
+    /// Draft `k` tokens, then put the head's cache back.
+    ///
+    /// The head is walked over tokens that may be rejected, so the
+    /// snapshot/restore pair around the chain is what keeps this an
+    /// observation rather than a commitment. Returns the drafted token
+    /// ids, argmax-sampled through the target's own (replicated)
+    /// `lm_head`.
+    pub fn mtp_draft(&mut self, first_token: u32, start_pos: usize, k: usize) -> Result<Vec<u32>> {
+        anyhow::ensure!(k > 0, "mtp_draft: k must be positive");
+        let mut hidden = self
+            .last_hidden
+            .as_ref()
+            .context("mtp_draft: no hidden state retained; run the target forward first")?
+            .clone();
+        let (_, have, _) = hidden.dims3()?;
+        if have > 1 {
+            hidden = hidden.i((.., have - 1.., ..))?;
+        }
+        let device = self.base.device.clone();
+        let rotary = Arc::clone(&self.base.rotary);
+
+        let snap = {
+            let head = self
+                .mtp
+                .as_ref()
+                .context("mtp_draft: no draft head on this rank")?;
+            head.snapshot_kv()?
+        };
+
+        let mut drafted = Vec::with_capacity(k);
+        let mut token = first_token;
+        let mut cur = hidden;
+        let mut err: Option<anyhow::Error> = None;
+        for step in 0..k {
+            let ids = match Tensor::new(&[token], &device).and_then(|t| t.unsqueeze(0)) {
+                Ok(t) => t,
+                Err(e) => {
+                    err = Some(e.into());
+                    break;
+                }
+            };
+            let out = (|| -> Result<Tensor> {
+                let embeds = self.base.embed_tokens.forward(&ids)?;
+                let (cos, sin) = rotary.cos_sin_at(&[start_pos + step])?;
+                let head = self
+                    .mtp
+                    .as_mut()
+                    .context("mtp_draft: no draft head on this rank")?;
+                head.forward(&embeds, &cur, None, &cos, &sin)
+            })();
+            let out = match out {
+                Ok(t) => t,
+                Err(e) => {
+                    err = Some(e);
+                    break;
+                }
+            };
+            let next = match self.lm_head.forward(&out).and_then(|logits| {
+                logits
+                    .i((0, 0, ..))?
+                    .to_dtype(DType::F32)?
+                    .argmax(candle_core::D::Minus1)?
+                    .to_scalar::<u32>()
+            }) {
+                Ok(t) => t,
+                Err(e) => {
+                    err = Some(e.into());
+                    break;
+                }
+            };
+            drafted.push(next);
+            token = next;
+            cur = out;
+        }
+
+        // Restore even when a step failed: a half-walked head is worse
+        // than no draft, because every later round inherits it.
+        if let Some(head) = self.mtp.as_mut() {
+            head.restore_kv(&snap)?;
+        }
+        match err {
+            Some(e) => Err(e),
+            None => Ok(drafted),
+        }
+    }
+
+    /// Walk the head over one token the target actually committed, so
+    /// its cache stays level with the target's.
+    pub fn mtp_advance(&mut self, token: u32, pos: usize) -> Result<()> {
+        let hidden = self
+            .last_hidden
+            .as_ref()
+            .context("mtp_advance: no hidden state retained")?
+            .clone();
+        let (_, have, _) = hidden.dims3()?;
+        let hidden = if have > 1 {
+            hidden.i((.., have - 1.., ..))?
+        } else {
+            hidden
+        };
+        let device = self.base.device.clone();
+        let rotary = Arc::clone(&self.base.rotary);
+        let ids = Tensor::new(&[token], &device)?.unsqueeze(0)?;
+        let embeds = self.base.embed_tokens.forward(&ids)?;
+        let (cos, sin) = rotary.cos_sin_at(&[pos])?;
+        let head = self
+            .mtp
+            .as_mut()
+            .context("mtp_advance: no draft head on this rank")?;
+        head.forward(&embeds, &hidden, None, &cos, &sin)?;
+        Ok(())
+    }
+
+    pub fn forward_multi(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
+        let hidden = self.base.forward(input, offset)?;
+        self.last_hidden = Some(hidden.clone());
+        hidden.apply(&self.lm_head)
+    }
+
     /// Lockstep batched decode step (#98): `(B, 1)` input, per-row
     /// positions, padding mask from
     /// [`TpQwen3_5Model::batch_decode_mask`]. Returns `(B, 1, vocab)`.

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

@@ -52,6 +52,21 @@ impl WorkerModel {
         }
     }
 
+    /// Logits at every position (#96). This rank discards them; what
+    /// matters is the identical collective sequence.
+    fn forward_multi(
+        &mut self,
+        input: &candle_core::Tensor,
+        offset: usize,
+    ) -> candle_core::Result<candle_core::Tensor> {
+        match self {
+            WorkerModel::Qwen3_5(m) => m.forward_multi(input, offset),
+            WorkerModel::Qwen3(_) => Err(candle_core::Error::Msg(
+                "forward_multi: the tp qwen3 arch has no multi-position logits path".into(),
+            )),
+        }
+    }
+
     /// Chunked image prefill on this rank. Only the vision-capable
     /// `qwen3_5` arch has a replicated tower; the dense `qwen3` arch
     /// errors. The returned logits are discarded by the caller (the
@@ -269,6 +284,11 @@ impl WorkerState {
                 tokens,
                 offset,
             } => self.handle_generate_step(&model_id, tokens, offset),
+            WorkerRequest::GenerateStepMulti {
+                model_id,
+                tokens,
+                offset,
+            } => self.handle_generate_step_multi(&model_id, tokens, offset),
             WorkerRequest::GenerateStepWithImages {
                 model_id,
                 tokens,
@@ -594,6 +614,80 @@ impl WorkerState {
         }
     }
 
+    /// Multi-position forward (#96). Identical collectives to
+    /// `handle_generate_step` — the only difference is that the model
+    /// keeps every position's logits instead of the last. This rank
+    /// discards them either way; what matters is that it issues the
+    /// same `AllReduce` sequence, or the leader's forward never
+    /// completes.
+    #[cfg(feature = "cuda")]
+    fn handle_generate_step_multi(
+        &mut self,
+        model_id: &str,
+        tokens: Vec<u32>,
+        offset: usize,
+    ) -> WorkerResponse {
+        use candle_core::Tensor;
+
+        let Some(model) = self.models.get_mut(model_id) else {
+            return WorkerResponse::Error {
+                kind: "model_not_loaded".into(),
+                message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
+            };
+        };
+        let device = model.device().clone();
+        let input = match Tensor::new(tokens.as_slice(), &device).and_then(|t| t.unsqueeze(0)) {
+            Ok(t) => t,
+            Err(e) => {
+                return WorkerResponse::Error {
+                    kind: "forward_failed".into(),
+                    message: format!("build input tensor: {e}"),
+                };
+            }
+        };
+        let start = std::time::Instant::now();
+        tracing::debug!(
+            rank = self.config.rank,
+            model = %model_id,
+            tokens = tokens.len(),
+            offset,
+            "worker GenerateStepMulti: forward starting"
+        );
+        if let Err(e) = model.forward_multi(&input, offset) {
+            tracing::warn!(
+                rank = self.config.rank,
+                model = %model_id,
+                elapsed_ms = start.elapsed().as_millis(),
+                error = %e,
+                "worker GenerateStepMulti: forward failed"
+            );
+            return WorkerResponse::Error {
+                kind: "forward_failed".into(),
+                message: format!("TP multi-position forward: {e}"),
+            };
+        }
+        tracing::debug!(
+            rank = self.config.rank,
+            model = %model_id,
+            elapsed_ms = start.elapsed().as_millis(),
+            "worker GenerateStepMulti: forward done"
+        );
+        WorkerResponse::GenerateStepOk
+    }
+
+    #[cfg(not(feature = "cuda"))]
+    fn handle_generate_step_multi(
+        &mut self,
+        _model_id: &str,
+        _tokens: Vec<u32>,
+        _offset: usize,
+    ) -> WorkerResponse {
+        WorkerResponse::Error {
+            kind: "cuda_feature_not_enabled".into(),
+            message: "GenerateStepMulti requires --features cuda".into(),
+        }
+    }
+
     /// One lockstep batched decode step on this rank (#98). Mirrors
     /// the leader's `forward_logits_batch`: derives per-row positions
     /// and the padding mask locally from the broadcast geometry, runs

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

@@ -158,6 +158,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
         PreflightError::EmptyRepo { .. } => "empty_repo",
         PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
         PreflightError::QuantNotFound { .. } => "quant_not_found",
+        PreflightError::Unsupported { .. } => "unsupported_load_option",
     }
 }
 

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/src/wire/openai_responses.rs

@@ -219,9 +219,6 @@ pub fn request_to_chat(req: ResponsesRequest) -> Result<ChatCompletionRequest, T
         }
     }
 
-    // Only the newest turn's thinking is context; the rest is noise.
-    prune_stale_reasoning(&mut messages);
-
     // Carry the caller's extension fields across the hop (#277).
     //
     // This used to start from an empty map and insert exactly one key,
@@ -455,54 +452,6 @@ fn merge_tool_call_into(prev: &mut ChatMessage, msg: &mut ChatMessage) -> bool {
     true
 }
 
-/// Keep replayed reasoning on the most recent assistant turn only.
-///
-/// #277 restored the round-trip, and then over-corrected: a client that
-/// replays every turn verbatim — which is what an agentic harness does —
-/// had the model's entire thinking history rendered back into the
-/// prompt, because Qwen3.8's template defaults `preserve_thinking` to
-/// true. Measured on a live session: **62% of the assistant content
-/// being replayed was prior think blocks** (95,311 characters of
-/// thinking against 467 characters of actual output), and the model's
-/// per-turn reasoning collapsed to tens of tokens while it made
-/// increasingly shallow passes at the task.
-///
-/// Reasoning models are trained on histories that carry *final content*,
-/// not the deliberation behind it — Qwen's own multi-turn convention
-/// strips prior think blocks. But the newest block is different: when a
-/// turn was cut off mid-thought and the caller says "continue", that
-/// thinking is the work being resumed, which is what #277 exists to
-/// preserve.
-///
-/// So: newest assistant turn keeps its reasoning, everything older
-/// loses it. Note the template's own `preserve_thinking: false` is
-/// exactly backwards for this — it keeps thinking on turns *after* the
-/// last user message (the agentic case, where it is noise) and drops it
-/// on turns *before* (the continuation case, where it is the point).
-fn prune_stale_reasoning(messages: &mut [ChatMessage]) {
-    let newest = messages
-        .iter()
-        .rposition(|m| m.role == "assistant" && m.extra.get("reasoning_content").is_some());
-    let Some(newest) = newest else { return };
-    let mut dropped = 0usize;
-    for (i, m) in messages.iter_mut().enumerate() {
-        if i == newest || m.role != "assistant" {
-            continue;
-        }
-        if let Some(obj) = m.extra.as_object_mut()
-            && obj.remove("reasoning_content").is_some()
-        {
-            dropped += 1;
-        }
-    }
-    if dropped > 0 {
-        tracing::debug!(
-            dropped,
-            "responses: replaying reasoning for the newest assistant turn only"
-        );
-    }
-}
-
 /// The text of a replayed reasoning item, under either spelling.
 ///
 /// `summary` is what neuron emits (`summary_text` parts) and `content`
@@ -1934,14 +1883,28 @@ mod tests {
         );
     }
 
-    /// An agentic harness replays every turn verbatim. Rendering all of
-    /// that thinking back into the prompt made the context mostly
-    /// self-quotation — 62% of replayed assistant content on a live
-    /// session — and the model started making shallow passes. Only the
-    /// newest turn's reasoning is context; the rest is noise the model
-    /// was never trained to read.
+    /// Prior-turn reasoning reaches the model; the model's own chat
+    /// template decides what to do with it (#329).
+    ///
+    /// The translator used to keep reasoning on the newest assistant
+    /// turn and delete it from the rest. That was a third policy, owned
+    /// by us, that no model vendor implements: every vendor's rule is
+    /// tool-loop aware (keep thinking within the current tool-calling
+    /// episode, drop it across user turns), which is exactly what
+    /// Qwen's `last_query_index` does and what Anthropic requires
+    /// inside a tool-use turn.
+    ///
+    /// It also made the prompt non-append-only: the block the model
+    /// generated into the KV cache was deleted on the next turn, so
+    /// turn N's token sequence was no longer a prefix of turn N+1's.
+    /// Measured across the change on a live fleet, prefix-cache reuse
+    /// on this surface went 53.9% -> 0.5%.
+    ///
+    /// Policy now lives where it is expressible: the template, through
+    /// `chat_template_kwargs.preserve_thinking`, with precedence
+    /// request > operator > the template's own default.
     #[test]
-    fn only_the_newest_turns_reasoning_is_replayed() {
+    fn prior_turn_reasoning_is_left_for_the_template_to_judge() {
         let raw = r#"{
             "model": "Qwen/Qwen3.8-27B",
             "input": [
@@ -1962,10 +1925,9 @@ mod tests {
             .collect();
         assert_eq!(
             carried,
-            ["newer plan"],
-            "older turns must keep their content but lose their thinking"
+            ["old plan", "newer plan"],
+            "every replayed turn's thinking must survive translation"
         );
-        // The turns themselves survive — only the thinking is pruned.
         assert_eq!(
             chat.messages
                 .iter()
@@ -1975,8 +1937,29 @@ mod tests {
         );
     }
 
-    /// The continuation case still works: a turn cut off mid-thought is
-    /// the newest assistant turn, so its reasoning is what survives.
+    /// The caller's own policy still reaches the template unaltered —
+    /// this is the toggle #329 specifies, and the one a client can
+    /// actually set (pi sends it verbatim via `samplingParams`).
+    #[test]
+    fn a_callers_preserve_thinking_reaches_the_template_kwargs() {
+        let raw = r#"{
+            "model": "Qwen/Qwen3.8-27B",
+            "input": [{"type": "message", "role": "user", "content": "hi"}],
+            "chat_template_kwargs": {"preserve_thinking": false}
+        }"#;
+        let req: ResponsesRequest = serde_json::from_str(raw).expect("parse");
+        let chat = request_to_chat(req).expect("translate");
+        assert_eq!(
+            chat.extra
+                .get("chat_template_kwargs")
+                .and_then(|k| k.get("preserve_thinking")),
+            Some(&serde_json::Value::Bool(false))
+        );
+    }
+
+    /// The continuation case (#277): a turn cut off mid-thought is
+    /// replayed as reasoning with no assistant turn after it, and that
+    /// thinking — the work being resumed — must reach the model.
     #[test]
     fn a_truncated_turn_keeps_its_reasoning_through_a_continue() {
         let raw = r#"{

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

@@ -415,11 +415,15 @@ fn qwen4_exp_case(name: &str, max_abs_bound: f32) {
     let mut model = Qwen4ExpForCausalLM::load(
         &cfg,
         DType::F32,
-        None,
         &device,
         &vb,
-        std::slice::from_ref(&weights),
-        None,
+        &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");
 
@@ -428,7 +432,11 @@ fn qwen4_exp_case(name: &str, max_abs_bound: f32) {
     let ids = expected["input_ids"].to_dtype(DType::U32).unwrap();
     let reference: Vec<f32> = expected["logits"].flatten_all().unwrap().to_vec1().unwrap();
 
-    let logits = model.forward(&ids, 0).expect("forward");
+    // 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()
@@ -465,4 +473,46 @@ fn qwen4_exp_case(name: &str, max_abs_bound: f32) {
         "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,
     }
 }
 

helexa.ai/content/docs/2-operating/03-neurons.md

@@ -358,6 +358,20 @@ Whether narrowing replay *helps* the model is unmeasured. Set it in both
 `/v1/models` and logged at load so a session can be attributed to the
 arm it ran under.
 
+A caller can override it per request with
+`chat_template_kwargs.preserve_thinking`, on `/v1/chat/completions` and
+`/v1/responses` alike — precedence is request > operator > template
+default.
+
+The setting also decides whether a multi-turn prompt is append-only, and
+therefore whether the prefix cache can work at all. Replay is
+byte-stable: the think block the model generated is re-rendered
+unchanged, so turn N's tokens stay a prefix of turn N+1's. Anything that
+deletes a prior turn's reasoning rewrites the middle of the prompt and
+invalidates the cache from that point. Measured on this fleet when the
+Responses path pruned prior reasoning unconditionally, reuse fell from
+53.9% to 0.5% and stayed there (#329).
+
 **`reasoning_budget.answer_reserve_tokens`** (default `4096`) holds back
 part of the caller's output budget for the answer:
 

script/generate_laya_fixture.py

@@ -0,0 +1,153 @@
+"""Generate a tiny Laya parity fixture from the upstream reference (#336).
+
+Builds the `laya` SDK's own `DecisionModel` (laya/common.py) around a
+tiny ModernBERT, then records its outputs for a right-padded batch.
+
+Writes, into --out:
+  rl_agent_config.json   the decision-model config, as neuron parses it
+  encoder/config.json    the ModernBERT config, saved by transformers
+  model.safetensors      the reference's own weights, float32
+  tokenizer.json         a placeholder word-level tokenizer (the fixture
+                         is scored on token ids; loading needs a file)
+  expected.json          the batch (ids, markers, question types) and the
+                         reference's per-row scorer logits and act softmax
+
+The shape is chosen so every mechanism is exercised at toy size:
+
+- four encoder layers with a global layer on each side of two sliding
+  ones, and a sliding window (local_attention = 8, so +-4 positions)
+  far shorter than the 40-token row, so a window bug changes the result;
+- two different rotary bases for global and sliding layers, so using
+  the wrong one for either kind is visible;
+- rows of different lengths, so padding is present and must be masked;
+- a one-option row, which takes the reference's single-marker branch
+  of the act features;
+- every question type, so the type embedding is indexed three ways.
+
+Weights are initialised well above transformers' default scale so
+attention is sharp rather than near-uniform, and norm weights and biases
+are perturbed away from their ones/zeros init: at init a missing norm
+and a present one are indistinguishable.
+
+Usage, in a venv holding the `laya` SDK and CPU torch:
+
+    python script/generate_laya_fixture.py \
+        --out crates/neuron/src/harness/testdata/laya/tiny
+"""
+
+import argparse
+import json
+import os
+
+import torch
+from safetensors.torch import save_file
+from transformers import ModernBertConfig, AutoModel
+
+from laya.common import DecisionModel
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--out", required=True)
+ap.add_argument("--seed", type=int, default=7)
+args = ap.parse_args()
+torch.manual_seed(args.seed)
+os.makedirs(os.path.join(args.out, "encoder"), exist_ok=True)
+
+ecfg = ModernBertConfig(
+    vocab_size=97,
+    hidden_size=32,
+    num_hidden_layers=4,
+    num_attention_heads=2,
+    intermediate_size=48,
+    max_position_embeddings=128,
+    local_attention=8,
+    global_attn_every_n_layers=3,
+    pad_token_id=0,
+    bos_token_id=1,
+    eos_token_id=2,
+    cls_token_id=1,
+    sep_token_id=2,
+    norm_eps=1e-5,
+    hidden_activation="gelu",
+    # Well above the 0.02 default: at 0.02 attention is near-uniform and
+    # a window or rotary bug moves the logits only in the sixth decimal.
+    initializer_range=0.3,
+    rope_parameters={
+        "full_attention": {"rope_theta": 1000.0, "rope_type": "default"},
+        "sliding_attention": {"rope_theta": 300.0, "rope_type": "default"},
+    },
+)
+encoder = AutoModel.from_config(ecfg, attn_implementation="sdpa")
+model = DecisionModel(encoder, head_layers=2, n_act=2).eval()
+
+with torch.no_grad():
+    for name, p in model.named_parameters():
+        if "norm" in name or name.startswith("scorer.0"):
+            p.add_(0.3 * torch.randn_like(p))
+        elif name.startswith("type_emb"):
+            p.mul_(3.0)
+
+ecfg.save_pretrained(os.path.join(args.out, "encoder"))
+rl_cfg = {
+    "encoder": "tiny-modernbert",
+    "head_layers": 2,
+    "max_len": 64,
+    "head_max_len": 24,
+    "act_costs": {"escalate": 0.5},
+    "amp_dtype": "bf16",
+    "temperature": [1.0, 1.0, 1.0],
+    "temperature_by_options": {},
+}
+with open(os.path.join(args.out, "rl_agent_config.json"), "w") as f:
+    json.dump(rl_cfg, f, indent=1)
+# A placeholder: the fixture is scored on token ids, but a checkpoint
+# directory is loaded with its tokenizer, so one must exist.
+with open(os.path.join(args.out, "tokenizer.json"), "w") as f:
+    json.dump({"version": "1.0", "truncation": None, "padding": None, "added_tokens": [],
+               "normalizer": None, "pre_tokenizer": {"type": "Whitespace"},
+               "post_processor": None, "decoder": None,
+               "model": {"type": "WordLevel", "vocab": {"[UNK]": 0, "a": 1, "b": 2},
+                         "unk_token": "[UNK]"}}, f)
+
+state = {k: v.detach().contiguous().float() for k, v in model.state_dict().items()}
+save_file(state, os.path.join(args.out, "model.safetensors"))
+
+# (length, marker positions, qtype). Ids avoid 0 (pad) inside a row.
+rows = [
+    (40, [3, 9, 17], 0),
+    (23, [4, 11], 1),
+    (9, [2], 2),
+    (31, [5, 7, 12, 20, 26], 0),
+]
+seqs = []
+for n, markers, qtype in rows:
+    ids = torch.randint(1, ecfg.vocab_size, (n,)).tolist()
+    seqs.append({"input_ids": ids, "markers": markers, "qtype": qtype})
+
+b = len(seqs)
+L = max(len(s["input_ids"]) for s in seqs)
+K = max(len(s["markers"]) for s in seqs)
+ids = torch.zeros((b, L), dtype=torch.long)
+att = torch.zeros((b, L), dtype=torch.long)
+mpos = torch.zeros((b, K), dtype=torch.long)
+mmask = torch.zeros((b, K), dtype=torch.bool)
+for i, s in enumerate(seqs):
+    n = len(s["input_ids"])
+    ids[i, :n] = torch.tensor(s["input_ids"])
+    att[i, :n] = 1
+    k = len(s["markers"])
+    mpos[i, :k] = torch.tensor(s["markers"])
+    mmask[i, :k] = True
+qtype = torch.tensor([s["qtype"] for s in seqs])
+
+with torch.no_grad():
+    logits, act = model(ids, att, mpos, mmask, qtype)
+act = torch.softmax(act.float(), -1)
+
+for i, s in enumerate(seqs):
+    k = len(s["markers"])
+    s["logits"] = logits[i, :k].tolist()
+    s["act"] = act[i].tolist()
+
+with open(os.path.join(args.out, "expected.json"), "w") as f:
+    json.dump({"seed": args.seed, "seqs": seqs}, f, indent=1)
+print("wrote", args.out, "params:", sum(v.numel() for v in state.values()))

script/laya-reference.py

@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+"""Record Laya reference outputs for neuron's parity tests (#334).
+
+Runs the upstream `laya` SDK (the reference implementation, not the
+`rl_common.py` shipped in the model repo, which has diverged) over a fixed
+corpus and writes, per case:
+
+  - the request exactly as a `/v1/systemone` client would send it
+  - per question: `input_ids`, `[MASK]` marker positions, question type and
+    the option-collapse stats, straight out of `Agent._encode_state`
+  - the raw scorer logits and act probabilities out of `Agent._forward`,
+    before any temperature is applied
+  - the final Jev-shaped response (`answers`, `usage`, `routing`)
+
+A separate section records `laya.lang.analyse` on a language corpus, for
+the router port.
+
+The reference runs on CPU in float32 so the numbers are deterministic and
+free of autocast noise; neuron's reduced-precision paths are compared
+against them with a tolerance.
+
+Usage (in a venv with `pip install laya==<pinned>` and CPU torch):
+
+    HF_ENDPOINT=... HF_TOKEN=... python script/laya-reference.py \
+        --out crates/neuron/src/harness/testdata/laya
+
+Point HF_ENDPOINT at whichever Hugging Face-compatible registry you fetch
+weights through. Python's huggingface_hub may additionally need
+SSL_CERT_FILE (for a private CA) and HF_HUB_DISABLE_XET=1 (for registries
+without Xet support).
+"""
+import argparse
+import json
+import os
+import sys
+
+import laya
+import numpy as np
+import torch
+from laya import Router
+from laya.agent import Agent
+from laya.lang import analyse
+from laya.serve import _resolve_model
+
+LONG_TEXT = (
+    "Hello team, following up on the invoice discrepancy from last quarter. "
+    "The amounts on INV-2291 and INV-2307 do not match the purchase order, "
+    "and our finance department needs a corrected copy before month end. "
+) * 12
+
+CASES = [
+    # --- the three primitives, default budgets
+    {
+        "name": "choice_basic",
+        "state": "I was charged twice this month, I want my money back",
+        "questions": {
+            "queue": {"type": "choice", "instructions": "Which team?",
+                      "criteria": {"billing": "billing and refunds", "tech": "login and app issues",
+                                   "other": "everything else"}},
+        },
+    },
+    {
+        "name": "score_basic",
+        "state": "This is the THIRD time I've asked. Fix it today or I'm cancelling.",
+        "questions": {
+            "urgency": {"type": "score", "instructions": "How urgent?",
+                        "criteria": ["calm", "firm", "angry", "furious"]},
+        },
+    },
+    {
+        "name": "noul_default_criteria",
+        "state": "Please refund the duplicate charge on my card.",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+    },
+    {
+        "name": "noul_custom_criteria_and_labels",
+        "state": "Can you reset my password? I'm locked out.",
+        "questions": {
+            "locked": {"type": "noul", "instructions": "Is the user locked out of their account?",
+                       "criteria": {"true": "the user cannot sign in", "false": "the user can sign in"},
+                       "labels": {"true": "locked", "false": "not locked"}},
+        },
+    },
+    # --- several questions in one request: the shared-state, one-forward path
+    {
+        "name": "multi_question",
+        "state": "I was charged twice this month, I want my money back",
+        "questions": {
+            "queue": {"type": "choice", "instructions": "Which team?",
+                      "criteria": {"billing": "billing and refunds", "tech": "login and app issues",
+                                   "other": "everything else"}},
+            "urgency": {"type": "score", "instructions": "How urgent?",
+                        "criteria": ["calm", "firm", "angry", "furious"]},
+            "refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
+            "tone": {"type": "choice", "instructions": "Tone?", "criteria": ["polite", "rude", "neutral"]},
+            "spam": {"type": "noul", "instructions": "Is this spam?"},
+        },
+    },
+    # --- every temperature bucket
+    {
+        "name": "choice_2",
+        "state": "The app crashes when I open settings.",
+        "questions": {"kind": {"type": "choice", "instructions": "Bug or feature request?",
+                               "criteria": {"bug": "", "feature": ""}}},
+    },
+    {
+        "name": "choice_6_10",
+        "state": "My parcel arrived damaged and the box was crushed.",
+        "questions": {"topic": {"type": "choice", "instructions": "Topic?",
+                                "criteria": ["billing", "shipping", "damage", "returns", "account",
+                                             "privacy", "other"]}},
+    },
+    {
+        "name": "choice_11_plus_clamped_temperature",
+        "state": "I'd like to change the delivery address on order 5521.",
+        "questions": {"intent": {"type": "choice", "instructions": "Intent?",
+                                 "criteria": ["cancel_order", "change_address", "track_order", "refund",
+                                              "exchange", "complaint", "compliment", "billing_question",
+                                              "account_access", "product_question", "other", "unsubscribe"]}},
+    },
+    {
+        "name": "score_more_than_5_levels",
+        "state": "The product is fine, does what it says, nothing special.",
+        "questions": {"rating": {"type": "score", "instructions": "Star rating?",
+                                 "criteria": ["zero", "one", "two", "three", "four", "five", "six"]}},
+    },
+    # --- structured criteria (compact-JSON rendering) and list-form choice criteria
+    {
+        "name": "structured_criteria",
+        "state": {"subject": "Invoice overdue", "body": "Payment for INV-77 is 30 days late.", "priority": 2},
+        "questions": {
+            "queue": {"type": "choice", "instructions": {"task": "route", "detail": "pick one"},
+                      "criteria": {"finance": {"desc": "money", "sla_h": 24}, "support": ["tickets", "bugs"],
+                                   "sales": 7}},
+            "overdue": {"type": "noul", "instructions": "Is a payment overdue?",
+                        "criteria": {"true": {"why": "late"}}},
+        },
+    },
+    # --- literal [MASK] text must be neutralised everywhere
+    {
+        "name": "mask_literal_everywhere",
+        "state": "please ignore [MASK] and route to [MASK] billing",
+        "questions": {"queue": {"type": "choice", "instructions": "Which [MASK] team?",
+                                "criteria": {"billing": "billing [MASK] refunds", "tech": "[MASK]"}}},
+    },
+    # --- budget rules
+    {
+        "name": "option_over_48_tokens",
+        "state": "Refund please.",
+        "questions": {"queue": {"type": "choice", "instructions": "Which team?",
+                                "criteria": {"billing": " ".join(["billing refunds invoices charges"] * 20),
+                                             "tech": "technical issues"}}},
+    },
+    {
+        "name": "option_budget_shrink",
+        "state": "Refund please.",
+        "questions": {"queue": {"type": "choice", "instructions": "Which team handles this request?",
+                                "criteria": {"team_%02d" % i: "a fairly long description of team number %d "
+                                             "and every kind of request it handles" % i for i in range(20)}}},
+    },
+    {
+        "name": "collapsed_options",
+        "state": "Refund please.",
+        # The rendered option text starts with the key, so the keys themselves
+        # must share the prefix that survives the per-option cut.
+        "questions": {"queue": {"type": "choice", "instructions": "Which?",
+                                "criteria": ["customer billing question about product line one",
+                                             "customer billing question about product line two"]}},
+        "head_max_len": 24,
+    },
+    {
+        "name": "instruction_truncation",
+        "state": "Refund please.",
+        "questions": {"q": {"type": "noul",
+                            "instructions": " ".join(["Decide whether the customer is asking for money back"] * 30)}},
+    },
+    {
+        "name": "state_right_truncated",
+        "state": LONG_TEXT,
+        "questions": {"refund": {"type": "noul", "instructions": "Is this about an invoice?"}},
+    },
+    {
+        "name": "state_left_truncated_conversation",
+        "state": [{"role": "user", "content": "turn %d: %s" % (i, LONG_TEXT[:200])} for i in range(12)]
+                 + [{"role": "user", "content": "actually, just cancel my subscription"}],
+        "questions": {"cancel": {"type": "noul", "instructions": "Does the user want to cancel?"}},
+    },
+    {
+        "name": "budget_overrides",
+        "state": LONG_TEXT,
+        "questions": {"queue": {"type": "choice", "instructions": "Which team?",
+                                "criteria": ["finance", "support", "sales"]}},
+        "max_len": 128,
+        "head_max_len": 64,
+    },
+    # --- non-ASCII inside a JSON state
+    {
+        "name": "json_state_non_ascii",
+        "state": {"from": "Zoë", "text": "Die Rechnung ist falsch — bitte korrigieren. 请尽快处理"},
+        "questions": {"queue": {"type": "choice", "instructions": "Which team?",
+                                "criteria": ["billing", "tech", "other"]}},
+    },
+    # --- routing to the multilingual checkpoint
+    {
+        "name": "ml_devanagari",
+        "state": "मुझे इस महीने दो बार शुल्क लिया गया, मुझे पैसे वापस चाहिए",
+        "questions": {"queue": {"type": "choice", "instructions": "Which team?",
+                                "criteria": {"billing": "billing and refunds", "tech": "login and app issues",
+                                             "other": "everything else"}}},
+    },
+    {
+        "name": "ml_khmer",
+        "state": "ខ្ញុំត្រូវបានគិតប្រាក់ពីរដងក្នុងខែនេះ ខ្ញុំចង់បានប្រាក់វិញ",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+    },
+    {
+        "name": "ml_cjk",
+        "state": "这个月我被扣了两次钱，我要退款",
+        "questions": {"urgency": {"type": "score", "instructions": "How urgent?",
+                                  "criteria": ["calm", "firm", "angry", "furious"]}},
+    },
+    {
+        "name": "ml_arabic",
+        "state": "تم خصم المبلغ مرتين هذا الشهر، أريد استرداد أموالي",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+    },
+    {
+        "name": "ml_latin_non_english",
+        "state": "Me han cobrado dos veces este mes y quiero que me devuelvan el dinero",
+        "questions": {"queue": {"type": "choice", "instructions": "Which team?",
+                                "criteria": ["billing", "tech", "other"]}},
+    },
+    {
+        "name": "ml_mixed_script",
+        "state": "Order #5521 — 注文をキャンセルしたいです. Please help.",
+        "questions": {"cancel": {"type": "noul", "instructions": "Does the user want to cancel?"}},
+    },
+    # --- explicit checkpoint selection through `model`
+    {
+        "name": "model_pinned_multilingual",
+        "state": "I was charged twice this month, I want my money back",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+        "model": "multilingual",
+    },
+    {
+        "name": "model_pinned_typed_decisions",
+        "state": "I was charged twice this month, I want my money back",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+        "model": "typed-decisions",
+    },
+    {
+        "name": "model_foreign_name_routes_by_default",
+        "state": "I was charged twice this month, I want my money back",
+        "questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
+        "model": "jev-1",
+    },
+]
+
+LANG_CASES = [
+    "I was charged twice this month, I want my money back",
+    "Me han cobrado dos veces este mes",
+    "Die Rechnung ist falsch, bitte korrigieren",
+    "Je voudrais annuler ma commande",
+    "मुझे पैसे वापस चाहिए",
+    "ខ្ញុំចង់បានប្រាក់វិញ",
+    "这个月我被扣了两次钱",
+    "注文をキャンセルしたいです",
+    "تم خصم المبلغ مرتين",
+    "Мне дважды списали деньги",
+    "Order #5521 — 注文をキャンセルしたいです. Please help.",
+    "ok",
+    "",
+    "12345 67890",
+    "café résumé naïve",
+    {"subject": "Invoice", "body": "Bitte zahlen Sie bis Freitag"},
+    ["hello", "bonjour", "hola"],
+]
+
+
+class Capture:
+    """Wraps the two Agent internals the fixtures need, per call."""
+
+    def __init__(self):
+        self.items = None
+        self.logits = None
+        self.act = None
+        self.agent = None
+        orig_encode, orig_forward = Agent._encode_state, Agent._forward
+        cap = self
+
+        def encode(agent, *a, **kw):
+            items = orig_encode(agent, *a, **kw)
+            cap.items, cap.agent = items, agent
+            return items
+
+        def forward(agent, b):
+            logits, act = orig_forward(agent, b)
+            cap.logits, cap.act = logits, act
+            return logits, act
+
+        Agent._encode_state, Agent._forward = encode, forward
+
+    def reset(self):
+        self.items = self.logits = self.act = self.agent = None
+
+
+def to_py(x):
+    if isinstance(x, np.ndarray):
+        return x.tolist()
+    if isinstance(x, (np.floating, np.integer)):
+        return x.item()
+    if isinstance(x, dict):
+        return {k: to_py(v) for k, v in x.items()}
+    if isinstance(x, (list, tuple)):
+        return [to_py(v) for v in x]
+    return x
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", required=True)
+    args = ap.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+    torch.manual_seed(0)
+    torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))
+
+    cap = Capture()
+    router = Router(device="cpu", preload=True, max_loaded=3)
+    header = {
+        "laya_version": laya.__version__,
+        "torch_version": torch.__version__,
+        "device": "cpu",
+        "dtype": "float32",
+    }
+
+    cases, checkpoints = [], {}
+    for case in CASES:
+        req = {"state": case["state"], "questions": case["questions"]}
+        kw = {}
+        for key in ("model", "max_len", "head_max_len"):
+            if key in case:
+                req[key] = case[key]
+                if key != "model":
+                    kw[key] = case[key]
+        cap.reset()
+        # The HTTP server, not the Router, maps a foreign model name (a Jev id)
+        # to "let the router choose"; go through the same function.
+        resp = router.predict(case["state"], case["questions"], model=_resolve_model(case.get("model")), **kw)
+        agent = cap.agent
+        rows = []
+        for j, (qid, it) in enumerate(zip(case["questions"], cap.items)):
+            k = len(it["markers"])
+            rows.append({
+                "qid": qid,
+                "input_ids": list(it["ids"]),
+                "markers": list(it["markers"]),
+                "qtype": it["qtype"],
+                "options": it.get("options"),
+                "logits": cap.logits[j, :k].tolist(),
+                "act": cap.act[j].tolist(),
+            })
+        ckpt = resp["routing"]["model"]
+        if ckpt not in checkpoints:
+            checkpoints[ckpt] = {
+                "repo": resp["routing"].get("repo"),
+                "revision": getattr(agent, "revision", None),
+                "cfg": {k: agent.cfg.get(k) for k in ("max_len", "head_max_len", "temperature",
+                                                       "temperature_by_options")},
+                "special_ids": {"cls": agent.tok.cls_token_id, "sep": agent.tok.sep_token_id,
+                                "mask": agent.tok.mask_token_id, "pad": agent.tok.pad_token_id},
+            }
+        cases.append({"name": case["name"], "request": req, "items": rows, "response": to_py(resp)})
+        print("%-40s -> %-16s %d question(s)" % (case["name"], ckpt, len(rows)), file=sys.stderr)
+
+    lang = [{"state": s, "analysis": to_py(analyse(s))} for s in LANG_CASES]
+
+    header["checkpoints"] = checkpoints
+    with open(os.path.join(args.out, "reference.json"), "w") as f:
+        json.dump({"header": header, "cases": cases}, f, ensure_ascii=False, indent=1)
+    with open(os.path.join(args.out, "lang.json"), "w") as f:
+        json.dump({"header": header, "cases": lang}, f, ensure_ascii=False, indent=1)
+
+
+if __name__ == "__main__":
+    main()