doc(qwen4_exp): settle the MTP wiring, and correct a norm this got wrong

§8 said every 10240-wide norm in the model is grouped over 2560 and
named mtp.pre_fc_norm_hidden among them. It is not. vLLM applies a
plain GemmaRMSNorm over the whole 10240 there, while using its
GroupedGemmaRMSNorm for the hyper-connection — both classes exist in
that tree, so the distinction is deliberate. Grouping it would rescale
the draft head's input in four independent pieces and surface as a poor
acceptance rate rather than as an error, which is the kind of wrong this
document exists to prevent.

§9 now carries the wiring instead of describing it as clear and leaving
it. Three parts of it are not recoverable from the tensor shapes: the
norm above; that the hidden state is the pre-final-mixer multi-stream
rather than the 2560 the LM head sees; and that the head emits both, the
collapsed stream for the head and the multi stream for its own next
step.

It also separates the wiring from the accept rule, which §9 had
conflated. The accept rule is greedy prefix match or the Leviathan
scheme and is not specific to this architecture.

Read from vllm/models/qwen4_exp/nvidia/mtp.py at 76ba321.

Refs #307, #313.

.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

@@ -334,6 +334,53 @@ jobs:
             [ "${desired}" = "$(cat "${conf}" 2>/dev/null || true)" ] || config_changed=1
           fi
 
+          # ── 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) ──
           # `rpm -q` prints "package X is not installed" on STDOUT and exits
           # 1, so `|| echo not-installed` never fires — it captures the
@@ -375,6 +422,17 @@ jobs:
             echo "applied ${conf}: NEURON_MAX_PROMPT_TOKENS=${max_prompt_tokens}"
           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
             sudo /usr/bin/systemctl stop neuron.service
           fi

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

@@ -15,9 +15,36 @@ use std::path::{Path, PathBuf};
 pub const DEFAULT_SOURCE_SCHEME: &str = "huggingface";
 
 /// Endpoint URL for the default huggingface source, used when no
-/// `[harness.candle.sources.huggingface]` is configured.
+/// `[harness.candle.sources.huggingface]` is configured and no
+/// `HF_ENDPOINT` is set in the environment.
 pub const DEFAULT_HF_ENDPOINT: &str = "https://huggingface.co";
 
+/// The ecosystem-standard environment variable for redirecting HF
+/// traffic at a mirror.
+pub const HF_ENDPOINT_ENV: &str = "HF_ENDPOINT";
+
+/// Resolve the endpoint for the synthesised `huggingface` source.
+///
+/// `hf-hub` only consults `HF_ENDPOINT` from `ApiBuilder::from_env()`,
+/// and we build with `ApiBuilder::new()` and then set the endpoint
+/// explicitly from config — so without this the variable is read by
+/// nobody and an operator who sets it gets canonical huggingface.co
+/// anyway, silently. Honouring it here keeps that from being a trap.
+///
+/// Precedence is explicit config > `HF_ENDPOINT` > `huggingface.co`;
+/// this function covers the last two, since an operator-supplied
+/// `[harness.candle.sources.huggingface]` block is never overwritten.
+///
+/// Blank or whitespace-only values count as unset, and a trailing
+/// slash is trimmed — `hf-hub` interpolates `{endpoint}/{repo}/...`,
+/// so `https://rf.internal/` would otherwise produce a doubled slash.
+pub fn resolve_hf_endpoint(from_env: Option<String>) -> String {
+    match from_env {
+        Some(raw) if !raw.trim().is_empty() => raw.trim().trim_end_matches('/').to_string(),
+        _ => DEFAULT_HF_ENDPOINT.to_string(),
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct NeuronConfig {
     #[serde(default = "default_port")]
@@ -67,9 +94,14 @@ pub struct CandleHarnessConfig {
     /// directory.
     ///
     /// When absent or missing the `huggingface` key, the loader
-    /// synthesises a `huggingface` entry pointing at
-    /// `https://huggingface.co` with `hf_cache` (above) as its
-    /// cache_dir. This keeps single-source configs ergonomic.
+    /// synthesises a `huggingface` entry with `hf_cache` (above) as its
+    /// cache_dir, pointing at `HF_ENDPOINT` if that is set in the
+    /// environment and `https://huggingface.co` otherwise. This keeps
+    /// single-source configs ergonomic, and lets a whole host be
+    /// redirected at a mirror with one environment variable.
+    ///
+    /// An entry written here always wins over `HF_ENDPOINT` — an
+    /// operator who has typed a sources block means it.
     #[serde(default)]
     pub sources: HashMap<String, SourceConfig>,
 
@@ -450,10 +482,22 @@ impl CandleHarnessConfig {
     /// (operator-typed) config can still be serialized back to TOML
     /// for diagnostics.
     pub fn effective_sources(&self) -> HashMap<String, SourceConfig> {
+        self.effective_sources_with_hf_endpoint(resolve_hf_endpoint(
+            std::env::var(HF_ENDPOINT_ENV).ok(),
+        ))
+    }
+
+    /// [`Self::effective_sources`] with the huggingface endpoint passed
+    /// in, so the precedence rules are testable without mutating
+    /// process-wide environment state from a parallel test runner.
+    pub fn effective_sources_with_hf_endpoint(
+        &self,
+        hf_endpoint: String,
+    ) -> HashMap<String, SourceConfig> {
         let mut out = self.sources.clone();
         out.entry(DEFAULT_SOURCE_SCHEME.to_string())
             .or_insert_with(|| SourceConfig {
-                endpoint: DEFAULT_HF_ENDPOINT.to_string(),
+                endpoint: hf_endpoint,
                 auth_env: Some("HF_TOKEN".to_string()),
                 cache_dir: self.hf_cache.clone(),
             });
@@ -598,4 +642,63 @@ mod tests {
         };
         assert_eq!(cfg.effective_default_source(), "helexa");
     }
+
+    #[test]
+    fn hf_endpoint_env_redirects_the_synthesised_source() {
+        let cfg = CandleHarnessConfig::default();
+        let sources = cfg.effective_sources_with_hf_endpoint(resolve_hf_endpoint(Some(
+            "https://rf.internal".to_string(),
+        )));
+        assert_eq!(sources["huggingface"].endpoint, "https://rf.internal");
+        // The token still comes from HF_TOKEN, so a mirror requiring
+        // auth works without further config.
+        assert_eq!(sources["huggingface"].auth_env.as_deref(), Some("HF_TOKEN"));
+    }
+
+    /// Explicit config beats the environment. An operator who has
+    /// written a sources block means it.
+    #[test]
+    fn explicit_source_config_wins_over_hf_endpoint_env() {
+        let cfg = CandleHarnessConfig {
+            sources: HashMap::from([(
+                "huggingface".to_string(),
+                SourceConfig {
+                    endpoint: "https://pinned.example.org".into(),
+                    ..Default::default()
+                },
+            )]),
+            ..Default::default()
+        };
+        let sources = cfg.effective_sources_with_hf_endpoint(resolve_hf_endpoint(Some(
+            "https://rf.internal".to_string(),
+        )));
+        assert_eq!(
+            sources["huggingface"].endpoint,
+            "https://pinned.example.org"
+        );
+    }
+
+    #[test]
+    fn hf_endpoint_falls_back_to_canonical_when_unset_or_blank() {
+        assert_eq!(resolve_hf_endpoint(None), DEFAULT_HF_ENDPOINT);
+        assert_eq!(
+            resolve_hf_endpoint(Some(String::new())),
+            DEFAULT_HF_ENDPOINT
+        );
+        assert_eq!(resolve_hf_endpoint(Some("   ".into())), DEFAULT_HF_ENDPOINT);
+    }
+
+    /// hf-hub interpolates `{endpoint}/{repo}/resolve/...`, so a
+    /// trailing slash would produce a doubled one in every URL.
+    #[test]
+    fn hf_endpoint_trailing_slash_is_trimmed() {
+        assert_eq!(
+            resolve_hf_endpoint(Some("https://rf.internal/".into())),
+            "https://rf.internal"
+        );
+        assert_eq!(
+            resolve_hf_endpoint(Some("  https://rf.internal//  ".into())),
+            "https://rf.internal"
+        );
+    }
 }

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

@@ -21,3 +21,4 @@
 //! the pattern set by `tp_qwen3.rs`.
 
 pub mod qwen3_5;
+pub mod qwen4_exp;

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,8 +17,11 @@ pub struct Qwen3_5MLP {
 }
 
 impl Qwen3_5MLP {
-    /// Construct directly from pre-built projections (MoE-block tests).
-    #[cfg(test)]
+    /// Construct directly from pre-built projections.
+    ///
+    /// `qwen4_exp` ships its routed experts as fused 3D tensors rather
+    /// than per-expert modules, so its loader slices them and builds the
+    /// experts here — same SwiGLU, different storage.
     pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
         Self {
             gate_proj,

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

@@ -52,6 +52,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: Vec<Qwen3_5MLP>,
+        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,

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

@@ -1,4 +1,4 @@
-//! Norm primitives for Qwen3-Next.
+//! Norm primitives for Qwen3-Next, shared with `qwen4_exp`.
 //!
 //! Two reasons we can't reuse `candle_nn::RmsNorm` directly:
 //!
@@ -11,13 +11,25 @@
 //!    introduced the `(1 + w)` form to recover from the zero-init.
 //!
 //! 2. **Gated variant.** The linear-attention layer post-normalises
-//!    its output by an RMSNorm *gated* with a per-element SiLU on
-//!    a sibling `z` projection — fused for numerical reasons (the
-//!    norm's float32 promotion has to happen before the SiLU
+//!    its output by an RMSNorm *gated* with a per-element nonlinearity
+//!    on a sibling `z` projection — fused for numerical reasons (the
+//!    norm's float32 promotion has to happen before the gate
 //!    multiply). Not a single existing candle op.
 //!
 //! Both ops accept inputs in any compute dtype; promotion to f32 for
 //! the variance calculation matches the Python reference.
+//!
+//! `qwen4_exp` (`doc/qwen4_exp-port-spec.md`) reuses both with two
+//! additions, because the conventions are otherwise identical:
+//!
+//! - **`group_size`** on the plain norm, for the 10240-wide norms whose
+//!   four hidden-size slices normalise independently.
+//! - **`OutputGate`** on the gated norm, because that arch sets
+//!   `output_gate_type = "sigmoid"` where Qwen3.6 falls back to SiLU.
+//!
+//! Both default to the Qwen3.6 behaviour, so existing callers — the
+//! decoder, the attention q/k norms, and the TP analogues in
+//! `harness/tp/tp_qwen3_5.rs` — are unaffected.
 
 use anyhow::{Context, Result};
 use candle_core::{D, Module, Tensor};
@@ -39,16 +51,52 @@ pub fn l2norm(x: &Tensor, eps: f32) -> candle_core::Result<Tensor> {
 
 /// Qwen3-Next's RMSNorm. Stores the raw weight tensor; forward applies
 /// `(1.0 + weight) * x_normed`.
+///
+/// With `group_size` set, the last axis is split into
+/// `size / group_size` groups that are normalised **independently**
+/// before the scale is applied across the whole axis — upstream's
+/// `Qwen4ExpTextRMSNorm(dim, group_size=...)`. `qwen4_exp` uses this
+/// for every 10240-wide norm (`hc_norm`, PLE's three norms,
+/// `mtp.pre_fc_norm_hidden`), all with `group_size = hidden_size`, i.e.
+/// four independent normalisations rather than one over 10240. A single
+/// norm over the full axis compiles, runs, and is quietly wrong — see
+/// `doc/qwen4_exp-port-spec.md` §8.
 pub struct Qwen3_5RmsNorm {
     weight: Tensor,
     eps: f32,
     size: usize,
+    /// `None` = one normalisation over the whole last axis (Qwen3.6).
+    group_size: Option<usize>,
 }
 
 impl Qwen3_5RmsNorm {
     /// Load `weight` from the ShardedVarBuilder. `vb` should already be
     /// `.pp(...)`-ed to the norm's tensor prefix.
     pub fn load(vb: &ShardedVarBuilder, size: usize, eps: f64) -> Result<Self> {
+        Self::load_inner(vb, size, eps, None)
+    }
+
+    /// Grouped variant: normalise each `group_size`-wide slice of the
+    /// last axis independently.
+    pub fn load_grouped(
+        vb: &ShardedVarBuilder,
+        size: usize,
+        group_size: usize,
+        eps: f64,
+    ) -> Result<Self> {
+        anyhow::ensure!(
+            group_size > 0 && size.is_multiple_of(group_size),
+            "grouped RMSNorm needs size ({size}) divisible by group_size ({group_size})"
+        );
+        Self::load_inner(vb, size, eps, Some(group_size))
+    }
+
+    fn load_inner(
+        vb: &ShardedVarBuilder,
+        size: usize,
+        eps: f64,
+        group_size: Option<usize>,
+    ) -> Result<Self> {
         let weight = vb
             .get(size, "weight")
             .with_context(|| format!("load '{}/weight'", vb.prefix()))?;
@@ -56,9 +104,23 @@ impl Qwen3_5RmsNorm {
             weight,
             eps: eps as f32,
             size,
+            group_size,
         })
     }
 
+    /// Direct constructor — used by unit tests that build a norm
+    /// without going through a VarBuilder.
+    #[cfg(test)]
+    pub(crate) fn from_weight(weight: Tensor, eps: f64, group_size: Option<usize>) -> Self {
+        let size = weight.dims()[0];
+        Self {
+            weight,
+            eps: eps as f32,
+            size,
+            group_size,
+        }
+    }
+
     pub fn size(&self) -> usize {
         self.size
     }
@@ -68,8 +130,28 @@ impl Module for Qwen3_5RmsNorm {
     fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
         let dtype = x.dtype();
         let x_f32 = x.to_dtype(candle_core::DType::F32)?;
-        let var = x_f32.sqr()?.mean_keepdim(D::Minus1)?;
-        let normed = x_f32.broadcast_mul(&(var + self.eps as f64)?.sqrt()?.recip()?)?;
+        // Grouped: reshape the last axis into (groups, group_size),
+        // normalise over group_size, then flatten back. The scale below
+        // still applies across the full axis.
+        let normed = match self.group_size {
+            Some(gs) => {
+                let rank = x_f32.rank();
+                let mut shape = x_f32.dims().to_vec();
+                let last = shape.pop().expect("tensor has at least one axis");
+                shape.push(last / gs);
+                shape.push(gs);
+                let grouped = x_f32.reshape(shape)?;
+                let var = grouped.sqr()?.mean_keepdim(D::Minus1)?;
+                let normed = grouped.broadcast_mul(&(var + self.eps as f64)?.sqrt()?.recip()?)?;
+                // rank-1 is the groups axis; flattening from there
+                // restores the original rank.
+                normed.flatten_from(rank - 1)?
+            }
+            None => {
+                let var = x_f32.sqr()?.mean_keepdim(D::Minus1)?;
+                x_f32.broadcast_mul(&(var + self.eps as f64)?.sqrt()?.recip()?)?
+            }
+        };
         // Promote weight to f32 and shift by 1.0 *before* multiplying.
         // Doing the (1 + w) operation in fp16 lands at -inf for the
         // bottom-of-range weights at load time.
@@ -90,10 +172,49 @@ pub struct Qwen3_5RmsNormGated {
     weight: Tensor,
     eps: f32,
     size: usize,
+    activation: OutputGate,
+}
+
+/// Which nonlinearity the output gate applies. Qwen3.6 leaves
+/// `output_gate_type` unset and falls back to `hidden_act` (SiLU);
+/// `qwen4_exp` sets it to `sigmoid` explicitly. Inheriting SiLU there
+/// would be wrong on all 36 linear-attention layers, and wrong in a way
+/// that still produces fluent output.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum OutputGate {
+    Silu,
+    Sigmoid,
+}
+
+impl OutputGate {
+    /// Map a config's `output_gate_type` (or `hidden_act` fallback).
+    pub fn from_name(name: &str) -> Result<Self> {
+        match name {
+            "silu" => Ok(Self::Silu),
+            "sigmoid" => Ok(Self::Sigmoid),
+            other => anyhow::bail!("unsupported output_gate_type '{other}'"),
+        }
+    }
+
+    fn apply(self, gate: &Tensor) -> candle_core::Result<Tensor> {
+        match self {
+            Self::Silu => candle_nn::ops::silu(gate),
+            Self::Sigmoid => candle_nn::ops::sigmoid(gate),
+        }
+    }
 }
 
 impl Qwen3_5RmsNormGated {
     pub fn load(vb: &ShardedVarBuilder, size: usize, eps: f64) -> Result<Self> {
+        Self::load_with_gate(vb, size, eps, OutputGate::Silu)
+    }
+
+    pub fn load_with_gate(
+        vb: &ShardedVarBuilder,
+        size: usize,
+        eps: f64,
+        activation: OutputGate,
+    ) -> Result<Self> {
         let weight = vb
             .get(size, "weight")
             .with_context(|| format!("load '{}/weight'", vb.prefix()))?;
@@ -101,6 +222,7 @@ impl Qwen3_5RmsNormGated {
             weight,
             eps: eps as f32,
             size,
+            activation,
         })
     }
 
@@ -113,6 +235,7 @@ impl Qwen3_5RmsNormGated {
             weight,
             eps: eps as f32,
             size,
+            activation: OutputGate::Silu,
         }
     }
 
@@ -128,11 +251,11 @@ impl Qwen3_5RmsNormGated {
         let normed = x_f32.broadcast_mul(&(var + self.eps as f64)?.sqrt()?.recip()?)?;
         let w = self.weight.to_dtype(candle_core::DType::F32)?;
         let out = normed.broadcast_mul(&w)?;
-        // SiLU on the float32 gate, multiply back into the normed
+        // Activate the float32 gate, multiply back into the normed
         // tensor, then cast to the model dtype.
         let g = gate.to_dtype(candle_core::DType::F32)?;
-        let silu_gate = candle_nn::ops::silu(&g)?;
-        (out * silu_gate)?.to_dtype(dtype)
+        let activated = self.activation.apply(&g)?;
+        (out * activated)?.to_dtype(dtype)
     }
 }
 
@@ -158,4 +281,98 @@ mod tests {
         let v: Vec<f32> = out.to_vec1().unwrap();
         assert!(v.iter().all(|x| x.is_finite()));
     }
+
+    /// Grouped norm over `[1, 1, 4]` with `group_size = 2`: each pair is
+    /// normalised on its own RMS, so the second pair — an order of
+    /// magnitude larger — comes back at the same scale as the first.
+    /// Ungrouped, the shared RMS would leave the pairs far apart, which
+    /// is exactly the silent failure this guards.
+    #[test]
+    fn grouped_rmsnorm_normalises_each_group_independently() {
+        let dev = Device::Cpu;
+        let x = Tensor::new(&[[[3.0_f32, 4.0, 30.0, 40.0]]], &dev).unwrap();
+        let w = Tensor::zeros(4, candle_core::DType::F32, &dev).unwrap();
+
+        let grouped = Qwen3_5RmsNorm::from_weight(w.clone(), 1e-6, Some(2));
+        let got: Vec<f32> = grouped
+            .forward(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // RMS([3,4]) = sqrt(12.5); RMS([30,40]) = sqrt(1250).
+        let a = (12.5_f32).sqrt();
+        let b = (1250.0_f32).sqrt();
+        let want = [3.0 / a, 4.0 / a, 30.0 / b, 40.0 / b];
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-4, "grouped: got {got:?} want {want:?}");
+        }
+        // Both groups land at the same magnitude — the whole point.
+        assert!((got[0] - got[2]).abs() < 1e-4);
+    }
+
+    #[test]
+    fn ungrouped_rmsnorm_is_unchanged_by_the_group_option() {
+        let dev = Device::Cpu;
+        let x = Tensor::new(&[[[3.0_f32, 4.0, 30.0, 40.0]]], &dev).unwrap();
+        let w = Tensor::zeros(4, candle_core::DType::F32, &dev).unwrap();
+
+        let plain = Qwen3_5RmsNorm::from_weight(w, 1e-6, None);
+        let got: Vec<f32> = plain
+            .forward(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // One RMS over all four: sqrt((9+16+900+1600)/4).
+        let rms = ((9.0_f32 + 16.0 + 900.0 + 1600.0) / 4.0).sqrt();
+        let want = [3.0 / rms, 4.0 / rms, 30.0 / rms, 40.0 / rms];
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-4, "plain: got {got:?} want {want:?}");
+        }
+    }
+
+    /// The grouped reshape must not disturb the leading axes.
+    #[test]
+    fn grouped_rmsnorm_preserves_shape() {
+        let dev = Device::Cpu;
+        let x = Tensor::rand(0f32, 1f32, (2, 3, 8), &dev).unwrap();
+        let w = Tensor::zeros(8, candle_core::DType::F32, &dev).unwrap();
+        let n = Qwen3_5RmsNorm::from_weight(w, 1e-6, Some(4));
+        assert_eq!(n.forward(&x).unwrap().dims(), &[2, 3, 8]);
+    }
+
+    #[test]
+    fn output_gate_sigmoid_differs_from_silu() {
+        let dev = Device::Cpu;
+        let g = Tensor::new(&[1.0_f32, -2.0, 0.5], &dev).unwrap();
+
+        let sig: Vec<f32> = OutputGate::Sigmoid.apply(&g).unwrap().to_vec1().unwrap();
+        let silu: Vec<f32> = OutputGate::Silu.apply(&g).unwrap().to_vec1().unwrap();
+
+        // sigmoid(1) = 0.7311; silu(1) = 1 * 0.7311 = 0.7311 — equal at
+        // x = 1, so check a point where they must diverge.
+        assert!(
+            (sig[1] - 0.11920292).abs() < 1e-5,
+            "sigmoid(-2) = {}",
+            sig[1]
+        );
+        assert!(
+            (silu[1] + 0.23840584).abs() < 1e-5,
+            "silu(-2) = {}",
+            silu[1]
+        );
+        assert!((sig[1] - silu[1]).abs() > 0.3);
+    }
+
+    #[test]
+    fn output_gate_rejects_unknown_names() {
+        assert!(OutputGate::from_name("silu").is_ok());
+        assert!(OutputGate::from_name("sigmoid").is_ok());
+        assert!(OutputGate::from_name("gelu").is_err());
+    }
 }

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,302 @@
+//! `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 super::config::TextConfig;
+use super::full_attn::Attention;
+use super::hyper::HyperConnection;
+use super::ple::PleBlock;
+use super::{linear_attn, moe};
+
+/// 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,
+        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()
+                )
+            })?;
+        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}' for layer {layer_idx} (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 cfg.ple_layers().contains(&layer_idx) {
+            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, &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(())
+    }
+
+    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,353 @@
+//! `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)?)
+    }
+
+    /// 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

@@ -0,0 +1,336 @@
+//! Hyper-connections: `qwen4_exp`'s replacement for the pre-norm
+//! residual stream.
+//!
+//! Every architecture we serve today carries one residual stream of
+//! `hidden_size` and wraps each sublayer as
+//! `x = x + sublayer(norm(x))`. `qwen4_exp` carries **four** streams,
+//! concatenated into a single `hc_count * hidden_size` tensor, and
+//! there is 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.
+//!
+//! Per sublayer:
+//!
+//! ```text
+//! h [.., 10240]
+//!   │
+//!   ├─ hn   = hc_norm(h)                       grouped over 4x2560
+//!   ├─ w    = sigmoid(up(silu(down(hn) / 4)))  10240 -> 320 -> 10240
+//!   ├─ x    = mean over the 4 streams of (w * hn)        -> [.., 2560]
+//!   ├─ inj  = 2 * sigmoid(block_inject(hn) / 4)          -> [.., 4]
+//!   │
+//!   │  y = sublayer(x)                                   -> [.., 2560]
+//!   │
+//!   └─ h' = h + (y (x) inj).flatten(-2)                  -> [.., 10240]
+//! ```
+//!
+//! Three details are easy to get subtly wrong, and all three produce
+//! fluent-but-worse output rather than a crash, so each is pinned by a
+//! test below:
+//!
+//! 1. Both gates divide by `hc_count` **before** the nonlinearity.
+//! 2. The stream mix is a `mean`, not a `sum`.
+//! 3. The tensor carried into the residual add is the **un-normalised**
+//!    `h`. Only the mixing and injection paths see `hn`.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §1.
+
+use anyhow::{Context, Result};
+use candle_core::{D, Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm;
+
+/// What a hyper-connection hands its sublayer.
+pub struct HyperSplit {
+    /// The mixed, normalised input for the sublayer: `[.., hidden]`.
+    pub input: Tensor,
+    /// The **un-normalised** stream tensor to add back: `[.., hidden*hc]`.
+    pub residual: Tensor,
+    /// Per-stream injection weights: `[.., hc]`. `None` on the final
+    /// mixer, which has no `block_inject_weight` and never rejoins.
+    pub inject: Option<Tensor>,
+}
+
+/// `Qwen4ExpTextGatedResidual`.
+pub struct HyperConnection {
+    hc_norm: Qwen3_5RmsNorm,
+    mix_down: Linear,
+    mix_up: Linear,
+    /// Absent on the top-level `hyper_connection_mixer`, which collapses
+    /// the streams for the LM head instead of feeding a sublayer.
+    block_inject: Option<Linear>,
+    hidden_size: usize,
+    hc_count: usize,
+}
+
+impl HyperConnection {
+    pub fn load(
+        vb: &ShardedVarBuilder,
+        hidden_size: usize,
+        hc_count: usize,
+        hc_lowrank: usize,
+        eps: f64,
+        use_combine: bool,
+    ) -> Result<Self> {
+        let wide = hidden_size * hc_count;
+        let hc_norm = Qwen3_5RmsNorm::load_grouped(&vb.pp("hc_norm"), wide, hidden_size, eps)
+            .with_context(|| format!("load '{}/hc_norm'", vb.prefix()))?;
+        let mix_down = linear(vb, "input_mix_weight_down", wide, hc_lowrank)?;
+        let mix_up = linear(vb, "input_mix_weight_up", hc_lowrank, wide)?;
+        let block_inject = if use_combine {
+            Some(linear(vb, "block_inject_weight", wide, hc_count)?)
+        } else {
+            None
+        };
+        Ok(Self {
+            hc_norm,
+            mix_down,
+            mix_up,
+            block_inject,
+            hidden_size,
+            hc_count,
+        })
+    }
+
+    /// 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> {
+        let wide = self.hidden_size * self.hc_count;
+        debug_assert_eq!(h.dims()[h.rank() - 1], wide);
+
+        let hn = self.hc_norm.forward(h)?;
+        let hc = self.hc_count as f64;
+
+        // silu(down(hn) / hc) -> sigmoid(up(.)) : a low-rank per-stream,
+        // per-channel gate over the normalised streams.
+        let w = self.mix_down.forward(&hn)?;
+        let w = candle_nn::ops::silu(&(w / hc)?)?;
+        let w = candle_nn::ops::sigmoid(&self.mix_up.forward(&w)?)?;
+
+        // Mean over the stream axis — NOT a sum.
+        let input = (unflatten_streams(&w, self.hc_count, self.hidden_size)?
+            * unflatten_streams(&hn, self.hc_count, self.hidden_size)?)?
+        .mean(D::Minus2)?;
+
+        let inject = match &self.block_inject {
+            Some(bi) => {
+                let raw = (bi.forward(&hn)? / hc)?;
+                Some((candle_nn::ops::sigmoid(&raw)? * 2.0)?)
+            }
+            None => None,
+        };
+
+        Ok(HyperSplit {
+            input,
+            residual: h.clone(),
+            inject,
+        })
+    }
+
+    /// Scatter a sublayer's output back across the streams and add it to
+    /// the residual.
+    pub fn combine(
+        &self,
+        residual: &Tensor,
+        block_out: &Tensor,
+        inject: &Tensor,
+    ) -> candle_core::Result<Tensor> {
+        // outer product: [.., 1, hidden] * [.., hc, 1] -> [.., hc, hidden]
+        let scattered = block_out
+            .unsqueeze(D::Minus2)?
+            .broadcast_mul(&inject.unsqueeze(D::Minus1)?)?;
+        residual.add(&scattered.flatten_from(D::Minus2)?)
+    }
+
+    /// The top-level collapse before the LM head: mix the streams down
+    /// to one `hidden_size` vector and discard the rest. There is no
+    /// separate final norm in this architecture; `hc_norm` is it.
+    pub fn collapse(&self, h: &Tensor) -> candle_core::Result<Tensor> {
+        Ok(self.split(h)?.input)
+    }
+}
+
+fn unflatten_streams(x: &Tensor, hc: usize, hidden: usize) -> candle_core::Result<Tensor> {
+    let mut shape = x.dims().to_vec();
+    shape.pop();
+    shape.push(hc);
+    shape.push(hidden);
+    x.reshape(shape)
+}
+
+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::Device;
+
+    /// Build a hyper-connection whose three projections are all zero.
+    /// That makes the whole gate analytic:
+    ///   down(hn) = 0 -> silu(0) = 0 -> up(0) = 0 -> sigmoid(0) = 0.5
+    ///   block_inject(hn) = 0 -> 2 * sigmoid(0) = 1.0
+    /// so the mix weight is exactly 0.5 everywhere and every stream is
+    /// 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 {
+        HyperConnection::zeroed_for_test(hidden, hc, lowrank, true)
+    }
+
+    #[test]
+    fn split_matches_hand_computed_forward() {
+        let dev = Device::Cpu;
+        let hc = zeroed(2, 2, 2);
+        // Two streams whose magnitudes differ 10x, so the grouped norm
+        // has to bring them to the same scale before mixing.
+        let h = Tensor::new(&[[[3.0f32, 4.0, 30.0, 40.0]]], &dev).unwrap();
+
+        let split = hc.split(&h).unwrap();
+
+        // hn = [3,4]/rms([3,4]) ++ [30,40]/rms([30,40]) — identical pairs.
+        let a = (12.5f32).sqrt();
+        // mix weight is 0.5, and the two normalised streams are equal,
+        // so mean over streams leaves them unchanged.
+        let want = [0.5 * 3.0 / a, 0.5 * 4.0 / a];
+        let got: Vec<f32> = split.input.flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(got.len(), 2);
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-4, "input: got {got:?} want {want:?}");
+        }
+
+        let inj: Vec<f32> = split
+            .inject
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        assert_eq!(inj.len(), 2);
+        for v in &inj {
+            assert!(
+                (v - 1.0).abs() < 1e-5,
+                "inject should be 2*sigmoid(0) = 1, got {inj:?}"
+            );
+        }
+    }
+
+    /// The residual handed back for the skip connection is the raw
+    /// input, not the normalised one. Getting this wrong still produces
+    /// fluent output, so assert it exactly.
+    #[test]
+    fn residual_is_the_unnormalised_input() {
+        let dev = Device::Cpu;
+        let hc = zeroed(2, 2, 2);
+        let h = Tensor::new(&[[[3.0f32, 4.0, 30.0, 40.0]]], &dev).unwrap();
+        let got: Vec<f32> = hc
+            .split(&h)
+            .unwrap()
+            .residual
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        assert_eq!(got, vec![3.0, 4.0, 30.0, 40.0]);
+    }
+
+    #[test]
+    fn combine_scatters_the_block_output_across_streams() {
+        let dev = Device::Cpu;
+        let hc = zeroed(2, 2, 2);
+        let residual = Tensor::new(&[[[3.0f32, 4.0, 30.0, 40.0]]], &dev).unwrap();
+        let block_out = Tensor::new(&[[[1.0f32, 2.0]]], &dev).unwrap();
+        // Distinct per-stream weights, so a transposed outer product
+        // would give a different answer.
+        let inject = Tensor::new(&[[[10.0f32, 100.0]]], &dev).unwrap();
+
+        let got: Vec<f32> = hc
+            .combine(&residual, &block_out, &inject)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // stream 0 gets 10*[1,2]; stream 1 gets 100*[1,2].
+        assert_eq!(
+            got,
+            vec![3.0 + 10.0, 4.0 + 20.0, 30.0 + 100.0, 40.0 + 200.0]
+        );
+    }
+
+    /// A round trip through split + combine must preserve the stream
+    /// width, or the next layer sees the wrong shape.
+    #[test]
+    fn split_combine_preserves_stream_width() {
+        let dev = Device::Cpu;
+        let hc = zeroed(4, 4, 3);
+        let h = Tensor::rand(0f32, 1f32, (2, 5, 16), &dev).unwrap();
+        let split = hc.split(&h).unwrap();
+        assert_eq!(split.input.dims(), &[2, 5, 4]);
+        assert_eq!(split.inject.as_ref().unwrap().dims(), &[2, 5, 4]);
+        let out = hc
+            .combine(
+                &split.residual,
+                &split.input,
+                split.inject.as_ref().unwrap(),
+            )
+            .unwrap();
+        assert_eq!(out.dims(), &[2, 5, 16]);
+    }
+
+    /// The final mixer has no block_inject_weight and collapses to one
+    /// hidden-size vector for the LM head.
+    #[test]
+    fn collapse_has_no_injection_and_narrows_to_hidden() {
+        let dev = Device::Cpu;
+        let mut hc = zeroed(4, 4, 3);
+        hc.block_inject = None;
+        let h = Tensor::rand(0f32, 1f32, (1, 3, 16), &dev).unwrap();
+        assert!(hc.split(&h).unwrap().inject.is_none());
+        assert_eq!(hc.collapse(&h).unwrap().dims(), &[1, 3, 4]);
+    }
+}

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

@@ -0,0 +1,43 @@
+//! `qwen4_exp` — the architecture behind `Qwen/Qwen3.8-Flash-Next`.
+//!
+//! A Qwen 4 experimental generation rather than a Qwen3 variant. The
+//! hybrid skeleton (36 linear-attention layers to 12 full-attention,
+//! `full_attention_interval: 4`) and the linear-attention tensor layout
+//! are shared with [`super::qwen3_5`], so much of that implementation
+//! carries over. What is new is the residual structure — four streams
+//! joined by hyper-connections, with no layernorms of the usual kind —
+//! plus QSA sparse attention and PLE hashed n-gram embeddings.
+//!
+//! The written specification lives at `doc/qwen4_exp-port-spec.md`;
+//! it is the reference for every dimension and dataflow here, and it
+//! records what was measured from the checkpoint rather than inferred.
+//!
+//! 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.
+//! - [`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 ple;
+pub mod qsa;

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

@@ -0,0 +1,606 @@
+//! `qwen4_exp` end to end: embeddings in, logits out.
+//!
+//! ```text
+//! input_ids ─ embed_tokens ─► x [B,T,2560]
+//!                              │  repeated hc_count times
+//!                              ▼
+//!                            h [B,T,10240]   four identical streams
+//!                              │
+//!                            48 decoder layers   (PLE joins at layer 1)
+//!                              │
+//!                            hyper_connection_mixer   collapses 4 → 1
+//!                              ▼
+//!                            lm_head
+//! ```
+//!
+//! Three things about this shape are unlike every other model we serve,
+//! and all three are visible in the checkpoint rather than inferred:
+//!
+//! 1. **The residual between layers is `hidden_size * hc_count` wide**,
+//!    initialised as four copies of the token embedding.
+//! 2. **There is no `model.norm`.** The index has
+//!    `hyper_connection_mixer.{hc_norm, input_mix_weight_down,
+//!    input_mix_weight_up}` and nothing else before `lm_head`, so the
+//!    mixer's own `hc_norm` is the only normalisation the logits see.
+//!    It also has **no `block_inject_weight`** — it collapses the
+//!    streams rather than scattering back into them.
+//! 3. **The n-gram embedding is computed once per step, not per
+//!    layer.** It depends only on the input ids, and exactly one layer
+//!    consumes it.
+//!
+//! The weights hang off two different roots: everything structural is
+//! under `model.language_model`, but `lm_head` is at the top level.
+//!
+//! See `doc/qwen4_exp-port-spec.md`.
+
+use anyhow::{Context, Result, ensure};
+use candle_core::{DType, Device, IndexOp, Module, Tensor};
+use candle_nn::var_builder::ShardedVarBuilder;
+use candle_nn::{Embedding, Linear};
+use std::sync::Arc;
+
+use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
+
+use super::config::{Config, TextConfig};
+use super::decoder::DecoderLayer;
+use super::hyper::HyperConnection;
+use super::ple::{NGramHasher, NGramTable, ShardedNGramTable};
+
+/// The hashed n-gram lookup, hoisted to the model because it is a
+/// function of the input ids alone.
+struct NGramEmbedding {
+    hasher: NGramHasher,
+    table: Box<dyn NGramTable>,
+    heads: usize,
+    head_dim: usize,
+    /// The previous `ngram_size - 1` ids per batch row, so a decode step
+    /// hashes the same n-grams a prefill would. Upstream keeps this as a
+    /// `(B, 2)` "conv state" on the PLE layer; here it is what it is —
+    /// token ids, not activations.
+    carried: Vec<Vec<i64>>,
+}
+
+impl NGramEmbedding {
+    /// `[B, T, heads * head_dim]` for this step's ids.
+    fn forward(&mut self, ids: &[Vec<i64>], device: &Device, dtype: DType) -> Result<Tensor> {
+        let seq_len = ids.first().map_or(0, Vec::len);
+        if self.carried.len() != ids.len() {
+            self.carried = vec![Vec::new(); ids.len()];
+        }
+        let mut rows = Vec::with_capacity(ids.len());
+        for (row, step) in ids.iter().enumerate() {
+            let mut history = self.carried[row].clone();
+            history.extend_from_slice(step);
+            let per_position = self.hasher.rows(&history, step.len())?;
+            let flat: Vec<i64> = per_position.into_iter().flatten().collect();
+            rows.push(
+                self.table
+                    .gather(&flat)?
+                    .reshape((step.len(), self.heads * self.head_dim))?,
+            );
+            // Keep only what the next step's shifts can still reach.
+            let keep = self.hasher.context_len().min(history.len());
+            self.carried[row] = history[history.len() - keep..].to_vec();
+        }
+        let stacked = Tensor::stack(&rows, 0)?
+            .to_dtype(dtype)?
+            .to_device(device)?;
+        debug_assert_eq!(stacked.dims()[1], seq_len);
+        Ok(stacked)
+    }
+
+    fn reset(&mut self) {
+        for row in &mut self.carried {
+            row.clear();
+        }
+    }
+}
+
+pub struct Qwen4ExpForCausalLM {
+    embed_tokens: Embedding,
+    layers: Vec<DecoderLayer>,
+    /// `hyper_connection_mixer` — `use_combine = false`, so it collapses
+    /// rather than scatters, and there is no final norm beside it.
+    mixer: HyperConnection,
+    lm_head: Linear,
+    ngram: Option<NGramEmbedding>,
+    rotary: Arc<RotaryEmbedding>,
+    hc_count: usize,
+    device: Device,
+    dtype: DType,
+}
+
+impl Qwen4ExpForCausalLM {
+    pub fn load(
+        cfg: &Config,
+        dtype: DType,
+        device: &Device,
+        vb: &ShardedVarBuilder,
+    ) -> Result<Self> {
+        let text = &cfg.text_config;
+        // `RotaryEmbedding` reads a qwen3_5 config; the rope block is the
+        // same type and the same values (spec §7), so the reuse is real
+        // rather than a coincidence — see `config::tests`.
+        let rotary = Arc::new(rotary_for(text, dtype, device)?);
+
+        let root = vb.pp("model").pp("language_model");
+        let embed_weight = root
+            .pp("embed_tokens")
+            .get((text.vocab_size, text.hidden_size), "weight")
+            .context("load 'model.language_model.embed_tokens.weight'")?;
+        let embed_tokens = Embedding::new(embed_weight, text.hidden_size);
+
+        let layers_vb = root.pp("layers");
+        let mut layers = Vec::with_capacity(text.num_hidden_layers);
+        for i in 0..text.num_hidden_layers {
+            layers.push(
+                DecoderLayer::load(text, rotary.clone(), i, &layers_vb.pp(i))
+                    .with_context(|| format!("load decoder layer {i}"))?,
+            );
+        }
+
+        let mixer = HyperConnection::load(
+            &root.pp("hyper_connection_mixer"),
+            text.hidden_size,
+            text.hc_count,
+            text.hc_lowrank,
+            text.rms_norm_eps,
+            // No block_inject_weight in the checkpoint: this one ends
+            // the streams rather than feeding a sublayer.
+            false,
+        )
+        .context("load hyper_connection_mixer")?;
+
+        let lm_head_weight = vb
+            .pp("lm_head")
+            .get((text.vocab_size, text.hidden_size), "weight")
+            .context("load 'lm_head.weight'")?;
+
+        let ngram = match text.ple_layers().first() {
+            Some(layer) => Some(load_ngram(text, &layers_vb.pp(*layer).pp("ple"))?),
+            None => None,
+        };
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            mixer,
+            lm_head: Linear::new(lm_head_weight, None),
+            ngram,
+            rotary,
+            hc_count: text.hc_count,
+            device: device.clone(),
+            dtype,
+        })
+    }
+
+    /// `input_ids` is `(B, T)`; `offset` is the sequence position the
+    /// first of them sits at. Returns logits `(B, T, vocab)`.
+    pub fn forward(&mut self, input_ids: &Tensor, offset: usize) -> Result<Tensor> {
+        let (batch, seq_len) = input_ids.dims2()?;
+
+        // Four identical streams — the residual this architecture
+        // carries between layers, not one.
+        let x = self.embed_tokens.forward(input_ids)?;
+        let mut h = x.repeat((1, 1, self.hc_count))?;
+
+        let (cos, sin) = self.rotary.plain_cos_sin(offset, seq_len)?;
+        let mask = if seq_len > 1 {
+            Some(self.causal_mask(batch, seq_len, offset)?)
+        } else {
+            None
+        };
+
+        // Computed once: it is a function of the ids, and one layer
+        // consumes it.
+        let ngram = match &mut self.ngram {
+            Some(ngram) => {
+                let ids = ids_to_i64(input_ids)?;
+                Some(ngram.forward(&ids, &self.device, self.dtype)?)
+            }
+            None => None,
+        };
+
+        for (i, layer) in self.layers.iter_mut().enumerate() {
+            h = layer
+                .forward(&h, ngram.as_ref(), mask.as_ref(), &cos, &sin, offset)
+                .with_context(|| format!("layer {i}"))?;
+        }
+
+        // The mixer's hc_norm is the only normalisation before the head.
+        let x = self.mixer.collapse(&h)?;
+        Ok(self.lm_head.forward(&x)?)
+    }
+
+    pub fn clear_kv_cache(&mut self) -> Result<()> {
+        for layer in &mut self.layers {
+            layer.clear_kv_cache()?;
+        }
+        if let Some(ngram) = &mut self.ngram {
+            ngram.reset();
+        }
+        Ok(())
+    }
+
+    fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> Result<Tensor> {
+        let minf = f32::NEG_INFINITY;
+        let mask: Vec<f32> = (0..tgt)
+            .flat_map(|i| (0..(tgt + offset)).map(move |j| if j <= i + offset { 0. } else { minf }))
+            .collect();
+        Ok(
+            Tensor::from_vec(mask, (b, 1, tgt, tgt + offset), &self.device)?
+                .to_dtype(self.dtype)?,
+        )
+    }
+}
+
+fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbedding> {
+    let vb = ple_vb.pp("ple_embedding");
+    let heads = cfg.ngram_heads();
+    let head_dim = cfg.ngram_head_dim();
+
+    // The three derived buffers ship as I64. Reading them through a
+    // bf16 VarBuilder would silently round multipliers of order 1e13
+    // and prime vocab sizes of 2e7 into nonsense, so the dtype is
+    // forced rather than inherited.
+    let ints = vb.to_dtype(DType::I64);
+    let multipliers = int_vec(&ints, "layer_multipliers", cfg.ngram_size)?;
+    let vocab_sizes = int_vec(&ints, "ngram_heads_vocab_sizes", heads)?;
+    let offsets = int_vec(&ints, "ngram_heads_offsets", heads)?;
+
+    // Total rows: the last head's offset plus its own vocab, rounded up
+    // the way the checkpoint pads it.
+    let declared: i64 = offsets[heads - 1] + vocab_sizes[heads - 1];
+    let pad = cfg.make_ngram_vocab_size_divisible_by.max(1) as i64;
+    let rows = (declared + pad - 1) / pad * pad;
+    let rows = rows as usize;
+    ensure!(
+        rows.is_multiple_of(cfg.split_ngram_parts),
+        "{rows} padded rows do not divide into {} shards",
+        cfg.split_ngram_parts
+    );
+
+    let table = ShardedNGramTable::load(
+        &vb.pp("ngram_embedding"),
+        cfg.split_ngram_parts,
+        rows,
+        head_dim,
+    )?;
+
+    Ok(NGramEmbedding {
+        hasher: NGramHasher::new(
+            cfg.ngram_size,
+            cfg.heads_per_ngram,
+            multipliers,
+            vocab_sizes,
+            offsets,
+            cfg.eos_token_id,
+        )?,
+        table: Box::new(table),
+        heads,
+        head_dim,
+        carried: Vec::new(),
+    })
+}
+
+fn int_vec(vb: &ShardedVarBuilder, name: &str, len: usize) -> Result<Vec<i64>> {
+    let t = vb
+        .get(len, name)
+        .with_context(|| format!("load '{}/{name}'", vb.prefix()))?;
+    Ok(t.to_dtype(DType::I64)?.to_vec1::<i64>()?)
+}
+
+fn ids_to_i64(input_ids: &Tensor) -> Result<Vec<Vec<i64>>> {
+    let ids = input_ids.to_dtype(DType::I64)?;
+    let (batch, _) = ids.dims2()?;
+    let mut out = Vec::with_capacity(batch);
+    for row in 0..batch {
+        out.push(ids.i(row)?.to_vec1::<i64>()?);
+    }
+    Ok(out)
+}
+
+/// Build the rotary tables from this architecture's config.
+///
+/// The rotary reads only `head_dim`, `max_position_embeddings` and the
+/// `rope_parameters` block, and this checkpoint declares all three
+/// identically to Qwen3.6 (spec §7) — so `qwen3_5`'s module is reused
+/// as-is, through the parts it actually needs rather than through a
+/// fabricated config.
+fn rotary_for(cfg: &TextConfig, dtype: DType, device: &Device) -> Result<RotaryEmbedding> {
+    RotaryEmbedding::from_parts(
+        dtype,
+        cfg.head_dim,
+        cfg.max_position_embeddings,
+        &cfg.rope_parameters,
+        device,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_nn::var_builder::ShardedSafeTensors;
+    use std::collections::HashMap;
+
+    /// A model small enough to build in a test and complete enough to
+    /// exercise every loader: two layers so both attention flavours
+    /// appear, PLE on layer 1 because `ple_layer_ids` is one-indexed,
+    /// and a sharded n-gram table with real prime vocab sizes.
+    const TINY: &str = r#"{
+        "model_type": "qwen4_exp",
+        "tie_word_embeddings": false,
+        "text_config": {
+            "vocab_size": 32, "hidden_size": 8, "num_hidden_layers": 2,
+            "num_attention_heads": 2, "num_key_value_heads": 1, "head_dim": 4,
+            "max_position_embeddings": 64, "rms_norm_eps": 1e-6,
+            "hidden_act": "silu", "eos_token_id": 3,
+            "full_attention_interval": 2,
+            "rope_parameters": {"rope_theta": 10000.0, "partial_rotary_factor": 0.5},
+            "hc_count": 2, "hc_lowrank": 4,
+            "ple_layer_ids": [2], "ple_embed_dim": 8, "ple_conv_kernel_size": 2,
+            "ngram_size": 3, "heads_per_ngram": 2,
+            "ngram_vocab_size_base": 10,
+            "make_ngram_vocab_size_divisible_by": 4,
+            "split_ngram_parts": 2,
+            "indexer_n_heads": 2, "indexer_kv_heads": 1, "indexer_head_dim": 4,
+            "indexer_budget": 8, "indexer_compress_ratio": 2,
+            "linear_num_key_heads": 1, "linear_key_head_dim": 4,
+            "linear_num_value_heads": 2, "linear_value_head_dim": 4,
+            "linear_conv_kernel_dim": 2,
+            "mamba_ssm_dtype": "float32", "output_gate_type": "sigmoid",
+            "num_experts": 2, "num_experts_per_tok": 1,
+            "moe_intermediate_size": 4, "shared_expert_intermediate_size": 4
+        }
+    }"#;
+
+    /// Per-head vocab sizes and their running offsets, as the checkpoint
+    /// would ship them: four distinct primes, cumulative offsets, padded
+    /// to a multiple of `make_ngram_vocab_size_divisible_by`.
+    const VOCABS: [i64; 4] = [11, 13, 17, 19];
+    const OFFSETS: [i64; 4] = [0, 11, 24, 41];
+    const TABLE_ROWS: usize = 60; // 41 + 19, already a multiple of 4
+
+    fn tensors(cfg: &TextConfig) -> HashMap<String, Tensor> {
+        let dev = Device::Cpu;
+        let mut t: HashMap<String, Tensor> = HashMap::new();
+        // Small non-zero values: zeros would hide a transposed load, and
+        // large ones make the MoE softmax degenerate.
+        let mut seed = 0u32;
+        let mut rand = |shape: Vec<usize>| {
+            let n: usize = shape.iter().product();
+            let v: Vec<f32> = (0..n)
+                .map(|i| {
+                    seed = seed
+                        .wrapping_mul(1664525)
+                        .wrapping_add(1013904223 + i as u32);
+                    ((seed >> 16) as f32 / 32768.0 - 1.0) * 0.1
+                })
+                .collect();
+            Tensor::from_vec(v, shape, &dev).unwrap()
+        };
+
+        let (h, hc) = (cfg.hidden_size, cfg.hc_count);
+        let wide = h * hc;
+        t.insert("lm_head.weight".to_string(), rand(vec![cfg.vocab_size, h]));
+        let root = "model.language_model";
+        t.insert(
+            format!("{root}.embed_tokens.weight"),
+            rand(vec![cfg.vocab_size, h]),
+        );
+        t.insert(
+            format!("{root}.hyper_connection_mixer.hc_norm.weight"),
+            rand(vec![wide]),
+        );
+        t.insert(
+            format!("{root}.hyper_connection_mixer.input_mix_weight_down.weight"),
+            rand(vec![cfg.hc_lowrank, wide]),
+        );
+        t.insert(
+            format!("{root}.hyper_connection_mixer.input_mix_weight_up.weight"),
+            rand(vec![wide, cfg.hc_lowrank]),
+        );
+
+        for i in 0..cfg.num_hidden_layers {
+            let l = format!("{root}.layers.{i}");
+            for which in ["attn_hyper_connection", "mlp_hyper_connection"] {
+                t.insert(format!("{l}.{which}.hc_norm.weight"), rand(vec![wide]));
+                t.insert(
+                    format!("{l}.{which}.input_mix_weight_down.weight"),
+                    rand(vec![cfg.hc_lowrank, wide]),
+                );
+                t.insert(
+                    format!("{l}.{which}.input_mix_weight_up.weight"),
+                    rand(vec![wide, cfg.hc_lowrank]),
+                );
+                t.insert(
+                    format!("{l}.{which}.block_inject_weight.weight"),
+                    rand(vec![hc, wide]),
+                );
+            }
+
+            if cfg.is_full_attention(i) {
+                let a = format!("{l}.self_attn");
+                let heads = cfg.num_attention_heads * cfg.head_dim;
+                t.insert(format!("{a}.q_proj.weight"), rand(vec![heads * 2, h]));
+                t.insert(
+                    format!("{a}.k_proj.weight"),
+                    rand(vec![cfg.num_key_value_heads * cfg.head_dim, h]),
+                );
+                t.insert(
+                    format!("{a}.v_proj.weight"),
+                    rand(vec![cfg.num_key_value_heads * cfg.head_dim, h]),
+                );
+                t.insert(format!("{a}.o_proj.weight"), rand(vec![h, heads]));
+                t.insert(format!("{a}.q_norm.weight"), rand(vec![cfg.head_dim]));
+                t.insert(format!("{a}.k_norm.weight"), rand(vec![cfg.head_dim]));
+                let fused = (cfg.indexer_n_heads + cfg.indexer_kv_heads) * cfg.indexer_head_dim;
+                t.insert(
+                    format!("{a}.indexer.index_qk_proj.weight"),
+                    rand(vec![fused, h]),
+                );
+                t.insert(
+                    format!("{a}.indexer.q_layernorm.weight"),
+                    rand(vec![cfg.indexer_head_dim]),
+                );
+                t.insert(
+                    format!("{a}.indexer.k_layernorm.weight"),
+                    rand(vec![cfg.indexer_head_dim]),
+                );
+            } else {
+                let n = format!("{l}.linear_attn");
+                let key_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim;
+                let value_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim;
+                let conv_dim = key_dim * 2 + value_dim;
+                t.insert(format!("{n}.in_proj_qkv.weight"), rand(vec![conv_dim, h]));
+                t.insert(format!("{n}.in_proj_z.weight"), rand(vec![value_dim, h]));
+                t.insert(
+                    format!("{n}.in_proj_a.weight"),
+                    rand(vec![cfg.linear_num_value_heads, h]),
+                );
+                t.insert(
+                    format!("{n}.in_proj_b.weight"),
+                    rand(vec![cfg.linear_num_value_heads, h]),
+                );
+                t.insert(
+                    format!("{n}.conv1d.weight"),
+                    rand(vec![conv_dim, 1, cfg.linear_conv_kernel_dim]),
+                );
+                t.insert(
+                    format!("{n}.dt_bias"),
+                    rand(vec![cfg.linear_num_value_heads]),
+                );
+                t.insert(format!("{n}.A_log"), rand(vec![cfg.linear_num_value_heads]));
+                t.insert(
+                    format!("{n}.norm.weight"),
+                    rand(vec![cfg.linear_value_head_dim]),
+                );
+                t.insert(format!("{n}.out_proj.weight"), rand(vec![h, value_dim]));
+            }
+
+            let m = format!("{l}.mlp");
+            t.insert(format!("{m}.gate.weight"), rand(vec![cfg.num_experts, h]));
+            t.insert(
+                format!("{m}.experts.gate_up_proj"),
+                rand(vec![cfg.num_experts, cfg.moe_intermediate_size * 2, h]),
+            );
+            t.insert(
+                format!("{m}.experts.down_proj"),
+                rand(vec![cfg.num_experts, h, cfg.moe_intermediate_size]),
+            );
+            for p in ["gate_proj", "up_proj"] {
+                t.insert(
+                    format!("{m}.shared_expert.{p}.weight"),
+                    rand(vec![cfg.shared_expert_intermediate_size, h]),
+                );
+            }
+            t.insert(
+                format!("{m}.shared_expert.down_proj.weight"),
+                rand(vec![h, cfg.shared_expert_intermediate_size]),
+            );
+            t.insert(format!("{m}.shared_expert_gate.weight"), rand(vec![1, h]));
+
+            if cfg.ple_layers().contains(&i) {
+                let p = format!("{l}.ple");
+                t.insert(format!("{p}.key_proj.weight"), rand(vec![wide, h]));
+                t.insert(format!("{p}.value_proj.weight"), rand(vec![h, h]));
+                for norm in ["norm_key", "norm_query", "norm_conv"] {
+                    t.insert(format!("{p}.{norm}.weight"), rand(vec![wide]));
+                }
+                t.insert(
+                    format!("{p}.conv1d.weight"),
+                    rand(vec![wide, 1, cfg.ple_conv_kernel_size]),
+                );
+                let e = format!("{p}.ple_embedding");
+                // The three derived buffers, as I64 — the dtype is the
+                // point of this fixture.
+                t.insert(
+                    format!("{e}.layer_multipliers"),
+                    Tensor::from_vec(vec![7i64, 13, 29], 3, &dev).unwrap(),
+                );
+                t.insert(
+                    format!("{e}.ngram_heads_vocab_sizes"),
+                    Tensor::from_vec(VOCABS.to_vec(), VOCABS.len(), &dev).unwrap(),
+                );
+                t.insert(
+                    format!("{e}.ngram_heads_offsets"),
+                    Tensor::from_vec(OFFSETS.to_vec(), OFFSETS.len(), &dev).unwrap(),
+                );
+                let per_shard = TABLE_ROWS / cfg.split_ngram_parts;
+                for s in 0..cfg.split_ngram_parts {
+                    t.insert(
+                        format!("{e}.ngram_embedding.shard_{s}.weight"),
+                        rand(vec![per_shard, cfg.ngram_head_dim()]),
+                    );
+                }
+            }
+        }
+        t
+    }
+
+    fn build() -> (tempfile::TempDir, Qwen4ExpForCausalLM) {
+        let cfg = Config::from_config_json(TINY).unwrap();
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&tensors(&cfg.text_config), &path).unwrap();
+        let vb =
+            unsafe { ShardedSafeTensors::var_builder(&[&path], DType::F32, &Device::Cpu).unwrap() };
+        let model = Qwen4ExpForCausalLM::load(&cfg, DType::F32, &Device::Cpu, &vb).unwrap();
+        (dir, model)
+    }
+
+    /// Every loader in the architecture, against a checkpoint laid out
+    /// the way the real one is. Tensor names and shapes are the failure
+    /// mode a maths test cannot reach, and finding them on a 360 GB
+    /// download is the expensive way.
+    #[test]
+    fn the_whole_model_loads_and_produces_logits() {
+        let (_dir, mut model) = build();
+        let ids = Tensor::from_vec(vec![5u32, 9, 3, 12, 7], (1, 5), &Device::Cpu).unwrap();
+
+        let logits = model.forward(&ids, 0).unwrap();
+        assert_eq!(logits.dims(), &[1, 5, 32]);
+        let values: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(
+            values.iter().all(|v| v.is_finite()),
+            "a wrong stream width or a dead mask row shows up as NaN"
+        );
+    }
+
+    /// A decode step continues the prefill: the caches, the rotary
+    /// offset and PLE's carried ids all have to agree, and any one of
+    /// them being wrong still returns a tensor of the right shape.
+    #[test]
+    fn a_decode_step_follows_a_prefill() {
+        let (_dir, mut model) = build();
+        let prompt = Tensor::from_vec(vec![5u32, 9, 3, 12], (1, 4), &Device::Cpu).unwrap();
+        model.forward(&prompt, 0).unwrap();
+
+        let next = Tensor::from_vec(vec![7u32], (1, 1), &Device::Cpu).unwrap();
+        let logits = model.forward(&next, 4).unwrap();
+        assert_eq!(logits.dims(), &[1, 1, 32]);
+        let values: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(values.iter().all(|v| v.is_finite()));
+
+        // And clearing lets the same model serve a fresh request.
+        model.clear_kv_cache().unwrap();
+        assert!(model.forward(&prompt, 0).is_ok());
+    }
+
+    /// The n-gram buffers are I64 in a checkpoint whose activations are
+    /// not. Loading them through the model's dtype would round a prime
+    /// vocab size and a multiplier into nonsense, and the addressing
+    /// would still return rows.
+    #[test]
+    fn the_ngram_buffers_survive_a_non_integer_model_dtype() {
+        let (_dir, model) = build();
+        let ngram = model.ngram.as_ref().expect("layer 1 carries PLE");
+        assert_eq!(ngram.heads, 4);
+        assert_eq!(ngram.head_dim, 2);
+        // The table was sized from offsets + vocab sizes read as I64;
+        // a bf16 round-trip of 41 + 19 would not land on 60.
+        assert_eq!(ngram.table.head_dim(), 2);
+        assert!(ngram.table.gather(&[(TABLE_ROWS - 1) as i64]).is_ok());
+        assert!(ngram.table.gather(&[TABLE_ROWS as i64]).is_err());
+    }
+}

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

@@ -0,0 +1,188 @@
+//! `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.
+//!
+//! The fused layout is the better one for a grouped GEMM, and keeping
+//! it would be the faster path; that is an algorithm change, and the
+//! spec is explicit that this stage is not it. Sparsity is extreme —
+//! 10 of 512 with an intermediate of 640 — so the routing, not the
+//! GEMM shape, is what dominates today.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §6.
+
+use anyhow::{Context, Result};
+use candle_core::{IndexOp, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::mlp::Qwen3_5MLP;
+use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
+
+use super::config::TextConfig;
+
+/// `vb` should be `.pp(...)`-ed to the layer's `mlp` prefix.
+pub fn load(cfg: &TextConfig, 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()))?;
+    let experts = split_fused_experts(&gate_up, &down, 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,
+    ))
+}
+
+/// Slice `(E, 2I, H)` and `(E, H, I)` into `E` SwiGLU experts.
+///
+/// Gate first, up second — see the module note.
+pub(crate) fn split_fused_experts(
+    gate_up: &Tensor,
+    down: &Tensor,
+    intermediate: usize,
+) -> Result<Vec<Qwen3_5MLP>> {
+    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
+    );
+
+    let mut out = Vec::with_capacity(experts);
+    for e in 0..experts {
+        let fused_e = gate_up.i(e)?;
+        out.push(Qwen3_5MLP::from_weights(
+            Linear::new(fused_e.narrow(0, 0, intermediate)?.contiguous()?, None),
+            Linear::new(
+                fused_e
+                    .narrow(0, intermediate, intermediate)?
+                    .contiguous()?,
+                None,
+            ),
+            Linear::new(down.i(e)?.contiguous()?, None),
+        ));
+    }
+    Ok(out)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device, Module};
+
+    /// One expert, two hidden dims, one intermediate. The gate row picks
+    /// the first input channel and the up row picks the second, so
+    /// swapping them swaps which channel is squashed by the SiLU — and
+    /// the two answers differ by more than any tolerance.
+    #[test]
+    fn the_gate_is_the_first_half_of_the_fused_tensor() {
+        let dev = Device::Cpu;
+        // gate_up: [[gate row], [up row]] = [[1, 0], [0, 1]]
+        let gate_up = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (1, 2, 2), &dev).unwrap();
+        // down: (E=1, H=2, I=1)
+        let down = Tensor::from_vec(vec![2.0f32, 3.0], (1, 2, 1), &dev).unwrap();
+
+        let experts = split_fused_experts(&gate_up, &down, 1).unwrap();
+        assert_eq!(experts.len(), 1);
+
+        let x = Tensor::from_vec(vec![1.0f32, 2.0], (1, 2), &dev).unwrap();
+        let got: Vec<f32> = experts[0]
+            .forward(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // gate(x) = 1, up(x) = 2  ->  silu(1) * 2  ->  down
+        let silu1 = 1.0f32 / (1.0 + (-1.0f32).exp());
+        let want = [2.0 * silu1 * 2.0, 3.0 * silu1 * 2.0];
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-5, "got {got:?} want {want:?}");
+        }
+
+        // The swapped reading: silu(2) * 1, which is a different model.
+        let silu2 = 2.0f32 / (1.0 + (-2.0f32).exp());
+        assert!(
+            (got[0] - 2.0 * silu2).abs() > 1e-3,
+            "this is silu(up) * gate, not silu(gate) * up: {got:?}"
+        );
+    }
+
+    #[test]
+    fn a_fused_tensor_of_the_wrong_width_is_rejected() {
+        let dev = Device::Cpu;
+        let gate_up = Tensor::zeros((2, 6, 4), DType::F32, &dev).unwrap();
+        let down = Tensor::zeros((2, 4, 3), DType::F32, &dev).unwrap();
+        assert!(split_fused_experts(&gate_up, &down, 3).is_ok());
+        // 2 x 4 != 6
+        assert!(split_fused_experts(&gate_up, &down, 4).is_err());
+        // expert counts must agree
+        let down_short = Tensor::zeros((1, 4, 3), DType::F32, &dev).unwrap();
+        assert!(split_fused_experts(&gate_up, &down_short, 3).is_err());
+    }
+}

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

@@ -0,0 +1,986 @@
+//! PLE — per-layer embedding over hashed n-grams.
+//!
+//! `qwen4_exp` puts one of these on zero-indexed layer **1**
+//! (`ple_layer_ids: [2]` is one-indexed upstream), and it accounts for
+//! **28.44% of the model's parameters** — a 320,001,536-row table that
+//! each token touches 16 rows of.
+//!
+//! PLE is three separable pieces. [`NGramHasher`] is the *addressing*:
+//! token ids to row indices. [`PleBlock`] is the *consumption*: the
+//! gated mix of the gathered embedding into the four residual streams.
+//! Between them sits the gather, behind the [`NGramTable`] trait —
+//! *where* those bytes live is #310's decision, and the trait is the
+//! seam so that decision changes one implementation and nothing else.
+//! [`ShardedNGramTable`] is the naive reading of the checkpoint, which
+//! is the baseline that issue measures against rather than its
+//! answer.
+//!
+//! ## Geometry
+//!
+//! `ngram_heads = (ngram_size - 1) * heads_per_ngram` — for the shipped
+//! config, `(3 - 1) * 8 = 16` heads of `ple_embed_dim / 16 = 160` dims.
+//! There is **no unigram head**: with `ngram_size = 3` the orders are
+//! bigram and trigram only. `layer_multipliers` still has three entries
+//! because a trigram needs three shift multipliers.
+//!
+//! ## The hash
+//!
+//! For n-gram order *n*, over heads `[(n-2)*heads_per_ngram, +8)`:
+//!
+//! ```text
+//! mixed = tok[0] * m[0]
+//! for p in 1..n:  mixed ^= tok[p] * m[p]
+//! row[head] = (mixed mod vocab_sizes[head]) + offsets[head]
+//! ```
+//!
+//! where `tok[s]` is the id stream shifted right by `s`.
+//!
+//! Three things worth stating precisely:
+//!
+//! 1. **The arithmetic is int64 and stays non-negative by
+//!    construction.** Upstream derives `multiplier_max` as
+//!    `i64::MAX / vocab_size` specifically so that `token_id *
+//!    multiplier` cannot overflow: the largest possible product is
+//!    9,223,334,893,764,784,449 against an `i64::MAX` of
+//!    9,223,372,036,854,775,807. Every product therefore has its sign
+//!    bit clear, so their XOR does too and `mixed` is never negative.
+//!    The products still overflow `i32` by four orders of magnitude, so
+//!    the width matters.
+//! 2. **The modulo is nonetheless Euclidean.** Given (1) it agrees with
+//!    a truncated `%` on every value this checkpoint can produce, but
+//!    torch's `remainder` follows Python's sign convention, and a
+//!    checkpoint whose constants were derived differently would
+//!    silently produce negative row indices under Rust's `%`.
+//!    `rem_euclid` makes the addressing correct rather than
+//!    incidentally correct.
+//! 3. **Shifts do not cross an EOS.** Positions closer to the start of
+//!    their EOS-delimited segment than `shift` read `eos_token_id`
+//!    rather than borrowing the previous document's tokens. Invisible
+//!    on a single-turn prompt, wrong on anything batched or
+//!    multi-document.
+//!
+//! The three derived quantities — `layer_multipliers`,
+//! `ngram_heads_vocab_sizes`, `ngram_heads_offsets` — all ship in the
+//! checkpoint as I64 buffers, so the prime search and the splitmix64
+//! multiplier derivation are load-time trivia we deliberately do not
+//! reimplement. See `doc/qwen4_exp-port-spec.md` §2.
+
+use anyhow::{Context, Result, ensure};
+use candle_core::{D, Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm;
+
+/// Addressing for the n-gram table.
+pub struct NGramHasher {
+    ngram_size: usize,
+    heads_per_ngram: usize,
+    /// One per shift, `len == ngram_size`.
+    multipliers: Vec<i64>,
+    /// Per-head, `len == ngram_heads()`.
+    head_vocab_sizes: Vec<i64>,
+    head_offsets: Vec<i64>,
+    eos_token_id: i64,
+}
+
+impl NGramHasher {
+    /// Build from the checkpoint's own buffers.
+    pub fn new(
+        ngram_size: usize,
+        heads_per_ngram: usize,
+        multipliers: Vec<i64>,
+        head_vocab_sizes: Vec<i64>,
+        head_offsets: Vec<i64>,
+        eos_token_id: i64,
+    ) -> Result<Self> {
+        ensure!(ngram_size >= 2, "ngram_size must be >= 2, got {ngram_size}");
+        ensure!(
+            multipliers.len() == ngram_size,
+            "expected {ngram_size} layer_multipliers, got {}",
+            multipliers.len()
+        );
+        let heads = (ngram_size - 1) * heads_per_ngram;
+        ensure!(
+            head_vocab_sizes.len() == heads && head_offsets.len() == heads,
+            "expected {heads} per-head vocab sizes and offsets, got {} and {}",
+            head_vocab_sizes.len(),
+            head_offsets.len()
+        );
+        ensure!(
+            head_vocab_sizes.iter().all(|v| *v > 0),
+            "per-head vocab sizes must be positive"
+        );
+        Ok(Self {
+            ngram_size,
+            heads_per_ngram,
+            multipliers,
+            head_vocab_sizes,
+            head_offsets,
+            eos_token_id,
+        })
+    }
+
+    /// `(ngram_size - 1) * heads_per_ngram` — 16 for the shipped config.
+    pub fn ngram_heads(&self) -> usize {
+        (self.ngram_size - 1) * self.heads_per_ngram
+    }
+
+    /// How many previous ids must be carried across a decode step for
+    /// the shifts to be correct: `ngram_size - 1`.
+    pub fn context_len(&self) -> usize {
+        self.ngram_size - 1
+    }
+
+    /// Row indices for the last `seq_len` positions of `history`.
+    ///
+    /// `history` is the carried context followed by this step's ids, so
+    /// `history.len() >= seq_len`. Returns one row per head per position,
+    /// flattened as `[position][head]`.
+    pub fn rows(&self, history: &[i64], seq_len: usize) -> Result<Vec<Vec<i64>>> {
+        ensure!(
+            history.len() >= seq_len,
+            "history ({}) shorter than seq_len ({seq_len})",
+            history.len()
+        );
+        // tok[s] for s in 0..ngram_size, each segment-aware.
+        let shifted: Vec<Vec<i64>> = (0..self.ngram_size)
+            .map(|s| self.shift_right_ignore_eos(history, s))
+            .collect();
+
+        // Transpose to [position][shift] so each position's n-gram is a
+        // contiguous little slice rather than a stride across vectors.
+        let start = history.len() - seq_len;
+        let per_position =
+            (start..history.len()).map(|pos| shifted.iter().map(|s| s[pos]).collect::<Vec<i64>>());
+
+        let mut out = Vec::with_capacity(seq_len);
+        for toks in per_position {
+            let mut row = Vec::with_capacity(self.ngram_heads());
+            for order in 2..=self.ngram_size {
+                // Products are bounded by construction: upstream caps the
+                // multiplier at i64::MAX / vocab_size, so a token id times
+                // its multiplier cannot overflow. wrapping_mul keeps that
+                // a documented assumption rather than a debug panic.
+                let mut mixed = toks[0].wrapping_mul(self.multipliers[0]);
+                for (tok, mult) in toks.iter().zip(&self.multipliers).take(order).skip(1) {
+                    mixed ^= tok.wrapping_mul(*mult);
+                }
+                let head_start = (order - 2) * self.heads_per_ngram;
+                let head_end = head_start + self.heads_per_ngram;
+                let vocabs = &self.head_vocab_sizes[head_start..head_end];
+                let offsets = &self.head_offsets[head_start..head_end];
+                for (vocab, offset) in vocabs.iter().zip(offsets) {
+                    // Euclidean, not truncated — see the module note.
+                    row.push(mixed.rem_euclid(*vocab) + *offset);
+                }
+            }
+            out.push(row);
+        }
+        Ok(out)
+    }
+
+    /// Shift ids right by `shift` without crossing an EOS boundary.
+    /// Positions with fewer than `shift` predecessors inside their own
+    /// segment read `eos_token_id`.
+    fn shift_right_ignore_eos(&self, ids: &[i64], shift: usize) -> Vec<i64> {
+        if shift == 0 {
+            return ids.to_vec();
+        }
+        let mut out = Vec::with_capacity(ids.len());
+        // Position of the most recent EOS strictly before i.
+        let mut previous_eos: i64 = -1;
+        for (i, &id) in ids.iter().enumerate() {
+            let segment_start = previous_eos + 1;
+            let position_in_segment = i as i64 - segment_start;
+            let source = i as i64 - shift as i64;
+            let valid = position_in_segment >= shift as i64 && source >= 0;
+            out.push(if valid {
+                ids[source as usize]
+            } else {
+                self.eos_token_id
+            });
+            if id == self.eos_token_id {
+                previous_eos = i as i64;
+            }
+        }
+        out
+    }
+}
+
+/// Where the n-gram table's rows come from.
+///
+/// 320,001,536 rows of 160, and every token needs sixteen of them —
+/// a few KB per token out of 27 GB quantised. That ratio is why the
+/// table can live off-device at almost no cost, and why #310 is a
+/// residency decision rather than an algorithm. Nothing above this
+/// trait depends on the answer.
+pub trait NGramTable: Send + Sync {
+    /// Gather `rows`, **in the order given**, as
+    /// `[rows.len(), head_dim]`.
+    fn gather(&self, rows: &[i64]) -> Result<Tensor>;
+
+    /// Width of one row — `ple_embed_dim / ngram_heads`, 160.
+    fn head_dim(&self) -> usize;
+}
+
+/// The checkpoint's 128 shards, held wherever they were loaded.
+///
+/// The shards are a **flat split of one concatenated table**, not a
+/// split per head: shard `s` holds rows `[s * rows_per_shard,
+/// (s+1) * rows_per_shard)`, so a global row index from the hash
+/// addresses straight into it. They are kept separate rather than
+/// concatenated because joining 128 tensors of 2,500,012 x 160 would
+/// need the whole table free a second time.
+///
+/// This is the naive residency — the bytes end up wherever the
+/// VarBuilder put them, which for a 27 GB table is exactly the problem
+/// #310 exists to solve. It is the control that issue measures its
+/// host- and NVMe-resident variants against.
+pub struct ShardedNGramTable {
+    shards: Vec<Tensor>,
+    rows_per_shard: usize,
+    head_dim: usize,
+}
+
+impl ShardedNGramTable {
+    /// `vb` should be `.pp(...)`-ed to `ple_embedding.ngram_embedding`.
+    pub fn load(
+        vb: &ShardedVarBuilder,
+        parts: usize,
+        rows: usize,
+        head_dim: usize,
+    ) -> Result<Self> {
+        ensure!(parts > 0, "split_ngram_parts must be > 0");
+        ensure!(
+            rows.is_multiple_of(parts),
+            "{rows} rows do not divide evenly into {parts} shards"
+        );
+        let rows_per_shard = rows / parts;
+        let mut shards = Vec::with_capacity(parts);
+        for s in 0..parts {
+            shards.push(
+                vb.pp(format!("shard_{s}"))
+                    .get((rows_per_shard, head_dim), "weight")
+                    .with_context(|| format!("load '{}/shard_{s}/weight'", vb.prefix()))?,
+            );
+        }
+        Ok(Self {
+            shards,
+            rows_per_shard,
+            head_dim,
+        })
+    }
+
+    #[cfg(test)]
+    pub(crate) fn from_shards(shards: Vec<Tensor>) -> Result<Self> {
+        ensure!(!shards.is_empty(), "need at least one shard");
+        let (rows_per_shard, head_dim) = shards[0].dims2()?;
+        Ok(Self {
+            shards,
+            rows_per_shard,
+            head_dim,
+        })
+    }
+
+    fn shard_of(&self, row: i64) -> usize {
+        row as usize / self.rows_per_shard
+    }
+}
+
+impl NGramTable for ShardedNGramTable {
+    fn gather(&self, rows: &[i64]) -> Result<Tensor> {
+        let total = (self.shards.len() * self.rows_per_shard) as i64;
+        if let Some(bad) = rows.iter().find(|r| **r < 0 || **r >= total) {
+            anyhow::bail!("row {bad} is outside the table's 0..{total}");
+        }
+        if rows.is_empty() {
+            let device = self.shards[0].device();
+            return Ok(Tensor::zeros(
+                (0, self.head_dim),
+                self.shards[0].dtype(),
+                device,
+            )?);
+        }
+
+        // Visit each shard once: sort the requests by shard, gather, then
+        // undo the permutation. Getting that inverse wrong hands every
+        // token somebody else's n-gram, which is not a crash.
+        let mut order: Vec<usize> = (0..rows.len()).collect();
+        order.sort_by_key(|&i| (self.shard_of(rows[i]), i));
+
+        let device = self.shards[0].device();
+        let mut parts = Vec::new();
+        let mut i = 0;
+        while i < order.len() {
+            let shard = self.shard_of(rows[order[i]]);
+            let mut within = Vec::new();
+            while i < order.len() && self.shard_of(rows[order[i]]) == shard {
+                within.push((rows[order[i]] as usize - shard * self.rows_per_shard) as u32);
+                i += 1;
+            }
+            let idx = Tensor::from_vec(within.clone(), within.len(), device)?;
+            parts.push(self.shards[shard].index_select(&idx, 0)?);
+        }
+        let gathered = Tensor::cat(&parts, 0)?;
+
+        let mut inverse = vec![0u32; rows.len()];
+        for (position, &original) in order.iter().enumerate() {
+            inverse[original] = position as u32;
+        }
+        let inverse = Tensor::from_vec(inverse, rows.len(), device)?;
+        Ok(gathered.index_select(&inverse, 0)?)
+    }
+
+    fn head_dim(&self) -> usize {
+        self.head_dim
+    }
+}
+
+/// The consumption half of PLE: `Qwen4ExpTextPLE` minus the table.
+///
+/// Given the gathered n-gram embedding `e` and the four-stream residual
+/// `h`, this produces the `hidden * hc_count` tensor that layer 1 adds
+/// back into `h`. It is a gated read: the embedding proposes a value,
+/// the residual streams decide per-stream how much of it to take.
+///
+/// ```text
+/// k    = norm_key(key_proj(e))      -> [.., hc, hidden]
+/// q    = norm_query(h)              -> [.., hc, hidden]
+/// v    = value_proj(e)              -> [.., hidden]
+/// g    = (k * q).sum(-1) / sqrt(hidden)               [.., hc, 1]
+/// g    = sign(g) * sqrt(max(|g|, 1e-6))               signed sqrt
+/// gv   = sigmoid(g) * v                               [.., hc, hidden]
+/// out  = gv.flatten(-2) + short_conv(norm_conv(gv.flatten(-2)))
+/// ```
+///
+/// Two things here are not what they look like:
+///
+/// 1. **`short_conv` is dilated.** `dilation = ngram_size = 3`, so its
+///    four taps reach 9 positions back, not 3, and the cached left
+///    context is 9 wide rather than `kernel_size`. A plain causal conv
+///    runs, trains nothing, and mixes the wrong neighbours — the test
+///    below moves a single position and asserts where it lands.
+/// 2. **The conv is a residual branch, not the output.** `gv` passes
+///    through whether or not the conv contributes.
+///
+/// All three norms are grouped over `hidden_size` (§8): four
+/// independent normalisations of a 10240-wide vector, not one.
+/// See `doc/qwen4_exp-port-spec.md` §2.
+pub struct PleBlock {
+    key_proj: Linear,
+    value_proj: Linear,
+    norm_key: Qwen3_5RmsNorm,
+    norm_query: Qwen3_5RmsNorm,
+    norm_conv: Qwen3_5RmsNorm,
+    /// Depthwise: `(hidden * hc_count, 1, kernel_size)`.
+    conv1d_weight: Tensor,
+    /// Rolling left context for the dilated conv,
+    /// `(B, hidden * hc_count, (kernel_size - 1) * dilation)`. Upstream
+    /// keeps this as `conv_states[1]` on layer 1.
+    conv_state: Option<Tensor>,
+    hidden_size: usize,
+    hc_count: usize,
+    kernel_size: usize,
+    dilation: usize,
+}
+
+impl PleBlock {
+    /// `vb` should be `.pp(...)`-ed to the layer's `ple` prefix.
+    pub fn load(
+        vb: &ShardedVarBuilder,
+        hidden_size: usize,
+        hc_count: usize,
+        kernel_size: usize,
+        dilation: usize,
+        eps: f64,
+    ) -> Result<Self> {
+        ensure!(kernel_size >= 1, "ple conv kernel_size must be >= 1");
+        ensure!(dilation >= 1, "ple conv dilation must be >= 1");
+        let wide = hidden_size * hc_count;
+        let conv1d_weight = vb
+            .pp("conv1d")
+            .get((wide, 1, kernel_size), "weight")
+            .with_context(|| format!("load '{}/conv1d/weight'", vb.prefix()))?;
+        Ok(Self {
+            key_proj: linear_no_bias(vb, "key_proj", hidden_size, wide)?,
+            value_proj: linear_no_bias(vb, "value_proj", hidden_size, hidden_size)?,
+            norm_key: Qwen3_5RmsNorm::load_grouped(&vb.pp("norm_key"), wide, hidden_size, eps)?,
+            norm_query: Qwen3_5RmsNorm::load_grouped(&vb.pp("norm_query"), wide, hidden_size, eps)?,
+            norm_conv: Qwen3_5RmsNorm::load_grouped(&vb.pp("norm_conv"), wide, hidden_size, eps)?,
+            conv1d_weight,
+            conv_state: None,
+            hidden_size,
+            hc_count,
+            kernel_size,
+            dilation,
+        })
+    }
+
+    /// `h`: `[B, L, hidden * hc]` — the residual streams.
+    /// `ngram_embed`: `[B, L, hidden]` — the gathered n-gram rows.
+    /// Returns `[B, L, hidden * hc]`, to be **added** to `h`.
+    pub fn forward(&mut self, h: &Tensor, ngram_embed: &Tensor) -> candle_core::Result<Tensor> {
+        let (b, l, wide) = h.dims3()?;
+        debug_assert_eq!(wide, self.hidden_size * self.hc_count);
+        let streams = (b, l, self.hc_count, self.hidden_size);
+
+        let k = self
+            .norm_key
+            .forward(&self.key_proj.forward(ngram_embed)?)?
+            .reshape(streams)?;
+        let q = self.norm_query.forward(h)?.reshape(streams)?;
+        let v = self.value_proj.forward(ngram_embed)?;
+
+        // Per-stream agreement between the proposed key and the stream's
+        // own state, scaled like an attention logit.
+        let g = (k.mul(&q)?.sum_keepdim(D::Minus1)? / (self.hidden_size as f64).sqrt())?;
+        // Signed sqrt: compresses the magnitude without losing the sign,
+        // so a strongly disagreeing stream still shuts the gate.
+        let g = g.sign()?.mul(&g.abs()?.maximum(1e-6)?.sqrt()?)?;
+
+        let gv = candle_nn::ops::sigmoid(&g)?
+            .broadcast_mul(&v.unsqueeze(D::Minus2)?)?
+            .flatten_from(D::Minus2)?;
+
+        let conv = self.short_conv(&self.norm_conv.forward(&gv)?)?;
+        gv + conv
+    }
+
+    /// Depthwise dilated causal conv + SiLU, with rolling left context.
+    ///
+    /// Unlike the GatedDeltaNet short conv this cannot use
+    /// `run_causal_conv1d`: that path (and the fused cuda kernel behind
+    /// it) is hard-wired to `dilation = 1`.
+    fn short_conv(&mut self, x: &Tensor) -> candle_core::Result<Tensor> {
+        let (b, l, wide) = x.dims3()?;
+        let reach = self.left_context();
+        // (B, wide, L) for the conv.
+        let x = x.transpose(1, 2)?.contiguous()?;
+
+        let prepended = match self.conv_state.take() {
+            Some(prev) => Tensor::cat(&[&prev, &x], 2)?,
+            None => x,
+        };
+        let prep_len = prepended.dim(2)?;
+
+        self.conv_state = Some(if prep_len >= reach {
+            prepended.narrow(2, prep_len - reach, reach)?.contiguous()?
+        } else {
+            let pad = Tensor::zeros(
+                (b, wide, reach - prep_len),
+                prepended.dtype(),
+                prepended.device(),
+            )?;
+            Tensor::cat(&[&pad, &prepended], 2)?
+        });
+
+        // Pad both sides by the full reach, then keep the left-aligned
+        // window: output t sees inputs t - j*dilation only.
+        let out = prepended.conv1d(&self.conv1d_weight, reach, 1, self.dilation, wide)?;
+        let out = candle_nn::ops::silu(&out.narrow(2, 0, prep_len)?)?;
+        out.narrow(2, prep_len - l, l)?
+            .transpose(1, 2)?
+            .contiguous()
+    }
+
+    /// How many positions of left context the dilated conv reaches back
+    /// over — `(kernel_size - 1) * dilation`, 9 for the shipped config.
+    /// This, not `kernel_size`, is the cached state width.
+    pub fn left_context(&self) -> usize {
+        (self.kernel_size - 1) * self.dilation
+    }
+
+    /// Drop the rolling conv context. Call alongside a KV-cache clear;
+    /// the carried n-gram ids ([`NGramHasher::context_len`]) reset with
+    /// it.
+    pub fn clear_state(&mut self) {
+        self.conv_state = None;
+    }
+}
+
+fn linear_no_bias(
+    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::*;
+
+    const EOS: i64 = 248044;
+
+    /// The shipped geometry, with the checkpoint's real derivation for
+    /// the head vocab sizes (primes just above 20,000,000) approximated
+    /// by distinct nearby primes — only their distinctness and
+    /// positivity matter to the addressing.
+    fn hasher(ngram_size: usize, heads_per_ngram: usize) -> NGramHasher {
+        let heads = (ngram_size - 1) * heads_per_ngram;
+        let vocabs: Vec<i64> = (0..heads).map(|i| 20_000_003 + 2 * i as i64).collect();
+        let mut offsets = Vec::with_capacity(heads);
+        let mut acc = 0i64;
+        for v in &vocabs {
+            offsets.push(acc);
+            acc += v;
+        }
+        // Realistic magnitude: odd, ~1e13, as _build_layer_multipliers
+        // produces.
+        let mults: Vec<i64> = (0..ngram_size)
+            .map(|i| 12_345_678_901_235 + 2 * i as i64)
+            .collect();
+        NGramHasher::new(ngram_size, heads_per_ngram, mults, vocabs, offsets, EOS).unwrap()
+    }
+
+    #[test]
+    fn geometry_matches_the_shipped_config() {
+        let h = hasher(3, 8);
+        assert_eq!(h.ngram_heads(), 16, "(ngram_size - 1) * heads_per_ngram");
+        assert_eq!(h.context_len(), 2);
+    }
+
+    /// Every row must land inside its own head's slice of the table.
+    /// This is what a truncated remainder would break, and it is also
+    /// the invariant that keeps a gather in bounds.
+    #[test]
+    fn rows_stay_within_their_head_slice() {
+        let h = hasher(3, 8);
+        let ids: Vec<i64> = (0..64).map(|i| (i * 3907) % 248_320).collect();
+        let rows = h.rows(&ids, ids.len()).unwrap();
+        assert_eq!(rows.len(), ids.len());
+        for row in &rows {
+            assert_eq!(row.len(), 16);
+            for (head, &r) in row.iter().enumerate() {
+                let lo = h.head_offsets[head];
+                let hi = lo + h.head_vocab_sizes[head];
+                assert!(r >= lo && r < hi, "head {head}: {r} outside [{lo}, {hi})");
+            }
+        }
+    }
+
+    /// Upstream's multiplier bound guarantees products fit i64. Assert
+    /// it rather than trusting the comment, since the whole hash relies
+    /// on the sign bit staying clear.
+    #[test]
+    fn largest_possible_product_does_not_overflow_i64() {
+        let vocab_size: i64 = 248_320;
+        let multiplier_max = i64::MAX / vocab_size;
+        let largest_multiplier = 2 * (multiplier_max / 2 - 1) + 1;
+        let product = (vocab_size - 1).checked_mul(largest_multiplier);
+        assert!(product.is_some(), "token id * multiplier must fit i64");
+        assert!(product.unwrap() >= 0, "and stay non-negative");
+    }
+
+    /// A shift must not read across an EOS: the token after a document
+    /// boundary sees EOS as its predecessor, not the previous
+    /// document's last token.
+    #[test]
+    fn shifts_do_not_cross_an_eos_boundary() {
+        let h = hasher(3, 8);
+        //        0    1    2     3(eos)  4    5
+        let ids = [11i64, 12, 13, EOS, 21, 22];
+
+        let by_one = h.shift_right_ignore_eos(&ids, 1);
+        // Position 4 is the first of a new segment — its predecessor
+        // must be EOS, not the boundary token itself.
+        assert_eq!(
+            by_one[4], EOS,
+            "first token of a segment has no predecessor"
+        );
+        assert_eq!(by_one[5], 21, "second token reads within its segment");
+        assert_eq!(by_one[0], EOS, "very first position has no predecessor");
+        assert_eq!(by_one[1], 11);
+
+        let by_two = h.shift_right_ignore_eos(&ids, 2);
+        assert_eq!(by_two[4], EOS);
+        assert_eq!(by_two[5], EOS, "only one predecessor inside the segment");
+        assert_eq!(by_two[2], 11);
+    }
+
+    #[test]
+    fn shift_zero_is_the_identity() {
+        let h = hasher(3, 8);
+        let ids = [11i64, 12, EOS, 21];
+        assert_eq!(h.shift_right_ignore_eos(&ids, 0), ids.to_vec());
+    }
+
+    /// The bigram heads see two shifts, the trigram heads three, so a
+    /// token whose grandparent differs must change only the trigram
+    /// half of the row.
+    #[test]
+    fn bigram_heads_ignore_the_third_token_but_trigram_heads_do_not() {
+        let h = hasher(3, 4);
+        // Same last two tokens, different one before that.
+        let a = [7i64, 8, 9];
+        let b = [99i64, 8, 9];
+        let ra = &h.rows(&a, 1).unwrap()[0];
+        let rb = &h.rows(&b, 1).unwrap()[0];
+
+        assert_eq!(&ra[0..4], &rb[0..4], "bigram heads depend on 2 tokens only");
+        assert_ne!(&ra[4..8], &rb[4..8], "trigram heads must see the third");
+    }
+
+    /// Carried context makes a decode step agree with the prefill that
+    /// would have produced the same position.
+    #[test]
+    fn carried_context_reproduces_the_prefill_row() {
+        let h = hasher(3, 8);
+        let full = [11i64, 12, 13, 14];
+        let prefill = h.rows(&full, full.len()).unwrap();
+        // Decode step for the last token, carrying ngram_size - 1 ids.
+        let step = h.rows(&full[1..], 1).unwrap();
+        assert_eq!(step[0], prefill[3]);
+    }
+
+    #[test]
+    fn rejects_buffers_that_disagree_with_the_geometry() {
+        // 3 multipliers expected for ngram_size 3.
+        assert!(
+            NGramHasher::new(3, 2, vec![1, 3], vec![7; 4], vec![0; 4], EOS).is_err(),
+            "wrong multiplier count must be rejected at load"
+        );
+        // (3-1)*2 = 4 heads expected.
+        assert!(
+            NGramHasher::new(3, 2, vec![1, 3, 5], vec![7; 3], vec![0; 3], EOS).is_err(),
+            "wrong head count must be rejected at load"
+        );
+        assert!(
+            NGramHasher::new(3, 2, vec![1, 3, 5], vec![0; 4], vec![0; 4], EOS).is_err(),
+            "a zero vocab size would divide by zero"
+        );
+    }
+
+    // ---- consumption ----
+
+    use candle_core::{DType, Device};
+
+    /// A `PleBlock` whose projections and norms are supplied directly,
+    /// so each test can make the parts it is not measuring analytic.
+    /// All three norm weights are zero, which makes the `(1 + w)` scale
+    /// exactly 1 and leaves a pure RMS normalisation.
+    fn block(
+        hidden: usize,
+        hc: usize,
+        kernel_size: usize,
+        dilation: usize,
+        key_w: Tensor,
+        value_w: Tensor,
+        conv_w: Tensor,
+    ) -> PleBlock {
+        let dev = Device::Cpu;
+        let wide = hidden * hc;
+        let norm = || {
+            Qwen3_5RmsNorm::from_weight(
+                Tensor::zeros(wide, DType::F32, &dev).unwrap(),
+                1e-6,
+                Some(hidden),
+            )
+        };
+        PleBlock {
+            key_proj: Linear::new(key_w, None),
+            value_proj: Linear::new(value_w, None),
+            norm_key: norm(),
+            norm_query: norm(),
+            norm_conv: norm(),
+            conv1d_weight: conv_w,
+            conv_state: None,
+            hidden_size: hidden,
+            hc_count: hc,
+            kernel_size,
+            dilation,
+        }
+    }
+
+    fn eye(n: usize) -> Tensor {
+        let dev = Device::Cpu;
+        let mut v = vec![0.0f32; n * n];
+        for i in 0..n {
+            v[i * n + i] = 1.0;
+        }
+        Tensor::from_vec(v, (n, n), &dev).unwrap()
+    }
+
+    /// The gate, end to end, against arithmetic done by hand.
+    ///
+    /// With identity projections and zero norm weights, `k` and `q` both
+    /// normalise to all-ones over four dims, so the dot product is 4 and
+    /// the scale divides it by sqrt(4) = 2. The signed sqrt then takes
+    /// that 2 to sqrt(2) — the step a plain `g` would skip, and 2 vs
+    /// sqrt(2) is 0.88 vs 0.80 through the sigmoid, which is exactly the
+    /// kind of difference that reads as "slightly worse model".
+    #[test]
+    fn gate_matches_hand_computed_forward() {
+        let dev = Device::Cpu;
+        let hidden = 4;
+        let mut ple = block(
+            hidden,
+            1,
+            2,
+            3,
+            eye(hidden),
+            eye(hidden),
+            // Zero conv weight: silu(0) = 0, so the residual branch
+            // contributes nothing and `out` is the gate alone.
+            Tensor::zeros((hidden, 1, 2), DType::F32, &dev).unwrap(),
+        );
+
+        let e = Tensor::new(&[[[1.0f32, 1.0, 1.0, 1.0]]], &dev).unwrap();
+        // Magnitude is irrelevant — norm_query normalises it away, which
+        // is why the gate is bounded no matter how large the streams get.
+        let h = Tensor::new(&[[[2.0f32, 2.0, 2.0, 2.0]]], &dev).unwrap();
+
+        let got: Vec<f32> = ple
+            .forward(&h, &e)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let want = 1.0 / (1.0 + (-(2.0f32.sqrt())).exp());
+        assert_eq!(got.len(), hidden);
+        for g in &got {
+            assert!(
+                (g - want).abs() < 1e-5,
+                "gate: got {got:?} want {want} in every channel"
+            );
+        }
+    }
+
+    /// The conv is a residual branch. Zeroing the key projection shuts
+    /// the gate to sigmoid(0) = 0.5 and the output is then exactly half
+    /// the value projection, broadcast across every stream — which also
+    /// pins the flatten order (stream-major, hidden-minor).
+    #[test]
+    fn value_is_broadcast_across_streams_and_survives_a_dead_conv() {
+        let dev = Device::Cpu;
+        let (hidden, hc) = (2, 2);
+        let mut ple = block(
+            hidden,
+            hc,
+            2,
+            3,
+            Tensor::zeros((hidden * hc, hidden), DType::F32, &dev).unwrap(),
+            eye(hidden),
+            Tensor::zeros((hidden * hc, 1, 2), DType::F32, &dev).unwrap(),
+        );
+
+        let e = Tensor::new(&[[[3.0f32, 5.0]]], &dev).unwrap();
+        let h = Tensor::new(&[[[1.0f32, 1.0, 7.0, 7.0]]], &dev).unwrap();
+
+        let got: Vec<f32> = ple
+            .forward(&h, &e)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        let want = [1.5f32, 2.5, 1.5, 2.5];
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-5, "got {got:?} want {want:?}");
+        }
+    }
+
+    /// The conv reaches `(kernel_size - 1) * dilation` back, not
+    /// `kernel_size - 1`. A single impulse is placed at position 0 and
+    /// the far tap is the only live weight, so the response must appear
+    /// at position `dilation`. With the dilation dropped it would land
+    /// at position 1 — the whole difference between mixing an n-gram's
+    /// neighbours and mixing the token next door.
+    #[test]
+    fn short_conv_is_dilated_not_merely_causal() {
+        let dev = Device::Cpu;
+        let (wide, dilation) = (1, 3);
+        // Taps are [far, near]: output t = w[0]*x[t - dilation] + w[1]*x[t].
+        let conv_w = Tensor::from_vec(vec![1.0f32, 0.0], (wide, 1, 2), &dev).unwrap();
+        let mut ple = block(1, 1, 2, dilation, eye(1), eye(1), conv_w);
+
+        let x = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 0.0, 0.0], (1, 5, wide), &dev).unwrap();
+        let got: Vec<f32> = ple
+            .short_conv(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let silu1 = 1.0f32 / (1.0 + (-1.0f32).exp());
+        let want = [0.0, 0.0, 0.0, silu1, 0.0];
+        for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
+            assert!(
+                (g - w).abs() < 1e-5,
+                "position {i}: got {got:?} want {want:?}"
+            );
+        }
+        assert_eq!(ple.left_context(), dilation);
+    }
+
+    /// Decoding one token at a time must give the same answer as running
+    /// the whole sequence at once — the property the rolling state
+    /// exists for, and the one that breaks silently when the cached
+    /// context is `kernel_size` wide instead of the full reach.
+    #[test]
+    fn decode_step_by_step_matches_prefill() {
+        let dev = Device::Cpu;
+        let (wide, dilation) = (2, 3);
+        let conv_w = Tensor::from_vec(vec![0.5f32, -0.25, 1.5, 0.75], (wide, 1, 2), &dev).unwrap();
+        let mk = || block(2, 1, 2, dilation, eye(2), eye(2), conv_w.clone());
+
+        let seq: Vec<f32> = (0..16).map(|i| (i as f32 * 0.37).sin()).collect();
+        let x = Tensor::from_vec(seq.clone(), (1, 8, wide), &dev).unwrap();
+
+        let mut whole = mk();
+        let want: Vec<f32> = whole
+            .short_conv(&x)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let mut stepwise = mk();
+        let mut got = Vec::new();
+        for t in 0..8 {
+            let step = x.narrow(1, t, 1).unwrap().contiguous().unwrap();
+            got.extend(
+                stepwise
+                    .short_conv(&step)
+                    .unwrap()
+                    .flatten_all()
+                    .unwrap()
+                    .to_vec1::<f32>()
+                    .unwrap(),
+            );
+        }
+
+        assert_eq!(got.len(), want.len());
+        for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
+            assert!((g - w).abs() < 1e-5, "position {i}: got {g} want {w}");
+        }
+
+        // And a cleared block starts over rather than reading the tail of
+        // the previous request.
+        stepwise.clear_state();
+        let first = stepwise
+            .short_conv(&x.narrow(1, 0, 1).unwrap().contiguous().unwrap())
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1::<f32>()
+            .unwrap();
+        assert!(
+            (first[0] - want[0]).abs() < 1e-5,
+            "clear_state should return the block to a fresh prefill"
+        );
+    }
+
+    // ---- the table ----
+
+    /// Rows come back in the order asked for, not the order they were
+    /// fetched in. Two shards of three rows, requested out of order and
+    /// across the boundary: the shard-major gather returns them sorted,
+    /// so the inverse permutation is the whole test. Getting it wrong
+    /// hands every token somebody else's n-gram, silently.
+    #[test]
+    fn a_gather_preserves_the_order_it_was_asked_for() {
+        let dev = Device::Cpu;
+        // Row r holds [r, r] so a row is identifiable by its contents.
+        let shard = |base: f32| {
+            Tensor::from_vec(
+                vec![base, base, base + 1.0, base + 1.0, base + 2.0, base + 2.0],
+                (3, 2),
+                &dev,
+            )
+            .unwrap()
+        };
+        let table = ShardedNGramTable::from_shards(vec![shard(0.0), shard(3.0)]).unwrap();
+        assert_eq!(table.head_dim(), 2);
+
+        let got: Vec<f32> = table
+            .gather(&[4, 1, 3, 0, 5])
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        assert_eq!(got, vec![4.0, 4.0, 1.0, 1.0, 3.0, 3.0, 0.0, 0.0, 5.0, 5.0]);
+    }
+
+    /// Repeated rows are legitimate — sixteen heads can collide — and
+    /// must each come back.
+    #[test]
+    fn a_repeated_row_is_returned_once_per_request() {
+        let dev = Device::Cpu;
+        let shard = Tensor::from_vec(vec![0.0f32, 1.0, 2.0, 3.0], (2, 2), &dev).unwrap();
+        let table = ShardedNGramTable::from_shards(vec![shard]).unwrap();
+        let got: Vec<f32> = table
+            .gather(&[1, 1, 0])
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        assert_eq!(got, vec![2.0, 3.0, 2.0, 3.0, 0.0, 1.0]);
+    }
+
+    /// A row index past the end is a hashing bug upstream, and must not
+    /// become a panic in the middle of a forward.
+    #[test]
+    fn a_row_outside_the_table_is_an_error() {
+        let dev = Device::Cpu;
+        let shard = Tensor::zeros((2, 2), DType::F32, &dev).unwrap();
+        let table = ShardedNGramTable::from_shards(vec![shard]).unwrap();
+        assert!(table.gather(&[0, 1]).is_ok());
+        assert!(table.gather(&[2]).is_err());
+        assert!(table.gather(&[-1]).is_err());
+        assert_eq!(table.gather(&[]).unwrap().dims(), &[0, 2]);
+    }
+
+    /// The hasher and the table meet here: sixteen row indices per
+    /// position, gathered and flattened into the 2560-wide embedding the
+    /// block consumes. The widths have to agree or the projections do
+    /// not fit.
+    #[test]
+    fn hashed_rows_gather_into_one_embedding_per_position() {
+        let dev = Device::Cpu;
+        let heads = 4;
+        let head_dim = 2;
+        let rows = 64;
+        let shard = Tensor::from_vec(
+            (0..rows * head_dim).map(|v| v as f32).collect::<Vec<_>>(),
+            (rows, head_dim),
+            &dev,
+        )
+        .unwrap();
+        let table = ShardedNGramTable::from_shards(vec![shard]).unwrap();
+
+        let hasher = hasher(3, heads / 2);
+        let ids = [7i64, 11, 13];
+        let per_position = hasher.rows(&ids, ids.len()).unwrap();
+        assert_eq!(per_position.len(), 3);
+        assert_eq!(per_position[0].len(), heads);
+
+        let flat: Vec<i64> = per_position
+            .iter()
+            .flatten()
+            .map(|r| r.rem_euclid(rows as i64))
+            .collect();
+        let gathered = table.gather(&flat).unwrap();
+        assert_eq!(gathered.dims(), &[ids.len() * heads, head_dim]);
+        // Reshaped per position, that is the ple_embed_dim-wide row.
+        assert_eq!(
+            gathered
+                .reshape((ids.len(), heads * head_dim))
+                .unwrap()
+                .dims(),
+            &[3, heads * head_dim]
+        );
+    }
+}

crates/neuron/src/harness/candle.rs

@@ -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(),
         }
     }
 
@@ -1332,8 +1347,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"];
+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 +1409,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
@@ -2802,6 +2828,30 @@ 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")?
+                    };
+                    let model = super::arch::qwen4_exp::model::Qwen4ExpForCausalLM::load(
+                        &cfg,
+                        dtype,
+                        &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
@@ -8113,6 +8163,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/tp/tp_qwen3_5.rs

@@ -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,
         )?;

doc/qwen4_exp-port-spec.md

@@ -0,0 +1,618 @@
+# `qwen4_exp` port spec
+
+Reference specification for `Qwen/Qwen3.8-Flash-Next`, written so that
+the decoder can be implemented without re-reading the 131 KB upstream
+Python. Every dimension in this document was read from the checkpoint's
+own safetensors headers or from the pinned reference source; nothing is
+inferred from the model card.
+
+Companion to issue #308 (child A of epic #307).
+
+## Provenance
+
+| what | pinned value |
+|---|---|
+| reference source | `huggingface/transformers` `src/transformers/models/qwen4_exp/`, commit `fc5c5bde8e656dad91cbf34e61940d984b1c7b91` ("Add Qwen4Exp model", #48337), read at `main` = `5f8ab9bb53ec9e0c9329153d18bd825ff1db80f9` |
+| checkpoint | `Qwen/Qwen3.8-Flash-Next`, revision `de4b8e4d43b917e7706784d8bb445c9af86a3540` |
+| declared | `model_type: qwen4_exp`, `architectures: [Qwen4ExpForConditionalGeneration]`, `transformers_version: 5.8.0.dev0` |
+
+The checkpoint ships **no remote code**, so the transformers package is
+the only authority. Note that `modeling_qwen4_exp.py` has already been
+amended once since the model landed (commit `83d46aa2`, "Make gated
+delta rule more explicit") — re-diff before trusting this document for
+the linear-attention path.
+
+## Measured parameter split
+
+Read from all 131 shard headers (`total_size` 359,999,963,128 B;
+1,658 tensors; 1,655 BF16 + 3 I64). This **replaces** the estimate in
+#307 — the shape of the estimate was right, the residual bucket was
+understated because MTP was counted as dense.
+
+| component | params | share | BF16 | tensors |
+|---|---:|---:|---:|---:|
+| routed experts | 120,795,955,200 | 67.11% | 241.6 GB | 96 |
+| PLE n-gram tables | 51,200,245,760 | 28.44% | 102.4 GB | 128 |
+| MTP head | 2,607,150,848 | 1.45% | 5.2 GB | 31 |
+| linear attention | 2,086,510,464 | 1.16% | 4.2 GB | 324 |
+| hyper-connections | 640,624,640 | 0.36% | 1.3 GB | 387 |
+| `embed_tokens` | 635,699,200 | 0.35% | 1.3 GB | 1 |
+| `lm_head` | 635,699,200 | 0.35% | 1.3 GB | 1 |
+| full attention | 597,694,464 | 0.33% | 1.2 GB | 72 |
+| vision tower | 448,931,056 | 0.25% | 0.9 GB | 333 |
+| shared experts | 236,052,480 | 0.13% | 0.5 GB | 192 |
+| MoE routers | 62,914,560 | 0.03% | 0.1 GB | 48 |
+| PLE consumption (conv1d, projections, norms) | 32,839,683 | 0.02% | 0.1 GB | 7 |
+| QSA indexers | 19,663,872 | 0.01% | 0.04 GB | 36 |
+| **total** | **179,999,981,459** | | **360.0 GB** | **1,658** |
+
+The offload plan in #307 survives contact with the measurement:
+
+- **sparse-access, offloadable: 171.996 B (95.55%)** — routed experts
+  (10 of 512 per token) plus PLE tables (a few rows per token, one
+  layer).
+- **always-resident text path: 4.94 B (2.74%)** — everything above
+  minus MTP and vision. 9.9 GB at BF16.
+- MTP (2.61 B) is only resident if speculative decoding is on (#313);
+  vision (0.45 B) only if the vision tower is loaded (#314).
+
+## Skeleton
+
+```
+hidden_size            2560          num_hidden_layers   48
+head_dim               256           num_attention_heads 24
+num_key_value_heads    2             vocab_size          248,320
+max_position_embeddings              262,144
+layer_types            36 linear_attention + 12 full_attention,
+                       full_attention_interval 4 → layers 3,7,11,…,47
+num_experts            512           num_experts_per_tok 10
+moe_intermediate_size  640           shared_expert_intermediate_size 640
+norm_topk_prob         true (default; absent from config.json)
+hc_count               4             hc_lowrank          320
+ple_layer_ids          [2]  (one-indexed → zero-indexed layer 1)
+rms_norm_eps           1e-6          hidden_act          silu
+output_gate_type       sigmoid       mamba_ssm_dtype     float32
+```
+
+`layer_types` is authoritative; `full_attention_interval` is a hint. The
+same convention as `arch/qwen3_5/`.
+
+### Top-level dataflow
+
+```
+input_ids ──► embed_tokens ──► x [B,T,2560]
+                                 │
+                    repeat(1,1,hc_count) ──► h [B,T,10240]     4 identical streams
+                                 │
+        ┌────────────────────────▼───────────────────────────────┐
+        │ for layer in 0..48:                                    │
+        │     if layer == 1:  h = h + PLE(h, input_ids)          │  [B,T,10240]
+        │     h = decoder_layer(h)                               │
+        └────────────────────────┬───────────────────────────────┘
+                                 │
+                 hyper_connection_mixer(h)   (use_combine=False)
+                                 │
+                                 ▼  [B,T,2560]
+                             lm_head
+```
+
+Two things to note, because both differ from every arch we serve:
+
+1. **The inter-layer residual is 10240 wide, not 2560.** Four parallel
+   residual streams, initialised as four copies of the token embedding.
+2. **There is no final `model.norm`.** The `hc_norm` inside
+   `hyper_connection_mixer` is the only normalisation before `lm_head`.
+   Do not go looking for a missing tensor — the checkpoint has
+   `hyper_connection_mixer.{hc_norm,input_mix_weight_down,input_mix_weight_up}`
+   and no `norm.weight`.
+
+### Decoder layer
+
+There is **no `input_layernorm` and no `post_attention_layernorm`**. The
+hyper-connection block absorbs both roles: it normalises, mixes the four
+streams down to one 2560 vector for the sublayer, and scatters the
+sublayer's output back across the four streams.
+
+```
+h [B,T,10240]
+  │
+  ├─ (layer 1 only) h = h + PLE(h, ple_input_ids)
+  │
+  ├─ (x, h_saved, inj) = attn_hyper_connection(h)      x:[.,2560] inj:[.,4]
+  │   x = linear_attn(x)  |  x = self_attn(x)          → [.,2560]
+  │   h = h_saved + (x ⊗ inj).flatten(-2)              → [.,10240]
+  │
+  └─ (x, h_saved, inj) = mlp_hyper_connection(h)
+      x = sparse_moe(x)                                → [.,2560]
+      h = h_saved + (x ⊗ inj).flatten(-2)              → [.,10240]
+```
+
+`x ⊗ inj` is the outer product `x.unsqueeze(-2) * inj.unsqueeze(-1)`,
+giving `[.,4,2560]` — the sublayer output scaled by a per-stream weight
+before being added into each stream.
+
+## Component specs
+
+### 1. Hyper-connections (`Qwen4ExpTextGatedResidual`)
+
+Tensors, per instance:
+
+```
+hc_norm.weight              (10240,)
+input_mix_weight_down.weight  (320, 10240)
+input_mix_weight_up.weight    (10240, 320)
+block_inject_weight.weight    (4, 10240)     # absent on the final mixer
+```
+
+Instances: `layers.{i}.attn_hyper_connection`,
+`layers.{i}.mlp_hyper_connection` for all 48 layers, plus
+`model.language_model.hyper_connection_mixer` and
+`mtp.hyper_connection_mixer` (both `use_combine=False`, i.e. no
+`block_inject_weight`). 387 + 3 tensors.
+
+Forward, exactly:
+
+```
+hn   = hc_norm(h)                                   # grouped RMSNorm, see §8
+w    = silu(input_mix_weight_down(hn) / hc_count)   # [.,320]
+w    = sigmoid(input_mix_weight_up(w))              # [.,10240]
+w    = w.unflatten(-1, (4, 2560))
+x    = (w * hn.unflatten(-1, (4, 2560))).mean(dim=-2)     # [.,2560]
+inj  = 2 * sigmoid(block_inject_weight(hn) / hc_count)    # [.,4]
+return x, h, inj                                    # note: h, NOT hn
+```
+
+Three details that are easy to get wrong and produce fluent-but-worse
+output rather than a crash:
+
+- The **division by `hc_count` (4)** happens inside both gates, before
+  the nonlinearity.
+- The mix uses `.mean()` over the four streams, not `.sum()`.
+- The residual returned for the skip connection is the **un-normalised**
+  `h`. Only the mixing and injection paths see `hn`.
+
+### 2. PLE — per-layer embedding over hashed n-grams
+
+Applied at zero-indexed **layer 1** (`ple_layer_ids: [2]` is
+one-indexed; upstream computes
+`config.ple_layer_ids.index(layer_idx + 1)`, and the checkpoint
+confirms it — the tensors live at
+`model.language_model.layers.1.ple.*`).
+
+#### Table geometry
+
+```
+ngram_size            3        heads_per_ngram   8
+ngram_heads = (ngram_size - 1) * heads_per_ngram = 16
+ple_embed_dim         2560  →  head_dim_per_ngram = 2560 / 16 = 160
+ngram_vocab_size_base 20,000,000
+make_ngram_vocab_size_divisible_by 128
+split_ngram_parts     128
+```
+
+Head *h* (0..15) gets vocab size = the (h+1)-th prime strictly greater
+than 19,999,999; offsets are the running cumulative sum. Total padded up
+to a multiple of 128 gives **320,001,536 rows × 160 dims**, sharded as
+128 × `(2500012, 160)` tensors named
+`ple_embedding.ngram_embedding.shard_{0..127}.weight`. The shards are a
+**flat split of the concatenated table**, not a per-head split —
+concatenate along dim 0 in shard-index order and the global row index
+from the hash indexes straight into it.
+
+Do **not** reimplement the prime search or the splitmix64 multiplier
+derivation. All three derived quantities ship in the checkpoint as I64
+buffers and should simply be loaded:
+
+```
+ple_embedding.ngram_heads_offsets      (16,)
+ple_embedding.ngram_heads_vocab_sizes  (16,)
+ple_embedding.layer_multipliers        (3,)
+```
+
+#### Hashing
+
+Let `tok[s]` be the input ids shifted right by `s` (s = 0,1,2) and
+`m = layer_multipliers`. For n-gram order *n* ∈ {2, 3}, occupying heads
+`[(n-2)*8, (n-2)*8 + 8)`:
+
+```
+mixed = tok[0] * m[0]
+for p in 1..n:  mixed ^= tok[p] * m[p]        # XOR, int64
+row[head] = (mixed mod vocab_sizes[head]) + offsets[head]
+```
+
+so each token produces 16 row indices, gathers 16 × 160 values and
+flattens to a 2560-wide embedding. Arithmetic is int64 and the
+multipliers are ~10^13, so the products overflow int32 by a wide margin
+— this must be done in 64-bit.
+
+There is **no unigram head**: with `ngram_size = 3` the loop runs
+n = 2,3 only. `layer_multipliers` has three entries because the trigram
+needs three shift multipliers.
+
+#### Shifting is segment-aware
+
+`_shift_right_ignore_eos` does not shift across an EOS token. Within
+each EOS-delimited segment, positions closer to the segment start than
+`shift` get filled with `eos_token_id` (248044) instead of borrowing
+the previous document's tokens. Additionally
+`Qwen4ExpTextModel.forward` replaces padded positions in
+`ple_input_ids` with EOS before the layer is called. Getting this wrong
+is invisible on a single-turn prompt and wrong on anything batched or
+multi-document.
+
+#### Consumption
+
+```
+key_proj.weight    (10240, 2560)     value_proj.weight (2560, 2560)
+norm_key.weight    (10240,)          norm_query.weight (10240,)
+norm_conv.weight   (10240,)          conv1d.weight     (10240, 1, 4)
+```
+
+```
+e     = ngram_lookup(input_ids)                             # [.,2560]
+k     = norm_key(key_proj(e)).unflatten(-1, (4, 2560))
+v     = value_proj(e)                                       # [.,2560]
+q     = norm_query(h).unflatten(-1, (4, 2560))              # h is the 10240 stream
+g     = (k * q).sum(-1, keepdim=True) / sqrt(2560)          # [.,4,1]
+g     = sign(g) * sqrt(clamp_min(|g|, 1e-6))                # signed sqrt
+gv    = sigmoid(g) * v.unsqueeze(-2)                        # [.,4,2560]
+out   = gv.flatten(-2) + short_conv(norm_conv(gv.flatten(-2)))
+```
+
+`short_conv` is a **depthwise, dilated** Conv1d: `groups = 10240`,
+`kernel_size = ple_conv_kernel_size = 4`, **`dilation = ngram_size = 3`**,
+no bias, followed by SiLU. Its receptive field is
+`(4 - 1) * 3 = 9` positions of left context, which is also the cached
+conv-state length. The dilation is the part most likely to be dropped
+by accident — it is not a plain causal conv.
+
+The three norms are grouped RMSNorm (§8) over 10240 in groups of 2560.
+
+#### Cache state
+
+PLE needs two pieces of per-request state, both stored upstream as
+extra "conv states" on layer 1:
+
+| slot | shape | contents |
+|---|---|---|
+| `conv_states[1]` | `(B, 10240, 9)` | short-conv left context |
+| `conv_states[2]` | `(B, 2)` | previous `ngram_size - 1` input **ids**, EOS-initialised |
+
+The second is token ids, not activations. Reset both when a request's
+KV cache is cleared.
+
+### 3. QSA sparse attention indexer
+
+On the 12 full-attention layers and the MTP layer.
+
+```
+indexer_n_heads 4   indexer_kv_heads 1   indexer_head_dim 128
+indexer_budget 2048 indexer_compress_ratio 4
+block_topk = indexer_budget / indexer_compress_ratio = 512
+
+index_qk_proj.weight (640, 2560)   # 4*128 query + 1*128 key, fused
+q_layernorm.weight   (128,)
+k_layernorm.weight   (128,)
+```
+
+`indexer_budget` is in **tokens**; 512 *blocks* of 4 tokens are
+selected. `index_qk_proj` is a fused projection: the first 512 output
+channels are the 4 query heads, the last 128 are the single index key.
+
+Per query position:
+
+1. `q = q_layernorm(q); q = RoPE(q, pos)` — RoPE at the query's own
+   position, using the same partial-rotary cos/sin as the main
+   attention (64 of 128 dims rotated).
+2. The raw index key (**pre-norm, pre-RoPE**) is appended to a
+   **separate indexer KV cache**, one head × 128 dims.
+3. Take the causally-visible key positions. Group them into
+   `floor(n_visible / 4)` complete blocks of 4 **consecutive visible
+   positions**.
+4. Pool: `mean` over the 4 keys **in float32**, then `k_layernorm`, then
+   RoPE at the **first token position of the block**.
+5. Score: `relu(q @ pooled_k^T).sum(over the 4 index heads) / sqrt(128)`.
+   The ReLU is before the head sum — a signed sum is a different model.
+6. Select the top `min(512, n_blocks)` blocks; expand back to their
+   4 token positions each.
+7. **Unconditionally append the tail** — the `n_visible mod 4` positions
+   that did not form a complete block are always attended, never
+   scored.
+8. The selection becomes a mask that is ANDed (bool) or added (float)
+   onto the causal mask before the main attention.
+
+Consequences for the port:
+
+- **The full KV cache is still materialised for every position.** #307's
+  24 KiB/token arithmetic holds. QSA bounds the *attention compute*, not
+  the cache.
+- **There is a second cache.** 12 layers × 1 head × 128 dims. At BF16
+  that is 3 KiB/token on top of the main 24 KiB/token — #307's KV table
+  omitted it. llama.cpp's own accounting on a 2-GPU box reports
+  6144 MB main + 2304 MB indexer at the full 262,144 window, i.e. the
+  indexer is a real ~27% surcharge, not a rounding error.
+- Below the budget (fewer than 2048 visible tokens) selection is
+  effectively a no-op: every complete block gets selected and the tail
+  is appended, so the result is dense causal attention. **Short-prompt
+  parity against a dense implementation is therefore a valid test**, and
+  is the cheapest correctness gate available.
+- The reference is a Python double loop over batch × query position. It
+  is a correctness oracle only — a usable implementation must batch the
+  block pooling and score in one pass, and gather rather than build a
+  dense `[T, kv_len]` mask.
+
+### 4. Full attention
+
+```
+q_proj.weight (12288, 2560)   # 24 heads × 256 × 2 — query and gate
+k_proj.weight   (512, 2560)   # 2 kv heads × 256
+v_proj.weight   (512, 2560)
+o_proj.weight (2560, 6144)
+q_norm.weight   (256,)  k_norm.weight (256,)
+```
+
+Mechanically identical to `arch/qwen3_5/full_attn.rs`: GQA 24:2, per-head
+RMSNorm on q and k, a sigmoid output gate carried in the widened
+`q_proj`, scaling `head_dim^-0.5`.
+
+The gate split is **per head, not a flat halving**: reshape to
+`(..., n_heads, 2 * head_dim)` and chunk on the last axis, so within
+each head the first 256 channels are the query and the last 256 are the
+gate. Splitting the 12288 vector in half is a different and wrong
+permutation.
+
+The gate is applied to the attention output **before** `o_proj`:
+`attn_out = attn_out * sigmoid(gate)`.
+
+### 5. Linear attention (GatedDeltaNet)
+
+```
+linear_num_key_heads 16   linear_key_head_dim   128
+linear_num_value_heads 48 linear_value_head_dim 128
+linear_conv_kernel_dim 4  → conv_dim = 16*128*2 + 48*128 = 10240
+
+in_proj_qkv.weight (10240, 2560)   in_proj_z.weight (6144, 2560)
+in_proj_a.weight      (48, 2560)   in_proj_b.weight   (48, 2560)
+conv1d.weight    (10240, 1, 4)     A_log (48,)  dt_bias (48,)
+norm.weight  (128,)                out_proj.weight (2560, 6144)
+```
+
+**Tensor names and layout are identical to Qwen3.6's split (non-fused)
+form**, which `arch/qwen3_5/linear_attn.rs` already loads — it branches
+on the presence of `in_proj_qkvz.weight` and takes the split path
+otherwise. Same delta rule, same conv, same state shapes.
+
+One real difference: the output gated RMSNorm's activation is
+`config.output_gate_type or config.hidden_act`, and this config sets
+**`output_gate_type: sigmoid`**. Our `Qwen3_5RmsNormGated` hardcodes
+SiLU. This must become configurable, or the linear-attention output is
+wrong on all 36 layers.
+
+`mamba_ssm_dtype: float32` — upstream carries the recurrent state in
+f32. We default the equivalent GatedDeltaNet state to a bf16 round-trip
+(#284, `NEURON_GDN_STATE_F32` off). **A port must not silently inherit
+our default against upstream's explicit choice**: default this arch to
+f32 state and measure the cost, rather than assuming #284's tradeoff
+transfers.
+
+### 6. MoE
+
+```
+mlp.gate.weight            (512, 2560)
+mlp.experts.gate_up_proj   (512, 1280, 2560)
+mlp.experts.down_proj      (512, 2560, 640)
+mlp.shared_expert.{gate,up}_proj.weight (640, 2560)
+mlp.shared_expert.down_proj.weight      (2560, 640)
+mlp.shared_expert_gate.weight (1, 2560)
+```
+
+```
+shared = sigmoid(shared_expert_gate(x)) * shared_expert(x)
+probs  = softmax(gate(x), dim=-1, dtype=f32)      # over all 512
+w, idx = topk(probs, 10)
+w      = w / w.sum(-1, keepdim=True)              # norm_topk_prob = true
+out    = sum_k w_k * expert_{idx_k}(x)  +  shared
+```
+
+Routing arithmetic matches `arch/qwen3_5/moe.rs` exactly. The **storage
+layout does not**: this checkpoint ships experts as fused 3D tensors,
+where ours loads per-expert `mlp.experts.{i}.{gate,up,down}_proj`. In
+`gate_up_proj` the first 640 output rows are the gate and the last 640
+are the up projection. The fused layout is the better one for a grouped
+GEMM; the loader needs a new branch, not a new algorithm.
+
+Sparsity is extreme: 10 of 512, `moe_intermediate_size` only 640.
+Per token, 10 × (1280 + 2560) × 2560 ≈ 98.3 M params ≈ 2.36 B across
+48 layers — the number that drives every placement decision in #309.
+
+### 7. Rotary — interleaved M-RoPE
+
+```
+rope_theta 1e7    partial_rotary_factor 0.25    mrope_section [11,11,10]
+rotary_dim = head_dim * 0.25 = 64  →  32 inverse frequencies
+```
+
+`sum(mrope_section) = 32` = the number of frequencies. The interleave
+assigns frequency index *i* to axis T, H or W:
+
+```
+freqs = freqs[0]                      # start from temporal
+freqs[..., 1:33:3] = freqs_H[..., 1:33:3]     # 11 indices: 1,4,…,31
+freqs[..., 2:30:3] = freqs_W[..., 2:30:3]     # 10 indices: 2,5,…,29
+```
+
+i.e. the layout is `T,H,W,T,H,W,…` with temporal keeping the remainder,
+rather than three contiguous chunks. Rotation is `rotate_half` over the
+**first 64 of 256** head dims; the remaining 192 pass through unrotated.
+
+**This is byte-for-byte the configuration `arch/qwen3_5/rope.rs`
+already implements** (`head_dim 256`, `partial_rotary_factor 0.25`,
+`mrope_section [11,11,10]`, `mrope_interleaved true`) — the module's
+own tests pin exactly these values. Verify the index sets against the
+slices above, but expect this item to be free.
+
+`rope_theta` differs between the main model (1e7) and the MTP layer's
+config block (also 1e7 here — `mtp.rope_theta: 10000000`); do not
+assume they always agree.
+
+### 8. Norms
+
+Two conventions coexist in the same file, and mixing them up is silent:
+
+| variant | scale | used by |
+|---|---|---|
+| `Qwen4ExpTextRMSNorm` | `(1.0 + weight)` | everything except the GDN output |
+| `Qwen4ExpTextRMSNormGated` | `weight` | GatedDeltaNet output gate only |
+
+`arch/qwen3_5/rmsnorm.rs` already implements both with these exact
+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. 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
+`rmsnorm.rs`.
+
+### 9. MTP head
+
+```
+mtp.pre_fc_norm_embedding.weight (2560,)        # plain RMSNorm
+mtp.pre_fc_norm_hidden.weight    (10240,)       # grouped, group_size 2560
+mtp.fc_embedding.weight       (2560, 2560)
+mtp.fc_hidden.weight          (2560, 2560)
+mtp.hyper_connection_mixer.*                    # use_combine = False
+mtp.layers.0.*                                  # a complete full_attention
+                                                # decoder layer: 512-expert
+                                                # MoE, shared expert, QSA
+                                                # indexer, both hyper-
+                                                # connections
+```
+
+`mtp_num_hidden_layers: 1`, `mtp.hybrid: true`, `mtp.layer_types:
+["full_attention"]`, `mtp_use_dedicated_embeddings: false` (so it reuses
+`embed_tokens` and `lm_head`).
+
+**Transformers does not implement the MTP head.**
+`Qwen4ExpPreTrainedModel._keys_to_ignore_on_load_unexpected = [r"^mtp.*"]`
+— 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
+part of #310's problem.
+
+### 10. Vision tower
+
+```
+depth 27   hidden_size 1152   intermediate_size 4304   num_heads 16
+patch_size 16   spatial_merge_size 2   temporal_patch_size 2
+out_hidden_size 2560   num_position_embeddings 2304
+hidden_act gelu_pytorch_tanh   deepstack_visual_indexes []
+```
+
+448.9 M params. Smaller than the 27B's tower and with deepstack merging
+disabled. Out of scope for the text path (#314); `language_model_only`
+is a config knob.
+
+## Reuse boundary against `arch/qwen3_5/`
+
+| piece | verdict |
+|---|---|
+| `rope.rs` | **reuse unchanged.** Identical `head_dim`/`partial_rotary_factor`/`mrope_section`/`mrope_interleaved`. Verify the interleave index sets, expect no change. |
+| `rmsnorm.rs` | **extend.** Both `(1+w)` and gated conventions already correct. Add `group_size`; make the gated variant's activation configurable (sigmoid here, SiLU in Qwen3.6). |
+| `linear_attn.rs` | **reuse, two deltas.** Names and layout already match the split form. Switch the output-gate activation to sigmoid; default the recurrent state to f32 per `mamba_ssm_dtype`. |
+| `full_attn.rs` | **extend.** Same GQA + q/k norm + widened-`q_proj` output gate. Add the QSA indexer and its separate cache. |
+| `moe.rs` | **extend loader only.** Routing arithmetic identical; add a fused-3D `experts.{gate_up,down}_proj` branch beside the per-expert one. |
+| `mlp.rs` | **reuse** for the shared expert. |
+| `decoder.rs` | **rewrite.** Hyper-connections replace both layernorms and the residual is 4× wide. |
+| `mod.rs` (config, load, forward) | **new sibling arch.** Config field vocabulary overlaps heavily; the model-level plumbing (10240 stream, no final norm, PLE at layer 1, MTP) does not. |
+| `snapshot.rs` | **extend.** Batched decode needs snapshot support for the new per-request state: PLE's two conv slots and the QSA indexer cache, on top of the existing conv + recurrent state. |
+| `vision.rs` | **likely reuse with a new config.** Deferred (#314). |
+
+Net: the linear-attention half of the model — 36 of 48 layers and the
+piece that took the most work in `qwen3_5` — carries over close to
+unchanged. The genuinely new maths is hyper-connections (§1), PLE (§2)
+and QSA (§3).
+
+## Correctness gates, in the order they should be built
+
+Each of these fails loudly, unlike the model as a whole, which fails
+fluently.
+
+1. **Grouped RMSNorm** against a hand-computed 4×2560 case.
+2. **Hyper-connection block** in isolation: feed a known `h`, check
+   `x`, `inj` and the reconstructed `h` against a torch trace. Do this
+   before stacking anything on it — a wrong stream mix still produces
+   plausible text.
+3. **PLE hashing** against the reference for a short id sequence,
+   including one that crosses an EOS, in int64.
+4. **QSA below budget** — a prompt under 2048 tokens must give
+   bit-comparable logits to dense causal attention through the same
+   layer.
+5. **QSA above budget** against the reference double loop on a small
+   synthetic case (short `indexer_budget`, so the selection is actually
+   exercised).
+6. Full-model logits parity on a short prompt against a reference trace.
+
+## Open questions this document does not settle
+
+- **The accept rule and proposal depth for the MTP head** (§9). Needs
+  llama.cpp #27836 or vLLM, not transformers.
+- **Whether QSA can engage flash-attn at all.** The selection is a
+  gather, our flavour builds are per-sm, and flash-attn is only wired on
+  some (#95). Measure; do not assume.
+- **Whether the f32 recurrent state (§5) costs anything measurable**
+  here, given 36 linear layers rather than Qwen3.6's proportion.
+- **How `kv_budget_mb` should account for the indexer cache** (§3) — it
+  is a second, differently-shaped allocation and `derive_limit` does not
+  know about it. Feeds #310 and #315.

neuron.example.toml

@@ -45,9 +45,25 @@ name = "candle"
 # default_source = "huggingface"
 
 # Per-scheme source endpoints. Each scheme maps to an HF-compatible
-# registry. The `huggingface` source is auto-synthesised pointing at
-# `https://huggingface.co` when omitted; declare it explicitly here to
-# override the endpoint, auth env, or cache dir.
+# registry. The `huggingface` source is auto-synthesised when omitted;
+# declare it explicitly here to override the endpoint, auth env, or
+# cache dir.
+#
+# Endpoint resolution for the synthesised source (first hit wins):
+#   1. `endpoint` in an explicit block below.
+#   2. `HF_ENDPOINT` env var — the same convention the Python
+#      `huggingface_hub` library uses, so a whole host can be pointed
+#      at a mirror with one systemd drop-in and no config edit:
+#        /etc/systemd/system/neuron.service.d/local.conf:
+#          [Service]
+#          Environment=HF_ENDPOINT=https://rf.internal
+#          Environment=HF_TOKEN=...
+#   3. `https://huggingface.co`.
+#
+# Note the mirror must speak the HF-compatible wire format (siblings
+# listing at `/api/models/{org}/{name}`, blobs at
+# `/{org}/{name}/resolve/{rev}/{path}`), and if it requires auth, the
+# token comes from `auth_env` (default `HF_TOKEN`).
 #
 # [harness.candle.sources.huggingface]
 # endpoint = "https://huggingface.co"

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