fix(examples): anchor logit_trace chunk 0 with the BOS token
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Found during PR #1779 (the Laguna reference comparison).
## Problem
`examples/logit_trace.rs` encodes the corpus without special tokens and prepends the tokenizer's BOS to every chunk except chunk 0; with `PREFILL > 0`, chunk 0 also gets no context pass. Chunk 0 therefore starts with no anchor at every width. On a model trained behind a document separator this collapses the first chunk: on Laguna XS 2.1 at width 256, chunk 0 reads perplexity 955 in the float32 reference and 999 in mlxcel, against about 17.8 in a BOS-anchored chunk, and it held 42 of the 60 top-1 disagreements (PR #1779). Both arms mirror the gap, so A/B comparisons stay valid, but a pooled perplexity or disagreement figure reads like a defect.
`examples/perplexity.rs` has the same `c == 0` skip (`:87`) but encodes with `encode(&text, true)` (`:50-59`), so its chunk 0 already starts at the tokenizer's BOS. It needs no change.
## Evidence
- `examples/logit_trace.rs:191-192` (`encode(text.as_str(), false)`), `:200-207` (the "Same BOS anchoring as `examples/perplexity`" comment and `bos_prefix`), `:223-233` (`if prefill > 0 { ... } else if c == 0 || bos_prefix.is_empty()`), `:245-256` (context pass only when `ctx_from < start`, never true for chunk 0).
- `scripts/laguna_oracle_trace.py` mirrors it on purpose: docstring `:44-49`, `plan_chunks` `:184-205`.
## Proposed fix
- Without `PREFILL`: prepend `bos_prefix` to chunk 0 as for later chunks, with `target_offset = bos_prefix.len()`.
- With `PREFILL`: give chunk 0 a context pass of `bos_prefix` alone, so the traced forward keeps the chunk's width and `target_offset` 0. Skip it when `bos_prefix` is empty.
- Keep `encode(text, false)`: encoding with special tokens would shift every chunk boundary, not just chunk 0.
- Add a `#` metadata line recording the anchoring. `scripts/compare_logit_traces.py` reads only the `model` metadata key (`:43-63`, `:107-108`), so parsing is unaffected.
- Update `plan_chunks` and its docstring in `scripts/laguna_oracle_trace.py` in the same PR.
- Note in `docs/benchmarks.md` (`:407`, "Judging a change that moves the numbers") that older traces are not comparable at chunk 0, and refresh the unanchored-window figures in the Laguna entry of `docs/supported-models.md:55`.
## Acceptance criteria
- [ ] Chunk 0 carries the BOS at every width, with and without `PREFILL`, whenever `bos_prefix` is non-empty; with an empty `bos_prefix` the data rows are unchanged.
- [ ] Laguna XS 2.1 at width 256: chunk-0 perplexity drops into the BOS-anchored range (about 17.8 in PR #1779).
- [ ] mlxcel and `laguna_oracle_trace.py` still agree on every `(chunk, pos, target)`.
## Verification
```bash
cargo build --release --features metal,accelerate --example logit_trace
./target/release/examples/logit_trace models/mlx/laguna-xs-2.1-nvfp4 CORPUS.txt 256 2 8 0 > mlxcel.tsv
python3 scripts/laguna_oracle_trace.py models/mlx/laguna-xs-2.1-nvfp4 CORPUS.txt 256 2 8 0 > oracle.tsv
python3 scripts/compare_logit_traces.py oracle.tsv mlxcel.tsv
```
Repeat both traces with `8 80 8 512` (width 8 behind 512 tokens) for the prefill path. Clippy and `cargo fmt --all -- --check` stay clean.
Contributor guide
Research direction
Start in examples/logit_trace.rs at the cited encoding, BOS, and context-pass sections, then compare the chunk plan with scripts/laguna_oracle_trace.py and its docstring. Run both verification commands for normal and prefill paths, then update the cited benchmark and supported-model documentation. Done means chunk 0 is BOS-anchored, empty-BOS rows are unchanged, both traces agree, and formatting and Clippy pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust
- Domain
- performance, testing, tooling
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100