# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC # SPDX-License-Identifier: Apache-2.0 """Trace capture, replay, and persistent-resource ownership.""" from __future__ import annotations import ctypes from collections.abc import Callable from dataclasses import dataclass, field from typing import Any from loguru import logger from ttnn.tools import trace_allocation_tracker import ttnn from models.common.llm_runtime.program_compiler import ( ProgramCompiler, ProgramKey, signature_digest, validate_sha256_digest, ) from models.common.llm_runtime.tensor_resources import ( TensorResourceOrphan, attach_cleanup_failures, best_effort_deallocate_owned_tensors, raise_cleanup_failures, release_orphans, ) _TRACE_KEY_DOMAIN = "tttv2.llm-runtime.trace" _TRACE_KEY_SCHEMA_VERSION = 1 @dataclass(frozen=True) class TraceKey: """Full content digest for one operation-produced trace signature.""" digest: str def __post_init__(self) -> None: validate_sha256_digest(self.digest, "trace") @classmethod def from_signature(cls, signature: Any) -> "TraceKey": return cls(signature_digest(_TRACE_KEY_DOMAIN, _TRACE_KEY_SCHEMA_VERSION, signature)) @dataclass class PersistentInputs: """Trace-owned persistent replay inputs opaque to public runtime APIs.""" values: Any @dataclass(frozen=True) class InputRefreshPolicy: every_replay: tuple[str, ...] = () full_on_batch_reset: bool = True full_on_graph_switch: bool = True full_without_device_feedback: bool = True refresh_page_table_on_change: bool = True @dataclass(frozen=True) class RefreshDecision: full: bool page_table: bool fields: tuple[str, ...] @dataclass class TraceArtifact: trace_id: int persistent_inputs: PersistentInputs outputs: Any refresh_policy: InputRefreshPolicy trace_released: bool = False deallocated_tensor_ids: set[int] = field(default_factory=set, repr=False) @dataclass(frozen=True) class TraceCapturePlan: """Operation-produced specification for one trace-capable compiled program.""" program_key: ProgramKey trace_signature: Any operation: str prepare_inputs: Callable[[], PersistentInputs | Any] capture: Callable[[PersistentInputs], Any] refresh_policy: InputRefreshPolicy = InputRefreshPolicy() schema_fingerprint: Any = None prepare_workspace: Callable[[], Any] | None = None workspace_fingerprint: Any = None prime: Callable[[PersistentInputs], Any] | None = None release_prime_output: Callable[[Any], list[BaseException]] | None = None def __post_init__(self) -> None: if self.operation not in ("prefill", "decode"): raise ValueError(f"Unsupported trace operation: {self.operation!r}") if (self.prime is None) is not (self.release_prime_output is None): raise ValueError("trace capture prime and output releaser must be configured together") @dataclass class TraceRecord: signature: Any operation: str artifact: TraceArtifact | None = None @dataclass class TraceAliasRecord: """Program-local postprocess state kept outside the shared hidden trace.""" trace_key: TraceKey workspace_fingerprint: Any prepare_workspace: Callable[[], Any] | None = field(default=None, repr=False) workspace: Any = None deallocated_tensor_ids: set[int] = field(default_factory=set, repr=False) class TraceCompiler: """Register, capture, replay, and release traces for compiled programs. ``TracedExecutor.compile_*`` first compiles an eager program and calls `register_capture_plan`. `WarmupCoordinator` calls `capture_all` only after the complete configured program set exists. Forward execution then calls `replay` with operation-owned refresh logic. The compiler owns trace artifacts and persistent inputs, while the composed ``ProgramCompiler`` remains the sole program registry. """ def __init__(self, program_compiler: ProgramCompiler): if not isinstance(program_compiler, ProgramCompiler): raise TypeError("program_compiler must be a ProgramCompiler") self.program_compiler = program_compiler self.mesh_device = program_compiler.mesh_device self._traces: dict[TraceKey, TraceRecord] = {} self._plans: dict[TraceKey, TraceCapturePlan] = {} self._program_to_trace: dict[ProgramKey, TraceKey] = {} self._aliases: dict[ProgramKey, TraceAliasRecord] = {} self._rollback_orphans: list[TensorResourceOrphan] = [] self._capture_in_progress = False self._activated = False self._released = False self._previous_replay_key: TraceKey | None = None self._replay_count = 0 self._replay_counts = {"prefill": 0, "decode": 0} # Public API @property def trace_active(self) -> bool: return self._activated @property def replay_count(self) -> int: """Return successfully submitted trace replays across all operations.""" return self._replay_count @property def replay_counts(self) -> dict[str, int]: """Return a snapshot of successfully submitted replays by operation.""" return dict(self._replay_counts) @property def trace_count(self) -> int: """Return the number of semantic hidden traces in the registry.""" return len(self._traces) @property def trace_association_count(self) -> int: """Return the number of compiled-program aliases associated to traces.""" return len(self._program_to_trace) def registered_coverage(self, operation: str) -> tuple[tuple[TraceKey, Any], ...]: """Return registered trace keys/signatures for one operation.""" if operation not in ("prefill", "decode"): raise ValueError(f"Unsupported trace operation: {operation!r}") return tuple( (trace_key, record.signature) for trace_key, record in self._traces.items() if record.operation == operation ) def get(self, key: TraceKey) -> TraceRecord | None: """Return the record needed to finish an operation-specific replay.""" return self._traces.get(key) def trace_key_for_program(self, program_key: ProgramKey) -> TraceKey | None: """Return the registered trace association for one compiled program.""" return self._program_to_trace.get(program_key) def workspace_for_program(self, program_key: ProgramKey) -> Any: """Return one alias's postprocess workspace after capture allocation.""" alias = self._aliases.get(program_key) if alias is None: raise RuntimeError(f"Program key {program_key.digest} has no trace alias workspace") return alias.workspace def register_capture_plan(self, plan: TraceCapturePlan) -> TraceKey: """Validate a compiled source and register one explicit trace association.""" self._ensure_live() if self._capture_in_progress or self._activated: raise RuntimeError("Cannot register trace capture plans during capture or after trace activation") self.program_compiler.require_compiled(plan.program_key) trace_key = TraceKey.from_signature(plan.trace_signature) existing_association = self._program_to_trace.get(plan.program_key) if existing_association is not None and existing_association != trace_key: raise ValueError(f"Program key {plan.program_key.digest} already has a different trace association") existing_alias = self._aliases.get(plan.program_key) if existing_alias is not None and existing_alias.workspace_fingerprint != plan.workspace_fingerprint: raise ValueError( f"Program key {plan.program_key.digest} was registered with a different workspace fingerprint" ) record = self._traces.get(trace_key) if record is None: record = TraceRecord( signature=plan.trace_signature, operation=plan.operation, ) self._traces[trace_key] = record self._plans[trace_key] = plan else: if record.signature != plan.trace_signature: raise RuntimeError(f"Trace key collision for digest {trace_key.digest}: retained signature differs") if record.operation != plan.operation: raise RuntimeError(f"Trace key collision for digest {trace_key.digest}: operation differs") if self._plans[trace_key].refresh_policy != plan.refresh_policy: raise ValueError(f"Trace key {trace_key.digest} was registered with a different refresh policy") if self._plans[trace_key].schema_fingerprint != plan.schema_fingerprint: raise ValueError(f"Trace key {trace_key.digest} was registered with a different schema fingerprint") self._program_to_trace[plan.program_key] = trace_key if existing_alias is None: self._aliases[plan.program_key] = TraceAliasRecord( trace_key=trace_key, workspace_fingerprint=plan.workspace_fingerprint, prepare_workspace=plan.prepare_workspace, ) return trace_key def capture_all(self) -> None: """Allocate every persistent input before beginning the first capture.""" self._ensure_live() if self._activated: return if self._capture_in_progress: raise RuntimeError("Trace capture is already in progress") if not self._plans: return if self.program_compiler.compile_orphan_count: raise RuntimeError("Cannot capture while unreleased compile outputs remain") prepared: dict[TraceKey, tuple[PersistentInputs, TraceCapturePlan]] = {} captured_keys: set[TraceKey] = set() self._capture_in_progress = True try: for trace_key, plan in self._plans.items(): self.program_compiler.require_compiled(plan.program_key) values = plan.prepare_inputs() persistent = values if isinstance(values, PersistentInputs) else PersistentInputs(values) prepared[trace_key] = (persistent, plan) for program_key, alias in self._aliases.items(): if alias.prepare_workspace is not None: alias.workspace = alias.prepare_workspace() capture_order = sorted( prepared, key=lambda trace_key: self._traces[trace_key].operation == "prefill", ) for trace_key in capture_order: persistent, plan = prepared[trace_key] record = self._traces[trace_key] # Program signatures intentionally describe padded trace # identity, not active-row cardinality. Selected operation # plans therefore prime their exact persistent-input body # immediately before capturing that same body. No unrelated # trace can perturb allocator/program state between the prime # and ``begin_trace_capture``. if plan.prime is not None: prime_output = None try: prime_output = plan.prime(persistent) ttnn.synchronize_device(self.mesh_device) except BaseException as primary: cleanup_failures = plan.release_prime_output(prime_output) try: ttnn.synchronize_device(self.mesh_device) except BaseException as error: cleanup_failures.append(error) attach_cleanup_failures(primary, cleanup_failures) raise release_failures = plan.release_prime_output(prime_output) try: ttnn.synchronize_device(self.mesh_device) except BaseException as error: release_failures.append(error) if release_failures: raise_cleanup_failures(release_failures) logger.info(f"Primed {plan.operation} trace capture body: signature={plan.trace_signature!r}") self.program_compiler.set_trace_capture_in_progress(True) # Whatever the model allocates between begin/end_trace_capture belongs to the trace # being recorded and must stay allocated for replay; recording N traces means capture # N runs while 1..N-1 are live, which ordering cannot avoid. Acknowledge the window # (no-op unless TT_METAL_TRACE_ALLOC_TRACKING=1), as tt_transformers' generator does. with trace_allocation_tracker.corruptible_allocation_scope(self.mesh_device): trace_id = ttnn.begin_trace_capture(self.mesh_device, cq_id=0) outputs = None capture_ended = False try: outputs = plan.capture(persistent) ttnn.end_trace_capture(self.mesh_device, trace_id, cq_id=0) capture_ended = True ttnn.synchronize_device(self.mesh_device) except BaseException as primary: cleanup_failures = [] if not capture_ended: try: ttnn.end_trace_capture(self.mesh_device, trace_id, cq_id=0) except BaseException as error: cleanup_failures.append(error) record.artifact = TraceArtifact( trace_id=trace_id, persistent_inputs=persistent, outputs=outputs, refresh_policy=plan.refresh_policy, ) captured_keys.add(trace_key) cleanup_failures.extend(self._release_trace(record)) attach_cleanup_failures(primary, cleanup_failures) raise record.artifact = TraceArtifact( trace_id=trace_id, persistent_inputs=persistent, outputs=outputs, refresh_policy=plan.refresh_policy, ) logger.info(f"Captured {plan.operation} trace: signature={plan.trace_signature!r}") captured_keys.add(trace_key) self.program_compiler.set_trace_capture_in_progress(False) ttnn.synchronize_device(self.mesh_device) self._capture_in_progress = False self.program_compiler.set_trace_capture_in_progress(False) self._activated = True self.program_compiler.set_trace_active(True) _trim_host_allocator() except BaseException as primary: cleanup_failures = self._release_trace_resources() cleanup_failures.extend(self._release_alias_workspaces()) for trace_key, (persistent, _) in prepared.items(): if trace_key in captured_keys: continue orphan = TensorResourceOrphan(persistent.values) orphan_failures = best_effort_deallocate_owned_tensors( orphan.values, orphan.deallocated_tensor_ids, ) cleanup_failures.extend(orphan_failures) if orphan_failures: self._rollback_orphans.append(orphan) self._activated = ( bool(self._rollback_orphans) or any(record.artifact is not None for record in self._traces.values()) or any(alias.workspace is not None for alias in self._aliases.values()) ) self._capture_in_progress = False self.program_compiler.set_trace_capture_in_progress(False) self.program_compiler.set_trace_active(self._activated) attach_cleanup_failures(primary, cleanup_failures) raise def replay( self, program_key: ProgramKey, refresh_inputs: Callable[[TraceArtifact, RefreshDecision], None], *, reset_batch: bool = False, device_feedback_enabled: bool = False, feedback_compatible: bool = False, page_table_changed: bool = False, ) -> Any: """Refresh persistent inputs and enqueue one non-blocking trace replay.""" self._ensure_live() self.program_compiler.require_compiled(program_key) trace_key = self._program_to_trace.get(program_key) if trace_key is None: raise RuntimeError(f"Program key {program_key.digest} has no trace association") record = self._traces[trace_key] artifact = record.artifact if artifact is None: raise RuntimeError(f"Trace key {trace_key.digest} has not been captured") policy = artifact.refresh_policy switched = self._previous_replay_key != trace_key full = ( (policy.full_on_batch_reset and reset_batch) or (policy.full_on_graph_switch and switched) or (policy.full_without_device_feedback and not (device_feedback_enabled and feedback_compatible)) ) decision = RefreshDecision( full=full, page_table=policy.refresh_page_table_on_change and page_table_changed, fields=policy.every_replay, ) refresh_inputs(artifact, decision) ttnn.execute_trace(self.mesh_device, artifact.trace_id, cq_id=0, blocking=False) self._replay_count += 1 self._replay_counts[record.operation] += 1 self._previous_replay_key = trace_key return artifact.outputs def cleanup(self) -> None: """Release traces and persistent inputs, then reopen the program gate.""" if self._released: return failures = self._release_trace_resources() failures.extend(self._release_alias_workspaces()) failures.extend(release_orphans(self._rollback_orphans)) if failures: self._activated = True self.program_compiler.set_trace_active(True) error = RuntimeError(f"Failed to release {len(failures)} trace resource(s)") attach_cleanup_failures(error, failures) raise error from failures[0] self._capture_in_progress = False self._activated = False self._previous_replay_key = None self.program_compiler.set_trace_capture_in_progress(False) self.program_compiler.set_trace_active(False) self._released = True # Private implementation def _release_trace_resources(self) -> list[BaseException]: failures: list[BaseException] = [] for record in self._traces.values(): failures.extend(self._release_trace(record)) return failures def _release_alias_workspaces(self) -> list[BaseException]: failures: list[BaseException] = [] for alias in self._aliases.values(): if alias.workspace is None: continue alias_failures = best_effort_deallocate_owned_tensors( alias.workspace, alias.deallocated_tensor_ids, ) failures.extend(alias_failures) if not alias_failures: alias.workspace = None return failures def _release_trace(self, record: TraceRecord) -> list[BaseException]: artifact = record.artifact if artifact is None: return [] if not artifact.trace_released: try: ttnn.release_trace(self.mesh_device, artifact.trace_id) except BaseException as error: return [error] artifact.trace_released = True failures = best_effort_deallocate_owned_tensors( (artifact.persistent_inputs.values, artifact.outputs), artifact.deallocated_tensor_ids, ) if failures: return failures record.artifact = None return [] def _ensure_live(self) -> None: if self._released: raise RuntimeError("TraceCompiler has been released") def _trim_host_allocator() -> None: """Return released trace-capture staging arenas to the OS when supported.""" try: malloc_trim = ctypes.CDLL(None).malloc_trim except (AttributeError, OSError): return malloc_trim.argtypes = (ctypes.c_size_t,) malloc_trim.restype = ctypes.c_int malloc_trim(0)