skippy_cache/
identity.rs

1use skippy_protocol::StageConfig;
2
3pub const NATIVE_KV_RUNTIME_ABI_VERSION: &str = "stage-abi-0.1/native-kv-page-v1";
4pub const NATIVE_KV_DTYPE: &str = "ggml-native-kv";
5const NATIVE_KV_LAYER_CONTIGUOUS_LAYOUT: i32 = 4;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct PrefixIdentity {
9    pub prefix_hash: String,
10    pub page_id: String,
11    pub token_start: u64,
12    pub token_count: u64,
13}
14
15pub fn prefix_identity(
16    config: &StageConfig,
17    token_start: u64,
18    token_ids: &[i32],
19) -> PrefixIdentity {
20    prefix_identity_with_namespace(config, token_start, token_ids, None)
21}
22
23pub fn prefix_identity_with_namespace(
24    config: &StageConfig,
25    token_start: u64,
26    token_ids: &[i32],
27    cache_namespace: Option<&str>,
28) -> PrefixIdentity {
29    let token_count = token_ids.len() as u64;
30    let prefix_hash = prefix_hash_with_namespace(config, token_start, token_ids, cache_namespace);
31    let page_id = page_id(config, token_start, token_count, &prefix_hash);
32    PrefixIdentity {
33        prefix_hash,
34        page_id,
35        token_start,
36        token_count,
37    }
38}
39
40pub fn prefix_hash(config: &StageConfig, token_start: u64, token_ids: &[i32]) -> String {
41    prefix_hash_with_namespace(config, token_start, token_ids, None)
42}
43
44pub fn prefix_hash_with_namespace(
45    config: &StageConfig,
46    token_start: u64,
47    token_ids: &[i32],
48    cache_namespace: Option<&str>,
49) -> String {
50    let mut hasher = blake3::Hasher::new();
51    hasher.update(config.model_id.as_bytes());
52    hasher.update(config.topology_id.as_bytes());
53    hasher.update(config.stage_id.as_bytes());
54    hasher.update(&config.stage_index.to_le_bytes());
55    hasher.update(&config.layer_start.to_le_bytes());
56    hasher.update(&config.layer_end.to_le_bytes());
57    hasher.update(NATIVE_KV_RUNTIME_ABI_VERSION.as_bytes());
58    hasher.update(&NATIVE_KV_LAYER_CONTIGUOUS_LAYOUT.to_le_bytes());
59    hasher.update(NATIVE_KV_DTYPE.as_bytes());
60    hasher.update(format!("ctx:{}", config.ctx_size).as_bytes());
61    if let Some(cache_namespace) = cache_namespace {
62        hasher.update(b"openai-cache-namespace-v1");
63        hasher.update(cache_namespace.as_bytes());
64    }
65    hasher.update(&token_start.to_le_bytes());
66    for token_id in token_ids {
67        hasher.update(token_id.to_le_bytes().as_slice());
68    }
69    format!("blake3:{}", hasher.finalize().to_hex())
70}
71
72pub fn page_id(
73    config: &StageConfig,
74    token_start: u64,
75    token_count: u64,
76    prefix_hash: &str,
77) -> String {
78    let digest = prefix_hash
79        .strip_prefix("blake3:")
80        .or_else(|| prefix_hash.strip_prefix("sha256:"))
81        .unwrap_or(prefix_hash);
82    let short = digest.get(..16).unwrap_or(digest);
83    format!(
84        "{}:{}:{}:{}:{}",
85        config.stage_id, token_start, token_count, config.layer_start, short
86    )
87}
88
89pub fn activation_page_id(page_id: &str, activation_width: i32) -> String {
90    format!("act:{}:w{}", page_id, activation_width.max(0))
91}
skippy_cache/identity.rs

skippy_cache/
identity.rs
