diff --git a/crates/core/src/engine/common/sequence.rs b/crates/core/src/engine/common/sequence.rs index b20813a2e3..5ae4076775 100644 --- a/crates/core/src/engine/common/sequence.rs +++ b/crates/core/src/engine/common/sequence.rs @@ -5,6 +5,22 @@ use crate::engine::common::hashing::{ BlockHash, SequenceHash, XXH3_SEED, compute_block_hash_for_tokens, compute_next_sequence_hash, }; use rand::random; +use uuid::Uuid; +use xxhash_rust::xxh3::xxh3_64; + +/// Synthetic token for output position `output_ordinal` of a request that has +/// no planned token there. +/// +/// The token depends only on the request ID and position, so replaying a +/// request under the same ID reproduces its output tokens and KV block hashes, +/// while distinct requests receive independent streams. Hashing a fixed +/// little-endian encoding keeps values stable across platforms and toolchains. +pub(crate) fn synthetic_output_token(request_id: Uuid, output_ordinal: usize) -> u32 { + let mut bytes = [0; 24]; + bytes[..16].copy_from_slice(request_id.as_bytes()); + bytes[16..].copy_from_slice(&(output_ordinal as u64).to_le_bytes()); + xxh3_64(&bytes) as u32 +} #[derive(Debug)] struct FlatTokens { @@ -103,6 +119,7 @@ impl BlockIdentity { /// positional lineage, or allocation bookkeeping. #[derive(Debug)] pub(crate) struct RequestSequence { + request_id: Uuid, tokens: FlatTokens, block_size: usize, max_output_tokens: usize, @@ -117,6 +134,7 @@ pub(crate) struct RequestSequence { impl RequestSequence { #[allow(clippy::too_many_arguments)] pub(crate) fn new( + request_id: Uuid, tokens: Vec, max_output_tokens: usize, output_capacity_hint: usize, @@ -159,6 +177,7 @@ impl RequestSequence { } let sequence = Self { + request_id, tokens: FlatTokens::new(tokens, output_capacity_hint, block_size, emit_token_ids), block_size, max_output_tokens, @@ -203,6 +222,7 @@ impl RequestSequence { /// Append the next planned or synthetic token. /// + /// Positions without a planned token derive one from the request ID. /// Returns the token and whether this append opened a new partial block. pub(crate) fn generate_token(&mut self) -> (u32, bool) { assert!( @@ -213,7 +233,7 @@ impl RequestSequence { .planned_output_ids .as_ref() .and_then(|ids| ids.get(self.generated_tokens).copied()) - .unwrap_or_else(random::); + .unwrap_or_else(|| synthetic_output_token(self.request_id, self.generated_tokens)); let opened_partial = self.len().is_multiple_of(self.block_size); self.tokens.push(token); self.generated_tokens += 1; @@ -347,6 +367,18 @@ impl RequestSequence { mod tests { use super::*; + #[test] + fn synthetic_output_tokens_match_independent_xxh3_values() { + // xxh3_64(request UUID bytes ‖ ordinal as u64 LE) truncated to u32, computed + // with an independent XXH3 implementation (Python `xxhash`). Pinning the + // values keeps them stable across processes, toolchains, and platforms. + let tokens: Vec = (0..4) + .map(|ordinal| synthetic_output_token(Uuid::from_u128(1), ordinal)) + .collect(); + assert_eq!(tokens, [2515880452, 3600266899, 2716793611, 2249174135]); + assert_eq!(synthetic_output_token(Uuid::from_u128(2), 0), 1927195132); + } + fn sequence( tokens: Vec, max_output_tokens: usize, @@ -355,6 +387,7 @@ mod tests { planned_output_ids: Option>, ) -> (RequestSequence, Vec) { RequestSequence::new( + Uuid::nil(), tokens, max_output_tokens, max_output_tokens, @@ -452,8 +485,17 @@ mod tests { #[test] fn completion_footprint_is_bounded_by_requested_output() { - let (sequence, identities) = - RequestSequence::new(vec![0, 1, 2, 3, 4], 100, 3, 4, false, false, false, None); + let (sequence, identities) = RequestSequence::new( + Uuid::nil(), + vec![0, 1, 2, 3, 4], + 100, + 3, + 4, + false, + false, + false, + None, + ); assert_eq!(sequence.current_known_blocks(), 2); assert_eq!(sequence.to_completion_blocks(), 27); assert_eq!( diff --git a/crates/core/src/engine/kv_manager/g1_manager.rs b/crates/core/src/engine/kv_manager/g1_manager.rs index 243eec9f0e..681bb003bc 100644 --- a/crates/core/src/engine/kv_manager/g1_manager.rs +++ b/crates/core/src/engine/kv_manager/g1_manager.rs @@ -402,6 +402,7 @@ mod tests { let mut manager = G1Manager::new_with_event_sink(8, 4, KvEventPublishers::new(Some(sink)), 0); let (mut sequence, identities) = RequestSequence::new( + owner, (0..8).collect(), 4, 4, diff --git a/crates/core/src/engine/kv_manager/state_cache_tests.rs b/crates/core/src/engine/kv_manager/state_cache_tests.rs index 0c319461fb..0a1bd8b69e 100644 --- a/crates/core/src/engine/kv_manager/state_cache_tests.rs +++ b/crates/core/src/engine/kv_manager/state_cache_tests.rs @@ -30,6 +30,7 @@ fn request( output: usize, ) -> (RequestSequence, BlockRequestLease) { let (sequence, identities) = RequestSequence::new( + owner, (0..tokens).map(|value| value as u32).collect(), output, output, @@ -895,8 +896,9 @@ fn kda_partial_page_can_match_when_tokens_after_the_hit_diverge() { manager.finish_lease(owner, lease); let mut tokens: Vec = (0..19).collect(); tokens[16] = 999; - let (mut query, ids) = RequestSequence::new(tokens.clone(), 0, 0, 6, true, true, false, None); let owner = Uuid::from_u128(603); + let (mut query, ids) = + RequestSequence::new(owner, tokens.clone(), 0, 0, 6, true, true, false, None); let mut reader = BlockRequestLease::new(owner, ids); reader.configure_prefix_hashes(&tokens, 2); assert_eq!( diff --git a/crates/core/src/engine/kv_manager/vllm_backend.rs b/crates/core/src/engine/kv_manager/vllm_backend.rs index 33d3f91971..f2073d9fe0 100644 --- a/crates/core/src/engine/kv_manager/vllm_backend.rs +++ b/crates/core/src/engine/kv_manager/vllm_backend.rs @@ -1747,7 +1747,8 @@ mod tests { emit_token_ids: bool, ) -> (RequestSequence, BlockRequestLease) { let tokens = (0..hashes.len() * 4).map(|token| token as u32).collect(); - let (sequence, _) = RequestSequence::new(tokens, 0, 0, 4, true, true, emit_token_ids, None); + let (sequence, _) = + RequestSequence::new(owner, tokens, 0, 0, 4, true, true, emit_token_ids, None); let identities = hashes .iter() .copied() @@ -1986,8 +1987,17 @@ mod tests { // blocks plus one partial tail occupy four physical blocks. let first = Uuid::from_u128(1); let first_tokens = (0..15).collect::>(); - let (mut first_sequence, first_identities) = - RequestSequence::new(first_tokens, 2, 2, 4, true, true, false, Some(vec![15, 16])); + let (mut first_sequence, first_identities) = RequestSequence::new( + first, + first_tokens, + 2, + 2, + 4, + true, + true, + false, + Some(vec![15, 16]), + ); let mut first_lease = BlockRequestLease::new(first, first_identities); ready(manager.allocate_lease(first, &mut first_lease, 15, 0)); manager.finalize_lease_computed_prefix(first, &mut first_sequence, &mut first_lease, 0, 15); @@ -2032,7 +2042,7 @@ mod tests { let mut second_tokens = (0..10).collect::>(); second_tokens.extend([100, 101, 102, 103]); let (mut second_sequence, second_identities) = - RequestSequence::new(second_tokens, 0, 0, 4, true, true, false, None); + RequestSequence::new(second, second_tokens, 0, 0, 4, true, true, false, None); let mut second_lease = BlockRequestLease::new(second, second_identities); let prefill = manager.get_lease_prefill_cost(&second_sequence, &second_lease); assert_eq!(prefill.cached_tokens, 8); diff --git a/crates/core/src/engine/protocol.rs b/crates/core/src/engine/protocol.rs index ffc5e940a3..48ff50d661 100644 --- a/crates/core/src/engine/protocol.rs +++ b/crates/core/src/engine/protocol.rs @@ -18,6 +18,8 @@ pub struct Request { /// Requested output length. pub max_output_tokens: usize, /// Optional exact output IDs. Its length overrides `max_output_tokens`. + /// When absent, each output token is derived from `request_id` and its + /// output position, so reusing a request ID reproduces the same stream. #[serde(default, skip_serializing_if = "Option::is_none")] pub output_token_ids: Option>, } diff --git a/crates/core/src/engine/scheduler/sglang/request.rs b/crates/core/src/engine/scheduler/sglang/request.rs index 074f5dfacd..3f513019f8 100644 --- a/crates/core/src/engine/scheduler/sglang/request.rs +++ b/crates/core/src/engine/scheduler/sglang/request.rs @@ -2,17 +2,16 @@ // SPDX-License-Identifier: Apache-2.0 use std::collections::VecDeque; -use std::collections::hash_map::DefaultHasher; use std::ops::{Deref, RangeFull}; use crate::engine::scheduler::queue_metrics::{QueueStats, QueuedLength}; -use std::hash::{Hash, Hasher}; use uuid::Uuid; #[cfg(test)] use crate::engine::cache::radix_cache::KvPageId; use crate::engine::common::protocols::DirectRequest; +use crate::engine::common::sequence::synthetic_output_token; use crate::engine::kv_manager::sglang_backend::RadixRequestLease; #[derive(Debug)] @@ -126,10 +125,7 @@ impl SglangRequest { return *token_id; } - let mut hasher = DefaultHasher::new(); - self.uuid.hash(&mut hasher); - self.output_len().hash(&mut hasher); - hasher.finish() as u32 + synthetic_output_token(self.uuid, self.output_len()) } pub(super) fn append_output_token(&mut self, token: u32, block_size: usize) { diff --git a/crates/core/src/engine/scheduler/vllm/host_offload.rs b/crates/core/src/engine/scheduler/vllm/host_offload.rs index 319b1c4bff..aab1b84a24 100644 --- a/crates/core/src/engine/scheduler/vllm/host_offload.rs +++ b/crates/core/src/engine/scheduler/vllm/host_offload.rs @@ -694,7 +694,7 @@ mod tests { tokens: Vec, ) -> (RequestSequence, BlockRequestLease, VllmHostRequestState) { let (sequence, identities) = - RequestSequence::new(tokens, 0, 0, 4, true, false, false, None); + RequestSequence::new(owner, tokens, 0, 0, 4, true, false, false, None); let lease = BlockRequestLease::new(owner, identities); let host = VllmHostRequestState::new(&sequence, &lease, 4); (sequence, lease, host) diff --git a/crates/core/src/engine/scheduler/vllm/request.rs b/crates/core/src/engine/scheduler/vllm/request.rs index e1894e552d..4c6153c4f6 100644 --- a/crates/core/src/engine/scheduler/vllm/request.rs +++ b/crates/core/src/engine/scheduler/vllm/request.rs @@ -26,6 +26,7 @@ impl RequestKvState { planned_output_ids: Option>, ) -> Self { let (sequence, identities) = RequestSequence::new( + owner, tokens, max_output_tokens, output_capacity_hint, diff --git a/crates/core/src/engine/trace.rs b/crates/core/src/engine/trace.rs index 5a0694b5d0..2753ed96b6 100644 --- a/crates/core/src/engine/trace.rs +++ b/crates/core/src/engine/trace.rs @@ -31,7 +31,9 @@ pub fn g1_parent_chain_events(block_size: usize) -> Vec { let prompt_len = computed_after - 2; let prompt_len_u32 = u32::try_from(prompt_len).expect("ordering regression prompt length must fit in u32"); + let owner = Uuid::from_u128(1); let (mut sequence, identities) = RequestSequence::new( + owner, (0..prompt_len_u32).collect(), 2, 2, @@ -41,7 +43,6 @@ pub fn g1_parent_chain_events(block_size: usize) -> Vec { false, Some(vec![prompt_len_u32, prompt_len_u32 + 1]), ); - let owner = Uuid::from_u128(1); let mut lease = BlockRequestLease::new(owner, identities); let (events, sink) = capture_kv_event_sink(); let mut manager = diff --git a/crates/core/src/replay/capture.rs b/crates/core/src/replay/capture.rs index afe3116f00..f32d7a9cb5 100644 --- a/crates/core/src/replay/capture.rs +++ b/crates/core/src/replay/capture.rs @@ -11,7 +11,8 @@ pub const CANONICAL_SELECTOR_SEED: u64 = 0xd1a0_5eed; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum ReplayDeterminism { - /// Preserve normal runtime randomness. + /// Preserve normal runtime randomness: fresh random request IDs, so synthetic + /// outputs vary between runs, and no canonical selector seed. #[default] Random, /// Use ordinal request UUIDs and the canonical selector seed. diff --git a/crates/core/tests/artifacts.rs b/crates/core/tests/artifacts.rs index e6c9023e86..0693d0d7e0 100644 --- a/crates/core/tests/artifacts.rs +++ b/crates/core/tests/artifacts.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use aisimulate_core::engine::{ - Backend, EngineConfig, KvEvent, NativeHostOffloadConfig, TimingModelConfig, + Backend, EngineConfig, KvEvent, KvEventData, NativeHostOffloadConfig, TimingModelConfig, }; use aisimulate_core::replay::loadgen::{SessionTrace, Trace, TurnTrace, WorkloadDriver}; use aisimulate_core::replay::{ @@ -261,6 +261,98 @@ fn cache_events_advertise_reusable_prefixes_only_when_caching_is_enabled() { } } +#[test] +fn canonical_replay_reproduces_unplanned_output_tokens_and_kv_hashes() { + let output_tokens = |artifacts: &ReplayArtifacts, request_index: usize| -> Vec { + let request_id = artifacts.requests[request_index].request_id; + artifacts + .outputs + .iter() + .filter(|output| output.request_id == request_id) + .filter_map(|output| output.token_id) + .collect() + }; + let mut first_request_streams = Vec::new(); + for backend in [Backend::Vllm, Backend::Trtllm, Backend::Sglang] { + let run = || { + let mut replay_spec = spec(backend, 1, 1); + let mut config: ReplayEngineConfig = + serde_json::from_value(replay_spec.engine.clone()).unwrap(); + config.rank.emit_kv_events = true; + config.rank.emit_kv_token_ids = true; + replay_spec.engine = serde_json::to_value(config).unwrap(); + // Identical length-only requests carry no output plan, so every + // generated token and output-block hash is engine-synthesized. + replay_spec.requests = (0..2) + .map(|index| { + let mut request = + replay_request(&format!("unplanned-{index}"), 0.0, (0..6).collect()); + request.output_tokens = 9; + request + }) + .collect(); + Replayer::new(replay_spec, ReplayEngineFactory::new()) + .unwrap() + .with_capture_options(ReplayCaptureOptions { + determinism: ReplayDeterminism::CanonicalV1, + ..ReplayCaptureOptions::default() + }) + .run_with_artifacts(ReplayArtifactKvEventVisibility::Native) + .unwrap() + .1 + }; + let first = run(); + let second = run(); + for request_index in 0..2 { + let tokens = output_tokens(&first, request_index); + assert_eq!(tokens.len(), 9, "{backend:?}"); + assert_eq!( + tokens, + output_tokens(&second, request_index), + "{backend:?} request {request_index}" + ); + } + assert_ne!( + output_tokens(&first, 0), + output_tokens(&first, 1), + "{backend:?}: distinct requests must not share synthetic output blocks" + ); + assert_eq!(kv_parts(&first), kv_parts(&second), "{backend:?}"); + + // Block 2 of request 0 (block_size 4, 6 prompt tokens) holds only generated + // tokens; it must be published with a hash derived from those tokens. + let request_tokens: Vec = (0..6).chain(output_tokens(&first, 0)).collect(); + let output_block = &request_tokens[8..12]; + let stored = first + .kv_events + .iter() + .filter_map(|event| match &event.event.data { + KvEventData::Stored(stored) => Some(&stored.blocks), + KvEventData::Removed { .. } => None, + }) + .flatten() + .find(|block| block.token_ids.as_deref() == Some(output_block)) + .unwrap_or_else(|| panic!("{backend:?}: no stored block holds generated tokens")); + let bytes: Vec = output_block + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect(); + assert_eq!( + stored.tokens_hash, + xxhash_rust::xxh3::xxh3_64_with_seed(&bytes, 1337), + "{backend:?}" + ); + first_request_streams.push(output_tokens(&first, 0)); + } + // Every engine derives unplanned tokens from the same request identity. + assert!( + first_request_streams + .windows(2) + .all(|pair| pair[0] == pair[1]), + "engines diverged: {first_request_streams:?}" + ); +} + #[test] fn capped_passes_respect_visibility_boundaries() { let capped = |visibility| {