From d9198795bfdc5c1920a1a3bcb3086908aac752be Mon Sep 17 00:00:00 2001 From: Evandro Fonseca Junior Date: Mon, 30 Mar 2026 17:28:39 -0300 Subject: [PATCH] fix(quick-1): bail on NULL KV cache and improve prefill error diagnostics - engine.rs: new_session() now returns Err when llama_get_memory returns NULL instead of silently continuing, preventing decode corruption on 2nd request - session.rs: decode_prefill error includes token count for easier diagnosis --- crates/inference-core/src/engine.rs | 25 ++++++++++++--- crates/inference-core/src/session.rs | 47 +++++++++++++++------------- 2 files changed, 45 insertions(+), 27 deletions(-) diff --git a/crates/inference-core/src/engine.rs b/crates/inference-core/src/engine.rs index c41e421..14edab7 100644 --- a/crates/inference-core/src/engine.rs +++ b/crates/inference-core/src/engine.rs @@ -263,10 +263,6 @@ impl InferenceEngine { info!("Running engine warmup (dummy decode)…"); unsafe { - // Mark context as warmup so llama.cpp skips perf counters - if let Ok(f) = std::panic::catch_unwind(|| {}) { - let _ = f; - } (self.lib.set_warmup)(self.ctx.ptr, true); // set_warmup may not exist in all builds; call decode with 1 token let dummy_token: llama_sys::llama_token = 1; // BOS-ish token @@ -279,7 +275,11 @@ impl InferenceEngine { } let memory = (self.lib.get_memory)(self.ctx.ptr); - (self.lib.memory_clear)(memory, false); + if memory.is_null() { + warn!("llama_get_memory returned NULL during warmup — KV cache not cleared"); + } else { + (self.lib.memory_clear)(memory, false); + } } info!("Warmup complete"); @@ -295,6 +295,21 @@ impl InferenceEngine { prompt: &str, sampler_cfg: SamplerConfig, ) -> Result { + // Clear KV cache from any previous session so the new prefill starts at position 0. + // `true` forces a full clear including M-RoPE position metadata; `false` only clears + // cell occupancy flags and leaves position data intact, which causes the + // "inconsistent sequence positions" error on models that use M-RoPE (e.g. Qwen). + unsafe { + let memory = (self.lib.get_memory)(self.ctx.ptr); + if memory.is_null() { + tracing::error!("llama_get_memory returned NULL — cannot clear KV cache, refusing to create session to avoid decode corruption"); + anyhow::bail!("KV cache clear failed: llama_get_memory returned NULL"); + } else { + (self.lib.memory_clear)(memory, true); + tracing::info!("engine.kv_cache_cleared"); + } + } + let vocab = unsafe { (self.lib.model_get_vocab)(self._model.ptr) }; GenerationSession::with_vocab( Arc::clone(&self.lib), diff --git a/crates/inference-core/src/session.rs b/crates/inference-core/src/session.rs index b8d4c04..8488117 100644 --- a/crates/inference-core/src/session.rs +++ b/crates/inference-core/src/session.rs @@ -52,18 +52,6 @@ struct PrefillChunkPlan { unsafe impl Send for GenerationSession {} impl GenerationSession { - /// Tokenize prompt and run prefill decode. Returns a session ready to generate. - pub fn new( - _lib: Arc, - _ctx: *mut llama_context, - _sampler_cfg: SamplerConfig, - ) -> Result { - // Get vocab from context's model (via a workaround: we pass vocab separately) - // For now we'll get it indirectly — see Engine::new_session which passes vocab. - // This constructor signature will be refined; use the vocab variant below. - bail!("Use GenerationSession::with_vocab() instead") - } - /// Full constructor called by InferenceEngine::new_session(). pub fn with_vocab( lib: Arc, @@ -159,8 +147,10 @@ impl GenerationSession { let prefill_ms = prefill_start.elapsed().as_secs_f64() * 1000.0; let prefill_tps = n_prompt_tokens as f64 / (prefill_ms / 1000.0); info!( - "Prefill finished in {:.2} ms ({:.2} t/s)", - prefill_ms, prefill_tps + n_prompt_tokens, + prefill_ms = format_args!("{:.2}", prefill_ms), + prefill_tps = format_args!("{:.2}", prefill_tps), + "session.prefill_complete" ); // create session state @@ -224,6 +214,14 @@ impl GenerationSession { self.n_decoded += 1; self.tokens_generated += 1; + if self.tokens_generated % 100 == 0 { + tracing::debug!( + tokens_generated = self.tokens_generated, + max_tokens = self.max_tokens, + "session.generation_progress" + ); + } + // Detokenize let token_text = self.detokenize_token(new_token); @@ -249,14 +247,18 @@ impl GenerationSession { self.tokens_per_sec = self.tokens_generated as f64 / elapsed; } } + let stop_reason = if self.tokens_generated >= self.max_tokens { + "max_tokens" + } else { + "eos_or_stop_sequence" + }; info!( tokens_generated = self.tokens_generated, - tps = format!("{:.1}", self.tokens_per_sec), - prefill_ms = self - .prefill_ms - .map(|v| format!("{:.1}", v)) - .unwrap_or_default(), - "Generation complete" + max_tokens = self.max_tokens, + tps = self.tokens_per_sec, + prefill_ms = self.prefill_ms.unwrap_or(0.0), + stop_reason, + "session.generation_complete" ); } @@ -347,9 +349,10 @@ fn decode_prefill( if ret != 0 { bail!( - "llama_decode (prefill) returned {} while decoding prompt chunk {:?}", + "llama_decode (prefill) returned {} for chunk {:?} — likely KV cache not cleared or context window ({} tokens) exceeded by prompt", ret, - plan.range + plan.range, + tokens.len() ); } }