test(qwen4_exp): logits parity against the upstream reference (#323)

The architecture was ported from a written spec and every test asserts
it does what we read that spec to mean. This compares it to the
implementation it was ported from: max_abs 0.000002, cosine 1.000000,
same argmax, through all four novel blocks composed.

Unlike the other fixtures here it needs no snapshot and does not
self-skip. It is 44 KB of random weights in the real architecture — 4
layers, 3 linear-attention and 1 QSA, PLE on layer 1 — generated *by*
transformers, so it is checked in and runs on every push. Generating
the weights from the reference rather than by hand means a tensor name
we cannot find is a finding about our loader, not a fixture that
drifted.

Two properties of the fixture are load-bearing and both were learned by
getting them wrong first.

The weights are uniform(-0.8, 0.8) rather than the (-0.1, 0.1) I
started with. Near the origin silu(x) ~ x/2, so the SwiGLU is nearly
symmetric: with small weights, reading the fused expert's gate and up
halves the wrong way round — the one error moe.rs's own comment warns
about — moved the logits by 2.4e-4 and passed. At this scale the same
mutation moves them by 2.7 and takes the argmax with it. A fixture that
cannot express a defect is not evidence about it.

The tolerance is 1e-5 against an observed 2e-6, rather than a
comfortable 1e-3. The same reasoning: at 1e-3 the gate/up swap passed.
A bound wide enough to feel safe is wide enough to hide the thing it
was written for.

The generator refuses to emit a fixture whose QSA selection rests on a
tie. Blocks score exactly 0.0 whenever relu clamps every head, which is
common with random weights, and when the k-th and (k+1)-th scores are
equal it is `topk`'s ordering of equal elements that picks the survivor
— which torch does not specify and whose CPU and CUDA kernels need not
agree. Seeds 0, 1 and 4 tie; seed 2 does not. This is the one place we
genuinely differ from the reference, and it is benign: a tie means the
indexer scored both blocks equally worthless. The fixture avoids the
situation rather than pretending it is resolved.

On failure the test breaks the error down per position before
asserting. Localised means a discrete choice went differently, spread
means the arithmetic drifted; knowing which cost an hour the first
time.

Also carries the `qwen_sparse_attention` layer-type alias, found while
standing the reference up.

Refs #307, #311, #312, #323.

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

Cargo.lock

@@ -3147,6 +3147,7 @@ dependencies = [
  "half",
  "hf-hub",
  "image",
+ "memmap2",
  "minijinja",
  "minijinja-contrib",
  "rayon",

crates/neuron/Cargo.toml

@@ -104,6 +104,12 @@ minijinja-contrib = { version = "2", features = ["pycompat"] }
 # tp `fused_load` module to read per-rank slices of fused QKV tensors
 # without materialising the full tensor on device.
 safetensors = "0.8"
+# The qwen4_exp PLE n-gram table is 102.4 GB of BF16 across 128 shards
+# (#310). It is a pure gather of ~5 KB per token, so it is read straight
+# out of the checkpoint by mmap rather than made resident anywhere —
+# the page cache decides what stays warm, and almost all of the table is
+# cold at any instant.
+memmap2 = "0.9"
 # Vision capability for Qwen3.6 (Stage A of the vision plan in
 # doc/vision-qwen3_6-spec.md). `image` decodes PNG/JPEG/etc from
 # the bytes embedded in `data:image/...;base64,...` content parts;

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

@@ -98,39 +98,59 @@ pub(crate) struct QuantizedExpert {
 /// `quantized tensor must have their last dim divisible by block size
 /// [2560, 640] 256`, having read 131 shards to get there.
 ///
-/// Substituting a 32-block type for that half is what llama.cpp does
-/// with the same constraint, and it costs almost nothing here: `Q4_0`
-/// and `Q4K` are both 4.5 bpw, so the memory arithmetic the placement
-/// decisions rest on is unchanged. The quality is a little worse per
-/// bit — `Q4_0` has one scale per 32 values where `Q4K` has a
-/// two-level hierarchy — which is a real trade, but it applies to one
-/// of the three expert matrices and the alternative is not loading.
+/// **This table is llama.cpp's**, from `tensor_type_fallback` in
+/// `src/llama-quant.cpp`, and the principle behind it is the part
+/// worth stating: **every fallback spends bits rather than saving
+/// them.** A k-quant carries a two-level scale hierarchy that the
+/// legacy 32-block types do not, so it is better than a legacy type at
+/// equal width — dropping `Q4K` (4.5 bpw) to `Q4_0` (also 4.5) would
+/// hold the memory constant and take the quality hit silently.
+/// Upstream declines that trade and pays a bit instead:
 ///
-/// Errors rather than falls back further if the row will not even
-/// divide 32: that is a config we do not understand, and silently
-/// leaving those weights dense would be a memory surprise, not a
-/// mercy.
+/// | requested | fallback | bpw |
+/// |---|---|---|
+/// | `Q2K`, `Q3K` | `Q4_0` | 2.63 / 3.44 -> 4.5 |
+/// | `Q4K` | `Q5_0` | 4.5 -> 5.5 |
+/// | `Q5K` | `Q5_1` | 5.5 -> 6.0 |
+/// | `Q6K`, `Q8K` | `Q8_0` | 6.56 -> 8.5 |
+///
+/// For `qwen4_exp` that is not free: `down_proj` is exactly a third of
+/// the 120.8 B routed parameters, so `Q5_0` rather than `Q4_0` costs
+/// **+5.0 GB** across the expert bank (67.9 -> 73.0 GB). We take it
+/// anyway. The cheaper table was an unmeasured divergence in the one
+/// direction that flatters the constraint we are straining against,
+/// and this matrix is the one a mixed k-quant scheme normally treats
+/// as quality-critical. If the 5 GB turns out to buy nothing, #315 is
+/// where that gets measured and reversed on evidence.
+///
+/// A row that divides 32 either — vanishingly rare, and impossible for
+/// this checkpoint since 640 and 2560 are both multiples of 32 — falls
+/// back to F16 as upstream does, loudly. It is a 3.5x blowup on that
+/// tensor, so it must never be quiet.
 pub(crate) fn quant_for_row(requested: GgmlDType, row: usize) -> candle_core::Result<GgmlDType> {
     if row.is_multiple_of(requested.block_size()) {
         return Ok(requested);
     }
     let fallback = match requested {
-        GgmlDType::Q2K | GgmlDType::Q3K | GgmlDType::Q4K => GgmlDType::Q4_0,
-        GgmlDType::Q5K => GgmlDType::Q5_0,
+        GgmlDType::Q2K | GgmlDType::Q3K => GgmlDType::Q4_0,
+        GgmlDType::Q4K => GgmlDType::Q5_0,
+        GgmlDType::Q5K => GgmlDType::Q5_1,
         GgmlDType::Q6K | GgmlDType::Q8K => GgmlDType::Q8_0,
-        // Already a 32-block type: there is nothing coarser to fall
-        // back to, so this is the error case below.
+        // Already a 32-block type (or not a quant at all): there is no
+        // smaller block to demote to, so the F16 check below decides.
         other => other,
     };
     if row.is_multiple_of(fallback.block_size()) {
         return Ok(fallback);
     }
-    candle_core::bail!(
-        "cannot quantise a row of {row} elements as {requested:?} (block \
-         {}) or {fallback:?} (block {}): the row divides neither",
-        requested.block_size(),
-        fallback.block_size(),
-    )
+    tracing::warn!(
+        ?requested,
+        ?fallback,
+        row,
+        "row divides neither the requested block size nor its 32-wide \
+         fallback; storing this tensor as F16 — 3.5x the bytes"
+    );
+    Ok(GgmlDType::F16)
 }
 
 impl Experts {
@@ -708,16 +728,25 @@ mod tests {
         );
         assert_eq!(
             quant_for_row(GgmlDType::Q4K, 640).unwrap(),
-            GgmlDType::Q4_0,
+            GgmlDType::Q5_0,
             "down reduces over moe_intermediate_size, which is not"
         );
+        // llama.cpp's `tensor_type_fallback` table, which spends bits
+        // rather than saving them: a legacy 32-block type is worse than
+        // a k-quant at equal width, so upstream pays one to avoid the
+        // quality cliff. Q4_0 above would be 4.5 bpw like Q4K and would
+        // save 5.0 GB across the expert bank — that is the trade we
+        // declined, and this is where it is pinned.
+        assert_eq!(quant_for_row(GgmlDType::Q5K, 640).unwrap(), GgmlDType::Q5_1);
+        assert_eq!(quant_for_row(GgmlDType::Q2K, 640).unwrap(), GgmlDType::Q4_0);
+        assert_eq!(quant_for_row(GgmlDType::Q3K, 640).unwrap(), GgmlDType::Q4_0);
         assert_eq!(quant_for_row(GgmlDType::Q6K, 640).unwrap(), GgmlDType::Q8_0);
         assert_eq!(
             quant_for_row(GgmlDType::Q8_0, 640).unwrap(),
             GgmlDType::Q8_0
         );
-        // 100 divides neither 256 nor 32.
-        assert!(quant_for_row(GgmlDType::Q4K, 100).is_err());
+        // 100 divides neither 256 nor 32: F16, as upstream does.
+        assert_eq!(quant_for_row(GgmlDType::Q4K, 100).unwrap(), GgmlDType::F16);
 
         // And the whole path runs at those widths.
         let (n_experts, hidden, inter) = (2usize, 2560usize, 640usize);

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

@@ -197,10 +197,18 @@ impl TextConfig {
             .map(|d| d.eq_ignore_ascii_case("float32") || d.eq_ignore_ascii_case("f32"))
     }
 
+    /// Both spellings upstream uses. The shipped checkpoint's
+    /// `layer_types` says `full_attention` (36 linear, 12 full), but
+    /// `Qwen4ExpTextConfig` *derives* the array with
+    /// `qwen_sparse_attention` when the checkpoint omits it. Matching
+    /// only the first is a silent misclassification: every layer would
+    /// read as linear attention, and a model that is 3:1 linear would
+    /// look almost right until its full-attention tensors were not
+    /// found.
     pub fn is_full_attention(&self, layer: usize) -> bool {
         self.layer_types
             .get(layer)
-            .is_some_and(|t| t == "full_attention")
+            .is_some_and(|t| t == "full_attention" || t == "qwen_sparse_attention")
     }
 
     fn fill_layer_types(&mut self) -> Result<()> {
@@ -474,6 +482,35 @@ mod tests {
     }
 
     /// A one-indexed field that reads 0 is a config written against the
+    /// Upstream has two names for the same layer, and we must read both.
+    ///
+    /// The shipped `config.json` spells the attention layers
+    /// `full_attention`. `Qwen4ExpTextConfig` derives them as
+    /// `qwen_sparse_attention` when the checkpoint omits the array —
+    /// observed directly from `transformers` `main`, which reports
+    /// `['linear_attention', ..., 'qwen_sparse_attention']` for a config
+    /// carrying only `full_attention_interval`.
+    ///
+    /// Reading only the shipped spelling would classify every layer as
+    /// linear on such a config. That is not a load error at the point it
+    /// happens — it is a model built with the wrong 12 layers, which
+    /// fails later and elsewhere.
+    #[test]
+    fn both_upstream_spellings_of_a_full_attention_layer_are_read() {
+        let mut cfg = Config::from_config_json(SHIPPED).unwrap().text_config;
+
+        cfg.layer_types = vec!["linear_attention".into(), "full_attention".into()];
+        assert!(!cfg.is_full_attention(0));
+        assert!(cfg.is_full_attention(1), "the shipped spelling");
+
+        cfg.layer_types = vec!["linear_attention".into(), "qwen_sparse_attention".into()];
+        assert!(!cfg.is_full_attention(0));
+        assert!(
+            cfg.is_full_attention(1),
+            "the derived spelling must classify as full attention too"
+        );
+    }
+
     /// wrong convention, not a layer.
     #[test]
     fn rejects_a_zero_ple_layer_id() {

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

@@ -45,7 +45,7 @@ use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
 use super::config::{Config, TextConfig};
 use super::decoder::DecoderLayer;
 use super::hyper::HyperConnection;
-use super::ple::{NGramHasher, NGramTable, ShardedNGramTable};
+use super::ple::{MmapNGramTable, NGramHasher, NGramTable, ShardedNGramTable};
 use crate::harness::arch::snapshot::{KvCacheSnapshot, PleSnapshot};
 
 /// The hashed n-gram lookup, hoisted to the model because it is a
@@ -129,12 +129,18 @@ impl Qwen4ExpForCausalLM {
     /// keeps the checkpoint's precision, which for this architecture is
     /// useful only on a model small enough to hold: 241.6 GB of BF16
     /// experts fit neither 64 GB of VRAM nor 123 GB of host RAM.
+    /// `safetensors_paths` is the checkpoint's own file set, used only
+    /// to mmap the PLE n-gram table where it already lies (#310). Empty
+    /// falls back to [`ShardedNGramTable`], which puts 102.4 GB
+    /// wherever `vb` points — the oracle, not a residency to serve
+    /// from.
     pub fn load(
         cfg: &Config,
         dtype: DType,
         quant: Option<GgmlDType>,
         device: &Device,
         vb: &ShardedVarBuilder,
+        safetensors_paths: &[std::path::PathBuf],
     ) -> Result<Self> {
         let text = &cfg.text_config;
         // `RotaryEmbedding` reads a qwen3_5 config; the rope block is the
@@ -150,6 +156,25 @@ impl Qwen4ExpForCausalLM {
         let embed_tokens = Embedding::new(embed_weight, text.hidden_size);
 
         let layers_vb = root.pp("layers");
+
+        // Before the layers, not after them. The table is mapped rather
+        // than made resident (#310), so its position in the load order
+        // costs nothing — but everything it can get wrong is a name or
+        // an extent, and those are worth finding in the second before
+        // the expert loop rather than the half-hour after it. Loading
+        // it last meant a checkpoint whose shard names had moved would
+        // report that only once 48 layers of experts had been
+        // quantised, or — since they do not currently fit — never.
+        let ngram = match text.ple_layers().first() {
+            Some(layer) => Some(load_ngram(
+                text,
+                &layers_vb.pp(*layer).pp("ple"),
+                safetensors_paths,
+                device,
+            )?),
+            None => None,
+        };
+
         let mut layers = Vec::with_capacity(text.num_hidden_layers);
         for i in 0..text.num_hidden_layers {
             layers.push(
@@ -175,11 +200,6 @@ impl Qwen4ExpForCausalLM {
             .get((text.vocab_size, text.hidden_size), "weight")
             .context("load 'lm_head.weight'")?;
 
-        let ngram = match text.ple_layers().first() {
-            Some(layer) => Some(load_ngram(text, &layers_vb.pp(*layer).pp("ple"))?),
-            None => None,
-        };
-
         Ok(Self {
             embed_tokens,
             layers,
@@ -356,7 +376,12 @@ impl Qwen4ExpForCausalLM {
     }
 }
 
-fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbedding> {
+fn load_ngram(
+    cfg: &TextConfig,
+    ple_vb: &ShardedVarBuilder,
+    safetensors_paths: &[std::path::PathBuf],
+    device: &Device,
+) -> Result<NGramEmbedding> {
     let vb = ple_vb.pp("ple_embedding");
     let heads = cfg.ngram_heads();
     let head_dim = cfg.ngram_head_dim();
@@ -382,12 +407,28 @@ fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbed
         cfg.split_ngram_parts
     );
 
-    let table = ShardedNGramTable::load(
-        &vb.pp("ngram_embedding"),
-        cfg.split_ngram_parts,
-        rows,
-        head_dim,
-    )?;
+    // The shard names are whatever this VarBuilder would have asked
+    // for, so the prefix is taken from it rather than rebuilt from the
+    // layer index — the two cannot disagree that way, and getting it
+    // wrong is a 102.4 GB read of the wrong thing.
+    let ngram_vb = vb.pp("ngram_embedding");
+    let table: Box<dyn NGramTable> = if safetensors_paths.is_empty() {
+        Box::new(ShardedNGramTable::load(
+            &ngram_vb,
+            cfg.split_ngram_parts,
+            rows,
+            head_dim,
+        )?)
+    } else {
+        Box::new(MmapNGramTable::open(
+            safetensors_paths,
+            &ngram_vb.prefix(),
+            cfg.split_ngram_parts,
+            rows,
+            head_dim,
+            device,
+        )?)
+    };
 
     Ok(NGramEmbedding {
         hasher: NGramHasher::new(
@@ -398,7 +439,7 @@ fn load_ngram(cfg: &TextConfig, ple_vb: &ShardedVarBuilder) -> Result<NGramEmbed
             offsets,
             cfg.eos_token_id,
         )?,
-        table: Box::new(table),
+        table,
         heads,
         head_dim,
         carried: Vec::new(),
@@ -757,7 +798,18 @@ mod tests {
         candle_core::safetensors::save(&tensors(&cfg.text_config), &path).unwrap();
         let vb =
             unsafe { ShardedSafeTensors::var_builder(&[&path], DType::F32, &Device::Cpu).unwrap() };
-        let model = Qwen4ExpForCausalLM::load(&cfg, DType::F32, None, &Device::Cpu, &vb).unwrap();
+        // Passing the path exercises the mmap table (#310), which is
+        // the residency production uses — the whole suite then runs
+        // against the gather that actually serves, not the oracle.
+        let model = Qwen4ExpForCausalLM::load(
+            &cfg,
+            DType::F32,
+            None,
+            &Device::Cpu,
+            &vb,
+            std::slice::from_ref(&path),
+        )
+        .unwrap();
         (dir, model)
     }
 

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

@@ -66,7 +66,7 @@
 //! reimplement. See `doc/qwen4_exp-port-spec.md` §2.
 
 use anyhow::{Context, Result, ensure};
-use candle_core::{D, Module, Tensor};
+use candle_core::{D, DType, Device, Module, Tensor};
 use candle_nn::Linear;
 use candle_nn::var_builder::ShardedVarBuilder;
 
@@ -337,6 +337,216 @@ impl NGramTable for ShardedNGramTable {
     }
 }
 
+/// The table read straight out of the checkpoint by mmap (#310).
+///
+/// The shipping residency. Nothing is made resident anywhere: the 128
+/// shard tensors are mapped where they already lie in the checkpoint's
+/// safetensors files, and a gather copies the requested rows out of
+/// the mapping. The page cache decides what stays warm, which for this
+/// access pattern is the right decision-maker — a token needs sixteen
+/// rows of 160 values, about 5 KB, out of 102.4 GB, so almost the whole
+/// table is cold at any instant and nothing is served by making it
+/// otherwise.
+///
+/// This is what makes the model loadable at all. [`ShardedNGramTable`]
+/// puts those 102.4 GB wherever the `VarBuilder` points, which for a
+/// CUDA load is a 32 GB card; it stays as the oracle this is diffed
+/// against, not as a residency anyone would choose.
+///
+/// Mapping the checkpoint's own tensors in place — rather than building
+/// a separate artifact — is FreeToken's `source_from_safetensors`
+/// trick, and the reason it works here is that the shards are already
+/// exactly what we want: contiguous, equal-stride, row-major, one
+/// extent each. There is nothing to convert and nothing to keep in
+/// sync.
+///
+/// Unlike the sharded gather this needs no sort-and-permute: rows are
+/// copied in the order asked for, so the inverse permutation that has
+/// to be right there cannot be wrong here.
+pub struct MmapNGramTable {
+    /// One mapping per checkpoint file that holds shards. Kept for the
+    /// table's lifetime because every `gather` reads through them.
+    maps: Vec<memmap2::Mmap>,
+    /// Byte position of each shard's row 0: which mapping, and where in
+    /// it. Resolved once at load so a gather is arithmetic.
+    shards: Vec<ShardExtent>,
+    rows_per_shard: usize,
+    head_dim: usize,
+    dtype: DType,
+    row_bytes: usize,
+    device: Device,
+}
+
+/// Where one shard's bytes live.
+struct ShardExtent {
+    map: usize,
+    offset: usize,
+}
+
+impl MmapNGramTable {
+    /// `prefix` is the full tensor-name prefix of the shards, e.g.
+    /// `model.language_model.layers.1.ple.ple_embedding.ngram_embedding`
+    /// — the shards are then `{prefix}.shard_{s}.weight`.
+    ///
+    /// `files` is the checkpoint's safetensors set; shards may be
+    /// spread across any of them (in the real checkpoint, 33 of the
+    /// 131 files carry them), so each is opened and asked what it has.
+    pub fn open(
+        files: &[std::path::PathBuf],
+        prefix: &str,
+        parts: usize,
+        rows: usize,
+        head_dim: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        ensure!(parts > 0, "split_ngram_parts must be > 0");
+        ensure!(
+            rows.is_multiple_of(parts),
+            "{rows} rows do not divide evenly into {parts} shards"
+        );
+        let rows_per_shard = rows / parts;
+
+        let mut maps: Vec<memmap2::Mmap> = Vec::new();
+        let mut found: Vec<Option<ShardExtent>> = (0..parts).map(|_| None).collect();
+        let mut dtype: Option<DType> = None;
+
+        for path in files {
+            let file = std::fs::File::open(path)
+                .with_context(|| format!("open {} for the n-gram table", path.display()))?;
+            // SAFETY: the same contract candle's own safetensors
+            // loading takes — the checkpoint is not modified while we
+            // serve from it. A truncated or rewritten file underneath a
+            // live mapping is a corrupt-model problem, not one this can
+            // defend against.
+            let map = unsafe { memmap2::Mmap::map(&file) }
+                .with_context(|| format!("mmap {}", path.display()))?;
+            let (header_len, meta) =
+                safetensors::SafeTensors::read_metadata(&map).map_err(|e| {
+                    anyhow::anyhow!("read safetensors metadata from {}: {e}", path.display())
+                })?;
+            // `data_offsets` are relative to the start of the data
+            // section, which follows the 8-byte header length and the
+            // header itself.
+            let data_start = 8 + header_len;
+
+            let mut used = false;
+            for (s, slot) in found.iter_mut().enumerate() {
+                if slot.is_some() {
+                    continue;
+                }
+                let name = format!("{prefix}.shard_{s}.weight");
+                let Some(info) = meta.info(&name) else {
+                    continue;
+                };
+                ensure!(
+                    info.shape == [rows_per_shard, head_dim],
+                    "'{name}' is {:?}, expected [{rows_per_shard}, {head_dim}]",
+                    info.shape
+                );
+                let dt = st_dtype(info.dtype)
+                    .with_context(|| format!("'{name}' has a dtype the gather cannot read"))?;
+                match dtype {
+                    None => dtype = Some(dt),
+                    Some(seen) => ensure!(
+                        seen == dt,
+                        "shard dtypes disagree: {seen:?} then {dt:?} at '{name}'"
+                    ),
+                }
+                *slot = Some(ShardExtent {
+                    map: maps.len(),
+                    offset: data_start + info.data_offsets.0,
+                });
+                used = true;
+            }
+            if used {
+                maps.push(map);
+            }
+        }
+
+        let dtype = dtype.context("no n-gram shard was found in any checkpoint file")?;
+        let shards = found
+            .into_iter()
+            .enumerate()
+            .map(|(s, e)| e.with_context(|| format!("'{prefix}.shard_{s}.weight' is in no file")))
+            .collect::<Result<Vec<_>>>()?;
+
+        let row_bytes = head_dim * dtype.size_in_bytes();
+        // Every extent must actually contain its rows: a header that
+        // disagrees with the file length would otherwise surface as a
+        // panic inside a gather, mid-request.
+        for (s, extent) in shards.iter().enumerate() {
+            let end = extent.offset + rows_per_shard * row_bytes;
+            ensure!(
+                end <= maps[extent.map].len(),
+                "shard {s} claims bytes {}..{end} of a {}-byte file",
+                extent.offset,
+                maps[extent.map].len()
+            );
+        }
+
+        tracing::info!(
+            parts,
+            rows,
+            head_dim,
+            ?dtype,
+            files = maps.len(),
+            gib = (rows * row_bytes) as f64 / (1024.0 * 1024.0 * 1024.0),
+            "ple: n-gram table mapped from the checkpoint, not made resident (#310)"
+        );
+
+        Ok(Self {
+            maps,
+            shards,
+            rows_per_shard,
+            head_dim,
+            dtype,
+            row_bytes,
+            device: device.clone(),
+        })
+    }
+}
+
+impl NGramTable for MmapNGramTable {
+    fn gather(&self, rows: &[i64]) -> Result<Tensor> {
+        let total = (self.shards.len() * self.rows_per_shard) as i64;
+        if let Some(bad) = rows.iter().find(|r| **r < 0 || **r >= total) {
+            anyhow::bail!("row {bad} is outside the table's 0..{total}");
+        }
+        let mut out = vec![0u8; rows.len() * self.row_bytes];
+        for (slot, &row) in rows.iter().enumerate() {
+            let shard = row as usize / self.rows_per_shard;
+            let within = row as usize % self.rows_per_shard;
+            let extent = &self.shards[shard];
+            let from = extent.offset + within * self.row_bytes;
+            out[slot * self.row_bytes..(slot + 1) * self.row_bytes]
+                .copy_from_slice(&self.maps[extent.map][from..from + self.row_bytes]);
+        }
+        Ok(Tensor::from_raw_buffer(
+            &out,
+            self.dtype,
+            &[rows.len(), self.head_dim],
+            &self.device,
+        )?)
+    }
+
+    fn head_dim(&self) -> usize {
+        self.head_dim
+    }
+}
+
+/// safetensors' dtype vocabulary is wider than candle's. Only the
+/// float widths a checkpoint would store an embedding table in are
+/// accepted; anything else is a checkpoint we do not understand rather
+/// than one to guess at.
+fn st_dtype(dtype: safetensors::Dtype) -> Result<DType> {
+    Ok(match dtype {
+        safetensors::Dtype::BF16 => DType::BF16,
+        safetensors::Dtype::F16 => DType::F16,
+        safetensors::Dtype::F32 => DType::F32,
+        other => anyhow::bail!("unsupported n-gram table dtype {other:?}"),
+    })
+}
+
 /// The consumption half of PLE: `Qwen4ExpTextPLE` minus the table.
 ///
 /// Given the gathered n-gram embedding `e` and the four-stream residual
@@ -903,6 +1113,88 @@ mod tests {
 
     // ---- the table ----
 
+    /// The mmap table against the device-resident one, row for row.
+    ///
+    /// This is the differential #310 asks for: [`ShardedNGramTable`] is
+    /// the oracle — obviously correct, and unusable at 102.4 GB — and
+    /// [`MmapNGramTable`] is what serves. They must agree on every row
+    /// of every shard, in whatever order they are asked, or the model
+    /// quietly reads somebody else's n-gram.
+    ///
+    /// The row list is chosen to break the things that can break: shard
+    /// boundaries in both directions, the first and last row of the
+    /// table, a repeat, and an order that is not sorted (the oracle
+    /// sorts internally and inverts the permutation; the mmap path does
+    /// not, so a disagreement here is exactly that inverse being wrong
+    /// on one side).
+    #[test]
+    fn the_mapped_table_agrees_with_the_resident_one() {
+        let dev = Device::Cpu;
+        let (parts, rows_per_shard, head_dim) = (4usize, 5usize, 3usize);
+        let prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding";
+
+        // Row r holds [r, r + 0.5, r + 0.25] — distinct per row, and
+        // not a pattern a stride mistake would reproduce.
+        let mut tensors = std::collections::HashMap::new();
+        for s in 0..parts {
+            let mut v = Vec::with_capacity(rows_per_shard * head_dim);
+            for r in 0..rows_per_shard {
+                let global = (s * rows_per_shard + r) as f32;
+                v.extend_from_slice(&[global, global + 0.5, global + 0.25]);
+            }
+            tensors.insert(
+                format!("{prefix}.shard_{s}.weight"),
+                Tensor::from_vec(v, (rows_per_shard, head_dim), &dev).unwrap(),
+            );
+        }
+
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("model.safetensors");
+        candle_core::safetensors::save(&tensors, &path).unwrap();
+
+        let vb = unsafe {
+            candle_nn::var_builder::ShardedSafeTensors::var_builder(&[&path], DType::F32, &dev)
+                .unwrap()
+        };
+        let total = parts * rows_per_shard;
+        let oracle = ShardedNGramTable::load(&vb.pp(prefix), parts, total, head_dim).unwrap();
+        let mapped = MmapNGramTable::open(
+            std::slice::from_ref(&path),
+            prefix,
+            parts,
+            total,
+            head_dim,
+            &dev,
+        )
+        .unwrap();
+
+        assert_eq!(mapped.head_dim(), oracle.head_dim());
+
+        let asked: Vec<i64> = vec![19, 0, 5, 4, 9, 10, 19, 1, 15, 14, 7, 7];
+        let flat = |t: Tensor| -> Vec<f32> { t.flatten_all().unwrap().to_vec1().unwrap() };
+        let from_oracle = flat(oracle.gather(&asked).unwrap());
+        let from_mapped = flat(mapped.gather(&asked).unwrap());
+        assert_eq!(
+            from_mapped, from_oracle,
+            "the mapped gather disagrees with the resident one"
+        );
+        // And that both are right, rather than wrong together.
+        assert_eq!(&from_mapped[..3], &[19.0, 19.5, 19.25]);
+        assert_eq!(&from_mapped[3..6], &[0.0, 0.5, 0.25]);
+
+        // Every row of the table, in order, so no shard is untouched.
+        let all: Vec<i64> = (0..total as i64).collect();
+        assert_eq!(
+            flat(mapped.gather(&all).unwrap()),
+            flat(oracle.gather(&all).unwrap())
+        );
+
+        // The empty gather and the out-of-range guard behave alike.
+        assert_eq!(mapped.gather(&[]).unwrap().dims(), &[0, head_dim]);
+        assert!(mapped.gather(&[total as i64]).is_err());
+        assert!(mapped.gather(&[-1]).is_err());
+    }
+
     /// Rows come back in the order asked for, not the order they were
     /// fetched in. Two shards of three rows, requested out of order and
     /// across the boundary: the shard-major gather returns them sorted,

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

@@ -134,7 +134,13 @@ impl BlockSelector {
     /// 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()?;
+        // f32 for the same reason `pool` is: this is a *selection*
+        // function. See the note there and on the batched path in
+        // `Indexer::positions_from`.
+        let per_head = q
+            .to_dtype(DType::F32)?
+            .matmul(&pooled.to_dtype(DType::F32)?.t()?)?
+            .relu()?;
         per_head.sum(0)? / (self.head_dim as f64).sqrt()
     }
 
@@ -505,15 +511,30 @@ impl Indexer {
         // 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()?;
+        //
+        // In f32, matching upstream, which casts both operands before
+        // the matmul (`torch.matmul(q.float(), block_key_states.float()
+        // .transpose(-1, -2))`). It is not a rounding nicety: these
+        // scores feed a top-k over as many as 65,536 blocks, and the
+        // output is the *set of positions a query may attend*. Two
+        // blocks whose f32 scores differ by less than a bf16 tick —
+        // 0.03125 near 4.0 — compare equal in bf16 and the tie-break
+        // decides instead, so bf16 does not perturb the answer, it
+        // changes which past the model sees. Same argument as `pool`,
+        // which we already got right.
+        let q_rows = q
+            .reshape((batch, self.n_heads * seq_len, self.head_dim))?
+            .to_dtype(DType::F32)?;
+        let keys = blocks
+            .squeeze(1)?
+            .transpose(1, 2)?
+            .contiguous()?
+            .to_dtype(DType::F32)?;
         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 scores = (per_head.sum(1)? / (self.head_dim as f64).sqrt())?.to_vec3::<f32>()?;
 
         let mut out = Vec::with_capacity(batch);
         for row in scores.iter().take(batch) {
@@ -649,6 +670,63 @@ mod tests {
     /// 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
+    /// Scoring is f32 even when the cache is bf16, because it decides
+    /// *which past a query sees*, not merely by how much.
+    ///
+    /// Upstream casts both operands before the matmul
+    /// (`torch.matmul(q.float(), block_key_states.float().T)`); we did
+    /// not, and ran the matmul, the relu and the sum over heads in the
+    /// model dtype.
+    ///
+    /// The construction is the whole argument. Two blocks, four heads,
+    /// per-head contributions of 1.0 and — for one head of one block —
+    /// 1.0078125, which is 1 + 2^-7 and therefore exact in bf16. Their
+    /// sums are 4.0 and 4.0078125, and bf16's tick near 4.0 is 2^-5 =
+    /// 0.03125, so the second rounds back onto the first. The inputs
+    /// survive the dtype; the *comparison* does not. In bf16 the two
+    /// blocks tie and `select`'s index tie-break picks the wrong one,
+    /// silently, for as long as the model is served.
+    #[test]
+    fn scoring_is_f32_even_when_the_cache_is_bf16() {
+        let dev = Device::Cpu;
+        // Budget 4 over 4-wide blocks keeps exactly one block, so the
+        // comparison decides the answer rather than being absorbed by a
+        // budget large enough to take both.
+        let qsa = BlockSelector::new(4, 2, 4, 4).unwrap();
+
+        // Block 0 reads dimension 0 of each head, block 1 reads
+        // dimension 1, so the two scores are independent sums.
+        let pooled = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], (2, 2), &dev)
+            .unwrap()
+            .to_dtype(DType::BF16)
+            .unwrap();
+        let q = Tensor::from_vec(
+            vec![
+                1.0f32, 1.0, // head 0: block 0 <- 1.0, block 1 <- 1.0
+                1.0, 1.0, // head 1
+                1.0, 1.0, // head 2
+                1.0, 1.0078125, // head 3: block 1 gets the extra tick
+            ],
+            (4, 2),
+            &dev,
+        )
+        .unwrap()
+        .to_dtype(DType::BF16)
+        .unwrap();
+
+        let scores: Vec<f32> = qsa.scores(&q, &pooled).unwrap().to_vec1().unwrap();
+        assert!(
+            scores[1] > scores[0],
+            "bf16 scoring lost the distinction between the blocks: {scores:?}"
+        );
+        // And the selection that rests on it picks the right block.
+        assert_eq!(
+            qsa.select(&scores, 8).unwrap(),
+            vec![4, 5, 6, 7],
+            "the higher-scoring block should be the one kept"
+        );
+    }
+
     /// this asserts the value rather than the dtype.
     #[test]
     fn pooling_accumulates_in_f32() {

crates/neuron/src/harness/candle.rs

@@ -2863,6 +2863,7 @@ impl CandleHarness {
                         isq,
                         &device_for_load,
                         &sharded_vb,
+                        &safetensors_paths,
                     )
                     .context("build qwen4_exp model")?;
                     Ok(ModelArch::Qwen4Exp(Box::new(model)))

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

@@ -1059,6 +1059,7 @@ fn load_dense_inner(
                 quant,
                 device,
                 &sharded_vb,
+                safetensors_paths,
             )
             .context("build qwen4_exp model")?;
             Ok(ModelArch::Qwen4Exp(Box::new(model)))

crates/neuron/tests/fixtures/numerical/README.md

@@ -12,6 +12,59 @@ HuggingFace reference" claim to checked-in numbers.
 | `qwen3_5-0.8b-text` | Qwen/Qwen3.5-0.8B | text (>64-token prompt → chunked GDN prefill) | f32 | `text_logits_match_reference` |
 | `qwen3_5-0.8b-vision` | Qwen/Qwen3.5-0.8B | 448×448 synthetic image + prompt | f32 | `vision_tower_and_logits_match_reference` |
 | `qwen3_6-27b-text` | Qwen/Qwen3.6-27B | text | bf16 | manual (see below) |
+| `qwen4_exp-tiny` | random weights, real architecture | text, 12 tokens | f32 | `qwen4_exp_logits_match_reference` |
+
+## `qwen4_exp-tiny` is different, and runs in CI
+
+The others replay a real snapshot and self-skip without
+`NEURON_REF_MODEL_PATH`. This one is 44 KB of random weights in the
+real architecture — 4 layers (3 linear-attention, 1 QSA), PLE on layer
+1, 2 experts — generated *by* `transformers` itself via
+[`script/generate_qwen4_exp_fixture.py`](../../../../../script/generate_qwen4_exp_fixture.py).
+So it is checked in, needs no weights, and runs on every push.
+
+That matters because the whole `qwen4_exp` architecture was ported from
+a written spec (#308) and tested against our own reading of it. This is
+the only thing in the tree that compares it to the implementation it
+was ported from (#323).
+
+Two properties of the fixture are load-bearing, and both were learned
+by getting them wrong:
+
+- **Weights are `uniform(-0.8, 0.8)`, not `uniform(-0.1, 0.1)`.** Near
+  the origin `silu(x) ~ x/2`, so the SwiGLU is nearly symmetric and
+  reading the fused expert's gate and up halves the wrong way round
+  moved the logits by 2.4e-4 — inside any tolerance worth writing. At
+  the larger scale the same mutation moves them by 2.7.
+- **The generator refuses to emit a fixture containing a QSA selection
+  tie.** A block scores exactly 0.0 whenever relu clamps every head, so
+  ties are common with random weights; when the k-th and (k+1)-th
+  scores are equal, which block is dropped is decided by `topk`'s
+  ordering of equal elements, which torch does not specify and whose
+  CPU and CUDA kernels need not agree. Seeds 0, 1 and 4 produce ties;
+  seed 2 (the default) does not. See "the tie-break" below.
+
+Regenerate with the reference from `transformers` **main** — the 5.9.0
+release does not carry `qwen4_exp`:
+
+```sh
+python -m venv --system-site-packages .venv
+.venv/bin/pip install --no-deps "transformers @ git+https://github.com/huggingface/transformers@main"
+.venv/bin/pip install "tokenizers>=0.23.1,<0.24" "safetensors>=0.8.0"
+.venv/bin/python script/generate_qwen4_exp_fixture.py \
+    --config crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny/config.json \
+    --out crates/neuron/tests/fixtures/numerical/qwen4_exp-tiny
+```
+
+### The tie-break
+
+Where the reference and this port genuinely differ: at an exact tie in
+the QSA block scores, `topk` keeps one block and we keep another. Ours
+is deterministic (lower block index wins); torch's is unspecified. A
+tie means every head's dot product was clamped to zero for both blocks
+— the indexer scored them equally worthless — so either choice is
+defensible, and there is no stable target to match. The fixture avoids
+the situation rather than pretending it is resolved.
 
 ## Running the comparison
 

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

@@ -0,0 +1,50 @@
+{
+  "model_type": "qwen4_exp",
+  "tie_word_embeddings": false,
+  "text_config": {
+    "vocab_size": 32,
+    "hidden_size": 8,
+    "num_hidden_layers": 4,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "head_dim": 4,
+    "max_position_embeddings": 64,
+    "rms_norm_eps": 1e-06,
+    "hidden_act": "silu",
+    "eos_token_id": 3,
+    "full_attention_interval": 4,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "partial_rotary_factor": 0.5,
+      "rope_type": "default"
+    },
+    "hc_count": 2,
+    "hc_lowrank": 4,
+    "ple_layer_ids": [
+      2
+    ],
+    "ple_embed_dim": 8,
+    "ple_conv_kernel_size": 2,
+    "ngram_size": 3,
+    "heads_per_ngram": 2,
+    "ngram_vocab_size_base": 10,
+    "make_ngram_vocab_size_divisible_by": 4,
+    "split_ngram_parts": 2,
+    "indexer_n_heads": 2,
+    "indexer_kv_heads": 1,
+    "indexer_head_dim": 4,
+    "indexer_budget": 8,
+    "indexer_compress_ratio": 2,
+    "linear_num_key_heads": 1,
+    "linear_key_head_dim": 4,
+    "linear_num_value_heads": 2,
+    "linear_value_head_dim": 4,
+    "linear_conv_kernel_dim": 2,
+    "mamba_ssm_dtype": "float32",
+    "output_gate_type": "sigmoid",
+    "num_experts": 2,
+    "num_experts_per_tok": 1,
+    "moe_intermediate_size": 4,
+    "shared_expert_intermediate_size": 4
+  }
+}

crates/neuron/tests/numerical_reference.rs

@@ -339,3 +339,94 @@ fn vision_tower_and_logits_match_reference() {
     );
     assert_eq!(cr.argmax_ours, cr.argmax_ref, "chunked argmax vs reference");
 }
+
+/// `qwen4_exp` logits against the upstream reference (#323).
+///
+/// Unlike the fixtures above this one needs no `NEURON_REF_MODEL_PATH`
+/// and no snapshot: the model is 44 KB of random weights in the real
+/// architecture, generated *by* `transformers` itself
+/// (`script/generate_qwen4_exp_fixture.py`), so it is checked in and
+/// this runs on every push. A tensor name we cannot find is therefore a
+/// finding about our loader rather than a fixture that drifted.
+///
+/// f32 on both sides, for the reason in this file's header: cross-dtype
+/// comparison measures bf16 rounding, not implementation error.
+///
+/// What it covers that the unit tests cannot: the four novel blocks
+/// composed. Hyper-connections carrying four residual streams with no
+/// layernorm of the usual kind, the PLE n-gram gather on layer 1, QSA
+/// on layer 3, and the final mixer collapse with no `model.norm` before
+/// `lm_head` — every one of which was written from prose and tested
+/// against our own reading of it.
+#[test]
+fn qwen4_exp_logits_match_reference() {
+    use candle_nn::var_builder::ShardedSafeTensors;
+    use neuron::harness::arch::qwen4_exp::config::Config as Qwen4ExpConfig;
+    use neuron::harness::arch::qwen4_exp::model::Qwen4ExpForCausalLM;
+
+    let fixture = fixture_root().join("qwen4_exp-tiny");
+    let weights = fixture.join("model.safetensors");
+    let device = Device::Cpu;
+
+    let cfg = Qwen4ExpConfig::from_config_json(
+        &std::fs::read_to_string(fixture.join("config.json")).unwrap(),
+    )
+    .expect("parse the fixture config");
+
+    let vb = unsafe { ShardedSafeTensors::var_builder(&[&weights], DType::F32, &device).unwrap() };
+    let mut model = Qwen4ExpForCausalLM::load(
+        &cfg,
+        DType::F32,
+        None,
+        &device,
+        &vb,
+        std::slice::from_ref(&weights),
+    )
+    .expect("load the fixture through the production loader");
+
+    let expected = candle_core::safetensors::load(fixture.join("expected.safetensors"), &device)
+        .expect("read the reference output");
+    let ids = expected["input_ids"].to_dtype(DType::U32).unwrap();
+    let reference: Vec<f32> = expected["logits"].flatten_all().unwrap().to_vec1().unwrap();
+
+    let logits = model.forward(&ids, 0).expect("forward");
+    let ours: Vec<f32> = logits
+        .to_dtype(DType::F32)
+        .unwrap()
+        .flatten_all()
+        .unwrap()
+        .to_vec1()
+        .unwrap();
+
+    let c = compare(&ours, &reference);
+    eprintln!(
+        "qwen4_exp: max_abs={:.6} cosine={:.6} argmax ours={} ref={}",
+        c.max_abs, c.cosine, c.argmax_ours, c.argmax_ref
+    );
+    // On failure, break the error down by position before asserting.
+    // Whether the divergence is confined to one position or spread
+    // across all of them is the first thing worth knowing: localised
+    // means a discrete choice went differently (an expert, a QSA
+    // block), spread means the arithmetic drifted. Finding that out
+    // cost an hour the first time.
+    if c.max_abs >= 1e-5 {
+        let v = cfg.text_config.vocab_size;
+        for t in 0..(ours.len() / v) {
+            let pc = compare(&ours[t * v..(t + 1) * v], &reference[t * v..(t + 1) * v]);
+            eprintln!(
+                "  pos {t:2}: max_abs={:.6} cosine={:.6}",
+                pc.max_abs, pc.cosine
+            );
+        }
+    }
+    assert_eq!(c.argmax_ours, c.argmax_ref, "argmax token diverged");
+    // f32 against f32 is *exact* here — observed max_abs 0.000000 — so
+    // the bound is set just off the floor rather than at a comfortable
+    // distance. That is deliberate: at 1e-3 this test still passed with
+    // the fused expert's gate and up halves read the wrong way round
+    // (max_abs 2.4e-4), which is the single error the loader's own
+    // comment warns about. A tolerance wide enough to be safe is wide
+    // enough to hide the defect it was written for.
+    assert!(c.cosine > 0.999_999, "cosine {:.6} too low", c.cosine);
+    assert!(c.max_abs < 1e-5, "max abs diff {:.6} too high", c.max_abs);
+}

script/generate_qwen4_exp_fixture.py

@@ -0,0 +1,116 @@
+"""Generate a qwen4_exp parity fixture from the upstream reference (#323).
+
+Writes, into --out:
+  config.json          the tiny config, as both implementations parse it
+  model.safetensors    the reference's own weights, with the PLE n-gram
+                       table split into `split_ngram_parts` shards on
+                       dim 0 the way the released checkpoint ships it
+  expected.safetensors input_ids and the reference's logits
+
+The weights are the reference's, so a name we cannot find is a finding
+rather than a fixture bug. Norm weights are perturbed away from their
+zero init on purpose: at exactly zero, `(1 + w)` and `w` differ, but
+`(1 + w)` and `1` do not, and the whole point is to tell those apart.
+"""
+
+import argparse, json, os
+import torch
+from safetensors.torch import save_file
+from transformers.models.qwen4_exp import Qwen4ExpConfig, Qwen4ExpForConditionalGeneration
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--config", required=True, help="the tiny config.json both implementations parse")
+ap.add_argument("--out", required=True, help="fixture directory to write")
+# Seed 2 is tie-free; 0, 1 and 4 are not. See the tie check below.
+ap.add_argument("--seed", type=int, default=2)
+ap.add_argument("--tokens", type=int, default=12)
+args = ap.parse_args()
+
+torch.manual_seed(args.seed)
+torch.use_deterministic_algorithms(True)
+
+raw = json.load(open(args.config))
+cfg = Qwen4ExpConfig(**raw)
+tc = cfg.text_config
+
+model = Qwen4ExpForConditionalGeneration(cfg).eval()
+
+# Every parameter gets a non-degenerate value. Zeros hide a transposed
+# load. But *small* values hide more than that: silu(x) ~ x/2 near the
+# origin, so with weights in +-0.1 the SwiGLU is close enough to
+# symmetric that reading gate and up the wrong way round changes the
+# logits by 2e-4 — inside any tolerance you would write. The scale here
+# is chosen to put activations where silu is actually curved, which is
+# what makes the ordering observable, while staying below the point
+# where the MoE router saturates and top-k stops discriminating.
+with torch.no_grad():
+    for name, p in model.named_parameters():
+        if "visual" in name or "vision" in name:
+            continue
+        p.copy_(torch.empty_like(p).uniform_(-0.8, 0.8))
+
+ids = torch.randint(0, tc.vocab_size, (1, args.tokens), dtype=torch.long)
+# Keep eos out of the prompt: it starts a new n-gram segment, and the
+# first fixture should not conflate segment handling with everything
+# else. A separate fixture exercises it deliberately.
+ids[ids == tc.eos_token_id] = (tc.eos_token_id + 1) % tc.vocab_size
+
+# The QSA selection must not rest on a tie.
+#
+# A block scores exactly 0.0 whenever relu clamps every head's dot
+# product, which is common enough that a random fixture hits it. When
+# the k-th and (k+1)-th scores are equal, which block gets dropped is
+# decided by `topk`'s ordering of equal elements — and torch does not
+# specify that, nor does its CPU kernel agree with its CUDA one. A
+# fixture containing such a tie measures that unspecified detail rather
+# than our arithmetic, so we refuse to emit one.
+orig_topk = torch.Tensor.topk
+ties = []
+def _topk(self, k, *a, **kw):
+    if self.dim() == 1 and self.numel() <= 64 and k < self.numel():
+        ordered = sorted(self.detach().tolist(), reverse=True)
+        if ordered[k - 1] == ordered[k]:
+            ties.append((self.detach().tolist(), k))
+    return orig_topk(self, k, *a, **kw)
+torch.Tensor.topk = _topk
+with torch.no_grad():
+    out = model(input_ids=ids, use_cache=False)
+torch.Tensor.topk = orig_topk
+
+if ties:
+    raise SystemExit(
+        f"seed {args.seed} produces {len(ties)} QSA selection tie(s), e.g. "
+        f"scores={[round(v, 6) for v in ties[0][0]]} k={ties[0][1]}. "
+        "Pick another --seed: the fixture must isolate the arithmetic."
+    )
+
+os.makedirs(args.out, exist_ok=True)
+
+sd = {k: v for k, v in model.state_dict().items() if "visual" not in k and "vision" not in k}
+
+# Split the runtime n-gram embedding back into the checkpoint's shard
+# layout: dim 0, ascending, per upstream's own note that "the
+# checkpoints for it are sharded on dim0" and loading concatenates them.
+parts = tc.split_ngram_parts
+for key in [k for k in sd if k.endswith("ngram_embedding.weight")]:
+    table = sd.pop(key)
+    rows = table.shape[0]
+    assert rows % parts == 0, f"{rows} rows do not divide into {parts} shards"
+    per = rows // parts
+    base = key[: -len(".weight")]
+    for s in range(parts):
+        sd[f"{base}.shard_{s}.weight"] = table[s * per : (s + 1) * per].clone()
+
+save_file({k: v.contiguous() for k, v in sd.items()}, os.path.join(args.out, "model.safetensors"))
+save_file(
+    {"input_ids": ids.to(torch.int64), "logits": out.logits.to(torch.float32).contiguous()},
+    os.path.join(args.out, "expected.safetensors"),
+)
+with open(os.path.join(args.out, "config.json"), "w") as f:
+    json.dump(raw, f, indent=2)
+
+print(f"tokens      {args.tokens}")
+print(f"tensors     {len(sd)}")
+print(f"logits      {tuple(out.logits.shape)}")
+print(f"logits absmax {out.logits.abs().max().item():.6f}")
+print(f"finite      {bool(torch.isfinite(out.logits).all())}")