From d34dadd570e21968cfad08352ef8074a99942572 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 10:03:34 +1000 Subject: [PATCH 01/18] feat(skippy): add run-ahead verify-window admission - WireCondition gains an exponential jitter component plus probabilistic burst stalls so benches can model contended links (Wi-Fi, WAN) instead of a constant-latency pipe; new --downstream-wire-jitter-ms / --downstream-wire-stall-ms / --downstream-wire-stall-p flags and MESH_LLM_BENCH_DOWNSTREAM_WIRE_{JITTER_MS,STALL_MS,STALL_P} envs. - VerifyWindowScheduler gains a run-ahead mode: admission bounded by a speculative-token budget (verify_window.runahead_max_tokens) instead of a fixed window count, capped at the native checkpoint-retention bound. Config plumbed as verify_window_runahead_tokens through model config, schema, validation, and the skippy resolver. --- crates/mesh-llm-config/src/model.rs | 9 ++ .../control_behavior/speculative.rs | 3 +- .../src/model/built_in_schema/declarations.rs | 4 + .../mesh-llm-config/src/model_validation.rs | 8 +- .../src/inference/skippy/mod.rs | 45 ++++-- .../inference/skippy/resolver/speculative.rs | 13 ++ crates/skippy-protocol/src/lib.rs | 3 +- crates/skippy-protocol/src/validation.rs | 4 + .../src/binary_transport/options.rs | 10 +- .../src/binary_transport/wire.rs | 147 +++++++++++++++++- crates/skippy-server/src/cli.rs | 18 +++ .../src/frontend/decode_scheduler.rs | 144 +++++++++++++---- .../src/frontend/embedded_generation.rs | 17 +- .../src/frontend/native_mtp/verify_window.rs | 7 +- .../skippy-server/src/frontend/speculative.rs | 16 +- 15 files changed, 391 insertions(+), 57 deletions(-) diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index 128b723732..5d5f8705cb 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -562,6 +562,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, } @@ -613,6 +614,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()) @@ -686,6 +688,8 @@ struct SpeculativeConfigRaw { #[serde(default)] verify_window_pipeline_depth: Option, #[serde(default)] + verify_window_runahead_tokens: Option, + #[serde(default)] spec_default: Option, } @@ -730,6 +734,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, }) @@ -792,6 +797,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 daebdb289e..127fed5df3 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 d5b6a027bf..8b62e4f882 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, @@ -691,6 +691,12 @@ 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"), + 1, + u32::try_from(MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS).expect("runahead limit fits u32"), ) } 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 6ef94c4a57..11fadef6cf 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -89,26 +89,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)] @@ -1473,9 +1483,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 ); } @@ -1483,7 +1496,7 @@ mod tests { #[test] fn benchmark_wire_delay_rejects_invalid_values() { for value in ["-1", "NaN", "inf", "not-a-number"] { - assert!(parse_benchmark_downstream_wire_delay_ms(value).is_err()); + assert!(parse_benchmark_downstream_wire_value("test", value).is_err()); } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index 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/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index b054a59c20..ea054df872 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -29,7 +29,8 @@ pub use messages::{ StateImportMessage, StopMessage, TokenReplyMessage, }; pub use validation::{ - MAX_STAGE_FRAME_BYTES, MAX_VERIFY_WINDOW_PIPELINE_DEPTH, SCHEMA_VERSION, STAGE_ALPN_V2, + 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, diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index e28ce8dbcd..1a158a5c5d 100644 --- a/crates/skippy-protocol/src/validation.rs +++ b/crates/skippy-protocol/src/validation.rs @@ -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 { diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index f8177f7f3f..d27bccab0d 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() { @@ -260,6 +265,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/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index b5567b2e4e..94a6765389 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -1,27 +1,92 @@ -use std::{io, thread, time::Duration}; +use std::{ + io, + sync::atomic::{AtomicU64, Ordering}, + thread, + time::Duration, +}; use anyhow::{Result, bail}; use skippy_protocol::binary::{StageWireMessage, write_stage_message}; +/// Process-wide sample counter so conditioned writes draw a deterministic +/// pseudo-random sequence per process without threading RNG state through the +/// `Copy` condition value. +static WIRE_SAMPLE_COUNTER: AtomicU64 = AtomicU64::new(0); + +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, } 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"); + } + 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; + } + Duration::from_secs_f64(delay_ms / 1000.0) } fn sleep_for(&self, message: &StageWireMessage) { @@ -40,6 +105,20 @@ impl WireCondition { } } +/// Deterministic uniform sample in [0, 1) via splitmix64 over a process-wide +/// counter. Not cryptographic; just reproducible-enough conditioning for +/// benches and tests. +fn next_uniform_sample() -> f64 { + let index = WIRE_SAMPLE_COUNTER.fetch_add(1, Ordering::Relaxed); + let mut state = index.wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 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 +} + pub(crate) fn write_stage_message_conditioned( writer: impl io::Write, message: &StageWireMessage, @@ -76,10 +155,72 @@ 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_COUNTER.load(Ordering::Relaxed); + let _ = condition.propagation_delay(); + assert_eq!(WIRE_SAMPLE_COUNTER.load(Ordering::Relaxed), before); + } + + #[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..501319a1ae 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,7 +235,13 @@ 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 { @@ -264,6 +305,7 @@ impl VerifyWindowScheduler { &mut self, base_position: usize, decode_step: usize, + token_count: usize, ) -> OpenAiResult { if !self.has_capacity() { return Err(OpenAiError::backend( @@ -279,11 +321,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 +346,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 +356,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 +439,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 +461,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 +480,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 +500,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 +514,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 +539,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 +587,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 +606,45 @@ 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: one more window may still open. + assert!(scheduler.has_capacity()); + let _third = scheduler.open(106, 2, 48).unwrap(); + assert!(!scheduler.has_capacity()); + assert!(scheduler.open(154, 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, 144); + 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()); + } } diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 089dcb5dd1..6aeca46acd 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -820,7 +820,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(); @@ -1056,8 +1064,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 { 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/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index dfe2927d85..2b9081812e 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; @@ -81,6 +81,8 @@ impl NgramProposerKind { /// Longest suffix match window, and upper bound for a suffix proposer's `max_ngram`. pub const SUFFIX_NGRAM_MAX_WINDOW: usize = 64; + + /// N-gram proposer kind and its match-length and draft-length bounds. #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(deny_unknown_fields)] @@ -108,6 +110,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 +137,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 +211,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(()) } From df357f4e2f97d64019f3189dd5b24edba34dc497 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 10:14:11 +1000 Subject: [PATCH 02/18] feat(skippy): discard stale run-ahead windows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On divergence the driver now sends DiscardStaleWindows (window-id range in the token sideband) down the chain. Each stage connection gains a reader thread that parses inbound messages ahead of execution and records discard ranges in a shared registry the moment they are read, so buffered stale verify windows are answered with an empty PredictedTokens reply instead of being executed. Middle stages forward the discard and keep executing (their forwarded activations must stay valid); the final stage — which carries the sampling head — skips. Sent only in run-ahead mode, so fixed-depth setups keep today's wire behavior. --- crates/skippy-protocol/src/binary/types.rs | 11 ++ .../src/binary_transport/binary_messaging.rs | 1 + .../binary_messaging/connection.rs | 74 ++++++++- .../binary_messaging/message_receive.rs | 99 ++++++++---- .../binary_messaging/stale_discard.rs | 150 ++++++++++++++++++ .../src/binary_transport/stage_execution.rs | 1 + .../src/frontend/decode_scheduler.rs | 4 + .../src/frontend/embedded_execution.rs | 51 +++++- .../src/frontend/embedded_generation.rs | 35 +++- .../frontend/embedded_generation/lifecycle.rs | 18 +++ .../src/frontend/wire_messages.rs | 33 ++++ 11 files changed, 441 insertions(+), 36 deletions(-) create mode 100644 crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index da0a0e3f37..3b33aa3429 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -36,6 +36,7 @@ pub enum WireMessageKind { DecodeLightCtx = 9, VerifyWindow = 21, RetireVerifyWindow = 22, + DiscardStaleWindows = 23, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -79,6 +80,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) } @@ -125,6 +134,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")), } } @@ -348,6 +358,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-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index aaca99a49b..f5e19bd7f9 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -40,6 +40,7 @@ mod message_receive; mod prefill_recording; pub(in crate::binary_transport) mod reply; mod session_lifecycle; +mod stale_discard; mod session_tracker; mod summary; mod telemetry; 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 1824e7687e..993e9797d8 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,8 @@ 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::stale_discard::StaleDiscardRegistry; 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}; @@ -156,6 +157,13 @@ fn handle_binary_connection_messages( let mut request_summary = BinaryRequestSummary::default(); let mut prediction_return_streams: BTreeMap<(u64, u64), TcpStream> = BTreeMap::new(); let mut next_message = Some(first_message); + let discard_registry = Arc::new(StaleDiscardRegistry::default()); + let inbound_reader = spawn_message_reader( + upstream, + input_activation_width, + max_inflight.max(1), + discard_registry.clone(), + )?; let mut async_forwarder = if async_prefill_forward || max_inflight > 1 { downstream .as_ref() @@ -171,10 +179,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, + let Some(mut message) = inbound_reader.next( next_message.take(), pending_prefill_replies, request_summary.message_count, @@ -222,6 +227,33 @@ fn handle_binary_connection_messages( continue; } + if message.kind.is_stale_window_discard() { + // The reader thread already recorded the range; middle stages + // forward it so downstream stages can skip their buffered stale + // windows too. No reply is expected. + if let Some(downstream) = downstream.as_mut() { + if let Some(forwarder) = async_forwarder.as_mut() { + forwarder + .send( + message, + wire_dtype, + downstream_wire_condition, + BTreeMap::new(), + ) + .context("forward stale window discard downstream")?; + } else { + write_stage_message_conditioned( + &mut *downstream, + &message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward stale window discard downstream")?; + } + } + continue; + } + if message.kind.is_verify_retirement() { handle_verify_retirement( iteration_scheduler, @@ -306,6 +338,38 @@ fn handle_binary_connection_messages( bail!("binary stage state does not match message kind"); } + if message.kind == WireMessageKind::VerifyWindow + && downstream.is_none() + && discard_registry.is_discarded( + message.request_id, + message.session_id, + message.state.seq_id, + ) + { + // A discard raced ahead of this buffered stale window: answer with + // an empty prediction set instead of executing it. The driver's + // stale drain only uses the window id for FIFO bookkeeping. + let reply = StageReply { + kind: WireReplyKind::PredictedTokens, + predicted: message.state.current_token, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: reply_window_for_message(&message), + stats: StageReplyStats::default(), + }; + if let Some(return_stream) = + prediction_return_streams.get_mut(&(message.request_id, message.session_id)) + { + direct_return::send_direct_prediction_return(return_stream, reply) + .context("send discarded verify window reply")?; + } else { + send_stage_reply(&mut *upstream, reply) + .context("send discarded verify window reply")?; + } + request_summary.message_count += 1; + continue; + } + let requires_predicted = message.kind.requires_predicted_reply(); let early_prefill_ack = message.kind.is_prefill() && !requires_predicted; let mut upstream_reply_start_unix_nanos = None; diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs index 40a376a8d8..a9980e0652 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,9 +1,13 @@ -use super::ConnectionWorkerControl; use anyhow::{Context, Result}; use skippy_protocol::binary::{StageWireMessage, read_stage_message}; use std::io; use std::net::TcpStream; +use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc; +use std::thread; + +use super::stale_discard::StaleDiscardRegistry; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); @@ -11,35 +15,72 @@ 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, + +/// 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. +pub(super) struct InboundMessageReader { + receiver: mpsc::Receiver>, +} + +pub(super) fn spawn_message_reader( + upstream: &TcpStream, activation_width: i32, - first_message: Option, - pending_prefill_replies: usize, - observed_message_count: usize, -) -> Result> { - if first_message.is_some() { - return Ok(first_message); - } - 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(upstream, activation_width) { - Ok(message) => Ok(Some(message)), - Err(error) - if error.kind() == io::ErrorKind::UnexpectedEof - && pending_prefill_replies == 0 - && observed_message_count == 0 => - { - Ok(None) + capacity: usize, + registry: Arc, +) -> Result { + let mut reader = upstream + .try_clone() + .context("clone upstream stream for inbound message reader")?; + let (sender, receiver) = mpsc::sync_channel(capacity.max(1)); + thread::spawn(move || { + loop { + match read_stage_message(&mut reader, activation_width) { + Ok(message) => { + if message.kind.is_stale_window_discard() { + registry.record_message(&message); + } + if sender.send(Ok(message)).is_err() { + return; + } + } + Err(error) => { + let _ = sender.send(Err(error)); + return; + } + } + } + }); + Ok(InboundMessageReader { receiver }) +} + +impl InboundMessageReader { + /// Mirrors `receive_next_message`'s 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); + } + match self.receiver.recv() { + Ok(Ok(message)) => 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), } - Err(error) => Err(error).context("read binary stage message"), } } diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs new file mode 100644 index 0000000000..35ec9ea58f --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs @@ -0,0 +1,150 @@ +use std::collections::VecDeque; +use std::sync::Mutex; + +use skippy_protocol::binary::StageWireMessage; + +/// Bounds so a misbehaving upstream cannot grow the registry without limit. +const MAX_TRACKED_REQUESTS: usize = 32; +const MAX_RANGES_PER_REQUEST: usize = 8; + +#[derive(Debug)] +struct RequestDiscards { + request_id: u64, + session_id: u64, + ranges: VecDeque<(i32, i32)>, +} + +/// Discarded verify-window id ranges, recorded by the connection's reader +/// thread the moment a `DiscardStaleWindows` message is read and consulted by +/// the executor before running each buffered verify window. This is what lets +/// a divergence cancel the stale run-ahead tail instead of executing it. +#[derive(Debug, Default)] +pub(super) struct StaleDiscardRegistry { + requests: Mutex>, +} + +impl StaleDiscardRegistry { + /// Records the range carried by a `DiscardStaleWindows` message + /// (`tokens = [min_window_id, max_window_id]`). Malformed messages are + /// ignored: a discard is an optimization, never a correctness dependency. + pub(super) fn record_message(&self, message: &StageWireMessage) { + let (Some(&min_id), Some(&max_id)) = (message.tokens.first(), message.tokens.get(1)) else { + return; + }; + if min_id > max_id { + return; + } + self.record(message.request_id, message.session_id, min_id, max_id); + } + + pub(super) fn record(&self, request_id: u64, session_id: u64, min_id: i32, max_id: i32) { + let mut requests = self.requests.lock().expect("stale discard lock poisoned"); + if let Some(entry) = requests + .iter_mut() + .find(|entry| entry.request_id == request_id && entry.session_id == session_id) + { + if entry.ranges.len() >= MAX_RANGES_PER_REQUEST { + entry.ranges.pop_front(); + } + entry.ranges.push_back((min_id, max_id)); + return; + } + if requests.len() >= MAX_TRACKED_REQUESTS { + requests.pop_front(); + } + let mut ranges = VecDeque::new(); + ranges.push_back((min_id, max_id)); + requests.push_back(RequestDiscards { + request_id, + session_id, + ranges, + }); + } + + pub(super) fn is_discarded(&self, request_id: u64, session_id: u64, window_id: i32) -> bool { + let requests = self.requests.lock().expect("stale discard lock poisoned"); + requests + .iter() + .filter(|entry| entry.request_id == request_id && entry.session_id == session_id) + .any(|entry| { + entry + .ranges + .iter() + .any(|&(min_id, max_id)| (min_id..=max_id).contains(&window_id)) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn records_and_matches_ranges_per_request() { + let registry = StaleDiscardRegistry::default(); + registry.record(7, 9, 10, 14); + + assert!(registry.is_discarded(7, 9, 10)); + assert!(registry.is_discarded(7, 9, 14)); + assert!(!registry.is_discarded(7, 9, 15)); + assert!(!registry.is_discarded(7, 9, 9)); + // Other requests and sessions are unaffected. + assert!(!registry.is_discarded(8, 9, 12)); + assert!(!registry.is_discarded(7, 10, 12)); + } + + #[test] + fn range_count_is_bounded_per_request() { + let registry = StaleDiscardRegistry::default(); + for index in 0..(MAX_RANGES_PER_REQUEST as i32 + 4) { + registry.record(1, 1, index * 10, index * 10 + 1); + } + // The oldest ranges were evicted; the newest still match. + assert!(!registry.is_discarded(1, 1, 0)); + let newest = (MAX_RANGES_PER_REQUEST as i32 + 3) * 10; + assert!(registry.is_discarded(1, 1, newest)); + } + + #[test] + fn tracked_request_count_is_bounded() { + let registry = StaleDiscardRegistry::default(); + for request in 0..(MAX_TRACKED_REQUESTS as u64 + 4) { + registry.record(request, 1, 0, 10); + } + assert!(!registry.is_discarded(0, 1, 5)); + assert!(registry.is_discarded(MAX_TRACKED_REQUESTS as u64 + 3, 1, 5)); + } + + #[test] + fn malformed_discard_messages_are_ignored() { + use skippy_protocol::binary::{StageStateHeader, WireActivationDType, WireMessageKind}; + let registry = StaleDiscardRegistry::default(); + let mut message = StageWireMessage { + kind: WireMessageKind::DiscardStaleWindows, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new( + WireMessageKind::DiscardStaleWindows, + WireActivationDType::F32, + ), + request_id: 1, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }; + registry.record_message(&message); + assert!(!registry.is_discarded(1, 1, 0)); + + message.tokens = vec![9, 3]; + registry.record_message(&message); + assert!(!registry.is_discarded(1, 1, 5)); + + message.tokens = vec![3, 9]; + registry.record_message(&message); + assert!(registry.is_discarded(1, 1, 5)); + } +} diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index 5f6ed785d1..045cc38a11 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/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index 501319a1ae..f903b686bf 100644 --- a/crates/skippy-server/src/frontend/decode_scheduler.rs +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -248,6 +248,10 @@ impl VerifyWindowScheduler { self.config.depth() } + 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; diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index bef9a4ea09..fcb1e4dcb9 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -16,7 +16,9 @@ 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; @@ -102,6 +104,45 @@ impl StageOpenAiBackend { Ok(()) } + /// Sends a stale-window discard downstream without waiting for the + /// write receipt: the message queues behind the already-dispatched stale + /// windows, and blocking here would stall recovery for the whole stale + /// tail's wire time. + pub(super) fn discard_stale_windows( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + async_forwarder: Option<&mut AsyncForwarder>, + discard: StaleWindowDiscard, + ) -> OpenAiResult<()> { + let message = discard_stale_windows_message( + request.wire_dtype, + discard.request_id, + discard.session_id, + discard.min_window_id, + discard.max_window_id, + )?; + if let Some(forwarder) = async_forwarder { + forwarder + .send( + message, + request.wire_dtype, + request.downstream_wire_condition, + self.openai_attrs(request.ids), + ) + .map_err(openai_backend_error)?; + } else { + write_stage_message_conditioned( + downstream, + &message, + request.wire_dtype, + request.downstream_wire_condition, + ) + .map_err(openai_io_error)?; + } + Ok(()) + } + pub(super) fn execute_embedded_stage_message( &self, request: &EmbeddedStageZeroGeneration<'_>, @@ -877,3 +918,11 @@ mod tests { writer.join().unwrap(); } } + +/// 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, +} diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 6aeca46acd..f2bb31f090 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -10,7 +10,7 @@ use crate::binary_transport::{ AsyncForwarder, BinaryStageExecutionOptions, forwarded_stage_message, forwarded_stage_message_timed, run_binary_stage_message, write_stage_message_conditioned, }; -use crate::frontend::embedded_execution::VerifyRetirement; +use crate::frontend::embedded_execution::{StaleWindowDiscard, VerifyRetirement}; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window_with_threshold, propose_configured_ngram_tokens, @@ -38,6 +38,7 @@ 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, + stale_window_id_range, queued_active_tokens, refill_pipeline_ngram_candidates, speculation_after_prefix_restore, }; use openai_frontend::{OpenAiError, OpenAiResult}; @@ -1313,6 +1314,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(), @@ -1877,6 +1894,22 @@ 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, + }, + )?; + } while let Some(stale) = pipelined_windows.pop_front() { let stale_drain_timer = PhaseTimer::start(); let stale_reply = self.complete_dispatched_stage_message_direct( diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 78aae7b9fc..3f2dd143e6 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -194,6 +194,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/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 813e90a046..0e03eb814f 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -163,6 +163,39 @@ pub(super) fn embedded_verify_window_message( }) } +/// Invalidates verify windows `min_window_id..=max_window_id` for a request +/// after divergence. The downstream records the range at receive time so +/// buffered stale windows are answered with an empty reply instead of being +/// executed. The window-id range rides in `tokens`. +pub(super) fn discard_stale_windows_message( + wire_dtype: WireActivationDType, + request_id: u64, + session_id: u64, + min_window_id: i32, + max_window_id: i32, +) -> OpenAiResult { + if min_window_id > max_window_id { + return Err(OpenAiError::backend( + "stale window discard range must be non-empty", + )); + } + let kind = WireMessageKind::DiscardStaleWindows; + Ok(StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, wire_dtype), + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![min_window_id, max_window_id], + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }) +} + pub(super) fn retire_verify_window_message( request_id: u64, session_id: u64, From 0a39fce6d6adfdca5927defb307b1c22424886ce Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 16:57:09 +1000 Subject: [PATCH 03/18] fix(config): record verify_window_runahead_tokens in the defaults UI schema fixture --- .../fixtures/config_schema_defaults_ui_reference.json | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 9ccf0fb49c..e4e85168cf 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 @@ -791,6 +791,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", @@ -848,4 +855,4 @@ } } ] -} +} \ No newline at end of file From 08c267c9194b70f292ffe88f487f6f105da0734c Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 17:59:37 +1000 Subject: [PATCH 04/18] chore: rustfmt, move StaleWindowDiscard above the test module, regen console-print ratchet --- .../src/binary_transport/binary_messaging.rs | 2 +- .../binary_messaging/connection.rs | 2 +- .../binary_messaging/message_receive.rs | 1 - .../src/frontend/decode_scheduler.rs | 12 ++++++++--- .../src/frontend/embedded_execution.rs | 20 +++++++++---------- .../src/frontend/embedded_generation.rs | 2 +- .../skippy-server/src/frontend/speculative.rs | 2 -- 7 files changed, 21 insertions(+), 20 deletions(-) diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index f5e19bd7f9..e94bca1c45 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -40,8 +40,8 @@ mod message_receive; mod prefill_recording; pub(in crate::binary_transport) mod reply; mod session_lifecycle; -mod stale_discard; mod session_tracker; +mod stale_discard; mod summary; mod telemetry; 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 993e9797d8..fd85585f48 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -5,7 +5,6 @@ use super::control_messages::{ handle_verify_retirement, }; use super::message_receive::{next_connection_session_id, spawn_message_reader}; -use super::stale_discard::StaleDiscardRegistry; 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}; @@ -13,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; 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 a9980e0652..1fc396ebf7 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 @@ -15,7 +15,6 @@ pub(super) fn next_connection_session_id() -> u64 { BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed) } - /// 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 diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index f903b686bf..6266ce0d54 100644 --- a/crates/skippy-server/src/frontend/decode_scheduler.rs +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -647,8 +647,14 @@ mod tests { // 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()); + assert!( + scheduler + .open( + 10 + skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH, + 64, + 1 + ) + .is_err() + ); } } diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index fcb1e4dcb9..859a42de13 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -16,9 +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::{ - discard_stale_windows_message, 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; @@ -41,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, @@ -918,11 +924,3 @@ mod tests { writer.join().unwrap(); } } - -/// 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, -} diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index f2bb31f090..7e621fb460 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -38,8 +38,8 @@ 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, - stale_window_id_range, queued_active_tokens, refill_pipeline_ngram_candidates, speculation_after_prefix_restore, + stale_window_id_range, }; use openai_frontend::{OpenAiError, OpenAiResult}; use prefix_restore::EmbeddedPrefixRestore; diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index 2b9081812e..a9c7e16a5a 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -81,8 +81,6 @@ impl NgramProposerKind { /// Longest suffix match window, and upper bound for a suffix proposer's `max_ngram`. pub const SUFFIX_NGRAM_MAX_WINDOW: usize = 64; - - /// N-gram proposer kind and its match-length and draft-length bounds. #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(deny_unknown_fields)] From cf6409fb03e05614391fb14961d5f89c0a76358a Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sun, 23 Aug 2026 18:37:35 +1000 Subject: [PATCH 05/18] fix(skippy): gate discard frames and join inbound readers - STAGE_PROTOCOL_GENERATION 4 -> 5 with the matching stage-generation-5 feature token, so split planning excludes peers that cannot parse DiscardStaleWindows (kind 23). - verify_window_runahead_tokens validates 0..=MAX: zero is the documented fixed-depth sentinel and lets a model-level block turn inherited run-ahead off. Precedence test covers global 256 + model 0. - The inbound reader's channel now covers the whole admitted verify backlog (2 x MAX_VERIFY_WINDOW_PIPELINE_DEPTH) instead of max_inflight, so a DiscardStaleWindows behind a full backlog is read and recorded before the stale windows execute. Regression test feeds a 64-message backlog past a capacity-1 execution queue. - InboundMessageReader shuts the cloned socket down and joins its thread on drop, so a handler exiting while the peer holds the connection open no longer leaks a blocked thread and descriptor. --- .../mesh-llm-config/src/model_validation.rs | 4 +- .../src/inference/skippy/resolver/tests.rs | 43 +++++++ .../binary_messaging/message_receive.rs | 118 +++++++++++++++++- 3 files changed, 160 insertions(+), 5 deletions(-) diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index 8b62e4f882..fbf4300ad7 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -695,7 +695,9 @@ fn validate_verify_window_controls( validate_optional_u32_range( config.verify_window_runahead_tokens, &format!("{base_path}.verify_window_runahead_tokens"), - 1, + // 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-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/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs index 1fc396ebf7..7aa512d8d9 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,7 +1,7 @@ use anyhow::{Context, Result}; use skippy_protocol::binary::{StageWireMessage, read_stage_message}; use std::io; -use std::net::TcpStream; +use std::net::{Shutdown, TcpStream}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::mpsc; @@ -15,13 +15,38 @@ pub(super) fn next_connection_session_id() -> u64 { BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed) } +/// 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: the reader never blocks on a stale window while a +/// `DiscardStaleWindows` for it is still unread in the socket. +pub(super) const INBOUND_LOOKAHEAD_MESSAGES: usize = + 2 * skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; + /// Reads upstream messages on a dedicated thread so the executor can run a /// buffered message while later ones are already parsed. This is what lets a /// `DiscardStaleWindows` control message take effect before the buffered /// stale verify windows behind it get executed: the reader records the /// discard range in the shared registry the moment it reads the message. +/// +/// Dropping the reader shuts the socket down and joins the thread, so a +/// handler that exits on a local error while the peer keeps its side open +/// does not leak a blocked thread and its cloned descriptor. pub(super) struct InboundMessageReader { receiver: mpsc::Receiver>, + stream: TcpStream, + thread: Option>, +} + +impl Drop for InboundMessageReader { + fn drop(&mut self) { + // Unblock a pending `read_stage_message`; errors here only mean the + // socket is already closed. + let _ = self.stream.shutdown(Shutdown::Both); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } } pub(super) fn spawn_message_reader( @@ -33,8 +58,11 @@ pub(super) fn spawn_message_reader( let mut reader = upstream .try_clone() .context("clone upstream stream for inbound message reader")?; - let (sender, receiver) = mpsc::sync_channel(capacity.max(1)); - thread::spawn(move || { + let stream = upstream + .try_clone() + .context("clone upstream stream for inbound reader shutdown")?; + let (sender, receiver) = mpsc::sync_channel(capacity.max(INBOUND_LOOKAHEAD_MESSAGES)); + let thread = thread::spawn(move || { loop { match read_stage_message(&mut reader, activation_width) { Ok(message) => { @@ -52,7 +80,11 @@ pub(super) fn spawn_message_reader( } } }); - Ok(InboundMessageReader { receiver }) + Ok(InboundMessageReader { + receiver, + stream, + thread: Some(thread), + }) } impl InboundMessageReader { @@ -83,3 +115,81 @@ impl InboundMessageReader { } } } + +#[cfg(test)] +mod tests { + use super::*; + use skippy_protocol::binary::{ + StageStateHeader, WireActivationDType, WireMessageKind, write_stage_message, + }; + use std::net::TcpListener; + use std::time::{Duration, Instant}; + + fn control_message(kind: WireMessageKind, tokens: Vec) -> StageWireMessage { + StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, WireActivationDType::F32), + request_id: 7, + session_id: 9, + sampling: None, + chat_sampling_metadata: None, + tokens, + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("local addr"); + let client = TcpStream::connect(address).expect("connect"); + let (server, _) = listener.accept().expect("accept"); + (client, server) + } + + #[test] + fn discard_is_recorded_behind_a_backlog_larger_than_the_execution_queue() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + // Execution queue of one; the reader must still look past a full + // admitted backlog without anything being dequeued. + let reader = spawn_message_reader(&upstream, 4, 1, registry.clone()).expect("spawn"); + + for _ in 0..skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH { + let stale = control_message(WireMessageKind::Stop, Vec::new()); + write_stage_message(&mut peer, &stale, WireActivationDType::F32).expect("write"); + } + let discard = control_message(WireMessageKind::DiscardStaleWindows, vec![3, 9]); + write_stage_message(&mut peer, &discard, WireActivationDType::F32).expect("write"); + + let deadline = Instant::now() + Duration::from_secs(5); + while !registry.is_discarded(7, 9, 5) { + assert!( + Instant::now() < deadline, + "discard must be recorded while the stale backlog is still queued" + ); + thread::sleep(Duration::from_millis(5)); + } + drop(reader); + } + + #[test] + fn dropping_the_reader_joins_the_thread_while_the_peer_stays_open() { + let (peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("drop must shut the socket down and join the blocked reader thread"); + drop(peer); + } +} From 600c8e1165f4589dc3c4a92b207a5e7e942683de Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Tue, 25 Aug 2026 18:16:30 +1000 Subject: [PATCH 06/18] fix(skippy): enforce run-ahead budgets and bound shutdown - The scheduler enforces the run-ahead token budget from the second in-flight window on (admissible_window_tokens); the caller clamps its chunk width to the remaining budget and waits for a retirement instead of planning a chunk the budget cannot fit. A first window wider than the whole budget still opens so a narrow budget cannot stall a request. - InboundMessageReader::drop disconnects the channel receiver before the socket shutdown and join: a reader blocked in send on a full lookahead queue is not woken by the shutdown alone. Regression test fills the queue before dropping. - WireCondition rejects delay/jitter/stall inputs beyond one simulated hour and clamps the sampled delay, keeping Duration::from_secs_f64 in its domain. - Remaining generation-4 prose in the README and design docs now names generation 5. --- .../binary_messaging/message_receive.rs | 37 ++++++++++++++-- .../src/binary_transport/wire.rs | 18 +++++++- .../src/frontend/decode_scheduler.rs | 43 +++++++++++++++++-- .../src/frontend/embedded_generation.rs | 15 ++++++- 4 files changed, 104 insertions(+), 9 deletions(-) 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 7aa512d8d9..f8883429e0 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 @@ -33,13 +33,16 @@ pub(super) const INBOUND_LOOKAHEAD_MESSAGES: usize = /// handler that exits on a local error while the peer keeps its side open /// does not leak a blocked thread and its cloned descriptor. pub(super) struct InboundMessageReader { - receiver: mpsc::Receiver>, + receiver: Option>>, stream: TcpStream, thread: Option>, } 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. + drop(self.receiver.take()); // Unblock a pending `read_stage_message`; errors here only mean the // socket is already closed. let _ = self.stream.shutdown(Shutdown::Both); @@ -81,7 +84,7 @@ pub(super) fn spawn_message_reader( } }); Ok(InboundMessageReader { - receiver, + receiver: Some(receiver), stream, thread: Some(thread), }) @@ -99,7 +102,11 @@ impl InboundMessageReader { if first_message.is_some() { return Ok(first_message); } - match self.receiver.recv() { + let receiver = self + .receiver + .as_ref() + .expect("inbound receiver present until drop"); + match receiver.recv() { Ok(Ok(message)) => Ok(Some(message)), Ok(Err(error)) if error.kind() == io::ErrorKind::UnexpectedEof @@ -176,6 +183,30 @@ mod tests { drop(reader); } + #[test] + fn dropping_the_reader_completes_while_the_lookahead_channel_is_full() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + // Fill the lookahead queue and leave the reader blocked in `send`. + for _ in 0..(INBOUND_LOOKAHEAD_MESSAGES + 4) { + let message = control_message(WireMessageKind::Stop, Vec::new()); + write_stage_message(&mut peer, &message, WireActivationDType::F32).expect("write"); + } + thread::sleep(Duration::from_millis(100)); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("dropping the receiver must unblock the queued send before the join"); + drop(peer); + } + #[test] fn dropping_the_reader_joins_the_thread_while_the_peer_stays_open() { let (peer, upstream) = connected_pair(); diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index 94a6765389..0083b01693 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -24,6 +24,11 @@ pub struct WireCondition { 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) @@ -46,6 +51,15 @@ impl WireCondition { stall_ms: f64, stall_p: f64, ) -> Result { + 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"); + } + } if !delay_ms.is_finite() || delay_ms < 0.0 { bail!("downstream wire delay must be finite and non-negative"); } @@ -86,7 +100,9 @@ impl WireCondition { if self.stall_p > 0.0 && next_uniform_sample() < self.stall_p { delay_ms += self.stall_ms; } - Duration::from_secs_f64(delay_ms / 1000.0) + // 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) { diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index 6266ce0d54..3e8840e23d 100644 --- a/crates/skippy-server/src/frontend/decode_scheduler.rs +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -248,6 +248,20 @@ impl VerifyWindowScheduler { 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() } @@ -316,6 +330,11 @@ impl VerifyWindowScheduler { "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 @@ -624,16 +643,18 @@ mod tests { 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: one more window may still open. + // 96 tokens in flight, budget 100: only 4 more tokens fit. assert!(scheduler.has_capacity()); - let _third = scheduler.open(106, 2, 48).unwrap(); + 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(154, 3, 1).is_err()); + 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, 144); + assert_eq!(scheduler.stats().max_in_flight_tokens, 100); assert_eq!(scheduler.stats().runahead_max_tokens, 100); } @@ -657,4 +678,18 @@ mod tests { .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_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 7e621fb460..4f9a78fe72 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1029,6 +1029,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 { @@ -1050,7 +1056,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(); From 07123c58e3268d31e831394677aaa0fce46e5401 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 16:20:58 +1000 Subject: [PATCH 07/18] fix(skippy): serialize discard writes and bound read-ahead - AsyncForwarder joins its writer thread on drop, so no queued frame is still being written when the request returns its lane and a teardown Stop goes out through another clone of the same socket; the teardown discard also flushes explicitly so write errors surface there. The mid-generation discard still does not wait, since everything behind it is queued on the same forwarder and stays ordered. - The inbound lookahead queue is bounded by bytes as well as message count: 128 wide activation frames would otherwise retain many GiB. - Wire conditioning draws its jitter sequence from a per-thread counter instead of a process-global one, so parallel tests and per-lane conditioning stop depending on scheduler interleaving. - bandwidth_delay clamps the serialization delay the same way the propagation delay is clamped, so a near-zero mbps cannot panic Duration::from_secs_f64. --- .../binary_messaging/async_forwarder.rs | 70 +++++++++++++- .../binary_messaging/message_receive.rs | 32 ++++++- .../src/binary_transport/wire.rs | 92 ++++++++++++++----- .../src/frontend/embedded_generation.rs | 9 ++ 4 files changed, 176 insertions(+), 27 deletions(-) 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..f338dd3529 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,48 @@ mod tests { } } + /// A delayed discard must be fully written before a teardown `Stop` that + /// goes out through a different clone of the same socket, otherwise the + /// two frames interleave and poison a lane that is handed back for reuse. + #[test] + fn a_delayed_discard_lands_before_a_teardown_stop_on_another_clone() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let mut client = TcpStream::connect(address).unwrap(); + let (mut server, _) = listener.accept().unwrap(); + let telemetry = Telemetry::new(None, 1, prefix_cache_test_config(), TelemetryLevel::Off); + let mut forwarder = AsyncForwarder::new(&client, telemetry, 8).unwrap(); + // 250ms of simulated propagation: without the drop-time join, the + // teardown write below wins the race and the frames interleave. + let condition = WireCondition::new(250.0, None).unwrap(); + + forwarder + .send( + message(WireMessageKind::DiscardStaleWindows, 11), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap(); + drop(forwarder); + + write_stage_message_after_propagation( + &mut client, + &message(WireMessageKind::Stop, 22), + WireActivationDType::F32, + WireCondition::new(0.0, None).unwrap(), + ) + .unwrap(); + + let first = read_stage_message(&mut server, 4).unwrap(); + let second = read_stage_message(&mut server, 4).unwrap(); + + assert_eq!(first.kind, WireMessageKind::DiscardStaleWindows); + assert_eq!(first.pos_start, 11); + assert_eq!(second.kind, WireMessageKind::Stop); + assert_eq!(second.pos_start, 22); + } + #[test] fn retirement_receipt_orders_all_prior_verify_writes() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs index f8883429e0..92cfb51666 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 @@ -3,9 +3,10 @@ use skippy_protocol::binary::{StageWireMessage, read_stage_message}; use std::io; use std::net::{Shutdown, TcpStream}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc; use std::thread; +use std::time::Duration; use super::stale_discard::StaleDiscardRegistry; @@ -23,6 +24,14 @@ pub(super) fn next_connection_session_id() -> u64 { pub(super) const INBOUND_LOOKAHEAD_MESSAGES: usize = 2 * skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; +/// Byte ceiling for the same queue. The message count alone bounds nothing +/// useful for memory: a full queue of wide activation frames would retain +/// many gigabytes. When the parsed backlog reaches this many bytes the reader +/// stops reading ahead until the executor drains it; control messages are +/// small, so a discard still overtakes an activation backlog well before the +/// ceiling matters. +pub(super) const INBOUND_LOOKAHEAD_BYTES: usize = 256 * 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 @@ -36,6 +45,7 @@ pub(super) struct InboundMessageReader { receiver: Option>>, stream: TcpStream, thread: Option>, + queued_bytes: Arc, } impl Drop for InboundMessageReader { @@ -65,13 +75,22 @@ pub(super) fn spawn_message_reader( .try_clone() .context("clone upstream stream for inbound reader shutdown")?; let (sender, receiver) = mpsc::sync_channel(capacity.max(INBOUND_LOOKAHEAD_MESSAGES)); + let queued_bytes = Arc::new(AtomicUsize::new(0)); + let reader_queued_bytes = queued_bytes.clone(); let 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 { + thread::sleep(Duration::from_millis(1)); + } match read_stage_message(&mut reader, activation_width) { Ok(message) => { if message.kind.is_stale_window_discard() { registry.record_message(&message); } + let message_bytes = message.estimated_wire_bytes(); + reader_queued_bytes.fetch_add(message_bytes, Ordering::AcqRel); if sender.send(Ok(message)).is_err() { return; } @@ -87,6 +106,7 @@ pub(super) fn spawn_message_reader( receiver: Some(receiver), stream, thread: Some(thread), + queued_bytes, }) } @@ -107,7 +127,15 @@ impl InboundMessageReader { .as_ref() .expect("inbound receiver present until drop"); match receiver.recv() { - Ok(Ok(message)) => Ok(Some(message)), + Ok(Ok(message)) => { + self.queued_bytes.fetch_sub( + message + .estimated_wire_bytes() + .min(self.queued_bytes.load(Ordering::Acquire)), + Ordering::AcqRel, + ); + Ok(Some(message)) + } Ok(Err(error)) if error.kind() == io::ErrorKind::UnexpectedEof && pending_prefill_replies == 0 diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index 0083b01693..e6ad45387b 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -1,18 +1,8 @@ -use std::{ - io, - sync::atomic::{AtomicU64, Ordering}, - thread, - time::Duration, -}; +use std::{cell::Cell, io, thread, time::Duration}; use anyhow::{Result, bail}; use skippy_protocol::binary::{StageWireMessage, write_stage_message}; -/// Process-wide sample counter so conditioned writes draw a deterministic -/// pseudo-random sequence per process without threading RNG state through the -/// `Copy` condition value. -static WIRE_SAMPLE_COUNTER: AtomicU64 = AtomicU64::new(0); - const WIRE_SAMPLE_SEED: u64 = 0x9E37_79B9_7F4A_7C15; #[derive(Clone, Copy, Debug)] @@ -110,22 +100,47 @@ impl WireCondition { self.sleep_for_bandwidth(message); } + /// Serialization delay for `bytes` on this link. A near-zero `mbps` + /// makes the quotient arbitrarily large, so it is clamped to the same + /// bound as the propagation delay: `Duration::from_secs_f64` panics on an + /// out-of-domain value. + pub(crate) fn bandwidth_delay(&self, bytes: usize) -> Duration { + let Some(mbps) = self.mbps else { + return Duration::ZERO; + }; + let seconds = bytes as f64 / (mbps * 125_000.0); + if !seconds.is_finite() || seconds <= 0.0 { + return Duration::ZERO; + } + Duration::from_secs_f64((seconds * 1000.0).min(MAX_SIMULATED_DELAY_MS) / 1000.0) + } + fn sleep_for_bandwidth(&self, message: &StageWireMessage) { - let bandwidth_seconds = self - .mbps - .map(|mbps| message.estimated_wire_bytes() as f64 / (mbps * 125_000.0)) - .unwrap_or(0.0); - if bandwidth_seconds > 0.0 { - thread::sleep(Duration::from_secs_f64(bandwidth_seconds)); + let delay = self.bandwidth_delay(message.estimated_wire_bytes()); + if !delay.is_zero() { + thread::sleep(delay); } } } -/// Deterministic uniform sample in [0, 1) via splitmix64 over a process-wide +thread_local! { + /// Per-thread draw index. A process-global counter makes each thread's + /// sequence depend on how the scheduler interleaves the others, which + /// leaves parallel tests and per-lane conditioning irreproducible; each + /// thread drawing its own sequence keeps a single conditioned stream + /// deterministic. + static WIRE_SAMPLE_INDEX: Cell = const { Cell::new(0) }; +} + +/// Deterministic uniform sample in [0, 1) via splitmix64 over a per-thread /// counter. Not cryptographic; just reproducible-enough conditioning for /// benches and tests. fn next_uniform_sample() -> f64 { - let index = WIRE_SAMPLE_COUNTER.fetch_add(1, Ordering::Relaxed); + let index = WIRE_SAMPLE_INDEX.with(|counter| { + let index = counter.get(); + counter.set(index.wrapping_add(1)); + index + }); let mut state = index.wrapping_mul(0x2545_F491_4F6C_DD1D) ^ WIRE_SAMPLE_SEED; state ^= state >> 30; state = state.wrapping_mul(0xBF58_476D_1CE4_E5B9); @@ -195,9 +210,44 @@ mod tests { #[test] fn constant_condition_never_draws_samples() { let condition = WireCondition::new(3.0, None).unwrap(); - let before = WIRE_SAMPLE_COUNTER.load(Ordering::Relaxed); + let before = WIRE_SAMPLE_INDEX.with(Cell::get); let _ = condition.propagation_delay(); - assert_eq!(WIRE_SAMPLE_COUNTER.load(Ordering::Relaxed), before); + assert_eq!(WIRE_SAMPLE_INDEX.with(Cell::get), before); + } + + #[test] + fn each_thread_draws_its_own_deterministic_sequence() { + 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::>() + }; + // Two threads that each draw from index 0 see the same sequence, so a + // conditioned stream no longer depends on how other threads interleave. + let first = thread::spawn(sample_three).join().expect("first thread"); + let second = thread::spawn(sample_three).join().expect("second thread"); + + assert_eq!(first, second); + } + + #[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] diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 4f9a78fe72..634d2bfebc 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1922,6 +1922,15 @@ impl StageOpenAiBackend { 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(); From da7c38f9537996a316dadea7cb3cb2d4c94257bb Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 16:42:16 +1000 Subject: [PATCH 08/18] fix(skippy): unblock readers parked at the byte ceiling The backoff loop only observed the byte counter, so a reader waiting on an executor that is going away would spin past both the receiver drop and the socket shutdown and block Drop's join. Drop now sets a stop flag the loop checks, and the counter decrement saturates so it cannot wrap the reader into a permanent park. --- .../binary_messaging/message_receive.rs | 57 ++++++++++++++++--- 1 file changed, 50 insertions(+), 7 deletions(-) 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 92cfb51666..61d29deb25 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 @@ -3,7 +3,7 @@ use skippy_protocol::binary::{StageWireMessage, read_stage_message}; use std::io; use std::net::{Shutdown, TcpStream}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc; use std::thread; use std::time::Duration; @@ -46,12 +46,17 @@ pub(super) struct InboundMessageReader { stream: TcpStream, thread: Option>, queued_bytes: Arc, + stopped: Arc, } impl Drop for InboundMessageReader { fn drop(&mut self) { // Disconnect the channel first: a reader blocked in `send` on a full // lookahead queue is not woken by the socket shutdown. + // Release a reader parked on the byte ceiling: it is waiting on the + // executor, which is not coming back, and neither the receiver drop + // nor the socket shutdown would wake it. + self.stopped.store(true, Ordering::Release); drop(self.receiver.take()); // Unblock a pending `read_stage_message`; errors here only mean the // socket is already closed. @@ -77,11 +82,16 @@ pub(super) fn spawn_message_reader( let (sender, receiver) = mpsc::sync_channel(capacity.max(INBOUND_LOOKAHEAD_MESSAGES)); let queued_bytes = Arc::new(AtomicUsize::new(0)); let reader_queued_bytes = queued_bytes.clone(); + let stopped = Arc::new(AtomicBool::new(false)); + let reader_stopped = stopped.clone(); let thread = thread::spawn(move || { loop { // Back off while the parsed backlog is over the byte ceiling; the // executor decrements as it takes messages off the queue. while reader_queued_bytes.load(Ordering::Acquire) >= INBOUND_LOOKAHEAD_BYTES { + if reader_stopped.load(Ordering::Acquire) { + return; + } thread::sleep(Duration::from_millis(1)); } match read_stage_message(&mut reader, activation_width) { @@ -107,6 +117,7 @@ pub(super) fn spawn_message_reader( stream, thread: Some(thread), queued_bytes, + stopped, }) } @@ -128,12 +139,14 @@ impl InboundMessageReader { .expect("inbound receiver present until drop"); match receiver.recv() { Ok(Ok(message)) => { - self.queued_bytes.fetch_sub( - message - .estimated_wire_bytes() - .min(self.queued_bytes.load(Ordering::Acquire)), - Ordering::AcqRel, - ); + // 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)) @@ -211,6 +224,36 @@ mod tests { drop(reader); } + #[test] + fn dropping_the_reader_completes_while_it_is_parked_on_the_byte_ceiling() { + let (mut peer, upstream) = connected_pair(); + let registry = Arc::new(StaleDiscardRegistry::default()); + let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + + // Park the reader: pretend the executor is holding the whole byte + // budget, so the backoff loop is the only thing running. + reader + .queued_bytes + .store(INBOUND_LOOKAHEAD_BYTES, Ordering::Release); + write_stage_message( + &mut peer, + &control_message(WireMessageKind::Stop, Vec::new()), + WireActivationDType::F32, + ) + .expect("write"); + thread::sleep(Duration::from_millis(50)); + + let (done, dropped) = mpsc::channel(); + thread::spawn(move || { + drop(reader); + let _ = done.send(()); + }); + dropped + .recv_timeout(Duration::from_secs(5)) + .expect("a reader parked on the byte ceiling must still be released on drop"); + drop(peer); + } + #[test] fn dropping_the_reader_completes_while_the_lookahead_channel_is_full() { let (mut peer, upstream) = connected_pair(); From ec078f7399918bb40bfd1536ae99658cffda02fe Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 23:26:36 +1000 Subject: [PATCH 09/18] fix(skippy): isolate jitter streams and bound read-ahead MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Salt each thread's draw index with a per-thread stream ordinal. The per-thread index alone handed every lane the identical sequence, so every writer thread took its burst stall on the same message index — a synchronized-loss model rather than the contended link the flag documents. uniform_sample is now pure in (stream, index), so both properties are tested directly: reproducible within a lane, distinct across lanes. - INBOUND_LOOKAHEAD_BYTES 256 MiB -> 32 MiB. Reading ahead moves frames into userspace, so this is the per-connection bound on what a peer can make the process buffer; a ~100-byte discard overtakes a 32 MiB backlog as reliably as a larger one. - The lookahead doc comment claimed the discard always overtakes the stale windows. With the byte gate it does not for wide frames, so it now states the real behaviour and the benign fallback. - DATA_FLOW.md documents DiscardStaleWindows, including the rule that middle stages forward and still execute while the final stage skips, and the section heading names generation 5. The README states that the standalone serve-binary path has no generation handshake, so its contract is that all stages upgrade together. - with_jitter checks finiteness before the magnitude bound, and the reader's EOF doc no longer references a deleted function. --- crates/skippy-server/README.md | 7 + .../binary_messaging/message_receive.rs | 29 +++-- .../src/binary_transport/wire.rs | 121 +++++++++++++----- docs/skippy/DATA_FLOW.md | 32 +++++ 4 files changed, 149 insertions(+), 40 deletions(-) diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 9a0d0e7b5f..3e15075ff0 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -133,6 +133,13 @@ deadline handling. older peers are rejected during split planning instead of being mixed into a generation-7 topology. Generation 6 is historical and is not accepted by the current serve binary. +- That rejection happens in mesh split planning. A manually wired + `serve-binary --downstream host:port` pair performs no generation + handshake, so **the contract for the standalone path is that all stages are + upgraded together**. Pointing a run-ahead coordinator at an older stage + binary is not degraded gracefully: the older peer rejects the + `DiscardStaleWindows` frame as an unknown message kind and drops the + request connection. - `serve-binary` accepts upstream protocol connections concurrently. Model execution remains serialized by the per-process runtime lock, but readiness, abandoned, or broken connections do not monopolize the listener and block the 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 61d29deb25..d2bd2aa963 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs @@ -19,18 +19,25 @@ pub(super) fn next_connection_session_id() -> u64 { /// 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: the reader never blocks on a stale window while a -/// `DiscardStaleWindows` for it is still unread in the socket. +/// 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. The message count alone bounds nothing -/// useful for memory: a full queue of wide activation frames would retain -/// many gigabytes. When the parsed backlog reaches this many bytes the reader -/// stops reading ahead until the executor drains it; control messages are -/// small, so a discard still overtakes an activation backlog well before the -/// ceiling matters. -pub(super) const INBOUND_LOOKAHEAD_BYTES: usize = 256 * 1024 * 1024; +/// 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 @@ -122,8 +129,8 @@ pub(super) fn spawn_message_reader( } impl InboundMessageReader { - /// Mirrors `receive_next_message`'s EOF classification: a clean EOF before - /// any traffic is a normal connection close, anything else is an error. + /// 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, diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index e6ad45387b..8e1b999112 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -1,4 +1,10 @@ -use std::{cell::Cell, 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}; @@ -41,15 +47,6 @@ impl WireCondition { stall_ms: f64, stall_p: f64, ) -> Result { - 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"); - } - } if !delay_ms.is_finite() || delay_ms < 0.0 { bail!("downstream wire delay must be finite and non-negative"); } @@ -68,6 +65,17 @@ impl WireCondition { 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, @@ -123,25 +131,31 @@ impl WireCondition { } } +/// 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! { - /// Per-thread draw index. A process-global counter makes each thread's - /// sequence depend on how the scheduler interleaves the others, which - /// leaves parallel tests and per-lane conditioning irreproducible; each - /// thread drawing its own sequence keeps a single conditioned stream - /// deterministic. + /// 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 a per-thread -/// counter. Not cryptographic; just reproducible-enough conditioning for -/// benches and tests. -fn next_uniform_sample() -> f64 { - let index = WIRE_SAMPLE_INDEX.with(|counter| { - let index = counter.get(); - counter.set(index.wrapping_add(1)); - index - }); - let mut state = index.wrapping_mul(0x2545_F491_4F6C_DD1D) ^ WIRE_SAMPLE_SEED; +/// 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; @@ -150,6 +164,23 @@ fn next_uniform_sample() -> f64 { (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, @@ -216,19 +247,51 @@ mod tests { } #[test] - fn each_thread_draws_its_own_deterministic_sequence() { + 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::>() }; - // Two threads that each draw from index 0 see the same sequence, so a - // conditioned stream no longer depends on how other threads interleave. let first = thread::spawn(sample_three).join().expect("first thread"); let second = thread::spawn(sample_three).join().expect("second thread"); - assert_eq!(first, second); + assert_ne!(first, second, "per-lane writer threads must decorrelate"); } #[test] diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 2d6b536603..2cc93bf288 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -80,6 +80,38 @@ With the same four stages and 10 ms inter-stage delay, the no-spec decode hot path becomes `S0 -> S1 -> S2 -> S3 -> S0`: four hops, or about 40 ms before compute. That removes two serialized reply hops from every generated token. +## Stale Verify-Window Discard + +Generation 5 also adds the `DiscardStaleWindows` control frame (wire kind +23), which is what made the generation compatibility-breaking: a +generation-4 peer rejects the kind outright and drops the request +connection. + +Run-ahead admission dispatches verify windows before their predecessors are +verified, so a rejection strands every window queued behind the divergence. +Those windows are already on the wire. The coordinator sends one discard +naming a contiguous `[min_window_id, max_window_id]` range for the request, +and it travels the same ordered path as the windows it cancels, so it is +read after them. + +The handling rule differs by position in the chain, and this is the +non-obvious part: + +- **Middle stages forward the frame and still execute the stale windows.** + A stage that is not the final stage owns KV state its downstream neighbour + depends on; skipping its forward pass would desynchronize the chain. It + passes the discard along so the frame reaches the stage that can act on it. +- **The final/readout stage skips the named windows.** It is the only stage + whose output is thrown away by a discard, so skipping there is what + actually saves the work — the expensive readout and the direct prediction + return back to stage 0. + +The receiving stage records the range the moment it parses the frame, ahead +of executing the backlog queued in front of it, which is what lets the skip +take effect before the stale tail runs. See `INBOUND_LOOKAHEAD_BYTES` for +the case where a wide-frame backlog outruns that read-ahead and the tail +executes anyway. + ## Relative Sizes | Flow | Size | From bad6c10d556f8f68c2e75d6b789045ba72175a84 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 2 Sep 2026 10:18:17 +1000 Subject: [PATCH 10/18] fix(skippy): update discard generation and cap bandwidth delays Rebased onto main, which has moved three times under this branch. Two things it changed that this commit answers for: - Main is now at stage protocol generation 7, and generation 7 does not carry wire kind 23. So the generation boundary this branch adds has to move again: a generation-7 peer advertises current support, clears split planning, and still rejects the discard frame as an unknown message kind, which is the mixed-version teardown the review blocked on. Bumped to generation 8 (`stage-generation-8`), keeping main's wording for the rest of the generation contract and adding kind 23 to it. - Main split the connection's activation width into input and output halves (#1585). The inbound reader parses upstream frames, so it takes the input width. Review fix in the same commit: - `WireCondition::bandwidth_delay` returned `Duration::ZERO` when a positive but tiny rate made the quotient overflow to infinity, serving the slowest configurable link as an unmetered one. Only NaN and non-positive quotients yield no delay now; an infinite quotient takes the `MAX_SIMULATED_DELAY_MS` cap like any other oversized delay. Regression test covers the rate that produces the overflow. - Added the lane-reuse half of the teardown-ordering test: after a delayed discard and a teardown `Stop`, the next request's frame follows them intact on the same socket, so a lane returned to the pool cannot be corrupted by a half-written frame. `verify_window_runahead_tokens` also carries its website config-reference row, wiring-manifest entry, and reverse-audit entry, which main now requires of every config field. --- crates/mesh-llm-config/src/wiring_status.rs | 7 +++ .../src/mesh/tests/stage_control.rs | 2 +- .../src/protocol/convert.rs | 4 +- .../src/protocol/tests/announcements.rs | 10 +-- .../src/runtime/local_split/loading.rs | 2 +- crates/skippy-protocol/README.md | 4 +- crates/skippy-protocol/src/lib.rs | 25 +++++--- crates/skippy-protocol/src/validation.rs | 6 +- crates/skippy-server/README.md | 13 ++-- .../src/binary_transport/binary_messaging.rs | 2 +- .../binary_messaging/async_forwarder.rs | 62 ++++++++++++++++++- .../binary_messaging/connection.rs | 13 ++-- .../binary_messaging/message_receive.rs | 48 ++++++++++---- .../binary_messaging/stale_discard.rs | 7 +-- .../src/binary_transport/wire.rs | 42 +++++++++++-- .../src/frontend/embedded_execution.rs | 3 - .../src/frontend/wire_messages.rs | 3 +- docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md | 3 +- docs/design/TESTING.md | 4 +- docs/design/message_protocol.md | 2 +- docs/skippy/DATA_FLOW.md | 13 ++-- tools/xtask/data/console_print_allowlist.json | 12 ++-- website/src/docs/pages/config-reference.md | 1 + website/src/docs/pages/testing.md | 2 +- 24 files changed, 202 insertions(+), 88 deletions(-) diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index e994fe9dd6..573c7bf1eb 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/mesh/tests/stage_control.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/stage_control.rs index 1c88103a1b..d196cb59d4 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 @@ -120,7 +120,7 @@ async fn stage_control_bundle_gate_rejects_legacy_peer() -> Result<()> { assert!( error .to_string() - .contains("does not advertise the required generation-7 control bundle"), + .contains("does not advertise the required generation-8 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 1d8c66153a..c87fffd2c7 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_V7.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8.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_V7, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, 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 f17fb08626..ba63cf74e3 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs @@ -84,7 +84,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_V7)); + == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8)); assert_eq!( proto_pa .owner_attestation @@ -438,13 +438,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_V7, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, 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_V7, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1, ] @@ -466,7 +466,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-8 bundle" ); } } @@ -483,7 +483,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_V7 + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8 .to_string(), ], }, 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 e764c6bbfe..286546e870 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 @@ -1327,7 +1327,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 8 requires graph-observed boundary descriptors"); assert!( error .to_string() diff --git a/crates/skippy-protocol/README.md b/crates/skippy-protocol/README.md index a1df1518ca..ff55a2f18e 100644 --- a/crates/skippy-protocol/README.md +++ b/crates/skippy-protocol/README.md @@ -37,11 +37,11 @@ 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 7 requires mesh-subprotocol control, list-valued status responses, +Generation 8 requires mesh-subprotocol control, list-valued status responses, and strict local-content identity as one fail-closed capability bundle. The dedicated `skippy-stage/2` ALPN accepts activation transport only. diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index ea054df872..8a1945a86c 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -30,12 +30,11 @@ pub use messages::{ }; pub use validation::{ 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, + 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_V7, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, STAGE_SUBPROTOCOL_MAJOR, STAGE_SUBPROTOCOL_NAME, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -71,7 +70,7 @@ mod tests { stage_control_response, }; use super::{ - STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V7, + STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -80,7 +79,7 @@ mod tests { #[test] fn stage_protocol_generation_feature_names_current_generation() { assert_eq!( - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V7, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8, format!("stage-generation-{STAGE_PROTOCOL_GENERATION}") ); } @@ -366,14 +365,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: 6 }) - )); + Err(StageFrameError::BadGeneration { + got: previous_generation + }) + ); } #[test] diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index 1a158a5c5d..8a0e43904d 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 = 7; +pub const STAGE_PROTOCOL_GENERATION: u32 = 8; /// 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_V7: &str = "stage-generation-7"; +pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8: &str = "stage-generation-8"; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V7; + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V8; 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"; diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 3e15075ff0..e80c860f51 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -15,7 +15,7 @@ mesh/openai-frontend; diagnostic and benchmark clients may connect directly to the first stage. The full request/reply path is tip-to-tip: token IDs enter at the driver-facing -tip, and activations flow through the stage chain. Stage protocol generation 7 +tip, and activations flow through the stage chain. Stage protocol generation 8 is a compatibility-breaking contract: prediction-bearing replies return directly from the final/readout tip to the driver-facing stage instead of being relayed back through intermediate stages. Middle-out is the prefill optimization @@ -126,13 +126,16 @@ deadline handling. ## Notes - `serve-binary` is the tuned binary stage-to-stage path. -- `serve-binary` participates in the breaking generation-7 stage protocol. - Stage compatibility requires the complete `stage-generation-7` control, +- `serve-binary` participates in the breaking generation-8 stage protocol. + Stage compatibility requires the complete `stage-generation-8` control, status-list, and strict-content-identity bundle; direct prediction return and exact verify-checkpoint retirement are part of that generation's contract, so older peers are rejected during split planning instead of being mixed into a - generation-7 topology. Generation 6 is historical and is not accepted by the + generation-8 topology. Generation 6 is historical and is not accepted by the current serve binary. + `DiscardStaleWindows` (wire kind 23) is part of that contract too, and is + what makes this generation breaking for peers that predate it: they parse + every other current frame and reject kind 23 as an unknown message kind. - That rejection happens in mesh split planning. A manually wired `serve-binary --downstream host:port` pair performs no generation handshake, so **the contract for the standalone path is that all stages are @@ -155,7 +158,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. + generation-8 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 e94bca1c45..59e07a93f5 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -506,7 +506,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 f338dd3529..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 @@ -254,7 +254,6 @@ mod tests { forwarder .send( message(WireMessageKind::DiscardStaleWindows, 11), - WireActivationDType::F32, condition, BTreeMap::new(), ) @@ -264,7 +263,6 @@ mod tests { write_stage_message_after_propagation( &mut client, &message(WireMessageKind::Stop, 22), - WireActivationDType::F32, WireCondition::new(0.0, None).unwrap(), ) .unwrap(); @@ -278,6 +276,66 @@ mod tests { 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 fd85585f48..5fc95e3965 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -91,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 = @@ -145,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<()> { @@ -163,6 +163,7 @@ fn handle_binary_connection_messages( input_activation_width, max_inflight.max(1), discard_registry.clone(), + worker_control, )?; let mut async_forwarder = if async_prefill_forward || max_inflight > 1 { downstream @@ -234,18 +235,12 @@ fn handle_binary_connection_messages( if let Some(downstream) = downstream.as_mut() { if let Some(forwarder) = async_forwarder.as_mut() { forwarder - .send( - message, - wire_dtype, - downstream_wire_condition, - BTreeMap::new(), - ) + .send(message, downstream_wire_condition, BTreeMap::new()) .context("forward stale window discard downstream")?; } else { write_stage_message_conditioned( &mut *downstream, &message, - wire_dtype, downstream_wire_condition, ) .context("forward stale window discard downstream")?; 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 d2bd2aa963..a2615274eb 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 @@ -8,6 +8,7 @@ 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); @@ -79,6 +80,7 @@ pub(super) fn spawn_message_reader( activation_width: i32, capacity: usize, registry: Arc, + worker_control: Arc, ) -> Result { let mut reader = upstream .try_clone() @@ -96,11 +98,24 @@ pub(super) fn spawn_message_reader( // 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) { + if reader_stopped.load(Ordering::Acquire) || worker_control.is_shutting_down() { return; } thread::sleep(Duration::from_millis(1)); } + // Reads stay interruptible by shutdown. `TcpStream::shutdown` on a + // tracked clone does not unblock an in-flight `read` on Windows + // (#1538), and moving the read onto this thread must not lose + // that: peek under a short timeout and end the reader when the + // worker is shutting down. + match worker_control.wait_for_readable(&reader) { + Ok(true) => {} + Ok(false) => return, + Err(error) => { + let _ = sender.send(Err(error)); + return; + } + } match read_stage_message(&mut reader, activation_width) { Ok(message) => { if message.kind.is_stale_window_discard() { @@ -174,18 +189,22 @@ impl InboundMessageReader { #[cfg(test)] mod tests { use super::*; - use skippy_protocol::binary::{ - StageStateHeader, WireActivationDType, WireMessageKind, write_stage_message, - }; + use skippy_protocol::binary::{StageStateHeader, WireMessageKind, write_stage_message}; 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 control_message(kind: WireMessageKind, tokens: Vec) -> StageWireMessage { StageWireMessage { kind, pos_start: 0, token_count: 0, - state: StageStateHeader::new(kind, WireActivationDType::F32), + state: StageStateHeader::new(kind), request_id: 7, session_id: 9, sampling: None, @@ -211,14 +230,15 @@ mod tests { let registry = Arc::new(StaleDiscardRegistry::default()); // Execution queue of one; the reader must still look past a full // admitted backlog without anything being dequeued. - let reader = spawn_message_reader(&upstream, 4, 1, registry.clone()).expect("spawn"); + let reader = spawn_message_reader(&upstream, 4, 1, registry.clone(), test_worker_control()) + .expect("spawn"); for _ in 0..skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH { let stale = control_message(WireMessageKind::Stop, Vec::new()); - write_stage_message(&mut peer, &stale, WireActivationDType::F32).expect("write"); + write_stage_message(&mut peer, &stale).expect("write"); } let discard = control_message(WireMessageKind::DiscardStaleWindows, vec![3, 9]); - write_stage_message(&mut peer, &discard, WireActivationDType::F32).expect("write"); + write_stage_message(&mut peer, &discard).expect("write"); let deadline = Instant::now() + Duration::from_secs(5); while !registry.is_discarded(7, 9, 5) { @@ -235,7 +255,8 @@ mod tests { fn dropping_the_reader_completes_while_it_is_parked_on_the_byte_ceiling() { let (mut peer, upstream) = connected_pair(); let registry = Arc::new(StaleDiscardRegistry::default()); - let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + let reader = + spawn_message_reader(&upstream, 4, 1, registry, test_worker_control()).expect("spawn"); // Park the reader: pretend the executor is holding the whole byte // budget, so the backoff loop is the only thing running. @@ -245,7 +266,6 @@ mod tests { write_stage_message( &mut peer, &control_message(WireMessageKind::Stop, Vec::new()), - WireActivationDType::F32, ) .expect("write"); thread::sleep(Duration::from_millis(50)); @@ -265,12 +285,13 @@ mod tests { fn dropping_the_reader_completes_while_the_lookahead_channel_is_full() { let (mut peer, upstream) = connected_pair(); let registry = Arc::new(StaleDiscardRegistry::default()); - let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + let reader = + spawn_message_reader(&upstream, 4, 1, registry, test_worker_control()).expect("spawn"); // Fill the lookahead queue and leave the reader blocked in `send`. for _ in 0..(INBOUND_LOOKAHEAD_MESSAGES + 4) { let message = control_message(WireMessageKind::Stop, Vec::new()); - write_stage_message(&mut peer, &message, WireActivationDType::F32).expect("write"); + write_stage_message(&mut peer, &message).expect("write"); } thread::sleep(Duration::from_millis(100)); @@ -289,7 +310,8 @@ mod tests { fn dropping_the_reader_joins_the_thread_while_the_peer_stays_open() { let (peer, upstream) = connected_pair(); let registry = Arc::new(StaleDiscardRegistry::default()); - let reader = spawn_message_reader(&upstream, 4, 1, registry).expect("spawn"); + let reader = + spawn_message_reader(&upstream, 4, 1, registry, test_worker_control()).expect("spawn"); let (done, dropped) = mpsc::channel(); thread::spawn(move || { 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 index 35ec9ea58f..e63ff37c0f 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/stale_discard.rs @@ -117,16 +117,13 @@ mod tests { #[test] fn malformed_discard_messages_are_ignored() { - use skippy_protocol::binary::{StageStateHeader, WireActivationDType, WireMessageKind}; + 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, - WireActivationDType::F32, - ), + state: StageStateHeader::new(WireMessageKind::DiscardStaleWindows), request_id: 1, session_id: 1, sampling: None, diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index 8e1b999112..51d9102caa 100644 --- a/crates/skippy-server/src/binary_transport/wire.rs +++ b/crates/skippy-server/src/binary_transport/wire.rs @@ -108,19 +108,27 @@ impl WireCondition { self.sleep_for_bandwidth(message); } - /// Serialization delay for `bytes` on this link. A near-zero `mbps` - /// makes the quotient arbitrarily large, so it is clamped to the same - /// bound as the propagation delay: `Duration::from_secs_f64` panics on an - /// out-of-domain value. + /// 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_finite() || seconds <= 0.0 { + if seconds.is_nan() || seconds <= 0.0 { return Duration::ZERO; } - Duration::from_secs_f64((seconds * 1000.0).min(MAX_SIMULATED_DELAY_MS) / 1000.0) + 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) { @@ -313,6 +321,28 @@ mod tests { ); } + #[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(); diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index 859a42de13..e9c36c2322 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -122,7 +122,6 @@ impl StageOpenAiBackend { discard: StaleWindowDiscard, ) -> OpenAiResult<()> { let message = discard_stale_windows_message( - request.wire_dtype, discard.request_id, discard.session_id, discard.min_window_id, @@ -132,7 +131,6 @@ impl StageOpenAiBackend { forwarder .send( message, - request.wire_dtype, request.downstream_wire_condition, self.openai_attrs(request.ids), ) @@ -141,7 +139,6 @@ impl StageOpenAiBackend { write_stage_message_conditioned( downstream, &message, - request.wire_dtype, request.downstream_wire_condition, ) .map_err(openai_io_error)?; diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 0e03eb814f..9fd9516d90 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -168,7 +168,6 @@ pub(super) fn embedded_verify_window_message( /// buffered stale windows are answered with an empty reply instead of being /// executed. The window-id range rides in `tokens`. pub(super) fn discard_stale_windows_message( - wire_dtype: WireActivationDType, request_id: u64, session_id: u64, min_window_id: i32, @@ -184,7 +183,7 @@ pub(super) fn discard_stale_windows_message( kind, pos_start: 0, token_count: 0, - state: StageStateHeader::new(kind, wire_dtype), + state: StageStateHeader::new(kind), request_id, session_id, sampling: None, diff --git a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md index cd785a9c78..acd3c95f95 100644 --- a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md +++ b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md @@ -101,7 +101,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 d684736018..e7f61db5ce 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -914,8 +914,8 @@ cached and a worker does not: to open `skippy-stage/2`, then Skippy artifact-transfer stream 0x03, to fetch only its assigned package files before the normal HF fallback path. - Current/released mixed mesh: a released coordinator without the complete - `stage-generation-7` control/status/content-identity bundle and - `direct-prediction-return` support must not be selected for a generation-7 + `stage-generation-8` control/status/content-identity bundle and + `direct-prediction-return` support must not be selected for a generation-8 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 8195b4fe05..82cb1753db 100644 --- a/docs/design/message_protocol.md +++ b/docs/design/message_protocol.md @@ -448,7 +448,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 2cc93bf288..7419f53393 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -40,11 +40,11 @@ activation links and then crossed three reply links before stage 0 could emit the token. On a topology with a fixed 10 ms delay per inter-stage hop, the reply chain alone makes the hot path six hops, or about 60 ms before compute. -## Generation 7 Direct Prediction Return and Verify Retirement +## Generation 8 Direct Prediction Return and Verify Retirement -Stage protocol generation 7 is a compatibility-breaking change. A peer is stage +Stage protocol generation 8 is a compatibility-breaking change. A peer is stage compatible only when it advertises both `skippy-stage/2` and -the complete `stage-generation-7` control/status/content-identity bundle. +the complete `stage-generation-8` control/status/content-identity bundle. 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,10 +82,9 @@ compute. That removes two serialized reply hops from every generated token. ## Stale Verify-Window Discard -Generation 5 also adds the `DiscardStaleWindows` control frame (wire kind -23), which is what made the generation compatibility-breaking: a -generation-4 peer rejects the kind outright and drops the request -connection. +Generation 8 adds the `DiscardStaleWindows` control frame (wire kind 23), +which is what makes this generation compatibility-breaking: a pre-generation-8 +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. diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 93a8ddaaf1..9a5842cc27 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4251,27 +4251,27 @@ ], "crates/skippy-server/src/binary_transport/binary_messaging.rs": [ { - "line": 380, + "line": 381, "macro_name": "eprintln!" }, { - "line": 390, + "line": 391, "macro_name": "println!" }, { - "line": 425, + "line": 426, "macro_name": "eprintln!" }, { - "line": 446, + "line": 447, "macro_name": "eprintln!" }, { - "line": 454, + "line": 455, "macro_name": "eprintln!" }, { - "line": 519, + "line": 520, "macro_name": "eprintln!" } ], diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index c958acae4a..1f4f10ff9d 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -261,6 +261,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 From 48692e3194206ba96c2c75935e28b5818a5630f8 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 2 Sep 2026 15:04:38 +1000 Subject: [PATCH 11/18] fix(skippy): interrupt readers stalled mid-frame MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `wait_for_readable` returns as soon as one byte is readable and clears the socket timeout before the framed read, so a peer that sends a frame prefix and then stalls leaves the reader blocked inside `read_exact` past the readable check. `Drop` shut down a *different* clone of the socket and then joined, and shutting down a cloned handle does not interrupt a read pending on another one on Windows (#1538) — the same property the readable wait exists to preserve. Because `Drop` joins, that hangs the dropping handler rather than merely leaking a thread. The reader thread and the handle now share one `TcpStream` through an `Arc`, and `Drop` shuts down that exact handle. `&TcpStream` implements `Read`, so the framed read runs on the shared handle with no other change to the read path. Regression test `dropping_the_reader_completes_while_a_read_is_stalled_mid_frame` writes a 4-byte frame prefix and stalls, so the reader is committed to the framed read rather than parked in the readable wait, then asserts the drop completes. The existing peer-stays-open test sends no bytes and only covers the idle case. Note the Windows CI lane builds the host but does not run this suite, so the test guards the behavior rather than proving it there. --- .../binary_messaging/message_receive.rs | 84 +++++++++++++++---- 1 file changed, 69 insertions(+), 15 deletions(-) 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 a2615274eb..99342d3a3f 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 @@ -49,9 +49,17 @@ pub(super) const INBOUND_LOOKAHEAD_BYTES: usize = 32 * 1024 * 1024; /// 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: TcpStream, + stream: Arc, thread: Option>, queued_bytes: Arc, stopped: Arc, @@ -66,8 +74,10 @@ impl Drop for InboundMessageReader { // nor the socket shutdown would wake it. self.stopped.store(true, Ordering::Release); drop(self.receiver.take()); - // Unblock a pending `read_stage_message`; errors here only mean the - // socket is already closed. + // 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(); @@ -82,12 +92,12 @@ pub(super) fn spawn_message_reader( registry: Arc, worker_control: Arc, ) -> Result { - let mut reader = upstream - .try_clone() - .context("clone upstream stream for inbound message reader")?; - let stream = upstream - .try_clone() - .context("clone upstream stream for inbound reader shutdown")?; + let 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(); @@ -103,11 +113,12 @@ pub(super) fn spawn_message_reader( } thread::sleep(Duration::from_millis(1)); } - // Reads stay interruptible by shutdown. `TcpStream::shutdown` on a - // tracked clone does not unblock an in-flight `read` on Windows - // (#1538), and moving the read onto this thread must not lose - // that: peek under a short timeout and end the reader when the - // worker is shutting down. + // 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, @@ -116,7 +127,9 @@ pub(super) fn spawn_message_reader( return; } } - match read_stage_message(&mut reader, activation_width) { + // `&TcpStream` implements `Read`, so the framed read runs on the + // shared handle that `Drop` can shut down. + match read_stage_message(&mut &*reader, activation_width) { Ok(message) => { if message.kind.is_stale_window_discard() { registry.record_message(&message); @@ -190,6 +203,7 @@ impl InboundMessageReader { 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}; @@ -306,6 +320,46 @@ mod tests { 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_message_reader(&upstream, 4, 1, registry, test_worker_control()).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(); From d4038f32d6aafb58ddaeb08976bc3d9a06a23223 Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 08:39:32 +1000 Subject: [PATCH 12/18] docs(skippy): name the combined generation 10 contract --- .../src/inference/skippy/package.rs | 4 ++-- crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs | 10 +++++----- .../mesh-llm-host-runtime/src/runtime/local_package.rs | 2 +- .../mesh-llm-host-runtime/src/runtime/local_split.rs | 4 ++-- .../src/runtime/local_split/loading.rs | 2 +- .../src/runtime/stage_admission.rs | 4 ++-- crates/skippy-protocol/src/admission.rs | 4 ++-- crates/skippy-protocol/src/config.rs | 2 +- 8 files changed, 16 insertions(+), 16 deletions(-) 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/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/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 f7b85f2ed2..f30220ad84 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs @@ -542,7 +542,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( @@ -554,7 +554,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 a4232a2239..f04d3809f6 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 @@ -1319,7 +1319,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 8 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/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/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, From 7753a8c1032d19a21832b0eba2273ec7254b42da Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 12:46:45 +1000 Subject: [PATCH 13/18] fix(skippy): resolve run-ahead review follow-ups --- .../src/prompt_cli/binary_repl.rs | 61 +++++++++++++++-- crates/skippy-prompt/src/prompt_cli/mod.rs | 2 +- .../binary_messaging/connection.rs | 67 ++++++++++--------- .../src/frontend/embedded_generation.rs | 37 +--------- crates/skippy-server/src/frontend/prefill.rs | 35 ++++++++++ tools/xtask/data/console_print_allowlist.json | 38 +++++------ 6 files changed, 148 insertions(+), 92 deletions(-) diff --git a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs index 6c33c40615..f338058993 100644 --- a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs +++ b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs @@ -13,12 +13,22 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { let tokenizer_path = materialized_tokenizer .as_deref() .unwrap_or(requested_tokenizer_path); + let tokenizer_layer_start = tokenizer_layer_start(&args, materialized_tokenizer.is_some()); + let tokenizer_layer_end = tokenizer_layer_end(&args, materialized_tokenizer.is_some()); + let tokenizer_load_mode = tokenizer_load_mode(&args, materialized_tokenizer.is_some()); + let tokenizer_resident_tensor_names = resident_tensor_names_for_direct_load( + tokenizer_path, + tokenizer_load_mode, + tokenizer_layer_start, + tokenizer_layer_end, + args.ctx_size, + )?; let tokenizer = StageModel::open( tokenizer_path, &RuntimeConfig { stage_index: 0, - layer_start: tokenizer_layer_start(&args, materialized_tokenizer.is_some()), - layer_end: tokenizer_layer_end(&args, materialized_tokenizer.is_some()), + layer_start: tokenizer_layer_start, + layer_end: tokenizer_layer_end, ctx_size: args.ctx_size, lane_count: 1, n_batch: None, @@ -39,7 +49,7 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { cache_type_k: GGML_TYPE_F16, cache_type_v: GGML_TYPE_F16, flash_attn_type: skippy_runtime::FlashAttentionType::Auto, - load_mode: tokenizer_load_mode(&args, materialized_tokenizer.is_some()), + load_mode: tokenizer_load_mode, projector_path: None, projector_use_gpu: None, media_marker: None, @@ -51,7 +61,7 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { include_output: false, mtp_source: MtpSource::Disabled, filter_tensors_on_load: true, - resident_tensor_names: Vec::new(), + resident_tensor_names: tokenizer_resident_tensor_names, checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, checkpoint_imatrix: None, checkpoint_imatrix_sha256: None, @@ -80,6 +90,13 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { .as_deref() .is_some_and(|path| path != args.model_path.as_path()) { + let resident_tensor_names = resident_tensor_names_for_direct_load( + &args.model_path, + RuntimeLoadMode::RuntimeSlice, + 0, + 1, + args.ctx_size, + )?; let model = StageModel::open( &args.model_path, &RuntimeConfig { @@ -118,7 +135,7 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { include_output: false, mtp_source: MtpSource::Disabled, filter_tensors_on_load: true, - resident_tensor_names: Vec::new(), + resident_tensor_names, checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, checkpoint_imatrix: None, checkpoint_imatrix_sha256: None, @@ -364,6 +381,40 @@ fn tokenizer_load_mode(args: &BinaryReplArgs, materialized_package: bool) -> Run } } +fn resident_tensor_names_for_direct_load( + model_path: &Path, + load_mode: RuntimeLoadMode, + layer_start: u32, + layer_end: u32, + ctx_size: u32, +) -> Result> { + match load_mode { + RuntimeLoadMode::RuntimeSlice => plan_gguf_stage_resident_tensor_names( + model_path, + &[(layer_start, layer_end)], + ctx_size, + 1, + ) + .context("derive tokenizer resident tensor closure")? + .into_iter() + .next() + .context("tokenizer resident tensor plan is empty"), + RuntimeLoadMode::ArtifactSlice | RuntimeLoadMode::LayerPackage => { + let mut names = ModelInfo::open(model_path) + .with_context(|| format!("open tokenizer tensor inventory {}", model_path.display()))? + .tensors() + .context("read tokenizer tensor inventory")? + .into_iter() + .map(|tensor| tensor.name) + .collect::>(); + names.sort(); + names.dedup(); + anyhow::ensure!(!names.is_empty(), "tokenizer tensor inventory is empty"); + Ok(names) + } + } +} + fn tokenizer_load_mode_label(args: &BinaryReplArgs, materialized_package: bool) -> String { if materialized_package { "artifact-slice(materialized layer package)".to_string() diff --git a/crates/skippy-prompt/src/prompt_cli/mod.rs b/crates/skippy-prompt/src/prompt_cli/mod.rs index f4beb2d374..03e0554fd9 100644 --- a/crates/skippy-prompt/src/prompt_cli/mod.rs +++ b/crates/skippy-prompt/src/prompt_cli/mod.rs @@ -33,7 +33,7 @@ use skippy_runtime::{ ChatTemplateMessage, ChatTemplateOptions, GGML_TYPE_F16, ModelInfo, MtpSource, RuntimeConfig, RuntimeLoadMode, StageModel, StageSession, package::{PackageStageRequest, inspect_layer_package, materialize_layer_package}, - restore_native_logs, suppress_native_logs, + plan_gguf_stage_resident_tensor_names, restore_native_logs, suppress_native_logs, }; use skippy_topology::{ BoundaryDecision, NodeSpec, PlannerPolicy, TopologyPlanRequest, dense_attention_layers, 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 cf30f28ee8..92fb8dddac 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -412,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, @@ -423,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; @@ -556,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 { @@ -612,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, @@ -620,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/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 634d2bfebc..b93b2f13eb 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -30,7 +30,7 @@ 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; @@ -46,23 +46,6 @@ 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, @@ -2008,21 +1991,3 @@ impl StageOpenAiBackend { 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/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/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 2bc7c99079..8d498149e0 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -3807,79 +3807,79 @@ ], "crates/skippy-prompt/src/prompt_cli/binary_repl.rs": [ { - "line": 64, + "line": 74, "macro_name": "eprintln!" }, { - "line": 131, + "line": 148, "macro_name": "eprintln!" }, { - "line": 150, + "line": 167, "macro_name": "eprintln!" }, { - "line": 157, + "line": 174, "macro_name": "eprintln!" }, { - "line": 164, + "line": 181, "macro_name": "eprintln!" }, { - "line": 171, + "line": 188, "macro_name": "eprintln!" }, { - "line": 173, + "line": 190, "macro_name": "eprintln!" }, { - "line": 175, + "line": 192, "macro_name": "eprintln!" }, { - "line": 176, + "line": 193, "macro_name": "eprintln!" }, { - "line": 182, + "line": 199, "macro_name": "eprintln!" }, { - "line": 185, + "line": 202, "macro_name": "eprintln!" }, { - "line": 188, + "line": 205, "macro_name": "eprintln!" }, { - "line": 193, + "line": 210, "macro_name": "eprintln!" }, { - "line": 197, + "line": 214, "macro_name": "eprintln!" }, { - "line": 257, + "line": 274, "macro_name": "eprintln!" }, { - "line": 262, + "line": 279, "macro_name": "eprintln!" }, { - "line": 278, + "line": 295, "macro_name": "eprintln!" }, { - "line": 281, + "line": 298, "macro_name": "eprintln!" }, { - "line": 342, + "line": 359, "macro_name": "eprintln!" } ], From bba22db28a0223c7639d0f181ac3b3a76158d6c9 Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 13:03:44 +1000 Subject: [PATCH 14/18] fix(ci): authorize runtime-event gate model --- .../manifests/competitive-benchmark.json | 2 +- ci/model-artifacts/manifests/hf-download-smoke.json | 2 +- ci/model-artifacts/manifests/openai-smoke.json | 2 +- .../manifests/product-integration-smoke.json | 2 +- ci/model-artifacts/manifests/product-smoke.json | 2 +- ci/model-artifacts/manifests/radix-cache.json | 2 +- .../manifests/safetensors-runtime-smoke.json | 2 +- .../manifests/scripted-binary-smoke.json | 2 +- ci/model-artifacts/manifests/sdk-smoke.json | 2 +- ci/model-artifacts/manifests/skippy-ci-smoke.json | 4 +++- ci/model-artifacts/manifests/skippy-correctness.json | 2 +- ci/model-artifacts/manifests/skippy-parity.json | 2 +- ci/model-artifacts/registry.json | 2 ++ scripts/tests/test_runtime_events_native_gate.py | 11 +++++++++++ 14 files changed, 27 insertions(+), 12 deletions(-) diff --git a/ci/model-artifacts/manifests/competitive-benchmark.json b/ci/model-artifacts/manifests/competitive-benchmark.json index 3504c73a95..d92195f1e6 100644 --- a/ci/model-artifacts/manifests/competitive-benchmark.json +++ b/ci/model-artifacts/manifests/competitive-benchmark.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "competitive-benchmark", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "family-llama", diff --git a/ci/model-artifacts/manifests/hf-download-smoke.json b/ci/model-artifacts/manifests/hf-download-smoke.json index 05aee70ae5..4cd7986dc9 100644 --- a/ci/model-artifacts/manifests/hf-download-smoke.json +++ b/ci/model-artifacts/manifests/hf-download-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "hf-download-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-q4-download", diff --git a/ci/model-artifacts/manifests/openai-smoke.json b/ci/model-artifacts/manifests/openai-smoke.json index 9c3ccdf388..46b4ca6478 100644 --- a/ci/model-artifacts/manifests/openai-smoke.json +++ b/ci/model-artifacts/manifests/openai-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "openai-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-q4-download", diff --git a/ci/model-artifacts/manifests/product-integration-smoke.json b/ci/model-artifacts/manifests/product-integration-smoke.json index fe015820ed..637c170327 100644 --- a/ci/model-artifacts/manifests/product-integration-smoke.json +++ b/ci/model-artifacts/manifests/product-integration-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "product-integration-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "family-granite-hybrid", diff --git a/ci/model-artifacts/manifests/product-smoke.json b/ci/model-artifacts/manifests/product-smoke.json index 30347719d1..5893680fdb 100644 --- a/ci/model-artifacts/manifests/product-smoke.json +++ b/ci/model-artifacts/manifests/product-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "product-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-q8-inference", diff --git a/ci/model-artifacts/manifests/radix-cache.json b/ci/model-artifacts/manifests/radix-cache.json index 8cedbcdb42..033b5df0e6 100644 --- a/ci/model-artifacts/manifests/radix-cache.json +++ b/ci/model-artifacts/manifests/radix-cache.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "radix-cache", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "family-deepseek2", diff --git a/ci/model-artifacts/manifests/safetensors-runtime-smoke.json b/ci/model-artifacts/manifests/safetensors-runtime-smoke.json index f333747029..df174bd325 100644 --- a/ci/model-artifacts/manifests/safetensors-runtime-smoke.json +++ b/ci/model-artifacts/manifests/safetensors-runtime-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "safetensors-runtime-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-safetensors", diff --git a/ci/model-artifacts/manifests/scripted-binary-smoke.json b/ci/model-artifacts/manifests/scripted-binary-smoke.json index e62ddf0cdd..bcd387b5e2 100644 --- a/ci/model-artifacts/manifests/scripted-binary-smoke.json +++ b/ci/model-artifacts/manifests/scripted-binary-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "scripted-binary-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-q8-inference", diff --git a/ci/model-artifacts/manifests/sdk-smoke.json b/ci/model-artifacts/manifests/sdk-smoke.json index c55578838d..f226facdb5 100644 --- a/ci/model-artifacts/manifests/sdk-smoke.json +++ b/ci/model-artifacts/manifests/sdk-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "sdk-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "smollm2-q8-inference", diff --git a/ci/model-artifacts/manifests/skippy-ci-smoke.json b/ci/model-artifacts/manifests/skippy-ci-smoke.json index 560e2b7989..1e02cd08e4 100644 --- a/ci/model-artifacts/manifests/skippy-ci-smoke.json +++ b/ci/model-artifacts/manifests/skippy-ci-smoke.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "skippy-ci-smoke", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "family-qwen3-dense", @@ -20,6 +20,8 @@ "llama-bump", "manual-full", "nightly", + "pull-request", + "main", "manual" ], "repo": "Qwen/Qwen3-0.6B-GGUF", diff --git a/ci/model-artifacts/manifests/skippy-correctness.json b/ci/model-artifacts/manifests/skippy-correctness.json index 267a98421c..cd04059f3b 100644 --- a/ci/model-artifacts/manifests/skippy-correctness.json +++ b/ci/model-artifacts/manifests/skippy-correctness.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "skippy-correctness", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "qwen3-q8-correctness", diff --git a/ci/model-artifacts/manifests/skippy-parity.json b/ci/model-artifacts/manifests/skippy-parity.json index 0e76600138..09b3bdb606 100644 --- a/ci/model-artifacts/manifests/skippy-parity.json +++ b/ci/model-artifacts/manifests/skippy-parity.json @@ -2,7 +2,7 @@ "schema_version": 1, "manifest_kind": "test-model-artifacts", "suite": "skippy-parity", - "registry_sha256": "0fb272556fa31f2ab91d06c8265af82aa531c2001dbcbe6c3af3dab941140b3c", + "registry_sha256": "dbd9c20f1d7b0176f31eda393eb9aa7cb8af989399f83758650a450a0be1d268", "artifacts": [ { "id": "family-deepseek2", diff --git a/ci/model-artifacts/registry.json b/ci/model-artifacts/registry.json index df92a39206..848b9ac3f4 100644 --- a/ci/model-artifacts/registry.json +++ b/ci/model-artifacts/registry.json @@ -88,6 +88,8 @@ "llama-bump", "manual-full", "nightly", + "pull-request", + "main", "manual" ], "capability_tags": [ diff --git a/scripts/tests/test_runtime_events_native_gate.py b/scripts/tests/test_runtime_events_native_gate.py index 7ebffa597e..7126841266 100644 --- a/scripts/tests/test_runtime_events_native_gate.py +++ b/scripts/tests/test_runtime_events_native_gate.py @@ -275,6 +275,17 @@ def test_the_gate_selects_one_artifact_from_the_shared_manifest(self) -> None: ) self.assertEqual(step["with"]["model_artifact_id"], "family-qwen3-dense") + def test_the_gate_model_is_authorized_for_pr_and_main_ci(self) -> None: + manifest = yaml.safe_load( + (ROOT / "ci/model-artifacts/manifests/skippy-ci-smoke.json").read_text( + encoding="utf-8" + ) + ) + artifact = next( + row for row in manifest["artifacts"] if row["id"] == "family-qwen3-dense" + ) + self.assertTrue({"pull-request", "main"}.issubset(artifact["cadences"])) + def test_evidence_is_uploaded_even_when_the_gate_fails(self) -> None: """The evidence file is how a failure is diagnosed, so it must survive one.""" From a4d8b414ff69345f0f005d833a7b33a9ed39a21c Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 13:20:06 +1000 Subject: [PATCH 15/18] fix(ci): preserve native gate evidence path --- scripts/ci-runtime-events-native-gate.sh | 5 ++++ .../tests/test_runtime_events_native_gate.py | 26 ++++++++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/scripts/ci-runtime-events-native-gate.sh b/scripts/ci-runtime-events-native-gate.sh index 8ab3eeb31f..3f104069c6 100755 --- a/scripts/ci-runtime-events-native-gate.sh +++ b/scripts/ci-runtime-events-native-gate.sh @@ -74,6 +74,11 @@ if [[ ! -s "$MODEL_PATH" ]]; then fi mkdir -p "$(dirname "$EVIDENCE_FILE")" +# Cargo may run integration-test binaries from a package directory. Resolve +# the caller's evidence path before invoking it so both the test and this +# wrapper always refer to the same file. +EVIDENCE_DIR="$(cd "$(dirname "$EVIDENCE_FILE")" && pwd)" +EVIDENCE_FILE="$EVIDENCE_DIR/$(basename "$EVIDENCE_FILE")" # Start from an empty file so the assertion below reads THIS run's markers, # never a previous run's left behind by a warm workspace. : >"$EVIDENCE_FILE" diff --git a/scripts/tests/test_runtime_events_native_gate.py b/scripts/tests/test_runtime_events_native_gate.py index 7126841266..d8f007f8ae 100644 --- a/scripts/tests/test_runtime_events_native_gate.py +++ b/scripts/tests/test_runtime_events_native_gate.py @@ -82,6 +82,7 @@ def run_gate( bundle: str | None = None, model: str | None = None, evidence_seed: str | None = None, + relative_evidence: bool = False, ) -> subprocess.CompletedProcess[str]: stub_bin = root / "stub-bin" stub_bin.mkdir(exist_ok=True) @@ -101,6 +102,7 @@ def run_gate( evidence = root / "evidence.txt" if evidence_seed is not None: evidence.write_text(evidence_seed, encoding="utf-8") + evidence_argument = evidence.name if relative_evidence else str(evidence) return subprocess.run( [ @@ -111,7 +113,7 @@ def run_gate( "--model", model, "--evidence", - str(evidence), + evidence_argument, ], capture_output=True, text=True, @@ -120,6 +122,7 @@ def run_gate( **os.environ, "PATH": f"{stub_bin}{os.pathsep}{os.environ['PATH']}", }, + cwd=root, ) EXECUTES = ( @@ -140,6 +143,27 @@ def test_a_gate_that_executed_passes(self) -> None: self.assertEqual(result.returncode, 0, result.stdout + result.stderr) self.assertIn("executed", result.stdout) + def test_relative_evidence_path_survives_cargo_working_directory_change(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + captured = root / "evidence-path.txt" + result = self.run_gate( + root, + cargo_body=( + "#!/usr/bin/env bash\n" + f'printf \'%s\\n\' "$MESH_LLM_RUNTIME_EVENTS_EVIDENCE_FILE" ' + f"> {captured}\n" + "cd /\n" + 'printf \'executed\\n\' ' + '>> "$MESH_LLM_RUNTIME_EVENTS_EVIDENCE_FILE"\n' + ), + relative_evidence=True, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + evidence_path = Path(captured.read_text(encoding="utf-8").strip()) + self.assertTrue(evidence_path.is_absolute()) + self.assertEqual(evidence_path.resolve(), (root / "evidence.txt").resolve()) + def test_a_blocked_gate_fails_even_though_the_test_exits_zero(self) -> None: """The whole reason the script checks the marker. From 0771b041ea4c0f8d39f8d48f4e3f23c1851ff2b8 Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 13:50:35 +1000 Subject: [PATCH 16/18] refactor(skippy): bound embedded generation module --- .../src/frontend/embedded_generation.rs | 34 +++----------- .../frontend/embedded_generation/lifecycle.rs | 44 ++++++++++++++++++- tools/xtask/data/console_print_allowlist.json | 2 +- 3 files changed, 51 insertions(+), 29 deletions(-) diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 9fc1932a1c..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; @@ -11,7 +10,6 @@ use crate::binary_transport::{ forwarded_stage_message_timed, run_binary_stage_message, write_stage_message_conditioned, }; use crate::frontend::embedded_execution::{StaleWindowDiscard, VerifyRetirement}; -use crate::frontend::generation_receipt::GenerationLifecycleState; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window_with_threshold, propose_configured_ngram_tokens, @@ -36,9 +34,10 @@ use crate::frontend::{ }; 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, }; @@ -68,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(); @@ -2008,10 +1991,7 @@ 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) diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 3f2dd143e6..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 diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index dc9f36a30a..bf81d4e646 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4405,7 +4405,7 @@ ], "crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs": [ { - "line": 159, + "line": 201, "macro_name": "eprintln!" } ], From 1e2f5b9553ef56839d8adb8e0d7d9a30b7296142 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 15:59:49 +1000 Subject: [PATCH 17/18] fix(ci): reconcile model registry after main sync --- ci/model-artifacts/registry.json | 2 -- 1 file changed, 2 deletions(-) diff --git a/ci/model-artifacts/registry.json b/ci/model-artifacts/registry.json index 7ef2792034..2b499466c2 100644 --- a/ci/model-artifacts/registry.json +++ b/ci/model-artifacts/registry.json @@ -84,8 +84,6 @@ "llama-bump", "manual-full", "nightly", - "pull-request", - "main", "manual" ], "capability_tags": [ From a2acf7f145d56f8ec932c47df2fca0b5d2d02c58 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 16:08:52 +1000 Subject: [PATCH 18/18] chore(ci): refresh console print ratchet after main sync --- tools/xtask/data/console_print_allowlist.json | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) 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!" } ],