Problem
Every OpenXLA generation path returns the terminating EOS id as if it were generated
output. The eager MLX paths drop it. The result is a visible control token in user
content and a completion-token count that is one too high.
Reproduced on GB10 (driver 580.159.03, CUDA 13.0, MLX pin b7c3dd6d27f4,
MLX_CUDA_ARCHITECTURES=121, MLX_ENABLE_TF32=0), release profile, checkpoint
molmo2-4b, prompt Answer with one word only. What color is snow?, -n 12:
| path |
returned content |
reported tokens |
| eager MLX CLI (default backend) |
White |
1 |
MLXCEL_BACKEND=xla CLI |
White<|im_end|> |
2 |
The same defect shows on the XLA server: a text-only request to a Molmo2 XLA server
returns content Red<|im_end|> with finish_reason=stop, where the eager server
returns Red.
This is not model-specific and not server-specific. It affects every model served or
generated through the XLA backend, on both the CLI and the OpenAI-compatible server.
Root cause
The eager path treats the terminating id as a control token. BatchScheduler checks it
before recording or detokenizing, so it never reaches generated_tokens or the
incremental detokenizer:
// src/server/batch/scheduler.rs
if seq.merged_eos.contains(&token_val) {
// transition to Finished(Stop)
continue;
}
seq.generated_tokens.push(token_val);
Every XLA path records or emits the token before testing it:
src/lib/mlxcel-xla/src/batch.rs:686 — XlaBatchEngine::pump admission: pushes
EngineEvent::Token for the prefill argmax, then computes finish_reason.
src/lib/mlxcel-xla/src/batch.rs:763 — XlaBatchEngine::pump decode: pushes
EngineEvent::Token, then checks finish_reason.
src/lib/mlxcel-xla/src/lib.rs:420 — generate_streaming_greedy: out.push(next)
before the EOS test.
src/lib/mlxcel-xla/src/lib.rs:586 — the prepared-input greedy loop, same shape.
src/lib/mlxcel-xla/src/batch.rs — XlaReferenceEngine::generate, same shape.
Because XlaServeWorker counts one generated token per EngineEvent::Token
(src/server/batch/xla_worker.rs:328) and detokenizes whatever it receives, the engine
handing over the EOS token is enough to produce both symptoms.
Note the admission path is already inconsistent with itself: when the prefill argmax is
EOS it emits Token and Finished, while the slot is discarded without ever decoding.
Expected contract
The terminating EOS id is a control token, never output. On every XLA path:
- it must not appear in returned tokens or streamed content,
- it must not count toward
completion_tokens,
finish_reason must still be stop,
- a token that ends generation by hitting the length cap is real output and must still
be emitted.
Required implementation
- Make the emit/finish decision one explicit, unit-testable step shared by both
pump sites, so Length still emits and Stop does not.
- Apply the same ordering to the three single-sequence greedy loops.
- Keep
finish_reason's existing "EOS wins over Length" precedence.
Validation
- Unit tests, no device needed: EOS suppressed and reported as
Stop; a cap-reached
non-EOS token emitted and reported as Length; EOS as the very first (prefill) token
emits nothing and finishes.
- Real checkpoint on CUDA:
MLXCEL_BACKEND=xla CLI output and token count match the
eager CLI byte for byte on a prompt that terminates on EOS.
- Real checkpoint on CUDA: XLA server non-streaming and streaming content carries no
EOS text, finish_reason=stop, and completion_tokens matches the eager server.
Acceptance criteria
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 and is not addressed here.
Found while qualifying #916 against #871's "text-only behaviour remains token-exact"
criterion. Refs #566, #932.
Problem
Every OpenXLA generation path returns the terminating EOS id as if it were generated
output. The eager MLX paths drop it. The result is a visible control token in user
content and a completion-token count that is one too high.
Reproduced on GB10 (driver 580.159.03, CUDA 13.0, MLX pin
b7c3dd6d27f4,MLX_CUDA_ARCHITECTURES=121,MLX_ENABLE_TF32=0), release profile, checkpointmolmo2-4b, promptAnswer with one word only. What color is snow?,-n 12:WhiteMLXCEL_BACKEND=xlaCLIWhite<|im_end|>The same defect shows on the XLA server: a text-only request to a Molmo2 XLA server
returns content
Red<|im_end|>withfinish_reason=stop, where the eager serverreturns
Red.This is not model-specific and not server-specific. It affects every model served or
generated through the XLA backend, on both the CLI and the OpenAI-compatible server.
Root cause
The eager path treats the terminating id as a control token.
BatchSchedulerchecks itbefore recording or detokenizing, so it never reaches
generated_tokensor theincremental detokenizer:
Every XLA path records or emits the token before testing it:
src/lib/mlxcel-xla/src/batch.rs:686—XlaBatchEngine::pumpadmission: pushesEngineEvent::Tokenfor the prefill argmax, then computesfinish_reason.src/lib/mlxcel-xla/src/batch.rs:763—XlaBatchEngine::pumpdecode: pushesEngineEvent::Token, then checksfinish_reason.src/lib/mlxcel-xla/src/lib.rs:420—generate_streaming_greedy:out.push(next)before the EOS test.
src/lib/mlxcel-xla/src/lib.rs:586— the prepared-input greedy loop, same shape.src/lib/mlxcel-xla/src/batch.rs—XlaReferenceEngine::generate, same shape.Because
XlaServeWorkercounts one generated token perEngineEvent::Token(
src/server/batch/xla_worker.rs:328) and detokenizes whatever it receives, the enginehanding over the EOS token is enough to produce both symptoms.
Note the admission path is already inconsistent with itself: when the prefill argmax is
EOS it emits
TokenandFinished, while the slot is discarded without ever decoding.Expected contract
The terminating EOS id is a control token, never output. On every XLA path:
completion_tokens,finish_reasonmust still bestop,be emitted.
Required implementation
pumpsites, soLengthstill emits andStopdoes not.finish_reason's existing "EOS wins over Length" precedence.Validation
Stop; a cap-reachednon-EOS token emitted and reported as
Length; EOS as the very first (prefill) tokenemits nothing and finishes.
MLXCEL_BACKEND=xlaCLI output and token count match theeager CLI byte for byte on a prompt that terminates on EOS.
EOS text,
finish_reason=stop, andcompletion_tokensmatches the eager server.Acceptance criteria
finish_reasonvalues are unchanged (stopfor EOS,lengthfor the cap).Out of scope
SampleParamscarries no per-request stop token ids andXlaBatchEnginebuilds itsEOS set from the checkpoint config alone, so request-level
stop_token_idsare nothonoured by the XLA engine at all. That is a separate gap and is not addressed here.
Found while qualifying #916 against #871's "text-only behaviour remains token-exact"
criterion. Refs #566, #932.