src/pqc_kv_cache/entry.py
2.2 KB · 72 lines · python Raw
1 """KV cache entry data structures."""
2
3 from __future__ import annotations
4
5 from dataclasses import asdict, dataclass
6 from typing import Any
7
8
9 @dataclass(frozen=True)
10 class EntryMetadata:
11 """Non-secret metadata about a KV cache entry."""
12
13 tenant_id: str
14 session_id: str
15 layer_idx: int
16 position: int # token position in sequence
17 token_id: int = -1 # optional vocab id (for debugging; not required)
18 kv_role: str = "both" # "key" | "value" | "both"
19
20 def to_dict(self) -> dict[str, Any]:
21 return asdict(self)
22
23
24 @dataclass
25 class KVCacheEntry:
26 """Plaintext KV cache entry - only exists inside TenantSession scope."""
27
28 metadata: EntryMetadata
29 key_tensor_bytes: bytes # raw bytes of the K vector for this position
30 value_tensor_bytes: bytes # raw bytes of the V vector for this position
31
32 def plaintext_size(self) -> int:
33 return len(self.key_tensor_bytes) + len(self.value_tensor_bytes)
34
35
36 @dataclass
37 class EncryptedEntry:
38 """AES-256-GCM encrypted KV cache entry."""
39
40 metadata: EntryMetadata
41 nonce: str # hex
42 ciphertext: str # hex (contains both K and V concatenated)
43 key_len: int # bytes of K portion (V starts after)
44 sequence_number: int
45
46 def to_dict(self) -> dict[str, Any]:
47 return {
48 "metadata": self.metadata.to_dict(),
49 "nonce": self.nonce,
50 "ciphertext": self.ciphertext,
51 "key_len": self.key_len,
52 "sequence_number": self.sequence_number,
53 }
54
55 @classmethod
56 def from_dict(cls, data: dict[str, Any]) -> EncryptedEntry:
57 meta = data["metadata"]
58 return cls(
59 metadata=EntryMetadata(
60 tenant_id=meta["tenant_id"],
61 session_id=meta["session_id"],
62 layer_idx=int(meta["layer_idx"]),
63 position=int(meta["position"]),
64 token_id=int(meta.get("token_id", -1)),
65 kv_role=meta.get("kv_role", "both"),
66 ),
67 nonce=data["nonce"],
68 ciphertext=data["ciphertext"],
69 key_len=int(data["key_len"]),
70 sequence_number=int(data["sequence_number"]),
71 )
72