From 96ac5071934f87df67e86dbca9b450643068cf08 Mon Sep 17 00:00:00 2001 From: chrishayuk Date: Thu, 12 Nov 2026 01:34:17 +0000 Subject: [PATCH] fix(inference): a GPU decode step that fails mid-generation is an error When the backend returned nothing for a decode step, or a step produced no candidate tokens (how a failed command buffer is reported), the decode loop stopped and returned the tokens it had. The caller built the result with error: None, so a generation cut short by a GPU fault was reported as a complete one. The loop now returns the reason with its outcome, and the result carries it as an error alongside the tokens produced before the failure. The mock GPU backend can be told to fail after a number of steps. --- .../layer_graph/generate/gpu/decode_loop.rs | 88 +++++++++++++++++++ .../src/layer_graph/generate/gpu/mod.rs | 2 +- .../src/layer_graph/generate/gpu/tests.rs | 51 +++++++++++ crates/larql-inference/src/test_utils/q4k.rs | 19 +++- 4 files changed, 158 insertions(+), 2 deletions(-) diff --git a/crates/larql-inference/src/layer_graph/generate/gpu/decode_loop.rs b/crates/larql-inference/src/layer_graph/generate/gpu/decode_loop.rs index 95d919531..71fdc8890 100644 --- a/crates/larql-inference/src/layer_graph/generate/gpu/decode_loop.rs +++ b/crates/larql-inference/src/layer_graph/generate/gpu/decode_loop.rs @@ -12,15 +12,33 @@ use crate::layer_graph::generate::eos::EosConfig; use crate::layer_graph::generate::lm_head::lm_head_topk_with_policy; use crate::layer_graph::generate::policy::GenerationRuntimeConfig; use crate::layer_graph::generate::sampling::Sampler; +use crate::layer_graph::generate::types::GenerateError; use crate::model::ModelWeights; use larql_compute::prelude::*; use larql_compute::FullPipelineLayer; +/// The error for a decode step that failed part-way through a generation. +/// +/// The loop used to stop here and return what it had, and the caller +/// reported that truncated output as a success. +fn step_failed(step: usize, produced: usize, what: &str) -> GenerateError { + GenerateError::Other { + reason: format!( + "generation stopped at decode step {step}: {what} ({produced} token(s) were \ + generated by the loop before it)" + ), + } +} + /// Aggregated output of the decode-loop phase. pub(super) struct DecodeLoopOutcome { /// `(text, prob)` per generated token (excluding the first, which the /// caller already produced from prefill). pub tokens: Vec<(String, f64)>, + /// Why the loop stopped, when it stopped because a step failed. `None` + /// for a loop that reached `max_tokens` or an end-of-sequence token. + /// The tokens produced before the failure are still in `tokens`. + pub error: Option, /// Per-step wall time in ms. pub decode_ms: Vec, pub t_embed: f64, @@ -74,6 +92,7 @@ where { let mut tokens: Vec<(String, f64)> = Vec::with_capacity(crate::generation_capacity::generation_capacity(max_tokens)); + let mut error: Option = None; let mut decode_ms: Vec = Vec::with_capacity(crate::generation_capacity::generation_capacity(max_tokens)); @@ -199,6 +218,13 @@ where "[profile] step={step} — GPU decode returned None; stopping generation" ); } + // An early stop, and not a successful one: the output is + // shorter than the model would have made it. + error = Some(step_failed( + step, + tokens.len(), + "the backend returned no output", + )); break; }; @@ -247,6 +273,11 @@ where "[profile] step={step} — no candidates (empty head, or GPU step failed); break" ); } + error = Some(step_failed( + step, + tokens.len(), + "the step produced no candidate tokens", + )); break; }; let detok_ms = t4.elapsed().as_secs_f64() * 1000.0; @@ -296,6 +327,7 @@ where DecodeLoopOutcome { tokens, + error, decode_ms, t_embed, t_gpu, @@ -633,6 +665,62 @@ mod tests { assert_eq!(outcome.tokens.len(), callback_count); } + /// A backend that fails a step part-way stops the loop with an error, + /// and keeps the tokens produced before it. The loop used to stop and + /// return those tokens as a complete, successful generation. + #[test] + fn a_failed_decode_step_is_an_error_with_the_tokens_before_it() { + use crate::layer_graph::generate::detok::Detokenizer; + use crate::layer_graph::generate::eos::EosConfig; + use crate::layer_graph::generate::policy::GenerationRuntimeConfig; + use crate::layer_graph::generate::sampling::{Sampler, SamplingConfig}; + + /// Decode steps the backend completes before it fails. + const GOOD_STEPS: usize = 2; + const MAX_TOKENS: usize = 6; + + let run = |backend: MockGpuBackend| { + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + let tokenizer = crate::test_utils::make_test_tokenizer(weights.vocab_size); + let mut sampler = Sampler::new(SamplingConfig::greedy()); + let runtime = GenerationRuntimeConfig::default(); + let mut detok = Detokenizer::new(&tokenizer); + let mut generated_ids: Vec = vec![0]; + run_decode_loop( + &weights, + &tokenizer, + &index, + &backend, + &[], + weights.hidden_size, + weights.intermediate_size, + weights.arch.norm_weight_offset(), + 5, + &runtime, + &mut sampler, + &mut detok, + &EosConfig::empty(), + &mut generated_ids, + /*current_token_id=*/ 0, + MAX_TOKENS, + None, + &mut |_id, _text, _prob| {}, + ) + }; + + // The control: a backend that never fails runs to max_tokens with + // no error, so the error below is the failure and nothing else. + let whole = run(MockGpuBackend::new()); + assert!(whole.error.is_none(), "{:?}", whole.error); + assert_eq!(whole.tokens.len(), MAX_TOKENS - 1); + + let cut_short = run(MockGpuBackend::failing_after(GOOD_STEPS)); + assert_eq!(cut_short.tokens.len(), GOOD_STEPS); + let error = cut_short.error.expect("a failed step must be reported"); + assert!(error.to_string().contains("decode step"), "{error}"); + } + /// `profile_decode=true` exercises the per-step diagnostic blocks /// (lines 97-112, 134-136, 162-164, 187-220). All stderr writes, no /// behaviour change. diff --git a/crates/larql-inference/src/layer_graph/generate/gpu/mod.rs b/crates/larql-inference/src/layer_graph/generate/gpu/mod.rs index 9fd0b50ef..1a1702bf0 100644 --- a/crates/larql-inference/src/layer_graph/generate/gpu/mod.rs +++ b/crates/larql-inference/src/layer_graph/generate/gpu/mod.rs @@ -476,7 +476,7 @@ where detok_ms_total: outcome.t_detok, dequant_ms_total: 0.0, }, - error: None, + error: outcome.error, } } diff --git a/crates/larql-inference/src/layer_graph/generate/gpu/tests.rs b/crates/larql-inference/src/layer_graph/generate/gpu/tests.rs index ffae1d7d2..9d957832e 100644 --- a/crates/larql-inference/src/layer_graph/generate/gpu/tests.rs +++ b/crates/larql-inference/src/layer_graph/generate/gpu/tests.rs @@ -266,3 +266,54 @@ fn diag_compare_cpu_topk_runs_against_synthetic_weights() { &h_1d, ); } + +/// A GPU that fails a decode step part-way through a generation is an +/// error on the result, with the tokens produced before it. The public +/// entry point used to return those tokens with no error: a truncated +/// generation reported as a complete one. +#[test] +fn a_gpu_decode_step_that_fails_mid_generation_is_reported() { + use crate::test_utils::{ + make_test_q4k_vindex, make_test_q4k_weights, make_test_tokenizer, MockGpuBackend, + }; + /// Decode steps the backend completes before it fails. + const GOOD_STEPS: usize = 2; + const MAX_TOKENS: usize = 6; + + let run = |backend: MockGpuBackend| { + let mut weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + let tokenizer = make_test_tokenizer(weights.vocab_size); + let cached = CachedLayerGraph::from_residuals(vec![]); + let num_layers = weights.num_layers; + generate( + &mut weights, + &tokenizer, + &[0u32, 1, 2], + MAX_TOKENS, + &index, + &backend, + &cached, + 0..num_layers, + ) + }; + + // The control: the same request on a backend that does not fail. + let whole = run(MockGpuBackend::new()); + assert!(whole.error.is_none(), "{:?}", whole.error); + assert_eq!(whole.tokens.len(), MAX_TOKENS); + + let cut_short = run(MockGpuBackend::failing_after(GOOD_STEPS)); + assert!( + cut_short.tokens.len() < MAX_TOKENS, + "the failing backend must cut the generation short" + ); + assert!( + !cut_short.tokens.is_empty(), + "the tokens before the failure are kept" + ); + let error = cut_short + .error + .expect("a generation cut short by a failed step must carry an error"); + assert!(error.to_string().contains("decode step"), "{error}"); +} diff --git a/crates/larql-inference/src/test_utils/q4k.rs b/crates/larql-inference/src/test_utils/q4k.rs index 67fd4aadd..3c04f73af 100644 --- a/crates/larql-inference/src/test_utils/q4k.rs +++ b/crates/larql-inference/src/test_utils/q4k.rs @@ -461,6 +461,9 @@ impl Q4KTestFixtures { pub struct MockGpuBackend { inner: larql_compute::CpuBackend, kv_len: std::sync::atomic::AtomicUsize, + /// `decode_token` answers `None` once this many calls have succeeded. + /// `usize::MAX` never fails. + decode_steps_before_failure: usize, } impl Default for MockGpuBackend { @@ -474,6 +477,16 @@ impl MockGpuBackend { Self { inner: larql_compute::CpuBackend, kv_len: std::sync::atomic::AtomicUsize::new(0), + decode_steps_before_failure: usize::MAX, + } + } + + /// A backend whose `decode_token` succeeds `steps` times and then + /// answers `None`: a GPU that fails part-way through a generation. + pub fn failing_after(steps: usize) -> Self { + Self { + decode_steps_before_failure: steps, + ..Self::new() } } } @@ -529,8 +542,12 @@ impl larql_compute::DecodeBackend for MockGpuBackend { hidden: usize, _inter: usize, ) -> Option> { - self.kv_len + let succeeded = self + .kv_len .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if succeeded >= self.decode_steps_before_failure { + return None; + } Some(vec![0.0f32; hidden]) }