From eef5b07b9d828637cb3199e6ac1b346d89292830 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Sun, 27 Sep 2026 02:39:08 -0700 Subject: [PATCH 1/3] fix: derive vLLM synthetic output tokens from request identity The vLLM and TRT-LLM scheduler path filled output positions without a planned token with unseeded random u32 values. Those tokens determine output-block hashes, so two otherwise identical canonical replays emitted different output token IDs and KV events. Derive each missing position from the request UUID and output ordinal, matching the inputs SGLang already uses. RequestSequence now carries its request ID. XXH3 over a fixed little-endian encoding keeps the values stable across platforms and toolchains and costs no more than the thread RNG it replaces. Signed-off-by: PeaBrane --- crates/core/src/engine/common/sequence.rs | 36 +++++++++++- .../core/src/engine/kv_manager/g1_manager.rs | 1 + .../engine/kv_manager/state_cache_tests.rs | 4 +- .../src/engine/kv_manager/vllm_backend.rs | 18 ++++-- .../src/engine/scheduler/vllm/host_offload.rs | 2 +- .../core/src/engine/scheduler/vllm/request.rs | 1 + crates/core/src/engine/trace.rs | 3 +- crates/core/tests/artifacts.rs | 55 +++++++++++++++++++ 8 files changed, 110 insertions(+), 10 deletions(-) diff --git a/crates/core/src/engine/common/sequence.rs b/crates/core/src/engine/common/sequence.rs index b20813a2e3..dee90f6795 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. +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; @@ -355,6 +375,7 @@ mod tests { planned_output_ids: Option>, ) -> (RequestSequence, Vec) { RequestSequence::new( + Uuid::nil(), tokens, max_output_tokens, max_output_tokens, @@ -452,8 +473,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 80d794461d..b679239faf 100644 --- a/crates/core/src/engine/kv_manager/g1_manager.rs +++ b/crates/core/src/engine/kv_manager/g1_manager.rs @@ -394,6 +394,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 1542d89898..58a978b1b6 100644 --- a/crates/core/src/engine/kv_manager/vllm_backend.rs +++ b/crates/core/src/engine/kv_manager/vllm_backend.rs @@ -1709,7 +1709,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() @@ -1948,8 +1949,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); @@ -1994,7 +2004,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/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/tests/artifacts.rs b/crates/core/tests/artifacts.rs index e6c9023e86..4a39c443dd 100644 --- a/crates/core/tests/artifacts.rs +++ b/crates/core/tests/artifacts.rs @@ -261,6 +261,61 @@ 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() + }; + for backend in [Backend::Vllm, Backend::Trtllm, Backend::Sglang] { + let run = || { + let mut replay_spec = spec(backend, 1, 1); + // 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!(!first.kv_events.is_empty(), "{backend:?}"); + assert_eq!(kv_parts(&first), kv_parts(&second), "{backend:?}"); + } +} + #[test] fn capped_passes_respect_visibility_boundaries() { let capped = |visibility| { From aefbf7808192e99180d47a23b0c930c014c9a14c Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 28 Sep 2026 22:20:01 -0700 Subject: [PATCH 2/3] fix: derive SGLang synthetic output tokens from the shared helper SGLang hashed (request uuid, ordinal) with std's DefaultHasher, whose algorithm is not guaranteed stable across Rust releases, so canonical SGLang output-block hashes could shift with a toolchain upgrade. Route SGLang through the same xxh3-based synthetic_output_token helper as vLLM/TRT-LLM so every engine derives unplanned output tokens identically and stably. Signed-off-by: PeaBrane --- crates/core/src/engine/common/sequence.rs | 2 +- crates/core/src/engine/scheduler/sglang/request.rs | 9 ++------- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/crates/core/src/engine/common/sequence.rs b/crates/core/src/engine/common/sequence.rs index dee90f6795..2804f70308 100644 --- a/crates/core/src/engine/common/sequence.rs +++ b/crates/core/src/engine/common/sequence.rs @@ -15,7 +15,7 @@ use xxhash_rust::xxh3::xxh3_64; /// 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. -fn synthetic_output_token(request_id: Uuid, output_ordinal: usize) -> u32 { +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()); diff --git a/crates/core/src/engine/scheduler/sglang/request.rs b/crates/core/src/engine/scheduler/sglang/request.rs index b68e45dedc..d480283482 100644 --- a/crates/core/src/engine/scheduler/sglang/request.rs +++ b/crates/core/src/engine/scheduler/sglang/request.rs @@ -1,14 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::hash_map::DefaultHasher; -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)] @@ -120,10 +118,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) { From ba96a1a0ca90932c3508625671e9ce233ba78ed1 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 28 Sep 2026 22:31:01 -0700 Subject: [PATCH 3/3] test: pin synthetic output tokens and hash a generated block Pin synthetic_output_token to independently computed XXH3 values, require every engine to emit the same unplanned stream for one request identity, and check that a stored block holding only generated tokens carries the hash of those tokens. Document that absent output_token_ids derive from request_id and position. Signed-off-by: PeaBrane --- crates/core/src/engine/common/sequence.rs | 12 +++++++ crates/core/src/engine/protocol.rs | 2 ++ crates/core/src/replay/capture.rs | 3 +- crates/core/tests/artifacts.rs | 41 +++++++++++++++++++++-- 4 files changed, 55 insertions(+), 3 deletions(-) diff --git a/crates/core/src/engine/common/sequence.rs b/crates/core/src/engine/common/sequence.rs index 2804f70308..5ae4076775 100644 --- a/crates/core/src/engine/common/sequence.rs +++ b/crates/core/src/engine/common/sequence.rs @@ -367,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, diff --git a/crates/core/src/engine/protocol.rs b/crates/core/src/engine/protocol.rs index 12b3e6aebb..535c129644 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/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 4a39c443dd..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::{ @@ -272,9 +272,15 @@ fn canonical_replay_reproduces_unplanned_output_tokens_and_kv_hashes() { .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) @@ -311,9 +317,40 @@ fn canonical_replay_reproduces_unplanned_output_tokens_and_kv_hashes() { output_tokens(&first, 1), "{backend:?}: distinct requests must not share synthetic output blocks" ); - assert!(!first.kv_events.is_empty(), "{backend:?}"); 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]