File size: 6,737 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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Parallel decisions over one exact shared state using a native prefix cache."""
from __future__ import annotations
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,
# This value occurs after the extracted evidence boundary.
"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"]
if prompt.count(payload) != 1:
raise ValueError("Cannot locate the unmodified evidence payload in the chat template")
evidence = json.dumps({"evidence": state}, ensure_ascii=False)[:-1]
if not payload.startswith(evidence):
raise ValueError("Evidence serialization changed")
text = prompt[: prompt.index(payload)] + evidence
# Appending JSON punctuation can merge with the final boundary token.
return tokenizer.encode(text, add_special_tokens=False)[:-1]
def _suffix_layout(sequences: list[list[int]], prefix_length: int, pad_id: int):
if not sequences or any(not sequence for sequence in sequences):
raise ValueError("Every decision needs a nonempty suffix")
width = max(map(len, sequences))
ids, masks, positions, ends = [], [], [], []
for sequence in sequences:
padding = width - len(sequence)
ids.append(sequence + [pad_id] * padding)
masks.append([1] * (prefix_length + len(sequence)) + [0] * padding)
positions.append(list(range(prefix_length, prefix_length + len(sequence))) + [0] * padding)
ends.append(len(sequence) - 1)
return {"input_ids": ids, "attention_mask": masks, "position_ids": positions}, ends
def score_shared(model, tokenizer, rows: list[dict], metadata: dict, max_tokens: int = 4096):
"""Return all option distributions together after one state prefill."""
import torch
if not rows or any(row["state"] != rows[0]["state"] for row in rows[1:]):
raise ValueError("Shared scoring requires one nonempty exact state")
if len({row["id"] for row in rows}) != len(rows):
raise ValueError("Decision IDs must be unique")
started = time.perf_counter()
encoded = [encode_prompt(tokenizer, row, max_tokens) for row in rows]
prefix = _state_prefix(tokenizer, rows[0]["state"])
if not prefix or any(ids[: len(prefix)] != prefix or len(ids) <= len(prefix) for ids, _, _ in encoded):
raise ValueError("The fixed state prefix does not match every full prompt")
pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
if pad is None:
raise ValueError("Tokenizer requires a padding or EOS token")
layout, ends = _suffix_layout([ids[len(prefix) :] for ids, _, _ in encoded], len(prefix), pad)
selected_positions = sorted(set(ends))
encode_seconds = time.perf_counter() - started
device = next(model.parameters()).device
sync = lambda: torch.cuda.synchronize(device) if device.type == "cuda" else None
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-position logits needed by shared scoring")
model.eval()
with torch.inference_mode():
sync()
mark = time.perf_counter()
output = model(
input_ids=torch.tensor([prefix], dtype=torch.long, device=device),
attention_mask=torch.ones((1, len(prefix)), dtype=torch.long, device=device),
use_cache=True,
return_dict=True,
logits_to_keep=1,
)
cache = output.past_key_values
del output
sync()
prefill_seconds = time.perf_counter() - mark
if cache is None or cache.get_seq_length() != len(prefix):
raise RuntimeError("Invalid native prefix cache")
if not callable(getattr(cache, "reorder_cache", None)):
raise RuntimeError("Native cache does not support duplicate branch selection")
sync()
mark = time.perf_counter()
cache.reorder_cache(torch.zeros(len(rows), dtype=torch.long, device=device))
sync()
replicate_seconds = time.perf_counter() - mark
inputs = {key: torch.tensor(value, dtype=torch.long, device=device) for key, value in layout.items()}
sync()
mark = time.perf_counter()
output = model(
**inputs,
past_key_values=cache,
use_cache=True,
return_dict=True,
logits_to_keep=torch.tensor(selected_positions, dtype=torch.long, device=device),
)
sync()
suffix_seconds = time.perf_counter() - mark
results = []
for index, (row, (ids, slots, prompt_hash)) in enumerate(zip(rows, encoded)):
vocabulary = output.logits[index, selected_positions.index(ends[index]), :].float()
selected = vocabulary[slots].cpu().tolist()
results.append(
{
"id": row["id"],
"option_ids": [option["id"] for option in row["options"]],
"probabilities": softmax(selected),
"option_logits": selected,
"input_tokens": len(ids),
"prompt_sha256": prompt_hash,
"prompt_version": PROMPT_VERSION,
"model": {**metadata, "serving_config": "native-state-prefix-parallel-v1"},
"readout": "native selected suffix-position logits",
"probability_status": "conditional option score; uncalibrated as decision confidence",
}
)
del output, cache
sync()
timing = {
"total_seconds": time.perf_counter() - started,
"encode_seconds": encode_seconds,
"prefix_tokens": len(prefix),
"prefill_seconds": prefill_seconds,
"replicate_seconds": replicate_seconds,
"suffix_forward_seconds": suffix_seconds,
"batch_size": len(rows),
"true_suffix_tokens": sum(len(ids) - len(prefix) for ids, _, _ in encoded),
"padded_suffix_tokens": len(rows) * len(layout["input_ids"][0]),
}
return results, timing
|