"""Fail-closed KV-cache arithmetic and content-addressed capacity evidence. The original ``ModelShape``, ``cache_bytes``, ``gib``, and ``max_sequences`` API remains available. The capacity-plan records add the identities and reservation evidence required to use the arithmetic operationally. """ from __future__ import annotations from dataclasses import asdict, dataclass, is_dataclass from decimal import Decimal, InvalidOperation, ROUND_CEILING, ROUND_FLOOR import hashlib import json import math import re from typing import Any, Iterable GIB = 1024**3 MAX_LAYERS = 512 MAX_HEADS = 512 MAX_HEAD_DIM = 65_536 MAX_TOKENS = 16_777_216 MAX_SEQUENCES = 1_000_000 MAX_REQUESTS = 100_000 MAX_SAFE_BYTES = 2**63 - 1 SUPPORTED_BYTES_PER_ELEMENT = frozenset({0.5, 1.0, 2.0, 4.0, 8.0}) _PRECISION_BYTES = { "int4-packed-cache-v1": 0.5, "int8-cache-v1": 1.0, "bf16-cache-v1": 2.0, "fp16-cache-v1": 2.0, "fp32-cache-v1": 4.0, "fp64-cache-v1": 8.0, } _IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/@=+-]{0,159}$") _DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$") def _identifier(name: str, value: object) -> str: if type(value) is not str or not _IDENTIFIER.fullmatch(value): raise ValueError(f"{name} must be a stable identifier") return value def _digest(name: str, value: object) -> str: if type(value) is not str or not _DIGEST.fullmatch(value): raise ValueError(f"{name} must be a sha256 digest") return value def _owner(name: str, value: object) -> str: checked = _identifier(name, value) if not checked.startswith("team:"): raise ValueError(f"{name} must identify an accountable team") return checked def _integer(name: str, value: object, minimum: int, maximum: int) -> int: if type(value) is not int or not minimum <= value <= maximum: raise ValueError(f"{name} must be an integer in [{minimum}, {maximum}]") return value def _finite_decimal(name: str, value: object) -> Decimal: if isinstance(value, bool) or not isinstance(value, (int, float, Decimal)): raise ValueError(f"{name} must be a real number") if isinstance(value, float) and not math.isfinite(value): raise ValueError(f"{name} must be finite") try: checked = Decimal(str(value)) except (InvalidOperation, ValueError) as error: raise ValueError(f"{name} must be finite") from error if not checked.is_finite(): raise ValueError(f"{name} must be finite") return checked def _canonical(value: Any) -> Any: if is_dataclass(value) and not isinstance(value, type): return _canonical(asdict(value)) if isinstance(value, dict): return {key: _canonical(item) for key, item in sorted(value.items())} if isinstance(value, tuple): return [_canonical(item) for item in value] return value def _content_id(prefix: str, payload: dict[str, Any]) -> str: encoded = json.dumps( _canonical(payload), sort_keys=True, separators=(",", ":"), allow_nan=False ).encode("utf-8") return f"{prefix}@sha256:{hashlib.sha256(encoded).hexdigest()}" def _checked_byte_count(name: str, value: Decimal, *, rounding: str) -> int: if not value.is_finite() or value < 0: raise OverflowError(f"{name} produced an invalid byte count") rounded = value.to_integral_value(rounding=rounding) if rounded > MAX_SAFE_BYTES: raise OverflowError(f"{name} exceeds the supported byte range") return int(rounded) @dataclass(frozen=True) class ModelShape: layers: int kv_heads: int head_dim: int bytes_per_element: float def validate(self) -> None: _integer("layers", self.layers, 1, MAX_LAYERS) _integer("kv_heads", self.kv_heads, 1, MAX_HEADS) _integer("head_dim", self.head_dim, 1, MAX_HEAD_DIM) precision = _finite_decimal("bytes_per_element", self.bytes_per_element) if float(precision) not in SUPPORTED_BYTES_PER_ELEMENT: raise ValueError("bytes_per_element is not a supported cache precision") @property def content_id(self) -> str: self.validate() return _content_id("kv-model-shape", asdict(self)) def _raw_cache_decimal(shape: ModelShape, tokens: int, sequences: int) -> Decimal: if type(shape) is not ModelShape: raise TypeError("shape must be a concrete ModelShape") shape.validate() _integer("tokens", tokens, 1, MAX_TOKENS) _integer("sequences", sequences, 1, MAX_SEQUENCES) raw = ( Decimal(sequences) * Decimal(tokens) * Decimal(shape.layers) * 2 * Decimal(shape.kv_heads) * Decimal(shape.head_dim) * _finite_decimal("bytes_per_element", shape.bytes_per_element) ) _checked_byte_count("raw cache", raw, rounding=ROUND_CEILING) return raw def cache_bytes( shape: ModelShape, tokens: int, sequences: int = 1, overhead: float = 1.0, ) -> float: """Return bounded cache bytes while preserving the original helper API.""" factor = _finite_decimal("overhead", overhead) if not Decimal("1") <= factor <= Decimal("4"): raise ValueError("overhead must be in [1, 4]") value = _raw_cache_decimal(shape, tokens, sequences) * factor _checked_byte_count("cache bytes", value, rounding=ROUND_CEILING) return float(value) def gib(value: float) -> float: checked = _finite_decimal("value", value) if (checked != 0 and checked < Decimal("1e-12")) or checked < 0 or checked > MAX_SAFE_BYTES: raise ValueError("value is outside the supported byte range") return float(checked / Decimal(GIB)) def max_sequences( shape: ModelShape, tokens: int, available_gib: float, overhead: float = 1.0, ) -> int: available = _finite_decimal("available_gib", available_gib) if not Decimal("1e-12") <= available <= Decimal(MAX_SAFE_BYTES) / Decimal(GIB): raise ValueError("available_gib is outside the supported range") per_sequence = _finite_decimal( "per_sequence", cache_bytes(shape, tokens, overhead=overhead) ) capacity = available * Decimal(GIB) result = int((capacity / per_sequence).to_integral_value(rounding=ROUND_FLOOR)) if result > MAX_SEQUENCES: raise OverflowError("sequence capacity exceeds the supported count") return result @dataclass(frozen=True) class CapacityContract: contract_version: str model_id: str model_version: str architecture_content_digest: str shape: ModelShape cache_precision_id: str scheduler_id: str scheduler_version: str allocator_id: str allocator_version: str reservation_policy_version: str device_pool_id: str device_capacity_bytes: int reserved_non_cache_bytes: int block_tokens: int overhead_basis_points: int maximum_requests: int maximum_tokens_per_request: int numeric_convention: str model_owner: str scheduler_owner: str capacity_owner: str def __post_init__(self) -> None: for field in ( "contract_version", "model_id", "model_version", "cache_precision_id", "scheduler_id", "scheduler_version", "allocator_id", "allocator_version", "reservation_policy_version", "device_pool_id", ): _identifier(field, getattr(self, field)) _digest("architecture_content_digest", self.architecture_content_digest) for field in ("model_owner", "scheduler_owner", "capacity_owner"): _owner(field, getattr(self, field)) if type(self.shape) is not ModelShape: raise TypeError("shape must be a concrete ModelShape") self.shape.validate() if self.cache_precision_id not in _PRECISION_BYTES: raise ValueError("cache_precision_id is unsupported") if float(self.shape.bytes_per_element) != _PRECISION_BYTES[self.cache_precision_id]: raise ValueError("cache precision identity does not match model shape") _integer("device_capacity_bytes", self.device_capacity_bytes, 1, MAX_SAFE_BYTES) _integer( "reserved_non_cache_bytes", self.reserved_non_cache_bytes, 0, self.device_capacity_bytes - 1, ) _integer("block_tokens", self.block_tokens, 1, MAX_TOKENS) _integer("overhead_basis_points", self.overhead_basis_points, 10_000, 40_000) _integer("maximum_requests", self.maximum_requests, 1, MAX_REQUESTS) _integer( "maximum_tokens_per_request", self.maximum_tokens_per_request, 1, MAX_TOKENS, ) if self.block_tokens > self.maximum_tokens_per_request: raise ValueError("block_tokens exceeds maximum_tokens_per_request") if self.numeric_convention != "integer-bytes-ceil-block-bps-v1": raise ValueError("unsupported numeric_convention") @property def available_cache_bytes(self) -> int: return self.device_capacity_bytes - self.reserved_non_cache_bytes @property def model_content_id(self) -> str: return _content_id( "kv-model", { "model_id": self.model_id, "model_version": self.model_version, "architecture_content_digest": self.architecture_content_digest, "shape": self.shape, "cache_precision_id": self.cache_precision_id, }, ) @property def content_id(self) -> str: return _content_id("kv-capacity-contract", asdict(self)) @dataclass(frozen=True) class ReservationRequest: request_id: str request_version: str tenant_scope: str prompt_tokens: int maximum_new_tokens: int sequences: int reserved_tokens_per_sequence: int request_content_id: str @classmethod def capture( cls, *, contract: CapacityContract, request_id: str, request_version: str, tenant_scope: str, prompt_tokens: int, maximum_new_tokens: int, sequences: int, reserved_tokens_per_sequence: int, ) -> "ReservationRequest": if type(contract) is not CapacityContract: raise TypeError("contract must be a concrete CapacityContract") _validate_request_values( contract, prompt_tokens, maximum_new_tokens, sequences, reserved_tokens_per_sequence, ) payload = { "request_id": request_id, "request_version": request_version, "tenant_scope": tenant_scope, "prompt_tokens": prompt_tokens, "maximum_new_tokens": maximum_new_tokens, "sequences": sequences, "reserved_tokens_per_sequence": reserved_tokens_per_sequence, } return cls( **payload, request_content_id=_content_id("kv-reservation-request", payload), ) def __post_init__(self) -> None: for field in ( "request_id", "request_version", "tenant_scope", "request_content_id", ): _identifier(field, getattr(self, field)) @dataclass(frozen=True) class ReservationEvidence: evidence_id: str captured_at: str source_id: str source_version: str contract_content_id: str model_content_id: str scheduler_id: str scheduler_version: str allocator_id: str allocator_version: str cache_precision_id: str device_pool_id: str reported_available_cache_bytes: int requests: tuple[ReservationRequest, ...] evidence_owner: str evidence_content_id: str @classmethod def capture( cls, *, contract: CapacityContract, evidence_id: str, captured_at: str, source_id: str, source_version: str, reported_available_cache_bytes: int, requests: Iterable[ReservationRequest], evidence_owner: str, ) -> "ReservationEvidence": if type(contract) is not CapacityContract: raise TypeError("contract must be a concrete CapacityContract") frozen_requests = tuple(requests) if not 1 <= len(frozen_requests) <= contract.maximum_requests: raise ValueError("request count is outside the contract") if not all(type(item) is ReservationRequest for item in frozen_requests): raise TypeError("requests must contain concrete ReservationRequest records") _integer( "reported_available_cache_bytes", reported_available_cache_bytes, 1, MAX_SAFE_BYTES, ) payload = { "evidence_id": evidence_id, "captured_at": captured_at, "source_id": source_id, "source_version": source_version, "contract_content_id": contract.content_id, "model_content_id": contract.model_content_id, "scheduler_id": contract.scheduler_id, "scheduler_version": contract.scheduler_version, "allocator_id": contract.allocator_id, "allocator_version": contract.allocator_version, "cache_precision_id": contract.cache_precision_id, "device_pool_id": contract.device_pool_id, "reported_available_cache_bytes": reported_available_cache_bytes, "requests": frozen_requests, "evidence_owner": evidence_owner, } return cls( **payload, evidence_content_id=_content_id("kv-reservation-evidence", payload), ) def __post_init__(self) -> None: for field in ( "evidence_id", "captured_at", "source_id", "source_version", "contract_content_id", "model_content_id", "scheduler_id", "scheduler_version", "allocator_id", "allocator_version", "cache_precision_id", "device_pool_id", "evidence_content_id", ): _identifier(field, getattr(self, field)) _owner("evidence_owner", self.evidence_owner) if type(self.requests) is not tuple: raise TypeError("requests must be an immutable tuple") @dataclass(frozen=True) class CapacityPlan: contract_content_id: str model_content_id: str evidence_content_id: str request_count: int sequence_count: int reserved_tokens: int required_cache_bytes: int available_cache_bytes: int remaining_cache_bytes: int decision: str plan_content_id: str claim: str def _validate_request_values( contract: CapacityContract, prompt_tokens: object, maximum_new_tokens: object, sequences: object, reserved_tokens_per_sequence: object, ) -> None: prompt = _integer( "prompt_tokens", prompt_tokens, 1, contract.maximum_tokens_per_request ) growth = _integer( "maximum_new_tokens", maximum_new_tokens, 1, contract.maximum_tokens_per_request, ) _integer("sequences", sequences, 1, MAX_SEQUENCES) reserved = _integer( "reserved_tokens_per_sequence", reserved_tokens_per_sequence, 1, contract.maximum_tokens_per_request, ) if prompt + growth > contract.maximum_tokens_per_request: raise ValueError("request shape exceeds maximum_tokens_per_request") if reserved < prompt + growth: raise ValueError("reservation does not cover prompt plus maximum growth") def _request_required_bytes( contract: CapacityContract, request: ReservationRequest ) -> tuple[int, int]: rounded_tokens = ( (request.reserved_tokens_per_sequence + contract.block_tokens - 1) // contract.block_tokens ) * contract.block_tokens if rounded_tokens > contract.maximum_tokens_per_request: raise ValueError("block-rounded reservation exceeds the request token cap") raw = _raw_cache_decimal(contract.shape, rounded_tokens, request.sequences) required = raw * Decimal(contract.overhead_basis_points) / Decimal(10_000) return rounded_tokens * request.sequences, _checked_byte_count( "request reservation", required, rounding=ROUND_CEILING ) def build_capacity_plan( contract: CapacityContract, evidence: ReservationEvidence ) -> CapacityPlan: """Revalidate reservations and return a content-addressed admission plan.""" if type(contract) is not CapacityContract: raise TypeError("contract must be a concrete CapacityContract") if type(evidence) is not ReservationEvidence: raise TypeError("evidence must be concrete ReservationEvidence") contract.__post_init__() evidence.__post_init__() scope = ( evidence.contract_content_id, evidence.model_content_id, evidence.scheduler_id, evidence.scheduler_version, evidence.allocator_id, evidence.allocator_version, evidence.cache_precision_id, evidence.device_pool_id, ) expected = ( contract.content_id, contract.model_content_id, contract.scheduler_id, contract.scheduler_version, contract.allocator_id, contract.allocator_version, contract.cache_precision_id, contract.device_pool_id, ) if scope != expected: raise ValueError("reservation evidence scope does not match the contract") if evidence.reported_available_cache_bytes != contract.available_cache_bytes: raise ValueError("reported cache capacity does not match the contract") if not 1 <= len(evidence.requests) <= contract.maximum_requests: raise ValueError("request count is outside the contract") seen_requests: set[tuple[str, str]] = set() request_bytes: list[int] = [] reserved_tokens = 0 sequence_count = 0 for request in evidence.requests: if type(request) is not ReservationRequest: raise TypeError("evidence contains a non-concrete request") request.__post_init__() identity = (request.request_id, request.request_version) if identity in seen_requests: raise ValueError("duplicate request identity") seen_requests.add(identity) _validate_request_values( contract, request.prompt_tokens, request.maximum_new_tokens, request.sequences, request.reserved_tokens_per_sequence, ) payload = { field: getattr(request, field) for field in ( "request_id", "request_version", "tenant_scope", "prompt_tokens", "maximum_new_tokens", "sequences", "reserved_tokens_per_sequence", ) } if request.request_content_id != _content_id( "kv-reservation-request", payload ): raise ValueError("request content identity does not match contents") tokens, required = _request_required_bytes(contract, request) reserved_tokens += tokens sequence_count += request.sequences if reserved_tokens > MAX_TOKENS * MAX_SEQUENCES: raise OverflowError("aggregate reserved-token count is unsupported") if sequence_count > MAX_SEQUENCES: raise OverflowError("aggregate sequence count is unsupported") request_bytes.append(required) evidence_payload = { field: getattr(evidence, field) for field in ( "evidence_id", "captured_at", "source_id", "source_version", "contract_content_id", "model_content_id", "scheduler_id", "scheduler_version", "allocator_id", "allocator_version", "cache_precision_id", "device_pool_id", "reported_available_cache_bytes", "requests", "evidence_owner", ) } if evidence.evidence_content_id != _content_id( "kv-reservation-evidence", evidence_payload ): raise ValueError("evidence content identity does not match contents") required_cache_bytes = sum(request_bytes) if required_cache_bytes > MAX_SAFE_BYTES: raise OverflowError("aggregate cache reservation exceeds the supported range") remaining = contract.available_cache_bytes - required_cache_bytes decision = "FIT" if remaining >= 0 else "EXCEEDS_CAPACITY" plan_payload = { "contract_content_id": contract.content_id, "model_content_id": contract.model_content_id, "evidence_content_id": evidence.evidence_content_id, "request_count": len(evidence.requests), "sequence_count": sequence_count, "reserved_tokens": reserved_tokens, "required_cache_bytes": required_cache_bytes, "available_cache_bytes": contract.available_cache_bytes, "remaining_cache_bytes": remaining, "decision": decision, "claim": "CAPACITY_ARITHMETIC_ONLY_REQUIRES_LOAD_VALIDATION", } return CapacityPlan( **plan_payload, plan_content_id=_content_id("kv-capacity-plan", plan_payload), ) ILLUSTRATIVE_CONTRACT = CapacityContract( contract_version="kv-capacity-plan-v2", model_id="decoder:academy-8b", model_version="model-v8", architecture_content_digest="sha256:" + "c" * 64, shape=ModelShape(layers=32, kv_heads=8, head_dim=128, bytes_per_element=2), cache_precision_id="bf16-cache-v1", scheduler_id="continuous-batcher", scheduler_version="scheduler-v4", allocator_id="paged-kv-allocator", allocator_version="allocator-v3", reservation_policy_version="reserve-prompt-plus-max-output-v2", device_pool_id="pool:a100-80gb", device_capacity_bytes=80 * GIB, reserved_non_cache_bytes=40 * GIB, block_tokens=16, overhead_basis_points=11_500, maximum_requests=64, maximum_tokens_per_request=16_384, numeric_convention="integer-bytes-ceil-block-bps-v1", model_owner="team:model-platform", scheduler_owner="team:inference", capacity_owner="team:serving-sre", ) ILLUSTRATIVE_REQUESTS = ( ReservationRequest.capture( contract=ILLUSTRATIVE_CONTRACT, request_id="request-001", request_version="request-v1", tenant_scope="tenant:academy-a", prompt_tokens=6_000, maximum_new_tokens=2_192, sequences=8, reserved_tokens_per_sequence=8_192, ), ReservationRequest.capture( contract=ILLUSTRATIVE_CONTRACT, request_id="request-002", request_version="request-v1", tenant_scope="tenant:academy-b", prompt_tokens=3_000, maximum_new_tokens=1_096, sequences=8, reserved_tokens_per_sequence=4_096, ), ) ILLUSTRATIVE_EVIDENCE = ReservationEvidence.capture( contract=ILLUSTRATIVE_CONTRACT, evidence_id="reservation-window-001", captured_at="2026-08-25T00:00:00Z", source_id="scheduler:reservation-log", source_version="log-schema-v3", reported_available_cache_bytes=ILLUSTRATIVE_CONTRACT.available_cache_bytes, requests=ILLUSTRATIVE_REQUESTS, evidence_owner="team:serving-sre", ) def format_example() -> str: plan = build_capacity_plan(ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_EVIDENCE) return "\n".join( ( "example=illustrative_only", f"contract_version={ILLUSTRATIVE_CONTRACT.contract_version}", f"model_content_id={plan.model_content_id}", f"evidence_content_id={plan.evidence_content_id}", f"requests={plan.request_count}", f"sequences={plan.sequence_count}", f"reserved_tokens={plan.reserved_tokens}", f"required_cache_gib={gib(plan.required_cache_bytes):.3f}", f"available_cache_gib={gib(plan.available_cache_bytes):.3f}", f"decision={plan.decision}", f"plan_content_id={plan.plan_content_id}", f"claim={plan.claim}", ) ) if __name__ == "__main__": print(format_example())