feat(qwen4_exp): the decoder layer, where four streams replace the norms

The composition. Every other architecture we serve wraps a sublayer as
x = x + sublayer(norm(x)) over one residual stream; this one carries
four, 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, which is why looking for the
missing tensor is a dead end.

The spec's warning about this layer is the reason for its shape: a wrong
stream mix still produces fluent text. So the plumbing is a free
function taking the sublayer as a closure, and can be checked against a
known function without building an attention block. Two properties are
asserted there, because both are invisible downstream:

The sublayer is handed the mixed 2560 vector, not the 10240 stream —
every sublayer in the model is built for the narrow one, so the wide
tensor would be caught, but the narrow-but-wrong one would not.

And a sublayer that returns zero leaves the stream exactly as it was,
not as its normalised form. That single assertion is what separates a
correct residual from one that quietly re-normalises the model at every
layer, which is a plausible reading of a block whose only norm lives
inside the hyper-connection.

PLE is added into the stream before the attention hyper-connection
rather than merged into it, and only on the layers ple_layer_ids names
— one-indexed, so layer 1. A layer that carries the block and is handed
no gathered embedding says so rather than continuing without it.

Refs #307, #312.

.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

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

@@ -165,10 +165,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 +225,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,350 @@
+//! `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)?,
+            indexer: Indexer::load(vb, 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,40 @@
+//! `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.
+//! - [`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 moe;
+pub mod ple;
+pub mod qsa;

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,751 @@
+//! 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 itself, which is where the table lives
+//! and therefore belongs to #310 — at 27 GB quantised it is not
+//! device-resident, and nothing above depends on how it is fetched.
+//!
+//! ## 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
+    }
+}
+
+/// 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"
+        );
+    }
+}

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

@@ -0,0 +1,955 @@
+//! QSA — the sparse-attention indexer's block selection.
+//!
+//! The 12 full-attention layers (and the MTP layer) do not attend over
+//! every visible position. A small side-channel — 4 query heads and a
+//! single key head of 128 dims, projected by `index_qk_proj` — scores
+//! *blocks* of 4 consecutive visible positions, and only the best 512
+//! blocks are attended. That is a 2048-token budget over a 262,144-token
+//! window, which is why decode degrades only ~20% from 0 to 64k.
+//!
+//! [`BlockSelector`] is the selection arithmetic: pooling visible keys
+//! into blocks, scoring them, and turning the scores into a position
+//! set. [`Indexer`] wraps it in the layer's own machinery — the fused
+//! `index_qk_proj`, the two layer norms, the rotation, and the second
+//! KV cache the epic's per-token arithmetic omitted. The mask itself
+//! stays with the attention layer.
+//!
+//! ```text
+//! keys [n_visible, 128]  (raw: pre-norm, pre-RoPE, straight off the cache)
+//!   │
+//!   ├─ complete blocks of 4 ─► mean in f32 ─► k_layernorm ─► RoPE(block start)
+//!   │                                                          │
+//!   │   q = RoPE(q_layernorm(q), pos)   [4, 128] ──────────────┤
+//!   │                                                          ▼
+//!   ├─ score = relu(q @ pooled^T).sum(heads) / sqrt(128)   [n_blocks]
+//!   ├─ top-512 blocks ─► expand each to its 4 positions
+//!   └─ + the n_visible mod 4 tail positions, never scored, always kept
+//! ```
+//!
+//! Two of those steps are the ones to get right:
+//!
+//! 1. **The ReLU is before the head sum.** Summing signed scores lets a
+//!    head that strongly rejects a block cancel one that wants it, which
+//!    is a different model that still produces text.
+//! 2. **The tail is unconditional.** The `n_visible mod 4` most recent
+//!    positions never form a complete block, are never scored, and are
+//!    always attended — including the query's own position, which is
+//!    why dropping them is not a subtle degradation but a broken
+//!    causal chain.
+//!
+//! Below the budget the selection is a no-op — every complete block is
+//! selected and the tail is appended, so the result is dense causal
+//! attention. That makes short-prompt parity against a dense
+//! implementation an exact test rather than a tolerance, and it is the
+//! cheapest correctness gate this architecture offers.
+//!
+//! See `doc/qwen4_exp-port-spec.md` §3.
+
+use anyhow::{Context, Result, ensure};
+use candle_core::{D, DType, IndexOp, Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use crate::harness::arch::qwen3_5::rmsnorm::Qwen3_5RmsNorm;
+use crate::harness::arch::qwen3_5::rope::{RotaryEmbedding, apply_partial_rotary};
+
+use super::config::TextConfig as Qwen4ExpTextConfig;
+
+/// The indexer's geometry, from `config.json`.
+pub struct BlockSelector {
+    n_heads: usize,
+    head_dim: usize,
+    /// Positions per block — `indexer_compress_ratio`, 4.
+    block_size: usize,
+    /// `indexer_budget` in **tokens**, not blocks.
+    budget: usize,
+}
+
+impl BlockSelector {
+    pub fn new(n_heads: usize, head_dim: usize, block_size: usize, budget: usize) -> Result<Self> {
+        ensure!(n_heads >= 1, "indexer_n_heads must be >= 1");
+        ensure!(head_dim >= 1, "indexer_head_dim must be >= 1");
+        ensure!(block_size >= 1, "indexer_compress_ratio must be >= 1");
+        ensure!(
+            budget.is_multiple_of(block_size),
+            "indexer_budget ({budget}) must be a whole number of blocks of {block_size}"
+        );
+        Ok(Self {
+            n_heads,
+            head_dim,
+            block_size,
+            budget,
+        })
+    }
+
+    /// How many blocks may be selected: `indexer_budget /
+    /// indexer_compress_ratio`, 512 for the shipped config.
+    pub fn block_topk(&self) -> usize {
+        self.budget / self.block_size
+    }
+
+    /// Complete blocks over `n_visible` causally visible positions. The
+    /// remainder is the tail, and is not a block.
+    pub fn n_blocks(&self, n_visible: usize) -> usize {
+        n_visible / self.block_size
+    }
+
+    /// The position each block starts at — where its pooled key is
+    /// RoPE'd, **not** the position it is scored against.
+    pub fn block_starts(&self, n_visible: usize) -> Vec<usize> {
+        (0..self.n_blocks(n_visible))
+            .map(|b| b * self.block_size)
+            .collect()
+    }
+
+    /// Mean-pool the visible keys into blocks.
+    ///
+    /// `keys`: `[n_visible, head_dim]`, raw off the indexer cache. The
+    /// tail is dropped — it is never pooled and never scored.
+    ///
+    /// The mean is taken in f32 whatever the cache dtype, and the result
+    /// **stays f32**. A bf16 accumulation of four keys loses the small
+    /// ones outright, and casting the f32 mean straight back is no
+    /// better: one bf16 tick near 0.25 is 0.002, wider than the
+    /// difference the mean exists to carry. The caller normalises and
+    /// RoPEs this in f32 and casts once, at the matmul.
+    pub fn pool(&self, keys: &Tensor) -> candle_core::Result<Tensor> {
+        let (n_visible, dim) = keys.dims2()?;
+        debug_assert_eq!(dim, self.head_dim);
+        let blocks = self.n_blocks(n_visible);
+        if blocks == 0 {
+            return Tensor::zeros((0, dim), DType::F32, keys.device());
+        }
+        keys.narrow(0, 0, blocks * self.block_size)?
+            .to_dtype(DType::F32)?
+            .reshape((blocks, self.block_size, dim))?
+            .mean(1)
+    }
+
+    /// Score every block for one query position.
+    ///
+    /// `q`: `[n_heads, head_dim]`, already layer-normed and RoPE'd at
+    /// the query's own position. `pooled`: `[n_blocks, head_dim]`,
+    /// already layer-normed and RoPE'd at each block's start.
+    /// Returns `[n_blocks]`.
+    pub fn scores(&self, q: &Tensor, pooled: &Tensor) -> candle_core::Result<Tensor> {
+        debug_assert_eq!(q.dims2()?, (self.n_heads, self.head_dim));
+        let per_head = q.matmul(&pooled.t()?)?.relu()?;
+        per_head.sum(0)? / (self.head_dim as f64).sqrt()
+    }
+
+    /// Turn block scores into the visible positions the main attention
+    /// may see, ascending.
+    ///
+    /// Ties are broken towards the earlier block, which matters only for
+    /// reproducibility — the reference's `topk` does the same.
+    pub fn select(&self, scores: &[f32], n_visible: usize) -> Result<Vec<usize>> {
+        let blocks = self.n_blocks(n_visible);
+        ensure!(
+            scores.len() == blocks,
+            "expected {blocks} block scores for {n_visible} visible positions, got {}",
+            scores.len()
+        );
+
+        let mut order: Vec<usize> = (0..blocks).collect();
+        let keep = self.block_topk().min(blocks);
+        if keep < blocks {
+            order.sort_by(|a, b| {
+                scores[*b]
+                    .partial_cmp(&scores[*a])
+                    .unwrap_or(std::cmp::Ordering::Equal)
+                    .then(a.cmp(b))
+            });
+            order.truncate(keep);
+            order.sort_unstable();
+        }
+
+        let tail_start = blocks * self.block_size;
+        let mut positions = Vec::with_capacity(keep * self.block_size + (n_visible - tail_start));
+        for b in order {
+            let start = b * self.block_size;
+            positions.extend(start..start + self.block_size);
+        }
+        // The tail never competes for the budget. It holds the query's
+        // own position.
+        positions.extend(tail_start..n_visible);
+        Ok(positions)
+    }
+
+    /// Whether the selection can bind at all at this depth. Below the
+    /// budget it cannot, and the layer is free to skip scoring entirely
+    /// and attend densely — same answer, no indexer work.
+    pub fn is_dense(&self, n_visible: usize) -> bool {
+        self.n_blocks(n_visible) <= self.block_topk()
+    }
+}
+
+/// The indexer as a layer: projection, norms, rotation, and its own KV
+/// cache, wrapped around [`BlockSelector`].
+///
+/// ```text
+/// hidden [B,L,2560] ─ index_qk_proj ─┬─ q  [B,L,4,128] ─ q_layernorm ─ RoPE(own pos)
+///                                    └─ k  [B,L,128] ─────────────────► cache (raw)
+///                                                                        │
+///        pooled blocks ◄─ RoPE(block start) ◄─ k_layernorm ◄─ mean(f32) ◄┘
+/// ```
+///
+/// Three orderings matter and none of them are guessable from the
+/// shapes:
+///
+/// 1. **The cached key is raw** — pre-norm and pre-RoPE. Normalising or
+///    rotating before the cache would rotate each key at its own
+///    position, but a pooled block is rotated once at the block's
+///    *first* position. Caching post-RoPE quietly makes every block's
+///    position wrong except its first member's.
+/// 2. **Pooling comes before the norm**, not after: `mean` then
+///    `k_layernorm`, so the block's key is normalised as a block.
+/// 3. **`index_qk_proj` is fused, and the split is not a halving.** The
+///    first `n_heads * head_dim` output channels are the queries and
+///    the last `kv_heads * head_dim` are the single index key — 512 and
+///    128 of 640, not 320 and 320.
+///
+/// This is the second KV cache: 12 layers x 1 head x 128 dims, about
+/// 3 KiB/token at bf16 on top of the main 24 KiB. llama.cpp measures
+/// 2304 MB of it beside 6144 MB of main cache at the full 262k window,
+/// so it is a ~27% surcharge and belongs in any `kv_budget_mb`
+/// arithmetic (#310, #315).
+pub struct Indexer {
+    /// Fused `index_qk_proj`: `[n_heads * head_dim + kv_heads * head_dim, hidden]`.
+    qk_proj: Linear,
+    q_layernorm: Qwen3_5RmsNorm,
+    k_layernorm: Qwen3_5RmsNorm,
+    selector: BlockSelector,
+    n_heads: usize,
+    kv_heads: usize,
+    head_dim: usize,
+    /// Raw index keys, `(B, T, kv_heads * head_dim)`. Pre-norm,
+    /// pre-RoPE — see the note above.
+    key_cache: Option<Tensor>,
+}
+
+impl Indexer {
+    /// `vb` should be `.pp(...)`-ed to the attention layer's prefix, so
+    /// the tensors resolve as `index_qk_proj`, `q_layernorm`,
+    /// `k_layernorm`. Every dimension comes from the config rather than
+    /// the caller, so a layer cannot be built against a geometry the
+    /// checkpoint does not declare.
+    pub fn load(vb: &ShardedVarBuilder, cfg: &Qwen4ExpTextConfig) -> Result<Self> {
+        let (n_heads, kv_heads) = (cfg.indexer_n_heads, cfg.indexer_kv_heads);
+        let head_dim = cfg.indexer_head_dim;
+        let fused = (n_heads + kv_heads) * head_dim;
+        let weight = vb
+            .pp("index_qk_proj")
+            .get((fused, cfg.hidden_size), "weight")
+            .with_context(|| format!("load '{}/index_qk_proj/weight'", vb.prefix()))?;
+        Ok(Self {
+            qk_proj: Linear::new(weight, None),
+            q_layernorm: Qwen3_5RmsNorm::load(&vb.pp("q_layernorm"), head_dim, cfg.rms_norm_eps)?,
+            k_layernorm: Qwen3_5RmsNorm::load(&vb.pp("k_layernorm"), head_dim, cfg.rms_norm_eps)?,
+            selector: BlockSelector::new(
+                n_heads,
+                head_dim,
+                cfg.indexer_compress_ratio,
+                cfg.indexer_budget,
+            )?,
+            n_heads,
+            kv_heads,
+            head_dim,
+            key_cache: None,
+        })
+    }
+
+    /// Build an indexer from parts, for tests in sibling modules that
+    /// have no VarBuilder to load from. The norms carry zero weight, so
+    /// `(1 + w)` leaves a plain RMS normalisation.
+    #[cfg(test)]
+    pub(crate) fn from_parts(
+        qk_proj: Linear,
+        selector: BlockSelector,
+        n_heads: usize,
+        kv_heads: usize,
+        head_dim: usize,
+    ) -> Self {
+        let norm = || {
+            Qwen3_5RmsNorm::from_weight(
+                Tensor::zeros(head_dim, DType::F32, &candle_core::Device::Cpu).unwrap(),
+                1e-6,
+                None,
+            )
+        };
+        Self {
+            qk_proj,
+            q_layernorm: norm(),
+            k_layernorm: norm(),
+            selector,
+            n_heads,
+            kv_heads,
+            head_dim,
+            key_cache: None,
+        }
+    }
+
+    pub fn selector(&self) -> &BlockSelector {
+        &self.selector
+    }
+
+    /// Drop the indexer's cache. Must happen with the main KV clear —
+    /// a stale indexer cache selects blocks of another request's text.
+    pub fn clear_cache(&mut self) {
+        self.key_cache = None;
+    }
+
+    /// How many positions the indexer has seen.
+    pub fn cached_len(&self) -> usize {
+        self.key_cache.as_ref().map_or(0, |c| c.dims()[1])
+    }
+
+    /// Split the fused projection into normed queries and raw keys.
+    ///
+    /// Returns `q` as `(B, n_heads, L, head_dim)` — normed, unrotated —
+    /// and `k` as `(B, L, kv_heads * head_dim)`, untouched.
+    fn project(&self, hidden: &Tensor) -> candle_core::Result<(Tensor, Tensor)> {
+        let (b, l, _) = hidden.dims3()?;
+        let qk = self.qk_proj.forward(hidden)?;
+        let q_width = self.n_heads * self.head_dim;
+        let q = qk
+            .narrow(D::Minus1, 0, q_width)?
+            .reshape((b, l, self.n_heads, self.head_dim))?;
+        let q = self
+            .q_layernorm
+            .forward(&q)?
+            .transpose(1, 2)?
+            .contiguous()?;
+        let k = qk
+            .narrow(D::Minus1, q_width, self.kv_heads * self.head_dim)?
+            .contiguous()?;
+        Ok((q, k))
+    }
+
+    /// Pool the cache into blocks, norm them, and rotate each at its own
+    /// block's first position. `(B, 1, n_blocks, head_dim)`.
+    fn block_keys(
+        &self,
+        cache: &Tensor,
+        rope: &RotaryEmbedding,
+        total: usize,
+    ) -> candle_core::Result<Tensor> {
+        let b = cache.dims()[0];
+        let starts = self.selector.block_starts(total);
+        let mut pooled = Vec::with_capacity(b);
+        for row in 0..b {
+            pooled.push(self.selector.pool(&cache.i(row)?)?);
+        }
+        let pooled = Tensor::stack(&pooled, 0)?.to_dtype(cache.dtype())?;
+        let normed = self.k_layernorm.forward(&pooled)?.unsqueeze(1)?;
+        let (cos, sin) = rope.cos_sin_at(&starts)?;
+        apply_partial_rotary(&normed, &cos, &sin, rope.rotary_dim())
+    }
+
+    /// The additive attention mask for this step: `(B, 1, L, past + L)`,
+    /// `0` where the query may attend and `-inf` where it may not.
+    ///
+    /// Same shape and convention as `qwen3_5`'s own causal mask, so it
+    /// goes where that one goes — and it is already causal, since the
+    /// selection only ever offers positions at or behind the query.
+    ///
+    /// **This mask must not reach the flash-attn path.** That path takes
+    /// `attn_mask.is_some()` as "apply causal masking" and never reads
+    /// the tensor, so a QSA mask handed to it is silently replaced by a
+    /// dense causal one — the model runs, the budget does nothing, and
+    /// the only symptom is that long context costs what it always did.
+    /// Whether QSA can use flash-attn at all is an open question in the
+    /// spec (a gather, not a mask, is what the kernel would need);
+    /// until it is answered, a layer building this must take the eager
+    /// path.
+    ///
+    /// The dense `[L, S]` form is the correctness-first one. The spec is
+    /// explicit that a serving implementation should gather the selected
+    /// positions instead of materialising a mask over the whole window —
+    /// at 262k that is a 68 G-element tensor per layer.
+    pub fn attention_mask(
+        &mut self,
+        hidden: &Tensor,
+        rope: &RotaryEmbedding,
+        past_len: usize,
+        dtype: DType,
+    ) -> Result<Tensor> {
+        let (batch, seq_len, _) = hidden.dims3()?;
+        let visible = self.visible_positions(hidden, rope, past_len)?;
+        Ok(dense_mask(
+            &visible,
+            batch,
+            seq_len,
+            past_len + seq_len,
+            hidden.device(),
+            dtype,
+        )?)
+    }
+
+    /// The mask this step needs, or `None` when the selection cannot
+    /// bind at this depth.
+    ///
+    /// This is the serving entry point. Below the budget the mask would
+    /// be the plain causal mask exactly, so `None` lets the caller keep
+    /// the one it already has — and keep the flash path, which an
+    /// additive mask forecloses. The cache is advanced either way.
+    ///
+    /// The saving is not just the mask: below the budget the pooling,
+    /// the rotation of the block keys and the scoring are all skipped,
+    /// because their answer is known.
+    pub fn binding_mask(
+        &mut self,
+        hidden: &Tensor,
+        rope: &RotaryEmbedding,
+        past_len: usize,
+        dtype: DType,
+    ) -> Result<Option<Tensor>> {
+        let (batch, seq_len, _) = hidden.dims3()?;
+        let (q, cache) = self.advance(hidden, past_len)?;
+        let total = cache.dims()[1];
+        if self.selector.is_dense(total) {
+            return Ok(None);
+        }
+        let visible = self.positions_from(&q, &cache, rope, past_len)?;
+        Ok(Some(dense_mask(
+            &visible,
+            batch,
+            seq_len,
+            total,
+            hidden.device(),
+            dtype,
+        )?))
+    }
+
+    /// The positions each query may attend, as `[batch][query]`.
+    ///
+    /// `past_len` is the query's sequence offset, so query `i` sits at
+    /// absolute position `past_len + i` and sees `past_len + i + 1`
+    /// positions.
+    ///
+    /// The pooling and scoring are batched — one matmul covers every
+    /// query against every block — because a block's pooled key does not
+    /// depend on which query is looking at it. Only the top-k runs per
+    /// query, over the prefix of blocks that query can see. The
+    /// reference is a double loop over batch x query position; it is an
+    /// oracle, not a design.
+    pub fn visible_positions(
+        &mut self,
+        hidden: &Tensor,
+        rope: &RotaryEmbedding,
+        past_len: usize,
+    ) -> Result<Vec<Vec<Vec<usize>>>> {
+        let (q, cache) = self.advance(hidden, past_len)?;
+        self.positions_from(&q, &cache, rope, past_len)
+    }
+
+    /// Project this step and extend the cache. Every entry point goes
+    /// through here, because a step that skipped it would leave the
+    /// indexer's cache missing keys that later blocks pool.
+    fn advance(&mut self, hidden: &Tensor, past_len: usize) -> Result<(Tensor, Tensor)> {
+        let (_, seq_len, _) = hidden.dims3()?;
+        let (q, k) = self.project(hidden)?;
+        let cache = match self.key_cache.take() {
+            Some(prev) => Tensor::cat(&[&prev, &k], 1)?,
+            None => k,
+        };
+        let total = cache.dims()[1];
+        ensure!(
+            total == past_len + seq_len,
+            "indexer cache holds {total} positions but the queries start at {past_len} \
+             and run for {seq_len} — the cache and the main KV have diverged"
+        );
+        self.key_cache = Some(cache.clone());
+        Ok((q, cache))
+    }
+
+    fn positions_from(
+        &self,
+        q: &Tensor,
+        cache: &Tensor,
+        rope: &RotaryEmbedding,
+        past_len: usize,
+    ) -> Result<Vec<Vec<Vec<usize>>>> {
+        let (batch, _, seq_len, _) = q.dims4()?;
+        let total = cache.dims()[1];
+
+        // Queries rotate at their own positions, which are contiguous.
+        let (cos, sin) = rope.plain_cos_sin(past_len, seq_len)?;
+        let q = apply_partial_rotary(q, &cos, &sin, rope.rotary_dim())?;
+
+        let n_blocks = self.selector.n_blocks(total);
+        if n_blocks == 0 {
+            // Nothing is poolable yet; every query is all tail.
+            return Ok((0..batch)
+                .map(|_| (0..seq_len).map(|i| (0..=past_len + i).collect()).collect())
+                .collect());
+        }
+        let blocks = self.block_keys(cache, rope, total)?;
+
+        // relu(q . k) summed over heads, for every (query, block) pair
+        // in one batched matmul: every head and query shares the same
+        // block keys, so fold them into the row axis rather than
+        // broadcasting the keys across heads.
+        let q_rows = q.reshape((batch, self.n_heads * seq_len, self.head_dim))?;
+        let keys = blocks.squeeze(1)?.transpose(1, 2)?.contiguous()?;
+        let per_head = q_rows
+            .matmul(&keys)?
+            .reshape((batch, self.n_heads, seq_len, n_blocks))?
+            .relu()?;
+        let scores = (per_head.sum(1)? / (self.head_dim as f64).sqrt())?
+            .to_dtype(DType::F32)?
+            .to_vec3::<f32>()?;
+
+        let mut out = Vec::with_capacity(batch);
+        for row in scores.iter().take(batch) {
+            let mut per_query = Vec::with_capacity(seq_len);
+            for (i, all_blocks) in row.iter().enumerate().take(seq_len) {
+                let visible = past_len + i + 1;
+                // Only blocks wholly behind this query exist for it.
+                let visible_blocks = self.selector.n_blocks(visible);
+                per_query.push(
+                    self.selector
+                        .select(&all_blocks[..visible_blocks], visible)?,
+                );
+            }
+            out.push(per_query);
+        }
+        Ok(out)
+    }
+}
+
+/// Expand per-query position sets into the additive `(B, 1, L, S)` mask
+/// `qwen3_5`'s attention already takes.
+fn dense_mask(
+    visible: &[Vec<Vec<usize>>],
+    batch: usize,
+    seq_len: usize,
+    total: usize,
+    device: &candle_core::Device,
+    dtype: DType,
+) -> candle_core::Result<Tensor> {
+    let mut mask = vec![f32::NEG_INFINITY; batch * seq_len * total];
+    for (b, row) in visible.iter().enumerate() {
+        for (i, positions) in row.iter().enumerate() {
+            let base = (b * seq_len + i) * total;
+            for p in positions {
+                mask[base + p] = 0.0;
+            }
+        }
+    }
+    Tensor::from_vec(mask, (batch, 1, seq_len, total), device)?.to_dtype(dtype)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Device;
+
+    /// The shipped geometry, shrunk where the numbers do not matter:
+    /// 4 heads of 8 dims, blocks of 4, and a budget of 3 blocks so the
+    /// selection actually has to choose.
+    fn small() -> BlockSelector {
+        BlockSelector::new(4, 8, 4, 12).unwrap()
+    }
+
+    fn shipped() -> BlockSelector {
+        BlockSelector::new(4, 128, 4, 2048).unwrap()
+    }
+
+    #[test]
+    fn shipped_geometry_selects_512_blocks() {
+        let qsa = shipped();
+        assert_eq!(qsa.block_topk(), 512);
+        // 2048 tokens of budget over the full window.
+        assert_eq!(qsa.block_topk() * 4, 2048);
+    }
+
+    /// The cheapest gate in the architecture: below the budget, the
+    /// selection returns every visible position, so a short prompt must
+    /// be bit-comparable to dense causal attention.
+    #[test]
+    fn below_budget_selection_is_dense() {
+        let qsa = shipped();
+        for n_visible in [1usize, 3, 4, 7, 2047, 2048] {
+            assert!(qsa.is_dense(n_visible), "{n_visible} should be dense");
+            let scores = vec![0.0f32; qsa.n_blocks(n_visible)];
+            let got = qsa.select(&scores, n_visible).unwrap();
+            assert_eq!(
+                got,
+                (0..n_visible).collect::<Vec<_>>(),
+                "{n_visible} visible positions should all be attended"
+            );
+        }
+        // One block past the budget it can bind.
+        assert!(!qsa.is_dense(2052));
+    }
+
+    /// The tail is not scored and not budgeted. Here every block is
+    /// beaten out of a full budget, and the three tail positions —
+    /// including the query's own — must still come back.
+    #[test]
+    fn the_tail_is_always_attended() {
+        let qsa = small();
+        // 19 visible: 4 complete blocks (16 positions) + a 3-position tail.
+        let n_visible = 19;
+        assert_eq!(qsa.n_blocks(n_visible), 4);
+        // Block 2 is worthless; with topk 3 it is the one dropped.
+        let scores = [0.9f32, 0.8, 0.1, 0.7];
+        let got = qsa.select(&scores, n_visible).unwrap();
+        let want: Vec<usize> = (0..8).chain(12..19).collect();
+        assert_eq!(got, want);
+        assert!(got.contains(&(n_visible - 1)), "the query's own position");
+    }
+
+    /// ReLU before the head sum. Head 0 likes block 0 a little; head 1
+    /// hates it a lot. A signed sum ranks block 0 last; the reference
+    /// ranks it first, because rejection is clamped to indifference.
+    #[test]
+    fn relu_precedes_the_head_sum() {
+        let dev = Device::Cpu;
+        let qsa = BlockSelector::new(2, 1, 4, 8).unwrap();
+        let q = Tensor::from_vec(vec![1.0f32, 1.0], (2, 1), &dev).unwrap();
+        // block 0: +2 for head 0, -5 for head 1.  block 1: +1, +1.
+        let pooled = Tensor::from_vec(vec![2.0f32, 1.0], (2, 1), &dev).unwrap();
+        let scores: Vec<f32> = qsa.scores(&q, &pooled).unwrap().to_vec1().unwrap();
+        // Both heads share the query here, so the point is made by the
+        // per-head clamp: 2 + 2 = 4 beats 1 + 1 = 2.
+        assert!(scores[0] > scores[1]);
+
+        // Now give head 1 the opposite sign and check the clamp holds.
+        let q = Tensor::from_vec(vec![1.0f32, -2.5], (2, 1), &dev).unwrap();
+        let scores: Vec<f32> = qsa.scores(&q, &pooled).unwrap().to_vec1().unwrap();
+        // head 0: relu(2) = 2, head 1: relu(-5) = 0  ->  2
+        // signed would be 2 - 5 = -3, which loses to block 1's -1.5.
+        assert!(
+            (scores[0] - 2.0).abs() < 1e-6,
+            "expected the negative head clamped away, got {scores:?}"
+        );
+        assert!(scores[0] > scores[1], "got {scores:?}");
+    }
+
+    /// Pooling accumulates in f32 even when the cache is bf16, and keeps
+    /// the result there. Four keys whose sum a bf16 accumulator cannot
+    /// represent: 1.0 + 3 x 0.001 rounds straight back to 1.0, so a bf16
+    /// mean is exactly 0.25 — and so is an f32 mean cast back to bf16,
+    /// because the true 0.2507 sits below the first bf16 tick above
+    /// 0.25. Both mistakes land on the same wrong number, which is why
+    /// this asserts the value rather than the dtype.
+    #[test]
+    fn pooling_accumulates_in_f32() {
+        let dev = Device::Cpu;
+        let qsa = BlockSelector::new(4, 1, 4, 8).unwrap();
+        let keys = Tensor::from_vec(vec![1.0f32, 0.001, 0.001, 0.001], (4, 1), &dev)
+            .unwrap()
+            .to_dtype(DType::BF16)
+            .unwrap();
+        let pooled_t = qsa.pool(&keys).unwrap();
+        assert_eq!(pooled_t.dtype(), DType::F32);
+        let pooled: Vec<f32> = pooled_t.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(
+            (pooled[0] - 0.25).abs() > 1e-4,
+            "a bf16 accumulation, or an f32 one cast back, gives exactly \
+             0.25; got {pooled:?}"
+        );
+        assert!((pooled[0] - 0.2507).abs() < 1e-3, "got {pooled:?}");
+    }
+
+    /// The tail is dropped from pooling, and a prompt shorter than one
+    /// block pools nothing at all.
+    #[test]
+    fn pooling_drops_the_tail() {
+        let dev = Device::Cpu;
+        let qsa = small();
+        let keys = Tensor::zeros((11, 8), DType::BF16, &dev).unwrap();
+        assert_eq!(qsa.pool(&keys).unwrap().dims(), &[2, 8]);
+        assert_eq!(qsa.block_starts(11), vec![0, 4]);
+
+        let short = Tensor::zeros((3, 8), DType::F32, &dev).unwrap();
+        assert_eq!(qsa.pool(&short).unwrap().dims(), &[0, 8]);
+        assert!(qsa.block_starts(3).is_empty());
+    }
+
+    #[test]
+    fn selection_is_ascending_and_rejects_a_wrong_score_count() {
+        let qsa = small();
+        let got = qsa.select(&[0.1f32, 0.9, 0.5, 0.7, 0.2], 23).unwrap();
+        assert!(got.windows(2).all(|w| w[0] < w[1]), "got {got:?}");
+        // blocks 1, 3, 2 win; block 0 and 4 lose; tail 20..23 stays.
+        assert_eq!(
+            got,
+            vec![4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 21, 22]
+        );
+        assert!(qsa.select(&[0.1f32], 23).is_err());
+    }
+
+    // ---- the indexer as a layer ----
+
+    use crate::harness::arch::qwen3_5::TextConfig;
+
+    /// A rope whose `head_dim` matches the indexer's, with half of it
+    /// rotating — the shipped model rotates 64 of the indexer's 128,
+    /// which is the same fraction as the main attention's 64 of 256.
+    fn rope_for(head_dim: usize) -> RotaryEmbedding {
+        let json = format!(
+            r#"{{"vocab_size": 8, "hidden_size": 8, "intermediate_size": 8,
+                 "num_hidden_layers": 1, "num_attention_heads": 1,
+                 "num_key_value_heads": 1, "head_dim": {head_dim},
+                 "max_position_embeddings": 128, "rms_norm_eps": 1e-6,
+                 "rope_parameters": {{"rope_theta": 10000.0,
+                                      "partial_rotary_factor": 0.5}}}}"#
+        );
+        let cfg: TextConfig = serde_json::from_str(&json).unwrap();
+        RotaryEmbedding::new(DType::F32, &cfg, &Device::Cpu).unwrap()
+    }
+
+    /// Deterministic, well-spread values — the selection is a top-k, so
+    /// tied scores would make a test flaky rather than wrong.
+    fn spread(n: usize, seed: f32) -> Vec<f32> {
+        (0..n)
+            .map(|i| ((i as f32 * 0.7 + seed) * 1.3).sin() + 0.1 * (i as f32).cos())
+            .collect()
+    }
+
+    fn indexer(
+        hidden_size: usize,
+        n_heads: usize,
+        head_dim: usize,
+        block_size: usize,
+        budget: usize,
+    ) -> Indexer {
+        let dev = Device::Cpu;
+        let fused = (n_heads + 1) * head_dim;
+        let qk =
+            Tensor::from_vec(spread(fused * hidden_size, 0.3), (fused, hidden_size), &dev).unwrap();
+        Indexer::from_parts(
+            Linear::new(qk, None),
+            BlockSelector::new(n_heads, head_dim, block_size, budget).unwrap(),
+            n_heads,
+            1,
+            head_dim,
+        )
+    }
+
+    fn hidden(batch: usize, seq: usize, width: usize, seed: f32) -> Tensor {
+        Tensor::from_vec(
+            spread(batch * seq * width, seed),
+            (batch, seq, width),
+            &Device::Cpu,
+        )
+        .unwrap()
+    }
+
+    /// `index_qk_proj` is fused 4:1, not halved. With 4 query heads and
+    /// 1 key head of 8 dims the split is 32 and 8 of 40 — a halving
+    /// would take 20 and 20 and put half a query head in the key.
+    #[test]
+    fn the_fused_projection_splits_by_head_count_not_in_half() {
+        let (hidden_size, n_heads, head_dim) = (8, 4, 8);
+        let ix = indexer(hidden_size, n_heads, head_dim, 4, 16);
+        let h = hidden(1, 3, hidden_size, 1.1);
+
+        let (q, k) = ix.project(&h).unwrap();
+        assert_eq!(q.dims(), &[1, n_heads, 3, head_dim]);
+        assert_eq!(k.dims(), &[1, 3, head_dim]);
+
+        // The key is the LAST head_dim channels of the 40-wide fused
+        // output, and it is untouched by the norms.
+        let fused = ix.qk_proj.forward(&h).unwrap();
+        let want: Vec<f32> = fused
+            .narrow(D::Minus1, n_heads * head_dim, head_dim)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        let got: Vec<f32> = k.flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(got, want);
+    }
+
+    /// The cached key is raw: pre-norm and pre-RoPE. A block is rotated
+    /// once at its first position, so a cache of per-position-rotated
+    /// keys would put three of every four keys at the wrong angle — and
+    /// still return a plausible selection.
+    #[test]
+    fn the_cache_holds_raw_keys() {
+        let (hidden_size, head_dim) = (8, 8);
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 4096);
+        let h = hidden(1, 6, hidden_size, 2.2);
+
+        let (_, raw) = ix.project(&h).unwrap();
+        ix.visible_positions(&h, &rope_for(head_dim), 0).unwrap();
+
+        let cached: Vec<f32> = ix
+            .key_cache
+            .as_ref()
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        let want: Vec<f32> = raw.flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(cached, want, "the cache must not be normed or rotated");
+        assert_eq!(ix.cached_len(), 6);
+
+        ix.clear_cache();
+        assert_eq!(ix.cached_len(), 0);
+    }
+
+    /// Gate 4, through the whole layer rather than the bare selector:
+    /// below the budget every query sees every position it causally
+    /// could, so the layer is exactly dense causal attention.
+    #[test]
+    fn below_budget_every_query_sees_its_whole_past() {
+        let (hidden_size, head_dim) = (8, 8);
+        // Budget of 512 blocks, as shipped — 11 positions cannot reach it.
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 2048);
+        let h = hidden(2, 11, hidden_size, 3.3);
+
+        let got = ix.visible_positions(&h, &rope_for(head_dim), 0).unwrap();
+        assert_eq!(got.len(), 2);
+        for row in &got {
+            assert_eq!(row.len(), 11);
+            for (i, positions) in row.iter().enumerate() {
+                assert_eq!(
+                    *positions,
+                    (0..=i).collect::<Vec<_>>(),
+                    "query {i} should attend its whole past"
+                );
+            }
+        }
+    }
+
+    /// Decoding a token must select what a prefill covering the same
+    /// text would select for that position. This is the property the
+    /// cache and the `past_len` arithmetic exist for, and it is tested
+    /// with a budget small enough that the selection actually binds —
+    /// at 1 block of 4 from 9 positions, two blocks compete.
+    #[test]
+    fn a_decode_step_selects_what_the_prefill_would() {
+        let (hidden_size, head_dim) = (8, 8);
+        let rope = rope_for(head_dim);
+        let h = hidden(1, 9, hidden_size, 4.4);
+
+        let mut whole = indexer(hidden_size, 4, head_dim, 4, 4);
+        let want = whole.visible_positions(&h, &rope, 0).unwrap()[0][8].clone();
+
+        let mut split = indexer(hidden_size, 4, head_dim, 4, 4);
+        let prefill = h.narrow(1, 0, 8).unwrap().contiguous().unwrap();
+        split.visible_positions(&prefill, &rope, 0).unwrap();
+        let step = h.narrow(1, 8, 1).unwrap().contiguous().unwrap();
+        let got = split.visible_positions(&step, &rope, 8).unwrap()[0][0].clone();
+
+        assert_eq!(got, want);
+        // And the budget really did bind, or this proves nothing.
+        assert!(
+            got.len() < 9,
+            "expected a bound selection, got all {} positions",
+            got.len()
+        );
+        // The tail is 9 mod 4 = 1 position: the query's own.
+        assert!(got.contains(&8));
+    }
+
+    /// A cache that has drifted from the main KV selects blocks of some
+    /// other request's text. Fail loudly instead.
+    #[test]
+    fn a_cache_out_of_step_with_the_main_kv_is_an_error() {
+        let (hidden_size, head_dim) = (8, 8);
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 2048);
+        let h = hidden(1, 2, hidden_size, 5.5);
+        let err = ix
+            .visible_positions(&h, &rope_for(head_dim), 3)
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains("diverged"), "got: {err}");
+    }
+
+    /// Below the budget the mask must be the plain causal mask, value
+    /// for value — the consumable form of gate 4. If it is not, the
+    /// layer is not a drop-in for dense attention on short prompts and
+    /// the cheap parity test is worthless.
+    #[test]
+    fn below_budget_the_mask_is_exactly_causal() {
+        let (hidden_size, head_dim) = (8, 8);
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 2048);
+        let (batch, seq) = (2, 7);
+        let h = hidden(batch, seq, hidden_size, 6.6);
+
+        let got = ix
+            .attention_mask(&h, &rope_for(head_dim), 0, DType::F32)
+            .unwrap();
+        assert_eq!(got.dims(), &[batch, 1, seq, seq]);
+
+        let want: Vec<f32> = (0..seq)
+            .flat_map(|i| (0..seq).map(move |j| if j <= i { 0.0 } else { f32::NEG_INFINITY }))
+            .collect();
+        let got: Vec<f32> = got.i(0).unwrap().flatten_all().unwrap().to_vec1().unwrap();
+        assert_eq!(got, want);
+    }
+
+    /// Above the budget the mask is strictly sparser than causal and
+    /// still causal.
+    ///
+    /// It does **not** always keep the query's own position, which is
+    /// surprising enough to pin rather than discover. The tail — the
+    /// positions that never formed a complete block — is what carries
+    /// the diagonal, and when the visible count is an exact multiple of
+    /// the block size there is no tail at all: the query's own position
+    /// sits inside the last complete block and competes for the budget
+    /// like any other. So the diagonal is guaranteed for three visible
+    /// counts in four and merely likely for the fourth.
+    ///
+    /// That follows from step 7 of the spec as written. It is worth a
+    /// reference trace to confirm, because the alternative reading —
+    /// that the current block is always kept — is also a sensible
+    /// design and the two differ only every fourth position.
+    #[test]
+    fn above_budget_the_mask_is_sparse_but_still_causal() {
+        let (hidden_size, head_dim) = (8, 8);
+        let mut ix = indexer(hidden_size, 4, head_dim, 4, 4);
+        let seq = 12;
+        let h = hidden(1, seq, hidden_size, 7.7);
+
+        let m: Vec<f32> = ix
+            .attention_mask(&h, &rope_for(head_dim), 0, DType::F32)
+            .unwrap()
+            .i(0)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        let block = 4;
+        let mut sparser = false;
+        for i in 0..seq {
+            for j in (i + 1)..seq {
+                assert!(m[i * seq + j].is_infinite(), "query {i} peeked at {j}");
+            }
+            // A tail exists whenever the visible count is not a whole
+            // number of blocks, and the tail always holds the diagonal.
+            if (i + 1) % block != 0 {
+                assert_eq!(
+                    m[i * seq + i],
+                    0.0,
+                    "query {i} has a tail, so it must see itself"
+                );
+            }
+            let visible = (0..=i).filter(|j| m[i * seq + j] == 0.0).count();
+            if visible < i + 1 {
+                sparser = true;
+            }
+        }
+        assert!(sparser, "a budget of one block over 12 positions must bind");
+
+        // The exact-multiple case: 8 visible positions is two whole
+        // blocks and no tail, so position 7 is only attended if its own
+        // block won the budget.
+        let own_block_kept = m[7 * seq + 7] == 0.0;
+        let earlier_block_kept = m[7 * seq] == 0.0;
+        assert!(
+            own_block_kept != earlier_block_kept,
+            "with a one-block budget exactly one of the two blocks wins"
+        );
+    }
+}

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

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