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); +}