skippy-server 0.76.1

Embedded Skippy staged runtime server
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use skippy_protocol::StageConfig;
use skippy_scheduler::{CacheAffinity, StageCacheAffinity};

use super::{KvStageIntegration, PrefillKvIdentity, StagePrefixCachePayload};

impl KvStageIntegration {
    /// Inspect cache value for scheduling without mutating LRU recency.
    pub fn peek_cache_affinity(
        &self,
        config: &StageConfig,
        identities: &[PrefillKvIdentity],
    ) -> CacheAffinity {
        if !self.should_lookup() {
            return CacheAffinity::default();
        }
        let radix = match self.radix.try_lock() {
            Ok(radix) => radix,
            Err(std::sync::TryLockError::WouldBlock) => return CacheAffinity::default(),
            Err(std::sync::TryLockError::Poisoned(poisoned)) => poisoned.into_inner(),
        };
        let cache_epoch = radix.epoch();
        let matched_tokens = identities
            .iter()
            .filter_map(|identity| match self.payload {
                StagePrefixCachePayload::ResidentKv => radix
                    .peek_resident(&identity.namespace, &identity.token_ids)
                    .map(|hit| hit.matched_tokens),
                StagePrefixCachePayload::KvRecurrent | StagePrefixCachePayload::FullState => radix
                    .peek_recurrent(&identity.namespace, &identity.token_ids)
                    .map(|hit| hit.matched_tokens),
                StagePrefixCachePayload::Disabled => None,
            })
            .filter(|matched_tokens| self.meets_shared_prefix_min_tokens(*matched_tokens))
            .max()
            .unwrap_or(0);
        if matched_tokens == 0 {
            return CacheAffinity::default();
        }
        // Layer count is a deterministic first-order proxy for work saved on
        // heterogeneous stages. The policy type supports measured weights once
        // stage timing calibration is available.
        let prefill_cost_per_token =
            u64::from(config.layer_end.saturating_sub(config.layer_start).max(1));
        CacheAffinity::from_stage(StageCacheAffinity {
            stage_index: config.stage_index,
            matched_tokens,
            prefill_cost_per_token,
            restore_cost: 0,
            cache_epoch,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{Arc, mpsc},
        time::{Duration, Instant},
    };

    use skippy_protocol::{LoadMode, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload};

    use super::*;

    fn test_config() -> StageConfig {
        StageConfig {
            run_id: "run".to_string(),
            topology_id: "topology".to_string(),
            model_id: "hugging-quants/Llama-3.2-1B-Instruct-GGUF:Q4_K_M".to_string(),
            package_ref: None,
            manifest_sha256: None,
            source_model_path: None,
            source_model_sha256: None,
            source_model_bytes: None,
            materialized_path: None,
            materialized_pinned: false,
            model_path: None,
            projector_path: None,
            stage_id: "stage-0".to_string(),
            stage_index: 0,
            layer_start: 0,
            layer_end: 4,
            ctx_size: 512,
            lane_count: 2,
            n_batch: None,
            n_ubatch: None,
            n_gpu_layers: 0,
            mmap: None,
            mlock: false,
            repack: false,
            op_offload: None,
            no_host_buffer: false,
            check_tensors: false,
            direct_io: false,
            main_gpu: None,
            split_mode: skippy_protocol::SplitMode::Auto,
            cache_type_k: "f16".to_string(),
            cache_type_v: "f16".to_string(),
            flash_attn_type: Default::default(),
            kv_offload: None,
            kv_unified: None,
            swa_full: None,
            cache_idle_slots: None,
            filter_tensors_on_load: false,
            resident_tensor_names: Vec::new(),
            selected_device: None,
            kv_cache: Some(StageKvCacheConfig {
                mode: StageKvCacheMode::LookupRecord,
                payload: StageKvCachePayload::ResidentKv,
                max_entries: 8,
                max_bytes: 0,
                min_tokens: 64,
                shared_prefix_stride_tokens: 32,
                shared_prefix_record_limit: 2,
            }),
            native_mtp_enabled: true,
            load_mode: LoadMode::RuntimeSlice,
            bind_addr: "127.0.0.1:0".to_string(),
            upstream: None,
            downstream: None,
            ..StageConfig::default()
        }
    }

    #[test]
    fn busy_radix_returns_cold_affinity_without_waiting() {
        let config = test_config();
        let integration =
            KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
                .unwrap()
                .unwrap();
        let radix = Arc::clone(&integration.radix);
        let (locked_tx, locked_rx) = mpsc::channel();
        let (release_tx, release_rx) = mpsc::channel();
        let holder = std::thread::spawn(move || {
            let _guard = radix.lock().unwrap();
            locked_tx.send(()).unwrap();
            release_rx.recv().unwrap();
        });
        locked_rx.recv().unwrap();

        let started = Instant::now();
        assert_eq!(
            integration.peek_cache_affinity(&config, &[]),
            CacheAffinity::default()
        );
        assert!(started.elapsed() < Duration::from_millis(100));

        release_tx.send(()).unwrap();
        holder.join().unwrap();
    }
}