diff --git a/crates/skippy-runtime/src/tests.rs b/crates/skippy-runtime/src/tests.rs index 100143e9e..d96a80d1d 100644 --- a/crates/skippy-runtime/src/tests.rs +++ b/crates/skippy-runtime/src/tests.rs @@ -8,7 +8,13 @@ mod tests { GGML_TYPE_F16, ModelInfo, NativeMtpDraft, RuntimeConfig, RuntimeLoadMode, SamplingConfig, StageModel, StageSession, Status, TensorRole, format_skippy_error, }; - use std::{env, path::PathBuf}; + use std::{ + env, + path::PathBuf, + time::{Duration, Instant}, + }; + + const TOOL_CALLS_JSON: &str = r#"[{"type":"function","function":{"name":"execute_bash","description":"Run a command.","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]"#; fn correctness_model() -> Option { env::var_os("SKIPPY_CORRECTNESS_MODEL").map(PathBuf::from) @@ -47,12 +53,19 @@ mod tests { } fn open_correctness_model(model_path: &PathBuf) -> anyhow::Result { + open_correctness_model_with_context(model_path, 256) + } + + fn open_correctness_model_with_context( + model_path: &PathBuf, + ctx_size: u32, + ) -> anyhow::Result { let layer_end = infer_layer_end(model_path)?; let config = RuntimeConfig { stage_index: 0, layer_start: 0, layer_end, - ctx_size: 256, + ctx_size, lane_count: 1, n_batch: None, n_ubatch: None, @@ -74,6 +87,13 @@ mod tests { StageModel::open(model_path, &config) } + fn tool_call_template_options() -> ChatTemplateJsonOptions { + ChatTemplateJsonOptions { + tools_json: Some(TOOL_CALLS_JSON.to_string()), + ..ChatTemplateJsonOptions::default() + } + } + #[test] fn chat_template_applies_when_model_is_configured() -> anyhow::Result<()> { let Some(model_path) = correctness_model() else { @@ -226,13 +246,7 @@ mod tests { let model = open_correctness_model(&model_path)?; let rendered = model.apply_chat_template_json( r#"[{"role":"user","content":"Call execute_bash."}]"#, - ChatTemplateJsonOptions { - tools_json: Some( - r#"[{"type":"function","function":{"name":"execute_bash","description":"Run a command.","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]"# - .to_string(), - ), - ..ChatTemplateJsonOptions::default() - }, + tool_call_template_options(), )?; let metadata: Value = serde_json::from_str(&rendered.metadata_json)?; assert!( @@ -275,11 +289,14 @@ mod tests { prompt_token_count, Some(&sampling), )?; - let serial_predictions = verify_inputs - .iter() - .copied() - .map(|token| serial.decode_step_sampled(token, Some(&sampling))) - .collect::>>()?; + let mut serial_predictions = Vec::with_capacity(verify_inputs.len()); + for (index, token) in verify_inputs.iter().copied().enumerate() { + let predicted = serial.decode_step_sampled(token, Some(&sampling))?; + serial_predictions.push(predicted); + if index + 1 < verify_inputs.len() && predicted != verify_inputs[index + 1] { + break; + } + } let serial_token_count = serial.token_count(); let serial_native_position = serial.native_position()?; drop(serial); @@ -291,15 +308,182 @@ mod tests { prompt_token_count, Some(&sampling), )?; - let batched_predictions = - batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?; - - assert_eq!(batched_predictions, serial_predictions); + let batched_predictions = batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?; + assert_eq!( + batched_predictions, serial_predictions, + "batched verification must stop at the first target mismatch" + ); + batched.trim_session(serial_token_count)?; assert_eq!(batched.token_count(), serial_token_count); assert_eq!(batched.native_position()?, serial_native_position); Ok(()) } + #[test] + fn long_resident_tool_context_preserves_grammar_and_native_mtp_acceptance() -> anyhow::Result<()> + { + const MIN_RESIDENT_TOKENS: usize = 8_192; + const CONTEXT_SIZE: u32 = 10_240; + // Resident prefixes reserve IDs immediately after the active lane IDs. + // A single-lane runtime uses `3`, matching the state-handoff harness. + const RESIDENT_PREFIX_ID: i32 = 3; + + let Some(model_path) = correctness_model() else { + eprintln!("skipping: SKIPPY_CORRECTNESS_MODEL is not set"); + return Ok(()); + }; + let model = open_correctness_model_with_context(&model_path, CONTEXT_SIZE)?; + + let resident_sentence = + "The resident tool context records a completed command result cwd workspace. "; + let sentence_tokens = model.tokenize(resident_sentence, false)?; + assert!( + !sentence_tokens.is_empty(), + "resident context sentence must tokenize" + ); + let mut resident_sentence_count = MIN_RESIDENT_TOKENS / sentence_tokens.len() + 1; + let (rendered, prompt_tokens) = loop { + let resident_context = resident_sentence.repeat(resident_sentence_count); + let rendered = model.apply_chat_template_json( + &format!( + r#"[{{"role":"user","content":"Call execute_bash after this resident context: {resident_context}"}}]"# + ), + tool_call_template_options(), + )?; + let prompt_tokens = model.tokenize(&rendered.prompt, true)?; + if prompt_tokens.len() >= MIN_RESIDENT_TOKENS { + break (rendered, prompt_tokens); + } + let shortfall = MIN_RESIDENT_TOKENS - prompt_tokens.len(); + resident_sentence_count += + shortfall * resident_sentence_count / prompt_tokens.len() + 1; + }; + let metadata: Value = serde_json::from_str(&rendered.metadata_json)?; + assert_eq!( + metadata.get("grammar_lazy").and_then(Value::as_bool), + Some(true), + "tool grammar must wait for its trigger" + ); + + assert!( + prompt_tokens.len() >= MIN_RESIDENT_TOKENS, + "expected at least {MIN_RESIDENT_TOKENS} resident tokens, got {}", + prompt_tokens.len() + ); + assert!( + prompt_tokens.len() < CONTEXT_SIZE as usize, + "resident prompt must leave room for tool sampling" + ); + let prompt_prefix = &prompt_tokens[..prompt_tokens.len() - 1]; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + let last_prompt_token = *prompt_tokens.last().expect("checked nonempty prompt"); + let sampling = SamplingConfig { + enabled: true, + temperature: 0.0, + top_p: 0.95, + top_k: 40, + min_p: 0.05, + ..SamplingConfig::default() + }; + + let mut prefix_owner = model.create_session()?; + prefix_owner.prefill_chunked(prompt_prefix)?; + prefix_owner.save_prefix(RESIDENT_PREFIX_ID, prompt_prefix.len() as u64)?; + drop(prefix_owner); + + let mut verify_inputs = vec![last_prompt_token]; + verify_inputs.extend(model.tokenize("", false)?); + verify_inputs.extend(model.tokenize( + "execute_bashcommandpwd", + false, + )?); + + let mut serial = + model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?; + serial.configure_chat_sampling( + &rendered.metadata_json, + prompt_token_count, + Some(&sampling), + )?; + let mut serial_predictions = Vec::with_capacity(verify_inputs.len()); + for (index, token) in verify_inputs.iter().copied().enumerate() { + let predicted = serial.decode_step_sampled(token, Some(&sampling))?; + serial_predictions.push(predicted); + if index + 1 < verify_inputs.len() && predicted != verify_inputs[index + 1] { + break; + } + } + let serial_token_count = serial.token_count(); + let serial_native_position = serial.native_position()?; + drop(serial); + + let mut batched = + model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?; + batched.configure_chat_sampling( + &rendered.metadata_json, + prompt_token_count, + Some(&sampling), + )?; + let batched_predictions = batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?; + assert_eq!( + batched_predictions, serial_predictions, + "resident-KV verification must stop at the first tool-grammar mismatch" + ); + batched.trim_session(serial_token_count)?; + assert_eq!(batched.token_count(), serial_token_count); + assert_eq!(batched.native_position()?, serial_native_position); + drop(batched); + + let mut native_mtp = + model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?; + native_mtp.configure_chat_sampling( + &rendered.metadata_json, + prompt_token_count, + Some(&sampling), + )?; + let decode_started = Instant::now(); + let (predicted, draft) = + native_mtp.decode_step_sampled_mtp(last_prompt_token, Some(&sampling), 4)?; + let resident_decode_elapsed = decode_started.elapsed(); + assert!( + resident_decode_elapsed < Duration::from_secs(30), + "sampling after an 8k resident KV prefix took {resident_decode_elapsed:?}" + ); + drop(native_mtp); + + if let Some(draft) = draft { + let mut target = + model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?; + target.configure_chat_sampling( + &rendered.metadata_json, + prompt_token_count, + Some(&sampling), + )?; + let mut target_inputs = vec![last_prompt_token, predicted]; + target_inputs.extend(&draft.token_ids); + let target_predictions = + target.verify_tokens_sampled(&target_inputs, Some(&sampling))?; + assert_eq!( + target_predictions.first(), + Some(&predicted), + "resident target decode must agree with the native-MTP source token" + ); + let accepted_draft_tokens = target_predictions + .iter() + .skip(1) + .zip(&draft.token_ids) + .take_while(|(target, draft)| target == draft) + .count(); + assert!( + accepted_draft_tokens > 0, + "native MTP must accept a draft token after the resident tool context; draft={:?}, target={target_predictions:?}", + draft.token_ids + ); + } + + Ok(()) + } + #[test] fn stage_session_exposes_non_frame_native_mtp_decode_api() { type DecodeStepSampledMtp = fn( diff --git a/crates/skippy-server/src/frontend/linear_proposal.rs b/crates/skippy-server/src/frontend/linear_proposal.rs index e0dfa7aa3..ee0702f8a 100644 --- a/crates/skippy-server/src/frontend/linear_proposal.rs +++ b/crates/skippy-server/src/frontend/linear_proposal.rs @@ -171,7 +171,8 @@ pub struct LinearProposalReceipt { pub accepted_proposal_tokens: usize, /// Target tokens committed to the response stream. pub committed_tokens: Box<[i32]>, - /// Prediction observed for every verification row. + /// Authoritative prediction prefix through the full-accept boundary or + /// first mismatch. Rejected branch-conditioned suffixes are not sampled. pub verification_row_predictions: Box<[i32]>, /// Prefix length of row predictions that remained canonical. pub canonical_prediction_count: usize, diff --git a/crates/skippy-server/src/frontend/native_mtp/hybrid.rs b/crates/skippy-server/src/frontend/native_mtp/hybrid.rs index 4f029fba5..685fb9c08 100644 --- a/crates/skippy-server/src/frontend/native_mtp/hybrid.rs +++ b/crates/skippy-server/src/frontend/native_mtp/hybrid.rs @@ -343,17 +343,14 @@ where F: FnMut(i32) -> OpenAiResult, { let required_predictions = proposal_tokens.len().saturating_add(1); - if predicted_tokens.len() < required_predictions { - return Err(OpenAiError::backend(format!( - "native MTP verify window returned too few tokens: got {} expected {}", - predicted_tokens.len(), - required_predictions - ))); - } - let mut accepted_proposal_tokens = 0usize; for (index, proposal_token) in proposal_tokens.iter().enumerate() { - let predicted = predicted_tokens[index]; + let Some(&predicted) = predicted_tokens.get(index) else { + return Err(OpenAiError::backend(format!( + "native MTP verify window ended before a decision: got {} predictions after accepting {accepted_proposal_tokens} proposal tokens", + predicted_tokens.len() + ))); + }; let commit_count = index + 1; if predicted != *proposal_token { return Ok(NativeMtpVerifyWindowDecision { @@ -373,6 +370,13 @@ where } } + if predicted_tokens.len() < required_predictions { + return Err(OpenAiError::backend(format!( + "native MTP verify window omitted the boundary token after accepting every proposal token: got {} expected {required_predictions}", + predicted_tokens.len() + ))); + } + Ok(NativeMtpVerifyWindowDecision { accepted_proposal_tokens, commit_count: required_predictions.min(max_new_tokens.saturating_sub(generated_len)), @@ -633,11 +637,26 @@ mod tests { #[test] fn verify_window_commits_the_target_correction_after_rejection() { let decision = - classify_native_mtp_verify_window(&[11, 12], &[11, 42, 99], 0, 8, |_| Ok(false)) - .unwrap(); + classify_native_mtp_verify_window(&[11, 12], &[11, 42], 0, 8, |_| Ok(false)).unwrap(); assert_eq!(decision.accepted_proposal_tokens, 1); assert_eq!(decision.commit_count, 2); assert!(decision.rejected); } + + #[test] + fn verify_window_rejects_a_prefix_that_ends_before_any_decision() { + let error = + classify_native_mtp_verify_window(&[11, 12], &[11], 0, 8, |_| Ok(false)).unwrap_err(); + + assert!(error.to_string().contains("ended before a decision")); + } + + #[test] + fn verify_window_requires_a_boundary_after_full_acceptance() { + let error = classify_native_mtp_verify_window(&[11, 12], &[11, 12], 0, 8, |_| Ok(false)) + .unwrap_err(); + + assert!(error.to_string().contains("omitted the boundary token")); + } } diff --git a/third_party/llama.cpp/patches/0066-Stop-chat-grammar-verification-at-the-first-mismatch.patch b/third_party/llama.cpp/patches/0066-Stop-chat-grammar-verification-at-the-first-mismatch.patch new file mode 100644 index 000000000..88bac328f --- /dev/null +++ b/third_party/llama.cpp/patches/0066-Stop-chat-grammar-verification-at-the-first-mismatch.patch @@ -0,0 +1,37 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 5 Aug 2026 15:33:46 +1000 +Subject: [PATCH] Stop chat grammar verification at the first mismatch + +--- + src/skippy.cpp | 12 +++++++++++- + 1 file changed, 11 insertions(+), 1 deletion(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 6942f08a4..cfc0bac8f 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -6631,12 +6631,22 @@ enum skippy_status skippy_verify_tokens_frame_sampled( + return SKIPPY_STATUS_RUNTIME_ERROR; + } + session->token_history.resize(history_size_before_verification); ++ const bool stop_after_first_mismatch = session->grammar_sampler != nullptr; ++ bool proposal_mismatched = false; + for (size_t i = 0; i < token_count; ++i) { + skippy_record_tokens(session, &token_ids[i], 1); + const int32_t logits_index = static_cast(i); + output_tokens[i] = skippy_sample_token_ith(session, sampling, logits_index); ++ if (stop_after_first_mismatch && i + 1 < token_count && output_tokens[i] != token_ids[i + 1]) { ++ // Later rows are conditioned on a rejected token and are not ++ // authoritative predictions. Stop advancing the sampler and ++ // grammar state through that rejected suffix. ++ proposal_mismatched = true; ++ *out_token_count = i + 1; ++ break; ++ } + } +- if (out_mtp_draft != nullptr && token_count > 0) { ++ if (out_mtp_draft != nullptr && token_count > 0 && !proposal_mismatched) { + status = skippy_mtp_propose_next( + session, + output_tokens[token_count - 1],