From ffa41ae963849822e19f6bfb73691ea4c7ff96f7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 30 Jul 2026 19:59:38 +0900 Subject: [PATCH] fix(xla): stop emitting the terminating EOS token as output Every OpenXLA path handed the terminating EOS id to its consumers as if it were generated output, while the eager MLX paths drop it. On GB10 the eager CLI returns `White` (1 token) and the XLA CLI returned `White<|im_end|>` (2 tokens); the XLA server showed the same with finish_reason=stop and completion_tokens one too high. 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: XlaBatchEngine::pump pushed EngineEvent::Token before computing finish_reason at both the admission and decode sites, and the four single-sequence greedy loops plus XlaReferenceEngine::generate pushed into their output vector before the EOS test. XlaServeWorker counts one generated token per Token event and detokenizes whatever it receives, so the same defect produced both the leaked text and the inflated count. Introduce commit_token, which returns both whether a sampled token is client-visible and why the sequence ended, and route both pump sites through it. Collapse the four greedy loops onto one drive_greedy_loop so the contract is stated once, and keep that loop free of the session so it is unit-testable without a device. A token that ends generation by reaching the cap is still real output: it is collected and it is still handed to the streaming callback, which is why the callback is invoked before the budget is tested rather than after. finish_reason keeps its existing precedence of EOS over Length. The decode-width metric now counts a Stop finish as the one sampled token whose id was withheld, so suppressing the id does not silently drop a decode step from /metrics. Closes #963 Refs #566, #932, #916 --- src/lib/mlxcel-xla/src/batch.rs | 123 +++++++++++++++--- src/lib/mlxcel-xla/src/lib.rs | 187 ++++++++++++++++++++------- src/server/batch/xla_worker.rs | 30 ++++- src/server/batch/xla_worker_tests.rs | 73 +++++++++++ 4 files changed, 349 insertions(+), 64 deletions(-) diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index b08cb4b2..7361cd32 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -112,7 +112,8 @@ impl From for XlaAdmissionError { /// called [`XlaBatchEngine::cancel`] already knows), so it is not a finish reason. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FinishReason { - /// An EOS token was emitted. + /// An EOS token was sampled. The id itself is a control token and is not + /// emitted as a [`EngineEvent::Token`] (issue #963). Stop, /// The token budget (`max_new_tokens`) was reached. Length, @@ -233,6 +234,35 @@ fn finish_reason(produced: usize, cap: usize, last: i32, eos: &[i32]) -> Option< } } +/// What committing one sampled token means for its slot. +#[derive(Debug, PartialEq, Eq)] +struct TokenCommit { + /// Whether the token is client-visible output. + emit: bool, + /// Why the sequence ended, if it did. + finish: Option, +} + +/// Decide, for a slot that has now produced `produced` tokens against `cap`, +/// whether `last` is output and whether it ends the sequence. +/// +/// The terminating EOS id is a control token, not output. The eager +/// `BatchScheduler` drops it before it reaches `generated_tokens` or the +/// incremental detokenizer, so the XLA engine must not hand it to consumers +/// either: `XlaServeWorker` detokenizes every `EngineEvent::Token` it receives +/// and counts one generated token per event, so emitting the EOS id both leaks +/// it into content (`White<|im_end|>`) and reports one token too many +/// (issue #963). +/// +/// A token that ends the sequence by reaching `cap` is real output and is still +/// emitted. EOS still wins over Length when a token is both. +fn commit_token(produced: usize, cap: usize, last: i32, eos: &[i32]) -> TokenCommit { + TokenCommit { + emit: !eos.contains(&last), + finish: finish_reason(produced, cap, last, eos), + } +} + /// Backend-neutral scheduling state: `B_max` slots plus a FIFO admission queue. /// Holds no IREE handles, so its admit/evict/cancel bookkeeping is unit-testable /// without a device. [`XlaBatchEngine`] owns one of these next to the IREE engine @@ -683,10 +713,13 @@ impl XlaBatchEngine { if needs_history { history.push(first); } - events.push(EngineEvent::Token { - req_id: p.req_id, - token: first, - }); + let commit = commit_token(1, p.cap, first, &eos); + if commit.emit { + events.push(EngineEvent::Token { + req_id: p.req_id, + token: first, + }); + } let slot = Slot { req_id: p.req_id, adapter_mode: p.input.adapter_mode(), @@ -699,7 +732,7 @@ impl XlaBatchEngine { rng, history, }; - if let Some(reason) = finish_reason(slot.produced, slot.cap, first, &eos) { + if let Some(reason) = commit.finish { // Finished at its first token: leave the slot free for the next admit. events.push(EngineEvent::Finished { req_id: p.req_id, @@ -759,9 +792,11 @@ impl XlaBatchEngine { slot.cache_len += 1; slot.produced += 1; let req_id = slot.req_id; - let done = finish_reason(slot.produced, slot.cap, nt, &eos); - events.push(EngineEvent::Token { req_id, token: nt }); - if let Some(reason) = done { + let commit = commit_token(slot.produced, slot.cap, nt, &eos); + if commit.emit { + events.push(EngineEvent::Token { req_id, token: nt }); + } + if let Some(reason) = commit.finish { events.push(EngineEvent::Finished { req_id, reason }); *slot_opt = None; } @@ -1232,6 +1267,10 @@ impl XlaReferenceEngine { /// is prefilled in full and its argmax is the first token, then decode advances /// from there. /// + /// The terminating EOS id is a control token and is not returned, so this + /// reference sequence is directly comparable with an eager MLX one + /// (issue #963). + /// /// # Errors /// /// Propagates prefill / decode failures. @@ -1249,15 +1288,16 @@ impl XlaReferenceEngine { if max_new_tokens == 0 { return Ok(Vec::new()); } - let first = self.inner.prefill_first(prompt)?; - let mut out = vec![first]; + let mut out = Vec::with_capacity(max_new_tokens); let mut cache_len = prompt.len() as i32; - let mut cur = first; - while out.len() < max_new_tokens && !eos.contains(&cur) { - let nt = self.inner.decode(cur, cache_len)?; - out.push(nt); + let mut cur = self.inner.prefill_first(prompt)?; + while !eos.contains(&cur) { + out.push(cur); + if out.len() >= max_new_tokens { + break; + } + cur = self.inner.decode(cur, cache_len)?; cache_len += 1; - cur = nt; } Ok(out) } @@ -1287,6 +1327,57 @@ mod tests { assert_eq!(finish_reason(1, 1, 7, &EOS), Some(FinishReason::Length)); } + #[test] + fn commit_token_withholds_the_terminating_eos_id() { + // The EOS id ends the sequence and is not output: emitting it would leak + // `<|im_end|>` into content and count one token too many (issue #963). + assert_eq!( + commit_token(3, 8, 42, &EOS), + TokenCommit { + emit: false, + finish: Some(FinishReason::Stop), + } + ); + // Even as the very first (prefill) token, so an immediately-finished + // request streams nothing at all. + assert_eq!( + commit_token(1, 8, 42, &EOS), + TokenCommit { + emit: false, + finish: Some(FinishReason::Stop), + } + ); + // And even when it is also the cap-th token, where EOS wins over Length. + assert_eq!( + commit_token(8, 8, 42, &EOS), + TokenCommit { + emit: false, + finish: Some(FinishReason::Stop), + } + ); + } + + #[test] + fn commit_token_emits_a_length_terminated_final_token() { + // A non-EOS token that exhausts the budget is real output: it must be + // emitted, unlike the EOS id above. + assert_eq!( + commit_token(8, 8, 7, &EOS), + TokenCommit { + emit: true, + finish: Some(FinishReason::Length), + } + ); + // Mid-sequence: emitted, not finished. + assert_eq!( + commit_token(3, 8, 7, &EOS), + TokenCommit { + emit: true, + finish: None, + } + ); + } + fn g() -> SampleParams { SampleParams::greedy() } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b132826..2be9febd 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -400,7 +400,7 @@ impl XlaInferenceSession { prompt_tokens: &[i32], max_new_tokens: usize, eos_token_ids: &[i32], - mut on_token: F, + on_token: F, ) -> Result, String> { if prompt_tokens.is_empty() { return Err("XLA generation requires a non-empty prompt".to_string()); @@ -413,18 +413,42 @@ impl XlaInferenceSession { // and avoids double-counting the last prompt token in the cache. let split = prompt_tokens.len() - 1; self.prefill(&prompt_tokens[..split])?; - let mut current = prompt_tokens[split]; - let mut out = Vec::with_capacity(max_new_tokens); - for _ in 0..max_new_tokens { - let next = self.decode_step(current)?; - out.push(next); - let keep_going = on_token(next); - if !keep_going || eos_token_ids.contains(&next) { - break; - } - current = next; + if max_new_tokens == 0 { + return Ok(Vec::new()); } - Ok(out) + let first = self.decode_step(prompt_tokens[split])?; + self.drive_greedy(first, max_new_tokens, eos_token_ids, on_token) + } + + /// Advance greedy decode from `first` (a prefill's argmax) and collect the + /// generated tokens. + /// + /// Applies the shared terminating-token contract: the EOS id that ends the + /// sequence is a control token, so it is neither collected nor handed to + /// `on_token`. This mirrors the eager `BatchScheduler`, which drops the EOS + /// id before it reaches `generated_tokens` or the incremental detokenizer; + /// without it the CLI renders `White<|im_end|>` where eager renders `White`, + /// and reports one generated token too many (issue #963). A token that ends + /// generation by reaching `max_new_tokens` is real output and is collected. + /// + /// Every collected token is handed to `on_token` exactly once, including the + /// one that exhausts `max_new_tokens`, so a streaming consumer never loses the + /// final piece. `on_token` returning `false` stops generation with that token + /// still collected. The order of the two break conditions matters: testing the + /// budget first would short-circuit past the callback for the last token. + /// + /// `max_new_tokens` is assumed non-zero; callers return early on zero before + /// spending a prefill. + fn drive_greedy bool>( + &mut self, + first: i32, + max_new_tokens: usize, + eos_token_ids: &[i32], + on_token: F, + ) -> Result, String> { + drive_greedy_loop(first, max_new_tokens, eos_token_ids, on_token, |token| { + self.decode_step(token) + }) } /// Seed KV from a complete token prompt and return the prefill argmax. @@ -521,14 +545,7 @@ impl XlaInferenceSession { ) .map_err(|error| error.to_string())?; let first = self.prefill_deepstack_prepared(request)?; - let mut out = Vec::with_capacity(max_new_tokens); - out.push(first); - let mut current = first; - while out.len() < max_new_tokens && !eos_token_ids.contains(¤t) { - current = self.decode_step(current)?; - out.push(current); - } - Ok(out) + self.drive_greedy(first, max_new_tokens, eos_token_ids, |_| true) } /// Greedy generation seeded by Gemma3n post-scale embeddings and dense PLE. @@ -548,14 +565,7 @@ impl XlaInferenceSession { ) .map_err(|error| error.to_string())?; let first = self.prefill_gemma3n_prepared(request)?; - let mut out = Vec::with_capacity(max_new_tokens); - out.push(first); - let mut current = first; - while out.len() < max_new_tokens && !eos_token_ids.contains(¤t) { - current = self.decode_step(current)?; - out.push(current); - } - Ok(out) + self.drive_greedy(first, max_new_tokens, eos_token_ids, |_| true) } /// Streaming greedy generation from prepared embeddings. The embeddings @@ -567,7 +577,7 @@ impl XlaInferenceSession { prepared: &PreparedPrefill, max_new_tokens: usize, eos_token_ids: &[i32], - mut on_token: F, + on_token: F, ) -> Result, String> { if max_new_tokens == 0 { return Ok(Vec::new()); @@ -575,21 +585,7 @@ impl XlaInferenceSession { validate_request_capacity(prepared.sequence_len, max_new_tokens, self.context_capacity) .map_err(|error| error.to_string())?; let first = self.prefill_prepared(prepared)?; - let mut out = Vec::with_capacity(max_new_tokens); - out.push(first); - if !on_token(first) || eos_token_ids.contains(&first) { - return Ok(out); - } - let mut current = first; - while out.len() < max_new_tokens { - let next = self.decode_step(current)?; - out.push(next); - if !on_token(next) || eos_token_ids.contains(&next) { - break; - } - current = next; - } - Ok(out) + self.drive_greedy(first, max_new_tokens, eos_token_ids, on_token) } } @@ -656,6 +652,34 @@ impl InferenceSession for XlaInferenceSession { } } +/// The greedy drive loop shared by every single-sequence generator, with `decode` +/// standing in for `decode_step`. +/// +/// Split out from [`XlaInferenceSession::drive_greedy`] so the terminating-token +/// contract of issue #963 is unit-testable without an IREE device. +fn drive_greedy_loop( + first: i32, + max_new_tokens: usize, + eos_token_ids: &[i32], + mut on_token: F, + mut decode: D, +) -> Result, String> +where + F: FnMut(i32) -> bool, + D: FnMut(i32) -> Result, +{ + let mut out = Vec::with_capacity(max_new_tokens); + let mut current = first; + while !eos_token_ids.contains(¤t) { + out.push(current); + if !on_token(current) || out.len() >= max_new_tokens { + break; + } + current = decode(current)?; + } + Ok(out) +} + // These scaffold tests cover the without-`iree` behavior (load records inputs; // the token primitives report NOT_WIRED). Under `--features iree`, `load` // constructs a real engine from a real model directory, so the fake-path fixture @@ -804,3 +828,78 @@ mod tests { assert!(!super::llava_diagnostic_device_memory_note("local-task").contains("GB10")); } } + +// The terminating-token contract (issue #963) is device-independent, so these run +// under both the plain and the `iree` build. +#[cfg(test)] +mod greedy_contract_tests { + use super::drive_greedy_loop; + + /// Drive the loop over a scripted token stream, returning what was collected + /// and, separately, what the streaming callback actually saw. + fn drive( + first: i32, + max_new_tokens: usize, + eos: &[i32], + script: &[i32], + ) -> (Vec, Vec) { + let mut streamed = Vec::new(); + let mut next = script.iter().copied(); + let out = drive_greedy_loop( + first, + max_new_tokens, + eos, + |token| { + streamed.push(token); + true + }, + |_| next.next().ok_or_else(|| "script exhausted".to_string()), + ) + .expect("scripted drive"); + (out, streamed) + } + + #[test] + fn withholds_the_terminating_eos_and_never_streams_it() { + // 7, 8, then EOS: the EOS id ends the sequence without being collected or + // streamed, so the result is comparable with an eager MLX sequence. + let (out, streamed) = drive(7, 8, &[42], &[8, 42]); + assert_eq!(out, vec![7, 8]); + assert_eq!(streamed, vec![7, 8]); + } + + #[test] + fn returns_nothing_when_the_prefill_token_is_eos() { + let (out, streamed) = drive(42, 8, &[42], &[]); + assert!(out.is_empty()); + assert!(streamed.is_empty()); + } + + #[test] + fn streams_the_token_that_exhausts_the_budget() { + // Regression guard: testing the budget before invoking the callback would + // short-circuit past it and silently drop the last streamed piece. + let (out, streamed) = drive(7, 3, &[42], &[8, 9]); + assert_eq!(out, vec![7, 8, 9]); + assert_eq!(streamed, vec![7, 8, 9], "the final token must still stream"); + } + + #[test] + fn a_false_callback_stops_generation_keeping_that_token() { + let mut seen = Vec::new(); + let mut next = [8, 9].iter().copied(); + let out = drive_greedy_loop( + 7, + 8, + &[42], + |token| { + seen.push(token); + token != 8 + }, + |_| next.next().ok_or_else(|| "script exhausted".to_string()), + ) + .expect("scripted drive"); + assert_eq!(out, vec![7, 8]); + assert_eq!(seen, vec![7, 8]); + } +} diff --git a/src/server/batch/xla_worker.rs b/src/server/batch/xla_worker.rs index 81105bd6..32894efc 100644 --- a/src/server/batch/xla_worker.rs +++ b/src/server/batch/xla_worker.rs @@ -178,6 +178,31 @@ struct PendingAudioState { start: Instant, } +/// How many tokens one engine step actually sampled, which is the step's decode +/// width for `/metrics`. +/// +/// This is not simply the number of `Token` events: a sequence that ends on EOS +/// still sampled a token, but the engine withholds that id because it is a +/// control token rather than output (issue #963). Each `Finished` with +/// [`XlaFinishReason::Stop`] therefore stands for exactly one sampled-but- +/// withheld token, while a `Length` finish always accompanies its emitted +/// `Token`. Counting both keeps the throughput metric measuring device work. +fn sampled_token_count(events: &[EngineEvent]) -> usize { + events + .iter() + .filter(|event| { + matches!( + event, + EngineEvent::Token { .. } + | EngineEvent::Finished { + reason: XlaFinishReason::Stop, + .. + } + ) + }) + .count() +} + /// Server-side worker that serves requests through the OpenXLA continuous-batching /// engine. Built and run on a single worker thread (see /// `model_worker::spawn_xla_model_worker`). @@ -474,10 +499,7 @@ impl BatchEngine for XlaServeWorker { Ok(events) => { // Each `Token` event is one token produced this step across the // active batch, so the count is the step's decode width. - let decoded = events - .iter() - .filter(|e| matches!(e, EngineEvent::Token { .. })) - .count(); + let decoded = sampled_token_count(&events); if decoded > 0 { self.batch_observability.record_decode_step(decoded); } diff --git a/src/server/batch/xla_worker_tests.rs b/src/server/batch/xla_worker_tests.rs index 39307d3d..c861e4ec 100644 --- a/src/server/batch/xla_worker_tests.rs +++ b/src/server/batch/xla_worker_tests.rs @@ -666,3 +666,76 @@ fn pending_preprocess_poll_timeout_does_not_shutdown_worker() { assert!(!worker.shutdown); drop(request_tx); } + +#[test] +fn stop_finish_without_a_token_event_returns_empty_content_and_no_completion_tokens() { + // The engine withholds the terminating EOS id (issue #963), so a sequence + // that ends immediately on EOS reaches the worker as a bare `Finished`. + // Nothing may be streamed, `completion_tokens` must stay 0, and the finish + // reason must still be `stop`. + let mut worker = worker(None); + let (response_tx, response_rx) = mpsc::channel(); + + worker.admit( + String::new(), + Some(vec![10, 11]), + options(8), + Vec::new(), + Vec::new(), + Vec::new(), + media(0, 0, 0), + response_tx, + Arc::new(AtomicBool::new(false)), + ); + let req_id = *worker.states.keys().next().expect("one admitted request"); + + worker.dispatch(vec![EngineEvent::Finished { + req_id, + reason: XlaFinishReason::Stop, + }]); + + let result = match response_rx.recv_timeout(Duration::from_secs(1)).unwrap() { + GenerateEvent::Done(result) => result, + GenerateEvent::Token(text) | GenerateEvent::TokenWithLogprobs(text, _) => { + panic!("the withheld EOS id was streamed as content: {text:?}") + } + GenerateEvent::Error(error) => panic!("unexpected generation failure: {error}"), + }; + assert_eq!(result.text, ""); + assert_eq!(result.completion_tokens, 0); + assert_eq!(result.finish_reason, "stop"); + assert!(worker.states.is_empty()); +} + +#[test] +fn sampled_token_count_charges_a_withheld_eos_to_the_decode_step() { + // Suppressing the EOS id must not shrink the decode-width metric: the + // engine still ran that step on device. + let one_token_then_eos = vec![ + EngineEvent::Token { + req_id: 0, + token: 7, + }, + EngineEvent::Finished { + req_id: 1, + reason: XlaFinishReason::Stop, + }, + ]; + assert_eq!(sampled_token_count(&one_token_then_eos), 2); + + // A length finish accompanies its own emitted token, so it is not counted + // twice. + let length_terminated = vec![ + EngineEvent::Token { + req_id: 0, + token: 7, + }, + EngineEvent::Finished { + req_id: 0, + reason: XlaFinishReason::Length, + }, + ]; + assert_eq!(sampled_token_count(&length_terminated), 1); + + assert_eq!(sampled_token_count(&[]), 0); +}