lablup / lablup/mlxcel

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

Open
#1,785 0 comments 0 reactions 0 assignees View on GitHub
area:benchmark priority:low status:ready type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.