fix(inference): a GPU decode step that fails mid-generation is an error

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<GenerateError>,
     /// Per-step wall time in ms.
     pub decode_ms: Vec<f64>,
     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<GenerateError> = None;
     let mut decode_ms: Vec<f64> =
         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<u32> = 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.

crates/larql-inference/src/layer_graph/generate/gpu/mod.rs

@@ -485,7 +485,7 @@ where
             detok_ms_total: outcome.t_detok,
             dequant_ms_total: 0.0,
         },
-        error: None,
+        error: outcome.error,
     }
 }
 

crates/larql-inference/src/layer_graph/generate/gpu/tests.rs

@@ -267,6 +267,57 @@ fn diag_compare_cpu_topk_runs_against_synthetic_weights() {
     );
 }
 
+/// 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}");
+}
+
 // ── The CPU fallback, through the streaming entry point ─────────────────
 //
 // `CpuBackend` has no fused Q4 pipeline, so `generate_streaming` hands the

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<Vec<f32>> {
-        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])
     }