feat(neuron): qwen3_next MoE FFN block, single-GPU path + HF parity (#92)

F1 slice 2 — the MoE block itself, CPU/single-GPU:

- arch/qwen3_5/moe.rs: Qwen3_5MoeBlock — top-k router (upstream
  softmax-then-topk order, renorm iff norm_topk_prob), per-expert
  SwiGLU (reusing Qwen3_5MLP at moe_intermediate_size), and the
  always-on shared expert mixed via sigmoid(shared_expert_gate).
  Correctness-first host-side scatter dispatch; the fused grouped-GEMM
  path is slice 4 behind the same forward signature.
- decoder.rs: MlpKind { Dense, Moe } dispatch on layer_uses_moe,
  mirroring AttentionKind.
- linear_attn.rs: fused-checkpoint support — qwen3_next stores
  in_proj_qkvz / in_proj_ba interleaved per key-head group (upstream
  fix_query_key_value_ordering layout); split_fused_qkvz/ba
  de-interleave once at load into the contiguous [Q|K|V] + Z / B + A
  layout the forward path (incl. the conv channels) already uses.
  Auto-detected via contains_tensor, so Qwen3.6 checkpoints are
  untouched.
- mod.rs: text_weight_prefix() — qwen3_next checkpoints put the text
  core at `model.*`, Qwen3.6 at `model.language_model.*`; the slice-1
  single-GPU MoE guard is removed (TP guard stays until slice 3).

Validation:
- qwen3_next_parity integration test replays a committed tiny
  random-weight HF Qwen3NextForCausalLM checkpoint (generated by
  script/dump_qwen3_next_tiny.py on beast: transformers 5.9.0,
  torch 2.9.1) through neuron's full load path: max_abs 0.000000,
  cosine 1.00000000 at f32 — exact parity, pinning the config
  normalisation, weight prefix, qkvz/ba de-interleave, hybrid layer
  interleaving, and the whole MoE block against upstream.
- Unit tests: scatter forward vs per-token dense reference,
  no-shared-expert/no-renorm behaviour, fused-split round-trip, and a
  flat-layout end-to-end structural load.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

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

@@ -22,6 +22,7 @@ use super::TextConfig;
 use super::full_attn::Qwen3_5Attention;
 use super::linear_attn::GatedDeltaNet;
 use super::mlp::Qwen3_5MLP;
+use super::moe::Qwen3_5MoeBlock;
 use super::rmsnorm::Qwen3_5RmsNorm;
 use super::rope::RotaryEmbedding;
 use super::snapshot::LayerKvSnapshot;
@@ -35,10 +36,27 @@ enum AttentionKind {
     Linear(GatedDeltaNet),
 }
 
+/// The FFN slot: dense SwiGLU (Qwen3.6) or the high-sparsity MoE block
+/// (qwen3_next 80B-A3B family, #92), selected per layer by
+/// [`TextConfig::layer_uses_moe`].
+enum MlpKind {
+    Dense(Qwen3_5MLP),
+    Moe(Qwen3_5MoeBlock),
+}
+
+impl Module for MlpKind {
+    fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
+        match self {
+            MlpKind::Dense(mlp) => mlp.forward(x),
+            MlpKind::Moe(moe) => moe.forward(x),
+        }
+    }
+}
+
 pub struct Qwen3_5DecoderLayer {
     input_layernorm: Qwen3_5RmsNorm,
     post_attention_layernorm: Qwen3_5RmsNorm,
-    mlp: Qwen3_5MLP,
+    mlp: MlpKind,
     attention: AttentionKind,
 }
 
@@ -73,7 +91,11 @@ impl Qwen3_5DecoderLayer {
             ),
         };
 
-        let mlp = Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?;
+        let mlp = if cfg.layer_uses_moe(layer_idx) {
+            MlpKind::Moe(Qwen3_5MoeBlock::load(cfg, &vb.pp("mlp"))?)
+        } else {
+            MlpKind::Dense(Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?)
+        };
         let input_layernorm =
             Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
         let post_attention_layernorm = Qwen3_5RmsNorm::load(

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

@@ -139,10 +139,42 @@ impl GatedDeltaNet {
         let conv_dim = key_dim * 2 + value_dim;
 
         // ----- Linear projections (all `bias=False` in the reference). -----
-        let in_proj_qkv = load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?;
-        let in_proj_z = load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?;
-        let in_proj_b = load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?;
-        let in_proj_a = load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?;
+        // Two checkpoint layouts exist for the input projections:
+        // - Qwen3.6 (qwen3_5): separate `in_proj_qkv` / `in_proj_z` /
+        //   `in_proj_b` / `in_proj_a`, with qkv stored as contiguous
+        //   [Q | K | V] blocks — loads directly.
+        // - Qwen3-Next 80B-A3B (qwen3_next, #92): fused `in_proj_qkvz`
+        //   + `in_proj_ba`, **interleaved per key-head group** (see
+        //   `split_fused_qkvz`/`split_fused_ba`) — de-interleaved once
+        //   at load into the same contiguous layout, so the forward
+        //   path (incl. the conv over [Q|K|V] channels) is unchanged.
+        let (in_proj_qkv, in_proj_z, in_proj_b, in_proj_a) =
+            if vb.contains_tensor("in_proj_qkvz.weight") {
+                let qkvz = vb
+                    .pp("in_proj_qkvz")
+                    .get((2 * key_dim + 2 * value_dim, cfg.hidden_size), "weight")
+                    .with_context(|| format!("load '{}/in_proj_qkvz/weight'", vb.prefix()))?;
+                let ba = vb
+                    .pp("in_proj_ba")
+                    .get((2 * num_v_heads, cfg.hidden_size), "weight")
+                    .with_context(|| format!("load '{}/in_proj_ba/weight'", vb.prefix()))?;
+                let (qkv_w, z_w) =
+                    split_fused_qkvz(&qkvz, num_k_heads, num_v_heads, head_k_dim, head_v_dim)?;
+                let (b_w, a_w) = split_fused_ba(&ba, num_k_heads, num_v_heads)?;
+                (
+                    Linear::new(qkv_w, None),
+                    Linear::new(z_w, None),
+                    Linear::new(b_w, None),
+                    Linear::new(a_w, None),
+                )
+            } else {
+                (
+                    load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?,
+                    load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?,
+                    load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?,
+                    load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?,
+                )
+            };
         let out_proj = load_linear_no_bias(vb, "out_proj", value_dim, cfg.hidden_size)?;
 
         // ----- Conv1d weight (depthwise, bias=False). -----
@@ -889,6 +921,57 @@ fn load_linear_no_bias(
     Ok(Linear::new(weight, None))
 }
 
+/// De-interleave a fused `in_proj_qkvz.weight` (qwen3_next layout, #92)
+/// into a contiguous `[Q | K | V]` qkv weight plus a `Z` weight.
+///
+/// The fused rows are grouped **per key head**: for each of the
+/// `num_k_heads` groups (`r = num_v_heads / num_k_heads`, group stride
+/// `s = 2*head_k + 2*head_v*r`), the group holds
+/// `[q (head_k) | k (head_k) | v (head_v*r) | z (head_v*r)]` — the
+/// reshape in upstream `fix_query_key_value_ordering`
+/// `(num_k_heads, 2*head_k + 2*head_v*num_v/num_k)`. Concatenating the
+/// per-group regions restores the global-contiguous layout the rest of
+/// this module (incl. the conv over `[Q|K|V]` channels) expects.
+fn split_fused_qkvz(
+    qkvz: &Tensor,
+    num_k_heads: usize,
+    num_v_heads: usize,
+    head_k_dim: usize,
+    head_v_dim: usize,
+) -> Result<(Tensor, Tensor)> {
+    let r = num_v_heads / num_k_heads;
+    let stride = 2 * head_k_dim + 2 * head_v_dim * r;
+    let (mut qs, mut ks, mut vs, mut zs) = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
+    for g in 0..num_k_heads {
+        let base = g * stride;
+        qs.push(qkvz.narrow(0, base, head_k_dim)?);
+        ks.push(qkvz.narrow(0, base + head_k_dim, head_k_dim)?);
+        vs.push(qkvz.narrow(0, base + 2 * head_k_dim, head_v_dim * r)?);
+        zs.push(qkvz.narrow(0, base + 2 * head_k_dim + head_v_dim * r, head_v_dim * r)?);
+    }
+    let parts: Vec<Tensor> = qs.into_iter().chain(ks).chain(vs).collect();
+    let qkv = Tensor::cat(&parts, 0)?.contiguous()?;
+    let z = Tensor::cat(&zs, 0)?.contiguous()?;
+    Ok((qkv, z))
+}
+
+/// De-interleave a fused `in_proj_ba.weight` (qwen3_next layout, #92)
+/// into per-v-head `b` (beta) and `a` (decay) weights. Same per-key-head
+/// grouping as [`split_fused_qkvz`]: each group holds `[b (r) | a (r)]`
+/// rows, `r = num_v_heads / num_k_heads`.
+fn split_fused_ba(ba: &Tensor, num_k_heads: usize, num_v_heads: usize) -> Result<(Tensor, Tensor)> {
+    let r = num_v_heads / num_k_heads;
+    let (mut bs, mut r#as) = (Vec::new(), Vec::new());
+    for g in 0..num_k_heads {
+        let base = g * 2 * r;
+        bs.push(ba.narrow(0, base, r)?);
+        r#as.push(ba.narrow(0, base + r, r)?);
+    }
+    let b = Tensor::cat(&bs, 0)?.contiguous()?;
+    let a = Tensor::cat(&r#as, 0)?.contiguous()?;
+    Ok((b, a))
+}
+
 /// Numerically-stable `softplus(x) = ln(1 + exp(x))`. Matches PyTorch's
 /// `F.softplus` default (beta=1, threshold=20: for large positive x,
 /// returns x as-is to avoid overflow in the exp).
@@ -1186,4 +1269,115 @@ mod tests {
         let v: Vec<f32> = y.flatten_all().unwrap().to_vec1().unwrap();
         assert!(v.iter().all(|x| x.is_finite()));
     }
+
+    /// Interleave known per-head Q/K/V/Z (and B/A) rows into the fused
+    /// qwen3_next layout, split, and expect the original contiguous
+    /// blocks back. Layout under test: per key-head group g,
+    /// `[q_g | k_g | v_g | z_g]` with r = num_v/num_k value heads per
+    /// group (upstream `fix_query_key_value_ordering`).
+    #[test]
+    fn split_fused_qkvz_and_ba_roundtrip() {
+        let dev = Device::Cpu;
+        let (num_k, num_v, head_k, head_v, hidden) = (2usize, 4usize, 3usize, 2usize, 5usize);
+        let r = num_v / num_k;
+
+        // Distinct constant per logical row so any mis-slicing shows.
+        let row = |tag: f32| Tensor::full(tag, (1, hidden), &dev).unwrap();
+        let mut fused_rows: Vec<Tensor> = Vec::new();
+        let (mut q_rows, mut k_rows, mut v_rows, mut z_rows) =
+            (Vec::new(), Vec::new(), Vec::new(), Vec::new());
+        for g in 0..num_k {
+            let base = 1000.0 * (g as f32 + 1.0);
+            for i in 0..head_k {
+                let t = row(base + i as f32);
+                fused_rows.push(t.clone());
+                q_rows.push(t);
+            }
+            for i in 0..head_k {
+                let t = row(base + 100.0 + i as f32);
+                fused_rows.push(t.clone());
+                k_rows.push(t);
+            }
+            for i in 0..head_v * r {
+                let t = row(base + 200.0 + i as f32);
+                fused_rows.push(t.clone());
+                v_rows.push(t);
+            }
+            for i in 0..head_v * r {
+                let t = row(base + 300.0 + i as f32);
+                fused_rows.push(t.clone());
+                z_rows.push(t);
+            }
+        }
+        let fused = Tensor::cat(&fused_rows, 0).unwrap();
+        let expected_qkv = Tensor::cat(
+            &q_rows
+                .iter()
+                .chain(k_rows.iter())
+                .chain(v_rows.iter())
+                .cloned()
+                .collect::<Vec<_>>(),
+            0,
+        )
+        .unwrap();
+        let expected_z = Tensor::cat(&z_rows, 0).unwrap();
+
+        let (qkv, z) = split_fused_qkvz(&fused, num_k, num_v, head_k, head_v).unwrap();
+        assert_eq!(qkv.dims(), &[2 * num_k * head_k + num_v * head_v, hidden]);
+        let diff_qkv: f32 = (qkv - expected_qkv)
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar()
+            .unwrap();
+        let diff_z: f32 = (z - expected_z)
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar()
+            .unwrap();
+        assert_eq!(diff_qkv, 0.0);
+        assert_eq!(diff_z, 0.0);
+
+        // ba: per group, [b (r rows) | a (r rows)].
+        let mut ba_rows = Vec::new();
+        let (mut b_rows, mut a_rows) = (Vec::new(), Vec::new());
+        for g in 0..num_k {
+            let base = 10.0 * (g as f32 + 1.0);
+            for i in 0..r {
+                let t = row(base + i as f32);
+                ba_rows.push(t.clone());
+                b_rows.push(t);
+            }
+            for i in 0..r {
+                let t = row(base + 5.0 + i as f32);
+                ba_rows.push(t.clone());
+                a_rows.push(t);
+            }
+        }
+        let ba = Tensor::cat(&ba_rows, 0).unwrap();
+        let (b, a) = split_fused_ba(&ba, num_k, num_v).unwrap();
+        let diff_b: f32 = (b - Tensor::cat(&b_rows, 0).unwrap())
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar()
+            .unwrap();
+        let diff_a: f32 = (a - Tensor::cat(&a_rows, 0).unwrap())
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar()
+            .unwrap();
+        assert_eq!(diff_b, 0.0);
+        assert_eq!(diff_a, 0.0);
+    }
 }

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

@@ -17,12 +17,32 @@ pub struct Qwen3_5MLP {
 }
 
 impl Qwen3_5MLP {
+    /// Construct directly from pre-built projections (MoE-block tests).
+    #[cfg(test)]
+    pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
+        Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+        }
+    }
+
     pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
-        let h = cfg.hidden_size;
-        let i = cfg.intermediate_size;
-        let gate_proj = load_linear_no_bias(vb, "gate_proj", h, i)?;
-        let up_proj = load_linear_no_bias(vb, "up_proj", h, i)?;
-        let down_proj = load_linear_no_bias(vb, "down_proj", i, h)?;
+        Self::load_with_dims(vb, cfg.hidden_size, cfg.intermediate_size)
+    }
+
+    /// Load with explicit dims — the MoE block (#92) reuses this SwiGLU
+    /// shape for routed experts (`moe_intermediate_size`) and the shared
+    /// expert (`shared_expert_intermediate_size`), both narrower than
+    /// the dense `intermediate_size`.
+    pub fn load_with_dims(
+        vb: &ShardedVarBuilder,
+        hidden: usize,
+        intermediate: usize,
+    ) -> Result<Self> {
+        let gate_proj = load_linear_no_bias(vb, "gate_proj", hidden, intermediate)?;
+        let up_proj = load_linear_no_bias(vb, "up_proj", hidden, intermediate)?;
+        let down_proj = load_linear_no_bias(vb, "down_proj", intermediate, hidden)?;
         Ok(Self {
             gate_proj,
             up_proj,

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

@@ -76,6 +76,7 @@ pub mod decoder;
 pub mod full_attn;
 pub mod linear_attn;
 pub mod mlp;
+pub mod moe;
 pub mod rmsnorm;
 pub mod rope;
 pub mod snapshot;
@@ -460,17 +461,20 @@ pub struct Qwen3_5Model {
 }
 
 impl Qwen3_5Model {
-    pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
+    /// `text_prefix` is where the text core lives in the checkpoint:
+    /// - Qwen3.6 (multimodal, `model_type = "qwen3_5"`):
+    ///   `model.language_model` — sibling to `model.visual.*` (the
+    ///   vision tower) and top-level `lm_head` / `mtp.*`.
+    /// - Qwen3-Next-80B-A3B (text-only, `model_type = "qwen3_next"`):
+    ///   plain `model`.
+    ///
+    /// [`Qwen3_5ForCausalLM::new`] picks by `Config::model_type` via
+    /// [`text_weight_prefix`].
+    pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder, text_prefix: &str) -> Result<Self> {
         let dtype = vb.dtype();
         let device = vb.device().clone();
 
-        // Qwen3-Next is a multimodal architecture whose text core lives
-        // under `model.language_model.*` — sibling to `model.visual.*`
-        // (the vision tower) and to top-level `lm_head` / `mtp.*`.
-        // Every text-side tensor in the safetensors files is under
-        // this prefix; we ignore the vision and MTP weights for
-        // language-model inference.
-        let text_vb = vb.pp("model.language_model");
+        let text_vb = vb.pp(text_prefix);
 
         let embed_vb = text_vb.pp("embed_tokens");
         let embed_weight = embed_vb
@@ -489,17 +493,6 @@ impl Qwen3_5Model {
             );
         }
 
-        // MoE FFN wiring is F1 (#92) slice 2 — fail with a clear message
-        // rather than a cryptic missing-tensor error from the dense MLP
-        // loader ('mlp.gate_proj/weight not found').
-        if cfg.num_experts > 0 {
-            anyhow::bail!(
-                "config declares a MoE FFN (num_experts={}) but the qwen3_5 MoE \
-                 block is not implemented yet (#92); only dense-FFN checkpoints load",
-                cfg.num_experts
-            );
-        }
-
         let vb_l = text_vb.pp("layers");
         let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
         for i in 0..cfg.num_hidden_layers {
@@ -746,10 +739,20 @@ pub struct Qwen3_5ForCausalLM {
     image_token_id: Option<u32>,
 }
 
+/// Checkpoint prefix of the text core for a given `model_type` — see
+/// [`Qwen3_5Model::load`].
+pub fn text_weight_prefix(model_type: &str) -> &'static str {
+    if model_type == MODEL_TYPE_NEXT {
+        "model"
+    } else {
+        "model.language_model"
+    }
+}
+
 impl Qwen3_5ForCausalLM {
     pub fn new(config: Config, vb: ShardedVarBuilder) -> Result<Self> {
         let cfg = &config.text_config;
-        let base = Qwen3_5Model::load(cfg, &vb)?;
+        let base = Qwen3_5Model::load(cfg, &vb, text_weight_prefix(&config.model_type))?;
         let lm_head = if cfg.tie_word_embeddings {
             Linear::new(base.embed_weight().clone(), None)
         } else {
@@ -1137,6 +1140,129 @@ mod tests {
         assert!(t.layer_uses_moe(47));
     }
 
+    /// End-to-end structural check for the qwen3_next path (#92): a
+    /// tiny random-weight checkpoint in the **flat** layout (`model.*`
+    /// prefix, fused `in_proj_qkvz`/`in_proj_ba`, per-expert MoE
+    /// tensors, shared expert) loads through `Config::from_config_json`
+    /// and `Qwen3_5ForCausalLM::new`, producing finite logits of the
+    /// right shape. Numerical parity vs HF is pinned separately by the
+    /// `qwen3_next_parity` fixture test.
+    #[test]
+    fn tiny_qwen3_next_checkpoint_loads_and_forwards() {
+        use candle_core::Device;
+        use std::collections::HashMap;
+
+        let raw = r#"{
+            "model_type": "qwen3_next",
+            "vocab_size": 32, "hidden_size": 8, "intermediate_size": 16,
+            "num_hidden_layers": 2, "num_attention_heads": 2,
+            "num_key_value_heads": 1, "head_dim": 4,
+            "max_position_embeddings": 64, "rms_norm_eps": 1e-6,
+            "full_attention_interval": 2,
+            "linear_num_value_heads": 4, "linear_num_key_heads": 2,
+            "linear_key_head_dim": 4, "linear_value_head_dim": 4,
+            "linear_conv_kernel_dim": 4,
+            "num_experts": 4, "num_experts_per_tok": 2,
+            "moe_intermediate_size": 4,
+            "shared_expert_intermediate_size": 4,
+            "norm_topk_prob": true
+        }"#;
+        let cfg = Config::from_config_json(raw).expect("parse tiny qwen3_next config");
+        assert_eq!(cfg.text_config.layer_types[0], "linear_attention");
+        assert_eq!(cfg.text_config.layer_types[1], "full_attention");
+
+        let dev = Device::Cpu;
+        let randn = |shape: &[usize]| Tensor::randn(0f32, 0.1f32, shape, &dev).unwrap();
+        let ones = |shape: &[usize]| Tensor::ones(shape, DType::F32, &dev).unwrap();
+        let mut t: HashMap<String, Tensor> = HashMap::new();
+
+        let (h, vocab) = (8usize, 32usize);
+        t.insert("model.embed_tokens.weight".into(), randn(&[vocab, h]));
+        t.insert("lm_head.weight".into(), randn(&[vocab, h]));
+        t.insert("model.norm.weight".into(), ones(&[h]));
+
+        let moe = |t: &mut HashMap<String, Tensor>, p: &str| {
+            t.insert(format!("{p}.gate.weight"), randn(&[4, h]));
+            for e in 0..4 {
+                t.insert(format!("{p}.experts.{e}.gate_proj.weight"), randn(&[4, h]));
+                t.insert(format!("{p}.experts.{e}.up_proj.weight"), randn(&[4, h]));
+                t.insert(format!("{p}.experts.{e}.down_proj.weight"), randn(&[h, 4]));
+            }
+            t.insert(
+                format!("{p}.shared_expert.gate_proj.weight"),
+                randn(&[4, h]),
+            );
+            t.insert(format!("{p}.shared_expert.up_proj.weight"), randn(&[4, h]));
+            t.insert(
+                format!("{p}.shared_expert.down_proj.weight"),
+                randn(&[h, 4]),
+            );
+            t.insert(format!("{p}.shared_expert_gate.weight"), randn(&[1, h]));
+        };
+
+        // Layer 0: linear_attention with the FUSED qwen3_next input
+        // projections. key_dim = 2*4 = 8, value_dim = 4*4 = 16 →
+        // qkvz rows = 2*8 + 2*16 = 48, ba rows = 2*4 = 8, conv_dim = 32.
+        let l0 = "model.layers.0";
+        t.insert(
+            format!("{l0}.linear_attn.in_proj_qkvz.weight"),
+            randn(&[48, h]),
+        );
+        t.insert(
+            format!("{l0}.linear_attn.in_proj_ba.weight"),
+            randn(&[8, h]),
+        );
+        t.insert(
+            format!("{l0}.linear_attn.conv1d.weight"),
+            randn(&[32, 1, 4]),
+        );
+        t.insert(format!("{l0}.linear_attn.dt_bias"), randn(&[4]));
+        t.insert(format!("{l0}.linear_attn.A_log"), randn(&[4]));
+        t.insert(format!("{l0}.linear_attn.norm.weight"), ones(&[4]));
+        t.insert(format!("{l0}.linear_attn.out_proj.weight"), randn(&[h, 16]));
+        t.insert(format!("{l0}.input_layernorm.weight"), ones(&[h]));
+        t.insert(format!("{l0}.post_attention_layernorm.weight"), ones(&[h]));
+        moe(&mut t, &format!("{l0}.mlp"));
+
+        // Layer 1: full_attention (output-gated: q_proj is 2×).
+        let l1 = "model.layers.1";
+        t.insert(
+            format!("{l1}.self_attn.q_proj.weight"),
+            randn(&[2 * 2 * 4, h]),
+        );
+        t.insert(format!("{l1}.self_attn.k_proj.weight"), randn(&[4, h]));
+        t.insert(format!("{l1}.self_attn.v_proj.weight"), randn(&[4, h]));
+        t.insert(format!("{l1}.self_attn.o_proj.weight"), randn(&[h, 8]));
+        t.insert(format!("{l1}.self_attn.q_norm.weight"), ones(&[4]));
+        t.insert(format!("{l1}.self_attn.k_norm.weight"), ones(&[4]));
+        t.insert(format!("{l1}.input_layernorm.weight"), ones(&[h]));
+        t.insert(format!("{l1}.post_attention_layernorm.weight"), ones(&[h]));
+        moe(&mut t, &format!("{l1}.mlp"));
+
+        let dir = tempfile::tempdir().expect("tempdir");
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&t, &path).expect("save safetensors");
+        // SAFETY: mmap of a file this test just wrote; nothing mutates it.
+        let vb = unsafe {
+            candle_nn::var_builder::ShardedSafeTensors::var_builder(
+                std::slice::from_ref(&path),
+                DType::F32,
+                &dev,
+            )
+            .expect("build ShardedVarBuilder")
+        };
+
+        let mut model = Qwen3_5ForCausalLM::new(cfg, vb).expect("load tiny qwen3_next checkpoint");
+        let input = Tensor::new(&[1u32, 5, 9], &dev)
+            .unwrap()
+            .unsqueeze(0)
+            .unwrap();
+        let logits = model.forward(&input, 0).expect("forward");
+        assert_eq!(logits.dims(), &[1, 1, vocab]);
+        let v: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
+        assert!(v.iter().all(|x| x.is_finite()), "logits must be finite");
+    }
+
     /// `mlp_only_layers` and `decoder_sparse_step` gate `layer_uses_moe`
     /// per the upstream convention.
     #[test]

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

@@ -0,0 +1,336 @@
+//! High-sparsity MoE FFN block for the qwen3_next family (#92).
+//!
+//! Qwen3-Next-80B-A3B replaces the dense SwiGLU in (almost) every
+//! decoder layer with `Qwen3NextSparseMoeBlock`: a top-k router over
+//! `num_experts` small SwiGLU experts, plus an always-on **shared
+//! expert** mixed in through a per-token sigmoid gate:
+//!
+//! ```text
+//! probs   = softmax(gate(x))                 # over ALL experts, f32
+//! w, idx  = topk(probs, num_experts_per_tok)
+//! w       = w / sum(w)                       # iff norm_topk_prob
+//! routed  = Σ_j w_j · expert_{idx_j}(x)
+//! shared  = sigmoid(shared_expert_gate(x)) · shared_expert(x)
+//! y       = routed + shared
+//! ```
+//!
+//! Routing follows the upstream softmax-then-topk order (NOT
+//! topk-then-softmax — the renormalisation only equals softmax over
+//! the selected logits when `norm_topk_prob` is on, and the reference
+//! renormalises the *global* softmax values).
+//!
+//! ## Dispatch strategy
+//!
+//! This is the correctness-first implementation: a host-side scatter
+//! loop over the experts that actually received tokens (the pattern
+//! candle-transformers' `Qwen3SparseMoeBlock` uses). Batch-1 decode
+//! touches `num_experts_per_tok` experts per layer; prefill batches
+//! per-expert token groups. The fused grouped-GEMM path (slice 4)
+//! replaces the loop behind the same `forward` signature.
+
+use anyhow::{Context, Result};
+use candle_core::{DType, Module, Tensor};
+use candle_nn::Linear;
+use candle_nn::var_builder::ShardedVarBuilder;
+
+use super::TextConfig;
+use super::mlp::Qwen3_5MLP;
+
+pub struct Qwen3_5MoeBlock {
+    /// Router: `(num_experts, hidden)`, checkpoint name `mlp.gate`.
+    gate: Linear,
+    /// Routed experts, checkpoint names `mlp.experts.{i}.{gate,up,down}_proj`.
+    experts: Vec<Qwen3_5MLP>,
+    /// Always-on expert, `mlp.shared_expert.*`. `None` when the config
+    /// declares no shared expert (Qwen3-30B-A3B style).
+    shared_expert: Option<Qwen3_5MLP>,
+    /// Per-token sigmoid mix for the shared expert: `(1, hidden)`,
+    /// checkpoint name `mlp.shared_expert_gate`.
+    shared_expert_gate: Option<Linear>,
+    num_experts_per_tok: usize,
+    norm_topk_prob: bool,
+}
+
+impl Qwen3_5MoeBlock {
+    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,
+            "MoE block needs num_experts ({}), num_experts_per_tok ({}) and \
+             moe_intermediate_size ({}) all > 0",
+            cfg.num_experts,
+            cfg.num_experts_per_tok,
+            cfg.moe_intermediate_size,
+        );
+        anyhow::ensure!(
+            cfg.num_experts_per_tok <= cfg.num_experts,
+            "num_experts_per_tok ({}) exceeds num_experts ({})",
+            cfg.num_experts_per_tok,
+            cfg.num_experts,
+        );
+
+        let h = cfg.hidden_size;
+
+        let gate_weight = vb
+            .pp("gate")
+            .get((cfg.num_experts, h), "weight")
+            .with_context(|| format!("load '{}/gate/weight'", vb.prefix()))?;
+        let gate = Linear::new(gate_weight, None);
+
+        let experts_vb = vb.pp("experts");
+        let mut experts = Vec::with_capacity(cfg.num_experts);
+        for i in 0..cfg.num_experts {
+            experts.push(
+                Qwen3_5MLP::load_with_dims(&experts_vb.pp(i), h, cfg.moe_intermediate_size)
+                    .with_context(|| format!("load expert {i}"))?,
+            );
+        }
+
+        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(Self {
+            gate,
+            experts,
+            shared_expert,
+            shared_expert_gate,
+            num_experts_per_tok: cfg.num_experts_per_tok,
+            norm_topk_prob: cfg.norm_topk_prob,
+        })
+    }
+}
+
+impl Module for Qwen3_5MoeBlock {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let (b, l, hidden) = xs.dims3()?;
+        let xs_flat = xs.reshape(((), hidden))?;
+        let n_tokens = b * l;
+
+        // Router probabilities in f32 (reference uses float softmax
+        // regardless of activations dtype).
+        let router_logits = self.gate.forward(&xs_flat)?;
+        let probs = candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?;
+
+        // Top-k selection: descending argsort, take the first k. The
+        // renormalisation (iff norm_topk_prob) divides by the sum of
+        // the selected global-softmax values.
+        let sorted = probs.arg_sort_last_dim(false)?;
+        let topk_idx = sorted
+            .narrow(1, 0, self.num_experts_per_tok)?
+            .contiguous()?;
+        let mut topk_w = probs.gather(&topk_idx, 1)?;
+        if self.norm_topk_prob {
+            let denom = topk_w.sum_keepdim(1)?;
+            topk_w = topk_w.broadcast_div(&denom)?;
+        }
+
+        // Host-side scatter: token row lists per expert. Cheap relative
+        // to the expert GEMMs; replaced by grouped-GEMM in slice 4.
+        let idx_host: Vec<Vec<u32>> = topk_idx.to_vec2()?;
+        let w_host: Vec<Vec<f32>> = topk_w.to_vec2()?;
+        let mut tokens_for: Vec<Vec<u32>> = vec![Vec::new(); self.experts.len()];
+        let mut weights_for: Vec<Vec<f32>> = vec![Vec::new(); self.experts.len()];
+        for t in 0..n_tokens {
+            for j in 0..self.num_experts_per_tok {
+                let e = idx_host[t][j] as usize;
+                tokens_for[e].push(t as u32);
+                weights_for[e].push(w_host[t][j]);
+            }
+        }
+
+        let mut ys = xs_flat.zeros_like()?;
+        for (e, expert) in self.experts.iter().enumerate() {
+            if tokens_for[e].is_empty() {
+                continue;
+            }
+            let rows = Tensor::new(tokens_for[e].as_slice(), xs.device())?;
+            let picked = xs_flat.index_select(&rows, 0)?;
+            let out = expert.forward(&picked)?;
+            let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
+                .to_dtype(out.dtype())?
+                .reshape(((), 1))?;
+            ys = ys.index_add(&rows, &out.broadcast_mul(&w)?, 0)?;
+        }
+
+        if let (Some(shared), Some(gate)) = (&self.shared_expert, &self.shared_expert_gate) {
+            let mix = candle_nn::ops::sigmoid(&gate.forward(&xs_flat)?)?;
+            let shared_out = shared.forward(&xs_flat)?.broadcast_mul(&mix)?;
+            ys = (ys + shared_out)?;
+        }
+
+        ys.reshape((b, l, hidden))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Device;
+
+    fn randn(shape: &[usize]) -> Tensor {
+        Tensor::randn(0f32, 0.5f32, shape, &Device::Cpu).unwrap()
+    }
+
+    fn rand_mlp(hidden: usize, inter: usize) -> Qwen3_5MLP {
+        Qwen3_5MLP::from_weights(
+            Linear::new(randn(&[inter, hidden]), None),
+            Linear::new(randn(&[inter, hidden]), None),
+            Linear::new(randn(&[hidden, inter]), None),
+        )
+    }
+
+    /// The batched scatter forward must equal a per-token dense
+    /// reference: route each token independently (host softmax → top-k
+    /// → renorm), run its selected experts one by one, and mix in the
+    /// shared expert through the sigmoid gate. Catches indexing,
+    /// weighting, and renormalisation bugs in the scatter path.
+    #[test]
+    fn scatter_forward_matches_per_token_reference() {
+        let (hidden, inter, n_exp, top_k) = (8, 4, 6, 2);
+
+        let block = Qwen3_5MoeBlock {
+            gate: Linear::new(randn(&[n_exp, hidden]), None),
+            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            shared_expert: Some(rand_mlp(hidden, inter)),
+            shared_expert_gate: Some(Linear::new(randn(&[1, hidden]), None)),
+            num_experts_per_tok: top_k,
+            norm_topk_prob: true,
+        };
+
+        let (b, l) = (2, 3);
+        let xs = randn(&[b, l, hidden]);
+        let got = block.forward(&xs).unwrap();
+        assert_eq!(got.dims(), &[b, l, hidden]);
+
+        let xs_flat = xs.reshape(((), hidden)).unwrap();
+        let logits: Vec<Vec<f32>> = block.gate.forward(&xs_flat).unwrap().to_vec2().unwrap();
+        let got_flat: Vec<Vec<f32>> = got.reshape(((), hidden)).unwrap().to_vec2().unwrap();
+
+        for t in 0..b * l {
+            // Host-side softmax over all experts, then top-k + renorm.
+            let max = logits[t].iter().cloned().fold(f32::MIN, f32::max);
+            let exps: Vec<f32> = logits[t].iter().map(|v| (v - max).exp()).collect();
+            let sum: f32 = exps.iter().sum();
+            let probs: Vec<f32> = exps.iter().map(|e| e / sum).collect();
+            let mut order: Vec<usize> = (0..n_exp).collect();
+            order.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());
+            let selected = &order[..top_k];
+            let denom: f32 = selected.iter().map(|&e| probs[e]).sum();
+
+            let row = xs_flat.narrow(0, t, 1).unwrap();
+            let mut expect = vec![0f32; hidden];
+            for &e in selected {
+                let w = probs[e] / denom;
+                let out: Vec<f32> = block.experts[e]
+                    .forward(&row)
+                    .unwrap()
+                    .flatten_all()
+                    .unwrap()
+                    .to_vec1()
+                    .unwrap();
+                for (acc, o) in expect.iter_mut().zip(out) {
+                    *acc += w * o;
+                }
+            }
+            let gate_v: f32 = block
+                .shared_expert_gate
+                .as_ref()
+                .unwrap()
+                .forward(&row)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1::<f32>()
+                .unwrap()[0];
+            let mix = 1.0 / (1.0 + (-gate_v).exp());
+            let shared: Vec<f32> = block
+                .shared_expert
+                .as_ref()
+                .unwrap()
+                .forward(&row)
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .to_vec1()
+                .unwrap();
+            for (acc, s) in expect.iter_mut().zip(shared) {
+                *acc += mix * s;
+            }
+
+            for (i, (&g, &e)) in got_flat[t].iter().zip(expect.iter()).enumerate() {
+                assert!(
+                    (g - e).abs() < 1e-4,
+                    "token {t} dim {i}: got {g}, expected {e}"
+                );
+            }
+        }
+    }
+
+    /// Without a shared expert (Qwen3-30B-A3B shape) the block is pure
+    /// routed output; without norm_topk_prob the raw global-softmax
+    /// weights apply (they do NOT sum to 1 across the selected k).
+    #[test]
+    fn no_shared_expert_and_no_renorm() {
+        let (hidden, inter, n_exp) = (4, 2, 3);
+        let block = Qwen3_5MoeBlock {
+            gate: Linear::new(randn(&[n_exp, hidden]), None),
+            experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
+            shared_expert: None,
+            shared_expert_gate: None,
+            num_experts_per_tok: 1,
+            norm_topk_prob: false,
+        };
+        let xs = randn(&[1, 1, hidden]);
+        let got: Vec<f32> = block
+            .forward(&xs)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+
+        // Reference: the argmax expert's output scaled by its raw
+        // softmax probability.
+        let logits: Vec<f32> = block
+            .gate
+            .forward(&xs.reshape(((), hidden)).unwrap())
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        let max = logits.iter().cloned().fold(f32::MIN, f32::max);
+        let exps: Vec<f32> = logits.iter().map(|v| (v - max).exp()).collect();
+        let sum: f32 = exps.iter().sum();
+        let best = (0..n_exp)
+            .max_by(|&a, &b| exps[a].partial_cmp(&exps[b]).unwrap())
+            .unwrap();
+        let w = exps[best] / sum;
+        let out: Vec<f32> = block.experts[best]
+            .forward(&xs.reshape(((), hidden)).unwrap())
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
+        for (i, (&g, &o)) in got.iter().zip(out.iter()).enumerate() {
+            assert!(
+                (g - w * o).abs() < 1e-5,
+                "dim {i}: got {g}, expected {}",
+                w * o
+            );
+        }
+    }
+}

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

@@ -208,7 +208,7 @@ mod tests {
             )
             .expect("build ShardedVarBuilder")
         };
-        Qwen3_5Model::load(cfg, &vb).expect("load tiny qwen3_5 model")
+        Qwen3_5Model::load(cfg, &vb, "model.language_model").expect("load tiny qwen3_5 model")
     }
 
     fn forward_tokens(model: &mut Qwen3_5Model, tokens: &[u32], offset: usize) -> Vec<f32> {

crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/config.json

@@ -0,0 +1,56 @@
+{
+  "architectures": [
+    "Qwen3NextForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "decoder_sparse_step": 1,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 64,
+  "initializer_range": 0.02,
+  "intermediate_size": 128,
+  "layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_key_head_dim": 16,
+  "linear_num_key_heads": 2,
+  "linear_num_value_heads": 4,
+  "linear_value_head_dim": 16,
+  "max_position_embeddings": 512,
+  "mlp_only_layers": [],
+  "model_type": "qwen3_next",
+  "moe_intermediate_size": 32,
+  "norm_topk_prob": true,
+  "num_attention_heads": 4,
+  "num_experts": 16,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 2,
+  "output_router_logits": false,
+  "pad_token_id": null,
+  "partial_rotary_factor": 0.25,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "partial_rotary_factor": 0.25,
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "router_aux_loss_coef": 0.001,
+  "shared_expert_intermediate_size": 32,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.9.0",
+  "use_cache": true,
+  "vocab_size": 512
+}

crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.9.0",
+  "use_cache": true
+}

crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/manifest.json

@@ -0,0 +1,114 @@
+{
+  "case": "qwen3_next-tiny",
+  "seed": 92,
+  "token_ids": [
+    3,
+    10,
+    17,
+    24,
+    31,
+    38,
+    45,
+    52,
+    59,
+    66,
+    73,
+    80,
+    87,
+    94,
+    101,
+    108,
+    115,
+    122,
+    129,
+    136,
+    143,
+    150,
+    157,
+    164,
+    171,
+    178,
+    185,
+    192,
+    199,
+    206,
+    213,
+    220,
+    227,
+    234,
+    241,
+    248,
+    255,
+    262,
+    269,
+    276,
+    283,
+    290,
+    297,
+    304,
+    311,
+    318,
+    325,
+    332,
+    339,
+    346,
+    353,
+    360,
+    367,
+    374,
+    381,
+    388,
+    395,
+    402,
+    409,
+    416,
+    423,
+    430,
+    437,
+    444,
+    451,
+    458,
+    465,
+    472,
+    479,
+    486,
+    493,
+    500,
+    507,
+    2,
+    9,
+    16,
+    23,
+    30,
+    37,
+    44,
+    51,
+    58,
+    65,
+    72,
+    79,
+    86,
+    93,
+    100,
+    107,
+    114,
+    121,
+    128,
+    135,
+    142,
+    149,
+    156
+  ],
+  "files": {
+    "logits": {
+      "file": "logits.f32",
+      "shape": [
+        512
+      ]
+    }
+  },
+  "versions": {
+    "transformers": "5.9.0",
+    "torch": "2.9.1+cu128"
+  }
+}

crates/neuron/tests/qwen3_next_parity.rs

@@ -0,0 +1,122 @@
+//! Numerical parity for the qwen3_next path (#92) against the HF
+//! transformers reference, via the tiny self-contained fixture
+//! generated by `script/dump_qwen3_next_tiny.py`.
+//!
+//! The fixture directory carries the WHOLE checkpoint (tiny
+//! random-weight `Qwen3NextForCausalLM`: config.json +
+//! model.safetensors, a few hundred KB) plus the reference
+//! final-position logits, so this test needs no snapshot, no env var,
+//! and runs in CI. It pins: flat-config normalisation, the `model.*`
+//! weight prefix, the fused `in_proj_qkvz`/`in_proj_ba` de-interleave,
+//! hybrid full/linear layer interleaving, and the MoE block (routing,
+//! per-expert SwiGLU, shared expert + sigmoid gate).
+//!
+//! Self-skips (with a loud eprintln) while the fixture has not been
+//! generated yet — regeneration instructions in the script docstring.
+
+use candle_core::{DType, Device, Tensor};
+use serde::Deserialize;
+use std::path::{Path, PathBuf};
+
+#[derive(Deserialize)]
+struct Manifest {
+    token_ids: Vec<u32>,
+    files: std::collections::HashMap<String, FileEntry>,
+}
+
+#[derive(Deserialize)]
+struct FileEntry {
+    file: String,
+    shape: Vec<usize>,
+}
+
+fn fixture_dir() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/numerical/qwen3_next-tiny")
+}
+
+fn read_f32(path: &Path) -> Vec<f32> {
+    let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
+    bytes
+        .chunks_exact(4)
+        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
+        .collect()
+}
+
+#[test]
+fn tiny_qwen3_next_logits_match_hf_reference() {
+    let dir = fixture_dir();
+    let manifest_path = dir.join("manifest.json");
+    if !manifest_path.exists() {
+        eprintln!(
+            "SKIP qwen3_next parity: fixture not generated yet — run \
+             script/dump_qwen3_next_tiny.py --out {} on a host with \
+             torch + transformers>=4.57",
+            dir.display()
+        );
+        return;
+    }
+
+    let manifest: Manifest =
+        serde_json::from_str(&std::fs::read_to_string(&manifest_path).expect("read manifest"))
+            .expect("parse manifest");
+    let logits_entry = &manifest.files["logits"];
+    let reference = read_f32(&dir.join(&logits_entry.file));
+    assert_eq!(
+        reference.len(),
+        logits_entry.shape.iter().product::<usize>()
+    );
+
+    let config_json = std::fs::read_to_string(dir.join("config.json")).expect("read config.json");
+    let cfg = neuron::harness::arch::qwen3_5::Config::from_config_json(&config_json)
+        .expect("normalise qwen3_next config");
+
+    let dev = Device::Cpu;
+    let st_path = dir.join("model.safetensors");
+    // SAFETY: mmap of committed fixture files; nothing mutates them.
+    let vb = unsafe {
+        candle_nn::var_builder::ShardedSafeTensors::var_builder(
+            std::slice::from_ref(&st_path),
+            DType::F32,
+            &dev,
+        )
+        .expect("build ShardedVarBuilder over fixture checkpoint")
+    };
+    let mut model = neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM::new(cfg, vb)
+        .expect("load tiny qwen3_next checkpoint through neuron");
+
+    let input = Tensor::new(manifest.token_ids.as_slice(), &dev)
+        .unwrap()
+        .unsqueeze(0)
+        .unwrap();
+    let logits = model.forward(&input, 0).expect("forward");
+    let got: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
+    assert_eq!(got.len(), reference.len());
+
+    // f32-vs-f32 through an 8-layer doll-house model: agreement should
+    // be tight (the qwen3_5 text fixtures observe max_abs ≈ 0.000,
+    // cosine ≈ 1.0). Thresholds sit far above rounding noise and far
+    // below any real wiring bug (a swapped de-interleave region, a
+    // topk-before-softmax, a missing shared-expert gate all blow past
+    // them instantly).
+    let max_abs = got
+        .iter()
+        .zip(&reference)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0f32, f32::max);
+    let dot: f64 = got
+        .iter()
+        .zip(&reference)
+        .map(|(a, b)| (*a as f64) * (*b as f64))
+        .sum();
+    let na: f64 = got.iter().map(|a| (*a as f64).powi(2)).sum::<f64>().sqrt();
+    let nb: f64 = reference
+        .iter()
+        .map(|b| (*b as f64).powi(2))
+        .sum::<f64>()
+        .sqrt();
+    let cosine = dot / (na * nb);
+
+    eprintln!("qwen3_next parity: max_abs={max_abs:.6} cosine={cosine:.8}");
+    assert!(max_abs < 1e-3, "max abs diff {max_abs} exceeds 1e-3");
+    assert!(cosine > 0.9999, "cosine {cosine} below 0.9999");
+}

script/dump_qwen3_next_tiny.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""Synthesize a tiny qwen3_next parity fixture (#92).
+
+Unlike script/dump_reference.py (which replays a real HF snapshot and
+therefore needs the weights on disk), this builds a TINY random-weight
+`Qwen3NextForCausalLM` from scratch, saves the checkpoint INTO the
+fixture directory, runs the reference forward on fixed token ids, and
+dumps the final-position logits. The whole fixture (weights included)
+is a few hundred KB, so it is committed and the companion Rust test
+(crates/neuron/tests/qwen3_next_parity.rs) runs everywhere — no env
+var, no snapshot, CI included.
+
+What this pins, exactly: neuron's qwen3_next wiring against upstream —
+flat config normalisation, the `model.*` weight prefix, the
+per-key-head-group de-interleave of the fused `in_proj_qkvz` /
+`in_proj_ba` projections, hybrid layer interleaving, and the MoE block
+(softmax→top-k→renorm routing, per-expert SwiGLU, shared expert +
+sigmoid gate). The full-size 80B checkpoint differs only in dimensions.
+
+The config mirrors the real 80B's *shape decisions* at doll-house
+scale: interval-4 hybrid, 8 layers (so two full-attention layers),
+every layer MoE (decoder_sparse_step 1), 16 experts / top-4 + shared
+expert, partial rotary 0.25, 2 KV heads.
+
+Usage (host with torch + transformers>=4.57, e.g. beast):
+  python3 script/dump_qwen3_next_tiny.py \
+      --out crates/neuron/tests/fixtures/numerical/qwen3_next-tiny
+
+Regenerate whenever the transformers reference implementation changes;
+record the transformers version from the manifest in the commit
+message.
+"""
+
+import argparse
+import json
+import os
+import struct
+
+# ---------------------------------------------------------------------------
+# Compat shim (same as dump_reference.py): transformers 5.9 constructs
+# kernels-hub repository objects at import time without the
+# revision/version that kernels 0.15 requires. The hub kernels are
+# never used here; the constructors just must not throw.
+os.environ.setdefault("USE_HUB_KERNELS", "NO")
+try:
+    import kernels.layer.func as _kf
+    import kernels.layer.layer as _kl
+
+    def _patch(cls):
+        orig = cls.__init__
+
+        def patched(self, *a, **kw):
+            if "revision" not in kw and "version" not in kw:
+                kw["revision"] = "main"
+            orig(self, *a, **kw)
+
+        cls.__init__ = patched
+
+    _patch(_kl.LayerRepository)
+    _patch(_kf.FuncRepository)
+except Exception:  # noqa: BLE001 — older/newer kernels may not need it
+    pass
+# ---------------------------------------------------------------------------
+
+import torch  # noqa: E402
+
+# Fixed input: 96 token ids (> 64 so neuron's chunked delta-rule
+# prefill path is exercised), deterministic arithmetic sequence folded
+# into the tiny vocab.
+TOKEN_IDS = [(7 * i + 3) % 512 for i in range(96)]
+SEED = 92
+
+
+def tiny_config():
+    from transformers import Qwen3NextConfig
+
+    return Qwen3NextConfig(
+        vocab_size=512,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=8,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=512,
+        partial_rotary_factor=0.25,
+        rope_theta=10000000,
+        rms_norm_eps=1e-6,
+        tie_word_embeddings=False,
+        full_attention_interval=4,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=16,
+        linear_num_key_heads=2,
+        linear_num_value_heads=4,
+        linear_value_head_dim=16,
+        decoder_sparse_step=1,
+        mlp_only_layers=[],
+        moe_intermediate_size=32,
+        norm_topk_prob=True,
+        num_experts=16,
+        num_experts_per_tok=4,
+        shared_expert_intermediate_size=32,
+    )
+
+
+def write_f32(path, tensor):
+    data = tensor.detach().to(torch.float32).cpu().contiguous().reshape(-1)
+    with open(path, "wb") as f:
+        f.write(struct.pack(f"<{data.numel()}f", *data.tolist()))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", required=True, help="fixture directory to write")
+    args = ap.parse_args()
+
+    import transformers
+    from transformers import Qwen3NextForCausalLM
+
+    os.makedirs(args.out, exist_ok=True)
+
+    torch.manual_seed(SEED)
+    cfg = tiny_config()
+    model = Qwen3NextForCausalLM(cfg).to(torch.float32).eval()
+
+    # Save the checkpoint into the fixture itself (config.json +
+    # model.safetensors) — the Rust test loads neuron's implementation
+    # from exactly these files.
+    model.save_pretrained(args.out, safe_serialization=True)
+
+    ids = torch.tensor([TOKEN_IDS], dtype=torch.long)
+    with torch.no_grad():
+        out = model(input_ids=ids)
+    logits = out.logits[0, -1]  # final position, (vocab,)
+
+    write_f32(f"{args.out}/logits.f32", logits)
+    manifest = {
+        "case": "qwen3_next-tiny",
+        "seed": SEED,
+        "token_ids": TOKEN_IDS,
+        "files": {"logits": {"file": "logits.f32", "shape": [cfg.vocab_size]}},
+        "versions": {
+            "transformers": transformers.__version__,
+            "torch": torch.__version__,
+        },
+    }
+    with open(f"{args.out}/manifest.json", "w") as f:
+        json.dump(manifest, f, indent=2)
+
+    print(f"fixture written to {args.out}")
+    print(f"  transformers {transformers.__version__}, torch {torch.__version__}")
+    print(f"  logits[:4] = {logits[:4].tolist()}")
+
+
+if __name__ == "__main__":
+    main()