aisimulate_core/engine/
handoff.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Runtime-neutral identities and timing inputs for prefill-to-decode handoff.
5//!
6//! These value types are owned by the `aisimulate-core::engine` module because native
7//! schedulers must retain them with their KV ownership state. Replay owns the
8//! handoff coordinator, ordering, and virtual transfer event.
9//!
10use serde::{Deserialize, Serialize};
11use uuid::Uuid;
12
13use crate::engine::config::WorkerType;
14
15/// Stable identifier for one replay-local prefill-to-decode handoff attempt.
16///
17/// The caller allocates the UUID identity. Engines only retain and echo it
18/// in lifecycle effects, so creating an engine does not require randomness or
19/// a Dynamo runtime.
20#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
21#[serde(transparent)]
22pub struct HandoffId(Uuid);
23
24impl HandoffId {
25    /// Construct an identity from the caller-owned replay coordinator.
26    pub const fn new(value: Uuid) -> Self {
27        Self(value)
28    }
29
30    /// Return the caller-owned UUID identity.
31    pub const fn get(self) -> Uuid {
32        self.0
33    }
34}
35
36impl From<Uuid> for HandoffId {
37    fn from(value: Uuid) -> Self {
38        Self(value)
39    }
40}
41
42impl From<HandoffId> for Uuid {
43    fn from(value: HandoffId) -> Self {
44        value.0
45    }
46}
47
48/// Prompt footprint used when the Replayer models KV transfer time.
49#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51pub enum TransferTimingMode {
52    /// Transfer time is based on every prompt token.
53    #[default]
54    FullPrompt,
55    /// Transfer time is based only on prompt tokens missing at the destination.
56    DestinationMissing,
57}
58
59/// Source-provided inputs for calculating KV transfer delay.
60///
61/// Schedulers publish this value when source KV ownership becomes held.
62/// The Replayer combines it with the destination's missing-token observation
63/// and schedules the virtual transfer completion.
64#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
65pub struct HandoffTransferTiming {
66    /// Which prompt footprint contributes to transfer time.
67    pub mode: TransferTimingMode,
68    /// Full source prompt length.
69    pub full_prompt_tokens: usize,
70    /// Modeled KV bytes occupied by one prompt token.
71    pub kv_bytes_per_token: Option<usize>,
72    /// Modeled transfer bandwidth in decimal gigabytes per second.
73    pub bandwidth_gb_s: Option<f64>,
74}
75
76impl HandoffTransferTiming {
77    /// Calculate transfer delay in milliseconds.
78    ///
79    /// Returns `None` when the source did not provide a complete timing model
80    /// or when bandwidth is non-positive. In that case the Replayer may apply
81    /// its configured fallback delay.
82    pub fn delay_ms(self, destination_missing_tokens: usize) -> Option<f64> {
83        let tokens = match self.mode {
84            TransferTimingMode::FullPrompt => self.full_prompt_tokens,
85            TransferTimingMode::DestinationMissing => destination_missing_tokens,
86        };
87        let (Some(bytes_per_token), Some(bandwidth_gb_s)) =
88            (self.kv_bytes_per_token, self.bandwidth_gb_s)
89        else {
90            return None;
91        };
92        if bandwidth_gb_s <= 0.0 {
93            return None;
94        }
95
96        Some(tokens as f64 * bytes_per_token as f64 / (bandwidth_gb_s * 1e9) * 1000.0)
97    }
98
99    /// Calculate delay using the full prompt irrespective of `mode`.
100    pub fn full_prompt_delay_ms(self) -> Option<f64> {
101        Self {
102            mode: TransferTimingMode::FullPrompt,
103            ..self
104        }
105        .delay_ms(0)
106    }
107}
108
109/// Compute the client-visible prefill-to-decode handoff delay.
110///
111/// A delay is exposed only for a completed prefill-worker request. Aggregated
112/// and decode workers do not cross a prefill/decode transfer boundary.
113pub fn prefill_handoff_delay_ms(
114    worker_type: WorkerType,
115    completed: bool,
116    num_input_tokens: usize,
117    bandwidth_gb_s: Option<f64>,
118    kv_bytes_per_token: Option<usize>,
119) -> Option<f64> {
120    if worker_type != WorkerType::Prefill || !completed {
121        return None;
122    }
123    HandoffTransferTiming {
124        mode: TransferTimingMode::FullPrompt,
125        full_prompt_tokens: num_input_tokens,
126        kv_bytes_per_token,
127        bandwidth_gb_s,
128    }
129    .full_prompt_delay_ms()
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn handoff_id_round_trips_caller_owned_value() {
138        let value = Uuid::from_u128(17);
139        let handoff_id = HandoffId::new(value);
140        assert_eq!(handoff_id.get(), value);
141    }
142
143    #[test]
144    fn transfer_delay_uses_selected_prompt_footprint() {
145        let timing = HandoffTransferTiming {
146            mode: TransferTimingMode::DestinationMissing,
147            full_prompt_tokens: 100,
148            kv_bytes_per_token: Some(1_000),
149            bandwidth_gb_s: Some(1.0),
150        };
151
152        assert_eq!(timing.delay_ms(20), Some(0.02));
153        assert_eq!(timing.full_prompt_delay_ms(), Some(0.1));
154    }
155
156    #[test]
157    fn incomplete_or_non_positive_timing_model_has_no_delay() {
158        let timing = HandoffTransferTiming {
159            mode: TransferTimingMode::FullPrompt,
160            full_prompt_tokens: 100,
161            kv_bytes_per_token: None,
162            bandwidth_gb_s: Some(1.0),
163        };
164        assert_eq!(timing.delay_ms(0), None);
165
166        let timing = HandoffTransferTiming {
167            kv_bytes_per_token: Some(1_000),
168            bandwidth_gb_s: Some(0.0),
169            ..timing
170        };
171        assert_eq!(timing.delay_ms(0), None);
172    }
173
174    #[test]
175    fn prefill_handoff_delay_requires_completed_prefill_work() {
176        let args = (128, Some(1.0), Some(1_000_000));
177        assert_eq!(
178            prefill_handoff_delay_ms(WorkerType::Prefill, true, args.0, args.1, args.2,),
179            Some(128.0)
180        );
181        assert_eq!(
182            prefill_handoff_delay_ms(WorkerType::Prefill, false, args.0, args.1, args.2,),
183            None
184        );
185        assert_eq!(
186            prefill_handoff_delay_ms(WorkerType::Decode, true, args.0, args.1, args.2,),
187            None
188        );
189    }
190}
aisimulate_core/engine/handoff.rs

aisimulate_core/engine/
handoff.rs
