diff --git a/Cargo.lock b/Cargo.lock index e5757e33f7..181f591d97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4046,6 +4046,7 @@ dependencies = [ "serial_test", "sha2 0.10.9", "skippy-coordinator", + "skippy-ffi", "skippy-protocol", "skippy-runtime", "skippy-server", diff --git a/crates/mesh-llm-cli/src/lib.rs b/crates/mesh-llm-cli/src/lib.rs index c31cc27609..ed952c7dd6 100644 --- a/crates/mesh-llm-cli/src/lib.rs +++ b/crates/mesh-llm-cli/src/lib.rs @@ -12,6 +12,7 @@ pub use mesh_llm_events::LogFormat; pub use parser::{ AuthCommand, BinaryFlavor, Cli, Command, ConfigCommand, DiscoveryScope, DoctorCommand, GpuCommand, MeshDiscoveryMode, MeshGuardrailCliMode, NormalizedRuntimeArgs, PluginCommand, - RuntimeSurface, SkillAgentArg, SkillCommand, TrustCommand, TrustPolicy, - legacy_runtime_surface_warning, normalize_runtime_surface_args, validate_discovery_mode_args, + RuntimeSurface, SkillAgentArg, SkillCommand, SpeculativeNgramProposerCli, TrustCommand, + TrustPolicy, legacy_runtime_surface_warning, normalize_runtime_surface_args, + validate_discovery_mode_args, }; diff --git a/crates/mesh-llm-cli/src/parser.rs b/crates/mesh-llm-cli/src/parser.rs index 80d06db13d..a74ef94514 100644 --- a/crates/mesh-llm-cli/src/parser.rs +++ b/crates/mesh-llm-cli/src/parser.rs @@ -395,6 +395,21 @@ pub enum MeshGuardrailCliMode { Enforce, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +pub enum SpeculativeNgramProposerCli { + Simple, + Cache, +} + +impl SpeculativeNgramProposerCli { + pub fn as_str(self) -> &'static str { + match self { + Self::Simple => "simple", + Self::Cache => "cache", + } + } +} + impl MeshGuardrailCliMode { pub const fn as_str(self) -> &'static str { match self { @@ -547,6 +562,70 @@ pub struct Cli { #[arg(long, hide = true)] pub no_draft: bool, + /// Override the package speculative decoding strategy for this invocation. + #[arg(long, hide = true)] + pub speculative_strategy: Option, + + /// Override the N-gram proposer kind for this invocation. + #[arg(long, value_enum, hide = true)] + pub speculative_ngram_proposer: Option, + + /// Minimum matching N-gram length for a direct N-gram proposer. + #[arg(long, hide = true)] + pub speculative_ngram_min: Option, + + /// Maximum matching N-gram length for a direct N-gram proposer. + #[arg(long, hide = true)] + pub speculative_ngram_max: Option, + + /// Cap N-gram tokens proposed in one verify window. + #[arg(long, hide = true)] + pub speculative_ngram_max_proposal_tokens: Option, + + /// Initial N-gram extension length for a composite MTP strategy. + #[arg(long, hide = true)] + pub speculative_extension_initial_tokens: Option, + + /// Maximum N-gram extension length for a composite MTP strategy. + #[arg(long, hide = true)] + pub speculative_extension_max_tokens: Option, + + /// Consecutive weak extensions before the composite strategy backs off. + #[arg(long, hide = true)] + pub speculative_extension_tail_backoff_proposals: Option, + + /// Native MTP rejection cooldown in generated tokens. + #[arg(long, hide = true)] + pub speculative_native_mtp_reject_cooldown_tokens: Option, + + /// Suppress native MTP drafts while its rejection cooldown is active. + #[arg(long, hide = true)] + pub speculative_native_mtp_suppress_cooldown_drafts: bool, + + /// Keep native MTP drafts during its rejection cooldown. + #[arg( + long, + hide = true, + conflicts_with = "speculative_native_mtp_suppress_cooldown_drafts" + )] + pub speculative_native_mtp_allow_cooldown_drafts: bool, + + /// Maximum native MTP drafts suppressed by a cooldown. + #[arg(long, hide = true)] + pub speculative_native_mtp_suppress_cooldown_draft_limit: Option, + + /// Minimum tokens to include in a pipelined verify window. + #[arg(long, hide = true)] + pub speculative_verify_window_min_tokens: Option, + + /// Maximum tokens to include in a pipelined verify window. + #[arg(long, hide = true)] + pub speculative_verify_window_max_tokens: Option, + + /// Number of in-flight pipelined verify windows. + #[arg(long, hide = true)] + pub speculative_verify_window_pipeline_depth: Option, + /// Force tensor split even if the model fits on one node. #[arg(long, hide = true)] pub split: bool, @@ -1508,6 +1587,39 @@ mod tests { ); } + #[test] + fn serve_parses_speculative_decode_overrides() { + let normalized = normalize_runtime_surface_args([ + "mesh-llm", + "serve", + "--speculative-strategy", + "mtp-cache", + "--speculative-ngram-proposer", + "cache", + "--speculative-ngram-min", + "2", + "--speculative-ngram-max", + "6", + "--speculative-extension-max-tokens", + "8", + "--speculative-native-mtp-allow-cooldown-drafts", + "--speculative-verify-window-pipeline-depth", + "3", + ]); + let cli = Cli::try_parse_from(normalized.normalized).expect("clap parse"); + + assert_eq!(cli.speculative_strategy.as_deref(), Some("mtp-cache")); + assert_eq!( + cli.speculative_ngram_proposer, + Some(SpeculativeNgramProposerCli::Cache) + ); + assert_eq!(cli.speculative_ngram_min, Some(2)); + assert_eq!(cli.speculative_ngram_max, Some(6)); + assert_eq!(cli.speculative_extension_max_tokens, Some(8)); + assert!(cli.speculative_native_mtp_allow_cooldown_drafts); + assert_eq!(cli.speculative_verify_window_pipeline_depth, Some(3)); + } + #[test] fn legacy_runtime_surface_warning_for_top_level_serve_flags() { let normalized = diff --git a/crates/mesh-llm-config/src/lib.rs b/crates/mesh-llm-config/src/lib.rs index ffcce98170..12ea2ab8f9 100644 --- a/crates/mesh-llm-config/src/lib.rs +++ b/crates/mesh-llm-config/src/lib.rs @@ -38,8 +38,8 @@ pub use validate::{ mod tests { use super::{ ConfigStore, GpuAssignment, LocalServingNodeConfig, MeshConfig, ModelRuntimeKind, - built_in_config_schema, canonicalize_built_in_config_identifier, parse_config_toml, - validate_config, + SpeculativeConfig, built_in_config_schema, canonicalize_built_in_config_identifier, + parse_config_toml, validate_config, }; use std::collections::{BTreeMap, BTreeSet}; use std::fs; @@ -55,6 +55,31 @@ mod tests { assert!(config.models.is_empty()); } + #[test] + fn speculative_config_precedence_keeps_lower_layer_fields() { + let defaults = SpeculativeConfig { + strategy: Some("mtp-cache".to_string()), + verify_window_pipeline_depth: Some(2), + ..Default::default() + }; + let model = SpeculativeConfig { + ngram_max_proposal_tokens: Some(6), + ..Default::default() + }; + let overrides = SpeculativeConfig { + strategy: Some("mtp".to_string()), + verify_window_pipeline_depth: Some(3), + ..Default::default() + }; + + let resolved = + SpeculativeConfig::with_precedence(Some(&overrides), Some(&model), Some(&defaults)); + + assert_eq!(resolved.strategy.as_deref(), Some("mtp")); + assert_eq!(resolved.ngram_max_proposal_tokens, Some(6)); + assert_eq!(resolved.verify_window_pipeline_depth, Some(3)); + } + #[test] fn plugin_startup_config_round_trips_from_toml() { let config: MeshConfig = toml::from_str( @@ -731,7 +756,6 @@ gpu_id = "pci:0000:65:00.0" "models", "plugins", "settings", - "strategy", ]; let mut total = 0usize; diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index c02a5d2969..b5a326d2ef 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -559,10 +559,80 @@ pub struct SpeculativeConfig { pub draft_cache_type_v: Option, pub ngram_min: Option, pub ngram_max: Option, + pub ngram_proposer: Option, + pub ngram_max_proposal_tokens: Option, + pub extension_initial_tokens: Option, + pub extension_max_tokens: Option, + pub extension_tail_backoff_proposals: Option, + pub native_mtp_reject_cooldown_tokens: Option, + pub native_mtp_suppress_cooldown_drafts: Option, + pub native_mtp_suppress_cooldown_draft_limit: Option, + pub verify_window_min_tokens: Option, + pub verify_window_max_tokens: Option, + pub verify_window_pipeline_depth: Option, pub spec_default: Option, pub(crate) legacy_draft_model_path_used: bool, } +impl SpeculativeConfig { + /// Resolves the three supported policy layers without discarding fields + /// that are not overridden by a more specific layer. + pub fn with_precedence( + overrides: Option<&Self>, + model: Option<&Self>, + defaults: Option<&Self>, + ) -> Self { + macro_rules! pick { + ($field:ident) => { + overrides + .and_then(|config| config.$field.clone()) + .or_else(|| model.and_then(|config| config.$field.clone())) + .or_else(|| defaults.and_then(|config| config.$field.clone())) + }; + } + + Self { + strategy: pick!(strategy), + mode: pick!(mode), + draft_model: pick!(draft_model), + draft_hf_repo: pick!(draft_hf_repo), + draft_hf_file: pick!(draft_hf_file), + draft_selection_policy: pick!(draft_selection_policy), + pairing_fault: pick!(pairing_fault), + draft_max_tokens: pick!(draft_max_tokens), + draft_min_tokens: pick!(draft_min_tokens), + draft_acceptance_threshold: pick!(draft_acceptance_threshold), + draft_split_probability: pick!(draft_split_probability), + draft_gpu_layers: pick!(draft_gpu_layers), + draft_device: pick!(draft_device), + draft_threads: pick!(draft_threads), + draft_cache_type_k: pick!(draft_cache_type_k), + draft_cache_type_v: pick!(draft_cache_type_v), + ngram_min: pick!(ngram_min), + ngram_max: pick!(ngram_max), + ngram_proposer: pick!(ngram_proposer), + ngram_max_proposal_tokens: pick!(ngram_max_proposal_tokens), + extension_initial_tokens: pick!(extension_initial_tokens), + extension_max_tokens: pick!(extension_max_tokens), + extension_tail_backoff_proposals: pick!(extension_tail_backoff_proposals), + native_mtp_reject_cooldown_tokens: pick!(native_mtp_reject_cooldown_tokens), + native_mtp_suppress_cooldown_drafts: pick!(native_mtp_suppress_cooldown_drafts), + native_mtp_suppress_cooldown_draft_limit: pick!( + native_mtp_suppress_cooldown_draft_limit + ), + verify_window_min_tokens: pick!(verify_window_min_tokens), + verify_window_max_tokens: pick!(verify_window_max_tokens), + verify_window_pipeline_depth: pick!(verify_window_pipeline_depth), + spec_default: pick!(spec_default), + legacy_draft_model_path_used: overrides + .filter(|config| config.draft_model.is_some()) + .or_else(|| model.filter(|config| config.draft_model.is_some())) + .or_else(|| defaults.filter(|config| config.draft_model.is_some())) + .is_some_and(|config| config.legacy_draft_model_path_used), + } + } +} + /// Raw deserialization helper that accepts both `draft_model` and the legacy /// `draft_model_path` key. The public `SpeculativeConfig` is constructed from /// this after detecting which key was used. @@ -608,6 +678,28 @@ struct SpeculativeConfigRaw { #[serde(default)] ngram_max: Option, #[serde(default)] + ngram_proposer: Option, + #[serde(default)] + ngram_max_proposal_tokens: Option, + #[serde(default)] + extension_initial_tokens: Option, + #[serde(default)] + extension_max_tokens: Option, + #[serde(default)] + extension_tail_backoff_proposals: Option, + #[serde(default)] + native_mtp_reject_cooldown_tokens: Option, + #[serde(default)] + native_mtp_suppress_cooldown_drafts: Option, + #[serde(default)] + native_mtp_suppress_cooldown_draft_limit: Option, + #[serde(default)] + verify_window_min_tokens: Option, + #[serde(default)] + verify_window_max_tokens: Option, + #[serde(default)] + verify_window_pipeline_depth: Option, + #[serde(default)] spec_default: Option, } @@ -643,6 +735,17 @@ impl<'de> Deserialize<'de> for SpeculativeConfig { draft_cache_type_v: raw.draft_cache_type_v, ngram_min: raw.ngram_min, ngram_max: raw.ngram_max, + ngram_proposer: raw.ngram_proposer, + ngram_max_proposal_tokens: raw.ngram_max_proposal_tokens, + extension_initial_tokens: raw.extension_initial_tokens, + extension_max_tokens: raw.extension_max_tokens, + extension_tail_backoff_proposals: raw.extension_tail_backoff_proposals, + native_mtp_reject_cooldown_tokens: raw.native_mtp_reject_cooldown_tokens, + native_mtp_suppress_cooldown_drafts: raw.native_mtp_suppress_cooldown_drafts, + native_mtp_suppress_cooldown_draft_limit: raw.native_mtp_suppress_cooldown_draft_limit, + verify_window_min_tokens: raw.verify_window_min_tokens, + verify_window_max_tokens: raw.verify_window_max_tokens, + verify_window_pipeline_depth: raw.verify_window_pipeline_depth, spec_default: raw.spec_default, legacy_draft_model_path_used: legacy_used, }) @@ -656,7 +759,7 @@ impl Serialize for SpeculativeConfig { { use serde::ser::SerializeMap; - let mut map = serializer.serialize_map(Some(21))?; + let mut map = serializer.serialize_map(Some(32))?; map.serialize_entry("strategy", &self.strategy)?; map.serialize_entry("mode", &self.mode)?; if self.legacy_draft_model_path_used { @@ -684,6 +787,32 @@ impl Serialize for SpeculativeConfig { map.serialize_entry("draft_cache_type_v", &self.draft_cache_type_v)?; map.serialize_entry("ngram_min", &self.ngram_min)?; map.serialize_entry("ngram_max", &self.ngram_max)?; + map.serialize_entry("ngram_proposer", &self.ngram_proposer)?; + map.serialize_entry("ngram_max_proposal_tokens", &self.ngram_max_proposal_tokens)?; + map.serialize_entry("extension_initial_tokens", &self.extension_initial_tokens)?; + map.serialize_entry("extension_max_tokens", &self.extension_max_tokens)?; + map.serialize_entry( + "extension_tail_backoff_proposals", + &self.extension_tail_backoff_proposals, + )?; + map.serialize_entry( + "native_mtp_reject_cooldown_tokens", + &self.native_mtp_reject_cooldown_tokens, + )?; + map.serialize_entry( + "native_mtp_suppress_cooldown_drafts", + &self.native_mtp_suppress_cooldown_drafts, + )?; + map.serialize_entry( + "native_mtp_suppress_cooldown_draft_limit", + &self.native_mtp_suppress_cooldown_draft_limit, + )?; + map.serialize_entry("verify_window_min_tokens", &self.verify_window_min_tokens)?; + map.serialize_entry("verify_window_max_tokens", &self.verify_window_max_tokens)?; + map.serialize_entry( + "verify_window_pipeline_depth", + &self.verify_window_pipeline_depth, + )?; map.serialize_entry("spec_default", &self.spec_default)?; map.end() } diff --git a/crates/mesh-llm-config/src/model/built_in_schema.rs b/crates/mesh-llm-config/src/model/built_in_schema.rs index f4d6dff0b6..e8c8cc1d34 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema.rs @@ -603,6 +603,7 @@ fn skippy_settings(prefix: &str) -> Vec { fn speculative_settings(prefix: &str) -> Vec { vec![ + basic_setting(&format!("{prefix}.strategy"), ConfigValueSchema::String), basic_setting( &format!("{prefix}.mode"), string_enum(["auto", "disabled", "draft", "ngram"]), @@ -665,6 +666,50 @@ fn speculative_settings(prefix: &str) -> Vec { ), basic_setting(&format!("{prefix}.ngram_min"), ConfigValueSchema::Integer), basic_setting(&format!("{prefix}.ngram_max"), ConfigValueSchema::Integer), + basic_setting( + &format!("{prefix}.ngram_proposer"), + string_enum(["simple", "cache"]), + ), + basic_setting( + &format!("{prefix}.ngram_max_proposal_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.extension_initial_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.extension_max_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.extension_tail_backoff_proposals"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.native_mtp_reject_cooldown_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.native_mtp_suppress_cooldown_drafts"), + ConfigValueSchema::Boolean, + ), + basic_setting( + &format!("{prefix}.native_mtp_suppress_cooldown_draft_limit"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.verify_window_min_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.verify_window_max_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting( + &format!("{prefix}.verify_window_pipeline_depth"), + ConfigValueSchema::Integer, + ), basic_setting(&format!("{prefix}.spec_default"), bool_or_auto_schema()), ] } diff --git a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs index adec7bd982..e05b28faea 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs @@ -11,7 +11,7 @@ pub(super) fn apply_speculative_behavior( suffix: &str, ) { match suffix { - "mode" | "draft_selection_policy" | "pairing_fault" | "spec_default" => { + "strategy" | "mode" | "draft_selection_policy" | "pairing_fault" | "spec_default" => { set_static_options(setting); } "draft_model" => { @@ -68,6 +68,18 @@ pub(super) fn apply_speculative_behavior( push_range_constraint(setting, Some(format!("{prefix}.ngram_min")), None); push_mode_dependency(setting, prefix, "ngram", suffix); } + "ngram_proposer" | "ngram_max_proposal_tokens" => { + push_mode_dependency(setting, prefix, "ngram", suffix); + } + "extension_initial_tokens" + | "extension_max_tokens" + | "extension_tail_backoff_proposals" + | "native_mtp_reject_cooldown_tokens" + | "native_mtp_suppress_cooldown_drafts" + | "native_mtp_suppress_cooldown_draft_limit" + | "verify_window_min_tokens" + | "verify_window_max_tokens" + | "verify_window_pipeline_depth" => {} _ => {} } } diff --git a/crates/mesh-llm-config/src/validate.rs b/crates/mesh-llm-config/src/validate.rs index 2396354cbe..08587931da 100644 --- a/crates/mesh-llm-config/src/validate.rs +++ b/crates/mesh-llm-config/src/validate.rs @@ -878,11 +878,9 @@ fn validate_skippy(config: &SkippyConfig, base_path: &str) -> DiagnosticResult { } fn validate_speculative(config: &SpeculativeConfig, base_path: &str) -> DiagnosticResult { - validate_optional_enum( - config.strategy.as_deref(), - &["auto", "disabled", "mtp"], - &format!("{base_path}.strategy"), - )?; + if let Some(strategy) = config.strategy.as_deref() { + validate_non_empty(strategy, &format!("{base_path}.strategy"))?; + } validate_optional_enum( config.mode.as_deref(), &["auto", "disabled", "draft", "ngram"], @@ -982,6 +980,7 @@ fn validate_speculative(config: &SpeculativeConfig, base_path: &str) -> Diagnost format!("{base_path}.ngram_max must be greater than or equal to {base_path}.ngram_min"), )); } + validate_speculative_proposer_controls(config, base_path)?; validate_bool_or_auto( config.spec_default.as_ref(), &format!("{base_path}.spec_default"), @@ -1001,6 +1000,109 @@ fn validate_speculative(config: &SpeculativeConfig, base_path: &str) -> Diagnost Ok(()) } +fn validate_speculative_proposer_controls( + config: &SpeculativeConfig, + base_path: &str, +) -> DiagnosticResult { + validate_optional_enum( + config.ngram_proposer.as_deref(), + &["simple", "cache"], + &format!("{base_path}.ngram_proposer"), + )?; + validate_optional_u32_range( + config.ngram_max_proposal_tokens, + &format!("{base_path}.ngram_max_proposal_tokens"), + 1, + 10_000_000, + )?; + validate_extension_controls(config, base_path)?; + validate_native_mtp_controls(config, base_path)?; + validate_verify_window_controls(config, base_path) +} + +fn validate_extension_controls(config: &SpeculativeConfig, base_path: &str) -> DiagnosticResult { + validate_optional_u32_range( + config.extension_initial_tokens, + &format!("{base_path}.extension_initial_tokens"), + 1, + 10_000_000, + )?; + validate_optional_u32_range( + config.extension_max_tokens, + &format!("{base_path}.extension_max_tokens"), + 1, + 10_000_000, + )?; + if let (Some(initial), Some(max)) = + (config.extension_initial_tokens, config.extension_max_tokens) + && initial > max + { + return Err(validation_diagnostic( + &format!("{base_path}.extension_initial_tokens"), + format!( + "{base_path}.extension_initial_tokens must be less than or equal to {base_path}.extension_max_tokens" + ), + )); + } + validate_optional_u32_range( + config.extension_tail_backoff_proposals, + &format!("{base_path}.extension_tail_backoff_proposals"), + 0, + 10_000_000, + ) +} + +fn validate_native_mtp_controls(config: &SpeculativeConfig, base_path: &str) -> DiagnosticResult { + validate_optional_u32_range( + config.native_mtp_reject_cooldown_tokens, + &format!("{base_path}.native_mtp_reject_cooldown_tokens"), + 0, + 10_000_000, + )?; + validate_optional_u32_range( + config.native_mtp_suppress_cooldown_draft_limit, + &format!("{base_path}.native_mtp_suppress_cooldown_draft_limit"), + 0, + 10_000_000, + ) +} + +fn validate_verify_window_controls( + config: &SpeculativeConfig, + base_path: &str, +) -> DiagnosticResult { + validate_optional_u32_range( + config.verify_window_min_tokens, + &format!("{base_path}.verify_window_min_tokens"), + 1, + 10_000_000, + )?; + validate_optional_u32_range( + config.verify_window_max_tokens, + &format!("{base_path}.verify_window_max_tokens"), + 1, + 10_000_000, + )?; + if let (Some(min), Some(max)) = ( + config.verify_window_min_tokens, + config.verify_window_max_tokens, + ) && min > max + { + return Err(validation_diagnostic( + &format!("{base_path}.verify_window_min_tokens"), + format!( + "{base_path}.verify_window_min_tokens must be less than or equal to {base_path}.verify_window_max_tokens" + ), + )); + } + validate_optional_u32_range( + config.verify_window_pipeline_depth, + &format!("{base_path}.verify_window_pipeline_depth"), + 1, + 1_024, + ) +} + fn validate_request_defaults(config: &RequestDefaultsConfig, base_path: &str) -> DiagnosticResult { validate_optional_u32_range( config.max_tokens, @@ -1885,7 +1987,7 @@ gpu_id = "metal:0" } #[test] - fn speculative_strategy_rejects_unknown_values() { + fn speculative_strategy_allows_package_declared_names() { let config: MeshConfig = toml::from_str( r#" [defaults.speculative] @@ -1894,17 +1996,9 @@ strategy = "mystery-oracle" ) .expect("config should parse before validation"); - let diagnostics = validate_config_diagnostics(&config); - assert_eq!(diagnostics.len(), 1); - assert_eq!( - diagnostics[0].path.as_ref().map(ConfigPath::render), - Some("defaults.speculative.strategy".to_string()) - ); - assert!( - diagnostics[0] - .message - .contains("defaults.speculative.strategy must be one of") - ); + validate_config(&config) + .expect("package strategy names are validated after package resolution"); + assert!(validate_config_diagnostics(&config).is_empty()); } #[test] @@ -1928,7 +2022,7 @@ strategy = "native-mtp-n1" } #[test] - fn speculative_strategy_native_mtp_n1_raw_value_is_invalid() { + fn speculative_strategy_raw_name_is_deferred_to_package_resolution() { let config = MeshConfig { defaults: Some(ModelConfigDefaults { speculative: Some(SpeculativeConfig { @@ -1940,17 +2034,9 @@ strategy = "native-mtp-n1" ..MeshConfig::default() }; - let diagnostics = validate_config_diagnostics(&config); - assert_eq!(diagnostics.len(), 1); - assert_eq!( - diagnostics[0].path.as_ref().map(ConfigPath::render), - Some("defaults.speculative.strategy".to_string()) - ); - assert!( - diagnostics[0] - .message - .contains("defaults.speculative.strategy must be one of") - ); + validate_config(&config) + .expect("package strategy names are validated after package resolution"); + assert!(validate_config_diagnostics(&config).is_empty()); } #[test] diff --git a/crates/mesh-llm-host-runtime/Cargo.toml b/crates/mesh-llm-host-runtime/Cargo.toml index 45034e06d4..26a04cadbb 100644 --- a/crates/mesh-llm-host-runtime/Cargo.toml +++ b/crates/mesh-llm-host-runtime/Cargo.toml @@ -55,6 +55,7 @@ openai-frontend = { path = "../openai-frontend", version = "0.72.1" } skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } skippy-coordinator = { path = "../skippy-coordinator", version = "0.72.1" } skippy-runtime = { path = "../skippy-runtime", version = "0.72.1" } +skippy-ffi = { path = "../skippy-ffi", version = "0.72.1", default-features = false } skippy-server = { path = "../skippy-server", version = "0.72.1" } skippy-topology = { path = "../skippy-topology", version = "0.72.1" } iroh = "1.0.0" diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 3e3f7edf94..a6616a5eb6 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -185,6 +185,14 @@ impl SkippyTelemetryOptions { level: TelemetryLevel::Debug, } } + + pub(crate) fn summary(metrics_otlp_grpc: String) -> Self { + Self { + metrics_otlp_grpc: Some(metrics_otlp_grpc), + queue_capacity: 1024, + level: TelemetryLevel::Summary, + } + } } pub(crate) fn default_skippy_openai_guardrails() -> OpenAiGuardrailsConfig { @@ -419,6 +427,7 @@ fn embedded_openai_args_from( speculative_window: embedded_args.speculative_window, adaptive_speculative_window: embedded_args.adaptive_speculative_window, draft_n_gpu_layers: embedded_args.draft_n_gpu_layers, + speculative: embedded_args.speculative, ngram_min: embedded_args.ngram_min, ngram_max: embedded_args.ngram_max, native_mtp_enabled: embedded_args.native_mtp_enabled, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs index 689b851d16..125d6dfa1f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs @@ -8,6 +8,7 @@ use std::{ use anyhow::{Context, Result}; use serde::Serialize; use sha2::{Digest, Sha256}; +use skippy_ffi::TensorRole; use skippy_runtime::package::PackageGenerationInfo; use super::hash_cache::{self, SidecarDigestCache}; @@ -87,6 +88,13 @@ pub fn synthetic_direct_gguf_package( source_model_path.display() ); let source_model_bytes = source_files.iter().map(|file| file.bytes).sum(); + let layer_weight_bytes = direct_gguf_layer_weight_bytes(&source_files, compact.layer_count) + .with_context(|| { + format!( + "inspect GGUF tensor weights {}", + source_model_path.display() + ) + })?; let source_model_sha256 = aggregate_source_sha256(&source_files); @@ -113,7 +121,7 @@ pub fn synthetic_direct_gguf_package( source_model_sha256, source_model_bytes, source_files, - layer_weight_bytes: Vec::new(), + layer_weight_bytes, layer_count: compact.layer_count, activation_width: compact.embedding_size, tensor_count, @@ -311,6 +319,88 @@ fn gguf_tensor_count(path: &Path) -> Result { read_gguf_count(&mut reader, version) } +fn direct_gguf_layer_weight_bytes( + source_files: &[SkippyPackageSourceFile], + layer_count: u32, +) -> Result> { + let mut tensors = Vec::new(); + for source_file in source_files { + let info = match skippy_runtime::ModelInfo::open(&source_file.path) { + Ok(info) => info, + Err(error) => { + tracing::debug!( + path = %source_file.path.display(), + error = %error, + "GGUF tensor layout unavailable; using capacity-based split planning" + ); + return Ok(Vec::new()); + } + }; + tensors.extend( + info.tensors() + .with_context(|| format!("read GGUF tensors {}", source_file.path.display()))?, + ); + } + Ok(layer_weight_bytes_from_tensors(&tensors, layer_count)) +} + +fn layer_weight_bytes_from_tensors( + tensors: &[skippy_runtime::TensorInfo], + layer_count: u32, +) -> Vec { + let Ok(layer_count) = usize::try_from(layer_count) else { + return Vec::new(); + }; + if layer_count == 0 { + return Vec::new(); + } + + let mut weights = vec![0_u64; layer_count]; + let mut shared_bytes = 0_u64; + let mut seen = std::collections::BTreeSet::new(); + + for tensor in tensors { + if !seen.insert(tensor.name.as_str()) { + continue; + } + let bytes = tensor.byte_size; + match tensor.layer_index { + Some(layer) if (layer as usize) < layer_count => { + weights[layer as usize] = weights[layer as usize].saturating_add(bytes); + } + // Native MTP blocks are appended after the trunk's declared layer + // count and must stay with the final stage that owns logits. + Some(_) => { + let last = weights.len() - 1; + weights[last] = weights[last].saturating_add(bytes); + } + None => match tensor.role { + TensorRole::Embedding => { + weights[0] = weights[0].saturating_add(bytes); + } + TensorRole::FinalNorm | TensorRole::Output => { + let last = weights.len() - 1; + weights[last] = weights[last].saturating_add(bytes); + } + TensorRole::Unknown + | TensorRole::Metadata + | TensorRole::Tokenizer + | TensorRole::Layer => { + shared_bytes = shared_bytes.saturating_add(bytes); + } + }, + } + } + + // Metadata is loaded at every stage but is normally tiny. Split it between + // endpoints so total model weight stays conserved without biasing a middle + // stage in multi-node plans. + weights[0] = weights[0].saturating_add(shared_bytes.div_ceil(2)); + let last = weights.len() - 1; + weights[last] = weights[last].saturating_add(shared_bytes / 2); + weights +} + fn read_u32_le(reader: &mut impl Read) -> Result { let mut bytes = [0u8; 4]; reader.read_exact(&mut bytes).context("read u32")?; @@ -463,6 +553,7 @@ fn required_layer_package_activation_width( #[cfg(test)] mod tests { use super::*; + use skippy_runtime::TensorInfo; #[test] fn synthetic_manifest_identity_is_stable_and_metadata_sensitive() { @@ -764,4 +855,35 @@ mod tests { info.layers[1].layer_index = 1; assert_eq!(layer_weight_bytes_from_info(&info), vec![30, 40]); } + + #[test] + fn direct_gguf_weights_charge_native_mtp_block_to_final_stage() { + let tensors = vec![ + tensor("token_embd.weight", None, TensorRole::Embedding, 5), + tensor("blk.0.attn_norm.weight", Some(0), TensorRole::Layer, 10), + tensor("blk.1.attn_norm.weight", Some(1), TensorRole::Layer, 10), + tensor("blk.2.nextn.eh_proj.weight", Some(2), TensorRole::Layer, 7), + tensor("output_norm.weight", None, TensorRole::FinalNorm, 1), + tensor("output.weight", None, TensorRole::Output, 9), + tensor("general.alignment", None, TensorRole::Metadata, 3), + ]; + + assert_eq!(layer_weight_bytes_from_tensors(&tensors, 2), vec![17, 28]); + } + + fn tensor( + name: &str, + layer_index: Option, + role: TensorRole, + byte_size: u64, + ) -> TensorInfo { + TensorInfo { + name: name.to_string(), + layer_index, + role, + ggml_type: 0, + byte_size, + element_count: byte_size, + } + } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs index f6e80a12a6..5444c61b2f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs @@ -3,31 +3,149 @@ use super::*; use crate::inference::skippy::SkippyTelemetryOptions; use skippy_protocol::LoadMode; use skippy_runtime::package::{ - PackageGenerationInfo, PackageSpeculativeDecodingInfo, PackageSpeculativeStrategyInfo, - PackageWindowPolicyInfo, + PackageExtensionPolicyInfo, PackageGenerationInfo, PackageSpeculativeDecodingInfo, + PackageSpeculativeProposerInfo, PackageSpeculativeStrategyInfo, PackageWindowPolicyInfo, }; use std::collections::BTreeMap; fn native_mtp_generation() -> PackageGenerationInfo { + let mut proposers = BTreeMap::new(); + proposers.insert( + "mtp".to_string(), + PackageSpeculativeProposerInfo { + proposer_type: "native-mtp".to_string(), + prediction_depth: Some(1), + layer_indices: vec![46], + ngram_min: None, + ngram_max: None, + max_proposal_tokens: None, + history_scope: None, + }, + ); let mut strategies = BTreeMap::new(); strategies.insert( "mtp".to_string(), PackageSpeculativeStrategyInfo { strategy_type: "native-mtp".to_string(), - prediction_depth: Some(1), - layer_indices: vec![46], + prediction_depth: None, + layer_indices: Vec::new(), window_policy: Some(PackageWindowPolicyInfo { default: "fixed".to_string(), initial_window: 1, min_window: 1, max_window: 1, }), + proposer: Some("mtp".to_string()), + primary: None, + extender: None, + extension_policy: None, }, ); PackageGenerationInfo { speculative_decoding: Some(PackageSpeculativeDecodingInfo { default: "mtp".to_string(), + proposers, + strategies, + }), + } +} + +fn native_mtp_cache_generation() -> PackageGenerationInfo { + let mut proposers = BTreeMap::new(); + proposers.insert( + "mtp".to_string(), + PackageSpeculativeProposerInfo { + proposer_type: "native-mtp".to_string(), + prediction_depth: Some(1), + layer_indices: vec![46], + ngram_min: None, + ngram_max: None, + max_proposal_tokens: None, + history_scope: None, + }, + ); + proposers.insert( + "cache".to_string(), + PackageSpeculativeProposerInfo { + proposer_type: "ngram-cache".to_string(), + prediction_depth: None, + layer_indices: Vec::new(), + ngram_min: Some(2), + ngram_max: Some(4), + max_proposal_tokens: Some(10), + history_scope: Some("request".to_string()), + }, + ); + let mut strategies = BTreeMap::new(); + strategies.insert( + "mtp-cache".to_string(), + PackageSpeculativeStrategyInfo { + strategy_type: "composite".to_string(), + prediction_depth: None, + layer_indices: Vec::new(), + window_policy: Some(PackageWindowPolicyInfo { + default: "adaptive".to_string(), + initial_window: 2, + min_window: 1, + max_window: 6, + }), + proposer: None, + primary: Some("mtp".to_string()), + extender: Some("cache".to_string()), + extension_policy: Some(PackageExtensionPolicyInfo { + initial_tokens: 2, + max_tokens: 8, + tail_backoff_proposals: 5, + }), + }, + ); + PackageGenerationInfo { + speculative_decoding: Some(PackageSpeculativeDecodingInfo { + default: "mtp-cache".to_string(), + proposers, + strategies, + }), + } +} + +fn ngram_cache_generation() -> PackageGenerationInfo { + let mut proposers = BTreeMap::new(); + proposers.insert( + "cache".to_string(), + PackageSpeculativeProposerInfo { + proposer_type: "ngram-cache".to_string(), + prediction_depth: None, + layer_indices: Vec::new(), + ngram_min: Some(2), + ngram_max: Some(4), + max_proposal_tokens: Some(6), + history_scope: Some("request".to_string()), + }, + ); + let mut strategies = BTreeMap::new(); + strategies.insert( + "ngram-cache".to_string(), + PackageSpeculativeStrategyInfo { + strategy_type: "ngram-cache".to_string(), + prediction_depth: None, + layer_indices: Vec::new(), + window_policy: Some(PackageWindowPolicyInfo { + default: "fixed".to_string(), + initial_window: 6, + min_window: 1, + max_window: 6, + }), + proposer: Some("cache".to_string()), + primary: None, + extender: None, + extension_policy: None, + }, + ); + PackageGenerationInfo { + speculative_decoding: Some(PackageSpeculativeDecodingInfo { + default: "ngram-cache".to_string(), + proposers, strategies, }), } @@ -185,6 +303,284 @@ fn speculative_strategy_auto_uses_package_native_mtp_default() { assert_eq!(openai.native_mtp_min_tokens, 0); } +#[test] +fn package_composite_strategy_resolves_native_mtp_with_cache_extension() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "mtp-cache" +ngram_max_proposal_tokens = 4 +extension_max_tokens = 8 +verify_window_pipeline_depth = 1 + +[[models]] +model = "meshllm/GLM-4.7-Flash-MTP-GGUF" + +[models.speculative] +ngram_max_proposal_tokens = 9 +extension_max_tokens = 7 +verify_window_pipeline_depth = 2 +"#, + ); + assert_eq!( + mesh_config + .defaults + .as_ref() + .and_then(|defaults| defaults.speculative.as_ref()) + .and_then(|speculative| speculative.ngram_max_proposal_tokens), + Some(4) + ); + assert_eq!( + mesh_config + .models + .first() + .and_then(|model| model.speculative.as_ref()) + .and_then(|speculative| speculative.ngram_max_proposal_tokens), + Some(9) + ); + let model_file = temp_model_file(); + let generation = native_mtp_cache_generation(); + + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: Some(&generation), + }) + .expect("package composite strategy should resolve"); + + assert!(resolved.speculative.native_mtp_enabled); + assert_eq!( + resolved.speculative.decode.effective_strategy, + "native-mtp+ngram-cache" + ); + let ngram = resolved + .speculative + .decode + .ngram + .as_ref() + .expect("cache proposer should resolve"); + assert_eq!(ngram.kind, skippy_server::NgramProposerKind::Cache); + assert_eq!(ngram.min_ngram, 2); + assert_eq!(ngram.max_ngram, 4); + assert_eq!(ngram.max_proposal_tokens, 9); + let extension = resolved + .speculative + .decode + .extension + .as_ref() + .expect("extension policy should resolve"); + assert_eq!(extension.initial_tokens, 2); + assert_eq!(extension.max_tokens, 7); + assert_eq!(extension.tail_backoff_proposals, 5); + assert_eq!(resolved.speculative.decode.verify_window.min_tokens, 1); + assert_eq!(resolved.speculative.decode.verify_window.max_tokens, 6); + assert_eq!(resolved.speculative.decode.verify_window.pipeline_depth, 2); +} + +#[test] +fn package_cache_strategy_uses_the_declared_verify_window() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "ngram-cache" +"#, + ); + let model_file = temp_model_file(); + let generation = ngram_cache_generation(); + + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: Some(&generation), + }) + .expect("package cache strategy should resolve"); + + assert!(!resolved.speculative.native_mtp_enabled); + let openai = resolved + .to_embedded_openai_args(4096, true) + .expect("package cache strategy should build OpenAI args"); + assert_eq!(openai.speculative_window, 6); + assert_eq!(openai.ngram_min, 2); + assert_eq!(openai.ngram_max, 6); + assert_eq!( + openai.speculative.ngram.as_ref().map(|ngram| ngram.kind), + Some(skippy_server::NgramProposerKind::Cache) + ); +} + +#[test] +fn direct_native_mtp_can_use_a_request_local_cache_extension() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "mtp" +ngram_proposer = "cache" +ngram_min = 2 +ngram_max = 4 +ngram_max_proposal_tokens = 6 +"#, + ); + let model_file = temp_model_file_with_tensor_names(&["blk.23.nextn.eh_proj.weight"], None); + + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .expect("direct native MTP with cache extension should resolve"); + + assert!(resolved.speculative.native_mtp_enabled); + assert_eq!( + resolved.speculative.decode.effective_strategy, + "native-mtp+ngram-cache" + ); + let extension = resolved + .speculative + .decode + .extension + .as_ref() + .expect("direct cache strategy should synthesize an extension plan"); + assert_eq!(extension.initial_tokens, 2); + assert_eq!(extension.max_tokens, 6); + let openai = resolved + .to_embedded_openai_args(4096, true) + .expect("direct cache strategy should build OpenAI args"); + assert!(openai.native_mtp_enabled); + assert_eq!(openai.ngram_min, 2); + assert_eq!(openai.ngram_max, 6); + assert_eq!( + openai.speculative.ngram.as_ref().map(|ngram| ngram.kind), + Some(skippy_server::NgramProposerKind::Cache) + ); +} + +#[test] +fn direct_cache_strategy_rejects_an_unsupported_cache_window() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "ngram-cache" +ngram_proposer = "cache" +ngram_min = 2 +ngram_max = 5 +ngram_max_proposal_tokens = 6 +"#, + ); + let model_file = temp_model_file(); + + let error = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .expect_err("cache windows above the llama.cpp limit must be rejected"); + + assert!( + error + .to_string() + .contains("must not exceed llama.cpp limit 4") + ); +} + +#[test] +fn direct_cache_strategy_resolves_a_request_local_cache_proposer() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "ngram-cache" +ngram_min = 2 +ngram_max = 4 +ngram_max_proposal_tokens = 6 +"#, + ); + let model_file = temp_model_file(); + + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .expect("direct cache strategy should resolve"); + + assert!(!resolved.speculative.native_mtp_enabled); + assert_eq!( + resolved.speculative.decode.effective_strategy, + "ngram-cache" + ); + let ngram = resolved + .speculative + .decode + .ngram + .as_ref() + .expect("direct cache strategy should select an N-gram proposer"); + assert_eq!(ngram.kind, skippy_server::NgramProposerKind::Cache); + assert_eq!(ngram.min_ngram, 2); + assert_eq!(ngram.max_ngram, 4); + assert_eq!(ngram.max_proposal_tokens, 6); +} + +#[test] +fn direct_native_mtp_can_use_a_simple_ngram_extension() { + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "mtp" +ngram_proposer = "simple" +ngram_min = 2 +ngram_max = 6 +ngram_max_proposal_tokens = 6 +"#, + ); + let model_file = temp_model_file_with_tensor_names(&["blk.23.nextn.eh_proj.weight"], None); + + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/GLM-4.7-Flash-MTP-GGUF", + model_path: model_file.path(), + model_bytes: 4 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .expect("direct native MTP with simple N-gram extension should resolve"); + + assert!(resolved.speculative.native_mtp_enabled); + assert_eq!( + resolved.speculative.decode.effective_strategy, + "native-mtp+ngram-simple" + ); + assert!(resolved.speculative.decode.extension.is_some()); + let openai = resolved + .to_embedded_openai_args(4096, true) + .expect("direct simple strategy should build OpenAI args"); + assert_eq!(openai.ngram_min, 2); + assert_eq!(openai.ngram_max, 6); + assert_eq!( + openai.speculative.ngram.as_ref().map(|ngram| ngram.kind), + Some(skippy_server::NgramProposerKind::Simple) + ); +} + #[test] fn speculative_strategy_native_mtp_rejects_direct_gguf_without_proven_support() { let mesh_config = parse_config( diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index c3421bb719..68d9a9a0cf 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs @@ -4,7 +4,14 @@ use crate::models::find_model_path; use anyhow::{Result, bail}; use mesh_llm_system::util::validate_draft_min_max; use model_artifact::gguf::{scan_gguf_compact_meta, scan_gguf_tensor_names_any}; -use skippy_runtime::package::{PackageGenerationInfo, PackageSpeculativeDecodingInfo}; +use skippy_runtime::package::{ + PackageExtensionPolicyInfo, PackageGenerationInfo, PackageSpeculativeDecodingInfo, + PackageSpeculativeProposerInfo, PackageSpeculativeStrategyInfo, PackageWindowPolicyInfo, +}; +use skippy_server::{ + NativeMtpProposalConfig, NgramExtensionConfig, NgramProposalConfig, NgramProposerKind, + SpeculativeDecodeConfig, VerifyWindowConfig, +}; use skippy_topology::infer_family_capability; use super::support::{pick_owned, pick_string, pick_string_owned}; @@ -49,23 +56,13 @@ pub(super) fn resolve_speculative_config( global_config.and_then(|config| config.strategy.as_deref()), Some("auto"), ); - let native_mtp_enabled = match strategy.as_str() { - "auto" => { - auto_defaults_enabled - && package_generation_or_direct_default_supports_native_mtp( - package_generation, - model_path, - ) - } - "mtp" => { - if !supports_native_mtp { - bail!("skippy speculative.strategy = \"mtp\" requires proven native MTP support"); - } - true - } - "disabled" => false, - _ => bail!("skippy speculative.strategy must be auto, disabled, or mtp"), - }; + let (strategy, native_mtp_enabled) = resolve_native_mtp_strategy( + strategy, + auto_defaults_enabled, + supports_native_mtp, + package_generation, + model_path, + )?; let mode = pick_string_owned( model_config.and_then(|config| config.mode.as_deref()), global_config.and_then(|config| config.mode.as_deref()), @@ -134,7 +131,7 @@ pub(super) fn resolve_speculative_config( draft_model_path = None; } Ok(ResolvedSpeculativeConfig { - strategy, + strategy: strategy.clone(), native_mtp_enabled, mode, draft_model_path, @@ -145,9 +142,492 @@ pub(super) fn resolve_speculative_config( draft_n_gpu_layers, ngram_min, ngram_max, + decode: resolve_decode_config(DecodeResolutionInput { + requested_strategy: &strategy, + native_mtp_enabled, + draft_max_tokens: effective_draft_max_tokens, + draft_min_tokens, + legacy_ngram_min: ngram_min, + legacy_ngram_max: ngram_max, + model_config, + global_config, + package_generation, + })?, }) } +fn resolve_native_mtp_strategy( + strategy: String, + auto_defaults_enabled: bool, + supports_native_mtp: bool, + package_generation: Option<&PackageGenerationInfo>, + model_path: &Path, +) -> Result<(String, bool)> { + let native_mtp_enabled = match strategy.as_str() { + "auto" => { + auto_defaults_enabled + && package_generation_or_direct_default_supports_native_mtp( + package_generation, + model_path, + ) + } + "mtp" => { + if !supports_native_mtp { + bail!("skippy speculative.strategy = \"mtp\" requires proven native MTP support"); + } + true + } + "ngram-simple" | "ngram-cache" => false, + "disabled" => false, + package_strategy if package_strategy_exists(package_generation, package_strategy) => { + let speculative = package_generation + .and_then(|generation| generation.speculative_decoding.as_ref()) + .expect("checked package strategy exists"); + strategy_uses_native_mtp(speculative, package_strategy) + } + _ => bail!( + "skippy speculative.strategy must be auto, disabled, mtp, or a strategy declared by model-package.json" + ), + }; + Ok((strategy, native_mtp_enabled)) +} + +struct DecodeResolutionInput<'a> { + requested_strategy: &'a str, + native_mtp_enabled: bool, + draft_max_tokens: u32, + draft_min_tokens: u32, + legacy_ngram_min: u32, + legacy_ngram_max: u32, + model_config: Option<&'a SpeculativeConfig>, + global_config: Option<&'a SpeculativeConfig>, + package_generation: Option<&'a PackageGenerationInfo>, +} + +fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result { + let mut config = package_decode_config(input.requested_strategy, input.package_generation)? + .unwrap_or_else(SpeculativeDecodeConfig::default); + config.requested_strategy = input.requested_strategy.to_string(); + + if input.native_mtp_enabled { + config.native_mtp.enabled = true; + config.native_mtp.max_draft_tokens = input.draft_max_tokens.max(1) as usize; + config.native_mtp.min_draft_tokens = input.draft_min_tokens as usize; + } + if config.native_mtp.enabled && config.effective_strategy == "disabled" { + config.effective_strategy = "native-mtp".to_string(); + } + + let ngram_min = pick_optional_u32( + input.model_config.and_then(|config| config.ngram_min), + input.global_config.and_then(|config| config.ngram_min), + ) + .unwrap_or(input.legacy_ngram_min); + let ngram_max = pick_optional_u32( + input.model_config.and_then(|config| config.ngram_max), + input.global_config.and_then(|config| config.ngram_max), + ) + .unwrap_or(input.legacy_ngram_max); + let ngram_proposer = pick_owned( + input + .model_config + .and_then(|config| config.ngram_proposer.clone()), + input + .global_config + .and_then(|config| config.ngram_proposer.clone()), + ); + let ngram_max_proposal_tokens = pick_optional_u32( + input + .model_config + .and_then(|config| config.ngram_max_proposal_tokens), + input + .global_config + .and_then(|config| config.ngram_max_proposal_tokens), + ); + if config.ngram.is_some() || ngram_min > 0 || ngram_max > 0 || ngram_proposer.is_some() { + let existing = config.ngram.as_ref(); + let min_ngram = nonzero_or( + ngram_min, + existing.map_or(0, |ngram| ngram.min_ngram as u32), + ); + let max_ngram = nonzero_or( + ngram_max, + existing.map_or(0, |ngram| ngram.max_ngram as u32), + ); + if min_ngram == 0 || max_ngram == 0 || min_ngram > max_ngram { + bail!("skippy speculative N-gram proposer requires 0 < ngram_min <= ngram_max"); + } + let kind = match ngram_proposer.as_deref() { + Some("cache") => NgramProposerKind::Cache, + Some("simple") => NgramProposerKind::Simple, + None => existing.map_or_else( + || match input.requested_strategy { + "ngram-cache" => NgramProposerKind::Cache, + _ => NgramProposerKind::Simple, + }, + |ngram| ngram.kind, + ), + Some(_) => unreachable!("validated by mesh configuration"), + }; + let max_proposal_tokens = ngram_max_proposal_tokens + .map(|value| value as usize) + .unwrap_or_else(|| { + existing.map_or(max_ngram as usize, |ngram| ngram.max_proposal_tokens) + }); + config.ngram = Some(NgramProposalConfig { + kind, + min_ngram: min_ngram as usize, + max_ngram: max_ngram as usize, + max_proposal_tokens, + }); + if config.effective_strategy == "disabled" { + config.effective_strategy = ngram_effective_strategy(kind).to_string(); + } + } + + if config.native_mtp.enabled + && let Some(ngram) = config.ngram.as_ref() + { + config.effective_strategy = match ngram.kind { + NgramProposerKind::Simple => "native-mtp+ngram-simple", + NgramProposerKind::Cache => "native-mtp+ngram-cache", + } + .to_string(); + if config.extension.is_none() { + config.extension = Some(NgramExtensionConfig { + initial_tokens: ngram.max_proposal_tokens.clamp(1, 2), + max_tokens: ngram.max_proposal_tokens, + tail_backoff_proposals: 0, + }); + } + } + + let extension_initial = pick_optional_u32( + input + .model_config + .and_then(|config| config.extension_initial_tokens), + input + .global_config + .and_then(|config| config.extension_initial_tokens), + ); + let extension_max = pick_optional_u32( + input + .model_config + .and_then(|config| config.extension_max_tokens), + input + .global_config + .and_then(|config| config.extension_max_tokens), + ); + let extension_backoff = pick_optional_u32( + input + .model_config + .and_then(|config| config.extension_tail_backoff_proposals), + input + .global_config + .and_then(|config| config.extension_tail_backoff_proposals), + ); + if extension_initial.is_some() || extension_max.is_some() || extension_backoff.is_some() { + let Some(extension) = config.extension.as_mut() else { + bail!( + "skippy speculative extension controls require native MTP and an N-gram proposer" + ); + }; + if let Some(value) = extension_initial { + extension.initial_tokens = value as usize; + } + if let Some(value) = extension_max { + extension.max_tokens = value as usize; + } + if let Some(value) = extension_backoff { + extension.tail_backoff_proposals = value as usize; + } + } + if config.extension.is_some() && (!config.native_mtp.enabled || config.ngram.is_none()) { + bail!("skippy speculative extension requires both native MTP and an N-gram proposer"); + } + + config.native_mtp.reject_cooldown_tokens = pick_optional_u32( + input + .model_config + .and_then(|config| config.native_mtp_reject_cooldown_tokens), + input + .global_config + .and_then(|config| config.native_mtp_reject_cooldown_tokens), + ) + .map_or(config.native_mtp.reject_cooldown_tokens, |value| { + value as usize + }); + config.native_mtp.suppress_cooldown_drafts = pick_owned( + input + .model_config + .and_then(|config| config.native_mtp_suppress_cooldown_drafts), + input + .global_config + .and_then(|config| config.native_mtp_suppress_cooldown_drafts), + ) + .unwrap_or(config.native_mtp.suppress_cooldown_drafts); + config.native_mtp.suppress_cooldown_draft_limit = pick_optional_u32( + input + .model_config + .and_then(|config| config.native_mtp_suppress_cooldown_draft_limit), + input + .global_config + .and_then(|config| config.native_mtp_suppress_cooldown_draft_limit), + ) + .map_or(config.native_mtp.suppress_cooldown_draft_limit, |value| { + value as usize + }); + config.verify_window.min_tokens = pick_optional_u32( + input + .model_config + .and_then(|config| config.verify_window_min_tokens), + input + .global_config + .and_then(|config| config.verify_window_min_tokens), + ) + .map_or(config.verify_window.min_tokens, |value| value as usize); + config.verify_window.max_tokens = pick_optional_u32( + input + .model_config + .and_then(|config| config.verify_window_max_tokens), + input + .global_config + .and_then(|config| config.verify_window_max_tokens), + ) + .map_or(config.verify_window.max_tokens, |value| value as usize); + config.verify_window.pipeline_depth = pick_optional_u32( + input + .model_config + .and_then(|config| config.verify_window_pipeline_depth), + input + .global_config + .and_then(|config| config.verify_window_pipeline_depth), + ) + .map_or(config.verify_window.pipeline_depth, |value| value as usize); + if config.verify_window.min_tokens > config.verify_window.max_tokens { + bail!("skippy speculative verify window requires min_tokens <= max_tokens"); + } + config.validate()?; + Ok(config) +} + +fn package_decode_config( + requested_strategy: &str, + package_generation: Option<&PackageGenerationInfo>, +) -> Result> { + let Some(speculative) = + package_generation.and_then(|generation| generation.speculative_decoding.as_ref()) + else { + return Ok(None); + }; + let strategy_name = if requested_strategy == "auto" { + speculative.default.as_str() + } else { + requested_strategy + }; + let Some(strategy) = speculative.strategies.get(strategy_name) else { + return Ok(None); + }; + let mut native_mtp = None; + let mut ngram = None; + match strategy.strategy_type.as_str() { + "native-mtp" => { + native_mtp = Some(native_proposer_config( + strategy + .proposer + .as_deref() + .and_then(|name| speculative.proposers.get(name)), + strategy, + )?); + } + "ngram-simple" | "ngram-cache" => { + ngram = Some(ngram_proposer_config( + strategy + .proposer + .as_deref() + .and_then(|name| speculative.proposers.get(name)), + strategy.strategy_type.as_str(), + )?); + } + "composite" => { + let primary = strategy + .primary + .as_deref() + .and_then(|name| speculative.proposers.get(name)) + .ok_or_else(|| anyhow::anyhow!("package speculative strategy {strategy_name} has no native MTP primary proposer"))?; + let extender = strategy + .extender + .as_deref() + .and_then(|name| speculative.proposers.get(name)) + .ok_or_else(|| { + anyhow::anyhow!( + "package speculative strategy {strategy_name} has no N-gram extender" + ) + })?; + native_mtp = Some(native_proposer_config(Some(primary), strategy)?); + ngram = Some(ngram_proposer_config( + Some(extender), + extender.proposer_type.as_str(), + )?); + } + other => bail!("package speculative strategy {strategy_name} has unsupported type {other}"), + } + let extension = strategy.extension_policy.as_ref().map(extension_config); + let verify_window = strategy + .window_policy + .as_ref() + .map(verify_window_config) + .unwrap_or(VerifyWindowConfig { + min_tokens: 1, + max_tokens: 4, + pipeline_depth: 1, + }); + let effective_strategy = match (native_mtp.is_some(), ngram.as_ref().map(|value| value.kind)) { + (true, Some(NgramProposerKind::Simple)) => "native-mtp+ngram-simple", + (true, Some(NgramProposerKind::Cache)) => "native-mtp+ngram-cache", + (true, None) => "native-mtp", + (false, Some(kind)) => ngram_effective_strategy(kind), + (false, None) => "disabled", + }; + Ok(Some(SpeculativeDecodeConfig { + requested_strategy: requested_strategy.to_string(), + effective_strategy: effective_strategy.to_string(), + native_mtp: native_mtp.unwrap_or(NativeMtpProposalConfig { + enabled: false, + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }), + ngram, + extension, + verify_window, + })) +} + +fn native_proposer_config( + proposer: Option<&PackageSpeculativeProposerInfo>, + legacy_strategy: &PackageSpeculativeStrategyInfo, +) -> Result { + let (proposer_type, prediction_depth, layer_indices) = proposer.map_or( + ( + legacy_strategy.strategy_type.as_str(), + legacy_strategy.prediction_depth, + legacy_strategy.layer_indices.as_slice(), + ), + |proposer| { + ( + proposer.proposer_type.as_str(), + proposer.prediction_depth, + proposer.layer_indices.as_slice(), + ) + }, + ); + if proposer_type != "native-mtp" || prediction_depth != Some(1) || layer_indices.is_empty() { + bail!("package native MTP proposer is not valid for the embedded runtime"); + } + Ok(NativeMtpProposalConfig { + enabled: true, + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }) +} + +fn ngram_proposer_config( + proposer: Option<&PackageSpeculativeProposerInfo>, + expected_type: &str, +) -> Result { + let proposer = proposer + .ok_or_else(|| anyhow::anyhow!("package N-gram strategy must reference a proposer"))?; + let kind = match proposer.proposer_type.as_str() { + "ngram-simple" => NgramProposerKind::Simple, + "ngram-cache" => NgramProposerKind::Cache, + other => bail!("package N-gram proposer has unsupported type {other}"), + }; + if expected_type != "composite" && expected_type != proposer.proposer_type { + bail!("package N-gram strategy type does not match its proposer"); + } + let min_ngram = proposer + .ngram_min + .ok_or_else(|| anyhow::anyhow!("package N-gram proposer has no ngram_min"))?; + let max_ngram = proposer + .ngram_max + .ok_or_else(|| anyhow::anyhow!("package N-gram proposer has no ngram_max"))?; + let max_proposal_tokens = proposer.max_proposal_tokens.unwrap_or(max_ngram); + Ok(NgramProposalConfig { + kind, + min_ngram: min_ngram as usize, + max_ngram: max_ngram as usize, + max_proposal_tokens: max_proposal_tokens as usize, + }) +} + +fn extension_config(policy: &PackageExtensionPolicyInfo) -> NgramExtensionConfig { + NgramExtensionConfig { + initial_tokens: policy.initial_tokens as usize, + max_tokens: policy.max_tokens as usize, + tail_backoff_proposals: policy.tail_backoff_proposals as usize, + } +} + +fn verify_window_config(policy: &PackageWindowPolicyInfo) -> VerifyWindowConfig { + VerifyWindowConfig { + min_tokens: policy.min_window as usize, + max_tokens: policy.max_window as usize, + pipeline_depth: 1, + } +} + +fn ngram_effective_strategy(kind: NgramProposerKind) -> &'static str { + match kind { + NgramProposerKind::Simple => "ngram-simple", + NgramProposerKind::Cache => "ngram-cache", + } +} + +fn nonzero_or(value: u32, default: u32) -> u32 { + if value == 0 { default } else { value } +} + +fn pick_optional_u32(model: Option, global: Option) -> Option { + pick_owned(model, global) +} + +fn package_strategy_exists(generation: Option<&PackageGenerationInfo>, strategy: &str) -> bool { + generation + .and_then(|generation| generation.speculative_decoding.as_ref()) + .is_some_and(|speculative| speculative.strategies.contains_key(strategy)) +} + +fn strategy_uses_native_mtp( + speculative: &PackageSpeculativeDecodingInfo, + strategy_name: &str, +) -> bool { + let Some(strategy) = speculative.strategies.get(strategy_name) else { + return false; + }; + match strategy.strategy_type.as_str() { + "native-mtp" => strategy + .proposer + .as_deref() + .and_then(|name| speculative.proposers.get(name)) + .map_or( + strategy.prediction_depth == Some(1) && !strategy.layer_indices.is_empty(), + |proposer| proposer.proposer_type == "native-mtp", + ), + "composite" => strategy + .primary + .as_deref() + .and_then(|name| speculative.proposers.get(name)) + .is_some_and(|proposer| proposer.proposer_type == "native-mtp"), + _ => false, + } +} + fn reject_unsupported_speculative_runtime_fields( model_config: Option<&SpeculativeConfig>, global_config: Option<&SpeculativeConfig>, @@ -280,16 +760,7 @@ fn package_generation_supports_default_native_mtp( ) -> bool { generation .and_then(|generation| generation.speculative_decoding.as_ref()) - .is_some_and(|speculative| { - speculative - .strategies - .get(&speculative.default) - .is_some_and(|strategy| { - strategy.strategy_type == "native-mtp" - && strategy.prediction_depth == Some(1) - && !strategy.layer_indices.is_empty() - }) - }) + .is_some_and(|speculative| strategy_uses_native_mtp(speculative, &speculative.default)) } fn package_generation_supports_native_mtp(generation: Option<&PackageGenerationInfo>) -> bool { @@ -299,11 +770,7 @@ fn package_generation_supports_native_mtp(generation: Option<&PackageGenerationI } fn speculative_supports_native_mtp(speculative: &PackageSpeculativeDecodingInfo) -> bool { - speculative.strategies.get("mtp").is_some_and(|strategy| { - strategy.strategy_type == "native-mtp" - && strategy.prediction_depth == Some(1) - && !strategy.layer_indices.is_empty() - }) + strategy_uses_native_mtp(speculative, "mtp") } fn direct_gguf_supports_native_mtp(model_path: &Path) -> bool { diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index 725a840c34..3bd1c87805 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -8,7 +8,8 @@ use anyhow::{Result, bail}; use openai_frontend::OpenAiHookPolicy; use skippy_protocol::{LoadMode, StageConfig, StageKvCacheConfig, StageKvCachePayload}; use skippy_server::{ - EmbeddedOpenAiArgs, EmbeddedOpenAiRequestDefaults, EmbeddedRuntimeOptions, telemetry::Telemetry, + EmbeddedOpenAiArgs, EmbeddedOpenAiRequestDefaults, EmbeddedRuntimeOptions, + NativeMtpProposalConfig, SpeculativeDecodeConfig, telemetry::Telemetry, }; use super::super::{ @@ -246,32 +247,27 @@ impl ResolvedSkippyConfig { } else { None }, - ngram_min: if mode == "ngram" { - self.speculative.ngram_min as usize - } else { - 0 - }, - ngram_max: if mode == "ngram" { - self.speculative.ngram_max as usize - } else { - 0 - }, + speculative: self.speculative_decode_config(), + ngram_min: self + .speculative + .decode + .ngram + .as_ref() + .map_or(0, |ngram| ngram.min_ngram), + ngram_max: self + .speculative + .decode + .ngram + .as_ref() + .map_or(0, |ngram| ngram.max_proposal_tokens), native_mtp_enabled: self.speculative.native_mtp_enabled, native_mtp_draft_model_path: if self.speculative.native_mtp_enabled { self.speculative.draft_model_path.clone() } else { None }, - native_mtp_max_tokens: if self.speculative.native_mtp_enabled { - self.speculative.draft_max_tokens as usize - } else { - 0 - }, - native_mtp_min_tokens: if self.speculative.native_mtp_enabled { - self.speculative.draft_min_tokens as usize - } else { - 0 - }, + native_mtp_max_tokens: self.speculative.decode.native_mtp.max_draft_tokens, + native_mtp_min_tokens: self.speculative.decode.native_mtp.min_draft_tokens, activation_width, wire_dtype: self.skippy.activation_wire_dtype.into(), reply_credit_limit: None, @@ -301,7 +297,9 @@ impl ResolvedSkippyConfig { fn speculative_mode_for_embedded(&self, _staged: bool) -> &'static str { if self.speculative.mode == "draft" && self.speculative.draft_model_path.is_some() { "draft" - } else if self.speculative.mode == "ngram" && self.speculative.ngram_min > 0 { + } else if self.speculative.decode.ngram.is_some() + && !self.speculative.decode.native_mtp.enabled + { "ngram" } else { "disabled" @@ -311,7 +309,14 @@ impl ResolvedSkippyConfig { fn speculative_window_for_embedded(&self, mode: &str) -> usize { match mode { "draft" => self.speculative.draft_max_tokens as usize, - "ngram" => self.speculative.ngram_max as usize, + "ngram" => self + .speculative + .decode + .ngram + .as_ref() + .map_or(self.speculative.ngram_max as usize, |ngram| { + ngram.max_proposal_tokens + }), _ => 0, } } @@ -356,6 +361,12 @@ impl ResolvedSkippyConfig { } } +impl ResolvedSkippyConfig { + fn speculative_decode_config(&self) -> SpeculativeDecodeConfig { + self.speculative.decode.clone() + } +} + impl ResolvedEmbeddedOpenAiArgs { pub(crate) fn direct_single_stage_defaults( model_id: String, @@ -379,6 +390,26 @@ impl ResolvedEmbeddedOpenAiArgs { speculative_window: 0, adaptive_speculative_window: false, draft_n_gpu_layers: None, + speculative: SpeculativeDecodeConfig { + native_mtp: NativeMtpProposalConfig { + enabled: native_mtp_enabled, + max_draft_tokens: if native_mtp_enabled { + DEFAULT_NATIVE_MTP_MAX_TOKENS + } else { + 1 + }, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }, + effective_strategy: if native_mtp_enabled { + "native-mtp".to_string() + } else { + "disabled".to_string() + }, + ..SpeculativeDecodeConfig::default() + }, ngram_min: 0, ngram_max: 0, native_mtp_enabled, @@ -419,6 +450,26 @@ impl ResolvedEmbeddedOpenAiArgs { speculative_window: 0, adaptive_speculative_window: false, draft_n_gpu_layers: None, + speculative: SpeculativeDecodeConfig { + native_mtp: NativeMtpProposalConfig { + enabled: native_mtp_enabled, + max_draft_tokens: if native_mtp_enabled { + DEFAULT_NATIVE_MTP_MAX_TOKENS + } else { + 1 + }, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }, + effective_strategy: if native_mtp_enabled { + "native-mtp".to_string() + } else { + "disabled".to_string() + }, + ..SpeculativeDecodeConfig::default() + }, ngram_min: 0, ngram_max: 0, native_mtp_enabled, @@ -462,6 +513,7 @@ impl ResolvedEmbeddedOpenAiArgs { speculative_window: self.speculative_window, adaptive_speculative_window: self.adaptive_speculative_window, draft_n_gpu_layers: self.draft_n_gpu_layers, + speculative: self.speculative, ngram_min: self.ngram_min, ngram_max: self.ngram_max, native_mtp_enabled: self.native_mtp_enabled, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index 128e68cbb4..dd0f3c1740 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -2,7 +2,7 @@ use std::path::{Path, PathBuf}; use skippy_protocol::{FlashAttentionType, StageKvCacheMode, StageKvCachePayload}; use skippy_runtime::package::PackageGenerationInfo; -use skippy_server::EmbeddedOpenAiRequestDefaults; +use skippy_server::{EmbeddedOpenAiRequestDefaults, SpeculativeDecodeConfig}; use super::super::StageWireDType; use crate::plugin::{MeshConfig, ReasoningBudget, ReasoningEnabled, RequestDefaultsConfig}; @@ -102,6 +102,7 @@ pub(crate) struct ResolvedSpeculativeConfig { pub(crate) draft_n_gpu_layers: Option, pub(crate) ngram_min: u32, pub(crate) ngram_max: u32, + pub(crate) decode: SpeculativeDecodeConfig, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -157,6 +158,7 @@ pub(crate) struct ResolvedEmbeddedOpenAiArgs { pub(crate) speculative_window: usize, pub(crate) adaptive_speculative_window: bool, pub(crate) draft_n_gpu_layers: Option, + pub(crate) speculative: SpeculativeDecodeConfig, pub(crate) ngram_min: usize, pub(crate) ngram_max: usize, pub(crate) native_mtp_enabled: bool, diff --git a/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs b/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs index c6f102d41f..0fbf8bf118 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs @@ -458,6 +458,16 @@ pub struct StageRuntimeStatus { pub shutdown_generation: u64, } +/// Classifies a stage status-refresh failure so that only a definitive +/// "peer answered but has no such stage" outcome marks the stage Failed. A +/// transient failure (unreachable peer, timeout) retains the last-known status +/// instead of tearing down a healthy split on a momentary blip. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum StageStatusRefreshFailure { + MissingFromRuntime, + Transient, +} + #[derive(Clone, Debug, Default)] pub(crate) struct StageTopologyState { pub(crate) topologies: HashMap, @@ -551,14 +561,22 @@ impl StageTopologyState { pub(crate) fn record_status_refresh_failure( &mut self, status: &StageRuntimeStatus, - error: String, + failure: StageStatusRefreshFailure, ) { + // A transient refresh failure (peer briefly unreachable, request + // timeout) must NOT mark the stage Failed — that discards a still-valid + // last-known status and can wrongly tear down a healthy split on a + // momentary blip. Only a definitive "missing from runtime" signal, where + // the peer answered but has no such stage, marks the stage Failed. + if failure == StageStatusRefreshFailure::Transient { + return; + } self.record_status(stage_runtime_status_from_snapshot( status.node_id, stage_snapshot_from_runtime_status( status, crate::inference::skippy::StageRuntimeState::Failed, - Some(error), + Some("stage status missing from runtime".to_string()), ), )); } @@ -727,73 +745,110 @@ impl Node { pub async fn refresh_stage_runtime_statuses(&self, timeout: std::time::Duration) { let active_statuses = self.stage_topologies.lock().await.active_statuses(); for status in active_statuses { - if status.stage_index == 0 { - continue; - } - let Some(peer_id) = status.node_id else { - continue; - }; - let filter = crate::inference::skippy::StageStatusFilter { - topology_id: Some(status.topology_id.clone()), - run_id: Some(status.run_id.clone()), - stage_id: Some(status.stage_id.clone()), - }; - let refresh = async { - if peer_id == self.endpoint.id() { - self.query_local_stage_status(filter) - .await - .map(crate::inference::skippy::StageControlResponse::Status) - } else { - self.send_stage_control( - peer_id, - crate::inference::skippy::StageControlRequest::Status(filter), - ) + self.refresh_stage_runtime_status(status, timeout).await; + } + } + + async fn refresh_stage_runtime_status( + &self, + status: StageRuntimeStatus, + timeout: std::time::Duration, + ) { + if status.stage_index == 0 { + return; + } + let Some(peer_id) = status.node_id else { + return; + }; + let filter = crate::inference::skippy::StageStatusFilter { + topology_id: Some(status.topology_id.clone()), + run_id: Some(status.run_id.clone()), + stage_id: Some(status.stage_id.clone()), + }; + let refresh = async { + if peer_id == self.endpoint.id() { + self.query_local_stage_status(filter) .await - } - }; - match tokio::time::timeout(timeout, refresh).await { - Ok(Ok(crate::inference::skippy::StageControlResponse::Status(statuses))) => { - if statuses.is_empty() { - self.stage_topologies - .lock() - .await - .record_status_refresh_failure( - &status, - "stage status missing from runtime".to_string(), - ); - } else { - for status in statuses { - self.record_stage_status(Some(peer_id), status).await; - } - } - } - Ok(Ok(crate::inference::skippy::StageControlResponse::Ready(ready))) => { - self.record_stage_status(Some(peer_id), ready.status).await; - } - Ok(Ok(_)) => {} - Ok(Err(error)) => { - self.stage_topologies - .lock() - .await - .record_status_refresh_failure(&status, error.to_string()); - } - Err(_) => { + .map(crate::inference::skippy::StageControlResponse::Status) + } else { + self.send_stage_control( + peer_id, + crate::inference::skippy::StageControlRequest::Status(filter), + ) + .await + } + }; + match tokio::time::timeout(timeout, refresh).await { + Ok(Ok(response)) => { + self.record_stage_status_refresh(&status, peer_id, response) + .await; + } + Ok(Err(error)) => { + self.record_transient_stage_status_refresh_failure(&status, peer_id, Some(&error)) + .await; + } + Err(_) => { + self.record_transient_stage_status_refresh_failure(&status, peer_id, None) + .await; + } + } + } + + async fn record_stage_status_refresh( + &self, + status: &StageRuntimeStatus, + peer_id: EndpointId, + response: crate::inference::skippy::StageControlResponse, + ) { + match response { + crate::inference::skippy::StageControlResponse::Status(statuses) => { + if statuses.is_empty() { self.stage_topologies .lock() .await .record_status_refresh_failure( - &status, - "stage status refresh timed out".to_string(), + status, + StageStatusRefreshFailure::MissingFromRuntime, ); - tracing::debug!( - topology_id = %status.topology_id, - run_id = %status.run_id, - stage_id = %status.stage_id, - peer = %peer_id.fmt_short(), - "stage status refresh timed out; marking stage failed" - ); + return; } + for status in statuses { + self.record_stage_status(Some(peer_id), status).await; + } + } + crate::inference::skippy::StageControlResponse::Ready(ready) => { + self.record_stage_status(Some(peer_id), ready.status).await; } + _ => {} + } + } + + async fn record_transient_stage_status_refresh_failure( + &self, + status: &StageRuntimeStatus, + peer_id: EndpointId, + error: Option<&anyhow::Error>, + ) { + self.stage_topologies + .lock() + .await + .record_status_refresh_failure(status, StageStatusRefreshFailure::Transient); + match error { + Some(error) => tracing::debug!( + topology_id = %status.topology_id, + run_id = %status.run_id, + stage_id = %status.stage_id, + peer = %peer_id.fmt_short(), + error = %error, + "stage status refresh failed; retaining last stage status" + ), + None => tracing::debug!( + topology_id = %status.topology_id, + run_id = %status.run_id, + stage_id = %status.stage_id, + peer = %peer_id.fmt_short(), + "stage status refresh timed out; retaining last stage status" + ), } } diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests/requirements.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/requirements.rs index f945c4a3c6..c0bbb41fd3 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests/requirements.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests/requirements.rs @@ -922,7 +922,7 @@ fn active_stage_refresh_marks_missing_stage_failed() { } #[test] -fn active_stage_refresh_timeout_marks_cached_stage_failed() { +fn active_stage_missing_from_runtime_marks_cached_stage_failed() { let node_id = EndpointId::from(SecretKey::from_bytes(&[0x43; 32]).public()); let mut state = StageTopologyState::default(); state.record_status(test_stage_status( @@ -934,7 +934,10 @@ fn active_stage_refresh_timeout_marks_cached_stage_failed() { )); let cached = state.active_statuses().into_iter().next().unwrap(); - state.record_status_refresh_failure(&cached, "stage status refresh timed out".to_string()); + state.record_status_refresh_failure( + &cached, + crate::mesh::stage_transport::StageStatusRefreshFailure::MissingFromRuntime, + ); let status = state.runtime_statuses().into_iter().next().unwrap(); assert_eq!( @@ -943,8 +946,37 @@ fn active_stage_refresh_timeout_marks_cached_stage_failed() { ); assert_eq!( status.error.as_deref(), - Some("stage status refresh timed out") + Some("stage status missing from runtime") + ); +} + +#[test] +fn active_stage_refresh_timeout_retains_cached_stage_status() { + // A transient refresh failure (timeout / unreachable peer) must retain the + // last-known Ready status rather than tearing down a healthy split on a + // momentary blip. + let node_id = EndpointId::from(SecretKey::from_bytes(&[0x43; 32]).public()); + let mut state = StageTopologyState::default(); + state.record_status(test_stage_status( + node_id, + "stage-1", + 1, + "127.0.0.1:51234", + crate::inference::skippy::StageRuntimeState::Ready, + )); + let cached = state.active_statuses().into_iter().next().unwrap(); + + state.record_status_refresh_failure( + &cached, + crate::mesh::stage_transport::StageStatusRefreshFailure::Transient, + ); + + let status = state.runtime_statuses().into_iter().next().unwrap(); + assert_eq!( + status.state, + crate::inference::skippy::StageRuntimeState::Ready ); + assert_eq!(status.error, None); } #[test] diff --git a/crates/mesh-llm-host-runtime/src/plugin/mod.rs b/crates/mesh-llm-host-runtime/src/plugin/mod.rs index 404cd6e135..8c733c225e 100644 --- a/crates/mesh-llm-host-runtime/src/plugin/mod.rs +++ b/crates/mesh-llm-host-runtime/src/plugin/mod.rs @@ -49,16 +49,16 @@ pub use self::config::ExternalPluginSpec; pub(crate) use self::config::{ BoolOrAuto, HardwareConfig, IntegerOrString, ModelConfigDefaults, ModelFitConfig, MultimodalConfig, ReasoningBudget, ReasoningEnabled, RequestDefaultsConfig, SkippyConfig, - SpeculativeConfig, StringOrStringList, ThroughputConfig, + StringOrStringList, ThroughputConfig, }; #[allow(unused_imports)] pub use self::config::{ ConfigEditor, ConfigStore, GpuAssignment, GpuConfig, LocalServingNodeConfig, MeshConfig, MeshRequirementsConfig, ModelConfigEditor, ModelConfigEntry, ModelDefaultsEditor, ModelRuntimeKind, OwnerControlConfig, PluginConfigEditor, PluginConfigEntry, PluginHostMode, - PluginStartupConfig, PluginWebUiPreference, ResolvedPlugins, TelemetryConfig, - TelemetryMetricsConfig, bundled_cli_plugin_spec, config_path, config_to_toml, load_config, - parse_config_toml, resolve_plugins, validate_config_file, + PluginStartupConfig, PluginWebUiPreference, ResolvedPlugins, SpeculativeConfig, + TelemetryConfig, TelemetryMetricsConfig, bundled_cli_plugin_spec, config_path, config_to_toml, + load_config, parse_config_toml, resolve_plugins, validate_config_file, }; #[cfg(test)] pub(crate) use self::config::{ diff --git a/crates/mesh-llm-host-runtime/src/runtime/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/mod.rs index ffa5b8a947..25b72979d1 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/mod.rs @@ -3488,20 +3488,9 @@ async fn prepare_runtime_startup( .iter() .map(|model| model.resolved_path.clone()) .collect(); - let update_check_paths = resolved_models.clone(); - match tokio::task::spawn_blocking(move || { - models::warn_about_updates_for_paths(&update_check_paths); - }) - .await - { - Ok(()) => {} - Err(err) => { - let _ = emit_event(OutputEvent::Warning { - message: format!("Could not join Hugging Face update check task: {err}"), - context: None, - }); - } - } + spawn_advisory_startup_task(move || { + models::warn_about_updates_for_paths(&resolved_models); + }); let requested_model_names = startup_models .iter() @@ -3514,6 +3503,11 @@ async fn prepare_runtime_startup( })) } +// Snapshot update checks are advisory. Serving must not wait on Hub reachability. +fn spawn_advisory_startup_task(task: impl FnOnce() + Send + 'static) { + std::mem::drop(tokio::task::spawn_blocking(task)); +} + pub(crate) async fn run() -> Result<()> { initialize_runtime_entrypoint()?; run_runtime_cli(RuntimeOptions::default(), None, None, None).await @@ -3613,8 +3607,9 @@ async fn run_runtime_cli( autoupdate::check_for_update(crate::BUILD_VERSION).await; } - let config = plugin::load_config(options.config.as_deref())?; + let mut config = plugin::load_config(options.config.as_deref())?; apply_runtime_config_options(&mut options, &config); + apply_runtime_cli_speculative_overrides(&mut config, options.speculative_overrides.as_ref()); let startup_mesh_creation_state = resolve_startup_mesh_creation_state(&options, &config)?; let cli_has_explicit_models = cli_has_explicit_models(&options); let has_config_models = !config.models.is_empty(); @@ -3674,6 +3669,35 @@ fn apply_runtime_config_options(options: &mut RuntimeOptions, config: &plugin::M options.listen_all |= config.runtime.listen_all; } +fn apply_runtime_cli_speculative_overrides( + config: &mut plugin::MeshConfig, + overrides: Option<&plugin::SpeculativeConfig>, +) { + let Some(overrides) = overrides else { + return; + }; + let defaults = config + .defaults + .as_ref() + .and_then(|defaults| defaults.speculative.as_ref()) + .cloned(); + + let resolved_defaults = + plugin::SpeculativeConfig::with_precedence(Some(overrides), None, defaults.as_ref()); + config + .defaults + .get_or_insert_with(plugin::ModelConfigDefaults::default) + .speculative = Some(resolved_defaults); + + for model in &mut config.models { + model.speculative = Some(plugin::SpeculativeConfig::with_precedence( + Some(overrides), + model.speculative.as_ref(), + defaults.as_ref(), + )); + } +} + #[cfg(test)] fn runtime_options_for_test(args: &[&str]) -> RuntimeOptions { let mut options = RuntimeOptions::default(); @@ -5890,18 +5914,17 @@ async fn attach_local_release_attestation(node: &mesh::Node) -> Result<()> { } fn skippy_telemetry_options(options: &RuntimeOptions) -> skippy::SkippyTelemetryOptions { - if !options.debug { - return skippy::SkippyTelemetryOptions::off(); + let endpoint = options + .skippy_metrics_otlp_grpc + .as_deref() + .map(str::trim) + .filter(|endpoint| !endpoint.is_empty()) + .map(str::to_owned); + match (endpoint, options.debug) { + (Some(endpoint), true) => skippy::SkippyTelemetryOptions::debug(Some(endpoint)), + (Some(endpoint), false) => skippy::SkippyTelemetryOptions::summary(endpoint), + (None, _) => skippy::SkippyTelemetryOptions::off(), } - - skippy::SkippyTelemetryOptions::debug( - options - .skippy_metrics_otlp_grpc - .as_deref() - .map(str::trim) - .filter(|endpoint| !endpoint.is_empty()) - .map(str::to_owned), - ) } fn configure_run_auto_process_state( @@ -9206,6 +9229,28 @@ mod tests { } } + #[tokio::test] + async fn advisory_startup_task_does_not_block_runtime_startup() { + let started = std::sync::Arc::new(tokio::sync::Notify::new()); + let completed = std::sync::Arc::new(AtomicBool::new(false)); + let task_started = std::sync::Arc::clone(&started); + let task_completed = std::sync::Arc::clone(&completed); + + spawn_advisory_startup_task(move || { + task_started.notify_one(); + std::thread::sleep(Duration::from_millis(100)); + task_completed.store(true, Ordering::Release); + }); + + tokio::time::timeout(Duration::from_secs(1), started.notified()) + .await + .expect("advisory task should be scheduled"); + assert!( + !completed.load(Ordering::Acquire), + "startup must not wait for the advisory task" + ); + } + #[test] fn noq_proto_tracing_messages_use_transport_context() { let message = "2026-06-11T03:49:18.033043Z WARN noq_proto::connection: err=LastOpenPath failed closing path"; @@ -9216,6 +9261,37 @@ mod tests { assert_eq!(context.as_deref(), Some("transport")); } + #[test] + fn cli_speculative_overrides_take_precedence_without_dropping_model_tuning() { + let mut config: plugin::MeshConfig = toml::from_str( + r#" +[defaults.speculative] +strategy = "mtp-cache" +verify_window_pipeline_depth = 2 + +[[models]] +model = "test/model" + +[models.speculative] +ngram_max_proposal_tokens = 6 +"#, + ) + .expect("config parses"); + let mut overrides = plugin::SpeculativeConfig::default(); + overrides.strategy = Some("mtp".to_string()); + overrides.verify_window_pipeline_depth = Some(3); + + apply_runtime_cli_speculative_overrides(&mut config, Some(&overrides)); + + let model = config.models[0] + .speculative + .as_ref() + .expect("model speculative config is resolved"); + assert_eq!(model.strategy.as_deref(), Some("mtp")); + assert_eq!(model.ngram_max_proposal_tokens, Some(6)); + assert_eq!(model.verify_window_pipeline_depth, Some(3)); + } + #[test] fn routed_tracing_messages_strip_ansi_sequences() { let formatted = "\u{1b}[2m2026-06-11T03:49:18.033043Z\u{1b}[0m \u{1b}[33m WARN\u{1b}[0m"; @@ -10688,6 +10764,41 @@ mod tests { assert!(metrics.contains(&hardware::Metric::IsSoc)); } + #[test] + fn skippy_telemetry_endpoint_enables_summary_without_debug() { + let options = RuntimeOptions { + skippy_metrics_otlp_grpc: Some("http://127.0.0.1:14317".to_string()), + ..RuntimeOptions::default() + }; + + let telemetry = skippy_telemetry_options(&options); + + assert_eq!( + telemetry.metrics_otlp_grpc.as_deref(), + Some("http://127.0.0.1:14317") + ); + assert_eq!( + telemetry.level, + skippy_server::telemetry::TelemetryLevel::Summary + ); + } + + #[test] + fn skippy_telemetry_debug_keeps_debug_level_when_endpoint_is_set() { + let options = RuntimeOptions { + debug: true, + skippy_metrics_otlp_grpc: Some("http://127.0.0.1:14317".to_string()), + ..RuntimeOptions::default() + }; + + let telemetry = skippy_telemetry_options(&options); + + assert_eq!( + telemetry.level, + skippy_server::telemetry::TelemetryLevel::Debug + ); + } + #[test] fn pinned_gpu_startup_preflight_cli_models_bypass_config_gpu_id() { let options = runtime_options_for_test(&["mesh-llm", "--model", "Qwen3-8B-Q4_K_M"]); diff --git a/crates/mesh-llm-host-runtime/src/runtime/options.rs b/crates/mesh-llm-host-runtime/src/runtime/options.rs index 20b2918c2f..ef53ecfa68 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/options.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/options.rs @@ -5,6 +5,7 @@ use mesh_llm_events::LogFormat; use crate::crypto::TrustPolicy; use crate::discovery::MeshDiscoveryMode; +use crate::plugin::SpeculativeConfig; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum RuntimeSurface { @@ -56,6 +57,7 @@ pub struct RuntimeOptions { pub draft: Option, pub draft_max: u16, pub no_draft: bool, + pub speculative_overrides: Option, pub split: bool, pub ctx_size: Option, pub max_vram: Option, @@ -121,6 +123,7 @@ impl Default for RuntimeOptions { draft: None, draft_max: 8, no_draft: false, + speculative_overrides: None, split: false, ctx_size: None, max_vram: None, diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json index f9353261be..af06512f8b 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json @@ -819,6 +819,27 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.extension_initial_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.extension_max_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.extension_tail_backoff_proposals", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.mode", "support": "supported", @@ -826,6 +847,27 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.native_mtp_reject_cooldown_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.native_mtp_suppress_cooldown_draft_limit", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.native_mtp_suppress_cooldown_drafts", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.ngram_max", "support": "supported", @@ -833,6 +875,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.ngram_max_proposal_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.ngram_min", "support": "supported", @@ -840,6 +889,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.ngram_proposer", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.pairing_fault", "support": "supported", @@ -854,6 +910,34 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.strategy", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.verify_window_max_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.verify_window_min_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, + { + "canonical_path": "defaults.speculative.verify_window_pipeline_depth", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.throughput.continuous_batching", "support": "supported", diff --git a/crates/mesh-llm/src/lib.rs b/crates/mesh-llm/src/lib.rs index a90df5cc5a..d967980699 100644 --- a/crates/mesh-llm/src/lib.rs +++ b/crates/mesh-llm/src/lib.rs @@ -131,6 +131,7 @@ fn runtime_help_text() -> Option { } fn runtime_options_from_cli(cli: mesh_llm_cli::Cli) -> mesh_llm_host_runtime::RuntimeOptions { + let speculative_overrides = speculative_overrides_from_cli(&cli); mesh_llm_host_runtime::RuntimeOptions { log_format: cli.log_format, debug: cli.debug, @@ -166,6 +167,7 @@ fn runtime_options_from_cli(cli: mesh_llm_cli::Cli) -> mesh_llm_host_runtime::Ru draft: cli.draft, draft_max: cli.draft_max, no_draft: cli.no_draft, + speculative_overrides, split: cli.split, ctx_size: cli.ctx_size, max_vram: cli.max_vram, @@ -195,6 +197,37 @@ fn runtime_options_from_cli(cli: mesh_llm_cli::Cli) -> mesh_llm_host_runtime::Ru } } +fn speculative_overrides_from_cli( + cli: &mesh_llm_cli::Cli, +) -> Option { + let suppress_cooldown_drafts = if cli.speculative_native_mtp_allow_cooldown_drafts { + Some(false) + } else { + cli.speculative_native_mtp_suppress_cooldown_drafts + .then_some(true) + }; + let mut overrides = mesh_llm_host_runtime::plugin::SpeculativeConfig::default(); + overrides.strategy = cli.speculative_strategy.clone(); + overrides.ngram_proposer = cli + .speculative_ngram_proposer + .map(mesh_llm_cli::SpeculativeNgramProposerCli::as_str) + .map(str::to_string); + overrides.ngram_min = cli.speculative_ngram_min; + overrides.ngram_max = cli.speculative_ngram_max; + overrides.ngram_max_proposal_tokens = cli.speculative_ngram_max_proposal_tokens; + overrides.extension_initial_tokens = cli.speculative_extension_initial_tokens; + overrides.extension_max_tokens = cli.speculative_extension_max_tokens; + overrides.extension_tail_backoff_proposals = cli.speculative_extension_tail_backoff_proposals; + overrides.native_mtp_reject_cooldown_tokens = cli.speculative_native_mtp_reject_cooldown_tokens; + overrides.native_mtp_suppress_cooldown_drafts = suppress_cooldown_drafts; + overrides.native_mtp_suppress_cooldown_draft_limit = + cli.speculative_native_mtp_suppress_cooldown_draft_limit; + overrides.verify_window_min_tokens = cli.speculative_verify_window_min_tokens; + overrides.verify_window_max_tokens = cli.speculative_verify_window_max_tokens; + overrides.verify_window_pipeline_depth = cli.speculative_verify_window_pipeline_depth; + (overrides != Default::default()).then_some(overrides) +} + fn command_uses_machine_output(command: Option<&mesh_llm_cli::Command>) -> bool { matches!( command, @@ -279,6 +312,7 @@ fn map_trust_policy( #[cfg(test)] mod cli_entrypoint_tests { + use clap::Parser; use std::ffi::OsString; #[test] @@ -332,4 +366,32 @@ mod cli_entrypoint_tests { OsString::from("--help"), ])); } + + #[test] + fn cli_ngram_override_reaches_runtime_config() { + let normalized = mesh_llm_cli::normalize_runtime_surface_args([ + "mesh-llm", + "serve", + "--speculative-strategy", + "mtp", + "--speculative-ngram-proposer", + "cache", + "--speculative-ngram-min", + "2", + "--speculative-ngram-max", + "6", + "--speculative-ngram-max-proposal-tokens", + "5", + ]); + let cli = mesh_llm_cli::Cli::try_parse_from(normalized.normalized).expect("CLI parses"); + + let config = super::speculative_overrides_from_cli(&cli) + .expect("speculative flags produce an override"); + + assert_eq!(config.strategy.as_deref(), Some("mtp")); + assert_eq!(config.ngram_proposer.as_deref(), Some("cache")); + assert_eq!(config.ngram_min, Some(2)); + assert_eq!(config.ngram_max, Some(6)); + assert_eq!(config.ngram_max_proposal_tokens, Some(5)); + } } diff --git a/crates/openai-frontend/src/completions.rs b/crates/openai-frontend/src/completions.rs index 1bc6977f0e..387dfdca9c 100644 --- a/crates/openai-frontend/src/completions.rs +++ b/crates/openai-frontend/src/completions.rs @@ -140,6 +140,8 @@ pub struct CompletionResponse { pub model: String, pub choices: Vec, pub usage: Usage, + #[serde(skip_serializing_if = "Option::is_none")] + pub timings: Option>, } impl CompletionResponse { @@ -165,8 +167,29 @@ impl CompletionResponse { finish_reason: Some(finish_reason), }], usage, + timings: None, } } + + pub fn with_timings(mut self, timings: Option>) -> Self { + self.timings = timings; + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn completion_response_serializes_optional_timings() { + let response = CompletionResponse::new("model", "text", Usage::new(1, 1)) + .with_timings(Some(BTreeMap::from([("draft_n".to_string(), json!(2))]))); + + let value = serde_json::to_value(response).unwrap(); + assert_eq!(value["timings"]["draft_n"], json!(2)); + } } #[derive(Debug, Clone, Serialize, PartialEq, Eq)] diff --git a/crates/skippy-bench/src/chat_corpus.rs b/crates/skippy-bench/src/chat_corpus.rs index ebdde6c751..5c6022135c 100644 --- a/crates/skippy-bench/src/chat_corpus.rs +++ b/crates/skippy-bench/src/chat_corpus.rs @@ -1,4 +1,5 @@ use std::{ + collections::BTreeMap, fs, io::{BufRead, BufReader}, path::{Path, PathBuf}, @@ -78,6 +79,7 @@ struct ChatCorpusResult { total_tokens: Option, finish_reason: Option, output_chars: usize, + timings: Option>, error: Option, api_error_code: Option, } @@ -99,6 +101,9 @@ struct ChatCorpusSummary { total_wall_ms: f64, completion_tok_s: Option, total_tok_s: Option, + drafted_tokens: u64, + accepted_draft_tokens: u64, + draft_acceptance: Option, } pub fn chat_corpus(args: ChatCorpusArgs) -> Result<()> { @@ -330,6 +335,7 @@ fn parse_json_response( .map(str::chars) .map(Iterator::count) .unwrap_or_default(), + timings: response_timings(&value), error: None, api_error_code: None, } @@ -370,6 +376,7 @@ fn parse_stream_response( let mut prompt_tokens = None; let mut total_tokens = None; let mut finish_reason = None; + let mut timings = None; let mut error = None; let mut api_error_code_value = None; @@ -419,6 +426,9 @@ fn parse_stream_response( .and_then(Value::as_str) .map(ToOwned::to_owned); } + if let Some(response_timings) = response_timings(&value) { + timings = Some(response_timings); + } if let Some(usage) = value.get("usage").filter(|usage| !usage.is_null()) { completion_tokens = usage_u64(usage, "completion_tokens"); prompt_tokens = usage_u64(usage, "prompt_tokens"); @@ -440,6 +450,7 @@ fn parse_stream_response( total_tokens, finish_reason, output_chars, + timings, error, api_error_code: api_error_code_value, } @@ -466,6 +477,7 @@ fn error_result( total_tokens: None, finish_reason: None, output_chars: 0, + timings: None, error: Some(error), api_error_code, } @@ -622,6 +634,14 @@ fn summarize(results: &[ChatCorpusResult], total_wall_ms: f64) -> ChatCorpusSumm .iter() .filter_map(|result| result.total_tokens) .sum::(); + let drafted_tokens = results + .iter() + .filter_map(|result| timing_u64(result, "draft_n")) + .sum::(); + let accepted_draft_tokens = results + .iter() + .filter_map(|result| timing_u64(result, "draft_n_accepted")) + .sum::(); elapsed.sort_by(f64::total_cmp); ttft.sort_by(f64::total_cmp); ChatCorpusSummary { @@ -640,9 +660,21 @@ fn summarize(results: &[ChatCorpusResult], total_wall_ms: f64) -> ChatCorpusSumm total_wall_ms, completion_tok_s: rate(completion_tokens, total_wall_ms), total_tok_s: rate(total_tokens, total_wall_ms), + drafted_tokens, + accepted_draft_tokens, + draft_acceptance: (drafted_tokens > 0) + .then(|| accepted_draft_tokens as f64 / drafted_tokens as f64), } } +fn response_timings(value: &Value) -> Option> { + serde_json::from_value(value.get("timings")?.clone()).ok() +} + +fn timing_u64(result: &ChatCorpusResult, key: &str) -> Option { + result.timings.as_ref()?.get(key)?.as_u64() +} + fn mean(values: &[f64]) -> Option { (!values.is_empty()).then(|| values.iter().sum::() / values.len() as f64) } @@ -812,4 +844,43 @@ mod tests { assert_eq!(body["user"], "repo:turns"); } + + #[test] + fn summary_aggregates_canonical_draft_counters() { + let results = vec![ + result_with_timings(20, 12), + result_with_timings(5, 4), + result_with_timings(0, 0), + ]; + + let summary = summarize(&results, 1_000.0); + + assert_eq!(summary.drafted_tokens, 25); + assert_eq!(summary.accepted_draft_tokens, 16); + assert_eq!(summary.draft_acceptance, Some(0.64)); + } + + fn result_with_timings(drafted: u64, accepted: u64) -> ChatCorpusResult { + ChatCorpusResult { + sequence: 0, + prompt_id: None, + category: None, + length_bucket: None, + session_id: "session".to_string(), + prompt_chars: 1, + elapsed_ms: 1.0, + ttft_ms: None, + completion_tokens: Some(1), + prompt_tokens: Some(1), + total_tokens: Some(2), + finish_reason: Some("stop".to_string()), + output_chars: 1, + timings: Some(BTreeMap::from([ + ("draft_n".to_string(), json!(drafted)), + ("draft_n_accepted".to_string(), json!(accepted)), + ])), + error: None, + api_error_code: None, + } + } } diff --git a/crates/skippy-bench/src/cli.rs b/crates/skippy-bench/src/cli.rs index 0584d122d4..d4c95893c4 100644 --- a/crates/skippy-bench/src/cli.rs +++ b/crates/skippy-bench/src/cli.rs @@ -20,8 +20,8 @@ pub enum CommandKind { LocalSplitBinary(LocalSplitBinaryArgs), LocalSplitCompare(LocalSplitCompareArgs), LocalSplitChainBinary(LocalSplitChainBinaryArgs), - #[command(name = "verify-span-local")] - VerifySpanLocal(VerifySpanLocalArgs), + #[command(name = "verify-window-local")] + VerifyWindowLocal(VerifyWindowLocalArgs), #[command(name = "chat-corpus")] ChatCorpus(ChatCorpusArgs), #[command(name = "token-lengths")] @@ -208,7 +208,7 @@ pub struct TokenLengthsArgs { } #[derive(Parser)] -pub struct VerifySpanLocalArgs { +pub struct VerifyWindowLocalArgs { #[arg(long)] pub model_path: PathBuf, #[arg(long, default_value_t = 48)] @@ -635,10 +635,10 @@ mod tests { } #[test] - fn parses_verify_span_local_command() { + fn parses_verify_window_local_command() { let cli = Cli::try_parse_from([ "skippy-bench", - "verify-span-local", + "verify-window-local", "--model-path", "/tmp/model.gguf", "--layer-end", @@ -652,8 +652,8 @@ mod tests { ]) .unwrap(); - let CommandKind::VerifySpanLocal(args) = cli.command else { - panic!("expected verify-span-local subcommand"); + let CommandKind::VerifyWindowLocal(args) = cli.command else { + panic!("expected verify-window-local subcommand"); }; assert_eq!(args.model_path, PathBuf::from("/tmp/model.gguf")); @@ -665,10 +665,10 @@ mod tests { } #[test] - fn parses_verify_span_local_split_layer() { + fn parses_verify_window_local_split_layer() { let cli = Cli::try_parse_from([ "skippy-bench", - "verify-span-local", + "verify-window-local", "--model-path", "/tmp/model.gguf", "--split-layer", @@ -676,8 +676,8 @@ mod tests { ]) .unwrap(); - let CommandKind::VerifySpanLocal(args) = cli.command else { - panic!("expected verify-span-local subcommand"); + let CommandKind::VerifyWindowLocal(args) = cli.command else { + panic!("expected verify-window-local subcommand"); }; assert_eq!(args.split_layer, Some(24)); diff --git a/crates/skippy-bench/src/evals.rs b/crates/skippy-bench/src/evals.rs index f2b4ca2684..806b538cd0 100644 --- a/crates/skippy-bench/src/evals.rs +++ b/crates/skippy-bench/src/evals.rs @@ -462,9 +462,10 @@ mod tests { doctor::preflight_eval_run, registry::{definition, selected_evals}, run::{ - fill_client_rates, resolved_harness_commit, speed_bench_metrics, - speed_bench_output_path, swe_bench_pro_metrics, swe_bench_pro_output_path, - telemetry_or_unavailable, terminal_bench_metrics, terminal_bench_output_path, + fill_client_rates, resolved_harness_commit, run_artifacts, speed_bench_metrics, + speed_bench_output_path, speed_bench_response_timings_path, swe_bench_pro_metrics, + swe_bench_pro_output_path, telemetry_or_unavailable, terminal_bench_metrics, + terminal_bench_output_path, }, sync::existing_repo_sync_steps, }; @@ -538,6 +539,14 @@ mod tests { .envs .contains(&("SKIPPY_BENCH_BASE_URL".to_string(), args.base_url.clone())) ); + assert!( + command.envs.contains(&( + "SKIPPY_BENCH_RESPONSE_TIMINGS_PATH".to_string(), + speed_bench_response_timings_path(&run_dir) + .display() + .to_string(), + )) + ); assert!(!command.display().contains("test")); assert!( command @@ -547,6 +556,24 @@ mod tests { let launcher = fs::read_to_string(run_dir.join("raw/speed-bench-auth.py")).unwrap(); assert!(launcher.contains("request_origin(url) == benchmark_origin")); assert!(launcher.contains("headers.setdefault(\"Authorization\"")); + assert!(launcher.contains("capture_response_timings")); + assert!(launcher.contains("response.get(\"timings\")")); + assert!(launcher.contains("safe_timings")); + let _ = fs::remove_dir_all(run_dir); + } + + #[test] + fn speed_bench_reports_the_response_timing_artifact() { + let run_dir = temp_run_dir("speed-artifacts"); + let artifacts = run_artifacts(definition(EvalId::SpeedBench), &run_dir); + + assert!(artifacts.iter().any(|artifact| { + artifact.kind == "speed-bench-response-timings" + && artifact.path + == speed_bench_response_timings_path(&run_dir) + .display() + .to_string() + })); let _ = fs::remove_dir_all(run_dir); } diff --git a/crates/skippy-bench/src/evals/adapters/speed_bench.rs b/crates/skippy-bench/src/evals/adapters/speed_bench.rs index 53cb19fdf0..a102364821 100644 --- a/crates/skippy-bench/src/evals/adapters/speed_bench.rs +++ b/crates/skippy-bench/src/evals/adapters/speed_bench.rs @@ -1,9 +1,14 @@ -use super::super::{run::speed_bench_output_path, *}; +use super::super::{ + run::{speed_bench_output_path, speed_bench_response_timings_path}, + *, +}; const AUTH_LAUNCHER: &str = r#"from __future__ import annotations import os import runpy import sys +import json +import threading from urllib.parse import urlparse import requests @@ -27,6 +32,32 @@ def authorized_request(self, method, url, **kwargs): return original_request(self, method, url, **kwargs) requests.sessions.Session.request = authorized_request + +timings_path = os.environ.get("SKIPPY_BENCH_RESPONSE_TIMINGS_PATH") +timings_lock = threading.Lock() +original_json = requests.models.Response.json + +def capture_response_timings(self, *args, **kwargs): + response = original_json(self, *args, **kwargs) + timings = response.get("timings") if isinstance(response, dict) else None + if ( + timings_path + and isinstance(timings, dict) + and not getattr(self, "_skippy_timings_captured", False) + ): + # Preserve only scalar timing counters; never copy request or response content. + safe_timings = { + key: value + for key, value in timings.items() + if isinstance(key, str) and isinstance(value, (bool, int, float)) + } + with timings_lock: + with open(timings_path, "a", encoding="utf-8") as output: + output.write(json.dumps({"timings": safe_timings}, sort_keys=True) + "\\n") + self._skippy_timings_captured = True + return response + +requests.models.Response.json = capture_response_timings script = sys.argv.pop(1) sys.argv[0] = script runpy.run_path(script, run_name="__main__") @@ -80,6 +111,12 @@ pub(in crate::evals) fn speed_bench_command( ) .env("UV_CACHE_DIR", cache_root.join("uv").display().to_string()) .env("SKIPPY_BENCH_BASE_URL", args.base_url.clone()) + .env( + "SKIPPY_BENCH_RESPONSE_TIMINGS_PATH", + speed_bench_response_timings_path(run_dir) + .display() + .to_string(), + ) .secret_env("SKIPPY_BENCH_API_KEY", args.api_key.clone()); Ok(command) } diff --git a/crates/skippy-bench/src/evals/run.rs b/crates/skippy-bench/src/evals/run.rs index 6f5ec0c05a..01f393f05d 100644 --- a/crates/skippy-bench/src/evals/run.rs +++ b/crates/skippy-bench/src/evals/run.rs @@ -128,7 +128,7 @@ pub(super) fn resolved_harness_commit( Ok(Some(commit)) } -fn run_artifacts(definition: EvalDefinition, run_dir: &Path) -> Vec { +pub(super) fn run_artifacts(definition: EvalDefinition, run_dir: &Path) -> Vec { let mut artifacts = vec![ RunArtifact { kind: "stdout", @@ -140,10 +140,18 @@ fn run_artifacts(definition: EvalDefinition, run_dir: &Path) -> Vec }, ]; match definition.id { - EvalId::SpeedBench => artifacts.push(RunArtifact { - kind: "speed-bench-json", - path: speed_bench_output_path(run_dir).display().to_string(), - }), + EvalId::SpeedBench => artifacts.extend([ + RunArtifact { + kind: "speed-bench-json", + path: speed_bench_output_path(run_dir).display().to_string(), + }, + RunArtifact { + kind: "speed-bench-response-timings", + path: speed_bench_response_timings_path(run_dir) + .display() + .to_string(), + }, + ]), EvalId::SweBenchPro => artifacts.extend([ RunArtifact { kind: "swe-bench-pro-sweagent-results", @@ -417,6 +425,10 @@ pub(super) fn speed_bench_output_path(run_dir: &Path) -> PathBuf { run_dir.join("raw/speed-bench.json") } +pub(super) fn speed_bench_response_timings_path(run_dir: &Path) -> PathBuf { + run_dir.join("raw/speed-bench-response-timings.jsonl") +} + pub(super) fn swe_bench_pro_output_path(run_dir: &Path) -> PathBuf { run_dir.join("raw/swe-bench-pro/eval/eval_results.json") } diff --git a/crates/skippy-bench/src/main.rs b/crates/skippy-bench/src/main.rs index dc3ada3e9e..c814fcfb04 100644 --- a/crates/skippy-bench/src/main.rs +++ b/crates/skippy-bench/src/main.rs @@ -8,7 +8,7 @@ mod model_identity; mod support; mod telemetry_report; mod token_lengths; -mod verify_span_local; +mod verify_window_local; use anyhow::Result; use clap::Parser; @@ -23,7 +23,7 @@ use crate::{ local_split_binary, local_split_chain_binary, local_split_compare, local_split_inprocess, }, token_lengths::token_lengths, - verify_span_local::verify_span_local, + verify_window_local::verify_window_local, }; fn main() -> Result<()> { @@ -33,7 +33,7 @@ fn main() -> Result<()> { CommandKind::LocalSplitBinary(args) => local_split_binary(args), CommandKind::LocalSplitCompare(args) => local_split_compare(args), CommandKind::LocalSplitChainBinary(args) => local_split_chain_binary(args), - CommandKind::VerifySpanLocal(args) => verify_span_local(args), + CommandKind::VerifyWindowLocal(args) => verify_window_local(args), CommandKind::ChatCorpus(args) => chat_corpus(args), CommandKind::TokenLengths(args) => token_lengths(args), CommandKind::FocusedRuntime(args) => focused_runtime(args), diff --git a/crates/skippy-bench/src/verify_span_local.rs b/crates/skippy-bench/src/verify_window_local.rs similarity index 97% rename from crates/skippy-bench/src/verify_span_local.rs rename to crates/skippy-bench/src/verify_window_local.rs index b0046fa96b..f47a195d76 100644 --- a/crates/skippy-bench/src/verify_span_local.rs +++ b/crates/skippy-bench/src/verify_window_local.rs @@ -11,7 +11,7 @@ use skippy_runtime::{ parse_cache_type, }; -use crate::cli::VerifySpanLocalArgs; +use crate::cli::VerifyWindowLocalArgs; #[derive(Debug, Serialize)] struct TimingStats { @@ -35,7 +35,7 @@ struct TimingShape { } #[derive(Debug, Serialize)] -struct VerifySpanLocalReport { +struct VerifyWindowLocalReport { mode: &'static str, model_path: PathBuf, layer_end: u32, @@ -90,7 +90,7 @@ struct SplitTimingDiagnostics { serial_stage1: TimingShape, } -pub fn verify_span_local(args: VerifySpanLocalArgs) -> Result<()> { +pub fn verify_window_local(args: VerifyWindowLocalArgs) -> Result<()> { validate_args(&args)?; let output = args.output.clone(); let full = run_full_model_samples(&args)?; @@ -115,7 +115,7 @@ pub fn verify_span_local(args: VerifySpanLocalArgs) -> Result<()> { Ok(()) } -fn validate_args(args: &VerifySpanLocalArgs) -> Result<()> { +fn validate_args(args: &VerifyWindowLocalArgs) -> Result<()> { if args.layer_end == 0 { bail!("layer_end must be greater than zero"); } @@ -130,7 +130,7 @@ fn validate_args(args: &VerifySpanLocalArgs) -> Result<()> { Ok(()) } -fn full_runtime_config(args: &VerifySpanLocalArgs) -> Result { +fn full_runtime_config(args: &VerifyWindowLocalArgs) -> Result { Ok(RuntimeConfig { stage_index: 0, layer_start: 0, @@ -162,7 +162,7 @@ struct FullModelSamples { samples: SampleSet, } -fn run_full_model_samples(args: &VerifySpanLocalArgs) -> Result { +fn run_full_model_samples(args: &VerifyWindowLocalArgs) -> Result { let config = full_runtime_config(args)?; let model = StageModel::open(&args.model_path, &config) .with_context(|| format!("failed to open {}", args.model_path.display()))?; @@ -364,7 +364,7 @@ fn measure_batched( let start = Instant::now(); let prediction = session .verify_tokens_frame_sampled(verify_tokens, Some(&SamplingConfig::default()), None, 0) - .context("batched width-2 VerifySpan failed")? + .context("batched width-2 VerifyWindow failed")? .0; Ok((start.elapsed(), prediction)) } @@ -383,7 +383,7 @@ fn measure_serial( } fn run_split_inprocess_samples( - args: &VerifySpanLocalArgs, + args: &VerifyWindowLocalArgs, split_layer: u32, tokens: &[i32], verify_tokens: &[i32], @@ -416,7 +416,7 @@ fn run_split_inprocess_samples( } fn split_runtime_configs( - args: &VerifySpanLocalArgs, + args: &VerifyWindowLocalArgs, split_layer: u32, ) -> Result<(RuntimeConfig, RuntimeConfig)> { let cache_type_k = parse_cache_type(&args.cache_type_k)?; @@ -567,11 +567,11 @@ fn measure_split_batched( let stage0_start = Instant::now(); let (_stage0_prediction, boundary) = session0 .verify_tokens_frame_sampled(verify_tokens, Some(&SamplingConfig::default()), None, 0) - .context("in-process split stage 0 VerifySpan failed")?; + .context("in-process split stage 0 VerifyWindow failed")?; let stage0 = stage0_start.elapsed(); let boundary_payload_bytes = boundary.payload.len(); if boundary_payload_bytes == 0 { - bail!("in-process split stage 0 produced an empty VerifySpan activation frame"); + bail!("in-process split stage 0 produced an empty VerifyWindow activation frame"); } let stage1_start = Instant::now(); @@ -582,7 +582,7 @@ fn measure_split_batched( Some(&boundary), 0, ) - .context("in-process split stage 1 VerifySpan failed")? + .context("in-process split stage 1 VerifyWindow failed")? .0; let stage1 = stage1_start.elapsed(); Ok(SplitSample { @@ -758,16 +758,16 @@ impl SampleSet { } fn build_report( - args: VerifySpanLocalArgs, + args: VerifyWindowLocalArgs, full: FullModelSamples, split_inprocess_width2: Option, -) -> Result { +) -> Result { let batched_width2 = timing_stats(&full.samples.batched)?; let serial_two_decode_mtp_n1 = timing_stats(&full.samples.serial)?; let batched_avg = batched_width2.avg_us; let serial_avg = serial_two_decode_mtp_n1.avg_us; - Ok(VerifySpanLocalReport { - mode: "verify-span-local", + Ok(VerifyWindowLocalReport { + mode: "verify-window-local", model_path: args.model_path, layer_end: args.layer_end, split_layer: args.split_layer, diff --git a/crates/skippy-correctness/src/runner/native_mtp.rs b/crates/skippy-correctness/src/runner/native_mtp.rs index 982648f17b..591cc4106e 100644 --- a/crates/skippy-correctness/src/runner/native_mtp.rs +++ b/crates/skippy-correctness/src/runner/native_mtp.rs @@ -70,33 +70,21 @@ pub(crate) fn native_mtp_verification_satisfies_requirement( pub(crate) fn native_mtp_sideband_report(reply: &StageReply) -> NativeMtpSidebandReport { let authoritative_token = reply.predicted_tokens.first().copied(); - let advertised_count = reply - .predicted_tokens - .get(1) - .copied() - .and_then(|value| usize::try_from(value).ok()) - .unwrap_or(0); - let available_draft_count = reply.predicted_tokens.len().saturating_sub(3); - let draft_token_count = advertised_count.min(available_draft_count); - let draft_start = 2; - let draft_end = draft_start + draft_token_count; let draft_tokens = reply - .predicted_tokens - .get(draft_start..draft_end) - .unwrap_or(&[]) - .to_vec(); + .native_mtp_draft + .as_ref() + .map_or_else(Vec::new, |draft| draft.token_ids.clone()); let proposal_compute_us = reply - .predicted_tokens - .get(draft_end) - .copied() - .map(|value| i64::from(value.max(0))); + .native_mtp_draft + .as_ref() + .map(|draft| draft.proposal_compute_us.max(0)); NativeMtpSidebandReport { - sideband_present: !draft_tokens.is_empty(), + sideband_present: reply.native_mtp_draft.is_some(), predicted_token_count: reply.predicted_tokens.len(), authoritative_matches_reply: authoritative_token .is_none_or(|token| token == reply.predicted), authoritative_token, - draft_token_count, + draft_token_count: draft_tokens.len(), draft_tokens, proposal_compute_us, } @@ -203,20 +191,26 @@ pub(crate) fn emit_report(report: &T, report_out: Option<&Path>) - #[cfg(test)] mod tests { use super::*; - use skippy_protocol::binary::{StageReplyStats, WireReplyKind}; + use skippy_protocol::binary::{StageNativeMtpDraft, StageReplyStats, WireReplyKind}; - fn predicted_reply(predicted: i32, predicted_tokens: Vec) -> StageReply { + fn predicted_reply( + predicted: i32, + predicted_tokens: Vec, + native_mtp_draft: Option, + ) -> StageReply { StageReply { kind: WireReplyKind::PredictedToken, predicted, predicted_tokens, + native_mtp_draft, + window: Default::default(), stats: StageReplyStats::default(), } } #[test] fn native_mtp_report_treats_plain_authoritative_token_as_no_draft() { - let report = native_mtp_sideband_report(&predicted_reply(11, vec![11])); + let report = native_mtp_sideband_report(&predicted_reply(11, vec![11], None)); assert!(!report.sideband_present); assert_eq!(report.predicted_token_count, 1); @@ -228,11 +222,18 @@ mod tests { } #[test] - fn native_mtp_report_extracts_draft_sideband() { - let report = native_mtp_sideband_report(&predicted_reply(11, vec![11, 2, 12, 13, 34])); + fn native_mtp_report_extracts_typed_draft() { + let report = native_mtp_sideband_report(&predicted_reply( + 11, + vec![11], + Some(StageNativeMtpDraft { + token_ids: vec![12, 13], + proposal_compute_us: 34, + }), + )); assert!(report.sideband_present); - assert_eq!(report.predicted_token_count, 5); + assert_eq!(report.predicted_token_count, 1); assert!(report.authoritative_matches_reply); assert_eq!(report.authoritative_token, Some(11)); assert_eq!(report.draft_token_count, 2); @@ -241,8 +242,15 @@ mod tests { } #[test] - fn native_mtp_report_flags_authoritative_sideband_mismatch() { - let report = native_mtp_sideband_report(&predicted_reply(11, vec![10, 1, 12, 34])); + fn native_mtp_report_flags_authoritative_prediction_mismatch() { + let report = native_mtp_sideband_report(&predicted_reply( + 11, + vec![10], + Some(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: 34, + }), + )); assert!(report.sideband_present); assert!(!report.authoritative_matches_reply); @@ -252,15 +260,29 @@ mod tests { #[test] fn native_mtp_report_clamps_negative_proposal_time() { - let report = native_mtp_sideband_report(&predicted_reply(11, vec![11, 1, 12, -34])); + let report = native_mtp_sideband_report(&predicted_reply( + 11, + vec![11], + Some(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: -34, + }), + )); assert_eq!(report.proposal_compute_us, Some(0)); } #[test] fn native_mtp_requirement_can_require_draft_presence() { - let no_draft = native_mtp_sideband_report(&predicted_reply(11, vec![11])); - let draft = native_mtp_sideband_report(&predicted_reply(11, vec![11, 1, 12, 34])); + let no_draft = native_mtp_sideband_report(&predicted_reply(11, vec![11], None)); + let draft = native_mtp_sideband_report(&predicted_reply( + 11, + vec![11], + Some(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: 34, + }), + )); let optional = NativeMtpRequirement { require_draft: false, }; @@ -275,7 +297,14 @@ mod tests { #[test] fn native_mtp_verification_report_accepts_matching_second_target() { - let first = native_mtp_sideband_report(&predicted_reply(11, vec![11, 2, 12, 13, 34])); + let first = native_mtp_sideband_report(&predicted_reply( + 11, + vec![11], + Some(StageNativeMtpDraft { + token_ids: vec![12, 13], + proposal_compute_us: 34, + }), + )); let report = native_mtp_verification_report(true, &first, Some(12), Some(12), Some(9)) .expect("verification report"); @@ -299,7 +328,14 @@ mod tests { #[test] fn native_mtp_verification_report_rejects_mismatched_draft_without_failing_byte_identity() { - let first = native_mtp_sideband_report(&predicted_reply(11, vec![11, 1, 12, 34])); + let first = native_mtp_sideband_report(&predicted_reply( + 11, + vec![11], + Some(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: 34, + }), + )); let report = native_mtp_verification_report(true, &first, Some(13), Some(13), Some(9)) .expect("verification report"); @@ -320,7 +356,7 @@ mod tests { #[test] fn native_mtp_verification_requirement_fails_when_required_draft_is_missing() { - let first = native_mtp_sideband_report(&predicted_reply(11, vec![11])); + let first = native_mtp_sideband_report(&predicted_reply(11, vec![11], None)); let report = native_mtp_verification_report(true, &first, Some(13), Some(13), Some(9)) .expect("required verification report"); diff --git a/crates/skippy-ffi/README.md b/crates/skippy-ffi/README.md index 2a139d43e5..8f7153c787 100644 --- a/crates/skippy-ffi/README.md +++ b/crates/skippy-ffi/README.md @@ -52,11 +52,10 @@ same Rust crate. ## ABI Contract -The staged ABI is versioned as `0.1.26`. The patch header in -`third_party/llama.cpp/patches/0083-skippy-add-model-open-runtime-events-ABI.patch` -and the Rust constants in `crates/skippy-ffi/src/lib.rs` are the source of -truth, so keep this README aligned with those files instead of treating it as -canonical prose. +The staged ABI is versioned as `0.1.32`. The patch header in +`third_party/llama.cpp/patches/` and the Rust constants in +`crates/skippy-ffi/src/lib.rs` are the source of truth, so keep this README +aligned with those files instead of treating it as canonical prose. Version `0` is still experimental, so callers should treat the ABI as feature-probed rather than permanently stable. @@ -67,6 +66,13 @@ The runtime-event additions are part of that `0.1.26` bump. They add versioned argument on the `_with_events` model-open entrypoints instead of extending `RuntimeConfig`. +Version `0.1.31` adds `skippy_ngram_simple_draft`, a stateless adapter for +llama.cpp's upstream self-speculative `ngram-simple` proposer. + +Version `0.1.32` adds a request-owned adapter for llama.cpp's upstream +`ngram-cache` proposer. Callers reset or append only target-committed tokens; +an optional provisional continuation is read-only input to drafting. + The Rust FFI layer binds `skippy_abi_features`. This README records ABI intent and compatibility expectations only; higher-level gating belongs in `skippy-runtime` or later tasks. @@ -144,6 +150,9 @@ read it directly: | `CHAT_SAMPLING_GRAMMAR` | `1 << 22` | Session-local llama.cpp grammar-constrained sampling from chat template metadata | | `BACKEND_DEVICES` | `1 << 23` | Backend-device capability reporting | | `RUNTIME_EVENTS` | `1 << 24` | `_with_events` model-open entrypoints and runtime-event callbacks | +| `NATIVE_MTP_N1` | `1 << 25` | Typed, non-frame native MTP draft sideband | +| `NGRAM_SIMPLE_DRAFT` | `1 << 26` | llama.cpp upstream `ngram-simple` proposal over accepted token history | +| `NGRAM_CACHE_DRAFT` | `1 << 27` | Stateful request-local llama.cpp `ngram-cache` proposer | Runtime-event compatibility expectations are narrow on purpose: @@ -213,6 +222,15 @@ hook currently bound by this crate. | `skippy_prefill_chunk_frame` | Prefills a token chunk using `ActivationDesc` plus payload buffers. | | `skippy_decode_step_frame_sampled` | Decodes one token with activation-frame I/O and `SamplingConfig`. | +### Self-speculative proposal + +| Function | Purpose | +| --- | --- | +| `skippy_ngram_simple_draft` | Calls llama.cpp's upstream `ngram-simple` proposer with accepted history and returns only the proposed continuation. It owns no persistent cache state. | +| `skippy_ngram_cache_create` / `free` | Allocates or releases a request-owned cache handle. | +| `skippy_ngram_cache_reset` / `append` | Rebuilds or extends the cache with target-committed token history only. | +| `skippy_ngram_cache_draft` | Drafts after committed history and an optional non-mutating continuation prefix, such as a native MTP candidate. | + ### Token and chat helpers | Function | Purpose | diff --git a/crates/skippy-ffi/build.rs b/crates/skippy-ffi/build.rs index 76c285d6a3..f976e6e434 100644 --- a/crates/skippy-ffi/build.rs +++ b/crates/skippy-ffi/build.rs @@ -218,7 +218,7 @@ fn main() { fn default_build_dir(workspace_root: &std::path::Path, target: &str) -> std::path::PathBuf { let suffix = default_backend(target); - workspace_root.join(format!(".deps/llama-build/build-stage-abi-{suffix}")) + workspace_root.join(format!(".deps/llama-build/build-stage-abi-static-{suffix}")) } fn default_backend(target: &str) -> &'static str { diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index aa644918bd..b1798ebf86 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,9 +1,11 @@ pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 30; +pub const ABI_VERSION_PATCH: u32 = 32; pub const FEATURE_BACKEND_DEVICES: u64 = 1 << 23; pub const FEATURE_RUNTIME_EVENTS: u64 = 1 << 24; pub const FEATURE_NATIVE_MTP_N1: u64 = 1 << 25; +pub const FEATURE_NGRAM_SIMPLE_DRAFT: u64 = 1 << 26; +pub const FEATURE_NGRAM_CACHE_DRAFT: u64 = 1 << 27; #[repr(C)] #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -239,6 +241,11 @@ pub struct Model { _private: [u8; 0], } +#[repr(C)] +pub struct NgramCache { + _private: [u8; 0], +} + #[repr(C)] pub struct Session { _private: [u8; 0], @@ -791,6 +798,12 @@ mod dynamic { llama_model_quantize_default_params() -> LlamaModelQuantizeParams; llama_model_quantize(fname_inp: *const c_char, fname_out: *const c_char, params: *const LlamaModelQuantizeParams) -> u32; skippy_error_free(error: *mut Error); + skippy_ngram_simple_draft(token_ids: *const i32, token_count: usize, sampled_token: i32, ngram_size: u16, max_draft_tokens: u16, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; + skippy_ngram_cache_create(ngram_min: u16, ngram_max: u16, out_cache: *mut *mut NgramCache, out_error: *mut *mut Error) -> Status; + skippy_ngram_cache_free(cache: *mut NgramCache); + skippy_ngram_cache_reset(cache: *mut NgramCache, token_ids: *const i32, token_count: usize, out_error: *mut *mut Error) -> Status; + skippy_ngram_cache_append(cache: *mut NgramCache, token_ids: *const i32, token_count: usize, out_error: *mut *mut Error) -> Status; + skippy_ngram_cache_draft(cache: *mut NgramCache, continuation_prefix: *const i32, continuation_prefix_count: usize, max_draft_tokens: u16, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_backend_device_count(out_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_backend_device_at(index: usize, out_device: *mut BackendDevice, out_error: *mut *mut Error) -> Status; skippy_model_open(path: *const c_char, config: *const RuntimeConfig, out_model: *mut *mut Model, out_error: *mut *mut Error) -> Status; @@ -1175,6 +1188,52 @@ unsafe extern "C" { pub fn skippy_error_free(error: *mut Error); + pub fn skippy_ngram_simple_draft( + token_ids: *const i32, + token_count: usize, + sampled_token: i32, + ngram_size: u16, + max_draft_tokens: u16, + output_tokens: *mut i32, + output_token_capacity: usize, + out_token_count: *mut usize, + out_error: *mut *mut Error, + ) -> Status; + + pub fn skippy_ngram_cache_create( + ngram_min: u16, + ngram_max: u16, + out_cache: *mut *mut NgramCache, + out_error: *mut *mut Error, + ) -> Status; + + pub fn skippy_ngram_cache_free(cache: *mut NgramCache); + + pub fn skippy_ngram_cache_reset( + cache: *mut NgramCache, + token_ids: *const i32, + token_count: usize, + out_error: *mut *mut Error, + ) -> Status; + + pub fn skippy_ngram_cache_append( + cache: *mut NgramCache, + token_ids: *const i32, + token_count: usize, + out_error: *mut *mut Error, + ) -> Status; + + pub fn skippy_ngram_cache_draft( + cache: *mut NgramCache, + continuation_prefix: *const i32, + continuation_prefix_count: usize, + max_draft_tokens: u16, + output_tokens: *mut i32, + output_token_capacity: usize, + out_token_count: *mut usize, + out_error: *mut *mut Error, + ) -> Status; + pub fn skippy_backend_device_count(out_count: *mut usize, out_error: *mut *mut Error) -> Status; diff --git a/crates/skippy-model-package/src/main.rs b/crates/skippy-model-package/src/main.rs index 60985d0bf6..927f8940a7 100644 --- a/crates/skippy-model-package/src/main.rs +++ b/crates/skippy-model-package/src/main.rs @@ -231,9 +231,29 @@ struct PackageGeneration { #[derive(Debug, Deserialize, Serialize)] struct PackageSpeculativeDecoding { default: String, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + proposers: BTreeMap, strategies: BTreeMap, } +#[derive(Debug, Deserialize, Serialize)] +struct PackageSpeculativeProposer { + #[serde(rename = "type")] + proposer_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + prediction_depth: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + layer_indices: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + ngram_min: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + ngram_max: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + max_proposal_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + history_scope: Option, +} + #[derive(Debug, Deserialize, Serialize)] struct PackageSpeculativeStrategy { #[serde(rename = "type")] @@ -244,6 +264,21 @@ struct PackageSpeculativeStrategy { layer_indices: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] window_policy: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + proposer: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + primary: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + extender: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + extension_policy: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +struct PackageExtensionPolicy { + initial_tokens: u32, + max_tokens: u32, + tail_backoff_proposals: u32, } #[derive(Debug, Deserialize, Serialize)] @@ -1771,6 +1806,19 @@ fn package_generation(tensors: &[TensorInfo]) -> Option { let strategy_id = "mtp".to_string(); let mut strategies = BTreeMap::new(); + let mut proposers = BTreeMap::new(); + proposers.insert( + strategy_id.clone(), + PackageSpeculativeProposer { + proposer_type: "native-mtp".to_string(), + prediction_depth: Some(1), + layer_indices: mtp_layers.clone(), + ngram_min: None, + ngram_max: None, + max_proposal_tokens: None, + history_scope: None, + }, + ); strategies.insert( strategy_id.clone(), PackageSpeculativeStrategy { @@ -1783,12 +1831,17 @@ fn package_generation(tensors: &[TensorInfo]) -> Option { min_window: 1, max_window: 1, }), + proposer: Some(strategy_id.clone()), + primary: None, + extender: None, + extension_policy: None, }, ); Some(PackageGeneration { speculative_decoding: Some(PackageSpeculativeDecoding { default: strategy_id, + proposers, strategies, }), }) @@ -2112,11 +2165,19 @@ mod tests { .speculative_decoding .expect("MTP generation should configure speculative decoding"); assert_eq!(speculative.default, "mtp"); + let proposer = speculative + .proposers + .get("mtp") + .expect("native MTP proposer should be present"); + assert_eq!(proposer.proposer_type, "native-mtp"); + assert_eq!(proposer.prediction_depth, Some(1)); + assert_eq!(proposer.layer_indices, vec![47]); let strategy = speculative .strategies .get("mtp") .expect("default strategy should be present"); assert_eq!(strategy.strategy_type, "native-mtp"); + assert_eq!(strategy.proposer.as_deref(), Some("mtp")); assert_eq!(strategy.prediction_depth, Some(1)); assert_eq!(strategy.layer_indices, vec![47]); let window = strategy diff --git a/crates/skippy-model-package/src/preflight.rs b/crates/skippy-model-package/src/preflight.rs index 8635d852de..df34d5423e 100644 --- a/crates/skippy-model-package/src/preflight.rs +++ b/crates/skippy-model-package/src/preflight.rs @@ -85,9 +85,30 @@ pub(crate) struct PreflightGeneration { #[derive(Debug, Serialize)] pub(crate) struct PreflightSpeculativeDecoding { pub default: String, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub proposers: Vec, pub strategies: Vec, } +#[derive(Debug, Serialize)] +pub(crate) struct PreflightSpeculativeProposer { + pub name: String, + #[serde(rename = "type")] + pub proposer_type: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub prediction_depth: Option, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub layer_indices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub ngram_min: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ngram_max: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_proposal_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub history_scope: Option, +} + #[derive(Debug, Serialize)] pub(crate) struct PreflightSpeculativeStrategy { pub name: String, @@ -99,6 +120,21 @@ pub(crate) struct PreflightSpeculativeStrategy { pub layer_indices: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub window_policy: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub proposer: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub primary: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extender: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extension_policy: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct PreflightExtensionPolicy { + pub initial_tokens: u32, + pub max_tokens: u32, + pub tail_backoff_proposals: u32, } #[derive(Debug, Serialize)] @@ -150,9 +186,29 @@ struct PackageGeneration { struct PackageSpeculativeDecoding { default: String, #[serde(default)] + proposers: BTreeMap, + #[serde(default)] strategies: BTreeMap, } +#[derive(Debug, Deserialize)] +struct PackageSpeculativeProposer { + #[serde(rename = "type")] + proposer_type: String, + #[serde(default)] + prediction_depth: Option, + #[serde(default)] + layer_indices: Vec, + #[serde(default)] + ngram_min: Option, + #[serde(default)] + ngram_max: Option, + #[serde(default)] + max_proposal_tokens: Option, + #[serde(default)] + history_scope: Option, +} + #[derive(Debug, Deserialize)] struct PackageSpeculativeStrategy { #[serde(rename = "type")] @@ -163,6 +219,21 @@ struct PackageSpeculativeStrategy { layer_indices: Vec, #[serde(default)] window_policy: Option, + #[serde(default)] + proposer: Option, + #[serde(default)] + primary: Option, + #[serde(default)] + extender: Option, + #[serde(default)] + extension_policy: Option, +} + +#[derive(Debug, Deserialize)] +struct PackageExtensionPolicy { + initial_tokens: u32, + max_tokens: u32, + tail_backoff_proposals: u32, } #[derive(Debug, Deserialize)] @@ -428,14 +499,53 @@ fn validate_generation( "add the default strategy entry or point default at an existing strategy", ); } + for (name, proposer) in &speculative.proposers { + validate_speculative_proposer(name, proposer, layer_count, report); + } for (name, strategy) in &speculative.strategies { - validate_speculative_strategy(name, strategy, layer_count, report); + validate_speculative_strategy(name, strategy, &speculative.proposers, layer_count, report); + } +} + +fn validate_speculative_proposer( + name: &str, + proposer: &PackageSpeculativeProposer, + layer_count: u32, + report: &mut PackagePreflightReport, +) { + if name.trim().is_empty() { + report.error( + "empty_speculative_proposer_name", + "generation.speculative_decoding proposer names must not be empty", + Some("model-package.json".to_string()), + "use a stable non-empty proposer id such as mtp or ngram-cache", + ); + } + match proposer.proposer_type.as_str() { + "native-mtp" => validate_native_mtp_parts( + name, + proposer.prediction_depth, + &proposer.layer_indices, + layer_count, + report, + ), + "ngram-simple" | "ngram-cache" => validate_ngram_proposer(name, proposer, report), + _ => report.error( + "unsupported_speculative_proposer_type", + format!( + "speculative proposer {name} has unsupported type {}", + proposer.proposer_type + ), + Some("model-package.json".to_string()), + "use native-mtp, ngram-simple, or ngram-cache", + ), } } fn validate_speculative_strategy( name: &str, strategy: &PackageSpeculativeStrategy, + proposers: &BTreeMap, layer_count: u32, report: &mut PackagePreflightReport, ) { @@ -455,21 +565,185 @@ fn validate_speculative_strategy( "set a supported strategy type such as native-mtp", ); } + if let Some(proposer) = &strategy.proposer { + validate_proposer_reference(name, "proposer", proposer, proposers, report); + } if strategy.strategy_type == "native-mtp" { - validate_native_mtp_strategy(name, strategy, layer_count, report); + validate_native_mtp_strategy_proposer_or_inline( + name, + strategy, + proposers, + layer_count, + report, + ); + } + if matches!( + strategy.strategy_type.as_str(), + "ngram-simple" | "ngram-cache" + ) { + validate_ngram_strategy_proposer_type(name, strategy, proposers, report); + } + if strategy.strategy_type == "composite" { + validate_composite_strategy(name, strategy, proposers, report); + } + if let Some(policy) = &strategy.extension_policy { + validate_extension_policy(name, policy, report); } if let Some(window) = &strategy.window_policy { validate_window_policy(name, window, report); } } +fn validate_native_mtp_strategy_proposer_or_inline( + strategy_name: &str, + strategy: &PackageSpeculativeStrategy, + proposers: &BTreeMap, + layer_count: u32, + report: &mut PackagePreflightReport, +) { + let Some(proposer_name) = strategy.proposer.as_deref() else { + validate_native_mtp_strategy(strategy_name, strategy, layer_count, report); + return; + }; + let Some(proposer) = proposers.get(proposer_name) else { + return; + }; + if proposer.proposer_type != "native-mtp" { + report.error( + "native_mtp_strategy_proposer_type_mismatch", + format!( + "native MTP speculative strategy {strategy_name} references proposer {proposer_name} with type {}", + proposer.proposer_type + ), + Some("model-package.json".to_string()), + "set proposer to a declared native-mtp proposer", + ); + } +} + +fn validate_ngram_strategy_proposer_type( + strategy_name: &str, + strategy: &PackageSpeculativeStrategy, + proposers: &BTreeMap, + report: &mut PackagePreflightReport, +) { + let Some(proposer_name) = strategy.proposer.as_deref() else { + report.error( + "missing_ngram_strategy_proposer", + format!("N-gram speculative strategy {strategy_name} must declare a proposer"), + Some("model-package.json".to_string()), + "set proposer to a declared ngram-simple or ngram-cache proposer", + ); + return; + }; + let Some(proposer) = proposers.get(proposer_name) else { + return; + }; + if proposer.proposer_type != strategy.strategy_type { + report.error( + "ngram_strategy_proposer_type_mismatch", + format!( + "N-gram speculative strategy {strategy_name} type {} does not match proposer {proposer_name} type {}", + strategy.strategy_type, proposer.proposer_type + ), + Some("model-package.json".to_string()), + "make the strategy type match its referenced N-gram proposer", + ); + } +} + +fn validate_proposer_reference( + strategy_name: &str, + field: &str, + proposer_name: &str, + proposers: &BTreeMap, + report: &mut PackagePreflightReport, +) { + if !proposers.contains_key(proposer_name) { + report.error( + "missing_speculative_proposer", + format!("speculative strategy {strategy_name} references missing {field} proposer {proposer_name}"), + Some("model-package.json".to_string()), + "declare the referenced proposer under generation.speculative_decoding.proposers", + ); + } +} + +fn validate_composite_strategy( + name: &str, + strategy: &PackageSpeculativeStrategy, + proposers: &BTreeMap, + report: &mut PackagePreflightReport, +) { + let Some(primary) = strategy.primary.as_deref() else { + report.error( + "missing_composite_primary", + format!("composite speculative strategy {name} must declare primary"), + Some("model-package.json".to_string()), + "set primary to a declared native-mtp proposer", + ); + return; + }; + let Some(extender) = strategy.extender.as_deref() else { + report.error( + "missing_composite_extender", + format!("composite speculative strategy {name} must declare extender"), + Some("model-package.json".to_string()), + "set extender to a declared ngram-simple or ngram-cache proposer", + ); + return; + }; + validate_proposer_reference(name, "primary", primary, proposers, report); + validate_proposer_reference(name, "extender", extender, proposers, report); + if proposers + .get(primary) + .is_some_and(|proposer| proposer.proposer_type != "native-mtp") + { + report.error( + "invalid_composite_primary_type", + format!("composite speculative strategy {name} primary {primary} must be native-mtp"), + Some("model-package.json".to_string()), + "set primary to a native-mtp proposer", + ); + } + if proposers.get(extender).is_some_and(|proposer| { + !matches!( + proposer.proposer_type.as_str(), + "ngram-simple" | "ngram-cache" + ) + }) { + report.error( + "invalid_composite_extender_type", + format!("composite speculative strategy {name} extender {extender} must be an N-gram proposer"), + Some("model-package.json".to_string()), + "set extender to an ngram-simple or ngram-cache proposer", + ); + } +} + fn validate_native_mtp_strategy( name: &str, strategy: &PackageSpeculativeStrategy, layer_count: u32, report: &mut PackagePreflightReport, ) { - if strategy.prediction_depth != Some(1) { + validate_native_mtp_parts( + name, + strategy.prediction_depth, + &strategy.layer_indices, + layer_count, + report, + ); +} + +fn validate_native_mtp_parts( + name: &str, + prediction_depth: Option, + layer_indices: &[u32], + layer_count: u32, + report: &mut PackagePreflightReport, +) { + if prediction_depth != Some(1) { report.error( "unsupported_native_mtp_prediction_depth", format!("native MTP strategy {name} must use prediction_depth 1"), @@ -477,7 +751,7 @@ fn validate_native_mtp_strategy( "rebuild the package with the mtp policy supported by this runtime", ); } - if strategy.layer_indices.is_empty() { + if layer_indices.is_empty() { report.error( "missing_native_mtp_layers", format!("native MTP strategy {name} must declare MTP layer_indices"), @@ -485,7 +759,7 @@ fn validate_native_mtp_strategy( "rebuild the package from a GGUF with native MTP tensors", ); } - for layer_index in &strategy.layer_indices { + for layer_index in layer_indices { if *layer_index >= layer_count { report.error( "native_mtp_layer_out_of_range", @@ -499,6 +773,75 @@ fn validate_native_mtp_strategy( } } +fn validate_ngram_proposer( + name: &str, + proposer: &PackageSpeculativeProposer, + report: &mut PackagePreflightReport, +) { + let min = proposer.ngram_min.unwrap_or_default(); + let max = proposer.ngram_max.unwrap_or_default(); + if min == 0 || max == 0 || min > max { + report.error( + "invalid_ngram_proposer_window", + format!("N-gram proposer {name} must set ngram_min and ngram_max with 1 <= min <= max"), + Some("model-package.json".to_string()), + "set positive ngram_min and ngram_max values with min less than or equal to max", + ); + } + if proposer.max_proposal_tokens.unwrap_or_default() == 0 { + report.error( + "invalid_ngram_proposer_max_tokens", + format!("N-gram proposer {name} must set max_proposal_tokens greater than zero"), + Some("model-package.json".to_string()), + "set max_proposal_tokens to a positive value", + ); + } + if proposer.proposer_type == "ngram-cache" + && max as usize > skippy_runtime::NGRAM_CACHE_MAX_NGRAM + { + report.error( + "unsupported_ngram_cache_max_window", + format!( + "N-gram cache proposer {name} ngram_max {max} exceeds llama.cpp limit {}", + skippy_runtime::NGRAM_CACHE_MAX_NGRAM + ), + Some("model-package.json".to_string()), + format!( + "set ngram_max to at most {} while keeping max_proposal_tokens independent", + skippy_runtime::NGRAM_CACHE_MAX_NGRAM + ), + ); + } + if proposer.proposer_type == "ngram-cache" + && proposer.history_scope.as_deref() != Some("request") + { + report.error( + "invalid_ngram_cache_history_scope", + format!("N-gram cache proposer {name} must set history_scope to request"), + Some("model-package.json".to_string()), + "set history_scope to request; shared cache history is not supported", + ); + } +} + +fn validate_extension_policy( + name: &str, + policy: &PackageExtensionPolicy, + report: &mut PackagePreflightReport, +) { + if policy.initial_tokens == 0 + || policy.max_tokens == 0 + || policy.initial_tokens > policy.max_tokens + { + report.error( + "invalid_extension_policy_tokens", + format!("speculative strategy {name} extension_policy must satisfy 1 <= initial_tokens <= max_tokens"), + Some("model-package.json".to_string()), + "set positive initial_tokens and max_tokens with initial_tokens no larger than max_tokens", + ); + } +} + fn validate_window_policy( name: &str, window: &PackageWindowPolicy, @@ -558,6 +901,20 @@ fn preflight_speculative_decoding( ) -> PreflightSpeculativeDecoding { PreflightSpeculativeDecoding { default: speculative.default.clone(), + proposers: speculative + .proposers + .iter() + .map(|(name, proposer)| PreflightSpeculativeProposer { + name: name.clone(), + proposer_type: proposer.proposer_type.clone(), + prediction_depth: proposer.prediction_depth, + layer_indices: proposer.layer_indices.clone(), + ngram_min: proposer.ngram_min, + ngram_max: proposer.ngram_max, + max_proposal_tokens: proposer.max_proposal_tokens, + history_scope: proposer.history_scope.clone(), + }) + .collect(), strategies: speculative .strategies .iter() @@ -567,6 +924,16 @@ fn preflight_speculative_decoding( prediction_depth: strategy.prediction_depth, layer_indices: strategy.layer_indices.clone(), window_policy: strategy.window_policy.as_ref().map(preflight_window_policy), + proposer: strategy.proposer.clone(), + primary: strategy.primary.clone(), + extender: strategy.extender.clone(), + extension_policy: strategy.extension_policy.as_ref().map(|policy| { + PreflightExtensionPolicy { + initial_tokens: policy.initial_tokens, + max_tokens: policy.max_tokens, + tail_backoff_proposals: policy.tail_backoff_proposals, + } + }), }) .collect(), } @@ -1262,6 +1629,274 @@ mod tests { fs::remove_dir_all(dir).unwrap(); } + #[test] + fn preflight_accepts_request_local_ngram_cache_composite_strategy() { + let dir = unique_test_dir("ngram-cache-composite"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "mtp-cache", + "proposers": { + "mtp": { + "type": "native-mtp", + "prediction_depth": 1, + "layer_indices": [1] + }, + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 4, + "history_scope": "request" + } + }, + "strategies": { + "mtp-cache": { + "type": "composite", + "primary": "mtp", + "extender": "cache", + "extension_policy": { + "initial_tokens": 2, + "max_tokens": 4, + "tail_backoff_proposals": 6 + } + } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(report.valid, "{:?}", report.issues); + let speculative = report + .generation + .and_then(|generation| generation.speculative_decoding) + .expect("generation strategy should be reported"); + assert_eq!(speculative.proposers.len(), 2); + assert_eq!(speculative.strategies[0].primary.as_deref(), Some("mtp")); + assert_eq!(speculative.strategies[0].extender.as_deref(), Some("cache")); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_accepts_complete_native_and_ngram_strategy_matrix() { + let dir = unique_test_dir("complete-speculative-matrix"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "mtp", + "proposers": { + "mtp": { + "type": "native-mtp", + "prediction_depth": 1, + "layer_indices": [1] + }, + "simple": { + "type": "ngram-simple", + "ngram_min": 2, + "ngram_max": 6, + "max_proposal_tokens": 6 + }, + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 6, + "history_scope": "request" + } + }, + "strategies": { + "mtp": { + "type": "native-mtp", + "proposer": "mtp" + }, + "ngram-simple": { + "type": "ngram-simple", + "proposer": "simple" + }, + "ngram-cache": { + "type": "ngram-cache", + "proposer": "cache" + }, + "mtp-simple": { + "type": "composite", + "primary": "mtp", + "extender": "simple", + "extension_policy": { + "initial_tokens": 2, + "max_tokens": 6, + "tail_backoff_proposals": 2 + } + }, + "mtp-cache": { + "type": "composite", + "primary": "mtp", + "extender": "cache", + "extension_policy": { + "initial_tokens": 2, + "max_tokens": 6, + "tail_backoff_proposals": 2 + } + } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(report.valid, "{:?}", report.issues); + let strategies = report + .generation + .and_then(|generation| generation.speculative_decoding) + .expect("generation strategies should be reported") + .strategies; + assert_eq!(strategies.len(), 5); + assert!( + strategies + .iter() + .any(|strategy| strategy.name == "mtp-simple") + ); + assert!( + strategies + .iter() + .any(|strategy| strategy.name == "mtp-cache") + ); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_rejects_ngram_strategy_with_mismatched_proposer_type() { + let dir = unique_test_dir("ngram-strategy-type-mismatch"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "simple", + "proposers": { + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 4, + "history_scope": "request" + } + }, + "strategies": { + "simple": { "type": "ngram-simple", "proposer": "cache" } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(!report.valid); + assert_issue(&report, "ngram_strategy_proposer_type_mismatch"); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_rejects_native_mtp_strategy_with_ngram_proposer() { + let dir = unique_test_dir("native-mtp-strategy-type-mismatch"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "mtp", + "proposers": { + "simple": { + "type": "ngram-simple", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 4 + } + }, + "strategies": { + "mtp": { "type": "native-mtp", "proposer": "simple" } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(!report.valid); + assert_issue(&report, "native_mtp_strategy_proposer_type_mismatch"); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_rejects_shared_ngram_cache_history() { + let dir = unique_test_dir("ngram-cache-shared-history"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "cache", + "proposers": { + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 4, + "history_scope": "shared" + } + }, + "strategies": { + "cache": { "type": "ngram-cache", "proposer": "cache" } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(!report.valid); + assert_issue(&report, "invalid_ngram_cache_history_scope"); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_rejects_ngram_cache_window_above_llama_limit() { + let dir = unique_test_dir("ngram-cache-max-window"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "cache", + "proposers": { + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 5, + "max_proposal_tokens": 6, + "history_scope": "request" + } + }, + "strategies": { + "cache": { "type": "ngram-cache", "proposer": "cache" } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(!report.valid); + assert_issue(&report, "unsupported_ngram_cache_max_window"); + fs::remove_dir_all(dir).unwrap(); + } + #[test] fn preflight_rejects_native_mtp_layer_out_of_range() { let dir = unique_test_dir("native-mtp-out-of-range"); diff --git a/crates/skippy-prompt/src/prompt_cli/generation.rs b/crates/skippy-prompt/src/prompt_cli/generation.rs index eaf01404fb..e5407e5228 100644 --- a/crates/skippy-prompt/src/prompt_cli/generation.rs +++ b/crates/skippy-prompt/src/prompt_cli/generation.rs @@ -368,10 +368,9 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { .max(adaptive_window); let decode_index = generated.len(); let verify_inputs = verify_inputs_for_proposals(current, &draft_tokens); - let reply = send_verify_span( + let reply = send_verify_window( stream, wire_dtype, - prompt_index, request_id, wire_session_id, token_ids.len(), @@ -386,7 +385,7 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { speculative_stats.observe_primary_verify(&reply, verify_inputs.len()); reply_stats.merge(reply.stats); first_time_to_token_ms.get_or_insert_with(|| elapsed_ms(wall_started)); - let decision = classify_verify_span( + let decision = classify_verify_window( &draft_tokens, &reply.predicted_tokens, generated.len(), @@ -442,10 +441,9 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { speculative_stats.recovery_decode_elapsed_ms += repair.elapsed_ms; } else { let repair_inputs = &verify_inputs[..repair_input_count]; - let repair = send_verify_span( + let repair = send_verify_window( stream, wire_dtype, - prompt_index, request_id, wire_session_id, token_ids.len(), @@ -469,13 +467,13 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { speculative_stats.recovery_reverify_write_ms += repair.write_ms; speculative_stats.recovery_reverify_wait_ms += repair.wait_ms; speculative_stats.recovery_reverify_compute_us += - repair.stats.verify_span_compute_us; + repair.stats.verify_window_compute_us; speculative_stats.recovery_reverify_forward_write_us += - repair.stats.verify_span_forward_write_us; + repair.stats.verify_window_forward_write_us; speculative_stats.recovery_reverify_downstream_wait_us += - repair.stats.verify_span_downstream_wait_us; + repair.stats.verify_window_downstream_wait_us; speculative_stats.recovery_reverify_stage_count += - repair.stats.verify_span_stage_count; + repair.stats.verify_window_stage_count; } } let mut reached_eog = false; diff --git a/crates/skippy-prompt/src/prompt_cli/speculative.rs b/crates/skippy-prompt/src/prompt_cli/speculative.rs index 52d5ff09b6..894c381e30 100644 --- a/crates/skippy-prompt/src/prompt_cli/speculative.rs +++ b/crates/skippy-prompt/src/prompt_cli/speculative.rs @@ -48,23 +48,23 @@ struct SpeculativeStats { } impl SpeculativeStats { - fn observe_primary_verify(&mut self, reply: &VerifySpanReply, token_count: usize) { + fn observe_primary_verify(&mut self, reply: &VerifyWindowReply, token_count: usize) { self.primary_verify_requests += 1; self.primary_verify_tokens += token_count; self.primary_verify_elapsed_ms += reply.elapsed_ms; self.primary_verify_write_ms += reply.write_ms; self.primary_verify_wait_ms += reply.wait_ms; - self.primary_verify_compute_us += reply.stats.verify_span_compute_us; - self.primary_verify_forward_write_us += reply.stats.verify_span_forward_write_us; - self.primary_verify_downstream_wait_us += reply.stats.verify_span_downstream_wait_us; - self.primary_verify_total_us += reply.stats.verify_span_total_us; - self.primary_verify_stage_count += reply.stats.verify_span_stage_count; + self.primary_verify_compute_us += reply.stats.verify_window_compute_us; + self.primary_verify_forward_write_us += reply.stats.verify_window_forward_write_us; + self.primary_verify_downstream_wait_us += reply.stats.verify_window_downstream_wait_us; + self.primary_verify_total_us += reply.stats.verify_window_total_us; + self.primary_verify_stage_count += reply.stats.verify_window_stage_count; self.checkpoint_ms += us_to_ms(reply.stats.checkpoint_total_us); } fn observe_verify_decision( &mut self, - decision: VerifySpanDecision, + decision: VerifyWindowDecision, adaptive_window: &mut usize, adaptive_enabled: bool, max_speculative_window: usize, @@ -75,7 +75,7 @@ impl SpeculativeStats { } match decision.kind { - VerifySpanDecisionKind::FullAccept => { + VerifyWindowDecisionKind::FullAccept => { self.full_accept_windows += 1; self.grow_adaptive_window( adaptive_window, @@ -83,10 +83,10 @@ impl SpeculativeStats { max_speculative_window, ); } - VerifySpanDecisionKind::AcceptedStop => { + VerifyWindowDecisionKind::AcceptedStop => { self.accepted_stop_windows += 1; } - VerifySpanDecisionKind::TailReject => { + VerifyWindowDecisionKind::TailReject => { self.observe_reject(decision); self.tail_reject_windows += 1; self.grow_adaptive_window( @@ -95,13 +95,13 @@ impl SpeculativeStats { max_speculative_window, ); } - VerifySpanDecisionKind::EarlyReject => { + VerifyWindowDecisionKind::EarlyReject => { self.observe_reject(decision); self.early_reject_windows += 1; self.repair_required_windows += 1; self.shrink_adaptive_window(adaptive_window, adaptive_enabled, decision); } - VerifySpanDecisionKind::EarlyRejectStop => { + VerifyWindowDecisionKind::EarlyRejectStop => { self.observe_reject(decision); self.early_reject_windows += 1; self.early_reject_stop_windows += 1; @@ -109,7 +109,7 @@ impl SpeculativeStats { } } - fn observe_reject(&mut self, decision: VerifySpanDecision) { + fn observe_reject(&mut self, decision: VerifyWindowDecision) { if let Some(repair_input_count) = decision.repair_input_count { self.rejected_windows += 1; self.first_reject_position_sum += repair_input_count; @@ -132,7 +132,7 @@ impl SpeculativeStats { &mut self, adaptive_window: &mut usize, adaptive_enabled: bool, - decision: VerifySpanDecision, + decision: VerifyWindowDecision, ) { if !adaptive_enabled { return; @@ -162,7 +162,7 @@ fn verify_inputs_for_proposals(current: i32, proposals: &[i32]) -> Vec { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum VerifySpanDecisionKind { +enum VerifyWindowDecisionKind { FullAccept, AcceptedStop, TailReject, @@ -171,46 +171,46 @@ enum VerifySpanDecisionKind { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct VerifySpanDecision { - kind: VerifySpanDecisionKind, +struct VerifyWindowDecision { + kind: VerifyWindowDecisionKind, accepted_before_reject: usize, repair_input_count: Option, commit_count: usize, } -impl VerifySpanDecision { +impl VerifyWindowDecision { fn rejected(self) -> bool { matches!( self.kind, - VerifySpanDecisionKind::TailReject - | VerifySpanDecisionKind::EarlyReject - | VerifySpanDecisionKind::EarlyRejectStop + VerifyWindowDecisionKind::TailReject + | VerifyWindowDecisionKind::EarlyReject + | VerifyWindowDecisionKind::EarlyRejectStop ) } fn requires_repair(self) -> bool { - self.kind == VerifySpanDecisionKind::EarlyReject + self.kind == VerifyWindowDecisionKind::EarlyReject } #[cfg(test)] fn tail_reject(self) -> bool { - self.kind == VerifySpanDecisionKind::TailReject + self.kind == VerifyWindowDecisionKind::TailReject } } -fn classify_verify_span( +fn classify_verify_window( draft_tokens: &[i32], predicted_tokens: &[i32], generated_len: usize, max_new_tokens: usize, mut token_is_eog: F, -) -> Result +) -> Result where F: FnMut(i32) -> Result, { if predicted_tokens.len() < draft_tokens.len() { bail!( - "verify span returned too few tokens: got {} expected {}", + "verify window returned too few tokens: got {} expected {}", predicted_tokens.len(), draft_tokens.len() ); @@ -226,8 +226,8 @@ where if accepted { accepted_before_reject += 1; if (reached_eog || reached_limit) && commit_count < draft_tokens.len() { - return Ok(VerifySpanDecision { - kind: VerifySpanDecisionKind::AcceptedStop, + return Ok(VerifyWindowDecision { + kind: VerifyWindowDecisionKind::AcceptedStop, accepted_before_reject, repair_input_count: None, commit_count, @@ -238,13 +238,13 @@ where let repair_input_count = accepted_before_reject + 1; let kind = if repair_input_count == draft_tokens.len() { - VerifySpanDecisionKind::TailReject + VerifyWindowDecisionKind::TailReject } else if reached_eog || reached_limit { - VerifySpanDecisionKind::EarlyRejectStop + VerifyWindowDecisionKind::EarlyRejectStop } else { - VerifySpanDecisionKind::EarlyReject + VerifyWindowDecisionKind::EarlyReject }; - return Ok(VerifySpanDecision { + return Ok(VerifyWindowDecision { kind, accepted_before_reject, repair_input_count: Some(repair_input_count), @@ -252,8 +252,8 @@ where }); } - Ok(VerifySpanDecision { - kind: VerifySpanDecisionKind::FullAccept, + Ok(VerifyWindowDecision { + kind: VerifyWindowDecisionKind::FullAccept, accepted_before_reject, repair_input_count: None, commit_count, diff --git a/crates/skippy-prompt/src/prompt_cli/tests.rs b/crates/skippy-prompt/src/prompt_cli/tests.rs index b2c4e9e18b..864eaabde0 100644 --- a/crates/skippy-prompt/src/prompt_cli/tests.rs +++ b/crates/skippy-prompt/src/prompt_cli/tests.rs @@ -56,13 +56,13 @@ mod speculative_tests { } #[test] - fn classify_verify_span_full_accept() { + fn classify_verify_window_full_accept() { let decision = - classify_verify_span(&[10, 11, 12], &[10, 11, 12], 0, 16, |_| Ok(false)).unwrap(); + classify_verify_window(&[10, 11, 12], &[10, 11, 12], 0, 16, |_| Ok(false)).unwrap(); assert_eq!( decision, - VerifySpanDecision { - kind: VerifySpanDecisionKind::FullAccept, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::FullAccept, accepted_before_reject: 3, repair_input_count: None, commit_count: 3, @@ -73,13 +73,13 @@ mod speculative_tests { } #[test] - fn classify_verify_span_tail_reject_keeps_state() { + fn classify_verify_window_tail_reject_keeps_state() { let decision = - classify_verify_span(&[10, 11, 12], &[10, 11, 42], 0, 16, |_| Ok(false)).unwrap(); + classify_verify_window(&[10, 11, 12], &[10, 11, 42], 0, 16, |_| Ok(false)).unwrap(); assert_eq!( decision, - VerifySpanDecision { - kind: VerifySpanDecisionKind::TailReject, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::TailReject, accepted_before_reject: 2, repair_input_count: Some(3), commit_count: 3, @@ -91,14 +91,14 @@ mod speculative_tests { } #[test] - fn classify_verify_span_early_reject_requires_repair() { + fn classify_verify_window_early_reject_requires_repair() { let decision = - classify_verify_span(&[10, 11, 12, 13], &[10, 42, 77, 88], 0, 16, |_| Ok(false)) + classify_verify_window(&[10, 11, 12, 13], &[10, 42, 77, 88], 0, 16, |_| Ok(false)) .unwrap(); assert_eq!( decision, - VerifySpanDecision { - kind: VerifySpanDecisionKind::EarlyReject, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::EarlyReject, accepted_before_reject: 1, repair_input_count: Some(2), commit_count: 2, @@ -110,14 +110,14 @@ mod speculative_tests { } #[test] - fn classify_verify_span_accepted_eog_stops_without_growing_window() { + fn classify_verify_window_accepted_eog_stops_without_growing_window() { let decision = - classify_verify_span(&[10, 99, 12], &[10, 99, 12], 0, 16, |token| Ok(token == 99)) + classify_verify_window(&[10, 99, 12], &[10, 99, 12], 0, 16, |token| Ok(token == 99)) .unwrap(); assert_eq!( decision, - VerifySpanDecision { - kind: VerifySpanDecisionKind::AcceptedStop, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::AcceptedStop, accepted_before_reject: 2, repair_input_count: None, commit_count: 2, @@ -128,13 +128,13 @@ mod speculative_tests { } #[test] - fn classify_verify_span_early_reject_at_limit_does_not_repair() { + fn classify_verify_window_early_reject_at_limit_does_not_repair() { let decision = - classify_verify_span(&[10, 11, 12], &[10, 42, 77], 2, 4, |_| Ok(false)).unwrap(); + classify_verify_window(&[10, 11, 12], &[10, 42, 77], 2, 4, |_| Ok(false)).unwrap(); assert_eq!( decision, - VerifySpanDecision { - kind: VerifySpanDecisionKind::EarlyRejectStop, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::EarlyRejectStop, accepted_before_reject: 1, repair_input_count: Some(2), commit_count: 2, @@ -146,11 +146,11 @@ mod speculative_tests { } #[test] - fn classify_verify_span_requires_complete_predictions() { - let err = classify_verify_span(&[10, 11, 12], &[10, 11], 0, 16, |_| Ok(false)).unwrap_err(); + fn classify_verify_window_requires_complete_predictions() { + let err = classify_verify_window(&[10, 11, 12], &[10, 11], 0, 16, |_| Ok(false)).unwrap_err(); assert!( err.to_string() - .contains("verify span returned too few tokens"), + .contains("verify window returned too few tokens"), "{err:#}" ); } @@ -160,8 +160,8 @@ mod speculative_tests { let mut stats = SpeculativeStats::default(); let mut adaptive_window = 4; stats.observe_verify_decision( - VerifySpanDecision { - kind: VerifySpanDecisionKind::FullAccept, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::FullAccept, accepted_before_reject: 4, repair_input_count: None, commit_count: 4, @@ -182,8 +182,8 @@ mod speculative_tests { let mut stats = SpeculativeStats::default(); let mut adaptive_window = 4; stats.observe_verify_decision( - VerifySpanDecision { - kind: VerifySpanDecisionKind::AcceptedStop, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::AcceptedStop, accepted_before_reject: 2, repair_input_count: None, commit_count: 2, @@ -193,8 +193,8 @@ mod speculative_tests { 8, ); stats.observe_verify_decision( - VerifySpanDecision { - kind: VerifySpanDecisionKind::EarlyRejectStop, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::EarlyRejectStop, accepted_before_reject: 1, repair_input_count: Some(2), commit_count: 2, @@ -217,8 +217,8 @@ mod speculative_tests { let mut stats = SpeculativeStats::default(); let mut adaptive_window = 6; stats.observe_verify_decision( - VerifySpanDecision { - kind: VerifySpanDecisionKind::EarlyReject, + VerifyWindowDecision { + kind: VerifyWindowDecisionKind::EarlyReject, accepted_before_reject: 1, repair_input_count: Some(2), commit_count: 2, diff --git a/crates/skippy-prompt/src/prompt_cli/wire_messages.rs b/crates/skippy-prompt/src/prompt_cli/wire_messages.rs index 3d7098ba09..07ef33fdf3 100644 --- a/crates/skippy-prompt/src/prompt_cli/wire_messages.rs +++ b/crates/skippy-prompt/src/prompt_cli/wire_messages.rs @@ -4,7 +4,7 @@ struct DecodeStepReply { elapsed_ms: f64, } -struct VerifySpanReply { +struct VerifyWindowReply { predicted_tokens: Vec, stats: StageReplyStats, write_ms: f64, @@ -65,10 +65,9 @@ fn send_decode_step( } #[allow(clippy::too_many_arguments)] -fn send_verify_span( +fn send_verify_window( stream: &mut TcpStream, wire_dtype: WireActivationDType, - prompt_index: usize, request_id: u64, session_id: u64, prompt_token_count: usize, @@ -76,13 +75,13 @@ fn send_verify_span( decode_index: usize, tokens: &[i32], checkpoint: bool, -) -> Result { +) -> Result { if tokens.is_empty() { - bail!("verify span requires at least one token"); + bail!("verify window requires at least one token"); } let verify_started = Instant::now(); - let mut state = StageStateHeader::new(WireMessageKind::VerifySpan, wire_dtype); - state.seq_id = i32::try_from(prompt_index).context("prompt index exceeds i32")?; + let mut state = StageStateHeader::new(WireMessageKind::VerifyWindow, wire_dtype); + state.seq_id = i32::try_from(decode_index).context("decode step exceeds i32")?; state.prompt_token_count = i32::try_from(prompt_token_count).context("prompt token count exceeds i32")?; state.decode_step = i32::try_from(decode_index).context("decode step exceeds i32")?; @@ -92,9 +91,9 @@ fn send_verify_span( state.flags |= state_flags::SKIP_VERIFY_CHECKPOINT; } let message = StageWireMessage { - kind: WireMessageKind::VerifySpan, - pos_start: i32::try_from(pos_start).context("verify span position exceeds i32")?, - token_count: i32::try_from(tokens.len()).context("verify span exceeds i32")?, + kind: WireMessageKind::VerifyWindow, + pos_start: i32::try_from(pos_start).context("verify window position exceeds i32")?, + token_count: i32::try_from(tokens.len()).context("verify window exceeds i32")?, state, request_id, session_id, @@ -107,14 +106,14 @@ fn send_verify_span( }; let write_started = Instant::now(); write_stage_message(&mut *stream, &message, wire_dtype) - .with_context(|| format!("send verify span at decode step {decode_index}"))?; + .with_context(|| format!("send verify window at decode step {decode_index}"))?; let write_ms = elapsed_ms(write_started); let wait_started = Instant::now(); let reply = recv_reply(&mut *stream) - .with_context(|| format!("receive verify span {decode_index} reply"))?; + .with_context(|| format!("receive verify window {decode_index} reply"))?; ensure_reply_kind(&reply, WireReplyKind::PredictedTokens)?; let wait_ms = elapsed_ms(wait_started); - Ok(VerifySpanReply { + Ok(VerifyWindowReply { predicted_tokens: reply.predicted_tokens, stats: reply.stats, write_ms, @@ -417,8 +416,8 @@ fn print_stats(stats: Stats) { stats.reply_stats.restore_prefill_drained_replies ); } - if stats.reply_stats.verify_span_total_us > 0 { - let verify_total_ms = us_to_ms(stats.reply_stats.verify_span_total_us); + if stats.reply_stats.verify_window_total_us > 0 { + let verify_total_ms = us_to_ms(stats.reply_stats.verify_window_total_us); let verify_tok_s = if verify_total_ms > 0.0 { 1000.0 * stats.speculative_stats.draft_tokens as f64 / verify_total_ms } else { @@ -427,26 +426,26 @@ fn print_stats(stats: Stats) { eprintln!( " spec verify_breakdown_ms total={:.2} compute={:.2} forward={:.2} downstream_wait={:.2} stages={} proposed_tok_s={:.2}", verify_total_ms, - us_to_ms(stats.reply_stats.verify_span_compute_us), - us_to_ms(stats.reply_stats.verify_span_forward_write_us), - us_to_ms(stats.reply_stats.verify_span_downstream_wait_us), - stats.reply_stats.verify_span_stage_count, + us_to_ms(stats.reply_stats.verify_window_compute_us), + us_to_ms(stats.reply_stats.verify_window_forward_write_us), + us_to_ms(stats.reply_stats.verify_window_downstream_wait_us), + stats.reply_stats.verify_window_stage_count, verify_tok_s ); - let protocol_avg_span = if stats.reply_stats.verify_span_request_count > 0 { - stats.reply_stats.verify_span_token_count as f64 - / stats.reply_stats.verify_span_request_count as f64 + let protocol_avg_span = if stats.reply_stats.verify_window_request_count > 0 { + stats.reply_stats.verify_window_token_count as f64 + / stats.reply_stats.verify_window_request_count as f64 } else { 0.0 }; eprintln!( " spec verify_batch_stats protocol_requests={} protocol_tokens={} max_span={} avg_span={:.2} checkpointed_requests={} skip_checkpoint_requests={}", - stats.reply_stats.verify_span_request_count, - stats.reply_stats.verify_span_token_count, - stats.reply_stats.verify_span_max_tokens, + stats.reply_stats.verify_window_request_count, + stats.reply_stats.verify_window_token_count, + stats.reply_stats.verify_window_max_tokens, protocol_avg_span, - stats.reply_stats.verify_span_checkpointed_requests, - stats.reply_stats.verify_span_skip_checkpoint_requests + stats.reply_stats.verify_window_checkpointed_requests, + stats.reply_stats.verify_window_skip_checkpoint_requests ); } if stats.speculative_stats.recovery_reverify_elapsed_ms > 0.0 { diff --git a/crates/skippy-protocol/src/binary/codec.rs b/crates/skippy-protocol/src/binary/codec.rs index 59739a9d6f..65985871de 100644 --- a/crates/skippy-protocol/src/binary/codec.rs +++ b/crates/skippy-protocol/src/binary/codec.rs @@ -4,8 +4,9 @@ use super::{ MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES, MAX_STAGE_DECODED_ACTIVATION_BYTES, MAX_STAGE_LOGIT_BIAS, MAX_STAGE_PREDICTED_TOKENS, MAX_STAGE_SIDEBAND_VALUES, MAX_STAGE_STATE_IMPORT_BYTES, READY_MAGIC, STAGE_STATE_VERSION, - StageLogitBias, StageReply, StageReplyStats, StageSamplingConfig, StageStateHeader, - StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + StageLogitBias, StageNativeMtpDraft, StageReply, StageReplyStats, StageReplyWindow, + StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, + WireReplyKind, activation::{ activation_decoded_f32_bytes_with_state_flags, activation_wire_bytes_with_state_flags, }, @@ -29,10 +30,17 @@ pub fn send_reply_ack(mut writer: impl Write) -> io::Result<()> { } pub fn send_reply_ack_with_stats(mut writer: impl Write, stats: StageReplyStats) -> io::Result<()> { - write_i32(&mut writer, WireReplyKind::Ack as i32)?; - write_i32(&mut writer, 0)?; - write_i32(&mut writer, 0)?; - write_reply_stats(&mut writer, stats) + send_reply_message( + &mut writer, + &StageReply { + kind: WireReplyKind::Ack, + predicted: 0, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: StageReplyWindow::default(), + stats, + }, + ) } pub fn send_reply_predicted(mut writer: impl Write, predicted: i32) -> io::Result<()> { @@ -44,7 +52,13 @@ pub fn send_reply_predicted_with_stats( predicted: i32, stats: StageReplyStats, ) -> io::Result<()> { - send_reply_predicted_with_tokens_and_stats(&mut writer, predicted, &[predicted], stats) + send_reply_predicted_with_tokens_window_and_stats( + &mut writer, + predicted, + &[predicted], + StageReplyWindow::default(), + stats, + ) } pub fn send_reply_predicted_with_tokens_and_stats( @@ -52,20 +66,34 @@ pub fn send_reply_predicted_with_tokens_and_stats( predicted: i32, predicted_tokens: &[i32], stats: StageReplyStats, +) -> io::Result<()> { + send_reply_predicted_with_tokens_window_and_stats( + &mut writer, + predicted, + predicted_tokens, + StageReplyWindow::default(), + stats, + ) +} + +pub fn send_reply_predicted_with_tokens_window_and_stats( + mut writer: impl Write, + predicted: i32, + predicted_tokens: &[i32], + window: StageReplyWindow, + stats: StageReplyStats, ) -> io::Result<()> { if predicted_tokens.len() > MAX_STAGE_PREDICTED_TOKENS { return Err(invalid_input("too many predicted tokens")); } - write_i32(&mut writer, WireReplyKind::PredictedToken as i32)?; - write_i32(&mut writer, predicted)?; - write_i32( + write_reply_header( &mut writer, - i32::try_from(predicted_tokens.len()) - .map_err(|_| invalid_input("too many predicted tokens"))?, + WireReplyKind::PredictedToken, + predicted, + predicted_tokens, + window, )?; - for token in predicted_tokens { - write_i32(&mut writer, *token)?; - } + write_native_mtp_draft(&mut writer, None)?; write_reply_stats(&mut writer, stats) } @@ -73,24 +101,51 @@ pub fn send_reply_predicted_tokens_with_stats( mut writer: impl Write, predicted_tokens: &[i32], stats: StageReplyStats, +) -> io::Result<()> { + send_reply_predicted_tokens_with_window_and_stats( + &mut writer, + predicted_tokens, + StageReplyWindow::default(), + stats, + ) +} + +pub fn send_reply_predicted_tokens_with_window_and_stats( + mut writer: impl Write, + predicted_tokens: &[i32], + window: StageReplyWindow, + stats: StageReplyStats, ) -> io::Result<()> { if predicted_tokens.len() > MAX_STAGE_PREDICTED_TOKENS { return Err(invalid_input("too many predicted tokens")); } let predicted = predicted_tokens.first().copied().unwrap_or(0); - write_i32(&mut writer, WireReplyKind::PredictedTokens as i32)?; - write_i32(&mut writer, predicted)?; - write_i32( + write_reply_header( &mut writer, - i32::try_from(predicted_tokens.len()) - .map_err(|_| invalid_input("too many predicted tokens"))?, + WireReplyKind::PredictedTokens, + predicted, + predicted_tokens, + window, )?; - for token in predicted_tokens { - write_i32(&mut writer, *token)?; - } + write_native_mtp_draft(&mut writer, None)?; write_reply_stats(&mut writer, stats) } +pub fn send_reply_message(mut writer: impl Write, reply: &StageReply) -> io::Result<()> { + if reply.predicted_tokens.len() > MAX_STAGE_PREDICTED_TOKENS { + return Err(invalid_input("too many predicted tokens")); + } + write_reply_header( + &mut writer, + reply.kind, + reply.predicted, + &reply.predicted_tokens, + reply.window, + )?; + write_native_mtp_draft(&mut writer, reply.native_mtp_draft.as_ref())?; + write_reply_stats(&mut writer, reply.stats) +} + pub fn recv_reply(mut reader: impl Read) -> io::Result { let kind = WireReplyKind::try_from(read_i32(&mut reader)?)?; let predicted = read_i32(&mut reader)?; @@ -104,15 +159,84 @@ pub fn recv_reply(mut reader: impl Read) -> io::Result { for _ in 0..predicted_count { predicted_tokens.push(read_i32(&mut reader)?); } + let window = read_reply_window(&mut reader)?; + let native_mtp_draft = read_native_mtp_draft(&mut reader)?; let stats = read_reply_stats(&mut reader)?; Ok(StageReply { kind, predicted, predicted_tokens, + native_mtp_draft, + window, stats, }) } +fn write_reply_header( + mut writer: impl Write, + kind: WireReplyKind, + predicted: i32, + predicted_tokens: &[i32], + window: StageReplyWindow, +) -> io::Result<()> { + write_i32(&mut writer, kind as i32)?; + write_i32(&mut writer, predicted)?; + write_i32( + &mut writer, + i32::try_from(predicted_tokens.len()) + .map_err(|_| invalid_input("too many predicted tokens"))?, + )?; + for token in predicted_tokens { + write_i32(&mut writer, *token)?; + } + write_reply_window(&mut writer, window) +} + +fn write_native_mtp_draft( + mut writer: impl Write, + draft: Option<&StageNativeMtpDraft>, +) -> io::Result<()> { + let Some(draft) = draft else { + return write_i32(&mut writer, 0); + }; + if draft.token_ids.len() > MAX_STAGE_PREDICTED_TOKENS { + return Err(invalid_input("too many native MTP draft tokens")); + } + write_i32(&mut writer, 1)?; + write_i32( + &mut writer, + i32::try_from(draft.token_ids.len()) + .map_err(|_| invalid_input("too many native MTP draft tokens"))?, + )?; + for token in &draft.token_ids { + write_i32(&mut writer, *token)?; + } + write_i64(&mut writer, draft.proposal_compute_us) +} + +fn read_native_mtp_draft(mut reader: impl Read) -> io::Result> { + match read_i32(&mut reader)? { + 0 => Ok(None), + 1 => { + let token_count = checked_i32_len( + read_i32(&mut reader)?, + MAX_STAGE_PREDICTED_TOKENS, + "negative native MTP draft token count", + "native MTP draft token count exceeds maximum", + )?; + let mut token_ids = Vec::with_capacity(token_count); + for _ in 0..token_count { + token_ids.push(read_i32(&mut reader)?); + } + Ok(Some(StageNativeMtpDraft { + token_ids, + proposal_compute_us: read_i64(&mut reader)?, + })) + } + _ => Err(invalid_data("unknown native MTP draft reply marker")), + } +} + pub fn write_stage_message( mut writer: impl Write, message: &StageWireMessage, @@ -463,6 +587,20 @@ fn read_reply_stats(mut reader: impl Read) -> io::Result { Ok(reply_stats_from_fields(fields)) } +fn write_reply_window(mut writer: impl Write, window: StageReplyWindow) -> io::Result<()> { + write_i32(&mut writer, window.window_id)?; + write_i32(&mut writer, window.accepted_len)?; + write_i32(&mut writer, window.correction_token) +} + +fn read_reply_window(mut reader: impl Read) -> io::Result { + Ok(StageReplyWindow { + window_id: read_i32(&mut reader)?, + accepted_len: read_i32(&mut reader)?, + correction_token: read_i32(&mut reader)?, + }) +} + fn reply_stats_fields(stats: StageReplyStats) -> [i64; REPLY_STATS_FIELD_COUNT] { [ stats.kv_lookup_hits, @@ -488,16 +626,16 @@ fn reply_stats_fields(stats: StageReplyStats) -> [i64; REPLY_STATS_FIELD_COUNT] stats.restore_downstream_wait_us, stats.restore_total_us, stats.restore_prefill_drained_replies, - stats.verify_span_compute_us, - stats.verify_span_forward_write_us, - stats.verify_span_downstream_wait_us, - stats.verify_span_total_us, - stats.verify_span_stage_count, - stats.verify_span_request_count, - stats.verify_span_token_count, - stats.verify_span_max_tokens, - stats.verify_span_checkpointed_requests, - stats.verify_span_skip_checkpoint_requests, + stats.verify_window_compute_us, + stats.verify_window_forward_write_us, + stats.verify_window_downstream_wait_us, + stats.verify_window_total_us, + stats.verify_window_stage_count, + stats.verify_window_request_count, + stats.verify_window_token_count, + stats.verify_window_max_tokens, + stats.verify_window_checkpointed_requests, + stats.verify_window_skip_checkpoint_requests, stats.prefill_edge_write_us_max, stats.prefill_edge_wait_us_max, stats.prefill_edge_total_us_max, @@ -532,16 +670,16 @@ fn reply_stats_from_fields(fields: [i64; REPLY_STATS_FIELD_COUNT]) -> StageReply restore_downstream_wait_us: fields[20], restore_total_us: fields[21], restore_prefill_drained_replies: fields[22], - verify_span_compute_us: fields[23], - verify_span_forward_write_us: fields[24], - verify_span_downstream_wait_us: fields[25], - verify_span_total_us: fields[26], - verify_span_stage_count: fields[27], - verify_span_request_count: fields[28], - verify_span_token_count: fields[29], - verify_span_max_tokens: fields[30], - verify_span_checkpointed_requests: fields[31], - verify_span_skip_checkpoint_requests: fields[32], + verify_window_compute_us: fields[23], + verify_window_forward_write_us: fields[24], + verify_window_downstream_wait_us: fields[25], + verify_window_total_us: fields[26], + verify_window_stage_count: fields[27], + verify_window_request_count: fields[28], + verify_window_token_count: fields[29], + verify_window_max_tokens: fields[30], + verify_window_checkpointed_requests: fields[31], + verify_window_skip_checkpoint_requests: fields[32], prefill_edge_write_us_max: fields[33], prefill_edge_wait_us_max: fields[34], prefill_edge_total_us_max: fields[35], @@ -561,6 +699,16 @@ fn write_i32(mut writer: impl Write, value: i32) -> io::Result<()> { writer.write_all(&value.to_le_bytes()) } +fn read_i64(mut reader: impl Read) -> io::Result { + let mut bytes = [0_u8; 8]; + reader.read_exact(&mut bytes)?; + Ok(i64::from_le_bytes(bytes)) +} + +fn write_i64(mut writer: impl Write, value: i64) -> io::Result<()> { + writer.write_all(&value.to_le_bytes()) +} + fn read_u32(mut reader: impl Read) -> io::Result { let mut bytes = [0_u8; 4]; reader.read_exact(&mut bytes)?; diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index 0eaac71cf8..5f25017251 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -9,9 +9,10 @@ pub use activation::{ }; pub use codec::{ read_stage_message, recv_ready, recv_reply, send_ready, send_reply_ack, - send_reply_ack_with_stats, send_reply_predicted, send_reply_predicted_tokens_with_stats, + send_reply_ack_with_stats, send_reply_message, send_reply_predicted, + send_reply_predicted_tokens_with_stats, send_reply_predicted_tokens_with_window_and_stats, send_reply_predicted_with_stats, send_reply_predicted_with_tokens_and_stats, - write_stage_message, + send_reply_predicted_with_tokens_window_and_stats, write_stage_message, }; pub use types::{ ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_RWKV7_V_FIRST, LLAMA_TOKEN_NULL, @@ -19,10 +20,11 @@ pub use types::{ MAX_STAGE_DECODED_ACTIVATION_BYTES, MAX_STAGE_LOGIT_BIAS, MAX_STAGE_PREDICTED_TOKENS, MAX_STAGE_SIDEBAND_VALUES, MAX_STAGE_STATE_IMPORT_BYTES, READY_MAGIC, STAGE_LOGIT_BIAS_WIRE_BYTES, STAGE_SAMPLING_CONFIG_BASE_BYTES, STAGE_STATE_HEADER_BYTES, - STAGE_STATE_VERSION, STAGE_WIRE_FIXED_HEADER_BYTES, StageLogitBias, StageReply, - StageReplyStats, StageRequestEpoch, StageSamplingConfig, StageStateHeader, StageWireMessage, - WireActivationDType, WireMessageKind, WireReplyKind, WireStagePhase, - activation_frame_flags_from_state_flags, activation_state_flags_from_frame_flags, state_flags, + STAGE_STATE_VERSION, STAGE_WIRE_FIXED_HEADER_BYTES, StageLogitBias, StageNativeMtpDraft, + StageReply, StageReplyStats, StageReplyWindow, StageRequestEpoch, StageSamplingConfig, + StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + WireStagePhase, activation_frame_flags_from_state_flags, + activation_state_flags_from_frame_flags, state_flags, }; pub(crate) fn invalid_data(message: &'static str) -> std::io::Error { @@ -103,6 +105,26 @@ mod tests { assert_eq!(reply.kind, WireReplyKind::PredictedToken); assert_eq!(reply.predicted, 42); assert_eq!(reply.predicted_tokens, vec![42]); + assert_eq!(reply.native_mtp_draft, None); + } + + #[test] + fn reply_round_trips_typed_native_mtp_draft() { + let reply = StageReply { + kind: WireReplyKind::PredictedToken, + predicted: 42, + predicted_tokens: vec![42], + native_mtp_draft: Some(StageNativeMtpDraft { + token_ids: vec![43, 44], + proposal_compute_us: 12_345, + }), + window: StageReplyWindow::default(), + stats: StageReplyStats::default(), + }; + let mut bytes = Vec::new(); + send_reply_message(&mut bytes, &reply).unwrap(); + + assert_eq!(recv_reply(Cursor::new(bytes)).unwrap(), reply); } #[test] @@ -150,6 +172,29 @@ mod tests { assert_eq!(reply.predicted_tokens, vec![1, 2, 3]); } + #[test] + fn reply_window_metadata_round_trips() { + let mut bytes = Vec::new(); + send_reply_predicted_tokens_with_window_and_stats( + &mut bytes, + &[1, 2, 3], + StageReplyWindow { + window_id: 42, + accepted_len: 2, + correction_token: 9, + }, + StageReplyStats::default(), + ) + .unwrap(); + let reply = recv_reply(Cursor::new(bytes)).unwrap(); + + assert_eq!(reply.kind, WireReplyKind::PredictedTokens); + assert_eq!(reply.predicted_tokens, vec![1, 2, 3]); + assert_eq!(reply.window.window_id, 42); + assert_eq!(reply.window.accepted_len, 2); + assert_eq!(reply.window.correction_token, 9); + } + #[test] fn reply_rejects_predicted_token_count_over_limit() { let mut bytes = Vec::new(); @@ -231,6 +276,71 @@ mod tests { assert_eq!(sampling.logit_bias[0].bias, -50.0); } + #[test] + fn verify_window_message_round_trips_window_metadata() { + let mut state = + StageStateHeader::new(WireMessageKind::VerifyWindow, WireActivationDType::F32); + state.seq_id = 42; + state.prompt_token_count = 128; + state.decode_step = 7; + state.current_token = 1001; + state.flags |= state_flags::SKIP_VERIFY_CHECKPOINT; + let message = StageWireMessage { + kind: WireMessageKind::VerifyWindow, + pos_start: 135, + token_count: 4, + state, + request_id: 7, + session_id: 11, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![1001, 1002, 1003, 1004], + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }; + + let mut bytes = Vec::new(); + write_stage_message(&mut bytes, &message, WireActivationDType::F32).unwrap(); + let decoded = read_stage_message(Cursor::new(bytes), 2).unwrap(); + + assert_eq!(decoded.kind, WireMessageKind::VerifyWindow); + assert_eq!(decoded.verify_window_id(), Some(42)); + assert_eq!(decoded.verify_window_base_position(), Some(135)); + assert_eq!(decoded.verify_window_token_count(), Some(4)); + assert_eq!(decoded.tokens, vec![1001, 1002, 1003, 1004]); + assert_eq!(decoded.state.decode_step, 7); + assert_ne!(decoded.state.flags & state_flags::SKIP_VERIFY_CHECKPOINT, 0); + } + + #[test] + fn stage_message_rejects_old_state_version() { + let mut state = + StageStateHeader::new(WireMessageKind::DecodeEmbd, WireActivationDType::F32); + state.version = STAGE_STATE_VERSION - 1; + let bytes = stage_frame_prefix(WireMessageKind::DecodeEmbd, 1, 0, 0, state); + + assert_invalid_data( + read_stage_message(Cursor::new(bytes), 2), + "unsupported stage state version", + ); + } + + #[test] + fn stage_message_rejects_legacy_kind_10() { + let mut bytes = Vec::new(); + push_i32(&mut bytes, 10); + push_i32(&mut bytes, 0); + push_i32(&mut bytes, 1); + push_i32(&mut bytes, 0); + push_i32(&mut bytes, 0); + + assert_invalid_data( + read_stage_message(Cursor::new(bytes), 2), + "unknown stage message kind", + ); + } + #[test] fn stage_message_estimates_full_wire_transfer_bytes() { let message = StageWireMessage { diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index b11bc4f702..6de1b589ff 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -5,7 +5,8 @@ use super::{ invalid_data, }; -pub const STAGE_STATE_VERSION: i32 = 6; +// v8 adds the typed native-MTP reply section. Stage peers must be upgraded together. +pub const STAGE_STATE_VERSION: i32 = 8; pub const MAX_STAGE_LOGIT_BIAS: usize = 256; pub const MAX_STAGE_PREDICTED_TOKENS: usize = 262_144; pub const MAX_STAGE_SIDEBAND_VALUES: usize = 1_048_576; @@ -53,7 +54,7 @@ pub enum WireMessageKind { StateImport = 7, DecodeReadout = 8, DecodeLightCtx = 9, - VerifySpan = 10, + VerifyWindow = 21, CheckpointSession = 11, RestoreSession = 12, StateExport = 13, @@ -85,7 +86,7 @@ impl WireMessageKind { Self::DecodeEmbd | Self::DecodeReadout | Self::DecodeLightCtx - | Self::VerifySpan + | Self::VerifyWindow | Self::PrefillFinalEmbd | Self::DecodeReplayFinalEmbd ) @@ -134,7 +135,6 @@ impl TryFrom for WireMessageKind { 7 => Ok(Self::StateImport), 8 => Ok(Self::DecodeReadout), 9 => Ok(Self::DecodeLightCtx), - 10 => Ok(Self::VerifySpan), 11 => Ok(Self::CheckpointSession), 12 => Ok(Self::RestoreSession), 13 => Ok(Self::StateExport), @@ -145,6 +145,7 @@ impl TryFrom for WireMessageKind { 18 => Ok(Self::TryRestorePrefillDecode), 19 => Ok(Self::TrimSession), 20 => Ok(Self::PredictionReturnOpen), + 21 => Ok(Self::VerifyWindow), _ => Err(invalid_data("unknown stage message kind")), } } @@ -395,6 +396,18 @@ pub struct StageWireMessage { } impl StageWireMessage { + pub fn verify_window_id(&self) -> Option { + (self.kind == WireMessageKind::VerifyWindow).then_some(self.state.seq_id) + } + + pub fn verify_window_base_position(&self) -> Option { + (self.kind == WireMessageKind::VerifyWindow).then_some(self.pos_start) + } + + pub fn verify_window_token_count(&self) -> Option { + (self.kind == WireMessageKind::VerifyWindow).then_some(self.token_count) + } + pub fn estimated_wire_bytes(&self) -> usize { let sampling_bytes = self.sampling.as_ref().map_or(0, |sampling| { STAGE_SAMPLING_CONFIG_BASE_BYTES @@ -537,10 +550,41 @@ impl StageWireMessage { pub struct StageReply { pub kind: WireReplyKind, pub predicted: i32, + /// Target-model predictions only. Native MTP proposals are carried separately. pub predicted_tokens: Vec, + pub native_mtp_draft: Option, + pub window: StageReplyWindow, pub stats: StageReplyStats, } +/// A native MTP proposal associated with a stage reply. +/// +/// This deliberately has its own reply field rather than sharing the target +/// prediction vector. Consumers must never mistake proposal metadata for a +/// target prediction while verifying a composite speculative window. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StageNativeMtpDraft { + pub token_ids: Vec, + pub proposal_compute_us: i64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StageReplyWindow { + pub window_id: i32, + pub accepted_len: i32, + pub correction_token: i32, +} + +impl Default for StageReplyWindow { + fn default() -> Self { + Self { + window_id: 0, + accepted_len: 0, + correction_token: LLAMA_TOKEN_NULL, + } + } +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct StageReplyStats { pub kv_lookup_hits: i64, @@ -566,16 +610,16 @@ pub struct StageReplyStats { pub restore_downstream_wait_us: i64, pub restore_total_us: i64, pub restore_prefill_drained_replies: i64, - pub verify_span_compute_us: i64, - pub verify_span_forward_write_us: i64, - pub verify_span_downstream_wait_us: i64, - pub verify_span_total_us: i64, - pub verify_span_stage_count: i64, - pub verify_span_request_count: i64, - pub verify_span_token_count: i64, - pub verify_span_max_tokens: i64, - pub verify_span_checkpointed_requests: i64, - pub verify_span_skip_checkpoint_requests: i64, + pub verify_window_compute_us: i64, + pub verify_window_forward_write_us: i64, + pub verify_window_downstream_wait_us: i64, + pub verify_window_total_us: i64, + pub verify_window_stage_count: i64, + pub verify_window_request_count: i64, + pub verify_window_token_count: i64, + pub verify_window_max_tokens: i64, + pub verify_window_checkpointed_requests: i64, + pub verify_window_skip_checkpoint_requests: i64, pub prefill_edge_write_us_max: i64, pub prefill_edge_wait_us_max: i64, pub prefill_edge_total_us_max: i64, @@ -609,18 +653,18 @@ impl StageReplyStats { self.restore_downstream_wait_us += other.restore_downstream_wait_us; self.restore_total_us += other.restore_total_us; self.restore_prefill_drained_replies += other.restore_prefill_drained_replies; - self.verify_span_compute_us += other.verify_span_compute_us; - self.verify_span_forward_write_us += other.verify_span_forward_write_us; - self.verify_span_downstream_wait_us += other.verify_span_downstream_wait_us; - self.verify_span_total_us += other.verify_span_total_us; - self.verify_span_stage_count += other.verify_span_stage_count; - self.verify_span_request_count += other.verify_span_request_count; - self.verify_span_token_count += other.verify_span_token_count; - self.verify_span_max_tokens = self - .verify_span_max_tokens - .max(other.verify_span_max_tokens); - self.verify_span_checkpointed_requests += other.verify_span_checkpointed_requests; - self.verify_span_skip_checkpoint_requests += other.verify_span_skip_checkpoint_requests; + self.verify_window_compute_us += other.verify_window_compute_us; + self.verify_window_forward_write_us += other.verify_window_forward_write_us; + self.verify_window_downstream_wait_us += other.verify_window_downstream_wait_us; + self.verify_window_total_us += other.verify_window_total_us; + self.verify_window_stage_count += other.verify_window_stage_count; + self.verify_window_request_count += other.verify_window_request_count; + self.verify_window_token_count += other.verify_window_token_count; + self.verify_window_max_tokens = self + .verify_window_max_tokens + .max(other.verify_window_max_tokens); + self.verify_window_checkpointed_requests += other.verify_window_checkpointed_requests; + self.verify_window_skip_checkpoint_requests += other.verify_window_skip_checkpoint_requests; self.prefill_edge_write_us_max = self .prefill_edge_write_us_max .max(other.prefill_edge_write_us_max); @@ -680,16 +724,16 @@ impl StageReplyStats { && self.restore_downstream_wait_us == 0 && self.restore_total_us == 0 && self.restore_prefill_drained_replies == 0 - && self.verify_span_compute_us == 0 - && self.verify_span_forward_write_us == 0 - && self.verify_span_downstream_wait_us == 0 - && self.verify_span_total_us == 0 - && self.verify_span_stage_count == 0 - && self.verify_span_request_count == 0 - && self.verify_span_token_count == 0 - && self.verify_span_max_tokens == 0 - && self.verify_span_checkpointed_requests == 0 - && self.verify_span_skip_checkpoint_requests == 0 + && self.verify_window_compute_us == 0 + && self.verify_window_forward_write_us == 0 + && self.verify_window_downstream_wait_us == 0 + && self.verify_window_total_us == 0 + && self.verify_window_stage_count == 0 + && self.verify_window_request_count == 0 + && self.verify_window_token_count == 0 + && self.verify_window_max_tokens == 0 + && self.verify_window_checkpointed_requests == 0 + && self.verify_window_skip_checkpoint_requests == 0 && self.prefill_edge_observation_count == 0 } } diff --git a/crates/skippy-runtime/src/lib.rs b/crates/skippy-runtime/src/lib.rs index 2d71491c50..41860973d9 100644 --- a/crates/skippy-runtime/src/lib.rs +++ b/crates/skippy-runtime/src/lib.rs @@ -26,9 +26,11 @@ use tokio::sync::mpsc; mod devices; mod native_mtp; +mod ngram; pub mod package; mod runtime_events; pub use native_mtp::NativeMtpDraft; +pub use ngram::{Cache as NgramCache, NGRAM_CACHE_MAX_NGRAM, simple_draft as ngram_simple_draft}; pub const MAX_LOGIT_BIAS: usize = 256; pub const GGML_TYPE_F16: u32 = 1; diff --git a/crates/skippy-runtime/src/ngram.rs b/crates/skippy-runtime/src/ngram.rs new file mode 100644 index 0000000000..0955929546 --- /dev/null +++ b/crates/skippy-runtime/src/ngram.rs @@ -0,0 +1,210 @@ +use std::ptr::{self, NonNull}; + +use anyhow::{Context, Result, bail}; + +/// llama.cpp's stateful N-gram cache supports match windows up to four tokens. +/// The proposal length remains independently configurable. +pub const NGRAM_CACHE_MAX_NGRAM: usize = 4; + +/// Uses llama.cpp's ngram-simple proposer against accepted history, including +/// the current sampled token as the final item in `history`. +pub fn simple_draft( + history: &[i32], + ngram_size: usize, + max_draft_tokens: usize, +) -> Result> { + if ngram_size == 0 || max_draft_tokens == 0 || history.len() < 2 { + return Ok(Vec::new()); + } + let output_limit = max_draft_tokens; + let search_draft_tokens = max_draft_tokens.max(ngram_size); + let ngram_size = u16::try_from(ngram_size).context("ngram size exceeds llama.cpp limit")?; + let search_draft_tokens = u16::try_from(search_draft_tokens) + .context("N-gram draft search limit exceeds llama.cpp limit")?; + let (sampled_token, token_ids) = history + .split_last() + .expect("history length is checked above"); + let mut output_tokens = vec![0_i32; usize::from(search_draft_tokens)]; + let mut output_token_count = 0_usize; + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_ngram_simple_draft( + token_ids.as_ptr(), + token_ids.len(), + *sampled_token, + ngram_size, + search_draft_tokens, + output_tokens.as_mut_ptr(), + output_tokens.len(), + &mut output_token_count, + &mut error, + ) + }; + super::ensure_ok(status, error)?; + if output_token_count > output_tokens.len() { + bail!("llama.cpp N-gram proposer exceeded its requested draft limit"); + } + output_tokens.truncate(output_token_count.min(output_limit)); + Ok(output_tokens) +} + +/// Stateful adapter for llama.cpp's cache-based N-gram proposer. +/// +/// Callers must feed only target-committed history through [`Self::reset`] or +/// [`Self::append`]. `draft_after` may include provisional tokens, but native +/// state is not mutated while producing that candidate. +pub struct Cache { + raw: NonNull, +} + +impl Cache { + pub fn new(ngram_min: usize, ngram_max: usize) -> Result { + if ngram_min == 0 || ngram_min > ngram_max { + bail!("cache N-gram proposer requires 0 < ngram_min <= ngram_max"); + } + if ngram_max > NGRAM_CACHE_MAX_NGRAM { + bail!( + "cache N-gram proposer ngram_max {ngram_max} exceeds llama.cpp limit {NGRAM_CACHE_MAX_NGRAM}" + ); + } + let ngram_min = u16::try_from(ngram_min).context("cache N-gram minimum exceeds limit")?; + let ngram_max = u16::try_from(ngram_max).context("cache N-gram maximum exceeds limit")?; + let mut raw = ptr::null_mut(); + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_ngram_cache_create(ngram_min, ngram_max, &mut raw, &mut error) + }; + super::ensure_ok(status, error)?; + let raw = NonNull::new(raw).context("llama.cpp created a null N-gram cache")?; + Ok(Self { raw }) + } + + pub fn reset(&mut self, history: &[i32]) -> Result<()> { + self.update(history, true) + } + + pub fn append(&mut self, committed_tokens: &[i32]) -> Result<()> { + if committed_tokens.is_empty() { + return Ok(()); + } + self.update(committed_tokens, false) + } + + pub fn draft_after( + &mut self, + continuation_prefix: &[i32], + max_draft_tokens: usize, + ) -> Result> { + if max_draft_tokens == 0 { + return Ok(Vec::new()); + } + let max_draft_tokens = + u16::try_from(max_draft_tokens).context("cache N-gram draft limit exceeds limit")?; + let mut output_tokens = vec![0_i32; usize::from(max_draft_tokens)]; + let mut output_token_count = 0_usize; + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_ngram_cache_draft( + self.raw.as_ptr(), + continuation_prefix.as_ptr(), + continuation_prefix.len(), + max_draft_tokens, + output_tokens.as_mut_ptr(), + output_tokens.len(), + &mut output_token_count, + &mut error, + ) + }; + super::ensure_ok(status, error)?; + if output_token_count > output_tokens.len() { + bail!("llama.cpp cache N-gram proposer exceeded its requested draft limit"); + } + output_tokens.truncate(output_token_count); + Ok(output_tokens) + } + + fn update(&mut self, tokens: &[i32], reset: bool) -> Result<()> { + let mut error = ptr::null_mut(); + let status = unsafe { + if reset { + skippy_ffi::skippy_ngram_cache_reset( + self.raw.as_ptr(), + tokens.as_ptr(), + tokens.len(), + &mut error, + ) + } else { + skippy_ffi::skippy_ngram_cache_append( + self.raw.as_ptr(), + tokens.as_ptr(), + tokens.len(), + &mut error, + ) + } + }; + super::ensure_ok(status, error) + } +} + +impl Drop for Cache { + fn drop(&mut self) { + unsafe { skippy_ffi::skippy_ngram_cache_free(self.raw.as_ptr()) }; + } +} + +#[cfg(test)] +mod tests { + use super::{Cache, NGRAM_CACHE_MAX_NGRAM, simple_draft}; + + #[test] + fn drafts_the_continuation_from_the_latest_matching_ngram() { + let history = [1, 2, 3, 4, 9, 2, 3, 4]; + + assert_eq!(simple_draft(&history, 2, 2).unwrap(), vec![9, 2]); + } + + #[test] + fn draft_output_budget_is_independent_from_match_length() { + let history = [ + 99, 1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 13, 14, 15, 16, 17, 18, 77, 1, 2, 3, 4, 5, 6, 7, 8, + ]; + + assert_eq!(simple_draft(&history, 8, 2).unwrap(), vec![11, 12]); + } + + #[test] + fn respects_zero_limits_without_entering_the_native_abi() { + assert!(simple_draft(&[1, 2, 1, 2], 0, 4).unwrap().is_empty()); + assert!(simple_draft(&[1, 2, 1, 2], 1, 0).unwrap().is_empty()); + } + + #[test] + fn cache_drafts_from_committed_history_and_never_mutates_for_a_prefix() { + let mut cache = Cache::new(2, 2).unwrap(); + cache.reset(&[1, 2, 3, 1, 2, 3, 1, 2]).unwrap(); + + assert_eq!(cache.draft_after(&[], 2).unwrap(), vec![3, 1]); + assert_eq!(cache.draft_after(&[9], 2).unwrap(), Vec::::new()); + assert_eq!(cache.draft_after(&[], 2).unwrap(), vec![3, 1]); + } + + #[test] + fn cache_append_extends_the_committed_history() { + let mut cache = Cache::new(2, 2).unwrap(); + cache.reset(&[1, 9, 7, 1, 9, 7, 1]).unwrap(); + + assert_eq!(cache.draft_after(&[9], 1).unwrap(), vec![7]); + cache.append(&[9, 7, 1]).unwrap(); + assert_eq!(cache.draft_after(&[9], 1).unwrap(), vec![7]); + } + + #[test] + fn cache_rejects_match_windows_above_the_llama_limit() { + let error = match Cache::new(2, NGRAM_CACHE_MAX_NGRAM + 1) { + Ok(_) => panic!("must reject max > 4"), + Err(error) => error, + }; + + assert!(error.to_string().contains("exceeds llama.cpp limit 4")); + } +} diff --git a/crates/skippy-runtime/src/package.rs b/crates/skippy-runtime/src/package.rs index bba53cc0b2..c4e20e5e40 100644 --- a/crates/skippy-runtime/src/package.rs +++ b/crates/skippy-runtime/src/package.rs @@ -66,15 +66,38 @@ pub struct PackageGenerationInfo { #[derive(Debug, Clone, Eq, PartialEq)] pub struct PackageSpeculativeDecodingInfo { pub default: String, + pub proposers: BTreeMap, pub strategies: BTreeMap, } +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct PackageSpeculativeProposerInfo { + pub proposer_type: String, + pub prediction_depth: Option, + pub layer_indices: Vec, + pub ngram_min: Option, + pub ngram_max: Option, + pub max_proposal_tokens: Option, + pub history_scope: Option, +} + #[derive(Debug, Clone, Eq, PartialEq)] pub struct PackageSpeculativeStrategyInfo { pub strategy_type: String, pub prediction_depth: Option, pub layer_indices: Vec, pub window_policy: Option, + pub proposer: Option, + pub primary: Option, + pub extender: Option, + pub extension_policy: Option, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct PackageExtensionPolicyInfo { + pub initial_tokens: u32, + pub max_tokens: u32, + pub tail_backoff_proposals: u32, } #[derive(Debug, Clone, Eq, PartialEq)] @@ -213,9 +236,29 @@ struct PackageGeneration { struct PackageSpeculativeDecoding { default: String, #[serde(default)] + proposers: BTreeMap, + #[serde(default)] strategies: BTreeMap, } +#[derive(Debug, Deserialize)] +struct PackageSpeculativeProposer { + #[serde(rename = "type")] + proposer_type: String, + #[serde(default)] + prediction_depth: Option, + #[serde(default)] + layer_indices: Vec, + #[serde(default)] + ngram_min: Option, + #[serde(default)] + ngram_max: Option, + #[serde(default)] + max_proposal_tokens: Option, + #[serde(default)] + history_scope: Option, +} + #[derive(Debug, Deserialize)] struct PackageSpeculativeStrategy { #[serde(rename = "type")] @@ -226,6 +269,21 @@ struct PackageSpeculativeStrategy { layer_indices: Vec, #[serde(default)] window_policy: Option, + #[serde(default)] + proposer: Option, + #[serde(default)] + primary: Option, + #[serde(default)] + extender: Option, + #[serde(default)] + extension_policy: Option, +} + +#[derive(Debug, Deserialize)] +struct PackageExtensionPolicy { + initial_tokens: u32, + max_tokens: u32, + tail_backoff_proposals: u32, } #[derive(Debug, Deserialize)] @@ -556,6 +614,11 @@ fn package_speculative_decoding_info( ) -> PackageSpeculativeDecodingInfo { PackageSpeculativeDecodingInfo { default: speculative.default, + proposers: speculative + .proposers + .into_iter() + .map(|(name, proposer)| (name, package_speculative_proposer_info(proposer))) + .collect(), strategies: speculative .strategies .into_iter() @@ -564,6 +627,20 @@ fn package_speculative_decoding_info( } } +fn package_speculative_proposer_info( + proposer: PackageSpeculativeProposer, +) -> PackageSpeculativeProposerInfo { + PackageSpeculativeProposerInfo { + proposer_type: proposer.proposer_type, + prediction_depth: proposer.prediction_depth, + layer_indices: proposer.layer_indices, + ngram_min: proposer.ngram_min, + ngram_max: proposer.ngram_max, + max_proposal_tokens: proposer.max_proposal_tokens, + history_scope: proposer.history_scope, + } +} + fn package_speculative_strategy_info( strategy: PackageSpeculativeStrategy, ) -> PackageSpeculativeStrategyInfo { @@ -579,6 +656,16 @@ fn package_speculative_strategy_info( min_window: window.min_window, max_window: window.max_window, }), + proposer: strategy.proposer, + primary: strategy.primary, + extender: strategy.extender, + extension_policy: strategy + .extension_policy + .map(|policy| PackageExtensionPolicyInfo { + initial_tokens: policy.initial_tokens, + max_tokens: policy.max_tokens, + tail_backoff_proposals: policy.tail_backoff_proposals, + }), } } diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 6b46f88f57..61e765c0ae 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -153,7 +153,7 @@ unload or replan. `--openai-draft-model-path`, `--openai-speculative-window`, and `--openai-adaptive-speculative-window`. The draft model runs locally in the stage0 process as a complete model without stage tensor filtering, and - proposal windows are verified through the existing staged `VerifySpan` binary + proposal windows are verified through the existing staged `VerifyWindow` binary request, so acceptance, rejection, checkpoint, restore, draft-propose, and recovery costs are visible on OpenAI-path spans. The draft runner is single-session guarded; use this first as a depth-1 measurement knob before diff --git a/crates/skippy-server/src/binary_transport.rs b/crates/skippy-server/src/binary_transport.rs index 81832b6370..3760010939 100644 --- a/crates/skippy-server/src/binary_transport.rs +++ b/crates/skippy-server/src/binary_transport.rs @@ -27,11 +27,11 @@ use skippy_metrics::{attr, metric}; use skippy_protocol::{ MessageBase, SCHEMA_VERSION, StageConfig, StageTopology, binary::{ - READY_MAGIC, StageReply, StageReplyStats, StageSamplingConfig, StageStateHeader, - StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, - activation_frame_flags_from_state_flags, read_stage_message, recv_reply, send_ready, - send_reply_ack, send_reply_ack_with_stats, send_reply_predicted_tokens_with_stats, - send_reply_predicted_with_tokens_and_stats, state_flags, + LLAMA_TOKEN_NULL, READY_MAGIC, StageNativeMtpDraft, StageReply, StageReplyStats, + StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, + WireMessageKind, WireReplyKind, activation_frame_flags_from_state_flags, + read_stage_message, recv_reply, send_ready, send_reply_ack, send_reply_ack_with_stats, + send_reply_message, state_flags, }, }; use skippy_runtime::{ @@ -48,6 +48,7 @@ pub(crate) mod forwarding; mod kv_eviction; mod options; mod preconnect; +mod restore_prefill_decode; mod socket; mod wire; @@ -63,6 +64,9 @@ use self::kv_eviction::{ }; pub use self::options::{BinaryStageOptions, EmbeddedOpenAiStageOptions, parse_wire_dtype}; use self::preconnect::spawn_downstream_preconnector; +use self::restore_prefill_decode::handle_binary_restore_prefill_decode_control; +#[cfg(test)] +use self::restore_prefill_decode::restore_prefill_decode_as_decode_message; use self::socket::*; pub use self::wire::WireCondition; pub(crate) use self::wire::write_stage_message_conditioned; @@ -278,9 +282,21 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R speculative_window: openai_options.speculative_window, adaptive_speculative_window: openai_options.adaptive_speculative_window, draft_n_gpu_layers: openai_options.draft_n_gpu_layers, - ngram_min: openai_options.ngram_min, - ngram_max: openai_options.ngram_max, - native_mtp_enabled, + speculative: openai_options.speculative.clone(), + ngram_min: openai_options + .speculative + .ngram + .as_ref() + .filter(|ngram| ngram.kind == frontend::NgramProposerKind::Simple) + .map_or(0, |ngram| ngram.min_ngram), + ngram_max: openai_options + .speculative + .ngram + .as_ref() + .filter(|ngram| ngram.kind == frontend::NgramProposerKind::Simple) + .map_or(0, |ngram| ngram.max_proposal_tokens.min(ngram.max_ngram)), + native_mtp_enabled: native_mtp_enabled + && openai_options.speculative.native_mtp.enabled, native_mtp_draft_model_path: None, native_mtp_max_tokens: openai_options.native_mtp_max_tokens, native_mtp_min_tokens: openai_options.native_mtp_min_tokens, @@ -717,6 +733,8 @@ fn handle_binary_connection( reset_start_unix_nanos, reset_end_unix_nanos, ); + prediction_return_streams.remove(&(message.request_id, message.session_id)); + prediction_return_sinks.remove(message.request_id, message.session_id); send_reply_ack_with_stats(&mut *upstream, stop_stats).context("send stop ACK")?; continue; } @@ -884,6 +902,7 @@ fn handle_binary_connection( activation_width, control_started, control_stats, + prediction_return_sinks, &mut prediction_return_streams, downstream_connect_timeout_secs, native_mtp_enabled, @@ -1043,7 +1062,7 @@ fn handle_binary_connection( let mut decode_batch_wait_ms = 0.0; let input_activation_bytes = message.activation.len(); let mut proactive_eviction = None; - let (predicted_token, predicted_tokens, output, compute_ms) = if restored_prefill { + let (predicted_token, mut predicted_tokens, output, compute_ms) = if restored_prefill { let now = now_unix_nanos() as u64; compute_start_unix_nanos = now; compute_end_unix_nanos = now; @@ -1064,7 +1083,7 @@ fn handle_binary_connection( } else { elapsed_ms(input_decode_started) }; - if message.kind == WireMessageKind::VerifySpan + if message.kind == WireMessageKind::VerifyWindow && (message.state.flags & state_flags::SKIP_VERIFY_CHECKPOINT) == 0 { let checkpoint_started = Instant::now(); @@ -1072,7 +1091,7 @@ fn handle_binary_connection( let mut runtime = runtime.lock().expect("runtime lock poisoned"); runtime .checkpoint_session(&session_key) - .context("checkpoint binary stage session before verify span")?; + .context("checkpoint binary stage session before verify window")?; } let checkpoint_us = elapsed_us(checkpoint_started); record_session_control_timing( @@ -1511,19 +1530,20 @@ fn handle_binary_connection( ); message_reply_stats.merge(pending_reply_stats); pending_reply_stats = StageReplyStats::default(); - record_verify_span_timing( + record_verify_window_timing( &mut message_reply_stats, &message, compute_ms, forward_write_ms, downstream_wait_ms, ); - let reply_kind = if message.kind == WireMessageKind::VerifySpan { + let reply_kind = if message.kind == WireMessageKind::VerifyWindow { WireReplyKind::PredictedTokens } else { WireReplyKind::PredictedToken }; - let predicted_token_count = if message.kind == WireMessageKind::VerifySpan { + let native_mtp_draft = split_native_mtp_reply(&message, &mut predicted_tokens)?; + let predicted_token_count = if message.kind == WireMessageKind::VerifyWindow { predicted_tokens.len() } else { predicted_tokens.len().max(1) @@ -1531,10 +1551,13 @@ fn handle_binary_connection( let reply_start_unix_nanos = now_unix_nanos() as u64; upstream_reply_start_unix_nanos.get_or_insert(reply_start_unix_nanos); let reply_started = Instant::now(); + let reply_window = reply_window_for_message(&message, &predicted_tokens); let reply = StageReply { kind: reply_kind, predicted: predicted_token, predicted_tokens, + native_mtp_draft, + window: reply_window, stats: message_reply_stats, }; if let Some(return_stream) = @@ -1605,24 +1628,24 @@ fn handle_binary_connection( let message_end_unix_nanos = now_unix_nanos() as u64; let message_elapsed_ms = elapsed_ms(message_started); - let verify_span_pre_compute_ms = if message.kind == WireMessageKind::VerifySpan { + let verify_window_pre_compute_ms = if message.kind == WireMessageKind::VerifyWindow { nanos_delta_ms(message_start_unix_nanos, compute_start_unix_nanos) } else { 0.0 }; - let verify_span_post_compute_ms = if message.kind == WireMessageKind::VerifySpan { + let verify_window_post_compute_ms = if message.kind == WireMessageKind::VerifyWindow { nanos_delta_ms(compute_end_unix_nanos, message_end_unix_nanos) } else { 0.0 }; - let verify_span_pre_reply_ms = if message.kind == WireMessageKind::VerifySpan { + let verify_window_pre_reply_ms = if message.kind == WireMessageKind::VerifyWindow { upstream_reply_start_unix_nanos .map(|reply_start| nanos_delta_ms(compute_end_unix_nanos, reply_start)) .unwrap_or(0.0) } else { 0.0 }; - let verify_span_after_reply_ms = if message.kind == WireMessageKind::VerifySpan { + let verify_window_after_reply_ms = if message.kind == WireMessageKind::VerifyWindow { upstream_reply_end_unix_nanos .map(|reply_end| nanos_delta_ms(reply_end, message_end_unix_nanos)) .unwrap_or(0.0) @@ -1651,10 +1674,10 @@ fn handle_binary_connection( session_auto_align_count, session_auto_align_ms, session_auto_align_trimmed_tokens, - verify_span_pre_compute_ms, - verify_span_post_compute_ms, - verify_span_pre_reply_ms, - verify_span_after_reply_ms, + verify_window_pre_compute_ms, + verify_window_post_compute_ms, + verify_window_pre_reply_ms, + verify_window_after_reply_ms, upstream_message_wait_ms: recv_read_ms, }); @@ -1829,6 +1852,51 @@ fn native_mtp_prediction_tokens(predicted: i32, draft: Option) - tokens } +/// Converts the temporary llama-stage sideband into the typed stage reply field. +/// +/// The C ABI still returns the proposal as a trailer. The network boundary is +/// authoritative: consumers receive only target predictions plus a separate +/// native-MTP draft. A malformed trailer is rejected instead of being exposed +/// as target-model output. +fn split_native_mtp_reply( + message: &StageWireMessage, + prediction_tokens: &mut Vec, +) -> Result> { + let sideband_offset = match message.kind { + WireMessageKind::DecodeEmbd + | WireMessageKind::DecodeReadout + | WireMessageKind::DecodeLightCtx + | WireMessageKind::DecodeReplayEmbd + | WireMessageKind::DecodeReplayFinalEmbd => 1, + WireMessageKind::VerifyWindow => message.tokens.len(), + _ => return Ok(None), + }; + if prediction_tokens.len() <= sideband_offset { + return Ok(None); + } + let draft_token_count = usize::try_from(prediction_tokens[sideband_offset]) + .map_err(|_| anyhow!("negative native MTP draft token count"))?; + let draft_start = sideband_offset + 1; + let draft_end = draft_start + .checked_add(draft_token_count) + .ok_or_else(|| anyhow!("native MTP draft token count overflow"))?; + let compute_index = draft_end; + if prediction_tokens.len() != compute_index + 1 { + bail!( + "malformed native MTP sideband: expected {} values, got {}", + compute_index + 1, + prediction_tokens.len() + ); + } + let proposal_compute_us = i64::from(prediction_tokens[compute_index].max(0)); + let token_ids = prediction_tokens[draft_start..draft_end].to_vec(); + prediction_tokens.truncate(sideband_offset); + Ok((!token_ids.is_empty()).then_some(StageNativeMtpDraft { + token_ids, + proposal_compute_us, + })) +} + fn binary_auto_align_session_enabled() -> bool { truthy_env(env::var(AUTO_ALIGN_SESSION_ENV).ok().as_deref()) } @@ -1850,7 +1918,7 @@ fn message_allows_session_auto_align(message: &StageWireMessage) -> bool { WireMessageKind::DecodeEmbd | WireMessageKind::DecodeReadout | WireMessageKind::DecodeLightCtx - | WireMessageKind::VerifySpan + | WireMessageKind::VerifyWindow ) } @@ -1858,14 +1926,6 @@ fn message_pos_start_as_token_count(message: &StageWireMessage) -> Option { u64::try_from(message.pos_start).ok() } -#[cfg(test)] -fn native_mtp_enabled_from(value: Option<&str>) -> bool { - !matches!( - value.map(str::trim).map(str::to_ascii_lowercase).as_deref(), - Some("0" | "false" | "off" | "disable" | "disabled" | "no") - ) -} - pub(crate) fn stage_output_activation_capacity( config: &StageConfig, token_count: i32, @@ -2082,32 +2142,32 @@ fn record_prefill_edge_transport( ); } -fn record_verify_span_timing( +fn record_verify_window_timing( stats: &mut StageReplyStats, message: &StageWireMessage, compute_ms: f64, forward_write_ms: f64, downstream_wait_ms: f64, ) { - if message.kind != WireMessageKind::VerifySpan { + if message.kind != WireMessageKind::VerifyWindow { return; } let compute_us = ms_to_us(compute_ms); let forward_write_us = ms_to_us(forward_write_ms); let downstream_wait_us = ms_to_us(downstream_wait_ms); let token_count = i64::from(message.token_count.max(0)); - stats.verify_span_compute_us += compute_us; - stats.verify_span_forward_write_us += forward_write_us; - stats.verify_span_downstream_wait_us += downstream_wait_us; - stats.verify_span_total_us += compute_us + forward_write_us + downstream_wait_us; - stats.verify_span_stage_count += 1; - stats.verify_span_request_count += 1; - stats.verify_span_token_count += token_count; - stats.verify_span_max_tokens = stats.verify_span_max_tokens.max(token_count); + stats.verify_window_compute_us += compute_us; + stats.verify_window_forward_write_us += forward_write_us; + stats.verify_window_downstream_wait_us += downstream_wait_us; + stats.verify_window_total_us += compute_us + forward_write_us + downstream_wait_us; + stats.verify_window_stage_count += 1; + stats.verify_window_request_count += 1; + stats.verify_window_token_count += token_count; + stats.verify_window_max_tokens = stats.verify_window_max_tokens.max(token_count); if (message.state.flags & state_flags::SKIP_VERIFY_CHECKPOINT) == 0 { - stats.verify_span_checkpointed_requests += 1; + stats.verify_window_checkpointed_requests += 1; } else { - stats.verify_span_skip_checkpoint_requests += 1; + stats.verify_window_skip_checkpoint_requests += 1; } } @@ -2398,6 +2458,9 @@ fn configure_prediction_return_stream( prediction_return_sinks: &PredictionReturnSinks, prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, ) { + if prediction_return_streams.contains_key(&(request_id, session_id)) { + return; + } match prediction_return_sinks.take_wait(request_id, session_id, Duration::from_millis(250)) { Ok(Some(stream)) => { prediction_return_streams.insert((request_id, session_id), stream); @@ -2429,257 +2492,32 @@ fn configure_prediction_return_stream( } } -#[allow(clippy::too_many_arguments)] -fn handle_binary_restore_prefill_decode_control( - config: &StageConfig, - topology: Option<&StageTopology>, - runtime: &Arc>, - kv: Option<&Arc>, - telemetry: &Telemetry, - session_id: &str, - wire_session_id: u64, - mut message: StageWireMessage, - downstream: Option<&mut TcpStream>, - wire_dtype: WireActivationDType, - downstream_wire_condition: WireCondition, - activation_width: i32, - control_started: Instant, - mut control_stats: StageReplyStats, - prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, - downstream_connect_timeout_secs: u64, - native_mtp_enabled: bool, -) -> Result<()> { - let (prefix_tokens, current_token) = restore_decode_sideband(&message)?; - let local = maybe_prefix_cache_control( - config, - runtime, - kv, - telemetry, - session_id, - &message, - prefix_tokens, - ); - control_stats.merge(local.stats); - if !local.hit { - let mut attrs = binary_message_attrs(config, wire_session_id, &message); - attrs.insert("skippy.kv.control_hit".to_string(), json!(false)); - attrs.insert( - "llama_stage.elapsed_ms".to_string(), - json!(elapsed_ms(control_started)), - ); - telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); - send_one_off_direct_return( - config, - topology, - &message, - wire_dtype, - downstream_connect_timeout_secs, - StageReply { - kind: WireReplyKind::Ack, - predicted: 0, - predicted_tokens: Vec::new(), - stats: control_stats, - }, - ) - .context("send restore-decode miss direct ACK")?; - return Ok(()); - } - - let input = input_activation_frame(config, topology, &mut message, activation_width)?; - let decode_message = restore_prefill_decode_as_decode_message(&message, current_token); - let compute_started = Instant::now(); - let (predicted_token, output, runtime_lock_wait_ms, runtime_lock_hold_ms, proactive_eviction) = { - let lock_started = Instant::now(); - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - let runtime_lock_wait_ms = elapsed_ms(lock_started); - let lock_hold_started = Instant::now(); - if let Some(metadata) = message.chat_sampling_metadata.as_deref() { - let sampling = runtime_sampling_config(message.sampling.as_ref()); - runtime - .configure_chat_sampling( - session_id, - metadata, - message.state.prompt_token_count.max(0) as u64, - sampling.as_ref(), - ) - .context("configure restore-decode chat sampling")?; - } - let proactive_eviction = evict_binary_resident_prefix_for_decode( - &mut runtime, - kv, - session_id, - BinaryProactiveEvictionPlan { - required: true, - ensure_session_before_eviction: false, - }, - )?; - let (predicted, _, output) = run_binary_stage_message( - &mut runtime, - session_id, - &decode_message, - &[current_token], - input.as_ref(), - BinaryStageExecutionOptions::new( - downstream.is_none(), - stage_output_activation_capacity( - config, - decode_message.token_count, - activation_width, - )?, - native_mtp_enabled, - ), - ) - .context("execute restore-decode stage message")?; - ( - predicted, - output, - runtime_lock_wait_ms, - elapsed_ms(lock_hold_started), - proactive_eviction, - ) - }; - let compute_ms = elapsed_ms(compute_started); - emit_binary_proactive_eviction(telemetry, &proactive_eviction); - - if let Some(downstream) = downstream { - let forwarded = - forwarded_stage_message_timed(config, &message, &output, wire_dtype, activation_width) - .context("forward restore-decode activation")?; - write_stage_message_conditioned( - &mut *downstream, - &forwarded.message, - wire_dtype, - downstream_wire_condition, - ) - .context("forward restore-decode downstream")?; - let mut attrs = binary_message_attrs(config, wire_session_id, &message); - attrs.insert("skippy.kv.control_hit".to_string(), json!(true)); - attrs.insert( - "llama_stage.elapsed_ms".to_string(), - json!(elapsed_ms(control_started)), - ); - attrs.insert("llama_stage.compute_ms".to_string(), json!(compute_ms)); - attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - json!(runtime_lock_wait_ms), - ); - attrs.insert( - "llama_stage.runtime_lock_hold_ms".to_string(), - json!(runtime_lock_hold_ms), - ); - proactive_eviction.insert_attrs(&mut attrs); - attrs.insert( - "llama_stage.forward_activation_bytes".to_string(), - json!(forwarded.message.activation.len()), - ); - attrs.insert( - "llama_stage.activation_encode_ms".to_string(), - json!(forwarded.activation_encode_ms), - ); - telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); - let downstream_reply = - recv_reply(&mut *downstream).context("restore-decode downstream reply")?; - if downstream_reply.kind != WireReplyKind::PredictedToken { - bail!( - "restore-decode expected downstream PredictedToken, got {:?}", - downstream_reply.kind - ); - } - control_stats.merge(downstream_reply.stats); - send_one_off_direct_return( - config, - topology, - &message, - wire_dtype, - downstream_connect_timeout_secs, - StageReply { - stats: control_stats, - ..downstream_reply - }, - ) - .context("relay restore-decode direct predicted reply")?; - return Ok(()); - } - - { - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - let record = maybe_record_binary_full_prefill( - config, - &mut runtime, - kv, - telemetry, - session_id, - &message, - message.tokens.as_slice(), - ); - add_binary_record_stats(&mut control_stats, config, &record); - } - let mut attrs = binary_message_attrs(config, wire_session_id, &message); - attrs.insert("skippy.kv.control_hit".to_string(), json!(true)); - attrs.insert( - "llama_stage.elapsed_ms".to_string(), - json!(elapsed_ms(control_started)), - ); - attrs.insert("llama_stage.compute_ms".to_string(), json!(compute_ms)); - attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - json!(runtime_lock_wait_ms), - ); - attrs.insert( - "llama_stage.runtime_lock_hold_ms".to_string(), - json!(runtime_lock_hold_ms), - ); - proactive_eviction.insert_attrs(&mut attrs); - telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); - let return_stream = prediction_return_streams - .get_mut(&(message.request_id, message.session_id)) - .ok_or_else(|| anyhow!("missing direct prediction return stream"))?; - direct_return::send_direct_prediction_return( - return_stream, - StageReply { - kind: WireReplyKind::PredictedToken, - predicted: predicted_token, - predicted_tokens: vec![predicted_token], - stats: control_stats, - }, - ) - .context("send restore-decode direct predicted reply")?; - Ok(()) +fn send_stage_reply(stream: &mut TcpStream, reply: StageReply) -> Result<()> { + send_reply_message(stream, &reply).context("send stage reply") } -fn send_one_off_direct_return( - config: &StageConfig, - topology: Option<&StageTopology>, +fn reply_window_for_message( message: &StageWireMessage, - wire_dtype: WireActivationDType, - downstream_connect_timeout_secs: u64, - reply: StageReply, -) -> Result<()> { - let mut stream = direct_return::open_prediction_return_stream( - config, - topology, - message.request_id, - message.session_id, - wire_dtype, - downstream_connect_timeout_secs, - )?; - direct_return::send_direct_prediction_return(&mut stream, reply) -} - -fn send_stage_reply(stream: &mut TcpStream, reply: StageReply) -> Result<()> { - match reply.kind { - WireReplyKind::PredictedToken => send_reply_predicted_with_tokens_and_stats( - stream, - reply.predicted, - &reply.predicted_tokens, - reply.stats, - ) - .context("send predicted-token reply"), - WireReplyKind::PredictedTokens => { - send_reply_predicted_tokens_with_stats(stream, &reply.predicted_tokens, reply.stats) - .context("send predicted-tokens reply") + predicted_tokens: &[i32], +) -> skippy_protocol::binary::StageReplyWindow { + if message.kind == WireMessageKind::VerifyWindow { + let accepted_len = message + .tokens + .iter() + .skip(1) + .zip(predicted_tokens) + .take_while(|(input, predicted)| input == predicted) + .count(); + skippy_protocol::binary::StageReplyWindow { + window_id: message.state.seq_id, + accepted_len: i32::try_from(accepted_len).unwrap_or(i32::MAX), + correction_token: predicted_tokens + .get(accepted_len) + .copied() + .unwrap_or(LLAMA_TOKEN_NULL), } - WireReplyKind::Ack => send_reply_ack_with_stats(stream, reply.stats).context("send ACK"), + } else { + Default::default() } } @@ -2691,36 +2529,6 @@ fn emit_binary_proactive_eviction(telemetry: &Telemetry, eviction: &BinaryProact } } -fn restore_decode_sideband(message: &StageWireMessage) -> Result<(&[i32], i32)> { - let Some((¤t, prefix_tokens)) = message.tokens.split_last() else { - bail!("restore-decode message requires prefix tokens plus current token"); - }; - if prefix_tokens.is_empty() { - bail!("restore-decode message requires non-empty prefix tokens"); - } - Ok((prefix_tokens, current)) -} - -fn restore_prefill_decode_as_decode_message( - message: &StageWireMessage, - current_token: i32, -) -> StageWireMessage { - let mut decode = message.clone(); - decode.kind = WireMessageKind::DecodeEmbd; - decode.token_count = 1; - decode.tokens = vec![current_token]; - decode.positions.clear(); - decode.activation.clear(); - decode.raw_bytes.clear(); - decode.state.phase = StageStateHeader::new( - WireMessageKind::DecodeEmbd, - message.state.dtype().unwrap_or(WireActivationDType::F32), - ) - .phase; - decode.state.current_token = current_token; - decode -} - #[allow(clippy::too_many_arguments)] fn maybe_lookup_binary_prefill( config: &StageConfig, @@ -3319,21 +3127,21 @@ struct BinaryRequestSummary { session_auto_align_count: usize, session_auto_align_ms: f64, session_auto_align_trimmed_tokens: u64, - verify_span_count: usize, - verify_span_session_auto_align_count: usize, - verify_span_session_auto_align_ms: f64, - verify_span_session_auto_align_trimmed_tokens: u64, - verify_span_token_count: u64, - verify_span_max_tokens: u64, - verify_span_compute_ms: f64, - verify_span_input_activation_decode_ms: f64, - verify_span_runtime_lock_hold_ms: f64, - verify_span_upstream_reply_ms: f64, - verify_span_pre_compute_ms: f64, - verify_span_post_compute_ms: f64, - verify_span_pre_reply_ms: f64, - verify_span_after_reply_ms: f64, - verify_span_upstream_message_wait_ms: f64, + verify_window_count: usize, + verify_window_session_auto_align_count: usize, + verify_window_session_auto_align_ms: f64, + verify_window_session_auto_align_trimmed_tokens: u64, + verify_window_token_count: u64, + verify_window_max_tokens: u64, + verify_window_compute_ms: f64, + verify_window_input_activation_decode_ms: f64, + verify_window_runtime_lock_hold_ms: f64, + verify_window_upstream_reply_ms: f64, + verify_window_pre_compute_ms: f64, + verify_window_post_compute_ms: f64, + verify_window_pre_reply_ms: f64, + verify_window_after_reply_ms: f64, + verify_window_upstream_message_wait_ms: f64, reply_stats: StageReplyStats, } @@ -3359,10 +3167,10 @@ struct BinaryMessageObservation<'a> { session_auto_align_count: usize, session_auto_align_ms: f64, session_auto_align_trimmed_tokens: u64, - verify_span_pre_compute_ms: f64, - verify_span_post_compute_ms: f64, - verify_span_pre_reply_ms: f64, - verify_span_after_reply_ms: f64, + verify_window_pre_compute_ms: f64, + verify_window_post_compute_ms: f64, + verify_window_pre_reply_ms: f64, + verify_window_after_reply_ms: f64, upstream_message_wait_ms: f64, } @@ -3517,25 +3325,26 @@ impl BinaryRequestSummary { self.session_auto_align_trimmed_tokens = self .session_auto_align_trimmed_tokens .saturating_add(observation.session_auto_align_trimmed_tokens); - if message.kind == WireMessageKind::VerifySpan { + if message.kind == WireMessageKind::VerifyWindow { let token_count = message.token_count.max(0) as u64; - self.verify_span_count += 1; - self.verify_span_token_count = self.verify_span_token_count.saturating_add(token_count); - self.verify_span_max_tokens = self.verify_span_max_tokens.max(token_count); - self.verify_span_session_auto_align_count += observation.session_auto_align_count; - self.verify_span_session_auto_align_ms += observation.session_auto_align_ms; - self.verify_span_session_auto_align_trimmed_tokens = self - .verify_span_session_auto_align_trimmed_tokens + self.verify_window_count += 1; + self.verify_window_token_count = + self.verify_window_token_count.saturating_add(token_count); + self.verify_window_max_tokens = self.verify_window_max_tokens.max(token_count); + self.verify_window_session_auto_align_count += observation.session_auto_align_count; + self.verify_window_session_auto_align_ms += observation.session_auto_align_ms; + self.verify_window_session_auto_align_trimmed_tokens = self + .verify_window_session_auto_align_trimmed_tokens .saturating_add(observation.session_auto_align_trimmed_tokens); - self.verify_span_compute_ms += observation.compute_ms; - self.verify_span_input_activation_decode_ms += observation.input_activation_decode_ms; - self.verify_span_runtime_lock_hold_ms += observation.runtime_lock_hold_ms; - self.verify_span_upstream_reply_ms += observation.upstream_reply_ms; - self.verify_span_pre_compute_ms += observation.verify_span_pre_compute_ms; - self.verify_span_post_compute_ms += observation.verify_span_post_compute_ms; - self.verify_span_pre_reply_ms += observation.verify_span_pre_reply_ms; - self.verify_span_after_reply_ms += observation.verify_span_after_reply_ms; - self.verify_span_upstream_message_wait_ms += observation.upstream_message_wait_ms; + self.verify_window_compute_ms += observation.compute_ms; + self.verify_window_input_activation_decode_ms += observation.input_activation_decode_ms; + self.verify_window_runtime_lock_hold_ms += observation.runtime_lock_hold_ms; + self.verify_window_upstream_reply_ms += observation.upstream_reply_ms; + self.verify_window_pre_compute_ms += observation.verify_window_pre_compute_ms; + self.verify_window_post_compute_ms += observation.verify_window_post_compute_ms; + self.verify_window_pre_reply_ms += observation.verify_window_pre_reply_ms; + self.verify_window_after_reply_ms += observation.verify_window_after_reply_ms; + self.verify_window_upstream_message_wait_ms += observation.upstream_message_wait_ms; } self.reply_stats.merge(observation.reply_stats); } @@ -3662,115 +3471,115 @@ impl BinaryRequestSummary { ); } attrs.insert( - "skippy.verify_span_count".to_string(), - json!(self.verify_span_count), + "skippy.verify_window_count".to_string(), + json!(self.verify_window_count), ); attrs.insert( - "skippy.verify_span_token_count".to_string(), - json!(self.verify_span_token_count), + "skippy.verify_window_token_count".to_string(), + json!(self.verify_window_token_count), ); attrs.insert( - "skippy.verify_span_max_tokens".to_string(), - json!(self.verify_span_max_tokens), + "skippy.verify_window_max_tokens".to_string(), + json!(self.verify_window_max_tokens), ); attrs.insert( - "skippy.verify_span_session_auto_align_count".to_string(), - json!(self.verify_span_session_auto_align_count), + "skippy.verify_window_session_auto_align_count".to_string(), + json!(self.verify_window_session_auto_align_count), ); attrs.insert( - "skippy.verify_span_session_auto_align_ms".to_string(), - json!(self.verify_span_session_auto_align_ms), + "skippy.verify_window_session_auto_align_ms".to_string(), + json!(self.verify_window_session_auto_align_ms), ); attrs.insert( - "skippy.verify_span_session_auto_align_trimmed_tokens".to_string(), - json!(self.verify_span_session_auto_align_trimmed_tokens), + "skippy.verify_window_session_auto_align_trimmed_tokens".to_string(), + json!(self.verify_window_session_auto_align_trimmed_tokens), ); - if self.verify_span_session_auto_align_count > 0 { + if self.verify_window_session_auto_align_count > 0 { attrs.insert( - "skippy.verify_span_session_auto_align_ms_avg".to_string(), + "skippy.verify_window_session_auto_align_ms_avg".to_string(), json!( - self.verify_span_session_auto_align_ms - / self.verify_span_session_auto_align_count as f64 + self.verify_window_session_auto_align_ms + / self.verify_window_session_auto_align_count as f64 ), ); } attrs.insert( - "skippy.verify_span_pre_compute_ms".to_string(), - json!(self.verify_span_pre_compute_ms), + "skippy.verify_window_pre_compute_ms".to_string(), + json!(self.verify_window_pre_compute_ms), ); attrs.insert( - "skippy.verify_span_compute_ms".to_string(), - json!(self.verify_span_compute_ms), + "skippy.verify_window_compute_ms".to_string(), + json!(self.verify_window_compute_ms), ); attrs.insert( - "skippy.verify_span_input_activation_decode_ms".to_string(), - json!(self.verify_span_input_activation_decode_ms), + "skippy.verify_window_input_activation_decode_ms".to_string(), + json!(self.verify_window_input_activation_decode_ms), ); attrs.insert( - "skippy.verify_span_runtime_lock_hold_ms".to_string(), - json!(self.verify_span_runtime_lock_hold_ms), + "skippy.verify_window_runtime_lock_hold_ms".to_string(), + json!(self.verify_window_runtime_lock_hold_ms), ); attrs.insert( - "skippy.verify_span_upstream_reply_ms".to_string(), - json!(self.verify_span_upstream_reply_ms), + "skippy.verify_window_upstream_reply_ms".to_string(), + json!(self.verify_window_upstream_reply_ms), ); attrs.insert( - "skippy.verify_span_post_compute_ms".to_string(), - json!(self.verify_span_post_compute_ms), + "skippy.verify_window_post_compute_ms".to_string(), + json!(self.verify_window_post_compute_ms), ); attrs.insert( - "skippy.verify_span_pre_reply_ms".to_string(), - json!(self.verify_span_pre_reply_ms), + "skippy.verify_window_pre_reply_ms".to_string(), + json!(self.verify_window_pre_reply_ms), ); attrs.insert( - "skippy.verify_span_after_reply_ms".to_string(), - json!(self.verify_span_after_reply_ms), + "skippy.verify_window_after_reply_ms".to_string(), + json!(self.verify_window_after_reply_ms), ); attrs.insert( - "skippy.verify_span_upstream_message_wait_ms".to_string(), - json!(self.verify_span_upstream_message_wait_ms), + "skippy.verify_window_upstream_message_wait_ms".to_string(), + json!(self.verify_window_upstream_message_wait_ms), ); - if self.verify_span_count > 0 { - let verify_span_count = self.verify_span_count as f64; + if self.verify_window_count > 0 { + let verify_window_count = self.verify_window_count as f64; attrs.insert( - "skippy.verify_span_pre_compute_ms_avg".to_string(), - json!(self.verify_span_pre_compute_ms / verify_span_count), + "skippy.verify_window_pre_compute_ms_avg".to_string(), + json!(self.verify_window_pre_compute_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_compute_ms_avg".to_string(), - json!(self.verify_span_compute_ms / verify_span_count), + "skippy.verify_window_compute_ms_avg".to_string(), + json!(self.verify_window_compute_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_input_activation_decode_ms_avg".to_string(), - json!(self.verify_span_input_activation_decode_ms / verify_span_count), + "skippy.verify_window_input_activation_decode_ms_avg".to_string(), + json!(self.verify_window_input_activation_decode_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_runtime_lock_hold_ms_avg".to_string(), - json!(self.verify_span_runtime_lock_hold_ms / verify_span_count), + "skippy.verify_window_runtime_lock_hold_ms_avg".to_string(), + json!(self.verify_window_runtime_lock_hold_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_upstream_reply_ms_avg".to_string(), - json!(self.verify_span_upstream_reply_ms / verify_span_count), + "skippy.verify_window_upstream_reply_ms_avg".to_string(), + json!(self.verify_window_upstream_reply_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_tokens_avg".to_string(), - json!(self.verify_span_token_count as f64 / verify_span_count), + "skippy.verify_window_tokens_avg".to_string(), + json!(self.verify_window_token_count as f64 / verify_window_count), ); attrs.insert( - "skippy.verify_span_post_compute_ms_avg".to_string(), - json!(self.verify_span_post_compute_ms / verify_span_count), + "skippy.verify_window_post_compute_ms_avg".to_string(), + json!(self.verify_window_post_compute_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_pre_reply_ms_avg".to_string(), - json!(self.verify_span_pre_reply_ms / verify_span_count), + "skippy.verify_window_pre_reply_ms_avg".to_string(), + json!(self.verify_window_pre_reply_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_after_reply_ms_avg".to_string(), - json!(self.verify_span_after_reply_ms / verify_span_count), + "skippy.verify_window_after_reply_ms_avg".to_string(), + json!(self.verify_window_after_reply_ms / verify_window_count), ); attrs.insert( - "skippy.verify_span_upstream_message_wait_ms_avg".to_string(), - json!(self.verify_span_upstream_message_wait_ms / verify_span_count), + "skippy.verify_window_upstream_message_wait_ms_avg".to_string(), + json!(self.verify_window_upstream_message_wait_ms / verify_window_count), ); } let lookups = self.reply_stats.kv_lookup_hits + self.reply_stats.kv_lookup_misses; @@ -3958,7 +3767,7 @@ pub(crate) fn run_binary_stage_message( output, )) } - WireMessageKind::VerifySpan => { + WireMessageKind::VerifyWindow => { let sampling = runtime_sampling_config(message.sampling.as_ref()); let (predicted_tokens, output) = runtime.verify_frame_sampled( session_id, diff --git a/crates/skippy-server/src/binary_transport/direct_return.rs b/crates/skippy-server/src/binary_transport/direct_return.rs index cebeade3c7..4fcc38b6f4 100644 --- a/crates/skippy-server/src/binary_transport/direct_return.rs +++ b/crates/skippy-server/src/binary_transport/direct_return.rs @@ -17,9 +17,8 @@ use skippy_protocol::{ StageConfig, StageTopology, binary::{ StageReply, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, - WireReplyKind, read_stage_message, recv_ready, recv_reply, send_ready, - send_reply_ack_with_stats, send_reply_predicted_tokens_with_stats, - send_reply_predicted_with_tokens_and_stats, write_stage_message, + WireReplyKind, read_stage_message, recv_ready, recv_reply, send_ready, send_reply_message, + write_stage_message, }, }; @@ -215,12 +214,16 @@ impl PredictionReturnReceiver { } pub(crate) fn try_recv_expected(&self, expected: WireReplyKind) -> Result> { + self.try_recv_one_of(std::slice::from_ref(&expected)) + } + + pub(crate) fn try_recv_one_of(&self, expected: &[WireReplyKind]) -> Result> { let Some(reply) = self.try_recv()? else { return Ok(None); }; - if reply.kind != expected { + if !expected.contains(&reply.kind) { bail!( - "expected {expected:?} direct prediction return, got {:?}", + "expected one of {expected:?} from direct prediction return, got {:?}", reply.kind ); } @@ -285,6 +288,13 @@ impl PredictionReturnSinks { thread::sleep(Duration::from_millis(2)); } } + + pub(crate) fn remove(&self, request_id: u64, session_id: u64) { + let key = PredictionReturnKey::new(request_id, session_id); + if let Ok(mut streams) = self.streams.lock() { + streams.remove(&key); + } + } } /// Read timeout for the return-sink ready handshake. `recv_ready` is a blocking @@ -389,22 +399,7 @@ pub(crate) fn send_direct_prediction_return( stream: &mut TcpStream, reply: StageReply, ) -> Result<()> { - match reply.kind { - WireReplyKind::PredictedToken => send_reply_predicted_with_tokens_and_stats( - stream, - reply.predicted, - &reply.predicted_tokens, - reply.stats, - ) - .context("send direct predicted-token return"), - WireReplyKind::PredictedTokens => { - send_reply_predicted_tokens_with_stats(stream, &reply.predicted_tokens, reply.stats) - .context("send direct predicted-tokens return") - } - WireReplyKind::Ack => { - send_reply_ack_with_stats(stream, reply.stats).context("send direct ACK return") - } - } + send_reply_message(stream, &reply).context("send direct prediction return") } fn driver_stage_endpoint<'a>( @@ -487,7 +482,7 @@ mod tests { } #[test] - fn direct_prediction_return_preserves_predicted_token_sideband() { + fn direct_prediction_return_preserves_typed_native_mtp_draft() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); let mut client = TcpStream::connect(addr).unwrap(); @@ -496,7 +491,16 @@ mod tests { let reply = StageReply { kind: WireReplyKind::PredictedToken, predicted: 42, - predicted_tokens: vec![42, 43, 123], + predicted_tokens: vec![42], + native_mtp_draft: Some(skippy_protocol::binary::StageNativeMtpDraft { + token_ids: vec![43], + proposal_compute_us: 123, + }), + window: skippy_protocol::binary::StageReplyWindow { + window_id: 7, + accepted_len: 2, + correction_token: 123, + }, stats: Default::default(), }; send_direct_prediction_return(&mut server, reply).unwrap(); @@ -504,7 +508,17 @@ mod tests { let received = recv_reply(&mut client).unwrap(); assert_eq!(received.kind, WireReplyKind::PredictedToken); assert_eq!(received.predicted, 42); - assert_eq!(received.predicted_tokens, vec![42, 43, 123]); + assert_eq!(received.predicted_tokens, vec![42]); + assert_eq!( + received.native_mtp_draft, + Some(skippy_protocol::binary::StageNativeMtpDraft { + token_ids: vec![43], + proposal_compute_us: 123, + }) + ); + assert_eq!(received.window.window_id, 7); + assert_eq!(received.window.accepted_len, 2); + assert_eq!(received.window.correction_token, 123); } #[test] @@ -531,6 +545,32 @@ mod tests { assert_eq!(stream.peer_addr().unwrap(), client.local_addr().unwrap()); } + #[test] + fn prediction_return_sinks_remove_abandoned_streams() { + let request_id = 41; + let session_id = 43; + let sinks = PredictionReturnSinks::default(); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let client = TcpStream::connect(listener.local_addr().unwrap()).unwrap(); + let (server, _) = listener.accept().unwrap(); + + sinks + .insert_opened_sink( + prediction_return_open_message(request_id, session_id), + server, + ) + .unwrap(); + sinks.remove(request_id, session_id); + + assert!( + sinks + .take_wait(request_id, session_id, Duration::from_millis(1)) + .unwrap() + .is_none() + ); + drop(client); + } + fn poll_test_reply(receiver: &PredictionReturnReceiver, expected: WireReplyKind) -> StageReply { let started = std::time::Instant::now(); loop { diff --git a/crates/skippy-server/src/binary_transport/kv_eviction.rs b/crates/skippy-server/src/binary_transport/kv_eviction.rs index e3da6f95cd..6e023e44e4 100644 --- a/crates/skippy-server/src/binary_transport/kv_eviction.rs +++ b/crates/skippy-server/src/binary_transport/kv_eviction.rs @@ -82,7 +82,7 @@ pub(super) fn binary_proactive_eviction_required( | WireMessageKind::DecodeReplayFinalEmbd | WireMessageKind::DecodeReadout | WireMessageKind::DecodeLightCtx - | WireMessageKind::VerifySpan + | WireMessageKind::VerifyWindow ) } diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index 060700756c..6783b757c2 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -3,7 +3,12 @@ use std::{net::SocketAddr, path::PathBuf}; use anyhow::{Context, Result, bail}; use skippy_protocol::{StageConfig, StageTopology, binary::WireActivationDType}; -use crate::{cli::ServeBinaryArgs, config::load_json, telemetry::TelemetryLevel}; +use crate::{ + cli::ServeBinaryArgs, + config::load_json, + frontend::{NgramProposalConfig, NgramProposerKind, SpeculativeDecodeConfig}, + telemetry::TelemetryLevel, +}; use super::WireCondition; @@ -42,10 +47,9 @@ pub struct EmbeddedOpenAiStageOptions { pub speculative_window: usize, pub adaptive_speculative_window: bool, pub draft_n_gpu_layers: Option, - pub ngram_min: usize, - pub ngram_max: usize, pub native_mtp_max_tokens: usize, pub native_mtp_min_tokens: usize, + pub speculative: SpeculativeDecodeConfig, } impl BinaryStageOptions { @@ -72,6 +76,14 @@ impl BinaryStageOptions { None => None, }; let bind_addr = args.bind_addr.unwrap_or(config.bind_addr.parse()?); + let openai_speculative = args + .openai_speculative_config + .as_ref() + .map(load_json) + .transpose() + .context("load --openai-speculative-config")? + .unwrap_or_else(|| legacy_speculative_config(&args)); + openai_speculative.validate()?; let openai = args .openai_bind_addr .map(|bind_addr| EmbeddedOpenAiStageOptions { @@ -89,12 +101,11 @@ impl BinaryStageOptions { speculative_window: args.openai_speculative_window, adaptive_speculative_window: args.openai_adaptive_speculative_window, draft_n_gpu_layers: args.openai_draft_n_gpu_layers, - ngram_min: args.openai_ngram_min, - ngram_max: args.openai_ngram_max, native_mtp_max_tokens: 3, native_mtp_min_tokens: 0, + speculative: openai_speculative, }); - let native_mtp_enabled = native_mtp_enabled_from_env(); + let native_mtp_enabled = config.native_mtp_enabled; Ok(Self { config, topology, @@ -115,17 +126,18 @@ impl BinaryStageOptions { } } -fn native_mtp_enabled_from_env() -> bool { - !matches!( - std::env::var("SKIPPY_NATIVE_MTP_ENABLED") - .ok() - .map(|value| value.trim().to_ascii_lowercase()), - Some(value) - if matches!( - value.as_str(), - "0" | "false" | "off" | "disable" | "disabled" | "no" - ) - ) +fn legacy_speculative_config(args: &ServeBinaryArgs) -> SpeculativeDecodeConfig { + let mut config = SpeculativeDecodeConfig::default(); + if args.openai_ngram_min > 0 && args.openai_ngram_max > 0 { + config.effective_strategy = "ngram-simple".to_string(); + config.ngram = Some(NgramProposalConfig { + kind: NgramProposerKind::Simple, + min_ngram: args.openai_ngram_min, + max_ngram: args.openai_ngram_max, + max_proposal_tokens: args.openai_ngram_max, + }); + } + config } pub fn parse_wire_dtype(value: &str) -> Result { @@ -136,3 +148,140 @@ pub fn parse_wire_dtype(value: &str) -> Result { _ => bail!("unsupported activation wire dtype {value}"), } } + +#[cfg(test)] +mod tests { + use std::fs; + + use clap::Parser; + use skippy_protocol::{FlashAttentionType, LoadMode, StageConfig}; + + use super::*; + use crate::{ + cli::{Cli, Command}, + frontend::{NativeMtpProposalConfig, NgramExtensionConfig, VerifyWindowConfig}, + }; + + fn stage_config() -> StageConfig { + StageConfig { + run_id: "run".to_string(), + topology_id: "topology".to_string(), + model_id: "model".to_string(), + package_ref: None, + manifest_sha256: None, + source_model_path: None, + source_model_sha256: None, + source_model_bytes: None, + materialized_path: None, + materialized_pinned: false, + model_path: Some("/tmp/model.gguf".to_string()), + projector_path: None, + stage_id: "stage-0".to_string(), + stage_index: 0, + layer_start: 0, + layer_end: 4, + ctx_size: 512, + lane_count: 1, + n_batch: None, + n_ubatch: None, + n_gpu_layers: -1, + mmap: None, + mlock: false, + cache_type_k: "f16".to_string(), + cache_type_v: "f16".to_string(), + flash_attn_type: FlashAttentionType::Auto, + filter_tensors_on_load: true, + selected_device: None, + kv_cache: None, + native_mtp_enabled: true, + load_mode: LoadMode::RuntimeSlice, + bind_addr: "127.0.0.1:0".to_string(), + upstream: None, + downstream: None, + } + } + + fn cache_composite_plan() -> SpeculativeDecodeConfig { + SpeculativeDecodeConfig { + requested_strategy: "mtp-cache".to_string(), + effective_strategy: "native-mtp+ngram-cache".to_string(), + native_mtp: NativeMtpProposalConfig { + enabled: true, + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }, + ngram: Some(NgramProposalConfig { + kind: NgramProposerKind::Cache, + min_ngram: 2, + max_ngram: 4, + max_proposal_tokens: 6, + }), + extension: Some(NgramExtensionConfig { + initial_tokens: 2, + max_tokens: 6, + tail_backoff_proposals: 2, + }), + verify_window: VerifyWindowConfig { + min_tokens: 1, + max_tokens: 6, + pipeline_depth: 2, + }, + } + } + + #[test] + fn typed_speculative_plan_reaches_embedded_stage_without_policy_merging() { + let dir = tempfile::tempdir().expect("create temp directory"); + let stage_path = dir.path().join("stage.json"); + let plan_path = dir.path().join("speculative.json"); + fs::write( + &stage_path, + serde_json::to_vec(&stage_config()).expect("serialize stage config"), + ) + .expect("write stage config"); + let expected = cache_composite_plan(); + fs::write( + &plan_path, + serde_json::to_vec(&expected).expect("serialize speculative config"), + ) + .expect("write speculative config"); + + let cli = Cli::try_parse_from([ + "skippy-server", + "serve-binary", + "--config", + stage_path.to_str().expect("UTF-8 stage path"), + "--activation-width", + "2048", + "--openai-bind-addr", + "127.0.0.1:9337", + "--openai-speculative-config", + plan_path.to_str().expect("UTF-8 plan path"), + ]) + .expect("parse binary stage CLI"); + let Command::ServeBinary(args) = cli.command else { + panic!("expected serve-binary command"); + }; + + let options = BinaryStageOptions::from_cli_args(args).expect("resolve binary stage"); + let openai = options.openai.expect("embedded OpenAI configuration"); + + assert!(options.native_mtp_enabled); + assert_eq!(openai.speculative, expected); + } + + #[test] + fn cache_composite_plan_is_json_stable_for_stage_handoff() { + let plan = cache_composite_plan(); + let json = serde_json::to_value(&plan).expect("serialize speculative plan"); + + assert_eq!( + json["ngram"]["kind"], + serde_json::Value::String("cache".to_string()) + ); + assert_eq!(json["verify_window"]["pipeline_depth"], 2); + } +} diff --git a/crates/skippy-server/src/binary_transport/restore_prefill_decode.rs b/crates/skippy-server/src/binary_transport/restore_prefill_decode.rs new file mode 100644 index 0000000000..58e4403283 --- /dev/null +++ b/crates/skippy-server/src/binary_transport/restore_prefill_decode.rs @@ -0,0 +1,383 @@ +use super::*; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RestorePrefillDecodeRoute { + DirectMiss, + ForwardHit, + DirectHit, +} + +fn restore_prefill_decode_route( + local_hit: bool, + has_downstream: bool, +) -> RestorePrefillDecodeRoute { + match (local_hit, has_downstream) { + (false, _) => RestorePrefillDecodeRoute::DirectMiss, + (true, true) => RestorePrefillDecodeRoute::ForwardHit, + (true, false) => RestorePrefillDecodeRoute::DirectHit, + } +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_binary_restore_prefill_decode_control( + config: &StageConfig, + topology: Option<&StageTopology>, + runtime: &Arc>, + kv: Option<&Arc>, + telemetry: &Telemetry, + session_id: &str, + wire_session_id: u64, + mut message: StageWireMessage, + downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + activation_width: i32, + control_started: Instant, + mut control_stats: StageReplyStats, + prediction_return_sinks: &PredictionReturnSinks, + prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, + downstream_connect_timeout_secs: u64, + native_mtp_enabled: bool, +) -> Result<()> { + let has_downstream = downstream.is_some(); + if !has_downstream { + configure_prediction_return_stream( + config, + topology, + message.request_id, + message.session_id, + wire_dtype, + downstream_connect_timeout_secs, + prediction_return_sinks, + prediction_return_streams, + ); + } + + let (prefix_tokens, current_token) = restore_decode_sideband(&message)?; + let local = maybe_prefix_cache_control( + config, + runtime, + kv, + telemetry, + session_id, + &message, + prefix_tokens, + ); + control_stats.merge(local.stats); + let route = restore_prefill_decode_route(local.hit, has_downstream); + if route == RestorePrefillDecodeRoute::DirectMiss { + emit_restore_decode_control( + config, + telemetry, + wire_session_id, + &message, + control_started, + false, + None, + ); + send_restore_decode_direct_reply( + config, + topology, + &message, + wire_dtype, + downstream_connect_timeout_secs, + prediction_return_streams, + StageReply { + kind: WireReplyKind::Ack, + predicted: 0, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: Default::default(), + stats: control_stats, + }, + ) + .context("send restore-decode miss direct ACK")?; + return Ok(()); + } + + let input = input_activation_frame(config, topology, &mut message, activation_width)?; + let decode_message = restore_prefill_decode_as_decode_message(&message, current_token); + let compute_started = Instant::now(); + let (predicted_token, output, runtime_lock_wait_ms, runtime_lock_hold_ms, proactive_eviction) = { + let lock_started = Instant::now(); + let mut runtime = runtime.lock().expect("runtime lock poisoned"); + let runtime_lock_wait_ms = elapsed_ms(lock_started); + let lock_hold_started = Instant::now(); + if let Some(metadata) = message.chat_sampling_metadata.as_deref() { + let sampling = runtime_sampling_config(message.sampling.as_ref()); + runtime + .configure_chat_sampling( + session_id, + metadata, + message.state.prompt_token_count.max(0) as u64, + sampling.as_ref(), + ) + .context("configure restore-decode chat sampling")?; + } + let proactive_eviction = evict_binary_resident_prefix_for_decode( + &mut runtime, + kv, + session_id, + BinaryProactiveEvictionPlan { + required: true, + ensure_session_before_eviction: false, + }, + )?; + let (predicted, _, output) = run_binary_stage_message( + &mut runtime, + session_id, + &decode_message, + &[current_token], + input.as_ref(), + BinaryStageExecutionOptions::new( + route == RestorePrefillDecodeRoute::DirectHit, + stage_output_activation_capacity( + config, + decode_message.token_count, + activation_width, + )?, + native_mtp_enabled, + ), + ) + .context("execute restore-decode stage message")?; + ( + predicted, + output, + runtime_lock_wait_ms, + elapsed_ms(lock_hold_started), + proactive_eviction, + ) + }; + let compute_ms = elapsed_ms(compute_started); + emit_binary_proactive_eviction(telemetry, &proactive_eviction); + + if route == RestorePrefillDecodeRoute::ForwardHit { + let downstream = downstream.expect("forward route requires downstream stage"); + let forwarded = + forwarded_stage_message_timed(config, &message, &output, wire_dtype, activation_width) + .context("forward restore-decode activation")?; + write_stage_message_conditioned( + &mut *downstream, + &forwarded.message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward restore-decode downstream")?; + emit_restore_decode_control( + config, + telemetry, + wire_session_id, + &message, + control_started, + true, + Some(RestoreDecodeTiming { + compute_ms, + runtime_lock_wait_ms, + runtime_lock_hold_ms, + proactive_eviction: &proactive_eviction, + forwarded: Some(&forwarded), + }), + ); + return Ok(()); + } + + { + let mut runtime = runtime.lock().expect("runtime lock poisoned"); + let record = maybe_record_binary_full_prefill( + config, + &mut runtime, + kv, + telemetry, + session_id, + &message, + message.tokens.as_slice(), + ); + add_binary_record_stats(&mut control_stats, config, &record); + } + emit_restore_decode_control( + config, + telemetry, + wire_session_id, + &message, + control_started, + true, + Some(RestoreDecodeTiming { + compute_ms, + runtime_lock_wait_ms, + runtime_lock_hold_ms, + proactive_eviction: &proactive_eviction, + forwarded: None, + }), + ); + send_restore_decode_direct_reply( + config, + topology, + &message, + wire_dtype, + downstream_connect_timeout_secs, + prediction_return_streams, + StageReply { + kind: WireReplyKind::PredictedToken, + predicted: predicted_token, + predicted_tokens: vec![predicted_token], + native_mtp_draft: None, + window: Default::default(), + stats: control_stats, + }, + ) + .context("send restore-decode direct predicted reply") +} + +struct RestoreDecodeTiming<'a> { + compute_ms: f64, + runtime_lock_wait_ms: f64, + runtime_lock_hold_ms: f64, + proactive_eviction: &'a BinaryProactiveEviction, + forwarded: Option<&'a forwarding::ForwardedStageMessage>, +} + +#[allow(clippy::too_many_arguments)] +fn emit_restore_decode_control( + config: &StageConfig, + telemetry: &Telemetry, + wire_session_id: u64, + message: &StageWireMessage, + control_started: Instant, + hit: bool, + timing: Option>, +) { + let mut attrs = binary_message_attrs(config, wire_session_id, message); + attrs.insert("skippy.kv.control_hit".to_string(), json!(hit)); + attrs.insert( + "llama_stage.elapsed_ms".to_string(), + json!(elapsed_ms(control_started)), + ); + if let Some(timing) = timing { + attrs.insert( + "llama_stage.compute_ms".to_string(), + json!(timing.compute_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(timing.runtime_lock_wait_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + json!(timing.runtime_lock_hold_ms), + ); + timing.proactive_eviction.insert_attrs(&mut attrs); + if let Some(forwarded) = timing.forwarded { + attrs.insert( + "llama_stage.forward_activation_bytes".to_string(), + json!(forwarded.message.activation.len()), + ); + attrs.insert( + "llama_stage.activation_encode_ms".to_string(), + json!(forwarded.activation_encode_ms), + ); + } + } + telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); +} + +fn send_restore_decode_direct_reply( + config: &StageConfig, + topology: Option<&StageTopology>, + message: &StageWireMessage, + wire_dtype: WireActivationDType, + downstream_connect_timeout_secs: u64, + prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, + reply: StageReply, +) -> Result<()> { + if let Some(stream) = + prediction_return_streams.get_mut(&(message.request_id, message.session_id)) + { + return direct_return::send_direct_prediction_return(stream, reply); + } + send_one_off_direct_return( + config, + topology, + message, + wire_dtype, + downstream_connect_timeout_secs, + reply, + ) +} + +fn send_one_off_direct_return( + config: &StageConfig, + topology: Option<&StageTopology>, + message: &StageWireMessage, + wire_dtype: WireActivationDType, + downstream_connect_timeout_secs: u64, + reply: StageReply, +) -> Result<()> { + let mut stream = direct_return::open_prediction_return_stream( + config, + topology, + message.request_id, + message.session_id, + wire_dtype, + downstream_connect_timeout_secs, + )?; + direct_return::send_direct_prediction_return(&mut stream, reply) +} + +fn restore_decode_sideband(message: &StageWireMessage) -> Result<(&[i32], i32)> { + let Some((¤t, prefix_tokens)) = message.tokens.split_last() else { + bail!("restore-decode message requires prefix tokens plus current token"); + }; + if prefix_tokens.is_empty() { + bail!("restore-decode message requires non-empty prefix tokens"); + } + Ok((prefix_tokens, current)) +} + +pub(super) fn restore_prefill_decode_as_decode_message( + message: &StageWireMessage, + current_token: i32, +) -> StageWireMessage { + let mut decode = message.clone(); + decode.kind = WireMessageKind::DecodeEmbd; + decode.token_count = 1; + decode.tokens = vec![current_token]; + decode.positions.clear(); + decode.activation.clear(); + decode.raw_bytes.clear(); + decode.state.phase = StageStateHeader::new( + WireMessageKind::DecodeEmbd, + message.state.dtype().unwrap_or(WireActivationDType::F32), + ) + .phase; + decode.state.current_token = current_token; + decode +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn restore_decode_routes_misses_and_terminal_hits_directly() { + assert_eq!( + restore_prefill_decode_route(false, true), + RestorePrefillDecodeRoute::DirectMiss + ); + assert_eq!( + restore_prefill_decode_route(false, false), + RestorePrefillDecodeRoute::DirectMiss + ); + assert_eq!( + restore_prefill_decode_route(true, false), + RestorePrefillDecodeRoute::DirectHit + ); + } + + #[test] + fn restore_decode_forwards_intermediate_hits_without_waiting_for_a_lane_reply() { + assert_eq!( + restore_prefill_decode_route(true, true), + RestorePrefillDecodeRoute::ForwardHit + ); + } +} diff --git a/crates/skippy-server/src/binary_transport/tests.rs b/crates/skippy-server/src/binary_transport/tests.rs index 2f489eaaad..fa47e069b6 100644 --- a/crates/skippy-server/src/binary_transport/tests.rs +++ b/crates/skippy-server/src/binary_transport/tests.rs @@ -1,7 +1,7 @@ use super::{ binary_full_prefill_record_identities, decode_record_tokens_sideband, - is_decode_frame_batch_candidate, native_mtp_enabled_from, prepare_binary_stage_connection, - restore_prefill_decode_as_decode_message, token_sideband_or_fill, + is_decode_frame_batch_candidate, prepare_binary_stage_connection, reply_window_for_message, + restore_prefill_decode_as_decode_message, split_native_mtp_reply, token_sideband_or_fill, warm_downstream_preconnect_enabled_from, }; use std::{ @@ -54,16 +54,6 @@ fn accepted_binary_stage_connection_is_blocking() { drop(client.join().unwrap()); } -#[test] -fn native_mtp_enabled_flag_defaults_on_and_accepts_false_values() { - assert!(native_mtp_enabled_from(None)); - assert!(native_mtp_enabled_from(Some("1"))); - assert!(native_mtp_enabled_from(Some("true"))); - assert!(!native_mtp_enabled_from(Some("0"))); - assert!(!native_mtp_enabled_from(Some("false"))); - assert!(!native_mtp_enabled_from(Some(" disabled "))); -} - #[test] fn warm_preconnect_is_opt_in() { assert!(!warm_downstream_preconnect_enabled_from(None)); @@ -119,33 +109,74 @@ fn stale_warm_downstream_connection_is_replaced() { } #[test] -fn request_summary_tracks_verify_span_compute_ms() { +fn request_summary_tracks_verify_window_compute_ms() { let config = prefix_cache_test_config(); let mut summary = super::BinaryRequestSummary::default(); - let verify = test_message(WireMessageKind::VerifySpan, 2); + let verify = test_message(WireMessageKind::VerifyWindow, 2); let decode = test_message(WireMessageKind::DecodeEmbd, 1); summary.observe(summary_observation(&config, &verify, 12.5)); summary.observe(summary_observation(&config, &decode, 7.0)); - assert_eq!(summary.verify_span_count, 1); - assert_eq!(summary.verify_span_token_count, 2); - assert_eq!(summary.verify_span_max_tokens, 2); - assert_eq!(summary.verify_span_compute_ms, 12.5); - assert_eq!(summary.verify_span_input_activation_decode_ms, 1.25); - assert_eq!(summary.verify_span_runtime_lock_hold_ms, 2.5); - assert_eq!(summary.verify_span_upstream_reply_ms, 0.75); + assert_eq!(summary.verify_window_count, 1); + assert_eq!(summary.verify_window_token_count, 2); + assert_eq!(summary.verify_window_max_tokens, 2); + assert_eq!(summary.verify_window_compute_ms, 12.5); + assert_eq!(summary.verify_window_input_activation_decode_ms, 1.25); + assert_eq!(summary.verify_window_runtime_lock_hold_ms, 2.5); + assert_eq!(summary.verify_window_upstream_reply_ms, 0.75); assert_eq!(summary.compute_ms, 19.5); assert_eq!(summary.input_activation_decode_ms, 2.5); assert_eq!(summary.runtime_lock_hold_ms, 5.0); assert_eq!(summary.upstream_reply_ms, 1.5); } +#[test] +fn verify_window_reply_reports_accepted_prefix_and_correction() { + let mut message = test_message(WireMessageKind::VerifyWindow, 3); + message.state.seq_id = 42; + message.tokens = vec![10, 11, 12]; + + let reply = reply_window_for_message(&message, &[11, 99, 100]); + + assert_eq!(reply.window_id, 42); + assert_eq!(reply.accepted_len, 1); + assert_eq!(reply.correction_token, 99); +} + +#[test] +fn native_mtp_sideband_is_removed_from_verify_predictions() { + let mut message = test_message(WireMessageKind::VerifyWindow, 3); + message.tokens = vec![10, 11, 12]; + let mut predictions = vec![11, 12, 13, 2, 14, 15, 123]; + + let draft = split_native_mtp_reply(&message, &mut predictions).unwrap(); + + assert_eq!(predictions, vec![11, 12, 13]); + assert_eq!( + draft, + Some(skippy_protocol::binary::StageNativeMtpDraft { + token_ids: vec![14, 15], + proposal_compute_us: 123, + }) + ); +} + +#[test] +fn malformed_native_mtp_sideband_is_rejected() { + let message = test_message(WireMessageKind::DecodeEmbd, 1); + let mut predictions = vec![11, 2, 12]; + + let error = split_native_mtp_reply(&message, &mut predictions).unwrap_err(); + + assert!(error.to_string().contains("malformed native MTP sideband")); +} + #[test] fn request_summary_tracks_auto_align_totals() { let config = prefix_cache_test_config(); let mut summary = super::BinaryRequestSummary::default(); - let verify = test_message(WireMessageKind::VerifySpan, 2); + let verify = test_message(WireMessageKind::VerifyWindow, 2); let decode = test_message(WireMessageKind::DecodeEmbd, 1); let mut verify_observation = summary_observation(&config, &verify, 12.5); @@ -163,9 +194,9 @@ fn request_summary_tracks_auto_align_totals() { assert_eq!(summary.session_auto_align_count, 2); assert_eq!(summary.session_auto_align_ms, 2.0); assert_eq!(summary.session_auto_align_trimmed_tokens, 3); - assert_eq!(summary.verify_span_session_auto_align_count, 1); - assert_eq!(summary.verify_span_session_auto_align_ms, 0.75); - assert_eq!(summary.verify_span_session_auto_align_trimmed_tokens, 1); + assert_eq!(summary.verify_window_session_auto_align_count, 1); + assert_eq!(summary.verify_window_session_auto_align_ms, 0.75); + assert_eq!(summary.verify_window_session_auto_align_trimmed_tokens, 1); } #[test] @@ -342,10 +373,10 @@ fn summary_observation<'a>( session_auto_align_count: 0, session_auto_align_ms: 0.0, session_auto_align_trimmed_tokens: 0, - verify_span_pre_compute_ms: 0.25, - verify_span_post_compute_ms: 0.5, - verify_span_pre_reply_ms: 0.0, - verify_span_after_reply_ms: 0.0, + verify_window_pre_compute_ms: 0.25, + verify_window_post_compute_ms: 0.5, + verify_window_pre_reply_ms: 0.0, + verify_window_after_reply_ms: 0.0, upstream_message_wait_ms: 0.0, } } diff --git a/crates/skippy-server/src/cli.rs b/crates/skippy-server/src/cli.rs index 0d4182d447..f69cdfde91 100644 --- a/crates/skippy-server/src/cli.rs +++ b/crates/skippy-server/src/cli.rs @@ -135,6 +135,11 @@ pub struct ServeBinaryArgs { pub openai_ngram_min: usize, #[arg(long, default_value_t = 0)] pub openai_ngram_max: usize, + #[arg( + long, + help = "JSON file containing a complete resolved speculative decode plan. This replaces the legacy --openai-ngram-* tuning flags." + )] + pub openai_speculative_config: Option, } #[derive(Parser)] @@ -150,6 +155,11 @@ pub struct ServeOpenAiArgs { help = "Served model id to advertise and accept, for example org/repo:Q4_K_M. Defaults to config model_id." )] pub model_id: Option, + #[arg( + long, + help = "JSON file containing a complete resolved speculative decode plan." + )] + pub speculative_config: Option, #[arg(long, default_value_t = 16)] pub default_max_tokens: u32, #[arg( @@ -262,4 +272,43 @@ mod tests { }; assert_eq!(args.openai_guardrails, OpenAiGuardrailsCliMode::Enforce); } + + #[test] + fn standalone_commands_accept_resolved_speculative_config_files() { + let cli = Cli::try_parse_from([ + "skippy-server", + "serve-binary", + "--config", + "stage.json", + "--activation-width", + "2048", + "--openai-speculative-config", + "decode-plan.json", + ]) + .unwrap(); + let Command::ServeBinary(args) = cli.command else { + panic!("expected serve-binary command"); + }; + assert_eq!( + args.openai_speculative_config, + Some(PathBuf::from("decode-plan.json")) + ); + + let cli = Cli::try_parse_from([ + "skippy-server", + "serve-openai", + "--config", + "stage.json", + "--speculative-config", + "decode-plan.json", + ]) + .unwrap(); + let Command::ServeOpenAi(args) = cli.command else { + panic!("expected serve-openai command"); + }; + assert_eq!( + args.speculative_config, + Some(PathBuf::from("decode-plan.json")) + ); + } } diff --git a/crates/skippy-server/src/frontend.rs b/crates/skippy-server/src/frontend.rs index d6f3df5e7c..482aedca05 100644 --- a/crates/skippy-server/src/frontend.rs +++ b/crates/skippy-server/src/frontend.rs @@ -32,7 +32,7 @@ use openai_frontend::{ RetryExhaustionMode, StreamingGuardrailMode, Usage, apply_chat_hook_outcome, chat_mesh_hooks_enabled, }; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use skippy_metrics::attr as attr_key; @@ -73,6 +73,7 @@ use crate::{ mod admission; mod backend; mod decode_batcher; +mod decode_scheduler; mod embedded_execution; mod embedded_generation; mod generation_flow; @@ -93,6 +94,7 @@ mod wire_messages; use self::{ admission::{GenerationTokenBudget, GenerationTokenBudgetRequest}, decode_batcher::DecodeBatcher, + decode_scheduler::*, native_mtp::*, prefill::*, request::*, @@ -146,6 +148,7 @@ pub async fn serve_openai(args: ServeOpenAiArgs) -> Result<()> { if args.generation_concurrency == 0 { bail!("--generation-concurrency must be greater than zero"); } + let speculative = load_standalone_speculative_config(args.speculative_config.as_ref())?; let runtime = load_runtime(&config)?.ok_or_else(|| { anyhow!("serve-openai requires a stage config with model_path for tokenization and decode") @@ -199,11 +202,9 @@ pub async fn serve_openai(args: ServeOpenAiArgs) -> Result<()> { draft: None, speculative_window: 0, adaptive_speculative_window: false, - ngram_min: 0, - ngram_max: 0, - native_mtp_enabled: false, - native_mtp_max_tokens: 1, - native_mtp_min_tokens: 0, + ngram_min: standalone_simple_ngram_min(&speculative), + ngram_max: standalone_simple_ngram_max(&speculative), + speculative, generation_limit: Arc::new(Semaphore::new(args.generation_concurrency)), generation_queue_depth: Arc::new(AtomicUsize::new(0)), generation_queue_limit: args.generation_concurrency, @@ -227,6 +228,244 @@ pub async fn serve_openai(args: ServeOpenAiArgs) -> Result<()> { Ok(()) } +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct SpeculativeDecodeConfig { + pub requested_strategy: String, + pub effective_strategy: String, + pub native_mtp: NativeMtpProposalConfig, + pub ngram: Option, + pub extension: Option, + pub verify_window: VerifyWindowConfig, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct NativeMtpProposalConfig { + pub enabled: bool, + pub max_draft_tokens: usize, + pub min_draft_tokens: usize, + pub reject_cooldown_tokens: usize, + pub suppress_cooldown_drafts: bool, + pub suppress_cooldown_draft_limit: usize, +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum NgramProposerKind { + Simple, + Cache, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct NgramProposalConfig { + pub kind: NgramProposerKind, + pub min_ngram: usize, + pub max_ngram: usize, + pub max_proposal_tokens: usize, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct NgramExtensionConfig { + pub initial_tokens: usize, + pub max_tokens: usize, + pub tail_backoff_proposals: usize, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct VerifyWindowConfig { + pub min_tokens: usize, + pub max_tokens: usize, + pub pipeline_depth: usize, +} + +impl Default for SpeculativeDecodeConfig { + fn default() -> Self { + Self { + requested_strategy: "auto".to_string(), + effective_strategy: "disabled".to_string(), + native_mtp: NativeMtpProposalConfig { + enabled: false, + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }, + ngram: None, + extension: None, + verify_window: VerifyWindowConfig { + min_tokens: 1, + max_tokens: 4, + pipeline_depth: 1, + }, + } + } +} + +impl SpeculativeDecodeConfig { + pub fn validate(&self) -> Result<()> { + if self.requested_strategy.trim().is_empty() || self.effective_strategy.trim().is_empty() { + bail!("speculative decode strategies must not be empty"); + } + if self.native_mtp.min_draft_tokens > self.native_mtp.max_draft_tokens { + bail!("native MTP min_draft_tokens must not exceed max_draft_tokens"); + } + if let Some(ngram) = &self.ngram + && (ngram.min_ngram == 0 + || ngram.min_ngram > ngram.max_ngram + || ngram.max_proposal_tokens < ngram.min_ngram) + { + bail!( + "N-gram proposer requires 0 < min_ngram <= max_ngram and max_proposal_tokens >= min_ngram" + ); + } + if let Some(ngram) = &self.ngram + && ngram.kind == NgramProposerKind::Cache + && ngram.max_ngram > skippy_runtime::NGRAM_CACHE_MAX_NGRAM + { + bail!( + "cache N-gram proposer max_ngram must not exceed llama.cpp limit {}", + skippy_runtime::NGRAM_CACHE_MAX_NGRAM + ); + } + if self.extension.is_some() && (!self.native_mtp.enabled || self.ngram.is_none()) { + bail!("N-gram extension requires both native MTP and an N-gram proposer"); + } + if let Some(extension) = &self.extension + && (extension.initial_tokens == 0 + || extension.initial_tokens > extension.max_tokens + || extension.max_tokens == 0) + { + bail!("N-gram extension requires 0 < initial_tokens <= max_tokens"); + } + if self.verify_window.min_tokens == 0 + || self.verify_window.min_tokens > self.verify_window.max_tokens + || self.verify_window.pipeline_depth == 0 + { + bail!("verify window requires 0 < min_tokens <= max_tokens and pipeline_depth > 0"); + } + Ok(()) + } + + fn insert_telemetry_attrs(&self, attrs: &mut BTreeMap) { + attrs.insert( + "llama_stage.spec.requested_strategy".to_string(), + json!(self.requested_strategy), + ); + attrs.insert( + "llama_stage.spec.effective_strategy".to_string(), + json!(self.effective_strategy), + ); + } +} + +fn load_standalone_speculative_config(path: Option<&PathBuf>) -> Result { + let config = match path { + Some(path) => load_json(path) + .with_context(|| format!("load speculative decode config {}", path.display()))?, + None => SpeculativeDecodeConfig::default(), + }; + config.validate()?; + Ok(config) +} + +fn standalone_simple_ngram_min(config: &SpeculativeDecodeConfig) -> usize { + config + .ngram + .as_ref() + .filter(|ngram| ngram.kind == NgramProposerKind::Simple) + .map_or(0, |ngram| ngram.min_ngram) +} + +fn standalone_simple_ngram_max(config: &SpeculativeDecodeConfig) -> usize { + config + .ngram + .as_ref() + .filter(|ngram| ngram.kind == NgramProposerKind::Simple) + .map_or(0, |ngram| ngram.max_proposal_tokens.min(ngram.max_ngram)) +} + +#[cfg(test)] +mod standalone_speculative_config_tests { + use super::*; + + #[test] + fn standalone_speculative_config_rejects_invalid_composite_plan() { + let config = SpeculativeDecodeConfig { + extension: Some(NgramExtensionConfig { + initial_tokens: 2, + max_tokens: 4, + tail_backoff_proposals: 1, + }), + ..SpeculativeDecodeConfig::default() + }; + + let error = config.validate().expect_err("extension requires proposers"); + + assert!( + error + .to_string() + .contains("requires both native MTP and an N-gram proposer") + ); + } + + #[test] + fn standalone_speculative_config_round_trips_cache_composite() { + let config = SpeculativeDecodeConfig { + requested_strategy: "mtp-cache".to_string(), + effective_strategy: "native-mtp-cache".to_string(), + native_mtp: NativeMtpProposalConfig { + enabled: true, + max_draft_tokens: 2, + ..SpeculativeDecodeConfig::default().native_mtp + }, + ngram: Some(NgramProposalConfig { + kind: NgramProposerKind::Cache, + min_ngram: 2, + max_ngram: 4, + max_proposal_tokens: 6, + }), + extension: Some(NgramExtensionConfig { + initial_tokens: 2, + max_tokens: 6, + tail_backoff_proposals: 2, + }), + ..SpeculativeDecodeConfig::default() + }; + + let json = serde_json::to_string(&config).expect("serialize plan"); + let decoded: SpeculativeDecodeConfig = serde_json::from_str(&json).expect("parse plan"); + + assert_eq!(decoded, config); + decoded.validate().expect("valid composite plan"); + } + + #[test] + fn standalone_speculative_config_rejects_cache_windows_above_llama_limit() { + let config = SpeculativeDecodeConfig { + ngram: Some(NgramProposalConfig { + kind: NgramProposerKind::Cache, + min_ngram: 2, + max_ngram: skippy_runtime::NGRAM_CACHE_MAX_NGRAM + 1, + max_proposal_tokens: 6, + }), + ..SpeculativeDecodeConfig::default() + }; + + let error = config.validate().expect_err("cache max must be bounded"); + + assert!( + error + .to_string() + .contains("must not exceed llama.cpp limit 4") + ); + } +} + #[derive(Clone)] pub struct EmbeddedOpenAiArgs { pub bind_addr: SocketAddr, @@ -246,6 +485,7 @@ pub struct EmbeddedOpenAiArgs { pub speculative_window: usize, pub adaptive_speculative_window: bool, pub draft_n_gpu_layers: Option, + pub speculative: SpeculativeDecodeConfig, pub ngram_min: usize, pub ngram_max: usize, pub native_mtp_enabled: bool, @@ -577,9 +817,6 @@ pub fn embedded_openai_backend(args: EmbeddedOpenAiArgs) -> Result Result, generation_queue_depth: Arc, generation_queue_limit: usize, @@ -904,9 +1137,6 @@ enum OpenAiBackendMode { prefill_reply_credit_limit: usize, lane_pool: Option>, prediction_returns: Option>, - native_mtp_enabled: bool, - native_mtp_max_tokens: usize, - native_mtp_min_tokens: usize, }, } @@ -947,6 +1177,16 @@ struct PrefillTransportEstimate { /// lane's generation reads stay blocking. const LANE_READY_READ_TIMEOUT: Duration = Duration::from_secs(20); +/// Steady-state (post-warmup) lane reconnect deadline. A live split peer at +/// LAN/WAN RTT answers the ready handshake in milliseconds; only a dead or +/// wedged downstream stage waits this long. Keeping the steady-state deadline +/// short turns "a new request routed to a dead stage" into a fast error (~3s) +/// instead of the ~20s warmup deadline. The long `LANE_READY_READ_TIMEOUT` is +/// only needed during pool warmup, when the downstream stage may still be +/// loading its model; a mid-life reconnect to an already-serving mesh has no +/// such excuse for a multi-second silence. +const LANE_STEADY_CONNECT_TIMEOUT: Duration = Duration::from_secs(3); + impl PersistentStageLanePool { const PREFILL_TRANSPORT_EWMA_ALPHA: f64 = 0.25; @@ -970,7 +1210,9 @@ impl PersistentStageLanePool { }); let timer = PhaseTimer::start(); for _ in 0..capacity { - let lane = pool.connect_lane()?; + // Warmup: the downstream stage may still be loading its model, so + // allow the full ready deadline. + let lane = pool.connect_lane(LANE_READY_READ_TIMEOUT)?; pool.return_lane(lane); } let mut attrs = lifecycle_attrs(config); @@ -1000,9 +1242,21 @@ impl PersistentStageLanePool { .map_err(|_| OpenAiError::backend("persistent lane pool lock poisoned"))?; lanes.pop() }; - let lane = match lane { + // A pooled lane may be stale: if the downstream stage died while the + // lane was checked in, the cached TCP stream is dead but reusing it + // would block forever on the next generation read (the handshake + // read-timeout is deliberately cleared so long generations don't + // truncate). Probe liveness before reuse and drop a dead lane so we + // reconnect with the short steady-state deadline instead of hanging. + let live_pooled = lane.filter(|lane| lane_stream_is_live(&lane.stream)); + let lane = match live_pooled { Some(lane) => lane, - None => self.connect_lane().map_err(openai_backend_error)?, + // Steady-state reconnect: the mesh is already serving, so a healthy + // downstream answers fast. Use the short deadline so a request + // routed to a dead stage fails quickly instead of stalling ~20s. + None => self + .connect_lane(LANE_STEADY_CONNECT_TIMEOUT) + .map_err(openai_backend_error)?, }; let mut attrs = BTreeMap::from([ ( @@ -1111,7 +1365,10 @@ impl PersistentStageLanePool { "llama_stage.openai_downstream_retired_lane_id".to_string(), json!(retired_lane_id), ); - match self.connect_lane() { + // Replacing a retired lane on an already-serving mesh is a steady-state + // reconnect: a healthy downstream answers fast, a dead one should fail + // fast rather than stall. + match self.connect_lane(LANE_STEADY_CONNECT_TIMEOUT) { Ok(lane) => { attrs.insert( "llama_stage.openai_downstream_lane_id".to_string(), @@ -1145,16 +1402,16 @@ impl PersistentStageLanePool { } } - fn connect_lane(&self) -> Result { + fn connect_lane(&self, ready_timeout: Duration) -> Result { let lane_id = self.next_lane_id.fetch_add(1, Ordering::Relaxed); let timer = PhaseTimer::start(); // Single bounded attempt. `connect_binary_downstream` already retries // the TCP connect internally (its own attempt budget), and the ready - // handshake read below is bounded by `LANE_READY_READ_TIMEOUT`. An extra + // handshake read below is bounded by `ready_timeout`. An extra // outer retry here only multiplies the worst-case stall — a dead peer // would burn (connect budget × outer attempts) before failing (see PR // #1011 review). Bring-up wants a single, predictable deadline. - let stream = self.connect_lane_once(lane_id).inspect_err(|error| { + let stream = self.connect_lane_once(lane_id, ready_timeout).inspect_err(|error| { eprintln!( "openai downstream lane handshake failed: stage_id={} lane_id={lane_id}: {error:#}", self.config.stage_id, @@ -1187,10 +1444,12 @@ impl PersistentStageLanePool { /// Perform a single connect + ready-handshake attempt for one lane. /// - /// Returns the connected, ready `TcpStream` on success. Callers retry this - /// with a backoff so a transient short read on `recv_ready` does not fail - /// lane creation outright. - fn connect_lane_once(&self, lane_id: u64) -> Result { + /// Returns the connected, ready `TcpStream` on success. `ready_timeout` + /// bounds the handshake read: use `LANE_READY_READ_TIMEOUT` during pool + /// warmup (the downstream stage may still be loading) and the shorter + /// `LANE_STEADY_CONNECT_TIMEOUT` for mid-life reconnects on an + /// already-serving mesh, so a dead stage fails fast instead of stalling. + fn connect_lane_once(&self, lane_id: u64, ready_timeout: Duration) -> Result { let mut stream = connect_binary_downstream(&self.config, self.timeout_secs)? .ok_or_else(|| anyhow!("embedded stage0 has no downstream"))?; let local_addr = stream.local_addr().ok(); @@ -1201,24 +1460,7 @@ impl PersistentStageLanePool { ); send_client_ready_hello_if_enabled(&mut stream) .context("send persistent downstream lane client ready hello")?; - // Bound the ready handshake read. `recv_ready` is a blocking - // `read_exact`; without a timeout a stalled or half-dead downstream - // connection hangs lane creation forever (observed as a lane stuck in - // "waiting ready" that never completes and never errors). A single - // bounded read timeout turns the stall into an error. Both the set and - // the clear are propagated: if the set fails, `recv_ready` would be - // unbounded (defeating the fix); if the clear fails, the handshake - // timeout would leak into the persistent lane's later generation reads - // and truncate long generations. - stream - .set_read_timeout(Some(LANE_READY_READ_TIMEOUT)) - .context("set persistent downstream lane ready read timeout")?; - let ready = - recv_ready(&mut stream).context("persistent downstream lane did not become ready"); - stream - .set_read_timeout(None) - .context("clear persistent downstream lane ready read timeout")?; - ready?; + receive_persistent_lane_ready(&mut stream, ready_timeout)?; eprintln!( "openai downstream lane received ready: stage_id={} lane_id={lane_id} local={local_addr:?} peer={peer_addr:?}", self.config.stage_id @@ -1227,6 +1469,59 @@ impl PersistentStageLanePool { } } +/// Read the lane ready handshake with a bounded deadline. +/// +/// `recv_ready` is a blocking `read_exact`; without a read timeout a stalled or +/// half-dead downstream connection hangs lane creation forever (observed as a +/// lane stuck in "waiting ready" that never completes and never errors). Set the +/// deadline, read, then clear it so the persistent lane's later generation reads +/// stay blocking. Both the set and the clear are propagated: if the set fails, +/// `recv_ready` would be unbounded (defeating the fix); if the clear fails, the +/// handshake timeout would leak into generation reads and truncate long outputs. +fn receive_persistent_lane_ready(stream: &mut TcpStream, timeout: Duration) -> Result<()> { + stream + .set_read_timeout(Some(timeout)) + .context("set persistent downstream lane ready timeout")?; + let ready = recv_ready(&mut *stream).context("persistent downstream lane did not become ready"); + stream + .set_read_timeout(None) + .context("restore persistent downstream lane read timeout")?; + ready +} + +/// Cheap liveness probe for a pooled lane before reuse. +/// +/// A pooled lane whose downstream stage died is a dead TCP stream; reusing it +/// would block the next generation read forever (the handshake read-timeout is +/// cleared for pooled lanes so long generations don't truncate). A nonblocking +/// peek distinguishes cases without consuming data: +/// +/// - `Ok(0)` => peer sent EOF / closed => dead, discard. +/// - `Err(WouldBlock)` => connection open, no pending data => healthy, reuse. +/// - other `Err` (reset, etc.) => dead, discard. +/// +/// Between requests a healthy lane has no unread bytes, so a nonzero peek would +/// be unexpected protocol data; treat that as unsafe-to-reuse and discard too. +/// The blocking mode is always restored so the caller's reads are unaffected. +fn lane_stream_is_live(stream: &TcpStream) -> bool { + use std::io::ErrorKind; + if stream.set_nonblocking(true).is_err() { + return false; + } + let mut probe = [0u8; 1]; + let live = match stream.peek(&mut probe) { + Ok(0) => false, + Ok(_) => false, + Err(ref e) if e.kind() == ErrorKind::WouldBlock => true, + Err(_) => false, + }; + // Restore blocking mode; if we can't, the lane is not safe to hand back. + if stream.set_nonblocking(false).is_err() { + return false; + } + live +} + fn ewma(old: f64, sample: f64) -> f64 { old.mul_add( 1.0 - PersistentStageLanePool::PREFILL_TRANSPORT_EWMA_ALPHA, @@ -1314,7 +1609,6 @@ fn prompt_cache_retention_label(retention: openai_frontend::PromptCacheRetention } } -#[derive(Clone, Copy)] struct GenerationCacheStats { status: &'static str, cached_prompt_tokens: u32, @@ -1322,6 +1616,11 @@ struct GenerationCacheStats { suffix_prefill_tokens: u32, hit_kind: Option<&'static str>, native_mtp_stats: NativeMtpStats, + native_mtp_decode_telemetry: Option, + verify_window_pipeline_stats: Option, + speculative_stats: Option, + prompt_ms: f64, + predicted_ms: f64, } impl Default for GenerationCacheStats { @@ -1333,6 +1632,11 @@ impl Default for GenerationCacheStats { suffix_prefill_tokens: 0, hit_kind: None, native_mtp_stats: NativeMtpStats::default(), + native_mtp_decode_telemetry: None, + verify_window_pipeline_stats: None, + speculative_stats: None, + prompt_ms: 0.0, + predicted_ms: 0.0, } } } @@ -1763,6 +2067,19 @@ fn chat_response_from_generated_text( .with_timings(output.timings()) } +fn completion_response_from_generated_text( + model: String, + output: &GeneratedText, +) -> CompletionResponse { + CompletionResponse::new_with_reason( + model, + output.text.clone(), + output.usage(), + output.finish_reason, + ) + .with_timings(output.timings()) +} + fn parsed_chat_message_from_json( message_json: &str, request: &ChatCompletionRequest, @@ -1932,9 +2249,8 @@ struct LocalGeneration<'a> { max_tokens: u32, sampling: &'a SamplingConfig, chat_sampling_metadata: Option<&'a str>, + speculative: &'a SpeculativeDecodeConfig, native_mtp_enabled: bool, - native_mtp_max_tokens: usize, - native_mtp_min_tokens: usize, hook_request: Option, hook_runtime: Option, cancellation: Option<&'a openai_frontend::CancellationToken>, @@ -1953,11 +2269,11 @@ struct EmbeddedStageZeroGeneration<'a> { draft: Option>>, speculative_window: usize, adaptive_speculative_window: bool, + speculative: &'a SpeculativeDecodeConfig, ngram_min: usize, ngram_max: usize, native_mtp_enabled: bool, native_mtp_max_tokens: usize, - native_mtp_min_tokens: usize, prompt_token_ids: &'a [i32], max_tokens: u32, sampling: &'a SamplingConfig, @@ -2363,6 +2679,11 @@ where suffix_prefill_tokens: cache_stats.suffix_prefill_tokens, cache_hit_kind: cache_stats.hit_kind, native_mtp_stats: cache_stats.native_mtp_stats, + native_mtp_decode_telemetry: cache_stats.native_mtp_decode_telemetry, + verify_window_pipeline_stats: cache_stats.verify_window_pipeline_stats, + speculative_stats: cache_stats.speculative_stats, + prompt_ms: cache_stats.prompt_ms, + predicted_ms: cache_stats.predicted_ms, text: self.text, finish_reason: self.finish_reason, detokenize_ms: self.metrics.detokenize_ms, @@ -2381,6 +2702,11 @@ struct GeneratedText { suffix_prefill_tokens: u32, cache_hit_kind: Option<&'static str>, native_mtp_stats: NativeMtpStats, + native_mtp_decode_telemetry: Option, + verify_window_pipeline_stats: Option, + speculative_stats: Option, + prompt_ms: f64, + predicted_ms: f64, text: String, finish_reason: FinishReason, detokenize_ms: f64, @@ -2395,14 +2721,26 @@ impl GeneratedText { } fn timings(&self) -> Option> { - if !self.native_mtp_stats.enabled() { - return None; - } - let stats = self.native_mtp_stats; - Some(BTreeMap::from([ - ("draft_n".to_string(), json!(stats.drafted_tokens)), - ("draft_n_accepted".to_string(), json!(stats.accepted_tokens)), + let (drafted_tokens, accepted_tokens) = self + .native_mtp_decode_telemetry + .and_then(NativeMtpDecodeTelemetry::composite_proposal_totals) + .unwrap_or((stats.drafted_tokens, stats.accepted_tokens)); + let mut timings = BTreeMap::from([ + ("prompt_n".to_string(), json!(self.prompt_tokens)), + ("prompt_ms".to_string(), json!(self.prompt_ms)), + ( + "prompt_per_second".to_string(), + json!(tokens_per_second(self.prompt_tokens, self.prompt_ms)), + ), + ("predicted_n".to_string(), json!(self.completion_tokens)), + ("predicted_ms".to_string(), json!(self.predicted_ms)), + ( + "predicted_per_second".to_string(), + json!(tokens_per_second(self.completion_tokens, self.predicted_ms)), + ), + ("draft_n".to_string(), json!(drafted_tokens)), + ("draft_n_accepted".to_string(), json!(accepted_tokens)), ( "native_mtp_rejected".to_string(), json!(stats.rejected_tokens), @@ -2419,7 +2757,25 @@ impl GeneratedText { "native_mtp_verification_compute_us".to_string(), json!(stats.verification_compute_us), ), - ])) + ]); + if let Some(telemetry) = self.native_mtp_decode_telemetry { + telemetry.insert_response_timings(&mut timings); + } + if let Some(stats) = self.verify_window_pipeline_stats { + stats.insert_response_timings(&mut timings); + } + if let Some(stats) = self.speculative_stats.as_ref() { + stats.insert_response_timings(&mut timings); + } + Some(timings) + } +} + +fn tokens_per_second(token_count: u32, elapsed_ms: f64) -> f64 { + if elapsed_ms > 0.0 { + f64::from(token_count) * 1_000.0 / elapsed_ms + } else { + 0.0 } } diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index 9d2aa69db1..fe6dee836b 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -195,12 +195,7 @@ impl OpenAiBackend for StageOpenAiBackend { ) .await?; let response_timer = PhaseTimer::start(); - let response = CompletionResponse::new_with_reason( - request.model, - output.text.clone(), - output.usage(), - output.finish_reason, - ); + let response = completion_response_from_generated_text(request.model, &output); let mut response_attrs = self.openai_attrs(&ids); response_attrs.insert( "llama_stage.openai_operation".to_string(), diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs new file mode 100644 index 0000000000..e3a9567d69 --- /dev/null +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -0,0 +1,475 @@ +use std::collections::{BTreeMap, VecDeque}; + +use super::{OpenAiError, OpenAiResult}; + +const PIPELINE_PROFILE_MIN_OBSERVATIONS: usize = 8; +const PIPELINE_PROFILE_MAX_OBSERVATIONS: usize = 32; +const PIPELINE_PROFIT_MARGIN: f64 = 1.15; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct VerifyWindowPipelineConfig { + depth: usize, +} + +impl VerifyWindowPipelineConfig { + pub(super) fn new(depth: usize) -> Self { + Self { + depth: depth.max(1), + } + } + + pub(super) fn depth(self) -> usize { + self.depth + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub(super) struct VerifyWindowPipelineStats { + depth: usize, + direct_prediction_return: bool, + opened_windows: usize, + max_in_flight: usize, + stale_discarded: usize, + stale_drain_ms: f64, + policy_observed_windows: usize, + policy_continuation_windows: usize, + policy_permit_checks: usize, + policy_permits: usize, + policy_suppressed: usize, +} + +impl VerifyWindowPipelineStats { + pub(super) fn insert_response_timings(self, timings: &mut BTreeMap) { + timings.insert( + "verify_window_depth".to_string(), + serde_json::json!(self.depth), + ); + timings.insert( + "verify_window_direct_prediction_return".to_string(), + serde_json::json!(self.direct_prediction_return), + ); + timings.insert( + "verify_window_opened".to_string(), + serde_json::json!(self.opened_windows), + ); + timings.insert( + "verify_window_max_in_flight".to_string(), + serde_json::json!(self.max_in_flight), + ); + timings.insert( + "verify_window_stale_discarded".to_string(), + serde_json::json!(self.stale_discarded), + ); + timings.insert( + "verify_window_stale_drain_ms".to_string(), + serde_json::json!(self.stale_drain_ms), + ); + timings.insert( + "verify_window_policy_observed_windows".to_string(), + serde_json::json!(self.policy_observed_windows), + ); + timings.insert( + "verify_window_policy_continuation_windows".to_string(), + serde_json::json!(self.policy_continuation_windows), + ); + timings.insert( + "verify_window_policy_permit_checks".to_string(), + serde_json::json!(self.policy_permit_checks), + ); + timings.insert( + "verify_window_policy_permits".to_string(), + serde_json::json!(self.policy_permits), + ); + timings.insert( + "verify_window_policy_suppressed".to_string(), + serde_json::json!(self.policy_suppressed), + ); + } +} + +#[derive(Debug, Default)] +struct VerifyWindowWidthProfile { + observations: VecDeque, + continuation_windows: usize, + stage0_compute_ms: f64, + downstream_wait_ms: f64, +} + +#[derive(Debug, Clone, Copy)] +struct VerifyWindowProfileObservation { + continues: bool, + stage0_compute_ms: f64, + downstream_wait_ms: f64, +} + +impl VerifyWindowWidthProfile { + fn observe(&mut self, continues: bool, stage0_compute_ms: f64, downstream_wait_ms: f64) { + let observation = VerifyWindowProfileObservation { + continues, + stage0_compute_ms: stage0_compute_ms.max(0.0), + downstream_wait_ms: downstream_wait_ms.max(0.0), + }; + self.observations.push_back(observation); + self.continuation_windows = self + .continuation_windows + .saturating_add(usize::from(continues)); + self.stage0_compute_ms += observation.stage0_compute_ms; + self.downstream_wait_ms += observation.downstream_wait_ms; + if self.observations.len() > PIPELINE_PROFILE_MAX_OBSERVATIONS { + let expired = self + .observations + .pop_front() + .expect("profile exceeded its non-empty bound"); + self.continuation_windows = self + .continuation_windows + .saturating_sub(usize::from(expired.continues)); + self.stage0_compute_ms -= expired.stage0_compute_ms; + self.downstream_wait_ms -= expired.downstream_wait_ms; + } + } + + fn is_profitable(&self) -> bool { + if self.observations.len() < PIPELINE_PROFILE_MIN_OBSERVATIONS { + return false; + } + let observations = self.observations.len() as f64; + let continuation_rate = self.continuation_windows as f64 / observations; + let average_stage0_ms = self.stage0_compute_ms / observations; + let average_downstream_ms = self.downstream_wait_ms / observations; + let expected_overlap_ms = continuation_rate * average_downstream_ms; + let expected_stale_ms = + (1.0 - continuation_rate) * average_stage0_ms.max(average_downstream_ms); + expected_overlap_ms > expected_stale_ms * PIPELINE_PROFIT_MARGIN + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct VerifyWindow { + pub(super) id: i32, + pub(super) base_position: usize, + pub(super) decode_step: usize, +} + +#[derive(Debug)] +pub(super) struct VerifyWindowScheduler { + config: VerifyWindowPipelineConfig, + next_id: i32, + in_flight: VecDeque, + stats: VerifyWindowPipelineStats, + width_profiles: BTreeMap, +} + +impl VerifyWindowScheduler { + pub(super) fn new(config: VerifyWindowPipelineConfig) -> Self { + Self { + config, + next_id: 1, + in_flight: VecDeque::new(), + stats: VerifyWindowPipelineStats { + depth: config.depth(), + ..VerifyWindowPipelineStats::default() + }, + width_profiles: BTreeMap::new(), + } + } + + pub(super) fn has_capacity(&self) -> bool { + self.in_flight.len() < self.config.depth + } + + pub(super) fn depth(&self) -> usize { + self.config.depth() + } + + pub(super) fn mark_direct_prediction_return(&mut self) { + self.stats.direct_prediction_return = true; + } + + pub(super) fn observe_pipeline_profile( + &mut self, + width: usize, + continues: bool, + stage0_compute_ms: f64, + downstream_wait_ms: f64, + ) { + if width == 0 { + return; + } + self.stats.policy_observed_windows = self.stats.policy_observed_windows.saturating_add(1); + self.stats.policy_continuation_windows = self + .stats + .policy_continuation_windows + .saturating_add(usize::from(continues)); + self.width_profiles.entry(width).or_default().observe( + continues, + stage0_compute_ms, + downstream_wait_ms, + ); + } + + pub(super) fn has_profitable_pipeline_width(&self) -> bool { + self.config.depth() > 1 + && self + .width_profiles + .values() + .any(VerifyWindowWidthProfile::is_profitable) + } + + pub(super) fn permit_pipeline_width(&mut self, width: usize) -> bool { + self.stats.policy_permit_checks = self.stats.policy_permit_checks.saturating_add(1); + let permitted = self.config.depth() > 1 + && self + .width_profiles + .get(&width) + .is_some_and(VerifyWindowWidthProfile::is_profitable); + if permitted { + self.stats.policy_permits = self.stats.policy_permits.saturating_add(1); + } else { + self.stats.policy_suppressed = self.stats.policy_suppressed.saturating_add(1); + } + permitted + } + + pub(super) fn insert_policy_telemetry_attrs( + &self, + attrs: &mut BTreeMap, + ) { + attrs.insert( + "llama_stage.verify_window.pipeline_policy_observed_windows".to_string(), + serde_json::json!(self.stats.policy_observed_windows), + ); + attrs.insert( + "llama_stage.verify_window.pipeline_policy_continuation_windows".to_string(), + serde_json::json!(self.stats.policy_continuation_windows), + ); + attrs.insert( + "llama_stage.verify_window.pipeline_policy_permit_checks".to_string(), + serde_json::json!(self.stats.policy_permit_checks), + ); + attrs.insert( + "llama_stage.verify_window.pipeline_policy_permits".to_string(), + serde_json::json!(self.stats.policy_permits), + ); + attrs.insert( + "llama_stage.verify_window.pipeline_policy_suppressed".to_string(), + serde_json::json!(self.stats.policy_suppressed), + ); + attrs.insert( + "llama_stage.verify_window.pipeline_policy_profitable_widths".to_string(), + serde_json::json!( + self.width_profiles + .values() + .filter(|profile| profile.is_profitable()) + .count() + ), + ); + } + + pub(super) fn open( + &mut self, + base_position: usize, + decode_step: usize, + ) -> OpenAiResult { + if !self.has_capacity() { + return Err(OpenAiError::backend( + "verify window pipeline depth exceeded", + )); + } + let id = self.next_id; + self.next_id = self + .next_id + .checked_add(1) + .ok_or_else(|| OpenAiError::backend("verify window id overflow"))?; + let window = VerifyWindow { + id, + base_position, + decode_step, + }; + self.in_flight.push_back(window.clone()); + self.stats.opened_windows = self.stats.opened_windows.saturating_add(1); + self.stats.max_in_flight = self.stats.max_in_flight.max(self.in_flight.len()); + Ok(window) + } + + pub(super) fn complete_next(&mut self, reply_window_id: i32) -> OpenAiResult { + let Some(window) = self.in_flight.front() else { + return Err(OpenAiError::backend( + "verify window reply arrived with no in-flight window", + )); + }; + if window.id != reply_window_id { + return Err(OpenAiError::backend(format!( + "verify window reply out of order: got {reply_window_id}, expected {}", + window.id + ))); + } + Ok(self.in_flight.pop_front().expect("checked non-empty queue")) + } + + #[cfg(test)] + pub(super) fn discard_stale(&mut self) -> usize { + let discarded = self.in_flight.len(); + self.in_flight.clear(); + self.stats.stale_discarded = self.stats.stale_discarded.saturating_add(discarded); + discarded + } + + pub(super) fn record_stale_discarded(&mut self, count: usize, drain_ms: f64) { + self.stats.stale_discarded = self.stats.stale_discarded.saturating_add(count); + self.stats.stale_drain_ms += drain_ms; + } + + pub(super) fn in_flight_len(&self) -> usize { + self.in_flight.len() + } + + pub(super) fn stale_discard_count(&self) -> usize { + self.stats.stale_discarded + } + + pub(super) fn stats(&self) -> VerifyWindowPipelineStats { + self.stats + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bounds_depth_and_requires_fifo_reply_ids() { + let config = VerifyWindowPipelineConfig { depth: 2 }; + let mut scheduler = VerifyWindowScheduler::new(config); + let first = scheduler.open(10, 0).unwrap(); + let second = scheduler.open(11, 1).unwrap(); + + assert!(scheduler.open(12, 2).is_err()); + assert!(scheduler.complete_next(second.id).is_err()); + assert_eq!(scheduler.in_flight_len(), 2); + assert_eq!(scheduler.complete_next(first.id).unwrap(), first); + assert_eq!(scheduler.complete_next(second.id).unwrap(), second); + assert_eq!(first.id, 1); + assert_eq!(scheduler.stats().depth, 2); + assert_eq!(scheduler.stats().opened_windows, 2); + assert_eq!(scheduler.stats().max_in_flight, 2); + assert!(!scheduler.stats().direct_prediction_return); + } + + #[test] + fn discards_stale_windows_after_divergence() { + let config = VerifyWindowPipelineConfig { depth: 3 }; + let mut scheduler = VerifyWindowScheduler::new(config); + scheduler.open(10, 0).unwrap(); + scheduler.open(11, 1).unwrap(); + scheduler.open(12, 2).unwrap(); + + assert_eq!(scheduler.discard_stale(), 3); + assert_eq!(scheduler.stale_discard_count(), 3); + assert_eq!(scheduler.in_flight_len(), 0); + assert_eq!(scheduler.stats().stale_discarded, 3); + } + + #[test] + fn pipeline_policy_waits_for_enough_width_specific_evidence() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + for _ in 0..PIPELINE_PROFILE_MIN_OBSERVATIONS - 1 { + scheduler.observe_pipeline_profile(2, true, 20.0, 80.0); + } + assert!(!scheduler.has_profitable_pipeline_width()); + assert!(!scheduler.permit_pipeline_width(2)); + + scheduler.observe_pipeline_profile(2, true, 20.0, 80.0); + assert!(scheduler.has_profitable_pipeline_width()); + assert!(scheduler.permit_pipeline_width(2)); + } + + #[test] + fn pipeline_policy_suppresses_low_acceptance_local_work() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + for index in 0..PIPELINE_PROFILE_MIN_OBSERVATIONS { + scheduler.observe_pipeline_profile(2, index < 2, 31.0, 24.0); + } + + assert!(!scheduler.has_profitable_pipeline_width()); + assert!(!scheduler.permit_pipeline_width(2)); + assert_eq!(scheduler.stats().policy_permit_checks, 1); + assert_eq!(scheduler.stats().policy_permits, 0); + assert_eq!(scheduler.stats().policy_suppressed, 1); + } + + #[test] + fn pipeline_policy_profiles_each_verify_width_independently() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + for index in 0..PIPELINE_PROFILE_MIN_OBSERVATIONS { + scheduler.observe_pipeline_profile(1, index < 7, 20.0, 80.0); + scheduler.observe_pipeline_profile(2, index < 2, 31.0, 24.0); + } + + assert!(scheduler.has_profitable_pipeline_width()); + assert!(scheduler.permit_pipeline_width(1)); + assert!(!scheduler.permit_pipeline_width(2)); + } + + #[test] + fn pipeline_depth_one_never_admits_dependent_work() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 1 }); + for _ in 0..PIPELINE_PROFILE_MIN_OBSERVATIONS { + scheduler.observe_pipeline_profile(2, true, 20.0, 80.0); + } + + assert!(!scheduler.has_profitable_pipeline_width()); + assert!(!scheduler.permit_pipeline_width(2)); + } + + #[test] + fn pipeline_policy_adapts_when_recent_acceptance_changes() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + for _ in 0..PIPELINE_PROFILE_MAX_OBSERVATIONS { + scheduler.observe_pipeline_profile(2, true, 20.0, 80.0); + } + assert!(scheduler.permit_pipeline_width(2)); + + for _ in 0..PIPELINE_PROFILE_MAX_OBSERVATIONS { + scheduler.observe_pipeline_profile(2, false, 20.0, 80.0); + } + assert!(!scheduler.permit_pipeline_width(2)); + } + + #[test] + fn pipeline_policy_counters_are_exposed_in_response_timings() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + for _ in 0..PIPELINE_PROFILE_MIN_OBSERVATIONS { + scheduler.observe_pipeline_profile(2, true, 20.0, 80.0); + } + assert!(scheduler.permit_pipeline_width(2)); + let mut timings = BTreeMap::new(); + scheduler.stats().insert_response_timings(&mut timings); + + assert_eq!( + timings["verify_window_policy_observed_windows"], + serde_json::json!(PIPELINE_PROFILE_MIN_OBSERVATIONS) + ); + assert_eq!( + timings["verify_window_policy_continuation_windows"], + serde_json::json!(PIPELINE_PROFILE_MIN_OBSERVATIONS) + ); + assert_eq!(timings["verify_window_policy_permit_checks"], 1); + assert_eq!(timings["verify_window_policy_permits"], 1); + assert_eq!(timings["verify_window_policy_suppressed"], 0); + + let mut attrs = BTreeMap::new(); + scheduler.insert_policy_telemetry_attrs(&mut attrs); + assert_eq!( + attrs["llama_stage.verify_window.pipeline_policy_observed_windows"], + serde_json::json!(PIPELINE_PROFILE_MIN_OBSERVATIONS) + ); + assert_eq!( + attrs["llama_stage.verify_window.pipeline_policy_continuation_windows"], + serde_json::json!(PIPELINE_PROFILE_MIN_OBSERVATIONS) + ); + assert_eq!( + attrs["llama_stage.verify_window.pipeline_policy_profitable_widths"], + 1 + ); + } +} diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index 3e0bbaf098..2f75e30603 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -3,6 +3,14 @@ use super::*; const DIRECT_RETURN_FALLBACK_POLL: Duration = Duration::from_millis(10); const DIRECT_RETURN_FALLBACK_TIMEOUT: Duration = Duration::from_secs(300); +pub(super) struct DispatchedEmbeddedStage { + started: Instant, + stats: StageReplyStats, + execution: EmbeddedExecutionStats, + message_kind: WireMessageKind, + token_count: i32, +} + impl StageOpenAiBackend { pub(super) fn execute_embedded_stage_message( &self, @@ -13,7 +21,25 @@ impl StageOpenAiBackend { token_ids: &[i32], expected_reply: WireReplyKind, ) -> OpenAiResult { - let timer = PhaseTimer::start(); + let dispatched = self.dispatch_embedded_stage_message( + request, + downstream, + session_key, + message, + token_ids, + )?; + self.complete_dispatched_stage_message(request, downstream, dispatched, expected_reply) + } + + pub(super) fn dispatch_embedded_stage_message( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + session_key: &str, + message: &StageWireMessage, + token_ids: &[i32], + ) -> OpenAiResult { + let started = Instant::now(); let mut stats = StageReplyStats::default(); let stage0_timer = PhaseTimer::start(); let output = { @@ -24,7 +50,7 @@ impl StageOpenAiBackend { .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; let lock_wait_ms = lock_timer.elapsed_ms(); let hold_timer = PhaseTimer::start(); - if message.kind == WireMessageKind::VerifySpan + if message.kind == WireMessageKind::VerifyWindow && (message.state.flags & state_flags::SKIP_VERIFY_CHECKPOINT) == 0 { let checkpoint_timer = PhaseTimer::start(); @@ -34,9 +60,9 @@ impl StageOpenAiBackend { let checkpoint_us = ms_to_us(checkpoint_timer.elapsed_ms()); stats.checkpoint_local_us += checkpoint_us; stats.checkpoint_total_us += checkpoint_us; - stats.verify_span_checkpointed_requests += 1; - } else if message.kind == WireMessageKind::VerifySpan { - stats.verify_span_skip_checkpoint_requests += 1; + stats.verify_window_checkpointed_requests += 1; + } else if message.kind == WireMessageKind::VerifyWindow { + stats.verify_window_skip_checkpoint_requests += 1; } let output = run_binary_stage_message( &mut runtime, @@ -83,29 +109,10 @@ impl StageOpenAiBackend { ) .map_err(openai_io_error)?; let forward_write_ms = write_timer.elapsed_ms(); - let wait_timer = PhaseTimer::start(); - let reply = receive_embedded_stage_reply( - downstream, - request.prediction_return.as_ref(), - expected_reply, - )?; - let downstream_wait_ms = wait_timer.elapsed_ms(); - stats.merge(reply.stats); - if message.kind == WireMessageKind::VerifySpan { - stats.verify_span_compute_us += ms_to_us(stage0_compute_ms); - stats.verify_span_forward_write_us += ms_to_us(forward_write_ms); - stats.verify_span_downstream_wait_us += ms_to_us(downstream_wait_ms); - stats.verify_span_total_us += ms_to_us(timer.elapsed_ms()); - stats.verify_span_stage_count += 1; - stats.verify_span_request_count += 1; - stats.verify_span_token_count += i64::from(message.token_count.max(0)); - stats.verify_span_max_tokens = stats - .verify_span_max_tokens - .max(i64::from(message.token_count.max(0))); - } - Ok(EmbeddedStageExecution { - reply: StageReply { stats, ..reply }, - stats: EmbeddedExecutionStats { + Ok(DispatchedEmbeddedStage { + started, + stats, + execution: EmbeddedExecutionStats { stage0_compute_ms, runtime_lock_wait_ms: output.runtime_lock_wait_ms, runtime_lock_hold_ms: output.runtime_lock_hold_ms, @@ -113,9 +120,89 @@ impl StageOpenAiBackend { output_activation_bytes: output.output.payload.len(), forward_activation_bytes: forwarded.message.activation.len(), forward_write_ms, - downstream_wait_ms, + downstream_wait_ms: 0.0, }, - elapsed_ms: timer.elapsed_ms(), + message_kind: message.kind, + token_count: message.token_count, + }) + } + + pub(super) fn complete_dispatched_stage_message( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + dispatched: DispatchedEmbeddedStage, + expected_reply: WireReplyKind, + ) -> OpenAiResult { + self.complete_dispatched_stage_message_with_return( + request, + downstream, + dispatched, + expected_reply, + false, + ) + } + + pub(super) fn complete_dispatched_stage_message_direct( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + dispatched: DispatchedEmbeddedStage, + expected_reply: WireReplyKind, + ) -> OpenAiResult { + self.complete_dispatched_stage_message_with_return( + request, + downstream, + dispatched, + expected_reply, + true, + ) + } + + fn complete_dispatched_stage_message_with_return( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + mut dispatched: DispatchedEmbeddedStage, + expected_reply: WireReplyKind, + require_direct_return: bool, + ) -> OpenAiResult { + let wait_timer = PhaseTimer::start(); + let reply = if require_direct_return { + receive_direct_prediction_return(request.prediction_return.as_ref(), expected_reply)? + } else { + receive_embedded_stage_reply( + downstream, + request.prediction_return.as_ref(), + expected_reply, + )? + }; + dispatched.execution.downstream_wait_ms = wait_timer.elapsed_ms(); + dispatched.stats.merge(reply.stats); + if dispatched.message_kind == WireMessageKind::VerifyWindow { + dispatched.stats.verify_window_compute_us += + ms_to_us(dispatched.execution.stage0_compute_ms); + dispatched.stats.verify_window_forward_write_us += + ms_to_us(dispatched.execution.forward_write_ms); + dispatched.stats.verify_window_downstream_wait_us += + ms_to_us(dispatched.execution.downstream_wait_ms); + dispatched.stats.verify_window_total_us += + ms_to_us(dispatched.started.elapsed().as_secs_f64() * 1000.0); + dispatched.stats.verify_window_stage_count += 1; + dispatched.stats.verify_window_request_count += 1; + dispatched.stats.verify_window_token_count += i64::from(dispatched.token_count.max(0)); + dispatched.stats.verify_window_max_tokens = dispatched + .stats + .verify_window_max_tokens + .max(i64::from(dispatched.token_count.max(0))); + } + Ok(EmbeddedStageExecution { + reply: StageReply { + stats: dispatched.stats, + ..reply + }, + stats: dispatched.execution, + elapsed_ms: dispatched.started.elapsed().as_secs_f64() * 1000.0, }) } @@ -219,29 +306,29 @@ impl StageOpenAiBackend { downstream_wait_ms, }) } +} - pub(super) fn trim_embedded_stage_session_local( - &self, - session_key: &str, - token_count: usize, - ) -> OpenAiResult { - let timer = PhaseTimer::start(); - let local_timer = PhaseTimer::start(); +fn receive_direct_prediction_return( + prediction_return: Option<&PredictionReturnReceiver>, + expected_reply: WireReplyKind, +) -> OpenAiResult { + let prediction_return = prediction_return.ok_or_else(|| { + OpenAiError::backend("direct prediction return was required but is not configured") + })?; + let started = Instant::now(); + loop { + if let Some(reply) = prediction_return + .try_recv_expected(expected_reply) + .map_err(openai_backend_error)? { - let mut runtime = self - .runtime - .lock() - .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; - runtime - .trim_session(session_key, token_count as u64) - .map_err(openai_backend_error)?; + return Ok(reply); } - Ok(EmbeddedSessionControl { - elapsed_ms: timer.elapsed_ms(), - local_ms: local_timer.elapsed_ms(), - downstream_write_ms: 0.0, - downstream_wait_ms: 0.0, - }) + if started.elapsed() >= DIRECT_RETURN_FALLBACK_TIMEOUT { + return Err(OpenAiError::backend(format!( + "timed out waiting for {expected_reply:?} reply from direct prediction return" + ))); + } + std::thread::sleep(DIRECT_RETURN_FALLBACK_POLL); } } @@ -250,16 +337,33 @@ pub(crate) fn receive_embedded_stage_reply( prediction_return: Option<&PredictionReturnReceiver>, expected_reply: WireReplyKind, ) -> OpenAiResult { + receive_embedded_stage_reply_one_of( + downstream, + prediction_return, + std::slice::from_ref(&expected_reply), + ) +} + +pub(crate) fn receive_embedded_stage_reply_one_of( + downstream: &mut TcpStream, + prediction_return: Option<&PredictionReturnReceiver>, + expected_replies: &[WireReplyKind], +) -> OpenAiResult { + if expected_replies.is_empty() { + return Err(OpenAiError::backend( + "at least one expected stage reply kind is required", + )); + } let Some(prediction_return) = prediction_return else { - return receive_downstream_stage_reply(downstream, expected_reply); + return receive_downstream_stage_reply_one_of(downstream, expected_replies); }; - poll_direct_or_downstream_reply(downstream, prediction_return, expected_reply) + poll_direct_or_downstream_reply(downstream, prediction_return, expected_replies) } fn poll_direct_or_downstream_reply( downstream: &mut TcpStream, prediction_return: &PredictionReturnReceiver, - expected_reply: WireReplyKind, + expected_replies: &[WireReplyKind], ) -> OpenAiResult { let previous_timeout = downstream.read_timeout().map_err(openai_io_error)?; downstream @@ -268,7 +372,7 @@ fn poll_direct_or_downstream_reply( let started = Instant::now(); loop { if let Some(reply) = prediction_return - .try_recv_expected(expected_reply) + .try_recv_one_of(expected_replies) .map_err(openai_backend_error)? { restore_downstream_read_timeout(downstream, previous_timeout)?; @@ -276,12 +380,12 @@ fn poll_direct_or_downstream_reply( } if downstream_reply_available(downstream)? { restore_downstream_read_timeout(downstream, previous_timeout)?; - return receive_downstream_stage_reply(downstream, expected_reply); + return receive_downstream_stage_reply_one_of(downstream, expected_replies); } if started.elapsed() >= DIRECT_RETURN_FALLBACK_TIMEOUT { restore_downstream_read_timeout(downstream, previous_timeout)?; return Err(OpenAiError::backend(format!( - "timed out waiting for {expected_reply:?} reply from direct return or downstream" + "timed out waiting for one of {expected_replies:?} from direct return or downstream" ))); } } @@ -313,16 +417,101 @@ fn restore_downstream_read_timeout( .map_err(openai_io_error) } -fn receive_downstream_stage_reply( +fn receive_downstream_stage_reply_one_of( downstream: &mut TcpStream, - expected_reply: WireReplyKind, + expected_replies: &[WireReplyKind], ) -> OpenAiResult { let reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; - if reply.kind != expected_reply { + if !expected_replies.contains(&reply.kind) { return Err(OpenAiError::backend(format!( - "expected {expected_reply:?} reply from downstream, got {:?}", + "expected one of {expected_replies:?} from downstream, got {:?}", reply.kind ))); } Ok(reply) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn embedded_stage_reply_accepts_fused_restore_hits_and_misses_from_direct_return() { + assert_eq!( + receive_direct_reply_one_of( + WireReplyKind::PredictedToken, + &[WireReplyKind::PredictedToken, WireReplyKind::Ack], + ), + WireReplyKind::PredictedToken + ); + assert_eq!( + receive_direct_reply_one_of( + WireReplyKind::Ack, + &[WireReplyKind::PredictedToken, WireReplyKind::Ack], + ), + WireReplyKind::Ack + ); + } + + fn receive_direct_reply_one_of( + reply_kind: WireReplyKind, + expected_replies: &[WireReplyKind], + ) -> WireReplyKind { + let request_id = 17; + let session_id = 23; + let hub = Arc::new(PredictionReturnHub::default()); + let receiver = hub.register(request_id, session_id).unwrap(); + let (mut direct_client, direct_server) = tcp_pair(); + let hub_thread = { + let hub = hub.clone(); + std::thread::spawn(move || { + hub.handle_return_connection( + StageWireMessage { + kind: WireMessageKind::PredictionReturnOpen, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new( + WireMessageKind::PredictionReturnOpen, + WireActivationDType::F32, + ), + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }, + direct_server, + ) + }) + }; + skippy_protocol::binary::send_reply_message( + &mut direct_client, + &StageReply { + kind: reply_kind, + predicted: 0, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: Default::default(), + stats: StageReplyStats::default(), + }, + ) + .unwrap(); + let (mut downstream, _downstream_peer) = tcp_pair(); + let reply = + receive_embedded_stage_reply_one_of(&mut downstream, Some(&receiver), expected_replies) + .unwrap(); + drop(direct_client); + hub_thread.join().unwrap().unwrap(); + reply.kind + } + + fn tcp_pair() -> (TcpStream, TcpStream) { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let client = TcpStream::connect(listener.local_addr().unwrap()).unwrap(); + let (server, _) = listener.accept().unwrap(); + (client, server) + } +} diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 46e38fa1e0..edf9402033 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1,5 +1,17 @@ +use std::collections::VecDeque; + +use super::embedded_execution::DispatchedEmbeddedStage; use super::*; +struct PipelinedCompositeWindow { + window: VerifyWindow, + input_tokens: Vec, + proposal_tokens: Vec, + expected_free_target: Option, + native_mtp_token_count: usize, + dispatched: DispatchedEmbeddedStage, +} + impl StageOpenAiBackend { pub(super) fn generate_embedded_stage_zero_tokens( &self, @@ -13,9 +25,8 @@ impl StageOpenAiBackend { max_tokens: request.max_tokens, sampling: request.sampling, chat_sampling_metadata: request.chat_sampling_metadata, + speculative: request.speculative, native_mtp_enabled: request.native_mtp_enabled, - native_mtp_max_tokens: request.native_mtp_max_tokens, - native_mtp_min_tokens: request.native_mtp_min_tokens, hook_request: request.hook_request, hook_runtime: request.hook_runtime, cancellation: request.cancellation, @@ -34,6 +45,7 @@ impl StageOpenAiBackend { .as_ref() .ok_or_else(|| OpenAiError::backend("embedded stage 0 has no downstream lane pool"))?; let mut lane = lane_pool.checkout(request.ids)?; + let mut direct_prediction_return_opened = false; if let Some(prediction_return) = request.prediction_return.as_ref() { match crate::binary_transport::direct_return::open_downstream_prediction_return_stream( request.config, @@ -43,6 +55,7 @@ impl StageOpenAiBackend { ) { Ok(stream) => { prediction_return.attach_opened_stream(stream); + direct_prediction_return_opened = true; } Err(error) => { eprintln!( @@ -587,6 +600,7 @@ impl StageOpenAiBackend { request.activation_width, ), ); + cache_stats.prompt_ms = prefill_timer.elapsed_ms(); self.emit_openai_phase("stage.openai_prefill", prefill_timer, prefill_attrs); let message = generation_config_message( @@ -646,11 +660,14 @@ impl StageOpenAiBackend { )?; let mut fused_reached_stop = false; let mut native_mtp = NativeMtpVerifier::default(); - let native_mtp_options = NativeMtpDecodeOptions::from_env() - .with_window(request.native_mtp_max_tokens, request.native_mtp_min_tokens); + let native_mtp_options = NativeMtpDecodeOptions::from_config(request.speculative); let mut native_mtp_counters = NativeMtpDecodeCounters::default(); let mut native_mtp_reject_cooldown_remaining = 0usize; let mut native_mtp_suppress_cooldown_drafts_remaining = 0usize; + let mut ngram_sidecar_controller = NgramSidecarController::new( + native_mtp_options.ngram_initial_extension_tokens, + native_mtp_options.ngram_max_proposal_tokens, + ); if let Some(mut fused) = fused_first_decode.take() { current = fused.predicted; let mut fused_native_mtp_draft = fused.native_mtp_draft.take(); @@ -793,6 +810,7 @@ impl StageOpenAiBackend { } } } + let mut cached_ngram_proposer = CachedNgramProposer::from_config(request.speculative)?; let max_speculative_window = request.speculative_window.max(1); let mut adaptive_window = if request.adaptive_speculative_window { max_speculative_window.min(4) @@ -840,6 +858,28 @@ impl StageOpenAiBackend { } _ => None, }; + let mut verify_window_scheduler = VerifyWindowScheduler::new( + VerifyWindowPipelineConfig::new(request.speculative.verify_window.pipeline_depth), + ); + let composite_sidecar_enabled = + native_mtp_options.ngram_hybrid && draft_guard.is_none(); + let native_mtp_verify_windows_enabled = + (request.native_mtp_enabled || composite_sidecar_enabled) && draft_guard.is_none(); + let pipelined_decode_enabled = + composite_sidecar_enabled && verify_window_scheduler.depth() > 1; + if native_mtp_verify_windows_enabled && !direct_prediction_return_opened { + return Err(OpenAiError::backend( + "native MTP verify windows require direct prediction return", + )); + } + if native_mtp_verify_windows_enabled { + verify_window_scheduler.mark_direct_prediction_return(); + } + let mut pipelined_windows = VecDeque::new(); + let mut pipelined = None; + let mut pipelined_current = current; + let mut composite_proposal_buffer = None; + let mut adaptive_verify_window = AdaptiveVerifyWindow::new(native_mtp_options); for decode_step in decoded_tokens as u32..request.max_tokens { if fused_reached_stop { break; @@ -856,15 +896,83 @@ impl StageOpenAiBackend { let token_timer = PhaseTimer::start(); let native_mtp_remaining = (request.max_tokens as usize).saturating_sub(decoded_tokens); - let can_run_native_mtp_batched_verify = native_mtp_options.batched_verify + let mut pipeline_seed = None; + if pipelined_decode_enabled + && pipelined.is_none() + && pipelined_windows.is_empty() + && native_mtp_reject_cooldown_remaining == 0 + && verify_window_scheduler.has_profitable_pipeline_width() + { + let pending = request + .native_mtp_enabled + .then(|| native_mtp.take_pending_draft()) + .flatten(); + let native_mtp_origin = pending.as_ref().map(|draft| draft.origin); + let native_mtp_tokens = pending + .as_ref() + .map(|draft| { + draft + .tokens + .iter() + .copied() + .take(native_mtp_options.max_draft_tokens) + .take(native_mtp_remaining.saturating_sub(1)) + .collect::>() + }) + .unwrap_or_default(); + let native_mtp_tokens = + if native_mtp_tokens.len() >= native_mtp_options.min_draft_tokens { + native_mtp_tokens.as_slice() + } else { + &[] + }; + let proposal = CompositeProposalProvider::from_options(native_mtp_options) + .propose_with_ngram_extension( + native_mtp_tokens, + &context_tokens, + native_mtp_remaining, + ngram_sidecar_controller.extension_limit( + native_mtp_tokens, + native_mtp_remaining.saturating_sub(native_mtp_tokens.len()), + ), + cached_ngram_proposer.as_mut(), + )?; + match proposal.parallel_verify_width( + adaptive_verify_window.width(proposal.tokens().len()), + verify_window_scheduler.depth(), + ) { + Some(parallel_verify_width) + if verify_window_scheduler + .permit_pipeline_width(parallel_verify_width) => + { + pipeline_seed = Some(( + proposal, + if native_mtp_tokens.is_empty() { + None + } else { + native_mtp_origin + }, + parallel_verify_width, + )); + } + _ => { + if let Some(pending) = pending { + native_mtp.restore_pending_draft(pending); + } + } + } + } + if native_mtp_verify_windows_enabled + && (!pipelined_decode_enabled + || (pipelined.is_none() && pipeline_seed.is_none())) && native_mtp_reject_cooldown_remaining == 0 - && draft_guard.is_none() - && native_mtp_remaining >= 2; - let pending_native_mtp_draft = can_run_native_mtp_batched_verify + && native_mtp_remaining >= 2 + { + let pending_native_mtp_draft = (request.native_mtp_enabled + && composite_proposal_buffer.is_none()) .then(|| native_mtp.take_pending_draft()) .flatten(); - if let Some(pending_native_mtp_draft) = pending_native_mtp_draft { - match self.execute_native_mtp_batched_verify( + match self.execute_native_mtp_verify_window( &request, downstream, &session_key, @@ -873,7 +981,11 @@ impl StageOpenAiBackend { prefill_token_count, &wire_sampling, &native_mtp_options, + &mut verify_window_scheduler, pending_native_mtp_draft, + &mut composite_proposal_buffer, + &mut cached_ngram_proposer, + &mut adaptive_verify_window, &mut current, decode_step, &mut decoded_tokens, @@ -883,6 +995,7 @@ impl StageOpenAiBackend { &mut native_mtp_counters, &mut native_mtp_reject_cooldown_remaining, &mut native_mtp_suppress_cooldown_drafts_remaining, + &mut ngram_sidecar_controller, &mut decode_stage0_compute_ms, &mut decode_runtime_lock_wait_ms, &mut decode_runtime_lock_wait_max_ms, @@ -896,8 +1009,353 @@ impl StageOpenAiBackend { &mut decode_downstream_wait_ms, &mut on_token, )? { - BatchedVerifyControl::ReachedStop => break, - BatchedVerifyControl::Continue => continue, + NativeMtpVerifyWindowControl::ReachedStop => break, + NativeMtpVerifyWindowControl::Continue => continue, + NativeMtpVerifyWindowControl::NoProposal => {} + } + } + if pipelined_decode_enabled { + if let Some((proposal, origin, parallel_verify_width)) = pipeline_seed { + pipelined_current = current; + pipelined = Some(CompositeProposalPipeline::new( + proposal, + origin, + parallel_verify_width, + )); + } + if let Some(pipeline) = pipelined.as_mut() { + while verify_window_scheduler.has_capacity() + && decoded_tokens + + pipelined_windows + .iter() + .map(|window: &PipelinedCompositeWindow| { + window.input_tokens.len() + }) + .sum::() + < request.max_tokens as usize + { + let Some(planned) = pipeline.next_window( + adaptive_verify_window.width(pipeline.candidate_len()), + ) else { + break; + }; + let proposal_tokens = planned.proposal_tokens().to_vec(); + let expected_free_target = planned.expected_free_target(); + let native_mtp_token_count = planned.native_mtp_token_count(); + let offset = pipelined_windows + .iter() + .map(|window: &PipelinedCompositeWindow| window.input_tokens.len()) + .sum::(); + let window = verify_window_scheduler.open( + prefill_token_count + decoded_tokens + offset, + decoded_tokens + offset, + )?; + let mut input_tokens = Vec::with_capacity(proposal_tokens.len() + 1); + input_tokens.push(pipelined_current); + input_tokens.extend_from_slice(&proposal_tokens); + let message = embedded_verify_window_message( + request.wire_dtype, + VerifyWindowMessageArgs { + window_id: window.id, + request_id, + session_id, + prompt_token_count: request.prompt_token_ids.len(), + pos_start: window.base_position, + decode_step: window.decode_step, + tokens: &input_tokens, + sampling: wire_sampling.clone(), + checkpoint: false, + }, + )?; + let dispatched = self.dispatch_embedded_stage_message( + &request, + downstream, + &session_key, + &message, + &input_tokens, + )?; + pipelined_windows.push_back(PipelinedCompositeWindow { + window, + input_tokens, + proposal_tokens, + expected_free_target, + native_mtp_token_count, + dispatched, + }); + let Some(next_current) = expected_free_target else { + break; + }; + pipelined_current = next_current; + } + if let Some(window) = pipelined_windows.pop_front() { + let verify = self.complete_dispatched_stage_message_direct( + &request, + downstream, + window.dispatched, + WireReplyKind::PredictedTokens, + )?; + let completed = verify_window_scheduler + .complete_next(verify.reply.window.window_id)?; + if completed != window.window { + return Err(OpenAiError::backend( + "verify window scheduler lost FIFO state", + )); + } + let native_mtp_verify_decision = classify_native_mtp_verify_window( + &window.proposal_tokens, + &verify.reply.predicted_tokens, + decoded_tokens, + request.max_tokens as usize, + |token| token_is_eog_with_runtime(&self.runtime, token), + )?; + let fully_accepted_window = !native_mtp_verify_decision.rejected + && native_mtp_verify_decision.accepted_proposal_tokens + == window.proposal_tokens.len(); + let free_target_matches = + window.expected_free_target.is_none_or(|expected| { + verify + .reply + .predicted_tokens + .get(window.proposal_tokens.len()) + == Some(&expected) + }); + let pipeline_continues = fully_accepted_window && free_target_matches; + if window.expected_free_target.is_some() { + verify_window_scheduler.observe_pipeline_profile( + window.proposal_tokens.len(), + pipeline_continues, + verify.stats.stage0_compute_ms, + verify.stats.downstream_wait_ms, + ); + } + let accepted_candidate_tokens = native_mtp_verify_decision + .accepted_proposal_tokens + + usize::from( + fully_accepted_window + && window.expected_free_target.is_some() + && free_target_matches, + ); + if window.native_mtp_token_count > 0 { + let pipeline = pipelined.as_ref().expect("pipeline retained"); + let span = native_mtp.observe_taken_draft_span( + &window.proposal_tokens[..window.native_mtp_token_count], + &verify.reply.predicted_tokens, + ms_to_us(verify.elapsed_ms), + ); + for index in 0..span.accepted_count + usize::from(span.rejected) { + native_mtp_counters.observe_verify_window_verification( + pipeline.origin().expect("native MTP candidate has origin"), + index < span.accepted_count, + ); + } + } + speculative_stats.windows += 1; + speculative_stats.draft_tokens += window.proposal_tokens.len(); + speculative_stats.primary_verify_requests += 1; + speculative_stats.primary_verify_tokens += window.input_tokens.len(); + speculative_stats.primary_verify_elapsed_ms += verify.elapsed_ms; + speculative_stats.primary_verify_stage0_compute_ms += + verify.stats.stage0_compute_ms; + speculative_stats.primary_verify_runtime_lock_wait_ms += + verify.stats.runtime_lock_wait_ms; + speculative_stats.primary_verify_runtime_lock_hold_ms += + verify.stats.runtime_lock_hold_ms; + speculative_stats.primary_verify_activation_encode_ms += + verify.stats.activation_encode_ms; + speculative_stats.primary_verify_forward_write_ms += + verify.stats.forward_write_ms; + speculative_stats.primary_verify_downstream_wait_ms += + verify.stats.downstream_wait_ms; + speculative_stats.primary_verify_output_activation_bytes = + speculative_stats + .primary_verify_output_activation_bytes + .saturating_add(verify.stats.output_activation_bytes); + speculative_stats.primary_verify_forward_activation_bytes = + speculative_stats + .primary_verify_forward_activation_bytes + .saturating_add(verify.stats.forward_activation_bytes); + decode_stage0_compute_ms += verify.stats.stage0_compute_ms; + decode_runtime_lock_wait_ms += verify.stats.runtime_lock_wait_ms; + decode_runtime_lock_wait_max_ms = decode_runtime_lock_wait_max_ms + .max(verify.stats.runtime_lock_wait_ms); + decode_runtime_lock_hold_ms += verify.stats.runtime_lock_hold_ms; + decode_runtime_lock_hold_max_ms = decode_runtime_lock_hold_max_ms + .max(verify.stats.runtime_lock_hold_ms); + decode_runtime_lock_acquires += 1; + decode_forward_activation_encode_ms += + verify.stats.activation_encode_ms; + decode_output_activation_bytes = decode_output_activation_bytes + .saturating_add(verify.stats.output_activation_bytes); + decode_forward_activation_bytes = decode_forward_activation_bytes + .saturating_add(verify.stats.forward_activation_bytes); + decode_forward_write_ms += verify.stats.forward_write_ms; + decode_downstream_wait_ms += verify.stats.downstream_wait_ms; + if fully_accepted_window { + speculative_stats.accepted_tokens += accepted_candidate_tokens; + speculative_stats.full_accept_windows += 1; + let pipeline = pipelined.as_mut().expect("pipeline retained"); + pipeline.set_next_draft( + request.native_mtp_enabled, + verify + .reply + .native_mtp_draft + .clone() + .map(NativeMtpDraft::from_stage_draft), + ); + } else { + speculative_stats.rejected_tokens += 1; + speculative_stats.rejected_windows += 1; + speculative_stats.early_reject_windows += 1; + speculative_stats.first_reject_position_sum += 1; + } + pipelined + .as_mut() + .expect("pipeline retained") + .observe_accepted(accepted_candidate_tokens); + let mut reached_stop = false; + for token in verify + .reply + .predicted_tokens + .iter() + .copied() + .take(native_mtp_verify_decision.commit_count) + { + current = token; + decoded_tokens += 1; + exact_replay_tokens.push(current); + context_tokens.push(current); + if on_token(current)? == TokenControl::Stop + || decoded_tokens >= request.max_tokens as usize + { + reached_stop = true; + break; + } + } + let previous_verify_width = adaptive_verify_window.current_tokens(); + adaptive_verify_window.observe(pipeline_continues); + native_mtp_counters.observe_adaptive_verify_window( + window.proposal_tokens.len(), + previous_verify_width, + adaptive_verify_window.current_tokens(), + ); + if !pipeline_continues || reached_stop { + let stale_count = pipelined_windows.len(); + let stale_drain_timer = PhaseTimer::start(); + while let Some(stale) = pipelined_windows.pop_front() { + let stale_reply = self + .complete_dispatched_stage_message_direct( + &request, + downstream, + stale.dispatched, + WireReplyKind::PredictedTokens, + )?; + verify_window_scheduler + .complete_next(stale_reply.reply.window.window_id)?; + } + verify_window_scheduler.record_stale_discarded( + stale_count, + stale_drain_timer.elapsed_ms(), + ); + let pipeline = pipelined.take().expect("pipeline retained"); + if ngram_sidecar_controller.observe_tail_outcome( + pipeline.proposal(), + pipeline.accepted_tokens(), + native_mtp_options.ngram_tail_backoff_proposals, + ) { + native_mtp_counters.observe_ngram_tail_rejection(); + } + native_mtp_counters.observe_hybrid_proposal( + pipeline.proposal(), + pipeline.accepted_tokens(), + ); + native_mtp.clear_pending_draft(); + if native_mtp_verify_decision.rejected + && pipeline + .proposal() + .native_mtp_prefix_rejected(pipeline.accepted_tokens()) + && native_mtp_options.reject_cooldown_tokens > 0 + { + native_mtp_reject_cooldown_remaining = + native_mtp_options.reject_cooldown_tokens; + native_mtp_suppress_cooldown_drafts_remaining = + native_mtp_options.suppress_cooldown_draft_limit; + } + if native_mtp_verify_decision.rejected || stale_count > 0 { + let trim = self.trim_embedded_stage_session( + &request, + downstream, + &session_key, + request_id, + session_id, + prefill_token_count + decoded_tokens, + )?; + speculative_stats.recovery_ms += trim.elapsed_ms; + } + pipelined_current = current; + if reached_stop { + break; + } + } else if pipelined_windows.is_empty() + && pipelined + .as_ref() + .is_some_and(|pipeline| !pipeline.has_remaining_candidates()) + { + let mut pipeline = pipelined.take().expect("pipeline retained"); + let next_draft_available = pipeline.next_draft().is_some(); + ngram_sidecar_controller.observe_tail_outcome( + pipeline.proposal(), + pipeline.accepted_tokens(), + native_mtp_options.ngram_tail_backoff_proposals, + ); + native_mtp_counters.observe_hybrid_proposal( + pipeline.proposal(), + pipeline.accepted_tokens(), + ); + native_mtp_counters.observe_verify_next_draft( + next_draft_available, + next_draft_available, + ); + if let Some(next_draft) = pipeline.take_next_draft() { + native_mtp.observe_next_draft( + Some(next_draft), + NativeMtpDraftOrigin::VerifyNext, + ); + } + } + if self.telemetry.is_debug_enabled() { + let mut attrs = self.openai_attrs(request.ids); + attrs.insert( + "llama_stage.message_kind".to_string(), + json!("VerifyWindow"), + ); + attrs.insert( + "llama_stage.spec.proposal_source".to_string(), + json!("composite_mtp_ngram"), + ); + attrs.insert( + "llama_stage.verify_window_id".to_string(), + json!(window.window.id), + ); + attrs.insert( + "llama_stage.verify_window.accepted".to_string(), + json!(pipeline_continues), + ); + attrs.insert( + "llama_stage.verify_window.in_flight_after".to_string(), + json!(verify_window_scheduler.in_flight_len()), + ); + attrs.insert( + "llama_stage.verify_window.stale_discarded".to_string(), + json!(verify_window_scheduler.stale_discard_count()), + ); + self.emit_openai_phase( + "stage.openai_decode_verify_window", + token_timer, + attrs, + ); + } + continue; + } } } if draft_guard.is_some() || request.ngram_max > 0 { @@ -920,12 +1378,24 @@ impl StageOpenAiBackend { proposal_source = "draft-model"; } } + if let (true, Some(cache)) = + (draft_tokens.is_empty(), cached_ngram_proposer.as_mut()) + { + draft_tokens = cache.propose( + &context_tokens, + &[], + proposal_limit.min(request.ngram_max), + )?; + if !draft_tokens.is_empty() { + proposal_source = "ngram-cache"; + } + } if draft_tokens.is_empty() && request.ngram_max > 0 { draft_tokens = propose_ngram_tokens( &context_tokens, request.ngram_min, proposal_limit.min(request.ngram_max), - ); + )?; if !draft_tokens.is_empty() { proposal_source = "ngram"; } @@ -934,9 +1404,11 @@ impl StageOpenAiBackend { speculative_stats.draft_propose_ms += draft_propose_ms; if !draft_tokens.is_empty() { let verify_inputs = verify_inputs_for_proposals(current, &draft_tokens); - let message = embedded_verify_message( + let message = embedded_verify_window_message( request.wire_dtype, - VerifySpanMessageArgs { + VerifyWindowMessageArgs { + window_id: i32::try_from(decoded_tokens) + .map_err(|_| OpenAiError::backend("decode step exceeds i32"))?, request_id, session_id, prompt_token_count: request.prompt_token_ids.len(), @@ -997,7 +1469,7 @@ impl StageOpenAiBackend { decode_downstream_wait_ms += verify.stats.downstream_wait_ms; speculative_stats.checkpoint_ms += us_to_ms(verify.reply.stats.checkpoint_total_us); - let decision = classify_verify_span( + let decision = classify_verify_window( &draft_tokens, &verify.reply.predicted_tokens, decoded_tokens, @@ -1074,9 +1546,12 @@ impl StageOpenAiBackend { speculative_stats.recovery_decode_elapsed_ms += repair.elapsed_ms; } else { let repair_inputs = &verify_inputs[..repair_input_count]; - let repair_message = embedded_verify_message( + let repair_message = embedded_verify_window_message( request.wire_dtype, - VerifySpanMessageArgs { + VerifyWindowMessageArgs { + window_id: i32::try_from(decoded_tokens).map_err(|_| { + OpenAiError::backend("decode step exceeds i32") + })?, request_id, session_id, prompt_token_count: request.prompt_token_ids.len(), @@ -1149,8 +1624,10 @@ impl StageOpenAiBackend { let mut token_attrs = self.openai_attrs(request.ids); token_attrs .insert("llama_stage.decode_step".to_string(), json!(decode_step)); - token_attrs - .insert("llama_stage.message_kind".to_string(), json!("VerifySpan")); + token_attrs.insert( + "llama_stage.message_kind".to_string(), + json!("VerifyWindow"), + ); token_attrs.insert( "llama_stage.spec.windows".to_string(), json!(speculative_stats.windows), @@ -1330,7 +1807,10 @@ impl StageOpenAiBackend { let native_mtp_draft = if suppress_cooldown_draft { None } else { - NativeMtpDraft::from_prediction_tokens(&reply.predicted_tokens) + reply + .native_mtp_draft + .clone() + .map(NativeMtpDraft::from_stage_draft) }; if suppress_cooldown_draft { native_mtp.clear_pending_draft(); @@ -1342,7 +1822,7 @@ impl StageOpenAiBackend { current, ms_to_us(downstream_wait_ms), native_mtp_draft, - if native_mtp_counters.batched_verification_count() == 0 { + if native_mtp_counters.verify_window_verification_count() == 0 { NativeMtpDraftOrigin::InitialSerial } else { NativeMtpDraftOrigin::SerialAfterGap @@ -1424,6 +1904,35 @@ impl StageOpenAiBackend { break; } } + if !pipelined_windows.is_empty() { + let stale_count = pipelined_windows.len(); + let stale_drain_timer = PhaseTimer::start(); + while let Some(stale) = pipelined_windows.pop_front() { + let stale_reply = self.complete_dispatched_stage_message_direct( + &request, + downstream, + stale.dispatched, + WireReplyKind::PredictedTokens, + )?; + verify_window_scheduler.complete_next(stale_reply.reply.window.window_id)?; + } + verify_window_scheduler + .record_stale_discarded(stale_count, stale_drain_timer.elapsed_ms()); + let trim = self.trim_embedded_stage_session( + &request, + downstream, + &session_key, + request_id, + session_id, + prefill_token_count + decoded_tokens, + )?; + speculative_stats.recovery_ms += trim.elapsed_ms; + } + if let Some(pipeline) = pipelined.take() { + native_mtp_counters + .observe_hybrid_proposal(pipeline.proposal(), pipeline.accepted_tokens()); + native_mtp.clear_pending_draft(); + } let mut decode_attrs = self.openai_attrs(request.ids); decode_attrs.insert( "llama_stage.decode_token_count".to_string(), @@ -1481,11 +1990,22 @@ impl StageOpenAiBackend { "llama_stage.downstream_wait_ms".to_string(), json!(decode_downstream_wait_ms), ); + request + .speculative + .insert_telemetry_attrs(&mut decode_attrs); speculative_stats.insert_attrs(&mut decode_attrs); let native_mtp_stats = native_mtp.stats(); cache_stats.native_mtp_stats = native_mtp_stats; + cache_stats.native_mtp_decode_telemetry = Some(NativeMtpDecodeTelemetry::new( + native_mtp_options, + native_mtp_counters, + )); + cache_stats.verify_window_pipeline_stats = Some(verify_window_scheduler.stats()); + cache_stats.speculative_stats = Some(speculative_stats.clone()); + cache_stats.predicted_ms = decode_timer.elapsed_ms(); native_mtp_stats.insert_attrs(&mut decode_attrs); native_mtp_counters.insert_summary_attrs(&mut decode_attrs, native_mtp_options); + verify_window_scheduler.insert_policy_telemetry_attrs(&mut decode_attrs); self.emit_openai_summary("stage.openai_decode", decode_timer, decode_attrs); Ok(()) })(); diff --git a/crates/skippy-server/src/frontend/generation_flow.rs b/crates/skippy-server/src/frontend/generation_flow.rs index c74bfcd4ba..fdd5b5cd43 100644 --- a/crates/skippy-server/src/frontend/generation_flow.rs +++ b/crates/skippy-server/src/frontend/generation_flow.rs @@ -101,9 +101,9 @@ impl StageOpenAiBackend { max_tokens, sampling: &sampling, chat_sampling_metadata, - native_mtp_enabled: self.native_mtp_enabled, - native_mtp_max_tokens: self.native_mtp_max_tokens, - native_mtp_min_tokens: self.native_mtp_min_tokens, + speculative: &self.speculative, + native_mtp_enabled: self.config.native_mtp_enabled + && self.speculative.native_mtp.enabled, hook_request: hook_request.clone(), hook_runtime: hook_runtime.clone(), cancellation, @@ -120,9 +120,6 @@ impl StageOpenAiBackend { prefill_reply_credit_limit, lane_pool, prediction_returns, - native_mtp_enabled, - native_mtp_max_tokens, - native_mtp_min_tokens, } => self.generate_embedded_stage_zero_tokens( EmbeddedStageZeroGeneration { config: &config, @@ -140,11 +137,12 @@ impl StageOpenAiBackend { draft: self.draft.clone(), speculative_window: self.speculative_window, adaptive_speculative_window: self.adaptive_speculative_window, + speculative: &self.speculative, ngram_min: self.ngram_min, ngram_max: self.ngram_max, - native_mtp_enabled, - native_mtp_max_tokens, - native_mtp_min_tokens, + native_mtp_enabled: config.native_mtp_enabled + && self.speculative.native_mtp.enabled, + native_mtp_max_tokens: self.speculative.native_mtp.max_draft_tokens, prompt_token_ids: &prompt_token_ids, max_tokens, sampling: &sampling, diff --git a/crates/skippy-server/src/frontend/local_generation.rs b/crates/skippy-server/src/frontend/local_generation.rs index c24538d83a..37363c6e25 100644 --- a/crates/skippy-server/src/frontend/local_generation.rs +++ b/crates/skippy-server/src/frontend/local_generation.rs @@ -99,6 +99,7 @@ impl StageOpenAiBackend { "llama_stage.runtime_sessions_after", &runtime_sessions_after, ); + cache_stats.prompt_ms = prefill_timer.elapsed_ms(); self.emit_openai_phase("stage.openai_prefill", prefill_timer, attrs); } else if request.prompt_token_ids.len() > 1 { let prefill_timer = PhaseTimer::start(); @@ -454,6 +455,7 @@ impl StageOpenAiBackend { "llama_stage.runtime_sessions_after", &runtime_sessions_after, ); + cache_stats.prompt_ms = prefill_timer.elapsed_ms(); self.emit_openai_phase("stage.openai_prefill", prefill_timer, attrs); self.telemetry.emit( "stage.openai_kv_record_decision", @@ -510,8 +512,7 @@ impl StageOpenAiBackend { let generation_hooks_active = self.generation_hooks_active(&hook_request, hook_runtime.as_ref()); let emit_token_debug = self.telemetry.is_debug_enabled(); - let native_mtp_options = NativeMtpDecodeOptions::from_env() - .with_window(request.native_mtp_max_tokens, request.native_mtp_min_tokens); + let native_mtp_options = NativeMtpDecodeOptions::from_config(request.speculative); let mut native_mtp = NativeMtpVerifier::default(); let mut post_prefill_hook_checked = false; let mut last_mid_generation_hook_at = None; @@ -720,8 +721,10 @@ impl StageOpenAiBackend { stats, ); } + request.speculative.insert_telemetry_attrs(&mut attrs); let native_mtp_stats = native_mtp.stats(); cache_stats.native_mtp_stats = native_mtp_stats; + cache_stats.predicted_ms = decode_timer.elapsed_ms(); native_mtp_stats.insert_attrs(&mut attrs); self.emit_openai_summary("stage.openai_decode", decode_timer, attrs); Ok(()) diff --git a/crates/skippy-server/src/frontend/native_mtp/batched_verify.rs b/crates/skippy-server/src/frontend/native_mtp/batched_verify.rs deleted file mode 100644 index ffdae71cb3..0000000000 --- a/crates/skippy-server/src/frontend/native_mtp/batched_verify.rs +++ /dev/null @@ -1,347 +0,0 @@ -use std::net::TcpStream; - -use openai_frontend::{OpenAiError, OpenAiResult}; -use skippy_protocol::binary::WireReplyKind; - -use super::super::{ - EmbeddedSessionControl, EmbeddedStageZeroGeneration, NativeMtpDecodeCounters, - NativeMtpDecodeOptions, NativeMtpDraft, NativeMtpDraftOrigin, NativeMtpTrimAction, - NativeMtpVerifier, PendingNativeMtpDraft, PhaseTimer, StageOpenAiBackend, TokenControl, - VerifySpanMessageArgs, WireSamplingConfig, embedded_verify_message, ms_to_us, - native_mtp_trim_action, -}; - -/// Control signal returned after processing a batched native MTP verify step. -pub(in crate::frontend) enum BatchedVerifyControl { - /// The on_token callback returned Stop — outer loop should break. - ReachedStop, - /// Continue the outer decode loop normally. - Continue, -} - -impl StageOpenAiBackend { - #[allow(clippy::too_many_arguments)] - pub(in crate::frontend) fn execute_native_mtp_batched_verify( - &self, - request: &EmbeddedStageZeroGeneration<'_>, - downstream: &mut TcpStream, - session_key: &str, - request_id: u64, - session_id: u64, - prefill_token_count: usize, - wire_sampling: &Option, - native_mtp_options: &NativeMtpDecodeOptions, - pending_native_mtp_draft: PendingNativeMtpDraft, - current: &mut i32, - decode_step: u32, - // Mutable decode loop state - decoded_tokens: &mut usize, - context_tokens: &mut Vec, - exact_replay_tokens: &mut Vec, - native_mtp: &mut NativeMtpVerifier, - native_mtp_counters: &mut NativeMtpDecodeCounters, - native_mtp_reject_cooldown_remaining: &mut usize, - native_mtp_suppress_cooldown_drafts_remaining: &mut usize, - // Mutable decode accumulators - decode_stage0_compute_ms: &mut f64, - decode_runtime_lock_wait_ms: &mut f64, - decode_runtime_lock_wait_max_ms: &mut f64, - decode_runtime_lock_hold_ms: &mut f64, - decode_runtime_lock_hold_max_ms: &mut f64, - decode_runtime_lock_acquires: &mut usize, - decode_forward_activation_encode_ms: &mut f64, - decode_output_activation_bytes: &mut usize, - decode_forward_activation_bytes: &mut usize, - decode_forward_write_ms: &mut f64, - decode_downstream_wait_ms: &mut f64, - // Token emission callback - on_token: &mut impl FnMut(i32) -> OpenAiResult, - ) -> OpenAiResult { - let batched_token_timer = self.telemetry.is_debug_enabled().then(PhaseTimer::start); - let native_mtp_remaining = (request.max_tokens as usize).saturating_sub(*decoded_tokens); - let native_mtp_draft_tokens = pending_native_mtp_draft - .tokens - .into_iter() - .take(native_mtp_options.max_draft_tokens) - .take(native_mtp_remaining.saturating_sub(1)) - .collect::>(); - if native_mtp_draft_tokens.is_empty() - || native_mtp_draft_tokens.len() < native_mtp_options.min_draft_tokens - { - native_mtp.clear_pending_draft(); - return Ok(BatchedVerifyControl::Continue); - } - let native_mtp_draft_origin = pending_native_mtp_draft.origin; - let mut verify_inputs = Vec::with_capacity(native_mtp_draft_tokens.len() + 1); - verify_inputs.push(*current); - verify_inputs.extend(native_mtp_draft_tokens.iter().copied()); - let message = embedded_verify_message( - request.wire_dtype, - VerifySpanMessageArgs { - request_id, - session_id, - prompt_token_count: request.prompt_token_ids.len(), - pos_start: prefill_token_count + *decoded_tokens, - decode_step: *decoded_tokens, - tokens: &verify_inputs, - sampling: wire_sampling.clone(), - checkpoint: false, - }, - )?; - let verify = self.execute_embedded_stage_message( - request, - downstream, - session_key, - &message, - &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 target_token = verify.reply.predicted_tokens[0]; - let verify_next_mtp_draft = NativeMtpDraft::from_verify_prediction_tokens( - &verify.reply.predicted_tokens, - verify_inputs.len(), - ); - let span = native_mtp.observe_taken_draft_span( - &native_mtp_draft_tokens, - &verify.reply.predicted_tokens, - ms_to_us(verify.elapsed_ms), - ); - let native_mtp_decision = span.first_decision; - let accepted = !span.rejected; - let verified_draft_count = span.accepted_count + usize::from(span.rejected); - for index in 0..verified_draft_count { - native_mtp_counters - .observe_batched_verification(native_mtp_draft_origin, index < span.accepted_count); - } - let commit_token_count = span.accepted_count.saturating_add(1); - let consumed_positions = verify_inputs.len(); - let mut committed_positions = 0usize; - let mut reached_stop = false; - for token in verify - .reply - .predicted_tokens - .iter() - .copied() - .take(commit_token_count) - { - *current = token; - *decoded_tokens += 1; - committed_positions += 1; - exact_replay_tokens.push(*current); - context_tokens.push(*current); - if on_token(*current)? == TokenControl::Stop { - reached_stop = true; - break; - } - if *decoded_tokens >= request.max_tokens as usize { - break; - } - } - if !accepted && native_mtp_options.reject_cooldown_tokens > 0 { - *native_mtp_reject_cooldown_remaining = native_mtp_options.reject_cooldown_tokens; - *native_mtp_suppress_cooldown_drafts_remaining = - native_mtp_options.suppress_cooldown_draft_limit; - native_mtp.clear_pending_draft(); - } - let verify_next_mtp_draft_available = verify_next_mtp_draft.is_some(); - let verify_next_mtp_draft_adopted = accepted - && committed_positions == consumed_positions - && !reached_stop - && *decoded_tokens < request.max_tokens as usize - && verify_next_mtp_draft.is_some(); - native_mtp_counters.observe_verify_next_draft( - verify_next_mtp_draft_available, - verify_next_mtp_draft_adopted, - ); - if verify_next_mtp_draft_adopted { - native_mtp.observe_next_draft( - verify_next_mtp_draft.clone(), - NativeMtpDraftOrigin::VerifyNext, - ); - } - let mut trim_control: Option = None; - match native_mtp_trim_action(committed_positions, consumed_positions) { - 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 trim = if defer_trim { - let trim = - self.trim_embedded_stage_session_local(session_key, target_token_count)?; - native_mtp_counters.observe_deferred_reject_trim(trim.local_ms); - trim - } else { - self.trim_embedded_stage_session( - request, - downstream, - session_key, - request_id, - session_id, - target_token_count, - )? - }; - trim_control = Some(trim); - } - } - *decode_stage0_compute_ms += verify.stats.stage0_compute_ms; - *decode_runtime_lock_wait_ms += verify.stats.runtime_lock_wait_ms; - *decode_runtime_lock_wait_max_ms = - decode_runtime_lock_wait_max_ms.max(verify.stats.runtime_lock_wait_ms); - *decode_runtime_lock_hold_ms += verify.stats.runtime_lock_hold_ms; - *decode_runtime_lock_hold_max_ms = - decode_runtime_lock_hold_max_ms.max(verify.stats.runtime_lock_hold_ms); - *decode_runtime_lock_acquires += 1; - *decode_forward_activation_encode_ms += verify.stats.activation_encode_ms; - *decode_output_activation_bytes = - decode_output_activation_bytes.saturating_add(verify.stats.output_activation_bytes); - *decode_forward_activation_bytes = - decode_forward_activation_bytes.saturating_add(verify.stats.forward_activation_bytes); - *decode_forward_write_ms += verify.stats.forward_write_ms; - *decode_downstream_wait_ms += verify.stats.downstream_wait_ms; - - if let Some(batched_token_timer) = batched_token_timer { - let mut token_attrs = self.openai_attrs(request.ids); - token_attrs.insert( - "llama_stage.decode_step".to_string(), - serde_json::json!(decode_step), - ); - token_attrs.insert( - "llama_stage.message_kind".to_string(), - serde_json::json!("VerifySpan"), - ); - token_attrs.insert( - "llama_stage.native_mtp.batched_verification".to_string(), - serde_json::json!(true), - ); - token_attrs.insert( - "llama_stage.native_mtp.verification".to_string(), - serde_json::json!(native_mtp_decision.label()), - ); - token_attrs.insert( - "llama_stage.native_mtp.verify_elapsed_ms".to_string(), - serde_json::json!(verify.elapsed_ms), - ); - token_attrs.insert( - "llama_stage.native_mtp.draft_tokens".to_string(), - serde_json::json!(native_mtp_draft_tokens), - ); - token_attrs.insert( - "llama_stage.native_mtp.pending_origin".to_string(), - serde_json::json!(native_mtp_draft_origin.label()), - ); - token_attrs.insert( - "llama_stage.native_mtp.target_token".to_string(), - serde_json::json!(target_token), - ); - token_attrs.insert( - "llama_stage.native_mtp.accepted_count".to_string(), - serde_json::json!(span.accepted_count), - ); - token_attrs.insert( - "llama_stage.native_mtp.verify_next_draft_available".to_string(), - serde_json::json!(verify_next_mtp_draft_available), - ); - token_attrs.insert( - "llama_stage.native_mtp.verify_next_draft_adopted".to_string(), - serde_json::json!(verify_next_mtp_draft_adopted), - ); - if let Some(next_draft) = verify_next_mtp_draft.as_ref() { - token_attrs.insert( - "llama_stage.native_mtp.verify_next_draft_tokens".to_string(), - serde_json::json!(next_draft.tokens), - ); - token_attrs.insert( - "llama_stage.native_mtp.verify_next_draft_compute_us".to_string(), - serde_json::json!(next_draft.proposal_compute_us), - ); - } - token_attrs.insert( - "llama_stage.native_mtp.consumed_positions".to_string(), - serde_json::json!(consumed_positions), - ); - token_attrs.insert( - "llama_stage.native_mtp.committed_positions".to_string(), - serde_json::json!(committed_positions), - ); - token_attrs.insert( - "llama_stage.native_mtp.reject_cooldown_tokens".to_string(), - serde_json::json!(native_mtp_options.reject_cooldown_tokens), - ); - token_attrs.insert( - "llama_stage.native_mtp.reject_cooldown_remaining".to_string(), - serde_json::json!(*native_mtp_reject_cooldown_remaining), - ); - token_attrs.insert( - "llama_stage.native_mtp.defer_reject_trim".to_string(), - serde_json::json!(native_mtp_options.defer_reject_trim), - ); - if let Some(trim) = trim_control.as_ref() { - token_attrs.insert( - "llama_stage.native_mtp.trim_ms".to_string(), - serde_json::json!(trim.elapsed_ms), - ); - token_attrs.insert( - "llama_stage.native_mtp.trim_local_ms".to_string(), - serde_json::json!(trim.local_ms), - ); - token_attrs.insert( - "llama_stage.native_mtp.trim_downstream_write_ms".to_string(), - serde_json::json!(trim.downstream_write_ms), - ); - token_attrs.insert( - "llama_stage.native_mtp.trim_downstream_wait_ms".to_string(), - serde_json::json!(trim.downstream_wait_ms), - ); - } - token_attrs.insert( - "llama_stage.stage0_compute_ms".to_string(), - serde_json::json!(verify.stats.stage0_compute_ms), - ); - token_attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - serde_json::json!(verify.stats.runtime_lock_wait_ms), - ); - token_attrs.insert( - "llama_stage.runtime_lock_hold_ms".to_string(), - serde_json::json!(verify.stats.runtime_lock_hold_ms), - ); - token_attrs.insert( - "llama_stage.activation_encode_ms".to_string(), - serde_json::json!(verify.stats.activation_encode_ms), - ); - token_attrs.insert( - "llama_stage.forward_write_ms".to_string(), - serde_json::json!(verify.stats.forward_write_ms), - ); - token_attrs.insert( - "llama_stage.downstream_wait_ms".to_string(), - serde_json::json!(verify.stats.downstream_wait_ms), - ); - token_attrs.insert( - "llama_stage.output_activation_bytes".to_string(), - serde_json::json!(verify.stats.output_activation_bytes), - ); - token_attrs.insert( - "llama_stage.forward_activation_bytes".to_string(), - serde_json::json!(verify.stats.forward_activation_bytes), - ); - self.emit_openai_phase( - "stage.openai_native_mtp_verify", - batched_token_timer, - token_attrs, - ); - } - - if reached_stop { - return Ok(BatchedVerifyControl::ReachedStop); - } - Ok(BatchedVerifyControl::Continue) - } -} diff --git a/crates/skippy-server/src/frontend/native_mtp/decode.rs b/crates/skippy-server/src/frontend/native_mtp/decode.rs index 30e98a0c7c..4c147a4be8 100644 --- a/crates/skippy-server/src/frontend/native_mtp/decode.rs +++ b/crates/skippy-server/src/frontend/native_mtp/decode.rs @@ -2,51 +2,118 @@ use std::collections::BTreeMap; use serde_json::{Value, json}; -use super::{ - NativeMtpDraftOrigin, native_mtp_batched_verify_enabled, native_mtp_defer_reject_trim_enabled, - native_mtp_reject_cooldown_tokens, native_mtp_suppress_cooldown_draft_limit, - native_mtp_suppress_cooldown_drafts_enabled, -}; +use super::{NativeMtpDraftOrigin, NativeMtpHybridProposal}; +use crate::frontend::SpeculativeDecodeConfig; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(in crate::frontend) struct NativeMtpDecodeOptions { - pub(in crate::frontend) batched_verify: bool, pub(in crate::frontend) max_draft_tokens: usize, pub(in crate::frontend) min_draft_tokens: usize, pub(in crate::frontend) reject_cooldown_tokens: usize, - 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_initial_extension_tokens: usize, + pub(in crate::frontend) ngram_max_proposal_tokens: usize, + pub(in crate::frontend) ngram_tail_backoff_proposals: usize, + pub(in crate::frontend) verify_window_min_tokens: usize, + pub(in crate::frontend) verify_window_max_tokens: usize, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(in crate::frontend) enum NativeMtpTrimAction { + None, + FullSession, +} + +pub(in crate::frontend) fn native_mtp_trim_action( + committed_positions: usize, + consumed_positions: usize, +) -> NativeMtpTrimAction { + if committed_positions == consumed_positions { + NativeMtpTrimAction::None + } else { + NativeMtpTrimAction::FullSession + } } impl NativeMtpDecodeOptions { - pub(in crate::frontend) fn from_env() -> Self { + pub(in crate::frontend) fn from_config(config: &SpeculativeDecodeConfig) -> Self { Self { - batched_verify: native_mtp_batched_verify_enabled(), - max_draft_tokens: 1, - min_draft_tokens: 0, - reject_cooldown_tokens: native_mtp_reject_cooldown_tokens(), - 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(), + max_draft_tokens: config.native_mtp.max_draft_tokens.max(1), + min_draft_tokens: config + .native_mtp + .min_draft_tokens + .min(config.native_mtp.max_draft_tokens.max(1)), + reject_cooldown_tokens: config.native_mtp.reject_cooldown_tokens, + suppress_cooldown_drafts: config.native_mtp.suppress_cooldown_drafts, + suppress_cooldown_draft_limit: config.native_mtp.suppress_cooldown_draft_limit, + ngram_hybrid: config.extension.is_some() && config.ngram.is_some(), + ngram_size: config.ngram.as_ref().map_or(0, |ngram| ngram.min_ngram), + ngram_initial_extension_tokens: config + .extension + .as_ref() + .map_or(0, |extension| extension.initial_tokens), + ngram_max_proposal_tokens: config + .extension + .as_ref() + .map_or(0, |extension| extension.max_tokens), + ngram_tail_backoff_proposals: config + .extension + .as_ref() + .map_or(0, |extension| extension.tail_backoff_proposals), + verify_window_min_tokens: config.verify_window.min_tokens.max(1), + verify_window_max_tokens: config.verify_window.max_tokens.max(1), } } - pub(in crate::frontend) fn with_window( - mut self, - max_draft_tokens: usize, - min_draft_tokens: usize, - ) -> Self { - self.max_draft_tokens = max_draft_tokens.max(1); - self.min_draft_tokens = min_draft_tokens.min(self.max_draft_tokens); - self + pub(in crate::frontend) fn verify_window_bounds(self) -> (usize, usize) { + let min = self.verify_window_min_tokens.max(1); + (min, self.verify_window_max_tokens.max(min)) } } -#[derive(Debug, Default)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(in crate::frontend) struct AdaptiveVerifyWindow { + min_tokens: usize, + max_tokens: usize, + current_tokens: usize, +} + +impl AdaptiveVerifyWindow { + pub(in crate::frontend) fn new(options: NativeMtpDecodeOptions) -> Self { + let (min_tokens, max_tokens) = options.verify_window_bounds(); + Self { + min_tokens, + max_tokens, + current_tokens: max_tokens.min(2).max(min_tokens), + } + } + + pub(in crate::frontend) fn width(self, available_tokens: usize) -> usize { + self.current_tokens.min(available_tokens) + } + + pub(in crate::frontend) fn observe(&mut self, full_accept: bool) -> bool { + let previous = self.current_tokens; + if full_accept { + self.current_tokens = self.current_tokens.saturating_add(1).min(self.max_tokens); + } else { + self.current_tokens = self.current_tokens.saturating_sub(1).max(self.min_tokens); + } + self.current_tokens != previous + } + + pub(in crate::frontend) fn current_tokens(self) -> usize { + self.current_tokens + } +} + +#[derive(Clone, Copy, Debug, Default)] pub(in crate::frontend) struct NativeMtpDecodeCounters { suppressed_cooldown_draft_count: usize, - batched_verification_count: usize, + verify_window_verification_count: usize, initial_serial_verification_count: usize, initial_serial_accepted_count: usize, serial_after_gap_verification_count: usize, @@ -55,25 +122,42 @@ pub(in crate::frontend) struct NativeMtpDecodeCounters { verify_next_accepted_count: usize, verify_next_draft_available_count: usize, verify_next_draft_adopted_count: usize, - deferred_reject_trim_count: usize, - deferred_reject_trim_local_ms: f64, + hybrid_native_prefix_available_count: usize, + hybrid_ngram_continuation_available_count: usize, + hybrid_ngram_mtp_prefix_agreement_count: usize, + hybrid_ngram_mtp_prefix_disagreement_count: usize, + hybrid_proposal_token_count: usize, + hybrid_accepted_token_count: usize, + hybrid_accepted_tail_token_count: usize, + hybrid_native_mtp_token_count: usize, + hybrid_ngram_token_count: usize, + hybrid_pure_ngram_proposal_count: usize, + hybrid_accepted_native_mtp_token_count: usize, + hybrid_ngram_tail_rejection_count: usize, + hybrid_ngram_sidecar_backoff_count: usize, + adaptive_verify_window_count: usize, + adaptive_verify_window_width_sum: usize, + adaptive_verify_window_width_min: usize, + adaptive_verify_window_width_max: usize, + adaptive_verify_window_grow_count: usize, + adaptive_verify_window_shrink_count: usize, } impl NativeMtpDecodeCounters { - pub(in crate::frontend) fn batched_verification_count(&self) -> usize { - self.batched_verification_count + pub(in crate::frontend) fn verify_window_verification_count(&self) -> usize { + self.verify_window_verification_count } pub(in crate::frontend) fn observe_suppressed_cooldown_draft(&mut self) { self.suppressed_cooldown_draft_count += 1; } - pub(in crate::frontend) fn observe_batched_verification( + pub(in crate::frontend) fn observe_verify_window_verification( &mut self, origin: NativeMtpDraftOrigin, accepted: bool, ) { - self.batched_verification_count += 1; + self.verify_window_verification_count += 1; match origin { NativeMtpDraftOrigin::InitialSerial => { self.initial_serial_verification_count += 1; @@ -109,9 +193,51 @@ impl NativeMtpDecodeCounters { } } - pub(in crate::frontend) fn observe_deferred_reject_trim(&mut self, local_ms: f64) { - self.deferred_reject_trim_count += 1; - self.deferred_reject_trim_local_ms += local_ms; + pub(in crate::frontend) fn observe_hybrid_proposal( + &mut self, + proposal: &NativeMtpHybridProposal, + accepted_token_count: usize, + ) { + self.hybrid_native_prefix_available_count += + usize::from(proposal.native_mtp_token_count() > 0); + self.hybrid_ngram_continuation_available_count += + usize::from(proposal.ngram_span_available()); + self.hybrid_ngram_mtp_prefix_agreement_count += + usize::from(proposal.ngram_mtp_prefix_agreed()); + self.hybrid_ngram_mtp_prefix_disagreement_count += + usize::from(proposal.ngram_mtp_prefix_disagreed()); + self.hybrid_proposal_token_count += proposal.tokens().len(); + self.hybrid_accepted_token_count += accepted_token_count; + self.hybrid_accepted_tail_token_count += + accepted_token_count.saturating_sub(proposal.native_mtp_token_count()); + self.hybrid_accepted_native_mtp_token_count += + accepted_token_count.min(proposal.native_mtp_token_count()); + self.hybrid_native_mtp_token_count += proposal.native_mtp_token_count(); + self.hybrid_ngram_token_count += proposal.ngram_token_count(); + self.hybrid_pure_ngram_proposal_count += usize::from(proposal.is_pure_ngram()); + } + + pub(in crate::frontend) fn observe_ngram_tail_rejection(&mut self) { + self.hybrid_ngram_tail_rejection_count += 1; + self.hybrid_ngram_sidecar_backoff_count += 1; + } + + pub(in crate::frontend) fn observe_adaptive_verify_window( + &mut self, + width: usize, + previous_width: usize, + next_width: usize, + ) { + self.adaptive_verify_window_count += 1; + self.adaptive_verify_window_width_sum += width; + self.adaptive_verify_window_width_min = if self.adaptive_verify_window_width_min == 0 { + width + } else { + self.adaptive_verify_window_width_min.min(width) + }; + self.adaptive_verify_window_width_max = self.adaptive_verify_window_width_max.max(width); + self.adaptive_verify_window_grow_count += usize::from(next_width > previous_width); + self.adaptive_verify_window_shrink_count += usize::from(next_width < previous_width); } pub(in crate::frontend) fn insert_summary_attrs( @@ -131,10 +257,6 @@ impl NativeMtpDecodeCounters { "llama_stage.native_mtp.min_draft_tokens".to_string(), json!(options.min_draft_tokens), ); - attrs.insert( - "llama_stage.native_mtp.defer_reject_trim".to_string(), - json!(options.defer_reject_trim), - ); attrs.insert( "llama_stage.native_mtp.suppress_cooldown_drafts".to_string(), json!(options.suppress_cooldown_drafts), @@ -143,13 +265,41 @@ 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_initial_extension_tokens".to_string(), + json!(options.ngram_initial_extension_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.ngram_max_proposal_tokens".to_string(), + json!(options.ngram_max_proposal_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.ngram_tail_backoff_proposals".to_string(), + json!(options.ngram_tail_backoff_proposals), + ); + attrs.insert( + "llama_stage.native_mtp.verify_window_min_tokens".to_string(), + json!(options.verify_window_min_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.verify_window_max_tokens".to_string(), + json!(options.verify_window_max_tokens), + ); attrs.insert( "llama_stage.native_mtp.suppressed_cooldown_draft_count".to_string(), json!(self.suppressed_cooldown_draft_count), ); attrs.insert( - "llama_stage.native_mtp.batched_verification_count".to_string(), - json!(self.batched_verification_count), + "llama_stage.native_mtp.verify_window_verification_count".to_string(), + json!(self.verify_window_verification_count), ); attrs.insert( "llama_stage.native_mtp.initial_serial_verification_count".to_string(), @@ -184,47 +334,311 @@ impl NativeMtpDecodeCounters { json!(self.verify_next_draft_adopted_count), ); attrs.insert( - "llama_stage.native_mtp.deferred_reject_trim_count".to_string(), - json!(self.deferred_reject_trim_count), + "llama_stage.native_mtp.hybrid_native_prefix_available_count".to_string(), + json!(self.hybrid_native_prefix_available_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_continuation_available_count".to_string(), + json!(self.hybrid_ngram_continuation_available_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_mtp_prefix_agreement_count".to_string(), + json!(self.hybrid_ngram_mtp_prefix_agreement_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_mtp_prefix_disagreement_count".to_string(), + json!(self.hybrid_ngram_mtp_prefix_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), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_native_mtp_token_count".to_string(), + json!(self.hybrid_native_mtp_token_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_token_count".to_string(), + json!(self.hybrid_ngram_token_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_accepted_native_mtp_token_count".to_string(), + json!(self.hybrid_accepted_native_mtp_token_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_tail_rejection_count".to_string(), + json!(self.hybrid_ngram_tail_rejection_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_ngram_sidecar_backoff_count".to_string(), + json!(self.hybrid_ngram_sidecar_backoff_count), + ); + attrs.insert( + "llama_stage.native_mtp.hybrid_pure_ngram_proposal_count".to_string(), + json!(self.hybrid_pure_ngram_proposal_count), + ); + attrs.insert( + "llama_stage.native_mtp.adaptive_verify_window_count".to_string(), + json!(self.adaptive_verify_window_count), + ); + attrs.insert( + "llama_stage.native_mtp.adaptive_verify_window_width_sum".to_string(), + json!(self.adaptive_verify_window_width_sum), + ); + attrs.insert( + "llama_stage.native_mtp.adaptive_verify_window_width_min".to_string(), + json!(self.adaptive_verify_window_width_min), + ); + attrs.insert( + "llama_stage.native_mtp.adaptive_verify_window_width_max".to_string(), + json!(self.adaptive_verify_window_width_max), ); attrs.insert( - "llama_stage.native_mtp.deferred_reject_trim_local_ms".to_string(), - json!(self.deferred_reject_trim_local_ms), + "llama_stage.native_mtp.adaptive_verify_window_grow_count".to_string(), + json!(self.adaptive_verify_window_grow_count), + ); + attrs.insert( + "llama_stage.native_mtp.adaptive_verify_window_shrink_count".to_string(), + json!(self.adaptive_verify_window_shrink_count), + ); + } + + fn insert_response_timings(&self, timings: &mut BTreeMap) { + timings.insert( + "native_mtp_verify_window_verifications".to_string(), + json!(self.verify_window_verification_count), + ); + timings.insert( + "native_mtp_hybrid_native_prefix_available".to_string(), + json!(self.hybrid_native_prefix_available_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_continuation_available".to_string(), + json!(self.hybrid_ngram_continuation_available_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_mtp_prefix_agreements".to_string(), + json!(self.hybrid_ngram_mtp_prefix_agreement_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_mtp_prefix_disagreements".to_string(), + json!(self.hybrid_ngram_mtp_prefix_disagreement_count), + ); + timings.insert( + "native_mtp_hybrid_proposed_tokens".to_string(), + json!(self.hybrid_proposal_token_count), + ); + timings.insert( + "native_mtp_hybrid_accepted_tokens".to_string(), + json!(self.hybrid_accepted_token_count), + ); + timings.insert( + "native_mtp_hybrid_accepted_tail_tokens".to_string(), + json!(self.hybrid_accepted_tail_token_count), + ); + timings.insert( + "native_mtp_hybrid_native_tokens".to_string(), + json!(self.hybrid_native_mtp_token_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_tokens".to_string(), + json!(self.hybrid_ngram_token_count), + ); + timings.insert( + "native_mtp_hybrid_accepted_native_tokens".to_string(), + json!(self.hybrid_accepted_native_mtp_token_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_tail_rejections".to_string(), + json!(self.hybrid_ngram_tail_rejection_count), + ); + timings.insert( + "native_mtp_hybrid_ngram_sidecar_backoffs".to_string(), + json!(self.hybrid_ngram_sidecar_backoff_count), + ); + timings.insert( + "native_mtp_hybrid_pure_ngram_proposals".to_string(), + json!(self.hybrid_pure_ngram_proposal_count), + ); + timings.insert( + "native_mtp_adaptive_verify_windows".to_string(), + json!(self.adaptive_verify_window_count), + ); + timings.insert( + "native_mtp_adaptive_verify_window_width_sum".to_string(), + json!(self.adaptive_verify_window_width_sum), + ); + timings.insert( + "native_mtp_adaptive_verify_window_width_min".to_string(), + json!(self.adaptive_verify_window_width_min), + ); + timings.insert( + "native_mtp_adaptive_verify_window_width_max".to_string(), + json!(self.adaptive_verify_window_width_max), + ); + timings.insert( + "native_mtp_adaptive_verify_window_grows".to_string(), + json!(self.adaptive_verify_window_grow_count), + ); + timings.insert( + "native_mtp_adaptive_verify_window_shrinks".to_string(), + json!(self.adaptive_verify_window_shrink_count), ); } } +#[derive(Clone, Copy, Debug)] +pub(in crate::frontend) struct NativeMtpDecodeTelemetry { + options: NativeMtpDecodeOptions, + counters: NativeMtpDecodeCounters, +} + +impl NativeMtpDecodeTelemetry { + pub(in crate::frontend) fn new( + options: NativeMtpDecodeOptions, + counters: NativeMtpDecodeCounters, + ) -> Self { + Self { options, counters } + } + + pub(in crate::frontend) fn composite_proposal_totals(self) -> Option<(u64, u64)> { + (self.counters.hybrid_proposal_token_count > 0).then_some(( + self.counters.hybrid_proposal_token_count as u64, + self.counters.hybrid_accepted_token_count as u64, + )) + } + + pub(in crate::frontend) fn insert_response_timings( + self, + timings: &mut BTreeMap, + ) { + timings.insert( + "native_mtp_ngram_hybrid_enabled".to_string(), + json!(self.options.ngram_hybrid), + ); + timings.insert( + "native_mtp_ngram_size".to_string(), + json!(self.options.ngram_size), + ); + timings.insert( + "native_mtp_ngram_initial_extension_tokens".to_string(), + json!(self.options.ngram_initial_extension_tokens), + ); + timings.insert( + "native_mtp_ngram_max_proposal_tokens".to_string(), + json!(self.options.ngram_max_proposal_tokens), + ); + timings.insert( + "native_mtp_ngram_tail_backoff_proposals".to_string(), + json!(self.options.ngram_tail_backoff_proposals), + ); + timings.insert( + "native_mtp_verify_window_min_tokens".to_string(), + json!(self.options.verify_window_min_tokens), + ); + timings.insert( + "native_mtp_verify_window_max_tokens".to_string(), + json!(self.options.verify_window_max_tokens), + ); + self.counters.insert_response_timings(timings); + } +} + #[cfg(test)] mod tests { + use super::super::CompositeProposalProvider; use super::*; + fn options() -> NativeMtpDecodeOptions { + NativeMtpDecodeOptions { + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 2, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 2, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, + } + } + + fn composite_proposal() -> NativeMtpHybridProposal { + CompositeProposalProvider::from_options(options()).propose( + &[9], + &[1, 2, 3, 9, 1, 2, 3, 9, 1, 2, 3], + 4, + ) + } + #[test] - fn counters_track_batched_verification_by_origin() { + fn decode_options_preserve_configured_initial_extension_width() { + let config = SpeculativeDecodeConfig { + extension: Some(crate::frontend::NgramExtensionConfig { + initial_tokens: 3, + max_tokens: 7, + tail_backoff_proposals: 2, + }), + ngram: Some(crate::frontend::NgramProposalConfig { + kind: crate::frontend::NgramProposerKind::Cache, + min_ngram: 2, + max_ngram: 4, + max_proposal_tokens: 7, + }), + ..SpeculativeDecodeConfig::default() + }; + + let options = NativeMtpDecodeOptions::from_config(&config); + + assert_eq!(options.ngram_initial_extension_tokens, 3); + assert_eq!(options.ngram_max_proposal_tokens, 7); + } + + #[test] + fn counters_track_verify_window_verification_by_origin() { let mut counters = NativeMtpDecodeCounters::default(); - counters.observe_batched_verification(NativeMtpDraftOrigin::InitialSerial, true); - counters.observe_batched_verification(NativeMtpDraftOrigin::SerialAfterGap, false); - counters.observe_batched_verification(NativeMtpDraftOrigin::VerifyNext, true); + counters.observe_verify_window_verification(NativeMtpDraftOrigin::InitialSerial, true); + counters.observe_verify_window_verification(NativeMtpDraftOrigin::SerialAfterGap, false); + counters.observe_verify_window_verification(NativeMtpDraftOrigin::VerifyNext, true); counters.observe_verify_next_draft(true, false); counters.observe_verify_next_draft(true, true); counters.observe_suppressed_cooldown_draft(); - counters.observe_deferred_reject_trim(1.25); + counters.observe_hybrid_proposal(&composite_proposal(), 3); + counters.observe_ngram_tail_rejection(); + counters.observe_adaptive_verify_window(2, 2, 3); let mut attrs = BTreeMap::new(); counters.insert_summary_attrs( &mut attrs, NativeMtpDecodeOptions { - batched_verify: true, max_draft_tokens: 3, min_draft_tokens: 0, reject_cooldown_tokens: 6, - defer_reject_trim: true, suppress_cooldown_drafts: false, suppress_cooldown_draft_limit: 2, + ngram_hybrid: true, + ngram_size: 8, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 6, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, }, ); assert_eq!( - attrs.get("llama_stage.native_mtp.batched_verification_count"), + attrs.get("llama_stage.native_mtp.verify_window_verification_count"), Some(&json!(3)) ); assert_eq!( @@ -256,16 +670,146 @@ mod tests { Some(&json!(6)) ); assert_eq!( - attrs.get("llama_stage.native_mtp.defer_reject_trim"), + attrs.get("llama_stage.native_mtp.ngram_initial_extension_tokens"), + Some(&json!(2)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_accepted_tail_token_count"), + Some(&json!(2)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_accepted_native_mtp_token_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_ngram_mtp_prefix_agreement_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_ngram_mtp_prefix_disagreement_count"), + Some(&json!(0)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_ngram_tail_rejection_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.hybrid_ngram_sidecar_backoff_count"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.adaptive_verify_window_grow_count"), + Some(&json!(1)) + ); + } + + #[test] + fn short_candidate_does_not_look_like_adaptive_window_growth() { + let mut counters = NativeMtpDecodeCounters::default(); + counters.observe_adaptive_verify_window(1, 4, 4); + + let mut attrs = BTreeMap::new(); + counters.insert_summary_attrs(&mut attrs, options()); + + assert_eq!( + attrs.get("llama_stage.native_mtp.adaptive_verify_window_width_sum"), + Some(&json!(1)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.adaptive_verify_window_grow_count"), + Some(&json!(0)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.adaptive_verify_window_shrink_count"), + Some(&json!(0)) + ); + } + + #[test] + fn response_timings_show_hybrid_widening_evidence() { + let mut counters = NativeMtpDecodeCounters::default(); + counters.observe_verify_window_verification(NativeMtpDraftOrigin::InitialSerial, true); + counters.observe_hybrid_proposal(&composite_proposal(), 3); + counters.observe_adaptive_verify_window(2, 2, 3); + let telemetry = NativeMtpDecodeTelemetry::new( + NativeMtpDecodeOptions { + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 8, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 6, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, + }, + counters, + ); + + let mut timings = BTreeMap::new(); + telemetry.insert_response_timings(&mut timings); + + assert_eq!( + timings.get("native_mtp_ngram_hybrid_enabled"), Some(&json!(true)) ); assert_eq!( - attrs.get("llama_stage.native_mtp.deferred_reject_trim_count"), + timings.get("native_mtp_ngram_initial_extension_tokens"), + Some(&json!(2)) + ); + assert_eq!( + timings.get("native_mtp_hybrid_native_tokens"), + Some(&json!(1)) + ); + assert_eq!( + timings.get("native_mtp_hybrid_accepted_tail_tokens"), + Some(&json!(2)) + ); + assert_eq!( + timings.get("native_mtp_hybrid_accepted_native_tokens"), + Some(&json!(1)) + ); + assert_eq!( + timings.get("native_mtp_hybrid_ngram_mtp_prefix_agreements"), Some(&json!(1)) ); assert_eq!( - attrs.get("llama_stage.native_mtp.deferred_reject_trim_local_ms"), - Some(&json!(1.25)) + timings.get("native_mtp_adaptive_verify_window_grows"), + Some(&json!(1)) + ); + } + + #[test] + fn adaptive_verify_window_starts_at_two_then_grows_and_shrinks() { + let mut window = AdaptiveVerifyWindow::new(options()); + + assert_eq!(window.current_tokens(), 2); + assert_eq!(window.width(1), 1); + assert!(window.observe(true)); + assert_eq!(window.current_tokens(), 3); + assert!(window.observe(false)); + assert_eq!(window.current_tokens(), 2); + assert!(window.observe(false)); + assert_eq!(window.current_tokens(), 1); + assert!(!window.observe(false)); + } + + #[test] + fn composite_proposal_totals_include_pure_ngram_candidates() { + let mut counters = NativeMtpDecodeCounters::default(); + let proposal = CompositeProposalProvider::from_options(options()).propose( + &[], + &[0, 0, 2, 3, 9, 1, 7, 8, 2, 3], + 4, + ); + counters.observe_hybrid_proposal(&proposal, 4); + + assert_eq!( + NativeMtpDecodeTelemetry::new(options(), counters).composite_proposal_totals(), + Some((4, 4)) ); } } diff --git a/crates/skippy-server/src/frontend/native_mtp/draft.rs b/crates/skippy-server/src/frontend/native_mtp/draft.rs index 2f57f8d9d5..9ad72bd870 100644 --- a/crates/skippy-server/src/frontend/native_mtp/draft.rs +++ b/crates/skippy-server/src/frontend/native_mtp/draft.rs @@ -1,3 +1,5 @@ +use skippy_protocol::binary::StageNativeMtpDraft; + #[derive(Clone, Debug, PartialEq, Eq)] pub(in crate::frontend) struct NativeMtpDraft { pub(in crate::frontend) tokens: Vec, @@ -5,30 +7,11 @@ pub(in crate::frontend) struct NativeMtpDraft { } impl NativeMtpDraft { - pub(in crate::frontend) fn from_prediction_tokens(tokens: &[i32]) -> Option { - Self::from_sideband(tokens, 1) - } - - pub(in crate::frontend) fn from_verify_prediction_tokens( - tokens: &[i32], - verified_token_count: usize, - ) -> Option { - Self::from_sideband(tokens, verified_token_count) - } - - fn from_sideband(tokens: &[i32], offset: usize) -> Option { - let token_count = usize::try_from(*tokens.get(offset)?).ok()?; - if token_count == 0 { - return None; + pub(in crate::frontend) fn from_stage_draft(draft: StageNativeMtpDraft) -> Self { + Self { + tokens: draft.token_ids, + proposal_compute_us: draft.proposal_compute_us.max(0), } - let start = offset.saturating_add(1); - let end = start.checked_add(token_count)?; - let draft_tokens = tokens.get(start..end)?.to_vec(); - let proposal_compute_us = tokens.get(end).copied().unwrap_or_default(); - Some(Self { - tokens: draft_tokens, - proposal_compute_us: i64::from(proposal_compute_us.max(0)), - }) } } @@ -60,61 +43,37 @@ mod tests { use super::*; #[test] - fn parses_prediction_token_sideband() { + fn converts_typed_stage_draft() { assert_eq!( - NativeMtpDraft::from_prediction_tokens(&[11, 1, 12, 34]), - Some(NativeMtpDraft { + NativeMtpDraft::from_stage_draft(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: 34, + }), + NativeMtpDraft { tokens: vec![12], proposal_compute_us: 34, - }) - ); - assert_eq!( - NativeMtpDraft::from_prediction_tokens(&[11, 2, 34, 35, 567]), - Some(NativeMtpDraft { - tokens: vec![34, 35], - proposal_compute_us: 567, - }) + } ); - assert_eq!(NativeMtpDraft::from_prediction_tokens(&[11]), None); } #[test] - fn parses_verify_prediction_token_sideband_after_verified_tokens() { - assert_eq!( - NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 1, 12, 34], 2), - Some(NativeMtpDraft { - tokens: vec![12], - proposal_compute_us: 34, - }) - ); - assert_eq!( - NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 1, 12, -3], 2), - Some(NativeMtpDraft { - tokens: vec![12], - proposal_compute_us: 0, - }) - ); - assert_eq!( - NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 2, 12, 13, 567], 2), - Some(NativeMtpDraft { - tokens: vec![12, 13], - proposal_compute_us: 567, - }) - ); - assert_eq!( - NativeMtpDraft::from_verify_prediction_tokens(&[10, 11], 2), - None - ); + fn clamps_negative_typed_proposal_time() { + let draft = NativeMtpDraft::from_stage_draft(StageNativeMtpDraft { + token_ids: vec![12], + proposal_compute_us: -3, + }); + + assert_eq!(draft.proposal_compute_us, 0); } #[test] - fn pending_draft_keeps_origin_label() { + fn pending_draft_keeps_origin() { let pending = PendingNativeMtpDraft { tokens: vec![12, 13], origin: NativeMtpDraftOrigin::VerifyNext, }; assert_eq!(pending.tokens, vec![12, 13]); - assert_eq!(pending.origin.label(), "verify_next"); + assert_eq!(pending.origin, NativeMtpDraftOrigin::VerifyNext); } } diff --git a/crates/skippy-server/src/frontend/native_mtp/env.rs b/crates/skippy-server/src/frontend/native_mtp/env.rs deleted file mode 100644 index 54fbb698d0..0000000000 --- a/crates/skippy-server/src/frontend/native_mtp/env.rs +++ /dev/null @@ -1,94 +0,0 @@ -const BATCHED_VERIFY_ENV: &str = "SKIPPY_NATIVE_MTP_BATCHED_VERIFY"; -const REJECT_COOLDOWN_TOKENS_ENV: &str = "SKIPPY_NATIVE_MTP_REJECT_COOLDOWN_TOKENS"; -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"; - -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()) -} - -pub(in crate::frontend) fn native_mtp_reject_cooldown_tokens() -> usize { - parse_usize_env(REJECT_COOLDOWN_TOKENS_ENV, 0) -} - -pub(in crate::frontend) fn native_mtp_defer_reject_trim_enabled() -> bool { - truthy_env(std::env::var(DEFER_REJECT_TRIM_ENV).ok().as_deref()) -} - -pub(in crate::frontend) fn native_mtp_suppress_cooldown_drafts_enabled() -> bool { - truthy_env(std::env::var(SUPPRESS_COOLDOWN_DRAFTS_ENV).ok().as_deref()) -} - -pub(in crate::frontend) fn native_mtp_suppress_cooldown_draft_limit() -> usize { - parse_usize_env(SUPPRESS_COOLDOWN_DRAFT_LIMIT_ENV, 0) -} - -fn native_mtp_batched_verify_enabled_from(value: Option<&str>) -> bool { - !falsey_env(value) -} - -fn truthy_env(value: Option<&str>) -> bool { - matches!( - normalized_env(value).as_deref(), - Some("1" | "true" | "on" | "enable" | "enabled" | "yes") - ) -} - -fn falsey_env(value: Option<&str>) -> bool { - matches!( - normalized_env(value).as_deref(), - Some("0" | "false" | "off" | "disable" | "disabled" | "no") - ) -} - -fn normalized_env(value: Option<&str>) -> Option { - value.map(str::trim).map(str::to_ascii_lowercase) -} - -fn parse_usize_env(name: &str, default: usize) -> usize { - std::env::var(name) - .ok() - .and_then(|value| value.trim().parse::().ok()) - .unwrap_or(default) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn batched_verify_flag_defaults_on_and_accepts_false_values() { - assert!(native_mtp_batched_verify_enabled_from(None)); - assert!(native_mtp_batched_verify_enabled_from(Some("1"))); - assert!(native_mtp_batched_verify_enabled_from(Some("true"))); - assert!(!native_mtp_batched_verify_enabled_from(Some("0"))); - assert!(!native_mtp_batched_verify_enabled_from(Some("false"))); - assert!(!native_mtp_batched_verify_enabled_from(Some(" disabled "))); - } - - #[test] - fn truthy_env_accepts_enabled_aliases_only() { - for value in ["1", "true", " enabled ", "yes", "on"] { - assert!(truthy_env(Some(value)), "{value}"); - } - for value in [ - None, - Some("0"), - Some("false"), - Some("off"), - Some("disabled"), - ] { - assert!(!truthy_env(value), "{value:?}"); - } - } - - #[test] - fn numeric_options_default_when_absent() { - assert_eq!(parse_usize_env("SKIPPY_TEST_MISSING_REJECT_COOLDOWN", 0), 0); - assert_eq!( - parse_usize_env("SKIPPY_TEST_MISSING_SUPPRESS_COOLDOWN_LIMIT", 0), - 0 - ); - } -} 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..ccdddbfe80 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/hybrid.rs @@ -0,0 +1,779 @@ +use std::collections::VecDeque; + +use openai_frontend::{OpenAiError, OpenAiResult}; + +use super::NativeMtpDecodeOptions; +use crate::frontend::speculative::{CachedNgramProposer, propose_ngram_tokens}; + +const MIN_NGRAM_EXTENSION_TOKENS: usize = 2; + +/// Builds one speculative candidate from a native-MTP prefix and an optional +/// N-gram continuation. The N-gram proposal must independently predict the +/// native-MTP prefix before its remaining tokens may extend that prefix. +#[derive(Debug, Clone, Copy)] +pub(in crate::frontend) struct CompositeProposalProvider { + enabled: bool, + ngram_size: usize, + max_proposal_tokens: usize, +} + +impl CompositeProposalProvider { + pub(in crate::frontend) fn from_options(options: NativeMtpDecodeOptions) -> Self { + Self { + enabled: options.ngram_hybrid, + ngram_size: options.ngram_size, + max_proposal_tokens: options.ngram_max_proposal_tokens, + } + } +} + +impl CompositeProposalProvider { + #[cfg(test)] + pub(in crate::frontend) fn propose( + &self, + native_mtp_tokens: &[i32], + context_tokens: &[i32], + max_proposal_tokens: usize, + ) -> NativeMtpHybridProposal { + self.propose_with_ngram_extension( + native_mtp_tokens, + context_tokens, + max_proposal_tokens, + max_proposal_tokens, + None, + ) + .expect("llama.cpp N-gram proposal succeeds") + } + + pub(in crate::frontend) fn propose_with_ngram_extension( + &self, + native_mtp_tokens: &[i32], + context_tokens: &[i32], + max_proposal_tokens: usize, + max_ngram_extension_tokens: usize, + cached_ngram_proposer: Option<&mut CachedNgramProposer>, + ) -> OpenAiResult { + let native_mtp_tokens = + &native_mtp_tokens[..native_mtp_tokens.len().min(max_proposal_tokens)]; + if !self.enabled || self.max_proposal_tokens == 0 || max_ngram_extension_tokens == 0 { + return Ok(NativeMtpHybridProposal::from_native_mtp_tokens( + native_mtp_tokens.to_vec(), + )); + } + + let ngram_limit = max_proposal_tokens + .saturating_sub(native_mtp_tokens.len()) + .min(self.max_proposal_tokens) + .min(max_ngram_extension_tokens); + let ( + ngram_tokens, + ngram_span_available, + ngram_mtp_prefix_agreed, + ngram_mtp_prefix_disagreed, + ) = if let Some(cache) = cached_ngram_proposer { + // The cache sees only committed target history. Native MTP is + // an optional read-only continuation, so this returns the + // sidecar tail directly rather than trying to re-predict it. + let tail = cache.propose(context_tokens, native_mtp_tokens, ngram_limit)?; + let available = !tail.is_empty(); + (tail, available, false, false) + } else { + let candidates = if native_mtp_tokens.is_empty() { + propose_ngram_tokens(context_tokens, self.ngram_size, ngram_limit)? + } else { + // Preserve the #875 anchor rule for native MTP: the most + // recent earlier occurrence of the configured context + // N-gram supplies a complete span. Its leading tokens + // must agree with MTP before its remaining tokens may + // become the sidecar tail below. + propose_ngram_tokens( + context_tokens, + self.ngram_size, + native_mtp_tokens.len().saturating_add(ngram_limit), + )? + }; + if native_mtp_tokens.is_empty() { + let available = !candidates.is_empty(); + (candidates, available, false, false) + } else { + let available = candidates.len() > native_mtp_tokens.len(); + let agreed = available && candidates.starts_with(native_mtp_tokens); + let disagreed = available && !agreed; + let tail = if agreed { + candidates[native_mtp_tokens.len()..].to_vec() + } else { + Vec::new() + }; + (tail, available, agreed, disagreed) + } + }; + let ngram_tokens = if ngram_tokens.len() >= MIN_NGRAM_EXTENSION_TOKENS { + ngram_tokens + } else { + Vec::new() + }; + let mut tokens = native_mtp_tokens.to_vec(); + tokens.extend(ngram_tokens); + Ok(NativeMtpHybridProposal { + native_mtp_token_count: native_mtp_tokens.len(), + ngram_token_count: tokens.len().saturating_sub(native_mtp_tokens.len()), + tokens, + ngram_span_available, + ngram_mtp_prefix_agreed, + ngram_mtp_prefix_disagreed, + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpHybridProposal { + tokens: Vec, + native_mtp_token_count: usize, + ngram_token_count: usize, + ngram_span_available: bool, + ngram_mtp_prefix_agreed: bool, + ngram_mtp_prefix_disagreed: bool, +} + +impl NativeMtpHybridProposal { + pub(in crate::frontend) fn from_parts( + tokens: Vec, + native_mtp_token_count: usize, + ngram_span_available: bool, + ) -> Self { + let native_mtp_token_count = native_mtp_token_count.min(tokens.len()); + Self { + ngram_token_count: tokens.len().saturating_sub(native_mtp_token_count), + native_mtp_token_count, + tokens, + ngram_span_available, + ngram_mtp_prefix_agreed: native_mtp_token_count > 0 && ngram_span_available, + ngram_mtp_prefix_disagreed: false, + } + } + + pub(in crate::frontend) fn from_native_mtp_tokens(tokens: Vec) -> Self { + Self::from_parts(tokens, usize::MAX, false) + } + + pub(in crate::frontend) fn tokens(&self) -> &[i32] { + &self.tokens + } + + pub(in crate::frontend) fn native_mtp_token_count(&self) -> usize { + self.native_mtp_token_count + } + + pub(in crate::frontend) fn ngram_token_count(&self) -> usize { + self.ngram_token_count + } + + pub(in crate::frontend) fn is_pure_ngram(&self) -> bool { + self.native_mtp_token_count == 0 && self.ngram_token_count > 0 + } + + pub(in crate::frontend) fn ngram_span_available(&self) -> bool { + self.ngram_span_available + } + + pub(in crate::frontend) fn ngram_mtp_prefix_agreed(&self) -> bool { + self.ngram_mtp_prefix_agreed + } + + pub(in crate::frontend) fn ngram_mtp_prefix_disagreed(&self) -> bool { + self.ngram_mtp_prefix_disagreed + } + + /// A tail mismatch is not evidence that the native MTP prefix was bad. + /// Keep the native reject cooldown scoped to mismatches inside that prefix. + pub(in crate::frontend) fn native_mtp_prefix_rejected( + &self, + accepted_proposal_tokens: usize, + ) -> bool { + accepted_proposal_tokens < self.native_mtp_token_count + } + + /// A mismatch after the MTP prefix belongs to the optional N-gram sidecar. + /// It must not penalize native MTP, but it should temporarily stop extending + /// healthy MTP candidates with another unprofitable tail. + pub(in crate::frontend) fn ngram_tail_rejected(&self, accepted_proposal_tokens: usize) -> bool { + self.native_mtp_token_count > 0 + && self.ngram_token_count > 0 + && accepted_proposal_tokens >= self.native_mtp_token_count + && accepted_proposal_tokens < self.tokens.len() + } + + /// The first pipelined verify may consume a wider prefix, while a later + /// in-flight window safely consumes the remaining suffix. Reserve one + /// optimistic target plus one remaining candidate so depth actually buys + /// overlap for short MTP-plus-N-gram spans. + pub(in crate::frontend) fn parallel_verify_width( + &self, + adaptive_verify_width: usize, + pipeline_depth: usize, + ) -> Option { + if pipeline_depth < 2 || self.tokens.len() < 3 { + return None; + } + Some( + adaptive_verify_width + .min(self.tokens.len().saturating_sub(2)) + .max(1), + ) + } +} + +/// Request-local adaptive control for the N-gram extension. +/// +/// MTP is already a useful proposer on its own, so the N-gram sidecar starts +/// with the smallest useful tail. It widens only after the target accepted an +/// entire tail and resets after a mismatch. This keeps one accidental repeated +/// context from turning a healthy MTP candidate into an expensive long span. +/// Pure N-gram proposals retain the configured maximum because they have no +/// native prefix to protect. +#[derive(Debug)] +pub(in crate::frontend) struct NgramSidecarController { + remaining_proposals: usize, + initial_extension_tokens: usize, + current_extension_tokens: usize, + max_extension_tokens: usize, +} + +impl NgramSidecarController { + pub(in crate::frontend) fn new( + initial_extension_tokens: usize, + max_extension_tokens: usize, + ) -> Self { + let initial_extension_tokens = initial_extension_tokens.min(max_extension_tokens); + Self { + remaining_proposals: 0, + initial_extension_tokens, + current_extension_tokens: initial_extension_tokens, + max_extension_tokens, + } + } + + /// Returns the N-gram token budget for this proposal. A zero budget means + /// use the native MTP prefix alone while the sidecar cools down. + pub(in crate::frontend) fn extension_limit( + &mut self, + native_mtp_tokens: &[i32], + available_tokens: usize, + ) -> usize { + if native_mtp_tokens.is_empty() { + return available_tokens.min(self.max_extension_tokens); + } + if self.remaining_proposals > 0 { + self.remaining_proposals -= 1; + return 0; + } + self.current_extension_tokens.min(available_tokens) + } + + /// Applies the result of a completed composite candidate. Returns true + /// only when a rejected tail entered cooldown, so existing rejection + /// telemetry remains a direct count of sidecar failures. + pub(in crate::frontend) fn observe_tail_outcome( + &mut self, + proposal: &NativeMtpHybridProposal, + accepted_proposal_tokens: usize, + cooldown_proposals: usize, + ) -> bool { + if proposal.native_mtp_token_count() == 0 || proposal.ngram_token_count() == 0 { + return false; + } + if accepted_proposal_tokens >= proposal.tokens().len() { + self.current_extension_tokens = self + .current_extension_tokens + .saturating_add(1) + .min(self.max_extension_tokens); + return false; + } + if !proposal.ngram_tail_rejected(accepted_proposal_tokens) { + return false; + } + self.current_extension_tokens = self.initial_extension_tokens; + self.remaining_proposals = cooldown_proposals; + cooldown_proposals > 0 + } + + #[cfg(test)] + pub(in crate::frontend) fn remaining_proposals(&self) -> usize { + self.remaining_proposals + } + + #[cfg(test)] + pub(in crate::frontend) fn current_extension_tokens(&self) -> usize { + self.current_extension_tokens + } +} + +/// Holds the unverified portion of a composite proposal. A fully accepted +/// verify window may advance one additional target token, so that token is +/// removed from the buffer only when it agrees with the buffered candidate. +#[derive(Debug)] +pub(in crate::frontend) struct BufferedCompositeProposal { + proposal: NativeMtpHybridProposal, + remaining_tokens: VecDeque, + accepted_tokens: usize, +} + +impl BufferedCompositeProposal { + pub(in crate::frontend) fn new(proposal: NativeMtpHybridProposal) -> Self { + Self { + remaining_tokens: proposal.tokens.iter().copied().collect(), + proposal, + accepted_tokens: 0, + } + } + + pub(in crate::frontend) fn proposal(&self) -> &NativeMtpHybridProposal { + &self.proposal + } + + pub(in crate::frontend) fn verify_tokens(&self, width: usize) -> Vec { + self.remaining_tokens.iter().copied().take(width).collect() + } + + pub(in crate::frontend) fn expected_free_target(&self, width: usize) -> Option { + self.remaining_tokens.get(width).copied() + } + + pub(in crate::frontend) fn remaining_len(&self) -> usize { + self.remaining_tokens.len() + } + + pub(in crate::frontend) fn is_empty(&self) -> bool { + self.remaining_tokens.is_empty() + } + + pub(in crate::frontend) fn accepted_tokens(&self) -> usize { + self.accepted_tokens + } + + pub(in crate::frontend) fn native_mtp_prefix_rejected_after( + &self, + accepted_window_tokens: usize, + ) -> bool { + self.proposal + .native_mtp_prefix_rejected(self.accepted_tokens.saturating_add(accepted_window_tokens)) + } + + pub(in crate::frontend) fn accept_window( + &mut self, + verified_tokens: &[i32], + next_target_token: Option, + ) { + for expected in verified_tokens { + debug_assert_eq!(self.remaining_tokens.pop_front(), Some(*expected)); + } + self.accepted_tokens += verified_tokens.len(); + if let Some(next_target_token) = next_target_token { + if self.remaining_tokens.front() == Some(&next_target_token) { + self.remaining_tokens.pop_front(); + self.accepted_tokens += 1; + } else { + self.remaining_tokens.clear(); + } + } + } + + pub(in crate::frontend) fn reject_window(&mut self, accepted_tokens: usize) { + for _ in 0..accepted_tokens { + self.remaining_tokens + .pop_front() + .expect("accepted composite prefix must remain buffered"); + } + self.accepted_tokens += accepted_tokens; + self.remaining_tokens.clear(); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpVerifyWindowDecision { + 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 classify_native_mtp_verify_window( + 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 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 commit_count = index + 1; + if predicted != *proposal_token { + return Ok(NativeMtpVerifyWindowDecision { + accepted_proposal_tokens, + commit_count, + rejected: true, + }); + } + + accepted_proposal_tokens += 1; + if token_is_eog(predicted)? || generated_len + commit_count >= max_new_tokens { + return Ok(NativeMtpVerifyWindowDecision { + accepted_proposal_tokens, + commit_count, + rejected: false, + }); + } + } + + Ok(NativeMtpVerifyWindowDecision { + accepted_proposal_tokens, + commit_count: required_predictions.min(max_new_tokens.saturating_sub(generated_len)), + rejected: false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn context_with_upstream_span() -> [i32; 10] { + [0, 0, 2, 3, 9, 1, 7, 8, 2, 3] + } + + fn options() -> NativeMtpDecodeOptions { + NativeMtpDecodeOptions { + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 2, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 2, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, + } + } + + #[test] + fn appends_ngram_tail_after_native_mtp_prefix() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[9], &context_with_upstream_span(), 4); + + assert_eq!(proposal.tokens(), &[9, 1, 7, 8]); + assert_eq!(proposal.native_mtp_token_count(), 1); + assert_eq!(proposal.ngram_token_count(), 3); + assert!(proposal.ngram_span_available()); + assert!(proposal.ngram_mtp_prefix_agreed()); + assert!(!proposal.ngram_mtp_prefix_disagreed()); + } + + #[test] + fn match_length_does_not_starve_a_short_adaptive_extension() { + let mut options = options(); + options.ngram_size = 8; + options.ngram_max_proposal_tokens = 10; + let provider = CompositeProposalProvider::from_options(options); + let context = [ + 99, 1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 13, 14, 15, 16, 17, 18, 77, 1, 2, 3, 4, 5, 6, 7, 8, + ]; + + let proposal = provider + .propose_with_ngram_extension(&[11], &context, 10, 2, None) + .unwrap(); + + assert_eq!(proposal.tokens(), &[11, 12, 13]); + assert_eq!(proposal.native_mtp_token_count(), 1); + assert_eq!(proposal.ngram_token_count(), 2); + assert!(proposal.ngram_mtp_prefix_agreed()); + } + + #[test] + fn mtp_only_provider_preserves_native_proposals() { + let mut options = options(); + options.ngram_hybrid = false; + options.ngram_max_proposal_tokens = 0; + let provider = CompositeProposalProvider::from_options(options); + + let proposal = provider.propose(&[9, 10, 11], &[], 2); + + assert_eq!(proposal.tokens(), &[9, 10]); + assert_eq!(proposal.native_mtp_token_count(), 2); + assert_eq!(proposal.ngram_token_count(), 0); + } + + #[test] + fn ngram_limit_does_not_truncate_the_native_prefix() { + let mut options = options(); + options.ngram_max_proposal_tokens = 1; + let provider = CompositeProposalProvider::from_options(options); + + let proposal = provider + .propose_with_ngram_extension(&[9, 10, 11], &[], 3, 1, None) + .unwrap(); + + assert_eq!(proposal.tokens(), &[9, 10, 11]); + assert_eq!(proposal.native_mtp_token_count(), 3); + assert_eq!(proposal.ngram_token_count(), 0); + } + + #[test] + fn keeps_native_mtp_prefix_when_ngram_prefix_disagrees() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[8], &context_with_upstream_span(), 4); + + assert_eq!(proposal.tokens(), &[8]); + assert!(proposal.ngram_span_available()); + assert!(!proposal.ngram_mtp_prefix_agreed()); + assert!(proposal.ngram_mtp_prefix_disagreed()); + } + + #[test] + fn requires_every_native_mtp_token_to_match_before_extending() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[9, 1], &context_with_upstream_span(), 4); + + assert_eq!(proposal.tokens(), &[9, 1, 7, 8]); + assert_eq!(proposal.native_mtp_token_count(), 2); + assert_eq!(proposal.ngram_token_count(), 2); + assert!(proposal.ngram_mtp_prefix_agreed()); + } + + #[test] + fn falls_back_to_pure_ngram_without_native_mtp_tokens() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[], &context_with_upstream_span(), 4); + + assert_eq!(proposal.tokens(), &[9, 1, 7, 8]); + assert_eq!(proposal.native_mtp_token_count(), 0); + assert!(proposal.is_pure_ngram()); + assert!(proposal.ngram_span_available()); + } + + #[test] + fn retains_native_mtp_prefix_when_no_ngram_span_exists() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[9, 10], &[1, 2, 3, 4], 4); + + assert_eq!(proposal.tokens(), &[9, 10]); + assert_eq!(proposal.native_mtp_token_count(), 2); + assert_eq!(proposal.ngram_token_count(), 0); + assert!(!proposal.ngram_span_available()); + } + + #[test] + fn cache_extends_native_mtp_without_requiring_a_matching_prefix() { + let provider = CompositeProposalProvider::from_options(options()); + let mut cache = CachedNgramProposer::new(2, 2).unwrap(); + let context = [1, 9, 7, 1, 9, 7, 1]; + + let proposal = provider + .propose_with_ngram_extension(&[9], &context, 3, 2, Some(&mut cache)) + .unwrap(); + + assert_eq!(proposal.tokens(), &[9, 7, 1]); + assert_eq!(proposal.native_mtp_token_count(), 1); + assert_eq!(proposal.ngram_token_count(), 2); + assert!(proposal.ngram_span_available()); + assert!(!proposal.ngram_mtp_prefix_agreed()); + assert!(!proposal.ngram_mtp_prefix_disagreed()); + } + + #[test] + fn uses_a_single_prior_fixed_ngram_anchor_for_native_mtp() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[9], &context_with_upstream_span(), 4); + + assert_eq!(proposal.tokens(), &[9, 1, 7, 8]); + assert_eq!(proposal.native_mtp_token_count(), 1); + assert_eq!(proposal.ngram_token_count(), 3); + assert!(proposal.ngram_span_available()); + assert!(proposal.ngram_mtp_prefix_agreed()); + } + + #[test] + fn ignores_a_one_token_ngram_tail() { + let provider = CompositeProposalProvider::from_options(options()); + let proposal = provider.propose(&[9], &context_with_upstream_span(), 2); + + assert_eq!(proposal.tokens(), &[9]); + assert_eq!(proposal.native_mtp_token_count(), 1); + assert_eq!(proposal.ngram_token_count(), 0); + assert!(proposal.ngram_span_available()); + assert!(proposal.ngram_mtp_prefix_agreed()); + } + + #[test] + fn tail_rejection_does_not_count_as_native_mtp_rejection() { + let proposal = NativeMtpHybridProposal::from_parts(vec![9, 10, 11], 1, true); + + assert!(!proposal.native_mtp_prefix_rejected(1)); + assert!(proposal.native_mtp_prefix_rejected(0)); + assert!(proposal.ngram_tail_rejected(1)); + assert!(proposal.ngram_tail_rejected(2)); + assert!(!proposal.ngram_tail_rejected(3)); + } + + #[test] + fn tail_rejection_resets_and_backs_off_only_mtp_extensions() { + let proposal = NativeMtpHybridProposal::from_parts(vec![9, 10, 11], 1, true); + let mut controller = NgramSidecarController::new(2, 4); + + assert!(controller.observe_tail_outcome(&proposal, 1, 2)); + assert_eq!(controller.remaining_proposals(), 2); + assert_eq!(controller.current_extension_tokens(), 2); + assert_eq!(controller.extension_limit(&[9], 3), 0); + assert_eq!(controller.extension_limit(&[9], 3), 0); + assert_eq!(controller.extension_limit(&[9], 3), 2); + assert_eq!(controller.extension_limit(&[], 4), 4); + } + + #[test] + fn backoff_preserves_the_native_mtp_prefix() { + let provider = CompositeProposalProvider::from_options(options()); + let mut controller = NgramSidecarController::new(2, 4); + let rejected_tail = NativeMtpHybridProposal::from_parts(vec![9, 1, 2], 1, true); + assert!(controller.observe_tail_outcome(&rejected_tail, 1, 1)); + + let native_only = provider + .propose_with_ngram_extension( + &[9], + &context_with_upstream_span(), + 4, + controller.extension_limit(&[9], 3), + None, + ) + .unwrap(); + let pure_ngram = provider + .propose_with_ngram_extension( + &[], + &context_with_upstream_span(), + 4, + controller.extension_limit(&[], 4), + None, + ) + .unwrap(); + + assert_eq!(native_only.tokens(), &[9]); + assert_eq!(pure_ngram.tokens(), &[9, 1, 7, 8]); + } + + #[test] + fn fully_accepted_tail_grows_the_next_extension_budget() { + let proposal = NativeMtpHybridProposal::from_parts(vec![9, 1, 2], 1, true); + let mut controller = NgramSidecarController::new(3, 6); + + assert_eq!(controller.extension_limit(&[9], 5), 3); + assert!(!controller.observe_tail_outcome(&proposal, 3, 4)); + assert_eq!(controller.current_extension_tokens(), 4); + assert_eq!(controller.extension_limit(&[9], 5), 4); + } + + #[test] + fn caps_parallel_verify_width_to_the_available_candidate_depth() { + let too_shallow = NativeMtpHybridProposal::from_parts(vec![1, 2], 1, true); + let deep_enough = NativeMtpHybridProposal::from_parts(vec![1, 2, 3], 1, true); + let four_tokens = NativeMtpHybridProposal::from_parts(vec![1, 2, 3, 4], 1, true); + let wider = NativeMtpHybridProposal::from_parts(vec![1, 2, 3, 4, 5], 1, true); + + assert_eq!(too_shallow.parallel_verify_width(4, 2), None); + assert_eq!(deep_enough.parallel_verify_width(4, 2), Some(1)); + assert_eq!(four_tokens.parallel_verify_width(4, 2), Some(2)); + assert_eq!(wider.parallel_verify_width(4, 2), Some(3)); + assert_eq!(wider.parallel_verify_width(4, 1), None); + } + + #[test] + fn buffer_reuses_tail_only_when_target_advances_along_it() { + let mut buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 1, 2, 3], + 1, + true, + )); + + buffer.accept_window(&[9, 1], Some(2)); + assert_eq!(buffer.verify_tokens(4), vec![3]); + assert_eq!(buffer.accepted_tokens(), 3); + + buffer.reject_window(0); + assert!(buffer.is_empty()); + } + + #[test] + fn buffer_keeps_the_matching_prefix_when_the_tail_rejects() { + let mut buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 10, 11, 12], + 1, + true, + )); + + buffer.reject_window(3); + + assert!(buffer.is_empty()); + assert_eq!(buffer.accepted_tokens(), 3); + } + + #[test] + fn buffer_exposes_only_a_real_dependent_free_target() { + let buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 1, 2], + 1, + true, + )); + + assert_eq!(buffer.expected_free_target(1), Some(1)); + assert_eq!(buffer.expected_free_target(2), Some(2)); + assert_eq!(buffer.expected_free_target(3), None); + } + + #[test] + fn later_tail_rejection_does_not_reject_an_accepted_native_prefix() { + let mut buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 10, 11, 12], + 1, + true, + )); + + buffer.accept_window(&[9, 10], Some(11)); + + assert!(!buffer.native_mtp_prefix_rejected_after(0)); + } + + #[test] + fn verify_window_commits_the_extra_target_after_full_accept() { + let decision = + classify_native_mtp_verify_window(&[11, 12, 13], &[11, 12, 13, 14], 0, 8, |_| { + Ok(false) + }) + .unwrap(); + + assert_eq!(decision.accepted_proposal_tokens, 3); + assert_eq!(decision.commit_count, 4); + assert!(!decision.rejected); + } + + #[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(); + + assert_eq!(decision.accepted_proposal_tokens, 1); + assert_eq!(decision.commit_count, 2); + assert!(decision.rejected); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/mod.rs b/crates/skippy-server/src/frontend/native_mtp/mod.rs index bd49a35d89..ee0b0a9823 100644 --- a/crates/skippy-server/src/frontend/native_mtp/mod.rs +++ b/crates/skippy-server/src/frontend/native_mtp/mod.rs @@ -1,19 +1,21 @@ -mod batched_verify; mod decode; mod draft; -mod env; +mod hybrid; +mod pipeline; mod stats; -mod trim; mod verifier; +mod verify_window; -pub(super) use batched_verify::BatchedVerifyControl; -pub(super) use decode::{NativeMtpDecodeCounters, NativeMtpDecodeOptions}; +pub(super) use decode::{ + AdaptiveVerifyWindow, NativeMtpDecodeCounters, NativeMtpDecodeOptions, + NativeMtpDecodeTelemetry, NativeMtpTrimAction, native_mtp_trim_action, +}; 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_reject_cooldown_tokens, native_mtp_suppress_cooldown_draft_limit, - native_mtp_suppress_cooldown_drafts_enabled, +pub(super) use hybrid::{ + BufferedCompositeProposal, CompositeProposalProvider, NativeMtpHybridProposal, + NgramSidecarController, classify_native_mtp_verify_window, }; +pub(super) use pipeline::CompositeProposalPipeline; pub(super) use stats::{NativeMtpStats, NativeMtpVerification}; -pub(super) use trim::{NativeMtpTrimAction, native_mtp_trim_action}; pub(super) use verifier::NativeMtpVerifier; +pub(super) use verify_window::NativeMtpVerifyWindowControl; diff --git a/crates/skippy-server/src/frontend/native_mtp/pipeline.rs b/crates/skippy-server/src/frontend/native_mtp/pipeline.rs new file mode 100644 index 0000000000..ec30ffe7eb --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/pipeline.rs @@ -0,0 +1,234 @@ +use std::collections::VecDeque; + +use super::{NativeMtpDraft, NativeMtpDraftOrigin, NativeMtpHybridProposal}; + +/// The candidate portion of one dispatched asynchronous verify window. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(in crate::frontend) struct PipelinedCandidateWindow { + proposal_tokens: Vec, + expected_free_target: Option, + native_mtp_token_count: usize, +} + +impl PipelinedCandidateWindow { + pub(in crate::frontend) fn proposal_tokens(&self) -> &[i32] { + &self.proposal_tokens + } + + pub(in crate::frontend) fn expected_free_target(&self) -> Option { + self.expected_free_target + } + + pub(in crate::frontend) fn native_mtp_token_count(&self) -> usize { + self.native_mtp_token_count + } +} + +/// Owns a deeper composite candidate while asynchronous windows consume it. +/// Each planned window reserves the target's free-advance candidate as the +/// next window's optimistic current token, preventing duplicate KV positions. +#[derive(Debug)] +pub(in crate::frontend) struct CompositeProposalPipeline { + proposal: NativeMtpHybridProposal, + origin: Option, + candidates: VecDeque, + parallel_verify_width: usize, + dispatched_native_mtp_token_count: usize, + accepted_tokens: usize, + next_draft: Option, +} + +impl CompositeProposalPipeline { + pub(in crate::frontend) fn new( + proposal: NativeMtpHybridProposal, + origin: Option, + parallel_verify_width: usize, + ) -> Self { + Self { + candidates: proposal.tokens().iter().copied().collect(), + proposal, + origin, + parallel_verify_width: parallel_verify_width.max(1), + dispatched_native_mtp_token_count: 0, + accepted_tokens: 0, + next_draft: None, + } + } + + pub(in crate::frontend) fn next_window( + &mut self, + verify_width: usize, + ) -> Option { + let verify_width = verify_width + .min(self.parallel_verify_width) + .min(self.candidates.len()); + if verify_width == 0 { + return None; + } + let native_mtp_token_count = self + .proposal + .native_mtp_token_count() + .saturating_sub(self.dispatched_native_mtp_token_count) + .min(verify_width); + let proposal_tokens = self.candidates.drain(..verify_width).collect(); + self.dispatched_native_mtp_token_count += native_mtp_token_count; + Some(PipelinedCandidateWindow { + proposal_tokens, + expected_free_target: self.candidates.pop_front(), + native_mtp_token_count, + }) + } + + pub(in crate::frontend) fn proposal(&self) -> &NativeMtpHybridProposal { + &self.proposal + } + + pub(in crate::frontend) fn origin(&self) -> Option { + self.origin + } + + pub(in crate::frontend) fn has_remaining_candidates(&self) -> bool { + !self.candidates.is_empty() + } + + pub(in crate::frontend) fn candidate_len(&self) -> usize { + self.candidates.len() + } + + pub(in crate::frontend) fn observe_accepted(&mut self, count: usize) { + self.accepted_tokens += count; + } + + pub(in crate::frontend) fn accepted_tokens(&self) -> usize { + self.accepted_tokens + } + + pub(in crate::frontend) fn set_next_draft( + &mut self, + native_mtp_enabled: bool, + draft: Option, + ) { + self.next_draft = native_mtp_enabled.then_some(draft).flatten(); + } + + pub(in crate::frontend) fn next_draft(&self) -> Option<&NativeMtpDraft> { + self.next_draft.as_ref() + } + + pub(in crate::frontend) fn take_next_draft(&mut self) -> Option { + self.next_draft.take() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn proposal(tokens: Vec, native_mtp_tokens: usize) -> NativeMtpHybridProposal { + let ngram_span_available = native_mtp_tokens < tokens.len(); + NativeMtpHybridProposal::from_parts(tokens, native_mtp_tokens, ngram_span_available) + } + + #[test] + fn reserves_free_target_as_the_next_optimistic_current_token() { + let mut pipeline = CompositeProposalPipeline::new( + proposal(vec![9, 1, 2, 3, 4], 1), + Some(NativeMtpDraftOrigin::InitialSerial), + 2, + ); + + let first = pipeline.next_window(2).unwrap(); + assert_eq!(first.proposal_tokens(), &[9, 1]); + assert_eq!(first.expected_free_target(), Some(2)); + assert_eq!(first.native_mtp_token_count(), 1); + + let second = pipeline.next_window(2).unwrap(); + assert_eq!(second.proposal_tokens(), &[3, 4]); + assert_eq!(second.expected_free_target(), None); + assert_eq!(second.native_mtp_token_count(), 0); + } + + #[test] + fn supports_a_pure_ngram_candidate() { + let mut pipeline = CompositeProposalPipeline::new(proposal(vec![1, 2, 3], 0), None, 2); + + let window = pipeline.next_window(2).unwrap(); + assert_eq!(window.proposal_tokens(), &[1, 2]); + assert_eq!(window.expected_free_target(), Some(3)); + assert_eq!(window.native_mtp_token_count(), 0); + assert!(!pipeline.has_remaining_candidates()); + } + + #[test] + fn pure_ngram_pipeline_discards_verify_next_native_mtp_drafts() { + let mut pipeline = CompositeProposalPipeline::new(proposal(vec![1, 2, 3], 0), None, 2); + + pipeline.set_next_draft( + false, + Some(NativeMtpDraft { + tokens: vec![4], + proposal_compute_us: 12, + }), + ); + + assert!(pipeline.next_draft().is_none()); + } + + #[test] + fn caps_each_dispatched_window_to_the_parallel_width() { + let mut pipeline = CompositeProposalPipeline::new(proposal(vec![1, 2, 3], 0), None, 1); + + let window = pipeline.next_window(4).unwrap(); + + assert_eq!(window.proposal_tokens(), &[1]); + assert_eq!(window.expected_free_target(), Some(2)); + } + + #[test] + fn records_the_matching_prefix_of_a_rejected_window() { + let mut pipeline = CompositeProposalPipeline::new( + proposal(vec![9, 1, 2, 3], 1), + Some(NativeMtpDraftOrigin::InitialSerial), + 2, + ); + + let _ = pipeline.next_window(2).unwrap(); + pipeline.observe_accepted(1); + + assert_eq!(pipeline.accepted_tokens(), 1); + assert!( + pipeline + .proposal() + .ngram_tail_rejected(pipeline.accepted_tokens()) + ); + } + + #[test] + fn later_ngram_rejection_does_not_reject_an_accepted_native_prefix() { + let mut pipeline = CompositeProposalPipeline::new( + proposal(vec![9, 1, 2, 3, 4], 1), + Some(NativeMtpDraftOrigin::InitialSerial), + 2, + ); + + let first = pipeline.next_window(2).unwrap(); + assert_eq!(first.proposal_tokens(), &[9, 1]); + assert_eq!(first.expected_free_target(), Some(2)); + pipeline.observe_accepted(3); + + let second = pipeline.next_window(2).unwrap(); + assert_eq!(second.proposal_tokens(), &[3, 4]); + pipeline.observe_accepted(0); + + assert!( + !pipeline + .proposal() + .native_mtp_prefix_rejected(pipeline.accepted_tokens()) + ); + assert!( + pipeline + .proposal() + .ngram_tail_rejected(pipeline.accepted_tokens()) + ); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/trim.rs b/crates/skippy-server/src/frontend/native_mtp/trim.rs deleted file mode 100644 index 99d407acae..0000000000 --- a/crates/skippy-server/src/frontend/native_mtp/trim.rs +++ /dev/null @@ -1,39 +0,0 @@ -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(in crate::frontend) enum NativeMtpTrimAction { - None, - FullSession, -} - -pub(in crate::frontend) fn native_mtp_trim_action( - committed_positions: usize, - consumed_positions: usize, -) -> NativeMtpTrimAction { - if committed_positions < consumed_positions { - NativeMtpTrimAction::FullSession - } else { - NativeMtpTrimAction::None - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn rejected_or_partially_committed_spans_require_full_session_trim() { - assert_eq!( - native_mtp_trim_action(0, 2), - NativeMtpTrimAction::FullSession - ); - assert_eq!( - native_mtp_trim_action(1, 2), - NativeMtpTrimAction::FullSession - ); - } - - #[test] - fn fully_committed_spans_do_not_trim() { - assert_eq!(native_mtp_trim_action(2, 2), NativeMtpTrimAction::None); - assert_eq!(native_mtp_trim_action(3, 2), NativeMtpTrimAction::None); - } -} diff --git a/crates/skippy-server/src/frontend/native_mtp/verifier.rs b/crates/skippy-server/src/frontend/native_mtp/verifier.rs index 4f0307930c..57ee0fd720 100644 --- a/crates/skippy-server/src/frontend/native_mtp/verifier.rs +++ b/crates/skippy-server/src/frontend/native_mtp/verifier.rs @@ -25,6 +25,15 @@ impl NativeMtpVerifier { }) } + pub(in crate::frontend) fn restore_pending_draft(&mut self, pending: PendingNativeMtpDraft) { + debug_assert!(self.pending.is_none()); + debug_assert!(self.pending_tokens.is_empty()); + self.pending_tokens = pending.tokens; + self.pending = Some(PendingDraft { + origin: pending.origin, + }); + } + pub(in crate::frontend) fn clear_pending_draft(&mut self) { self.pending = None; self.pending_tokens.clear(); @@ -219,6 +228,20 @@ mod tests { ); } + #[test] + fn restoring_a_taken_draft_preserves_its_origin_without_counting_it_twice() { + let mut verifier = NativeMtpVerifier::default(); + verifier.observe_next_draft(Some(draft(12)), NativeMtpDraftOrigin::VerifyNext); + + let pending = verifier.take_pending_draft().unwrap(); + verifier.restore_pending_draft(pending); + let restored = verifier.take_pending_draft().unwrap(); + + assert_eq!(restored.tokens, vec![12]); + assert_eq!(restored.origin, NativeMtpDraftOrigin::VerifyNext); + assert_eq!(verifier.stats().drafted_tokens, 1); + } + #[test] fn matching_next_target_accepts_pending_draft() { let mut verifier = NativeMtpVerifier::default(); diff --git a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs new file mode 100644 index 0000000000..8a56d9662f --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs @@ -0,0 +1,593 @@ +use std::net::TcpStream; + +use openai_frontend::{OpenAiError, OpenAiResult}; +use skippy_protocol::binary::{StageNativeMtpDraft, WireReplyKind}; + +use super::super::{ + AdaptiveVerifyWindow, BufferedCompositeProposal, CachedNgramProposer, + CompositeProposalProvider, EmbeddedSessionControl, EmbeddedStageZeroGeneration, + NativeMtpDecodeCounters, NativeMtpDecodeOptions, NativeMtpDraft, NativeMtpDraftOrigin, + NativeMtpTrimAction, NativeMtpVerifier, NgramSidecarController, PendingNativeMtpDraft, + PhaseTimer, StageOpenAiBackend, TokenControl, VerifyWindowMessageArgs, VerifyWindowScheduler, + WireSamplingConfig, classify_native_mtp_verify_window, embedded_verify_window_message, + ms_to_us, native_mtp_trim_action, token_is_eog_with_runtime, +}; + +/// Control signal returned after processing a batched native MTP verify step. +pub(in crate::frontend) enum NativeMtpVerifyWindowControl { + /// The on_token callback returned Stop — outer loop should break. + ReachedStop, + /// Continue the outer decode loop normally. + Continue, + /// No native-MTP or N-gram candidate was available for this position. + NoProposal, +} + +impl StageOpenAiBackend { + #[allow(clippy::too_many_arguments)] + pub(in crate::frontend) fn execute_native_mtp_verify_window( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + session_key: &str, + request_id: u64, + session_id: u64, + prefill_token_count: usize, + wire_sampling: &Option, + native_mtp_options: &NativeMtpDecodeOptions, + verify_window_scheduler: &mut VerifyWindowScheduler, + pending_native_mtp_draft: Option, + proposal_buffer: &mut Option, + cached_ngram_proposer: &mut Option, + adaptive_verify_window: &mut AdaptiveVerifyWindow, + current: &mut i32, + decode_step: u32, + // Mutable decode loop state + decoded_tokens: &mut usize, + context_tokens: &mut Vec, + exact_replay_tokens: &mut Vec, + native_mtp: &mut NativeMtpVerifier, + native_mtp_counters: &mut NativeMtpDecodeCounters, + native_mtp_reject_cooldown_remaining: &mut usize, + native_mtp_suppress_cooldown_drafts_remaining: &mut usize, + ngram_sidecar_controller: &mut NgramSidecarController, + // Mutable decode accumulators + decode_stage0_compute_ms: &mut f64, + decode_runtime_lock_wait_ms: &mut f64, + decode_runtime_lock_wait_max_ms: &mut f64, + decode_runtime_lock_hold_ms: &mut f64, + decode_runtime_lock_hold_max_ms: &mut f64, + decode_runtime_lock_acquires: &mut usize, + decode_forward_activation_encode_ms: &mut f64, + decode_output_activation_bytes: &mut usize, + decode_forward_activation_bytes: &mut usize, + decode_forward_write_ms: &mut f64, + decode_downstream_wait_ms: &mut f64, + // Token emission callback + on_token: &mut impl FnMut(i32) -> OpenAiResult, + ) -> OpenAiResult { + let verify_window_timer = self.telemetry.is_debug_enabled().then(PhaseTimer::start); + let native_mtp_remaining = (request.max_tokens as usize).saturating_sub(*decoded_tokens); + let native_mtp_draft_origin = pending_native_mtp_draft.as_ref().map(|draft| draft.origin); + let native_mtp_draft_tokens = pending_native_mtp_draft + .as_ref() + .map(|draft| { + draft + .tokens + .iter() + .copied() + .take(native_mtp_options.max_draft_tokens) + .take(native_mtp_remaining.saturating_sub(1)) + .collect::>() + }) + .unwrap_or_default(); + if proposal_buffer.is_none() { + let native_mtp_tokens = + if native_mtp_draft_tokens.len() >= native_mtp_options.min_draft_tokens { + native_mtp_draft_tokens.as_slice() + } else { + &[] + }; + let proposal = CompositeProposalProvider::from_options(*native_mtp_options) + .propose_with_ngram_extension( + native_mtp_tokens, + context_tokens, + native_mtp_remaining.saturating_sub(1), + ngram_sidecar_controller.extension_limit( + native_mtp_tokens, + native_mtp_remaining.saturating_sub(native_mtp_tokens.len() + 1), + ), + cached_ngram_proposer.as_mut(), + )?; + if proposal.tokens().is_empty() { + return Ok(NativeMtpVerifyWindowControl::NoProposal); + } + *proposal_buffer = Some(BufferedCompositeProposal::new(proposal)); + } + let proposal_tokens = { + let buffer = proposal_buffer + .as_ref() + .expect("proposal buffer initialized"); + buffer.verify_tokens(adaptive_verify_window.width(buffer.remaining_len())) + }; + if proposal_tokens.is_empty() { + return Ok(NativeMtpVerifyWindowControl::NoProposal); + } + let verify_inputs = native_mtp_verify_window_inputs(*current, &proposal_tokens); + let window = + verify_window_scheduler.open(prefill_token_count + *decoded_tokens, *decoded_tokens)?; + let message = embedded_verify_window_message( + request.wire_dtype, + VerifyWindowMessageArgs { + window_id: window.id, + request_id, + session_id, + prompt_token_count: request.prompt_token_ids.len(), + pos_start: prefill_token_count + *decoded_tokens, + decode_step: *decoded_tokens, + tokens: &verify_inputs, + sampling: wire_sampling.clone(), + checkpoint: false, + }, + )?; + let verify = self.execute_embedded_stage_message( + request, + downstream, + session_key, + &message, + &verify_inputs, + WireReplyKind::PredictedTokens, + )?; + let completed = verify_window_scheduler.complete_next(verify.reply.window.window_id)?; + if completed != window { + return Err(OpenAiError::backend( + "verify window scheduler lost FIFO state", + )); + } + let native_mtp_verify_decision = classify_native_mtp_verify_window( + &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 verify_next_mtp_draft = next_native_mtp_draft( + request.native_mtp_enabled, + verify.reply.native_mtp_draft.clone(), + ); + let native_mtp_decision = (!native_mtp_draft_tokens.is_empty()).then(|| { + let span = native_mtp.observe_taken_draft_span( + &native_mtp_draft_tokens, + &verify.reply.predicted_tokens, + ms_to_us(verify.elapsed_ms), + ); + let verified_draft_count = span.accepted_count + usize::from(span.rejected); + for index in 0..verified_draft_count { + native_mtp_counters.observe_verify_window_verification( + native_mtp_draft_origin.expect("native MTP draft has origin"), + index < span.accepted_count, + ); + } + span.first_decision + }); + let commit_token_count = native_mtp_verify_decision.commit_count; + let consumed_positions = verify_inputs.len(); + let mut committed_positions = 0usize; + let mut reached_stop = false; + for token in verify + .reply + .predicted_tokens + .iter() + .copied() + .take(commit_token_count) + { + *current = token; + *decoded_tokens += 1; + committed_positions += 1; + exact_replay_tokens.push(*current); + context_tokens.push(*current); + if on_token(*current)? == TokenControl::Stop { + reached_stop = true; + break; + } + if *decoded_tokens >= request.max_tokens as usize { + break; + } + } + let fully_accepted_window = !native_mtp_verify_decision.rejected + && native_mtp_verify_decision.accepted_proposal_tokens == proposal_tokens.len() + && committed_positions == consumed_positions + && !reached_stop; + if let Some((profile_width, pipeline_continues)) = + proposal_buffer.as_ref().and_then(|buffer| { + prospective_pipeline_observation( + buffer, + adaptive_verify_window.width(buffer.remaining_len()), + verify_window_scheduler.depth(), + native_mtp_verify_decision.accepted_proposal_tokens, + &verify.reply.predicted_tokens, + ) + }) + { + verify_window_scheduler.observe_pipeline_profile( + profile_width, + pipeline_continues, + verify.stats.stage0_compute_ms, + verify.stats.downstream_wait_ms, + ); + } + let native_mtp_prefix_rejected = proposal_buffer.as_ref().is_some_and(|buffer| { + buffer.native_mtp_prefix_rejected_after( + native_mtp_verify_decision.accepted_proposal_tokens, + ) + }); + let (buffer_exhausted, accepted_proposal_tokens) = { + let buffer = proposal_buffer.as_mut().expect("proposal buffer retained"); + if fully_accepted_window { + buffer.accept_window( + &proposal_tokens, + verify + .reply + .predicted_tokens + .get(proposal_tokens.len()) + .copied(), + ); + } else { + buffer.reject_window(native_mtp_verify_decision.accepted_proposal_tokens); + } + let accepted_proposal_tokens = buffer.accepted_tokens(); + let buffer_exhausted = buffer.is_empty(); + (buffer_exhausted, accepted_proposal_tokens) + }; + let previous_verify_width = adaptive_verify_window.current_tokens(); + let window_adjusted = adaptive_verify_window.observe(fully_accepted_window); + native_mtp_counters.observe_adaptive_verify_window( + proposal_tokens.len(), + previous_verify_width, + adaptive_verify_window.current_tokens(), + ); + if native_mtp_verify_decision.rejected + && native_mtp_prefix_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 = + native_mtp_options.suppress_cooldown_draft_limit; + native_mtp.clear_pending_draft(); + } + let verify_next_mtp_draft_available = verify_next_mtp_draft.is_some(); + let verify_next_mtp_draft_adopted = buffer_exhausted + && fully_accepted_window + && *decoded_tokens < request.max_tokens as usize + && verify_next_mtp_draft.is_some(); + native_mtp_counters.observe_verify_next_draft( + verify_next_mtp_draft_available, + verify_next_mtp_draft_adopted, + ); + if verify_next_mtp_draft_adopted { + native_mtp.observe_next_draft( + verify_next_mtp_draft.clone(), + NativeMtpDraftOrigin::VerifyNext, + ); + } + if buffer_exhausted { + let buffer = proposal_buffer + .take() + .expect("empty proposal buffer retained"); + if ngram_sidecar_controller.observe_tail_outcome( + buffer.proposal(), + accepted_proposal_tokens, + native_mtp_options.ngram_tail_backoff_proposals, + ) { + native_mtp_counters.observe_ngram_tail_rejection(); + } + native_mtp_counters + .observe_hybrid_proposal(buffer.proposal(), buffer.accepted_tokens()); + } + let mut trim_control: Option = None; + match native_mtp_trim_action(committed_positions, consumed_positions) { + NativeMtpTrimAction::None => {} + NativeMtpTrimAction::FullSession => { + let target_token_count = prefill_token_count + *decoded_tokens; + let trim = self.trim_embedded_stage_session( + request, + downstream, + session_key, + request_id, + session_id, + target_token_count, + )?; + trim_control = Some(trim); + } + } + *decode_stage0_compute_ms += verify.stats.stage0_compute_ms; + *decode_runtime_lock_wait_ms += verify.stats.runtime_lock_wait_ms; + *decode_runtime_lock_wait_max_ms = + decode_runtime_lock_wait_max_ms.max(verify.stats.runtime_lock_wait_ms); + *decode_runtime_lock_hold_ms += verify.stats.runtime_lock_hold_ms; + *decode_runtime_lock_hold_max_ms = + decode_runtime_lock_hold_max_ms.max(verify.stats.runtime_lock_hold_ms); + *decode_runtime_lock_acquires += 1; + *decode_forward_activation_encode_ms += verify.stats.activation_encode_ms; + *decode_output_activation_bytes = + decode_output_activation_bytes.saturating_add(verify.stats.output_activation_bytes); + *decode_forward_activation_bytes = + decode_forward_activation_bytes.saturating_add(verify.stats.forward_activation_bytes); + *decode_forward_write_ms += verify.stats.forward_write_ms; + *decode_downstream_wait_ms += verify.stats.downstream_wait_ms; + + if let Some(verify_window_timer) = verify_window_timer { + let mut token_attrs = self.openai_attrs(request.ids); + token_attrs.insert( + "llama_stage.decode_step".to_string(), + serde_json::json!(decode_step), + ); + token_attrs.insert( + "llama_stage.message_kind".to_string(), + serde_json::json!("VerifyWindow"), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_window_batch".to_string(), + serde_json::json!(true), + ); + token_attrs.insert( + "llama_stage.native_mtp.verification".to_string(), + serde_json::json!(native_mtp_decision.map_or("ngram", |decision| decision.label())), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_elapsed_ms".to_string(), + serde_json::json!(verify.elapsed_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.draft_tokens".to_string(), + serde_json::json!(native_mtp_draft_tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.pending_origin".to_string(), + serde_json::json!( + native_mtp_draft_origin.map_or("ngram", NativeMtpDraftOrigin::label) + ), + ); + token_attrs.insert( + "llama_stage.native_mtp.target_token".to_string(), + serde_json::json!(target_token), + ); + token_attrs.insert( + "llama_stage.native_mtp.accepted_count".to_string(), + serde_json::json!(native_mtp_verify_decision.accepted_proposal_tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.hybrid_proposal_len".to_string(), + serde_json::json!(proposal_tokens.len()), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_window_width".to_string(), + serde_json::json!(proposal_tokens.len()), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_window_next_width".to_string(), + serde_json::json!(adaptive_verify_window.current_tokens()), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_window_adjusted".to_string(), + serde_json::json!(window_adjusted), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_available".to_string(), + serde_json::json!(verify_next_mtp_draft_available), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_adopted".to_string(), + serde_json::json!(verify_next_mtp_draft_adopted), + ); + if let Some(next_draft) = verify_next_mtp_draft.as_ref() { + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_tokens".to_string(), + serde_json::json!(next_draft.tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_compute_us".to_string(), + serde_json::json!(next_draft.proposal_compute_us), + ); + } + token_attrs.insert( + "llama_stage.native_mtp.consumed_positions".to_string(), + serde_json::json!(consumed_positions), + ); + token_attrs.insert( + "llama_stage.native_mtp.committed_positions".to_string(), + serde_json::json!(committed_positions), + ); + token_attrs.insert( + "llama_stage.native_mtp.reject_cooldown_tokens".to_string(), + serde_json::json!(native_mtp_options.reject_cooldown_tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.reject_cooldown_remaining".to_string(), + serde_json::json!(*native_mtp_reject_cooldown_remaining), + ); + if let Some(trim) = trim_control.as_ref() { + token_attrs.insert( + "llama_stage.native_mtp.trim_ms".to_string(), + serde_json::json!(trim.elapsed_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_local_ms".to_string(), + serde_json::json!(trim.local_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_downstream_write_ms".to_string(), + serde_json::json!(trim.downstream_write_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_downstream_wait_ms".to_string(), + serde_json::json!(trim.downstream_wait_ms), + ); + } + token_attrs.insert( + "llama_stage.stage0_compute_ms".to_string(), + serde_json::json!(verify.stats.stage0_compute_ms), + ); + token_attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + serde_json::json!(verify.stats.runtime_lock_wait_ms), + ); + token_attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + serde_json::json!(verify.stats.runtime_lock_hold_ms), + ); + token_attrs.insert( + "llama_stage.activation_encode_ms".to_string(), + serde_json::json!(verify.stats.activation_encode_ms), + ); + token_attrs.insert( + "llama_stage.forward_write_ms".to_string(), + serde_json::json!(verify.stats.forward_write_ms), + ); + token_attrs.insert( + "llama_stage.downstream_wait_ms".to_string(), + serde_json::json!(verify.stats.downstream_wait_ms), + ); + token_attrs.insert( + "llama_stage.output_activation_bytes".to_string(), + serde_json::json!(verify.stats.output_activation_bytes), + ); + token_attrs.insert( + "llama_stage.forward_activation_bytes".to_string(), + serde_json::json!(verify.stats.forward_activation_bytes), + ); + self.emit_openai_phase( + "stage.openai_native_mtp_verify", + verify_window_timer, + token_attrs, + ); + } + + if reached_stop { + return Ok(NativeMtpVerifyWindowControl::ReachedStop); + } + Ok(NativeMtpVerifyWindowControl::Continue) + } +} + +fn native_mtp_verify_window_inputs(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 +} + +fn next_native_mtp_draft( + native_mtp_enabled: bool, + stage_draft: Option, +) -> Option { + native_mtp_enabled + .then(|| stage_draft.map(NativeMtpDraft::from_stage_draft)) + .flatten() +} + +fn prospective_pipeline_observation( + buffer: &BufferedCompositeProposal, + adaptive_verify_width: usize, + pipeline_depth: usize, + accepted_proposal_tokens: usize, + predicted_tokens: &[i32], +) -> Option<(usize, bool)> { + if buffer.accepted_tokens() > 0 { + return None; + } + let width = buffer + .proposal() + .parallel_verify_width(adaptive_verify_width, pipeline_depth)?; + let expected_free_target = buffer.expected_free_target(width)?; + let continues = accepted_proposal_tokens >= width + && predicted_tokens.get(width) == Some(&expected_free_target); + Some((width, continues)) +} + +#[cfg(test)] +mod tests { + use super::{ + BufferedCompositeProposal, native_mtp_verify_window_inputs, next_native_mtp_draft, + prospective_pipeline_observation, + }; + use crate::frontend::NativeMtpHybridProposal; + use skippy_protocol::binary::StageNativeMtpDraft; + + #[test] + fn verify_window_inputs_include_every_native_mtp_proposal() { + assert_eq!(native_mtp_verify_window_inputs(10, &[11, 12]), [10, 11, 12]); + } + + #[test] + fn pure_ngram_verify_does_not_capture_a_native_mtp_draft() { + assert_eq!( + next_native_mtp_draft( + false, + Some(StageNativeMtpDraft { + token_ids: vec![11], + proposal_compute_us: 12, + }), + ), + None + ); + } + + #[test] + fn native_mtp_verify_captures_the_next_native_draft() { + let draft = next_native_mtp_draft( + true, + Some(StageNativeMtpDraft { + token_ids: vec![11], + proposal_compute_us: 12, + }), + ) + .expect("native MTP draft should be retained"); + + assert_eq!(draft.tokens, vec![11]); + assert_eq!(draft.proposal_compute_us, 12); + } + + #[test] + fn full_sync_window_profiles_the_narrower_parallel_width() { + let buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 1, 2, 3], + 1, + true, + )); + + assert_eq!( + prospective_pipeline_observation(&buffer, 4, 2, 4, &[9, 1, 2, 3, 4]), + Some((2, true)) + ); + } + + #[test] + fn later_sync_rejection_preserves_valid_parallel_prefix_evidence() { + let buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 1, 2, 3], + 1, + true, + )); + + assert_eq!( + prospective_pipeline_observation(&buffer, 4, 2, 3, &[9, 1, 2, 99, 4]), + Some((2, true)) + ); + } + + #[test] + fn mismatched_parallel_free_target_records_no_continuation() { + let buffer = BufferedCompositeProposal::new(NativeMtpHybridProposal::from_parts( + vec![9, 1, 2, 3], + 1, + true, + )); + + assert_eq!( + prospective_pipeline_observation(&buffer, 4, 2, 2, &[9, 1, 99, 3, 4]), + Some((2, false)) + ); + } +} diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index daaac52ec1..14c52bf98f 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -984,7 +984,11 @@ impl StageOpenAiBackend { .map_err(openai_io_error)?; let forward_write_ms = write_timer.elapsed_ms(); let wait_timer = PhaseTimer::start(); - let downstream_reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; + let downstream_reply = super::embedded_execution::receive_embedded_stage_reply_one_of( + downstream, + request.prediction_return.as_ref(), + &[WireReplyKind::PredictedToken, WireReplyKind::Ack], + )?; let downstream_wait_ms = wait_timer.elapsed_ms(); let downstream_missed = downstream_reply.kind != WireReplyKind::PredictedToken || downstream_reply.stats.kv_lookup_errors > 0 @@ -1032,9 +1036,10 @@ impl StageOpenAiBackend { Ok(Some(EmbeddedFusedFirstDecode { predicted: downstream_reply.predicted, predicted_tokens: vec![downstream_reply.predicted], - native_mtp_draft: NativeMtpDraft::from_prediction_tokens( - &downstream_reply.predicted_tokens, - ), + native_mtp_draft: downstream_reply + .native_mtp_draft + .clone() + .map(NativeMtpDraft::from_stage_draft), reply_stats, execution: EmbeddedExecutionStats { stage0_compute_ms, diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index 74b6b7822d..56cdc3cf5a 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -1,6 +1,6 @@ use super::*; -#[derive(Default)] +#[derive(Clone, Default)] pub(super) struct OpenAiSpeculativeStats { pub(super) windows: usize, pub(super) draft_tokens: usize, @@ -49,48 +49,95 @@ pub(super) struct OpenAiSpeculativeStats { pub(super) adaptive_window_enabled: bool, } -/// Upper bound on the n-gram suffix length scanned during speculative -/// proposal generation. Without a cap, `propose_ngram_tokens` performs a -/// nested scan (`match_len` × `candidate_start` × per-candidate slice -/// compare) over the full history, giving O(N³) worst case behavior. Real -/// n-gram repeats beyond a handful of tokens are vanishingly rare, so -/// bounding the outer match loop at this many tokens keeps the scan -/// tractable on long contexts without losing useful proposals. -const MAX_NGRAM_MATCH: usize = 32; - +/// Uses llama.cpp's ngram-simple self-speculative proposer. The accepted +/// history includes the current token; the upstream API keeps it separate +/// from the preceding history internally. pub(super) fn propose_ngram_tokens( history: &[i32], min_match_tokens: usize, max_proposed_tokens: usize, -) -> Vec { - if min_match_tokens == 0 || max_proposed_tokens == 0 || history.len() < min_match_tokens * 2 { - return Vec::new(); +) -> OpenAiResult> { + skippy_runtime::ngram_simple_draft(history, min_match_tokens, max_proposed_tokens) + .map_err(openai_backend_error) +} + +/// Request-local, cache-based N-gram proposer. It mirrors only committed +/// history into native state; speculative candidates remain read-only inputs. +pub(super) struct CachedNgramProposer { + cache: skippy_runtime::NgramCache, + committed_history: Vec, +} + +impl CachedNgramProposer { + pub(super) fn from_config(config: &SpeculativeDecodeConfig) -> OpenAiResult> { + let Some(ngram) = config.ngram.as_ref() else { + return Ok(None); + }; + if ngram.kind != NgramProposerKind::Cache { + return Ok(None); + } + Self::new(ngram.min_ngram, ngram.max_ngram).map(Some) } - let upper_match = (history.len() / 2).min(MAX_NGRAM_MATCH); - let start_match = min_match_tokens.min(upper_match); - for match_len in (start_match..=upper_match).rev() { - let suffix_start = history.len() - match_len; - let suffix = &history[suffix_start..]; - let latest_candidate_start = suffix_start.saturating_sub(match_len); - for candidate_start in (0..=latest_candidate_start).rev() { - let candidate_end = candidate_start + match_len; - if &history[candidate_start..candidate_end] != suffix { - continue; - } - let proposal_start = candidate_end; - let proposal_end = history.len().min(proposal_start + max_proposed_tokens); - if proposal_start < proposal_end { - return history[proposal_start..proposal_end].to_vec(); - } + + pub(super) fn new(ngram_min: usize, ngram_max: usize) -> OpenAiResult { + let cache = + skippy_runtime::NgramCache::new(ngram_min, ngram_max).map_err(openai_backend_error)?; + Ok(Self { + cache, + committed_history: Vec::new(), + }) + } + + pub(super) fn propose( + &mut self, + committed_history: &[i32], + continuation_prefix: &[i32], + max_proposed_tokens: usize, + ) -> OpenAiResult> { + self.sync(committed_history)?; + self.cache + .draft_after(continuation_prefix, max_proposed_tokens) + .map_err(openai_backend_error) + } + + fn sync(&mut self, committed_history: &[i32]) -> OpenAiResult<()> { + if committed_history.starts_with(&self.committed_history) { + let appended = &committed_history[self.committed_history.len()..]; + self.cache.append(appended).map_err(openai_backend_error)?; + } else { + self.cache + .reset(committed_history) + .map_err(openai_backend_error)?; } + self.committed_history.clear(); + self.committed_history.extend_from_slice(committed_history); + Ok(()) } - Vec::new() } impl OpenAiSpeculativeStats { + pub(super) fn insert_response_timings(&self, timings: &mut BTreeMap) { + timings.insert( + "verify_window_verify_elapsed_ms".to_string(), + json!(self.primary_verify_elapsed_ms), + ); + timings.insert( + "verify_window_stage0_compute_ms".to_string(), + json!(self.primary_verify_stage0_compute_ms), + ); + timings.insert( + "verify_window_forward_write_ms".to_string(), + json!(self.primary_verify_forward_write_ms), + ); + timings.insert( + "verify_window_downstream_wait_ms".to_string(), + json!(self.primary_verify_downstream_wait_ms), + ); + } + pub(super) fn observe_verify_decision( &mut self, - decision: VerifySpanDecision, + decision: VerifyWindowDecision, adaptive_window: &mut usize, adaptive_enabled: bool, max_speculative_window: usize, @@ -103,7 +150,7 @@ impl OpenAiSpeculativeStats { self.adaptive_window_min = nonzero_min(self.adaptive_window_min, *adaptive_window); self.adaptive_window_max_seen = self.adaptive_window_max_seen.max(*adaptive_window); match decision.kind { - VerifySpanDecisionKind::FullAccept => { + VerifyWindowDecisionKind::FullAccept => { self.full_accept_windows += 1; self.grow_adaptive_window( adaptive_window, @@ -111,10 +158,10 @@ impl OpenAiSpeculativeStats { max_speculative_window, ); } - VerifySpanDecisionKind::AcceptedStop => { + VerifyWindowDecisionKind::AcceptedStop => { self.accepted_stop_windows += 1; } - VerifySpanDecisionKind::TailReject => { + VerifyWindowDecisionKind::TailReject => { self.observe_reject(decision); self.tail_reject_windows += 1; self.grow_adaptive_window( @@ -123,13 +170,13 @@ impl OpenAiSpeculativeStats { max_speculative_window, ); } - VerifySpanDecisionKind::EarlyReject => { + VerifyWindowDecisionKind::EarlyReject => { self.observe_reject(decision); self.early_reject_windows += 1; self.repair_required_windows += 1; self.shrink_adaptive_window(adaptive_window, adaptive_enabled, decision); } - VerifySpanDecisionKind::EarlyRejectStop => { + VerifyWindowDecisionKind::EarlyRejectStop => { self.observe_reject(decision); self.early_reject_windows += 1; self.early_reject_stop_windows += 1; @@ -137,7 +184,7 @@ impl OpenAiSpeculativeStats { } } - pub(super) fn observe_reject(&mut self, decision: VerifySpanDecision) { + pub(super) fn observe_reject(&mut self, decision: VerifyWindowDecision) { if let Some(repair_input_count) = decision.repair_input_count { self.rejected_windows += 1; self.first_reject_position_sum += repair_input_count; @@ -160,7 +207,7 @@ impl OpenAiSpeculativeStats { &mut self, adaptive_window: &mut usize, adaptive_enabled: bool, - decision: VerifySpanDecision, + decision: VerifyWindowDecision, ) { if !adaptive_enabled { return; @@ -340,14 +387,35 @@ mod ngram_tests { fn proposes_tokens_after_latest_matching_suffix() { let history = [1, 2, 3, 4, 9, 2, 3, 4]; - assert_eq!(propose_ngram_tokens(&history, 2, 2), vec![9, 2]); + assert_eq!(propose_ngram_tokens(&history, 2, 2).unwrap(), vec![9, 2]); } #[test] fn returns_empty_without_enough_history() { - assert!(propose_ngram_tokens(&[1, 2, 3], 2, 4).is_empty()); - assert!(propose_ngram_tokens(&[1, 2, 1, 2], 0, 4).is_empty()); - assert!(propose_ngram_tokens(&[1, 2, 1, 2], 1, 0).is_empty()); + assert!(propose_ngram_tokens(&[1, 2, 3], 2, 4).unwrap().is_empty()); + assert!( + propose_ngram_tokens(&[1, 2, 1, 2], 0, 4) + .unwrap() + .is_empty() + ); + assert!( + propose_ngram_tokens(&[1, 2, 1, 2], 1, 0) + .unwrap() + .is_empty() + ); + } + + #[test] + fn cache_proposer_syncs_only_the_committed_prefix() { + let mut proposer = CachedNgramProposer::new(2, 2).unwrap(); + let history = [1, 2, 3, 1, 2, 3, 1, 2]; + + assert_eq!(proposer.propose(&history, &[], 2).unwrap(), vec![3, 1]); + assert_eq!( + proposer.propose(&history, &[9], 2).unwrap(), + Vec::::new() + ); + assert_eq!(proposer.propose(&history, &[], 2).unwrap(), vec![3, 1]); } } @@ -362,7 +430,7 @@ pub(super) fn verify_inputs_for_proposals(current: i32, proposals: &[i32]) -> Ve } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum VerifySpanDecisionKind { +pub(super) enum VerifyWindowDecisionKind { FullAccept, AcceptedStop, TailReject, @@ -371,41 +439,41 @@ pub(super) enum VerifySpanDecisionKind { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct VerifySpanDecision { - pub(super) kind: VerifySpanDecisionKind, +pub(super) struct VerifyWindowDecision { + pub(super) kind: VerifyWindowDecisionKind, pub(super) accepted_before_reject: usize, pub(super) repair_input_count: Option, pub(super) commit_count: usize, } -impl VerifySpanDecision { +impl VerifyWindowDecision { pub(super) fn rejected(self) -> bool { matches!( self.kind, - VerifySpanDecisionKind::TailReject - | VerifySpanDecisionKind::EarlyReject - | VerifySpanDecisionKind::EarlyRejectStop + VerifyWindowDecisionKind::TailReject + | VerifyWindowDecisionKind::EarlyReject + | VerifyWindowDecisionKind::EarlyRejectStop ) } pub(super) fn requires_repair(self) -> bool { - self.kind == VerifySpanDecisionKind::EarlyReject + self.kind == VerifyWindowDecisionKind::EarlyReject } } -pub(super) fn classify_verify_span( +pub(super) fn classify_verify_window( draft_tokens: &[i32], predicted_tokens: &[i32], generated_len: usize, max_new_tokens: usize, mut token_is_eog: F, -) -> OpenAiResult +) -> OpenAiResult where F: FnMut(i32) -> OpenAiResult, { if predicted_tokens.len() < draft_tokens.len() { return Err(OpenAiError::backend(format!( - "verify span returned too few tokens: got {} expected {}", + "verify window returned too few tokens: got {} expected {}", predicted_tokens.len(), draft_tokens.len() ))); @@ -421,8 +489,8 @@ where if accepted { accepted_before_reject += 1; if (reached_eog || reached_limit) && commit_count < draft_tokens.len() { - return Ok(VerifySpanDecision { - kind: VerifySpanDecisionKind::AcceptedStop, + return Ok(VerifyWindowDecision { + kind: VerifyWindowDecisionKind::AcceptedStop, accepted_before_reject, repair_input_count: None, commit_count, @@ -433,13 +501,13 @@ where let repair_input_count = accepted_before_reject + 1; let kind = if repair_input_count == draft_tokens.len() { - VerifySpanDecisionKind::TailReject + VerifyWindowDecisionKind::TailReject } else if reached_eog || reached_limit { - VerifySpanDecisionKind::EarlyRejectStop + VerifyWindowDecisionKind::EarlyRejectStop } else { - VerifySpanDecisionKind::EarlyReject + VerifyWindowDecisionKind::EarlyReject }; - return Ok(VerifySpanDecision { + return Ok(VerifyWindowDecision { kind, accepted_before_reject, repair_input_count: Some(repair_input_count), @@ -447,8 +515,8 @@ where }); } - Ok(VerifySpanDecision { - kind: VerifySpanDecisionKind::FullAccept, + Ok(VerifyWindowDecision { + kind: VerifyWindowDecisionKind::FullAccept, accepted_before_reject, repair_input_count: None, commit_count, diff --git a/crates/skippy-server/src/frontend/tests.rs b/crates/skippy-server/src/frontend/tests.rs index 089e1005d2..582962a084 100644 --- a/crates/skippy-server/src/frontend/tests.rs +++ b/crates/skippy-server/src/frontend/tests.rs @@ -4,7 +4,7 @@ use async_trait::async_trait; use std::io::Cursor; use std::{ env, fs, - net::SocketAddr, + net::{SocketAddr, TcpListener, TcpStream}, sync::{ Arc, Mutex, atomic::{AtomicU64, AtomicUsize, Ordering}, @@ -1096,6 +1096,11 @@ fn chat_response_from_parsed_message_separates_reasoning_content() { verification_compute_us: 200, ..NativeMtpStats::default() }, + native_mtp_decode_telemetry: None, + verify_window_pipeline_stats: None, + speculative_stats: None, + prompt_ms: 20.0, + predicted_ms: 100.0, text: "Checked facts first.Final answer.".to_string(), finish_reason: FinishReason::Stop, detokenize_ms: 0.0, @@ -1121,6 +1126,106 @@ fn chat_response_from_parsed_message_separates_reasoning_content() { ); assert_eq!(message.tool_calls, None); assert_eq!(response.choices[0].finish_reason, Some(FinishReason::Stop)); + + let completion = completion_response_from_generated_text("qwen".to_string(), &output); + let timings = completion + .timings + .as_ref() + .expect("completion native MTP timings"); + assert_eq!(timings.get("draft_n"), Some(&json!(7))); + assert_eq!(timings.get("draft_n_accepted"), Some(&json!(5))); + assert_eq!(timings.get("predicted_per_second"), Some(&json!(70.0))); +} + +#[test] +fn generated_text_timings_are_present_without_native_mtp() { + let output = GeneratedText { + prompt_tokens: 4, + completion_tokens: 8, + cache_status: "disabled", + cached_prompt_tokens: 0, + matched_prefix_tokens: 0, + suffix_prefill_tokens: 0, + cache_hit_kind: None, + native_mtp_stats: NativeMtpStats::default(), + native_mtp_decode_telemetry: None, + verify_window_pipeline_stats: None, + speculative_stats: None, + prompt_ms: 20.0, + predicted_ms: 100.0, + text: "Paris".to_string(), + finish_reason: FinishReason::Stop, + detokenize_ms: 0.0, + text_emit_ms: 0.0, + eog_check_ms: 0.0, + }; + + let timings = output.timings().expect("standard timings"); + assert_eq!(timings.get("draft_n"), Some(&json!(0))); + assert_eq!(timings.get("draft_n_accepted"), Some(&json!(0))); + assert_eq!(timings.get("prompt_per_second"), Some(&json!(200.0))); + assert_eq!(timings.get("predicted_per_second"), Some(&json!(80.0))); +} + +#[test] +fn generated_text_timings_prefer_composite_proposal_totals() { + let mut counters = NativeMtpDecodeCounters::default(); + let proposal = CompositeProposalProvider::from_options(NativeMtpDecodeOptions { + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 2, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 2, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, + }) + .propose(&[], &[0, 0, 2, 3, 9, 1, 7, 8, 2, 3], 4); + counters.observe_hybrid_proposal(&proposal, 4); + let output = GeneratedText { + prompt_tokens: 4, + completion_tokens: 8, + cache_status: "disabled", + cached_prompt_tokens: 0, + matched_prefix_tokens: 0, + suffix_prefill_tokens: 0, + cache_hit_kind: None, + native_mtp_stats: NativeMtpStats::default(), + native_mtp_decode_telemetry: Some(NativeMtpDecodeTelemetry::new( + NativeMtpDecodeOptions { + max_draft_tokens: 1, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + ngram_hybrid: true, + ngram_size: 2, + ngram_initial_extension_tokens: 2, + ngram_max_proposal_tokens: 4, + ngram_tail_backoff_proposals: 2, + verify_window_min_tokens: 1, + verify_window_max_tokens: 4, + }, + counters, + )), + verify_window_pipeline_stats: None, + speculative_stats: None, + prompt_ms: 20.0, + predicted_ms: 100.0, + text: "ok".to_string(), + finish_reason: FinishReason::Stop, + detokenize_ms: 0.0, + text_emit_ms: 0.0, + eog_check_ms: 0.0, + }; + + let timings = output.timings().expect("timings"); + assert_eq!(timings.get("draft_n"), Some(&json!(4))); + assert_eq!(timings.get("draft_n_accepted"), Some(&json!(4))); } #[test] @@ -1498,9 +1603,7 @@ fn local_openai_backend(config: StageConfig) -> Result { adaptive_speculative_window: false, ngram_min: 0, ngram_max: 0, - native_mtp_enabled: false, - native_mtp_max_tokens: 1, - native_mtp_min_tokens: 0, + speculative: SpeculativeDecodeConfig::default(), generation_limit: Arc::new(Semaphore::new(1)), generation_queue_depth: Arc::new(AtomicUsize::new(0)), generation_queue_limit: 1, @@ -1686,18 +1789,24 @@ async fn real_multimodal_split_smoke_when_fixture_is_set() -> Result<()> { prefill_reply_credit_limit: 0, lane_pool: Some(lane_pool), prediction_returns: None, - native_mtp_enabled: true, - native_mtp_max_tokens: 3, - native_mtp_min_tokens: 0, }, draft: None, speculative_window: 0, adaptive_speculative_window: false, ngram_min: 0, ngram_max: 0, - native_mtp_enabled: true, - native_mtp_max_tokens: 3, - native_mtp_min_tokens: 0, + speculative: SpeculativeDecodeConfig { + native_mtp: NativeMtpProposalConfig { + enabled: true, + max_draft_tokens: 3, + min_draft_tokens: 0, + reject_cooldown_tokens: 0, + suppress_cooldown_drafts: false, + suppress_cooldown_draft_limit: 0, + }, + effective_strategy: "native-mtp".to_string(), + ..SpeculativeDecodeConfig::default() + }, generation_limit: Arc::new(Semaphore::new(1)), generation_queue_depth: Arc::new(AtomicUsize::new(0)), generation_queue_limit: 1, @@ -2728,6 +2837,71 @@ fn prefill_transport_ewma_seeds_adaptive_ramp() { assert_eq!(planner.chunk_size_for(0, 512), 256); } +#[test] +fn persistent_lane_ready_handshake_times_out_for_silent_downstream() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + let (_stream, _) = listener.accept().unwrap(); + std::thread::sleep(Duration::from_millis(200)); + }); + let mut client = TcpStream::connect(address).unwrap(); + + let error = receive_persistent_lane_ready(&mut client, Duration::from_millis(25)).unwrap_err(); + + assert!( + error + .to_string() + .contains("persistent downstream lane did not become ready") + ); + server.join().unwrap(); +} + +#[test] +fn steady_state_lane_reconnect_deadline_is_much_shorter_than_warmup() { + // A mid-life reconnect to an already-serving mesh must fail fast so a new + // request routed to a dead stage errors quickly instead of stalling for the + // full warmup deadline. Guard the invariant that the steady-state deadline + // stays well under the warmup deadline. + assert!( + LANE_STEADY_CONNECT_TIMEOUT < LANE_READY_READ_TIMEOUT, + "steady-state reconnect deadline must be shorter than the warmup deadline" + ); + assert!( + LANE_STEADY_CONNECT_TIMEOUT <= Duration::from_secs(5), + "steady-state reconnect deadline must stay small enough to fail fast on a dead stage" + ); +} + +#[test] +fn steady_state_ready_handshake_times_out_fast_for_silent_downstream() { + // A downstream that accepts the TCP connection but never sends the ready + // frame must be bounded by the supplied deadline, not hang. This mirrors the + // silent-downstream case a dead split stage produces on a new request. + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + let (_stream, _) = listener.accept().unwrap(); + std::thread::sleep(Duration::from_millis(200)); + }); + let mut client = TcpStream::connect(address).unwrap(); + + let start = std::time::Instant::now(); + let error = receive_persistent_lane_ready(&mut client, Duration::from_millis(25)).unwrap_err(); + let elapsed = start.elapsed(); + + assert!( + error + .to_string() + .contains("persistent downstream lane did not become ready") + ); + assert!( + elapsed < Duration::from_secs(1), + "handshake read must fail within the supplied deadline, took {elapsed:?}" + ); + server.join().unwrap(); +} + #[test] fn model_matching_is_exact_for_mesh_style_ids() { ensure_requested_model( diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index e3bdb63f83..4998dbac97 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -121,7 +121,8 @@ impl ReusableDecodeMessage { } } -pub(super) struct VerifySpanMessageArgs<'a> { +pub(super) struct VerifyWindowMessageArgs<'a> { + pub(super) window_id: i32, pub(super) request_id: u64, pub(super) session_id: u64, pub(super) prompt_token_count: usize, @@ -132,17 +133,17 @@ pub(super) struct VerifySpanMessageArgs<'a> { pub(super) checkpoint: bool, } -pub(super) fn embedded_verify_message( +pub(super) fn embedded_verify_window_message( wire_dtype: WireActivationDType, - args: VerifySpanMessageArgs<'_>, + args: VerifyWindowMessageArgs<'_>, ) -> OpenAiResult { if args.tokens.is_empty() { return Err(OpenAiError::backend( - "verify span requires at least one token", + "verify window requires at least one token", )); } - let mut state = StageStateHeader::new(WireMessageKind::VerifySpan, wire_dtype); - state.seq_id = 0; + let mut state = StageStateHeader::new(WireMessageKind::VerifyWindow, wire_dtype); + state.seq_id = args.window_id; state.prompt_token_count = i32::try_from(args.prompt_token_count) .map_err(|_| OpenAiError::backend("prompt token count exceeds i32"))?; state.decode_step = i32::try_from(args.decode_step) @@ -153,11 +154,11 @@ pub(super) fn embedded_verify_message( state.flags |= state_flags::SKIP_VERIFY_CHECKPOINT; } Ok(StageWireMessage { - kind: WireMessageKind::VerifySpan, + kind: WireMessageKind::VerifyWindow, pos_start: i32::try_from(args.pos_start) - .map_err(|_| OpenAiError::backend("verify span position exceeds i32"))?, + .map_err(|_| OpenAiError::backend("verify window position exceeds i32"))?, token_count: i32::try_from(args.tokens.len()) - .map_err(|_| OpenAiError::backend("verify span exceeds i32"))?, + .map_err(|_| OpenAiError::backend("verify window exceeds i32"))?, state, request_id: args.request_id, session_id: args.session_id, diff --git a/crates/skippy-server/src/lib.rs b/crates/skippy-server/src/lib.rs index 7750e8a2ca..66d01d7fdb 100644 --- a/crates/skippy-server/src/lib.rs +++ b/crates/skippy-server/src/lib.rs @@ -27,7 +27,9 @@ pub use embedded::{ pub use frontend::{ CONTEXT_BUDGET_MAX_TOKENS, DEFAULT_EMBEDDED_MAX_TOKENS, EmbeddedOpenAiArgs, EmbeddedOpenAiBackend, EmbeddedOpenAiRequestDefaults, EmbeddedReasoningBudget, - EmbeddedReasoningEnabled, EmbeddedReasoningFormat, OpenAiGuardrailsConfig, - OpenAiGuardrailsStatus, OpenAiGuardrailsTarget, embedded_openai_backend, + EmbeddedReasoningEnabled, EmbeddedReasoningFormat, NativeMtpProposalConfig, + NgramExtensionConfig, NgramProposalConfig, NgramProposerKind, OpenAiGuardrailsConfig, + OpenAiGuardrailsStatus, OpenAiGuardrailsTarget, SpeculativeDecodeConfig, VerifyWindowConfig, + embedded_openai_backend, }; pub use skippy_protocol::StageConfig; diff --git a/docs/README.md b/docs/README.md index 399b2544a8..7d26c7daea 100644 --- a/docs/README.md +++ b/docs/README.md @@ -27,6 +27,7 @@ Use this hub to find project guides that are not owned by a single Rust crate. | [skippy/TOPOLOGY_PLANNER.md](skippy/TOPOLOGY_PLANNER.md) | Stage topology planning behavior | | [skippy/CONFIGURATION.md](skippy/CONFIGURATION.md) | Authoritative operator matrix for Skippy config keys and rejection boundaries | | [skippy/PROMPT_CACHE.md](skippy/PROMPT_CACHE.md) | OpenAI prompt-prefix cache behavior, defaults, telemetry, and benchmark flow | +| [skippy/PIPELINED_VERIFY_WINDOW.md](skippy/PIPELINED_VERIFY_WINDOW.md) | Native MTP, anchored N-gram extension, VerifyWindow protocol, pipeline behavior, and telemetry | | [skippy/DATA_FLOW.md](skippy/DATA_FLOW.md) | Stage data flow and transport details | | [skippy/LLAMA_PARITY.md](skippy/LLAMA_PARITY.md) | Remaining llama.cpp parity queue | | [specs/layer-package-repos.md](specs/layer-package-repos.md) | Manifest schema and package artifact rules | diff --git a/docs/USAGE.md b/docs/USAGE.md index 792b6c9378..9579c416a7 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -320,8 +320,8 @@ lifecycle_health_interval_ms = 5000 # health-check interval (ms) # --- Speculative decoding ------------------------------------------------ [defaults.speculative] -strategy = "auto" # auto disabled mtp -mode = "auto" # auto off draft ngram lookahead +strategy = "auto" # auto disabled mtp or a package strategy id +mode = "auto" # legacy draft-model mode: auto disabled draft ngram draft_selection_policy = "auto" # auto manual heuristic pairing_fault = "warn_disable" # warn_disable fail_open fail_closed draft_acceptance_threshold = 0.0 # 0.0 = use runtime default @@ -345,9 +345,23 @@ spec_default = "auto" # bool or "auto" # draft_cache_type_k = "q8_0" # draft_cache_type_v = "q8_0" -# N-gram speculative (when mode = "ngram") -# ngram_min = 1 -# ngram_max = 5 +# N-gram proposer and MTP extension. `simple` scans accepted history; +# `cache` is request-local and currently requires ngram_max <= 4. +# ngram_proposer = "cache" # simple cache +# ngram_min = 2 +# ngram_max = 4 +# ngram_max_proposal_tokens = 6 # output budget, separate from ngram_max +# extension_initial_tokens = 2 # requires native MTP plus an N-gram proposer +# extension_max_tokens = 6 +# extension_tail_backoff_proposals = 2 + +# Target VerifyWindow and native-MTP recovery controls +# verify_window_min_tokens = 1 +# verify_window_max_tokens = 6 +# verify_window_pipeline_depth = 2 +# native_mtp_reject_cooldown_tokens = 4 +# native_mtp_suppress_cooldown_drafts = true +# native_mtp_suppress_cooldown_draft_limit = 1 # --- Request defaults (merged at OpenAI frontend only) ------------------- [defaults.request_defaults] @@ -702,6 +716,51 @@ Config precedence: process launch until direct plugin use. This is useful for very slow legacy hosts or emulator-assisted startup paths. +## Speculative decode configuration + +Configure speculative decoding under `[defaults.speculative]` for all staged +models, or under `[models.speculative]` to override one configured model. CLI +flags have the highest precedence, followed by the selected model, then +`[defaults.speculative]`; package strategies supply the remaining declared +defaults. The resolved plan is validated once before Skippy starts. + +Set `strategy = "auto"` to use a package recommendation, `"disabled"` for +the no-speculation baseline, or `"mtp"` for native MTP. A package may also +publish stable names such as `mtp-cache`; that name is valid only for the +package that declares it. Direct GGUF serving can use `ngram-simple` or +`ngram-cache` when it supplies valid N-gram bounds. + +```toml +[[models]] +model = "meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M" + +[models.speculative] +strategy = "mtp" +ngram_proposer = "cache" +ngram_min = 2 +ngram_max = 4 +ngram_max_proposal_tokens = 6 +extension_initial_tokens = 2 +extension_max_tokens = 6 +extension_tail_backoff_proposals = 2 +verify_window_min_tokens = 1 +verify_window_max_tokens = 6 +verify_window_pipeline_depth = 2 +``` + +`ngram_min` and `ngram_max` determine the history match length. +`ngram_max_proposal_tokens` is separately the maximum continuation length. +The request-local `cache` proposer is limited to `ngram_max <= 4`; `simple` +searches accepted history instead. Extension settings require both native MTP +and an N-gram proposer. All combinations are verified together by the target, +so tuning these values changes speculative work, not output correctness. + +For package-authoring rules, see +[Layer Package Repositories](specs/layer-package-repos.md#generation-defaults). +For strategy diagrams, CLI overrides, and the VerifyWindow telemetry used to +evaluate a configuration, see +[Pipelined VerifyWindow Decode](skippy/PIPELINED_VERIFY_WINDOW.md). + ## Lemonade integration Use the `openai-endpoint` plugin to route requests to a local [Lemonade Server](https://lemonade-server.ai) through the same `http://localhost:9337/v1` API that mesh-llm exposes. diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index 12ee62ff87..0bbd7f0944 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -157,18 +157,23 @@ missing-key error. | Report section | Report setting name | Config key path | Priority | Owner module | Translation target | Supported modes | Live-apply behavior | Default source | Validation rule | Docs anchor | Test evidence | Notes | |---|---|---|---|---|---|---|---|---|---|---|---|---| -| 5.5 | Speculative strategy override | `speculative.strategy` | P1 | `resolver/speculative.rs` | native MTP stage config and embedded frontend config | single-stage, staged | restart/reload only | `auto`, using package/runtime defaults | enum `auto`, `disabled`, or `mtp`; `disabled` turns off native MTP for the configured scope | `#speculative-decoding` | resolver/config validation tests | package-level native MTP default override | -| 5.5 | Speculation mode/type | `speculative.mode` | P1 | `plugin/config.rs` | embedded frontend speculative config and draft planner | single-stage, staged | restart/reload only | non-speculative default unless operator enables draft or n-gram mode | enum must match certified speculative strategies; `draft` requires an explicit draft source and compatible pairing policy | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | high-level speculative decoding mode | +| 5.5 | Speculative strategy override | `speculative.strategy` | P1 | `resolver/speculative.rs` | resolved `SpeculativeDecodeConfig` for the embedded frontend | staged | restart/reload only | `auto`, using package/runtime defaults | `auto`, `disabled`, `mtp`, built-in `ngram-simple` / `ngram-cache` with valid bounds, or a selected package strategy id; `disabled` turns off target speculation | `#speculative-decoding` | resolver/config validation tests | CLI override wins, then the selected model, then defaults; a named package strategy must be declared by that package | +| 5.5 | Speculation mode/type | `speculative.mode` | P1 | `resolver/speculative.rs` | legacy draft-model and standalone N-gram compatibility selection | staged | restart/reload only | `auto` | `auto`, `disabled`, `draft`, or `ngram`; use `strategy` for package and native-MTP selection | `#speculative-decoding` | resolver/config validation tests | Kept for draft-model compatibility. New native-MTP and package configurations should use `strategy`. | | 5.5 | Draft model identifier / HF source | `speculative.draft_model
speculative.draft_hf_repo
speculative.draft_hf_file` | P1 | `plugin/config.rs` | draft model resolver | single-stage, staged | restart/reload only | unset until operator or planner chooses a draft model | set an HF-style identifier (e.g. `Qwen/Qwen3-0.6B:Q4_K_M`) or provide hf_repo plus hf_file for remote resolution; required when `speculative.mode = "draft"` | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | manual draft model source selection | | 5.5 | Draft model auto-selection | `speculative.draft_selection_policy` | P1 | `plugin/config.rs` | draft model resolver policy | single-stage, staged | restart/reload only | mesh policy default | enum must match supported draft-selection strategies | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | automatic draft pairing policy | | 5.5 | Pairing fault behavior | `speculative.pairing_fault` | P1 | `plugin/config.rs` | draft pairing safety handler | single-stage, staged | restart/reload only | safe fallback policy default | enum fail-open, fail-closed, or later certified policies only | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | what to do when draft pairing is invalid | -| 5.5 | Draft max tokens | `speculative.draft_max_tokens` | P1 | `plugin/config.rs` | embedded frontend speculative token window | single-stage, staged | restart/reload only | runtime default | integer >= 1; no documented executable relationship to draft_min_tokens today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | upper speculative proposal window | -| 5.5 | Draft min tokens | `speculative.draft_min_tokens` | P2 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime default | integer >= 1; any value above `0` currently fails closed and is not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | lower speculative proposal window; reserved until a real executable surface exists | +| 5.5 | Draft max tokens | `speculative.draft_max_tokens` | P1 | `resolver/speculative.rs` | native-MTP or draft-model proposal window | staged | restart/reload only | native MTP uses one token unless the package/setting provides another bound | non-negative integer; when set with `draft_min_tokens`, the minimum must not exceed the resolved maximum | `#speculative-decoding` | resolver/config validation tests | Applies to native MTP as well as a compatible external draft model. | +| 5.5 | Draft min tokens | `speculative.draft_min_tokens` | P2 | `resolver/speculative.rs` | native-MTP or draft-model minimum proposal window | staged | restart/reload only | `0` | non-negative integer no greater than the resolved maximum | `#speculative-decoding` | resolver/config validation tests | A lower bound, not an acceptance-rate control. | | 5.5 | Draft probability/accept threshold | `speculative.draft_acceptance_threshold` | P2 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime default | float in supported acceptance range, typically 0.0..1.0; any value above `0.0` currently fails closed and is not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | acceptance threshold for proposed draft tokens; reserved until a real executable surface exists | | 5.5 | Draft split probability | `speculative.draft_split_probability` | P2 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime default | float in supported probability range, typically 0.0..1.0; not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | advanced split probability control; schema-reserved until a real executable surface exists | | 5.5 | Draft GPU layers/device/threads | `speculative.draft_gpu_layers
speculative.draft_device
speculative.draft_threads` | P1 | `plugin/config.rs` | draft model runtime allocation fields | single-stage, staged | restart/reload only | planner or runtime default | `draft_gpu_layers` may propagate on the staged draft path; `draft_device` and `draft_threads` fail closed and are not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | separate allocation policy for draft model; keep `draft_gpu_layers` distinct from the rejected device/thread overrides | | 5.5 | Draft KV cache type | `speculative.draft_cache_type_k
speculative.draft_cache_type_v` | P1 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime draft defaults | enum values are currently rejected by the resolver and are not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | draft-model KV cache dtype overrides; reserve until an executable draft-runtime path exists | -| 5.5 | N-gram speculative min/max | `speculative.ngram_min
speculative.ngram_max` | P2 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime n-gram defaults | integers >= 1 and ngram_max >= ngram_min when both set; mode `ngram` itself fails closed, so the numeric bounds are not executable today | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | specialized n-gram speculative window sizing; reserve until a real executable surface exists | +| 5.5 | N-gram match range | `speculative.ngram_min
speculative.ngram_max` | P1 | `resolver/speculative.rs` | `NgramProposalConfig` | staged | restart/reload only | package proposer bounds, otherwise required for a direct N-gram plan | both values are required for a direct N-gram proposer and must satisfy `0 < min <= max`; request-local `cache` requires `max <= 4` | `#speculative-decoding` | resolver/config validation tests | Match range controls lookup context, not the number of proposed tokens. | +| 5.5 | N-gram proposer kind | `speculative.ngram_proposer` | P1 | `resolver/speculative.rs` | `NgramProposalConfig.kind` | staged | restart/reload only | package strategy kind, otherwise `simple` | `simple` or `cache` | `#speculative-decoding` | resolver/config validation tests | `simple` searches accepted history; `cache` owns request-local target-committed history and can extend an MTP prefix. | +| 5.5 | N-gram output budget | `speculative.ngram_max_proposal_tokens` | P1 | `resolver/speculative.rs` | `NgramProposalConfig.max_proposal_tokens` | staged | restart/reload only | package value, otherwise N-gram maximum | positive integer | `#speculative-decoding` | resolver/config validation tests | Independent from the N-gram match range. | +| 5.5 | MTP N-gram extension policy | `speculative.extension_initial_tokens
speculative.extension_max_tokens
speculative.extension_tail_backoff_proposals` | P1 | `resolver/speculative.rs` | `NgramExtensionConfig` | staged | restart/reload only | package extension policy or bounded runtime defaults | requires both native MTP and an N-gram proposer; config must satisfy the frontend plan validation | `#speculative-decoding` | resolver/config validation tests | Controls the adaptive suffix only. MTP and N-gram tokens remain one target-verified candidate. | +| 5.5 | Native MTP recovery | `speculative.native_mtp_reject_cooldown_tokens
speculative.native_mtp_suppress_cooldown_drafts
speculative.native_mtp_suppress_cooldown_draft_limit` | P2 | `resolver/speculative.rs` | `NativeMtpProposalConfig` | staged | restart/reload only | runtime defaults | non-negative integer bounds and boolean suppression flag | `#speculative-decoding` | resolver/config validation tests | Limits low-value MTP work after rejection without affecting target correctness. | +| 5.5 | VerifyWindow bounds and parallelism | `speculative.verify_window_min_tokens
speculative.verify_window_max_tokens
speculative.verify_window_pipeline_depth` | P1 | `resolver/speculative.rs` | `VerifyWindowConfig` | staged | restart/reload only | package policy or runtime defaults | `min_tokens <= max_tokens`; frontend validates the complete resolved plan | `#speculative-decoding` | resolver/config validation tests | Pipeline depth is the maximum number of target verification windows in flight. Request-local acceptance and stage timing can suppress dependent windows when expected overlap does not cover stale-work cost. It is not a separate draft source. | | 5.5 | Speculative default | `speculative.spec_default` | P2 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | runtime default | boolean or supported default-policy enum only; only `true` is rejected, while `false`/`auto`/unset are tolerated but do not produce executable behavior | `#speculative-decoding` | `.sisyphus/evidence/task-1-setting-matrix.txt`; negative: `.sisyphus/evidence/task-1-setting-matrix-error.txt` | raw default behavior switch for speculation; reserve until a real executable surface exists | ## Sampling and request defaults diff --git a/docs/skippy/DEAD_PEER_FAST_FAIL_PLAN.md b/docs/skippy/DEAD_PEER_FAST_FAIL_PLAN.md new file mode 100644 index 0000000000..cf047c7424 --- /dev/null +++ b/docs/skippy/DEAD_PEER_FAST_FAIL_PLAN.md @@ -0,0 +1,90 @@ +# Plan: fast-fail for new requests routed to a dead split stage + +## Problem (measured 2026-07-18) + +2-node WAN split, remote stage (vast) killed mid-serving: + +- **In-flight request:** fails fast, 502 in ~1.8s. Good — bounded by #1011 + (`open_return_sink_once` / `connect_lane_once` read timeouts). +- **New request after peer death:** coordinator still shows `peers:1, serving` + (stale peer state), routes to the dead stage, and hangs **~30s** before erroring. + +## Root cause (confirmed in code) + +Three timers interact; the split lane path uses the slow ones: + +1. `mesh/heartbeat.rs`: heartbeat runs **every 60s**, `failure_threshold = 2` + (direct) / `5` (relay-only). A dead direct peer is not declared down / evicted + for **~2 min**; relay peers ~5 min. Thresholds are deliberately lenient to + survive relay hiccups — do **not** tighten blindly. +2. New split requests trust that stale peer state and open lanes to the dead + stage. +3. `skippy-server/frontend.rs`: `connect_lane_once` bounds the ready handshake at + `LANE_READY_READ_TIMEOUT = 20s`, and `connect_binary_downstream` retries the + TCP connect `timeout_secs * 2` times × 500ms sleep. Combined worst case ≈ the + observed ~30s stall. + +## Fix (two layers; leave mesh timing alone) + +### Layer 1 — shorter first-attempt lane-open deadline (contained, unit-testable) + +`crates/skippy-server/src/frontend.rs` + +- Distinguish **cold-start / pool warmup** (where a longer wait is legitimate: + peer still loading) from **steady-state re-dial of a live split** (where a + healthy peer at 20ms RTT answers in ms). +- Add a separate, short deadline for steady-state `checkout()` reconnects, e.g. + `LANE_STEADY_CONNECT_TIMEOUT ≈ 3s`, distinct from the existing + `LANE_READY_READ_TIMEOUT = 20s` used during pool `new()`. +- Plumb it so `connect_lane_once` called from `checkout()` (mid-life) uses the + short deadline; pool construction keeps the long one. +- Effect: dead stage on a new request fails in ~3s, not ~30s. + +Risk: low. Only changes how long a *reconnect* waits. Cold-start unaffected. + +Test: unit test with a `TcpListener` that accepts but never sends ready (as in +`persistent_lane_ready_handshake_times_out_for_silent_downstream`) asserting the +steady-state path errors within the short bound. + +### Layer 2 — feed lane failures into `network/target_health.rs` (right layer, reuses existing machinery) + +`crates/mesh-llm-host-runtime/src/network/*` (routing side) + +- `target_health.rs` already exists: `record_outcome(... TargetHealthOutcome::Timeout | Unavailable ...)` + with a 30s base cooldown (`cooldown_for_failure`), and `eligible_candidates()` + / `strict_eligible_candidates()` for routing. +- When a split lane connect/handshake to a downstream stage fails, record a + `Timeout`/`Unavailable` outcome for that **stage target** so the next request + sees the cooldown and rejects immediately with a clear "stage unavailable" + instead of re-dialing the corpse. +- Requires wiring: the split coordinator's downstream-stage selection must + consult `target_health` before opening lanes, and report lane outcomes back + into it. Today `target_health` is used by the proxy/routing layer, not the + split lane pool — this is the missing connection. + +Risk: medium (touches routing/selection). Needs the live 2-node kill test to +confirm the cooldown actually engages on the split path and that a recovered peer +clears (via `record_reputation_success`). + +## Explicitly NOT doing (needs live multi-node validation first) + +- Tightening `heartbeat.rs` intervals / `failure_threshold`. Comments warn this + causes false-positive peer death on relay hiccups (Sydney↔Sydney relay spike to + 10s+ RTT). Out of scope for a fast-fail fix. +- Automatic re-planning onto a replacement worker. Bigger design conversation. + +## Validation gate (before committing either layer) + +Re-run the 2-node WAN kill test: + +1. Start split, confirm serving. +2. Fire a new request → baseline hang time. +3. Kill remote stage. +4. New request must fail in ≤ ~3s (Layer 1) with a clear stage-unavailable error + (Layer 2), not hang ~30s. +5. Restart remote stage → confirm the target cools back in and serving resumes. + +## Sequencing + +Bundle with the next dense-model / stage-count proof run (cached builds, ~\$0.30) +so both layers are validated live before landing. Do not commit blind. diff --git a/docs/skippy/PIPELINED_VERIFY_WINDOW.md b/docs/skippy/PIPELINED_VERIFY_WINDOW.md new file mode 100644 index 0000000000..1fe29d5e18 --- /dev/null +++ b/docs/skippy/PIPELINED_VERIFY_WINDOW.md @@ -0,0 +1,399 @@ +# Pipelined VerifyWindow Decode + +## Purpose + +This document describes Skippy's internal speculative-decode subsystem for +native multi-token prediction (MTP) and the optional MTP-anchored N-gram +extender. It covers the wire protocol, target-verification invariant, +asynchronous scheduling, operating modes, and diagnostic telemetry. + +The staged-runtime protocol deliberately has no compatibility path for the +retired synchronous `VerifySpan` message. Public mesh gossip and the +OpenAI-compatible API retain their normal compatibility guarantees. + +## Terms + +| Term | Meaning | +|---|---| +| Target | The full staged model, authoritative for every emitted token. | +| Native MTP | Model-provided typed draft attached to a target reply. GLM 4.7 Flash currently supplies a narrow `N+1` candidate. | +| N-gram sidecar | An upstream llama.cpp `ngram-simple` lookup or request-local `ngram-cache` proposer over target-committed tokens. | +| Composite proposal | Native-MTP prefix plus an optional N-gram suffix. | +| VerifyWindow | Versioned target request that verifies a candidate span at one session position. | +| Free target token | Target's next token after a fully verified span. | +| Stale window | Optimistic in-flight window invalidated by an earlier divergence. | + +## Safety Invariant + +MTP and N-gram are candidate sources, never authorities. The target verifies +each candidate sequence, and Skippy commits only the longest target-matching +prefix. A target correction is committed after a rejection. + +```mermaid +flowchart LR + MTP["Native MTP draft"] --> C["Composite candidate"] + NGRAM["N-gram continuation"] --> C + C --> W["VerifyWindow to target"] + W --> R{"Target result"} + R -->|"full accept"| A["Commit verified candidate\nand free target token"] + R -->|"partial accept"| P["Commit matching prefix\nthen target correction"] + R -->|"no proposal"| T["Ordinary target decode"] +``` + +This invariant also applies when multiple windows are in flight. + +## Wire Protocol + +`STAGE_STATE_VERSION` is `8`. `VerifyWindow` is wire message kind `21`; the +legacy kind `10` is rejected. An old/new staged-runtime pairing therefore fails +clearly instead of silently interpreting requests with different semantics. + +```mermaid +sequenceDiagram + participant S0 as "Stage 0 / OpenAI frontend" + participant ST as "Downstream stages + target" + Note over S0,ST: "Typed native MTP draft travels in target-reply sideband" + S0->>ST: "VerifyWindow(id, position, current + candidates)" + ST-->>S0: "PredictedTokens(window id, target tokens, next MTP draft)" + S0->>S0: "Classify longest matching candidate prefix" + S0->>S0: "Commit target-verified tokens only" +``` + +Pipelined decode requires direct prediction return. The target reply must reach +stage zero through the upstream-opened return sink; configuration fails when +that sink is unavailable. + +## Composite Proposals + +The sidecar extends native MTP; it never replaces it. Skippy uses upstream +llama.cpp proposers rather than a second Rust history scanner. `ngram-simple` +requires a historical continuation to begin with every MTP token before its +remaining tokens can become the sidecar tail. The request-local `ngram-cache` +instead reads directly after the provisional MTP prefix and returns only the +tail. + +```mermaid +flowchart TD + D["Receive typed MTP draft"] --> Q{"MTP tokens?"} + Q -->|"no"| P["Pure N-gram fallback\nfrom accepted context"] + Q -->|"yes"| H["Find latest earlier suffix match"] + H --> A{"Historical continuation\nstarts with full MTP prefix?"} + A -->|"yes"| X["Append useful suffix only"] + A -->|"no"| M["Keep MTP prefix only"] + P --> C["Composite proposal"] + X --> C + M --> C + C --> V["One VerifyWindow"] +``` + +For MTP `[a, b]`, a valid historical continuation must start `[a, b, ...]`. +The composite proposal becomes `[a, b, c, d]`, not two independent requests. +A one-token N-gram tail is discarded. A rejected tail does not count as an MTP +prefix rejection; it only backs off the sidecar. + +The cache is never shared between requests and is updated only after target +tokens commit. Drafting with `[a, b]` is read-only, so a rejected VerifyWindow +cannot affect a later lookup. This permits a cache tail to follow MTP even when +the cache would not independently predict `[a, b]`. + +## Adaptive Sidecar Policy + +The sidecar begins with the smallest useful tail. A fully accepted tail widens +the next tail by one token, up to the configured maximum. A rejected tail resets +the width and enters sidecar cooldown. With no MTP token, pure N-gram can use +the available N-gram budget. + +```mermaid +stateDiagram-v2 + [*] --> InitialTail + InitialTail --> WiderTail: "full tail accepted" + WiderTail --> WiderTail: "full tail accepted below maximum" + InitialTail --> Cooldown: "tail rejected" + WiderTail --> Cooldown: "tail rejected" + Cooldown --> NativeOnly: "cooldown proposal" + NativeOnly --> Cooldown: "cooldown remains" + NativeOnly --> InitialTail: "cooldown exhausted" +``` + +## Serial Native-MTP Mode + +Native MTP alone uses serial VerifyWindow processing. A window is opened, +verified, classified, and committed before the next window begins. + +```mermaid +sequenceDiagram + participant F as "Frontend" + participant T as "Target" + F->>T: "Window 41: current + MTP candidate" + T-->>F: "Window 41 reply" + F->>F: "Commit verified prefix" + F->>T: "Window 42: next candidate" + T-->>F: "Window 42 reply" +``` + +This is the native-MTP parity path. It is not decode parallelism by itself. + +## Pipelined Composite Mode + +`verify_window_pipeline_depth > 1` is a maximum rather than a command to keep +that many windows in flight. The request-local scheduler first measures full +acceptance, stage-zero compute time, and downstream wait by verify width. It +admits a dependent window only after enough observations show that expected +downstream overlap exceeds the cost of stale work. Otherwise the same composite +proposal uses one synchronous batched VerifyWindow. A deeper admitted proposal +is partitioned into FIFO windows. The target's free-advance candidate is +reserved as the next window's optimistic current token, preventing duplicate +KV positions. + +Profiles are independent by verify width and retain only recent observations. +An observation counts as a continuation only when the verified window and its +free target both match the buffered candidate. Admission compares expected +downstream overlap with the larger of local-compute or downstream stale-work +cost, including a safety margin. This lets WAN or downstream-heavy topologies +use configured depth while a stage-zero-heavy split remains on the profitable +synchronous batched path. + +```mermaid +sequenceDiagram + participant F as "Stage-zero frontend" + participant T as "Target" + Note over F: "Composite proposal: [m1, n1, n2, n3]" + F->>T: "Window 100 verifies [m1, n1]" + F->>T: "Window 101 verifies [n2, n3]" + T-->>F: "Window 100 reply" + F->>F: "Commit verified prefix" + T-->>F: "Window 101 reply" + F->>F: "Commit only if prefix remains valid" +``` + +Replies complete in FIFO window-id order. An earlier divergence invalidates +later optimistic windows. Skippy drains them, records them as stale, trims to +the committed target state, and resumes from the target correction. + +```mermaid +flowchart LR + W0["Earlier window\npartial accept"] --> C["Commit matching prefix\n+ correction"] + W0 --> D["Discard later windows\nas stale"] + D --> R["Trim/replay to target state"] + R --> N["Create candidate at corrected position"] +``` + +## Verification Outcomes + +| Target result | Committed output | Next action | +|---|---|---| +| Full accept | Candidate plus free target token where applicable | Continue; adaptive width may grow. | +| Tail rejection | MTP prefix and matching tail prefix, then correction | Back off sidecar only. | +| Prefix rejection | Matching prefix, then correction | Handle native MTP rejection and discard stale windows. | +| EOG | Verified prefix through EOG | Stop. | +| No candidate | Ordinary target token | Continue decode. | + +## Running On The Two-Host Lab + +Use the package-qualified model reference. The normal mesh runtime owns split +planning; do not replace it with a direct `gguf://` reference for this flow. + +The package owns a tested declarative default. `mesh-llm` resolves that package +plan once at launch, applies model-level settings before global defaults, and +passes the resulting typed configuration to `skippy-server`. The server does +not read `SKIPPY_NATIVE_MTP_*`, `SKIPPY_NGRAM_CACHE_*`, or +`SKIPPY_VERIFY_WINDOW_*` from its request hot path. Those variables are retired +from supported operation. + +### Package Strategy Shape + +`model-package.json` names reusable proposers and strategies. A GLM 4.7 Flash +package can expose native MTP plus a request-local cache sidecar as follows: + +```json +{ + "generation": { + "speculative_decoding": { + "default": "mtp-cache", + "proposers": { + "mtp": { + "type": "native-mtp", + "prediction_depth": 1, + "layer_indices": [47] + }, + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 10, + "history_scope": "request" + } + }, + "strategies": { + "mtp-cache": { + "type": "composite", + "primary": "mtp", + "extender": "cache", + "extension_policy": { + "initial_tokens": 2, + "max_tokens": 8, + "tail_backoff_proposals": 5 + } + } + } + } + } +} +``` + +Use the following stable strategy names when a package exposes the complete +native-MTP/N-gram menu: + +| Strategy | Composition | Benchmark condition | +|---|---|---| +| `mtp` | Native MTP proposer | MTP | +| `ngram-simple` | Pure prompt/history N-gram proposer | N-gram simple | +| `ngram-cache` | Pure request-local cache N-gram proposer | N-gram cache | +| `mtp-simple` | Native MTP primary plus simple N-gram tail | MTP + N-gram simple | +| `mtp-cache` | Native MTP primary plus request-local cache N-gram tail | MTP + N-gram cache | + +`disabled` is an operator control rather than a package strategy; it supplies +the no-MTP baseline. Every listed strategy is still target-verified. + +### Operator Configuration + +Choose a package strategy with `speculative.strategy`. `auto` uses the package +default; `disabled` turns speculation off; `mtp` preserves the direct native +MTP path. A named strategy such as `mtp-cache` is valid only when the selected +package declares it. Packages provide the recommended bounds and topology for +a model, while an explicit direct-GGUF configuration may select the built-in +simple or request-local cache N-gram proposer with its required bounds. + +```toml +[defaults.speculative] +strategy = "auto" + +[[models]] +model = "meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M" + +[models.speculative] +strategy = "mtp-cache" +ngram_max_proposal_tokens = 10 +extension_initial_tokens = 2 +extension_max_tokens = 8 +extension_tail_backoff_proposals = 5 +verify_window_min_tokens = 1 +verify_window_max_tokens = 6 +verify_window_pipeline_depth = 2 +``` + +### No MTP Baseline + +```bash +mesh-llm serve meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M --split --no-draft +``` + +Use `[models.speculative] strategy = "disabled"` to make this an explicit +baseline instead of relying on environment variables. + +### Native MTP Only + +```bash +mesh-llm serve meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M --split --no-draft +``` + +Use `[models.speculative] strategy = "mtp"` to force this control. + +### MTP With Cache-backed N-gram Extension + +```bash +mesh-llm serve meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M --split --no-draft +``` + +Use `[models.speculative] strategy = "mtp-cache"` with the bounded settings +above when the package declares that recommendation. For a direct GGUF, use +the built-in request-local cache proposer explicitly: + +```toml +[models.speculative] +strategy = "mtp" +ngram_proposer = "cache" +ngram_min = 2 +ngram_max = 4 +ngram_max_proposal_tokens = 6 +extension_max_tokens = 6 +``` + +With native MTP and an N-gram proposer present, mesh-llm creates the bounded +composite plan. The package remains the preferred way to publish tested values. + +### Invocation Overrides + +`mesh-llm serve` may temporarily override a package-selected strategy without +editing `config.toml`. CLI settings have highest precedence, then the selected +model entry, then `[defaults.speculative]`; unspecified CLI fields retain the +lower-layer value. Named package strategies remain package-declared. The CLI +may explicitly select the built-in simple or request-local cache proposer for +a direct GGUF only when it supplies valid N-gram bounds. + +```bash +mesh-llm serve meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M --split --no-draft \ + --speculative-strategy mtp \ + --speculative-ngram-proposer cache \ + --speculative-ngram-min 2 \ + --speculative-ngram-max 4 \ + --speculative-extension-max-tokens 8 \ + --speculative-verify-window-pipeline-depth 2 +``` + +The supported tuning flags are `--speculative-ngram-{min,max}`, +`--speculative-ngram-max-proposal-tokens`, +`--speculative-extension-{initial,max}-tokens`, +`--speculative-extension-tail-backoff-proposals`, +`--speculative-native-mtp-{reject-cooldown-tokens,suppress-cooldown-drafts,suppress-cooldown-draft-limit}`, +and `--speculative-verify-window-{min,max}-tokens` / `--speculative-verify-window-pipeline-depth`. +Use `--speculative-native-mtp-allow-cooldown-drafts` to explicitly override a +configured suppression policy to `false`. + +### Standalone Skippy Server + +`skippy-server` does not resolve layer-package recommendations. For isolated +stage-server operation it accepts a complete, already resolved JSON +`SpeculativeDecodeConfig` via `serve-binary --openai-speculative-config` or +`serve-openai --speculative-config`. The file is validated as one typed plan +before serving starts. This is intentionally not a second policy-merging path; +normal mesh serving always resolves the package and policy in `mesh-llm`. + +```mermaid +flowchart LR + C["SPEED-Bench client\non micstudio"] --> S0["micstudio :9337\nOpenAI frontend / stage 0"] + S0 --> L["Persistent direct-LAN\nbinary stage lanes"] + L --> S1["studio54\nstage 1"] + S1 --> S0 + S0 --> C +``` + +The normal planner currently selected `micstudio 0..47` and `studio54 47..48`. +Record layer ranges, direct RTT, lane count, context size, and binary commit +with every benchmark. That shape proves normal split serving but is not directly +comparable to historic 22/26 benchmark rows. + +## Telemetry And Interpretation + +The OpenAI response `timings` object provides aggregate evidence; debug +telemetry supplies per-window and per-stage detail. + +| Question | Counters | +|---|---| +| Is decode faster? | `predicted_per_second`, `predicted_n`, `predicted_ms` | +| Which plan actually ran? | `llama_stage.spec.requested_strategy`, `llama_stage.spec.effective_strategy` | +| Are proposals accepted? | `draft_n`, `draft_n_accepted` | +| Did the sidecar widen MTP? | `native_mtp_hybrid_native_tokens`, `native_mtp_hybrid_ngram_tokens`, `native_mtp_hybrid_proposed_tokens` | +| Did anchors agree? | `native_mtp_hybrid_ngram_mtp_prefix_agreements`, `native_mtp_hybrid_ngram_mtp_prefix_disagreements` | +| Were tails useful? | `native_mtp_hybrid_accepted_tail_tokens`, `native_mtp_hybrid_ngram_tail_rejections`, `native_mtp_hybrid_ngram_sidecar_backoffs` | +| Was it pipelined? | `verify_window_depth`, `verify_window_opened`, `verify_window_max_in_flight`, `verify_window_stale_discarded` | +| Why was depth used or suppressed? | `verify_window_policy_observed_windows`, `verify_window_policy_continuation_windows`, `verify_window_policy_profitable_widths`, `verify_window_policy_permit_checks`, `verify_window_policy_permits`, `verify_window_policy_suppressed` | +| Where was time spent? | `verify_window_downstream_wait_ms`, `verify_window_forward_write_ms`, `verify_window_stage0_compute_ms`, `verify_window_verify_elapsed_ms` | + +A useful hybrid run needs more than an increased `draft_n`: it needs accepted +tail tokens, high anchor agreement, bounded stale-window work, and completion +throughput higher than the native-MTP control. + +Pipeline-policy telemetry contains only bounded numeric counts and stage timing; +it does not include prompts, completions, token IDs, paths, endpoints, or node +identifiers. Debug OTLP export remains explicitly configured by the operator. diff --git a/docs/skippy/TOPOLOGY_PLANNER.md b/docs/skippy/TOPOLOGY_PLANNER.md index 465662151b..e5afa78dd6 100644 --- a/docs/skippy/TOPOLOGY_PLANNER.md +++ b/docs/skippy/TOPOLOGY_PLANNER.md @@ -262,3 +262,62 @@ contract: - exact recurrent-state transfer is an explicit opt-in policy, not the default; - exact activation wire defaults to `f16`; - `q8` is opt-in only when the relevant family/split has passed correctness. + +## Latency-Aware Placement: current behaviour and gaps + +This section records what the planner actually does today for network-aware +placement (verified against `crates/skippy-coordinator/src/topology.rs`, +`crates/skippy-topology/src/{lib,edge_order}.rs`, and the host-runtime call site +`crates/mesh-llm-host-runtime/src/inference/skippy/topology.rs`), and the gaps +that matter as the mesh grows to many nodes. + +### What works today + +- **Latency is a cost, not just an exclusion.** `node_package_score` subtracts + `rtt_ms × 16MiB` from a node's placement score, so higher-RTT nodes are + deprioritised but still usable. A relay-only peer (no direct path) is the one + case that is excluded outright. +- **The planner selects a node subset; it does not have to use every eligible + node.** `plan_topology` enumerates `node_count in minimum_nodes..=usable_nodes` + and, for each count, every node subset (`for_each_node_subset`). It can and + does leave nodes out. +- **Stage count is gated on a decode-TPOT target.** When any node has a measured + hop latency, planning is latency-aware: candidates are ranked + target-met-first, then by lowest `estimated_decode_network_ms_per_token` + (`latency_candidate_ordering` / `decode_tpot_target_met`, default target + `DEFAULT_TARGET_DECODE_TPOT_MS = 33`). A shallower split that meets the target + is preferred over a deeper split that does not — deeper is not chosen just + because it fits. + +### Gaps for the many-node / co-located future + +1. **No peer-to-peer RTT matrix in production.** `edge_order.rs` implements + adjacent-stage RTT ordering (exhaustive for ≤8 stages, greedy beyond) driven + by `StageEdgeSignal { source, target, rtt_ms }`. **The host runtime never + supplies edge signals** — it calls + `plan_package_aware_contiguous_with_signals` (no transport variant), so the + edge-matrix machinery is effectively dead in production. Placement sees only + each node's RTT *to the coordinator*, not node↔node RTT. +2. **The network estimate is a worst-case proxy, not a path sum.** + `estimate_decode_network_ms_per_token = max(per-node coordinator RTT) × + node_count`. It assumes every hop costs the slowest node's coordinator-RTT. + This is why **co-located nodes cannot be exploited**: two nodes close to each + other but each ~25 ms from the coordinator are modelled as ~25 ms hops, even + if their real A↔B hop is ~1 ms. Combined with vast.ai's NAT hairpin (two + containers behind one host cannot hole-punch and fall back to a ~200 ms + relay), co-located pipeline stages are currently not achievable there. +3. **`minimum_nodes` forces splitting in labs, but there is no first-class + "prefer fewer/none" cost knob beyond the TPOT gate.** For single-fit models + the TPOT gate already discourages needless splitting; for forced-split lab + runs `minimum_nodes` overrides it. Fine today, but a future many-node mesh + wants an explicit placement policy (e.g. "only add a stage if it improves + TPOT by X" and "never place on nodes above Y ms"). + +### Implications + +- Adding a real node↔node RTT matrix (wire `edge_signals` from measured + peer-to-peer RTT into `plan_package_aware_contiguous_with_transport`) is the + single change that would unlock co-location and correct many-node ordering. +- Until then, treat placement as **coordinator-RTT-aware only**: it will + deprioritise distant nodes and pick a sensible subset/stage-count, but it + cannot reason about which nodes are close to *each other*. diff --git a/docs/skippy/WAN_SPLIT_PERF.md b/docs/skippy/WAN_SPLIT_PERF.md new file mode 100644 index 0000000000..1a8a6380a6 --- /dev/null +++ b/docs/skippy/WAN_SPLIT_PERF.md @@ -0,0 +1,172 @@ +# WAN Split Performance Model + +This documents the throughput/latency model for Skippy stage-split serving over a +network, and the measured evidence behind it. Use it to predict whether a given +model + link + node count will be compute-bound or latency-bound, and therefore +whether adding stages helps or hurts. + +## Single-stream per-token cost + +For one in-flight request (generation concurrency = 1), the stages run +**serialized per token**: stage 0 computes its layers, forwards activations to +stage 1, ... the last stage produces the token, and the token/return path walks +back. Total compute across stages equals the whole model, so splitting does not +reduce single-stream compute — it only adds network hops. + +``` +TPOT ≈ C_total + (S - 1) · 2 · RTT + (S - 1) · P +``` + +- `TPOT` — time per output token (ms) +- `C_total` — compute time for all layers of one token (ms); ≈ solo single-GPU + decode ms/token, independent of `S` +- `S` — number of pipeline stages (nodes in the split) +- `RTT` — round-trip time between adjacent stages (ms) +- `2 · RTT` — activation-forward hop + token/return hop per inter-stage boundary +- `P` — per-boundary protocol/serialization overhead (ms) + +Single-stream throughput ceiling: + +``` +tok/s_max = 1000 / TPOT +``` + +### Compute-bound vs latency-bound + +Define the network share of a 2-stage split: + +``` +network_fraction = (2 · RTT + P) / TPOT +``` + +- `network_fraction` high (→1): **latency-bound**. Adding stages makes it worse + (each stage adds `2·RTT + P`). Minimize stage count; use speculation. +- `network_fraction` low (→0): **compute-bound**. `C_total` dominates and is not + reduced by splitting single-stream — but concurrent/batched load benefits from + pipeline overlap (see below). + +## Measured evidence (2026-07-18) + +2-node split, `meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M` (48 layers, MoE), +M5 Max (Metal, layers 0..25) ↔ RTX A5000 Melbourne (CUDA, layers 25..48), +direct iroh hole-punch, RTT ≈ 20 ms. + +| Quantity | Value | +| --- | ---: | +| Solo GPU decode (all layers, no network) `C_total` | 12.9 ms/token (77.3 tok/s) | +| 2-stage split observed `TPOT` | 57.8 ms/token (~17 tok/s) | +| Split overhead vs solo | +44.9 ms/token | +| `2 · RTT` at 20 ms | 40 ms | +| Implied protocol `P` | ~4.9 ms | + +Decomposition of the 57.8 ms/token: `12.9 (compute) + 40 (2·RTT) + 4.9 (protocol) += 57.8`. The model closes to the measured value. + +**Verdict for this workload: latency-bound.** Compute was ~22% of per-token time; +network round-trips were ~70%. A small MoE (few active params/token) has low +`C_total`, so `network_fraction ≈ 0.78`. + +### Consequence: adding a stage here would slow it down + +At RTT 20 ms, each extra stage adds ~44.9 ms/token. Going 2→3 stages: + +``` +TPOT(3) ≈ 12.9 + 2 · (2·20 + 4.9) = 12.9 + 89.8 = 102.7 ms/token (~9.7 tok/s) +``` + +Only justified if a 3rd stage is needed for memory residency, or under +concurrent load, or if `C_total` per stage is large (dense big models). + +## When does adding a stage help? + +### 1. Memory residency +If the model + KV cache does not fit in `S-1` nodes' VRAM, you must add stages. +This is a correctness constraint, not a speed choice. + +### 2. Concurrent / batched throughput (pipeline overlap) +With `N` concurrent requests and a pipeline of `S` stages, stages work on +different requests simultaneously. Aggregate throughput approaches: + +``` +throughput_agg ≈ min(N, S) · (1 / max_stage_compute) (network-hidden regime) +``` + +when compute per stage ≥ inter-stage latency, so hops overlap with compute. This +is the regime where more stages raise aggregate tok/s — but it needs enough +concurrency and enough per-stage compute to hide `2·RTT`. + +### 3. Compute-bound models (dense, large) +For a dense model, `C_total` is large and splitting across `S` stages makes each +stage's compute ≈ `C_total / S`. Adding a stage helps single-stream throughput +only while per-stage compute still dominates the added `2·RTT`: + +``` +add a stage if: C_total / (S·(S+1)) > 2·RTT + P +``` + +i.e. the compute saved per token by finer splitting must exceed the latency of +the new hop. Low-RTT links and heavy dense models satisfy this; WAN + small MoE +do not. + +## Speculation as the WAN lever + +Speculative decoding (native MTP + N-gram) commits `k` tokens per verify +round-trip, amortizing the fixed `2·RTT` boundary cost across accepted tokens: + +``` +TPOT_spec ≈ C_verify + (S-1)·2·RTT / k_accepted + (S-1)·P +``` + +where `k_accepted` is the mean accepted tokens per verify window. This is why the +measured MTP + N-gram gain grows on longer/coding output (higher `k_accepted`): ++13–19% over MTP-off at RTT 20 ms. + +## Planning checklist + +1. Estimate `C_total` from a solo run (or per-layer compute × layers). +2. Measure `RTT` between candidate nodes. +3. Compute `network_fraction` for `S=2`. +4. If latency-bound: use `S=2`, enable speculation, prefer lower-RTT peers. +5. If compute-bound or memory-forced: increase `S`, verify + `C_total/(S·(S+1)) > 2·RTT + P` still holds before adding each stage. +6. Under concurrency, prefer more stages for aggregate throughput once per-stage + compute ≥ `2·RTT`. + +## Speculative recovery cost over WAN (measured) + +Enabling N-gram speculation on a large MoE (MiniMax-M2.7, 2-stage M5↔Melbourne, +~18ms RTT, agentic OSL512) was **~40% slower** than no speculation (≈15 vs ≈17 +tok/s) despite a **high 0.89 token accept rate**. The cause is recovery cost on +rejected windows, not proposal quality. + +Measured (metrics-server, coordinator stage-0 spans): + +- accept_rate 0.89, accepted 660 / proposed 2214 +- full_accept_windows 45, **early_reject_windows 231**, rejected_windows 234 +- **recovery_restores 231**, recovery_ms ≈ 18,971, recovery_restore_downstream_wait_ms ≈ 4,620 +- **window_shrinks 0** despite 231 early rejects — the adaptive window stayed at + 12 the entire run + +Mechanism (verified in `frontend/embedded_generation.rs` + +`frontend/embedded_execution.rs`): each early-reject window pays **two serial +extra WAN round-trips** — `restore_embedded_stage_session` (write + wait for +ACK) then the repair `execute_embedded_stage_message` (write + wait for reply). +A rejected window therefore costs ~3 round-trips to commit what plain decode +commits in 1. Over WAN those round-trips dominate. + +The smoking gun is `window_shrinks 0`: the adaptive policy never narrowed the +window under a sustained reject storm, so it kept proposing deep, kept +rejecting, and kept paying the 3× round-trip recovery. + +Fix directions (not required for merge — speculation is opt-in and off by +default; this only affects deployments that explicitly enable N-gram): + +1. Adaptive window must actually shrink on early-reject (it did not here). +2. Fuse restore+repair into a single round-trip instead of two serial waits. +3. For single-token repair (`repair_input_count == 1`), avoid the separate + restore round-trip entirely (piggyback the correction). + +Takeaway for large-model MTP work: over WAN, speculation only pays off when the +reject/recovery rate stays low. Native MTP (high acceptance) plus a +reject-responsive adaptive window is the regime to target; blind deep N-gram +proposal on a latency-bound split is counterproductive. diff --git a/docs/skippy/family/qwen-results.md b/docs/skippy/family/qwen-results.md index d9178c0bea..31fc99c65b 100644 --- a/docs/skippy/family/qwen-results.md +++ b/docs/skippy/family/qwen-results.md @@ -663,7 +663,7 @@ Per-window timing: | draft adaptive max `8` | `365.43/812.61ms` | `153.54/524.55ms` | `188.98/197.73ms` | Interpretation: speculation is doing the thing we wanted: it replaces roughly -`195` single-token staged decode turns with tens of `VerifySpan` requests, and +`195` single-token staged decode turns with tens of `VerifyWindow` requests, and the draft is usually right. The limiting cost is now local draft decode time on stage0, not target verification acceptance. Larger windows reduce target verify count but increase draft cost and rejection cost. The next useful step diff --git a/docs/skippy/speculative_decoding.md b/docs/skippy/speculative_decoding.md index 7bbef4b562..6b5fd237e8 100644 --- a/docs/skippy/speculative_decoding.md +++ b/docs/skippy/speculative_decoding.md @@ -9,7 +9,7 @@ in the family result documents under this directory. N-gram speculative decoding is implemented and useful, especially for repeated coding/editing sessions. It is model-free: the pool observes accepted target tokens, proposes continuations when a context suffix repeats, and the staged -target verifies every proposed token through `VerifySpan`. +target verifies every proposed token through `VerifyWindow`. Current policy: @@ -48,7 +48,7 @@ Open items: - Make repair decisions cost-aware, not only confidence/window-size aware. - Preserve the tail-reject fast path. -- Avoid repair `VerifySpan` when a normal decode step is cheaper. +- Avoid repair `VerifyWindow` when a normal decode step is cheaper. - Track repair cost by task type, not only globally. ### Pool Policy And Lifetime diff --git a/docs/specs/layer-package-repos.md b/docs/specs/layer-package-repos.md index 7ad2fdbd04..766b57358d 100644 --- a/docs/specs/layer-package-repos.md +++ b/docs/specs/layer-package-repos.md @@ -275,6 +275,7 @@ configured. When present, it may declare `speculative_decoding` defaults: - `default`: the strategy id the package recommends for this distribution. +- `proposers`: a map of reusable proposal-source ids to configuration. - `strategies`: a map of strategy id to strategy configuration. The current native MTP strategy shape is: @@ -293,6 +294,77 @@ The current native MTP strategy shape is: } ``` +New packages SHOULD name the native MTP source under `proposers` and reference +it from the strategy. This also permits a package to add an N-gram sidecar +without changing the MTP source: + +```json +{ + "default": "mtp-cache", + "proposers": { + "mtp": { + "type": "native-mtp", + "prediction_depth": 1, + "layer_indices": [47] + }, + "cache": { + "type": "ngram-cache", + "ngram_min": 2, + "ngram_max": 4, + "max_proposal_tokens": 10, + "history_scope": "request" + } + }, + "strategies": { + "mtp-cache": { + "type": "composite", + "primary": "mtp", + "extender": "cache", + "extension_policy": { + "initial_tokens": 2, + "max_tokens": 8, + "tail_backoff_proposals": 5 + } + } + } +} +``` + +Supported proposer types are `native-mtp`, `ngram-simple`, and +`ngram-cache`. An `ngram-cache` proposer MUST use `history_scope: "request"` +and `ngram_max` no greater than `4`, llama.cpp's current cache match-window +limit. It contains only target-committed history for one request and is never shared +between users or sessions. A `composite` strategy MUST use a `native-mtp` +primary and an N-gram extender. Its `extension_policy` bounds the adaptive +tail; every combined candidate is still verified by one target VerifyWindow. + +The package schema separates a proposer match length from its output budget: + +| Field | Applies to | Requirement | +|---|---|---| +| `prediction_depth` | `native-mtp` | Must be `1` for the current native MTP runtime. | +| `layer_indices` | `native-mtp` | Must identify the package layers that contain the model's NextN/MTP tensors. | +| `ngram_min` / `ngram_max` | N-gram proposers | Define the historical token match range. Both are required and `ngram_min <= ngram_max`. | +| `max_proposal_tokens` | N-gram proposers | Caps how many continuation tokens the proposer may return. It is independent of `ngram_max`. | +| `history_scope` | `ngram-cache` | Must be `"request"`; a cache proposer never observes another request's tokens. | +| `initial_tokens` / `max_tokens` | composite extension policy | Bound the adaptive N-gram tail after an MTP prefix. | +| `tail_backoff_proposals` | composite extension policy | Sets how many proposals to back off after an unhelpful tail. | + +`ngram-simple` looks up continuations in the accepted prompt/history. It can +be published for any compatible tokenizer. `ngram-cache` is a request-local +incremental lookup that can also start after a provisional MTP prefix; its +match window is intentionally capped at four tokens by the current llama.cpp +ABI. Neither N-gram proposer is authoritative: the target verifies the MTP +prefix and N-gram suffix together, then commits the accepted prefix only. + +Packages that expose the full product menu SHOULD use stable strategy ids: +`mtp`, `ngram-simple`, `ngram-cache`, `mtp-simple`, and `mtp-cache`. `mtp` +may reference a reusable `native-mtp` proposer instead of repeating its +prediction-depth and layer metadata. The N-gram-only names use `proposer`; the +two composite names use that MTP proposer as `primary` and the corresponding +N-gram proposer as `extender`. `disabled` is a runtime/operator baseline, not +a package strategy. + Native MTP strategy rules: - `type` MUST be `native-mtp`. @@ -309,12 +381,27 @@ Consumers that do not recognize a strategy type MUST ignore it unless it is the declared default for a request they are trying to serve. Operators may override the package recommendation in `config.toml` with -`speculative.strategy`. Supported values are `auto` (use package/runtime -defaults), `mtp` (force the current native MTP strategy), and -`disabled` (disable native MTP for the configured model/default scope). +`speculative.strategy`, or for one `mesh-llm serve` invocation with +`--speculative-strategy`. `auto` uses the package/runtime default, `mtp` forces +the direct native-MTP control, and `disabled` disables speculation. A named +package strategy such as `mtp-cache` is accepted only when the selected package +declares it. Precedence is CLI invocation, selected model entry, then global +defaults. These layers may bound N-gram proposal size, extension depth, +cooldown, and VerifyWindow depth. A named package strategy cannot be invented +outside its package. For direct GGUF operation, operators may explicitly select +the built-in `ngram-simple` or request-local `ngram-cache` proposer by supplying +valid N-gram bounds; mesh-llm constructs and validates that generic plan before +starting Skippy. `skippy-server` receives the resulting typed plan and does not +repeat this policy resolution. Operators may also pass the legacy `native-mtp-n1` value; the runtime normalizes it to `mtp` for backward compatibility. New configs should use `mtp`. +See [Speculative Decode Configuration](../skippy/PIPELINED_VERIFY_WINDOW.md#operator-configuration) +for the operator-side `config.toml` controls and CLI equivalents. Package +authors should publish conservative, tested values in the manifest; operators +can use configuration to choose a strategy or tighten its bounds without +changing the package's declared topology. + ## Layer Selection For a stage with `layer_start..layer_end`, consumers select: diff --git a/evals/skippy-coding-agent-loop.jsonl b/evals/skippy-coding-agent-loop.jsonl new file mode 100644 index 0000000000..fe384b4d22 --- /dev/null +++ b/evals/skippy-coding-agent-loop.jsonl @@ -0,0 +1 @@ +{"id":"coding-agent-loop","category":"coding_agent_loop","session_group":"skippy-proposer-matrix","prompt":"Continue this coding-agent task transcript and choose the next concrete action.\nRepository task: MoneyWidget decompress method breaks form validation with disabled fields. The task asks the agent to inspect /testbed/djmoney/forms/widgets.py, reproduce the issue, edit non-test source files only, and rerun validation. Current transcript shows the agent found /testbed/djmoney/forms/widgets.py and observed that MoneyWidget.decompress immediately returns [None, self.default_currency] before checking whether value is not None. Continue with the next concrete coding-agent action."} diff --git a/scripts/build-mac.sh b/scripts/build-mac.sh index 319a0147e0..8418b0d3b0 100755 --- a/scripts/build-mac.sh +++ b/scripts/build-mac.sh @@ -477,7 +477,16 @@ EOF echo " security find-identity -v -p codesigning" >&2 } -export LLAMA_STAGE_BUILD_DIR="${LLAMA_STAGE_BUILD_DIR:-${SKIPPY_LLAMA_BUILD_DIR:-$LLAMA_BUILD_ROOT/build-stage-abi-metal}}" +if [[ -z "${LLAMA_STAGE_BUILD_DIR:-}" && -n "${SKIPPY_LLAMA_BUILD_DIR:-}" ]]; then + export LLAMA_STAGE_BUILD_DIR="$SKIPPY_LLAMA_BUILD_DIR" +fi +if [[ -z "${LLAMA_STAGE_BUILD_DIR:-}" ]]; then + export LLAMA_STAGE_BUILD_DIR="$( + LLAMA_STAGE_BACKEND="${LLAMA_STAGE_BACKEND:-metal}" \ + LLAMA_STAGE_LINK_MODE=static \ + "$SCRIPT_DIR/build-llama.sh" --print-build-dir + )" +fi configure_lld_linker diff --git a/scripts/build-release.sh b/scripts/build-release.sh index 1baca3984d..9c120bf2c7 100755 --- a/scripts/build-release.sh +++ b/scripts/build-release.sh @@ -138,11 +138,11 @@ if [[ -z "${LLAMA_STAGE_BUILD_DIR:-}" && -n "${SKIPPY_LLAMA_BUILD_DIR:-}" ]]; th export LLAMA_STAGE_BUILD_DIR="$SKIPPY_LLAMA_BUILD_DIR" fi if [[ -z "${LLAMA_STAGE_BUILD_DIR:-}" ]]; then - if [[ "$BACKEND" == "cpu" ]]; then - export LLAMA_STAGE_BUILD_DIR="$LLAMA_BUILD_ROOT/build-stage-abi-static" - else - export LLAMA_STAGE_BUILD_DIR="$LLAMA_BUILD_ROOT/build-stage-abi-$BACKEND" - fi + export LLAMA_STAGE_BUILD_DIR="$( + LLAMA_STAGE_BACKEND="$BACKEND" \ + LLAMA_STAGE_LINK_MODE=static \ + "$SCRIPT_DIR/build-llama.sh" --print-build-dir + )" fi configure_lld_linker @@ -174,4 +174,8 @@ case "$BACKEND" in rocm) cargo_features+=(--features gpu-bench-hip) ;; esac stamp_build_version -(cd "$REPO_ROOT" && cargo build --release --locked -p mesh-llm "${cargo_features[@]}") +if ((${#cargo_features[@]})); then + (cd "$REPO_ROOT" && cargo build --release --locked -p mesh-llm "${cargo_features[@]}") +else + (cd "$REPO_ROOT" && cargo build --release --locked -p mesh-llm) +fi diff --git a/scripts/tests/test_build_release.py b/scripts/tests/test_build_release.py index 7eccbceee5..c44bf1bdea 100644 --- a/scripts/tests/test_build_release.py +++ b/scripts/tests/test_build_release.py @@ -46,7 +46,15 @@ def test_rocm_release_build_passes_hip_gpu_benchmark_feature_to_cargo(self) -> N self.assertIn("--features gpu-bench-hip", cargo_log) self.assertIn("--features dynamic-native-runtime", cargo_log) - def run_build_release_with_backend(self, backend: str) -> str: + def test_static_release_build_handles_an_empty_feature_list(self) -> None: + cargo_log = self.run_build_release_with_backend("metal", dynamic_native_runtime=False) + + self.assertIn("build --release --locked -p mesh-llm", cargo_log) + self.assertNotIn("--features", cargo_log) + + def run_build_release_with_backend( + self, backend: str, *, dynamic_native_runtime: bool = True + ) -> str: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) scripts_dir = tmp / "scripts" @@ -66,6 +74,22 @@ def run_build_release_with_backend(self, backend: str) -> str: echo "stub build-ui $*" """, ) + self.write_executable( + scripts_dir / "prepare-llama.sh", + """ + #!/usr/bin/env bash + set -euo pipefail + echo "stub prepare-llama $*" + """, + ) + self.write_executable( + scripts_dir / "build-llama.sh", + """ + #!/usr/bin/env bash + set -euo pipefail + echo "stub build-llama $*" + """, + ) self.write_executable( bin_dir / "uname", """ @@ -114,7 +138,9 @@ def run_build_release_with_backend(self, backend: str) -> str: { "CARGO_LOG": str(cargo_log), "LLAMA_STAGE_BACKEND": backend, - "MESH_LLM_DYNAMIC_NATIVE_RUNTIME": "1", + "MESH_LLM_DYNAMIC_NATIVE_RUNTIME": "1" + if dynamic_native_runtime + else "0", "PATH": f"{bin_dir}{os.pathsep}{env['PATH']}", } ) diff --git a/third_party/llama.cpp/patches/0017-Expose-upstream-ngram-simple-draft-ABI.patch b/third_party/llama.cpp/patches/0017-Expose-upstream-ngram-simple-draft-ABI.patch new file mode 100644 index 0000000000..faabda864c --- /dev/null +++ b/third_party/llama.cpp/patches/0017-Expose-upstream-ngram-simple-draft-ABI.patch @@ -0,0 +1,132 @@ +From 116925a64b0bee9ebd28681c41d557191a91c389 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 09:03:13 +1000 +Subject: [PATCH 17/17] Expose upstream ngram simple draft ABI + +--- + common/ngram-map.cpp | 38 ++++++++++++++++++++++++++++++++++++++ + include/skippy.h | 14 ++++++++++++++ + include/skippy/common.h | 3 ++- + src/skippy.cpp | 3 ++- + 4 files changed, 56 insertions(+), 2 deletions(-) + +diff --git a/common/ngram-map.cpp b/common/ngram-map.cpp +index d1cb4519..f6e0a372 100644 +--- a/common/ngram-map.cpp ++++ b/common/ngram-map.cpp +@@ -1,7 +1,9 @@ + #include "common.h" + #include "log.h" + #include "ngram-map.h" ++#include "../include/skippy.h" + ++#include + #include + #include + #include +@@ -111,6 +113,42 @@ llama_tokens common_ngram_simple_draft( + return draft_tokens; + } + ++extern "C" enum skippy_status skippy_ngram_simple_draft( ++ const llama_token * token_ids, ++ size_t token_count, ++ llama_token sampled_token, ++ uint16_t ngram_size, ++ uint16_t max_draft_tokens, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** /* out_error */) { ++ if (out_token_count == nullptr) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ *out_token_count = 0; ++ if (ngram_size == 0 || max_draft_tokens == 0) { ++ return SKIPPY_STATUS_OK; ++ } ++ if (token_ids == nullptr || token_count == 0) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ const auto draft_tokens = common_ngram_simple_draft( ++ { ngram_size, max_draft_tokens }, ++ { token_ids, token_ids + token_count }, ++ sampled_token); ++ *out_token_count = draft_tokens.size(); ++ if (draft_tokens.empty()) { ++ return SKIPPY_STATUS_OK; ++ } ++ if (output_tokens == nullptr || output_token_capacity < draft_tokens.size()) { ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ std::copy(draft_tokens.begin(), draft_tokens.end(), output_tokens); ++ return SKIPPY_STATUS_OK; ++} ++ + + // n-gram map + // +diff --git a/include/skippy.h b/include/skippy.h +index c5a4c4b7..f8511885 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -144,6 +144,20 @@ struct skippy_native_mtp_draft { + int64_t proposal_compute_us; + }; + ++// Runs llama.cpp's ngram-simple self-speculative proposer. `token_ids` is the ++// accepted token history excluding `sampled_token`; the proposer treats that ++// token as the current sample and returns only its proposed continuation. ++LLAMA_API enum skippy_status skippy_ngram_simple_draft( ++ const llama_token * token_ids, ++ size_t token_count, ++ llama_token sampled_token, ++ uint16_t ngram_size, ++ uint16_t max_draft_tokens, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_model_open( + const char * path, + const struct skippy_runtime_config * config, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index b7186bc9..c9c310c0 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 30 ++#define SKIPPY_ABI_VERSION_PATCH 31 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +@@ -55,6 +55,7 @@ enum skippy_feature { + SKIPPY_FEATURE_BACKEND_DEVICES = 1 << 23, + SKIPPY_FEATURE_RUNTIME_EVENTS = 1 << 24, + SKIPPY_FEATURE_NATIVE_MTP_N1 = 1 << 25, ++ SKIPPY_FEATURE_NGRAM_SIMPLE_DRAFT = 1 << 26, + }; + + enum skippy_status { +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 5f5596ce..1344df3a 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2458,7 +2458,8 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_EXTERNAL_MEDIA_PREFILL | + SKIPPY_FEATURE_RUNTIME_EVENTS | + SKIPPY_FEATURE_BACKEND_DEVICES | +- SKIPPY_FEATURE_NATIVE_MTP_N1; ++ SKIPPY_FEATURE_NATIVE_MTP_N1 | ++ SKIPPY_FEATURE_NGRAM_SIMPLE_DRAFT; + } + + void skippy_error_free(struct skippy_error * error) { +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0018-Expose-stateful-N-gram-cache-ABI.patch b/third_party/llama.cpp/patches/0018-Expose-stateful-N-gram-cache-ABI.patch new file mode 100644 index 0000000000..7120e7b25b --- /dev/null +++ b/third_party/llama.cpp/patches/0018-Expose-stateful-N-gram-cache-ABI.patch @@ -0,0 +1,258 @@ +From 45b9de4d71288fafc9edefd4676ae89bce851145 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 09:23:50 +1000 +Subject: [PATCH 18/18] Expose stateful N-gram cache ABI + +--- + common/ngram-cache.cpp | 136 ++++++++++++++++++++++++++++++++++++++++ + include/skippy.h | 35 +++++++++++ + include/skippy/common.h | 3 +- + src/skippy.cpp | 3 +- + 4 files changed, 175 insertions(+), 2 deletions(-) + +diff --git a/common/ngram-cache.cpp b/common/ngram-cache.cpp +index dce54b36..a0447210 100644 +--- a/common/ngram-cache.cpp ++++ b/common/ngram-cache.cpp +@@ -1,14 +1,150 @@ + #include "ngram-cache.h" + #include "common.h" + #include "log.h" ++#include "../include/skippy.h" + + #include + #include + #include + #include ++#include + #include + #include + ++struct skippy_ngram_cache { ++ uint16_t ngram_min; ++ uint16_t ngram_max; ++ std::vector history; ++ common_ngram_cache context; ++ common_ngram_cache dynamic; ++ common_ngram_cache static_cache; ++}; ++ ++static bool skippy_ngram_cache_valid_config(uint16_t ngram_min, uint16_t ngram_max) { ++ return ngram_min >= LLAMA_NGRAM_MIN && ngram_min <= ngram_max && ngram_max <= LLAMA_NGRAM_MAX; ++} ++ ++static enum skippy_status skippy_ngram_cache_update( ++ skippy_ngram_cache * cache, ++ const llama_token * token_ids, ++ size_t token_count, ++ bool reset) { ++ if (cache == nullptr || (token_count > 0 && token_ids == nullptr)) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (token_count == 0) { ++ if (reset) { ++ cache->history.clear(); ++ cache->context.clear(); ++ } ++ return SKIPPY_STATUS_OK; ++ } ++ if (token_count > static_cast(std::numeric_limits::max()) || ++ cache->history.size() > static_cast(std::numeric_limits::max()) - token_count) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (reset) { ++ cache->history.assign(token_ids, token_ids + token_count); ++ cache->context.clear(); ++ common_ngram_cache_update( ++ cache->context, ++ cache->ngram_min, ++ cache->ngram_max, ++ cache->history, ++ static_cast(token_count), ++ false); ++ return SKIPPY_STATUS_OK; ++ } ++ cache->history.insert(cache->history.end(), token_ids, token_ids + token_count); ++ common_ngram_cache_update( ++ cache->context, ++ cache->ngram_min, ++ cache->ngram_max, ++ cache->history, ++ static_cast(token_count), ++ false); ++ return SKIPPY_STATUS_OK; ++} ++ ++extern "C" enum skippy_status skippy_ngram_cache_create( ++ uint16_t ngram_min, ++ uint16_t ngram_max, ++ struct skippy_ngram_cache ** out_cache, ++ struct skippy_error ** /* out_error */) { ++ if (out_cache == nullptr || !skippy_ngram_cache_valid_config(ngram_min, ngram_max)) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ *out_cache = new skippy_ngram_cache{ ngram_min, ngram_max, {}, {}, {}, {} }; ++ return SKIPPY_STATUS_OK; ++} ++ ++extern "C" void skippy_ngram_cache_free(struct skippy_ngram_cache * cache) { ++ delete cache; ++} ++ ++extern "C" enum skippy_status skippy_ngram_cache_reset( ++ struct skippy_ngram_cache * cache, ++ const llama_token * token_ids, ++ size_t token_count, ++ struct skippy_error ** /* out_error */) { ++ return skippy_ngram_cache_update(cache, token_ids, token_count, true); ++} ++ ++extern "C" enum skippy_status skippy_ngram_cache_append( ++ struct skippy_ngram_cache * cache, ++ const llama_token * token_ids, ++ size_t token_count, ++ struct skippy_error ** /* out_error */) { ++ return skippy_ngram_cache_update(cache, token_ids, token_count, false); ++} ++ ++extern "C" enum skippy_status skippy_ngram_cache_draft( ++ struct skippy_ngram_cache * cache, ++ const llama_token * continuation_prefix, ++ size_t continuation_prefix_count, ++ uint16_t max_draft_tokens, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** /* out_error */) { ++ if (out_token_count == nullptr || cache == nullptr || ++ (continuation_prefix_count > 0 && continuation_prefix == nullptr)) { ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ *out_token_count = 0; ++ if (max_draft_tokens == 0 || cache->history.empty()) { ++ return SKIPPY_STATUS_OK; ++ } ++ ++ std::vector input = cache->history; ++ std::vector draft; ++ if (continuation_prefix_count == 0) { ++ draft.push_back(input.back()); ++ } else { ++ input.insert(input.end(), continuation_prefix, continuation_prefix + continuation_prefix_count); ++ draft.push_back(continuation_prefix[continuation_prefix_count - 1]); ++ } ++ common_ngram_cache_draft( ++ input, ++ draft, ++ max_draft_tokens, ++ cache->ngram_min, ++ cache->ngram_max, ++ cache->context, ++ cache->dynamic, ++ cache->static_cache); ++ const size_t drafted_count = draft.size() - 1; ++ *out_token_count = drafted_count; ++ if (drafted_count == 0) { ++ return SKIPPY_STATUS_OK; ++ } ++ if (output_tokens == nullptr || output_token_capacity < drafted_count) { ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ std::copy(draft.begin() + 1, draft.end(), output_tokens); ++ return SKIPPY_STATUS_OK; ++} ++ + void common_ngram_cache_update(common_ngram_cache & ngram_cache, int ngram_min, int ngram_max, + std::vector & inp, int nnew, bool print_progress) { + const int64_t t_start_ms = ggml_time_ms(); +diff --git a/include/skippy.h b/include/skippy.h +index f8511885..c9ecc349 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -48,6 +48,7 @@ struct skippy_model; + struct skippy_session; + struct skippy_model_info; + struct skippy_slice_plan; ++struct skippy_ngram_cache; + + struct skippy_runtime_config { + int32_t stage_index; +@@ -158,6 +159,40 @@ LLAMA_API enum skippy_status skippy_ngram_simple_draft( + size_t * out_token_count, + struct skippy_error ** out_error); + ++// Stateful, request-local adapter for llama.cpp's ngram-cache proposer. ++// The cache is updated only with committed history. `continuation_prefix` ++// supplies optional provisional tokens (such as native MTP output) used only ++// while drafting; it never mutates the cache. ++LLAMA_API enum skippy_status skippy_ngram_cache_create( ++ uint16_t ngram_min, ++ uint16_t ngram_max, ++ struct skippy_ngram_cache ** out_cache, ++ struct skippy_error ** out_error); ++ ++LLAMA_API void skippy_ngram_cache_free(struct skippy_ngram_cache * cache); ++ ++LLAMA_API enum skippy_status skippy_ngram_cache_reset( ++ struct skippy_ngram_cache * cache, ++ const llama_token * token_ids, ++ size_t token_count, ++ struct skippy_error ** out_error); ++ ++LLAMA_API enum skippy_status skippy_ngram_cache_append( ++ struct skippy_ngram_cache * cache, ++ const llama_token * token_ids, ++ size_t token_count, ++ struct skippy_error ** out_error); ++ ++LLAMA_API enum skippy_status skippy_ngram_cache_draft( ++ struct skippy_ngram_cache * cache, ++ const llama_token * continuation_prefix, ++ size_t continuation_prefix_count, ++ uint16_t max_draft_tokens, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_model_open( + const char * path, + const struct skippy_runtime_config * config, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index c9c310c0..771d1c78 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 31 ++#define SKIPPY_ABI_VERSION_PATCH 32 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +@@ -56,6 +56,7 @@ enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_EVENTS = 1 << 24, + SKIPPY_FEATURE_NATIVE_MTP_N1 = 1 << 25, + SKIPPY_FEATURE_NGRAM_SIMPLE_DRAFT = 1 << 26, ++ SKIPPY_FEATURE_NGRAM_CACHE_DRAFT = 1 << 27, + }; + + enum skippy_status { +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 1344df3a..46da7007 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2459,7 +2459,8 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_RUNTIME_EVENTS | + SKIPPY_FEATURE_BACKEND_DEVICES | + SKIPPY_FEATURE_NATIVE_MTP_N1 | +- SKIPPY_FEATURE_NGRAM_SIMPLE_DRAFT; ++ SKIPPY_FEATURE_NGRAM_SIMPLE_DRAFT | ++ SKIPPY_FEATURE_NGRAM_CACHE_DRAFT; + } + + void skippy_error_free(struct skippy_error * error) { +-- +2.54.0 (Apple Git-157) + diff --git a/website/src/docs/pages/CLI.md b/website/src/docs/pages/CLI.md index e7a0008f1a..0f464aea2f 100644 --- a/website/src/docs/pages/CLI.md +++ b/website/src/docs/pages/CLI.md @@ -208,6 +208,36 @@ Switches: - `--trust-policy `: override peer ownership trust policy. - `--trust-owner `: add trusted owner IDs on top of the local trust store. +### Speculative decoding overrides + +Advanced `serve` invocations can temporarily override a package or config-file +speculative decoding plan. CLI values have highest precedence; fields you omit +continue to come from the selected model, defaults, or model package. + +```bash +mesh-llm serve meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M --split --no-draft \ + --speculative-strategy mtp \ + --speculative-ngram-proposer cache \ + --speculative-ngram-min 2 \ + --speculative-ngram-max 4 \ + --speculative-ngram-max-proposal-tokens 6 \ + --speculative-extension-initial-tokens 2 \ + --speculative-extension-max-tokens 6 \ + --speculative-verify-window-pipeline-depth 2 +``` + +- `--speculative-strategy `: select `auto`, `disabled`, `mtp`, a built-in direct-GGUF N-gram strategy, or a strategy declared by the model package. +- `--speculative-ngram-proposer `: choose the N-gram implementation. The cache proposer is request-local. +- `--speculative-ngram-min ` / `--speculative-ngram-max `: set the history match bounds. +- `--speculative-ngram-max-proposal-tokens `: cap the N-gram continuation proposed at once. +- `--speculative-extension-initial-tokens ` / `--speculative-extension-max-tokens `: set the adaptive N-gram tail bounds when extending native MTP. +- `--speculative-extension-tail-backoff-proposals `: pause extension attempts after a rejected N-gram tail. +- `--speculative-native-mtp-reject-cooldown-tokens `: set the generated-token cooldown after native MTP rejection. +- `--speculative-native-mtp-suppress-cooldown-drafts`: suppress native drafts during cooldown; `--speculative-native-mtp-allow-cooldown-drafts` explicitly disables a configured suppression policy. +- `--speculative-native-mtp-suppress-cooldown-draft-limit `: cap the native drafts suppressed by one cooldown. +- `--speculative-verify-window-min-tokens ` / `--speculative-verify-window-max-tokens `: set adaptive verification window bounds. +- `--speculative-verify-window-pipeline-depth `: set the maximum in-flight asynchronous verification windows. + ## Commands ### `models`