aisimulate_core/replay/
handoff.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Replay-owned prefill/decode handoff ordering and virtual-transfer state.
5
6use crate::engine::Backend;
7pub use crate::engine::{HandoffId, HandoffTransferTiming};
8use anyhow::{Result, bail};
9use rustc_hash::FxHashMap;
10use serde::{Deserialize, Serialize};
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
13pub enum HandoffOrder {
14    SourceFirst,
15    DestinationFirst,
16}
17
18#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
19pub enum HandoffFact {
20    SourceHeld {
21        handoff_id: HandoffId,
22        transfer_timing: HandoffTransferTiming,
23    },
24    DestinationReserved {
25        handoff_id: HandoffId,
26        transferable_prompt_tokens: usize,
27    },
28    TransferCompleted {
29        handoff_id: HandoffId,
30    },
31    Failed {
32        handoff_id: HandoffId,
33    },
34    TimedOut {
35        handoff_id: HandoffId,
36    },
37    Canceled {
38        handoff_id: HandoffId,
39    },
40}
41
42impl HandoffFact {
43    fn handoff_id(&self) -> HandoffId {
44        match *self {
45            Self::SourceHeld { handoff_id, .. }
46            | Self::DestinationReserved { handoff_id, .. }
47            | Self::TransferCompleted { handoff_id }
48            | Self::Failed { handoff_id }
49            | Self::TimedOut { handoff_id }
50            | Self::Canceled { handoff_id } => handoff_id,
51        }
52    }
53}
54
55#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
56pub enum HandoffAction {
57    SubmitPrefill {
58        handoff_id: HandoffId,
59    },
60    ReserveDestination {
61        handoff_id: HandoffId,
62    },
63    StartTransfer {
64        handoff_id: HandoffId,
65        delay_ms: f64,
66    },
67    ActivateDestination {
68        handoff_id: HandoffId,
69    },
70    ReleaseSource {
71        handoff_id: HandoffId,
72    },
73    CancelSource {
74        handoff_id: HandoffId,
75    },
76    CancelDestination {
77        handoff_id: HandoffId,
78    },
79    Complete {
80        handoff_id: HandoffId,
81    },
82}
83
84#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
85pub struct HandoffActionId(u64);
86
87#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
88pub struct IssuedHandoffAction {
89    pub id: HandoffActionId,
90    pub action: HandoffAction,
91}
92
93#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
94pub enum HandoffActionOutcome {
95    Submitted,
96    Accepted,
97    Scheduled,
98    Applied,
99    Noop,
100    Failed(String),
101}
102
103#[derive(Clone, Copy, Debug, Eq, PartialEq)]
104enum CoordinatorMode {
105    Active,
106    CleaningUp,
107    Complete,
108}
109
110#[derive(Clone, Copy, Debug, Eq, PartialEq)]
111pub enum HandoffCompletion {
112    Success,
113    Canceled,
114}
115
116#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
117pub enum NormalizedHandoffEvent {
118    SourceHeld,
119    DestinationAccepted,
120    DestinationReserved,
121    DestinationActivated,
122    SourceReleased,
123    Completed,
124}
125
126#[doc(hidden)]
127#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
128pub struct NormalizedHandoffConformance {
129    pub engine_type: Backend,
130    pub order: HandoffOrder,
131    pub lifecycle: Vec<NormalizedHandoffEvent>,
132    pub source_output_tokens: usize,
133    pub destination_output_tokens: usize,
134    pub completed_requests: usize,
135    pub destination_stored: NormalizedStoredTiming,
136    pub source_drained: bool,
137    pub destination_drained: bool,
138    pub driver_drained: bool,
139}
140
141#[doc(hidden)]
142#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
143pub struct NormalizedStoredTiming {
144    pub before_activation: usize,
145    pub on_activation: usize,
146    pub repeated_activation_hashes_after_activation: usize,
147}
148
149impl NormalizedHandoffConformance {
150    #[doc(hidden)]
151    pub fn validate(&self) -> Result<()> {
152        let expected_order = match self.engine_type {
153            Backend::Vllm => HandoffOrder::SourceFirst,
154            Backend::Sglang => HandoffOrder::DestinationFirst,
155            Backend::Trtllm => bail!("TRT-LLM does not support destination handoff"),
156        };
157        if self.order != expected_order {
158            bail!(
159                "normalized handoff order mismatch: expected {expected_order:?}, got {:?}",
160                self.order
161            );
162        }
163        if self.lifecycle != expected_normalized_handoff(self.order) {
164            bail!(
165                "normalized handoff lifecycle mismatch: got {:?}",
166                self.lifecycle
167            );
168        }
169        if self.source_output_tokens != 1
170            || self.destination_output_tokens != 2
171            || self.completed_requests != 1
172        {
173            bail!("normalized handoff output/completion counts do not match the fixture");
174        }
175        if self.destination_stored.before_activation != 0
176            || self.destination_stored.on_activation == 0
177            || self
178                .destination_stored
179                .repeated_activation_hashes_after_activation
180                != 0
181        {
182            bail!("normalized handoff destination KV visibility is invalid");
183        }
184        if !self.source_drained || !self.destination_drained || !self.driver_drained {
185            bail!("normalized handoff did not drain");
186        }
187        Ok(())
188    }
189}
190
191pub fn expected_normalized_handoff(order: HandoffOrder) -> &'static [NormalizedHandoffEvent] {
192    use NormalizedHandoffEvent::*;
193    match order {
194        HandoffOrder::SourceFirst => &[
195            SourceHeld,
196            DestinationAccepted,
197            DestinationReserved,
198            DestinationActivated,
199            SourceReleased,
200            Completed,
201        ],
202        HandoffOrder::DestinationFirst => &[
203            DestinationAccepted,
204            DestinationReserved,
205            SourceHeld,
206            DestinationActivated,
207            SourceReleased,
208            Completed,
209        ],
210    }
211}
212
213#[derive(Default)]
214struct ActionJournal {
215    started: bool,
216    next_id: u64,
217    issued: FxHashMap<HandoffActionId, HandoffAction>,
218    outcomes: FxHashMap<HandoffActionId, HandoffActionOutcome>,
219}
220
221#[derive(Default)]
222struct SourceProgress {
223    submit_issued: bool,
224    submitted: bool,
225    held: bool,
226    transfer_timing: Option<HandoffTransferTiming>,
227    release_issued: bool,
228    cancel_issued: bool,
229    cleanup_done: bool,
230}
231
232#[derive(Default)]
233struct DestinationProgress {
234    reserve_issued: bool,
235    accepted: bool,
236    reserved: bool,
237    transferable_prompt_tokens: Option<usize>,
238    activation_issued: bool,
239    activation_applied: bool,
240    cancel_issued: bool,
241    cleanup_done: bool,
242}
243
244#[derive(Default)]
245struct TransferProgress {
246    issued: bool,
247    scheduled: bool,
248    completed: bool,
249}
250
251/// Pure state machine for one prefill-to-decode ownership handoff.
252pub struct HandoffCoordinatorCore {
253    handoff_id: HandoffId,
254    order: HandoffOrder,
255    mode: CoordinatorMode,
256    actions: ActionJournal,
257    source: SourceProgress,
258    destination: DestinationProgress,
259    transfer: TransferProgress,
260    completion: Option<HandoffCompletion>,
261    fallback_transfer_delay_ms: f64,
262}
263
264impl HandoffCoordinatorCore {
265    pub fn new(handoff_id: HandoffId, order: HandoffOrder) -> Self {
266        Self::new_with_fallback(handoff_id, order, 0.0)
267    }
268
269    pub fn new_with_fallback(
270        handoff_id: HandoffId,
271        order: HandoffOrder,
272        fallback_transfer_delay_ms: f64,
273    ) -> Self {
274        debug_assert!(fallback_transfer_delay_ms.is_finite());
275        debug_assert!(fallback_transfer_delay_ms >= 0.0);
276        Self {
277            handoff_id,
278            order,
279            mode: CoordinatorMode::Active,
280            actions: ActionJournal::default(),
281            source: SourceProgress::default(),
282            destination: DestinationProgress::default(),
283            transfer: TransferProgress::default(),
284            completion: None,
285            fallback_transfer_delay_ms,
286        }
287    }
288
289    pub fn start(&mut self) -> Result<Vec<IssuedHandoffAction>> {
290        if self.actions.started {
291            return Ok(Vec::new());
292        }
293        self.actions.started = true;
294        let action = match self.order {
295            HandoffOrder::SourceFirst => self.issue_submit_prefill(),
296            HandoffOrder::DestinationFirst => self.issue_reserve_destination(),
297        };
298        Ok(vec![action])
299    }
300
301    pub fn on_fact(&mut self, fact: HandoffFact) -> Result<Vec<IssuedHandoffAction>> {
302        self.validate_handoff(fact.handoff_id())?;
303        if self.mode != CoordinatorMode::Active {
304            return Ok(Vec::new());
305        }
306        match fact {
307            HandoffFact::SourceHeld {
308                transfer_timing, ..
309            } => {
310                if self.source.held {
311                    return Ok(Vec::new());
312                }
313                if !self.source.submitted {
314                    bail!("source held before prefill submission was acknowledged");
315                }
316                validate_transfer_timing(transfer_timing)?;
317                self.source.held = true;
318                self.source.transfer_timing = Some(transfer_timing);
319                self.advance_active()
320            }
321            HandoffFact::DestinationReserved {
322                transferable_prompt_tokens,
323                ..
324            } => {
325                if self.destination.reserved {
326                    return Ok(Vec::new());
327                }
328                if !self.destination.accepted {
329                    bail!("destination reserved before ownership was accepted");
330                }
331                self.destination.reserved = true;
332                self.destination.transferable_prompt_tokens = Some(transferable_prompt_tokens);
333                self.advance_active()
334            }
335            HandoffFact::TransferCompleted { .. } => {
336                if self.transfer.completed {
337                    return Ok(Vec::new());
338                }
339                if !self.transfer.scheduled {
340                    bail!("transfer completed before it was scheduled");
341                }
342                self.transfer.completed = true;
343                self.advance_active()
344            }
345            HandoffFact::Failed { .. }
346            | HandoffFact::TimedOut { .. }
347            | HandoffFact::Canceled { .. } => self.begin_cleanup(),
348        }
349    }
350
351    pub fn on_action_outcome(
352        &mut self,
353        action_id: HandoffActionId,
354        outcome: HandoffActionOutcome,
355    ) -> Result<Vec<IssuedHandoffAction>> {
356        if self.mode == CoordinatorMode::Complete {
357            return Ok(Vec::new());
358        }
359        let Some(action) = self.actions.issued.get(&action_id).copied() else {
360            bail!("unknown handoff action {action_id:?}");
361        };
362        if let Some(previous) = self.actions.outcomes.get(&action_id) {
363            if previous != &outcome {
364                bail!("conflicting outcome for handoff action {action_id:?}");
365            }
366            return Ok(Vec::new());
367        }
368        self.actions.outcomes.insert(action_id, outcome.clone());
369
370        if let HandoffActionOutcome::Failed(_) = outcome {
371            if matches!(
372                action,
373                HandoffAction::CancelSource { .. } | HandoffAction::CancelDestination { .. }
374            ) {
375                bail!("handoff cleanup action {action_id:?} failed");
376            }
377            return self.begin_cleanup();
378        }
379
380        match action {
381            HandoffAction::SubmitPrefill { .. } => {
382                require_outcome(&outcome, &[HandoffActionOutcome::Submitted])?;
383                self.source.submitted = true;
384            }
385            HandoffAction::ReserveDestination { .. } => {
386                require_outcome(&outcome, &[HandoffActionOutcome::Accepted])?;
387                self.destination.accepted = true;
388            }
389            HandoffAction::StartTransfer { .. } => {
390                require_outcome(&outcome, &[HandoffActionOutcome::Scheduled])?;
391                self.transfer.scheduled = true;
392            }
393            HandoffAction::ActivateDestination { .. } => {
394                require_outcome(&outcome, &[HandoffActionOutcome::Applied])?;
395                self.destination.activation_applied = true;
396            }
397            HandoffAction::ReleaseSource { .. } => {
398                require_outcome(
399                    &outcome,
400                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
401                )?;
402                self.source.cleanup_done = true;
403            }
404            HandoffAction::CancelSource { .. } => {
405                require_outcome(
406                    &outcome,
407                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
408                )?;
409                self.source.cleanup_done = true;
410            }
411            HandoffAction::CancelDestination { .. } => {
412                require_outcome(
413                    &outcome,
414                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
415                )?;
416                self.destination.cleanup_done = true;
417            }
418            HandoffAction::Complete { .. } => return Ok(Vec::new()),
419        }
420
421        match self.mode {
422            CoordinatorMode::Active => self.advance_active(),
423            CoordinatorMode::CleaningUp => self.advance_cleanup(),
424            CoordinatorMode::Complete => Ok(Vec::new()),
425        }
426    }
427
428    pub fn is_complete(&self) -> bool {
429        self.mode == CoordinatorMode::Complete
430    }
431
432    pub fn completion(&self) -> Option<HandoffCompletion> {
433        self.completion
434    }
435
436    fn advance_active(&mut self) -> Result<Vec<IssuedHandoffAction>> {
437        if self.order == HandoffOrder::SourceFirst
438            && self.source.held
439            && !self.destination.reserve_issued
440        {
441            return Ok(vec![self.issue_reserve_destination()]);
442        }
443        if self.order == HandoffOrder::DestinationFirst
444            && self.destination.reserved
445            && !self.source.submit_issued
446        {
447            return Ok(vec![self.issue_submit_prefill()]);
448        }
449        if self.source.held && self.destination.reserved && !self.transfer.issued {
450            self.transfer.issued = true;
451            let transfer_timing = self
452                .source
453                .transfer_timing
454                .expect("held source must retain transfer timing");
455            let transferable_prompt_tokens = self
456                .destination
457                .transferable_prompt_tokens
458                .expect("reserved destination must report its transferable footprint");
459            return Ok(vec![
460                self.issue(HandoffAction::StartTransfer {
461                    handoff_id: self.handoff_id,
462                    delay_ms: transfer_timing
463                        .delay_ms(transferable_prompt_tokens)
464                        .unwrap_or(self.fallback_transfer_delay_ms),
465                }),
466            ]);
467        }
468        if self.transfer.completed && !self.destination.activation_issued {
469            self.destination.activation_issued = true;
470            return Ok(vec![self.issue(HandoffAction::ActivateDestination {
471                handoff_id: self.handoff_id,
472            })]);
473        }
474        if self.destination.activation_applied && !self.source.release_issued {
475            self.source.release_issued = true;
476            return Ok(vec![self.issue(HandoffAction::ReleaseSource {
477                handoff_id: self.handoff_id,
478            })]);
479        }
480        if self.source.cleanup_done {
481            return Ok(vec![self.complete()]);
482        }
483        Ok(Vec::new())
484    }
485
486    fn begin_cleanup(&mut self) -> Result<Vec<IssuedHandoffAction>> {
487        if self.mode == CoordinatorMode::Complete {
488            return Ok(Vec::new());
489        }
490        self.mode = CoordinatorMode::CleaningUp;
491        self.advance_cleanup()
492    }
493
494    fn advance_cleanup(&mut self) -> Result<Vec<IssuedHandoffAction>> {
495        let mut actions = Vec::new();
496        if self.source.submit_issued && !self.source.cancel_issued && !self.source.cleanup_done {
497            self.source.cancel_issued = true;
498            actions.push(self.issue(HandoffAction::CancelSource {
499                handoff_id: self.handoff_id,
500            }));
501        }
502        if self.destination.reserve_issued
503            && !self.destination.cancel_issued
504            && !self.destination.cleanup_done
505        {
506            self.destination.cancel_issued = true;
507            actions.push(self.issue(HandoffAction::CancelDestination {
508                handoff_id: self.handoff_id,
509            }));
510        }
511        if actions.is_empty()
512            && (!self.source.submit_issued || self.source.cleanup_done)
513            && (!self.destination.reserve_issued || self.destination.cleanup_done)
514        {
515            actions.push(self.complete());
516        }
517        Ok(actions)
518    }
519
520    fn issue_submit_prefill(&mut self) -> IssuedHandoffAction {
521        self.source.submit_issued = true;
522        self.issue(HandoffAction::SubmitPrefill {
523            handoff_id: self.handoff_id,
524        })
525    }
526
527    fn issue_reserve_destination(&mut self) -> IssuedHandoffAction {
528        self.destination.reserve_issued = true;
529        self.issue(HandoffAction::ReserveDestination {
530            handoff_id: self.handoff_id,
531        })
532    }
533
534    fn complete(&mut self) -> IssuedHandoffAction {
535        self.completion = Some(if self.mode == CoordinatorMode::CleaningUp {
536            HandoffCompletion::Canceled
537        } else {
538            HandoffCompletion::Success
539        });
540        self.mode = CoordinatorMode::Complete;
541        let action = self.issue(HandoffAction::Complete {
542            handoff_id: self.handoff_id,
543        });
544        self.actions.issued.clear();
545        self.actions.outcomes.clear();
546        action
547    }
548
549    fn issue(&mut self, action: HandoffAction) -> IssuedHandoffAction {
550        let id = HandoffActionId(self.actions.next_id);
551        self.actions.next_id = self
552            .actions
553            .next_id
554            .checked_add(1)
555            .expect("handoff action ID overflow");
556        let previous = self.actions.issued.insert(id, action);
557        debug_assert!(previous.is_none());
558        IssuedHandoffAction { id, action }
559    }
560
561    fn validate_handoff(&self, handoff_id: HandoffId) -> Result<()> {
562        if handoff_id != self.handoff_id {
563            bail!("fact belongs to a different handoff");
564        }
565        Ok(())
566    }
567}
568
569pub fn validate_transfer_delay_ms(transfer_delay_ms: Option<f64>) -> Result<()> {
570    let Some(delay_ms) = transfer_delay_ms else {
571        return Ok(());
572    };
573    if !delay_ms.is_finite() || delay_ms < 0.0 {
574        bail!("invalid handoff transfer delay {delay_ms}");
575    }
576    Ok(())
577}
578
579pub fn validate_transfer_timing(transfer_timing: HandoffTransferTiming) -> Result<()> {
580    if let Some(bandwidth_gb_s) = transfer_timing.bandwidth_gb_s
581        && (!bandwidth_gb_s.is_finite() || bandwidth_gb_s < 0.0)
582    {
583        bail!("invalid handoff transfer bandwidth {bandwidth_gb_s}");
584    }
585    validate_transfer_delay_ms(transfer_timing.full_prompt_delay_ms())
586}
587
588fn require_outcome(outcome: &HandoffActionOutcome, allowed: &[HandoffActionOutcome]) -> Result<()> {
589    if allowed.contains(outcome) {
590        return Ok(());
591    }
592    bail!("invalid handoff action outcome {outcome:?}")
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598    use crate::engine::TransferTimingMode;
599    use uuid::Uuid;
600
601    fn start_transfer_delay(timing: HandoffTransferTiming, fallback_ms: f64) -> f64 {
602        let handoff_id = HandoffId::new(Uuid::from_u128(1));
603        let mut coordinator = HandoffCoordinatorCore::new_with_fallback(
604            handoff_id,
605            HandoffOrder::SourceFirst,
606            fallback_ms,
607        );
608        let submit = coordinator.start().unwrap().remove(0);
609        assert!(matches!(submit.action, HandoffAction::SubmitPrefill { .. }));
610        assert!(
611            coordinator
612                .on_action_outcome(submit.id, HandoffActionOutcome::Submitted)
613                .unwrap()
614                .is_empty()
615        );
616        let reserve = coordinator
617            .on_fact(HandoffFact::SourceHeld {
618                handoff_id,
619                transfer_timing: timing,
620            })
621            .unwrap()
622            .remove(0);
623        assert!(matches!(
624            reserve.action,
625            HandoffAction::ReserveDestination { .. }
626        ));
627        assert!(
628            coordinator
629                .on_action_outcome(reserve.id, HandoffActionOutcome::Accepted)
630                .unwrap()
631                .is_empty()
632        );
633        let transfer = coordinator
634            .on_fact(HandoffFact::DestinationReserved {
635                handoff_id,
636                transferable_prompt_tokens: 3,
637            })
638            .unwrap()
639            .remove(0);
640        match transfer.action {
641            HandoffAction::StartTransfer { delay_ms, .. } => delay_ms,
642            action => panic!("expected transfer action, got {action:?}"),
643        }
644    }
645
646    #[test]
647    fn missing_or_zero_bandwidth_uses_configured_fallback() {
648        let missing = HandoffTransferTiming {
649            mode: TransferTimingMode::DestinationMissing,
650            full_prompt_tokens: 10,
651            kv_bytes_per_token: None,
652            bandwidth_gb_s: None,
653        };
654        assert_eq!(start_transfer_delay(missing, 7.5), 7.5);
655
656        let zero_bandwidth = HandoffTransferTiming {
657            kv_bytes_per_token: Some(1024),
658            bandwidth_gb_s: Some(0.0),
659            ..missing
660        };
661        validate_transfer_timing(zero_bandwidth).unwrap();
662        assert_eq!(start_transfer_delay(zero_bandwidth, 7.5), 7.5);
663    }
664
665    #[test]
666    fn negative_bandwidth_is_rejected() {
667        let timing = HandoffTransferTiming {
668            mode: TransferTimingMode::FullPrompt,
669            full_prompt_tokens: 10,
670            kv_bytes_per_token: Some(1024),
671            bandwidth_gb_s: Some(-1.0),
672        };
673        assert!(validate_transfer_timing(timing).is_err());
674    }
675}
aisimulate_core/replay/handoff.rs

aisimulate_core/replay/
handoff.rs
