code/models/common/llm_runtime/prefill/README.md

Prefill Runtime

models/common/llm_runtime/prefill implements the common host-side prefill pipeline for TTTv2 language models. It plans logical requests, stages TT inputs, runs eager or traced prefill, postprocesses device outputs, restores the caller's row order, and releases invocation resources.

This document describes the internals of the prefill package. See the common LLM runtime README for model-executor composition, program and trace registries, paged-KV allocation, warmup, serving lifecycle, and cleanup order.

Scope

The package owns:

• immutable prefill request and chunk planning;
• regular-single, regular-batched, and chunked invocation geometry;
• eager program and trace identity construction;
• host and device input staging;
• trace capture inputs, replay refresh, and replay ownership;
• prefill postprocessing and device sampling state;
• streaming synchronized result collection; and
• retryable cleanup of prefill invocation transients.

The package does not choose a model, allocate the paged KV cache, select eager versus traced execution, own the program or trace registries, or adapt an inference-server API. Those decisions remain with the model executor and the common runtime components above this package.

Composition and construction

A model executor resolves one immutable config and constructs one prefill runtime per execution lane. PrefillRuntime then composes the focused package collaborators around the exact model, mesh, output reader, and page-table layout held by that config:

model + mesh + OutputReader + PageTableLayout + static prefill policy
  │
  └─> PrefillRuntimeConfig.resolve(...)
        │
        └─> PrefillRuntime(config)
              │
              ├─> PrefillInputStager(model, mesh)
              │
              ├─> PrefillPostprocessor(config)
              │
              ├─> PrefillResultAssembler(
              │     config,
              │     postprocessor,
              │     transient-release callback)
              │
              ├─> PrefillSequenceRunner(
              │     input stager,
              │     postprocessor,
              │     model-body callbacks,
              │     transient-release callback)
              │
              └─> PrefillTraceLifecycle(
                    input stager,
                    postprocessor,
                    model-body callbacks,
                    transient-release callback)

The collaborators share identities rather than constructing replacements:

• the input stager, postprocessor, sequence runner, and trace lifecycle use the exact objects composed by PrefillRuntime;
• eager and traced paths call the same runtime-owned model-body boundaries;
• eager invocation and trace replay finish through the same postprocessor;
• eager and traced results enter the same streaming result collector; and
• all transient release callbacks return to the runtime's retryable orphan owner.

configure_page_table_layout(...) is the one intentional late configuration step. Before allocation or execution, it installs a bounded immutable config replacement and forwards that exact replacement to config-dependent collaborators. It does not reconstruct the runtime or its execution helpers.

Data flow

One public prefill call follows four phases:

host request
  -> plan and classify                  [runtime.py, plan.py, signatures.py,
                                         postprocess.py]
  -> execute each prepared request      [sequence_runner.py or trace.py,
                                         inputs.py, postprocess.py]
  -> collect each invocation result     [result_collector.py]
  -> public logits or sampled tokens

PrefillRuntime is the package facade. Callers use its public methods rather than invoking the planner, input stager, sequence runner, postprocessor, or result collector directly.

Phase	Primary participants
Plan and classify	PrefillRuntime.prepare() coordinates _plan_prefill_requests(), PrefillPostprocessor.classify_sampling_path(), and the pure signature builders.
Stage inputs	PrefillInputStager prepares host values, allocates eager device inputs, or refreshes persistent trace inputs.
Execute eager	PrefillSequenceRunner stages and runs each PrefillChunk, releasing non-final outputs as it advances.
Execute traced	PrefillTraceLifecycle defines capture, refreshes each replay step, and records replay ownership.
Postprocess	PrefillPostprocessor turns the final hidden output into logits or device-sampled result values.
Collect	PrefillResultAssembler synchronizes readback, restores source rows, aggregates outputs, and releases invocation ownership.

1. Plan and classify once

PrefillRuntime.prepare(...) validates the request and produces one or more immutable PreparedPrefill values:

tokens + page table + lengths + cached positions + slots
  -> PrefillRequest
       -> one or more PrefillChunk values
  -> sampling path
  -> eager program signatures
  -> optional trace signature
  -> PreparedPrefill

The prepared value is the shared input to eager invocation, compilation, trace capture, and traced dispatch. Those paths do not re-plan the request or reclassify sampling.

The central immutable values are:

Value	Meaning
PrefillRequest	Logical rows, physical geometry, page tables, cached offsets, and ordered chunks for one invocation sequence
PrefillChunk	One model-body step, including its token slice, absolute position, page-table view, and final-token marker
PreparedPrefill	One request plus its sampling classification and program/trace identities
InvocationResult	A result value together with the transient resources that remain owned until collection

2. Execute eager or traced prefill

Eager execution runs the complete prepared sequence through PrefillSequenceRunner:

PreparedPrefill
  -> stage one chunk
  -> run the regular or chunked model body
  -> release a non-final chunk output
  -> repeat until the final-token chunk
  -> finish prefill postprocessing
  -> InvocationResult

A regular request is the one-step case. A cached or long request can contain multiple chunks, but it retains one ordered request-level sequence and one final result.

Traced execution captures a sampling-free hidden body and keeps program-alias-local postprocessing state separate:

PreparedPrefill
  -> capture plan
       -> persistent model inputs
       -> hidden-body trace
       -> alias-local replay workspace
  -> refresh inputs for each replay step
  -> replay each planned chunk
  -> finish only the final hidden output
  -> InvocationResult

Cached offsets and chunk starts are refreshed inputs. A fixed chunk trace can therefore serve multiple positions through a host-side replay loop without changing trace identity.

3. Postprocess device outputs

PrefillPostprocessor owns the boundary between model-body output and the result that will be read on the host. Its responsibilities include:

• selecting logits, forced-argmax, or top-k behavior;
• creating and refreshing K/P/T device values;
• selecting the logical last-token row or tile;
• allocating sampling and log-probability outputs; and
• preserving which tensors are borrowed, persistent, or invocation-owned.

Postprocessing does not restore public source-row order and does not perform the final synchronized host read. Those belong to result collection.

4. Collect results in source-row order

PrefillResultAssembler.assemble(...) consumes prepared invocation results as a stream. For each result it:

1. completes synchronized device-to-host readback;
2. normalizes logits or sampled tokens;
3. restores the original source-row mapping;
4. records per-row log probabilities when requested; and
5. releases that invocation's owned resources before advancing the stream.

Streaming is a correctness requirement for trace replay. Multiple replays can reuse one persistent output buffer, so a result must be read before the caller advances the generator and overwrites it.

Planning invariants

Logical rows and physical rows are distinct

source_rows preserves caller-visible order. active_batch_size records real rows, while padded_batch_size records the physical wave. Compatible active counts round up to a supported wave in {1, 2, 4, 8, 16, 32} without changing the logical result shape.

Program and trace identities use padded geometry rather than active row count. Requests with the same physical geometry can therefore share compiled and captured work.

Padding rows must not write KV

Regular paged-prefill tables use -1 as the skip sentinel for unallocated tails and padding rows. Token padding rows are zero. This makes unused physical rows inert while retaining one shared padded program identity.

Chunked full-request page tables retain the nonnegative filler required by the attention path. Fill-only chunk tables use -1 beyond their mapped prefix. These table roles are intentionally different.

Cached and chunk positions stay absolute

Planning preserves the cached-token offset and each chunk's absolute start. Last-token extraction is derived relative to the current invocation only at the staging or postprocessing boundary. Cached prefill is not a separate service or execution mode.

Planning fallback is not execution fallback

The planner can represent an ineligible batch as sequential single requests. Once a prepared request selects traced execution, missing program-to-trace coverage is an error; the traced path does not silently reinvoke eager execution.

Identity boundaries

Eager programs and traces have related but different identities:

• PrefillProgramSignature includes every material program choice, including operation variant, padded invocation geometry, page-table widths, sampling path, and any material last-token tile.
• PrefillTraceSignature identifies the static hidden body and persistent input geometry. Sampling state is deliberately excluded.
• capture-schema and workspace fingerprints allow one hidden trace family to keep distinct postprocessing workspaces for its program aliases.

Identity builders are pure. They consume an already-planned request and do not inspect mutable runtime state.

Ownership and failures

The runtime borrows the model, mesh device, output reader, and bound paged-KV context. Eager staging and postprocessing create invocation transients. Successful invocation transfers those resources in InvocationResult.owned to the result collector.

The sequence runner releases non-final chunk outputs immediately. On a primary failure it attempts to release everything acquired so far and attaches cleanup failures without replacing the primary error. Failed releases are retained by PrefillRuntime as retryable tensor-resource orphans and are retried during runtime cleanup.

Trace-owned persistent inputs and hidden outputs are not treated as eager invocation transients. PrefillReplayOwnership records the boundary between trace-owned, nested persistent, newly allocated, and replay-local values.

Module map

Module	Main responsibility
config.py	Resolve and validate static collaborators, capabilities, and geometry ceilings
plan.py	Build immutable logical requests, padded waves, page tables, and chunks
signatures.py	Classify prepared requests and construct pure eager/trace identities
inputs.py	Prepare host values, stage device inputs, and refresh replay inputs in place
sequence_runner.py	Run eager chunk sequences with failure-safe ownership transfer
trace.py	Define capture plans, persistent inputs, replay workspaces, refresh, and replay ownership
postprocess.py	Classify sampling and transform final hidden outputs into readable result values
result_collector.py	Read results synchronously, restore source rows, aggregate outputs, and release ownership
sampling_helpers.py	Provide stateless sampling normalization, slicing, and log-probability helpers
runtime.py	Expose the stable facade, compose collaborators, call model bodies, and retain cleanup orphans

The intended dependency direction is from the facade toward focused collaborators and from execution mechanics toward immutable plan/signature values. Focused modules must not import the model executor, execution-target selection, compiler registries, warmup coordination, or serving adapters.

Tests

The main contract suite is models/common/tests/llm_runtime/test_prefill_runtime.py. It covers planning, padding, page-table semantics, signatures, eager sequences, trace refresh, sampling, result collection, and cleanup failures.

models/common/tests/llm_runtime/test_execution.py covers the boundary between prepared prefill values and eager/traced execution, including strict trace coverage and streaming replay consumption. Warmup-specific prefill coverage is tested in models/common/tests/llm_runtime/test_warmup.py.