Nucleus-Resynthesis / runtime /src /resynthesis /base_forward_cache.py
Wl6adams's picture
Add portable Release 188 generation runtime
919fd68 verified
Raw
History Blame Contribute Delete
64.6 kB
"""Content-addressed frozen-parent prefill caches.
The frozen Resynthesis parent is identical across optimizer steps. When a decode arm
resets KV state and re-runs the same prompt prefill, this cache returns the
prior ``ResynthesisParentForward`` tensors from CPU-pinned storage instead of repeating
the 4M tiled parent forward.
Training uses the narrower ``FrozenBackbonePrefillPacket`` contract below.
Only the frozen decoder's final hidden state and tiled summaries are durable;
the additive head, historical RBO, and both parent/current Fabric graphs remain
live on every replay. Its disk lifecycle intentionally matches the packed
token sidecars: content-addressed object roots, one nonblocking ownership lock,
partial artifacts, a durable progress frontier, fsync, atomic promotion, and a
manifest published last.
"""
from __future__ import annotations
import fcntl
import hashlib
import json
import os
import stat
import threading
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final
import torch
FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1: Final = (
"nnf.resynthesis.frozen_backbone_prefill_authority.v1"
)
FROZEN_BACKBONE_AUTHORITY_SCHEMA: Final = (
"nnf.resynthesis.frozen_backbone_prefill_authority.v2"
)
FROZEN_BACKBONE_PACKET_SCHEMA: Final = (
"nnf.resynthesis.frozen_backbone_prefill_packet.v1"
)
FROZEN_BACKBONE_SIDECAR_SCHEMA: Final = (
"nnf.resynthesis.frozen_backbone_prefill_sidecar.v1"
)
FROZEN_BACKBONE_SIDECAR_PROGRESS_SCHEMA: Final = (
"nnf.resynthesis.frozen_backbone_prefill_sidecar.progress.v1"
)
_FROZEN_BACKBONE_DTYPES: Final = {
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}
_FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS: Final = frozenset(
{
"parentSourceBundleSha256",
"sourceHashDiagnosticOnly",
"sourceHashAffectsExecution",
}
)
_FROZEN_BACKBONE_VALIDATION_CACHE_ENTRIES: Final = 8_192
@dataclass(frozen=True)
class _FrozenBackboneFileStamp:
"""Kernel-owned identity observed at one immutable-file boundary."""
device: int
inode: int
mode: int
size: int
mtime_ns: int
ctime_ns: int
@dataclass(frozen=True)
class _FrozenBackboneManifestValidation:
"""Exact content proof reusable while every immutable file stamp is stable."""
manifest: dict[str, object]
manifest_sha256: str
artifact_sha256s: tuple[str, str, str]
file_stamps: tuple[
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
]
_LEGACY_FROZEN_BACKBONE_INDEX_LOCK: Final = threading.Lock()
_LEGACY_FROZEN_BACKBONE_INDEXES: OrderedDict[
Path,
dict[str, tuple[Path, ...]],
] = OrderedDict()
_FROZEN_BACKBONE_VALIDATION_CACHE_LOCK: Final = threading.Lock()
_FROZEN_BACKBONE_VALIDATION_CACHE: OrderedDict[
tuple[int, Path, str],
_FrozenBackboneManifestValidation,
] = OrderedDict()
def _is_sha256_hex(value: object) -> bool:
return (
isinstance(value, str)
and len(value) == 64
and all(character in "0123456789abcdef" for character in value)
)
def _canonical_json_bytes(value: object) -> bytes:
return json.dumps(
value,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=True,
).encode("utf-8")
def _string_object_mapping_boundary(
value: object,
) -> dict[str, object] | None:
"""Return a typed JSON-object mapping without coercing malformed keys."""
if not isinstance(value, dict):
return None
result: dict[str, object] = {}
for key, item in value.items():
if not isinstance(key, str):
return None
result[key] = item
return result
def _source_independent_authority_record_boundary(
value: object,
) -> dict[str, object] | None:
"""Normalize mutable source observations and the parent artifact alias."""
record = _string_object_mapping_boundary(value)
if record is None:
return None
for diagnostic_field in _FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS:
record.pop(diagnostic_field, None)
parent_artifact_sha256 = record.get("parentModelArtifactSha256")
if _is_sha256_hex(parent_artifact_sha256):
record["parentCheckpointId"] = (
f"resynthesis-native-parent:{parent_artifact_sha256}"
)
if record.get("schema") == FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1:
record["schema"] = FROZEN_BACKBONE_AUTHORITY_SCHEMA
return record
def _atomic_json_boundary(path: Path, payload: dict[str, object]) -> None:
"""Publish one external-I/O record atomically and fsync its directory."""
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
with temporary.open("wb") as handle:
handle.write(_canonical_json_bytes(payload))
handle.write(b"\n")
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, path)
directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
try:
os.fsync(directory_fd)
finally:
os.close(directory_fd)
def _file_sha256_boundary(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
while chunk := handle.read(8 * 1024 * 1024):
digest.update(chunk)
return digest.hexdigest()
def _tensor_storage_bytes_boundary(value: torch.Tensor) -> bytes:
"""Serialize one contiguous CPU tensor without dtype conversion."""
cpu = value.detach().to(device="cpu").contiguous()
return cpu.view(torch.uint8).numpy().tobytes()
def _write_all_fd_boundary(descriptor: int, payload: bytes) -> None:
"""Write one complete staged artifact at the external-I/O boundary."""
remaining = memoryview(payload)
while remaining:
written = os.write(descriptor, remaining)
if written < 1:
raise OSError("frozen-backbone staged artifact write made no progress")
remaining = remaining[written:]
def _synchronize_open_files_boundary(descriptors: tuple[int, ...]) -> None:
"""Flush one already-written outer wave as a filesystem transaction group."""
for descriptor in descriptors:
os.fdatasync(descriptor)
def _synchronize_directories_boundary(paths: tuple[Path, ...]) -> None:
"""Persist manifest-last renames from deepest object root to cache root."""
unique_paths = sorted(
{path.expanduser().resolve() for path in paths},
key=lambda path: (len(path.parts), str(path)),
reverse=True,
)
for path in unique_paths:
descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
try:
os.fsync(descriptor)
finally:
os.close(descriptor)
def frozen_backbone_prompt_mask_sha256_boundary(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
) -> str:
"""Hash exact prompt/mask geometry and bytes at the CPU staging boundary."""
if (
input_ids.device.type != "cpu"
or attention_mask.device.type != "cpu"
or input_ids.ndim != 2
or input_ids.shape[1] < 1
or attention_mask.shape != input_ids.shape
or input_ids.dtype not in {torch.int32, torch.int64}
or attention_mask.dtype not in {
torch.bool,
torch.int32,
torch.int64,
}
):
raise ValueError(
"frozen-backbone prompt/mask identity requires CPU [batch, sequence]"
)
digest = hashlib.sha256()
digest.update(b"nnf.resynthesis.frozen_backbone_prompt_mask.v1\x00")
for value in (input_ids, attention_mask):
digest.update(str(value.dtype).encode("ascii"))
digest.update(b"\x00")
digest.update(_canonical_json_bytes(tuple(value.shape)))
digest.update(b"\x00")
digest.update(_tensor_storage_bytes_boundary(value))
return digest.hexdigest()
def frozen_backbone_position_policy_sha256_boundary(
*,
final_hidden_adapter_id: str,
dual_chunk_pretrain_length: int,
dual_chunk_local_size: int,
pretrained_rope_band_tokens: int,
prefill_tile_tokens: int,
prefill_summaries_per_tile: int,
) -> str:
"""Bind every position/tile rule that can change frozen decoder features."""
policy = {
"schema": "nnf.resynthesis.frozen_backbone_position_policy.v1",
"finalHiddenAdapterId": final_hidden_adapter_id,
"dualChunkPretrainLength": dual_chunk_pretrain_length,
"dualChunkLocalSize": dual_chunk_local_size,
"pretrainedRopeBandTokens": pretrained_rope_band_tokens,
"prefillTileTokens": prefill_tile_tokens,
"prefillSummariesPerTile": prefill_summaries_per_tile,
}
if (
not final_hidden_adapter_id
or any(
isinstance(value, bool) or not isinstance(value, int) or value < 1
for value in (
dual_chunk_pretrain_length,
dual_chunk_local_size,
pretrained_rope_band_tokens,
prefill_tile_tokens,
prefill_summaries_per_tile,
)
)
):
raise ValueError("frozen-backbone position policy is malformed")
return hashlib.sha256(_canonical_json_bytes(policy)).hexdigest()
@dataclass(frozen=True)
class FrozenBackboneCacheAuthority:
"""Complete immutable identity for one prompt-only decoder feature object."""
qualified_work_id: str
prompt_sha256: str
prompt_mask_sha256: str
parent_checkpoint_id: str
parent_manifest_payload_sha256: str
parent_model_artifact_sha256: str
parent_source_bundle_sha256: str = field(compare=False)
position_policy_sha256: str
token_dtype: str
feature_dtype: str
hidden_size: int
summary_positions: int
def _record_boundary(
self,
*,
schema: str,
parent_checkpoint_id: str,
) -> dict[str, object]:
"""Serialize one validated cache-authority representation.
Source-tree hashes intentionally do not participate. Training follows
a moving source tree; launch/status receipts retain that observation,
while reusable frozen-parent features remain bound to the immutable
model artifact, checkpoint, prompt geometry, and decoder policy.
"""
hash_fields = (
self.prompt_sha256,
self.prompt_mask_sha256,
self.parent_manifest_payload_sha256,
self.parent_model_artifact_sha256,
self.position_policy_sha256,
)
if (
not self.qualified_work_id
or not parent_checkpoint_id
or not all(_is_sha256_hex(value) for value in hash_fields)
or self.token_dtype not in {"int32", "int64"}
or self.feature_dtype not in _FROZEN_BACKBONE_DTYPES
or isinstance(self.hidden_size, bool)
or self.hidden_size < 1
or isinstance(self.summary_positions, bool)
or self.summary_positions < 1
):
raise ValueError("frozen-backbone cache authority is malformed")
return {
"schema": schema,
"qualifiedWorkId": self.qualified_work_id,
"qualifiedWorkIdSha256": hashlib.sha256(
self.qualified_work_id.encode("utf-8")
).hexdigest(),
"promptSha256": self.prompt_sha256,
"promptMaskSha256": self.prompt_mask_sha256,
"parentCheckpointId": parent_checkpoint_id,
"parentManifestPayloadSha256": (
self.parent_manifest_payload_sha256
),
"parentModelArtifactSha256": self.parent_model_artifact_sha256,
"positionPolicySha256": self.position_policy_sha256,
"tokenDtype": self.token_dtype,
"featureDtype": self.feature_dtype,
"hiddenSize": self.hidden_size,
"summaryPositions": self.summary_positions,
"backboneOnly": True,
"additiveHeadCached": False,
"historicalRboCached": False,
"fabricCached": False,
"targetEnteredForward": False,
}
def record_boundary(self) -> dict[str, object]:
"""Return the canonical Resynthesis authority for new objects."""
return self._record_boundary(
schema=FROZEN_BACKBONE_AUTHORITY_SCHEMA,
parent_checkpoint_id=(
"resynthesis-native-parent:"
f"{self.parent_model_artifact_sha256}"
),
)
def legacy_v1_record_boundary(self) -> dict[str, object] | None:
"""Return the historical source-bound record for a direct legacy probe.
A malformed or changed source observation cannot block current cache
identity. When the reported digest still matches the historical
object, this record gives lookup a constant-time read-only fast path.
A per-root normalized manifest index handles all other legacy sources.
"""
if not _is_sha256_hex(self.parent_source_bundle_sha256):
return None
record = self._record_boundary(
schema=FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1,
parent_checkpoint_id=self.parent_checkpoint_id,
)
record["parentSourceBundleSha256"] = (
self.parent_source_bundle_sha256
)
return record
def legacy_v1_sha256_boundary(self) -> str | None:
record = self.legacy_v1_record_boundary()
if record is None:
return None
return hashlib.sha256(_canonical_json_bytes(record)).hexdigest()
def sha256_boundary(self) -> str:
return hashlib.sha256(
_canonical_json_bytes(self.record_boundary())
).hexdigest()
def digest_tensors_boundary(
self,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return typed row identity tensors for the active packet."""
cache_key_t = torch.tensor(
tuple(bytes.fromhex(self.sha256_boundary())),
dtype=torch.uint8,
).reshape(1, 32)
work_id_t = torch.tensor(
tuple(
hashlib.sha256(
self.qualified_work_id.encode("utf-8")
).digest()
),
dtype=torch.uint8,
).reshape(1, 32)
prompt_mask_t = torch.tensor(
tuple(bytes.fromhex(self.prompt_mask_sha256)),
dtype=torch.uint8,
).reshape(1, 32)
return cache_key_t, work_id_t, prompt_mask_t
@dataclass(frozen=True)
class FrozenBackbonePrefillPacket:
"""Raw frozen-decoder features; all additive/RBO/Fabric work stays live."""
final_hidden_t: torch.Tensor
summary_hidden_t: torch.Tensor
summary_mask_t: torch.Tensor
input_positions_t: torch.Tensor
cache_key_sha256_t: torch.Tensor
work_id_sha256_t: torch.Tensor
prompt_mask_sha256_t: torch.Tensor
def validate_boundary(self) -> None:
batch_size = self.final_hidden_t.shape[0]
if (
self.final_hidden_t.ndim != 3
or self.final_hidden_t.shape[1] != 1
or self.summary_hidden_t.ndim != 3
or self.summary_hidden_t.shape[0] != batch_size
or self.summary_hidden_t.shape[1] < 1
or self.summary_hidden_t.shape[2]
!= self.final_hidden_t.shape[2]
or self.summary_hidden_t.dtype != self.final_hidden_t.dtype
or self.summary_hidden_t.device != self.final_hidden_t.device
or self.summary_mask_t.shape
!= self.summary_hidden_t.shape[:2]
or self.summary_mask_t.dtype != torch.bool
or self.summary_mask_t.device != self.final_hidden_t.device
or self.input_positions_t.shape != (batch_size,)
or self.input_positions_t.dtype != torch.long
or self.input_positions_t.device != self.final_hidden_t.device
or any(
value.shape != (batch_size, 32)
or value.dtype != torch.uint8
or value.device != self.final_hidden_t.device
for value in (
self.cache_key_sha256_t,
self.work_id_sha256_t,
self.prompt_mask_sha256_t,
)
)
):
raise ValueError("frozen-backbone prefill packet geometry differs")
torch._assert_async(
self.summary_mask_t.all(),
"frozen-backbone replay requires exact unpadded summary geometry",
)
torch._assert_async(
self.input_positions_t.gt(0).all(),
"frozen-backbone input positions must be positive",
)
@classmethod
def from_features_boundary(
cls,
*,
authority: FrozenBackboneCacheAuthority,
final_hidden_t: torch.Tensor,
summary_hidden_t: torch.Tensor,
input_positions_t: torch.Tensor,
) -> FrozenBackbonePrefillPacket:
expected_feature_dtype = _FROZEN_BACKBONE_DTYPES.get(
authority.feature_dtype
)
if (
final_hidden_t.shape
!= (1, 1, authority.hidden_size)
or summary_hidden_t.shape
!= (
1,
authority.summary_positions,
authority.hidden_size,
)
or final_hidden_t.dtype != expected_feature_dtype
or summary_hidden_t.dtype != expected_feature_dtype
or input_positions_t.numel() != 1
):
raise ValueError(
"one frozen-backbone sidecar object geometry differs"
)
cache_key_t, work_id_t, prompt_mask_t = (
authority.digest_tensors_boundary()
)
device = final_hidden_t.device
packet = cls(
final_hidden_t=final_hidden_t,
summary_hidden_t=summary_hidden_t,
summary_mask_t=torch.ones(
summary_hidden_t.shape[:2],
dtype=torch.bool,
device=device,
),
input_positions_t=input_positions_t.reshape(1).to(
device=device,
dtype=torch.long,
),
cache_key_sha256_t=cache_key_t.to(device=device),
work_id_sha256_t=work_id_t.to(device=device),
prompt_mask_sha256_t=prompt_mask_t.to(device=device),
)
packet.validate_boundary()
return packet
@classmethod
def stack_boundary(
cls,
packets: tuple[FrozenBackbonePrefillPacket, ...],
) -> FrozenBackbonePrefillPacket:
if not packets:
raise ValueError("frozen-backbone packet stack is empty")
for packet in packets:
packet.validate_boundary()
summary_positions = {packet.summary_hidden_t.shape[1] for packet in packets}
hidden_sizes = {packet.final_hidden_t.shape[2] for packet in packets}
dtypes = {packet.final_hidden_t.dtype for packet in packets}
devices = {packet.final_hidden_t.device for packet in packets}
if (
len(summary_positions) != 1
or len(hidden_sizes) != 1
or len(dtypes) != 1
or len(devices) != 1
):
raise ValueError(
"frozen-backbone rows require one exact summary geometry"
)
stacked = cls(
final_hidden_t=torch.cat(
tuple(packet.final_hidden_t for packet in packets),
dim=0,
),
summary_hidden_t=torch.cat(
tuple(packet.summary_hidden_t for packet in packets),
dim=0,
),
summary_mask_t=torch.cat(
tuple(packet.summary_mask_t for packet in packets),
dim=0,
),
input_positions_t=torch.cat(
tuple(packet.input_positions_t for packet in packets),
dim=0,
),
cache_key_sha256_t=torch.cat(
tuple(packet.cache_key_sha256_t for packet in packets),
dim=0,
),
work_id_sha256_t=torch.cat(
tuple(packet.work_id_sha256_t for packet in packets),
dim=0,
),
prompt_mask_sha256_t=torch.cat(
tuple(packet.prompt_mask_sha256_t for packet in packets),
dim=0,
),
)
stacked.validate_boundary()
return stacked
def to_device_boundary(
self,
device: torch.device,
) -> FrozenBackbonePrefillPacket:
def move(value: torch.Tensor) -> torch.Tensor:
# A CPU capture may already live on the requested device. Force an
# isolated copy in that case so the sidecar I/O worker cannot observe
# a later parent forward reusing the producer's capture buffers.
return value.detach().to(
device=device,
non_blocking=device.type == "cuda",
copy=value.device == device,
)
packet = FrozenBackbonePrefillPacket(
final_hidden_t=move(self.final_hidden_t),
summary_hidden_t=move(self.summary_hidden_t),
summary_mask_t=move(self.summary_mask_t),
input_positions_t=move(self.input_positions_t),
cache_key_sha256_t=move(self.cache_key_sha256_t),
work_id_sha256_t=move(self.work_id_sha256_t),
prompt_mask_sha256_t=move(self.prompt_mask_sha256_t),
)
packet.validate_boundary()
return packet
@dataclass(frozen=True)
class _FrozenBackboneSidecarPaths:
root: Path
manifest: Path
progress: Path
lock: Path
final_hidden: Path
summary_hidden: Path
input_positions: Path
def artifact_paths_boundary(self) -> tuple[Path, Path, Path]:
return (
self.final_hidden,
self.summary_hidden,
self.input_positions,
)
def _frozen_backbone_file_stamp_boundary(
path: Path,
) -> _FrozenBackboneFileStamp | None:
"""Observe one regular file without reading or trusting its content."""
try:
observed = path.stat()
except FileNotFoundError:
return None
if not stat.S_ISREG(observed.st_mode):
return None
return _FrozenBackboneFileStamp(
device=observed.st_dev,
inode=observed.st_ino,
mode=observed.st_mode,
size=observed.st_size,
mtime_ns=observed.st_mtime_ns,
ctime_ns=observed.st_ctime_ns,
)
def _frozen_backbone_sidecar_file_stamps_boundary(
paths: _FrozenBackboneSidecarPaths,
) -> (
tuple[
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
]
| None
):
"""Return one complete immutable-object stamp, or no reusable object."""
manifest_stamp = _frozen_backbone_file_stamp_boundary(paths.manifest)
artifact_stamps = tuple(
_frozen_backbone_file_stamp_boundary(path)
for path in paths.artifact_paths_boundary()
)
if manifest_stamp is None or any(
artifact_stamp is None for artifact_stamp in artifact_stamps
):
return None
final_stamp, summary_stamp, positions_stamp = artifact_stamps
assert final_stamp is not None
assert summary_stamp is not None
assert positions_stamp is not None
return (
manifest_stamp,
final_stamp,
summary_stamp,
positions_stamp,
)
def _build_legacy_frozen_backbone_manifest_index_boundary(
cache_root: Path,
) -> dict[str, tuple[Path, ...]]:
"""Build one read-only normalized index of legacy source-bound manifests.
The index is read-only and contains manifest paths, not copied tensor
artifacts. Its key is the authority record after removing only code-source
diagnostics. Artifact geometry and hashes are deliberately validated later
by the selected sidecar reader, exactly as they are for current objects.
"""
objects_root = cache_root.expanduser().resolve() / "objects"
if not objects_root.is_dir():
return {}
indexed: dict[str, list[Path]] = {}
for manifest_path in sorted(objects_root.glob("*/*/manifest.json")):
object_root = manifest_path.parent
object_sha256 = object_root.name
if (
not _is_sha256_hex(object_sha256)
or object_root.parent.name != object_sha256[:2]
):
continue
try:
raw = json.loads(manifest_path.read_text(encoding="utf-8"))
except (OSError, UnicodeDecodeError, json.JSONDecodeError):
# An incomplete unrelated cache object is not a candidate. Exact
# candidates still pass the strict validator before any mmap.
continue
manifest = _string_object_mapping_boundary(raw)
if (
manifest is None
or manifest.get("schema") != FROZEN_BACKBONE_SIDECAR_SCHEMA
or manifest.get("authoritySha256") != object_sha256
):
continue
stored_authority = _string_object_mapping_boundary(
manifest.get("authority")
)
if (
stored_authority is None
or not any(
field in stored_authority
for field in _FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS
)
or hashlib.sha256(
_canonical_json_bytes(stored_authority)
).hexdigest()
!= object_sha256
):
continue
normalized = _source_independent_authority_record_boundary(
stored_authority
)
if normalized is None:
continue
normalized_sha256 = hashlib.sha256(
_canonical_json_bytes(normalized)
).hexdigest()
indexed.setdefault(normalized_sha256, []).append(manifest_path)
return {
authority_sha256: tuple(paths)
for authority_sha256, paths in indexed.items()
}
def _legacy_frozen_backbone_manifest_index_boundary(
cache_root: Path,
) -> dict[str, tuple[Path, ...]]:
"""Return one single-flight index per resolved root and process."""
resolved_root = cache_root.expanduser().resolve()
with _LEGACY_FROZEN_BACKBONE_INDEX_LOCK:
if resolved_root in _LEGACY_FROZEN_BACKBONE_INDEXES:
cached = _LEGACY_FROZEN_BACKBONE_INDEXES[resolved_root]
_LEGACY_FROZEN_BACKBONE_INDEXES.move_to_end(resolved_root)
return cached
built = _build_legacy_frozen_backbone_manifest_index_boundary(
resolved_root
)
_LEGACY_FROZEN_BACKBONE_INDEXES[resolved_root] = built
while len(_LEGACY_FROZEN_BACKBONE_INDEXES) > 16:
_LEGACY_FROZEN_BACKBONE_INDEXES.popitem(last=False)
return built
class FrozenBackboneFeatureSidecar:
"""One exact prompt feature object using packed-sidecar durability rules."""
@staticmethod
def _paths_for_authority_boundary(
*,
cache_root: Path,
authority_sha256: str,
feature_dtype: str,
) -> _FrozenBackboneSidecarPaths:
object_root = (
cache_root.expanduser().resolve()
/ "objects"
/ authority_sha256[:2]
/ authority_sha256
)
return _FrozenBackboneSidecarPaths(
root=object_root,
manifest=object_root / "manifest.json",
progress=object_root / "progress.json",
lock=object_root / ".writer.lock",
final_hidden=object_root
/ f"final_hidden.{feature_dtype}.bin",
summary_hidden=object_root
/ f"summary_hidden.{feature_dtype}.bin",
input_positions=object_root / "input_positions.i64le",
)
def __init__(
self,
*,
cache_root: Path,
authority: FrozenBackboneCacheAuthority,
) -> None:
self.authority = authority
self.authority_record = authority.record_boundary()
self.authority_sha256 = authority.sha256_boundary()
self.cache_root = cache_root.expanduser().resolve()
self.paths = self._paths_for_authority_boundary(
cache_root=self.cache_root,
authority_sha256=self.authority_sha256,
feature_dtype=authority.feature_dtype,
)
self.legacy_authority_record = (
authority.legacy_v1_record_boundary()
)
self.legacy_authority_sha256 = (
authority.legacy_v1_sha256_boundary()
)
self.legacy_paths = (
self._paths_for_authority_boundary(
cache_root=self.cache_root,
authority_sha256=self.legacy_authority_sha256,
feature_dtype=authority.feature_dtype,
)
if self.legacy_authority_sha256 is not None
else None
)
def _expected_elements_boundary(self) -> tuple[int, int, int]:
return (
self.authority.hidden_size,
self.authority.summary_positions * self.authority.hidden_size,
1,
)
def _artifact_record_boundary(
self,
path: Path,
*,
elements: int,
dtype_name: str,
) -> dict[str, object]:
stat = path.stat()
return {
"path": str(path),
"bytes": stat.st_size,
"elements": elements,
"dtype": dtype_name,
"sha256": _file_sha256_boundary(path),
"device": stat.st_dev,
"inode": stat.st_ino,
"mtimeNs": stat.st_mtime_ns,
}
def _validation_cache_key_boundary(
self,
*,
paths: _FrozenBackboneSidecarPaths,
) -> tuple[int, Path, str]:
"""Bind reusable validation to this process session and executable authority."""
return (
os.getpid(),
paths.manifest,
self.authority_sha256,
)
def _cached_manifest_validation_boundary(
self,
*,
paths: _FrozenBackboneSidecarPaths,
) -> dict[str, object] | None:
"""Reuse exact digests only while all kernel-owned file stamps remain stable."""
cache_key = self._validation_cache_key_boundary(paths=paths)
observed_stamps = _frozen_backbone_sidecar_file_stamps_boundary(paths)
with _FROZEN_BACKBONE_VALIDATION_CACHE_LOCK:
cached = _FROZEN_BACKBONE_VALIDATION_CACHE.get(cache_key)
if (
cached is not None
and observed_stamps is not None
and cached.file_stamps == observed_stamps
):
_FROZEN_BACKBONE_VALIDATION_CACHE.move_to_end(cache_key)
return cached.manifest
if cached is not None:
del _FROZEN_BACKBONE_VALIDATION_CACHE[cache_key]
return None
def _cache_manifest_validation_boundary(
self,
*,
paths: _FrozenBackboneSidecarPaths,
manifest: dict[str, object],
manifest_sha256: str,
artifact_sha256s: tuple[str, str, str],
file_stamps: tuple[
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
_FrozenBackboneFileStamp,
],
) -> None:
"""Retain a bounded process-session proof for later lookup boundaries."""
cache_key = self._validation_cache_key_boundary(paths=paths)
validation = _FrozenBackboneManifestValidation(
manifest=manifest,
manifest_sha256=manifest_sha256,
artifact_sha256s=artifact_sha256s,
file_stamps=file_stamps,
)
with _FROZEN_BACKBONE_VALIDATION_CACHE_LOCK:
_FROZEN_BACKBONE_VALIDATION_CACHE[cache_key] = validation
_FROZEN_BACKBONE_VALIDATION_CACHE.move_to_end(cache_key)
while (
len(_FROZEN_BACKBONE_VALIDATION_CACHE)
> _FROZEN_BACKBONE_VALIDATION_CACHE_ENTRIES
):
_FROZEN_BACKBONE_VALIDATION_CACHE.popitem(last=False)
def _validated_manifest_boundary(
self,
*,
paths: _FrozenBackboneSidecarPaths,
) -> dict[str, object]:
cached = self._cached_manifest_validation_boundary(paths=paths)
if cached is not None:
return cached
file_stamps_before = _frozen_backbone_sidecar_file_stamps_boundary(
paths
)
if file_stamps_before is None:
raise RuntimeError("frozen-backbone sidecar artifacts are absent")
manifest_bytes = paths.manifest.read_bytes()
raw = json.loads(manifest_bytes)
manifest = _string_object_mapping_boundary(raw)
if manifest is None:
raise RuntimeError("frozen-backbone sidecar manifest is not an object")
stored_authority = _string_object_mapping_boundary(
manifest.get("authority")
)
stored_authority_sha256 = manifest.get("authoritySha256")
normalized_authority = (
_source_independent_authority_record_boundary(stored_authority)
)
if (
manifest.get("schema") != FROZEN_BACKBONE_SIDECAR_SCHEMA
or not isinstance(stored_authority_sha256, str)
or not _is_sha256_hex(stored_authority_sha256)
or stored_authority is None
or hashlib.sha256(
_canonical_json_bytes(stored_authority)
).hexdigest()
!= stored_authority_sha256
or paths.root.name != stored_authority_sha256
or paths.root.parent.name != stored_authority_sha256[:2]
or normalized_authority != self.authority_record
or manifest.get("rows") != 1
or manifest.get("backboneOnly") is not True
or manifest.get("targetEnteredForward") is not False
or manifest.get("additiveHeadCached") is not False
or manifest.get("historicalRboCached") is not False
or manifest.get("fabricCached") is not False
):
raise RuntimeError("frozen-backbone sidecar authority differs")
artifacts = manifest.get("artifacts")
if not isinstance(artifacts, dict):
raise RuntimeError("frozen-backbone sidecar artifacts are absent")
expected_elements = self._expected_elements_boundary()
expected_paths = paths.artifact_paths_boundary()
expected_dtypes = (
self.authority.feature_dtype,
self.authority.feature_dtype,
"int64",
)
names = ("finalHidden", "summaryHidden", "inputPositions")
artifact_sha256s: list[str] = []
for name, path, elements, dtype_name in zip(
names,
expected_paths,
expected_elements,
expected_dtypes,
strict=True,
):
record = artifacts.get(name)
stat = path.stat() if path.is_file() else None
recorded_sha256 = (
record.get("sha256")
if isinstance(record, dict)
else None
)
observed_sha256 = (
_file_sha256_boundary(path)
if stat is not None
and _is_sha256_hex(recorded_sha256)
else None
)
if (
not isinstance(record, dict)
or record.get("path") != str(path)
or record.get("elements") != elements
or record.get("dtype") != dtype_name
or stat is None
or record.get("bytes") != stat.st_size
or record.get("device") != stat.st_dev
or record.get("inode") != stat.st_ino
or record.get("mtimeNs") != stat.st_mtime_ns
or recorded_sha256 != observed_sha256
):
raise RuntimeError(
f"frozen-backbone {name} artifact authority differs"
)
assert isinstance(recorded_sha256, str)
artifact_sha256s.append(recorded_sha256)
file_stamps_after = _frozen_backbone_sidecar_file_stamps_boundary(paths)
if file_stamps_after != file_stamps_before:
raise RuntimeError(
"frozen-backbone sidecar changed during content validation"
)
final_sha256, summary_sha256, positions_sha256 = artifact_sha256s
self._cache_manifest_validation_boundary(
paths=paths,
manifest=manifest,
manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(),
artifact_sha256s=(
final_sha256,
summary_sha256,
positions_sha256,
),
file_stamps=file_stamps_after,
)
return manifest
def _lookup_paths_boundary(
self,
*,
paths: _FrozenBackboneSidecarPaths,
) -> FrozenBackbonePrefillPacket | None:
if not paths.manifest.is_file():
return None
self._validated_manifest_boundary(paths=paths)
dtype = _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype]
final_elements, summary_elements, input_elements = (
self._expected_elements_boundary()
)
final_hidden_t = torch.from_file(
str(paths.final_hidden),
shared=False,
size=final_elements,
dtype=dtype,
).reshape(1, 1, self.authority.hidden_size)
summary_hidden_t = torch.from_file(
str(paths.summary_hidden),
shared=False,
size=summary_elements,
dtype=dtype,
).reshape(
1,
self.authority.summary_positions,
self.authority.hidden_size,
)
input_positions_t = torch.from_file(
str(paths.input_positions),
shared=False,
size=input_elements,
dtype=torch.long,
)
return FrozenBackbonePrefillPacket.from_features_boundary(
authority=self.authority,
final_hidden_t=final_hidden_t,
summary_hidden_t=summary_hidden_t,
input_positions_t=input_positions_t,
)
def _lookup_current_boundary(
self,
) -> FrozenBackbonePrefillPacket | None:
return self._lookup_paths_boundary(paths=self.paths)
def lookup_boundary(self) -> FrozenBackbonePrefillPacket | None:
"""Return a current or legacy exact object without rewriting artifacts."""
current = self._lookup_current_boundary()
if current is not None:
return current
if (
self.legacy_paths is None
or self.legacy_authority_record is None
or self.legacy_authority_sha256 is None
):
direct_legacy = None
else:
direct_legacy = self._lookup_paths_boundary(
paths=self.legacy_paths,
)
if direct_legacy is not None:
return direct_legacy
# A moving source tree may report a different or malformed diagnostic,
# so its old source-bound SHA cannot be derived. Build one read-only
# index per process/root and match only after removing source
# diagnostics. The selected object's manifest and artifact hashes are
# still validated exactly before mmap.
legacy_manifests = (
_legacy_frozen_backbone_manifest_index_boundary(
self.cache_root
).get(self.authority_sha256, ())
)
for manifest_path in legacy_manifests:
candidate = self._paths_for_authority_boundary(
cache_root=self.cache_root,
authority_sha256=manifest_path.parent.name,
feature_dtype=self.authority.feature_dtype,
)
if candidate.root == self.paths.root:
continue
legacy = self._lookup_paths_boundary(paths=candidate)
if legacy is not None:
return legacy
return None
def _progress_payload_boundary(
self,
*,
completed_rows: int,
) -> dict[str, object]:
final_elements, summary_elements, input_elements = (
self._expected_elements_boundary()
)
return {
"schema": FROZEN_BACKBONE_SIDECAR_PROGRESS_SCHEMA,
"authoritySha256": self.authority_sha256,
"authority": self.authority_record,
"completedRows": completed_rows,
"finalHiddenElements": final_elements if completed_rows else 0,
"summaryHiddenElements": summary_elements if completed_rows else 0,
"inputPositionElements": input_elements if completed_rows else 0,
"targetEnteredForward": False,
}
def _validated_progress_boundary(self) -> dict[str, object]:
if not self.paths.progress.is_file():
return self._progress_payload_boundary(completed_rows=0)
raw = json.loads(self.paths.progress.read_text(encoding="utf-8"))
if not isinstance(raw, dict):
raise RuntimeError("frozen-backbone progress is not an object")
expected_empty = self._progress_payload_boundary(completed_rows=0)
expected_complete = self._progress_payload_boundary(completed_rows=1)
if raw != expected_empty and raw != expected_complete:
raise RuntimeError("frozen-backbone progress authority differs")
return {
str(key): value
for key, value in raw.items()
}
@staticmethod
def _write_tensor_boundary(path: Path, value: torch.Tensor) -> None:
with path.open("wb") as handle:
handle.write(_tensor_storage_bytes_boundary(value))
handle.flush()
os.fsync(handle.fileno())
def _validate_packet_boundary(
self,
packet: FrozenBackbonePrefillPacket,
) -> None:
"""Validate one row before any staged or durable object mutation."""
packet.validate_boundary()
if (
packet.final_hidden_t.shape
!= (1, 1, self.authority.hidden_size)
or packet.summary_hidden_t.shape
!= (
1,
self.authority.summary_positions,
self.authority.hidden_size,
)
or packet.final_hidden_t.dtype
!= _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype]
):
raise ValueError("frozen-backbone packet does not match its authority")
expected_identity_t = self.authority.digest_tensors_boundary()
for actual_t, expected_t in zip(
(
packet.cache_key_sha256_t,
packet.work_id_sha256_t,
packet.prompt_mask_sha256_t,
),
expected_identity_t,
strict=True,
):
if not torch.equal(actual_t.detach().to(device="cpu"), expected_t):
raise ValueError(
"frozen-backbone packet identity differs from its authority"
)
def _manifest_payload_boundary(self) -> dict[str, object]:
"""Build one manifest only after all three final artifacts exist."""
final_elements, summary_elements, input_elements = (
self._expected_elements_boundary()
)
return {
"schema": FROZEN_BACKBONE_SIDECAR_SCHEMA,
"authoritySha256": self.authority_sha256,
"authority": self.authority_record,
"rows": 1,
"backboneOnly": True,
"additiveHeadCached": False,
"historicalRboCached": False,
"fabricCached": False,
"targetEnteredForward": False,
"artifacts": {
"finalHidden": self._artifact_record_boundary(
self.paths.final_hidden,
elements=final_elements,
dtype_name=self.authority.feature_dtype,
),
"summaryHidden": self._artifact_record_boundary(
self.paths.summary_hidden,
elements=summary_elements,
dtype_name=self.authority.feature_dtype,
),
"inputPositions": self._artifact_record_boundary(
self.paths.input_positions,
elements=input_elements,
dtype_name="int64",
),
},
}
@staticmethod
def store_batch_boundary(
captured: tuple[
tuple[
FrozenBackboneFeatureSidecar,
FrozenBackbonePrefillPacket,
],
...,
],
) -> int:
"""Durably publish one producer outer wave with manifest-last phases.
Every row keeps its independent content-addressed object and writer
lock. The optimization changes only when durability syscalls run:
artifact bytes for the complete outer wave are staged first, then
flushed as one filesystem transaction group; manifest bytes are staged
and flushed second; manifest renames are published last and their
object/prefix/cache directories are synchronized together. A failure
before the last phase leaves no manifest for that row, so lookup treats
it as an ordinary miss. The owning producer advances its durable cursor
only after this method returns.
"""
if not captured:
return 0
ordered = tuple(
sorted(
captured,
key=lambda row: str(row[0].paths.root),
)
)
if len({sidecar.paths.root for sidecar, _packet in ordered}) != len(
ordered
):
raise ValueError(
"frozen-backbone batch contains duplicate object authority"
)
for sidecar, packet in ordered:
sidecar._validate_packet_boundary(packet)
lock_descriptors: list[int] = []
artifact_descriptors: list[int] = []
manifest_descriptors: list[int] = []
staged_rows: list[
tuple[
FrozenBackboneFeatureSidecar,
tuple[Path, Path, Path],
]
] = []
manifest_promotions: list[tuple[Path, Path]] = []
try:
for sidecar, _packet in ordered:
sidecar.paths.root.mkdir(parents=True, exist_ok=True)
lock_descriptor = os.open(
sidecar.paths.lock,
os.O_RDWR | os.O_CREAT,
0o600,
)
fcntl.flock(lock_descriptor, fcntl.LOCK_EX)
lock_descriptors.append(lock_descriptor)
for sidecar, packet in ordered:
if sidecar._lookup_current_boundary() is not None:
continue
progress = sidecar._validated_progress_boundary()
completed_rows = progress["completedRows"]
if type(completed_rows) is not int:
raise RuntimeError(
"frozen-backbone progress row frontier is malformed"
)
final_paths = sidecar.paths.artifact_paths_boundary()
partial_paths: tuple[Path, Path, Path] = (
final_paths[0].with_name(
f"{final_paths[0].name}.partial"
),
final_paths[1].with_name(
f"{final_paths[1].name}.partial"
),
final_paths[2].with_name(
f"{final_paths[2].name}.partial"
),
)
for final_path, partial_path in zip(
final_paths,
partial_paths,
strict=True,
):
if final_path.exists() and partial_path.exists():
raise RuntimeError(
"frozen-backbone sidecar has duplicate "
"incomplete artifacts"
)
if final_path.exists():
os.replace(final_path, partial_path)
payloads = (
_tensor_storage_bytes_boundary(packet.final_hidden_t),
_tensor_storage_bytes_boundary(packet.summary_hidden_t),
_tensor_storage_bytes_boundary(packet.input_positions_t),
)
for partial_path, payload in zip(
partial_paths,
payloads,
strict=True,
):
descriptor = os.open(
partial_path,
os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
0o600,
)
artifact_descriptors.append(descriptor)
_write_all_fd_boundary(descriptor, payload)
staged_rows.append((sidecar, partial_paths))
_synchronize_open_files_boundary(tuple(artifact_descriptors))
for descriptor in artifact_descriptors:
os.close(descriptor)
artifact_descriptors.clear()
for sidecar, partial_paths in staged_rows:
for partial_path, final_path in zip(
partial_paths,
sidecar.paths.artifact_paths_boundary(),
strict=True,
):
os.replace(partial_path, final_path)
temporary_manifest = sidecar.paths.manifest.with_name(
f".{sidecar.paths.manifest.name}.batch.partial"
)
temporary_manifest.unlink(missing_ok=True)
descriptor = os.open(
temporary_manifest,
os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
0o600,
)
manifest_descriptors.append(descriptor)
_write_all_fd_boundary(
descriptor,
_canonical_json_bytes(
sidecar._manifest_payload_boundary()
)
+ b"\n",
)
manifest_promotions.append(
(temporary_manifest, sidecar.paths.manifest)
)
_synchronize_open_files_boundary(tuple(manifest_descriptors))
for descriptor in manifest_descriptors:
os.close(descriptor)
manifest_descriptors.clear()
for temporary_manifest, manifest_path in manifest_promotions:
os.replace(temporary_manifest, manifest_path)
for sidecar, _partial_paths in staged_rows:
sidecar.paths.progress.unlink(missing_ok=True)
_synchronize_directories_boundary(
tuple(
path
for sidecar, _partial_paths in staged_rows
for path in (
sidecar.paths.root,
sidecar.paths.root.parent,
sidecar.cache_root / "objects",
sidecar.cache_root,
)
)
)
for sidecar, _packet in ordered:
if sidecar._lookup_current_boundary() is None:
raise RuntimeError(
"frozen-backbone batch manifest published "
"without a readable packet"
)
return len(captured)
finally:
for descriptor in artifact_descriptors:
os.close(descriptor)
for descriptor in manifest_descriptors:
os.close(descriptor)
for descriptor in reversed(lock_descriptors):
fcntl.flock(descriptor, fcntl.LOCK_UN)
os.close(descriptor)
def store_boundary(
self,
packet: FrozenBackbonePrefillPacket,
) -> FrozenBackbonePrefillPacket:
"""Durably publish one row and return its mmap-backed representation."""
self._validate_packet_boundary(packet)
self.paths.root.mkdir(parents=True, exist_ok=True)
lock_fd = os.open(self.paths.lock, os.O_RDWR | os.O_CREAT, 0o600)
try:
fcntl.flock(lock_fd, fcntl.LOCK_EX)
# Admission already performed the one canonical/legacy read. The
# publisher rechecks only its canonical destination under lock;
# repeating the full compatibility lookup would add an avoidable
# manifest scan to every admitted miss.
existing = self._lookup_current_boundary()
if existing is not None:
return existing
progress = self._validated_progress_boundary()
completed_rows_value = progress["completedRows"]
if type(completed_rows_value) is not int:
raise RuntimeError(
"frozen-backbone progress row frontier is malformed"
)
completed_rows = completed_rows_value
final_elements, summary_elements, input_elements = (
self._expected_elements_boundary()
)
dtype = _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype]
expected_bytes = (
final_elements * torch.tensor([], dtype=dtype).element_size(),
summary_elements * torch.tensor([], dtype=dtype).element_size(),
input_elements
* torch.tensor([], dtype=torch.long).element_size(),
)
final_paths = self.paths.artifact_paths_boundary()
partial_paths = tuple(
path.with_name(f"{path.name}.partial")
for path in final_paths
)
if completed_rows == 0:
for final_path, partial_path in zip(
final_paths,
partial_paths,
strict=True,
):
if final_path.exists() and partial_path.exists():
raise RuntimeError(
"frozen-backbone sidecar has duplicate incomplete artifacts"
)
if final_path.exists():
os.replace(final_path, partial_path)
self._write_tensor_boundary(
partial_paths[0],
packet.final_hidden_t,
)
self._write_tensor_boundary(
partial_paths[1],
packet.summary_hidden_t,
)
self._write_tensor_boundary(
partial_paths[2],
packet.input_positions_t,
)
_atomic_json_boundary(
self.paths.progress,
self._progress_payload_boundary(completed_rows=1),
)
for final_path, partial_path, byte_count in zip(
final_paths,
partial_paths,
expected_bytes,
strict=True,
):
selected = (
final_path if final_path.is_file() else partial_path
)
if not selected.is_file():
raise RuntimeError(
"frozen-backbone durable progress lost an artifact"
)
with selected.open("r+b") as handle:
handle.truncate(byte_count)
handle.flush()
os.fsync(handle.fileno())
if selected == partial_path:
os.replace(partial_path, final_path)
directory_fd = os.open(
self.paths.root,
os.O_RDONLY | os.O_DIRECTORY,
)
try:
os.fsync(directory_fd)
finally:
os.close(directory_fd)
manifest = self._manifest_payload_boundary()
_atomic_json_boundary(self.paths.manifest, manifest)
self.paths.progress.unlink(missing_ok=True)
directory_fd = os.open(
self.paths.root,
os.O_RDONLY | os.O_DIRECTORY,
)
try:
os.fsync(directory_fd)
finally:
os.close(directory_fd)
loaded = self._lookup_current_boundary()
if loaded is None:
raise RuntimeError(
"frozen-backbone manifest published without a readable packet"
)
return loaded
finally:
fcntl.flock(lock_fd, fcntl.LOCK_UN)
os.close(lock_fd)
def base_forward_cache_enabled_boundary() -> bool:
"""Return whether the frozen-parent prefill cache is active."""
raw = os.environ.get("NNF_RESYNTHESIS_BASE_FORWARD_CACHE", "1")
return raw not in {"0", "false", "False", "no", "off"}
def base_forward_cache_capacity_boundary() -> int:
"""Return the LRU capacity for cached prefill results."""
raw = os.environ.get("NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES", "512")
try:
capacity = int(raw)
except ValueError as error:
raise RuntimeError(
"NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES must be an integer"
) from error
if capacity < 1:
raise RuntimeError(
"NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES must be positive"
)
return capacity
def base_forward_cache_allowed_boundary(
*,
input_mask: torch.Tensor | None,
use_past: bool,
parent_trainable: bool,
) -> bool:
"""Return whether a frozen-parent prefill may be served from cache.
Cache is restricted to batched training prefills that supply an explicit
attention mask. Autoregressive decode arms (no mask, or KV continuation)
must always run the live parent forward so ``past_key_values`` state is
populated for the next token.
"""
return (
base_forward_cache_enabled_boundary()
and input_mask is not None
and not use_past
and not parent_trainable
)
def base_forward_cache_key_boundary(
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None,
) -> tuple[str, str]:
"""Return full SHA-256 hex and its 16-hex display prefix for one prompt."""
if input_ids.ndim != 2 or input_ids.shape[1] < 1:
raise ValueError("base forward cache key requires [batch, sequence]")
digest = hashlib.sha256()
prompt = input_ids.detach().to(device="cpu", dtype=torch.long).contiguous()
digest.update(prompt.numpy().tobytes())
if attention_mask is not None:
if attention_mask.shape != input_ids.shape:
raise ValueError("base forward cache attention mask geometry differs")
mask = attention_mask.detach().to(device="cpu", dtype=torch.long).contiguous()
digest.update(mask.numpy().tobytes())
full_hex = digest.hexdigest()
return full_hex, full_hex[:16]
@dataclass(frozen=True)
class BaseForwardCacheEntry:
"""CPU-resident frozen-parent prefill tensors keyed by prompt identity."""
hidden: torch.Tensor
logits: torch.Tensor
parent_context_hidden: torch.Tensor
parent_expert_routes: torch.Tensor
parent_layer_routes: torch.Tensor
kv_prefix_positions: torch.Tensor
kv_new_positions: torch.Tensor
parent_prefill_hidden: torch.Tensor | None
parent_prefill_input_positions: torch.Tensor | None
sha16: str
@classmethod
def from_forward_boundary(
cls,
*,
hidden: torch.Tensor,
logits: torch.Tensor,
parent_context_hidden: torch.Tensor,
parent_expert_routes: torch.Tensor,
parent_layer_routes: torch.Tensor,
kv_prefix_positions: torch.Tensor,
kv_new_positions: torch.Tensor,
parent_prefill_hidden: torch.Tensor | None,
parent_prefill_input_positions: torch.Tensor | None,
sha16: str,
) -> BaseForwardCacheEntry:
"""Store detached CPU tensors, pinning when the host supports it."""
def _pin(value: torch.Tensor) -> torch.Tensor:
detached = value.detach().to(device="cpu").contiguous()
if detached.is_floating_point() or detached.dtype in {
torch.int32,
torch.int64,
torch.bool,
}:
try:
return detached.pin_memory()
except RuntimeError:
return detached
return detached
prefill_hidden = (
_pin(parent_prefill_hidden)
if isinstance(parent_prefill_hidden, torch.Tensor)
else None
)
prefill_positions = (
_pin(parent_prefill_input_positions)
if isinstance(parent_prefill_input_positions, torch.Tensor)
else None
)
return cls(
hidden=_pin(hidden),
logits=_pin(logits),
parent_context_hidden=_pin(parent_context_hidden),
parent_expert_routes=_pin(parent_expert_routes),
parent_layer_routes=_pin(parent_layer_routes),
kv_prefix_positions=_pin(kv_prefix_positions),
kv_new_positions=_pin(kv_new_positions),
parent_prefill_hidden=prefill_hidden,
parent_prefill_input_positions=prefill_positions,
sha16=sha16,
)
def to_device_boundary(self, device: torch.device) -> BaseForwardCacheEntry:
"""Materialize one cache entry on the active compute device."""
def _move(value: torch.Tensor) -> torch.Tensor:
return value.to(device=device, non_blocking=device.type == "cuda")
prefill_hidden = (
_move(self.parent_prefill_hidden)
if isinstance(self.parent_prefill_hidden, torch.Tensor)
else None
)
prefill_positions = (
_move(self.parent_prefill_input_positions)
if isinstance(self.parent_prefill_input_positions, torch.Tensor)
else None
)
return BaseForwardCacheEntry(
hidden=_move(self.hidden),
logits=_move(self.logits),
parent_context_hidden=_move(self.parent_context_hidden),
parent_expert_routes=_move(self.parent_expert_routes),
parent_layer_routes=_move(self.parent_layer_routes),
kv_prefix_positions=_move(self.kv_prefix_positions),
kv_new_positions=_move(self.kv_new_positions),
parent_prefill_hidden=prefill_hidden,
parent_prefill_input_positions=prefill_positions,
sha16=self.sha16,
)
class BaseForwardCache:
"""LRU cache for complete frozen-parent prefill outputs."""
def __init__(self, *, capacity: int) -> None:
if capacity < 1:
raise ValueError("base forward cache capacity must be positive")
self._capacity = capacity
self._entries: OrderedDict[str, BaseForwardCacheEntry] = OrderedDict()
self._hits = 0
self._misses = 0
@property
def hits(self) -> int:
return self._hits
@property
def misses(self) -> int:
return self._misses
def clear(self) -> None:
self._entries.clear()
def lookup(self, cache_key: str) -> BaseForwardCacheEntry | None:
entry = self._entries.get(cache_key)
if entry is None:
self._misses += 1
return None
self._hits += 1
self._entries.move_to_end(cache_key)
return entry
def store(self, cache_key: str, entry: BaseForwardCacheEntry) -> None:
if cache_key in self._entries:
self._entries.move_to_end(cache_key)
self._entries[cache_key] = entry
while len(self._entries) > self._capacity:
self._entries.popitem(last=False)
def telemetry_boundary(self) -> dict[str, object]:
return {
"schema": "nnf.resynthesis.base_forward_cache.v1",
"enabled": base_forward_cache_enabled_boundary(),
"capacity": self._capacity,
"entries": len(self._entries),
"hits": self._hits,
"misses": self._misses,
}