diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index bfe60b3d99..1bc374aaae 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -603,6 +603,8 @@ pub struct SpeculativeConfig { pub verify_window_min_tokens: Option, pub verify_window_max_tokens: Option, pub verify_window_pipeline_depth: Option, + pub verify_window_runahead_tokens: Option, + pub ngram_fallback: Option, pub spec_default: Option, pub(crate) legacy_draft_model_path_used: bool, } @@ -654,6 +656,8 @@ impl SpeculativeConfig { 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), + verify_window_runahead_tokens: pick!(verify_window_runahead_tokens), + ngram_fallback: pick!(ngram_fallback), spec_default: pick!(spec_default), legacy_draft_model_path_used: overrides .filter(|config| config.draft_model.is_some()) @@ -727,6 +731,10 @@ struct SpeculativeConfigRaw { #[serde(default)] verify_window_pipeline_depth: Option, #[serde(default)] + verify_window_runahead_tokens: Option, + #[serde(default)] + ngram_fallback: Option, + #[serde(default)] spec_default: Option, } @@ -771,6 +779,8 @@ impl<'de> Deserialize<'de> for SpeculativeConfig { 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, + verify_window_runahead_tokens: raw.verify_window_runahead_tokens, + ngram_fallback: raw.ngram_fallback, spec_default: raw.spec_default, legacy_draft_model_path_used: legacy_used, }) @@ -833,6 +843,11 @@ impl Serialize for SpeculativeConfig { "verify_window_pipeline_depth", &self.verify_window_pipeline_depth, )?; + map.serialize_entry( + "verify_window_runahead_tokens", + &self.verify_window_runahead_tokens, + )?; + map.serialize_entry("ngram_fallback", &self.ngram_fallback)?; map.serialize_entry("spec_default", &self.spec_default)?; map.end() } 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 bf09a8f8f4..4ec3059257 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 @@ -73,7 +73,9 @@ pub(super) fn apply_speculative_behavior( | "native_mtp_suppress_cooldown_draft_limit" | "verify_window_min_tokens" | "verify_window_max_tokens" - | "verify_window_pipeline_depth" => {} + | "verify_window_pipeline_depth" + | "verify_window_runahead_tokens" + | "ngram_fallback" => {} _ => {} } } diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 2095ae7f0e..5dc208ff48 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -659,6 +659,11 @@ fn speculative_settings(prefix: &str) -> Vec { &format!("{prefix}.verify_window_pipeline_depth"), ConfigValueSchema::Integer, ), + basic_setting( + &format!("{prefix}.verify_window_runahead_tokens"), + ConfigValueSchema::Integer, + ), + basic_setting(&format!("{prefix}.ngram_fallback"), ConfigValueSchema::String), basic_setting(&format!("{prefix}.spec_default"), bool_or_auto_schema()), ] } diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index 0505452fb0..2be79e0f69 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -9,7 +9,7 @@ use crate::model::{ SkippyConfig, SpeculativeConfig, StringOrStringList, merge_hardware, merge_model_fit, merge_multimodal, merge_throughput, }; -use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; +use skippy_protocol::{MAX_VERIFY_WINDOW_PIPELINE_DEPTH, MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS}; use crate::validation_support::{ looks_like_model_identifier, validate_allowed, validate_bool_or_auto, validate_hf_pair, @@ -633,7 +633,23 @@ fn validate_verify_window_controls( &format!("{base_path}.verify_window_pipeline_depth"), 1, u32::try_from(MAX_VERIFY_WINDOW_PIPELINE_DEPTH).expect("verify depth limit fits u32"), - ) + )?; + validate_optional_u32_range( + config.verify_window_runahead_tokens, + &format!("{base_path}.verify_window_runahead_tokens"), + // Zero is the documented fixed-depth sentinel, so a model-level block + // can switch run-ahead back off when the global defaults enable it. + 0, + u32::try_from(MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS).expect("runahead limit fits u32"), + )?; + if let Some(fallback) = config.ngram_fallback.as_deref() { + validate_allowed( + fallback, + &["draft", "none"], + &format!("{base_path}.ngram_fallback"), + )?; + } + Ok(()) } fn validate_request_defaults(config: &RequestDefaultsConfig, base_path: &str) -> DiagnosticResult { 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 140145e1df..9200f48e4a 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -79,26 +79,36 @@ pub(crate) use stage::{ pub(crate) use topology::{StageTopologyParticipant, plan_package_identity_topology}; const BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_DELAY_MS"; +const BENCH_DOWNSTREAM_WIRE_JITTER_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_JITTER_MS"; +const BENCH_DOWNSTREAM_WIRE_STALL_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_STALL_MS"; +const BENCH_DOWNSTREAM_WIRE_STALL_P_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_STALL_P"; fn benchmark_downstream_wire_condition() -> Result { - let delay_ms = match env::var(BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV) { - Ok(value) => parse_benchmark_downstream_wire_delay_ms(&value)?, - Err(env::VarError::NotPresent) => 0.0, + let delay_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV)?; + let jitter_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_JITTER_MS_ENV)?; + let stall_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_STALL_MS_ENV)?; + let stall_p = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_STALL_P_ENV)?; + WireCondition::with_jitter(delay_ms, None, jitter_ms, stall_ms, stall_p) +} + +fn parse_benchmark_wire_env(name: &'static str) -> Result { + match env::var(name) { + Ok(value) => parse_benchmark_downstream_wire_value(name, &value), + Err(env::VarError::NotPresent) => Ok(0.0), Err(env::VarError::NotUnicode(_)) => { - anyhow::bail!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be valid UTF-8") + anyhow::bail!("{name} must be valid UTF-8") } - }; - WireCondition::new(delay_ms, None) + } } -fn parse_benchmark_downstream_wire_delay_ms(value: &str) -> Result { - let delay_ms = value.parse::().with_context(|| { - format!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be a finite non-negative number") - })?; - if !delay_ms.is_finite() || delay_ms < 0.0 { - anyhow::bail!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be a finite non-negative number"); +fn parse_benchmark_downstream_wire_value(name: &str, value: &str) -> Result { + let parsed = value + .parse::() + .with_context(|| format!("{name} must be a finite non-negative number"))?; + if !parsed.is_finite() || parsed < 0.0 { + anyhow::bail!("{name} must be a finite non-negative number"); } - Ok(delay_ms) + Ok(parsed) } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -1345,9 +1355,12 @@ mod tests { #[test] fn benchmark_wire_delay_accepts_finite_non_negative_values() { - assert_eq!(parse_benchmark_downstream_wire_delay_ms("0").unwrap(), 0.0); assert_eq!( - parse_benchmark_downstream_wire_delay_ms("25.5").unwrap(), + parse_benchmark_downstream_wire_value("test", "0").unwrap(), + 0.0 + ); + assert_eq!( + parse_benchmark_downstream_wire_value("test", "25.5").unwrap(), 25.5 ); } @@ -1355,7 +1368,7 @@ mod tests { #[test] fn benchmark_wire_delay_rejects_invalid_values() { for value in ["-1", "NaN", "inf", "not-a-number"] { - assert!(parse_benchmark_downstream_wire_delay_ms(value).is_err()); + assert!(parse_benchmark_downstream_wire_value("test", value).is_err()); } } 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 6a17c09ddb..556ea314ec 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 @@ -136,6 +136,7 @@ pub(super) fn resolve_speculative_config( model_config, global_config, package_generation, + has_draft_model: draft_model_path.is_some(), })?; // A standalone N-gram plan (no native MTP, no draft model) runs in the // legacy `ngram` mode so the embedded frontend derives its window from the @@ -201,6 +202,7 @@ struct DecodeResolutionInput<'a> { model_config: Option<&'a SpeculativeConfig>, global_config: Option<&'a SpeculativeConfig>, package_generation: Option<&'a PackageGenerationInfo>, + has_draft_model: bool, } #[allow(clippy::too_many_lines)] @@ -397,9 +399,53 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result config.verify_window.max_tokens { bail!("skippy speculative verify window requires min_tokens <= max_tokens"); } + let ngram_fallback = pick_string( + input + .model_config + .and_then(|value| value.ngram_fallback.as_deref()), + input + .global_config + .and_then(|value| value.ngram_fallback.as_deref()), + None, + ); + config.ngram_fallback_draft = match ngram_fallback { + "draft" => { + if config.ngram.is_none() { + bail!("skippy speculative ngram_fallback = \"draft\" requires an N-gram strategy"); + } + // Both of these would otherwise start cleanly and never take the + // fallback path: the operator gets baseline behaviour and a + // telemetry counter stuck at zero, indistinguishable from a + // proposer that simply never missed. + if !input.has_draft_model { + bail!( + "skippy speculative ngram_fallback = \"draft\" requires speculative.draft_model" + ); + } + if config.verify_window.pipeline_depth <= 1 { + bail!( + "skippy speculative ngram_fallback = \"draft\" requires verify_window_pipeline_depth > 1; the classic serial draft loop is authoritative at depth 1" + ); + } + true + } + "none" | "" => false, + other => bail!("skippy speculative ngram_fallback must be draft or none, got {other}"), + }; config.validate()?; Ok(config) } @@ -479,6 +525,7 @@ fn package_decode_config( min_tokens: 1, max_tokens: 4, pipeline_depth: 1, + runahead_max_tokens: 0, }); let effective_strategy = match (native_mtp.enabled, ngram.as_ref().map(|value| value.kind)) { (true, Some(NgramProposerKind::Cache)) => "native-mtp+ngram-cache", @@ -494,6 +541,7 @@ fn package_decode_config( ngram, extension, verify_window, + ngram_fallback_draft: false, })) } @@ -564,6 +612,7 @@ fn verify_window_config(policy: &PackageWindowPolicyInfo) -> VerifyWindowConfig min_tokens: policy.min_window as usize, max_tokens: policy.max_window as usize, pipeline_depth: policy.pipeline_depth.unwrap_or(1) as usize, + runahead_max_tokens: 0, } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 86e5e8c948..b564aa9a7e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -1592,3 +1592,110 @@ verify_window_pipeline_depth = 2 assert_eq!(translated.max_proposal_tokens, 48); assert_eq!(args.speculative.verify_window.pipeline_depth, 2); } + +#[test] +fn model_level_zero_runahead_overrides_a_positive_global_default() { + use crate::plugin::SpeculativeConfig; + let global: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 2 +verify_window_runahead_tokens = 256 +"#, + ) + .expect("parse global speculative config"); + let model: SpeculativeConfig = toml::from_str( + r#" +verify_window_runahead_tokens = 0 +"#, + ) + .expect("parse model speculative config"); + let inherited = super::speculative::resolve_speculative_config( + None, + Some(&global), + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect("global run-ahead must resolve"); + assert_eq!(inherited.decode.verify_window.runahead_max_tokens, 256); + let overridden = super::speculative::resolve_speculative_config( + Some(&model), + Some(&global), + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect("model-level zero must resolve to fixed-depth mode"); + assert_eq!( + overridden.decode.verify_window.runahead_max_tokens, 0, + "Some(0) at the model level must win over the inherited positive default" + ); +} + +#[test] +fn draft_fallback_without_a_draft_model_is_rejected_not_silently_ignored() { + use crate::plugin::SpeculativeConfig; + let config: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 2 +ngram_fallback = "draft" +"#, + ) + .expect("parse speculative config"); + let error = super::speculative::resolve_speculative_config( + Some(&config), + None, + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect_err("a draft fallback with no draft model must not resolve"); + + assert!( + error + .to_string() + .contains("requires speculative.draft_model"), + "unexpected error: {error}" + ); +} + +#[test] +fn draft_fallback_at_pipeline_depth_one_is_rejected() { + use crate::plugin::SpeculativeConfig; + let config: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 1 +ngram_fallback = "draft" +draft_model = "meshllm/draft-model" +draft_max_tokens = 4 +"#, + ) + .expect("parse speculative config"); + let error = super::speculative::resolve_speculative_config( + Some(&config), + None, + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect_err("depth 1 keeps the serial draft loop authoritative"); + + assert!( + error + .to_string() + .contains("verify_window_pipeline_depth > 1"), + "unexpected error: {error}" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/protocol/convert.rs b/crates/mesh-llm-host-runtime/src/protocol/convert.rs index e928cb7931..c5c899b955 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/convert.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/convert.rs @@ -17,7 +17,7 @@ fn skippy_stage_subprotocols( let mut features = vec![skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL.to_string()]; if stage_protocol_generation_supported { features.push( - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5.to_string(), ); } if artifact_transfer_supported { @@ -50,7 +50,7 @@ fn supports_skippy_status_list(subprotocols: &[crate::proto::node::MeshSubprotoc fn supports_skippy_stage_generation(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { supports_skippy_stage_feature( subprotocols, - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, ) && supports_skippy_stage_feature( subprotocols, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, diff --git a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs index 691796c9de..c2af8166ac 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs @@ -82,7 +82,7 @@ fn owner_fields_roundtrip_through_proto_announcement() { .any(|feature| feature == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST) ); assert!(skippy.features.iter().any(|feature| feature - == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4)); + == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5)); assert_eq!( proto_pa .owner_attestation @@ -310,7 +310,7 @@ fn proto_announcement_without_stage_control_is_not_stage_compatible() { name: skippy_protocol::STAGE_SUBPROTOCOL_NAME.to_string(), major: skippy_protocol::STAGE_SUBPROTOCOL_MAJOR, features: vec![ - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5.to_string(), ], }], ..Default::default() 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 afccf958ef..d5bf2c2296 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 @@ -854,6 +854,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.ngram_fallback", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.ngram_max", "support": "supported", @@ -924,6 +931,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.verify_window_runahead_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.throughput.continuous_batching", "support": "supported", @@ -995,4 +1009,4 @@ } } ] -} +} \ No newline at end of file diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index 4541f7f9f0..376530fbaa 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -58,6 +58,7 @@ pub enum WireMessageKind { DecodeLightCtx = 9, VerifyWindow = 21, RetireVerifyWindow = 22, + DiscardStaleWindows = 23, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -101,6 +102,14 @@ impl WireMessageKind { matches!(self, Self::RetireVerifyWindow) } + /// Control message invalidating a contiguous range of not-yet-executed + /// verify windows after the driver detected divergence. Recorded at + /// message-receive time so buffered stale windows are skipped instead of + /// executed. + pub fn is_stale_window_discard(self) -> bool { + matches!(self, Self::DiscardStaleWindows) + } + pub fn is_generation_control(self) -> bool { matches!(self, Self::ConfigureGeneration) } @@ -147,6 +156,7 @@ impl TryFrom for WireMessageKind { 20 => Ok(Self::PredictionReturnOpen), 21 => Ok(Self::VerifyWindow), 22 => Ok(Self::RetireVerifyWindow), + 23 => Ok(Self::DiscardStaleWindows), _ => Err(invalid_data("unknown stage message kind")), } } @@ -329,6 +339,7 @@ impl StageStateHeader { WireMessageKind::StateImport | WireMessageKind::StateExport ) || kind.is_session_control() || kind.is_verify_retirement() + || kind.is_stale_window_discard() || kind.is_generation_control() { return true; diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index d2a7d72b56..6a7cd52906 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -29,11 +29,11 @@ pub use messages::{ StateImportMessage, StopMessage, TokenReplyMessage, }; pub use validation::{ - MAX_STAGE_FRAME_BYTES, MAX_VERIFY_WINDOW_PIPELINE_DEPTH, SCHEMA_VERSION, STAGE_ALPN_V2, - STAGE_PROTOCOL_GENERATION, STAGE_STREAM_ARTIFACT_TRANSFER, STAGE_STREAM_CONTROL, - STAGE_STREAM_TRANSPORT, STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER, + MAX_STAGE_FRAME_BYTES, MAX_VERIFY_WINDOW_PIPELINE_DEPTH, MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS, + SCHEMA_VERSION, STAGE_ALPN_V2, STAGE_PROTOCOL_GENERATION, STAGE_STREAM_ARTIFACT_TRANSFER, + STAGE_STREAM_CONTROL, STAGE_STREAM_TRANSPORT, STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER, STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, STAGE_SUBPROTOCOL_MAJOR, STAGE_SUBPROTOCOL_NAME, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -53,7 +53,7 @@ mod tests { stage_control_response, }; use super::{ - STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -62,7 +62,7 @@ mod tests { #[test] fn stage_protocol_generation_feature_names_current_generation() { assert_eq!( - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, format!("stage-generation-{STAGE_PROTOCOL_GENERATION}") ); } diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index 838a522c2a..2dc41cc467 100644 --- a/crates/skippy-protocol/src/validation.rs +++ b/crates/skippy-protocol/src/validation.rs @@ -6,13 +6,13 @@ pub const STAGE_ALPN_V2: &[u8] = b"skippy-stage/2"; pub const STAGE_SUBPROTOCOL_NAME: &str = "skippy-stage"; pub const STAGE_SUBPROTOCOL_MAJOR: u32 = 2; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL: &str = "stage-control"; -pub const STAGE_PROTOCOL_GENERATION: u32 = 4; +pub const STAGE_PROTOCOL_GENERATION: u32 = 5; /// Generation-scoped stage capability. A peer can advertise `stage-control` /// while still rejecting current-generation frames, so split planning gates on /// this exact token before sending current-generation control requests. -pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4: &str = "stage-generation-4"; +pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5: &str = "stage-generation-5"; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4; + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5; pub const STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER: &str = "artifact-transfer"; pub const STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST: &str = "status-list"; pub const STAGE_STREAM_CONTROL: u8 = 0x01; @@ -21,6 +21,10 @@ pub const STAGE_STREAM_ARTIFACT_TRANSFER: u8 = 0x03; pub const MAX_STAGE_FRAME_BYTES: usize = 8 * 1024 * 1024; /// Maximum number of unresolved verify windows covered by native checkpoints. pub const MAX_VERIFY_WINDOW_PIPELINE_DEPTH: usize = 64; +/// Sanity bound on the run-ahead speculative-token budget. Every in-flight +/// speculative token holds restorable recovery state downstream, so the budget +/// caps how much KV checkpoint memory one request can pin. +pub const MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS: usize = 4096; #[derive(Debug, Clone, PartialEq, Eq)] pub enum StageFrameError { diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 9859c0a5f5..4f07aff412 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -15,7 +15,7 @@ mesh/openai-frontend; diagnostic and benchmark clients may connect directly to the first stage. The full request/reply path is tip-to-tip: token IDs enter at the driver-facing -tip, and activations flow through the stage chain. Stage protocol generation 4 +tip, and activations flow through the stage chain. Stage protocol generation 5 is a compatibility-breaking contract: prediction-bearing replies return directly from the final/readout tip to the driver-facing stage instead of being relayed back through intermediate stages. Middle-out is the prefill optimization @@ -103,11 +103,18 @@ deadline handling. ## Notes - `serve-binary` is the tuned binary stage-to-stage path. -- `serve-binary` participates in the breaking generation-4 stage protocol. - Stage compatibility requires `stage-generation-4`; direct prediction return and - exact verify-checkpoint retirement are part of that generation's contract, so - older peers are rejected during split planning instead of being mixed into a - generation-4 topology. +- `serve-binary` participates in the breaking generation-5 stage protocol. + Stage compatibility requires `stage-generation-5`; direct prediction return, + exact verify-checkpoint retirement, and the `DiscardStaleWindows` control + frame are part of that generation's contract, so older peers are rejected + during split planning instead of being mixed into a generation-5 topology. +- That rejection happens in mesh split planning. A manually wired + `serve-binary --downstream host:port` pair performs no generation + handshake, so **the contract for the standalone path is that all stages are + upgraded together**. Pointing a run-ahead coordinator at a generation-4 + stage binary is not degraded gracefully: the older peer rejects the + `DiscardStaleWindows` frame as an unknown message kind and drops the + request connection. - `serve-binary` accepts upstream protocol connections concurrently. Model execution remains serialized by the per-process runtime lock, but readiness, abandoned, or broken connections do not monopolize the listener and block the @@ -123,7 +130,7 @@ deadline handling. `/v1/completions` using the shared `openai-frontend` crate for a local final/single-stage config with no downstream peer. Split serving uses embedded stage-0 OpenAI serving from `serve-binary --openai-bind-addr` because - generation-4 prediction returns flow directly from the final stage to stage 0. + generation-5 prediction returns flow directly from the final stage to stage 0. The older standalone `serve-openai --first-stage-addr` adapter is no longer supported. `--model-id` is the exact served model id to advertise and accept, for example `org/repo:Q4_K_M`; it is not parsed as stage topology. diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index d344d6e4d1..ff0989b72d 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -39,6 +39,7 @@ mod prefill_recording; pub(in crate::binary_transport) mod reply; mod session_lifecycle; mod session_tracker; +mod stale_discard; mod summary; mod telemetry; diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs b/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs index a64b50751c..c98f526c22 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs @@ -23,8 +23,24 @@ use std::time::Instant; const ASYNC_FORWARD_TERMINAL_TIMEOUT: Duration = Duration::from_secs(30); pub(crate) struct AsyncForwarder { - sender: mpsc::SyncSender, + sender: Option>, pending: VecDeque, + writer: Option>, +} + +impl Drop for AsyncForwarder { + /// Queued frames must not still be on the wire after the request that + /// owns them returns: a persistent lane is handed back for reuse, and a + /// teardown `Stop` written through another clone of the same socket + /// would interleave with a frame this forwarder is still writing. + /// Dropping the sender ends the writer loop once its queue drains, and + /// the join makes that ordering observable to the caller. + fn drop(&mut self) { + drop(self.sender.take()); + if let Some(writer) = self.writer.take() { + let _ = writer.join(); + } + } } pub(crate) struct AsyncForwardReceipt { @@ -54,10 +70,12 @@ impl AsyncForwarder { .set_write_timeout(Some(ASYNC_FORWARD_TERMINAL_TIMEOUT)) .context("set async activation forward write timeout")?; let (sender, receiver) = mpsc::sync_channel::(queue_capacity.max(1)); - thread::spawn(move || run_forwarder(&mut writer, &receiver, &telemetry)); + let writer_thread = + thread::spawn(move || run_forwarder(&mut writer, &receiver, &telemetry)); Ok(Self { - sender, + sender: Some(sender), pending: VecDeque::new(), + writer: Some(writer_thread), }) } @@ -83,6 +101,8 @@ impl AsyncForwarder { self.reap_completed()?; let (done, receiver) = mpsc::channel(); self.sender + .as_ref() + .ok_or_else(|| anyhow!("async activation forwarder stopped"))? .send(AsyncForwardJob { message, wire_dtype, @@ -114,7 +134,7 @@ impl AsyncForwarder { } } - pub(super) fn flush(&mut self) -> Result<()> { + pub(crate) fn flush(&mut self) -> Result<()> { while let Some(receiver) = self.pending.pop_front() { receiver.finish()?; } @@ -222,6 +242,48 @@ mod tests { } } + /// A delayed discard must be fully written before a teardown `Stop` that + /// goes out through a different clone of the same socket, otherwise the + /// two frames interleave and poison a lane that is handed back for reuse. + #[test] + fn a_delayed_discard_lands_before_a_teardown_stop_on_another_clone() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let mut client = TcpStream::connect(address).unwrap(); + let (mut server, _) = listener.accept().unwrap(); + let telemetry = Telemetry::new(None, 1, prefix_cache_test_config(), TelemetryLevel::Off); + let mut forwarder = AsyncForwarder::new(&client, telemetry, 8).unwrap(); + // 250ms of simulated propagation: without the drop-time join, the + // teardown write below wins the race and the frames interleave. + let condition = WireCondition::new(250.0, None).unwrap(); + + forwarder + .send( + message(WireMessageKind::DiscardStaleWindows, 11), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap(); + drop(forwarder); + + write_stage_message_after_propagation( + &mut client, + &message(WireMessageKind::Stop, 22), + WireActivationDType::F32, + WireCondition::new(0.0, None).unwrap(), + ) + .unwrap(); + + let first = read_stage_message(&mut server, 4).unwrap(); + let second = read_stage_message(&mut server, 4).unwrap(); + + assert_eq!(first.kind, WireMessageKind::DiscardStaleWindows); + assert_eq!(first.pos_start, 11); + assert_eq!(second.kind, WireMessageKind::Stop); + assert_eq!(second.pos_start, 22); + } + #[test] fn retirement_receipt_orders_all_prior_verify_writes() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs index 8340611883..9cb52623d6 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -3,7 +3,7 @@ use super::control_messages::{ handle_generation_control, handle_prefix_cache_control, handle_session_control, handle_stop, handle_verify_retirement, }; -use super::message_receive::{next_connection_session_id, receive_next_message}; +use super::message_receive::{next_connection_session_id, spawn_message_reader}; use super::reply::reply_window_for_message; use super::reply::send_stage_reply; use super::session_lifecycle::align_session_to_message; @@ -11,6 +11,7 @@ use super::session_tracker::{ ConnectionSessionTracker, combine_connection_and_cleanup_results, release_tracked_connection_sessions, }; +use super::stale_discard::StaleDiscardRegistry; use super::summary::BinaryMessageObservation; use super::summary::BinaryRequestSummary; use super::telemetry::UpstreamReplyWriteSpan; @@ -149,6 +150,13 @@ fn handle_binary_connection_messages( let mut request_summary = BinaryRequestSummary::default(); let mut prediction_return_streams: BTreeMap<(u64, u64), TcpStream> = BTreeMap::new(); let mut next_message = Some(first_message); + let discard_registry = Arc::new(StaleDiscardRegistry::default()); + let inbound_reader = spawn_message_reader( + upstream, + activation_width, + max_inflight.max(1), + discard_registry.clone(), + )?; let mut async_forwarder = if async_prefill_forward || max_inflight > 1 { downstream .as_ref() @@ -164,9 +172,7 @@ fn handle_binary_connection_messages( loop { let recv_start_unix_nanos = now_unix_nanos() as u64; let recv_started = Instant::now(); - let Some(mut message) = receive_next_message( - upstream, - activation_width, + let Some(mut message) = inbound_reader.next( next_message.take(), pending_prefill_replies, request_summary.message_count, @@ -215,6 +221,33 @@ fn handle_binary_connection_messages( continue; } + if message.kind.is_stale_window_discard() { + // The reader thread already recorded the range; middle stages + // forward it so downstream stages can skip their buffered stale + // windows too. No reply is expected. + if let Some(downstream) = downstream.as_mut() { + if let Some(forwarder) = async_forwarder.as_mut() { + forwarder + .send( + message, + wire_dtype, + downstream_wire_condition, + BTreeMap::new(), + ) + .context("forward stale window discard downstream")?; + } else { + write_stage_message_conditioned( + &mut *downstream, + &message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward stale window discard downstream")?; + } + } + continue; + } + if message.kind.is_verify_retirement() { handle_verify_retirement( iteration_scheduler, @@ -303,6 +336,38 @@ fn handle_binary_connection_messages( bail!("binary stage state does not match message kind"); } + if message.kind == WireMessageKind::VerifyWindow + && downstream.is_none() + && discard_registry.is_discarded( + message.request_id, + message.session_id, + message.state.seq_id, + ) + { + // A discard raced ahead of this buffered stale window: answer with + // an empty prediction set instead of executing it. The driver's + // stale drain only uses the window id for FIFO bookkeeping. + let reply = StageReply { + kind: WireReplyKind::PredictedTokens, + predicted: message.state.current_token, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: reply_window_for_message(&message), + stats: StageReplyStats::default(), + }; + if let Some(return_stream) = + prediction_return_streams.get_mut(&(message.request_id, message.session_id)) + { + direct_return::send_direct_prediction_return(return_stream, reply) + .context("send discarded verify window reply")?; + } else { + send_stage_reply(&mut *upstream, reply) + .context("send discarded verify window reply")?; + } + request_summary.message_count += 1; + continue; + } + let requires_predicted = message.kind.requires_predicted_reply(); let early_prefill_ack = message.kind.is_prefill() && !requires_predicted; let mut upstream_reply_start_unix_nanos = None; diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs index 8450c2c9a1..d2bd2aa963 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs @@ -1,8 +1,14 @@ use anyhow::{Context, Result}; use skippy_protocol::binary::{StageWireMessage, read_stage_message}; use std::io; -use std::net::TcpStream; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::net::{Shutdown, TcpStream}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; + +use super::stale_discard::StaleDiscardRegistry; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); @@ -10,25 +16,289 @@ pub(super) fn next_connection_session_id() -> u64 { BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed) } -pub(super) fn receive_next_message( - upstream: &mut TcpStream, +/// Messages the reader may hold parsed ahead of execution. The coordinator +/// admits at most `MAX_VERIFY_WINDOW_PIPELINE_DEPTH` verify windows per +/// request and retires each one with a control message, so this covers the +/// whole admitted backlog by count. +/// +/// A discard therefore usually overtakes the stale windows queued ahead of +/// it. It is not guaranteed to: `INBOUND_LOOKAHEAD_BYTES` binds first for +/// wide frames (a full 64-window backlog of `MAX_STAGE_FRAME_BYTES` frames is +/// far past the byte ceiling), and the reader then parks with the discard +/// still unread. That degrades to executing the stale tail — today's cost +/// without this path — rather than deadlocking, because the executor keeps +/// draining the queue. +pub(super) const INBOUND_LOOKAHEAD_MESSAGES: usize = + 2 * skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; + +/// Byte ceiling for the same queue, and the per-connection bound on what a +/// misbehaving peer can make this process buffer: reading ahead moves frames +/// out of the kernel socket buffer into userspace, so the message count alone +/// bounds nothing useful for memory. A `DiscardStaleWindows` frame is ~100 +/// bytes and overtakes a 32 MiB backlog exactly as reliably as a larger one, +/// so this is sized for the smallest backlog that preserves the property. +pub(super) const INBOUND_LOOKAHEAD_BYTES: usize = 32 * 1024 * 1024; + +/// Reads upstream messages on a dedicated thread so the executor can run a +/// buffered message while later ones are already parsed. This is what lets a +/// `DiscardStaleWindows` control message take effect before the buffered +/// stale verify windows behind it get executed: the reader records the +/// discard range in the shared registry the moment it reads the message. +/// +/// Dropping the reader shuts the socket down and joins the thread, so a +/// handler that exits on a local error while the peer keeps its side open +/// does not leak a blocked thread and its cloned descriptor. +pub(super) struct InboundMessageReader { + receiver: Option>>, + stream: TcpStream, + thread: Option>, + queued_bytes: Arc, + stopped: Arc, +} + +impl Drop for InboundMessageReader { + fn drop(&mut self) { + // Disconnect the channel first: a reader blocked in `send` on a full + // lookahead queue is not woken by the socket shutdown. + // Release a reader parked on the byte ceiling: it is waiting on the + // executor, which is not coming back, and neither the receiver drop + // nor the socket shutdown would wake it. + self.stopped.store(true, Ordering::Release); + drop(self.receiver.take()); + // Unblock a pending `read_stage_message`; errors here only mean the + // socket is already closed. + let _ = self.stream.shutdown(Shutdown::Both); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +pub(super) fn spawn_message_reader( + upstream: &TcpStream, activation_width: i32, - first_message: Option, - pending_prefill_replies: usize, - observed_message_count: usize, -) -> Result> { - if first_message.is_some() { - return Ok(first_message); - } - match read_stage_message(upstream, activation_width) { - Ok(message) => Ok(Some(message)), - Err(error) - if error.kind() == io::ErrorKind::UnexpectedEof - && pending_prefill_replies == 0 - && observed_message_count == 0 => - { - Ok(None) - } - Err(error) => Err(error).context("read binary stage message"), + capacity: usize, + registry: Arc, +) -> Result { + let mut reader = upstream + .try_clone() + .context("clone upstream stream for inbound message reader")?; + let stream = upstream + .try_clone() + .context("clone upstream stream for inbound reader shutdown")?; + let (sender, receiver) = mpsc::sync_channel(capacity.max(INBOUND_LOOKAHEAD_MESSAGES)); + let queued_bytes = Arc::new(AtomicUsize::new(0)); + let reader_queued_bytes = queued_bytes.clone(); + let stopped = Arc::new(AtomicBool::new(false)); + let reader_stopped = stopped.clone(); + let thread = thread::spawn(move || { + loop { + // Back off while the parsed backlog is over the byte ceiling; the + // executor decrements as it takes messages off the queue. + while reader_queued_bytes.load(Ordering::Acquire) >= INBOUND_LOOKAHEAD_BYTES { + if reader_stopped.load(Ordering::Acquire) { + return; + } + thread::sleep(Duration::from_millis(1)); + } + match read_stage_message(&mut reader, activation_width) { + Ok(message) => { + if message.kind.is_stale_window_discard() { + registry.record_message(&message); + } + let message_bytes = message.estimated_wire_bytes(); + reader_queued_bytes.fetch_add(message_bytes, Ordering::AcqRel); + if sender.send(Ok(message)).is_err() { + return; + } + } + Err(error) => { + let _ = sender.send(Err(error)); + return; + } + } + } + }); + Ok(InboundMessageReader { + receiver: Some(receiver), + stream, + thread: Some(thread), + queued_bytes, + stopped, + }) +} + +impl InboundMessageReader { + /// EOF classification: a clean EOF before any traffic is a normal + /// connection close, anything else is an error. + pub(super) fn next( + &self, + first_message: Option, + pending_prefill_replies: usize, + observed_message_count: usize, + ) -> Result> { + if first_message.is_some() { + return Ok(first_message); + } + let receiver = self + .receiver + .as_ref() + .expect("inbound receiver present until drop"); + match receiver.recv() { + Ok(Ok(message)) => { + // Saturating, not a load-then-subtract: the counter must + // never wrap, or the reader parks on the ceiling forever. + let bytes = message.estimated_wire_bytes(); + self.queued_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |queued| { + Some(queued.saturating_sub(bytes)) + }) + .ok(); + Ok(Some(message)) + } + Ok(Err(error)) + if error.kind() == io::ErrorKind::UnexpectedEof + && pending_prefill_replies == 0 + && observed_message_count == 0 => + { + Ok(None) + } + Ok(Err(error)) => Err(error).context("read binary stage message"), + // The reader thread is gone without a final error; treat it as a + // closed connection. + Err(_) => Ok(None), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use skippy_protocol::binary::{ + StageStateHeader, WireActivationDType, WireMessageKind, write_stage_message, + }; + use std::net::TcpListener; + use std::time::{Duration, Instant}; + + fn control_message(kind: WireMessageKind, tokens: Vec) -> StageWireMessage { + StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, WireActivationDType::F32), + request_id: 7, + session_id: 9, + sampling: None, + chat_sampling_metadata: None, + tokens, + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("local addr"); + let client = TcpStream::connect(address).expect("connect"); + let (server, _) = listener.accept().expect("accept"); + (client, server) + } + + #[test] + fn discard_is_recorded_behind_a_backlog_larger_than_the_execution_queue() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + // Execution queue of one; the reader must still look past a full + // admitted backlog without anything being dequeued. + let reader = spawn_message_reader(&upstream, 4, 1, registry.clone()).expect("spawn"); + + for _ in 0..skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH { + let stale = control_message(WireMessageKind::Stop, Vec::new()); + write_stage_message(&mut peer, &stale, WireActivationDType::F32).expect("write"); + } + let discard = control_message(WireMessageKind::DiscardStaleWindows, vec![3, 9]); + write_stage_message(&mut peer, &discard, WireActivationDType::F32).expect("write"); + + let deadline = Instant::now() + Duration::from_secs(5); + while !registry.is_discarded(7, 9, 5) { + assert!( + Instant::now() < deadline, + "discard must be recorded while the stale backlog is still queued" + ); + thread::sleep(Duration::from_millis(5)); + } + drop(reader); + } + + #[test] + fn dropping_the_reader_completes_while_it_is_parked_on_the_byte_ceiling() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + // Park the reader: pretend the executor is holding the whole byte + // budget, so the backoff loop is the only thing running. + reader + .queued_bytes + .store(INBOUND_LOOKAHEAD_BYTES, Ordering::Release); + write_stage_message( + &mut peer, + &control_message(WireMessageKind::Stop, Vec::new()), + WireActivationDType::F32, + ) + .expect("write"); + thread::sleep(Duration::from_millis(50)); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("a reader parked on the byte ceiling must still be released on drop"); + drop(peer); + } + + #[test] + fn dropping_the_reader_completes_while_the_lookahead_channel_is_full() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + // Fill the lookahead queue and leave the reader blocked in `send`. + for _ in 0..(INBOUND_LOOKAHEAD_MESSAGES + 4) { + let message = control_message(WireMessageKind::Stop, Vec::new()); + write_stage_message(&mut peer, &message, WireActivationDType::F32).expect("write"); + } + thread::sleep(Duration::from_millis(100)); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("dropping the receiver must unblock the queued send before the join"); + drop(peer); + } + + #[test] + fn dropping_the_reader_joins_the_thread_while_the_peer_stays_open() { + let (peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("drop must shut the socket down and join the blocked reader thread"); + drop(peer); } } diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs new file mode 100644 index 0000000000..35ec9ea58f --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs @@ -0,0 +1,150 @@ +use std::collections::VecDeque; +use std::sync::Mutex; + +use skippy_protocol::binary::StageWireMessage; + +/// Bounds so a misbehaving upstream cannot grow the registry without limit. +const MAX_TRACKED_REQUESTS: usize = 32; +const MAX_RANGES_PER_REQUEST: usize = 8; + +#[derive(Debug)] +struct RequestDiscards { + request_id: u64, + session_id: u64, + ranges: VecDeque<(i32, i32)>, +} + +/// Discarded verify-window id ranges, recorded by the connection's reader +/// thread the moment a `DiscardStaleWindows` message is read and consulted by +/// the executor before running each buffered verify window. This is what lets +/// a divergence cancel the stale run-ahead tail instead of executing it. +#[derive(Debug, Default)] +pub(super) struct StaleDiscardRegistry { + requests: Mutex>, +} + +impl StaleDiscardRegistry { + /// Records the range carried by a `DiscardStaleWindows` message + /// (`tokens = [min_window_id, max_window_id]`). Malformed messages are + /// ignored: a discard is an optimization, never a correctness dependency. + pub(super) fn record_message(&self, message: &StageWireMessage) { + let (Some(&min_id), Some(&max_id)) = (message.tokens.first(), message.tokens.get(1)) else { + return; + }; + if min_id > max_id { + return; + } + self.record(message.request_id, message.session_id, min_id, max_id); + } + + pub(super) fn record(&self, request_id: u64, session_id: u64, min_id: i32, max_id: i32) { + let mut requests = self.requests.lock().expect("stale discard lock poisoned"); + if let Some(entry) = requests + .iter_mut() + .find(|entry| entry.request_id == request_id && entry.session_id == session_id) + { + if entry.ranges.len() >= MAX_RANGES_PER_REQUEST { + entry.ranges.pop_front(); + } + entry.ranges.push_back((min_id, max_id)); + return; + } + if requests.len() >= MAX_TRACKED_REQUESTS { + requests.pop_front(); + } + let mut ranges = VecDeque::new(); + ranges.push_back((min_id, max_id)); + requests.push_back(RequestDiscards { + request_id, + session_id, + ranges, + }); + } + + pub(super) fn is_discarded(&self, request_id: u64, session_id: u64, window_id: i32) -> bool { + let requests = self.requests.lock().expect("stale discard lock poisoned"); + requests + .iter() + .filter(|entry| entry.request_id == request_id && entry.session_id == session_id) + .any(|entry| { + entry + .ranges + .iter() + .any(|&(min_id, max_id)| (min_id..=max_id).contains(&window_id)) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn records_and_matches_ranges_per_request() { + let registry = StaleDiscardRegistry::default(); + registry.record(7, 9, 10, 14); + + assert!(registry.is_discarded(7, 9, 10)); + assert!(registry.is_discarded(7, 9, 14)); + assert!(!registry.is_discarded(7, 9, 15)); + assert!(!registry.is_discarded(7, 9, 9)); + // Other requests and sessions are unaffected. + assert!(!registry.is_discarded(8, 9, 12)); + assert!(!registry.is_discarded(7, 10, 12)); + } + + #[test] + fn range_count_is_bounded_per_request() { + let registry = StaleDiscardRegistry::default(); + for index in 0..(MAX_RANGES_PER_REQUEST as i32 + 4) { + registry.record(1, 1, index * 10, index * 10 + 1); + } + // The oldest ranges were evicted; the newest still match. + assert!(!registry.is_discarded(1, 1, 0)); + let newest = (MAX_RANGES_PER_REQUEST as i32 + 3) * 10; + assert!(registry.is_discarded(1, 1, newest)); + } + + #[test] + fn tracked_request_count_is_bounded() { + let registry = StaleDiscardRegistry::default(); + for request in 0..(MAX_TRACKED_REQUESTS as u64 + 4) { + registry.record(request, 1, 0, 10); + } + assert!(!registry.is_discarded(0, 1, 5)); + assert!(registry.is_discarded(MAX_TRACKED_REQUESTS as u64 + 3, 1, 5)); + } + + #[test] + fn malformed_discard_messages_are_ignored() { + use skippy_protocol::binary::{StageStateHeader, WireActivationDType, WireMessageKind}; + let registry = StaleDiscardRegistry::default(); + let mut message = StageWireMessage { + kind: WireMessageKind::DiscardStaleWindows, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new( + WireMessageKind::DiscardStaleWindows, + WireActivationDType::F32, + ), + request_id: 1, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }; + registry.record_message(&message); + assert!(!registry.is_discarded(1, 1, 0)); + + message.tokens = vec![9, 3]; + registry.record_message(&message); + assert!(!registry.is_discarded(1, 1, 5)); + + message.tokens = vec![3, 9]; + registry.record_message(&message); + assert!(registry.is_discarded(1, 1, 5)); + } +} diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index 5512d1ef3c..903acb719e 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -64,8 +64,13 @@ impl BinaryStageOptions { bail!("--openai-prefill-chunk-size must be greater than zero"); } let wire_dtype = parse_wire_dtype(&args.activation_wire_dtype)?; - let downstream_wire_condition = - WireCondition::new(args.downstream_wire_delay_ms, args.downstream_wire_mbps)?; + let downstream_wire_condition = WireCondition::with_jitter( + args.downstream_wire_delay_ms, + args.downstream_wire_mbps, + args.downstream_wire_jitter_ms, + args.downstream_wire_stall_ms, + args.downstream_wire_stall_p, + )?; let config = load_json::(&args.config) .with_context(|| format!("load stage config {}", args.config.display()))?; let topology = match args.topology.as_ref() { @@ -231,7 +236,9 @@ mod tests { min_tokens: 1, max_tokens: 6, pipeline_depth: 2, + runahead_max_tokens: 0, }, + ngram_fallback_draft: false, } } diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index 0654f72355..a9eec4ed4b 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -721,6 +721,7 @@ pub(crate) fn run_binary_stage_message( | WireMessageKind::ConfigureGeneration | WireMessageKind::TrimSession | WireMessageKind::RetireVerifyWindow + | WireMessageKind::DiscardStaleWindows | WireMessageKind::ProbePrefill | WireMessageKind::RestorePrefill | WireMessageKind::TryRestorePrefill diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index 19e8c73914..d5dc2793db 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -1,27 +1,106 @@ -use std::{io, thread, time::Duration}; +use std::{ + cell::Cell, + io, + sync::atomic::{AtomicU64, Ordering}, + thread, + time::Duration, +}; use anyhow::{Result, bail}; use skippy_protocol::binary::{StageWireMessage, WireActivationDType, write_stage_message}; +const WIRE_SAMPLE_SEED: u64 = 0x9E37_79B9_7F4A_7C15; + #[derive(Clone, Copy, Debug)] pub struct WireCondition { delay_ms: f64, mbps: Option, + jitter_ms: f64, + stall_ms: f64, + stall_p: f64, } +/// Upper bound for each simulated delay component. An hour-long simulated +/// stall is already far beyond any useful wire model, and bounding the inputs +/// keeps the combined delay inside `Duration::from_secs_f64`'s domain. +const MAX_SIMULATED_DELAY_MS: f64 = 3_600_000.0; + impl WireCondition { pub fn new(delay_ms: f64, mbps: Option) -> Result { + Self::with_jitter(delay_ms, mbps, 0.0, 0.0, 0.0) + } + + /// A wire condition with a stochastic component, modeling jittery links + /// (Wi-Fi, congested WAN) instead of a constant-latency pipe: + /// + /// - `delay_ms`: fixed one-way propagation delay per message. + /// - `jitter_ms`: mean of an exponentially distributed extra delay added + /// per message (heavy-ish tail, like contention/retransmit variance). + /// - `stall_ms`/`stall_p`: with probability `stall_p` a message is hit by + /// an additional `stall_ms` burst stall (radio retry storms, channel + /// scans). Later messages queue behind it FIFO, which matches the + /// head-of-line blocking of an ordered transport. + pub fn with_jitter( + delay_ms: f64, + mbps: Option, + jitter_ms: f64, + stall_ms: f64, + stall_p: f64, + ) -> Result { if !delay_ms.is_finite() || delay_ms < 0.0 { bail!("downstream wire delay must be finite and non-negative"); } if mbps.is_some_and(|value| !value.is_finite() || value <= 0.0) { bail!("downstream wire mbps must be finite and greater than zero"); } - Ok(Self { delay_ms, mbps }) + if !jitter_ms.is_finite() || jitter_ms < 0.0 { + bail!("downstream wire jitter must be finite and non-negative"); + } + if !stall_ms.is_finite() || stall_ms < 0.0 { + bail!("downstream wire stall must be finite and non-negative"); + } + if !stall_p.is_finite() || !(0.0..=1.0).contains(&stall_p) { + bail!("downstream wire stall probability must be within [0, 1]"); + } + if stall_p > 0.0 && stall_ms == 0.0 { + bail!("downstream wire stall probability requires a stall duration"); + } + // After the finiteness checks, so an infinite input reports what is + // actually wrong with it rather than the magnitude bound. + for (value, name) in [ + (delay_ms, "delay"), + (jitter_ms, "jitter"), + (stall_ms, "stall"), + ] { + if value > MAX_SIMULATED_DELAY_MS { + bail!("downstream wire {name} must not exceed {MAX_SIMULATED_DELAY_MS} ms"); + } + } + Ok(Self { + delay_ms, + mbps, + jitter_ms, + stall_ms, + stall_p, + }) } + /// Samples the propagation delay for one message. With no stochastic + /// component configured this is the constant `delay_ms` and draws nothing + /// from the sample sequence. pub(crate) fn propagation_delay(&self) -> Duration { - Duration::from_secs_f64(self.delay_ms / 1000.0) + let mut delay_ms = self.delay_ms; + if self.jitter_ms > 0.0 { + // Inverse-CDF exponential sample with mean `jitter_ms`. + let uniform = next_uniform_sample(); + delay_ms += -self.jitter_ms * (1.0 - uniform).ln(); + } + if self.stall_p > 0.0 && next_uniform_sample() < self.stall_p { + delay_ms += self.stall_ms; + } + // The exponential jitter tail is unbounded, so clamp the combined + // delay: `Duration::from_secs_f64` panics on overflow. + Duration::from_secs_f64(delay_ms.min(MAX_SIMULATED_DELAY_MS) / 1000.0) } fn sleep_for(&self, message: &StageWireMessage) { @@ -29,17 +108,79 @@ impl WireCondition { self.sleep_for_bandwidth(message); } + /// Serialization delay for `bytes` on this link. A near-zero `mbps` + /// makes the quotient arbitrarily large, so it is clamped to the same + /// bound as the propagation delay: `Duration::from_secs_f64` panics on an + /// out-of-domain value. + pub(crate) fn bandwidth_delay(&self, bytes: usize) -> Duration { + let Some(mbps) = self.mbps else { + return Duration::ZERO; + }; + let seconds = bytes as f64 / (mbps * 125_000.0); + if !seconds.is_finite() || seconds <= 0.0 { + return Duration::ZERO; + } + Duration::from_secs_f64((seconds * 1000.0).min(MAX_SIMULATED_DELAY_MS) / 1000.0) + } + fn sleep_for_bandwidth(&self, message: &StageWireMessage) { - let bandwidth_seconds = self - .mbps - .map(|mbps| message.estimated_wire_bytes() as f64 / (mbps * 125_000.0)) - .unwrap_or(0.0); - if bandwidth_seconds > 0.0 { - thread::sleep(Duration::from_secs_f64(bandwidth_seconds)); + let delay = self.bandwidth_delay(message.estimated_wire_bytes()); + if !delay.is_zero() { + thread::sleep(delay); } } } +/// Hands each conditioned thread a distinct stream ordinal. A process-global +/// *draw* counter would make every stream depend on how the scheduler +/// interleaves the others; a per-thread draw index alone would hand every +/// thread the identical stream, which models synchronized loss rather than a +/// contended link. Salting the per-thread index with a per-thread ordinal +/// gives streams that are both reproducible and independent. +static WIRE_STREAM_ORDINALS: AtomicU64 = AtomicU64::new(0); + +thread_local! { + /// This thread's stream ordinal, claimed on first draw. + static WIRE_STREAM_ORDINAL: Cell> = const { Cell::new(None) }; + /// This thread's draw index within its stream. + static WIRE_SAMPLE_INDEX: Cell = const { Cell::new(0) }; +} + +/// Deterministic uniform sample in [0, 1) via splitmix64 over one stream's +/// draw index. Not cryptographic; just reproducible-enough conditioning for +/// benches and tests. Pure in `(stream, index)` so both properties the model +/// depends on — reproducibility within a lane, independence across lanes — +/// are directly testable. +fn uniform_sample(stream: u64, index: u64) -> f64 { + let mut state = index + .wrapping_mul(0x2545_F491_4F6C_DD1D) + .wrapping_add(stream.wrapping_mul(0x9E37_79B9_7F4A_7C15)) + ^ WIRE_SAMPLE_SEED; + state ^= state >> 30; + state = state.wrapping_mul(0xBF58_476D_1CE4_E5B9); + state ^= state >> 27; + state = state.wrapping_mul(0x94D0_49BB_1331_11EB); + state ^= state >> 31; + (state >> 11) as f64 / (1u64 << 53) as f64 +} + +fn next_uniform_sample() -> f64 { + let stream = WIRE_STREAM_ORDINAL.with(|ordinal| match ordinal.get() { + Some(stream) => stream, + None => { + let stream = WIRE_STREAM_ORDINALS.fetch_add(1, Ordering::Relaxed); + ordinal.set(Some(stream)); + stream + } + }); + let index = WIRE_SAMPLE_INDEX.with(|counter| { + let index = counter.get(); + counter.set(index.wrapping_add(1)); + index + }); + uniform_sample(stream, index) +} + pub(crate) fn write_stage_message_conditioned( writer: impl io::Write, message: &StageWireMessage, @@ -78,10 +219,139 @@ mod tests { } } + #[test] + fn wire_condition_rejects_invalid_jitter_and_stall_shapes() { + for jitter_ms in [-1.0, f64::NAN, f64::INFINITY] { + assert!(WireCondition::with_jitter(0.0, None, jitter_ms, 0.0, 0.0).is_err()); + } + for stall_ms in [-1.0, f64::NAN, f64::INFINITY] { + assert!(WireCondition::with_jitter(0.0, None, 0.0, stall_ms, 0.5).is_err()); + } + for stall_p in [-0.1, 1.1, f64::NAN] { + assert!(WireCondition::with_jitter(0.0, None, 0.0, 10.0, stall_p).is_err()); + } + assert!(WireCondition::with_jitter(0.0, None, 0.0, 0.0, 0.5).is_err()); + } + #[test] fn propagation_delay_is_exposed_without_bandwidth_serialization() { let condition = WireCondition::new(25.0, Some(100.0)).unwrap(); assert_eq!(condition.propagation_delay(), Duration::from_millis(25)); } + + #[test] + fn constant_condition_never_draws_samples() { + let condition = WireCondition::new(3.0, None).unwrap(); + let before = WIRE_SAMPLE_INDEX.with(Cell::get); + let _ = condition.propagation_delay(); + assert_eq!(WIRE_SAMPLE_INDEX.with(Cell::get), before); + } + + #[test] + fn a_stream_is_reproducible_from_its_ordinal_and_index() { + // Reproducibility: a lane replaying the same draws sees the same + // sequence, with no dependence on other threads' interleaving. + let first = (0..4) + .map(|index| uniform_sample(7, index)) + .collect::>(); + let second = (0..4) + .map(|index| uniform_sample(7, index)) + .collect::>(); + + assert_eq!(first, second); + } + + #[test] + fn separate_streams_are_independent_not_identical() { + // Independence: two lanes must not take their burst stalls on the + // same message index, which is a synchronized-loss model rather than + // the contended link the flag documents. + let lanes = (0..4) + .map(|stream| { + (0..8) + .map(|index| uniform_sample(stream, index)) + .collect::>() + }) + .collect::>(); + + for (left_index, left) in lanes.iter().enumerate() { + for right in lanes.iter().skip(left_index + 1) { + assert_ne!(left, right, "distinct streams must not share a sequence"); + } + } + } + + #[test] + fn each_thread_claims_its_own_stream() { + let condition = WireCondition::with_jitter(0.0, None, 5.0, 0.0, 0.0).unwrap(); + let sample_three = move || { + (0..3) + .map(|_| condition.propagation_delay()) + .collect::>() + }; + let first = thread::spawn(sample_three).join().expect("first thread"); + let second = thread::spawn(sample_three).join().expect("second thread"); + + assert_ne!(first, second, "per-lane writer threads must decorrelate"); + } + + #[test] + fn a_near_zero_rate_link_yields_a_bounded_bandwidth_delay() { + let condition = WireCondition::with_jitter(0.0, Some(f64::MIN_POSITIVE), 0.0, 0.0, 0.0) + .expect("a positive rate is accepted"); + + // Unclamped this overflows `Duration::from_secs_f64` and panics. + let delay = condition.bandwidth_delay(64 * 1024); + + assert_eq!( + delay, + Duration::from_secs_f64(MAX_SIMULATED_DELAY_MS / 1000.0) + ); + assert_eq!(condition.bandwidth_delay(0), Duration::ZERO); + assert_eq!( + WireCondition::new(1.0, None).unwrap().bandwidth_delay(4096), + Duration::ZERO + ); + } + + #[test] + fn jittered_condition_adds_a_bounded_positive_tail() { + let condition = WireCondition::with_jitter(2.0, None, 5.0, 0.0, 0.0).unwrap(); + let base = Duration::from_millis(2); + let mut above_base = 0usize; + for _ in 0..256 { + let sampled = condition.propagation_delay(); + assert!(sampled >= base); + // An exponential with mean 5ms virtually never exceeds 200ms; + // treat that as the sanity bound rather than an exact quantile. + assert!(sampled < base + Duration::from_millis(200)); + if sampled > base { + above_base += 1; + } + } + assert!(above_base > 200, "jitter should almost always add delay"); + } + + #[test] + fn stall_probability_gates_the_burst_component() { + let never = WireCondition::with_jitter(1.0, None, 0.0, 50.0, 0.0).unwrap(); + for _ in 0..64 { + assert_eq!(never.propagation_delay(), Duration::from_millis(1)); + } + + let always = WireCondition::with_jitter(1.0, None, 0.0, 50.0, 1.0).unwrap(); + for _ in 0..64 { + assert_eq!(always.propagation_delay(), Duration::from_millis(51)); + } + + let sometimes = WireCondition::with_jitter(0.0, None, 0.0, 50.0, 0.25).unwrap(); + let stalled = (0..512) + .filter(|_| sometimes.propagation_delay() >= Duration::from_millis(50)) + .count(); + assert!( + (32..480).contains(&stalled), + "stall rate {stalled}/512 is not plausibly 25%" + ); + } } diff --git a/crates/skippy-server/src/cli.rs b/crates/skippy-server/src/cli.rs index 030f4ab68c..39bb7c1234 100644 --- a/crates/skippy-server/src/cli.rs +++ b/crates/skippy-server/src/cli.rs @@ -78,6 +78,24 @@ pub struct ServeBinaryArgs { help = "Artificial downstream activation bandwidth cap in megabits per second." )] pub downstream_wire_mbps: Option, + #[arg( + long, + default_value_t = 0.0, + help = "Mean of an exponentially distributed extra per-message downstream delay in milliseconds (models link jitter)." + )] + pub downstream_wire_jitter_ms: f64, + #[arg( + long, + default_value_t = 0.0, + help = "Extra burst-stall delay in milliseconds applied with --downstream-wire-stall-p probability per message." + )] + pub downstream_wire_stall_ms: f64, + #[arg( + long, + default_value_t = 0.0, + help = "Probability in [0, 1] that a downstream message is hit by --downstream-wire-stall-ms." + )] + pub downstream_wire_stall_p: f64, #[arg(long, default_value_t = 60)] pub downstream_connect_timeout_secs: u64, #[arg( diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index ca04c54412..3e8840e23d 100644 --- a/crates/skippy-server/src/frontend/decode_scheduler.rs +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -7,23 +7,46 @@ use skippy_metrics::attr as attr_key; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct VerifyWindowPipelineConfig { depth: usize, + runahead_max_tokens: usize, } impl VerifyWindowPipelineConfig { pub(super) fn new(depth: usize) -> Self { Self { depth: depth.max(1), + runahead_max_tokens: 0, + } + } + + /// Run-ahead mode: admission is bounded by a speculative-token budget + /// instead of a fixed window count. The window count stays capped at the + /// native checkpoint-retention bound so downstream recovery state cannot + /// outgrow what the runtime can restore. + pub(super) fn with_runahead(max_tokens: usize) -> Self { + Self { + depth: skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH, + runahead_max_tokens: max_tokens.max(1), } } pub(super) fn depth(self) -> usize { self.depth } + + pub(super) fn runahead_max_tokens(self) -> usize { + self.runahead_max_tokens + } + + pub(super) fn is_runahead(self) -> bool { + self.runahead_max_tokens > 0 + } } #[derive(Debug, Clone, Default, PartialEq)] pub(super) struct VerifyWindowPipelineStats { depth: usize, + runahead_max_tokens: usize, + max_in_flight_tokens: usize, direct_prediction_return: bool, direct_prediction_return_upstream_opened: bool, direct_prediction_return_reverse_fallback: bool, @@ -156,6 +179,14 @@ impl VerifyWindowPipelineStats { "verify_window_occupancy_average_in_flight".to_string(), serde_json::json!(self.occupancy_average_in_flight), ); + timings.insert( + "verify_window_runahead_max_tokens".to_string(), + serde_json::json!(self.runahead_max_tokens), + ); + timings.insert( + "verify_window_max_in_flight_tokens".to_string(), + serde_json::json!(self.max_in_flight_tokens), + ); } } @@ -172,6 +203,7 @@ pub(super) struct VerifyWindow { pub(super) id: i32, pub(super) base_position: usize, pub(super) decode_step: usize, + pub(super) token_count: usize, } #[derive(Debug)] @@ -179,6 +211,7 @@ pub(super) struct VerifyWindowScheduler { config: VerifyWindowPipelineConfig, next_id: i32, in_flight: VecDeque, + in_flight_tokens: usize, stats: VerifyWindowPipelineStats, occupancy_ms_by_depth: Vec, occupancy_changed: Instant, @@ -190,8 +223,10 @@ impl VerifyWindowScheduler { config, next_id: 1, in_flight: VecDeque::new(), + in_flight_tokens: 0, stats: VerifyWindowPipelineStats { depth: config.depth(), + runahead_max_tokens: config.runahead_max_tokens(), ..VerifyWindowPipelineStats::default() }, occupancy_ms_by_depth: vec![0.0; config.depth().saturating_add(1)], @@ -200,13 +235,37 @@ impl VerifyWindowScheduler { } pub(super) fn has_capacity(&self) -> bool { - self.in_flight.len() < self.config.depth() + if self.in_flight.len() >= self.config.depth() { + return false; + } + if self.config.is_runahead() { + return self.in_flight_tokens < self.config.runahead_max_tokens(); + } + true } pub(super) fn depth(&self) -> usize { self.config.depth() } + /// Widest window (in input tokens, including an epoch-start boundary + /// token) that still fits the remaining run-ahead budget. Unbounded in + /// fixed-depth mode, and unbounded when idle so a budget narrower than + /// one window cannot stall a request: the hard bound applies from the + /// second in-flight window on. + pub(super) fn admissible_window_tokens(&self) -> usize { + if !self.config.is_runahead() || self.in_flight.is_empty() { + return usize::MAX; + } + self.config + .runahead_max_tokens() + .saturating_sub(self.in_flight_tokens) + } + + pub(super) fn is_runahead(&self) -> bool { + self.config.is_runahead() + } + pub(super) fn mark_direct_prediction_return(&mut self, upstream_opened: bool) { self.stats.direct_prediction_return = true; self.stats.direct_prediction_return_upstream_opened = upstream_opened; @@ -264,12 +323,18 @@ impl VerifyWindowScheduler { &mut self, base_position: usize, decode_step: usize, + token_count: usize, ) -> OpenAiResult { if !self.has_capacity() { return Err(OpenAiError::backend( "verify window pipeline depth exceeded", )); } + if token_count > self.admissible_window_tokens() { + return Err(OpenAiError::backend( + "verify window run-ahead token budget exceeded", + )); + } let id = self.next_id; self.next_id = self .next_id @@ -279,11 +344,15 @@ impl VerifyWindowScheduler { id, base_position, decode_step, + token_count, }; self.record_occupancy(); self.in_flight.push_back(window.clone()); + self.in_flight_tokens = self.in_flight_tokens.saturating_add(token_count); 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()); + self.stats.max_in_flight_tokens = + self.stats.max_in_flight_tokens.max(self.in_flight_tokens); Ok(window) } @@ -300,7 +369,9 @@ impl VerifyWindowScheduler { ))); } self.record_occupancy(); - Ok(self.in_flight.pop_front().expect("checked non-empty queue")) + let completed = self.in_flight.pop_front().expect("checked non-empty queue"); + self.in_flight_tokens = self.in_flight_tokens.saturating_sub(completed.token_count); + Ok(completed) } #[cfg(test)] @@ -308,6 +379,7 @@ impl VerifyWindowScheduler { let discarded = self.in_flight.len(); self.record_occupancy(); self.in_flight.clear(); + self.in_flight_tokens = 0; self.stats.stale_discarded = self.stats.stale_discarded.saturating_add(discarded); discarded } @@ -390,13 +462,13 @@ mod tests { #[test] fn records_preferred_and_reverse_direct_return_paths() { - let mut preferred = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + let mut preferred = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(2)); preferred.mark_direct_prediction_return(true); assert!(preferred.stats().direct_prediction_return); assert!(preferred.stats().direct_prediction_return_upstream_opened); assert!(!preferred.stats().direct_prediction_return_reverse_fallback); - let mut reverse = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + let mut reverse = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(2)); reverse.mark_direct_prediction_return(false); let mut timings = BTreeMap::new(); reverse.stats().insert_response_timings(&mut timings); @@ -412,12 +484,12 @@ mod tests { #[test] fn bounds_depth_and_requires_fifo_reply_ids() { - let config = VerifyWindowPipelineConfig { depth: 2 }; + let config = VerifyWindowPipelineConfig::new(2); let mut scheduler = VerifyWindowScheduler::new(config); - let first = scheduler.open(10, 0).unwrap(); - let second = scheduler.open(11, 1).unwrap(); + let first = scheduler.open(10, 0, 1).unwrap(); + let second = scheduler.open(11, 1, 1).unwrap(); - assert!(scheduler.open(12, 2).is_err()); + assert!(scheduler.open(12, 2, 1).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); @@ -431,13 +503,13 @@ mod tests { #[test] fn depth_nine_keeps_the_first_window_restorable() { - let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 9 }); + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(9)); let windows: Vec<_> = (0..9) - .map(|step| scheduler.open(10 + step, step).unwrap()) + .map(|step| scheduler.open(10 + step, step, 1).unwrap()) .collect(); assert_eq!(scheduler.in_flight_len(), 9); - assert!(scheduler.open(19, 9).is_err()); + assert!(scheduler.open(19, 9, 1).is_err()); assert_eq!(scheduler.stats().max_in_flight, 9); // The first window must still be completable after the ninth opens. @@ -451,11 +523,11 @@ mod tests { #[test] fn discards_stale_windows_after_divergence() { - let config = VerifyWindowPipelineConfig { depth: 3 }; + let config = VerifyWindowPipelineConfig::new(3); let mut scheduler = VerifyWindowScheduler::new(config); - scheduler.open(10, 0).unwrap(); - scheduler.open(11, 1).unwrap(); - scheduler.open(12, 2).unwrap(); + scheduler.open(10, 0, 1).unwrap(); + scheduler.open(11, 1, 1).unwrap(); + scheduler.open(12, 2, 1).unwrap(); assert_eq!(scheduler.discard_stale(), 3); assert_eq!(scheduler.stale_discard_count(), 3); @@ -465,10 +537,10 @@ mod tests { #[test] fn stale_recovery_tracks_marked_and_completed_work_separately() { - let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 3 }); - let first = scheduler.open(10, 0).unwrap(); - let second = scheduler.open(11, 1).unwrap(); - let third = scheduler.open(12, 2).unwrap(); + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(3)); + let first = scheduler.open(10, 0, 1).unwrap(); + let second = scheduler.open(11, 1, 1).unwrap(); + let third = scheduler.open(12, 2, 1).unwrap(); scheduler.complete_next(first.id).unwrap(); scheduler.mark_recovery_epoch(2); @@ -490,28 +562,28 @@ mod tests { #[test] fn configured_depth_is_the_fill_target() { - let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 8 }); + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(8)); assert!(scheduler.supports_pipelining(4)); assert_eq!(scheduler.depth(), 8); for position in 0..8 { - scheduler.open(100 + position, position).unwrap(); + scheduler.open(100 + position, position, 1).unwrap(); } assert!(!scheduler.has_capacity()); - assert!(scheduler.open(108, 8).is_err()); + assert!(scheduler.open(108, 8, 1).is_err()); } #[test] fn pipeline_depth_one_never_admits_dependent_work() { - let scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 1 }); + let scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(1)); assert!(!scheduler.supports_pipelining(2)); } #[test] fn fixed_fill_counters_are_exposed_in_response_timings() { - let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(2)); assert!(scheduler.supports_pipelining(2)); scheduler.record_horizon_refill(7); scheduler.record_horizon_refill(0); @@ -538,11 +610,11 @@ mod tests { #[test] fn occupancy_timings_measure_parallel_and_full_depth_time() { - let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 2 }); + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig::new(2)); scheduler.occupancy_changed = Instant::now() - std::time::Duration::from_millis(2); - let first = scheduler.open(10, 0).unwrap(); + let first = scheduler.open(10, 0, 1).unwrap(); scheduler.occupancy_changed = Instant::now() - std::time::Duration::from_millis(3); - let second = scheduler.open(11, 1).unwrap(); + let second = scheduler.open(11, 1, 1).unwrap(); scheduler.occupancy_changed = Instant::now() - std::time::Duration::from_millis(4); let stats = scheduler.stats(); @@ -557,4 +629,67 @@ mod tests { assert_eq!(scheduler.complete_next(first.id).unwrap(), first); assert_eq!(scheduler.complete_next(second.id).unwrap(), second); } + + #[test] + fn runahead_budget_bounds_admission_by_tokens() { + let mut scheduler = + VerifyWindowScheduler::new(VerifyWindowPipelineConfig::with_runahead(100)); + assert!(scheduler.supports_pipelining(4)); + assert_eq!( + scheduler.depth(), + skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH + ); + + let first = scheduler.open(10, 0, 48).unwrap(); + assert!(scheduler.has_capacity()); + let _second = scheduler.open(58, 1, 48).unwrap(); + // 96 tokens in flight, budget 100: only 4 more tokens fit. + assert!(scheduler.has_capacity()); + assert_eq!(scheduler.admissible_window_tokens(), 4); + assert!(scheduler.open(106, 2, 48).is_err()); + let _third = scheduler.open(106, 2, 4).unwrap(); + assert!(!scheduler.has_capacity()); + assert!(scheduler.open(110, 3, 1).is_err()); + + // Completing the head window frees its share of the budget. + assert_eq!(scheduler.complete_next(first.id).unwrap(), first); + assert!(scheduler.has_capacity()); + assert_eq!(scheduler.stats().max_in_flight_tokens, 100); + assert_eq!(scheduler.stats().runahead_max_tokens, 100); + } + + #[test] + fn runahead_window_count_stays_within_native_retention() { + let mut scheduler = + VerifyWindowScheduler::new(VerifyWindowPipelineConfig::with_runahead(1_000_000)); + for step in 0..skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH { + scheduler.open(10 + step, step, 1).unwrap(); + } + // The token budget is nowhere near spent, but the native checkpoint + // retention bound still caps the number of in-flight windows. + assert!(!scheduler.has_capacity()); + assert!( + scheduler + .open( + 10 + skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH, + 64, + 1 + ) + .is_err() + ); + } + + #[test] + fn a_first_window_wider_than_the_whole_budget_still_opens_when_idle() { + let mut scheduler = + VerifyWindowScheduler::new(VerifyWindowPipelineConfig::with_runahead(8)); + assert_eq!(scheduler.admissible_window_tokens(), usize::MAX); + let first = scheduler.open(10, 0, 32).unwrap(); + // Over budget: nothing else may open until the head retires. + assert!(!scheduler.has_capacity()); + assert_eq!(scheduler.admissible_window_tokens(), 0); + assert!(scheduler.open(42, 1, 1).is_err()); + assert_eq!(scheduler.complete_next(first.id).unwrap(), first); + assert!(scheduler.has_capacity()); + } } diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index fff3d21403..43e870dd1b 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -16,7 +16,7 @@ use crate::frontend::generation::stage_reply_timeout; use crate::frontend::util::ms_to_us; use crate::frontend::util::openai_backend_error; use crate::frontend::util::openai_io_error; -use crate::frontend::wire_messages::retire_verify_window_message; +use crate::frontend::wire_messages::{discard_stale_windows_message, retire_verify_window_message}; use crate::telemetry::now_unix_nanos; use openai_frontend::OpenAiError; use openai_frontend::OpenAiResult; @@ -39,6 +39,14 @@ const DIRECT_RETURN_FALLBACK_POLL: Duration = Duration::from_millis(10); // normal WAN verify traversal while remaining shorter than the HTTP client's // request timeout. +/// Identifies a contiguous stale verify-window range for one request. +pub(super) struct StaleWindowDiscard { + pub(super) request_id: u64, + pub(super) session_id: u64, + pub(super) min_window_id: i32, + pub(super) max_window_id: i32, +} + pub(super) struct VerifyRetirement { pub(super) request_id: u64, pub(super) session_id: u64, @@ -105,6 +113,45 @@ impl StageOpenAiBackend { Ok(()) } + /// Sends a stale-window discard downstream without waiting for the + /// write receipt: the message queues behind the already-dispatched stale + /// windows, and blocking here would stall recovery for the whole stale + /// tail's wire time. + pub(super) fn discard_stale_windows( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + async_forwarder: Option<&mut AsyncForwarder>, + discard: StaleWindowDiscard, + ) -> OpenAiResult<()> { + let message = discard_stale_windows_message( + request.wire_dtype, + discard.request_id, + discard.session_id, + discard.min_window_id, + discard.max_window_id, + )?; + if let Some(forwarder) = async_forwarder { + forwarder + .send( + message, + request.wire_dtype, + request.downstream_wire_condition, + self.openai_attrs(request.ids), + ) + .map_err(openai_backend_error)?; + } else { + write_stage_message_conditioned( + downstream, + &message, + request.wire_dtype, + request.downstream_wire_condition, + ) + .map_err(openai_io_error)?; + } + Ok(()) + } + pub(super) fn execute_embedded_stage_message( &self, request: &EmbeddedStageZeroGeneration<'_>, diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index b6adc416ef..42c17cd191 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -8,7 +8,7 @@ use crate::binary_transport::{ AsyncForwarder, BinaryStageExecutionOptions, forwarded_stage_message, forwarded_stage_message_timed, run_binary_stage_message, write_stage_message_conditioned, }; -use crate::frontend::embedded_execution::VerifyRetirement; +use crate::frontend::embedded_execution::{StaleWindowDiscard, VerifyRetirement}; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window, propose_configured_ngram_tokens, @@ -36,6 +36,7 @@ use lifecycle::{ compose_target_predictions, decode_uses_context_sideband, direct_prediction_return_path, mark_epoch_stale, open_upstream_prediction_return, pipelined_window_layout, queued_active_tokens, refill_pipeline_ngram_candidates, speculation_after_prefix_restore, + stale_window_id_range, }; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; @@ -794,19 +795,37 @@ impl StageOpenAiBackend { _ => None, }; let mut verify_window_scheduler = VerifyWindowScheduler::new( - VerifyWindowPipelineConfig::new(effective_speculative.verify_window.pipeline_depth), + if effective_speculative.verify_window.runahead_max_tokens > 0 { + VerifyWindowPipelineConfig::with_runahead( + effective_speculative.verify_window.runahead_max_tokens, + ) + } else { + VerifyWindowPipelineConfig::new( + effective_speculative.verify_window.pipeline_depth, + ) + }, ); + // Draft fallback keeps the draft model resident purely as a + // proposal source for the pipelined path when the N-gram proposer + // misses; the classic serial draft loop stays disabled. Pipelined + // paths only — at depth 1 the classic draft loop is strictly + // better. + let ngram_fallback_draft_enabled = effective_speculative.ngram_fallback_draft + && effective_speculative.ngram.is_some() + && draft_guard.is_some() + && verify_window_scheduler.depth() > 1; + let draft_blocks_pipeline = draft_guard.is_some() && !ngram_fallback_draft_enabled; let composite_sidecar_enabled = - native_mtp_options.ngram_hybrid && draft_guard.is_none(); + native_mtp_options.ngram_hybrid && !draft_blocks_pipeline; // A standalone N-gram plan (no native MTP, no draft model) can drive // the same verify-window pipeline; the single-window native-MTP path // stays composite-only, so standalone drafting still falls back to the // serial block at depth 1. let standalone_ngram_pipelining = !request.native_mtp_enabled && effective_speculative.ngram.is_some() - && draft_guard.is_none(); + && !draft_blocks_pipeline; let native_mtp_verify_windows_enabled = - (request.native_mtp_enabled || composite_sidecar_enabled) && draft_guard.is_none(); + (request.native_mtp_enabled || composite_sidecar_enabled) && !draft_blocks_pipeline; let pipelined_decode_enabled = (composite_sidecar_enabled || standalone_ngram_pipelining) && verify_window_scheduler.depth() > 1; @@ -888,6 +907,51 @@ impl StageOpenAiBackend { ), cached_ngram_proposer.as_mut(), )?; + let proposal = if proposal.tokens().len() < 2 + && ngram_fallback_draft_enabled + && pipelined_decode_enabled + // A fallback proposal is only worth its draft decode + // when at least two tokens still fit the remaining + // window; below that, fall through to the serial path + // rather than overshoot the budget by a token. + && native_mtp_remaining >= 2 + { + let fallback_timer = PhaseTimer::start(); + let draft = draft_guard + .as_deref_mut() + .expect("fallback requires a draft guard"); + // Propose from the token sync_to_context left + // unmaterialized rather than from `current`. The two + // agree today, but only by an invariant maintained + // across the whole decode loop; depending on it here + // would make a slip silent KV corruption instead of + // an error. + debug_assert_eq!(context_tokens.last(), Some(¤t)); + let Some(&propose_from) = context_tokens.last() else { + return Err(openai_backend_error(anyhow::anyhow!( + "draft fallback requires a non-empty context" + ))); + }; + draft + .sync_to_context(&context_tokens) + .map_err(openai_backend_error)?; + // The floor must not lift the budget back over the + // remaining window, so it is applied before the cap. + let budget = native_mtp_options + .ngram_max_proposal_tokens + .min(draft.window.max(1)) + .max(2) + .min(native_mtp_remaining); + let draft_tokens = draft + .propose(propose_from, budget) + .map_err(openai_backend_error)?; + speculative_stats.fallback_draft_proposals += 1; + speculative_stats.fallback_draft_tokens += draft_tokens.len(); + speculative_stats.fallback_draft_ms += fallback_timer.elapsed_ms(); + NativeMtpHybridProposal::from_parts(draft_tokens, 0, true) + } else { + proposal + }; if proposal.supports_positional_pipeline(verify_window_scheduler.depth()) && ngram_sidecar_controller.permit_pipeline_start() && verify_window_scheduler.supports_pipelining( @@ -988,6 +1052,12 @@ impl StageOpenAiBackend { && verify_window_scheduler.in_flight_len() < pipeline_in_flight_limit && decoded_tokens + queued_active_tokens(&pipelined_windows) < request.max_tokens as usize + // A window may need one budget token for its + // epoch-start boundary on top of the proposals, + // so wait for a retirement rather than plan a + // chunk the budget cannot fit. + && (verify_window_scheduler.in_flight_len() == 0 + || verify_window_scheduler.admissible_window_tokens() > 1) { let refill_threshold = chunk_width; if pipeline.candidate_len() < refill_threshold { @@ -998,18 +1068,52 @@ impl StageOpenAiBackend { ); let refill_budget = ngram_sidecar_controller.refill_limit(available_refill_tokens); - let appended = refill_pipeline_ngram_candidates( + let mut appended = refill_pipeline_ngram_candidates( pipeline, &context_tokens, &mut cached_ngram_proposer, refill_budget, )?; + if appended == 0 + && !pipeline.has_remaining_candidates() + && ngram_fallback_draft_enabled + && refill_budget >= 2 + { + let fallback_timer = PhaseTimer::start(); + let draft = draft_guard + .as_deref_mut() + .expect("fallback requires a draft guard"); + let mut sequence = context_tokens.clone(); + sequence.extend_from_slice(pipeline.optimistic_suffix()); + if let Some(&last) = sequence.last() { + draft + .sync_to_context(&sequence) + .map_err(openai_backend_error)?; + let budget = refill_budget.min(draft.window.max(1)); + let draft_tokens = draft + .propose(last, budget) + .map_err(openai_backend_error)?; + speculative_stats.fallback_draft_proposals += 1; + speculative_stats.fallback_draft_tokens += + draft_tokens.len(); + speculative_stats.fallback_draft_ms += + fallback_timer.elapsed_ms(); + appended = pipeline.append_ngram_candidates(&draft_tokens); + } + } verify_window_scheduler.record_horizon_refill(appended); } if !pipeline.has_remaining_candidates() { break; } - let Some(planned) = pipeline.next_chunk(chunk_width) else { + let budget_chunk_width = chunk_width + .min( + verify_window_scheduler + .admissible_window_tokens() + .saturating_sub(1), + ) + .max(1); + let Some(planned) = pipeline.next_chunk(budget_chunk_width) else { break; }; let proposal_tokens = planned.proposal_tokens().to_vec(); @@ -1024,8 +1128,11 @@ impl StageOpenAiBackend { current, &proposal_tokens, ); - let window = verify_window_scheduler - .open(layout.pos_start, layout.decode_step)?; + let window = verify_window_scheduler.open( + layout.pos_start, + layout.decode_step, + layout.input_tokens.len(), + )?; let input_tokens = layout.input_tokens; let message = embedded_verify_window_message( request.wire_dtype, @@ -1268,6 +1375,22 @@ impl StageOpenAiBackend { let stale_count = mark_epoch_stale(&mut pipelined_windows, pipeline_epoch); verify_window_scheduler.mark_recovery_epoch(stale_count); + if verify_window_scheduler.is_runahead() + && let Some((min_id, max_id)) = + stale_window_id_range(&pipelined_windows, pipeline_epoch) + { + self.discard_stale_windows( + &request, + downstream, + verify_window_forwarder.as_mut(), + StaleWindowDiscard { + request_id, + session_id, + min_window_id: min_id, + max_window_id: max_id, + }, + )?; + } let pipeline = pipelined.take().expect("pipeline retained"); if ngram_sidecar_controller.observe_tail_outcome( pipeline.proposal(), @@ -1361,7 +1484,7 @@ impl StageOpenAiBackend { continue; } } - if draft_guard.is_some() + if (draft_guard.is_some() && !ngram_fallback_draft_enabled) || (effective_speculative.ngram.is_some() && !pipelined_decode_enabled) { let remaining = (request.max_tokens as usize).saturating_sub(decoded_tokens); @@ -1828,6 +1951,31 @@ impl StageOpenAiBackend { if !pipelined_windows.is_empty() { let stale_count = mark_epoch_stale(&mut pipelined_windows, pipeline_epoch); verify_window_scheduler.mark_stale(stale_count); + if verify_window_scheduler.is_runahead() + && let Some((min_id, max_id)) = + stale_window_id_range(&pipelined_windows, pipeline_epoch) + { + self.discard_stale_windows( + &request, + downstream, + verify_window_forwarder.as_mut(), + StaleWindowDiscard { + request_id, + session_id, + min_window_id: min_id, + max_window_id: max_id, + }, + )?; + // Teardown discard: the request is ending and the lane + // goes back for reuse, so the discard must be fully on + // the wire before anything else writes to this socket. + // The mid-generation discard above deliberately does not + // wait, because everything behind it is queued on the + // same forwarder and stays ordered. + if let Some(forwarder) = verify_window_forwarder.as_mut() { + forwarder.flush().map_err(openai_backend_error)?; + } + } while let Some(stale) = pipelined_windows.pop_front() { let stale_drain_timer = PhaseTimer::start(); let stale_reply = self.complete_dispatched_stage_message_direct( diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index c724713847..80d00d526b 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -195,6 +195,24 @@ pub(super) fn can_seed_pipeline(windows: &VecDeque) -> windows.iter().all(|window| window.stale) } +/// Inclusive window-id range of the stale windows of `epoch`, if any. +pub(super) fn stale_window_id_range( + windows: &VecDeque, + epoch: u64, +) -> Option<(i32, i32)> { + let mut range: Option<(i32, i32)> = None; + for window in windows { + if window.epoch == epoch && window.stale { + let id = window.window.id; + range = Some(match range { + Some((min, max)) => (min.min(id), max.max(id)), + None => (id, id), + }); + } + } + range +} + pub(super) fn mark_epoch_stale( windows: &mut VecDeque, epoch: u64, diff --git a/crates/skippy-server/src/frontend/generation/draft_runner.rs b/crates/skippy-server/src/frontend/generation/draft_runner.rs index 76e53e5517..8cc8f678e1 100644 --- a/crates/skippy-server/src/frontend/generation/draft_runner.rs +++ b/crates/skippy-server/src/frontend/generation/draft_runner.rs @@ -20,6 +20,71 @@ pub(in crate::frontend) struct DraftRunner { pub(in crate::frontend) window: usize, pub(in crate::frontend) _model: StageModel, pub(in crate::frontend) session: StageSession, + /// Tokens currently materialized in the draft session's KV, maintained so + /// fallback proposals can advance incrementally instead of re-prefilling + /// the whole context on every call. + synced: DraftSyncState, +} + +/// What a sync to a given context requires of the draft session. Split out +/// from the session I/O because this bookkeeping is load-bearing for KV +/// correctness: claiming a prefix extension the session has not materialized +/// silently corrupts every later proposal. +#[derive(Debug, Eq, PartialEq)] +pub(in crate::frontend) enum DraftSyncPlan { + /// The session is already at the target; nothing to do. + AlreadySynced, + /// The synced tokens are a prefix of the target: prefill only the tail, + /// given here as a range into the target prefix. + Extend { from: usize, to: usize }, + /// The synced tokens diverge from the target: reset and prefill the + /// whole prefix. + Reset, +} + +/// Tokens the draft session has materialized, and the decisions derived from +/// them. +#[derive(Debug, Default)] +pub(in crate::frontend) struct DraftSyncState { + tokens: Vec, +} + +impl DraftSyncState { + /// The prefix a context implies: every token but the last, which is the + /// one a proposal decodes from. + fn target_len(context_tokens: &[i32]) -> usize { + context_tokens.len().saturating_sub(1) + } + + pub(in crate::frontend) fn plan(&self, context_tokens: &[i32]) -> DraftSyncPlan { + let target = &context_tokens[..Self::target_len(context_tokens)]; + if self.tokens.is_empty() || !target.starts_with(&self.tokens) { + return DraftSyncPlan::Reset; + } + if target.len() == self.tokens.len() { + return DraftSyncPlan::AlreadySynced; + } + DraftSyncPlan::Extend { + from: self.tokens.len(), + to: target.len(), + } + } + + fn record_extend(&mut self, delta: &[i32]) { + self.tokens.extend_from_slice(delta); + } + + fn record_reset(&mut self, prefix: &[i32]) { + self.tokens.clear(); + self.tokens.extend_from_slice(prefix); + } + + /// A proposal decodes from `current`, which the session materializes as + /// it steps — so it joins the synced prefix and the next sync can extend + /// instead of resetting. + fn record_proposal_step(&mut self, current: i32) { + self.tokens.push(current); + } } impl DraftRunner { @@ -70,19 +135,43 @@ impl DraftRunner { window, _model: model, session, + synced: DraftSyncState::default(), }) } pub(in crate::frontend) fn reset_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { self.session.reset().context("reset draft session")?; + self.synced.record_reset(&[]); if context_tokens.len() > 1 { + let prefix = &context_tokens[..context_tokens.len() - 1]; self.session - .prefill_chunk(&context_tokens[..context_tokens.len() - 1]) + .prefill_chunk(prefix) .context("prefill draft context")?; + self.synced.record_reset(prefix); } Ok(()) } + /// Brings the draft session to `context_tokens` (all but the last token + /// prefilled, ready to propose from the last). Extends incrementally when + /// the already-synced tokens are a prefix of the target — the common case + /// when prior fallback proposals were accepted — and falls back to a full + /// reset on divergence. + pub(in crate::frontend) fn sync_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { + match self.synced.plan(context_tokens) { + DraftSyncPlan::AlreadySynced => Ok(()), + DraftSyncPlan::Extend { from, to } => { + let delta = &context_tokens[from..to]; + self.session + .prefill_chunk(delta) + .context("advance draft context")?; + self.synced.record_extend(delta); + Ok(()) + } + DraftSyncPlan::Reset => self.reset_to_context(context_tokens), + } + } + pub(in crate::frontend) fn propose( &mut self, mut current: i32, @@ -90,10 +179,14 @@ impl DraftRunner { ) -> Result> { let mut tokens = Vec::with_capacity(max_tokens); for _ in 0..max_tokens { + // Record after the step succeeds: a failed decode must not leave + // the state claiming a token the session does not hold. + let stepped_from = current; current = self .session .decode_step(current) .context("draft decode step")?; + self.synced.record_proposal_step(stepped_from); tokens.push(current); } Ok(tokens) @@ -180,3 +273,93 @@ pub(in crate::frontend) fn model_layer_count(path: &Path) -> Result { .ok_or_else(|| anyhow!("could not infer layer count for {}", path.display()))?; Ok(layer_count) } + +#[cfg(test)] +mod tests { + use super::*; + + fn state(tokens: &[i32]) -> DraftSyncState { + DraftSyncState { + tokens: tokens.to_vec(), + } + } + + #[test] + fn an_empty_session_always_resets() { + assert_eq!(state(&[]).plan(&[1, 2, 3]), DraftSyncPlan::Reset); + // A context of one token has an empty prefix: still a reset, and + // `reset_to_context` then prefills nothing. + assert_eq!(state(&[]).plan(&[1]), DraftSyncPlan::Reset); + } + + #[test] + fn a_synced_prefix_extends_by_the_delta_only() { + // Synced [1, 2]; context [1, 2, 3, 4, 5] has prefix [1, 2, 3, 4]. + assert_eq!( + state(&[1, 2]).plan(&[1, 2, 3, 4, 5]), + DraftSyncPlan::Extend { from: 2, to: 4 } + ); + } + + #[test] + fn an_exactly_synced_prefix_is_a_no_op() { + // Synced [1, 2, 3]; context [1, 2, 3, 4] has prefix [1, 2, 3]. + assert_eq!( + state(&[1, 2, 3]).plan(&[1, 2, 3, 4]), + DraftSyncPlan::AlreadySynced + ); + } + + #[test] + fn divergence_resets_rather_than_extending() { + // Same length, different token: the KV past that point is wrong. + assert_eq!(state(&[1, 9]).plan(&[1, 2, 3, 4]), DraftSyncPlan::Reset); + // A rejected proposal leaves the session longer than the target. + assert_eq!(state(&[1, 2, 3, 4]).plan(&[1, 2, 3]), DraftSyncPlan::Reset); + } + + #[test] + fn proposal_steps_join_the_synced_prefix_so_the_next_sync_extends() { + let mut synced = state(&[1, 2]); + // Two proposal steps decoded from 3 then 4. + synced.record_proposal_step(3); + synced.record_proposal_step(4); + + // Both accepted, and the caller committed a fifth token: the session + // already holds [1, 2, 3, 4], so only [5] needs prefilling. + assert_eq!( + synced.plan(&[1, 2, 3, 4, 5, 6]), + DraftSyncPlan::Extend { from: 4, to: 5 } + ); + } + + #[test] + fn a_rejected_proposal_step_forces_a_reset() { + let mut synced = state(&[1, 2]); + synced.record_proposal_step(3); + // The verifier rejected 3 and committed 9 instead: the draft KV holds + // a token the target never accepted, so the prefix cannot be reused. + assert_eq!(synced.plan(&[1, 2, 9, 10]), DraftSyncPlan::Reset); + } + + #[test] + fn recording_a_reset_replaces_the_whole_prefix() { + let mut synced = state(&[1, 2, 3]); + synced.record_reset(&[7, 8]); + + assert_eq!(synced.plan(&[7, 8, 9]), DraftSyncPlan::AlreadySynced); + assert_eq!(synced.plan(&[1, 2, 3]), DraftSyncPlan::Reset); + } + + #[test] + fn an_extend_plan_indexes_the_context_the_caller_slices() { + // The plan's range must address `context_tokens` directly, since + // sync_to_context slices the context with it. + let context = [1, 2, 3, 4, 5]; + let DraftSyncPlan::Extend { from, to } = state(&[1, 2]).plan(&context) else { + panic!("a synced prefix must extend"); + }; + + assert_eq!(&context[from..to], &[3, 4]); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs index 7049364e80..22dcb61ddc 100644 --- a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs +++ b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs @@ -125,8 +125,11 @@ impl StageOpenAiBackend { 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 window = verify_window_scheduler.open( + prefill_token_count + *decoded_tokens, + *decoded_tokens, + verify_inputs.len(), + )?; let message = embedded_verify_window_message( request.wire_dtype, VerifyWindowMessageArgs { diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index 6abb8e61e1..9ce384dc0d 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -4,7 +4,7 @@ use openai_frontend::OpenAiResult; use serde::{Deserialize, Serialize}; use serde_json::Value; use serde_json::json; -use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; +use skippy_protocol::{MAX_VERIFY_WINDOW_PIPELINE_DEPTH, MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS}; use std::collections::BTreeMap; use std::path::PathBuf; use std::time::Instant; @@ -28,6 +28,11 @@ pub struct SpeculativeDecodeConfig { pub ngram: Option, pub extension: Option, pub verify_window: VerifyWindowConfig, + /// Propose from the configured draft model when the N-gram proposer + /// misses, instead of degrading to one token per round trip. Pipelined + /// paths only; requires a draft model. + #[serde(default)] + pub ngram_fallback_draft: bool, } /// Native multi-token-prediction draft settings. @@ -92,6 +97,12 @@ pub struct VerifyWindowConfig { pub min_tokens: usize, pub max_tokens: usize, pub pipeline_depth: usize, + /// Run-ahead speculative-token budget. Zero keeps the fixed + /// `pipeline_depth` window-count admission; a positive value switches the + /// scheduler to token-budget admission (windows stay capped at the native + /// checkpoint-retention bound). + #[serde(default)] + pub runahead_max_tokens: usize, } impl Default for SpeculativeDecodeConfig { @@ -113,7 +124,9 @@ impl Default for SpeculativeDecodeConfig { min_tokens: 1, max_tokens: 4, pipeline_depth: 1, + runahead_max_tokens: 0, }, + ngram_fallback_draft: false, } } } @@ -173,6 +186,14 @@ impl SpeculativeDecodeConfig { "verify window requires 0 < min_tokens <= max_tokens and 0 < pipeline_depth <= {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" ); } + if self.ngram_fallback_draft && self.ngram.is_none() { + bail!("ngram_fallback_draft requires an N-gram proposer to fall back from"); + } + if self.verify_window.runahead_max_tokens > MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS { + bail!( + "verify window runahead_max_tokens must not exceed {MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS}" + ); + } Ok(()) } @@ -367,6 +388,9 @@ mod standalone_speculative_config_tests { pub(super) struct OpenAiSpeculativeStats { pub(super) windows: usize, pub(super) draft_tokens: usize, + pub(super) fallback_draft_proposals: usize, + pub(super) fallback_draft_tokens: usize, + pub(super) fallback_draft_ms: f64, pub(super) accepted_tokens: usize, pub(super) rejected_tokens: usize, pub(super) full_accept_windows: usize, @@ -598,6 +622,18 @@ fn elapsed_us(started: Instant) -> u64 { impl OpenAiSpeculativeStats { pub(super) fn insert_response_timings(&self, timings: &mut BTreeMap) { timings.insert("speculative_windows".to_string(), json!(self.windows)); + timings.insert( + "speculative_fallback_draft_proposals".to_string(), + json!(self.fallback_draft_proposals), + ); + timings.insert( + "speculative_fallback_draft_tokens".to_string(), + json!(self.fallback_draft_tokens), + ); + timings.insert( + "speculative_fallback_draft_ms".to_string(), + json!(self.fallback_draft_ms), + ); timings.insert( "speculative_proposed_n".to_string(), json!(self.draft_tokens), diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 4fc48b42bc..e00a6670e4 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -174,6 +174,39 @@ pub(super) fn embedded_verify_window_message( }) } +/// Invalidates verify windows `min_window_id..=max_window_id` for a request +/// after divergence. The downstream records the range at receive time so +/// buffered stale windows are answered with an empty reply instead of being +/// executed. The window-id range rides in `tokens`. +pub(super) fn discard_stale_windows_message( + wire_dtype: WireActivationDType, + request_id: u64, + session_id: u64, + min_window_id: i32, + max_window_id: i32, +) -> OpenAiResult { + if min_window_id > max_window_id { + return Err(OpenAiError::backend( + "stale window discard range must be non-empty", + )); + } + let kind = WireMessageKind::DiscardStaleWindows; + Ok(StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, wire_dtype), + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![min_window_id, max_window_id], + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }) +} + pub(super) fn retire_verify_window_message( wire_dtype: WireActivationDType, request_id: u64, diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index c5fd451b7a..ef13dc5f4e 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -848,8 +848,8 @@ cached and a worker does not: to open `skippy-stage/2`, then Skippy artifact-transfer stream 0x03, to fetch only its assigned package files before the normal HF fallback path. - Current/released mixed mesh: a released coordinator without advertised - `skippy-stage/2` `artifact-transfer`, `stage-generation-4`, and - `direct-prediction-return` support must not be selected for a generation-4 + `skippy-stage/2` `artifact-transfer`, `stage-generation-5`, and + `direct-prediction-return` support must not be selected for a generation-5 split topology; the worker must fall back to local/HF package resolution. - Default public-mesh safety: with `MESH_LLM_ARTIFACT_TRANSFER` unset, the node must advertise no `artifact-transfer` feature, reject inbound artifact diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 9f78d34a17..1042d19b8f 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -40,11 +40,11 @@ activation links and then crossed three reply links before stage 0 could emit the token. On a topology with a fixed 10 ms delay per inter-stage hop, the reply chain alone makes the hot path six hops, or about 60 ms before compute. -## Generation 4 Direct Prediction Return and Verify Retirement +## Generation 5 Direct Prediction Return and Verify Retirement -Stage protocol generation 4 is a compatibility-breaking change. A peer is stage +Stage protocol generation 5 is a compatibility-breaking change. A peer is stage compatible only when it advertises both `skippy-stage/2` and -`stage-generation-4`. Prediction-bearing messages return +`stage-generation-5`. Prediction-bearing messages return directly from the final/readout stage to the driver-facing stage. Intermediate stages continue to forward activations and may handle cold-path control acknowledgments, but they are not part of the decode-token prediction return path. @@ -79,6 +79,38 @@ With the same four stages and 10 ms inter-stage delay, the no-spec decode hot path becomes `S0 -> S1 -> S2 -> S3 -> S0`: four hops, or about 40 ms before compute. That removes two serialized reply hops from every generated token. +## Stale Verify-Window Discard + +Generation 5 also adds the `DiscardStaleWindows` control frame (wire kind +23), which is what made the generation compatibility-breaking: a +generation-4 peer rejects the kind outright and drops the request +connection. + +Run-ahead admission dispatches verify windows before their predecessors are +verified, so a rejection strands every window queued behind the divergence. +Those windows are already on the wire. The coordinator sends one discard +naming a contiguous `[min_window_id, max_window_id]` range for the request, +and it travels the same ordered path as the windows it cancels, so it is +read after them. + +The handling rule differs by position in the chain, and this is the +non-obvious part: + +- **Middle stages forward the frame and still execute the stale windows.** + A stage that is not the final stage owns KV state its downstream neighbour + depends on; skipping its forward pass would desynchronize the chain. It + passes the discard along so the frame reaches the stage that can act on it. +- **The final/readout stage skips the named windows.** It is the only stage + whose output is thrown away by a discard, so skipping there is what + actually saves the work — the expensive readout and the direct prediction + return back to stage 0. + +The receiving stage records the range the moment it parses the frame, ahead +of executing the backlog queued in front of it, which is what lets the skip +take effect before the stale tail runs. See `INBOUND_LOOKAHEAD_BYTES` for +the case where a wide-frame backlog outruns that read-ahead and the tail +executes anyway. + ## Relative Sizes | Flow | Size | diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 001b401d61..7ab74f3f2b 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4205,27 +4205,27 @@ ], "crates/skippy-server/src/binary_transport/binary_messaging.rs": [ { - "line": 300, + "line": 301, "macro_name": "eprintln!" }, { - "line": 310, + "line": 311, "macro_name": "println!" }, { - "line": 333, + "line": 334, "macro_name": "eprintln!" }, { - "line": 353, + "line": 354, "macro_name": "eprintln!" }, { - "line": 361, + "line": 362, "macro_name": "eprintln!" }, { - "line": 415, + "line": 416, "macro_name": "eprintln!" } ],