fix(examples): anchor logit_trace chunk 0 with the BOS token

docs/benchmarks.md

@@ -573,6 +573,8 @@ python3 scripts/compare_logit_traces.py chain.tsv block.tsv --decided 2
 python3 scripts/compare_logit_traces.py chain.tsv verify.tsv --decided 2
 ```
 
+Every chunk starts at the tokenizer's BOS, chunk 0 included: without `PREFILL` the BOS is prepended to the traced forward, and with it the BOS leads the context pass (alone, for a chunk with no history), so the traced forward keeps the chunk's width. The trace records this in a `# bos_anchor` metadata line. Traces written before #1785 lack that line and left chunk 0 unanchored, which on a model trained behind a document separator collapses the first chunk (Laguna XS 2.1 read perplexity about 1000 there against about 18 in an anchored chunk), so do not compare an older trace with a newer one at chunk 0; regenerate both arms with the same build of the harness. `scripts/laguna_oracle_trace.py` follows the same rule.
+
 The metric to gate on is **disagreement on decided positions**: the fraction
 of positions where the reference's own top two were more than a stated gap
 apart and the candidate still picked something else. A position the reference

docs/supported-models.md

@@ -53,7 +53,7 @@ Implemented model families include:
 - StarCoder2, StableLM, SmolLM3, Baichuan, MiniCPM, MiniCPM3, MiniMax,
   Ministral3, Nemotron, Nemotron-NAS, Step 3.5, MiMo
 - Mellum / Mellum 2 (`mellum`, JetBrains code model: sliding/full hybrid attention driven by `layer_types`, with QK-RMSNorm and a sparse softmax-routed MoE (`norm_topk_prob`) in every layer. Full-attention layers use YaRN-scaled RoPE plus a standard `KVCache`; sliding-attention layers use default RoPE plus a `RotatingKVCache` windowed to `sliding_window`. Supports tied and untied LM heads. Validated against `JetBrains/Mellum2-12B-A2.5B-Base`.)
-- Laguna XS 2.1 / XS.2 / S 2.1 (`laguna`, Poolside code models: a pre-norm decoder with QK-RMSNorm, a hybrid of `full_attention` layers (YaRN over the first `partial_rotary_factor` share of each head, explicit `attention_factor`, dense `KVCache`) and `sliding_attention` layers (plain RoPE at base 10000, `RotatingKVCache` windowed to `sliding_window` 512), per-layer query-head counts from `num_attention_heads_per_layer`, a per-head softplus attention output gate (`g_proj`, `gating: "per-head"`), an optional per-head attention sink on sliding layers, a dense SwiGLU MLP on layer 0, and elsewhere a sigmoid-routed (or `sqrtsoftplus`) top-k MoE whose expert selection adds `e_score_correction_bias` while the weights stay bias-free, renormalized and scaled by `moe_routed_scaling_factor`, plus an always-on shared expert. The published NVFP4 checkpoints keep the experts and the shared expert in the `compressed-tensors` `nvfp4-pack-quantized` layout (`weight_packed` / `weight_scale` / `weight_global_scale`); the loader reinterprets the packed E2M1 bytes as MLX native NVFP4 words, keeps every E4M3 block scale byte for byte, and applies `1 / weight_global_scale` per expert after the routed matmul (the `SwitchLinear` `global_scale` sidecar) rather than folding it into the 3-mantissa-bit block scales. Pre-stacked MLX conversions (`switch_mlp.*`, a fused `gate_up_proj` is split) and per-expert bf16 exports load through the same path. The chat template emits the tokenizer's own BOS spelling (`〈|EOS|〉`, id 2, which is also an EOS) first, so the tokenize sites suppress `add_special_tokens` for a prompt that starts with the tokenizer's BOS string (`MlxcelTokenizer::prompt_carries_bos`); generation stops on ids 2 and 24 (`</assistant>`, a plain added token), and `<think>` / `</think>` are the reasoning markers. The YaRN block's `attention_factor` is used when the config carries one and derived from `factor` otherwise, which is what `transformers`' `_compute_yarn_parameters` does and therefore what the checkpoint's own `modeling_laguna.py` inherits; mlx-lm's port drops the key and always derives, so `Laguna-XS.2-4bit`, which declares `attention_factor: 1.0` where the derived value would be 1.3466, rotates differently under the two runtimes. XS 2.1 declares the derived value and agrees either way. Not wired for tensor parallel. Validated against `poolside/Laguna-XS-2.1-NVFP4` (native compressed-tensors NVFP4, 234 planes transcoded at load), `mlx-community/Laguna-XS.2-4bit` (pre-stacked affine, per-layer 8-bit router planes reconciled from the tensor shapes) and `poolside/Laguna-S-2.1-NVFP4` (48 layers, 256 experts at top-10). S 2.1 mixes expert formats in one checkpoint: layers 1 to 39 carry NVFP4 routed experts (117 planes), layers 40 to 47 keep their routed experts in bf16, and the shared experts are bf16 throughout, so a load exercises both the NVFP4 and the per-expert bf16 paths. It generates fluent Python with finite logits. Its load-time resident figure (49.4 GiB) counts only the transcoded planes, since the bf16 weights realize on the first forward, so plan for about 93 GiB and a 128 GB host. The external reference is the checkpoint's own `modeling_laguna.py`, run out of band in float32 by `scripts/laguna_oracle_trace.py`, which decodes the NVFP4 planes exactly and writes the `examples/logit_trace` format. Against it on XS 2.1, mlxcel (bf16) has no top-1 disagreement at a position the reference had decided (top-two gap of 2 or more): 0 of 75 at width 1 and 0 of 296 at width 8, both behind 512 tokens of context, and 0 of 107 at width 256 in a window that starts with the BOS token. A window without the `〈|EOS|〉` BOS degrades in the reference as well (perplexity about 950), and the only decided disagreement at width 256, 1 of 93, is in such a window. The reference implements neither `sqrtsoftplus` routing nor the sliding-layer sinks (it allocates `sink` but never reads it), and no published checkpoint enables either, so those two paths are covered by synthetic-config tests only.)
+- Laguna XS 2.1 / XS.2 / S 2.1 (`laguna`, Poolside code models: a pre-norm decoder with QK-RMSNorm, a hybrid of `full_attention` layers (YaRN over the first `partial_rotary_factor` share of each head, explicit `attention_factor`, dense `KVCache`) and `sliding_attention` layers (plain RoPE at base 10000, `RotatingKVCache` windowed to `sliding_window` 512), per-layer query-head counts from `num_attention_heads_per_layer`, a per-head softplus attention output gate (`g_proj`, `gating: "per-head"`), an optional per-head attention sink on sliding layers, a dense SwiGLU MLP on layer 0, and elsewhere a sigmoid-routed (or `sqrtsoftplus`) top-k MoE whose expert selection adds `e_score_correction_bias` while the weights stay bias-free, renormalized and scaled by `moe_routed_scaling_factor`, plus an always-on shared expert. The published NVFP4 checkpoints keep the experts and the shared expert in the `compressed-tensors` `nvfp4-pack-quantized` layout (`weight_packed` / `weight_scale` / `weight_global_scale`); the loader reinterprets the packed E2M1 bytes as MLX native NVFP4 words, keeps every E4M3 block scale byte for byte, and applies `1 / weight_global_scale` per expert after the routed matmul (the `SwitchLinear` `global_scale` sidecar) rather than folding it into the 3-mantissa-bit block scales. Pre-stacked MLX conversions (`switch_mlp.*`, a fused `gate_up_proj` is split) and per-expert bf16 exports load through the same path. The chat template emits the tokenizer's own BOS spelling (`〈|EOS|〉`, id 2, which is also an EOS) first, so the tokenize sites suppress `add_special_tokens` for a prompt that starts with the tokenizer's BOS string (`MlxcelTokenizer::prompt_carries_bos`); generation stops on ids 2 and 24 (`</assistant>`, a plain added token), and `<think>` / `</think>` are the reasoning markers. The YaRN block's `attention_factor` is used when the config carries one and derived from `factor` otherwise, which is what `transformers`' `_compute_yarn_parameters` does and therefore what the checkpoint's own `modeling_laguna.py` inherits; mlx-lm's port drops the key and always derives, so `Laguna-XS.2-4bit`, which declares `attention_factor: 1.0` where the derived value would be 1.3466, rotates differently under the two runtimes. XS 2.1 declares the derived value and agrees either way. Not wired for tensor parallel. Validated against `poolside/Laguna-XS-2.1-NVFP4` (native compressed-tensors NVFP4, 234 planes transcoded at load), `mlx-community/Laguna-XS.2-4bit` (pre-stacked affine, per-layer 8-bit router planes reconciled from the tensor shapes) and `poolside/Laguna-S-2.1-NVFP4` (48 layers, 256 experts at top-10). S 2.1 mixes expert formats in one checkpoint: layers 1 to 39 carry NVFP4 routed experts (117 planes), layers 40 to 47 keep their routed experts in bf16, and the shared experts are bf16 throughout, so a load exercises both the NVFP4 and the per-expert bf16 paths. It generates fluent Python with finite logits. Its load-time resident figure (49.4 GiB) counts only the transcoded planes, since the bf16 weights realize on the first forward, so plan for about 93 GiB and a 128 GB host. The external reference is the checkpoint's own `modeling_laguna.py`, run out of band in float32 by `scripts/laguna_oracle_trace.py`, which decodes the NVFP4 planes exactly and writes the `examples/logit_trace` format. Against it on XS 2.1, mlxcel (bf16) has no top-1 disagreement at a position the reference had decided (top-two gap of 2 or more): 0 of 75 at width 1 and 0 of 296 at width 8, both behind 512 tokens of context, and 0 of 107 at width 256 in a window that starts with the BOS token. Those traces left chunk 0 without the `〈|EOS|〉` BOS, which degrades the reference as well (perplexity about 950 against 17.8 anchored), and the only decided disagreement at width 256, 1 of 93, was in that window. Since #1785 both tools anchor every window, chunk 0 included. Re-measured that way on a GB10 (reference in float32 on CUDA, `docs/architecture.md` as the corpus): 0 of 250 decided positions disagree at width 256 and 0 of 317 at width 8 behind 512 tokens, and chunk 0 at width 256 reads perplexity 19.9 in both the reference and mlxcel, against 364 and 378 without the anchor. The reference implements neither `sqrtsoftplus` routing nor the sliding-layer sinks (it allocates `sink` but never reads it), and no published checkpoint enables either, so those two paths are covered by synthetic-config tests only.)
 - Inkling (`inkling_mm_model` and `inkling`): the text backbone combines hybrid sliding/global NoPE attention, learned banded relative-position bias and log-position temperature, four f32 short-convolution states per layer, dense-first SwiGLU, logsigmoid-softmax routed and shared experts, muP-scaled logits, and padded-vocabulary trimming. Image input is supported through both the CLI `--image` path and OpenAI-compatible server `image_url` content blocks when the checkpoint carries `vision_config` and `model.visual.*` weights. The HMLP tower maps row-major 40x40 RGB tiles duplicated across two temporal slots directly to text-width soft tokens; exact-width images intentionally receive a trailing all-padding tile column, and each rendered image placeholder expands to the number of tiles preprocessing actually produced before ordered feature scatter. Audio input is supported for checkpoints that also carry `audio_config` and `model.audio.*` weights through CLI `--audio` and server `input_audio`, alone or with still images. The bounded host path decodes, downmixes, and resamples WAV input to 16 kHz, produces reference-compatible dMel IDs, scatters image rows first and audio rows second, and keeps every multimodal request on classic prepared-embedding prefill. Video input is supported through CLI `--video` and server `video_url`: sampled frames default to 2 fps, odd sequences repeat their final frame, and at most 16 evenly spaced adjacent pairs are allocated request-wide without crossing clip boundaries. Each clip restarts its timestamp origin, and timestamps use actual probed source-frame time after sampling. Before decoding, the request is capped at eight clips, 32 unique selected frames, and 512 MiB of decoded pixel storage; only selected frames are decoded and retained. The second frame replaces temporal slot 1 only in the suffix of HMLP rows belonging to video tiles. Companion still images remain ahead of complete timestamped Inkling content parts inside the current user turn and retain duplicated temporal planes. Native chained MTP is available for B = 1 with the original `model.mtp.layers.*` tensors for standalone text targets and text-only requests on HMLP-wrapped checkpoints. Speculative decode uses the wrapper's `text` backbone, while image-, audio-, and video-bearing requests stay on the classic prepared-embedding prefill path before continuing through the shared text target. MTP shares the target embedding, final norm, and LM head while restoring and replaying exact KV plus four-convolution state after verification. The loader accepts original bf16/f32, native ModelOpt NVFP4, and pre-converted affine MLX text weights, plus dense or affine-quantized HMLP projections and dense or affine-quantized audio embeddings. Fused kernels, padded batching, and image-feature caching remain separate work. Validation covers f64-referenced dMel extraction and quantization, summed-channel audio embeddings, deterministic HMLP planning/folding, tiling/normalization, placeholder cardinality, image-then-audio scatter, current-user audio insertion with history, clip-local adjacent-frame selection, temporal suffix splicing, pre-decode resource caps, timestamp drift, raw-weight detection and normalization, the synthetic text graph, caches, sanitizer, chat markers, deterministic tiny-model greedy MTP parity, exact snapshot/replay, wrapper dispatch, and prepared-prefill preservation. The public 153.5 GB target, roughly 170.7 GB native NVFP4 checkpoint, and 4.5 GB MTP shard were not available on the validation host, so real image/audio/video answer quality, MTP throughput, and acceptance length remain unmeasured.)
 - TeleChat3 (`telechat3`, TeleAI: structurally a stock Llama decoder (RMSNorm, GQA, SwiGLU, untied `lm_head`), so what matters is the one thing that is not stock. `Telechat3-36B-Thinking` ships `rope_scaling.rope_type = "telechat3-yarn"` with `factor: 4.0`, `original_max_position_embeddings: 8192`, `beta_fast: 32` and `beta_slow: 1`. Upstream routes that into the same `YarnRoPE` it uses for `"yarn"` and `"deepseek_yarn"`: the vendor prefix names the checkpoint that produced the config, not a different algorithm, so the shared YaRN reader accepts all three spellings. This is why the family does not reuse the `llama3` decoder. `llama3::ModelArgs` reads its `rope_scaling` field since #1355, but only for the `default`, `linear` and `llama3` schemes, and it carries no `beta_fast` / `beta_slow`, so routing TeleChat3 through that path would warn once and then rotate at the unscaled base while every tensor kept the right shape; YaRN and default RoPE agree closely at small offsets and only diverge past `original_max_position_embeddings`, so no short-prompt output could expose it. TeleChat3 also keeps `rope_theta` at the top level of `config.json` while the shared reader looks for it inside the scaling block (defaulting to 500000, not this family's 1000000), so the base is injected into the block before the table is built; a table built at the wrong base is finite, correctly shaped and wrong. `attention_bias` puts bias terms on q/k/v/o and is plumbed because the family declares it, but the published 36B checkpoint sets it **false**, so it is not what distinguishes this family in practice. Like every other text family here it satisfies the maskless-prefill contract by routing a multi-token prefill with no caller mask through `causal_attention`. Not wired for tensor parallel, for the same reason it does not reuse the llama3 decoder: the TP runtime parses `config.json` straight into `llama3::ModelArgs`, which has no YaRN branch, so a sharded TeleChat3 would drop the scaling. Validation checkpoint: `mlx-community/Telechat3-36B-Thinking-4bit`.)
 - Apertus (`apertus`, Swiss AI: Llama-style dense transformer with an xIELU activation MLP (no gate), QK-norm, llama3 RoPE scaling, and untied embeddings)

examples/logit_trace.rs

@@ -89,6 +89,12 @@
 // traced, and the remaining rows are traced from a forward at the chunk's own
 // width. That is the shape a real verify runs at, long context and narrow
 // forward, which a bare small chunk does not reproduce.
+//
+// Every chunk, chunk 0 included, is anchored by the tokenizer's BOS: without
+// `PREFILL` it is prepended to the traced forward, with `PREFILL` it leads the
+// context pass (alone, for a chunk with no history). The `# bos_anchor`
+// metadata line records it; traces without that line left chunk 0 unanchored
+// and are not comparable with newer ones at chunk 0.
 
 #[path = "logit_trace/gemma4.rs"]
 mod gemma4;
@@ -229,7 +235,12 @@ fn main() -> Result<()> {
     // Same BOS anchoring as `examples/perplexity`, and for the same reason
     // (#686): the Gemma family collapses on a BOS-less window, so a
     // mid-corpus chunk scored raw measures the missing anchor rather than
-    // the change under test.
+    // the change under test. Every chunk carries it, chunk 0 included
+    // (#1785): the corpus is encoded without special tokens, so chunk 0
+    // starts unanchored too, and on Laguna XS 2.1 at width 256 that alone
+    // read as perplexity ~1000 against ~17.8 for an anchored chunk. The
+    // corpus stays encoded without special tokens because adding them would
+    // shift every chunk boundary, not just chunk 0's.
     let bos_prefix: Vec<i32> = tokenizer
         .encode("", true)
         .map(|ids| ids.into_iter().take(1).map(|t| t as i32).collect())
@@ -245,6 +256,19 @@ fn main() -> Result<()> {
         "# chunks\t{n_chunks}\tchunk_tokens\t{chunk_tokens}\ttopk\t{topk}\tprefill\t{prefill}"
     );
     println!("# columns\tchunk\tpos\ttarget\tnll\ttop_ids\ttop_logits");
+    // Traces written before #1785 left chunk 0 unanchored, so their chunk-0
+    // rows are not comparable with a trace that carries this line.
+    let (bos_anchor, anchored) = if bos_prefix.is_empty() {
+        ("none".to_string(), "none")
+    } else {
+        let ids = bos_prefix
+            .iter()
+            .map(i32::to_string)
+            .collect::<Vec<_>>()
+            .join(",");
+        (ids, "all")
+    };
+    println!("# bos_anchor\t{bos_anchor}\tchunks\t{anchored}");
 
     println!("# start_token\t{start_token}");
     if let Ok(mode) = std::env::var("MLXCEL_TRACE_GEMMA4") {
@@ -260,7 +284,7 @@ fn main() -> Result<()> {
             // The context pass already carried the anchor and the history, so
             // the traced forward is the continuation itself.
             (seg[..(l as usize - 1)].to_vec(), 0i32)
-        } else if c == 0 || bos_prefix.is_empty() {
+        } else if bos_prefix.is_empty() {
             (seg[..(l as usize - 1)].to_vec(), 0i32)
         } else {
             let mut with_bos = bos_prefix.clone();
@@ -279,12 +303,14 @@ fn main() -> Result<()> {
         let mut caches = model.make_caches();
         // Optional context pass. Its rows are discarded; it exists so the
         // traced forward runs at the chunk's width against a realistic
-        // context rather than against nothing.
+        // context rather than against nothing. A chunk with no history (chunk
+        // 0 at the corpus head) still gets the anchor alone, so the traced
+        // forward keeps the chunk's width and `target_offset` stays 0.
         if prefill > 0 && (!contiguous || c == 0) {
             let ctx_from = start.saturating_sub(prefill);
-            if ctx_from < start {
-                let mut ctx: Vec<i32> = bos_prefix.clone();
-                ctx.extend_from_slice(&ids[ctx_from..start]);
+            let mut ctx: Vec<i32> = bos_prefix.clone();
+            ctx.extend_from_slice(&ids[ctx_from..start]);
+            if !ctx.is_empty() {
                 let ctx_len = ctx.len() as i32;
                 if let Some(mode) = gemma4_mode {
                     mode.prefill(&model, &ctx, chunk_tokens)?;