use crate::binary_transport::BinaryStageExecutionOptions;
use crate::binary_transport::forwarded_stage_message_timed;
use crate::binary_transport::run_binary_stage_message;
use crate::binary_transport::stage_output_activation_capacity;
use crate::binary_transport::write_stage_message_conditioned;
use crate::frontend::NativeMtpDraft;
use crate::frontend::generation::ChainPrefixRestore;
use crate::frontend::generation::EmbeddedExecutionStats;
use crate::frontend::generation::EmbeddedFusedFirstDecode;
use crate::frontend::generation::EmbeddedStageZeroGeneration;
use crate::frontend::generation::MAX_EXACT_REPLAY_TOKENS;
use crate::frontend::generation::OpenAiGenerationIds;
use crate::frontend::generation::PhaseTimer;
use crate::frontend::generation::StageOpenAiBackend;
use crate::frontend::sampling_cache_key::{sampling_replay_safe, sampling_semantic_fingerprint};
use crate::frontend::util::openai_backend_error;
use crate::frontend::util::openai_io_error;
use crate::frontend::wire_messages::DecodeMessageArgs;
use crate::frontend::wire_messages::RestorePrefillDecodeMessageArgs;
use crate::frontend::wire_messages::embedded_decode_message;
use crate::frontend::wire_messages::embedded_prefix_cache_message;
use crate::frontend::wire_messages::embedded_restore_prefill_decode_message;
use crate::frontend::wire_messages::openai_stage_mask;
use crate::kv_integration::KvStageIntegration;
use crate::kv_integration::proactive_eviction_attrs;
use crate::kv_integration::proactive_eviction_error_kind;
use anyhow::Context;
use openai_frontend::OpenAiError;
use openai_frontend::OpenAiResult;
use serde_json::Value;
use serde_json::json;
use skippy_protocol::MessageBase;
use skippy_protocol::SCHEMA_VERSION;
use skippy_protocol::StageConfig;
use skippy_protocol::binary::StageReplyStats;
use skippy_protocol::binary::StageSamplingConfig as WireSamplingConfig;
use skippy_protocol::binary::StageWireMessage;
use skippy_protocol::binary::WireMessageKind;
use skippy_protocol::binary::WireReplyKind;
use skippy_protocol::binary::recv_reply;
use skippy_runtime::ActivationFrame;
use skippy_runtime::SamplingConfig;
use std::collections::BTreeMap;
use std::net::TcpStream;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(super) struct ChainPrefixCacheSavings {
pub(super) hit_stage_count: u32,
pub(super) stage0_activation_bytes_avoided: usize,
pub(super) interstage_activation_bytes_avoided_estimate: usize,
}
pub(super) fn chain_prefix_cache_savings(
stats: &StageReplyStats,
restored_tokens: usize,
activation_width: i32,
) -> ChainPrefixCacheSavings {
let hit_stage_count = prefix_cache_hit_stage_count(stats.kv_hit_stage_mask);
let stage0_activation_bytes_avoided =
estimated_activation_bytes(restored_tokens, activation_width);
let interstage_activation_bytes_avoided_estimate =
stage0_activation_bytes_avoided.saturating_mul(hit_stage_count.saturating_sub(1) as usize);
ChainPrefixCacheSavings {
hit_stage_count,
stage0_activation_bytes_avoided,
interstage_activation_bytes_avoided_estimate,
}
}
pub(super) fn insert_chain_prefix_cache_savings_attrs(
attrs: &mut BTreeMap<String, Value>,
savings: ChainPrefixCacheSavings,
) {
attrs.insert(
"skippy.kv.chain_cache_hit_stage_count".to_string(),
json!(savings.hit_stage_count),
);
attrs.insert(
"skippy.kv.chain_cache_stage0_activation_bytes_avoided".to_string(),
json!(savings.stage0_activation_bytes_avoided),
);
attrs.insert(
"skippy.kv.chain_cache_interstage_activation_bytes_avoided_estimate".to_string(),
json!(savings.interstage_activation_bytes_avoided_estimate),
);
}
fn prefix_cache_hit_stage_count(hit_stage_mask: i64) -> u32 {
if hit_stage_mask <= 0 {
return 0;
}
(hit_stage_mask as u64).count_ones()
}
fn estimated_activation_bytes(token_count: usize, activation_width: i32) -> usize {
let Ok(token_count) = i32::try_from(token_count) else {
return 0;
};
skippy_protocol::binary::activation_wire_bytes(token_count, activation_width).unwrap_or(0)
}
pub(super) fn request_allows_exact_replay(request: &EmbeddedStageZeroGeneration<'_>) -> bool {
request.draft.is_none() && sampling_replay_safe(request.sampling)
}
fn exact_replay_cache_key(
identity: &crate::kv_integration::PrefillKvIdentity,
sampling: &SamplingConfig,
chat_sampling_metadata: Option<&str>,
) -> String {
let fingerprint = sampling_semantic_fingerprint(sampling, chat_sampling_metadata);
format!("{}:replay:v2:{fingerprint}", identity.page_id)
}
fn record_replay_safe_first_token(
kv: &KvStageIntegration,
identity: &crate::kv_integration::PrefillKvIdentity,
predicted: i32,
sampling: &SamplingConfig,
chat_sampling_metadata: Option<&str>,
) -> bool {
if !sampling_replay_safe(sampling) {
return false;
}
let cache_key = exact_replay_cache_key(identity, sampling, chat_sampling_metadata);
kv.record_cached_first_token_with_key(&cache_key, identity, predicted)
}
fn lookup_replay_safe_first_token(
kv: &KvStageIntegration,
identity: &crate::kv_integration::PrefillKvIdentity,
sampling: &SamplingConfig,
chat_sampling_metadata: Option<&str>,
) -> Option<i32> {
if !sampling_replay_safe(sampling) {
return None;
}
let cache_key = exact_replay_cache_key(identity, sampling, chat_sampling_metadata);
kv.lookup_cached_first_token_with_key(&cache_key)
}
pub(super) fn stage0_prefill_record_identities(
kv: &KvStageIntegration,
config: &StageConfig,
base: &MessageBase,
token_start: u64,
token_ids: &[i32],
) -> Vec<crate::kv_integration::PrefillKvIdentity> {
kv.record_identities(config, base, token_start, token_ids)
}
pub(super) fn stage0_full_prefill_record_identities(
kv: &KvStageIntegration,
config: &StageConfig,
base: &MessageBase,
token_ids: &[i32],
) -> Vec<crate::kv_integration::PrefillKvIdentity> {
stage0_prefill_record_identities(kv, config, base, 0, token_ids)
}
pub(super) struct EmbeddedReplayCheckpointRecord<'a> {
pub(super) session_id: &'a str,
pub(super) ids: &'a OpenAiGenerationIds,
pub(super) prompt_token_ids: &'a [i32],
pub(super) checkpoint_token_ids: &'a [i32],
pub(super) predicted_tokens: &'a [i32],
pub(super) predicted: i32,
pub(super) sampling: &'a SamplingConfig,
pub(super) chat_sampling_metadata: Option<&'a str>,
}
impl StageOpenAiBackend {
pub(super) fn embedded_exact_checkpoint_boundary(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
prefill_tokens: &[i32],
restored_prefill_tokens: usize,
) -> Option<usize> {
let kv = self.kv.as_ref().filter(|kv| kv.payload_is_exact_state())?;
let shared_boundary = kv
.exact_shared_checkpoint_token_count(prefill_tokens.len() as u64)
.and_then(|token_count| usize::try_from(token_count).ok());
let chat_boundary = request
.recurrent_cache_prefix_token_ids
.filter(|tokens| {
!tokens.is_empty()
&& tokens.len() <= prefill_tokens.len()
&& prefill_tokens.starts_with(tokens)
})
.map(<[i32]>::len);
let valid = |boundary: &usize| {
*boundary > restored_prefill_tokens && *boundary < prefill_tokens.len()
};
shared_boundary
.filter(valid)
.or_else(|| chat_boundary.filter(valid))
}
pub(super) fn local_kv_message_base(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
) -> MessageBase {
MessageBase {
schema_version: SCHEMA_VERSION,
run_id: self.config.run_id.clone(),
request_id: ids.request_id_string(),
session_id: session_id.to_string(),
stage_id: "openai-local".to_string(),
stage_index: self.config.stage_index,
topology_id: self.config.topology_id.clone(),
model_id: Some(self.config.model_id.clone()),
tokenizer_id: None,
chat_template_id: ids.cache.namespace(),
seq: Some(ids.session_id),
}
}
pub(super) fn evict_embedded_stage0_resident_prefix(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
target_tokens: Option<u64>,
) -> OpenAiResult<()> {
let Some(kv) = self.kv.as_ref() else {
return Ok(());
};
let scheduler_kv = Arc::clone(kv);
let scheduler_session_id = session_id.to_string();
let eviction =
self.iteration_scheduler
.execute_runtime("embedded-prefix-evict", move |runtime| {
Ok((|| {
runtime
.ensure_session_active(&scheduler_session_id)
.context(
"activate embedded stage-0 session before resident-prefix eviction",
)?;
if let Some(target_tokens) = target_tokens {
scheduler_kv.evict_resident_prefix_for_tokens(
runtime,
&scheduler_session_id,
target_tokens,
)
} else {
scheduler_kv.evict_resident_prefix_for_decode_batch(
runtime,
&scheduler_session_id,
)
}
})())
})?;
let (status, error_kind, target_tokens, evicted_entries, evicted_tokens) = match &eviction {
Ok(eviction) => (
if eviction.evicted_entries > 0 {
"evicted"
} else {
"noop"
},
None,
eviction.target_tokens,
eviction.evicted_entries,
eviction.evicted_tokens,
),
Err(error) => (
"error",
Some(proactive_eviction_error_kind(error)),
target_tokens.unwrap_or_default(),
0,
0,
),
};
let mut attrs = self.openai_attrs(ids);
attrs.extend(proactive_eviction_attrs(
status,
error_kind,
target_tokens,
evicted_entries,
evicted_tokens,
));
if error_kind.is_some() || evicted_entries > 0 || evicted_tokens > 0 {
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
} else {
self.telemetry
.emit_debug("stage.openai_kv_record_decision", attrs);
}
eviction
.map(|_| ())
.map_err(|error| openai_backend_error(error.context("evict embedded stage-0 KV")))
}
pub(super) fn restore_embedded_stage0_prefill(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
token_start: u64,
token_ids: &[i32],
activation_width: i32,
) -> OpenAiResult<Option<ActivationFrame>> {
if token_start != 0 {
return Ok(None);
}
let Some(kv) = self.kv.as_ref() else {
return Ok(None);
};
let base = self.local_kv_message_base(session_id, ids);
let Some(activation) = kv.restore_resident_activation(
&self.config,
&base,
token_start,
token_ids,
activation_width,
) else {
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_activation_miss"),
);
attrs.insert("skippy.kv.token_start".to_string(), json!(token_start));
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
return Ok(None);
};
let scheduler_kv = Arc::clone(kv);
let scheduler_session_id = session_id.to_string();
let scheduler_identity = activation.identity.clone();
let scheduler_token_ids = token_ids.to_vec();
let restored = self.iteration_scheduler.execute_runtime(
"embedded-prefix-restore",
move |runtime| {
scheduler_kv
.restore_resident_prefix(
runtime,
&scheduler_session_id,
std::slice::from_ref(&scheduler_identity),
&scheduler_token_ids,
)
.map_err(openai_backend_error)
},
)?;
let Some(restored) = restored else {
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_activation_hit_kv_miss"),
);
attrs.insert(
"skippy.activation_cache.hit_page_id".to_string(),
json!(activation.page_id),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
return Ok(None);
};
if restored.token_count < token_ids.len() {
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_activation_hit_prefix_short"),
);
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(restored.token_count),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
return Ok(None);
}
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_activation_resident_hit"),
);
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(restored.token_count),
);
attrs.insert(
"skippy.activation_cache.hit_page_id".to_string(),
json!(activation.page_id),
);
attrs.insert(
"skippy.activation_cache.payload_bytes".to_string(),
json!(activation.payload_bytes),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
Ok(Some(activation.frame))
}
pub(super) fn record_embedded_stage0_prefill(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
token_start: u64,
token_ids: &[i32],
activation_width: i32,
output: &ActivationFrame,
) -> OpenAiResult<()> {
if token_start != 0 {
return Ok(());
}
let Some(kv) = self.kv.as_ref() else {
return Ok(());
};
let base = self.local_kv_message_base(session_id, ids);
let identities =
stage0_prefill_record_identities(kv, &self.config, &base, token_start, token_ids);
let record_candidate_count = identities.len();
let scheduler_kv = Arc::clone(kv);
let scheduler_session_id = session_id.to_string();
let scheduler_identities = identities.clone();
let scheduler_token_ids = token_ids.to_vec();
let resident_records =
self.iteration_scheduler
.execute_runtime("embedded-prefix-record", move |runtime| {
scheduler_identities
.iter()
.map(|identity| {
let token_count = identity
.identity
.token_count
.try_into()
.unwrap_or(usize::MAX)
.min(scheduler_token_ids.len());
scheduler_kv
.record_resident_prefix(
runtime,
&scheduler_session_id,
identity,
&scheduler_token_ids[..token_count],
)
.map_err(openai_backend_error)
})
.collect::<OpenAiResult<Vec<_>>>()
})?;
let activation_records = kv.record_resident_activation(
&self.config,
&base,
token_start,
token_ids,
activation_width,
output,
);
let mut recorded_any = false;
for record in resident_records.into_iter().flatten() {
recorded_any = true;
let mut attrs = self.openai_attrs(ids);
attrs.insert("skippy.kv.decision".to_string(), json!("stage0_record"));
attrs.insert(
"skippy.kv.record_candidates".to_string(),
json!(record_candidate_count),
);
attrs.insert("skippy.kv.token_start".to_string(), json!(token_start));
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
attrs.insert(
"skippy.kv.recorded_page_id".to_string(),
json!(record.page_id),
);
attrs.insert(
"skippy.kv.recorded_tokens".to_string(),
json!(record.token_count),
);
attrs.insert(
"skippy.kv.resident_seq_id".to_string(),
json!(record.seq_id),
);
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
}
for record in &activation_records {
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_activation_record"),
);
attrs.insert(
"skippy.kv.record_candidates".to_string(),
json!(record_candidate_count),
);
attrs.insert("skippy.kv.token_start".to_string(), json!(token_start));
attrs.insert(
"skippy.kv.token_count".to_string(),
json!(record.token_count),
);
attrs.insert(
"skippy.activation_cache.payload_bytes".to_string(),
json!(record.payload_bytes),
);
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
}
if activation_records.is_empty() && !recorded_any {
let mut attrs = self.openai_attrs(ids);
attrs.insert("skippy.kv.decision".to_string(), json!("stage0_record"));
attrs.insert(
"skippy.kv.record_candidates".to_string(),
json!(record_candidate_count),
);
attrs.insert("skippy.kv.token_start".to_string(), json!(token_start));
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
}
Ok(())
}
pub(super) fn record_embedded_stage0_full_prefill(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
token_ids: &[i32],
) -> OpenAiResult<bool> {
let Some(kv) = self.kv.as_ref() else {
return Ok(false);
};
if token_ids.is_empty() || !kv.should_record() {
return Ok(false);
}
let base = self.local_kv_message_base(session_id, ids);
let identities = stage0_full_prefill_record_identities(kv, &self.config, &base, token_ids);
let record_candidate_count = identities.len();
let scheduler_kv = Arc::clone(kv);
let scheduler_session_id = session_id.to_string();
let scheduler_token_ids = token_ids.to_vec();
let records = self.iteration_scheduler.execute_runtime(
"embedded-full-prefill-record",
move |runtime| {
identities
.iter()
.map(|identity| {
let token_count = identity
.identity
.token_count
.try_into()
.unwrap_or(usize::MAX)
.min(scheduler_token_ids.len());
scheduler_kv
.record_resident_prefix(
runtime,
&scheduler_session_id,
identity,
&scheduler_token_ids[..token_count],
)
.map_err(openai_backend_error)
})
.collect::<OpenAiResult<Vec<_>>>()
},
)?;
let exact_record_queued = kv.payload_is_exact_state()
&& self.enqueue_exact_state_record_at_tokens(
session_id,
ids,
token_ids.to_vec(),
"embedded_full_prefill_state",
);
let mut recorded_any = exact_record_queued;
for record in records.into_iter().flatten() {
recorded_any = true;
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_full_prefill_record"),
);
attrs.insert(
"skippy.kv.record_candidates".to_string(),
json!(record_candidate_count),
);
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
attrs.insert(
"skippy.kv.recorded_page_id".to_string(),
json!(record.page_id),
);
attrs.insert(
"skippy.kv.recorded_tokens".to_string(),
json!(record.token_count),
);
attrs.insert(
"skippy.kv.resident_seq_id".to_string(),
json!(record.seq_id),
);
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
}
if !recorded_any {
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_full_prefill_record"),
);
attrs.insert(
"skippy.kv.record_candidates".to_string(),
json!(record_candidate_count),
);
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
}
Ok(recorded_any)
}
pub(super) fn record_embedded_stage0_exact_checkpoint(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
checkpoint_tokens: &[i32],
) -> OpenAiResult<bool> {
let Some(kv) = self.kv.as_ref() else {
return Ok(false);
};
if checkpoint_tokens.is_empty() || !kv.payload_is_exact_state() || !kv.should_record() {
return Ok(false);
}
let scheduler_backend = self.clone();
let scheduler_session_id = session_id.to_string();
let scheduler_ids = ids.clone();
let scheduler_checkpoint_tokens = checkpoint_tokens.to_vec();
self.iteration_scheduler.execute_runtime(
"embedded-shared-exact-state-checkpoint",
move |runtime| {
Ok(scheduler_backend.record_exact_state_at_tokens(
runtime,
&scheduler_session_id,
&scheduler_ids,
&scheduler_checkpoint_tokens,
"embedded_shared_checkpoint",
))
},
)
}
pub(super) fn record_embedded_stage0_full_prompt_first_token(
&self,
session_id: &str,
ids: &OpenAiGenerationIds,
token_ids: &[i32],
predicted: i32,
sampling: &SamplingConfig,
chat_sampling_metadata: Option<&str>,
) -> OpenAiResult<bool> {
let Some(kv) = self.kv.as_ref() else {
return Ok(false);
};
if token_ids.is_empty() || !kv.should_record() {
return Ok(false);
}
let base = self.local_kv_message_base(session_id, ids);
let identity = kv.prefill_identity(&self.config, &base, 0, token_ids);
let recorded_state =
self.record_embedded_stage0_full_prefill(session_id, ids, token_ids)?;
let recorded_token = record_replay_safe_first_token(
kv,
&identity,
predicted,
sampling,
chat_sampling_metadata,
);
let mut attrs = self.openai_attrs(ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_full_prompt_first_token_record"),
);
attrs.insert("skippy.kv.token_count".to_string(), json!(token_ids.len()));
attrs.insert("skippy.kv.predicted_token".to_string(), json!(predicted));
attrs.insert(
"skippy.kv.recorded_page_id".to_string(),
json!(identity.page_id),
);
attrs.insert(
"skippy.kv.recorded_state".to_string(),
json!(recorded_state),
);
attrs.insert(
"skippy.kv.recorded_first_token".to_string(),
json!(recorded_token),
);
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
Ok(recorded_state || recorded_token)
}
pub(super) fn record_embedded_stage0_replay_checkpoint(
&self,
record: EmbeddedReplayCheckpointRecord<'_>,
) -> OpenAiResult<bool> {
let Some(kv) = self.kv.as_ref() else {
return Ok(false);
};
if record.prompt_token_ids.is_empty()
|| record.checkpoint_token_ids.is_empty()
|| record.predicted_tokens.len() >= MAX_EXACT_REPLAY_TOKENS
|| !kv.should_record()
{
return Ok(false);
}
let base = self.local_kv_message_base(record.session_id, record.ids);
let prompt_identity = kv.prefill_identity(&self.config, &base, 0, record.prompt_token_ids);
let replay_cache_key = exact_replay_cache_key(
&prompt_identity,
record.sampling,
record.chat_sampling_metadata,
);
let recorded_state = self.record_embedded_stage0_full_prefill(
record.session_id,
record.ids,
record.checkpoint_token_ids,
)?;
let recorded_replay = kv.record_cached_replay_tokens(
&replay_cache_key,
&prompt_identity,
record.predicted_tokens,
record.predicted,
MAX_EXACT_REPLAY_TOKENS,
);
let mut attrs = self.openai_attrs(record.ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("stage0_exact_replay_record"),
);
attrs.insert(
"skippy.kv.prompt_token_count".to_string(),
json!(record.prompt_token_ids.len()),
);
attrs.insert(
"skippy.kv.checkpoint_token_count".to_string(),
json!(record.checkpoint_token_ids.len()),
);
attrs.insert(
"skippy.kv.replay_token_count".to_string(),
json!(recorded_replay.unwrap_or(record.predicted_tokens.len())),
);
attrs.insert(
"skippy.kv.predicted_token".to_string(),
json!(record.predicted),
);
attrs.insert(
"skippy.kv.recorded_page_id".to_string(),
json!(prompt_identity.page_id),
);
attrs.insert(
"skippy.kv.recorded_state".to_string(),
json!(recorded_state),
);
attrs.insert(
"skippy.kv.recorded_replay".to_string(),
json!(recorded_replay.is_some()),
);
self.telemetry
.emit("stage.openai_kv_record_decision", attrs);
Ok(recorded_state || recorded_replay.is_some())
}
pub(super) fn try_restore_embedded_split_exact_replay(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
session_key: &str,
downstream: &mut TcpStream,
) -> OpenAiResult<Option<EmbeddedFusedFirstDecode>> {
let Some(kv) = self.kv.as_ref() else {
return Ok(None);
};
if request.prompt_token_ids.is_empty()
|| !kv.should_lookup()
|| !request_allows_exact_replay(request)
{
return Ok(None);
}
let timer = PhaseTimer::start();
let base = self.local_kv_message_base(session_key, request.ids);
let prompt_identity =
kv.prefill_identity(request.config, &base, 0, request.prompt_token_ids);
let replay_cache_key = exact_replay_cache_key(
&prompt_identity,
request.sampling,
request.chat_sampling_metadata,
);
let replay_tokens =
kv.lookup_cached_replay_tokens(&replay_cache_key, request.max_tokens as usize);
if replay_tokens.len() < 2 {
return Ok(None);
}
for replay_len in (2..=replay_tokens.len()).rev() {
let mut checkpoint_tokens = request.prompt_token_ids.to_vec();
checkpoint_tokens.extend_from_slice(&replay_tokens[..replay_len - 1]);
let Some(restore) = self.try_restore_embedded_split_prefill(
request,
session_key,
downstream,
&checkpoint_tokens,
)?
else {
continue;
};
if exact_replay_restore_is_partial(restore.restored_tokens, checkpoint_tokens.len()) {
self.drop_embedded_split_restore(request, session_key, downstream);
continue;
}
let replay = replay_tokens[..replay_len].to_vec();
let mut attrs = self.openai_attrs(request.ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("chain_exact_replay_hit"),
);
attrs.insert(
"skippy.kv.prompt_token_count".to_string(),
json!(request.prompt_token_ids.len()),
);
attrs.insert(
"skippy.kv.checkpoint_token_count".to_string(),
json!(checkpoint_tokens.len()),
);
attrs.insert(
"skippy.kv.replay_token_count".to_string(),
json!(replay.len()),
);
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(restore.restored_tokens),
);
attrs.insert(
"skippy.kv.hit_page_id".to_string(),
json!(prompt_identity.page_id),
);
attrs.insert(
"skippy.kv.lookup_hits".to_string(),
json!(restore.stats.kv_lookup_hits),
);
attrs.insert(
"skippy.kv.hit_stage_mask".to_string(),
json!(restore.stats.kv_hit_stage_mask),
);
insert_chain_prefix_cache_savings_attrs(
&mut attrs,
chain_prefix_cache_savings(
&restore.stats,
checkpoint_tokens.len(),
request.activation_width,
),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
return Ok(Some(EmbeddedFusedFirstDecode {
predicted: *replay.last().expect("checked replay length"),
predicted_tokens: replay,
native_mtp_draft: None,
reply_stats: restore.stats,
execution: EmbeddedExecutionStats::default(),
elapsed_ms: timer.elapsed_ms(),
token_phase: "exact-replay-cache",
message_kind: "TryRestorePrefill",
}));
}
Ok(None)
}
pub(super) fn try_restore_embedded_split_full_prompt_first_token(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
session_key: &str,
downstream: &mut TcpStream,
) -> OpenAiResult<Option<EmbeddedFusedFirstDecode>> {
let Some(kv) = self.kv.as_ref() else {
return Ok(None);
};
if request.prompt_token_ids.is_empty() || !kv.should_lookup() {
return Ok(None);
}
let timer = PhaseTimer::start();
let base = self.local_kv_message_base(session_key, request.ids);
let identity = kv.prefill_identity(request.config, &base, 0, request.prompt_token_ids);
let Some(predicted) = lookup_replay_safe_first_token(
kv,
&identity,
request.sampling,
request.chat_sampling_metadata,
) else {
return Ok(None);
};
let Some(restore) = self.try_restore_embedded_split_prefill(
request,
session_key,
downstream,
request.prompt_token_ids,
)?
else {
return Ok(None);
};
if restore.restored_tokens < request.prompt_token_ids.len() {
self.drop_embedded_split_restore(request, session_key, downstream);
return Ok(None);
}
let mut attrs = self.openai_attrs(request.ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("chain_full_prompt_first_token_hit"),
);
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(restore.restored_tokens),
);
attrs.insert("skippy.kv.predicted_token".to_string(), json!(predicted));
attrs.insert("skippy.kv.hit_page_id".to_string(), json!(identity.page_id));
attrs.insert(
"skippy.kv.lookup_hits".to_string(),
json!(restore.stats.kv_lookup_hits),
);
attrs.insert(
"skippy.kv.hit_stage_mask".to_string(),
json!(restore.stats.kv_hit_stage_mask),
);
insert_chain_prefix_cache_savings_attrs(
&mut attrs,
chain_prefix_cache_savings(
&restore.stats,
restore.restored_tokens,
request.activation_width,
),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
Ok(Some(EmbeddedFusedFirstDecode {
predicted,
predicted_tokens: vec![predicted],
native_mtp_draft: None,
reply_stats: restore.stats,
execution: EmbeddedExecutionStats::default(),
elapsed_ms: timer.elapsed_ms(),
token_phase: "full-prompt-cache",
message_kind: "TryRestorePrefill",
}))
}
pub(super) fn try_restore_embedded_split_prefill(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
session_key: &str,
downstream: &mut TcpStream,
prefill_tokens: &[i32],
) -> OpenAiResult<Option<ChainPrefixRestore>> {
let Some(kv) = self.kv.as_ref() else {
return Ok(None);
};
if prefill_tokens.is_empty() || !kv.should_lookup() {
return Ok(None);
}
let base = self.local_kv_message_base(session_key, request.ids);
let identities = kv.lookup_identities(request.config, &base, 0, prefill_tokens);
let mut restore_stats = StageReplyStats::default();
let scheduler_kv = Arc::clone(kv);
let scheduler_session_key = session_key.to_string();
let scheduler_prefill_tokens = prefill_tokens.to_vec();
let local_restore = self.iteration_scheduler.execute_runtime(
"embedded-split-prefix-restore",
move |runtime| match scheduler_kv
.restore_exact_state(runtime, &scheduler_session_key, &identities)
.map_err(openai_backend_error)?
{
Some(restored) => Ok(Some(restored.token_count)),
None => scheduler_kv
.restore_resident_prefix(
runtime,
&scheduler_session_key,
&identities,
&scheduler_prefill_tokens,
)
.map_err(openai_backend_error)
.map(|restored| restored.map(|restored| restored.token_count)),
},
)?;
let Some(local_restore) = local_restore else {
return Ok(None);
};
if local_restore == 0 {
return Ok(None);
}
let restored_tokens = local_restore.min(prefill_tokens.len());
restore_stats.kv_lookup_hits += 1;
restore_stats.kv_imported_pages += 1;
restore_stats.kv_imported_tokens += restored_tokens as i64;
restore_stats.kv_hit_stage_mask |= openai_stage_mask(request.config.stage_index);
let restore = embedded_prefix_cache_message(
WireMessageKind::TryRestorePrefill,
&prefill_tokens[..restored_tokens],
request.ids.request_id,
request.ids.session_id,
)?;
write_stage_message_conditioned(
&mut *downstream,
&restore,
request.downstream_wire_condition,
)
.map_err(openai_io_error)?;
let downstream_restore = recv_reply(&mut *downstream).map_err(openai_io_error)?;
if downstream_restore.kind != WireReplyKind::Ack {
return Err(OpenAiError::backend(format!(
"expected prefix try-restore ACK from downstream, got {:?}",
downstream_restore.kind
)));
}
restore_stats.merge(downstream_restore.stats);
if restore_stats.kv_lookup_errors > 0
|| restore_stats.kv_lookup_misses > 0
|| downstream_restore.stats.kv_lookup_hits == 0
{
self.drop_embedded_split_restore(request, session_key, downstream);
return Ok(None);
}
let mut attrs = self.openai_attrs(request.ids);
attrs.insert("skippy.kv.decision".to_string(), json!("chain_restore_hit"));
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(restored_tokens),
);
attrs.insert(
"skippy.kv.suffix_prefill_tokens".to_string(),
json!(prefill_tokens.len().saturating_sub(restored_tokens)),
);
attrs.insert(
"skippy.kv.lookup_hits".to_string(),
json!(restore_stats.kv_lookup_hits),
);
attrs.insert(
"skippy.kv.hit_stage_mask".to_string(),
json!(restore_stats.kv_hit_stage_mask),
);
insert_chain_prefix_cache_savings_attrs(
&mut attrs,
chain_prefix_cache_savings(&restore_stats, restored_tokens, request.activation_width),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
Ok(Some(ChainPrefixRestore {
restored_tokens,
stats: restore_stats,
}))
}
pub(super) fn try_restore_embedded_split_prefill_and_decode(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
session_key: &str,
downstream: &mut TcpStream,
prefill_tokens: &[i32],
current: i32,
wire_sampling: Option<WireSamplingConfig>,
) -> OpenAiResult<Option<EmbeddedFusedFirstDecode>> {
let Some(kv) = self.kv.as_ref() else {
return Ok(None);
};
if prefill_tokens.is_empty() || !kv.should_lookup() {
return Ok(None);
}
let timer = PhaseTimer::start();
let base = self.local_kv_message_base(session_key, request.ids);
let identity = kv.prefill_identity(request.config, &base, 0, prefill_tokens);
let mut reply_stats = StageReplyStats::default();
let stage0_timer = PhaseTimer::start();
let decode_message = embedded_decode_message(DecodeMessageArgs {
request_id: request.ids.request_id,
session_id: request.ids.session_id,
prompt_token_count: request.prompt_token_ids.len(),
pos_start: prefill_tokens.len(),
decode_step: 0,
current,
sampling: wire_sampling.clone(),
})?;
let output_capacity = stage_output_activation_capacity(
request.config,
decode_message.token_count,
request.activation_width,
)
.map_err(openai_backend_error)?;
let scheduler_kv = Arc::clone(kv);
let scheduler_session_key = session_key.to_string();
let scheduler_prefill_tokens = prefill_tokens.to_vec();
let scheduler_sampling = request.sampling.clone();
let scheduler_metadata = request.chat_sampling_metadata.map(str::to_string);
let prompt_token_count = request.prompt_token_ids.len();
let native_mtp_enabled = request.native_mtp_enabled;
let native_mtp_max_tokens = request.speculative.native_mtp.max_draft_tokens;
let scheduler_decode_message = decode_message.clone();
let outcome = self.iteration_scheduler.execute_runtime_timed(
"embedded-fused-prefix-decode",
move |runtime| {
let local_restore = scheduler_kv
.restore_resident_prefix(
runtime,
&scheduler_session_key,
std::slice::from_ref(&identity),
&scheduler_prefill_tokens,
)
.map_err(openai_backend_error)?;
let Some(local_restore) = local_restore else {
return Ok(None);
};
if local_restore.token_count < scheduler_prefill_tokens.len() {
let _ = runtime.drop_session_timed(&scheduler_session_key);
return Ok(None);
}
if let Some(metadata) = scheduler_metadata.as_deref() {
runtime
.configure_chat_sampling(
&scheduler_session_key,
metadata,
prompt_token_count as u64,
scheduler_sampling.enabled.then_some(&scheduler_sampling),
)
.map_err(openai_backend_error)?;
}
let output = run_binary_stage_message(
runtime,
&scheduler_session_key,
&scheduler_decode_message,
&[current],
None,
BinaryStageExecutionOptions::new(false, output_capacity, native_mtp_enabled)
.with_native_mtp_max_tokens(native_mtp_max_tokens),
)
.map_err(openai_backend_error)?
.2;
Ok(Some((local_restore.token_count, output)))
},
)?;
let token_runtime_lock_wait_ms = outcome.runtime_lock_wait_ms;
let token_runtime_lock_hold_ms = outcome.runtime_lock_hold_ms;
let Some((restored_token_count, output)) = outcome.value else {
return Ok(None);
};
reply_stats.kv_lookup_hits += 1;
reply_stats.kv_imported_pages += 1;
reply_stats.kv_imported_tokens += restored_token_count as i64;
reply_stats.kv_hit_stage_mask |= openai_stage_mask(request.config.stage_index);
let stage0_compute_ms = stage0_timer.elapsed_ms();
let fused_message =
embedded_restore_prefill_decode_message(RestorePrefillDecodeMessageArgs {
request_id: request.ids.request_id,
session_id: request.ids.session_id,
prompt_token_count: request.prompt_token_ids.len(),
pos_start: prefill_tokens.len(),
decode_step: 0,
prefix_tokens: prefill_tokens,
current,
sampling: wire_sampling,
chat_sampling_metadata: request.chat_sampling_metadata,
})?;
let forwarded = forwarded_stage_message_timed(
request.config,
&fused_message,
&output,
request.activation_width,
)
.map_err(openai_backend_error)?;
let write_timer = PhaseTimer::start();
write_stage_message_conditioned(
&mut *downstream,
&forwarded.message,
request.downstream_wire_condition,
)
.map_err(openai_io_error)?;
let forward_write_ms = write_timer.elapsed_ms();
let wait_timer = PhaseTimer::start();
let downstream_reply = super::embedded_execution::receive_embedded_stage_reply_one_of(
downstream,
request.prediction_return.as_ref(),
&[WireReplyKind::PredictedToken, WireReplyKind::Ack],
)?;
let downstream_wait_ms = wait_timer.elapsed_ms();
let downstream_missed = downstream_reply.kind != WireReplyKind::PredictedToken
|| downstream_reply.stats.kv_lookup_errors > 0
|| downstream_reply.stats.kv_lookup_misses > 0
|| downstream_reply.stats.kv_lookup_hits == 0;
reply_stats.merge(downstream_reply.stats);
if downstream_missed {
self.drop_embedded_split_restore(request, session_key, downstream);
return Ok(None);
}
let mut attrs = self.openai_attrs(request.ids);
attrs.insert(
"skippy.kv.decision".to_string(),
json!("chain_restore_decode_hit"),
);
attrs.insert(
"skippy.kv.restored_tokens".to_string(),
json!(prefill_tokens.len()),
);
attrs.insert(
"skippy.kv.lookup_hits".to_string(),
json!(reply_stats.kv_lookup_hits),
);
attrs.insert(
"skippy.kv.hit_stage_mask".to_string(),
json!(reply_stats.kv_hit_stage_mask),
);
insert_chain_prefix_cache_savings_attrs(
&mut attrs,
chain_prefix_cache_savings(
&reply_stats,
prefill_tokens.len(),
request.activation_width,
),
);
self.telemetry
.emit("stage.openai_kv_lookup_decision", attrs);
self.record_embedded_stage0_full_prompt_first_token(
session_key,
request.ids,
request.prompt_token_ids,
downstream_reply.predicted,
request.sampling,
request.chat_sampling_metadata,
)?;
Ok(Some(EmbeddedFusedFirstDecode {
predicted: downstream_reply.predicted,
predicted_tokens: vec![downstream_reply.predicted],
native_mtp_draft: downstream_reply
.native_mtp_draft
.clone()
.map(NativeMtpDraft::from_stage_draft),
reply_stats,
execution: EmbeddedExecutionStats {
stage0_compute_ms,
runtime_lock_wait_ms: token_runtime_lock_wait_ms,
runtime_lock_hold_ms: token_runtime_lock_hold_ms,
activation_encode_ms: forwarded.activation_encode_ms,
output_activation_bytes: output.payload.len(),
forward_activation_bytes: forwarded.message.activation.len(),
forward_write_ms,
downstream_wait_ms,
},
elapsed_ms: timer.elapsed_ms(),
token_phase: "fused-restore",
message_kind: "TryRestorePrefillDecode",
}))
}
pub(super) fn drop_embedded_split_restore(
&self,
request: &EmbeddedStageZeroGeneration<'_>,
session_key: &str,
downstream: &mut TcpStream,
) {
let scheduler_session_key = session_key.to_string();
let _ = self.iteration_scheduler.execute_runtime(
"embedded-split-restore-drop",
move |runtime| {
runtime
.drop_session_timed(&scheduler_session_key)
.map(|_| ())
.map_err(openai_backend_error)
},
);
let stop =
StageWireMessage::stop_with_identity(request.ids.request_id, request.ids.session_id);
if write_stage_message_conditioned(
&mut *downstream,
&stop,
request.downstream_wire_condition,
)
.is_ok()
{
let _ = recv_reply(&mut *downstream);
}
}
}
fn exact_replay_restore_is_partial(restored_tokens: usize, checkpoint_tokens: usize) -> bool {
restored_tokens < checkpoint_tokens
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chain_prefix_cache_savings_counts_confirmed_stage_hits() {
let stats = StageReplyStats {
kv_hit_stage_mask: openai_stage_mask(0)
| openai_stage_mask(1)
| openai_stage_mask(2)
| openai_stage_mask(3),
..Default::default()
};
let savings = chain_prefix_cache_savings(&stats, 256, 5120);
assert_eq!(savings.hit_stage_count, 4);
assert_eq!(savings.stage0_activation_bytes_avoided, 5_242_880);
assert_eq!(
savings.interstage_activation_bytes_avoided_estimate,
15_728_640
);
}
#[test]
fn chain_prefix_cache_savings_uses_f32_wire_size() {
let stats = StageReplyStats {
kv_hit_stage_mask: openai_stage_mask(0) | openai_stage_mask(1),
..Default::default()
};
let savings = chain_prefix_cache_savings(&stats, 256, 5120);
assert_eq!(savings.stage0_activation_bytes_avoided, 5_242_880);
assert_eq!(
savings.interstage_activation_bytes_avoided_estimate,
5_242_880
);
}
#[test]
fn exact_replay_rejects_a_shorter_restored_checkpoint() {
assert!(exact_replay_restore_is_partial(44_466, 44_467));
assert!(!exact_replay_restore_is_partial(44_467, 44_467));
}
#[test]
fn exact_replay_key_includes_ignore_eos_sampling_control() {
let identity = crate::kv_integration::PrefillKvIdentity {
identity: crate::kv_proto::PageIdentity {
token_count: 4,
..Default::default()
},
page_id: "page".to_string(),
namespace: "stage".to_string(),
token_ids: vec![1, 2, 3, 4],
};
let baseline = SamplingConfig {
enabled: true,
temperature: 0.0,
ignore_eos: false,
..Default::default()
};
let ignore_eos = SamplingConfig {
ignore_eos: true,
..baseline.clone()
};
assert_ne!(baseline, ignore_eos);
assert_ne!(
exact_replay_cache_key(&identity, &baseline, None),
exact_replay_cache_key(&identity, &ignore_eos, None),
"different greedy sampling semantics must not collide"
);
}
#[test]
fn first_token_cache_is_partitioned_and_skips_stochastic_requests() {
let config = StageConfig {
model_id: "hugging-quants/Llama-3.2-1B-Instruct-GGUF:Q4_K_M".to_string(),
stage_id: "stage-0".to_string(),
layer_end: 1,
ctx_size: 256,
lane_count: 1,
kv_cache: Some(skippy_protocol::StageKvCacheConfig {
mode: skippy_protocol::StageKvCacheMode::LookupRecord,
payload: skippy_protocol::StageKvCachePayload::ResidentKv,
max_entries: 8,
max_bytes: 0,
min_tokens: 1,
shared_prefix_stride_tokens: 1,
shared_prefix_record_limit: 1,
}),
..Default::default()
};
let cache = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
.unwrap()
.expect("resident cache");
let identity = crate::kv_integration::PrefillKvIdentity {
identity: crate::kv_proto::PageIdentity {
token_count: 4,
..Default::default()
},
page_id: "page".to_string(),
namespace: "stage".to_string(),
token_ids: vec![1, 2, 3, 4],
};
let first_sampling = SamplingConfig {
enabled: true,
seed: 1,
temperature: 0.0,
..Default::default()
};
let second_sampling = SamplingConfig {
ignore_eos: true,
..first_sampling.clone()
};
let stochastic_sampling = SamplingConfig {
temperature: 0.8,
..first_sampling.clone()
};
let first_key = exact_replay_cache_key(&identity, &first_sampling, None);
let second_key = exact_replay_cache_key(&identity, &second_sampling, None);
assert_ne!(
first_key, second_key,
"sampling controls must partition sampled first-token outputs"
);
assert!(sampling_replay_safe(&SamplingConfig::default()));
assert!(sampling_replay_safe(&first_sampling));
assert!(!sampling_replay_safe(&stochastic_sampling));
assert!(cache.record_cached_first_token_with_key(&first_key, &identity, 123));
assert_eq!(
cache.lookup_cached_first_token_with_key(&first_key),
Some(123),
"the same deterministic sampling semantics should reuse the token"
);
assert_eq!(
cache.lookup_cached_first_token_with_key(&second_key),
None,
"different sampling semantics must not reuse the token"
);
assert_eq!(
cache.lookup_cached_first_token_with_key(&exact_replay_cache_key(
&identity,
&stochastic_sampling,
None,
)),
None,
"stochastic requests must not reuse a deterministic first token"
);
}
#[test]
fn first_token_record_and_lookup_reject_rng_backed_zero_temperature_chains() {
let config = StageConfig {
model_id: "hugging-quants/Llama-3.2-1B-Instruct-GGUF:Q4_K_M".to_string(),
stage_id: "stage-0".to_string(),
layer_end: 1,
ctx_size: 256,
lane_count: 1,
kv_cache: Some(skippy_protocol::StageKvCacheConfig {
mode: skippy_protocol::StageKvCacheMode::LookupRecord,
payload: skippy_protocol::StageKvCachePayload::ResidentKv,
max_entries: 8,
max_bytes: 0,
min_tokens: 1,
shared_prefix_stride_tokens: 1,
shared_prefix_record_limit: 1,
}),
..Default::default()
};
let cache = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
.unwrap()
.expect("resident cache");
let identity = crate::kv_integration::PrefillKvIdentity {
identity: crate::kv_proto::PageIdentity {
token_count: 4,
..Default::default()
},
page_id: "page".to_string(),
namespace: "stage".to_string(),
token_ids: vec![1, 2, 3, 4],
};
let greedy = SamplingConfig {
enabled: true,
temperature: 0.0,
..Default::default()
};
let unsafe_configs = [
SamplingConfig {
mirostat_mode: 1,
..greedy.clone()
},
SamplingConfig {
samplers: vec!["top_k".to_string(), "top_p".to_string()],
..greedy.clone()
},
SamplingConfig {
xtc: skippy_runtime::XtcSamplingConfig {
probability: 0.5,
..greedy.xtc.clone()
},
..greedy.clone()
},
SamplingConfig {
dynatemp_range: 0.5,
..greedy
},
];
for (index, sampling) in unsafe_configs.iter().enumerate() {
let raw_key = exact_replay_cache_key(&identity, sampling, None);
assert!(!record_replay_safe_first_token(
&cache,
&identity,
100 + index as i32,
sampling,
None,
));
assert_eq!(
cache.lookup_cached_first_token_with_key(&raw_key),
None,
"an RNG-backed token must not be recorded"
);
assert!(cache.record_cached_first_token_with_key(
&raw_key,
&identity,
200 + index as i32,
));
assert_eq!(
lookup_replay_safe_first_token(&cache, &identity, sampling, None),
None,
"an RNG-backed request must not look up even a pre-existing token"
);
}
}
}