sipp-rs 0.1.0

Unified Rust library for extensible Sipp inference
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

use crate::runtime::config::KvReuseMode;
use crate::runtime::scheduler::BatchContributionKind;

use super::{unique_slot_first_use, InferenceRuntime};

impl InferenceRuntime {
    pub(super) fn capture_prefix_snapshots(
        &mut self,
        plan: &crate::runtime::scheduler::SharedBatchPlan,
    ) {
        let mut seen_slots: u64 = 0;
        for contribution in &plan.contributions {
            if contribution.kind != BatchContributionKind::Prefill
                || !unique_slot_first_use(&mut seen_slots, contribution.slot_index)
            {
                continue;
            }
            let Some(slot) = self.slot_scheduler.slots.get(contribution.slot_index) else {
                continue;
            };
            let Some(request) = slot.request() else {
                continue;
            };
            let Some(terminal_token_count) =
                decode_seed_snapshot_token_count(request.prompt_tokens.len())
            else {
                continue;
            };
            if slot.mirror.current_kv_tokens.len() < terminal_token_count {
                continue;
            }
            let snapshot_tokens = &slot.mirror.current_kv_tokens[..terminal_token_count];
            // Pure snapshot mode evicts the sequence at completion, so it must
            // materialize before finalization. Live+snapshot keeps the
            // sequence idle and can defer the expensive state readback. Both
            // paths store the prompt-minus-one identity; restore trims native
            // state captured after that boundary.
            if request.cache_mode == KvReuseMode::StateSnapshot {
                self.kv_cache.capture_prefix_snapshot(
                    &self.native_runtime,
                    self.model_fingerprint,
                    &request.context_key,
                    slot.seq_id,
                    snapshot_tokens,
                    terminal_token_count,
                );
            } else {
                self.kv_cache.queue_prefix_snapshot(
                    self.model_fingerprint,
                    &request.context_key,
                    slot.seq_id,
                    slot.lease_generation,
                    snapshot_tokens,
                    terminal_token_count,
                );
            }
        }
    }
}

pub(super) fn decode_seed_snapshot_token_count(prompt_len: usize) -> Option<usize> {
    prompt_len
        .checked_sub(1)
        .filter(|&token_count| token_count > 0)
}

#[cfg(test)]
#[path = "../../tests/runtime/inference_runtime/prefix_snapshots_tests.rs"]
mod prefix_snapshots_tests;