From 48aa9fc4e4de22014dac0be07ec6ee36378a714d Mon Sep 17 00:00:00 2001 From: James Dumay Date: Thu, 18 Jun 2026 19:22:08 +1000 Subject: [PATCH] Add native MTP ngram hybrid proposer --- crates/skippy-prompt/src/prompt_cli/args.rs | 61 ---- .../src/prompt_cli/binary_repl.rs | 24 -- crates/skippy-prompt/src/prompt_cli/draft.rs | 25 -- .../src/prompt_cli/generation.rs | 29 +- crates/skippy-prompt/src/prompt_cli/launch.rs | 37 --- crates/skippy-prompt/src/prompt_cli/mod.rs | 8 - .../skippy-prompt/src/prompt_cli/topology.rs | 13 - .../src/frontend/embedded_generation.rs | 89 +++++- .../src/frontend/native_mtp/decode.rs | 105 +++++++ .../src/frontend/native_mtp/env.rs | 15 + .../src/frontend/native_mtp/hybrid.rs | 278 ++++++++++++++++++ .../src/frontend/native_mtp/mod.rs | 6 + 12 files changed, 482 insertions(+), 208 deletions(-) create mode 100644 crates/skippy-server/src/frontend/native_mtp/hybrid.rs diff --git a/crates/skippy-prompt/src/prompt_cli/args.rs b/crates/skippy-prompt/src/prompt_cli/args.rs index b878f01777..4e51776b19 100644 --- a/crates/skippy-prompt/src/prompt_cli/args.rs +++ b/crates/skippy-prompt/src/prompt_cli/args.rs @@ -30,19 +30,10 @@ impl From for RuntimeLoadMode { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] -#[value(rename_all = "kebab-case")] -pub enum NgramProposalMode { - TransitionPool, - HistoryMatch, -} - #[derive(Parser)] pub struct PromptArgs { #[arg(long, default_value = "target/debug/metrics-server")] pub metrics_server_bin: PathBuf, - #[arg(long, default_value = "target/debug/ngram-pool-server")] - pub ngram_pool_server_bin: PathBuf, #[arg(long, default_value = "target/debug/skippy-server")] pub stage_server_bin: PathBuf, #[arg(long, default_value = "target/debug/skippy-model-package")] @@ -95,34 +86,10 @@ pub struct PromptArgs { pub speculative_window: usize, #[arg(long)] pub adaptive_speculative_window: bool, - #[arg(long)] - pub ngram_speculative: bool, - #[arg(long, value_enum, default_value = "transition-pool")] - pub ngram_proposal_mode: NgramProposalMode, - #[arg(long, default_value_t = 24)] - pub spec_ngram_size_n: usize, - #[arg(long, default_value_t = 1)] - pub ngram_history_min_hits: u32, #[arg(long, default_value_t = 12)] pub draft_min: usize, #[arg(long, default_value_t = 48)] pub draft_max: usize, - #[arg(long, default_value_t = DEFAULT_MIN_WINNER_COUNT)] - pub ngram_min_winner_count: u32, - #[arg(long, default_value_t = DEFAULT_MIN_CONFIDENCE)] - pub ngram_min_confidence: f32, - #[arg(long, default_value_t = DEFAULT_MIN_MARGIN)] - pub ngram_min_margin: u32, - #[arg(long, default_value_t = DEFAULT_CONFIDENCE_STEP)] - pub ngram_confidence_step: f32, - #[arg(long, default_value_t = DEFAULT_CONFIDENCE_STEP_TOKENS)] - pub ngram_confidence_step_tokens: usize, - #[arg(long, default_value_t = DEFAULT_MAX_CONFIDENCE)] - pub ngram_max_confidence: f32, - #[arg(long, default_value_t = DEFAULT_COUNT_STEP_TOKENS)] - pub ngram_count_step_tokens: usize, - #[arg(long, default_value_t = DEFAULT_MARGIN_STEP_TOKENS)] - pub ngram_margin_step_tokens: usize, #[arg(long, default_value = "lookup-record")] pub kv_mode: String, #[arg(long, default_value_t = 512)] @@ -155,8 +122,6 @@ pub struct PromptArgs { pub history_path: Option, #[arg(long)] pub session_id: Option, - #[arg(long)] - pub ngram_pool_uds_path: Option, #[arg(long, default_value_t = 80)] pub log_tail_lines: usize, #[arg(long)] @@ -210,34 +175,10 @@ pub struct BinaryReplArgs { pub speculative_window: usize, #[arg(long)] pub adaptive_speculative_window: bool, - #[arg(long)] - pub ngram_speculative: bool, - #[arg(long, value_enum, default_value = "transition-pool")] - pub ngram_proposal_mode: NgramProposalMode, - #[arg(long, default_value_t = 24)] - pub spec_ngram_size_n: usize, - #[arg(long, default_value_t = 1)] - pub ngram_history_min_hits: u32, #[arg(long, default_value_t = 12)] pub draft_min: usize, #[arg(long, default_value_t = 48)] pub draft_max: usize, - #[arg(long, default_value_t = DEFAULT_MIN_WINNER_COUNT)] - pub ngram_min_winner_count: u32, - #[arg(long, default_value_t = DEFAULT_MIN_CONFIDENCE)] - pub ngram_min_confidence: f32, - #[arg(long, default_value_t = DEFAULT_MIN_MARGIN)] - pub ngram_min_margin: u32, - #[arg(long, default_value_t = DEFAULT_CONFIDENCE_STEP)] - pub ngram_confidence_step: f32, - #[arg(long, default_value_t = DEFAULT_CONFIDENCE_STEP_TOKENS)] - pub ngram_confidence_step_tokens: usize, - #[arg(long, default_value_t = DEFAULT_MAX_CONFIDENCE)] - pub ngram_max_confidence: f32, - #[arg(long, default_value_t = DEFAULT_COUNT_STEP_TOKENS)] - pub ngram_count_step_tokens: usize, - #[arg(long, default_value_t = DEFAULT_MARGIN_STEP_TOKENS)] - pub ngram_margin_step_tokens: usize, #[arg(long, default_value_t = 60)] pub startup_timeout_secs: u64, #[arg(long, default_value_t = 30)] @@ -247,8 +188,6 @@ pub struct BinaryReplArgs { #[arg(long)] pub session_id: Option, #[arg(long)] - pub ngram_pool_uds_path: Option, - #[arg(long)] pub native_logs: bool, #[arg( long, diff --git a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs index 14fdf89287..56a6865301 100644 --- a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs +++ b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs @@ -141,27 +141,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { let default_session_id = args.session_id.clone().unwrap_or_else(default_session_id); let default_wire_session_id = stable_wire_id(&[default_session_id.as_bytes()]); eprintln!("session_id={default_session_id} wire_session_id={default_wire_session_id}"); - let mut ngram = if args.ngram_speculative { - eprintln!( - "ngram speculative enabled: mode={:?} n={} history_min_hits={} draft_min={} draft_max={} min_count={} min_confidence={:.2} min_margin={} confidence_step={:.2}/{} max_confidence={:.2} count_step={} margin_step={}", - args.ngram_proposal_mode, - args.spec_ngram_size_n, - args.ngram_history_min_hits, - args.draft_min, - args.draft_max, - args.ngram_min_winner_count, - args.ngram_min_confidence, - args.ngram_min_margin, - args.ngram_confidence_step, - args.ngram_confidence_step_tokens, - args.ngram_max_confidence, - args.ngram_count_step_tokens, - args.ngram_margin_step_tokens - ); - Some(NgramSource::open(&args, &default_session_id)?) - } else { - None - }; let interrupt = install_prompt_interrupt_handler()?; let mut history = PromptHistory::load(args.history_path.as_deref())?; let mut prompt_input = prompt_input(&history)?; @@ -198,7 +177,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { tokenizer: &tokenizer, chat_template_model: chat_template_model.as_ref(), draft: draft.as_mut(), - ngram: ngram.as_mut(), interrupt: &interrupt, wire_dtype, session_id: &prompt_session_id, @@ -263,7 +241,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { tokenizer: &tokenizer, chat_template_model: chat_template_model.as_ref(), draft: draft.as_mut(), - ngram: ngram.as_mut(), interrupt: &interrupt, wire_dtype, session_id: &default_session_id, @@ -283,7 +260,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { tokenizer: &tokenizer, chat_template_model: chat_template_model.as_ref(), draft: draft.as_mut(), - ngram: ngram.as_mut(), interrupt: &interrupt, wire_dtype, session_id: &default_session_id, diff --git a/crates/skippy-prompt/src/prompt_cli/draft.rs b/crates/skippy-prompt/src/prompt_cli/draft.rs index 90cf501523..21af2addc5 100644 --- a/crates/skippy-prompt/src/prompt_cli/draft.rs +++ b/crates/skippy-prompt/src/prompt_cli/draft.rs @@ -1,28 +1,3 @@ -struct NgramSource; - -impl NgramSource { - fn open(_args: &BinaryReplArgs, _session_id: &str) -> Result { - bail!("ngram speculative prompting is not imported into mesh-llm") - } - - fn observe_sequence(&mut self, _session_id: &str, _tokens: &[i32]) -> Result<()> { - Ok(()) - } - - fn observe_accepted(&mut self, _session_id: &str, _context_tokens: &[i32]) -> Result<()> { - Ok(()) - } - - fn propose( - &mut self, - _session_id: &str, - _context_tokens: &[i32], - _remaining: usize, - ) -> Result> { - Ok(Vec::new()) - } -} - struct DraftRunner { path: PathBuf, window: usize, diff --git a/crates/skippy-prompt/src/prompt_cli/generation.rs b/crates/skippy-prompt/src/prompt_cli/generation.rs index eaf01404fb..73e5437501 100644 --- a/crates/skippy-prompt/src/prompt_cli/generation.rs +++ b/crates/skippy-prompt/src/prompt_cli/generation.rs @@ -3,7 +3,6 @@ struct PromptRun<'a> { tokenizer: &'a StageModel, chat_template_model: Option<&'a StageModel>, draft: Option<&'a mut DraftRunner>, - ngram: Option<&'a mut NgramSource>, interrupt: &'a Arc, wire_dtype: skippy_protocol::binary::WireActivationDType, session_id: &'a str, @@ -19,7 +18,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { tokenizer, chat_template_model, mut draft, - mut ngram, interrupt, wire_dtype, session_id, @@ -280,16 +278,13 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { if let Some(draft) = draft.as_deref_mut() { draft.reset_to_context(&context_tokens)?; } - if let Some(ngram) = ngram.as_deref_mut() { - ngram.observe_sequence(session_id, &context_tokens)?; - } let max_speculative_window = args.speculative_window.max(1); let mut adaptive_window = if args.adaptive_speculative_window { max_speculative_window.min(4) } else { max_speculative_window }; - if draft.is_some() || ngram.is_some() { + if draft.is_some() { speculative_stats.adaptive_window_max = max_speculative_window; speculative_stats.adaptive_window_start = adaptive_window; speculative_stats.adaptive_window_enabled = args.adaptive_speculative_window; @@ -305,19 +300,11 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { let remaining = max_new_tokens - generated.len(); let proposal_limit = remaining.min(adaptive_window); - let draft_tokens = match ngram.as_deref_mut() { - Some(ngram) => ngram.propose(session_id, &context_tokens, proposal_limit)?, - None => Vec::new(), - }; - let draft_tokens = if draft_tokens.is_empty() { - match draft.as_deref_mut() { - Some(draft) if draft.window > 0 => { - draft.propose(current, proposal_limit.min(draft.window))? - } - _ => Vec::new(), + let draft_tokens = match draft.as_deref_mut() { + Some(draft) if draft.window > 0 => { + draft.propose(current, proposal_limit.min(draft.window))? } - } else { - draft_tokens + _ => Vec::new(), }; if draft_tokens.is_empty() { @@ -340,9 +327,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { current = reply.predicted; generated.push(current); context_tokens.push(current); - if let Some(ngram) = ngram.as_deref_mut() { - ngram.observe_accepted(session_id, &context_tokens)?; - } first_time_to_token_ms.get_or_insert_with(|| elapsed_ms(wall_started)); if tokenizer.token_is_eog(current)? { generation_reached_eog = true; @@ -483,9 +467,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { current = predicted; generated.push(current); context_tokens.push(current); - if let Some(ngram) = ngram.as_deref_mut() { - ngram.observe_accepted(session_id, &context_tokens)?; - } if tokenizer.token_is_eog(current)? { reached_eog = true; generation_reached_eog = true; diff --git a/crates/skippy-prompt/src/prompt_cli/launch.rs b/crates/skippy-prompt/src/prompt_cli/launch.rs index 1fb762a279..3fd86e126c 100644 --- a/crates/skippy-prompt/src/prompt_cli/launch.rs +++ b/crates/skippy-prompt/src/prompt_cli/launch.rs @@ -243,30 +243,6 @@ fn prompt_repl_launch(args: PromptArgs) -> Result<()> { ); children.push(ChildGuard::spawn(metrics)?); - let ngram_pool_uds_path = if args.ngram_speculative { - let socket_path = args - .ngram_pool_uds_path - .clone() - .unwrap_or_else(|| run_dir.join("ngram-pool.sock")); - let mut ngram_pool = Command::new(&args.ngram_pool_server_bin); - ngram_pool.args(["serve", "--uds-path", path_str(&socket_path)?]); - configure_process_log(&mut ngram_pool, &run_dir.join("ngram-pool-server.log"))?; - eprintln!( - "launch: starting ngram-pool-server socket={} log={}", - socket_path.display(), - run_dir.join("ngram-pool-server.log").display() - ); - children.push(ChildGuard::spawn(ngram_pool)?); - eprintln!( - "launch: waiting for ngram pool socket {}", - socket_path.display() - ); - wait_for_socket(&socket_path, args.startup_timeout_secs)?; - Some(socket_path) - } else { - None - }; - if remote { rsync_remote_stage_inputs(&args, &stages, &model_package_dir, hf_package_ref)?; eprintln!( @@ -356,20 +332,8 @@ fn prompt_repl_launch(args: PromptArgs) -> Result<()> { draft_model_path: args.draft_model_path, speculative_window: args.speculative_window, adaptive_speculative_window: args.adaptive_speculative_window, - ngram_speculative: args.ngram_speculative, - ngram_proposal_mode: args.ngram_proposal_mode, - spec_ngram_size_n: args.spec_ngram_size_n, - ngram_history_min_hits: args.ngram_history_min_hits, draft_min: args.draft_min, draft_max: args.draft_max, - ngram_min_winner_count: args.ngram_min_winner_count, - ngram_min_confidence: args.ngram_min_confidence, - ngram_min_margin: args.ngram_min_margin, - ngram_confidence_step: args.ngram_confidence_step, - ngram_confidence_step_tokens: args.ngram_confidence_step_tokens, - ngram_max_confidence: args.ngram_max_confidence, - ngram_count_step_tokens: args.ngram_count_step_tokens, - ngram_margin_step_tokens: args.ngram_margin_step_tokens, startup_timeout_secs: args.startup_timeout_secs, decode_timeout_secs: args.decode_timeout_secs, history_path: Some( @@ -377,7 +341,6 @@ fn prompt_repl_launch(args: PromptArgs) -> Result<()> { .unwrap_or_else(|| args.run_root.join("prompt-history.txt")), ), session_id: args.session_id, - ngram_pool_uds_path, native_logs: args.native_logs, raw_prompt: args.raw_prompt, no_think: args.no_think, diff --git a/crates/skippy-prompt/src/prompt_cli/mod.rs b/crates/skippy-prompt/src/prompt_cli/mod.rs index 8fee132d62..c51c8e58b3 100644 --- a/crates/skippy-prompt/src/prompt_cli/mod.rs +++ b/crates/skippy-prompt/src/prompt_cli/mod.rs @@ -41,14 +41,6 @@ use skippy_topology::{ dense_attention_layers, infer_family_capability, plan_contiguous_with_splits, }; -const DEFAULT_MIN_WINNER_COUNT: u32 = 2; -const DEFAULT_MIN_CONFIDENCE: f32 = 0.55; -const DEFAULT_MIN_MARGIN: u32 = 1; -const DEFAULT_CONFIDENCE_STEP: f32 = 0.0; -const DEFAULT_CONFIDENCE_STEP_TOKENS: usize = usize::MAX; -const DEFAULT_MAX_CONFIDENCE: f32 = 0.95; -const DEFAULT_COUNT_STEP_TOKENS: usize = usize::MAX; -const DEFAULT_MARGIN_STEP_TOKENS: usize = usize::MAX; const DEFAULT_MESH_CTX_SIZE: u32 = 4096; const DEFAULT_MESH_PROMPT_MAX_NEW_TOKENS: usize = 0; const PROMPT_EXACT_PREFIX_RESTORE_MIN_TOKENS: usize = 512; diff --git a/crates/skippy-prompt/src/prompt_cli/topology.rs b/crates/skippy-prompt/src/prompt_cli/topology.rs index a3f8254640..d934ab3e55 100644 --- a/crates/skippy-prompt/src/prompt_cli/topology.rs +++ b/crates/skippy-prompt/src/prompt_cli/topology.rs @@ -153,19 +153,6 @@ fn even_stage_ranges(stage_count: usize, layer_end: u32) -> Result Result<()> { - let deadline = Instant::now() + Duration::from_secs(timeout_secs.max(1)); - loop { - if socket_path.exists() { - return Ok(()); - } - if Instant::now() >= deadline { - bail!("timed out waiting for socket {}", socket_path.display()); - } - std::thread::sleep(std::time::Duration::from_millis(100)); - } -} - fn write_json(path: &Path, value: &serde_json::Value) -> Result<()> { fs::write(path, serde_json::to_vec_pretty(value)?) .with_context(|| format!("write {}", path.display())) diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 963464ae42..0bd9b57621 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -860,7 +860,17 @@ impl StageOpenAiBackend { self.telemetry.is_debug_enabled().then(PhaseTimer::start); let native_mtp_draft_token = pending_native_mtp_draft.token; let native_mtp_draft_origin = pending_native_mtp_draft.origin; - let verify_inputs = [current, native_mtp_draft_token]; + let max_native_mtp_proposals = native_mtp_remaining.saturating_sub(1).max(1); + let native_mtp_proposal = NativeMtpHybridProposal::from_anchor( + native_mtp_draft_token, + &context_tokens, + native_mtp_options, + max_native_mtp_proposals, + ); + let verify_inputs = native_mtp_verify_inputs_for_proposals( + current, + native_mtp_proposal.tokens(), + ); let message = embedded_verify_message( request.wire_dtype, VerifySpanMessageArgs { @@ -882,15 +892,15 @@ impl StageOpenAiBackend { &verify_inputs, WireReplyKind::PredictedTokens, )?; - if verify.reply.predicted_tokens.len() < verify_inputs.len() { - return Err(OpenAiError::backend(format!( - "native MTP verify span returned too few tokens: got {} expected {}", - verify.reply.predicted_tokens.len(), - verify_inputs.len() - ))); - } + let native_mtp_batched_decision = classify_native_mtp_batched_verify( + native_mtp_proposal.tokens(), + &verify.reply.predicted_tokens, + decoded_tokens, + request.max_tokens as usize, + |token| token_is_eog_with_runtime(&self.runtime, token), + )?; let target_token = verify.reply.predicted_tokens[0]; - let after_draft_token = verify.reply.predicted_tokens[1]; + let after_draft_token = verify.reply.predicted_tokens.get(1).copied(); let verify_next_mtp_draft = NativeMtpDraft::from_verify_prediction_tokens( &verify.reply.predicted_tokens, verify_inputs.len(), @@ -904,12 +914,24 @@ impl StageOpenAiBackend { matches!(native_mtp_decision, NativeMtpVerification::Accepted { .. }); native_mtp_counters .observe_batched_verification(native_mtp_draft_origin, accepted); - let commit_tokens = [target_token, after_draft_token]; - let commit_token_count = if accepted { 2 } else { 1 }; + native_mtp_counters.observe_hybrid_proposal( + native_mtp_proposal.ngram_span_available(), + native_mtp_proposal.ngram_anchor_agreed(), + native_mtp_proposal.ngram_anchor_disagreed(), + native_mtp_proposal.tokens().len(), + native_mtp_batched_decision.accepted_proposal_tokens, + ); + let commit_token_count = native_mtp_batched_decision.commit_count; let consumed_positions = verify_inputs.len(); let mut committed_positions = 0usize; let mut reached_stop = false; - for token in commit_tokens.into_iter().take(commit_token_count) { + for token in verify + .reply + .predicted_tokens + .iter() + .copied() + .take(commit_token_count) + { current = token; decoded_tokens += 1; committed_positions += 1; @@ -923,7 +945,9 @@ impl StageOpenAiBackend { break; } } - if !accepted && native_mtp_options.reject_cooldown_tokens > 0 { + if native_mtp_batched_decision.rejected + && native_mtp_options.reject_cooldown_tokens > 0 + { native_mtp_reject_cooldown_remaining = native_mtp_options.reject_cooldown_tokens; native_mtp_suppress_cooldown_drafts_remaining = @@ -931,7 +955,7 @@ impl StageOpenAiBackend { native_mtp.clear_pending_draft(); } let verify_next_mtp_draft_available = verify_next_mtp_draft.is_some(); - let verify_next_mtp_draft_adopted = accepted + let verify_next_mtp_draft_adopted = !native_mtp_batched_decision.rejected && committed_positions == consumed_positions && !reached_stop && decoded_tokens < request.max_tokens as usize @@ -951,8 +975,9 @@ impl StageOpenAiBackend { NativeMtpTrimAction::None => {} NativeMtpTrimAction::FullSession => { let target_token_count = prefill_token_count + decoded_tokens; - let defer_trim = - native_mtp_options.defer_reject_trim && !accepted && !reached_stop; + let defer_trim = native_mtp_options.defer_reject_trim + && native_mtp_batched_decision.rejected + && !reached_stop; let trim = if defer_trim { let trim = self.trim_embedded_stage_session_local( &session_key, @@ -1022,6 +1047,38 @@ impl StageOpenAiBackend { "llama_stage.native_mtp.after_draft_token".to_string(), json!(after_draft_token), ); + token_attrs.insert( + "llama_stage.native_mtp.ngram_hybrid".to_string(), + json!(native_mtp_options.ngram_hybrid), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_span_available".to_string(), + json!(native_mtp_proposal.ngram_span_available()), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_anchor_agreement".to_string(), + json!(native_mtp_proposal.ngram_anchor_agreed()), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_anchor_disagreement".to_string(), + json!(native_mtp_proposal.ngram_anchor_disagreed()), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_proposal_len".to_string(), + json!(native_mtp_proposal.tokens().len()), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_accepted_len".to_string(), + json!(native_mtp_batched_decision.accepted_proposal_tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_accepted_tail_len".to_string(), + json!( + native_mtp_batched_decision + .accepted_proposal_tokens + .saturating_sub(1) + ), + ); token_attrs.insert( "llama_stage.native_mtp.verify_next_draft_available".to_string(), json!(verify_next_mtp_draft_available), diff --git a/crates/skippy-server/src/frontend/native_mtp/decode.rs b/crates/skippy-server/src/frontend/native_mtp/decode.rs index 83bdcf441e..9703cc8c03 100644 --- a/crates/skippy-server/src/frontend/native_mtp/decode.rs +++ b/crates/skippy-server/src/frontend/native_mtp/decode.rs @@ -4,6 +4,7 @@ use serde_json::{Value, json}; use super::{ NativeMtpDraftOrigin, native_mtp_batched_verify_enabled, native_mtp_defer_reject_trim_enabled, + native_mtp_ngram_hybrid_enabled, native_mtp_ngram_max_proposal_tokens, native_mtp_ngram_size, native_mtp_reject_cooldown_tokens, native_mtp_suppress_cooldown_draft_limit, native_mtp_suppress_cooldown_drafts_enabled, }; @@ -15,6 +16,9 @@ pub(in crate::frontend) struct NativeMtpDecodeOptions { pub(in crate::frontend) defer_reject_trim: bool, pub(in crate::frontend) suppress_cooldown_drafts: bool, pub(in crate::frontend) suppress_cooldown_draft_limit: usize, + pub(in crate::frontend) ngram_hybrid: bool, + pub(in crate::frontend) ngram_size: usize, + pub(in crate::frontend) ngram_max_proposal_tokens: usize, } impl NativeMtpDecodeOptions { @@ -25,6 +29,9 @@ impl NativeMtpDecodeOptions { defer_reject_trim: native_mtp_defer_reject_trim_enabled(), suppress_cooldown_drafts: native_mtp_suppress_cooldown_drafts_enabled(), suppress_cooldown_draft_limit: native_mtp_suppress_cooldown_draft_limit(), + ngram_hybrid: native_mtp_ngram_hybrid_enabled(), + ngram_size: native_mtp_ngram_size(), + ngram_max_proposal_tokens: native_mtp_ngram_max_proposal_tokens(), } } } @@ -43,6 +50,13 @@ pub(in crate::frontend) struct NativeMtpDecodeCounters { verify_next_draft_adopted_count: usize, deferred_reject_trim_count: usize, deferred_reject_trim_local_ms: f64, + hybrid_anchor_available_count: usize, + hybrid_ngram_span_available_count: usize, + hybrid_anchor_agreement_count: usize, + hybrid_anchor_disagreement_count: usize, + hybrid_proposal_token_count: usize, + hybrid_accepted_token_count: usize, + hybrid_accepted_tail_token_count: usize, } impl NativeMtpDecodeCounters { @@ -100,6 +114,29 @@ impl NativeMtpDecodeCounters { self.deferred_reject_trim_local_ms += local_ms; } + pub(in crate::frontend) fn observe_hybrid_proposal( + &mut self, + ngram_span_available: bool, + ngram_anchor_agreed: bool, + ngram_anchor_disagreed: bool, + proposal_token_count: usize, + accepted_token_count: usize, + ) { + self.hybrid_anchor_available_count += 1; + if ngram_span_available { + self.hybrid_ngram_span_available_count += 1; + } + if ngram_anchor_agreed { + self.hybrid_anchor_agreement_count += 1; + } + if ngram_anchor_disagreed { + self.hybrid_anchor_disagreement_count += 1; + } + self.hybrid_proposal_token_count += proposal_token_count; + self.hybrid_accepted_token_count += accepted_token_count; + self.hybrid_accepted_tail_token_count += accepted_token_count.saturating_sub(1); + } + pub(in crate::frontend) fn insert_summary_attrs( &self, attrs: &mut BTreeMap, @@ -121,6 +158,18 @@ impl NativeMtpDecodeCounters { "llama_stage.native_mtp.suppress_cooldown_draft_limit".to_string(), json!(options.suppress_cooldown_draft_limit), ); + attrs.insert( + "llama_stage.native_mtp.ngram_hybrid".to_string(), + json!(options.ngram_hybrid), + ); + attrs.insert( + "llama_stage.native_mtp.ngram_size".to_string(), + json!(options.ngram_size), + ); + attrs.insert( + "llama_stage.native_mtp.ngram_max_proposal_tokens".to_string(), + json!(options.ngram_max_proposal_tokens), + ); attrs.insert( "llama_stage.native_mtp.suppressed_cooldown_draft_count".to_string(), json!(self.suppressed_cooldown_draft_count), @@ -169,6 +218,34 @@ impl NativeMtpDecodeCounters { "llama_stage.native_mtp.deferred_reject_trim_local_ms".to_string(), json!(self.deferred_reject_trim_local_ms), ); + attrs.insert( + "llama_stage.native_mtp.hybrid_anchor_available_count".to_string(), + json!(self.hybrid_anchor_available_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_span_available_count".to_string(), + json!(self.hybrid_ngram_span_available_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_anchor_agreement_count".to_string(), + json!(self.hybrid_anchor_agreement_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_anchor_disagreement_count".to_string(), + json!(self.hybrid_anchor_disagreement_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_proposal_token_count".to_string(), + json!(self.hybrid_proposal_token_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_accepted_token_count".to_string(), + json!(self.hybrid_accepted_token_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_accepted_tail_token_count".to_string(), + json!(self.hybrid_accepted_tail_token_count), + ); } } @@ -186,6 +263,7 @@ mod tests { counters.observe_verify_next_draft(true, true); counters.observe_suppressed_cooldown_draft(); counters.observe_deferred_reject_trim(1.25); + counters.observe_hybrid_proposal(true, true, false, 4, 3); let mut attrs = BTreeMap::new(); counters.insert_summary_attrs( @@ -196,6 +274,9 @@ mod tests { defer_reject_trim: true, suppress_cooldown_drafts: false, suppress_cooldown_draft_limit: 2, + ngram_hybrid: true, + ngram_size: 8, + ngram_max_proposal_tokens: 4, }, ); @@ -243,5 +324,29 @@ mod tests { attrs.get("llama_stage.native_mtp.deferred_reject_trim_local_ms"), Some(&json!(1.25)) ); + assert_eq!( + attrs.get("llama_stage.native_mtp.ngram_hybrid"), + Some(&json!(true)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_anchor_available_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_ngram_span_available_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_anchor_agreement_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_proposal_token_count"), + Some(&json!(4)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_accepted_tail_token_count"), + Some(&json!(2)) + ); } } diff --git a/crates/skippy-server/src/frontend/native_mtp/env.rs b/crates/skippy-server/src/frontend/native_mtp/env.rs index 54fbb698d0..f7d617880b 100644 --- a/crates/skippy-server/src/frontend/native_mtp/env.rs +++ b/crates/skippy-server/src/frontend/native_mtp/env.rs @@ -3,6 +3,9 @@ const REJECT_COOLDOWN_TOKENS_ENV: &str = "SKIPPY_NATIVE_MTP_REJECT_COOLDOWN_TOKE const DEFER_REJECT_TRIM_ENV: &str = "SKIPPY_NATIVE_MTP_DEFER_REJECT_TRIM"; const SUPPRESS_COOLDOWN_DRAFTS_ENV: &str = "SKIPPY_NATIVE_MTP_SUPPRESS_COOLDOWN_DRAFTS"; const SUPPRESS_COOLDOWN_DRAFT_LIMIT_ENV: &str = "SKIPPY_NATIVE_MTP_SUPPRESS_COOLDOWN_DRAFT_LIMIT"; +const NGRAM_HYBRID_ENV: &str = "SKIPPY_NATIVE_MTP_NGRAM_HYBRID"; +const NGRAM_SIZE_ENV: &str = "SKIPPY_NATIVE_MTP_NGRAM_SIZE"; +const NGRAM_MAX_PROPOSAL_TOKENS_ENV: &str = "SKIPPY_NATIVE_MTP_NGRAM_MAX_PROPOSAL_TOKENS"; pub(in crate::frontend) fn native_mtp_batched_verify_enabled() -> bool { native_mtp_batched_verify_enabled_from(std::env::var(BATCHED_VERIFY_ENV).ok().as_deref()) @@ -24,6 +27,18 @@ pub(in crate::frontend) fn native_mtp_suppress_cooldown_draft_limit() -> usize { parse_usize_env(SUPPRESS_COOLDOWN_DRAFT_LIMIT_ENV, 0) } +pub(in crate::frontend) fn native_mtp_ngram_hybrid_enabled() -> bool { + truthy_env(std::env::var(NGRAM_HYBRID_ENV).ok().as_deref()) +} + +pub(in crate::frontend) fn native_mtp_ngram_size() -> usize { + parse_usize_env(NGRAM_SIZE_ENV, 8) +} + +pub(in crate::frontend) fn native_mtp_ngram_max_proposal_tokens() -> usize { + parse_usize_env(NGRAM_MAX_PROPOSAL_TOKENS_ENV, 4) +} + fn native_mtp_batched_verify_enabled_from(value: Option<&str>) -> bool { !falsey_env(value) } diff --git a/crates/skippy-server/src/frontend/native_mtp/hybrid.rs b/crates/skippy-server/src/frontend/native_mtp/hybrid.rs new file mode 100644 index 0000000000..a1c6e71c88 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/hybrid.rs @@ -0,0 +1,278 @@ +use openai_frontend::{OpenAiError, OpenAiResult}; + +use super::NativeMtpDecodeOptions; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpHybridProposal { + tokens: Vec, + ngram_span_available: bool, + ngram_anchor_agreed: bool, + ngram_anchor_disagreed: bool, +} + +impl NativeMtpHybridProposal { + pub(in crate::frontend) fn from_anchor( + anchor: i32, + context_tokens: &[i32], + options: NativeMtpDecodeOptions, + max_proposal_tokens: usize, + ) -> Self { + let max_proposal_tokens = max_proposal_tokens.max(1); + if !options.ngram_hybrid || options.ngram_max_proposal_tokens == 0 { + return Self::anchor_only(anchor); + } + + let proposal_limit = max_proposal_tokens.min(options.ngram_max_proposal_tokens); + let ngram_tokens = + ngram_history_proposal(context_tokens, options.ngram_size, proposal_limit); + let ngram_span_available = !ngram_tokens.is_empty(); + let ngram_anchor_agreed = ngram_tokens.first().is_some_and(|token| *token == anchor); + let ngram_anchor_disagreed = ngram_span_available && !ngram_anchor_agreed; + if ngram_anchor_agreed { + return Self { + tokens: ngram_tokens, + ngram_span_available, + ngram_anchor_agreed, + ngram_anchor_disagreed, + }; + } + + Self { + tokens: vec![anchor], + ngram_span_available, + ngram_anchor_agreed, + ngram_anchor_disagreed, + } + } + + pub(in crate::frontend) fn tokens(&self) -> &[i32] { + &self.tokens + } + + pub(in crate::frontend) fn ngram_span_available(&self) -> bool { + self.ngram_span_available + } + + pub(in crate::frontend) fn ngram_anchor_agreed(&self) -> bool { + self.ngram_anchor_agreed + } + + pub(in crate::frontend) fn ngram_anchor_disagreed(&self) -> bool { + self.ngram_anchor_disagreed + } + + fn anchor_only(anchor: i32) -> Self { + Self { + tokens: vec![anchor], + ngram_span_available: false, + ngram_anchor_agreed: false, + ngram_anchor_disagreed: false, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpBatchedDecision { + pub(in crate::frontend) accepted_proposal_tokens: usize, + pub(in crate::frontend) commit_count: usize, + pub(in crate::frontend) rejected: bool, +} + +pub(in crate::frontend) fn native_mtp_verify_inputs_for_proposals( + current: i32, + proposals: &[i32], +) -> Vec { + let mut tokens = Vec::with_capacity(proposals.len().saturating_add(1)); + tokens.push(current); + tokens.extend_from_slice(proposals); + tokens +} + +pub(in crate::frontend) fn classify_native_mtp_batched_verify( + proposal_tokens: &[i32], + predicted_tokens: &[i32], + generated_len: usize, + max_new_tokens: usize, + mut token_is_eog: F, +) -> OpenAiResult +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 span 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 commit_count = index + 1; + let accepted = predicted == *proposal_token; + let reached_eog = token_is_eog(predicted)?; + let reached_limit = generated_len + commit_count >= max_new_tokens; + if !accepted { + return Ok(NativeMtpBatchedDecision { + accepted_proposal_tokens, + commit_count, + rejected: true, + }); + } + + accepted_proposal_tokens += 1; + if reached_eog || reached_limit { + return Ok(NativeMtpBatchedDecision { + accepted_proposal_tokens, + commit_count, + rejected: false, + }); + } + } + + let extra_commit_count = proposal_tokens.len().saturating_add(1); + Ok(NativeMtpBatchedDecision { + accepted_proposal_tokens, + commit_count: extra_commit_count.min(max_new_tokens.saturating_sub(generated_len)), + rejected: false, + }) +} + +fn ngram_history_proposal( + context_tokens: &[i32], + ngram_size: usize, + max_tokens: usize, +) -> Vec { + if ngram_size == 0 || max_tokens == 0 || context_tokens.len() <= ngram_size { + return Vec::new(); + } + + let suffix_start = context_tokens.len() - ngram_size; + let suffix = &context_tokens[suffix_start..]; + for candidate_start in (0..suffix_start).rev() { + let candidate_end = candidate_start + ngram_size; + if &context_tokens[candidate_start..candidate_end] != suffix { + continue; + } + let proposal_start = candidate_end; + let proposal_end = proposal_start + .saturating_add(max_tokens) + .min(context_tokens.len()); + if proposal_start < proposal_end { + return context_tokens[proposal_start..proposal_end].to_vec(); + } + } + + Vec::new() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn options() -> NativeMtpDecodeOptions { + NativeMtpDecodeOptions { + batched_verify: true, + reject_cooldown_tokens: 0, + defer_reject_trim: false, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 2, + ngram_max_proposal_tokens: 4, + } + } + + #[test] + fn hybrid_extends_when_ngram_first_token_matches_anchor() { + let proposal = + NativeMtpHybridProposal::from_anchor(3, &[1, 2, 3, 4, 5, 1, 2], options(), 4); + + assert_eq!(proposal.tokens(), &[3, 4, 5, 1]); + assert!(proposal.ngram_span_available()); + assert!(proposal.ngram_anchor_agreed()); + assert!(!proposal.ngram_anchor_disagreed()); + } + + #[test] + fn hybrid_keeps_only_anchor_when_ngram_first_token_disagrees() { + let proposal = + NativeMtpHybridProposal::from_anchor(9, &[1, 2, 3, 4, 5, 1, 2], options(), 4); + + assert_eq!(proposal.tokens(), &[9]); + assert!(proposal.ngram_span_available()); + assert!(!proposal.ngram_anchor_agreed()); + assert!(proposal.ngram_anchor_disagreed()); + } + + #[test] + fn hybrid_keeps_anchor_only_when_disabled() { + let mut options = options(); + options.ngram_hybrid = false; + + let proposal = NativeMtpHybridProposal::from_anchor(3, &[1, 2, 3, 4, 1, 2], options, 4); + + assert_eq!(proposal.tokens(), &[3]); + assert!(!proposal.ngram_span_available()); + } + + #[test] + fn verify_inputs_include_current_and_all_proposals() { + assert_eq!( + native_mtp_verify_inputs_for_proposals(10, &[11, 12, 13]), + vec![10, 11, 12, 13] + ); + } + + #[test] + fn classify_commits_extra_target_after_full_accept() { + let decision = + classify_native_mtp_batched_verify(&[11, 12], &[11, 12, 13], 0, 8, |_| Ok(false)) + .unwrap(); + + assert_eq!( + decision, + NativeMtpBatchedDecision { + accepted_proposal_tokens: 2, + commit_count: 3, + rejected: false, + } + ); + } + + #[test] + fn classify_commits_rejected_target_and_trims_rest() { + let decision = + classify_native_mtp_batched_verify(&[11, 12, 13], &[11, 42, 99, 100], 0, 8, |_| { + Ok(false) + }) + .unwrap(); + + assert_eq!( + decision, + NativeMtpBatchedDecision { + accepted_proposal_tokens: 1, + commit_count: 2, + rejected: true, + } + ); + } + + #[test] + fn classify_stops_without_extra_target_at_generation_limit() { + let decision = + classify_native_mtp_batched_verify(&[11, 12], &[11, 12, 13], 0, 2, |_| Ok(false)) + .unwrap(); + + assert_eq!( + decision, + NativeMtpBatchedDecision { + accepted_proposal_tokens: 2, + commit_count: 2, + rejected: false, + } + ); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/mod.rs b/crates/skippy-server/src/frontend/native_mtp/mod.rs index 589aeadc1b..7ce77b6d58 100644 --- a/crates/skippy-server/src/frontend/native_mtp/mod.rs +++ b/crates/skippy-server/src/frontend/native_mtp/mod.rs @@ -1,6 +1,7 @@ mod decode; mod draft; mod env; +mod hybrid; mod stats; mod trim; mod verifier; @@ -9,9 +10,14 @@ pub(super) use decode::{NativeMtpDecodeCounters, NativeMtpDecodeOptions}; pub(super) use draft::{NativeMtpDraft, NativeMtpDraftOrigin, PendingNativeMtpDraft}; pub(in crate::frontend) use env::{ native_mtp_batched_verify_enabled, native_mtp_defer_reject_trim_enabled, + native_mtp_ngram_hybrid_enabled, native_mtp_ngram_max_proposal_tokens, native_mtp_ngram_size, native_mtp_reject_cooldown_tokens, native_mtp_suppress_cooldown_draft_limit, native_mtp_suppress_cooldown_drafts_enabled, }; +pub(super) use hybrid::{ + NativeMtpHybridProposal, classify_native_mtp_batched_verify, + native_mtp_verify_inputs_for_proposals, +}; pub(super) use stats::{NativeMtpN1Stats, NativeMtpVerification}; pub(super) use trim::{NativeMtpTrimAction, native_mtp_trim_action}; pub(super) use verifier::NativeMtpN1Verifier;