# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC # SPDX-License-Identifier: Apache-2.0 import inspect from dataclasses import dataclass from types import SimpleNamespace import torch import models.common.llm_runtime.execution as execution_module import ttnn from models.common.llm_runtime.decode import DecodeRuntime from models.common.llm_runtime.decode import InvocationResult as DecodeInvocationResult from models.common.llm_runtime.execution import EagerExecutor, TracedExecutor from models.common.llm_runtime.prefill.result_collector import InvocationResult as PrefillInvocationResult from models.common.llm_runtime.prefill.runtime import PrefillRuntime from models.common.llm_runtime.program_compiler import ProgramCompiler from models.common.llm_runtime.trace_compiler import TraceCompiler @dataclass(frozen=True) class _Signature: operation: str variant: int @property def key_material(self): return (("operation", self.operation), ("variant", self.variant)) def _runtime(runtime_type, **methods): runtime = object.__new__(runtime_type) for name, method in methods.items(): setattr(runtime, name, method) return runtime def _compiler(monkeypatch): monkeypatch.setattr(ttnn, "synchronize_device", lambda mesh: None) return ProgramCompiler("mesh", lambda: object()) def _trace_compiler(program_compiler, *, mode="all"): return TraceCompiler(program_compiler) def _prepared_prefill(*, trace_eligible=True, signatures=None, name="regular"): if signatures is None: signatures = (_Signature("prefill", 1),) return SimpleNamespace( name=name, program_signatures=signatures, trace_eligible=trace_eligible, trace_signature=_Signature("prefill-trace", 1) if trace_eligible else None, ) def _prepared_decode(*, variant=1): return SimpleNamespace( variant=variant, device_feedback=True, reset_batch=False, page_table_changed=False, sampling_params=None, ) def test_execution_strategies_use_exact_identity_composition_without_type_frameworks(monkeypatch): prefill = _runtime(PrefillRuntime) decode = _runtime(DecodeRuntime) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=decode, program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) assert eager.prefill is prefill assert eager.decode is decode assert eager.program_compiler is program_compiler assert traced.eager_executor is eager assert traced.trace_compiler is trace_compiler assert EagerExecutor not in TracedExecutor.__mro__ assert EagerExecutor.__bases__ == (object,) assert TracedExecutor.__bases__ == (object,) source = inspect.getsource(execution_module) assert "Protocol" not in source assert "ABC" not in source assert "LightweightModule" not in source assert not hasattr(execution_module, "EagerExecutorConfig") assert not hasattr(execution_module, "TracedExecutorConfig") assert not hasattr(EagerExecutor, "cleanup") assert not hasattr(TracedExecutor, "cleanup") def test_execution_request_signatures_are_exact_and_aligned(): required = inspect.Parameter.empty positional = inspect.Parameter.POSITIONAL_OR_KEYWORD keyword_only = inspect.Parameter.KEYWORD_ONLY prefill_contract = [ ("self", positional, required), ("tokens", keyword_only, required), ("page_table", keyword_only, required), ("prompt_lens", keyword_only, None), ("start_pos", keyword_only, None), ("empty_slots", keyword_only, None), ("sampling_params", keyword_only, None), ("prompt_tokens", keyword_only, None), ("output_tokens", keyword_only, None), ("slot_remap", keyword_only, None), ] decode_contract = [ ("self", positional, required), ("tokens", keyword_only, required), ("start_pos", keyword_only, required), ("page_table", keyword_only, required), ("sampling_params", keyword_only, None), ("prompt_tokens", keyword_only, None), ("output_tokens", keyword_only, None), ("slot_remap", keyword_only, None), ("reset_batch", keyword_only, False), ] decode_forward_contract = [ *decode_contract, ("read_from_device", keyword_only, True), ] def parameter_contract(method): return [ (parameter.name, parameter.kind, parameter.default) for parameter in inspect.signature(method).parameters.values() ] for executor_type in (EagerExecutor, TracedExecutor): assert parameter_contract(executor_type.compile_prefill) == prefill_contract assert parameter_contract(executor_type.prefill_forward) == prefill_contract assert parameter_contract(executor_type.compile_decode) == decode_contract assert parameter_contract(executor_type.decode_forward) == decode_forward_contract assert parameter_contract(EagerExecutor._prepare_prefill) == prefill_contract assert parameter_contract(EagerExecutor._prepare_decode) == decode_contract for method_name in ("compile_prefill", "prefill_forward", "compile_decode", "decode_forward"): assert inspect.signature(getattr(EagerExecutor, method_name)) == inspect.signature( getattr(TracedExecutor, method_name) ) def test_execution_request_methods_reject_kv_cache(expect_error): prefill_fields = { "tokens": torch.zeros(1, 1), "page_table": torch.zeros(1, 1), } decode_fields = { "tokens": torch.zeros(1, 1), "start_pos": torch.zeros(1), "page_table": torch.zeros(1, 1), } for executor_type in (EagerExecutor, TracedExecutor): executor = object.__new__(executor_type) for method_name, fields in ( ("compile_prefill", prefill_fields), ("prefill_forward", prefill_fields), ("compile_decode", decode_fields), ("decode_forward", decode_fields), ): with expect_error(TypeError, "kv_cache"): getattr(executor, method_name)(**fields, kv_cache=object()) def test_traced_constructor_rejects_a_different_program_compiler(monkeypatch, expect_error): eager = EagerExecutor( prefill=_runtime(PrefillRuntime), decode=_runtime(DecodeRuntime), program_compiler=_compiler(monkeypatch), ) unrelated_trace_compiler = _trace_compiler(_compiler(monkeypatch)) with expect_error(ValueError, "compose eager.program_compiler"): TracedExecutor(eager=eager, trace_compiler=unrelated_trace_compiler) def test_eager_prefill_prepares_once_and_compiles_all_signatures_from_same_object(monkeypatch): prepared = _prepared_prefill( signatures=(_Signature("prefill", 1), _Signature("prefill", 2)), ) prepared_seen = [] prepare_calls = [] def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): prepare_calls.append( { "tokens": tokens, "page_table": page_table, "prompt_lens": prompt_lens, "start_pos": start_pos, "empty_slots": empty_slots, "sampling_params": sampling_params, } ) return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, invoke=lambda prepared, *, count_tokens=True: prepared_seen.append((prepared, count_tokens)) or PrefillInvocationResult(torch.zeros(1), ()), ) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=_compiler(monkeypatch)) tokens = torch.zeros(1, 1) page_table = torch.zeros(1, 1) prompt_lens = torch.tensor([1]) start_pos = torch.tensor([0]) empty_slots = [0] sampling_params = object() eager.compile_prefill( tokens=tokens, page_table=page_table, prompt_lens=prompt_lens, start_pos=start_pos, empty_slots=empty_slots, sampling_params=sampling_params, ) assert prepare_calls == [ { "tokens": tokens, "page_table": page_table, "prompt_lens": prompt_lens, "start_pos": start_pos, "empty_slots": empty_slots, "sampling_params": sampling_params, } ] assert prepared_seen == [(prepared, False), (prepared, False)] def test_traced_prefill_compile_does_not_interpret_request_eligibility(monkeypatch): prepared = _prepared_prefill(trace_eligible=False) identity_events = [] operation_plan = SimpleNamespace( signature=_Signature("prefill-trace", 1), prepare_inputs=lambda: (), capture=lambda persistent: torch.zeros(1), refresh_fields=("tokens",), prime=None, release_prime_output=None, ) def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): identity_events.append(("prepare", prepared)) return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, invoke=lambda prepared, *, count_tokens=True: identity_events.append(("invoke", prepared, count_tokens)) or PrefillInvocationResult(torch.zeros(1), ()), capture_plan=lambda prepared: identity_events.append(("capture_plan", prepared)) or operation_plan, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) registered = [] trace_compiler.register_capture_plan = registered.append traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) traced.compile_prefill(tokens=torch.zeros(1, 1), page_table=torch.zeros(1, 1)) assert [event[0] for event in identity_events] == ["prepare", "invoke", "capture_plan"] assert all(event[1] is prepared for event in identity_events) assert identity_events[1][2] is False assert len(registered) == 1 def test_traced_prefill_recompile_reuses_existing_trace_association(monkeypatch): prepared = _prepared_prefill() def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, invoke=lambda prepared, *, count_tokens=True: PrefillInvocationResult(torch.zeros(1), ()), capture_plan=lambda prepared: (_ for _ in ()).throw(AssertionError("capture plan rebuilt")), ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) trace_compiler.trace_key_for_program = lambda program_key: "existing-trace" trace_compiler.register_capture_plan = lambda plan: (_ for _ in ()).throw(AssertionError("plan registered")) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) traced.compile_prefill(tokens=torch.zeros(1, 1), page_table=torch.zeros(1, 1)) def test_traced_decode_recompile_reuses_existing_trace_association(monkeypatch): prepared = _prepared_decode() def prepare(*, tokens, start_pos, page_table, sampling_params=None, reset_batch=False): return prepared decode = _runtime( DecodeRuntime, prepare=prepare, program_signature=lambda prepared: _Signature("decode", 1), invoke=lambda prepared, *, device_feedback=False, count_tokens=True: DecodeInvocationResult( torch.zeros(1), (), False ), capture_plan=lambda prepared: (_ for _ in ()).throw(AssertionError("capture plan rebuilt")), ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=_runtime(PrefillRuntime), decode=decode, program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) trace_compiler.trace_key_for_program = lambda program_key: "existing-trace" trace_compiler.register_capture_plan = lambda plan: (_ for _ in ()).throw(AssertionError("plan registered")) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) traced.compile_decode(tokens=torch.zeros(1), start_pos=torch.zeros(1), page_table=torch.zeros(1, 1)) def test_execution_target_selection_is_external_to_traced_prefill(monkeypatch): prepared = _prepared_prefill(trace_eligible=True) invocations = [] def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, invoke=lambda prepared: invocations.append(prepared) or PrefillInvocationResult("eager", ()), assemble=lambda prepared_results, *, batch_size, sampling_params=None: prepared_results[0][1].value, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler, mode="decode_only") def replay( program_key, refresh_inputs, *, reset_batch=False, device_feedback_enabled=False, feedback_compatible=False, page_table_changed=False, ): raise AssertionError("trace replayed") trace_compiler.replay = replay traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) assert ( traced.eager_executor.prefill_forward( tokens=torch.zeros(1, 1), page_table=torch.zeros(1, 1), ) == "eager" ) assert invocations == [prepared] def test_prefill_replay_does_not_interpret_request_eligibility(monkeypatch): prepared = _prepared_prefill(trace_eligible=False) prepared.trace_signature = _Signature("prefill-trace", 1) persistent = object() hidden = object() identity_events = [] def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): identity_events.append(("prepare", prepared)) return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, refresh_trace=lambda prepared, persistent: identity_events.append(("refresh", prepared, persistent)), finish_trace=lambda prepared, hidden, persistent: identity_events.append( ("finish", prepared, hidden, persistent) ) or "traced", assemble=lambda prepared_results, *, batch_size, sampling_params=None: next(iter(prepared_results))[1], ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=persistent)) record = SimpleNamespace(artifact=artifact) trace_compiler.replay = ( lambda program_key, refresh_inputs, *, reset_batch=False, device_feedback_enabled=False, feedback_compatible=False, page_table_changed=False: refresh_inputs( artifact, object() ) or hidden ) trace_compiler.trace_key_for_program = lambda program_key: "trace-key" trace_compiler.get = lambda trace_key: record traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) result = traced.prefill_forward(tokens=torch.zeros(1, 1), page_table=torch.zeros(1, 1)) assert result == "traced" assert [event[0] for event in identity_events] == ["prepare", "refresh", "finish"] assert all(event[1] is prepared for event in identity_events) assert identity_events[1][2] is persistent assert identity_events[2][2:] == (hidden, persistent) def test_prefill_replay_is_consumed_before_shared_trace_output_is_overwritten(monkeypatch): prepared = ( _prepared_prefill(name="first"), _prepared_prefill(name="second"), ) persistent = {"output": None} def prepare(**kwargs): return prepared def refresh_trace(request, trace_inputs): trace_inputs["output"] = request.name def assemble(prepared_results, *, batch_size, sampling_params=None): return [result.value["output"] for _, result in prepared_results] prefill = _runtime( PrefillRuntime, prepare=prepare, refresh_trace=refresh_trace, finish_trace=lambda request, hidden, trace_inputs: PrefillInvocationResult(trace_inputs, ()), assemble=assemble, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=persistent)) trace_compiler.replay = lambda program_key, refresh_inputs, **kwargs: refresh_inputs(artifact, object()) or "hidden" trace_compiler.trace_key_for_program = lambda program_key: "trace-key" trace_compiler.get = lambda trace_key: SimpleNamespace(artifact=artifact) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) assert traced.prefill_forward(tokens=torch.zeros(2, 1), page_table=torch.zeros(2, 1)) == [ "first", "second", ] def test_prefill_missing_trace_artifact_is_an_error_without_eager_reinvocation(monkeypatch, expect_error): prepared = _prepared_prefill(trace_eligible=True) eager_invocations = [] def prepare( *, tokens, page_table, prompt_lens=None, start_pos=None, empty_slots=None, sampling_params=None, ): return (prepared,) prefill = _runtime( PrefillRuntime, prepare=prepare, invoke=lambda prepared: eager_invocations.append(prepared) or PrefillInvocationResult("eager", ()), refresh_trace=lambda prepared, persistent: None, assemble=lambda prepared_results, *, batch_size, sampling_params=None: next(iter(prepared_results))[1], ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=())) trace_compiler.replay = ( lambda program_key, refresh_inputs, *, reset_batch=False, device_feedback_enabled=False, feedback_compatible=False, page_table_changed=False: refresh_inputs( artifact, object() ) or "hidden" ) trace_compiler.trace_key_for_program = lambda program_key: "missing" trace_compiler.get = lambda trace_key: None traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) with expect_error(RuntimeError, "Required prefill trace") as exc_info: traced.prefill_forward(tokens=torch.zeros(1, 1), page_table=torch.zeros(1, 1)) assert eager_invocations == [] message = str(exc_info.value) for field in ( "operation=prefill", "trace_mode=all", "model=PrefillRuntime", "signature_material=", "signature_digest=", "configured_coverage=", "TraceConfig(mode='none')", ): assert field in message assert traced.coverage_miss_count == 1 def test_prefill_missing_trace_signature_is_rejected_before_replay(monkeypatch, expect_error): prepared = _prepared_prefill(trace_eligible=False) program_compiler = _compiler(monkeypatch) eager = EagerExecutor( prefill=_runtime(PrefillRuntime), decode=_runtime(DecodeRuntime), program_compiler=program_compiler, ) trace_compiler = _trace_compiler(program_compiler) trace_compiler.replay = lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("unexpected replay")) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) with expect_error(RuntimeError, "not trace-eligible"): traced._execute_prefill(prepared) def test_decode_missing_trace_artifact_reports_exact_strict_coverage(monkeypatch, expect_error): prepared = _prepared_decode() signature = _Signature("decode", 7) decode = _runtime( DecodeRuntime, config=SimpleNamespace(position_feedback_capable=True), program_signature=lambda value: signature, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=_runtime(PrefillRuntime), decode=decode, program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) trace_compiler.trace_key_for_program = lambda key: None traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler, trace_mode="decode_only") with expect_error(RuntimeError, "Required decode trace") as exc_info: traced._execute_decode(prepared) message = str(exc_info.value) for field in ( "operation=decode", "trace_mode=decode_only", "model=DecodeRuntime", "signature_material=", "signature_digest=", "program_key=", "trace_key=unavailable", "configured_coverage=", "TraceConfig(mode='none')", ): assert field in message assert traced.coverage_miss_count == 1 def test_prefill_whole_call_preflights_every_trace_before_first_replay(monkeypatch, expect_error): first = _prepared_prefill(signatures=(_Signature("prefill", 1),), name="first") second = _prepared_prefill(signatures=(_Signature("prefill", 2),), name="second") prefill = _runtime( PrefillRuntime, prepare=lambda **kwargs: (first, second), assemble=lambda results, **kwargs: tuple(results), ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) first_key = program_compiler.key_for(first.program_signatures[0]) second_key = program_compiler.key_for(second.program_signatures[0]) artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=())) trace_compiler.trace_key_for_program = lambda key: "first" if key == first_key else "second" trace_compiler.get = lambda key: SimpleNamespace(artifact=artifact) if key == "first" else None replays = [] trace_compiler.replay = lambda *args, **kwargs: replays.append(args) or "hidden" traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) with expect_error(RuntimeError, second_key.digest): traced.prefill_forward(tokens=torch.zeros(2, 1), page_table=torch.zeros(2, 1)) assert replays == [] def test_prefill_fixed_chunk_replays_every_step_and_finishes_only_final_hidden(monkeypatch): prepared = _prepared_prefill() chunks = ("chunk-0", "chunk-1", "chunk-2") prepared.request = SimpleNamespace(chunks=chunks) events = [] prefill = _runtime( PrefillRuntime, refresh_trace=lambda request, hidden_inputs, workspace, chunk: events.append(("refresh", chunk, workspace)), finish_trace=lambda request, hidden, persistent: events.append(("finish", hidden, persistent)) or PrefillInvocationResult(hidden, ()), ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) persistent = object() artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=persistent)) record = SimpleNamespace(artifact=artifact) trace_compiler.trace_key_for_program = lambda key: "trace" trace_compiler.get = lambda key: record trace_compiler.workspace_for_program = lambda key: persistent replayed = [] def replay(program_key, refresh_inputs, **kwargs): refresh_inputs(artifact, object()) replayed.append(program_key) return f"hidden-{len(replayed) - 1}" trace_compiler.replay = replay traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) result = traced._execute_prefill(prepared) assert result.value == "hidden-2" assert [event[:2] for event in events] == [ ("refresh", "chunk-0"), ("refresh", "chunk-1"), ("refresh", "chunk-2"), ("finish", "hidden-2"), ] assert all(event[-1] is persistent for event in events) def test_prefill_replay_emits_structured_serving_evidence(monkeypatch): class Request: source_rows = tuple(range(15)) padded_batch_size = 16 padded_sequence_length = 128 chunks = ("step",) monkeypatch.setattr(execution_module, "PrefillRequest", Request) signature = SimpleNamespace( operation_variant="regular-batched", key_material=(("operation_variant", "regular-batched"),), ) prepared = SimpleNamespace( request=Request(), program_signatures=(signature,), trace_signature=SimpleNamespace(key_material=(("padded_batch_size", 16),)), sampling_path="topk", sampling_params=object(), ) persistent = object() prefill = _runtime( PrefillRuntime, refresh_trace=lambda request, hidden_inputs, workspace, chunk: None, finish_trace=lambda request, hidden, workspace: PrefillInvocationResult(hidden, ()), assemble=lambda results, **kwargs: next(iter(results))[1], ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=prefill, decode=_runtime(DecodeRuntime), program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) program_key = program_compiler.key_for(signature) trace_key = SimpleNamespace(digest="1" * 64) artifact = SimpleNamespace(persistent_inputs=SimpleNamespace(values=persistent)) record = SimpleNamespace(artifact=artifact) trace_compiler.trace_key_for_program = lambda key: trace_key trace_compiler.workspace_for_program = lambda key: persistent trace_compiler.replay = lambda key, refresh, **kwargs: refresh(artifact, object()) or "hidden" traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) result = traced.execute_prepared_prefill( ((prepared, ((program_key, record),)),), batch_size=15, sampling_params=prepared.sampling_params, lane=3, ) assert result.value == "hidden" assert len(traced.recent_prefill_replay_evidence) == 1 evidence = traced.recent_prefill_replay_evidence[0] assert evidence.operation == "prefill" assert evidence.variant == "regular-batched" assert evidence.sampling_path == "topk" assert evidence.execution == "trace_replay" assert (evidence.active_batch_size, evidence.padded_batch_size) == (15, 16) assert evidence.padded_sequence_length == 128 assert (evidence.lane, evidence.rank) == (3, 3) assert evidence.program_key == program_key.digest assert evidence.trace_key == "1" * 64 assert evidence.replay_steps == 1 assert traced.runtime_summary() == { "eager_prefill_executions": 0, "semantic_program_count": 0, "rejected_post_activation_compile_attempts": 0, "ttnn_program_cache_count": None, "successful_trace_replays": 0, "trace_replays_by_operation": {"prefill": 0, "decode": 0}, "strict_coverage_misses": 0, "semantic_trace_count": 0, "trace_association_count": 0, } def test_decode_replay_prepares_once_and_uses_same_object_for_refresh_submission_and_consume(monkeypatch): prepared = _prepared_decode() events = [] def prepare(*, tokens, start_pos, page_table, sampling_params=None, reset_batch=False): events.append(("prepare", prepared, sampling_params, reset_batch)) return prepared decode = _runtime( DecodeRuntime, config=SimpleNamespace(position_feedback_capable=True), prepare=prepare, program_signature=lambda prepared: events.append(("signature", prepared)) or _Signature("decode", 1), refresh_trace=lambda artifact, prepared, decision: events.append(("refresh", prepared)), note_submitted=lambda prepared: events.append(("submitted", prepared)), consume=lambda result, *, read_from_device=True: events.append(("consume", result, read_from_device)) or result.value, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=_runtime(PrefillRuntime), decode=decode, program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) trace_key = SimpleNamespace(digest="2" * 64) trace_compiler.trace_key_for_program = lambda key: trace_key trace_compiler.get = lambda key: SimpleNamespace(artifact=object()) trace_compiler.replay = ( lambda program_key, refresh_inputs, *, reset_batch=False, device_feedback_enabled=False, feedback_compatible=False, page_table_changed=False: refresh_inputs( object(), object() ) or "token" ) traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) sampling_params = object() result = traced.decode_forward( tokens=torch.zeros(1, 1), start_pos=torch.zeros(1), page_table=torch.zeros(1, 1), sampling_params=sampling_params, reset_batch=True, read_from_device=False, ) assert result == "token" assert [event[0] for event in events] == ["prepare", "signature", "refresh", "submitted", "consume"] assert all(event[1] is prepared for event in events[:-1]) assert events[0][2:] == (sampling_params, True) assert isinstance(events[-1][1], DecodeInvocationResult) assert events[-1][1].owned is None assert events[-1][2] is False def test_explicit_eager_decode_delegates_once_and_execution_objects_do_not_cleanup(monkeypatch): prepared = _prepared_decode() calls = [] def prepare(*, tokens, start_pos, page_table, sampling_params=None, reset_batch=False): calls.append(("prepare", prepared, sampling_params, reset_batch)) return prepared decode = _runtime( DecodeRuntime, prepare=prepare, invoke=lambda prepared, *, device_feedback=False: calls.append(("invoke", prepared, device_feedback)) or DecodeInvocationResult("eager", (), False), consume=lambda result, *, read_from_device=True: calls.append(("consume", result.value, read_from_device)) or result.value, ) program_compiler = _compiler(monkeypatch) eager = EagerExecutor(prefill=_runtime(PrefillRuntime), decode=decode, program_compiler=program_compiler) trace_compiler = _trace_compiler(program_compiler) def replay( program_key, refresh_inputs, *, reset_batch=False, device_feedback_enabled=False, feedback_compatible=False, page_table_changed=False, ): raise AssertionError("trace replayed") trace_compiler.replay = replay traced = TracedExecutor(eager=eager, trace_compiler=trace_compiler) sampling_params = object() assert ( traced.eager_executor.decode_forward( tokens=torch.zeros(1, 1), start_pos=torch.zeros(1), page_table=torch.zeros(1, 1), sampling_params=sampling_params, reset_batch=True, read_from_device=False, ) == "eager" ) assert calls == [ ("prepare", prepared, sampling_params, True), ("invoke", prepared, False), ("consume", "eager", False), ]