diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index 434776daed..3c99437a87 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -601,6 +601,7 @@ 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 spec_default: Option, pub(crate) legacy_draft_model_path_used: bool, } @@ -652,6 +653,7 @@ 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), spec_default: pick!(spec_default), legacy_draft_model_path_used: overrides .filter(|config| config.draft_model.is_some()) @@ -725,6 +727,8 @@ struct SpeculativeConfigRaw { #[serde(default)] verify_window_pipeline_depth: Option, #[serde(default)] + verify_window_runahead_tokens: Option, + #[serde(default)] spec_default: Option, } @@ -769,6 +773,7 @@ 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, spec_default: raw.spec_default, legacy_draft_model_path_used: legacy_used, }) @@ -831,6 +836,10 @@ 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("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..0698071779 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,8 @@ 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" => {} _ => {} } } 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 ccd02db1ee..cd4fe39c3c 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 @@ -754,6 +754,10 @@ 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}.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 e18b504e77..6d4d5ac6a9 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, @@ -687,6 +687,14 @@ 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"), ) } diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index a924b901c5..58581d12eb 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -1257,6 +1257,13 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ reason: "", behavior: WiringBehavior::None, }, + WiringEntry { + path: "speculative.verify_window_runahead_tokens", + status: WiringStatus::Wired, + owner: "n/a", + reason: "", + behavior: WiringBehavior::None, + }, WiringEntry { path: "speculative.spec_default", status: WiringStatus::Wired, 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 d2193a555a..f343222492 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -102,26 +102,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)] @@ -1482,9 +1492,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 ); } @@ -1492,7 +1505,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/package.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs index 438de395cf..549dd56007 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs @@ -270,7 +270,7 @@ pub struct SkippyPackageSourceFile { /// Resolve a validated package-v2 directory into the existing host planning /// identity. The content-derived package ID remains in the v2 manifest and is -/// carried into split control by the generation-9 admission descriptor. +/// carried into split control by the generation-10 admission descriptor. pub fn identity_from_package_v2(package_dir: &Path) -> Result { let package_dir = package_dir.canonicalize().with_context(|| { format!( @@ -964,7 +964,7 @@ pub(crate) fn direct_gguf_source_paths(model_path: &Path) -> Result .collect() } -/// Build the source-complete metadata envelope required by generation-9 +/// Build the source-complete metadata envelope required by generation-10 /// planning directly from local GGUF shards. The envelope is in memory: the /// source files remain at their original paths and no package or layer shard /// is written. 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 b0aec41eaa..21d3d43019 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 @@ -463,6 +463,17 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result config.verify_window.max_tokens { bail!("skippy speculative verify window requires min_tokens <= max_tokens"); } @@ -545,6 +556,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", @@ -631,6 +643,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 11e120a260..ff18364e7b 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 @@ -2216,3 +2216,46 @@ placement = "auto" assert!(error.contains("defaults.hardware.placement")); } + +#[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" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs b/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs index ac7168a887..939a492f29 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs @@ -632,7 +632,7 @@ pub(super) fn stage_load_mode_from_proto(value: i32) -> anyhow::Result { - anyhow::bail!("unsupported generation-9 stage load mode {value}") + anyhow::bail!("unsupported generation-10 stage load mode {value}") } } } @@ -672,7 +672,7 @@ fn stage_activation_codec_from_proto( Ok(skippy_stage_proto::StageActivationCodec::S8RowF32RneV1) => { Ok(skippy_protocol::StageActivationCodec::S8RowF32RneV1) } - _ => anyhow::bail!("unsupported generation-9 activation codec {value}"), + _ => anyhow::bail!("unsupported generation-10 activation codec {value}"), } } @@ -694,7 +694,7 @@ fn stage_activation_codec_policy_from_proto( ) -> anyhow::Result { match skippy_stage_proto::StageActivationCodecPolicy::try_from(value) { Ok(skippy_stage_proto::StageActivationCodecPolicy::Unspecified) => { - anyhow::bail!("generation-9 activation codec policy must be explicit") + anyhow::bail!("generation-10 activation codec policy must be explicit") } Ok(skippy_stage_proto::StageActivationCodecPolicy::FixedV1) => { Ok(skippy_protocol::StageActivationCodecPolicy::Fixed) @@ -702,7 +702,7 @@ fn stage_activation_codec_policy_from_proto( Ok(skippy_stage_proto::StageActivationCodecPolicy::AutoLosslessV1) => { Ok(skippy_protocol::StageActivationCodecPolicy::AutoLosslessV1) } - _ => anyhow::bail!("unsupported generation-9 activation codec policy {value}"), + _ => anyhow::bail!("unsupported generation-10 activation codec policy {value}"), } } @@ -1007,7 +1007,7 @@ fn source_resolution_policy_from_proto(value: i32) -> anyhow::Result { Ok(skippy_stage_proto::SourceResolutionPolicy::Fallback) => Ok(false), Ok(skippy_stage_proto::SourceResolutionPolicy::LocalRequired) => Ok(true), Ok(skippy_stage_proto::SourceResolutionPolicy::Unspecified) | Err(_) => { - anyhow::bail!("unsupported generation-9 stage source resolution policy {value}") + anyhow::bail!("unsupported generation-10 stage source resolution policy {value}") } } } diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests/stage_control.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/stage_control.rs index e622a304f8..219649995b 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests/stage_control.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests/stage_control.rs @@ -150,7 +150,7 @@ async fn stage_control_bundle_gate_rejects_legacy_peer() -> Result<()> { assert!( error .to_string() - .contains("does not advertise the required generation-9 control bundle"), + .contains("does not advertise the required generation-10 control bundle"), "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 6bcd8e9b35..11610569ff 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/convert.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/convert.rs @@ -22,7 +22,7 @@ fn skippy_stage_subprotocols( } if stage_protocol_generation_supported { features.push( - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10.to_string(), ); } if artifact_transfer_supported { @@ -61,7 +61,7 @@ fn supports_local_gguf_content_id(subprotocols: &[crate::proto::node::MeshSubpro fn supports_skippy_stage_generation(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { let required_features = [ - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1, 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 e7921da3f2..c337269d5c 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs @@ -85,7 +85,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_V9)); + == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10)); assert_eq!( proto_pa .owner_attestation @@ -441,13 +441,13 @@ fn proto_announcement_without_required_generation_bundle_is_not_stage_compatible let peer_id = EndpointId::from(SecretKey::from_bytes(&[0xD0; 32]).public()); for missing in [ skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1, ] { let features = [ skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1, ] @@ -469,7 +469,7 @@ fn proto_announcement_without_required_generation_bundle_is_not_stage_compatible let (_, ann) = proto_ann_to_local(&proto_pa).expect("proto announcement should decode"); assert!( !ann.stage_protocol_generation_supported, - "missing {missing} must reject the generation-7 bundle" + "missing {missing} must reject the generation-10 bundle" ); } } @@ -486,7 +486,7 @@ fn partial_duplicate_stage_records_do_not_form_a_generation_bundle() { major: skippy_protocol::STAGE_SUBPROTOCOL_MAJOR, features: vec![ skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL.to_string(), - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9 + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10 .to_string(), ], }, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs index 9886d41eb1..4ed315437a 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs @@ -153,7 +153,7 @@ pub(super) async fn resolve_split_runtime_package( } anyhow::ensure!( model_path.is_file(), - "generation-9 split source must be a package-v2 directory or direct GGUF file: {}", + "generation-10 split source must be a package-v2 directory or direct GGUF file: {}", model_path.display() ); if local_source_required { diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs index 5d73f9ce6e..d33ad1cf5f 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs @@ -544,7 +544,7 @@ fn realize_split_stage_admissions( } else { anyhow::ensure!( model_path.is_file(), - "generation-9 direct-GGUF split source must be a local file: {}", + "generation-10 direct-GGUF split source must be a local file: {}", model_path.display() ); super::stage_admission::realize_direct_gguf_stage_admissions( @@ -556,7 +556,7 @@ fn realize_split_stage_admissions( "skippy-backend:auto:v1", ) } - .context("realize and admit generation-9 native stage chain") + .context("realize and admit generation-10 native stage chain") } async fn elect_split_start_coordinator( diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs index f27f090b32..cb58af5242 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs @@ -1339,7 +1339,7 @@ mod activation_boundary_tests { #[test] fn missing_graph_boundary_is_not_reconstructed_from_manifest_width() { let error = required_boundary(None, "stage-1", "input") - .expect_err("generation 7 requires graph-observed boundary descriptors"); + .expect_err("generation 10 requires graph-observed boundary descriptors"); assert!( error .to_string() diff --git a/crates/mesh-llm-host-runtime/src/runtime/stage_admission.rs b/crates/mesh-llm-host-runtime/src/runtime/stage_admission.rs index c5e520e85f..648c73cb4b 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/stage_admission.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/stage_admission.rs @@ -21,7 +21,7 @@ use skippy_package_format::{PackageManifest, Sidecar}; /// The planned admission expectations for one stage. /// /// This is carried by the planner on `RuntimeSliceStagePlan` from planning -/// time and mirrored into the generation-9 control protocol descriptor. +/// time and mirrored into the generation-10 control protocol descriptor. #[derive(Clone, Debug, Eq, PartialEq)] pub struct PlannedStageAdmission { /// Content-derived package identity (`sha256:...`). @@ -465,7 +465,7 @@ fn realize_native_stage_chain_from_manifest( } /// Realize, package-resolve, and admit every stage before a topology can be -/// published. Returned descriptors are canonical generation-9 wire values. +/// published. Returned descriptors are canonical generation-10 wire values. pub fn realize_stage_admissions( package_dir: &Path, ranges: &[(u32, u32)], 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 ac43057d6d..c84e9bbcdb 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 @@ -784,6 +784,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", @@ -841,4 +848,4 @@ } } ] -} +} \ No newline at end of file diff --git a/crates/skippy-protocol/README.md b/crates/skippy-protocol/README.md index 6e0e8bc894..2817a3a1d8 100644 --- a/crates/skippy-protocol/README.md +++ b/crates/skippy-protocol/README.md @@ -37,16 +37,16 @@ sequenceDiagram S0-->>D: PredictedToken ``` -Activation payloads dominate the wire path. Protocol generation 7 fixes every +Activation payloads dominate the wire path. Protocol generation 8 fixes every activation frame to raw little-endian `f32`; there is no negotiated or per-request activation dtype. -Generation 9 requires mesh-subprotocol control, list-valued status responses, -strict local-content identity, and canonical stage-admission descriptors as one -fail-closed capability bundle. Participants validate the descriptor while loading -and echo it when ready; any package, plan, range, tensor, sidecar, profile, backend, or -graph-configuration mismatch rejects the stage. The dedicated `skippy-stage/2` -ALPN accepts activation transport only. +Generation 10 requires mesh-subprotocol control, list-valued status responses, +strict local-content identity, canonical stage-admission descriptors, and stale +verify-window discard as one fail-closed capability bundle. Participants validate +the descriptor while loading and echo it when ready; any package, plan, range, +tensor, sidecar, profile, backend, or graph-configuration mismatch rejects the +stage. The dedicated `skippy-stage/2` ALPN accepts activation transport only. ## Responsibilities diff --git a/crates/skippy-protocol/src/admission.rs b/crates/skippy-protocol/src/admission.rs index 26e3aa30be..5b454a48fb 100644 --- a/crates/skippy-protocol/src/admission.rs +++ b/crates/skippy-protocol/src/admission.rs @@ -1,6 +1,6 @@ -//! Canonical generation-9 stage admission descriptors. +//! Canonical generation-10 stage admission descriptors. -/// Current descriptor schema carried by stage-control generation 9. +/// Current descriptor schema carried by stage-control generation 10. pub const STAGE_ADMISSION_DESCRIPTOR_VERSION: u32 = 1; #[derive(Clone, Debug, Eq, PartialEq)] diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index 03ef52b6a3..72f80cb7a8 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -38,6 +38,7 @@ pub enum WireMessageKind { DecodeLightCtx = 9, VerifyWindow = 21, RetireVerifyWindow = 22, + DiscardStaleWindows = 23, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -81,6 +82,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) } @@ -127,6 +136,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")), } } @@ -352,6 +362,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/config.rs b/crates/skippy-protocol/src/config.rs index 25afdc68d0..57332871c0 100644 --- a/crates/skippy-protocol/src/config.rs +++ b/crates/skippy-protocol/src/config.rs @@ -225,7 +225,7 @@ pub struct StageConfig { #[serde(default)] pub generation_signal_window: Option, /// Floating-point activation encoding for every downstream edge produced - /// by this stage. Generation 9 peers echo and bind this policy before the + /// by this stage. Generation 10 peers echo and bind this policy before the /// binary data plane starts. #[serde(default)] pub activation_codec: StageActivationCodec, diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 1cf0df9f74..23a7e5e4b8 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -35,12 +35,12 @@ 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_LOCAL_GGUF_CONTENT_ID_V1, STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, STAGE_SUBPROTOCOL_MAJOR, STAGE_SUBPROTOCOL_NAME, StageFrameError, validate_stage_admission_descriptor, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, @@ -206,7 +206,7 @@ mod tests { ); } use super::{ - STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, + STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, StageFrameError, validate_stage_admission_descriptor, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, @@ -216,7 +216,7 @@ mod tests { #[test] fn stage_protocol_generation_feature_names_current_generation() { assert_eq!( - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10, format!("stage-generation-{STAGE_PROTOCOL_GENERATION}") ); } @@ -483,14 +483,20 @@ mod tests { Err(StageFrameError::MissingStageControlCommand) )); + let previous_generation = STAGE_PROTOCOL_GENERATION - 1; let wrong_gen = StageControlRequest { - r#gen: STAGE_PROTOCOL_GENERATION - 1, + r#gen: previous_generation, ..frame }; - assert!(matches!( + // Compared against the computed previous generation rather than a + // literal, so the assertion keeps testing rejection of the previous + // generation across bumps instead of failing on the number. + assert_eq!( validate_stage_control_request(&wrong_gen), - Err(StageFrameError::BadGeneration { got }) if got == STAGE_PROTOCOL_GENERATION - 1 - )); + Err(StageFrameError::BadGeneration { + got: previous_generation, + }) + ); } #[test] diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index 5b745fb310..6921f97368 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 = 9; +pub const STAGE_PROTOCOL_GENERATION: u32 = 10; /// 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_V9: &str = "stage-generation-9"; +pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10: &str = "stage-generation-10"; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V9; + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V10; pub const STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER: &str = "artifact-transfer"; pub const STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST: &str = "status-list"; pub const STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1: &str = "local-gguf-content-id-v1"; @@ -22,6 +22,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 { @@ -115,23 +119,23 @@ impl std::fmt::Display for StageFrameError { StageFrameError::MissingStageAdmissionDescriptor => { write!( f, - "generation 9 stage load/status requires an admission descriptor" + "generation 10 stage load/status requires an admission descriptor" ) } StageFrameError::MissingLoadClaimHashes => { write!( f, - "generation 9 stage load requires participant and topology hashes" + "generation 10 stage load requires participant and topology hashes" ) } StageFrameError::InvalidActivationCodec { got } => { - write!(f, "unsupported generation-9 activation codec {got}") + write!(f, "unsupported generation-10 activation codec {got}") } StageFrameError::InvalidActivationCodecPolicy { got } => { - write!(f, "unsupported generation-9 activation codec policy {got}") + write!(f, "unsupported generation-10 activation codec policy {got}") } StageFrameError::InvalidTopologyStages(reason) => { - write!(f, "invalid generation-9 topology stage list: {reason}") + write!(f, "invalid generation-10 topology stage list: {reason}") } StageFrameError::InvalidStageAdmissionDescriptor(reason) => { write!(f, "invalid stage admission descriptor: {reason}") diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 76e5972eb5..580a9bc660 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -16,11 +16,11 @@ 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. Generation 7 introduced -direct prediction return from the final/readout tip to the driver-facing stage. -Generation 9 retains that path and requires mandatory canonical stage-admission -descriptors with exact participant echo before topology publication. Middle-out -is the prefill optimization inside that path, where internal boundary -activations are handed downstream while local compute advances. +direct prediction return from the final/readout tip to the driver-facing stage, +and generation 9 added canonical stage-admission descriptors. Generation 10 +retains both contracts and adds stale verify-window discard for run-ahead +execution. Middle-out is the prefill optimization inside that path, where +internal boundary activations are handed downstream while local compute advances. ```mermaid flowchart LR @@ -126,11 +126,13 @@ deadline handling. ## Notes - `serve-binary` is the tuned binary stage-to-stage path. -- `serve-binary` participates in the breaking generation-9 stage protocol. - Stage compatibility requires the complete `stage-generation-9` control, - status-list, strict-content-identity, and stage-admission bundle. Older peers, - including generation 7 peers, are rejected during split planning rather than - being mixed into a generation-9 topology. +- `serve-binary` participates in the breaking generation-10 stage protocol. + Stage compatibility requires the complete `stage-generation-10` control, + status-list, strict-content-identity, stage-admission, and stale-window-discard + bundle. Older peers are rejected during split planning rather than being mixed + into a generation-10 topology. A manually wired `serve-binary --downstream` + chain has no generation handshake, so every stage in that chain must be + upgraded together. - `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 @@ -146,7 +148,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-7 prediction returns flow directly from the final stage to stage 0. + 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 7d15876de8..9bee6d14e5 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -41,6 +41,7 @@ mod prefill_recording; pub(in crate::binary_transport) mod reply; mod session_lifecycle; mod session_tracker; +mod stale_discard; mod summary; mod telemetry; @@ -553,7 +554,7 @@ fn run_binary_stage( native_mtp_enabled, &prediction_return_sinks, session_ownership, - &task_control, + task_control.clone(), first_message, ) })() 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 c601354b97..04618b998f 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 @@ -22,8 +22,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 { @@ -52,10 +68,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), }) } @@ -79,6 +97,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, condition, @@ -109,7 +129,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()?; } @@ -216,6 +236,106 @@ 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), + condition, + BTreeMap::new(), + ) + .unwrap(); + drop(forwarder); + + write_stage_message_after_propagation( + &mut client, + &message(WireMessageKind::Stop, 22), + 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); + } + + /// The lane-reuse half of the same property: after a delayed discard and + /// the teardown `Stop`, the socket must be clean enough to serve the next + /// request. A frame left half-written by the torn-down forwarder would be + /// read as the next request's header, so this asserts both frames arrive + /// whole and in order and that the next request's traffic follows them + /// undisturbed on the same socket. + #[test] + fn a_reused_lane_carries_the_next_request_after_a_delayed_teardown() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let mut lane = 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(&lane, telemetry, 8).unwrap(); + let delayed = WireCondition::new(250.0, None).unwrap(); + let immediate = WireCondition::new(0.0, None).unwrap(); + + // First request: a discard still in flight when the request ends. + forwarder + .send( + message(WireMessageKind::DiscardStaleWindows, 31), + delayed, + BTreeMap::new(), + ) + .unwrap(); + drop(forwarder); + write_stage_message_after_propagation( + &mut lane, + &message(WireMessageKind::Stop, 32), + immediate, + ) + .unwrap(); + + // The pool hands the same socket to the next request. + write_stage_message_after_propagation( + &mut lane, + &message(WireMessageKind::VerifyWindow, 41), + immediate, + ) + .unwrap(); + + let frames = (0..3) + .map(|_| read_stage_message(&mut server, 4).unwrap()) + .collect::>(); + let observed = frames + .iter() + .map(|frame| (frame.kind, frame.pos_start)) + .collect::>(); + + assert_eq!( + observed, + vec![ + (WireMessageKind::DiscardStaleWindows, 31), + (WireMessageKind::Stop, 32), + (WireMessageKind::VerifyWindow, 41), + ], + "a lane returned to the pool must carry the next request intact" + ); + } + #[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 ed49851851..92fb8dddac 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -4,7 +4,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_target, record_session_auto_align}; @@ -12,6 +12,7 @@ use super::session_tracker::{ ConnectionSessionOwnership, 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; @@ -90,7 +91,7 @@ pub(super) fn handle_binary_connection( native_mtp_enabled: bool, prediction_return_sinks: &PredictionReturnSinks, session_ownership: Arc, - worker_control: &ConnectionWorkerControl, + worker_control: Arc, first_message: StageWireMessage, ) -> Result<()> { let mut session_tracker = @@ -144,7 +145,7 @@ fn handle_binary_connection_messages( downstream_connect_timeout_secs: u64, native_mtp_enabled: bool, prediction_return_sinks: &PredictionReturnSinks, - worker_control: &ConnectionWorkerControl, + worker_control: Arc, first_message: StageWireMessage, session_tracker: &mut ConnectionSessionTracker, ) -> Result<()> { @@ -156,6 +157,16 @@ 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, + input_activation_width, + config.activation_codec, + config.activation_codec_policy, + max_inflight.max(1), + discard_registry.clone(), + worker_control, + )?; let mut async_forwarder = if async_prefill_forward || max_inflight > 1 { downstream .as_ref() @@ -171,11 +182,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, - worker_control, - input_activation_width, - config, + let Some(mut message) = inbound_reader.next( next_message.take(), pending_prefill_replies, request_summary.message_count, @@ -223,6 +230,27 @@ 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, downstream_wire_condition, BTreeMap::new()) + .context("forward stale window discard downstream")?; + } else { + write_stage_message_conditioned( + &mut *downstream, + &message, + downstream_wire_condition, + ) + .context("forward stale window discard downstream")?; + } + } + continue; + } + if message.kind.is_verify_retirement() { handle_verify_retirement( iteration_scheduler, @@ -307,6 +335,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; @@ -352,9 +412,8 @@ fn handle_binary_connection_messages( let lookup_kv = kv.cloned(); let lookup_telemetry = telemetry.clone(); let lookup_session_key = session_key.clone(); - let lookup_message = message.clone(); let lookup_token_ids = token_ids.clone(); - let (auto_align, lookup_result) = iteration_scheduler + let (returned_message, auto_align, lookup_result) = iteration_scheduler .execute_runtime("binary-prefix-lookup", move |runtime| { let auto_align = align_session_to_target( runtime, @@ -363,21 +422,20 @@ fn handle_binary_connection_messages( align_target, ) .map_err(|error| openai_frontend::OpenAiError::backend(format!("{error:#}")))?; - Ok(( - auto_align, - maybe_lookup_binary_prefill( - &lookup_config, - runtime, - lookup_kv.as_ref(), - &lookup_telemetry, - &lookup_session_key, - &lookup_message, - &lookup_token_ids, - output_activation_width, - ), - )) + let lookup_result = maybe_lookup_binary_prefill( + &lookup_config, + runtime, + lookup_kv.as_ref(), + &lookup_telemetry, + &lookup_session_key, + &message, + &lookup_token_ids, + output_activation_width, + ); + Ok((message, auto_align, lookup_result)) }) .map_err(|error| anyhow::anyhow!(format!("{error:#}")))?; + message = returned_message; session_auto_align_count = auto_align.count; session_auto_align_ms = auto_align.elapsed_ms; session_auto_align_trimmed_tokens = auto_align.trimmed_tokens; @@ -496,12 +554,23 @@ fn handle_binary_connection_messages( let sample_prefill_final = message.kind == WireMessageKind::PrefillFinalEmbd && downstream.is_none(); let scheduler_session_key = session_key.clone(); - let scheduler_message = message.clone(); let scheduler_token_ids = executable_token_ids.to_vec(); let scheduler_kv = kv.cloned(); let scheduler_telemetry = telemetry.clone(); let align_in_compute = !lookup_needed; let collect_session_stats = telemetry.is_debug_enabled(); + let execute_context = format!( + "execute scheduler-owned binary stage message \ + kind={:?} pos_start={} token_count={} tokens={} \ + executable_tokens={} activation_bytes={}", + message.kind, + message.pos_start, + message.token_count, + message.tokens.len(), + executable_token_ids.len(), + input_activation_bytes, + ); + let scheduler_message = message; let outcome = iteration_scheduler .execute_runtime_timed("binary-stage-execute", move |runtime| { let auto_align = if align_in_compute { @@ -552,6 +621,7 @@ fn handle_binary_connection_messages( let sessions_after = collect_session_stats.then(|| runtime.session_stats()); Ok(( + scheduler_message, auto_align, sessions_before, sessions_after, @@ -560,24 +630,19 @@ fn handle_binary_connection_messages( )) }) .map_err(|error| anyhow::anyhow!(format!("{error:#}"))) - .with_context(|| { - format!( - "execute scheduler-owned binary stage message \ - kind={:?} pos_start={} token_count={} tokens={} \ - executable_tokens={} activation_bytes={}", - message.kind, - message.pos_start, - message.token_count, - message.tokens.len(), - executable_token_ids.len(), - input_activation_bytes, - ) - })?; + .with_context(|| execute_context)?; runtime_lock_wait_ms = outcome.runtime_lock_wait_ms; runtime_lock_hold_ms = outcome.runtime_lock_hold_ms; runtime_lock_acquires = 1; - let (auto_align, sessions_before, sessions_after, eviction, result) = - outcome.value; + let ( + returned_message, + auto_align, + sessions_before, + sessions_after, + eviction, + result, + ) = outcome.value; + message = returned_message; if align_in_compute { session_auto_align_count = auto_align.count; session_auto_align_ms = auto_align.elapsed_ms; 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 9423cbcd7f..0ea36e9505 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,12 +1,18 @@ -use super::ConnectionWorkerControl; use anyhow::{Context, Result}; use skippy_protocol::{ - StageConfig, + StageActivationCodec, StageActivationCodecPolicy, binary::{StageWireMessage, read_stage_message_for_codec_policy}, }; 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::ConnectionWorkerControl; +use super::stale_discard::StaleDiscardRegistry; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); @@ -14,41 +20,382 @@ 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, - worker_control: &ConnectionWorkerControl, +/// 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. +/// +/// The reader thread and this handle share one `TcpStream` rather than +/// holding separate clones. Shutdown has to interrupt the read that is +/// actually in flight: a peer can send a frame prefix and then stall, which +/// leaves the reader inside `read_exact` past the readable check, and on +/// Windows shutting down a *cloned* handle does not interrupt a read pending +/// on a different one (#1538). Since `Drop` joins the thread, that would hang +/// the dropping handler rather than merely leak a thread. +pub(super) struct InboundMessageReader { + receiver: Option>>, + stream: Arc, + 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` — including one stalled + // part-way through a frame — by shutting down the exact handle the + // reader thread is blocked on. 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, - config: &StageConfig, - first_message: Option, - pending_prefill_replies: usize, - observed_message_count: usize, -) -> Result> { - if first_message.is_some() { - return Ok(first_message); - } - if !worker_control - .wait_for_readable(upstream) - .context("wait for the next binary stage message")? - { - // Shutdown was requested while the connection was idle: end the - // worker cleanly instead of blocking in a read that a socket - // shutdown cannot interrupt on Windows (#1538). - return Ok(None); - } - match read_stage_message_for_codec_policy( - upstream, - activation_width, - config.activation_codec, - config.activation_codec_policy, - ) { - 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"), + activation_codec: StageActivationCodec, + activation_codec_policy: StageActivationCodecPolicy, + capacity: usize, + registry: Arc, + worker_control: Arc, +) -> Result { + let stream = Arc::new( + upstream + .try_clone() + .context("clone upstream stream for inbound message reader")?, + ); + let reader = stream.clone(); + 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) || worker_control.is_shutting_down() { + return; + } + thread::sleep(Duration::from_millis(1)); + } + // Wait for readability under a short timeout so an idle reader + // observes shutdown instead of parking in a read that a socket + // shutdown cannot interrupt on Windows (#1538). This only covers + // an idle socket; a peer that sends a frame prefix and stalls + // leaves the read below blocked mid-frame, which is why `Drop` + // shuts down this exact handle. + match worker_control.wait_for_readable(&reader) { + Ok(true) => {} + Ok(false) => return, + Err(error) => { + let _ = sender.send(Err(error)); + return; + } + } + // `&TcpStream` implements `Read`, so the framed read runs on the + // shared handle that `Drop` can shut down. + match read_stage_message_for_codec_policy( + &mut &*reader, + activation_width, + activation_codec, + activation_codec_policy, + ) { + 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, WireMessageKind, write_stage_message}; + use std::io::Write; + use std::net::TcpListener; + use std::time::{Duration, Instant}; + + /// A live worker control: reads stay interruptible by shutdown, and these + /// tests never request one. + fn test_worker_control() -> Arc { + Arc::new(ConnectionWorkerControl::default()) + } + + fn spawn_test_message_reader( + upstream: &TcpStream, + capacity: usize, + registry: Arc, + ) -> Result { + spawn_message_reader( + upstream, + 4, + StageActivationCodec::default(), + StageActivationCodecPolicy::default(), + capacity, + registry, + test_worker_control(), + ) + } + + fn control_message(kind: WireMessageKind, tokens: Vec) -> StageWireMessage { + StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind), + 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_test_message_reader(&upstream, 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).expect("write"); + } + let discard = control_message(WireMessageKind::DiscardStaleWindows, vec![3, 9]); + write_stage_message(&mut peer, &discard).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_test_message_reader(&upstream, 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()), + ) + .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_test_message_reader(&upstream, 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).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); + } + + /// The peer-stays-open case only exercises the readable wait: with no + /// bytes sent, the reader is parked in `wait_for_readable`, which polls + /// the shutdown flag. This is the harder case — a peer sends a frame + /// prefix and then stalls, so the reader is past the readable check and + /// blocked inside `read_exact` waiting for the rest of the frame. Only a + /// shutdown of the handle the read is pending on releases it, which is + /// why the reader thread and `Drop` share one `TcpStream` (#1538). + #[test] + fn dropping_the_reader_completes_while_a_read_is_stalled_mid_frame() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_test_message_reader(&upstream, 1, registry).expect("spawn"); + + // A frame prefix with no body behind it: enough to make the socket + // readable and commit the reader to the framed read, never enough to + // complete it. + let mut frame = Vec::new(); + write_stage_message( + &mut frame, + &control_message(WireMessageKind::Stop, Vec::new()), + ) + .expect("encode"); + peer.write_all(&frame[..4]).expect("write frame prefix"); + peer.flush().expect("flush"); + // Let the reader clear the readable check and block inside the frame. + thread::sleep(Duration::from_millis(100)); + + let (done, dropped) = mpsc::channel(); + let dropper = thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("drop must interrupt a read stalled mid-frame, not block in join"); + dropper.join().expect("dropper thread"); + 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_test_message_reader(&upstream, 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..e63ff37c0f --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs @@ -0,0 +1,147 @@ +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, WireMessageKind}; + let registry = StaleDiscardRegistry::default(); + let mut message = StageWireMessage { + kind: WireMessageKind::DiscardStaleWindows, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(WireMessageKind::DiscardStaleWindows), + 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 376e8b3f2e..546ad03665 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -73,8 +73,13 @@ impl BinaryStageOptions { { bail!("--openai-prefill-adaptive-target-ms must be finite and greater than zero"); } - 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() { @@ -261,6 +266,7 @@ mod tests { min_tokens: 1, max_tokens: 6, pipeline_depth: 2, + runahead_max_tokens: 0, }, ..SpeculativeDecodeConfig::default() } diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index f61a716961..db0c2de3b5 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -716,6 +716,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 b5567b2e4e..51d9102caa 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, 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,87 @@ impl WireCondition { self.sleep_for_bandwidth(message); } + /// Serialization delay for `bytes` on this link. A near-zero `mbps` makes + /// the quotient arbitrarily large — large enough to overflow to positive + /// infinity — so the delay is clamped to the same bound as the propagation + /// delay: `Duration::from_secs_f64` panics on an out-of-domain value. + /// + /// Only `NaN` and non-positive quotients yield no delay. An infinite + /// quotient is an arbitrarily slow link, not a free one, so it takes the + /// `MAX_SIMULATED_DELAY_MS` cap like any other oversized delay; returning + /// zero there would turn the slowest configurable link into the fastest. + 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_nan() || seconds <= 0.0 { + return Duration::ZERO; + } + let millis = seconds * 1000.0; + // `min` propagates the bound for an infinite left-hand side, so the + // cap covers the overflow case without a separate branch. + Duration::from_secs_f64(millis.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, @@ -76,10 +225,161 @@ 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 an_infinite_bandwidth_quotient_takes_the_cap_not_zero() { + // The smallest positive f64 rate makes the quotient overflow to + // positive infinity. That is the slowest link the flag can express, + // so it must take the cap; returning `Duration::ZERO` there would + // silently serve it as an unmetered link. + let rate = f64::from_bits(1); + let condition = WireCondition::with_jitter(0.0, Some(rate), 0.0, 0.0, 0.0) + .expect("a positive rate is accepted"); + assert!( + (1024_f64 / (rate * 125_000.0)).is_infinite(), + "this rate must produce an infinite quotient for the test to bite" + ); + + assert_eq!( + condition.bandwidth_delay(1024), + Duration::from_secs_f64(MAX_SIMULATED_DELAY_MS / 1000.0) + ); + // A zero-length frame still costs nothing: 0/inf is not a delay. + assert_eq!(condition.bandwidth_delay(0), 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 9a23dc7dd8..c0d82846df 100644 --- a/crates/skippy-server/src/cli.rs +++ b/crates/skippy-server/src/cli.rs @@ -75,6 +75,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 bef9a4ea09..e9c36c2322 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, @@ -102,6 +110,42 @@ 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( + discard.request_id, + discard.session_id, + discard.min_window_id, + discard.max_window_id, + )?; + if let Some(forwarder) = async_forwarder { + forwarder + .send( + message, + request.downstream_wire_condition, + self.openai_attrs(request.ids), + ) + .map_err(openai_backend_error)?; + } else { + write_stage_message_conditioned( + downstream, + &message, + 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 e85e7c4289..945df43626 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1,5 +1,4 @@ use std::collections::VecDeque; - mod fused_decode; mod lifecycle; mod prefix_restore; @@ -10,8 +9,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::generation_receipt::GenerationLifecycleState; +use crate::frontend::embedded_execution::{StaleWindowDiscard, VerifyRetirement}; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window_with_threshold, propose_configured_ngram_tokens, @@ -31,38 +29,23 @@ use crate::frontend::{ }, prefill::{ PrefillChunkObservation, drain_embedded_prefill_replies, drain_one_embedded_prefill_reply, - representative_prefill_compute_sample, + prefill_chunk_end, representative_prefill_compute_sample, }, }; use crate::telemetry::now_unix_nanos; use lifecycle::{ - DirectPredictionReturnPath, EmbeddedDecodeSummary, PipelinedCompositeWindow, can_seed_pipeline, - compose_target_predictions, decode_uses_context_sideband, direct_prediction_return_path, - mark_epoch_stale, open_upstream_prediction_return, pipelined_window_layout, + DirectPredictionReturnPath, EmbeddedDecodeSummary, PipelinedCompositeWindow, + begin_generation_lifecycle, can_seed_pipeline, compose_target_predictions, + decode_uses_context_sideband, direct_prediction_return_path, finish_generation_lifecycle, + lifecycle_on_token, 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 prefix_restore::EmbeddedPrefixRestore; use serde_json::json; use skippy_protocol::binary::{StageReplyStats, WireReplyKind, recv_reply}; -fn prefill_chunk_end( - pos_start: usize, - chunk_size: usize, - prefill_token_count: usize, - exact_checkpoint_boundary: Option, -) -> usize { - let mut end = pos_start - .saturating_add(chunk_size) - .min(prefill_token_count); - if let Some(boundary) = exact_checkpoint_boundary - && pos_start < boundary - { - end = end.min(boundary); - } - end -} - impl StageOpenAiBackend { pub(super) fn generate_embedded_stage_zero_tokens( &self, @@ -84,27 +67,11 @@ impl StageOpenAiBackend { let mut lane = lane_pool.checkout(request.ids)?; let direct_prediction_return_opened = open_upstream_prediction_return(&request); let mut cache_stats = GenerationCacheStats::default(); - let mut lifecycle = GenerationLifecycleState::new( - self.generation_lifecycle.as_ref(), - request.ids.request_id, - request.ids.session_id, - request.ids.agent_session_id.clone(), - request.ids.frontend_request_id, - request.prompt_token_ids, - ); - let request_started_at = request.ids.request_started_at; + let mut lifecycle = begin_generation_lifecycle(self, &request); let mut lifecycle_cancelled = false; let result = (|| { - let mut downstream_on_token = on_token; - let mut on_token = |token_id| { - lifecycle.commit(token_id, request_started_at.elapsed()); - let control = downstream_on_token(token_id)?; - if control == TokenControl::Stop { - lifecycle.mark_callback_stop(); - } - Ok(control) - }; + let mut on_token = lifecycle_on_token(&mut lifecycle, &request, on_token); let downstream = &mut lane.stream; let prefill_token_count = request.prompt_token_ids.len().saturating_sub(1); let prefill_timer = PhaseTimer::start(); @@ -841,7 +808,15 @@ 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, + ) + }, ); let composite_sidecar_enabled = native_mtp_options.ngram_hybrid && draft_guard.is_none(); @@ -1042,6 +1017,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 { @@ -1063,7 +1044,14 @@ impl StageOpenAiBackend { 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(); @@ -1078,8 +1066,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(VerifyWindowMessageArgs { @@ -1324,6 +1315,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(), @@ -1888,6 +1895,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( @@ -1959,30 +1991,9 @@ impl StageOpenAiBackend { Ok(()) })(); - if lifecycle_cancelled { - lifecycle.mark_cancelled(); - } - lifecycle.finish(result.is_ok()); + finish_generation_lifecycle(lifecycle, lifecycle_cancelled, result.is_ok()); self.finish_embedded_generation_session(&request, lane_pool, lane, &result, &session_key)?; result?; Ok(cache_stats) } } - -#[cfg(test)] -mod tests { - use super::prefill_chunk_end; - - #[test] - fn exact_checkpoint_splits_prefill_at_the_native_state_boundary() { - assert_eq!(prefill_chunk_end(0, 1024, 1400, Some(768)), 768); - assert_eq!(prefill_chunk_end(768, 1024, 1400, Some(768)), 1400); - } - - #[test] - fn exact_checkpoint_preserves_earlier_adaptive_chunks() { - assert_eq!(prefill_chunk_end(0, 256, 1400, Some(768)), 256); - assert_eq!(prefill_chunk_end(256, 256, 1400, Some(768)), 512); - assert_eq!(prefill_chunk_end(512, 512, 1400, Some(768)), 768); - } -} diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 78aae7b9fc..b1e9fecd56 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -1,4 +1,4 @@ -use std::{borrow::Cow, collections::VecDeque}; +use std::{borrow::Cow, collections::VecDeque, time::Instant}; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; @@ -13,10 +13,52 @@ use crate::frontend::{ EmbeddedStageZeroGeneration, GenerationCacheStats, LocalGeneration, PersistentStageLane, PersistentStageLanePool, PhaseTimer, StageOpenAiBackend, TokenControl, }, + generation_receipt::GenerationLifecycleState, speculative::{OpenAiSpeculativeStats, SpeculativeDecodeConfig}, util::{openai_backend_error, openai_io_error}, }; +pub(super) fn begin_generation_lifecycle( + backend: &StageOpenAiBackend, + request: &EmbeddedStageZeroGeneration<'_>, +) -> GenerationLifecycleState { + GenerationLifecycleState::new( + backend.generation_lifecycle.as_ref(), + request.ids.request_id, + request.ids.session_id, + request.ids.agent_session_id.clone(), + request.ids.frontend_request_id, + request.prompt_token_ids, + ) +} + +pub(super) fn lifecycle_on_token<'a>( + lifecycle: &'a mut GenerationLifecycleState, + request: &EmbeddedStageZeroGeneration<'_>, + mut on_token: impl FnMut(i32) -> OpenAiResult + 'a, +) -> impl FnMut(i32) -> OpenAiResult + 'a { + let request_started_at: Instant = request.ids.request_started_at; + move |token_id| { + lifecycle.commit(token_id, request_started_at.elapsed()); + let control = on_token(token_id)?; + if control == TokenControl::Stop { + lifecycle.mark_callback_stop(); + } + Ok(control) + } +} + +pub(super) fn finish_generation_lifecycle( + mut lifecycle: GenerationLifecycleState, + cancelled: bool, + succeeded: bool, +) { + if cancelled { + lifecycle.mark_cancelled(); + } + lifecycle.finish(succeeded); +} + /// Keeps the configured speculative plan unless a distributed prefix restore /// has already populated every stage's session. Pure N-gram verification after /// that restore currently races the restored session lifecycle, so only the @@ -194,6 +236,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/native_mtp/verify_window.rs b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs index 00a49bf10b..fadf32dfa2 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(VerifyWindowMessageArgs { window_id: window.id, request_id, diff --git a/crates/skippy-server/src/frontend/prefill.rs b/crates/skippy-server/src/frontend/prefill.rs index 8cea816e6f..81b35c9011 100644 --- a/crates/skippy-server/src/frontend/prefill.rs +++ b/crates/skippy-server/src/frontend/prefill.rs @@ -11,6 +11,23 @@ use skippy_protocol::binary::WireReplyKind; use skippy_protocol::binary::recv_reply; use std::net::TcpStream; +pub(super) fn prefill_chunk_end( + pos_start: usize, + chunk_size: usize, + prefill_token_count: usize, + exact_checkpoint_boundary: Option, +) -> usize { + let mut end = pos_start + .saturating_add(chunk_size) + .min(prefill_token_count); + if let Some(boundary) = exact_checkpoint_boundary + && pos_start < boundary + { + end = end.min(boundary); + } + end +} + #[derive(Clone, Debug, PartialEq, Eq)] pub(super) struct PrefillChunkSchedule { pub(super) sizes: Vec, @@ -366,3 +383,21 @@ pub(super) fn drain_embedded_prefill_replies( } Ok(drained) } + +#[cfg(test)] +mod tests { + use super::prefill_chunk_end; + + #[test] + fn exact_checkpoint_splits_prefill_at_the_native_state_boundary() { + assert_eq!(prefill_chunk_end(0, 1024, 1400, Some(768)), 768); + assert_eq!(prefill_chunk_end(768, 1024, 1400, Some(768)), 1400); + } + + #[test] + fn exact_checkpoint_preserves_earlier_adaptive_chunks() { + assert_eq!(prefill_chunk_end(0, 256, 1400, Some(768)), 256); + assert_eq!(prefill_chunk_end(256, 256, 1400, Some(768)), 512); + assert_eq!(prefill_chunk_end(512, 512, 1400, Some(768)), 768); + } +} diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index dfe2927d85..a9c7e16a5a 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; @@ -108,6 +108,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 { @@ -129,6 +135,7 @@ impl Default for SpeculativeDecodeConfig { min_tokens: 1, max_tokens: 4, pipeline_depth: 1, + runahead_max_tokens: 0, }, draft_acceptance_threshold: 0.0, draft_split_probability: 0.0, @@ -202,6 +209,11 @@ impl SpeculativeDecodeConfig { } skippy_runtime::parse_cache_type(&self.draft_cache_type_k)?; skippy_runtime::parse_cache_type(&self.draft_cache_type_v)?; + 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(()) } diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 813e90a046..9fd9516d90 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -163,6 +163,38 @@ 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( + 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), + 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( request_id: u64, session_id: u64, diff --git a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md index 5a37b93d37..0aa76681f0 100644 --- a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md +++ b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md @@ -102,7 +102,8 @@ schema row or stale manifest row from passing review: `speculative.native_mtp_suppress_cooldown_drafts`, `speculative.native_mtp_suppress_cooldown_draft_limit`, `speculative.verify_window_min_tokens`, `speculative.verify_window_max_tokens`, -`speculative.verify_window_pipeline_depth`, `speculative.spec_default`, +`speculative.verify_window_pipeline_depth`, +`speculative.verify_window_runahead_tokens`, `speculative.spec_default`, `request_defaults.max_tokens`, `request_defaults.stop`, `request_defaults.temperature`, `request_defaults.top_p`, `request_defaults.top_k`, `request_defaults.min_p`, `request_defaults.typical_p`, diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index 7bb15a14c3..7b757c9c95 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -892,8 +892,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 coordinator without the complete - `stage-generation-9` control/status/content-identity/admission bundle must not - be selected for a generation-9 split topology. Missing `artifact-transfer` + `stage-generation-10` control/status/content-identity/admission bundle must not + be selected for a generation-10 split topology. Missing `artifact-transfer` only prevents peer cache sourcing; the worker may still participate when local/HF package resolution provides an independent source. diff --git a/docs/design/message_protocol.md b/docs/design/message_protocol.md index 73739677ef..2d831e85ce 100644 --- a/docs/design/message_protocol.md +++ b/docs/design/message_protocol.md @@ -451,7 +451,7 @@ message MeshSubprotocolOpen { } ``` -Outbound transfer uses mesh stream `0x0d` (`STREAM_SUBPROTOCOL`), followed by a length-prefixed `MeshSubprotocolOpen { name: "skippy-stage", major: 2 }`, the Skippy-owned stream kind `0x03`, a length-prefixed `StageArtifactTransferRequest`, a length-prefixed `StageArtifactTransferResponse`, and raw artifact bytes when accepted. Generation 7 requires this mesh-subprotocol path for control and artifacts; the dedicated stage ALPN accepts activation transport only. +Outbound transfer uses mesh stream `0x0d` (`STREAM_SUBPROTOCOL`), followed by a length-prefixed `MeshSubprotocolOpen { name: "skippy-stage", major: 2 }`, the Skippy-owned stream kind `0x03`, a length-prefixed `StageArtifactTransferRequest`, a length-prefixed `StageArtifactTransferResponse`, and raw artifact bytes when accepted. Generation 8 requires this mesh-subprotocol path for control and artifacts; the dedicated stage ALPN accepts activation transport only. **Request:** ```proto diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 21b7fd53ed..e7de7332aa 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -40,13 +40,14 @@ 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 7 Direct Prediction Return and Verify Retirement +## Generation 10 Stage Admission and Verify Retirement -Generation 7 introduced direct prediction return. Generation 9 is the current +Generation 7 introduced direct prediction return, and generation 9 added +canonical stage-admission descriptors. Generation 10 is the current compatibility-breaking cutover: a peer is stage compatible only when it -advertises both `skippy-stage/2` and the complete `stage-generation-9` bundle. -Every load carries a canonical admission descriptor. The ready response echoes -that descriptor exactly before topology publication. +advertises both `skippy-stage/2` and the complete `stage-generation-10` bundle. +Every load carries a canonical admission descriptor, and the ready response +echoes it exactly before topology publication. 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, @@ -82,6 +83,37 @@ 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 10 adds the `DiscardStaleWindows` control frame (wire kind 23), +which makes this generation compatibility-breaking: a pre-generation-10 +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 8330df06fc..125c7915de 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4317,27 +4317,27 @@ ], "crates/skippy-server/src/binary_transport/binary_messaging.rs": [ { - "line": 424, + "line": 425, "macro_name": "eprintln!" }, { - "line": 434, + "line": 435, "macro_name": "println!" }, { - "line": 469, + "line": 470, "macro_name": "eprintln!" }, { - "line": 490, + "line": 491, "macro_name": "eprintln!" }, { - "line": 498, + "line": 499, "macro_name": "eprintln!" }, { - "line": 567, + "line": 568, "macro_name": "eprintln!" } ], @@ -4405,7 +4405,7 @@ ], "crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs": [ { - "line": 159, + "line": 201, "macro_name": "eprintln!" } ], diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index dc65d8f669..6f67a85fc5 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -263,6 +263,7 @@ configuration should use typed per-model `topology`; explicit `--model` and | `speculative.extension_max_tokens` | integer | N-gram output budget | both | model reload | wired (requires native MTP plus an N-gram proposer) | none | | `speculative.native_mtp_reject_cooldown_tokens`
`speculative.native_mtp_suppress_cooldown_drafts`
`speculative.native_mtp_suppress_cooldown_draft_limit` | integer / boolean | runtime defaults | both | model reload | wired | none | | `speculative.verify_window_min_tokens`
`speculative.verify_window_max_tokens`
`speculative.verify_window_pipeline_depth` | integer | package policy or runtime defaults; `min <= max` | both | model reload | wired | none | +| `speculative.verify_window_runahead_tokens` | integer | `0` (fixed-depth admission); `0..=4096`, where a positive budget admits verify windows by speculative-token budget instead of a fixed window count | both | model reload | wired (capped by the native checkpoint-retention bound of 64 windows) | none | | `speculative.spec_default` | bool-or-`auto` | `auto` | both | model reload | wired (`false` disables automatic speculation; `true`, `auto`, and omission enable supported automatic defaults) | none | ## Group 8: sampling, chat templates, reasoning, and request defaults diff --git a/website/src/docs/pages/testing.md b/website/src/docs/pages/testing.md index 69ff52a20f..4da516e4dd 100644 --- a/website/src/docs/pages/testing.md +++ b/website/src/docs/pages/testing.md @@ -484,7 +484,7 @@ cached and a worker does not: - Current/current mesh: the worker may use mesh `STREAM_SUBPROTOCOL` (0x0d) to open `skippy-stage/2`, then Skippy artifact-transfer stream 0x03, to fetch only its assigned package files before the normal HF fallback path. -- Mixed-generation mesh: peers missing any required generation-7 capability +- Mixed-generation mesh: peers missing any required generation-8 capability are not stage-compatible. Control and artifact streams on the dedicated `skippy-stage/2` ALPN are rejected rather than falling back to a legacy path. - Opt-out: when artifact transfer is disabled, the node must advertise no