Skip to content

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

Closed

Description

@inureyes

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 test(laguna): add a checkpoint-code reference trace and forward tests #1779).
  • mlxcel and laguna_oracle_trace.py still agree on every (chunk, pos, target).

Verification

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.

Activity

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:benchmarkBenchmark harness and performance measurement (bench_*.sh, /update-benchmarks)
    on Sep 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:benchmarkBenchmark harness and performance measurement (bench_*.sh, /update-benchmarks)priority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions