From afc2eb93eab450c8898d41e238b17f42f74e9577 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 10:03:34 +1000 Subject: [PATCH 01/15] Add run-ahead verify-window admission and wire jitter conditioning - 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. Co-Authored-By: Claude Opus 5 --- 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 bfe60b3d99..a9fb52bac6 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -603,6 +603,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, } @@ -654,6 +655,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()) @@ -727,6 +729,8 @@ struct SpeculativeConfigRaw { #[serde(default)] verify_window_pipeline_depth: Option, #[serde(default)] + verify_window_runahead_tokens: Option, + #[serde(default)] spec_default: Option, } @@ -771,6 +775,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, }) @@ -833,6 +838,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 2095ae7f0e..39003a11e7 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -659,6 +659,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 0505452fb0..de950ff04d 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -9,7 +9,7 @@ use crate::model::{ SkippyConfig, SpeculativeConfig, StringOrStringList, merge_hardware, merge_model_fit, merge_multimodal, merge_throughput, }; -use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; +use skippy_protocol::{MAX_VERIFY_WINDOW_PIPELINE_DEPTH, MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS}; use crate::validation_support::{ looks_like_model_identifier, validate_allowed, validate_bool_or_auto, validate_hf_pair, @@ -633,6 +633,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 140145e1df..9200f48e4a 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -79,26 +79,36 @@ pub(crate) use stage::{ pub(crate) use topology::{StageTopologyParticipant, plan_package_identity_topology}; const BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_DELAY_MS"; +const BENCH_DOWNSTREAM_WIRE_JITTER_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_JITTER_MS"; +const BENCH_DOWNSTREAM_WIRE_STALL_MS_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_STALL_MS"; +const BENCH_DOWNSTREAM_WIRE_STALL_P_ENV: &str = "MESH_LLM_BENCH_DOWNSTREAM_WIRE_STALL_P"; fn benchmark_downstream_wire_condition() -> Result { - let delay_ms = match env::var(BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV) { - Ok(value) => parse_benchmark_downstream_wire_delay_ms(&value)?, - Err(env::VarError::NotPresent) => 0.0, + let delay_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV)?; + let jitter_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_JITTER_MS_ENV)?; + let stall_ms = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_STALL_MS_ENV)?; + let stall_p = parse_benchmark_wire_env(BENCH_DOWNSTREAM_WIRE_STALL_P_ENV)?; + WireCondition::with_jitter(delay_ms, None, jitter_ms, stall_ms, stall_p) +} + +fn parse_benchmark_wire_env(name: &'static str) -> Result { + match env::var(name) { + Ok(value) => parse_benchmark_downstream_wire_value(name, &value), + Err(env::VarError::NotPresent) => Ok(0.0), Err(env::VarError::NotUnicode(_)) => { - anyhow::bail!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be valid UTF-8") + anyhow::bail!("{name} must be valid UTF-8") } - }; - WireCondition::new(delay_ms, None) + } } -fn parse_benchmark_downstream_wire_delay_ms(value: &str) -> Result { - let delay_ms = value.parse::().with_context(|| { - format!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be a finite non-negative number") - })?; - if !delay_ms.is_finite() || delay_ms < 0.0 { - anyhow::bail!("{BENCH_DOWNSTREAM_WIRE_DELAY_MS_ENV} must be a finite non-negative number"); +fn parse_benchmark_downstream_wire_value(name: &str, value: &str) -> Result { + let parsed = value + .parse::() + .with_context(|| format!("{name} must be a finite non-negative number"))?; + if !parsed.is_finite() || parsed < 0.0 { + anyhow::bail!("{name} must be a finite non-negative number"); } - Ok(delay_ms) + Ok(parsed) } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -1345,9 +1355,12 @@ mod tests { #[test] fn benchmark_wire_delay_accepts_finite_non_negative_values() { - assert_eq!(parse_benchmark_downstream_wire_delay_ms("0").unwrap(), 0.0); assert_eq!( - parse_benchmark_downstream_wire_delay_ms("25.5").unwrap(), + parse_benchmark_downstream_wire_value("test", "0").unwrap(), + 0.0 + ); + assert_eq!( + parse_benchmark_downstream_wire_value("test", "25.5").unwrap(), 25.5 ); } @@ -1355,7 +1368,7 @@ mod tests { #[test] fn benchmark_wire_delay_rejects_invalid_values() { for value in ["-1", "NaN", "inf", "not-a-number"] { - assert!(parse_benchmark_downstream_wire_delay_ms(value).is_err()); + assert!(parse_benchmark_downstream_wire_value("test", value).is_err()); } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index 6a17c09ddb..600f9d9e30 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 @@ -397,6 +397,17 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result config.verify_window.max_tokens { bail!("skippy speculative verify window requires min_tokens <= max_tokens"); } @@ -479,6 +490,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", @@ -564,6 +576,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 d2a7d72b56..89a16d5848 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_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index 838a522c2a..defc72bea4 100644 --- a/crates/skippy-protocol/src/validation.rs +++ b/crates/skippy-protocol/src/validation.rs @@ -21,6 +21,10 @@ pub const STAGE_STREAM_ARTIFACT_TRANSFER: u8 = 0x03; pub const MAX_STAGE_FRAME_BYTES: usize = 8 * 1024 * 1024; /// Maximum number of unresolved verify windows covered by native checkpoints. pub const MAX_VERIFY_WINDOW_PIPELINE_DEPTH: usize = 64; +/// Sanity bound on the run-ahead speculative-token budget. Every in-flight +/// speculative token holds restorable recovery state downstream, so the budget +/// caps how much KV checkpoint memory one request can pin. +pub const MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS: usize = 4096; #[derive(Debug, Clone, PartialEq, Eq)] pub enum StageFrameError { diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index 5512d1ef3c..e0c4b972d2 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -64,8 +64,13 @@ impl BinaryStageOptions { bail!("--openai-prefill-chunk-size must be greater than zero"); } let wire_dtype = parse_wire_dtype(&args.activation_wire_dtype)?; - let downstream_wire_condition = - WireCondition::new(args.downstream_wire_delay_ms, args.downstream_wire_mbps)?; + let downstream_wire_condition = WireCondition::with_jitter( + args.downstream_wire_delay_ms, + args.downstream_wire_mbps, + args.downstream_wire_jitter_ms, + args.downstream_wire_stall_ms, + args.downstream_wire_stall_p, + )?; let config = load_json::(&args.config) .with_context(|| format!("load stage config {}", args.config.display()))?; let topology = match args.topology.as_ref() { @@ -231,6 +236,7 @@ mod tests { min_tokens: 1, max_tokens: 6, pipeline_depth: 2, + runahead_max_tokens: 0, }, } } diff --git a/crates/skippy-server/src/binary_transport/wire.rs b/crates/skippy-server/src/binary_transport/wire.rs index 19e8c73914..1b81db3b47 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, WireActivationDType, 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, @@ -78,10 +157,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 030f4ab68c..39bb7c1234 100644 --- a/crates/skippy-server/src/cli.rs +++ b/crates/skippy-server/src/cli.rs @@ -78,6 +78,24 @@ pub struct ServeBinaryArgs { help = "Artificial downstream activation bandwidth cap in megabits per second." )] pub downstream_wire_mbps: Option, + #[arg( + long, + default_value_t = 0.0, + help = "Mean of an exponentially distributed extra per-message downstream delay in milliseconds (models link jitter)." + )] + pub downstream_wire_jitter_ms: f64, + #[arg( + long, + default_value_t = 0.0, + help = "Extra burst-stall delay in milliseconds applied with --downstream-wire-stall-p probability per message." + )] + pub downstream_wire_stall_ms: f64, + #[arg( + long, + default_value_t = 0.0, + help = "Probability in [0, 1] that a downstream message is hit by --downstream-wire-stall-ms." + )] + pub downstream_wire_stall_p: f64, #[arg(long, default_value_t = 60)] pub downstream_connect_timeout_secs: u64, #[arg( diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index ca04c54412..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 b6adc416ef..fc9b1c3081 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -794,7 +794,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(); @@ -1024,8 +1032,11 @@ impl StageOpenAiBackend { current, &proposal_tokens, ); - let window = verify_window_scheduler - .open(layout.pos_start, layout.decode_step)?; + let window = verify_window_scheduler.open( + layout.pos_start, + layout.decode_step, + layout.input_tokens.len(), + )?; let input_tokens = layout.input_tokens; let message = embedded_verify_window_message( request.wire_dtype, diff --git a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs index 7049364e80..22dcb61ddc 100644 --- a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs +++ b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs @@ -125,8 +125,11 @@ impl StageOpenAiBackend { return Ok(NativeMtpVerifyWindowControl::NoProposal); } let verify_inputs = native_mtp_verify_window_inputs(*current, &proposal_tokens); - let window = - verify_window_scheduler.open(prefill_token_count + *decoded_tokens, *decoded_tokens)?; + let window = verify_window_scheduler.open( + prefill_token_count + *decoded_tokens, + *decoded_tokens, + verify_inputs.len(), + )?; let message = embedded_verify_window_message( request.wire_dtype, VerifyWindowMessageArgs { diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index 6abb8e61e1..497a948ec4 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; @@ -65,6 +65,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)] @@ -92,6 +94,12 @@ pub struct VerifyWindowConfig { pub min_tokens: usize, pub max_tokens: usize, pub pipeline_depth: usize, + /// Run-ahead speculative-token budget. Zero keeps the fixed + /// `pipeline_depth` window-count admission; a positive value switches the + /// scheduler to token-budget admission (windows stay capped at the native + /// checkpoint-retention bound). + #[serde(default)] + pub runahead_max_tokens: usize, } impl Default for SpeculativeDecodeConfig { @@ -113,6 +121,7 @@ impl Default for SpeculativeDecodeConfig { min_tokens: 1, max_tokens: 4, pipeline_depth: 1, + runahead_max_tokens: 0, }, } } @@ -173,6 +182,11 @@ impl SpeculativeDecodeConfig { "verify window requires 0 < min_tokens <= max_tokens and 0 < pipeline_depth <= {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" ); } + 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 2fc65b3f75bee72d5fa28151b22d9e5cf2fa8dd5 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 10:14:11 +1000 Subject: [PATCH 02/15] Cancel the stale run-ahead tail with a wire-level discard 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. Co-Authored-By: Claude Opus 5 --- crates/skippy-protocol/src/binary/types.rs | 11 ++ .../src/binary_transport/binary_messaging.rs | 1 + .../binary_messaging/connection.rs | 73 ++++++++- .../binary_messaging/message_receive.rs | 88 +++++++--- .../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(+), 24 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 4541f7f9f0..376530fbaa 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -58,6 +58,7 @@ pub enum WireMessageKind { DecodeLightCtx = 9, VerifyWindow = 21, RetireVerifyWindow = 22, + DiscardStaleWindows = 23, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -101,6 +102,14 @@ impl WireMessageKind { matches!(self, Self::RetireVerifyWindow) } + /// Control message invalidating a contiguous range of not-yet-executed + /// verify windows after the driver detected divergence. Recorded at + /// message-receive time so buffered stale windows are skipped instead of + /// executed. + pub fn is_stale_window_discard(self) -> bool { + matches!(self, Self::DiscardStaleWindows) + } + pub fn is_generation_control(self) -> bool { matches!(self, Self::ConfigureGeneration) } @@ -147,6 +156,7 @@ impl TryFrom for WireMessageKind { 20 => Ok(Self::PredictionReturnOpen), 21 => Ok(Self::VerifyWindow), 22 => Ok(Self::RetireVerifyWindow), + 23 => Ok(Self::DiscardStaleWindows), _ => Err(invalid_data("unknown stage message kind")), } } @@ -329,6 +339,7 @@ impl StageStateHeader { WireMessageKind::StateImport | WireMessageKind::StateExport ) || kind.is_session_control() || kind.is_verify_retirement() + || kind.is_stale_window_discard() || kind.is_generation_control() { return true; diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index d344d6e4d1..bf810d3a4f 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -38,6 +38,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 8340611883..88ffd706b1 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -3,7 +3,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_message; @@ -149,6 +150,13 @@ fn handle_binary_connection_messages( let mut request_summary = BinaryRequestSummary::default(); let mut prediction_return_streams: BTreeMap<(u64, u64), TcpStream> = BTreeMap::new(); let mut next_message = Some(first_message); + let discard_registry = Arc::new(StaleDiscardRegistry::default()); + let inbound_reader = spawn_message_reader( + upstream, + activation_width, + max_inflight.max(1), + discard_registry.clone(), + )?; let mut async_forwarder = if async_prefill_forward || max_inflight > 1 { downstream .as_ref() @@ -164,9 +172,7 @@ fn handle_binary_connection_messages( loop { let recv_start_unix_nanos = now_unix_nanos() as u64; let recv_started = Instant::now(); - let Some(mut message) = receive_next_message( - upstream, - activation_width, + let Some(mut message) = inbound_reader.next( next_message.take(), pending_prefill_replies, request_summary.message_count, @@ -215,6 +221,33 @@ fn handle_binary_connection_messages( continue; } + if message.kind.is_stale_window_discard() { + // The reader thread already recorded the range; middle stages + // forward it so downstream stages can skip their buffered stale + // windows too. No reply is expected. + if let Some(downstream) = downstream.as_mut() { + if let Some(forwarder) = async_forwarder.as_mut() { + forwarder + .send( + message, + wire_dtype, + downstream_wire_condition, + BTreeMap::new(), + ) + .context("forward stale window discard downstream")?; + } else { + write_stage_message_conditioned( + &mut *downstream, + &message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward stale window discard downstream")?; + } + } + continue; + } + if message.kind.is_verify_retirement() { handle_verify_retirement( iteration_scheduler, @@ -303,6 +336,38 @@ fn handle_binary_connection_messages( bail!("binary stage state does not match message kind"); } + if message.kind == WireMessageKind::VerifyWindow + && downstream.is_none() + && discard_registry.is_discarded( + message.request_id, + message.session_id, + message.state.seq_id, + ) + { + // A discard raced ahead of this buffered stale window: answer with + // an empty prediction set instead of executing it. The driver's + // stale drain only uses the window id for FIFO bookkeeping. + let reply = StageReply { + kind: WireReplyKind::PredictedTokens, + predicted: message.state.current_token, + predicted_tokens: Vec::new(), + native_mtp_draft: None, + window: reply_window_for_message(&message), + stats: StageReplyStats::default(), + }; + if let Some(return_stream) = + prediction_return_streams.get_mut(&(message.request_id, message.session_id)) + { + direct_return::send_direct_prediction_return(return_stream, reply) + .context("send discarded verify window reply")?; + } else { + send_stage_reply(&mut *upstream, reply) + .context("send discarded verify window reply")?; + } + request_summary.message_count += 1; + continue; + } + let requires_predicted = message.kind.requires_predicted_reply(); let early_prefill_ack = message.kind.is_prefill() && !requires_predicted; let mut upstream_reply_start_unix_nanos = None; diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs index 8450c2c9a1..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 @@ -2,7 +2,12 @@ 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); @@ -10,25 +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, + +/// 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); - } - 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 0654f72355..a9eec4ed4b 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -721,6 +721,7 @@ pub(crate) fn run_binary_stage_message( | WireMessageKind::ConfigureGeneration | WireMessageKind::TrimSession | WireMessageKind::RetireVerifyWindow + | WireMessageKind::DiscardStaleWindows | WireMessageKind::ProbePrefill | WireMessageKind::RestorePrefill | WireMessageKind::TryRestorePrefill diff --git a/crates/skippy-server/src/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 fff3d21403..c15508f383 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; @@ -105,6 +107,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<'_>, @@ -652,3 +693,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 fc9b1c3081..44654f3e88 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -8,7 +8,7 @@ use crate::binary_transport::{ AsyncForwarder, BinaryStageExecutionOptions, forwarded_stage_message, forwarded_stage_message_timed, run_binary_stage_message, write_stage_message_conditioned, }; -use crate::frontend::embedded_execution::VerifyRetirement; +use crate::frontend::embedded_execution::{StaleWindowDiscard, VerifyRetirement}; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window, propose_configured_ngram_tokens, @@ -35,6 +35,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}; @@ -1279,6 +1280,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(), @@ -1839,6 +1856,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 c724713847..80d00d526b 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -195,6 +195,24 @@ pub(super) fn can_seed_pipeline(windows: &VecDeque) -> windows.iter().all(|window| window.stale) } +/// Inclusive window-id range of the stale windows of `epoch`, if any. +pub(super) fn stale_window_id_range( + windows: &VecDeque, + epoch: u64, +) -> Option<(i32, i32)> { + let mut range: Option<(i32, i32)> = None; + for window in windows { + if window.epoch == epoch && window.stale { + let id = window.window.id; + range = Some(match range { + Some((min, max)) => (min.min(id), max.max(id)), + None => (id, id), + }); + } + } + range +} + pub(super) fn mark_epoch_stale( windows: &mut VecDeque, epoch: u64, diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 4fc48b42bc..e00a6670e4 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -174,6 +174,39 @@ pub(super) fn embedded_verify_window_message( }) } +/// Invalidates verify windows `min_window_id..=max_window_id` for a request +/// after divergence. The downstream records the range at receive time so +/// buffered stale windows are answered with an empty reply instead of being +/// executed. The window-id range rides in `tokens`. +pub(super) fn discard_stale_windows_message( + wire_dtype: WireActivationDType, + request_id: u64, + session_id: u64, + min_window_id: i32, + max_window_id: i32, +) -> OpenAiResult { + if min_window_id > max_window_id { + return Err(OpenAiError::backend( + "stale window discard range must be non-empty", + )); + } + let kind = WireMessageKind::DiscardStaleWindows; + Ok(StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, wire_dtype), + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![min_window_id, max_window_id], + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }) +} + pub(super) fn retire_verify_window_message( wire_dtype: WireActivationDType, request_id: u64, From 7ecaba4de6edce67a6452a54048497d1360a1e29 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 16:57:09 +1000 Subject: [PATCH 03/15] config: record verify_window_runahead_tokens in the defaults UI schema fixture Co-Authored-By: Claude Opus 5 --- .../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 afccf958ef..a991f7875a 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 @@ -924,6 +924,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.verify_window_runahead_tokens", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.throughput.continuous_batching", "support": "supported", @@ -995,4 +1002,4 @@ } } ] -} +} \ No newline at end of file From 228b25a3805f8d5ef7b1edfbf77ea8ac52cc1579 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 17:59:37 +1000 Subject: [PATCH 04/15] chore: rustfmt, move StaleWindowDiscard above the test module, regen console-print ratchet Co-Authored-By: Claude Opus 5 --- crates/skippy-protocol/src/lib.rs | 5 ++--- .../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 -- tools/xtask/data/console_print_allowlist.json | 12 +++++------ 9 files changed, 29 insertions(+), 29 deletions(-) diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 89a16d5848..352611820a 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -30,9 +30,8 @@ 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_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, STAGE_SUBPROTOCOL_MAJOR, STAGE_SUBPROTOCOL_NAME, StageFrameError, diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index bf810d3a4f..ff0989b72d 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -38,8 +38,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 88ffd706b1..9cb52623d6 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,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_message; @@ -12,6 +11,7 @@ use super::session_tracker::{ ConnectionSessionTracker, combine_connection_and_cleanup_results, release_tracked_connection_sessions, }; +use super::stale_discard::StaleDiscardRegistry; use super::summary::BinaryMessageObservation; use super::summary::BinaryRequestSummary; use super::telemetry::UpstreamReplyWriteSpan; 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 c15508f383..43e870dd1b 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, @@ -693,11 +699,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 44654f3e88..bad2d0c909 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -35,8 +35,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 serde_json::json; diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index 497a948ec4..fd2a88df52 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -65,8 +65,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)] diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 001b401d61..7ab74f3f2b 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4205,27 +4205,27 @@ ], "crates/skippy-server/src/binary_transport/binary_messaging.rs": [ { - "line": 300, + "line": 301, "macro_name": "eprintln!" }, { - "line": 310, + "line": 311, "macro_name": "println!" }, { - "line": 333, + "line": 334, "macro_name": "eprintln!" }, { - "line": 353, + "line": 354, "macro_name": "eprintln!" }, { - "line": 361, + "line": 362, "macro_name": "eprintln!" }, { - "line": 415, + "line": 416, "macro_name": "eprintln!" } ], From dcd1fa4684cde92ca21013c8424fae88f77ed752 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sun, 23 Aug 2026 18:37:35 +1000 Subject: [PATCH 05/15] review: gate the discard wire kind on stage generation 5, allow runahead zero, decouple reader lookahead, join the reader on drop - 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. Co-Authored-By: Claude Opus 5 --- .../mesh-llm-config/src/model_validation.rs | 4 +- .../src/inference/skippy/resolver/tests.rs | 43 +++++++ .../src/protocol/convert.rs | 4 +- .../src/protocol/tests/announcements.rs | 4 +- crates/skippy-protocol/src/lib.rs | 6 +- crates/skippy-protocol/src/validation.rs | 6 +- crates/skippy-server/README.md | 2 +- .../binary_messaging/message_receive.rs | 118 +++++++++++++++++- docs/design/TESTING.md | 2 +- docs/skippy/DATA_FLOW.md | 2 +- 10 files changed, 173 insertions(+), 18 deletions(-) diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index de950ff04d..0487340639 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -637,7 +637,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 86e5e8c948..2d4e8d8592 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -1592,3 +1592,46 @@ verify_window_pipeline_depth = 2 assert_eq!(translated.max_proposal_tokens, 48); assert_eq!(args.speculative.verify_window.pipeline_depth, 2); } + +#[test] +fn model_level_zero_runahead_overrides_a_positive_global_default() { + use crate::plugin::SpeculativeConfig; + let global: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 2 +verify_window_runahead_tokens = 256 +"#, + ) + .expect("parse global speculative config"); + let model: SpeculativeConfig = toml::from_str( + r#" +verify_window_runahead_tokens = 0 +"#, + ) + .expect("parse model speculative config"); + let inherited = super::speculative::resolve_speculative_config( + None, + Some(&global), + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect("global run-ahead must resolve"); + assert_eq!(inherited.decode.verify_window.runahead_max_tokens, 256); + let overridden = super::speculative::resolve_speculative_config( + Some(&model), + Some(&global), + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect("model-level zero must resolve to fixed-depth mode"); + assert_eq!( + overridden.decode.verify_window.runahead_max_tokens, 0, + "Some(0) at the model level must win over the inherited positive default" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/protocol/convert.rs b/crates/mesh-llm-host-runtime/src/protocol/convert.rs index e928cb7931..c5c899b955 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/convert.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/convert.rs @@ -17,7 +17,7 @@ fn skippy_stage_subprotocols( let mut features = vec![skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL.to_string()]; if stage_protocol_generation_supported { features.push( - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5.to_string(), ); } if artifact_transfer_supported { @@ -50,7 +50,7 @@ fn supports_skippy_status_list(subprotocols: &[crate::proto::node::MeshSubprotoc fn supports_skippy_stage_generation(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { supports_skippy_stage_feature( subprotocols, - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, ) && supports_skippy_stage_feature( subprotocols, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, diff --git a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs index 691796c9de..c2af8166ac 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs @@ -82,7 +82,7 @@ fn owner_fields_roundtrip_through_proto_announcement() { .any(|feature| feature == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST) ); assert!(skippy.features.iter().any(|feature| feature - == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4)); + == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5)); assert_eq!( proto_pa .owner_attestation @@ -310,7 +310,7 @@ fn proto_announcement_without_stage_control_is_not_stage_compatible() { name: skippy_protocol::STAGE_SUBPROTOCOL_NAME.to_string(), major: skippy_protocol::STAGE_SUBPROTOCOL_MAJOR, features: vec![ - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5.to_string(), ], }], ..Default::default() diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 352611820a..6a7cd52906 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -33,7 +33,7 @@ pub use validation::{ SCHEMA_VERSION, STAGE_ALPN_V2, STAGE_PROTOCOL_GENERATION, STAGE_STREAM_ARTIFACT_TRANSFER, STAGE_STREAM_CONTROL, STAGE_STREAM_TRANSPORT, STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER, STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST, STAGE_SUBPROTOCOL_MAJOR, STAGE_SUBPROTOCOL_NAME, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -53,7 +53,7 @@ mod tests { stage_control_response, }; use super::{ - STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -62,7 +62,7 @@ mod tests { #[test] fn stage_protocol_generation_feature_names_current_generation() { assert_eq!( - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5, format!("stage-generation-{STAGE_PROTOCOL_GENERATION}") ); } diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index defc72bea4..2dc41cc467 100644 --- a/crates/skippy-protocol/src/validation.rs +++ b/crates/skippy-protocol/src/validation.rs @@ -6,13 +6,13 @@ pub const STAGE_ALPN_V2: &[u8] = b"skippy-stage/2"; pub const STAGE_SUBPROTOCOL_NAME: &str = "skippy-stage"; pub const STAGE_SUBPROTOCOL_MAJOR: u32 = 2; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL: &str = "stage-control"; -pub const STAGE_PROTOCOL_GENERATION: u32 = 4; +pub const STAGE_PROTOCOL_GENERATION: u32 = 5; /// Generation-scoped stage capability. A peer can advertise `stage-control` /// while still rejecting current-generation frames, so split planning gates on /// this exact token before sending current-generation control requests. -pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4: &str = "stage-generation-4"; +pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5: &str = "stage-generation-5"; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4; + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V5; pub const STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER: &str = "artifact-transfer"; pub const STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST: &str = "status-list"; pub const STAGE_STREAM_CONTROL: u8 = 0x01; diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 9859c0a5f5..6c20becb0d 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -104,7 +104,7 @@ deadline handling. - `serve-binary` is the tuned binary stage-to-stage path. - `serve-binary` participates in the breaking generation-4 stage protocol. - Stage compatibility requires `stage-generation-4`; direct prediction return and + Stage compatibility requires `stage-generation-5`; direct prediction return and exact verify-checkpoint retirement are part of that generation's contract, so older peers are rejected during split planning instead of being mixed into a generation-4 topology. 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); + } +} diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index c5fd451b7a..64447ea8ba 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -848,7 +848,7 @@ cached and a worker does not: to open `skippy-stage/2`, then Skippy artifact-transfer stream 0x03, to fetch only its assigned package files before the normal HF fallback path. - Current/released mixed mesh: a released coordinator without advertised - `skippy-stage/2` `artifact-transfer`, `stage-generation-4`, and + `skippy-stage/2` `artifact-transfer`, `stage-generation-5`, and `direct-prediction-return` support must not be selected for a generation-4 split topology; the worker must fall back to local/HF package resolution. - Default public-mesh safety: with `MESH_LLM_ARTIFACT_TRANSFER` unset, the node diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 9f78d34a17..7b269a6162 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -44,7 +44,7 @@ chain alone makes the hot path six hops, or about 60 ms before compute. Stage protocol generation 4 is a compatibility-breaking change. A peer is stage compatible only when it advertises both `skippy-stage/2` and -`stage-generation-4`. Prediction-bearing messages return +`stage-generation-5`. Prediction-bearing messages return directly from the final/readout stage to the driver-facing stage. Intermediate stages continue to forward activations and may handle cold-path control acknowledgments, but they are not part of the decode-token prediction return path. From c1110d95f4dfe0413a7acee4a4708457414d5dfa Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Tue, 25 Aug 2026 18:16:30 +1000 Subject: [PATCH 06/15] review: hard run-ahead budget past the first window, unblock reader drop on a full queue, bound simulated wire delays, finish generation-5 docs - 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. Co-Authored-By: Claude Opus 5 --- crates/skippy-server/README.md | 8 ++-- .../binary_messaging/message_receive.rs | 37 ++++++++++++++-- .../src/binary_transport/wire.rs | 18 +++++++- .../src/frontend/decode_scheduler.rs | 43 +++++++++++++++++-- .../src/frontend/embedded_generation.rs | 15 ++++++- docs/design/TESTING.md | 2 +- docs/skippy/DATA_FLOW.md | 2 +- 7 files changed, 110 insertions(+), 15 deletions(-) diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 6c20becb0d..93deb7a89f 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -15,7 +15,7 @@ mesh/openai-frontend; diagnostic and benchmark clients may connect directly to the first stage. The full request/reply path is tip-to-tip: token IDs enter at the driver-facing -tip, and activations flow through the stage chain. Stage protocol generation 4 +tip, and activations flow through the stage chain. Stage protocol generation 5 is a compatibility-breaking contract: prediction-bearing replies return directly from the final/readout tip to the driver-facing stage instead of being relayed back through intermediate stages. Middle-out is the prefill optimization @@ -103,11 +103,11 @@ deadline handling. ## Notes - `serve-binary` is the tuned binary stage-to-stage path. -- `serve-binary` participates in the breaking generation-4 stage protocol. +- `serve-binary` participates in the breaking generation-5 stage protocol. Stage compatibility requires `stage-generation-5`; direct prediction return and exact verify-checkpoint retirement are part of that generation's contract, so older peers are rejected during split planning instead of being mixed into a - generation-4 topology. + generation-5 topology. - `serve-binary` accepts upstream protocol connections concurrently. Model execution remains serialized by the per-process runtime lock, but readiness, abandoned, or broken connections do not monopolize the listener and block the @@ -123,7 +123,7 @@ deadline handling. `/v1/completions` using the shared `openai-frontend` crate for a local final/single-stage config with no downstream peer. Split serving uses embedded stage-0 OpenAI serving from `serve-binary --openai-bind-addr` because - generation-4 prediction returns flow directly from the final stage to stage 0. + generation-5 prediction returns flow directly from the final stage to stage 0. The older standalone `serve-openai --first-stage-addr` adapter is no longer supported. `--model-id` is the exact served model id to advertise and accept, for example `org/repo:Q4_K_M`; it is not parsed as stage topology. diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/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 1b81db3b47..18c363d4aa 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 bad2d0c909..4b0a6fb022 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -997,6 +997,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 { @@ -1018,7 +1024,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(); diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index 64447ea8ba..ef13dc5f4e 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -849,7 +849,7 @@ cached and a worker does not: fetch only its assigned package files before the normal HF fallback path. - Current/released mixed mesh: a released coordinator without advertised `skippy-stage/2` `artifact-transfer`, `stage-generation-5`, and - `direct-prediction-return` support must not be selected for a generation-4 + `direct-prediction-return` support must not be selected for a generation-5 split topology; the worker must fall back to local/HF package resolution. - Default public-mesh safety: with `MESH_LLM_ARTIFACT_TRANSFER` unset, the node must advertise no `artifact-transfer` feature, reject inbound artifact diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 7b269a6162..cf9aacc13a 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -42,7 +42,7 @@ chain alone makes the hot path six hops, or about 60 ms before compute. ## Generation 4 Direct Prediction Return and Verify Retirement -Stage protocol generation 4 is a compatibility-breaking change. A peer is stage +Stage protocol generation 5 is a compatibility-breaking change. A peer is stage compatible only when it advertises both `skippy-stage/2` and `stage-generation-5`. Prediction-bearing messages return directly from the final/readout stage to the driver-facing stage. Intermediate stages From 243d0503686bcca497c6a4715e0556f3efc0e530 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 16:20:58 +1000 Subject: [PATCH 07/15] review: serialize the discard writer with teardown, bound the reader by bytes, per-thread jitter - 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. Co-Authored-By: Claude Opus 5 --- .../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 a64b50751c..c98f526c22 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs @@ -23,8 +23,24 @@ use std::time::Instant; const ASYNC_FORWARD_TERMINAL_TIMEOUT: Duration = Duration::from_secs(30); pub(crate) struct AsyncForwarder { - sender: mpsc::SyncSender, + sender: Option>, pending: VecDeque, + writer: Option>, +} + +impl Drop for AsyncForwarder { + /// Queued frames must not still be on the wire after the request that + /// owns them returns: a persistent lane is handed back for reuse, and a + /// teardown `Stop` written through another clone of the same socket + /// would interleave with a frame this forwarder is still writing. + /// Dropping the sender ends the writer loop once its queue drains, and + /// the join makes that ordering observable to the caller. + fn drop(&mut self) { + drop(self.sender.take()); + if let Some(writer) = self.writer.take() { + let _ = writer.join(); + } + } } pub(crate) struct AsyncForwardReceipt { @@ -54,10 +70,12 @@ impl AsyncForwarder { .set_write_timeout(Some(ASYNC_FORWARD_TERMINAL_TIMEOUT)) .context("set async activation forward write timeout")?; let (sender, receiver) = mpsc::sync_channel::(queue_capacity.max(1)); - thread::spawn(move || run_forwarder(&mut writer, &receiver, &telemetry)); + let writer_thread = + thread::spawn(move || run_forwarder(&mut writer, &receiver, &telemetry)); Ok(Self { - sender, + sender: Some(sender), pending: VecDeque::new(), + writer: Some(writer_thread), }) } @@ -83,6 +101,8 @@ impl AsyncForwarder { self.reap_completed()?; let (done, receiver) = mpsc::channel(); self.sender + .as_ref() + .ok_or_else(|| anyhow!("async activation forwarder stopped"))? .send(AsyncForwardJob { message, wire_dtype, @@ -114,7 +134,7 @@ impl AsyncForwarder { } } - pub(super) fn flush(&mut self) -> Result<()> { + pub(crate) fn flush(&mut self) -> Result<()> { while let Some(receiver) = self.pending.pop_front() { receiver.finish()?; } @@ -222,6 +242,48 @@ mod tests { } } + /// A delayed discard must be fully written before a teardown `Stop` that + /// goes out through a different clone of the same socket, otherwise the + /// two frames interleave and poison a lane that is handed back for reuse. + #[test] + fn a_delayed_discard_lands_before_a_teardown_stop_on_another_clone() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let mut client = TcpStream::connect(address).unwrap(); + let (mut server, _) = listener.accept().unwrap(); + let telemetry = Telemetry::new(None, 1, prefix_cache_test_config(), TelemetryLevel::Off); + let mut forwarder = AsyncForwarder::new(&client, telemetry, 8).unwrap(); + // 250ms of simulated propagation: without the drop-time join, the + // teardown write below wins the race and the frames interleave. + let condition = WireCondition::new(250.0, None).unwrap(); + + forwarder + .send( + message(WireMessageKind::DiscardStaleWindows, 11), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap(); + drop(forwarder); + + write_stage_message_after_propagation( + &mut client, + &message(WireMessageKind::Stop, 22), + WireActivationDType::F32, + WireCondition::new(0.0, None).unwrap(), + ) + .unwrap(); + + let first = read_stage_message(&mut server, 4).unwrap(); + let second = read_stage_message(&mut server, 4).unwrap(); + + assert_eq!(first.kind, WireMessageKind::DiscardStaleWindows); + assert_eq!(first.pos_start, 11); + assert_eq!(second.kind, WireMessageKind::Stop); + assert_eq!(second.pos_start, 22); + } + #[test] fn retirement_receipt_orders_all_prior_verify_writes() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/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 18c363d4aa..98a295b40e 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, WireActivationDType, 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); @@ -197,9 +212,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 4b0a6fb022..8392232471 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1884,6 +1884,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 a8eeaa70c7196e49a75a0ed554888ed8c2044991 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 16:42:16 +1000 Subject: [PATCH 08/15] review: release a reader parked on the byte ceiling when the connection tears down 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. Co-Authored-By: Claude Opus 5 --- .../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 b60599e79f2547be12ee65786cd1307ea480e1a9 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 23:26:36 +1000 Subject: [PATCH 09/15] review: independent per-lane jitter streams, tighter read-ahead ceiling, document the discard frame 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. Co-Authored-By: Claude Opus 5 --- crates/skippy-server/README.md | 15 ++- .../binary_messaging/message_receive.rs | 29 +++-- .../src/binary_transport/wire.rs | 121 +++++++++++++----- docs/skippy/DATA_FLOW.md | 34 ++++- 4 files changed, 154 insertions(+), 45 deletions(-) diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 93deb7a89f..4f07aff412 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -104,10 +104,17 @@ deadline handling. - `serve-binary` is the tuned binary stage-to-stage path. - `serve-binary` participates in the breaking generation-5 stage protocol. - Stage compatibility requires `stage-generation-5`; 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-5 topology. + Stage compatibility requires `stage-generation-5`; direct prediction return, + exact verify-checkpoint retirement, and the `DiscardStaleWindows` control + frame are part of that generation's contract, so older peers are rejected + during split planning instead of being mixed into a generation-5 topology. +- That rejection happens in mesh split planning. A manually wired + `serve-binary --downstream host:port` pair performs no generation + handshake, so **the contract for the standalone path is that all stages are + upgraded together**. Pointing a run-ahead coordinator at a generation-4 + stage binary is not degraded gracefully: the older peer rejects the + `DiscardStaleWindows` frame as an unknown message kind and drops the + request connection. - `serve-binary` accepts upstream protocol connections concurrently. Model execution remains serialized by the per-process runtime lock, but readiness, abandoned, or broken connections do not monopolize the listener and block the 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 98a295b40e..d5dc2793db 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, WireActivationDType, 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, @@ -218,19 +249,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 cf9aacc13a..1042d19b8f 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -40,7 +40,7 @@ activation links and then crossed three reply links before stage 0 could emit the token. On a topology with a fixed 10 ms delay per inter-stage hop, the reply chain alone makes the hot path six hops, or about 60 ms before compute. -## Generation 4 Direct Prediction Return and Verify Retirement +## Generation 5 Direct Prediction Return and Verify Retirement Stage protocol generation 5 is a compatibility-breaking change. A peer is stage compatible only when it advertises both `skippy-stage/2` and @@ -79,6 +79,38 @@ With the same four stages and 10 ms inter-stage delay, the no-spec decode hot path becomes `S0 -> S1 -> S2 -> S3 -> S0`: four hops, or about 40 ms before compute. That removes two serialized reply hops from every generated token. +## Stale Verify-Window Discard + +Generation 5 also adds the `DiscardStaleWindows` control frame (wire kind +23), which is what made the generation compatibility-breaking: a +generation-4 peer rejects the kind outright and drops the request +connection. + +Run-ahead admission dispatches verify windows before their predecessors are +verified, so a rejection strands every window queued behind the divergence. +Those windows are already on the wire. The coordinator sends one discard +naming a contiguous `[min_window_id, max_window_id]` range for the request, +and it travels the same ordered path as the windows it cancels, so it is +read after them. + +The handling rule differs by position in the chain, and this is the +non-obvious part: + +- **Middle stages forward the frame and still execute the stale windows.** + A stage that is not the final stage owns KV state its downstream neighbour + depends on; skipping its forward pass would desynchronize the chain. It + passes the discard along so the frame reaches the stage that can act on it. +- **The final/readout stage skips the named windows.** It is the only stage + whose output is thrown away by a discard, so skipping there is what + actually saves the work — the expensive readout and the direct prediction + return back to stage 0. + +The receiving stage records the range the moment it parses the frame, ahead +of executing the backlog queued in front of it, which is what lets the skip +take effect before the stale tail runs. See `INBOUND_LOOKAHEAD_BYTES` for +the case where a wide-frame backlog outruns that read-ahead and the tail +executes anyway. + ## Relative Sizes | Flow | Size | From 5500159c6a6ead172e08ff6244630086f93c9f7b Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 10:39:28 +1000 Subject: [PATCH 10/15] Add draft-model fallback proposals for N-gram misses (ngram_fallback = "draft") When the suffix/cache proposer misses, the pipelined verify-window path now proposes from the configured draft model instead of degrading to one token per round trip: the draft session syncs incrementally to the committed context plus optimistic suffix (full re-prefill only on divergence), and its greedy rollout feeds the same candidate pipeline. Config: speculative.ngram_fallback = "draft" alongside a draft_model; the classic serial draft loop stays disabled while fallback drives the pipeline, and depth-1 setups keep classic behavior. Falls back opt-in, pipelined only. Co-Authored-By: Claude Opus 5 --- crates/mesh-llm-config/src/model.rs | 6 ++ .../control_behavior/speculative.rs | 3 +- .../src/model/built_in_schema/declarations.rs | 1 + .../mesh-llm-config/src/model_validation.rs | 10 ++- .../inference/skippy/resolver/speculative.rs | 22 ++++++ .../src/binary_transport/options.rs | 1 + .../src/frontend/embedded_generation.rs | 77 +++++++++++++++++-- .../src/frontend/generation/draft_runner.rs | 31 +++++++- .../skippy-server/src/frontend/speculative.rs | 24 ++++++ 9 files changed, 166 insertions(+), 9 deletions(-) diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index a9fb52bac6..1bc374aaae 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -604,6 +604,7 @@ pub struct SpeculativeConfig { pub verify_window_max_tokens: Option, pub verify_window_pipeline_depth: Option, pub verify_window_runahead_tokens: Option, + pub ngram_fallback: Option, pub spec_default: Option, pub(crate) legacy_draft_model_path_used: bool, } @@ -656,6 +657,7 @@ impl SpeculativeConfig { verify_window_max_tokens: pick!(verify_window_max_tokens), verify_window_pipeline_depth: pick!(verify_window_pipeline_depth), verify_window_runahead_tokens: pick!(verify_window_runahead_tokens), + ngram_fallback: pick!(ngram_fallback), spec_default: pick!(spec_default), legacy_draft_model_path_used: overrides .filter(|config| config.draft_model.is_some()) @@ -731,6 +733,8 @@ struct SpeculativeConfigRaw { #[serde(default)] verify_window_runahead_tokens: Option, #[serde(default)] + ngram_fallback: Option, + #[serde(default)] spec_default: Option, } @@ -776,6 +780,7 @@ impl<'de> Deserialize<'de> for SpeculativeConfig { verify_window_max_tokens: raw.verify_window_max_tokens, verify_window_pipeline_depth: raw.verify_window_pipeline_depth, verify_window_runahead_tokens: raw.verify_window_runahead_tokens, + ngram_fallback: raw.ngram_fallback, spec_default: raw.spec_default, legacy_draft_model_path_used: legacy_used, }) @@ -842,6 +847,7 @@ impl Serialize for SpeculativeConfig { "verify_window_runahead_tokens", &self.verify_window_runahead_tokens, )?; + map.serialize_entry("ngram_fallback", &self.ngram_fallback)?; map.serialize_entry("spec_default", &self.spec_default)?; map.end() } diff --git a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs index 0698071779..4ec3059257 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/speculative.rs @@ -74,7 +74,8 @@ pub(super) fn apply_speculative_behavior( | "verify_window_min_tokens" | "verify_window_max_tokens" | "verify_window_pipeline_depth" - | "verify_window_runahead_tokens" => {} + | "verify_window_runahead_tokens" + | "ngram_fallback" => {} _ => {} } } diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 39003a11e7..5dc208ff48 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -663,6 +663,7 @@ fn speculative_settings(prefix: &str) -> Vec { &format!("{prefix}.verify_window_runahead_tokens"), ConfigValueSchema::Integer, ), + basic_setting(&format!("{prefix}.ngram_fallback"), ConfigValueSchema::String), basic_setting(&format!("{prefix}.spec_default"), bool_or_auto_schema()), ] } diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index 0487340639..2be79e0f69 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -641,7 +641,15 @@ fn validate_verify_window_controls( // can switch run-ahead back off when the global defaults enable it. 0, u32::try_from(MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS).expect("runahead limit fits u32"), - ) + )?; + if let Some(fallback) = config.ngram_fallback.as_deref() { + validate_allowed( + fallback, + &["draft", "none"], + &format!("{base_path}.ngram_fallback"), + )?; + } + Ok(()) } fn validate_request_defaults(config: &RequestDefaultsConfig, base_path: &str) -> DiagnosticResult { diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index 600f9d9e30..ebd1824146 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 @@ -411,6 +411,27 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result config.verify_window.max_tokens { bail!("skippy speculative verify window requires min_tokens <= max_tokens"); } + let ngram_fallback = pick_string( + input + .model_config + .and_then(|value| value.ngram_fallback.as_deref()), + input + .global_config + .and_then(|value| value.ngram_fallback.as_deref()), + None, + ); + config.ngram_fallback_draft = match ngram_fallback { + "draft" => { + if config.ngram.is_none() { + bail!( + "skippy speculative ngram_fallback = \"draft\" requires an N-gram strategy" + ); + } + true + } + "none" | "" => false, + other => bail!("skippy speculative ngram_fallback must be draft or none, got {other}"), + }; config.validate()?; Ok(config) } @@ -506,6 +527,7 @@ fn package_decode_config( ngram, extension, verify_window, + ngram_fallback_draft: false, })) } diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index e0c4b972d2..903acb719e 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -238,6 +238,7 @@ mod tests { pipeline_depth: 2, runahead_max_tokens: 0, }, + ngram_fallback_draft: false, } } diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 8392232471..a5c00ef952 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -805,17 +805,28 @@ impl StageOpenAiBackend { ) }, ); + // Draft fallback keeps the draft model resident purely as a + // proposal source for the pipelined path when the N-gram proposer + // misses; the classic serial draft loop stays disabled. Pipelined + // paths only — at depth 1 the classic draft loop is strictly + // better. + let ngram_fallback_draft_enabled = effective_speculative.ngram_fallback_draft + && effective_speculative.ngram.is_some() + && draft_guard.is_some() + && verify_window_scheduler.depth() > 1; + let draft_blocks_pipeline = draft_guard.is_some() && !ngram_fallback_draft_enabled; let composite_sidecar_enabled = - native_mtp_options.ngram_hybrid && draft_guard.is_none(); + native_mtp_options.ngram_hybrid && !draft_blocks_pipeline; // A standalone N-gram plan (no native MTP, no draft model) can drive // the same verify-window pipeline; the single-window native-MTP path // stays composite-only, so standalone drafting still falls back to the // serial block at depth 1. let standalone_ngram_pipelining = !request.native_mtp_enabled && effective_speculative.ngram.is_some() - && draft_guard.is_none(); - let native_mtp_verify_windows_enabled = - (request.native_mtp_enabled || composite_sidecar_enabled) && draft_guard.is_none(); + && !draft_blocks_pipeline; + let native_mtp_verify_windows_enabled = (request.native_mtp_enabled + || composite_sidecar_enabled) + && !draft_blocks_pipeline; let pipelined_decode_enabled = (composite_sidecar_enabled || standalone_ngram_pipelining) && verify_window_scheduler.depth() > 1; @@ -897,6 +908,32 @@ impl StageOpenAiBackend { ), cached_ngram_proposer.as_mut(), )?; + let proposal = if proposal.tokens().len() < 2 + && ngram_fallback_draft_enabled + && pipelined_decode_enabled + { + let fallback_timer = PhaseTimer::start(); + let draft = draft_guard + .as_deref_mut() + .expect("fallback requires a draft guard"); + draft + .sync_to_context(&context_tokens) + .map_err(openai_backend_error)?; + let budget = native_mtp_options + .ngram_max_proposal_tokens + .min(draft.window.max(1)) + .min(native_mtp_remaining) + .max(2); + let draft_tokens = draft + .propose(current, budget) + .map_err(openai_backend_error)?; + speculative_stats.fallback_draft_proposals += 1; + speculative_stats.fallback_draft_tokens += draft_tokens.len(); + speculative_stats.fallback_draft_ms += fallback_timer.elapsed_ms(); + NativeMtpHybridProposal::from_parts(draft_tokens, 0, true) + } else { + proposal + }; if proposal.supports_positional_pipeline(verify_window_scheduler.depth()) && ngram_sidecar_controller.permit_pipeline_start() && verify_window_scheduler.supports_pipelining( @@ -1013,12 +1050,40 @@ impl StageOpenAiBackend { ); let refill_budget = ngram_sidecar_controller.refill_limit(available_refill_tokens); - let appended = refill_pipeline_ngram_candidates( + let mut appended = refill_pipeline_ngram_candidates( pipeline, &context_tokens, &mut cached_ngram_proposer, refill_budget, )?; + if appended == 0 + && !pipeline.has_remaining_candidates() + && ngram_fallback_draft_enabled + && refill_budget >= 2 + { + let fallback_timer = PhaseTimer::start(); + let draft = draft_guard + .as_deref_mut() + .expect("fallback requires a draft guard"); + let mut sequence = context_tokens.clone(); + sequence.extend_from_slice(pipeline.optimistic_suffix()); + if let Some(&last) = sequence.last() { + draft + .sync_to_context(&sequence) + .map_err(openai_backend_error)?; + let budget = refill_budget.min(draft.window.max(1)); + let draft_tokens = draft + .propose(last, budget) + .map_err(openai_backend_error)?; + speculative_stats.fallback_draft_proposals += 1; + speculative_stats.fallback_draft_tokens += + draft_tokens.len(); + speculative_stats.fallback_draft_ms += + fallback_timer.elapsed_ms(); + appended = + pipeline.append_ngram_candidates(&draft_tokens); + } + } verify_window_scheduler.record_horizon_refill(appended); } if !pipeline.has_remaining_candidates() { @@ -1402,7 +1467,7 @@ impl StageOpenAiBackend { continue; } } - if draft_guard.is_some() + if (draft_guard.is_some() && !ngram_fallback_draft_enabled) || (effective_speculative.ngram.is_some() && !pipelined_decode_enabled) { let remaining = (request.max_tokens as usize).saturating_sub(decoded_tokens); diff --git a/crates/skippy-server/src/frontend/generation/draft_runner.rs b/crates/skippy-server/src/frontend/generation/draft_runner.rs index 76e53e5517..fdbb539bf3 100644 --- a/crates/skippy-server/src/frontend/generation/draft_runner.rs +++ b/crates/skippy-server/src/frontend/generation/draft_runner.rs @@ -20,6 +20,10 @@ pub(in crate::frontend) struct DraftRunner { pub(in crate::frontend) window: usize, pub(in crate::frontend) _model: StageModel, pub(in crate::frontend) session: StageSession, + /// Tokens currently materialized in the draft session's KV, maintained so + /// fallback proposals can advance incrementally instead of re-prefilling + /// the whole context on every call. + synced: Vec, } impl DraftRunner { @@ -70,19 +74,43 @@ impl DraftRunner { window, _model: model, session, + synced: Vec::new(), }) } pub(in crate::frontend) fn reset_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { self.session.reset().context("reset draft session")?; + self.synced.clear(); if context_tokens.len() > 1 { + let prefix = &context_tokens[..context_tokens.len() - 1]; self.session - .prefill_chunk(&context_tokens[..context_tokens.len() - 1]) + .prefill_chunk(prefix) .context("prefill draft context")?; + self.synced.extend_from_slice(prefix); } Ok(()) } + /// Brings the draft session to `context_tokens` (all but the last token + /// prefilled, ready to propose from the last). Extends incrementally when + /// the already-synced tokens are a prefix of the target — the common case + /// when prior fallback proposals were accepted — and falls back to a full + /// reset on divergence. + pub(in crate::frontend) fn sync_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { + let target = &context_tokens[..context_tokens.len().saturating_sub(1)]; + if !self.synced.is_empty() && target.starts_with(&self.synced) { + let delta = &target[self.synced.len()..]; + if !delta.is_empty() { + self.session + .prefill_chunk(delta) + .context("advance draft context")?; + self.synced.extend_from_slice(delta); + } + return Ok(()); + } + self.reset_to_context(context_tokens) + } + pub(in crate::frontend) fn propose( &mut self, mut current: i32, @@ -90,6 +118,7 @@ impl DraftRunner { ) -> Result> { let mut tokens = Vec::with_capacity(max_tokens); for _ in 0..max_tokens { + self.synced.push(current); current = self .session .decode_step(current) diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index fd2a88df52..9ce384dc0d 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -28,6 +28,11 @@ pub struct SpeculativeDecodeConfig { pub ngram: Option, pub extension: Option, pub verify_window: VerifyWindowConfig, + /// Propose from the configured draft model when the N-gram proposer + /// misses, instead of degrading to one token per round trip. Pipelined + /// paths only; requires a draft model. + #[serde(default)] + pub ngram_fallback_draft: bool, } /// Native multi-token-prediction draft settings. @@ -121,6 +126,7 @@ impl Default for SpeculativeDecodeConfig { pipeline_depth: 1, runahead_max_tokens: 0, }, + ngram_fallback_draft: false, } } } @@ -180,6 +186,9 @@ impl SpeculativeDecodeConfig { "verify window requires 0 < min_tokens <= max_tokens and 0 < pipeline_depth <= {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" ); } + if self.ngram_fallback_draft && self.ngram.is_none() { + bail!("ngram_fallback_draft requires an N-gram proposer to fall back from"); + } if self.verify_window.runahead_max_tokens > MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS { bail!( "verify window runahead_max_tokens must not exceed {MAX_VERIFY_WINDOW_RUNAHEAD_TOKENS}" @@ -379,6 +388,9 @@ mod standalone_speculative_config_tests { pub(super) struct OpenAiSpeculativeStats { pub(super) windows: usize, pub(super) draft_tokens: usize, + pub(super) fallback_draft_proposals: usize, + pub(super) fallback_draft_tokens: usize, + pub(super) fallback_draft_ms: f64, pub(super) accepted_tokens: usize, pub(super) rejected_tokens: usize, pub(super) full_accept_windows: usize, @@ -610,6 +622,18 @@ fn elapsed_us(started: Instant) -> u64 { impl OpenAiSpeculativeStats { pub(super) fn insert_response_timings(&self, timings: &mut BTreeMap) { timings.insert("speculative_windows".to_string(), json!(self.windows)); + timings.insert( + "speculative_fallback_draft_proposals".to_string(), + json!(self.fallback_draft_proposals), + ); + timings.insert( + "speculative_fallback_draft_tokens".to_string(), + json!(self.fallback_draft_tokens), + ); + timings.insert( + "speculative_fallback_draft_ms".to_string(), + json!(self.fallback_draft_ms), + ); timings.insert( "speculative_proposed_n".to_string(), json!(self.draft_tokens), From 2998be022ef014e039def4cc25e093a612207da8 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 17:01:28 +1000 Subject: [PATCH 11/15] config: record ngram_fallback in the defaults UI schema fixture Co-Authored-By: Claude Opus 5 --- .../fixtures/config_schema_defaults_ui_reference.json | 7 +++++++ 1 file changed, 7 insertions(+) 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 a991f7875a..2f7f634575 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 @@ -931,6 +931,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.ngram_fallback", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.throughput.continuous_batching", "support": "supported", From bcc9a4d534fabb93e9784b7d53a97db2fb8601c1 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 17:09:34 +1000 Subject: [PATCH 12/15] config: keep the defaults UI schema fixture alphabetized The exported reference is sorted by canonical path; insert the new speculative entries in order. Co-Authored-By: Claude Opus 5 --- .../config_schema_defaults_ui_reference.json | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) 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 2f7f634575..d5bf2c2296 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json @@ -854,6 +854,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.speculative.ngram_fallback", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.speculative.ngram_max", "support": "supported", @@ -931,13 +938,6 @@ "kind": "built_in" } }, - { - "canonical_path": "defaults.speculative.ngram_fallback", - "support": "supported", - "source": { - "kind": "built_in" - } - }, { "canonical_path": "defaults.throughput.continuous_batching", "support": "supported", From 4a396621e5427189c8c0db080d7f476eecfd3c3d Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Sat, 22 Aug 2026 18:00:18 +1000 Subject: [PATCH 13/15] chore: rustfmt Co-Authored-By: Claude Opus 5 --- .../src/inference/skippy/resolver/speculative.rs | 4 +--- crates/skippy-server/src/frontend/embedded_generation.rs | 8 +++----- 2 files changed, 4 insertions(+), 8 deletions(-) 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 ebd1824146..9082152b2a 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 @@ -423,9 +423,7 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result { if config.ngram.is_none() { - bail!( - "skippy speculative ngram_fallback = \"draft\" requires an N-gram strategy" - ); + bail!("skippy speculative ngram_fallback = \"draft\" requires an N-gram strategy"); } true } diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index a5c00ef952..5e3e32a91c 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -824,9 +824,8 @@ impl StageOpenAiBackend { let standalone_ngram_pipelining = !request.native_mtp_enabled && effective_speculative.ngram.is_some() && !draft_blocks_pipeline; - let native_mtp_verify_windows_enabled = (request.native_mtp_enabled - || composite_sidecar_enabled) - && !draft_blocks_pipeline; + let native_mtp_verify_windows_enabled = + (request.native_mtp_enabled || composite_sidecar_enabled) && !draft_blocks_pipeline; let pipelined_decode_enabled = (composite_sidecar_enabled || standalone_ngram_pipelining) && verify_window_scheduler.depth() > 1; @@ -1080,8 +1079,7 @@ impl StageOpenAiBackend { draft_tokens.len(); speculative_stats.fallback_draft_ms += fallback_timer.elapsed_ms(); - appended = - pipeline.append_ngram_candidates(&draft_tokens); + appended = pipeline.append_ngram_candidates(&draft_tokens); } } verify_window_scheduler.record_horizon_refill(appended); From 226d92022de6541b2b2a4297561940f0017efd53 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 16:20:26 +1000 Subject: [PATCH 14/15] review: unit-test the draft sync bookkeeping and keep the fallback budget inside the window - DraftSyncState/DraftSyncPlan split the load-bearing synced-prefix decision out of the session I/O so it can be tested without a model: prefix extension, exact-sync no-op, divergence reset, proposal steps joining the prefix, and a rejected step forcing a reset. - The fallback budget applies its floor before the remaining-window cap, and the fallback is skipped when fewer than two tokens remain, so a proposal can no longer overshoot the window by a token. Co-Authored-By: Claude Opus 5 --- .../src/frontend/embedded_generation.rs | 11 +- .../src/frontend/generation/draft_runner.rs | 175 ++++++++++++++++-- 2 files changed, 172 insertions(+), 14 deletions(-) diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 5e3e32a91c..c6a29926e1 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -910,6 +910,11 @@ impl StageOpenAiBackend { let proposal = if proposal.tokens().len() < 2 && ngram_fallback_draft_enabled && pipelined_decode_enabled + // A fallback proposal is only worth its draft decode + // when at least two tokens still fit the remaining + // window; below that, fall through to the serial path + // rather than overshoot the budget by a token. + && native_mtp_remaining >= 2 { let fallback_timer = PhaseTimer::start(); let draft = draft_guard @@ -918,11 +923,13 @@ impl StageOpenAiBackend { draft .sync_to_context(&context_tokens) .map_err(openai_backend_error)?; + // The floor must not lift the budget back over the + // remaining window, so it is applied before the cap. let budget = native_mtp_options .ngram_max_proposal_tokens .min(draft.window.max(1)) - .min(native_mtp_remaining) - .max(2); + .max(2) + .min(native_mtp_remaining); let draft_tokens = draft .propose(current, budget) .map_err(openai_backend_error)?; diff --git a/crates/skippy-server/src/frontend/generation/draft_runner.rs b/crates/skippy-server/src/frontend/generation/draft_runner.rs index fdbb539bf3..b4e0156787 100644 --- a/crates/skippy-server/src/frontend/generation/draft_runner.rs +++ b/crates/skippy-server/src/frontend/generation/draft_runner.rs @@ -23,7 +23,68 @@ pub(in crate::frontend) struct DraftRunner { /// Tokens currently materialized in the draft session's KV, maintained so /// fallback proposals can advance incrementally instead of re-prefilling /// the whole context on every call. - synced: Vec, + synced: DraftSyncState, +} + +/// What a sync to a given context requires of the draft session. Split out +/// from the session I/O because this bookkeeping is load-bearing for KV +/// correctness: claiming a prefix extension the session has not materialized +/// silently corrupts every later proposal. +#[derive(Debug, Eq, PartialEq)] +pub(in crate::frontend) enum DraftSyncPlan { + /// The session is already at the target; nothing to do. + AlreadySynced, + /// The synced tokens are a prefix of the target: prefill only the tail, + /// given here as a range into the target prefix. + Extend { from: usize, to: usize }, + /// The synced tokens diverge from the target: reset and prefill the + /// whole prefix. + Reset, +} + +/// Tokens the draft session has materialized, and the decisions derived from +/// them. +#[derive(Debug, Default)] +pub(in crate::frontend) struct DraftSyncState { + tokens: Vec, +} + +impl DraftSyncState { + /// The prefix a context implies: every token but the last, which is the + /// one a proposal decodes from. + fn target_len(context_tokens: &[i32]) -> usize { + context_tokens.len().saturating_sub(1) + } + + pub(in crate::frontend) fn plan(&self, context_tokens: &[i32]) -> DraftSyncPlan { + let target = &context_tokens[..Self::target_len(context_tokens)]; + if self.tokens.is_empty() || !target.starts_with(&self.tokens) { + return DraftSyncPlan::Reset; + } + if target.len() == self.tokens.len() { + return DraftSyncPlan::AlreadySynced; + } + DraftSyncPlan::Extend { + from: self.tokens.len(), + to: target.len(), + } + } + + fn record_extend(&mut self, delta: &[i32]) { + self.tokens.extend_from_slice(delta); + } + + fn record_reset(&mut self, prefix: &[i32]) { + self.tokens.clear(); + self.tokens.extend_from_slice(prefix); + } + + /// A proposal decodes from `current`, which the session materializes as + /// it steps — so it joins the synced prefix and the next sync can extend + /// instead of resetting. + fn record_proposal_step(&mut self, current: i32) { + self.tokens.push(current); + } } impl DraftRunner { @@ -74,19 +135,19 @@ impl DraftRunner { window, _model: model, session, - synced: Vec::new(), + synced: DraftSyncState::default(), }) } pub(in crate::frontend) fn reset_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { self.session.reset().context("reset draft session")?; - self.synced.clear(); + self.synced.record_reset(&[]); if context_tokens.len() > 1 { let prefix = &context_tokens[..context_tokens.len() - 1]; self.session .prefill_chunk(prefix) .context("prefill draft context")?; - self.synced.extend_from_slice(prefix); + self.synced.record_reset(prefix); } Ok(()) } @@ -97,18 +158,18 @@ impl DraftRunner { /// when prior fallback proposals were accepted — and falls back to a full /// reset on divergence. pub(in crate::frontend) fn sync_to_context(&mut self, context_tokens: &[i32]) -> Result<()> { - let target = &context_tokens[..context_tokens.len().saturating_sub(1)]; - if !self.synced.is_empty() && target.starts_with(&self.synced) { - let delta = &target[self.synced.len()..]; - if !delta.is_empty() { + match self.synced.plan(context_tokens) { + DraftSyncPlan::AlreadySynced => Ok(()), + DraftSyncPlan::Extend { from, to } => { + let delta = &context_tokens[from..to]; self.session .prefill_chunk(delta) .context("advance draft context")?; - self.synced.extend_from_slice(delta); + self.synced.record_extend(delta); + Ok(()) } - return Ok(()); + DraftSyncPlan::Reset => self.reset_to_context(context_tokens), } - self.reset_to_context(context_tokens) } pub(in crate::frontend) fn propose( @@ -118,7 +179,7 @@ impl DraftRunner { ) -> Result> { let mut tokens = Vec::with_capacity(max_tokens); for _ in 0..max_tokens { - self.synced.push(current); + self.synced.record_proposal_step(current); current = self .session .decode_step(current) @@ -209,3 +270,93 @@ pub(in crate::frontend) fn model_layer_count(path: &Path) -> Result { .ok_or_else(|| anyhow!("could not infer layer count for {}", path.display()))?; Ok(layer_count) } + +#[cfg(test)] +mod tests { + use super::*; + + fn state(tokens: &[i32]) -> DraftSyncState { + DraftSyncState { + tokens: tokens.to_vec(), + } + } + + #[test] + fn an_empty_session_always_resets() { + assert_eq!(state(&[]).plan(&[1, 2, 3]), DraftSyncPlan::Reset); + // A context of one token has an empty prefix: still a reset, and + // `reset_to_context` then prefills nothing. + assert_eq!(state(&[]).plan(&[1]), DraftSyncPlan::Reset); + } + + #[test] + fn a_synced_prefix_extends_by_the_delta_only() { + // Synced [1, 2]; context [1, 2, 3, 4, 5] has prefix [1, 2, 3, 4]. + assert_eq!( + state(&[1, 2]).plan(&[1, 2, 3, 4, 5]), + DraftSyncPlan::Extend { from: 2, to: 4 } + ); + } + + #[test] + fn an_exactly_synced_prefix_is_a_no_op() { + // Synced [1, 2, 3]; context [1, 2, 3, 4] has prefix [1, 2, 3]. + assert_eq!( + state(&[1, 2, 3]).plan(&[1, 2, 3, 4]), + DraftSyncPlan::AlreadySynced + ); + } + + #[test] + fn divergence_resets_rather_than_extending() { + // Same length, different token: the KV past that point is wrong. + assert_eq!(state(&[1, 9]).plan(&[1, 2, 3, 4]), DraftSyncPlan::Reset); + // A rejected proposal leaves the session longer than the target. + assert_eq!(state(&[1, 2, 3, 4]).plan(&[1, 2, 3]), DraftSyncPlan::Reset); + } + + #[test] + fn proposal_steps_join_the_synced_prefix_so_the_next_sync_extends() { + let mut synced = state(&[1, 2]); + // Two proposal steps decoded from 3 then 4. + synced.record_proposal_step(3); + synced.record_proposal_step(4); + + // Both accepted, and the caller committed a fifth token: the session + // already holds [1, 2, 3, 4], so only [5] needs prefilling. + assert_eq!( + synced.plan(&[1, 2, 3, 4, 5, 6]), + DraftSyncPlan::Extend { from: 4, to: 5 } + ); + } + + #[test] + fn a_rejected_proposal_step_forces_a_reset() { + let mut synced = state(&[1, 2]); + synced.record_proposal_step(3); + // The verifier rejected 3 and committed 9 instead: the draft KV holds + // a token the target never accepted, so the prefix cannot be reused. + assert_eq!(synced.plan(&[1, 2, 9, 10]), DraftSyncPlan::Reset); + } + + #[test] + fn recording_a_reset_replaces_the_whole_prefix() { + let mut synced = state(&[1, 2, 3]); + synced.record_reset(&[7, 8]); + + assert_eq!(synced.plan(&[7, 8, 9]), DraftSyncPlan::AlreadySynced); + assert_eq!(synced.plan(&[1, 2, 3]), DraftSyncPlan::Reset); + } + + #[test] + fn an_extend_plan_indexes_the_context_the_caller_slices() { + // The plan's range must address `context_tokens` directly, since + // sync_to_context slices the context with it. + let context = [1, 2, 3, 4, 5]; + let DraftSyncPlan::Extend { from, to } = state(&[1, 2]).plan(&context) else { + panic!("a synced prefix must extend"); + }; + + assert_eq!(&context[from..to], &[3, 4]); + } +} From 2d47d287602094d447cc8bf5a5f6e02b6613a9a9 Mon Sep 17 00:00:00 2001 From: Daniel Winter-Wijntjes Date: Wed, 26 Aug 2026 23:40:19 +1000 Subject: [PATCH 15/15] review: propose from the synced token, reject silent no-op fallback configs, record steps after they land MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The seed path proposes from context_tokens.last() — the token sync_to_context deliberately leaves unmaterialized — instead of from the loop's `current`. The two agree today, but only by an invariant maintained across the whole decode loop, and a slip would be silent KV corruption rather than an error. A debug_assert keeps it visible. - ngram_fallback = "draft" now requires speculative.draft_model and a pipeline depth above 1. Either missing meant the stage started cleanly and never took the fallback path, leaving the operator a zero counter indistinguishable from a proposer that never missed. - propose records a step after decode_step succeeds, so a failed decode cannot leave the sync state claiming a token the session lacks. Co-Authored-By: Claude Opus 5 --- .../inference/skippy/resolver/speculative.rs | 16 +++++ .../src/inference/skippy/resolver/tests.rs | 64 +++++++++++++++++++ .../src/frontend/embedded_generation.rs | 14 +++- .../src/frontend/generation/draft_runner.rs | 5 +- 4 files changed, 97 insertions(+), 2 deletions(-) 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 9082152b2a..556ea314ec 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs @@ -136,6 +136,7 @@ pub(super) fn resolve_speculative_config( model_config, global_config, package_generation, + has_draft_model: draft_model_path.is_some(), })?; // A standalone N-gram plan (no native MTP, no draft model) runs in the // legacy `ngram` mode so the embedded frontend derives its window from the @@ -201,6 +202,7 @@ struct DecodeResolutionInput<'a> { model_config: Option<&'a SpeculativeConfig>, global_config: Option<&'a SpeculativeConfig>, package_generation: Option<&'a PackageGenerationInfo>, + has_draft_model: bool, } #[allow(clippy::too_many_lines)] @@ -425,6 +427,20 @@ fn resolve_decode_config(input: DecodeResolutionInput<'_>) -> Result 1; the classic serial draft loop is authoritative at depth 1" + ); + } true } "none" | "" => false, 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 2d4e8d8592..b564aa9a7e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -1635,3 +1635,67 @@ verify_window_runahead_tokens = 0 "Some(0) at the model level must win over the inherited positive default" ); } + +#[test] +fn draft_fallback_without_a_draft_model_is_rejected_not_silently_ignored() { + use crate::plugin::SpeculativeConfig; + let config: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 2 +ngram_fallback = "draft" +"#, + ) + .expect("parse speculative config"); + let error = super::speculative::resolve_speculative_config( + Some(&config), + None, + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect_err("a draft fallback with no draft model must not resolve"); + + assert!( + error + .to_string() + .contains("requires speculative.draft_model"), + "unexpected error: {error}" + ); +} + +#[test] +fn draft_fallback_at_pipeline_depth_one_is_rejected() { + use crate::plugin::SpeculativeConfig; + let config: SpeculativeConfig = toml::from_str( + r#" +strategy = "ngram-suffix" +ngram_proposer = "suffix" +ngram_min = 5 +ngram_max = 32 +verify_window_pipeline_depth = 1 +ngram_fallback = "draft" +draft_model = "meshllm/draft-model" +draft_max_tokens = 4 +"#, + ) + .expect("parse speculative config"); + let error = super::speculative::resolve_speculative_config( + Some(&config), + None, + "meshllm/test-model", + std::path::Path::new("/nonexistent/test-model.gguf"), + None, + ) + .expect_err("depth 1 keeps the serial draft loop authoritative"); + + assert!( + error + .to_string() + .contains("verify_window_pipeline_depth > 1"), + "unexpected error: {error}" + ); +} diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index c6a29926e1..42c17cd191 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -920,6 +920,18 @@ impl StageOpenAiBackend { let draft = draft_guard .as_deref_mut() .expect("fallback requires a draft guard"); + // Propose from the token sync_to_context left + // unmaterialized rather than from `current`. The two + // agree today, but only by an invariant maintained + // across the whole decode loop; depending on it here + // would make a slip silent KV corruption instead of + // an error. + debug_assert_eq!(context_tokens.last(), Some(¤t)); + let Some(&propose_from) = context_tokens.last() else { + return Err(openai_backend_error(anyhow::anyhow!( + "draft fallback requires a non-empty context" + ))); + }; draft .sync_to_context(&context_tokens) .map_err(openai_backend_error)?; @@ -931,7 +943,7 @@ impl StageOpenAiBackend { .max(2) .min(native_mtp_remaining); let draft_tokens = draft - .propose(current, budget) + .propose(propose_from, budget) .map_err(openai_backend_error)?; speculative_stats.fallback_draft_proposals += 1; speculative_stats.fallback_draft_tokens += draft_tokens.len(); diff --git a/crates/skippy-server/src/frontend/generation/draft_runner.rs b/crates/skippy-server/src/frontend/generation/draft_runner.rs index b4e0156787..8cc8f678e1 100644 --- a/crates/skippy-server/src/frontend/generation/draft_runner.rs +++ b/crates/skippy-server/src/frontend/generation/draft_runner.rs @@ -179,11 +179,14 @@ impl DraftRunner { ) -> Result> { let mut tokens = Vec::with_capacity(max_tokens); for _ in 0..max_tokens { - self.synced.record_proposal_step(current); + // Record after the step succeeds: a failed decode must not leave + // the state claiming a token the session does not hold. + let stepped_from = current; current = self .session .decode_step(current) .context("draft decode step")?; + self.synced.record_proposal_step(stepped_from); tokens.push(current); } Ok(tokens)