File size: 9,159 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
import pytest
import torch
import models.common.llm_runtime.program_compiler as program_compiler_module
import ttnn
from models.common.llm_runtime.program_compiler import OutputSpec, ProgramCompiler, ProgramKey
@dataclass(frozen=True)
class _Signature:
mode: str
batch: int
optional: int | None = None
runtime_token: int = field(default=0, compare=False)
@property
def key_material(self):
return (("mode", self.mode), ("batch", self.batch), ("optional", self.optional))
def _compiler(context=object()):
return ProgramCompiler("mesh", lambda: context)
def _patch_sync(monkeypatch, events=None):
events = [] if events is None else events
monkeypatch.setattr(ttnn, "synchronize_device", lambda mesh: events.append(("sync", mesh)))
return events
def test_program_key_has_stable_golden_digest_and_separate_trace_domain():
from models.common.llm_runtime.trace_compiler import TraceKey
signature = _Signature("decode", 32)
assert ProgramKey.from_signature(signature).digest == (
"c807e5574b34f3f83d386fc441580d1ae499aa534c8a426541dbeafdefa2ffdb"
)
assert TraceKey.from_signature(signature).digest == (
"27bea60de6b500f448e680d3423ad62efbd01294c254524e37c114ff8b5dff35"
)
assert ProgramKey.from_signature(signature).digest != TraceKey.from_signature(signature).digest
def test_program_key_changes_for_each_material_field_but_not_runtime_values():
base = _Signature("decode", 32)
assert ProgramKey.from_signature(base) != ProgramKey.from_signature(_Signature("prefill", 32))
assert ProgramKey.from_signature(base) != ProgramKey.from_signature(_Signature("decode", 16))
assert ProgramKey.from_signature(base) != ProgramKey.from_signature(_Signature("decode", 32, 1))
assert ProgramKey.from_signature(base) == ProgramKey.from_signature(_Signature("decode", 32, runtime_token=99))
def test_compiler_memoizes_program_key_by_canonical_material(monkeypatch):
calls = []
digest = program_compiler_module.signature_digest
monkeypatch.setattr(
program_compiler_module,
"signature_digest",
lambda domain, schema_version, signature: calls.append((domain, schema_version, signature))
or digest(domain, schema_version, signature),
)
compiler = _compiler()
first = compiler.key_for(_Signature("decode", 32))
second = compiler.key_for(_Signature("decode", 32, runtime_token=99))
assert second is first
assert len(calls) == 1
def test_compiler_program_key_memo_tracks_mutated_material(monkeypatch):
class MutableSignature:
def __init__(self):
self.value = 1
def key_material(self):
return (("value", self.value),)
calls = []
digest = program_compiler_module.signature_digest
monkeypatch.setattr(
program_compiler_module,
"signature_digest",
lambda domain, schema_version, signature: calls.append((domain, schema_version, signature))
or digest(domain, schema_version, signature),
)
compiler = _compiler()
signature = MutableSignature()
first = compiler.key_for(signature)
signature.value = 2
second = compiler.key_for(signature)
assert second != first
assert compiler.key_for(signature) is second
assert len(calls) == 2
def test_compiler_program_key_memo_preserves_tagged_type_identity():
compiler = _compiler()
keys = {
compiler.key_for(type("Signature", (), {"key_material": (("value", value),)})())
for value in (True, 1, 1.0, "1")
}
assert len(keys) == 4
@pytest.mark.parametrize("material", [{"unordered": "mapping"}, ["list"], torch.zeros(1)])
def test_program_key_rejects_noncanonical_material(material, expect_error):
class InvalidSignature:
key_material = (("invalid", material),)
with expect_error(TypeError, "key material"):
ProgramKey.from_signature(InvalidSignature())
def test_same_digest_with_different_retained_signature_is_rejected(monkeypatch, expect_error):
_patch_sync(monkeypatch)
compiler = _compiler()
forced_key = ProgramKey("0" * 64)
monkeypatch.setattr(compiler, "key_for", lambda signature: forced_key)
compiler.compile(_Signature("decode", 32), lambda context: torch.zeros(1))
with expect_error(RuntimeError, "collision"):
compiler.compile(_Signature("prefill", 1), lambda context: torch.zeros(1))
def test_compile_receives_exact_context_deduplicates_and_checks_output_contract(monkeypatch, expect_error):
events = _patch_sync(monkeypatch)
context = object()
compiler = _compiler(context)
calls = []
expected = OutputSpec((2, 3), torch.bfloat16)
program = compiler.compile(
_Signature("decode", 32),
lambda supplied: calls.append(supplied) or torch.zeros(2, 3, dtype=torch.bfloat16),
expected_output_spec=expected,
)
duplicate = compiler.compile(
_Signature("decode", 32),
lambda supplied: (_ for _ in ()).throw(AssertionError("duplicate invoked")),
expected_output_spec=expected,
)
assert program is duplicate
assert calls == [context]
assert program.signature == _Signature("decode", 32)
assert program.output_spec == expected
assert events == [("sync", "mesh"), ("sync", "mesh")]
with expect_error(ValueError, "different output contract"):
compiler.compile(
_Signature("decode", 32),
lambda supplied: torch.zeros(1),
expected_output_spec=OutputSpec((1,), torch.float32),
)
def test_compiled_program_snapshot_is_immutable_and_registry_authoritative(monkeypatch):
_patch_sync(monkeypatch)
compiler = _compiler()
first = compiler.compile(_Signature("prefill", 1), lambda _: torch.zeros(1))
snapshot = compiler.compiled_programs
second = compiler.compile(_Signature("decode", 32), lambda _: torch.zeros(1))
assert snapshot == (first,)
assert compiler.compiled_programs == (first, second)
def test_compile_uses_explicit_result_value_and_owned_release_selector(monkeypatch):
_patch_sync(monkeypatch)
class OwnedTensor:
pass
@dataclass
class InvocationResult:
value: torch.Tensor
owned: object
released = []
monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
monkeypatch.setattr(ttnn, "deallocate", released.append)
owned = OwnedTensor()
compiler = _compiler()
program = compiler.compile(
_Signature("prefill", 1),
lambda context: InvocationResult(torch.zeros(1, 2), owned),
output_spec=lambda result: OutputSpec.from_value(result.value),
release_output=lambda result: result.owned,
)
assert program.output_spec.shape == (1, 2)
assert released == [owned]
def test_compile_release_failure_is_retryable_and_blocks_further_compile(monkeypatch, expect_error):
_patch_sync(monkeypatch)
class OwnedTensor:
pass
first = OwnedTensor()
retry = OwnedTensor()
attempts = []
failure = RuntimeError("release once")
def deallocate(value):
attempts.append(value)
if value is retry and attempts.count(retry) == 1:
raise failure
monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
monkeypatch.setattr(ttnn, "deallocate", deallocate)
compiler = _compiler()
with expect_error(RuntimeError, "Failed to deallocate") as caught:
compiler.compile(
_Signature("decode", 32),
lambda context: (first, retry),
output_spec=lambda output: OutputSpec((1,), "dtype"),
)
assert caught.value.cleanup_failures == (failure,)
assert compiler.compile_orphan_count == 1
with expect_error(RuntimeError, "unreleased compile outputs"):
compiler.compile(_Signature("decode", 16), lambda context: torch.zeros(1))
compiler.cleanup()
assert attempts.count(first) == 1
assert attempts.count(retry) == 2
def test_compile_gate_distinguishes_capture_from_activation(monkeypatch, expect_error):
_patch_sync(monkeypatch)
compiler = _compiler()
compiler.set_trace_capture_in_progress(True)
with expect_error(RuntimeError, "capture is in progress"):
compiler.compile(_Signature("decode", 32), lambda context: torch.zeros(1))
assert compiler.post_activation_compile_rejections == 0
compiler.set_trace_capture_in_progress(False)
compiler.set_trace_active(True)
with expect_error(RuntimeError, "after trace activation"):
compiler.compile(_Signature("decode", 32), lambda context: torch.zeros(1))
assert compiler.post_activation_compile_rejections == 1
def test_cleanup_terminalizes_only_program_metadata(monkeypatch, expect_error):
_patch_sync(monkeypatch)
compiler = _compiler()
program = compiler.compile(_Signature("decode", 32), lambda context: torch.zeros(1))
compiler.cleanup()
compiler.cleanup()
with expect_error(RuntimeError, "released"):
compiler.require_compiled(program.key)
|