fix(xla): stop emitting the terminating EOS token as output

[inureyes]

Summary

The OpenXLA backend handed the terminating EOS id to its consumers as if it were
generated output. The eager MLX paths drop it. That produced a visible control token
in user-facing content and a completion-token count one too high, on both the CLI
and the OpenAI-compatible server, for every model served through this backend.

Measured on GB10 with molmo2-4b, prompt Answer with one word only. What color is snow?, -n 12:

path	content	tokens
eager MLX CLI	White	1
MLXCEL_BACKEND=xla CLI, before	White<|im_end|>	2
MLXCEL_BACKEND=xla CLI, after	White	1

What changed

The eager BatchScheduler tests merged_eos before recording or detokenizing, so
the id never reaches generated_tokens or the incremental detokenizer. Every XLA
path recorded or emitted first and tested after. This makes the ordering match.

• mlxcel-xla/src/batch.rs: new commit_token returns both whether a sampled token
  is client-visible and why the sequence ended, and both XlaBatchEngine::pump
  sites (admission and decode) use it. finish_reason's existing "EOS wins over
  Length" precedence is untouched, and a token that ends generation by reaching the
  cap is still emitted.
• mlxcel-xla/src/lib.rs: the four single-sequence greedy generators now share one
  drive_greedy driver that applies the same contract once, instead of four copies
  of a push-then-test loop.
• mlxcel-xla/src/batch.rs: XlaReferenceEngine::generate likewise, so a reference
  token sequence is directly comparable with an eager one.
• server/batch/xla_worker.rs: the decode-width metric is now
  sampled_token_count, which counts a Stop finish as the one sampled token whose
  id was withheld. Without this the throughput metric would silently lose one step
  per sequence.

Why this matters beyond the visible symptom

Under the reduction-order contract just decided for #932, a family PR is gated on
per-operator parity plus greedy token-exactness against MLX. Before this fix
every XLA family would have failed that gate on one trailing token, regardless of
its numerics, so this lands first.

Validation

Unit tests, no device required:

• commit_token_withholds_the_terminating_eos_id — suppressed and reported Stop,
  including as the very first (prefill) token and when it is also the cap-th token.
• commit_token_emits_a_length_terminated_final_token — a non-EOS token that
  exhausts the budget is still emitted and reported Length.
• stop_finish_without_a_token_event_returns_empty_content_and_no_completion_tokens
  — a bare Finished{Stop} yields empty content, completion_tokens=0, and
  finish_reason=stop.
• sampled_token_count_charges_a_withheld_eos_to_the_decode_step — the metric does
  not shrink.

cargo test -p mlxcel-xla --lib: 252 passed.
cargo test --release --lib --features xla-iree server::batch::xla_worker: 9 passed.
cargo fmt --all --check and git diff --check clean.

Real checkpoint on CUDA (GB10, driver 580.159.03, CUDA 13.0, MLX pin
b7c3dd6d27f4, MLX_CUDA_ARCHITECTURES=121, MLX_ENABLE_TF32=0, release):
the CLI table above, and the XLA server's non-streaming and streaming content.

Existing callers that pass an empty EOS set (tests/xla_prepared_prefill.rs,
tests/xla_gemma3n_prefill.rs) are unaffected: with no EOS id the loops emit
exactly max_new_tokens tokens and issue the same number of decode steps as before.

Out of scope

SampleParams carries no per-request stop token ids and XlaBatchEngine builds its
EOS set from the checkpoint config alone, so request-level stop_token_ids are not
honoured by the XLA engine at all. That is a separate gap, noted in the issue and
not addressed here.

Closes #963
Refs #566, #932, #916

[inureyes]

Real-checkpoint before/after on GB10

Both measurements use the same fixture, the same binary lineage and the same
host. The only difference is whether the three changed source files carry this
commit or its parent: I reverted exactly
mlxcel-xla/src/{batch.rs,lib.rs} and server/batch/xla_worker.rs to HEAD~1,
rebuilt --release --features cuda,xla-iree --bin mlxcel, measured, then restored
them and measured again. No result is carried over from another head or another
model.

Environment: GB10 / Grace-Blackwell, driver 580.159.03, CUDA 13.0, rustc 1.93.1,
MLX pin b7c3dd6d27f4, MLX_CUDA_ARCHITECTURES=121 (the shipped value, confirmed in
the live CMakeCache.txt), MLX_ENABLE_TF32=0, release profile, IREE
3.12.0rc20260721. Machine idle (loadavg 0.4) for both runs.

Prompt Answer with one word only. What color is snow?, max_tokens=12,
temperature=0.

consumer	before	after
eager MLX CLI (reference)	White, 1 token	White, 1 token
XLA CLI	White<|im_end|>, 2 tokens	White, 1 token
XLA server, non-streaming	White<|im_end|>, stop, completion_tokens=2	White, stop, completion_tokens=1
XLA server, streaming (assembled)	White<|im_end|>, stop	White, stop

After the fix the XLA CLI matches the eager CLI byte for byte and token for token,
the server carries no EOS text on either the streaming or the non-streaming path,
finish_reason stays stop, and completion_tokens drops from 2 to the correct 1.

Checkpoint note

The fixture is qwen2.5-1.5b-instruct-4bit plus a generation_config.json
containing only the two ids the checkpoint's own config.json already declares
(bos_token_id: 151643, eos_token_id: 151645); every other file is a hardlink to
the original. That file has to be added because both backends read EOS ids only
from generation_config.json (loading/mod.rs:314,
mlxcel-xla/src/lib.rs:227) and never fall back to config.json, which HF
transformers does. On the unmodified mlx-community repack, eager and XLA alike
run to the token cap and never reach an EOS at all, so that checkpoint cannot
exercise this path in either direction. That fallback gap is pre-existing, affects
both backends equally, and is not touched here; it is worth its own issue.

The originally reported symptom (Red<|im_end|> on a Molmo2 XLA server) cannot be
reproduced from main because Molmo2's XLA support lands with #916, not here. The
Molmo2 numbers in the issue were measured on that branch; the table above is the
main-based equivalent on a checkpoint main can actually serve.

[inureyes]

Amended before merge: a regression I introduced in the first push

Reviewing my own diff, drive_greedy had the two break conditions in the wrong
order:

if out.len() >= max_new_tokens || !on_token(current) { break; }

|| short-circuits, so once the token budget was exhausted the streaming callback
was never invoked for that final token. The original loop called on_token for
every token including the last, so a streaming consumer would have silently lost the
last piece. The CLI would not have shown it (it renders from the returned vector),
which is exactly the kind of thing that survives a casual test.

Fixed by invoking the callback first, and the ordering requirement is now stated in
the doc comment so it is not "tidied" back:

if !on_token(current) || out.len() >= max_new_tokens { break; }

To make this testable without a device I split the loop out of the session as
drive_greedy_loop, taking decode as a closure. Four new tests in
greedy_contract_tests run under both the plain and the iree build:

• withholds_the_terminating_eos_and_never_streams_it
• returns_nothing_when_the_prefill_token_is_eos
• streams_the_token_that_exhausts_the_budget (the regression guard)
• a_false_callback_stops_generation_keeping_that_token

cargo test -p mlxcel-xla --lib: 254 passed, 11 ignored.
cargo test --release --features xla-iree --test xla_prepared_prefill --test xla_gemma3n_prefill:
compiles and runs; all cases are #[ignore]d real-checkpoint gates, so this is
compile-level assurance only.
The unused import: vision::emit_vision_diagnostics warning is pre-existing on
main (verified by stashing this change).

The real-checkpoint before/after table above was measured with the callback in the
pre-amend order. That ordering affects only the streaming callback, not the returned
token vector or the server path (XlaServeWorker consumes EngineEvents from the
batch engine, not this loop), so the measured CLI and server results still hold. The
new unit tests cover the callback contract directly.