Skip to content

fix(xla): terminating EOS token leaks into generated output on every XLA path #963

Description

@inureyes

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:686XlaBatchEngine::pump admission: pushes
    EngineEvent::Token for the prefill argmax, then computes finish_reason.
  • src/lib/mlxcel-xla/src/batch.rs:763XlaBatchEngine::pump decode: pushes
    EngineEvent::Token, then checks finish_reason.
  • src/lib/mlxcel-xla/src/lib.rs:420generate_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.rsXlaReferenceEngine::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

  • No XLA path returns, streams, or counts the terminating EOS id.
  • A length-terminated final token is still emitted and counted.
  • finish_reason values are unchanged (stop for EOS, length for the cap).
  • XLA CLI matches eager CLI content and token count on the pinned fixture.
  • XLA server streaming and non-streaming content contain no EOS text.
  • Unit tests cover the emit/finish decision without a device.

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersarea:inferenceGeneration, sampling, decoding (incl. speculative, DRY)priority:highHigh 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