code/models/common/llm_runtime/prefill/trace.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Prefill trace capture, refresh, replay state, and ownership."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable

from models.common.llm_runtime.prefill.inputs import (
    PrefillDeviceInputs,
    PrefillInputStager,
    PrefillPositionInputs,
    allocate_device_tensors,
    copy_into_device_tensors,
)
from models.common.llm_runtime.prefill.plan import PrefillChunk
from models.common.llm_runtime.prefill.postprocess import (
    KPTSignature,
    PrefillPostprocessor,
    new_logprob_output,
    without_borrowed,
)
from models.common.llm_runtime.prefill.result_collector import InvocationResult
from models.common.llm_runtime.prefill.signatures import (
    PrefillTraceSignature,
    PreparedPrefill,
    capture_schema_fingerprint,
    workspace_fingerprint,
)
from models.common.llm_runtime.tensor_resources import attach_cleanup_failures


@dataclass(frozen=True)
class PrefillHiddenPersistentInputs:
    """Canonical trace-owned model inputs; deliberately sampling-free."""

    device_inputs: PrefillDeviceInputs

    def owned_tensor_values(self) -> tuple[Any, ...]:
        return self.device_inputs.model_values()


@dataclass
class PrefillReplayState:
    """Mutable program-alias-local postprocessing and sampling state."""

    position_inputs: PrefillPositionInputs
    kpt: tuple[Any, Any, Any] | None
    sampled_output: Any | None = None
    position_signature: int | None = None
    kpt_signature: KPTSignature = None

    def owned_tensor_values(self) -> tuple[Any, ...]:
        return self.position_inputs.values(), self.kpt, self.sampled_output


@dataclass(frozen=True)
class PrefillReplayOwnership:
    """Explicit ownership split for one hidden-trace postprocess result."""

    trace_owned_hidden_output: Any
    nested_persistent_output: Any | None
    new_logprob_output: Any | None
    replay_local_intermediates: tuple[Any, ...]


@dataclass(frozen=True)
class PrefillCapturePlan:
    """Operation hooks consumed by the trace compiler."""

    signature: PrefillTraceSignature
    prepare_inputs: Callable[[], PrefillHiddenPersistentInputs]
    capture: Callable[[PrefillHiddenPersistentInputs], Any]
    prepare_workspace: Callable[[], PrefillReplayState]
    schema_fingerprint: tuple[Any, ...]
    workspace_fingerprint: tuple[Any, ...]
    refresh_fields: tuple[str, ...] = ("tokens", "page_table", "last_token", "sampling")
    prime: Callable[[PrefillHiddenPersistentInputs], Any] | None = None
    release_prime_output: Callable[[Any], list[BaseException]] | None = None


@dataclass(frozen=True)
class PrefillTraceHooks:
    """Narrow input, model-body, postprocess, and cleanup collaborators."""

    input_stager: PrefillInputStager
    postprocessor: PrefillPostprocessor
    run_hidden_body: Callable[..., Any]
    run_chunk_hidden_body: Callable[..., Any]
    release_transient: Callable[[Any], list[BaseException]]
    trace_capture_prime_sequence_lengths: tuple[int, ...] = ()


class PrefillTraceLifecycle:
    """Own hidden-trace capture, replay refresh, and replay result ownership."""

    def __init__(self, *, hooks: PrefillTraceHooks) -> None:
        self.hooks = hooks

    def capture_plan(self, prepared: PreparedPrefill) -> PrefillCapturePlan:
        if prepared.trace_signature is None:
            raise ValueError("prepared prefill request has no configured trace family")

        def prepare_inputs() -> PrefillHiddenPersistentInputs:
            return self._prepare_hidden_persistent_inputs(prepared)

        def prepare_workspace() -> PrefillReplayState:
            return self._prepare_replay_state(prepared)

        def capture(persistent: PrefillHiddenPersistentInputs) -> Any:
            if prepared.request.uses_chunked_prefill:
                return self.hooks.run_chunk_hidden_body(
                    prepared,
                    prepared.request.chunks[0],
                    persistent.device_inputs,
                )
            # The captured padded identity must record every physical row so
            # replay does not depend on which active count registered it.
            return self.hooks.run_hidden_body(
                prepared.request,
                persistent.device_inputs,
                fill_rows=prepared.request.padded_batch_size,
            )

        should_prime = (
            prepared.request.kind == "batched"
            and prepared.trace_signature.padded_sequence_length in self.hooks.trace_capture_prime_sequence_lengths
        )
        return PrefillCapturePlan(
            signature=prepared.trace_signature,
            prepare_inputs=prepare_inputs,
            capture=capture,
            prepare_workspace=prepare_workspace,
            schema_fingerprint=capture_schema_fingerprint(prepared),
            workspace_fingerprint=workspace_fingerprint(
                prepared,
                sampling_output_rows=self.hooks.postprocessor.sampling_output_rows(prepared),
            ),
            prime=capture if should_prime else None,
            release_prime_output=self.hooks.release_transient if should_prime else None,
        )

    def refresh(
        self,
        prepared: PreparedPrefill,
        persistent: PrefillHiddenPersistentInputs,
        state: PrefillReplayState,
        chunk: PrefillChunk | None = None,
    ) -> None:
        request = prepared.request
        if request.uses_chunked_prefill:
            self._refresh_chunk_inputs(prepared, chunk or request.chunks[0], persistent, state)
            return
        relative_last = max(last - cached for last, cached in zip(request.last_token_indices, request.cached_tokens))
        # Rotary positions are fixed for regular trace families; only tokens
        # and the page table vary on every replay.
        self.hooks.input_stager.refresh_regular_device_inputs(request, persistent.device_inputs)
        # Host-sampling single requests now use the runtime last-token bounds as well (see
        # finish_regular_prefill), so their position inputs must follow the prompt on every replay.
        if not self.hooks.postprocessor.uses_static_q128_topk(request, prepared.sampling_path):
            if state.position_signature != relative_last:
                position_inputs = self.hooks.input_stager.prepare_position_inputs_host(
                    relative_last,
                    request.padded_sequence_length,
                )
                copy_into_device_tensors(position_inputs.values(), state.position_inputs.values())
                state.position_signature = relative_last
        state.kpt_signature = self.hooks.postprocessor.refresh_workspace_sampling(
            prepared,
            kpt=state.kpt,
            kpt_signature=state.kpt_signature,
        )

    def finish(
        self,
        prepared: PreparedPrefill,
        hidden: Any,
        state: PrefillReplayState,
    ) -> InvocationResult:
        replay_local: list[Any] = []
        output = self.hooks.postprocessor.finish_regular_prefill(
            prepared,
            hidden,
            state.kpt if prepared.sampling_path == "topk" else None,
            state.position_inputs,
            sampled_output=state.sampled_output,
            owned=replay_local,
        )
        new_logprob = new_logprob_output(output, state.sampled_output)
        sampled_output_alias = output[0] if state.sampled_output is not None else None
        caller_owned = without_borrowed(
            replay_local,
            (hidden, state.sampled_output, sampled_output_alias),
        )
        ownership = PrefillReplayOwnership(
            trace_owned_hidden_output=hidden,
            nested_persistent_output=state.sampled_output,
            new_logprob_output=new_logprob,
            replay_local_intermediates=without_borrowed(
                replay_local,
                (hidden, state.sampled_output, sampled_output_alias, new_logprob),
            ),
        )
        return InvocationResult(value=output, owned=caller_owned, replay_ownership=ownership)

    def _prepare_hidden_persistent_inputs(self, prepared: PreparedPrefill) -> PrefillHiddenPersistentInputs:
        request = prepared.request
        trace_inputs = self.hooks.input_stager.trace_inputs(request)
        host_inputs = self.hooks.input_stager.prepare_host_inputs(
            trace_inputs.tokens,
            request.page_table,
            start_pos=trace_inputs.start_pos,
            chunk_page_table=trace_inputs.chunk_page_table,
            chunk_start_idx=trace_inputs.chunk_start_idx,
            last_token_idx=max(request.last_token_indices),
        )
        device_inputs = None
        try:
            device_inputs = self.hooks.input_stager.stage_device_inputs(host_inputs)
            if request.uses_chunked_prefill:
                self.hooks.input_stager.copy_rotary_inputs(device_inputs)
        except BaseException as primary:
            failures = self.hooks.release_transient(device_inputs)
            attach_cleanup_failures(primary, failures)
            raise
        return PrefillHiddenPersistentInputs(device_inputs=device_inputs)

    def _prepare_replay_state(self, prepared: PreparedPrefill) -> PrefillReplayState:
        trace_inputs = self.hooks.input_stager.trace_inputs(prepared.request)
        position_inputs = None
        kpt = None
        sampled_output = None
        try:
            sampling_batch_size = self.hooks.postprocessor.sampling_output_rows(prepared)
            position_values = allocate_device_tensors(
                self.hooks.input_stager.prepare_position_inputs_host(
                    trace_inputs.relative_last,
                    trace_inputs.sequence_length,
                ).values(),
                mesh_device=self.hooks.input_stager.mesh_device,
            )
            position_inputs = PrefillPositionInputs(*position_values)
            kpt = self.hooks.postprocessor.make_device_kpt(
                self.hooks.postprocessor.prepared_sampling(prepared),
                sampling_batch_size,
                force_topk=prepared.sampling_path == "topk",
            )
            prepared_sampling = self.hooks.postprocessor.prepared_sampling(prepared)
            if prepared_sampling is not None:
                sampled_output = self.hooks.postprocessor.make_sampling_output(
                    self.hooks.postprocessor.sampling_output_rows(prepared)
                )
        except BaseException as primary:
            failures = self.hooks.release_transient((position_inputs, kpt, sampled_output))
            attach_cleanup_failures(primary, failures)
            raise
        kpt_signature = None
        if prepared_sampling is not None:
            kpt_signature = self.hooks.postprocessor.kpt_values(prepared_sampling, sampling_batch_size)
        return PrefillReplayState(
            position_inputs=position_inputs,
            kpt=kpt,
            sampled_output=sampled_output,
            position_signature=trace_inputs.relative_last,
            kpt_signature=kpt_signature,
        )

    def _refresh_chunk_inputs(
        self,
        prepared: PreparedPrefill,
        chunk: PrefillChunk,
        persistent: PrefillHiddenPersistentInputs,
        state: PrefillReplayState,
    ) -> None:
        request = prepared.request
        self.hooks.input_stager.refresh_chunk_device_inputs(request, chunk, persistent.device_inputs)

        final_chunk = request.chunks[-1]
        relative_last = (request.last_token_indices[0] - final_chunk.chunk_start_idx) % final_chunk.chunk_size
        if state.position_signature != relative_last:
            position_inputs = self.hooks.input_stager.prepare_position_inputs_host(
                relative_last,
                final_chunk.chunk_size,
            )
            copy_into_device_tensors(position_inputs.values(), state.position_inputs.values())
            state.position_signature = relative_last
        state.kpt_signature = self.hooks.postprocessor.refresh_workspace_sampling(
            prepared,
            kpt=state.kpt,
            kpt_signature=state.kpt_signature,
        )