fix(qwen4_exp): 640 is not a whole number of k-quant blocks

The first real cold load on beast got through routing, resolved all
131 shards, parsed the config, reached decoder layer 0 — and died
there:

    quantize experts to Q4K: quantized tensor must have their last dim
    divisible by block size [2560, 640] 256

ggml quantises along the last dimension and a row must be a whole
number of blocks. `gate_up` reduces over `hidden_size` = 2560, which
is ten blocks of 256. `down_proj` reduces over `moe_intermediate_size`
= 640, which is two and a half. So `quant = "q4k"` could never have
loaded this model, on any card, whatever else was fixed first.

It is not an orientation mistake — `[2560, 640]` is the `[out, in]`
layout a `QMatMul` wants, and 640 is genuinely the reduction. The
constraint is the number itself. llama.cpp meets the same wall and
answers it the same way: substitute a 32-block type for the row that
will not divide. Q4_0 and Q4K are both 4.5 bpw, so the memory
arithmetic every placement decision rests on is unchanged; the trade
is quality, on one of the three expert matrices, against not loading
at all. A row that divides neither 256 nor 32 is an error, because
silently leaving those weights dense would be a memory surprise rather
than a mercy.

The substitution is decided once per layer from the shapes rather than
per expert -- 512 identical lines a layer would bury the one fact
worth reading.

The test uses the checkpoint's own widths, 2560 and 640, scaled down
only in expert count. The fixture that missed this was 256 in both
dimensions -- the narrowest a k-quant accepts -- so every row divided
and the two halves were indistinguishable. A test whose fixture cannot
express the failure is not evidence about it.

Refs #307, #315, #320.

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

.gitea/workflows/build-prerelease.yml

@@ -564,6 +564,25 @@ jobs:
       # escalating to a slower uncached rebuild. The cache covers the
       # ~600-crate host-side dep tree (the bulk of the 10-14 min build),
       # shared across all three flavours, so even one run seeds the next.
+      # candle-flash-attn's build script clones CUTLASS from GitHub, and the
+      # runners are one VM per job so that cache starts empty every time.
+      # A burst of concurrent CUDA jobs is enough to start getting refused
+      # ("could not read Username for 'https://github.com'"), which fails a
+      # long job on a transient rate limit. Seed the cache first, with the
+      # retries the build script does not have. Both commands are
+      # best-effort: on failure the build proceeds exactly as it would have
+      # and reports the authoritative error itself.
+      #
+      # Only the flash-attn flavour pulls in candle-flash-attn, and only
+      # candle-flash-attn wants CUTLASS. Guard on that: unguarded, all
+      # three flavours raced for the same checkout at the same moment and
+      # two of them never used it, tripling a burst of exactly the request
+      # that gets refused.
+      - name: Seed the CUTLASS checkout
+        if: contains(matrix.cargo_features, 'flash-attn')
+        run: |
+          cargo fetch || true
+          script/seed-cutlass.sh
       - name: Build neuron with CUDA (${{ matrix.flavour }})
         run: |
           export PATH="${{ matrix.cuda_home }}/bin:${PATH}"
@@ -579,6 +598,15 @@ jobs:
           CUDA_COMPUTE_CAP: ${{ matrix.compute_cap }}
           CARGO_BUILD_JOBS: ${{ matrix.build_jobs }}
           NVCC_THREADS: ${{ matrix.nvcc_threads }}
+          # cudaforge shells out to `git clone` for CUTLASS if the seed
+          # step above did not populate the cache. Give that clone the
+          # same anonymous git environment the seed script uses: CUTLASS
+          # is public, and an inherited credential turns a 200 into a
+          # 401. Scoped to this step — actions/checkout writes global
+          # config and must not see these.
+          GIT_CONFIG_GLOBAL: /dev/null
+          GIT_CONFIG_SYSTEM: /dev/null
+          GIT_CONFIG_COUNT: "0"
 
       - name: Stage binary
         run: |

.gitea/workflows/ci.yml

@@ -290,6 +290,18 @@ jobs:
       # script (see build-prerelease.yml's neuron build for the same
       # pattern). It probes for sccache and, on a rustc SIGSEGV / OOM,
       # keeps the cache and fails fast rather than rebuilding uncached.
+      # candle-flash-attn's build script clones CUTLASS from GitHub, and the
+      # runners are one VM per job so that cache starts empty every time.
+      # A burst of concurrent CUDA jobs is enough to start getting refused
+      # ("could not read Username for 'https://github.com'"), which fails a
+      # long job on a transient rate limit. Seed the cache first, with the
+      # retries the build script does not have. Both commands are
+      # best-effort: on failure the build proceeds exactly as it would have
+      # and reports the authoritative error itself.
+      - name: Seed the CUTLASS checkout
+        run: |
+          cargo fetch || true
+          script/seed-cutlass.sh
       - name: cargo check --features cuda (sccache escalation)
         run: |
           # act launches the step shell without /etc/profile, so the
@@ -302,6 +314,16 @@ jobs:
           export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
           export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
           script/ci-cargo-escalate.sh cargo check -p neuron --features cuda,flash-attn --all-targets
+        env:
+          # cudaforge shells out to `git clone` for CUTLASS if the seed
+          # step above did not populate the cache. Give that clone the
+          # same anonymous git environment the seed script uses: CUTLASS
+          # is public, and an inherited credential turns a 200 into a
+          # 401. Scoped to this step — actions/checkout writes global
+          # config and must not see these.
+          GIT_CONFIG_GLOBAL: /dev/null
+          GIT_CONFIG_SYSTEM: /dev/null
+          GIT_CONFIG_COUNT: "0"
 
   srpm-cortex:
     name: Build cortex SRPM

.gitea/workflows/deploy.yml

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

CLAUDE.md

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

asset/sudoers.d/neuron-host.conf

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

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

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

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

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

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

@@ -43,6 +43,62 @@ fn flash_attn_enabled() -> bool {
     })
 }
 
+/// Which mask an attention call carries — and whether the flash kernel
+/// can express it.
+///
+/// The flash path encodes causality as a boolean flag and never reads a
+/// mask tensor. While the only mask this crate built was the causal
+/// one, `Option<&Tensor>` was enough: `Some` meant causal. It is no
+/// longer, because `qwen4_exp`'s QSA selection is an arbitrary additive
+/// mask — and handed to the flash path it would be silently replaced by
+/// a dense causal one. The model would run, the sparsity would do
+/// nothing, and the only symptom would be that long context cost what
+/// it always did.
+///
+/// So the distinction is a type rather than a comment: an
+/// [`AttnMask::Additive`] mask always takes the eager path.
+#[derive(Clone, Copy)]
+pub(crate) enum AttnMask<'a> {
+    /// Decode, or a single position — nothing to mask.
+    None,
+    /// The standard causal mask, which the flash kernel applies as a
+    /// flag rather than reading.
+    Causal(&'a Tensor),
+    /// A mask the kernel cannot express. Forces the eager path.
+    Additive(&'a Tensor),
+}
+
+impl<'a> AttnMask<'a> {
+    /// The historical `Option<&Tensor>` contract: present means causal.
+    pub(crate) fn causal(mask: Option<&'a Tensor>) -> Self {
+        match mask {
+            Some(m) => Self::Causal(m),
+            None => Self::None,
+        }
+    }
+
+    /// The tensor the eager path adds to the scores, if any.
+    fn tensor(&self) -> Option<&'a Tensor> {
+        match self {
+            Self::None => None,
+            Self::Causal(m) | Self::Additive(m) => Some(m),
+        }
+    }
+
+    /// Whether the flash kernel may be used at all for this mask.
+    /// Only the flash path asks, so only that build compiles it.
+    #[cfg(feature = "flash-attn")]
+    fn reaches_flash(&self) -> bool {
+        !matches!(self, Self::Additive(_))
+    }
+
+    /// The kernel's `causal` flag.
+    #[cfg(feature = "flash-attn")]
+    fn is_causal(&self) -> bool {
+        matches!(self, Self::Causal(_))
+    }
+}
+
 /// Attention core shared by the single-GPU and TP full-attention
 /// layers (#95): `(B, H, L, D)` query and `(B, H_kv, S, D)` key/value
 /// (post-KV-cache, NOT GQA-repeated) → `(B, H, L, D)` context.
@@ -56,19 +112,18 @@ fn flash_attn_enabled() -> bool {
 /// prefill continuation needs: a chunk of L new queries against
 /// `offset + L` cached keys masks correctly.
 ///
-/// INVARIANT: `attn_mask` is either `None` (decode / single position)
-/// or the standard causal mask — the only mask the qwen3_5 forward
-/// constructs. The flash path encodes it as `causal = attn_mask
-/// .is_some()`; a future non-causal mask must extend this signature,
-/// not silently pass through.
+/// The mask arrives as an [`AttnMask`] rather than an `Option`, because
+/// the flash path cannot express an arbitrary one and would discard it
+/// in silence — see that type.
 ///
 /// Falls back to the eager matmul→softmax→matmul everywhere else
-/// (CPU, f32, feature off, or `NEURON_FLASH_ATTN=0`).
+/// (CPU, f32, feature off, `NEURON_FLASH_ATTN=0`, or a mask the kernel
+/// cannot express).
 pub(crate) fn attention_context(
     q: &Tensor,
     k: &Tensor,
     v: &Tensor,
-    attn_mask: Option<&Tensor>,
+    attn_mask: AttnMask<'_>,
     num_kv_groups: usize,
     scale: f64,
 ) -> candle_core::Result<Tensor> {
@@ -84,6 +139,7 @@ pub(crate) fn attention_context(
         // byte-identical either way, so this split is purely a
         // performance routing decision.
         if flash_attn_enabled()
+            && attn_mask.reaches_flash()
             && q.dim(2)? > 1
             && q.device().is_cuda()
             && (dtype == DType::F16 || dtype == DType::BF16)
@@ -92,7 +148,7 @@ pub(crate) fn attention_context(
             let qf = q.transpose(1, 2)?.contiguous()?;
             let kf = k.transpose(1, 2)?.contiguous()?;
             let vf = v.transpose(1, 2)?.contiguous()?;
-            let causal = attn_mask.is_some();
+            let causal = attn_mask.is_causal();
             let ctx = candle_flash_attn::flash_attn(&qf, &kf, &vf, scale as f32, causal)?;
             return ctx.transpose(1, 2)?.contiguous();
         }
@@ -102,7 +158,7 @@ pub(crate) fn attention_context(
     let k = repeat_kv(k.clone(), num_kv_groups)?.contiguous()?;
     let v = repeat_kv(v.clone(), num_kv_groups)?.contiguous()?;
     let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
-    if let Some(m) = attn_mask {
+    if let Some(m) = attn_mask.tensor() {
         scores = scores.broadcast_add(m)?;
     }
     let probs = candle_nn::ops::softmax_last_dim(&scores)?;
@@ -222,7 +278,14 @@ impl Qwen3_5Attention {
         // 5+6. Attention core — FlashAttention when available, eager
         // GQA-repeat + masked softmax otherwise (#95).
         let scale = 1.0_f64 / (self.head_dim as f64).sqrt();
-        let ctx = attention_context(&q, &k, &v, attn_mask, self.num_kv_groups, scale)?; // (B, H, L, D)
+        let ctx = attention_context(
+            &q,
+            &k,
+            &v,
+            AttnMask::causal(attn_mask),
+            self.num_kv_groups,
+            scale,
+        )?; // (B, H, L, D)
 
         // 7. Reshape back, apply the output gate, project.
         let ctx = ctx

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

@@ -67,7 +67,7 @@ use candle_nn::var_builder::ShardedVarBuilder;
 #[cfg(test)]
 use super::RopeParameters;
 use super::TextConfig;
-use super::rmsnorm::{Qwen3_5RmsNormGated, l2norm};
+use super::rmsnorm::{OutputGate, Qwen3_5RmsNormGated, l2norm};
 
 /// Per-rank, per-layer state for the linear-attention block.
 ///
@@ -112,15 +112,69 @@ pub struct GatedDeltaNet {
     // `clear_kv_cache`; otherwise the state persists across forwards
     // and the per-token offset advances naturally.
     state: GatedDeltaNetState,
+    /// Whether the recurrent state is carried in f32 between forwards
+    /// (#284). Per layer rather than per process, because it is a
+    /// property of the checkpoint for architectures that declare
+    /// `mamba_ssm_dtype`.
+    state_f32: bool,
+}
+
+/// Everything a [`GatedDeltaNet`] needs that is not in its weights.
+///
+/// Split out of `TextConfig` so an architecture that is not `qwen3_5`
+/// can build one of these layers without owning a `qwen3_5` config —
+/// the tensor names and layout are shared, only the config type and two
+/// settings differ. `qwen4_exp` reuses 36 of its 48 layers this way.
+#[derive(Debug, Clone, Copy)]
+pub struct GatedDeltaNetParams {
+    pub hidden_size: usize,
+    pub num_value_heads: usize,
+    pub num_key_heads: usize,
+    pub key_head_dim: usize,
+    pub value_head_dim: usize,
+    pub conv_kernel_size: usize,
+    pub rms_norm_eps: f64,
+    /// The output gate's nonlinearity. Qwen3.6 leaves `output_gate_type`
+    /// unset and gets SiLU; `qwen4_exp` states sigmoid, and inheriting
+    /// SiLU there is wrong on every linear-attention layer.
+    pub output_gate: OutputGate,
+    /// Carry the recurrent state in f32 between forwards (#284) rather
+    /// than round-tripping it through the activation dtype. For
+    /// `qwen3_5` this follows `NEURON_GDN_STATE_F32`; an architecture
+    /// whose checkpoint states `mamba_ssm_dtype` should follow the
+    /// checkpoint instead.
+    pub state_f32: bool,
+}
+
+impl GatedDeltaNetParams {
+    /// The `qwen3_5` reading, preserving today's behaviour exactly:
+    /// SiLU regardless of `hidden_act`, and the process-wide flag.
+    pub fn from_qwen3_5(cfg: &TextConfig) -> Self {
+        Self {
+            hidden_size: cfg.hidden_size,
+            num_value_heads: cfg.linear_num_value_heads,
+            num_key_heads: cfg.linear_num_key_heads,
+            key_head_dim: cfg.linear_key_head_dim,
+            value_head_dim: cfg.linear_value_head_dim,
+            conv_kernel_size: cfg.linear_conv_kernel_dim,
+            rms_norm_eps: cfg.rms_norm_eps,
+            output_gate: OutputGate::Silu,
+            state_f32: gdn_state_f32(),
+        }
+    }
 }
 
 impl GatedDeltaNet {
     pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
-        let num_v_heads = cfg.linear_num_value_heads;
-        let num_k_heads = cfg.linear_num_key_heads;
-        let head_k_dim = cfg.linear_key_head_dim;
-        let head_v_dim = cfg.linear_value_head_dim;
-        let conv_kernel_size = cfg.linear_conv_kernel_dim;
+        Self::load_with_params(GatedDeltaNetParams::from_qwen3_5(cfg), vb)
+    }
+
+    pub fn load_with_params(cfg: GatedDeltaNetParams, vb: &ShardedVarBuilder) -> Result<Self> {
+        let num_v_heads = cfg.num_value_heads;
+        let num_k_heads = cfg.num_key_heads;
+        let head_k_dim = cfg.key_head_dim;
+        let head_v_dim = cfg.value_head_dim;
+        let conv_kernel_size = cfg.conv_kernel_size;
 
         if num_v_heads == 0 || num_k_heads == 0 {
             anyhow::bail!(
@@ -192,7 +246,12 @@ impl GatedDeltaNet {
             .with_context(|| format!("load '{}/A_log'", vb.prefix()))?;
 
         // ----- Output gated RMSNorm (per-head_v_dim). -----
-        let norm = Qwen3_5RmsNormGated::load(&vb.pp("norm"), head_v_dim, cfg.rms_norm_eps)?;
+        let norm = Qwen3_5RmsNormGated::load_with_gate(
+            &vb.pp("norm"),
+            head_v_dim,
+            cfg.rms_norm_eps,
+            cfg.output_gate,
+        )?;
 
         Ok(Self {
             in_proj_qkv,
@@ -213,6 +272,7 @@ impl GatedDeltaNet {
             conv_dim,
             conv_kernel_size,
             state: GatedDeltaNetState::default(),
+            state_f32: cfg.state_f32,
         })
     }
 
@@ -392,7 +452,7 @@ impl GatedDeltaNet {
         // by this point, and prefix snapshots deep-copy
         // (`snapshot_state`) precisely so they never share storage with
         // a buffer a later forward will overwrite.
-        self.state.recurrent_state = Some(if gdn_state_f32() {
+        self.state.recurrent_state = Some(if self.state_f32 {
             new_state
         } else {
             new_state.to_dtype(dtype)?
@@ -498,7 +558,7 @@ fn chunked_prefill_enabled() -> bool {
 ///
 /// Costs ~2× the recurrent slab: +37.7 MB per GPU per sequence on the
 /// deployed 27B under TP=2, so +302 MB/GPU at `max_in_flight = 8`.
-fn gdn_state_f32() -> bool {
+pub(crate) fn gdn_state_f32() -> bool {
     static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
     *ENABLED.get_or_init(|| {
         std::env::var("NEURON_GDN_STATE_F32")
@@ -1394,6 +1454,7 @@ mod tests {
             conv_dim,
             conv_kernel_size: cfg.linear_conv_kernel_dim,
             state: GatedDeltaNetState::default(),
+            state_f32: gdn_state_f32(),
         };
 
         let x = Tensor::ones(&[b, l, cfg.hidden_size], dtype, &dev).unwrap();

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

@@ -17,7 +17,13 @@ pub struct Qwen3_5MLP {
 }
 
 impl Qwen3_5MLP {
-    /// Construct directly from pre-built projections (MoE-block tests).
+    /// Construct directly from pre-built projections.
+    ///
+    /// Test-only again: `qwen4_exp` briefly built its 512 experts this
+    /// way before the fused tensors stayed banked (`Experts::Banked`),
+    /// which removed the last caller outside tests. The differential
+    /// test that proves banking neutral still needs it, to express the
+    /// same weights in the layout it is being compared against.
     #[cfg(test)]
     pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
         Self {

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

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

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

@@ -29,18 +29,218 @@
 //! replaces the loop behind the same `forward` signature.
 
 use anyhow::{Context, Result};
-use candle_core::{DType, Module, Tensor};
+use candle_core::quantized::{GgmlDType, QMatMul, QTensor};
+use candle_core::{DType, IndexOp, Module, Tensor};
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
 use super::TextConfig;
 use super::mlp::Qwen3_5MLP;
 
+/// How a layer's routed experts are stored.
+///
+/// Two shapes, because the right one depends on whether the model fits.
+///
+/// `PerExpert` is the Qwen3-Next layout: one module per expert, sliced
+/// out at load. Fine when the weights are device-resident anyway.
+///
+/// `Banked` keeps the checkpoint's own fused tensors whole and takes a
+/// view per routed expert. `qwen4_exp` needs it: slicing its 512
+/// experts into modules copies **241.6 GB** at BF16 — 5.03 GB per
+/// layer — which fits in neither 64 GB of VRAM nor 123 GB of host RAM,
+/// at any quantisation. It is also the layout a host-resident bank with
+/// a device-side gather wants (#318), so the fused form is where that
+/// work starts rather than something it would have to undo.
+///
+/// Banking costs nothing per call: a slice along dim 0 of a contiguous
+/// tensor is contiguous, and `x.matmul(w.t())` is what `Linear` does
+/// anyway, so the views feed the same GEMMs without a copy.
+pub(crate) enum Experts {
+    PerExpert(Vec<Qwen3_5MLP>),
+    /// Quantised in situ at load (#315). One `(gate, up, down)` triple
+    /// per expert.
+    ///
+    /// This is what makes `qwen4_exp` loadable at all: 120.8 B routed
+    /// parameters are 241.6 GB at BF16 and 67.9 GB at q4k, against
+    /// beast's 64 GB of VRAM and ~110 GB of usable host RAM. Nothing
+    /// about the arithmetic works until they shrink.
+    ///
+    /// Quantising *from* the banked tensor is what keeps the peak
+    /// bounded: one layer's fused BF16 tensor is 5.03 GB, and it is
+    /// dropped once its 512 experts are quantised, so the full 241.6 GB
+    /// never exists at once.
+    Quantized(Vec<QuantizedExpert>),
+    Banked {
+        /// `(num_experts, 2 * intermediate, hidden)` — gate rows first,
+        /// then up. Reversing them computes `silu(up) * gate`, which is
+        /// a different function of the same weights (#312).
+        gate_up: Tensor,
+        /// `(num_experts, hidden, intermediate)`
+        down: Tensor,
+        intermediate: usize,
+    },
+}
+
+/// One expert's three projections, quantised.
+pub(crate) struct QuantizedExpert {
+    gate: QMatMul,
+    up: QMatMul,
+    down: QMatMul,
+}
+
+/// The quantisation type actually usable for rows of `row` elements.
+///
+/// ggml quantises along the last dimension, and a row must be a whole
+/// number of blocks. The k-quants (`Q2K`..`Q8K`) block on 256; the
+/// older types block on 32. `qwen4_exp`'s `down_proj` reduces over
+/// `moe_intermediate_size = 640`, which is 2.5 blocks of 256 — so a
+/// `quant = "q4k"` load fails on the very first layer with
+/// `quantized tensor must have their last dim divisible by block size
+/// [2560, 640] 256`, having read 131 shards to get there.
+///
+/// Substituting a 32-block type for that half is what llama.cpp does
+/// with the same constraint, and it costs almost nothing here: `Q4_0`
+/// and `Q4K` are both 4.5 bpw, so the memory arithmetic the placement
+/// decisions rest on is unchanged. The quality is a little worse per
+/// bit — `Q4_0` has one scale per 32 values where `Q4K` has a
+/// two-level hierarchy — which is a real trade, but it applies to one
+/// of the three expert matrices and the alternative is not loading.
+///
+/// Errors rather than falls back further if the row will not even
+/// divide 32: that is a config we do not understand, and silently
+/// leaving those weights dense would be a memory surprise, not a
+/// mercy.
+pub(crate) fn quant_for_row(requested: GgmlDType, row: usize) -> candle_core::Result<GgmlDType> {
+    if row.is_multiple_of(requested.block_size()) {
+        return Ok(requested);
+    }
+    let fallback = match requested {
+        GgmlDType::Q2K | GgmlDType::Q3K | GgmlDType::Q4K => GgmlDType::Q4_0,
+        GgmlDType::Q5K => GgmlDType::Q5_0,
+        GgmlDType::Q6K | GgmlDType::Q8K => GgmlDType::Q8_0,
+        // Already a 32-block type: there is nothing coarser to fall
+        // back to, so this is the error case below.
+        other => other,
+    };
+    if row.is_multiple_of(fallback.block_size()) {
+        return Ok(fallback);
+    }
+    candle_core::bail!(
+        "cannot quantise a row of {row} elements as {requested:?} (block \
+         {}) or {fallback:?} (block {}): the row divides neither",
+        requested.block_size(),
+        fallback.block_size(),
+    )
+}
+
+impl Experts {
+    pub(crate) fn len(&self) -> usize {
+        match self {
+            Experts::PerExpert(v) => v.len(),
+            Experts::Banked { gate_up, .. } => gate_up.dims()[0],
+            Experts::Quantized(v) => v.len(),
+        }
+    }
+
+    /// Quantise a banked pair into per-expert triples.
+    ///
+    /// Takes the fused tensors by reference and drops nothing itself —
+    /// the caller owns the peak, and should hold one layer at a time.
+    ///
+    /// The two halves may not end up the same dtype: see
+    /// [`quant_for_row`]. `qwen4_exp`'s `down_proj` reduces over 640,
+    /// which no k-quant can block.
+    pub(crate) fn quantize_banked(
+        gate_up: &Tensor,
+        down: &Tensor,
+        intermediate: usize,
+        dtype: GgmlDType,
+    ) -> candle_core::Result<Self> {
+        let n = gate_up.dims()[0];
+        // Decided once, from the shapes, rather than per expert —
+        // every expert in a layer has the same two row widths, and
+        // 512 identical log lines per layer would bury the fact.
+        let gate_up_dtype = quant_for_row(dtype, gate_up.dims()[2])?;
+        let down_dtype = quant_for_row(dtype, down.dims()[2])?;
+        if gate_up_dtype != dtype || down_dtype != dtype {
+            tracing::info!(
+                requested = ?dtype,
+                gate_up_row = gate_up.dims()[2],
+                gate_up_dtype = ?gate_up_dtype,
+                down_row = down.dims()[2],
+                down_dtype = ?down_dtype,
+                "expert quantisation: a row width does not divide the requested \
+                 block size; substituting a 32-block type for that half"
+            );
+        }
+        let mut out = Vec::with_capacity(n);
+        for e in 0..n {
+            let gu = gate_up.i(e)?;
+            // The source must own its storage outright.
+            // `QTensor::quantize` reads `src.storage()` — the whole
+            // backing buffer, ignoring shape and offset — so a view into
+            // the fused tensor hands it every expert's bytes while the
+            // shape claims one, and candle panics on a block-count
+            // mismatch inside `from_float`. Neither `contiguous()` nor
+            // `copy()` saves you here: a slice along dim 0 is already
+            // contiguous, so both are free to hand back something still
+            // sharing the parent's buffer. Round-tripping through a Vec
+            // is the part that actually allocates, and it is also where
+            // the f32 the quantiser wants comes from.
+            let q = |t: Tensor, dtype: GgmlDType| -> candle_core::Result<QMatMul> {
+                let dims = t.dims().to_vec();
+                let vals = t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?;
+                let owned = Tensor::from_vec(vals, dims, t.device())?;
+                QMatMul::from_qtensor(QTensor::quantize(&owned, dtype)?)
+            };
+            out.push(QuantizedExpert {
+                gate: q(gu.narrow(0, 0, intermediate)?, gate_up_dtype)?,
+                up: q(gu.narrow(0, intermediate, intermediate)?, gate_up_dtype)?,
+                down: q(down.i(e)?, down_dtype)?,
+            });
+        }
+        Ok(Experts::Quantized(out))
+    }
+
+    /// Expert `e` applied to the rows routed to it.
+    pub(crate) fn forward_one(&self, e: usize, xs: &Tensor) -> candle_core::Result<Tensor> {
+        match self {
+            Experts::PerExpert(v) => v[e].forward(xs),
+            Experts::Banked {
+                gate_up,
+                down,
+                intermediate,
+            } => {
+                let gu = gate_up.i(e)?;
+                let gate_w = gu.narrow(0, 0, *intermediate)?;
+                let up_w = gu.narrow(0, *intermediate, *intermediate)?;
+                let lhs = candle_nn::ops::silu(&xs.matmul(&gate_w.t()?)?)?;
+                let rhs = xs.matmul(&up_w.t()?)?;
+                (lhs * rhs)?.matmul(&down.i(e)?.t()?)
+            }
+            Experts::Quantized(v) => {
+                let q = &v[e];
+                // The quantised kernels accumulate in f32 whatever the
+                // activation dtype, so the cast is explicit here rather
+                // than implied — and the result comes back in the
+                // caller's dtype so the routing arithmetic upstream
+                // does not have to know which layout it got.
+                let dtype = xs.dtype();
+                let xs32 = xs.to_dtype(DType::F32)?;
+                let lhs = candle_nn::ops::silu(&q.gate.forward(&xs32)?)?;
+                let rhs = q.up.forward(&xs32)?;
+                q.down.forward(&(lhs * rhs)?)?.to_dtype(dtype)
+            }
+        }
+    }
+}
+
 pub struct Qwen3_5MoeBlock {
     /// Router: `(num_experts, hidden)`, checkpoint name `mlp.gate`.
     gate: Linear,
-    /// Routed experts, checkpoint names `mlp.experts.{i}.{gate,up,down}_proj`.
-    experts: Vec<Qwen3_5MLP>,
+    /// Routed experts. Per-expert modules for `qwen3_5`; the
+    /// checkpoint's fused tensors for `qwen4_exp` — see [`Experts`].
+    experts: Experts,
     /// Always-on expert, `mlp.shared_expert.*`. `None` when the config
     /// declares no shared expert (Qwen3-30B-A3B style).
     shared_expert: Option<Qwen3_5MLP>,
@@ -52,6 +252,32 @@ pub struct Qwen3_5MoeBlock {
 }
 
 impl Qwen3_5MoeBlock {
+    /// Assemble from parts.
+    ///
+    /// The routing arithmetic is identical across the checkpoints that
+    /// use this block; only the storage of the experts differs.
+    /// `qwen4_exp` ships them as fused 3D tensors and slices them
+    /// itself, so it builds the block this way rather than through
+    /// [`Self::load`], which expects per-expert modules and a
+    /// `qwen3_5` config.
+    pub(crate) fn from_parts(
+        gate: Linear,
+        experts: Experts,
+        shared_expert: Option<Qwen3_5MLP>,
+        shared_expert_gate: Option<Linear>,
+        num_experts_per_tok: usize,
+        norm_topk_prob: bool,
+    ) -> Self {
+        Self {
+            gate,
+            experts,
+            shared_expert,
+            shared_expert_gate,
+            num_experts_per_tok,
+            norm_topk_prob,
+        }
+    }
+
     pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
         anyhow::ensure!(
             cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && cfg.moe_intermediate_size > 0,
@@ -103,7 +329,7 @@ impl Qwen3_5MoeBlock {
 
         Ok(Self {
             gate,
-            experts,
+            experts: Experts::PerExpert(experts),
             shared_expert,
             shared_expert_gate,
             num_experts_per_tok: cfg.num_experts_per_tok,
@@ -175,13 +401,13 @@ impl Module for Qwen3_5MoeBlock {
         )?;
 
         let mut ys = xs_flat.zeros_like()?;
-        for (e, expert) in self.experts.iter().enumerate() {
+        for e in 0..self.experts.len() {
             if tokens_for[e].is_empty() {
                 continue;
             }
             let rows = Tensor::new(tokens_for[e].as_slice(), xs.device())?;
             let picked = xs_flat.index_select(&rows, 0)?;
-            let out = expert.forward(&picked)?;
+            let out = self.experts.forward_one(e, &picked)?;
             let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
                 .to_dtype(out.dtype())?
                 .reshape(((), 1))?;
@@ -226,7 +452,7 @@ mod tests {
 
         let block = Qwen3_5MoeBlock {
             gate: Linear::new(randn(&[n_exp, hidden]), None),
-            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            experts: Experts::PerExpert((0..n_exp).map(|_| rand_mlp(hidden, inter)).collect()),
             shared_expert: Some(rand_mlp(hidden, inter)),
             shared_expert_gate: Some(Linear::new(randn(&[1, hidden]), None)),
             num_experts_per_tok: top_k,
@@ -257,8 +483,9 @@ mod tests {
             let mut expect = vec![0f32; hidden];
             for &e in selected {
                 let w = probs[e] / denom;
-                let out: Vec<f32> = block.experts[e]
-                    .forward(&row)
+                let out: Vec<f32> = block
+                    .experts
+                    .forward_one(e, &row)
                     .unwrap()
                     .flatten_all()
                     .unwrap()
@@ -310,7 +537,7 @@ mod tests {
         let (hidden, inter, n_exp) = (4, 2, 3);
         let block = Qwen3_5MoeBlock {
             gate: Linear::new(randn(&[n_exp, hidden]), None),
-            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            experts: Experts::PerExpert((0..n_exp).map(|_| rand_mlp(hidden, inter)).collect()),
             shared_expert: None,
             shared_expert_gate: None,
             num_experts_per_tok: 1,
@@ -342,8 +569,9 @@ mod tests {
             .max_by(|&a, &b| exps[a].partial_cmp(&exps[b]).unwrap())
             .unwrap();
         let w = exps[best] / sum;
-        let out: Vec<f32> = block.experts[best]
-            .forward(&xs.reshape(((), hidden)).unwrap())
+        let out: Vec<f32> = block
+            .experts
+            .forward_one(best, &xs.reshape(((), hidden)).unwrap())
             .unwrap()
             .flatten_all()
             .unwrap()
@@ -357,4 +585,207 @@ mod tests {
             );
         }
     }
+
+    /// Banking is a storage change, so the two layouts must compute the
+    /// same function. Same weights, expressed both ways, asserted
+    /// identical — otherwise "we only changed where the bytes live" is
+    /// a claim rather than a fact.
+    #[test]
+    fn banked_experts_equal_per_expert_modules() {
+        let (n_experts, hidden, inter) = (3usize, 4usize, 2usize);
+
+        // One source of truth for the weights, laid out both ways.
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+
+        let per_expert: Vec<Qwen3_5MLP> = (0..n_experts)
+            .map(|e| {
+                let gu = gate_up.i(e).unwrap();
+                Qwen3_5MLP::from_weights(
+                    Linear::new(gu.narrow(0, 0, inter).unwrap().contiguous().unwrap(), None),
+                    Linear::new(
+                        gu.narrow(0, inter, inter).unwrap().contiguous().unwrap(),
+                        None,
+                    ),
+                    Linear::new(down.i(e).unwrap().contiguous().unwrap(), None),
+                )
+            })
+            .collect();
+
+        let sliced = Experts::PerExpert(per_expert);
+        let banked = Experts::Banked {
+            gate_up,
+            down,
+            intermediate: inter,
+        };
+        assert_eq!(sliced.len(), banked.len());
+
+        let xs = randn(&[5, hidden]);
+        for e in 0..n_experts {
+            let a: Vec<f32> = sliced
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            let b: Vec<f32> = banked
+                .forward_one(e, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            assert_eq!(a.len(), b.len());
+            for (x, y) in a.iter().zip(b.iter()) {
+                assert!(
+                    (x - y).abs() < 1e-5,
+                    "expert {e} differs between layouts: {a:?} vs {b:?}"
+                );
+            }
+        }
+    }
+
+    /// And the halves are not interchangeable: reading `gate_up` the
+    /// other way round computes silu(up) * gate, which the banked path
+    /// would do just as silently as the sliced one.
+    #[test]
+    fn banked_reads_the_gate_before_the_up() {
+        let dev = Device::Cpu;
+        // gate row picks channel 0, up row picks channel 1.
+        let gate_up = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (1, 2, 2), &dev).unwrap();
+        let down = Tensor::from_vec(vec![1.0f32, 0.0], (1, 2, 1), &dev).unwrap();
+        let banked = Experts::Banked {
+            gate_up,
+            down,
+            intermediate: 1,
+        };
+        let xs = Tensor::from_vec(vec![1.0f32, 2.0], (1, 2), &dev).unwrap();
+        let got: Vec<f32> = banked
+            .forward_one(0, &xs)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // silu(gate.x) * (up.x) = silu(1) * 2
+        let want = (1.0f32 / (1.0 + (-1.0f32).exp())) * 2.0;
+        assert!((got[0] - want).abs() < 1e-5, "got {got:?} want {want}");
+        // the swapped reading would be silu(2) * 1
+        let swapped = (2.0f32 / (1.0 + (-2.0f32).exp())) * 1.0;
+        assert!(
+            (got[0] - swapped).abs() > 1e-3,
+            "this is silu(up) * gate: {got:?}"
+        );
+    }
+
+    /// In-situ quantisation has to approximate the dense expert, not
+    /// merely produce numbers of the right shape.
+    ///
+    /// Dimensions are 256 because the k-quants block on 256 elements —
+    /// a tiny fixture silently cannot be quantised at all, which is the
+    /// first thing that bites when this is tried on a toy model.
+    ///
+    /// `qwen4_exp`'s real expert widths, which no k-quant can block.
+    ///
+    /// `moe_intermediate_size` is 640, so `down_proj` is `[2560, 640]`
+    /// and its 640-element rows are 2.5 blocks of 256. The load died
+    /// on beast at decoder layer 0 with `quantized tensor must have
+    /// their last dim divisible by block size [2560, 640] 256`, after
+    /// resolving all 131 shards.
+    ///
+    /// The fixture that missed it was 256 wide in both dimensions —
+    /// the smallest width a k-quant accepts — so every row divided
+    /// and the orientation could not matter. This one uses the
+    /// checkpoint's own numbers, scaled down only in expert count.
+    #[test]
+    fn the_real_expert_widths_quantise_despite_a_640_wide_reduction() {
+        assert_eq!(
+            quant_for_row(GgmlDType::Q4K, 2560).unwrap(),
+            GgmlDType::Q4K,
+            "gate_up reduces over hidden_size, which is 10 whole blocks"
+        );
+        assert_eq!(
+            quant_for_row(GgmlDType::Q4K, 640).unwrap(),
+            GgmlDType::Q4_0,
+            "down reduces over moe_intermediate_size, which is not"
+        );
+        assert_eq!(quant_for_row(GgmlDType::Q6K, 640).unwrap(), GgmlDType::Q8_0);
+        assert_eq!(
+            quant_for_row(GgmlDType::Q8_0, 640).unwrap(),
+            GgmlDType::Q8_0
+        );
+        // 100 divides neither 256 nor 32.
+        assert!(quant_for_row(GgmlDType::Q4K, 100).is_err());
+
+        // And the whole path runs at those widths.
+        let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let quantised = Experts::quantize_banked(&gate_up, &down, inter, GgmlDType::Q4K)
+            .expect("the real widths must quantise");
+        let xs = randn(&[2, hidden]);
+        let got = quantised.forward_one(0, &xs).expect("forward");
+        assert_eq!(got.dims(), &[2, hidden]);
+        assert!(
+            got.flatten_all()
+                .unwrap()
+                .to_vec1::<f32>()
+                .unwrap()
+                .iter()
+                .all(|v| v.is_finite()),
+            "a substituted dtype must still produce finite output"
+        );
+    }
+
+    /// The assertion is relative error against the dense output, with a
+    /// noise baseline: a q8_0 expert should be close, a q4k expert
+    /// looser, and both far nearer the truth than an unrelated expert
+    /// is. Without that last comparison a "quantised" path that
+    /// returned anything smoothly wrong would pass.
+    #[test]
+    fn quantised_experts_approximate_the_dense_ones() {
+        let (n_experts, hidden, inter) = (2usize, 256usize, 256usize);
+        let gate_up = randn(&[n_experts, inter * 2, hidden]);
+        let down = randn(&[n_experts, hidden, inter]);
+        let banked = Experts::Banked {
+            gate_up: gate_up.clone(),
+            down: down.clone(),
+            intermediate: inter,
+        };
+        let xs = randn(&[4, hidden]);
+
+        let rel_err = |a: &[f32], b: &[f32]| -> f32 {
+            let num: f32 = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum();
+            let den: f32 = a.iter().map(|x| x * x).sum::<f32>().max(1e-12);
+            (num / den).sqrt()
+        };
+        let out = |e: &Experts, i: usize| -> Vec<f32> {
+            e.forward_one(i, &xs)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap()
+        };
+
+        let dense0 = out(&banked, 0);
+        // Baseline: a different expert on the same input. Any quantised
+        // reading must beat this by a wide margin.
+        let noise = rel_err(&dense0, &out(&banked, 1));
+
+        for (dtype, tol) in [(GgmlDType::Q8_0, 0.05f32), (GgmlDType::Q4K, 0.20f32)] {
+            let q = Experts::quantize_banked(&gate_up, &down, inter, dtype).unwrap();
+            assert_eq!(q.len(), n_experts);
+            let err = rel_err(&dense0, &out(&q, 0));
+            assert!(err < tol, "{dtype:?} relative error {err:.4} exceeds {tol}");
+            assert!(
+                err < noise / 2.0,
+                "{dtype:?} error {err:.4} is not clearly better than an unrelated \
+                 expert ({noise:.4}) — the quantised path may not be reading \
+                 this expert at all"
+            );
+        }
+    }
 }

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

@@ -23,7 +23,7 @@
 use anyhow::Result;
 use candle_core::{DType, Device, IndexOp, Tensor};
 
-use super::TextConfig;
+use super::{RopeParameters, TextConfig};
 
 #[derive(Debug, Clone)]
 pub struct RotaryEmbedding {
@@ -97,8 +97,30 @@ fn mrope_masks(
 
 impl RotaryEmbedding {
     pub fn new(dtype: DType, cfg: &TextConfig, dev: &Device) -> Result<Self> {
-        let head_dim = cfg.head_dim;
-        let rope = &cfg.rope_parameters;
+        Self::from_parts(
+            dtype,
+            cfg.head_dim,
+            cfg.max_position_embeddings,
+            &cfg.rope_parameters,
+            dev,
+        )
+    }
+
+    /// Build from the three things the rotary actually reads, rather
+    /// than from a whole config.
+    ///
+    /// `qwen4_exp` declares an identical `rope_parameters` block and the
+    /// same `head_dim` (spec §7), so it reuses this module unchanged —
+    /// but it owns a different config type, and fabricating a foreign
+    /// one to pass through here would be a lie the compiler could not
+    /// catch.
+    pub fn from_parts(
+        dtype: DType,
+        head_dim: usize,
+        max_seq_len: usize,
+        rope: &RopeParameters,
+        dev: &Device,
+    ) -> Result<Self> {
         let rotary_dim = (head_dim as f32 * rope.partial_rotary_factor) as usize;
         if !rotary_dim.is_multiple_of(2) {
             anyhow::bail!(
@@ -113,7 +135,6 @@ impl RotaryEmbedding {
                 rope.partial_rotary_factor
             );
         }
-        let max_seq_len = cfg.max_position_embeddings;
         let inv_freq: Vec<f32> = (0..rotary_dim)
             .step_by(2)
             .map(|i| 1f32 / rope.rope_theta.powf(i as f64 / rotary_dim as f64) as f32)
@@ -165,10 +186,23 @@ impl RotaryEmbedding {
     /// row, one decode token per step. [`Self::apply_cos_sin`] detects
     /// the rank-3 shape and broadcasts per row instead of per position.
     pub fn batch_cos_sin(&self, positions: &[usize]) -> candle_core::Result<(Tensor, Tensor)> {
+        let (cos, sin) = self.cos_sin_at(positions)?;
+        Ok((cos.unsqueeze(1)?, sin.unsqueeze(1)?))
+    }
+
+    /// cos/sin gathered at arbitrary, possibly non-contiguous positions,
+    /// shape `(positions.len(), rotary_dim/2)` — the same per-position
+    /// rank-2 form [`Self::plain_cos_sin`] returns, so it feeds the
+    /// sequence axis rather than a batch axis.
+    ///
+    /// `qwen4_exp`'s QSA indexer needs this: its pooled block keys are
+    /// rotated at the *first* position of each block, which is every
+    /// fourth position rather than a contiguous run.
+    pub fn cos_sin_at(&self, positions: &[usize]) -> candle_core::Result<(Tensor, Tensor)> {
         let idx: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
         let idx = Tensor::from_vec(idx, positions.len(), self.cos.device())?;
-        let cos = self.cos.index_select(&idx, 0)?.unsqueeze(1)?;
-        let sin = self.sin.index_select(&idx, 0)?.unsqueeze(1)?;
+        let cos = self.cos.index_select(&idx, 0)?;
+        let sin = self.sin.index_select(&idx, 0)?;
         Ok((cos, sin))
     }
 
@@ -212,38 +246,65 @@ impl RotaryEmbedding {
     ) -> candle_core::Result<(Tensor, Tensor)> {
         let (_, _, _seq_len, head_dim_in) = q.dims4()?;
         debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
-        let per_row = cos.rank() == 3;
-        let rope = |x: &Tensor| -> candle_core::Result<Tensor> {
-            if per_row {
-                rope_per_row(x, cos, sin)
-            } else {
-                candle_nn::rotary_emb::rope_slow(x, cos, sin)
-            }
-        };
-        if self.rotary_dim == self.head_dim {
-            let q_embed = rope(&q.contiguous()?)?;
-            let k_embed = rope(&k.contiguous()?)?;
-            Ok((q_embed, k_embed))
-        } else {
-            // Partial rotation: narrow → rotate → cat the untouched tail.
-            let tail = self.head_dim - self.rotary_dim;
-            let q_rot = q
-                .narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
-                .contiguous()?;
-            let q_pass = q.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
-            let k_rot = k
-                .narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
-                .contiguous()?;
-            let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
-            let q_rotated = rope(&q_rot)?;
-            let k_rotated = rope(&k_rot)?;
-            let q_embed =
-                Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
-            let k_embed =
-                Tensor::cat(&[&k_rotated, &k_pass.contiguous()?], candle_core::D::Minus1)?;
-            Ok((q_embed, k_embed))
-        }
+        Ok((
+            apply_partial_rotary(q, cos, sin, self.rotary_dim)?,
+            apply_partial_rotary(k, cos, sin, self.rotary_dim)?,
+        ))
     }
+
+    /// How many of each head's dims rotate. Heads narrower than
+    /// `head_dim` share the same rotation — `qwen4_exp`'s QSA indexer
+    /// keys are 128 wide and still rotate 64 dims against the same
+    /// cos/sin.
+    pub fn rotary_dim(&self) -> usize {
+        self.rotary_dim
+    }
+
+    /// The head width `apply_cos_sin` expects.
+    pub fn head_dim(&self) -> usize {
+        self.head_dim
+    }
+}
+
+/// Rotate the first `rotary_dim` dims of `x` and pass the rest through.
+///
+/// `x` is `(B, H, L, width)` for any `width >= rotary_dim`, and
+/// `cos`/`sin` are `(L, rotary_dim/2)` — or `(B, L, rotary_dim/2)` for
+/// the batched-decode path (#98), dispatched on rank.
+///
+/// Split out of [`RotaryEmbedding::apply_cos_sin`] because the rotation
+/// is a property of the *frequencies*, not of the head it is applied
+/// to: `qwen4_exp`'s sparse-attention indexer rotates 64 of its 128
+/// indexer dims using the same cos/sin the 256-wide attention heads
+/// use, so it needs the rotation without the head-width assertion.
+pub fn apply_partial_rotary(
+    x: &Tensor,
+    cos: &Tensor,
+    sin: &Tensor,
+    rotary_dim: usize,
+) -> candle_core::Result<Tensor> {
+    let width = x.dim(candle_core::D::Minus1)?;
+    debug_assert!(
+        width >= rotary_dim,
+        "rotary_dim {rotary_dim} exceeds head width {width}"
+    );
+    let per_row = cos.rank() == 3;
+    let rope = |t: &Tensor| -> candle_core::Result<Tensor> {
+        if per_row {
+            rope_per_row(t, cos, sin)
+        } else {
+            candle_nn::rotary_emb::rope_slow(t, cos, sin)
+        }
+    };
+    if rotary_dim == width {
+        return rope(&x.contiguous()?);
+    }
+    // Partial rotation: narrow → rotate → cat the untouched tail.
+    let rot = x
+        .narrow(candle_core::D::Minus1, 0, rotary_dim)?
+        .contiguous()?;
+    let pass = x.narrow(candle_core::D::Minus1, rotary_dim, width - rotary_dim)?;
+    Tensor::cat(&[&rope(&rot)?, &pass.contiguous()?], candle_core::D::Minus1)
 }
 
 /// GLM rotate-half (same convention as candle's private

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

@@ -0,0 +1,508 @@
+//! `qwen4_exp`'s `config.json`, and the geometry derived from it.
+//!
+//! Read from the checkpoint's own file
+//! (`Qwen/Qwen3.8-Flash-Next`, revision
+//! `de4b8e4d43b917e7706784d8bb445c9af86a3540`), which nests everything
+//! the text path needs under `text_config` alongside a `vision_config`
+//! this module ignores until #314.
+//!
+//! Most of the vocabulary overlaps [`super::super::qwen3_5`] — the same
+//! hybrid `layer_types`, the same GatedDeltaNet dimensions, the same
+//! nested `rope_parameters` — so the reused blocks read their settings
+//! from the same shapes. What is new is declared here: the four
+//! residual streams, the hashed n-gram table, and the sparse-attention
+//! indexer.
+//!
+//! Three of these fields say something other than what they appear to:
+//!
+//! 1. **`ple_layer_ids` is one-indexed.** The shipped `[2]` means
+//!    zero-indexed layer **1**. Upstream computes
+//!    `ple_layer_ids.index(layer_idx + 1)`; reading it as a plain layer
+//!    number puts 28% of the model's parameters on the wrong layer, and
+//!    the model still speaks.
+//! 2. **`output_gate_type: sigmoid`.** Qwen3.6 leaves this unset and
+//!    falls back to `hidden_act`, which is SiLU. Inheriting that here
+//!    is wrong on all 36 linear-attention layers.
+//! 3. **`mamba_ssm_dtype: float32`.** We default the equivalent
+//!    GatedDeltaNet state to a bf16 round-trip (#284). This checkpoint
+//!    makes an explicit choice and a port must not silently override
+//!    it with ours.
+//!
+//! There is no `intermediate_size`: every layer's FFN is the MoE, so
+//! there is no dense width to declare. See
+//! `doc/qwen4_exp-port-spec.md`.
+
+use anyhow::{Context, Result, ensure};
+use serde::Deserialize;
+
+use crate::harness::arch::qwen3_5::RopeParameters;
+use crate::harness::arch::qwen3_5::rmsnorm::OutputGate;
+
+/// `model_type` for the multimodal wrapper.
+pub const MODEL_TYPE: &str = "qwen4_exp";
+/// `model_type` of the nested `text_config`.
+pub const TEXT_MODEL_TYPE: &str = "qwen4_exp_text";
+
+/// The top-level config. `vision_config` is deliberately not modelled —
+/// serde drops it, and #314 owns the tower.
+#[derive(Debug, Clone, Deserialize)]
+pub struct Config {
+    #[serde(default)]
+    pub model_type: String,
+    pub text_config: TextConfig,
+    #[serde(default)]
+    pub tie_word_embeddings: bool,
+    #[serde(default)]
+    pub image_token_id: Option<u32>,
+    #[serde(default)]
+    pub video_token_id: Option<u32>,
+}
+
+impl Config {
+    pub fn from_config_json(json: &str) -> Result<Self> {
+        let mut cfg: Self = serde_json::from_str(json).context("parse qwen4_exp config.json")?;
+        cfg.text_config.fill_layer_types()?;
+        cfg.text_config.validate()?;
+        Ok(cfg)
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct TextConfig {
+    pub vocab_size: usize,
+    pub hidden_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: usize,
+    pub head_dim: usize,
+    pub max_position_embeddings: usize,
+    pub rms_norm_eps: f64,
+    /// Reused verbatim from `qwen3_5` — the nesting and every field in
+    /// it are identical, which is what makes `qwen3_5/rope.rs` reusable
+    /// unchanged (spec §7).
+    pub rope_parameters: RopeParameters,
+    #[serde(default = "default_hidden_act")]
+    pub hidden_act: String,
+    #[serde(default)]
+    pub tie_word_embeddings: bool,
+    /// 248044 for this checkpoint — the id PLE's segment-aware shifts
+    /// fill with, so it is load-bearing beyond stopping generation.
+    pub eos_token_id: i64,
+
+    /// Authoritative per-layer dispatch. Derived from
+    /// `full_attention_interval` if absent.
+    #[serde(default)]
+    pub layer_types: Vec<String>,
+    #[serde(default)]
+    pub full_attention_interval: Option<usize>,
+
+    // --- Hyper-connections (§1) ----------------------------------------
+    /// Parallel residual streams. The inter-layer residual is
+    /// `hidden_size * hc_count` wide.
+    pub hc_count: usize,
+    /// Bottleneck width of the stream-mixing gate.
+    pub hc_lowrank: usize,
+
+    // --- PLE (§2) ------------------------------------------------------
+    /// **One-indexed.** Use [`TextConfig::ple_layers`].
+    #[serde(default)]
+    pub ple_layer_ids: Vec<usize>,
+    pub ple_embed_dim: usize,
+    pub ple_conv_kernel_size: usize,
+    /// Also the short conv's dilation.
+    pub ngram_size: usize,
+    pub heads_per_ngram: usize,
+    pub ngram_vocab_size_base: i64,
+    pub make_ngram_vocab_size_divisible_by: usize,
+    /// Shards the table is stored in; a flat split, concatenated in
+    /// shard-index order.
+    pub split_ngram_parts: usize,
+
+    // --- QSA indexer (§3) ----------------------------------------------
+    pub indexer_n_heads: usize,
+    pub indexer_kv_heads: usize,
+    pub indexer_head_dim: usize,
+    /// In **tokens**, not blocks.
+    pub indexer_budget: usize,
+    /// Positions pooled into one scored block.
+    pub indexer_compress_ratio: usize,
+
+    // --- Linear attention (§5) -----------------------------------------
+    pub linear_num_key_heads: usize,
+    pub linear_key_head_dim: usize,
+    pub linear_num_value_heads: usize,
+    pub linear_value_head_dim: usize,
+    pub linear_conv_kernel_dim: usize,
+    /// `"float32"` here; see the module note.
+    #[serde(default)]
+    pub mamba_ssm_dtype: Option<String>,
+    /// `"sigmoid"` here, where Qwen3.6 falls back to `hidden_act`.
+    #[serde(default)]
+    pub output_gate_type: Option<String>,
+
+    // --- MoE (§6) ------------------------------------------------------
+    pub num_experts: usize,
+    pub num_experts_per_tok: usize,
+    pub moe_intermediate_size: usize,
+    pub shared_expert_intermediate_size: usize,
+    /// Absent from this `config.json`; upstream's default is `true`.
+    #[serde(default = "default_true")]
+    pub norm_topk_prob: bool,
+
+    // --- MTP (§9, #313) ------------------------------------------------
+    #[serde(default)]
+    pub mtp_num_hidden_layers: usize,
+}
+
+impl TextConfig {
+    /// Zero-indexed layers carrying a PLE block — `[2]` becomes `[1]`.
+    pub fn ple_layers(&self) -> Vec<usize> {
+        self.ple_layer_ids.iter().map(|id| id - 1).collect()
+    }
+
+    /// `(ngram_size - 1) * heads_per_ngram` — 16. There is no unigram
+    /// head.
+    pub fn ngram_heads(&self) -> usize {
+        (self.ngram_size - 1) * self.heads_per_ngram
+    }
+
+    /// Dims each n-gram head contributes to the 2560-wide embedding.
+    pub fn ngram_head_dim(&self) -> usize {
+        self.ple_embed_dim / self.ngram_heads()
+    }
+
+    /// The inter-layer residual width, `hidden_size * hc_count`.
+    pub fn stream_width(&self) -> usize {
+        self.hidden_size * self.hc_count
+    }
+
+    /// Blocks the QSA indexer may select.
+    pub fn block_topk(&self) -> usize {
+        self.indexer_budget / self.indexer_compress_ratio
+    }
+
+    /// Which nonlinearity the GatedDeltaNet output gate applies.
+    /// Falls back to `hidden_act` the way upstream does — but this
+    /// checkpoint states it, and the two disagree.
+    pub fn output_gate(&self) -> Result<OutputGate> {
+        OutputGate::from_name(self.output_gate_type.as_deref().unwrap_or(&self.hidden_act))
+    }
+
+    /// Whether to carry the GatedDeltaNet recurrent state in f32.
+    /// `mamba_ssm_dtype` is upstream's explicit choice; absent, we keep
+    /// our own default rather than inventing one.
+    pub fn ssm_state_is_f32(&self) -> Option<bool> {
+        self.mamba_ssm_dtype
+            .as_deref()
+            .map(|d| d.eq_ignore_ascii_case("float32") || d.eq_ignore_ascii_case("f32"))
+    }
+
+    pub fn is_full_attention(&self, layer: usize) -> bool {
+        self.layer_types
+            .get(layer)
+            .is_some_and(|t| t == "full_attention")
+    }
+
+    fn fill_layer_types(&mut self) -> Result<()> {
+        if !self.layer_types.is_empty() {
+            return Ok(());
+        }
+        let interval = self.full_attention_interval.unwrap_or(4);
+        ensure!(
+            interval > 0,
+            "full_attention_interval must be >= 1 to derive layer_types"
+        );
+        // Every `interval`-th layer counting from one, so with
+        // interval 4 the full-attention layers are 3, 7, … 47.
+        self.layer_types = (0..self.num_hidden_layers)
+            .map(|i| {
+                if (i + 1).is_multiple_of(interval) {
+                    "full_attention".to_string()
+                } else {
+                    "linear_attention".to_string()
+                }
+            })
+            .collect();
+        Ok(())
+    }
+
+    fn validate(&self) -> Result<()> {
+        ensure!(
+            self.layer_types.len() == self.num_hidden_layers,
+            "layer_types has {} entries for {} layers",
+            self.layer_types.len(),
+            self.num_hidden_layers
+        );
+        ensure!(self.hc_count >= 1, "hc_count must be >= 1");
+        ensure!(
+            self.ngram_size >= 2,
+            "ngram_size must be >= 2; {} leaves no n-gram heads",
+            self.ngram_size
+        );
+        let heads = self.ngram_heads();
+        ensure!(
+            heads > 0 && self.ple_embed_dim.is_multiple_of(heads),
+            "ple_embed_dim ({}) must divide evenly among {heads} n-gram heads",
+            self.ple_embed_dim
+        );
+        for id in &self.ple_layer_ids {
+            // One-indexed: 0 is not a layer, and N is the last one.
+            ensure!(
+                *id >= 1 && *id <= self.num_hidden_layers,
+                "ple_layer_ids is one-indexed; {id} is outside 1..={}",
+                self.num_hidden_layers
+            );
+        }
+        ensure!(
+            self.indexer_compress_ratio >= 1
+                && self
+                    .indexer_budget
+                    .is_multiple_of(self.indexer_compress_ratio),
+            "indexer_budget ({}) must be a whole number of blocks of {}",
+            self.indexer_budget,
+            self.indexer_compress_ratio
+        );
+        ensure!(
+            self.num_experts_per_tok <= self.num_experts,
+            "num_experts_per_tok ({}) exceeds num_experts ({})",
+            self.num_experts_per_tok,
+            self.num_experts
+        );
+        Ok(())
+    }
+}
+
+fn default_hidden_act() -> String {
+    "silu".to_string()
+}
+
+fn default_true() -> bool {
+    true
+}
+
+/// Whether a `config.json` declares this architecture, so the harness
+/// can route it here rather than to the `qwen3_5` sibling.
+pub fn is_qwen4_exp(json: &str) -> bool {
+    serde_json::from_str::<serde_json::Value>(json)
+        .ok()
+        .and_then(|v| {
+            v.get("model_type")
+                .and_then(|m| m.as_str())
+                .map(|m| m == MODEL_TYPE)
+        })
+        .unwrap_or(false)
+}
+
+/// The shipped `config.json`, verbatim apart from the vision block
+/// and the 48-entry `layer_types` (spelled out below in
+/// `SHIPPED_LAYER_TYPES` so the derivation can be tested against
+/// the real one). Values are from revision
+/// `de4b8e4d43b917e7706784d8bb445c9af86a3540`.
+#[cfg(test)]
+pub(crate) const SHIPPED: &str = r#"{
+    "architectures": ["Qwen4ExpForConditionalGeneration"],
+    "image_token_id": 248056,
+    "language_model_only": false,
+    "model_type": "qwen4_exp",
+    "tie_word_embeddings": false,
+    "video_token_id": 248057,
+    "vision_start_token_id": 248053,
+    "vision_end_token_id": 248054,
+    "vision_config": {"depth": 27, "hidden_size": 1152},
+    "text_config": {
+        "attention_bias": false,
+        "attention_dropout": 0.0,
+        "bos_token_id": 248044,
+        "dtype": "bfloat16",
+        "eos_token_id": 248044,
+        "full_attention_interval": 4,
+        "hc_count": 4,
+        "hc_lowrank": 320,
+        "head_dim": 256,
+        "heads_per_ngram": 8,
+        "hidden_act": "silu",
+        "hidden_size": 2560,
+        "indexer_budget": 2048,
+        "indexer_compress_ratio": 4,
+        "indexer_head_dim": 128,
+        "indexer_kv_heads": 1,
+        "indexer_n_heads": 4,
+        "initializer_range": 0.02,
+        "linear_conv_kernel_dim": 4,
+        "linear_key_head_dim": 128,
+        "linear_num_key_heads": 16,
+        "linear_num_value_heads": 48,
+        "linear_value_head_dim": 128,
+        "make_ngram_vocab_size_divisible_by": 128,
+        "mamba_ssm_dtype": "float32",
+        "max_position_embeddings": 262144,
+        "model_type": "qwen4_exp_text",
+        "moe_intermediate_size": 640,
+        "mtp": {"hybrid": true, "num_hidden_layers": 1, "rope_theta": 10000000},
+        "mtp_num_hidden_layers": 1,
+        "mtp_use_dedicated_embeddings": false,
+        "ngram_size": 3,
+        "ngram_vocab_size_base": 20000000,
+        "num_attention_heads": 24,
+        "num_experts": 512,
+        "num_experts_per_tok": 10,
+        "num_hidden_layers": 48,
+        "num_key_value_heads": 2,
+        "output_gate_type": "sigmoid",
+        "output_router_logits": false,
+        "pad_token_id": null,
+        "partial_rotary_factor": 0.25,
+        "ple_conv_kernel_size": 4,
+        "ple_embed_dim": 2560,
+        "ple_layer_ids": [2],
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "mrope_interleaved": true,
+            "mrope_section": [11, 11, 10],
+            "partial_rotary_factor": 0.25,
+            "rope_theta": 10000000,
+            "rope_type": "default"
+        },
+        "router_aux_loss_coef": 0.001,
+        "shared_expert_intermediate_size": 640,
+        "split_ngram_parts": 128,
+        "tie_word_embeddings": false,
+        "use_cache": true,
+        "vocab_size": 248320
+    }
+}"#;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The real file's `layer_types`, as shipped.
+    const SHIPPED_FULL_ATTENTION_LAYERS: [usize; 12] =
+        [3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 47];
+
+    fn shipped() -> Config {
+        Config::from_config_json(SHIPPED).unwrap()
+    }
+
+    #[test]
+    fn parses_the_shipped_config() {
+        let cfg = shipped();
+        assert_eq!(cfg.model_type, MODEL_TYPE);
+        assert!(!cfg.tie_word_embeddings);
+        let t = &cfg.text_config;
+        assert_eq!(t.hidden_size, 2560);
+        assert_eq!(t.num_hidden_layers, 48);
+        assert_eq!(t.vocab_size, 248_320);
+        assert_eq!(t.max_position_embeddings, 262_144);
+        assert_eq!(t.eos_token_id, 248_044);
+        // Four streams of 2560 — the residual every decoder block sees.
+        assert_eq!(t.stream_width(), 10_240);
+        assert_eq!(t.hc_lowrank, 320);
+        assert!(is_qwen4_exp(SHIPPED));
+    }
+
+    /// The single most expensive off-by-one available here: reading
+    /// `[2]` as a layer number rather than a one-indexed position puts
+    /// PLE on layer 2 instead of layer 1.
+    #[test]
+    fn ple_layer_ids_are_one_indexed() {
+        let t = &shipped().text_config;
+        assert_eq!(t.ple_layer_ids, vec![2], "as shipped, one-indexed");
+        assert_eq!(t.ple_layers(), vec![1], "zero-indexed, where it runs");
+    }
+
+    /// Sixteen heads of 160, and no unigram head — `ngram_size 3` runs
+    /// the bigram and trigram orders only.
+    #[test]
+    fn ngram_geometry_is_sixteen_heads_of_one_sixty() {
+        let t = &shipped().text_config;
+        assert_eq!(t.ngram_heads(), 16);
+        assert_eq!(t.ngram_head_dim(), 160);
+        assert_eq!(t.ngram_heads() * t.ngram_head_dim(), t.ple_embed_dim);
+        // The conv's dilation is the n-gram size, not its kernel.
+        assert_eq!(t.ple_conv_kernel_size, 4);
+        assert_eq!(t.ngram_size, 3);
+    }
+
+    /// The derivation must reproduce the file's own `layer_types`:
+    /// every fourth layer counting from one, so 3, 7, … 47.
+    #[test]
+    fn derived_layer_types_match_the_shipped_ones() {
+        let t = &shipped().text_config;
+        let full: Vec<usize> = (0..t.num_hidden_layers)
+            .filter(|i| t.is_full_attention(*i))
+            .collect();
+        assert_eq!(full, SHIPPED_FULL_ATTENTION_LAYERS);
+        assert_eq!(full.len(), 12);
+        assert_eq!(t.num_hidden_layers - full.len(), 36);
+    }
+
+    /// `output_gate_type` overrides `hidden_act`, and the two disagree:
+    /// inheriting SiLU here is wrong on all 36 linear-attention layers.
+    #[test]
+    fn the_output_gate_is_sigmoid_not_the_hidden_act() {
+        let t = &shipped().text_config;
+        assert_eq!(t.hidden_act, "silu");
+        assert_eq!(t.output_gate().unwrap(), OutputGate::Sigmoid);
+        assert_eq!(t.ssm_state_is_f32(), Some(true));
+    }
+
+    /// The reused rope module is only reusable if these four values are
+    /// what its masks were pinned against (spec §7).
+    #[test]
+    fn rope_parameters_are_the_ones_qwen3_5_already_implements() {
+        let t = &shipped().text_config;
+        let r = &t.rope_parameters;
+        assert_eq!(r.rope_theta, 10_000_000.0);
+        assert_eq!(r.partial_rotary_factor, 0.25);
+        assert_eq!(r.mrope_section, vec![11, 11, 10]);
+        assert!(r.mrope_interleaved);
+        // 64 of 256 head dims rotate; 32 frequencies, which is what
+        // mrope_section sums to.
+        assert_eq!((t.head_dim as f32 * r.partial_rotary_factor) as usize, 64);
+        assert_eq!(r.mrope_section.iter().sum::<usize>(), 32);
+    }
+
+    #[test]
+    fn qsa_geometry_selects_512_blocks_of_four() {
+        let t = &shipped().text_config;
+        assert_eq!(t.block_topk(), 512);
+        assert_eq!(t.indexer_compress_ratio, 4);
+        assert_eq!(t.indexer_n_heads, 4);
+        assert_eq!(t.indexer_kv_heads, 1);
+        assert_eq!(t.indexer_head_dim, 128);
+    }
+
+    /// A one-indexed field that reads 0 is a config written against the
+    /// wrong convention, not a layer.
+    #[test]
+    fn rejects_a_zero_ple_layer_id() {
+        let bad = SHIPPED.replace("\"ple_layer_ids\": [2]", "\"ple_layer_ids\": [0]");
+        let err = Config::from_config_json(&bad).unwrap_err().to_string();
+        assert!(err.contains("one-indexed"), "got: {err}");
+    }
+
+    #[test]
+    fn rejects_layer_types_that_do_not_cover_every_layer() {
+        let bad = SHIPPED.replace(
+            "\"num_hidden_layers\": 48",
+            "\"num_hidden_layers\": 48, \"layer_types\": [\"full_attention\"]",
+        );
+        let err = Config::from_config_json(&bad).unwrap_err().to_string();
+        assert!(err.contains("layer_types"), "got: {err}");
+    }
+
+    /// The harness dispatches on `model_type`, and the sibling archs
+    /// must not claim this checkpoint or be claimed by it.
+    #[test]
+    fn architecture_detection_is_exact() {
+        assert!(is_qwen4_exp(SHIPPED));
+        assert!(!is_qwen4_exp(r#"{"model_type": "qwen3_5"}"#));
+        assert!(!is_qwen4_exp(r#"{"model_type": "qwen3_next"}"#));
+        // The nested text_config's own model_type is a different string
+        // and must not be what we match on.
+        assert_ne!(MODEL_TYPE, TEXT_MODEL_TYPE);
+        assert!(!is_qwen4_exp(r#"{"model_type": "qwen4_exp_text"}"#));
+        assert!(!is_qwen4_exp("not json"));
+    }
+}

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

@@ -0,0 +1,389 @@
+//! `qwen4_exp`'s decoder layer — the composition, and the only part of
+//! it that is new.
+//!
+//! Every other architecture we serve wraps a sublayer as
+//! `x = x + sublayer(norm(x))` over one residual stream. This one
+//! carries **four**, concatenated into a `hidden_size * hc_count`
+//! tensor, and has **no `input_layernorm`, no
+//! `post_attention_layernorm` and no final `model.norm`** anywhere in
+//! the checkpoint. The `hc_norm` inside each hyper-connection does all
+//! three jobs.
+//!
+//! ```text
+//! h [B,T,10240]
+//!   ├─ (layer 1 only)  h = h + PLE(h, gathered n-grams)
+//!   ├─ attn_hyper_connection: split → linear_attn | self_attn → scatter
+//!   └─ mlp_hyper_connection:  split → sparse_moe               → scatter
+//! ```
+//!
+//! The spec's warning about this layer is worth repeating: a wrong
+//! stream mix still produces fluent text. So the plumbing lives in
+//! [`through`], which takes the sublayer as a closure and can be tested
+//! against a known function without constructing an attention block —
+//! and the two things it must get right are asserted there:
+//!
+//! 1. **The sublayer sees the mixed `hidden_size` vector**, not the
+//!    four-stream tensor. It is 2560 wide, and the sublayers are built
+//!    for exactly that.
+//! 2. **The residual added back is the un-normalised `h`.** Only the
+//!    mixing and injection paths see the normalised form.
+//!
+//! PLE is added into the stream *before* the attention
+//! hyper-connection, not merged into it — the layer's own blocks then
+//! run on the sum. See `doc/qwen4_exp-port-spec.md`.
+
+use anyhow::{Result, anyhow};
+use candle_core::{Module, Tensor};
+use candle_nn::var_builder::ShardedVarBuilder;
+use std::sync::Arc;
+
+use crate::harness::arch::qwen3_5::linear_attn::GatedDeltaNet;
+use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
+use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
+
+use candle_core::quantized::GgmlDType;
+
+use super::config::TextConfig;
+use super::full_attn::Attention;
+use super::hyper::HyperConnection;
+use super::ple::PleBlock;
+use super::{linear_attn, moe};
+use crate::harness::arch::snapshot::LayerKvSnapshot;
+
+/// The attention slot: one layer in four is full attention with the QSA
+/// indexer, the rest are the recurrent delta rule.
+enum Mixer {
+    Full(Box<Attention>),
+    Linear(Box<GatedDeltaNet>),
+}
+
+pub struct DecoderLayer {
+    attn_hc: HyperConnection,
+    mlp_hc: HyperConnection,
+    mixer: Mixer,
+    /// Every layer's FFN is the MoE; there is no dense width in this
+    /// config to build an alternative from.
+    mlp: Qwen3_5MoeBlock,
+    /// Present on exactly the layers `ple_layer_ids` names — layer 1,
+    /// zero-indexed, for the shipped checkpoint.
+    ple: Option<PleBlock>,
+}
+
+impl DecoderLayer {
+    pub fn load(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        layer_idx: usize,
+        quant: Option<GgmlDType>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let layer_type = cfg
+            .layer_types
+            .get(layer_idx)
+            .map(String::as_str)
+            .ok_or_else(|| {
+                anyhow!(
+                    "layer_types[{layer_idx}] missing (have {} entries)",
+                    cfg.layer_types.len()
+                )
+            })?;
+        Self::load_typed(
+            cfg,
+            rotary,
+            layer_type,
+            cfg.ple_layers().contains(&layer_idx),
+            quant,
+            vb,
+        )
+    }
+
+    /// Load a layer whose type the caller knows, rather than one looked
+    /// up in `layer_types`.
+    ///
+    /// The MTP head's layer (#313) needs this: it lives under
+    /// `mtp.layers.0` with its own single-entry type list, and it
+    /// carries no PLE wherever the main model puts one.
+    pub fn load_typed(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        layer_type: &str,
+        with_ple: bool,
+        quant: Option<GgmlDType>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let mixer = match layer_type {
+            "full_attention" => {
+                Mixer::Full(Box::new(Attention::load(cfg, rotary, &vb.pp("self_attn"))?))
+            }
+            "linear_attention" => {
+                Mixer::Linear(Box::new(linear_attn::load(cfg, &vb.pp("linear_attn"))?))
+            }
+            other => anyhow::bail!(
+                "unknown layer_type '{other}' (expected 'full_attention' or \
+                 'linear_attention')"
+            ),
+        };
+
+        let hc = |name: &str| {
+            HyperConnection::load(
+                &vb.pp(name),
+                cfg.hidden_size,
+                cfg.hc_count,
+                cfg.hc_lowrank,
+                cfg.rms_norm_eps,
+                // Decoder hyper-connections always scatter back; only
+                // the top-level mixer does not.
+                true,
+            )
+        };
+
+        let ple = if with_ple {
+            Some(PleBlock::load(
+                &vb.pp("ple"),
+                cfg.hidden_size,
+                cfg.hc_count,
+                cfg.ple_conv_kernel_size,
+                // The dilation is the n-gram size, not the kernel.
+                cfg.ngram_size,
+                cfg.rms_norm_eps,
+            )?)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            attn_hc: hc("attn_hyper_connection")?,
+            mlp_hc: hc("mlp_hyper_connection")?,
+            mixer,
+            mlp: moe::load(cfg, quant, &vb.pp("mlp"))?,
+            ple,
+        })
+    }
+
+    /// `h` is the four-stream residual, `[B, L, hidden * hc_count]`.
+    ///
+    /// `ple_embed` is the gathered n-gram embedding for this step,
+    /// `[B, L, hidden]` — required on the PLE layer and ignored
+    /// everywhere else.
+    pub fn forward(
+        &mut self,
+        h: &Tensor,
+        ple_embed: Option<&Tensor>,
+        causal_mask: Option<&Tensor>,
+        cos: &Tensor,
+        sin: &Tensor,
+        past_len: usize,
+    ) -> Result<Tensor> {
+        // Split the borrows so the sublayer closures can take `&mut`
+        // while the hyper-connections are read.
+        let Self {
+            attn_hc,
+            mlp_hc,
+            mixer,
+            mlp,
+            ple,
+        } = self;
+
+        let h = match ple {
+            Some(block) => {
+                let e = ple_embed.ok_or_else(|| {
+                    anyhow!(
+                        "this layer carries a PLE block but no gathered n-gram \
+                         embedding was supplied"
+                    )
+                })?;
+                (h + block.forward(h, e)?)?
+            }
+            None => h.clone(),
+        };
+
+        let h = through(attn_hc, &h, |x| match mixer {
+            Mixer::Full(attn) => attn.forward(x, causal_mask, cos, sin, past_len),
+            // The recurrent path's causality is in its state lifecycle,
+            // so it takes neither mask nor rotary.
+            Mixer::Linear(net) => Ok(net.forward(x)?),
+        })?;
+
+        through(mlp_hc, &h, |x| Ok(mlp.forward(x)?))
+    }
+
+    pub fn clear_kv_cache(&mut self) -> Result<()> {
+        match &mut self.mixer {
+            Mixer::Full(attn) => attn.clear_kv_cache(),
+            Mixer::Linear(net) => net.clear_kv_cache(),
+        }
+        if let Some(ple) = &mut self.ple {
+            ple.clear_state();
+        }
+        Ok(())
+    }
+
+    /// Capture this layer's attention-side state.
+    pub fn snapshot_kv(&self) -> candle_core::Result<LayerKvSnapshot> {
+        Ok(match &self.mixer {
+            Mixer::Full(attn) => {
+                let (kv, indexer_keys) = attn.snapshot_kv();
+                LayerKvSnapshot::FullSparse { kv, indexer_keys }
+            }
+            Mixer::Linear(net) => {
+                let (conv_state, recurrent_state) = net.snapshot_state()?;
+                LayerKvSnapshot::Linear {
+                    conv_state,
+                    recurrent_state,
+                }
+            }
+        })
+    }
+
+    /// Replace this layer's attention-side state. A variant mismatch
+    /// means the snapshot came from a different model, or from the same
+    /// model with a different `layer_types` — either way, restoring it
+    /// would put a recurrent state into an attention layer.
+    pub fn restore_kv(&mut self, snap: &LayerKvSnapshot) -> candle_core::Result<()> {
+        match (&mut self.mixer, snap) {
+            (Mixer::Full(attn), LayerKvSnapshot::FullSparse { kv, indexer_keys }) => {
+                attn.restore_kv(kv.as_ref(), indexer_keys.as_ref())
+            }
+            (
+                Mixer::Linear(net),
+                LayerKvSnapshot::Linear {
+                    conv_state,
+                    recurrent_state,
+                },
+            ) => net.restore_state(conv_state.as_ref(), recurrent_state.as_ref()),
+            _ => candle_core::bail!(
+                "restore_kv: snapshot layer kind does not match this layer's mixer kind"
+            ),
+        }
+    }
+
+    /// The PLE block's rolling conv context, if this layer carries one.
+    pub fn snapshot_ple(&self) -> candle_core::Result<Option<Tensor>> {
+        match &self.ple {
+            Some(block) => block.snapshot_state(),
+            None => Ok(None),
+        }
+    }
+
+    /// Restore the PLE conv context. Refuses a snapshot that disagrees
+    /// with this layer about whether PLE lives here at all.
+    pub fn restore_ple(&mut self, conv_state: Option<&Tensor>) -> candle_core::Result<()> {
+        match &mut self.ple {
+            Some(block) => block.restore_state(conv_state),
+            None if conv_state.is_none() => Ok(()),
+            None => candle_core::bail!(
+                "restore_ple: snapshot carries PLE state but this layer has no PLE block"
+            ),
+        }
+    }
+
+    pub fn has_ple(&self) -> bool {
+        self.ple.is_some()
+    }
+}
+
+/// Run one sublayer through a hyper-connection.
+///
+/// Split the four streams into the `hidden_size` vector the sublayer
+/// expects, call it, and scatter the result back across the streams
+/// with the per-stream injection weights. The tensor added back is the
+/// **un-normalised** `h`; only the mixing and injection paths see the
+/// normalised form.
+fn through<F>(hc: &HyperConnection, h: &Tensor, sublayer: F) -> Result<Tensor>
+where
+    F: FnOnce(&Tensor) -> Result<Tensor>,
+{
+    let split = hc.split(h)?;
+    let inject = split.inject.as_ref().ok_or_else(|| {
+        anyhow!("a decoder hyper-connection must carry block_inject_weight to scatter back")
+    })?;
+    let out = sublayer(&split.input)?;
+    Ok(hc.combine(&split.residual, &out, inject)?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+    use std::cell::RefCell;
+
+    /// With every projection zeroed the gate is exactly 0.5 and the
+    /// injection exactly 1, so `through` reduces to
+    /// `h + (sublayer(mean of normalised streams) broadcast to all)`.
+    fn hc(hidden: usize, hc_count: usize) -> HyperConnection {
+        HyperConnection::zeroed_for_test(hidden, hc_count, 2, true)
+    }
+
+    fn stream(values: &[f32]) -> Tensor {
+        Tensor::from_vec(values.to_vec(), (1, 1, values.len()), &Device::Cpu).unwrap()
+    }
+
+    /// The sublayer is handed the mixed `hidden_size` vector, not the
+    /// four-stream tensor. Feeding it the raw stream would be four
+    /// times too wide and every sublayer in the model is built for the
+    /// narrow one.
+    #[test]
+    fn the_sublayer_sees_one_stream_and_the_result_is_four() {
+        let seen = RefCell::new(Vec::new());
+        let h = stream(&[3.0, 4.0, 30.0, 40.0]);
+        let out = through(&hc(2, 2), &h, |x| {
+            seen.borrow_mut().extend_from_slice(x.dims());
+            Ok(x.zeros_like()?)
+        })
+        .unwrap();
+
+        assert_eq!(seen.into_inner(), vec![1, 1, 2], "sublayer input width");
+        assert_eq!(out.dims(), &[1, 1, 4], "the residual stays four streams");
+    }
+
+    /// A sublayer that returns zero must leave the stream exactly as it
+    /// was — not as its normalised form. This is the assertion that
+    /// separates a correct residual from one that quietly re-normalises
+    /// the model every layer.
+    #[test]
+    fn a_silent_sublayer_leaves_the_stream_untouched() {
+        let h = stream(&[3.0, 4.0, 30.0, 40.0]);
+        let out = through(&hc(2, 2), &h, |x| Ok(x.zeros_like()?)).unwrap();
+        let got: Vec<f32> = out.flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(got, vec![3.0, 4.0, 30.0, 40.0]);
+    }
+
+    /// The sublayer's output reaches every stream, scaled by that
+    /// stream's injection weight — 1.0 for all of them here.
+    #[test]
+    fn the_result_is_injected_into_every_stream() {
+        let h = stream(&[3.0, 4.0, 30.0, 40.0]);
+        let out = through(&hc(2, 2), &h, |x| Ok((x.ones_like()? * 7.0)?)).unwrap();
+        let got: Vec<f32> = out.flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(got, vec![10.0, 11.0, 37.0, 47.0]);
+    }
+
+    /// A hyper-connection with no `block_inject_weight` is the
+    /// top-level mixer, which collapses the streams instead of feeding
+    /// a sublayer. Using one here would silently drop the residual.
+    #[test]
+    fn a_mixer_cannot_stand_in_for_a_decoder_hyper_connection() {
+        let mixer = HyperConnection::zeroed_for_test(2, 2, 2, false);
+        let h = stream(&[1.0, 1.0, 1.0, 1.0]);
+        let err = through(&mixer, &h, |x| Ok(x.clone()))
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains("block_inject_weight"), "got: {err}");
+    }
+
+    #[test]
+    fn ple_layers_are_the_zero_indexed_ones() {
+        let cfg =
+            super::super::config::Config::from_config_json(super::super::config::SHIPPED).unwrap();
+        // The composition asks the config, and the config is one-indexed.
+        assert_eq!(cfg.text_config.ple_layers(), vec![1]);
+        assert!(!cfg.text_config.ple_layers().contains(&2));
+    }
+
+    #[test]
+    fn zeroed_stream_keeps_its_dtype_and_shape() {
+        let h = Tensor::zeros((2, 3, 8), DType::F32, &Device::Cpu).unwrap();
+        let out = through(&hc(2, 4), &h, |x| Ok(x.clone())).unwrap();
+        assert_eq!(out.dims(), &[2, 3, 8]);
+    }
+}

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

@@ -0,0 +1,381 @@
+//! `qwen4_exp`'s full-attention layer — one in every four.
+//!
+//! Mechanically this is `qwen3_5`'s: GQA 24:2, per-head `(1 + w)`
+//! RMSNorm on q and k, a sigmoid output gate carried in a widened
+//! `q_proj`, and `head_dim^-0.5` scaling. It is written here rather
+//! than reused because of the one thing that differs — the [`Indexer`]
+//! that decides which of the past this layer may look at — and because
+//! the two architectures' configs are different types.
+//!
+//! Two details are worth stating where the code can be read against
+//! them:
+//!
+//! 1. **The gate split is per head, not a halving.** `q_proj` is
+//!    `num_heads * head_dim * 2` wide; reshape to
+//!    `(.., num_heads, 2 * head_dim)` and take the halves *within* each
+//!    head. Splitting the flat 12288 vector down the middle gives a
+//!    different permutation — the first twelve heads' queries as the
+//!    query and the last twelve as the gate — which trains nothing and
+//!    reads as a weak model.
+//! 2. **QSA replaces the causal mask; it does not join it.** The
+//!    selection only ever offers positions at or behind the query, so
+//!    the returned mask is already causal. Below the budget the indexer
+//!    declines to produce one at all and this layer keeps the plain
+//!    causal mask — which is also what keeps the flash path available
+//!    for short prompts, since an additive mask forecloses it.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §3 and §4.
+
+use anyhow::{Context, Result};
+use candle_core::{Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::kv_cache::ConcatKvCache;
+use candle_nn::var_builder::ShardedVarBuilder;
+use std::sync::Arc;
+
+use crate::harness::arch::qwen3_5::full_attn::{AttnMask, attention_context};
+use crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm;
+use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
+
+use super::config::TextConfig;
+use super::qsa::Indexer;
+
+pub struct Attention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: Qwen3_5RmsNorm,
+    k_norm: Qwen3_5RmsNorm,
+    /// The sparse-attention side channel. Owns its own KV cache, which
+    /// must be cleared with this layer's.
+    indexer: Indexer,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    /// `num_heads * head_dim` — the attention output width, which is
+    /// not `hidden_size` here (6144 against 2560).
+    attn_width: usize,
+    rotary: Arc<RotaryEmbedding>,
+    kv_cache: ConcatKvCache,
+}
+
+impl Attention {
+    pub fn load(
+        cfg: &TextConfig,
+        rotary: Arc<RotaryEmbedding>,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let head_dim = cfg.head_dim;
+        let num_heads = cfg.num_attention_heads;
+        let num_kv_heads = cfg.num_key_value_heads;
+        anyhow::ensure!(
+            num_kv_heads > 0 && num_heads.is_multiple_of(num_kv_heads),
+            "num_attention_heads ({num_heads}) must be a positive multiple of \
+             num_key_value_heads ({num_kv_heads})"
+        );
+        let attn_width = num_heads * head_dim;
+        Ok(Self {
+            // Doubled output: query and gate, interleaved per head.
+            q_proj: linear(vb, "q_proj", cfg.hidden_size, attn_width * 2)?,
+            k_proj: linear(vb, "k_proj", cfg.hidden_size, num_kv_heads * head_dim)?,
+            v_proj: linear(vb, "v_proj", cfg.hidden_size, num_kv_heads * head_dim)?,
+            o_proj: linear(vb, "o_proj", attn_width, cfg.hidden_size)?,
+            q_norm: Qwen3_5RmsNorm::load(&vb.pp("q_norm"), head_dim, cfg.rms_norm_eps)?,
+            k_norm: Qwen3_5RmsNorm::load(&vb.pp("k_norm"), head_dim, cfg.rms_norm_eps)?,
+            // The indexer is a submodule of the attention block, not a
+            // sibling: `self_attn.indexer.index_qk_proj`, per the
+            // checkpoint's own index.
+            indexer: Indexer::load(&vb.pp("indexer"), cfg)?,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups: num_heads / num_kv_heads,
+            head_dim,
+            attn_width,
+            rotary,
+            kv_cache: ConcatKvCache::new(2),
+        })
+    }
+
+    /// `x` is the hyper-connection's mixed `hidden_size` input, not the
+    /// four-stream residual. `causal_mask` is the layer's usual mask,
+    /// used whenever the indexer declines to narrow it. `past_len` is
+    /// the sequence offset, and must match what the KV cache holds.
+    pub fn forward(
+        &mut self,
+        x: &Tensor,
+        causal_mask: Option<&Tensor>,
+        cos: &Tensor,
+        sin: &Tensor,
+        past_len: usize,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        // The indexer reads the same hidden states this layer does, and
+        // advances its own cache whether or not it produces a mask.
+        let selected = self
+            .indexer
+            .binding_mask(x, &self.rotary, past_len, x.dtype())?;
+
+        // q_proj is doubled: (query, gate) *within* each head.
+        let q_raw = self
+            .q_proj
+            .forward(x)?
+            .reshape((b, l, self.num_heads, self.head_dim * 2))?;
+        let q = q_raw.narrow(3, 0, self.head_dim)?;
+        let gate = q_raw
+            .narrow(3, self.head_dim, self.head_dim)?
+            .contiguous()?
+            .reshape((b, l, self.attn_width))?;
+
+        let q = self
+            .q_norm
+            .forward(&q.contiguous()?)?
+            .transpose(1, 2)?
+            .contiguous()?;
+        let k = self
+            .k_proj
+            .forward(x)?
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?;
+        let k = self
+            .k_norm
+            .forward(&k.contiguous()?)?
+            .transpose(1, 2)?
+            .contiguous()?;
+        let v = self
+            .v_proj
+            .forward(x)?
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?;
+
+        let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
+        let (k, v) = self.kv_cache.append(&k, &v)?;
+
+        // A QSA mask is additive and arbitrary, so it must not reach the
+        // flash kernel — the type says so, rather than a comment.
+        let mask = match &selected {
+            Some(m) => AttnMask::Additive(m),
+            None => AttnMask::causal(causal_mask),
+        };
+        let scale = 1.0_f64 / (self.head_dim as f64).sqrt();
+        let ctx = attention_context(&q, &k, &v, mask, self.num_kv_groups, scale)?;
+
+        let ctx = ctx
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((b, l, self.attn_width))?;
+        let gated = (ctx * candle_nn::ops::sigmoid(&gate)?)?;
+        Ok(self.o_proj.forward(&gated)?)
+    }
+
+    /// Capture both caches at one token boundary.
+    ///
+    /// The same reason `clear_kv_cache` drops them together: a
+    /// snapshot holding only the attention K/V restores a layer whose
+    /// indexer still holds the previous request's keys, and the
+    /// selection it computes from them is wrong without being invalid.
+    pub fn snapshot_kv(&self) -> (Option<(Tensor, Tensor)>, Option<Tensor>) {
+        let kv = match (self.kv_cache.k(), self.kv_cache.v()) {
+            (Some(k), Some(v)) => Some((k.clone(), v.clone())),
+            _ => None,
+        };
+        (kv, self.indexer.snapshot_keys())
+    }
+
+    /// Replace both caches from a snapshot.
+    pub fn restore_kv(
+        &mut self,
+        kv: Option<&(Tensor, Tensor)>,
+        indexer_keys: Option<&Tensor>,
+    ) -> candle_core::Result<()> {
+        self.kv_cache.reset();
+        if let Some((k, v)) = kv {
+            self.kv_cache.append(k, v)?;
+        }
+        self.indexer.restore_keys(indexer_keys);
+        Ok(())
+    }
+
+    /// Both caches, together. Dropping only one leaves the indexer
+    /// selecting blocks of a previous request's text.
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache.reset();
+        self.indexer.clear_cache();
+    }
+
+    pub fn indexer(&self) -> &Indexer {
+        &self.indexer
+    }
+}
+
+fn linear(vb: &ShardedVarBuilder, name: &str, in_dim: usize, out_dim: usize) -> Result<Linear> {
+    let weight = vb
+        .pp(name)
+        .get((out_dim, in_dim), "weight")
+        .with_context(|| format!("load '{}/{name}/weight'", vb.prefix()))?;
+    Ok(Linear::new(weight, None))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::harness::arch::qwen3_5::TextConfig as Qwen3_5TextConfig;
+    use crate::harness::arch::qwen4_exp::qsa::BlockSelector;
+    use candle_core::{DType, Device};
+
+    const HIDDEN: usize = 4;
+    const HEADS: usize = 2;
+    const HEAD_DIM: usize = 2;
+
+    fn rope() -> Arc<RotaryEmbedding> {
+        let json = format!(
+            r#"{{"vocab_size": 8, "hidden_size": {HIDDEN}, "intermediate_size": 8,
+                 "num_hidden_layers": 1, "num_attention_heads": {HEADS},
+                 "num_key_value_heads": 1, "head_dim": {HEAD_DIM},
+                 "max_position_embeddings": 64, "rms_norm_eps": 1e-6,
+                 "rope_parameters": {{"rope_theta": 10000.0,
+                                      "partial_rotary_factor": 1.0}}}}"#
+        );
+        let cfg: Qwen3_5TextConfig = serde_json::from_str(&json).unwrap();
+        Arc::new(RotaryEmbedding::new(DType::F32, &cfg, &Device::Cpu).unwrap())
+    }
+
+    /// Columns of a weight matrix, given as rows of `[out]` values for
+    /// input channel 0 and zeros elsewhere — so `W x` for `x = e0` is
+    /// exactly the values given.
+    fn from_first_column(values: &[f32], in_dim: usize) -> Linear {
+        let out = values.len();
+        let mut w = vec![0.0f32; out * in_dim];
+        for (row, v) in values.iter().enumerate() {
+            w[row * in_dim] = *v;
+        }
+        Linear::new(
+            Tensor::from_vec(w, (out, in_dim), &Device::Cpu).unwrap(),
+            None,
+        )
+    }
+
+    fn eye(n: usize) -> Linear {
+        let mut w = vec![0.0f32; n * n];
+        for i in 0..n {
+            w[i * n + i] = 1.0;
+        }
+        Linear::new(Tensor::from_vec(w, (n, n), &Device::Cpu).unwrap(), None)
+    }
+
+    /// A layer whose `q_proj` produces `[1,1, 2,2, 3,3, 4,4]` for the
+    /// unit input, so the per-head and the halved reading of the gate
+    /// disagree: per head the gate is `[2,2, 4,4]`, halved it would be
+    /// `[3,3, 4,4]`.
+    fn layer(budget: usize) -> Attention {
+        let attn_width = HEADS * HEAD_DIM;
+        Attention {
+            q_proj: from_first_column(&[1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0], HIDDEN),
+            k_proj: from_first_column(&[0.0; HEAD_DIM], HIDDEN),
+            v_proj: from_first_column(&[1.0; HEAD_DIM], HIDDEN),
+            o_proj: eye(attn_width),
+            q_norm: Qwen3_5RmsNorm::from_weight(
+                Tensor::zeros(HEAD_DIM, DType::F32, &Device::Cpu).unwrap(),
+                1e-6,
+                None,
+            ),
+            k_norm: Qwen3_5RmsNorm::from_weight(
+                Tensor::zeros(HEAD_DIM, DType::F32, &Device::Cpu).unwrap(),
+                1e-6,
+                None,
+            ),
+            indexer: Indexer::from_parts(
+                from_first_column(&[0.5; 2 * HEAD_DIM], HIDDEN),
+                BlockSelector::new(1, HEAD_DIM, 4, budget).unwrap(),
+                1,
+                1,
+                HEAD_DIM,
+            ),
+            num_heads: HEADS,
+            num_kv_heads: 1,
+            num_kv_groups: HEADS,
+            head_dim: HEAD_DIM,
+            attn_width,
+            rotary: rope(),
+            kv_cache: ConcatKvCache::new(2),
+        }
+    }
+
+    fn unit_input(seq: usize) -> Tensor {
+        let mut v = vec![0.0f32; seq * HIDDEN];
+        for i in 0..seq {
+            v[i * HIDDEN] = 1.0;
+        }
+        Tensor::from_vec(v, (1, seq, HIDDEN), &Device::Cpu).unwrap()
+    }
+
+    fn sigmoid(x: f32) -> f32 {
+        1.0 / (1.0 + (-x).exp())
+    }
+
+    /// The gate is the second half of each head, not the second half of
+    /// the vector. With one position the attention is the identity on
+    /// `v`, so the output is the gate alone — and the halved reading
+    /// gives sigmoid(3) where the correct one gives sigmoid(2).
+    #[test]
+    fn the_output_gate_splits_within_each_head() {
+        let mut attn = layer(2048);
+        let (cos, sin) = attn.rotary.plain_cos_sin(0, 1).unwrap();
+        let out: Vec<f32> = attn
+            .forward(&unit_input(1), None, &cos, &sin, 0)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let want = [sigmoid(2.0), sigmoid(2.0), sigmoid(4.0), sigmoid(4.0)];
+        for (g, w) in out.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-5, "got {out:?} want {want:?}");
+        }
+        // The halved reading would put sigmoid(3) in the first two.
+        assert!(
+            (out[0] - sigmoid(3.0)).abs() > 1e-3,
+            "this is the halved split, not the per-head one: {out:?}"
+        );
+    }
+
+    /// The indexer's cache and the main KV cache must advance together
+    /// across a prefill and a decode step, and be dropped together — a
+    /// surviving indexer cache selects blocks of the previous request.
+    #[test]
+    fn both_caches_advance_and_clear_together() {
+        let mut attn = layer(2048);
+        let (cos, sin) = attn.rotary.plain_cos_sin(0, 5).unwrap();
+        attn.forward(&unit_input(5), None, &cos, &sin, 0).unwrap();
+        assert_eq!(attn.indexer().cached_len(), 5);
+
+        let (cos, sin) = attn.rotary.plain_cos_sin(5, 1).unwrap();
+        attn.forward(&unit_input(1), None, &cos, &sin, 5).unwrap();
+        assert_eq!(attn.indexer().cached_len(), 6);
+        // Six positions is one block of four plus a tail — nowhere near
+        // a 2048-token budget, so the layer stayed on the causal path.
+        assert!(attn.indexer().selector().is_dense(6));
+
+        attn.clear_kv_cache();
+        assert_eq!(attn.indexer().cached_len(), 0);
+    }
+
+    /// A layer whose budget binds still produces a well-formed output —
+    /// the sparse path is exercised end to end, not just the mask.
+    #[test]
+    fn the_sparse_path_runs() {
+        let mut attn = layer(4);
+        let (cos, sin) = attn.rotary.plain_cos_sin(0, 12).unwrap();
+        let out = attn.forward(&unit_input(12), None, &cos, &sin, 0).unwrap();
+        assert_eq!(out.dims(), &[1, 12, HIDDEN]);
+        assert!(!attn.indexer().selector().is_dense(12));
+        let values: Vec<f32> = out.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(
+            values.iter().all(|v| v.is_finite()),
+            "a masked-out row would softmax to NaN: {values:?}"
+        );
+    }
+}

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

@@ -94,6 +94,43 @@ impl HyperConnection {
         })
     }
 
+    /// A hyper-connection whose three projections are all zero, for
+    /// tests in this module and its siblings.
+    ///
+    /// That makes the whole gate analytic: `down(hn) = 0` so
+    /// `silu(0) = 0`, `up(0) = 0` so `sigmoid(0) = 0.5`, and
+    /// `block_inject(hn) = 0` so `2 * sigmoid(0) = 1`. The mix weight is
+    /// exactly 0.5 everywhere and every stream is injected with weight
+    /// 1, so any error in the `/hc_count` divisions, the sigmoid
+    /// constants, or mean-versus-sum shows up as a wrong number rather
+    /// than a wrong shape.
+    #[cfg(test)]
+    pub(crate) fn zeroed_for_test(
+        hidden_size: usize,
+        hc_count: usize,
+        hc_lowrank: usize,
+        use_combine: bool,
+    ) -> Self {
+        use candle_core::{DType, Device};
+        let dev = Device::Cpu;
+        let wide = hidden_size * hc_count;
+        let z = |o: usize, i: usize| {
+            Linear::new(Tensor::zeros((o, i), DType::F32, &dev).unwrap(), None)
+        };
+        Self {
+            hc_norm: crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm::from_weight(
+                Tensor::zeros(wide, DType::F32, &dev).unwrap(),
+                1e-6,
+                Some(hidden_size),
+            ),
+            mix_down: z(hc_lowrank, wide),
+            mix_up: z(wide, hc_lowrank),
+            block_inject: use_combine.then(|| z(hc_count, wide)),
+            hidden_size,
+            hc_count,
+        }
+    }
+
     /// Split the stream tensor into a sublayer input plus what is needed
     /// to rejoin.
     pub fn split(&self, h: &Tensor) -> candle_core::Result<HyperSplit> {
@@ -171,7 +208,7 @@ fn linear(vb: &ShardedVarBuilder, name: &str, in_dim: usize, out_dim: usize) ->
 #[cfg(test)]
 mod tests {
     use super::*;
-    use candle_core::{DType, Device};
+    use candle_core::Device;
 
     /// Build a hyper-connection whose three projections are all zero.
     /// That makes the whole gate analytic:
@@ -181,23 +218,7 @@ mod tests {
     /// injected with weight 1. Any error in the /hc_count divisions, the
     /// sigmoid constants, or mean-vs-sum shows up as a wrong number.
     fn zeroed(hidden: usize, hc: usize, lowrank: usize) -> HyperConnection {
-        let dev = Device::Cpu;
-        let wide = hidden * hc;
-        let z = |o: usize, i: usize| {
-            Linear::new(Tensor::zeros((o, i), DType::F32, &dev).unwrap(), None)
-        };
-        HyperConnection {
-            hc_norm: Qwen3_5RmsNorm::from_weight(
-                Tensor::zeros(wide, DType::F32, &dev).unwrap(),
-                1e-6,
-                Some(hidden),
-            ),
-            mix_down: z(lowrank, wide),
-            mix_up: z(wide, lowrank),
-            block_inject: Some(z(hc, wide)),
-            hidden_size: hidden,
-            hc_count: hc,
-        }
+        HyperConnection::zeroed_for_test(hidden, hc, lowrank, true)
     }
 
     #[test]

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

@@ -0,0 +1,115 @@
+//! `qwen4_exp`'s linear-attention layer — 36 of 48.
+//!
+//! This is `qwen3_5`'s GatedDeltaNet with nothing changed. The tensor
+//! names, the split (non-fused) `in_proj_*` layout, the delta rule, the
+//! conv and the state shapes are all identical, which is why the
+//! largest single piece of the Qwen3.6 port carries over untouched and
+//! this module is a page rather than a file.
+//!
+//! What does not carry over is two settings, both of which our
+//! defaults would get wrong in the same quiet way:
+//!
+//! 1. **`output_gate_type: sigmoid`.** Qwen3.6 leaves the field unset
+//!    and the gated output norm falls back to SiLU. Inheriting that
+//!    here applies the wrong nonlinearity on every one of these 36
+//!    layers — fluent output, worse model.
+//! 2. **`mamba_ssm_dtype: float32`.** We carry the recurrent state
+//!    through a bf16 round-trip by default, deliberately (#284: the
+//!    flag is off because the #283 measurements were taken against
+//!    that baseline). This checkpoint states f32, and a port must not
+//!    silently answer an upstream choice with our own. So the
+//!    checkpoint wins here, and `NEURON_GDN_STATE_F32` remains the
+//!    default only where the config is silent.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §5.
+
+use anyhow::Result;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::linear_attn::{
+    GatedDeltaNet, GatedDeltaNetParams, gdn_state_f32,
+};
+
+use super::config::TextConfig;
+
+/// `vb` should be `.pp(...)`-ed to the layer's `linear_attn` prefix.
+pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<GatedDeltaNet> {
+    GatedDeltaNet::load_with_params(params(cfg)?, vb)
+}
+
+/// The settings this architecture asks for, as opposed to the ones
+/// `qwen3_5` would supply.
+pub(crate) fn params(cfg: &TextConfig) -> Result<GatedDeltaNetParams> {
+    Ok(GatedDeltaNetParams {
+        hidden_size: cfg.hidden_size,
+        num_value_heads: cfg.linear_num_value_heads,
+        num_key_heads: cfg.linear_num_key_heads,
+        key_head_dim: cfg.linear_key_head_dim,
+        value_head_dim: cfg.linear_value_head_dim,
+        conv_kernel_size: cfg.linear_conv_kernel_dim,
+        rms_norm_eps: cfg.rms_norm_eps,
+        output_gate: cfg.output_gate()?,
+        // The checkpoint's choice where it makes one; ours only where
+        // it is silent.
+        state_f32: cfg.ssm_state_is_f32().unwrap_or_else(gdn_state_f32),
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::harness::arch::qwen3_5::rmsnorm::OutputGate;
+    use crate::harness::arch::qwen4_exp::config::{Config, SHIPPED};
+
+    fn shipped() -> Config {
+        Config::from_config_json(SHIPPED).unwrap()
+    }
+
+    /// The layout `qwen3_5` already loads: 16 key heads and 48 value
+    /// heads of 128, a conv of 4. If these ever diverge the reuse claim
+    /// is void, so they are asserted rather than assumed.
+    #[test]
+    fn the_geometry_is_the_one_qwen3_5_already_loads() {
+        let p = params(&shipped().text_config).unwrap();
+        assert_eq!(p.hidden_size, 2560);
+        assert_eq!(p.num_key_heads, 16);
+        assert_eq!(p.num_value_heads, 48);
+        assert_eq!(p.key_head_dim, 128);
+        assert_eq!(p.value_head_dim, 128);
+        assert_eq!(p.conv_kernel_size, 4);
+        // conv_dim = 2 * 16 * 128 + 48 * 128 = 10240, the same width as
+        // the four-stream residual, which is a coincidence worth not
+        // reading anything into.
+        assert_eq!(
+            2 * p.num_key_heads * p.key_head_dim + p.num_value_heads * p.value_head_dim,
+            10240
+        );
+    }
+
+    /// The two deltas, together — this is the whole reason the module
+    /// exists.
+    #[test]
+    fn the_checkpoint_overrides_both_of_our_defaults() {
+        let p = params(&shipped().text_config).unwrap();
+        assert_eq!(
+            p.output_gate,
+            OutputGate::Sigmoid,
+            "inheriting hidden_act would give SiLU on all 36 layers"
+        );
+        assert!(
+            p.state_f32,
+            "mamba_ssm_dtype is float32; our #284 default is a bf16 round-trip"
+        );
+    }
+
+    /// Where the checkpoint says nothing, our default stands rather
+    /// than a guess — the fallback is the process-wide #284 flag, not a
+    /// hardcoded answer.
+    #[test]
+    fn a_silent_checkpoint_leaves_our_default_in_place() {
+        let json = SHIPPED.replace("\"mamba_ssm_dtype\": \"float32\",", "");
+        let cfg = Config::from_config_json(&json).unwrap();
+        assert_eq!(cfg.text_config.ssm_state_is_f32(), None);
+        assert_eq!(params(&cfg.text_config).unwrap().state_f32, gdn_state_f32());
+    }
+}

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

@@ -14,6 +14,33 @@
 //!
 //! Built in dependency order, so each piece can be parity-tested before
 //! anything stacks on it:
+//! - [`config`] — the checkpoint's own `config.json`, and what its
+//!   fields mean where they are not what they look like.
 //! - [`hyper`] — hyper-connections, the residual structure itself.
+//! - [`decoder`] — the composition: four residual streams, PLE on
+//!   layer 1, and no layernorm outside the hyper-connections.
+//! - [`full_attn`] — one layer in four, plus the indexer that
+//!   narrows what it may attend.
+//! - [`linear_attn`] — the other three layers in four, reused whole
+//!   from `qwen3_5` with two settings corrected.
+//! - [`model`] — embeddings in, logits out; four streams, one
+//!   n-gram lookup, and no final norm.
+//! - [`mtp`] — the shipped draft head: one more decoder layer, fed
+//!   the target's pre-mixer stream and the next token's embedding.
+//! - [`moe`] — the FFN every layer has, loaded from fused expert
+//!   tensors instead of 512 modules.
+//! - [`ple`] — the hashed n-gram table on layer 1: how its rows
+//!   are addressed, and how the gathered rows are consumed.
+//! - [`qsa`] — which blocks of the past a full-attention layer is
+//!   allowed to look at.
 
+pub mod config;
+pub mod decoder;
+pub mod full_attn;
 pub mod hyper;
+pub mod linear_attn;
+pub mod model;
+pub mod moe;
+pub mod mtp;
+pub mod ple;
+pub mod qsa;

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

@@ -0,0 +1,171 @@
+//! `qwen4_exp`'s MoE FFN — every layer's FFN, since there is no dense
+//! one.
+//!
+//! The routing is [`Qwen3_5MoeBlock`]'s, unchanged: softmax over all
+//! 512 in f32, top 10, renormalise, add a sigmoid-gated shared expert.
+//! What differs is **storage**. This checkpoint ships the routed
+//! experts as two fused 3D tensors rather than 512 modules:
+//!
+//! ```text
+//! mlp.experts.gate_up_proj   (512, 1280, 2560)   gate ++ up, stacked
+//! mlp.experts.down_proj      (512, 2560,  640)
+//! ```
+//!
+//! against our loader's `mlp.experts.{i}.{gate,up,down}_proj`. So this
+//! module is a loader, not an algorithm: it slices the fused tensors
+//! into the per-expert SwiGLUs the block already knows how to route.
+//!
+//! **In `gate_up_proj` the first `moe_intermediate_size` rows are the
+//! gate and the last are the up projection.** Reading them the other
+//! way round computes `silu(up) * gate` instead of `silu(gate) * up` —
+//! a different function of the same weights, which is exactly the kind
+//! of error that produces a model that talks and reasons slightly
+//! worse. It is asserted below rather than commented.
+//!
+//! ## Why the fused tensors stay fused
+//!
+//! Slicing them into 512 per-expert modules copies **241.6 GB** at
+//! BF16, 5.03 GB per layer. That fits in neither beast's 64 GB of VRAM
+//! nor its 123 GB of host RAM, and no quantisation rescues it: #309
+//! measured 59.4 GB of 63.7 used at 4.25 bpw with only *four* layers'
+//! experts offloaded. A model that cannot be held cannot be loaded.
+//!
+//! So the block takes `Experts::Banked` and views one expert per
+//! routed token instead. A slice along dim 0 of a contiguous tensor is
+//! contiguous, and `x.matmul(w.t())` is what `Linear` does anyway, so
+//! the views feed the same GEMMs with no copy. It is also the layout a
+//! host-resident bank with a device-side gather wants (#318).
+//!
+//! See `doc/qwen4_exp-port-spec.md` §6.
+
+use anyhow::{Context, Result};
+use candle_core::Tensor;
+use candle_core::quantized::GgmlDType;
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::mlp::Qwen3_5MLP;
+use crate::harness::arch::qwen3_5::moe::{Experts, Qwen3_5MoeBlock};
+
+use super::config::TextConfig;
+
+/// `vb` should be `.pp(...)`-ed to the layer's `mlp` prefix.
+///
+/// `quant` quantises the routed experts in situ (#315). Without it the
+/// experts stay banked at the checkpoint's precision, which is only
+/// useful for a model small enough to hold — 120.8 B routed parameters
+/// are 241.6 GB at BF16 and fit nowhere on this hardware, so a real
+/// load of `Qwen3.8-Flash-Next` always passes one.
+pub fn load(
+    cfg: &TextConfig,
+    quant: Option<GgmlDType>,
+    vb: &ShardedVarBuilder,
+) -> Result<Qwen3_5MoeBlock> {
+    let (h, inter) = (cfg.hidden_size, cfg.moe_intermediate_size);
+    anyhow::ensure!(
+        cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && inter > 0,
+        "MoE needs num_experts ({}), num_experts_per_tok ({}) and \
+         moe_intermediate_size ({inter}) all > 0",
+        cfg.num_experts,
+        cfg.num_experts_per_tok,
+    );
+
+    let gate = Linear::new(
+        vb.pp("gate")
+            .get((cfg.num_experts, h), "weight")
+            .with_context(|| format!("load '{}/gate/weight'", vb.prefix()))?,
+        None,
+    );
+
+    let experts_vb = vb.pp("experts");
+    let gate_up = experts_vb
+        .get((cfg.num_experts, inter * 2, h), "gate_up_proj")
+        .with_context(|| format!("load '{}/gate_up_proj'", experts_vb.prefix()))?;
+    let down = experts_vb
+        .get((cfg.num_experts, h, inter), "down_proj")
+        .with_context(|| format!("load '{}/down_proj'", experts_vb.prefix()))?;
+    check_fused_experts(&gate_up, &down, inter)?;
+    // The fused pair is the peak: 5.03 GB per layer at BF16. Quantising
+    // here and letting it drop at the end of this function is what keeps
+    // the whole 241.6 GB from ever existing.
+    let experts = match quant {
+        Some(dtype) => Experts::quantize_banked(&gate_up, &down, inter, dtype)
+            .with_context(|| format!("quantize experts to {dtype:?}"))?,
+        None => Experts::Banked {
+            gate_up,
+            down,
+            intermediate: inter,
+        },
+    };
+
+    let (shared_expert, shared_expert_gate) = if cfg.shared_expert_intermediate_size > 0 {
+        let shared = Qwen3_5MLP::load_with_dims(
+            &vb.pp("shared_expert"),
+            h,
+            cfg.shared_expert_intermediate_size,
+        )
+        .context("load shared_expert")?;
+        let gate_w = vb
+            .pp("shared_expert_gate")
+            .get((1, h), "weight")
+            .with_context(|| format!("load '{}/shared_expert_gate/weight'", vb.prefix()))?;
+        (Some(shared), Some(Linear::new(gate_w, None)))
+    } else {
+        (None, None)
+    };
+
+    Ok(Qwen3_5MoeBlock::from_parts(
+        gate,
+        experts,
+        shared_expert,
+        shared_expert_gate,
+        cfg.num_experts_per_tok,
+        cfg.norm_topk_prob,
+    ))
+}
+
+/// Validate the fused pair before it is banked.
+///
+/// The block will take views of these per routed expert, so a
+/// disagreement here surfaces as a shape error mid-forward on whichever
+/// token first routes to the offending expert — which is a worse place
+/// to find it than at load.
+pub(crate) fn check_fused_experts(
+    gate_up: &Tensor,
+    down: &Tensor,
+    intermediate: usize,
+) -> Result<()> {
+    let (experts, fused, _) = gate_up.dims3()?;
+    anyhow::ensure!(
+        fused == intermediate * 2,
+        "gate_up_proj is {fused} wide, expected 2 x {intermediate}"
+    );
+    anyhow::ensure!(
+        down.dims3()?.0 == experts,
+        "down_proj holds {} experts, gate_up_proj holds {experts}",
+        down.dims3()?.0
+    );
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+
+    /// The loader validates the fused pair up front. Without this the
+    /// first token to route to a bad expert fails mid-forward, which is
+    /// a worse place to learn the checkpoint disagrees with the config.
+    #[test]
+    fn a_fused_pair_that_disagrees_with_the_config_is_rejected() {
+        let dev = Device::Cpu;
+        let gate_up = Tensor::zeros((2, 6, 4), DType::F32, &dev).unwrap();
+        let down = Tensor::zeros((2, 4, 3), DType::F32, &dev).unwrap();
+        assert!(check_fused_experts(&gate_up, &down, 3).is_ok());
+        // 2 x 4 != 6
+        assert!(check_fused_experts(&gate_up, &down, 4).is_err());
+        // expert counts must agree
+        let down_short = Tensor::zeros((1, 4, 3), DType::F32, &dev).unwrap();
+        assert!(check_fused_experts(&gate_up, &down_short, 3).is_err());
+    }
+}

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

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

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

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

crates/neuron/src/harness/candle.rs

@@ -555,7 +555,7 @@ impl LoadedHandle {
 #[derive(Clone)]
 pub enum KvSnapshotRef {
     Worker(super::device_worker::KvSnapshotId),
-    Local(Arc<super::arch::qwen3_5::snapshot::KvCacheSnapshot>),
+    Local(Arc<super::arch::snapshot::KvCacheSnapshot>),
     #[cfg(feature = "cuda")]
     Tp(u64),
 }
@@ -1050,6 +1050,16 @@ pub enum ModelArch {
     // are real; forward bails "not implemented yet". See
     // `arch/qwen3_5.rs` for the open architecture work.
     Qwen3_5Dense(super::arch::qwen3_5::Qwen3_5ForCausalLM),
+
+    // Qwen3.8-Flash-Next (model_type "qwen4_exp") — a new architecture
+    // generation, not a Qwen3 variant: four residual streams joined by
+    // hyper-connections, QSA sparse attention, and a hashed n-gram
+    // table on one layer. Boxed for the same reason as LlamaDense.
+    //
+    // No snapshot or batched-decode support yet, so it falls through to
+    // the "no batched-decode support" arms and serves one request at a
+    // time. See `arch/qwen4_exp/` and the epic on #307.
+    Qwen4Exp(Box<super::arch::qwen4_exp::model::Qwen4ExpForCausalLM>),
 }
 
 impl ModelArch {
@@ -1066,6 +1076,7 @@ impl ModelArch {
             ModelArch::LlamaQuantized(m) => m.forward(input, offset)?,
             ModelArch::LlamaDense(m) => m.forward(input, offset)?,
             ModelArch::Qwen3_5Dense(m) => m.forward(input, offset)?,
+            ModelArch::Qwen4Exp(m) => m.forward(input, offset)?,
         };
         squeeze_to_vocab(&raw)
     }
@@ -1093,6 +1104,10 @@ impl ModelArch {
                 m.clear_kv_cache();
                 Ok(())
             }
+            // Clears the KV caches, the QSA indexer's separate cache and
+            // PLE's carried ids together — a survivor of any one of them
+            // reads the previous request's text.
+            ModelArch::Qwen4Exp(m) => m.clear_kv_cache(),
         }
     }
 
@@ -1101,15 +1116,16 @@ impl ModelArch {
     /// candle-transformers archs keep theirs private, so they stay on
     /// the clear-every-request path.
     pub fn supports_kv_snapshot(&self) -> bool {
-        matches!(self, ModelArch::Qwen3_5Dense(_))
+        matches!(self, ModelArch::Qwen3_5Dense(_) | ModelArch::Qwen4Exp(_))
     }
 
     /// Capture the live cache state as a prefix snapshot. See
     /// `arch/qwen3_5/snapshot.rs` for what a snapshot contains and the
     /// copy-semantics constraints.
-    pub fn snapshot_kv_cache(&self) -> Result<super::arch::qwen3_5::snapshot::KvCacheSnapshot> {
+    pub fn snapshot_kv_cache(&self) -> Result<super::arch::snapshot::KvCacheSnapshot> {
         match self {
             ModelArch::Qwen3_5Dense(m) => Ok(m.snapshot_kv_cache()?),
+            ModelArch::Qwen4Exp(m) => Ok(m.snapshot_kv_cache()?),
             _ => anyhow::bail!("snapshot_kv_cache: architecture has no snapshot support"),
         }
     }
@@ -1118,10 +1134,11 @@ impl ModelArch {
     /// restore-instead-of-clear half of prefix caching.
     pub fn restore_kv_cache(
         &mut self,
-        snap: &super::arch::qwen3_5::snapshot::KvCacheSnapshot,
+        snap: &super::arch::snapshot::KvCacheSnapshot,
     ) -> Result<()> {
         match self {
             ModelArch::Qwen3_5Dense(m) => Ok(m.restore_kv_cache(snap)?),
+            ModelArch::Qwen4Exp(m) => Ok(m.restore_kv_cache(snap)?),
             _ => anyhow::bail!("restore_kv_cache: architecture has no snapshot support"),
         }
     }
@@ -1332,8 +1349,14 @@ impl LlamaDense {
 /// value. New entries land alongside a new `ModelArch` variant + a
 /// dispatch branch in `load_arch_dense` (plus, for TP, a parallel
 /// pattern in `tp_qwen3.rs`).
-const DENSE_SUPPORTED_MODEL_TYPES: &[&str] =
-    &["llama", "qwen3", "qwen3_5", "qwen3_moe", "qwen3_next"];
+pub(crate) const DENSE_SUPPORTED_MODEL_TYPES: &[&str] = &[
+    "llama",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_moe",
+    "qwen3_next",
+    "qwen4_exp",
+];
 
 /// Pre-flight check the operator's `config.json` against the set of
 /// architectures the dense path actually knows how to build. Surfaces
@@ -1388,6 +1411,11 @@ pub(crate) fn check_dense_config_supported(config_json: &str, model_id: &str) ->
 /// families than the TP path because each TP-aware module is a real
 /// chunk of work (`tp_qwen3.rs` is the only one shipped today).
 #[cfg(feature = "cuda")]
+// `qwen4_exp` is deliberately absent: the single-GPU path loads it, but
+// there is no `tp_qwen4_exp.rs`, so a `tensor_parallel` request for it
+// is refused here rather than silently served from one device. This
+// list is `cfg(cuda)`, so no CPU-build test can cover it — the gate is
+// the list itself.
 const TP_SUPPORTED_MODEL_TYPES: &[&str] = &["qwen3", "qwen3_5", "qwen3_next"];
 
 /// TP-side counterpart to `check_dense_config_supported`. Gates the
@@ -2707,6 +2735,9 @@ impl CandleHarness {
             self.resolve_dense_files(spec, source_id).await?;
         let device_for_load = device.clone();
         let model_id_for_log = spec.model_id.clone();
+        // Resolved out here: `spec` is borrowed and does not outlive the
+        // blocking task.
+        let isq = super::quant::parse_quant_string(spec.quant.as_deref())?;
 
         let arch = tokio::task::spawn_blocking(move || -> Result<ModelArch> {
             let cfg_text = std::fs::read_to_string(&config_path).context("read config.json")?;
@@ -2802,6 +2833,40 @@ impl CandleHarness {
                         .context("build Qwen3-Next dense model")?;
                     Ok(ModelArch::Qwen3_5Dense(model))
                 }
+                "qwen4_exp" => {
+                    // Qwen3.8-Flash-Next. Like qwen3_5 it loads through
+                    // the sharded backend; unlike it, the config is a
+                    // separate type — the field vocabulary overlaps but
+                    // the model-level plumbing does not.
+                    let cfg = super::arch::qwen4_exp::config::Config::from_config_json(&cfg_text)
+                        .context("parse qwen4_exp config.json")?;
+                    let sharded_vb = unsafe {
+                        candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                            &safetensors_paths,
+                            dtype,
+                            &device_for_load,
+                        )
+                        .context("build ShardedVarBuilder for qwen4_exp")?
+                    };
+                    // ISQ is not optional for this architecture in
+                    // practice: without it the routed experts are
+                    // 241.6 GB. `quant` comes from the catalogue, and a
+                    // load without one will fail on memory rather than
+                    // on anything this code can say (#315).
+                    // ISQ is not really optional here: without it the
+                    // routed experts are 241.6 GB and the load fails on
+                    // memory rather than on anything this code can say.
+                    // The catalogue's `quant` is the knob (#315, #252).
+                    let model = super::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
+                        &cfg,
+                        dtype,
+                        isq,
+                        &device_for_load,
+                        &sharded_vb,
+                    )
+                    .context("build qwen4_exp model")?;
+                    Ok(ModelArch::Qwen4Exp(Box::new(model)))
+                }
                 other => {
                     // Defensive: `check_dense_config_supported` already
                     // gated on the supported set, so this branch is
@@ -2820,6 +2885,59 @@ impl CandleHarness {
         Ok((tokenizer_path, arch))
     }
 
+    /// Whether this spec's `quant` names a GGUF the repo actually has.
+    ///
+    /// `quant` is overloaded (#320): the TP path reads it as in-situ
+    /// quantisation of safetensors, while the single-GPU paths read it
+    /// as "select the GGUF whose filename contains this string". A repo
+    /// with no GGUF at all cannot mean the second, so rather than fail
+    /// with an error about a file that was never going to exist, the
+    /// same string is read the way the TP path reads it.
+    ///
+    /// This is deliberately narrow. A repo that *does* carry GGUFs keeps
+    /// the old behaviour exactly, including the old error when none of
+    /// them match the requested quant — that case is a typo, not an ISQ
+    /// request, and turning it into a silent 360 GB dense load would be
+    /// a worse answer than the error.
+    async fn quant_selects_a_gguf(
+        &self,
+        spec: &ModelSpec,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<bool> {
+        let Some(quant) = spec.quant.as_deref().filter(|q| !q.is_empty()) else {
+            return Ok(false);
+        };
+        let api = self.hf_api_for(&source_id.scheme)?;
+        let info = api
+            .model(source_id.repo_path())
+            .info()
+            .await
+            .with_context(|| format!("fetch HF repo info for {source_id}"))?;
+        let names: Vec<String> = info
+            .siblings
+            .iter()
+            .map(|s| s.rfilename.to_lowercase())
+            .collect();
+        let any_gguf = names.iter().any(|n| n.ends_with(".gguf"));
+        if any_gguf {
+            return Ok(true);
+        }
+        let any_safetensors = names.iter().any(|n| n.ends_with(".safetensors"));
+        if any_safetensors {
+            tracing::info!(
+                model = %source_id,
+                quant = %quant,
+                "no GGUF in this repo; reading `quant` as in-situ quantization of \
+                 the safetensors (#320)"
+            );
+            return Ok(false);
+        }
+        anyhow::bail!(
+            "repo {source_id} has neither a GGUF nor safetensors; `quant = {quant:?}` \
+             cannot be satisfied"
+        )
+    }
+
     /// Resolve a model spec to local GGUF and tokenizer file paths via
     /// hf-hub. Downloads on first use; subsequent calls are cached.
     async fn resolve_files(
@@ -3804,7 +3922,7 @@ impl Harness for CandleHarness {
             context_profile,
         ) = if let Some(w) = &worker {
             // CUDA path: resolve, then load in the worker.
-            if spec.quant.is_some() {
+            if self.quant_selects_a_gguf(spec, &source_id).await? {
                 let (gguf_path, tokenizer_path) = self.resolve_files(spec, &source_id).await?;
                 let handle = w
                     .load_gguf(gguf_path, spec.model_id.clone())
@@ -3828,36 +3946,38 @@ impl Harness for CandleHarness {
                 let (config_path, tokenizer_path, safetensors_paths) =
                     self.resolve_dense_files(spec, &source_id).await?;
                 let meta = VisionMeta::from_config_path(&config_path);
-                // Prefix snapshots (#11) exist only for the in-tree
-                // qwen3_5 arch — which serves BOTH its model_types
-                // (qwen3_5 and the flat-config qwen3_next MoE family);
-                // the worker holds the ModelArch so the async side
-                // decides from config.json instead.
-                let snapshot_capable = matches!(
-                    config_model_type(&config_path).as_deref(),
-                    Some(super::arch::qwen3_5::MODEL_TYPE)
-                        | Some(super::arch::qwen3_5::MODEL_TYPE_NEXT)
-                );
                 // Context-limit physics (#67): single-GPU → world_size 1.
-                // `None` for non-qwen3_5 dense archs.
+                // `None` for dense archs with no profile builder — they
+                // fall back to the static prompt cap.
                 let context_profile =
-                    super::context_limit::profile_from_qwen3_5_config(&config_path, 1);
-                let handle = w
-                    .load_dense(config_path, safetensors_paths, spec.model_id.clone())
+                    super::context_limit::profile_from_qwen3_5_config(&config_path, 1).or_else(
+                        || super::context_limit::profile_from_qwen4_exp_config(&config_path, 1),
+                    );
+                // In-situ quantisation (#320) reaches the worker load
+                // too, not just the CPU one — parsed here because
+                // `spec` does not cross the channel.
+                let isq = super::quant::parse_quant_string(spec.quant.as_deref())?;
+                let load = w
+                    .load_dense(config_path, safetensors_paths, spec.model_id.clone(), isq)
                     .await
                     .map_err(|e| anyhow::anyhow!("worker load_dense: {e:#}"))?;
+                // Prefix snapshots (#11): the built model answers for
+                // itself on the way back, rather than the async side
+                // re-deriving the answer from `config.json`. See
+                // `DenseLoad` — the re-derivation is what silently
+                // dropped batched decode for qwen4_exp.
                 (
                     tokenizer_path,
                     None,
-                    Some(handle),
+                    Some(load.handle),
                     meta,
-                    snapshot_capable,
+                    load.supports_kv_snapshot,
                     context_profile,
                 )
             }
         } else {
             // CPU path: legacy spawn_blocking + Arc<Mutex<ModelArch>>.
-            let (tokenizer_path, arch) = if spec.quant.is_some() {
+            let (tokenizer_path, arch) = if self.quant_selects_a_gguf(spec, &source_id).await? {
                 self.load_arch_gguf(spec, &source_id, &device).await?
             } else {
                 self.load_arch_dense(spec, &source_id, &device).await?
@@ -4675,6 +4795,13 @@ impl CandleHarness {
             image_token_id: vision_meta.image_token_id,
             image_grid_factor: vision_meta.image_grid_factor,
             spec: spec.clone(),
+            // Snapshot capability re-derived from `config.json`
+            // rather than asked of the loaded shard. Correct today —
+            // `TP_SUPPORTED_MODEL_TYPES` is exactly these two — but it
+            // is the same shape as the bug `DenseLoad` fixed on the
+            // single-GPU path, so whoever adds an arch here (#319)
+            // must add it in both places or lose batched decode
+            // silently.
             prefix_cache: self.new_prefix_cache(matches!(
                 config_model_type(&config_path).as_deref(),
                 Some(super::arch::qwen3_5::MODEL_TYPE)
@@ -6226,9 +6353,14 @@ struct VisionMeta {
 }
 
 /// Peek at `config.json` for its `model_type`. Best-effort — `None`
-/// on any read/parse error. The load path uses this to decide prefix-
-/// snapshot capability for worker-held models the async side can't
-/// inspect directly.
+/// on any read/parse error.
+///
+/// Only the TP load path still asks: the single-GPU path takes the
+/// arch's own answer off the load reply (`DenseLoad`) rather than
+/// re-deriving it here, which is what stopped the two from drifting.
+/// The TP path has the same hazard and not yet the same cure — see
+/// the note at its call site.
+#[cfg(feature = "cuda")]
 fn config_model_type(config_path: &std::path::Path) -> Option<String> {
     let text = std::fs::read_to_string(config_path).ok()?;
     let v: serde_json::Value = serde_json::from_str(&text).ok()?;
@@ -8113,6 +8245,22 @@ mod tests {
         check_dense_config_supported(cfg, "Qwen/Qwen3-1.7B").expect("qwen3 should pass");
     }
 
+    #[test]
+    fn check_dense_config_accepts_qwen4_exp() {
+        // Qwen3.8-Flash-Next is a new architecture generation rather
+        // than a Qwen3 variant, so it has its own model_type and its own
+        // arm in `load_arch_dense`. The gate and that match must agree —
+        // a type in the supported set with no arm reaches the defensive
+        // "unrouted" bail instead of loading.
+        let cfg = r#"{
+            "model_type": "qwen4_exp",
+            "architectures": ["Qwen4ExpForConditionalGeneration"],
+            "text_config": {"hidden_size": 2560}
+        }"#;
+        check_dense_config_supported(cfg, "Qwen/Qwen3.8-Flash-Next")
+            .expect("qwen4_exp should be in the supported set");
+    }
+
     #[test]
     fn check_dense_config_rejects_unsupported_arch_with_clear_message() {
         // Use a deliberately-fake model_type so this test stays

crates/neuron/src/harness/context_limit.rs

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

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

@@ -13,11 +13,13 @@
 //! ARCH model state in this state slab will gain a companion
 //! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
 
-use crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot;
+use crate::harness::arch::snapshot::KvCacheSnapshot;
 use crate::harness::candle::ModelArch;
 #[cfg(feature = "cuda")]
 use crate::harness::device_worker::jobs::TpHandle;
-use crate::harness::device_worker::jobs::{ArchHandle, ImageHandle, ImageInput, Job, KvSnapshotId};
+use crate::harness::device_worker::jobs::{
+    ArchHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
+};
 use crate::harness::image::ZImagePipeline;
 #[cfg(feature = "cuda")]
 use crate::harness::tp::TpLeaderModel;
@@ -131,17 +133,44 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
             } => {
                 let result = load_gguf_inner(&state.device, &gguf_path, &model_id)
                     .map(|arch| insert_arch(&mut state, Box::new(arch)));
+                // Same reasoning as `LoadImage`: release load-time
+                // transients before serving, and — on the failure
+                // path — give back whatever the partial load had
+                // allocated before it gave up. Without this a load
+                // that OOMs leaves its bytes in the pool, where the
+                // driver still counts them against the next load.
+                trim_device_pool(&state);
                 let _ = reply.send(result);
             }
             Job::LoadDense {
                 config_path,
                 safetensors_paths,
                 model_id,
+                quant,
                 reply,
             } => {
-                let result =
-                    load_dense_inner(&state.device, &config_path, &safetensors_paths, &model_id)
-                        .map(|arch| insert_arch(&mut state, Box::new(arch)));
+                let result = load_dense_inner(
+                    &state.device,
+                    &config_path,
+                    &safetensors_paths,
+                    &model_id,
+                    quant,
+                )
+                .map(|arch| {
+                    // Ask the model, once, while it is still in
+                    // hand. See `DenseLoad`.
+                    let supports_kv_snapshot = arch.supports_kv_snapshot();
+                    let handle = insert_arch(&mut state, Box::new(arch));
+                    DenseLoad {
+                        handle,
+                        supports_kv_snapshot,
+                    }
+                });
+                // As for `LoadGguf`. It matters more here: a dense
+                // load that runs out of VRAM part-way has typically
+                // allocated tens of gigabytes first, and the evicted
+                // model waiting to be restored needs them back.
+                trim_device_pool(&state);
                 let _ = reply.send(result);
             }
             Job::DropArch { handle, reply } => {
@@ -158,6 +187,14 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 // the end of the arm; calling drop() explicitly just
                 // makes the intent visible.
                 drop(removed);
+                // Return the freed blocks to the system rather than
+                // the stream-ordered pool, so the next load actually
+                // sees the VRAM — `DropImage` has done this since the
+                // benjy cold-swap, and an unload is an unload whatever
+                // modality asked for it.
+                if was_present {
+                    trim_device_pool(&state);
+                }
                 tracing::debug!(
                     device_index,
                     handle = handle.0,
@@ -374,6 +411,13 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 let was_present = removed.is_some();
                 drop(removed);
                 state.tp_kv_snapshots.retain(|(h, _), _| *h != handle);
+                // As `DropArch`: this is only the leader's shard —
+                // the other ranks free theirs in their own processes
+                // — but the leader's is the one sharing a card with
+                // whatever loads next.
+                if was_present {
+                    trim_device_pool(&state);
+                }
                 tracing::debug!(
                     device_index,
                     tp_handle = handle.0,
@@ -895,6 +939,7 @@ fn load_dense_inner(
     config_path: &std::path::Path,
     safetensors_paths: &[std::path::PathBuf],
     model_id: &str,
+    quant: Option<candle_core::quantized::GgmlDType>,
 ) -> anyhow::Result<ModelArch> {
     use anyhow::Context;
     use candle_core::DType;
@@ -988,6 +1033,36 @@ fn load_dense_inner(
                 .context("build Qwen3-Next dense model")?;
             Ok(ModelArch::Qwen3_5Dense(model))
         }
+        "qwen4_exp" => {
+            // Qwen3.8-Flash-Next. Mirrors the CPU path in
+            // `candle::load_arch_dense` — same config type, same
+            // sharded backend, same in-situ quantisation.
+            //
+            // ISQ is not optional for this architecture: the routed
+            // experts are 241.6 GB at BF16 (measured from the
+            // checkpoint's own headers), against 32.6 GB of card. A
+            // load without a `quant` will fail on memory rather than
+            // on anything this code can say (#315, #252).
+            let cfg = crate::harness::arch::qwen4_exp::config::Config::from_config_json(&cfg_text)
+                .context("parse qwen4_exp config.json")?;
+            let sharded_vb = unsafe {
+                candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                    safetensors_paths,
+                    dtype,
+                    device,
+                )
+                .context("build ShardedVarBuilder for qwen4_exp")?
+            };
+            let model = crate::harness::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
+                &cfg,
+                dtype,
+                quant,
+                device,
+                &sharded_vb,
+            )
+            .context("build qwen4_exp model")?;
+            Ok(ModelArch::Qwen4Exp(Box::new(model)))
+        }
         other => anyhow::bail!(
             "unrouted supported model_type '{other}' — \
              DENSE_SUPPORTED_MODEL_TYPES and load_dense_inner \
@@ -1045,7 +1120,7 @@ fn tp_load_shard_inner(
         "qwen3_5" | "qwen3_next" => {
             let cfg = crate::harness::tp::tp_qwen3_5::Config::from_config_json(config_json)
                 .context("parse Qwen3-Next Config JSON for leader load")?;
-            let quant_dtype = crate::harness::tp::worker::parse_quant_string(quant)?;
+            let quant_dtype = crate::harness::quant::parse_quant_string(quant)?;
             TpLeaderModel::Qwen3_5(crate::harness::tp::tp_qwen3_5::TpQwen3_5ForCausalLM::load(
                 cfg,
                 &vb,
@@ -1179,7 +1254,7 @@ fn tp_assemble_kv_batch(
         })?;
         pairs.push((snap, *len));
     }
-    let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
+    let batch = crate::harness::arch::snapshot::assemble_batch(&pairs)?;
     let model = tp_models
         .get_mut(&handle)
         .ok_or_else(|| anyhow::anyhow!("TpAssembleKvBatch: no model for handle {}", handle.0))?;
@@ -1212,9 +1287,8 @@ fn tp_extract_kv_rows(
         .snapshot_kv_cache()?;
     let mut total = 0u64;
     for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
-        let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
-            &live, row, prefix_len, padded_len, steps,
-        )?;
+        let snap =
+            crate::harness::arch::snapshot::extract_row(&live, row, prefix_len, padded_len, steps)?;
         total += snap.size_bytes();
         state.tp_kv_snapshots.insert((handle, id), snap);
     }
@@ -1337,7 +1411,7 @@ fn assemble_kv_batch(
         })?;
         pairs.push((snap, *len));
     }
-    let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
+    let batch = crate::harness::arch::snapshot::assemble_batch(&pairs)?;
     let arch = models
         .get_mut(&handle)
         .ok_or_else(|| anyhow::anyhow!("AssembleKvBatch: no model for handle {}", handle.0))?;
@@ -1363,9 +1437,8 @@ fn extract_kv_rows(
         .snapshot_kv_cache()?;
     let mut out = Vec::with_capacity(rows.len());
     for &(row, prefix_len) in rows {
-        let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
-            &live, row, prefix_len, padded_len, steps,
-        )?;
+        let snap =
+            crate::harness::arch::snapshot::extract_row(&live, row, prefix_len, padded_len, steps)?;
         let id = KvSnapshotId(state.next_kv_snapshot_id);
         state.next_kv_snapshot_id = state.next_kv_snapshot_id.wrapping_add(1);
         let bytes = snap.size_bytes();
@@ -1671,3 +1744,42 @@ fn drain_poisoned(job: Job, device_index: u32) {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Every architecture the preflight admits must have a loader arm
+    /// on **this** path, not only on the CPU one in
+    /// `candle::load_arch_dense`.
+    ///
+    /// There are two dense loaders — the worker's (every CUDA load)
+    /// and candle's (the CPU fallback) — and one list gating both.
+    /// `qwen4_exp` was added to the list and to the CPU loader, so it
+    /// passed preflight and then hit the "unrouted supported
+    /// model_type" bail here: on a GPU host the model was
+    /// unloadable, in about 200 ms, having read nothing. The bail's
+    /// own comment said the two must stay in sync; nothing checked.
+    ///
+    /// A bare `{"model_type": ...}` reaches the arm and then fails on
+    /// the typed parse, which is fine — this asserts only that an arm
+    /// exists to fail in.
+    #[test]
+    fn every_supported_dense_model_type_has_a_loader_arm() {
+        let dir = tempfile::tempdir().unwrap();
+        for model_type in crate::harness::candle::DENSE_SUPPORTED_MODEL_TYPES {
+            let path = dir.path().join(format!("{model_type}.json"));
+            std::fs::write(&path, format!(r#"{{"model_type": "{model_type}"}}"#)).unwrap();
+            let msg =
+                match load_dense_inner(&candle_core::Device::Cpu, &path, &[], model_type, None) {
+                    Ok(_) => panic!("a bare config must not build a model"),
+                    Err(e) => format!("{e:#}"),
+                };
+            assert!(
+                !msg.contains("unrouted"),
+                "`{model_type}` is in DENSE_SUPPORTED_MODEL_TYPES but load_dense_inner \
+                 has no arm for it: {msg}"
+            );
+        }
+    }
+}

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

@@ -19,6 +19,27 @@ use tokio::sync::oneshot;
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ArchHandle(pub u64);
 
+/// What a dense load hands back: the handle, plus what the arch it
+/// built says about itself.
+///
+/// The async side cannot ask the model anything — it lives in the
+/// worker's slab and never crosses the channel — so anything the
+/// serving path needs to know has to ride back on the reply. Prefix
+/// snapshot support (#11) is the first such fact, and it is here
+/// rather than re-derived from `config.json` because that is exactly
+/// how it went wrong: the async side matched a list of `model_type`
+/// strings that `ModelArch::supports_kv_snapshot` had since grown
+/// past, so `qwen4_exp` loaded with batched decode silently off.
+/// A capability the arch answers for itself cannot drift from the
+/// arch.
+#[derive(Debug, Clone, Copy)]
+pub struct DenseLoad {
+    pub handle: ArchHandle,
+    /// [`crate::harness::candle::ModelArch::supports_kv_snapshot`], as
+    /// answered by the model that was actually built.
+    pub supports_kv_snapshot: bool,
+}
+
 /// Opaque handle to a `TpLeaderModel` stored in the worker thread's
 /// state slab. Same shape as [`ArchHandle`] but in a separate
 /// namespace so the two slabs can coexist without ambiguity. Phase 3
@@ -107,7 +128,11 @@ pub enum Job {
         config_path: PathBuf,
         safetensors_paths: Vec<PathBuf>,
         model_id: String,
-        reply: oneshot::Sender<Result<ArchHandle>>,
+        /// In-situ quantisation (#320): the catalogue's `quant`,
+        /// already parsed. `None` loads at the checkpoint dtype.
+        /// Archs that cannot quantise in situ ignore it.
+        quant: Option<candle_core::quantized::GgmlDType>,
+        reply: oneshot::Sender<Result<DenseLoad>>,
     },
     /// Remove the model from the slab and drop it. The `Drop` runs on
     /// the worker thread so CUDA tensors release their memory on the
@@ -163,7 +188,7 @@ pub enum Job {
     /// state and install it as the model's live state (#98). `seqs`
     /// pairs each snapshot id with its true token length; attention
     /// K/V is right-padded to the batch max and `cat`ed on dim 0 (see
-    /// `arch::qwen3_5::snapshot::assemble_batch`). Replies with the
+    /// `arch::snapshot::assemble_batch`). Replies with the
     /// padded uniform KV length. The source snapshots remain stored —
     /// the caller drops them via `DropKvSnapshot` when the sequences
     /// leave the batch.

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

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

crates/neuron/src/harness/engine.rs

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

crates/neuron/src/harness/mod.rs

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

crates/neuron/src/harness/quant.rs

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

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

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

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

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

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

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

doc/qwen4_exp-port-spec.md

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

script/seed-cutlass.sh

@@ -0,0 +1,252 @@
+#!/bin/env bash
+#
+# Pre-seed the CUTLASS checkout that cudaforge would otherwise clone from
+# GitHub during every CUDA build.
+#
+# candle-flash-attn's build script calls `cudaforge`, which fetches CUTLASS
+# into `${CUDAFORGE_HOME:-$HOME/.cudaforge}/git/checkouts/cutlass-<commit[..16]>`.
+# If the cache starts empty — as it does on a fresh or single-use build
+# machine — every CUDA job clones CUTLASS afresh, and a clone GitHub
+# refuses fails the whole job late, after the 20 minutes of kernel
+# compilation that preceded it:
+#
+#   fatal: could not read Username for 'https://github.com': No such device
+#   Error: GitOperationFailed("git clone failed with status: exit status: 128")
+#
+# That message reads like a missing credential. It is the opposite: see
+# the note on anonymity below.
+#
+# cudaforge short-circuits when the cache directory already holds `include/`
+# at the right commit (dependency.rs, `fetch_with_lock`), so seeding it here
+# — with retries and backoff, which the build script has none of — removes
+# the fetch from the build's critical path.
+#
+# The fetch is forced ANONYMOUS, and probes the remote on every run.
+# CUTLASS is public, so no credential is needed or wanted; sending none
+# takes our side of the connection out of the picture, so a failure is
+# unambiguously not something this build machine attached to the request.
+# What has actually been observed is a 401 that nobody has yet traced to
+# a responder:
+#
+#   error: RPC failed; HTTP 401 curl 22 The requested URL returned error: 401
+#
+# Do not assume that is GitHub. GitHub rate-limits its API, not clones,
+# and a filtering proxy or a captive resolver on the path answers 401
+# just as readily. The diagnostics below print who replied and with what
+# headers on every build, so a refusal arrives with healthy samples
+# beside it rather than alone.
+#
+# Best effort by design: if seeding fails, the build behaves exactly as it
+# does today and cargo emits the authoritative error. This script never
+# fails a build that would otherwise have passed.
+set -uo pipefail
+
+# Read no git configuration at all for the duration of this script.
+# Clearing the individual keys instead (an empty credential.helper and an
+# empty http.<url>.extraheader) does neutralise an injected auth header,
+# but it leaves url.<...>.insteadOf rewrites and any config injected
+# through GIT_CONFIG_COUNT untouched — both of which can also redirect or
+# authenticate the fetch. Discarding the config files closes every vector
+# at once and needs no guess about which one a given machine carries;
+# this script has no configuration of its own to lose.
+# GIT_CONFIG_GLOBAL / GIT_CONFIG_SYSTEM require git >= 2.32.
+export GIT_CONFIG_GLOBAL=/dev/null
+export GIT_CONFIG_SYSTEM=/dev/null
+export GIT_CONFIG_COUNT=0
+
+# Never prompt: without this a 401 blocks on a username read instead of
+# failing, which turns a refused fetch into a confusing hang.
+export GIT_TERMINAL_PROMPT=0
+export GIT_ASKPASS=/bin/true
+
+CUTLASS_REPO="${CUTLASS_REPO:-https://github.com/NVIDIA/cutlass.git}"
+CACHE_ROOT="${CUDAFORGE_HOME:-$HOME/.cudaforge}/git/checkouts"
+
+# The commit is pinned in candle-flash-attn, not here — read it from the
+# source cargo actually resolved, so this cannot drift from the candle pin.
+BUILD_RS=$(find "${CARGO_HOME:-$HOME/.cargo}/git/checkouts" \
+  -path '*/candle-flash-attn/build.rs' -print -quit 2>/dev/null)
+if [[ -z "${BUILD_RS}" ]]; then
+  echo "seed-cutlass: candle-flash-attn not in the cargo git checkouts yet;" >&2
+  echo "              run 'cargo fetch' before this step. Skipping." >&2
+  exit 0
+fi
+
+COMMIT=$(sed -nE 's/^const CUTLASS_COMMIT: &str = "([0-9a-f]+)";/\1/p' "${BUILD_RS}")
+if [[ ! "${COMMIT}" =~ ^[0-9a-f]{40}$ ]]; then
+  echo "seed-cutlass: could not read CUTLASS_COMMIT from ${BUILD_RS}. Skipping." >&2
+  exit 0
+fi
+
+DEST="${CACHE_ROOT}/cutlass-${COMMIT:0:16}"
+echo "seed-cutlass: commit ${COMMIT}"
+echo "seed-cutlass: target ${DEST}"
+
+# Same short-circuit cudaforge applies, so a warm cache costs nothing.
+if [[ -d "${DEST}/include" ]] && [[ "$(git -C "${DEST}" rev-parse HEAD 2>/dev/null)" == "${COMMIT}" ]]; then
+  echo "seed-cutlass: already seeded"
+  exit 0
+fi
+
+mkdir -p "${CACHE_ROOT}" || exit 0
+rm -rf "${DEST}.partial"
+
+# ---------------------------------------------------------------------
+# Instrumentation. Two levels, both best effort: none of it may fail a
+# build, so every probe is bounded by a timeout and every failure is
+# swallowed.
+#
+#   probe_remote  runs on EVERY seeding run, successful ones included,
+#                 so consecutive builds form a comparable series. A
+#                 failure observed alone is nearly unreadable; the same
+#                 failure beside twenty healthy samples is evidence.
+#   diagnose      runs once, on the first failed attempt, adding the
+#                 slower checks only worth their time once something has
+#                 already gone wrong.
+#
+# What each field is for:
+#
+#   server, x-github-request-id  whether GitHub answered at all. A proxy
+#                                or captive portal returning 401 on its
+#                                behalf carries neither.
+#   x-github-edge-region, peer   which edge served us, so a fault
+#                                specific to one edge becomes visible.
+#   x-ratelimit-*                whether this egress address is being
+#                                throttled, measured rather than assumed.
+#   response body                distinguishes "Repository not found"
+#                                from any other refusal.
+#   tls issuer                   a re-signed certificate means something
+#                                is terminating TLS in the path.
+#   egress                       the identity any per-source limit is
+#                                applied to, and the only field that
+#                                makes failures on different runners
+#                                comparable. Costs one call to a third
+#                                party, which is why it is in the
+#                                failure path and not the common one.
+# ---------------------------------------------------------------------
+REFS_URL="${CUTLASS_REPO%.git}.git/info/refs?service=git-upload-pack"
+
+probe_remote() {
+  fmt='http=%{http_code} peer=%{remote_ip}:%{remote_port} tls_verify=%{ssl_verify_result}'
+  fmt="${fmt} connects=%{num_connects} redirects=%{num_redirects}"
+  fmt="${fmt} bytes=%{size_download} time=%{time_total}s"
+  # One request, not two: this runs on every build, and doubling the
+  # traffic to the endpoint under suspicion would corrupt what it
+  # measures. Headers land in a file so the summary and the identity
+  # come from the same response.
+  hdr=$(mktemp 2>/dev/null) || return 0
+  echo "seed-cutlass: probe $(curl -sS -o /dev/null -D "${hdr}" --max-time 20 \
+    -w "${fmt}" "${REFS_URL}" 2>&1 | tr '\n' ' ')" >&2
+  tr -d '\r' < "${hdr}" \
+    | grep -iE '^(HTTP/|server:|www-authenticate:|retry-after:|x-github-request-id:|x-github-edge-region:|x-ratelimit-)' \
+    | sed 's/^/seed-cutlass:   /' >&2 || true
+  rm -f "${hdr}"
+}
+
+diagnose() {
+  echo "seed-cutlass: --- diagnostics ---" >&2
+  echo "seed-cutlass: $(git --version 2>/dev/null || echo 'git --version failed')" \
+       "| $(uname -sr 2>/dev/null) | host $(hostname 2>/dev/null)" >&2
+
+  echo "seed-cutlass: dns $(getent ahosts github.com 2>/dev/null \
+    | awk '{print $1}' | sort -u | tr '\n' ' ')" >&2
+  echo "seed-cutlass: egress $(curl -sS --max-time 10 https://api.ipify.org 2>/dev/null \
+    || echo unknown)" >&2
+
+  proxies=$(env | grep -Ei '^(https?_proxy|all_proxy|no_proxy)=' | tr '\n' ' ')
+  echo "seed-cutlass: proxy env ${proxies:-none}" >&2
+
+  # `env -u` restores the real git configuration for this inspection
+  # only; the fetch itself still runs with config disabled.
+  relevant=$(env -u GIT_CONFIG_GLOBAL -u GIT_CONFIG_SYSTEM -u GIT_CONFIG_COUNT \
+    git config --list --show-origin 2>/dev/null \
+    | grep -Ei 'extraheader|insteadof|credential|proxy|askpass|http\.' | tr '\n' ' ')
+  echo "seed-cutlass: git config of interest ${relevant:-none}" >&2
+
+  echo "seed-cutlass: what this egress has left on the github api:" >&2
+  curl -sS -o /dev/null -D - --max-time 15 https://api.github.com/rate_limit 2>&1 \
+    | tr -d '\r' | grep -iE '^(HTTP/|x-ratelimit-)' \
+    | sed 's/^/seed-cutlass:   /' >&2 || true
+
+  echo "seed-cutlass: tls path:" >&2
+  curl -sSv -o /dev/null --max-time 20 "${REFS_URL}" 2>&1 \
+    | grep -iE 'subject:|issuer:|SSL connection using|ALPN, server accepted|subjectAltName' \
+    | sed 's/^[*] *//; s/^/seed-cutlass:   /' >&2 || true
+
+  echo "seed-cutlass: refusal body: $(curl -sS --max-time 20 "${REFS_URL}" 2>&1 \
+    | head -c 200 | tr '\n' ' ')" >&2
+
+  # Replay the fetch under git's own HTTP tracing. Everything above
+  # inspects a request of our choosing — a GET of the ref advertisement.
+  # The exchange that actually fails is the POST to git-upload-pack that
+  # follows it, and anything in the path may treat the two differently,
+  # so a clean GET proves very little. This shows which request was
+  # refused and what came back.
+  #
+  # It also settles, without inference, whether anything attached
+  # credentials: git logs an Authorization header here as <redacted>, so
+  # the line appearing at all means one was sent.
+  echo "seed-cutlass: traced replay of the real fetch:" >&2
+  rm -rf "${DEST}.trace"
+  if git init -q "${DEST}.trace" 2>/dev/null \
+    && git -C "${DEST}.trace" remote add origin "${CUTLASS_REPO}" 2>/dev/null; then
+    GIT_TRACE_CURL=1 GIT_TRACE_CURL_NO_DATA=1 GIT_TRACE_REDACT=1 \
+      git -C "${DEST}.trace" fetch --depth 1 origin "${COMMIT}" 2>&1 \
+      | grep -iE 'Send header: (GET|POST)|Recv header: (HTTP/|server:|www-authenticate:|x-github|content-type:|retry-after:)|Authorization|fatal:|error:' \
+      | sed 's/.*\(Send header\|Recv header\)/\1/; s/^/seed-cutlass:   /' \
+      | head -40 >&2 || true
+  fi
+  rm -rf "${DEST}.trace"
+
+  # Which wire protocol was refused? git >= 2.26 defaults to v2, where
+  # the ref listing is a POST to git-upload-pack carrying a
+  # Git-Protocol header; v0 is a plain GET of /info/refs first. libgit2
+  # — which cargo uses, and which has been observed reaching this same
+  # host seconds before git failed — still speaks v0. If v0 succeeds
+  # here while v2 fails, the fault is in how something on the path
+  # handles the newer protocol, not in access to the repository, and
+  # pinning protocol.version=0 for this fetch is the whole fix.
+  for v in 0 2; do
+    rm -rf "${DEST}.proto"
+    if git init -q "${DEST}.proto" 2>/dev/null \
+      && git -C "${DEST}.proto" remote add origin "${CUTLASS_REPO}" 2>/dev/null; then
+      if git -c protocol.version="${v}" -C "${DEST}.proto" \
+           fetch -q --depth 1 origin "${COMMIT}" 2>/dev/null; then
+        echo "seed-cutlass: protocol v${v}: OK" >&2
+      else
+        echo "seed-cutlass: protocol v${v}: FAILED" >&2
+      fi
+    fi
+    rm -rf "${DEST}.proto"
+  done
+
+  echo "seed-cutlass: --- end diagnostics ---" >&2
+}
+
+# Unconditional: one sample per run, so the failures have a baseline.
+probe_remote
+
+for attempt in 1 2 3 4 5; do
+  if git init -q "${DEST}.partial" \
+    && git -C "${DEST}.partial" remote add origin "${CUTLASS_REPO}" \
+    && git -C "${DEST}.partial" config core.sparseCheckout true \
+    && printf 'include/\ntools/util/include/\n' \
+         > "${DEST}.partial/.git/info/sparse-checkout" \
+    && git -C "${DEST}.partial" fetch -q --depth 1 origin "${COMMIT}" \
+    && git -C "${DEST}.partial" checkout -q FETCH_HEAD; then
+    rm -rf "${DEST}"
+    mv "${DEST}.partial" "${DEST}"
+    echo "seed-cutlass: seeded at attempt ${attempt}"
+    exit 0
+  fi
+  rm -rf "${DEST}.partial"
+  if [[ "${attempt}" == 1 ]]; then
+    diagnose
+  fi
+  backoff=$((attempt * 15))
+  echo "seed-cutlass: attempt ${attempt} failed; retrying in ${backoff}s" >&2
+  sleep "${backoff}"
+done
+
+echo "seed-cutlass: could not seed after 5 attempts; leaving it to the build" >&2
+exit 0