File size: 5,787 Bytes
7845694 | 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 | """Serialized decisions over one-state native prefix caches."""
from __future__ import annotations
import copy
import hashlib
import inspect
import json
import time
from .core import direct_messages, softmax
from .direct import PROMPT_VERSION, encode_prompt
def _state_prefix(tokenizer, state) -> list[int]:
row = {
"id": "prefix-only",
"state": state,
"question": "prefix boundary placeholder",
"options": [{"id": "yes", "description": "Yes"}, {"id": "no", "description": "No"}],
}
turns = direct_messages(row)
prompt = tokenizer.apply_chat_template(
turns, tokenize=False, add_generation_prompt=True, enable_thinking=False
)
payload = turns[-1]["content"]
evidence = json.dumps({"evidence": state}, ensure_ascii=False)[:-1]
if prompt.count(payload) != 1 or not payload.startswith(evidence):
raise ValueError("Cannot establish a deterministic evidence prefix")
text = prompt[: prompt.index(payload)] + evidence
return tokenizer.encode(text, add_special_tokens=False)[:-1]
def _cached_forward(model, inputs):
parameters = inspect.signature(model.forward).parameters
if "logits_to_keep" not in parameters and hasattr(model, "get_base_model"):
parameters = inspect.signature(model.get_base_model().forward).parameters
if "logits_to_keep" not in parameters:
raise RuntimeError("Model lacks selective last-position logits")
return model(**inputs, use_cache=True, return_dict=True, logits_to_keep=1)
class SerialPrefixScorer:
"""Cache the current state, then score independent copied suffix branches."""
def __init__(self, model, tokenizer, metadata: dict, max_tokens: int = 4096):
self.model = model
self.tokenizer = tokenizer
self.metadata = {**metadata, "serving_config": "native-state-prefix-cache-v1"}
self.max_tokens = max_tokens
self.device = next(model.parameters()).device
self.cache = None
self.state = None
self.prefix = None
def score(self, row: dict) -> dict:
import torch
started = time.perf_counter()
ids, slots, prompt_hash = encode_prompt(self.tokenizer, row, self.max_tokens)
hit = self.cache is not None and row["state"] == self.state
prefix = self.prefix if hit else _state_prefix(self.tokenizer, row["state"])
if not prefix or ids[: len(prefix)] != prefix or len(ids) <= len(prefix):
raise ValueError("State prefix does not match the full prompt")
sync = lambda: torch.cuda.synchronize(self.device) if self.device.type == "cuda" else None
prefill_seconds = 0.0
self.model.eval()
with torch.inference_mode():
if not hit:
self.cache = self.state = self.prefix = None
sync()
mark = time.perf_counter()
output = _cached_forward(
self.model,
{
"input_ids": torch.tensor([prefix], dtype=torch.long, device=self.device),
"attention_mask": torch.ones((1, len(prefix)), dtype=torch.long, device=self.device),
},
)
self.cache = output.past_key_values
del output
sync()
prefill_seconds = time.perf_counter() - mark
if self.cache is None or self.cache.get_seq_length() != len(prefix):
raise RuntimeError("Invalid native prefix cache")
self.state, self.prefix = row["state"], prefix
sync()
mark = time.perf_counter()
branch = copy.deepcopy(self.cache)
sync()
copy_seconds = time.perf_counter() - mark
inputs = {
"input_ids": torch.tensor([ids[len(prefix) :]], dtype=torch.long, device=self.device),
"attention_mask": torch.ones((1, len(ids)), dtype=torch.long, device=self.device),
"past_key_values": branch,
}
sync()
mark = time.perf_counter()
output = _cached_forward(self.model, inputs)
sync()
suffix_seconds = time.perf_counter() - mark
vocabulary = output.logits[0, -1, :].float()
selected_tensor = vocabulary[slots]
selected = selected_tensor.cpu().tolist()
result = {
"id": row["id"],
"option_ids": [option["id"] for option in row["options"]],
"probabilities": softmax(selected),
"option_logits": selected,
"answer_token_ids": slots,
"input_tokens": len(ids),
"prompt_sha256": prompt_hash,
"prompt_version": PROMPT_VERSION,
"model": self.metadata,
"readout": "native-state-prefix-cache-last-position",
"probability_status": "conditional option score; uncalibrated as decision confidence",
"cache_hit": hit,
"prefix_tokens": len(prefix),
"prefix_sha256": hashlib.sha256(json.dumps(prefix).encode()).hexdigest(),
"allowed_token_mass": float(
(selected_tensor.logsumexp(-1) - vocabulary.logsumexp(-1)).exp()
),
"full_vocab_argmax_id": int(vocabulary.argmax()),
"prefill_seconds": prefill_seconds,
"copy_seconds": copy_seconds,
"suffix_forward_seconds": suffix_seconds,
}
del output, branch
sync()
result["forward_seconds"] = prefill_seconds + suffix_seconds
result["total_seconds"] = time.perf_counter() - started
return result
|