aisimulate-core 0.12.0

Engine-neutral inference simulation, deterministic replay, and performance modeling
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::engine::common::protocols::{KvTransferTimingMode, WorkerType as CoreWorkerType};
use crate::engine::{
    HandoffTransferTiming, TransferTimingMode, WorkerType,
    prefill_handoff_delay_ms as native_prefill_handoff_delay_ms,
};

pub fn prefill_handoff_transfer_timing(
    num_input_tokens: usize,
    kv_transfer_bandwidth: Option<f64>,
    kv_bytes_per_token: Option<usize>,
    mode: KvTransferTimingMode,
) -> HandoffTransferTiming {
    HandoffTransferTiming {
        mode: match mode {
            KvTransferTimingMode::FullPrompt => TransferTimingMode::FullPrompt,
            KvTransferTimingMode::DestinationMissing => TransferTimingMode::DestinationMissing,
        },
        full_prompt_tokens: num_input_tokens,
        kv_bytes_per_token,
        bandwidth_gb_s: kv_transfer_bandwidth,
    }
}

/// Compute the modeled handoff delay after a prefill worker emits its terminal token.
///
/// NOTE: this intentionally does not model the internal prefill TTFT itself accurately, and the
/// exact prefill/decode boundary is backend dependent. For now we only care about decode-visible
/// TTFT, which is what the client observes, so modeling the delay as prefill-to-decode handoff is
/// good enough.
pub fn compute_prefill_handoff_delay_ms(
    worker_type: CoreWorkerType,
    completed: bool,
    num_input_tokens: usize,
    kv_transfer_bandwidth: Option<f64>,
    kv_bytes_per_token: Option<usize>,
) -> Option<f64> {
    let delay_ms = native_prefill_handoff_delay_ms(
        match worker_type {
            CoreWorkerType::Aggregated => WorkerType::Aggregated,
            CoreWorkerType::Prefill => WorkerType::Prefill,
            CoreWorkerType::Decode => WorkerType::Decode,
        },
        completed,
        num_input_tokens,
        kv_transfer_bandwidth,
        kv_bytes_per_token,
    );
    match delay_ms {
        Some(delay_ms) => {
            tracing::debug!(
                num_input_tokens,
                bandwidth_gb_s = kv_transfer_bandwidth,
                delay_ms = format!("{delay_ms:.2}"),
                "KV handoff delay for prefill completion"
            );
            Some(delay_ms)
        }
        None => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_prefill_handoff_delay_only_applies_to_completed_prefill() {
        let delay_ms = compute_prefill_handoff_delay_ms(
            CoreWorkerType::Prefill,
            true,
            128,
            Some(1.0),
            Some(1_000_000),
        )
        .expect("prefill completion should produce a handoff delay");
        assert!((delay_ms - 128.0).abs() < 1e-9);

        assert!(
            compute_prefill_handoff_delay_ms(
                CoreWorkerType::Prefill,
                false,
                128,
                Some(1.0),
                Some(1_000_000),
            )
            .is_none()
        );
        assert!(
            compute_prefill_handoff_delay_ms(
                CoreWorkerType::Decode,
                true,
                128,
                Some(1.0),
                Some(1_000_000),
            )
            .is_none()
        );
    }
}