File size: 12,239 Bytes
2415c4c | 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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Eager program readiness and compile-output resource ownership."""
from __future__ import annotations
import hashlib
import json
import math
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from typing import Any
import torch
import ttnn
from models.common.llm_runtime.tensor_resources import (
TensorResourceOrphan,
attach_cleanup_failures,
best_effort_deallocate_owned_tensors,
release_orphans,
)
_PROGRAM_KEY_DOMAIN = "tttv2.llm-runtime.program"
_PROGRAM_KEY_SCHEMA_VERSION = 1
@dataclass(frozen=True)
class ProgramKey:
"""Full content digest for one operation-produced program signature."""
digest: str
def __post_init__(self) -> None:
validate_sha256_digest(self.digest, "program")
@classmethod
def from_signature(cls, signature: Any) -> "ProgramKey":
return cls(signature_digest(_PROGRAM_KEY_DOMAIN, _PROGRAM_KEY_SCHEMA_VERSION, signature))
@dataclass(frozen=True)
class OutputSpec:
shape: tuple[int, ...]
dtype: Any
layout: Any = None
memory_config: Any = None
@classmethod
def from_value(cls, value: Any) -> "OutputSpec":
if isinstance(value, tuple):
if not value:
raise ValueError("Cannot derive an output specification from an empty tuple")
value = value[0]
if isinstance(value, torch.Tensor):
return cls(shape=tuple(value.shape), dtype=value.dtype)
if isinstance(value, ttnn.Tensor):
allocated = value.is_allocated() if hasattr(value, "is_allocated") else False
return cls(
shape=tuple(value.shape),
dtype=value.dtype,
layout=value.layout,
memory_config=value.spec.memory_config if allocated else None,
)
raise TypeError(f"Cannot derive an output specification from {type(value).__name__}")
@dataclass(frozen=True)
class CompiledProgram:
"""Program metadata retained independently from all trace state."""
key: ProgramKey
signature: Any
output_spec: OutputSpec
class ProgramCompiler:
"""Compile and retain the eager program set for one execution lane.
``EagerExecutor.compile_*`` derives an operation signature and calls
`compile`. Warmup registers the complete program set before
``TraceCompiler`` closes the compile gate. Forward execution then uses
`require_compiled` to reject unseen signatures after trace
activation. Trace artifacts remain in the separate trace registry.
"""
def __init__(self, mesh_device: Any, bound_cache_context: Callable[[], Any]):
self.mesh_device = mesh_device
self._bound_cache_context = bound_cache_context
self._programs: dict[ProgramKey, CompiledProgram] = {}
self._program_keys: dict[tuple[Any, ...], ProgramKey] = {}
self._compile_orphans: list[TensorResourceOrphan] = []
self._trace_capture_in_progress = False
self._trace_active = False
self._released = False
self._post_activation_compile_rejections = 0
# Public API
@property
def trace_capture_in_progress(self) -> bool:
return self._trace_capture_in_progress
@property
def trace_active(self) -> bool:
return self._trace_active
@property
def compile_orphan_count(self) -> int:
return len(self._compile_orphans)
@property
def compiled_programs(self) -> tuple[CompiledProgram, ...]:
"""Return an immutable snapshot of the authoritative program registry."""
return tuple(self._programs.values())
@property
def post_activation_compile_rejections(self) -> int:
"""Return unseen-program compile attempts rejected after activation."""
return self._post_activation_compile_rejections
def key_for(self, signature: Any) -> ProgramKey:
"""Return the stable program key for one operation signature."""
material = _canonical_value(_signature_key_material(signature))
key = self._program_keys.get(material)
if key is None:
key = ProgramKey.from_signature(signature)
self._program_keys[material] = key
return key
def compile(
self,
signature: Any,
invoke: Callable[[Any], Any],
*,
output_spec: Callable[[Any], OutputSpec] = OutputSpec.from_value,
release_output: Callable[[Any], Any] = lambda output: output,
expected_output_spec: OutputSpec | None = None,
) -> CompiledProgram:
"""Compile one signature and release its transient invocation output."""
self._ensure_live()
if self._compile_orphans:
raise RuntimeError("Cannot compile while unreleased compile outputs remain; clean up this compiler")
key = self.key_for(signature)
existing = self._programs.get(key)
if existing is not None:
_ensure_matching_signature(key, existing.signature, signature)
if expected_output_spec is not None and existing.output_spec != expected_output_spec:
raise ValueError(f"Program key {key.digest} was compiled with a different output contract")
return existing
if self._trace_capture_in_progress:
raise RuntimeError(f"Cannot compile uncompiled program key {key.digest} while trace capture is in progress")
if self._trace_active:
self._post_activation_compile_rejections += 1
raise RuntimeError(f"Cannot compile uncompiled program key {key.digest} after trace activation")
cache_context = self._bound_cache_context()
if cache_context is None:
raise RuntimeError("Paged KV cache must be allocated and bound before compilation")
output = invoke(cache_context)
owned_output = release_output(output)
try:
ttnn.synchronize_device(self.mesh_device)
spec = output_spec(output)
if expected_output_spec is not None and spec != expected_output_spec:
raise ValueError(f"Program key {key.digest} produced an unexpected output contract")
except BaseException as primary:
cleanup_failures = self._release_or_retain_compile_output(owned_output)
try:
ttnn.synchronize_device(self.mesh_device)
except BaseException as error:
cleanup_failures.append(error)
attach_cleanup_failures(primary, cleanup_failures)
raise
cleanup_failures = self._release_or_retain_compile_output(owned_output)
try:
ttnn.synchronize_device(self.mesh_device)
except BaseException as primary:
attach_cleanup_failures(primary, cleanup_failures)
raise
if cleanup_failures:
error = RuntimeError(f"Failed to deallocate {len(cleanup_failures)} compile output resource(s)")
attach_cleanup_failures(error, cleanup_failures)
raise error from cleanup_failures[0]
program = CompiledProgram(key=key, signature=signature, output_spec=spec)
self._programs[key] = program
return program
def require_compiled(self, key: ProgramKey, signature: Any | None = None) -> CompiledProgram:
"""Return registered program metadata or reject an unseen signature."""
self._ensure_live()
program = self._programs.get(key)
if program is None:
suffix = " after trace activation" if self._trace_active else ""
raise RuntimeError(f"Program key {key.digest} was not compiled{suffix}")
if signature is not None:
_ensure_matching_signature(key, program.signature, signature)
return program
def set_trace_capture_in_progress(self, value: bool) -> None:
self._ensure_live()
if not isinstance(value, bool):
raise TypeError("trace capture state must be bool")
if value and self._trace_active:
raise RuntimeError("Cannot begin trace capture after trace activation")
self._trace_capture_in_progress = value
def set_trace_active(self, value: bool) -> None:
self._ensure_live()
if not isinstance(value, bool):
raise TypeError("trace active state must be bool")
if value and self._trace_capture_in_progress:
raise RuntimeError("Trace capture must finish before trace activation")
self._trace_active = value
def cleanup(self) -> None:
"""Release retryable compile outputs and terminalize the registry."""
if self._released:
return
failures = release_orphans(self._compile_orphans)
if failures:
error = RuntimeError(f"Failed to release {len(failures)} compile output resource(s)")
attach_cleanup_failures(error, failures)
raise error from failures[0]
self._program_keys.clear()
self._trace_capture_in_progress = False
self._trace_active = False
self._released = True
# Private implementation
def _release_or_retain_compile_output(self, output: Any) -> list[BaseException]:
orphan = TensorResourceOrphan(output)
failures = best_effort_deallocate_owned_tensors(output, orphan.deallocated_tensor_ids)
if failures:
self._compile_orphans.append(orphan)
return failures
def _ensure_live(self) -> None:
if self._released:
raise RuntimeError("ProgramCompiler has been released")
# Public helpers
def signature_digest(domain: str, schema_version: int, signature: Any) -> str:
"""Return a stable SHA-256 digest for explicit signature key material."""
payload = (
("domain", domain),
("schema_version", schema_version),
("signature", _signature_key_material(signature)),
)
encoded = json.dumps(_canonical_value(payload), ensure_ascii=True, separators=(",", ":"))
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
def validate_sha256_digest(digest: str, domain: str) -> None:
"""Validate the full lowercase digest representation used by registry keys."""
if (
not isinstance(digest, str)
or len(digest) != 64
or any(character not in "0123456789abcdef" for character in digest)
):
raise ValueError(f"{domain} key digest must be a full lowercase SHA-256 hexadecimal digest")
# Private helpers
def _signature_key_material(signature: Any) -> tuple[Any, ...]:
"""Return the explicit tagged primitive tuple supplied by a signature."""
try:
material = signature.key_material
except AttributeError as error:
raise TypeError(f"{type(signature).__name__} must expose key_material") from error
if callable(material):
material = material()
if not isinstance(material, tuple):
raise TypeError("signature key_material must be a tuple")
return material
def _canonical_value(value: Any) -> Any:
if value is None:
return ("null",)
if isinstance(value, Enum):
if not isinstance(value.value, str):
raise TypeError("signature enum values must be stable strings")
return ("enum", value.value)
if isinstance(value, bool):
return ("bool", value)
if isinstance(value, int):
return ("int", str(value))
if isinstance(value, float):
if not math.isfinite(value):
raise ValueError("signature key material must not contain non-finite floats")
return ("float", value.hex())
if isinstance(value, str):
return ("str", value)
if isinstance(value, tuple):
return ("tuple", tuple(_canonical_value(item) for item in value))
raise TypeError(
"signature key material must contain only None, bool, int, finite float, str, stable enums, and tuples; "
f"got {type(value).__name__}"
)
def _ensure_matching_signature(key: ProgramKey, retained: Any, candidate: Any) -> None:
if retained != candidate:
raise RuntimeError(f"Program key collision for digest {key.digest}: retained signature differs")
|