Add files using upload-large-folder tool

README.md

@@ -0,0 +1,113 @@
+---
+tags:
+- blackhole
+- p300x2
+- tt-model-cache
+- tt-model-container
+- vllm-plugin
+license: apache-2.0
+pipeline_tag: automatic-speech-recognition
+base_model:
+- mistralai/Voxtral-Mini-4B-Realtime-2602
+---
+
+# voxtral-mini-4b-realtime-2602-p300x2
+
+Voxtral-Mini-4B-Realtime-2602, Mistral AI's realtime speech-to-text model (a 970M-parameter causal audio encoder feeding a 3.4B Ministral text decoder; 13 languages; one text token per 80 ms of audio), served on four Blackhole chips (2 x p300c, 1x4 ring, TP=4) through a TTNN autoport and the Tenstorrent vLLM plugin as an OpenAI-compatible `/v1/audio/transcriptions` endpoint.
+
+Runs on **p300x2** (mesh `P300x2`) — 131,072-token context, up to 32 concurrent sequences.
+
+Packaged and published with [tt-model-manager](https://github.com/tenstorrent/tt-model-manager) 0.1.0 (manifest schema 5.1).
+
+## At a glance
+
+| | |
+| --- | --- |
+| Architecture | 4B realtime ASR (970M causal audio encoder + 3.4B Ministral-3B decoder) |
+| Hardware | p300x2 |
+| Context | 131,072 tokens |
+| License | apache-2.0 |
+| Status | Experimental community bring-up (tt-model-bringup autoport, 11 stages, locally committed) |
+
+## Intended use
+
+**Direct use:** Speech-to-text transcription of mono speech audio (wav, flac, mp3 and other formats soundfile decodes; resampled to 16 kHz) through the OpenAI audio transcriptions API, one interactive user or up to 32 concurrent requests, in the 13 languages the checkpoint covers.
+
+**Out-of-scope use:** Chat, text generation, translation, speaker diarization, or audio question answering: the checkpoint is transcription-only and the server exposes only that task.
+
+## Quickstart
+
+```bash
+uv tool install tenstorrent   # once — the Tenstorrent CLI, `tt`
+tt model pull ndaly/Voxtral-Mini-4B-Realtime-2602-tt-p300x2
+tt serve ndaly/Voxtral-Mini-4B-Realtime-2602-tt-p300x2
+```
+
+`tt model pull` (or `tt-model pull --with-weights`) downloads the Docker image and the [`mistralai/Voxtral-Mini-4B-Realtime-2602`](https://huggingface.co/mistralai/Voxtral-Mini-4B-Realtime-2602) weights at `2769294da9567371363522aac9bbcfdd19447add` (into your HF cache; they are not in the image). `tt serve` (or `tt-model serve`) starts an OpenAI-compatible server on port 20000 (or the next free port, if that one is busy); the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs `Application startup complete`.
+
+Without tt-cli — tt-model alone does the whole job:
+
+```bash
+tt-model pull  ndaly/Voxtral-Mini-4B-Realtime-2602-tt-p300x2 --with-weights
+tt-model serve ndaly/Voxtral-Mini-4B-Realtime-2602-tt-p300x2
+```
+
+### What to expect
+- Hardware: a TT-QuietBox 2 or any host with 2 x p300c (4 Blackhole chips, 32 GB GDDR6 each), Docker, 1G hugepages.
+- Weights: one 8.9 GB safetensors file (plus tokenizer and config files) is downloaded into your Hugging Face cache on first `serve`.
+- First boot compiles kernels and captures the decode trace plus 32 per-slot prefill traces (several minutes); later boots reuse the
+  kernel cache under `~/.cache/tt-model/voxtral-mini-4b-realtime-2602-p300x2/`.
+- Measured (stage 10 serving run, 2026-09-23): 37.5 ms served TTFT and ~50 tokens/s per user for one user; 606 tokens/s aggregate
+  for 32 concurrent requests. Full evidence: `doc/optimized_vllm/README.md` and `doc/tti_release/RUN_NOTES.md` in the tt-metal autoport tree.
+### Use with your client
+Point an OpenAI-compatible client at `http://127.0.0.1:20000` (or the port `tt-model serve` printed) and call the audio
+transcriptions API with model id `mistralai/Voxtral-Mini-4B-Realtime-2602` (see Using it below).
+
+## Using it
+
+The server speaks the OpenAI API at `http://127.0.0.1:20000/v1` (or whichever port your serve command reported; chat completions, completions, and `/v1/models`). Pass `"model": "mistralai/Voxtral-Mini-4B-Realtime-2602"` — the weights id, not this package's name.
+
+The endpoint is `POST /v1/audio/transcriptions` (multipart form, OpenAI shape). With the server on the port `tt-model serve` printed
+(20000 by default):
+```bash
+curl -s http://127.0.0.1:20000/v1/audio/transcriptions \
+  -F model=mistralai/Voxtral-Mini-4B-Realtime-2602 -F file=@clip.wav
+# streaming (server-sent events, text deltas as they are decoded):
+curl -sN http://127.0.0.1:20000/v1/audio/transcriptions \
+  -F model=mistralai/Voxtral-Mini-4B-Realtime-2602 -F file=@clip.wav -F stream=true
+```
+Any OpenAI client works the same way (`client.audio.transcriptions.create(model=..., file=...)`). `tt-model curl` sends a chat
+completion and gets a 404 from this model: use the request above instead.
+
+## Expected performance
+
+Measured on p300x2 through the served endpoint (stage 10 of the bring-up, 2026-09-23): single user, one 5.28 s clip (78 completion tokens, greedy, streaming): 37.5 ms served time-to-first-token (the adapter's prefill wall), 20.14 ms per decode step (49.7 tokens/s per user), 1.59 s end to end; 32 concurrent requests: 606 tokens/s aggregate, 3.63 s median end to end, 32/32 complete. The generator's own traced decode is 19.55 ms per token. Accuracy: 5.64 % WER on LibriSpeech test-other (2,939 utterances, greedy, tt-inference-server release run through lmms-eval); transcripts identical to the Hugging Face reference on 8 of 10 qualitative clips (the other 2 differ by near-tie words); top-1 token agreement 0.995 against the bf16 reference at the selected precision (bfp4 + LoFi inner projections, bfp8 LM head and KV cache, 6.8 GB per chip).
+
+## Limitations
+
+Transcription only: the checkpoint supports no chat or text completion, so the server mounts `/v1/audio/transcriptions` (and `/v1/models`) but not `/v1/chat/completions`; `tt-model curl` cannot talk to it (see Using it). One TP=4 replica on exactly four chips (p300x2); no data parallelism, no other board validated. At most 32 concurrent sequences; 131,072-token context contract (the checkpoint's advertised context); `block_size` is fixed at 32 by the generator. Requires `--tokenizer-mode mistral` (tekken.json). Chunked prefill and prefix caching are disabled by the TT backend. Validated on utterances up to about 35 s (LibriSpeech) and a 30 s benchmark payload; longer audio runs through the same streaming path but was not measured. The vLLM stack is the tenstorrent/vllm fork's empty-target wheel plus its in-tree plugin with local commits, not the standalone vllm-tt-plugin.
+
+## Risks and safety considerations
+
+The bfp4 + LoFi precision policy drops one word onset on 1 of 2,939 LibriSpeech test-other utterances where the bf16 reference does not. Under concurrency, 1 of 2,939 utterances came back empty from a rotation-dependent near-tie onset (the Hugging Face reference also returns empty on the evaluator's exact encoding). The release accuracy gate compared against a placeholder reference because no published LibriSpeech test-other WER exists for this checkpoint. Transcripts are greedy by default; sampled decoding is supported but not evaluated for accuracy.
+
+## Licensing
+
+Weights under Mistral AI's Apache-2.0 licence; the TTNN port and serving code in `code/` are Apache-2.0 (tt-metal), as are the vLLM fork and its plugin inside the image.
+
+## Feedback
+
+Questions or problems with this package: open a discussion at https://huggingface.co/ndaly/Voxtral-Mini-4B-Realtime-2602-tt-p300x2/discussions — that is what reaches its author. A problem with the `tt` tooling itself: `tt report issue` (collects your environment and opens a prefilled issue against tenstorrent/tt-cli). Product feedback: support@tenstorrent.com.
+
+## Provenance
+
+The exact sources the image was built from — `code/` in this repo is byte-identical to the model code inside the image:
+
+| component | built from |
+| --- | --- |
+| tt-metal | a local checkout — commit not published *(dirty tree — the image includes uncommitted changes)* |
+| vLLM | `vllm-0.1.dev14195+g8c28fcecb.d20260916.empty-py3-none-any.whl` — a wheel the author built |
+| vllm-tt-plugin | a local checkout — commit not published *(dirty tree — the image includes uncommitted changes)* |
+| `code/` digest | `b26d652115a867ab` (sha256, first 16 hex digits) |
+| built | 2026-09-25T16:20:42+00:00 by tt-model 0.1.0 |
+

code/models/common/modules/moe/configs/deepseek_ocr.yaml

@@ -0,0 +1,28 @@
+# deepseek_ocr
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 1280
+select_experts_k: 6
+num_routed_experts: 64
+num_shared_experts: 2
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# deepseek_ocr = ungrouped softmax top-6, no bias. 64 experts → padded to the op's 256-face in
+# TTMoEGate.forward.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+gate_notes: "DeepSeek-OCR: softmax kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top-6→renorm (global Z cancels, no bias). Set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 896
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/deepseek_v3.yaml

@@ -0,0 +1,28 @@
+# deepseek_v3
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 7168
+select_experts_k: 8
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8)
+n_group: 8
+score_func: sigmoid
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only
+# (output weights stay unbiased). Explicit per-model flag (not inferred from score_func).
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/deepseek_v3_single_glx.yaml

@@ -0,0 +1,28 @@
+# deepseek_v3
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 7168
+select_experts_k: 8
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8)
+n_group: 8
+score_func: sigmoid
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only
+# (output weights stay unbiased). Explicit per-model flag (not inferred from score_func).
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/deepseek_v4_flash.yaml

@@ -0,0 +1,32 @@
+# deepseek_v4_flash
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 4096
+select_experts_k: 6
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# deepseek-v4: ungrouped (1 group) top-6 with the sqrtsoftplus score = sqrt(softplus(logit)) (applied
+# externally via ttnn.sqrt(ttnn.softplus(.)), no in-kernel op) + score_correction_bias for selection,
+# linear renormalization, scaled by 1.5. 256 experts → single 256-face (no combine). Same gate as
+# deepseek_v4_pro but 256 experts / scale 1.5 (pro is 384 / 2.5).
+n_group: 1
+score_func: sqrtsoftplus
+routed_scaling_factor: 1.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sqrtsoftplus scores for SELECTION
+# only (output weights stay unbiased). Explicit per-model flag (not inferred from score_func).
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/deepseek_v4_pro.yaml

@@ -0,0 +1,32 @@
+# deepseek_v4_pro
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 7168
+select_experts_k: 6
+num_routed_experts: 384
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# deepseek-v4: ungrouped top-6 with the sqrtsoftplus score = sqrt(softplus(logit)) (applied externally
+# via ttnn.sqrt(ttnn.softplus(.)), no in-kernel op) + score_correction_bias for selection, linear
+# renormalization, scaled by 2.5. 384 experts → the op's 512-combine path (pad 384→512, two 256-blocks →
+# global top-6), same combine path as kimi_k25.
+n_group: 1
+score_func: sqrtsoftplus
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sqrtsoftplus scores for SELECTION
+# only (output weights stay unbiased). Explicit per-model flag (not inferred from score_func).
+score_correction_bias: true
+
+compute:
+  intermediate_size: 3072
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/gemma_4_26b.yaml

@@ -0,0 +1,35 @@
+# gemma_4_26b
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 2816
+select_experts_k: 8
+num_routed_experts: 128
+num_shared_experts: 0
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Gemma4 router = linear → softmax(all) → top-8 → renormalize. TTMoEGate covers exactly that
+# (score_func=softmax ≡ softmax-front + renorm, the global Z cancels; routed_scaling_factor=1.0 → no
+# scalar scale). 128 experts → padded to the op's 256-face.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+
+# REMINDER — Gemma4's gate has steps TTMoEGate does NOT cover (caller must apply, see modular_gemma4.py
+# Gemma4TextRouter.forward):
+#   step 1: RMSNorm(hidden) (with_scale=False) — before this gate
+#   step 2: hidden *= self.scale[hidden] * hidden_size**-0.5 — a learned per-dim input scale × 1/√hidden
+#   step 6b: top_k_weights *= per_expert_scale[selected] — a learned PER-EXPERT output scale ([num_experts]),
+#            NOT a scalar, so it can't be folded into routed_scaling_factor.
+gate_notes: "Gemma4: TTMoEGate does steps 3-6a (linear→softmax→top8→renorm). MISSING upstream: RMSNorm + per-dim input scale (×1/√hidden); MISSING downstream: per-expert output scale per_expert_scale[selected] (not a scalar). SOFTMAX kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top8→renorm (global Z cancels, no bias); set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 704
+  activation_type: GELU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential

code/models/common/modules/moe/configs/glm5.yaml

@@ -0,0 +1,30 @@
+# glm5
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 6144
+select_experts_k: 8
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# GLM-5 is deepseek-arch but n_group=1 / topk_group=1 → the grouped path is a NO-OP (select 1 of 1 group =
+# all experts), i.e. UNGROUPED top-8 with sigmoid(logit)+score_correction_bias for selection, linear
+# renormalization, scaled by 2.5. 256 experts → single 256-face (no combine).
+n_group: 1
+score_func: sigmoid
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only.
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/glm_47.yaml

@@ -0,0 +1,31 @@
+# glm_47
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 5120
+select_experts_k: 8
+num_routed_experts: 160
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# GLM-4.7 is deepseek-arch but n_group=1 / topk_group=1 → the grouped path is a NO-OP (select 1 of 1 group
+# = all experts), i.e. UNGROUPED top-8 with sigmoid(logit)+score_correction_bias for selection, linear
+# renormalization, scaled by 2.5. Same gate as GLM-5 but 160 experts → padded to the op's 256-face
+# (phantom experts get a large-negative logit+bias so they rank last in the sigmoid selection).
+n_group: 1
+score_func: sigmoid
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only.
+score_correction_bias: true
+
+compute:
+  intermediate_size: 1536
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/gpt_oss.yaml

@@ -0,0 +1,33 @@
+# gpt_oss
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 2880
+select_experts_k: 4
+num_routed_experts: 128
+num_shared_experts: 0
+has_bias: true
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# GPT-OSS router (GptOssTopKRouter): logits = Wx + b → top-4 (by Wx+b) → softmax over the selected (Wx+b).
+# TTMoEGate's score_func=softmax is exactly that (rank by logit, then exp-over-selected), scale 1.0, and the
+# router LINEAR bias b is supplied as the gate projection bias (gate_proj_bias below → torch_gate_proj_bias),
+# so logits = Wx + b and b flows into BOTH selection and the softmax weights (matches the reference).
+# 128 experts → padded to the op's 256-face.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+# router LINEAR bias: GPT-OSS's GptOssTopKRouter does F.linear(x, W, b). TTMoEGate adds b via ttnn.linear's
+# bias (torch_gate_proj_bias) — distinct from the deepseek score-correction bias (which is selection-only).
+# (has_bias=true above is the separate EXPERT-MLP bias.)
+gate_proj_bias: true
+
+compute:
+  intermediate_size: 2880
+  activation_type: SWIGLU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential

code/models/common/modules/moe/configs/kimi_k25.yaml

@@ -0,0 +1,30 @@
+# kimi_k25
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 7168
+select_experts_k: 8
+num_routed_experts: 384
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Kimi-K2 is deepseek-arch but config.json has n_group=1 / topk_group=1 → the grouped path is a NO-OP
+# (1 group, select 1 of 1 = all experts), i.e. UNGROUPED top-8 with sigmoid(logit)+score_correction_bias
+# for selection, linear renormalization (norm_topk_prob), scaled by 2.827. 384 experts → 2-block combine.
+n_group: 1
+score_func: sigmoid
+routed_scaling_factor: 2.827
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only.
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/ling_1t.yaml

@@ -0,0 +1,29 @@
+# ling_1t
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 8192
+select_experts_k: 8
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# ling_1t is deepseek-style grouped routing: sigmoid(logit)+score_correction_bias for group selection,
+# linear renormalization of the selected experts, scaled by routed_scaling_factor.
+n_group: 8
+score_func: sigmoid
+routed_scaling_factor: 2.5
+# noaux_tc score-correction bias (e_score_correction_bias): added to the sigmoid scores for SELECTION only.
+score_correction_bias: true
+
+compute:
+  intermediate_size: 2048
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/models_table.md

@@ -0,0 +1,18 @@
+| **Model** | **hidden_size** | **moe_intermediate** | **shared_expert_interm** | **n_routed_experts** | **n_shared_experts** | **K (top-k)** | **activation** | **scoring_func** | **topk_method** | **n_group/topk_group** | **scaling_factor** | **expert_bias** | **router_bias** | **first_k_dense** | **num_layers** | **parallel dense** | **base_arch** |
+|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
+| **GPT-OSS 120B** | 2880 | 2880 | — | 128 | 0 | 4 | custom GELU-gated | softmax | simple | —/— | 1.0 | Yes | Yes | all MoE | 36 | No | GPT-OSS |
+| **DeepSeek V3** | 7168 | 2048 | 2048 | 256 | 1 | 8 | SiLU/SwiGLU | `sigmoid` | noaux_tc | 8/4 | 2.5 | No | correction | 3 | 61 | No | DS V3 |
+| **DS V4 Flash** | 4096 | 2048 | 2048 | 256 | 1 | 6 | SiLU/SwiGLU (swiglu_limit=10, routed only) | `sqrtsoftplus` (code) | noaux_tc; hash L0–2 (code) | —/— | 1.5 | No | correction (code) | 0 | 43 | No | DS V4 |
+| **DS V4 Pro** | 7168 | 3072 | 3072 | 384 | 1 | 6 | SiLU/SwiGLU (swiglu_limit=10, routed only) | `sqrtsoftplus` (code) | noaux_tc; hash L0–2 (code) | —/— | 2.5 | No | correction (code) | 0 | 61 | No | DS V4 |
+| **GLM-5** | 6144 | 2048 | 2048 | 256 | 1 | 8 | SiLU/SwiGLU | `sigmoid` | noaux_tc | 1/1 | 2.5 | No | correction | 3 | 78 | No | DS V3-like |
+| **Kimi K2.5** | 7168 | 2048 | 2048 | 384 | 1 | 8 | SiLU/SwiGLU | `sigmoid` | noaux_tc | 1/1 | 2.827 | No | correction | 1 | 61 | No | DS V3-like |
+| **Ling 1T** | 8192 | 2048 | 2048 | 256 | 1 | 8 | SiLU/SwiGLU | `sigmoid` | group_limited_topk (code) | 8/4 | 2.5 | No | correction (bias-enabled) | 4 | 80 | No | DS V3-like |
+| **GLM-4.7** | 5120 | 1536 | 1536 | 160 | 1 | 8 | SiLU/SwiGLU | `sigmoid` (code) | noaux_tc-style (code) | 1/1 | 2.5 | No | correction (code) | 3 | 92 | No | DS V3-like |
+| **Qwen3 235B** | 4096 | 1536 | — | 128 | 0 (no field) | 8 | SiLU/SwiGLU | softmax (code) | simple top-k (code) | —/— | — | No | No | all MoE | 94 | No | Qwen3 |
+| **Qwen3.5 397B** | 4096 | 1024 | 1024 | 512 | 1 (inferred) | 10 | SiLU/SwiGLU | softmax (code) | simple top-k (code) | —/— | — | No | No | all MoE | 60 | No | Qwen3.5 |
+| **Qwen3.5 35B** | 2048 | 512 | 512 | 256 | 1 (inferred) | 8 | SiLU/SwiGLU | softmax (code) | simple top-k (code) | —/— | — | No | No | all MoE | 40 | No | Qwen3.5 |
+| **Qwen3-Omni Thinker** | 2048 | 768 | — | 128 | 0 (inferred, size=0) | 8 | SiLU/SwiGLU | softmax (code) | simple top-k (code) | —/— | — | No | No | all MoE | 48 | No | Qwen3-Omni |
+| **Qwen3-Omni Talker** | 1024 | 384 | 768 | 128 | 1 (inferred, size=768) | 6 | SiLU/SwiGLU | softmax (code) | simple top-k (code) | —/— | — | No | No | all MoE | 20 | No | Qwen3-Omni |
+| **Mistral Large 3** | 7168 | 4096 | 4096 (from vLLM) | 128 | 1 | 4 | SiLU/SwiGLU (code) | softmax (code) | simple top-k (code) | 1/1 | 1.0 | No | No | 3 | 61 | No | Mistral |
+| **Gemma 4 26B** | 2816 | 704 | — (parallel dense) | 128 | 0 (parallel dense) | 8 | **GELU/SwiGLU** | softmax | simple+per_expert_scale | —/— | per-expert learned | No | No (has learned scale) | all MoE | 30 | **Yes** | Gemma4 |
+| **DS-OCR** | 1280 | 896 | 1792 | 64 | 2 | 6 | SiLU/SwiGLU | softmax (V2 default) | `greedy` | 1/1 | 1.0 (V2 default) | No | No | 1 | 12 | No | DS V2 |

code/models/common/modules/moe/configs/qwen35_35b.yaml

@@ -0,0 +1,26 @@
+# qwen35_35b
+mesh_shape: [16, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 2048
+select_experts_k: 8
+num_routed_experts: 256
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8)
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+gate_notes: "Qwen3.5 35B: softmax kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top-8→renorm (global Z cancels, no bias). Set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 512
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/qwen35_397b.yaml

@@ -0,0 +1,31 @@
+# qwen35_397b
+mesh_shape: [16, 8]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 4096
+select_experts_k: 10
+num_routed_experts: 512
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Qwen3.5 397B = linear → softmax(all) → top-10 → renormalize, no bias, no scale. UNSUPPORTED by the kernel
+# op (top-10 ∉ {4,6,8}, and 512 experts), so TTMoEGate takes the pure-ttnn FALLBACK (matmul → softmax(all)
+# → topk → linear renorm). softmax_position: pre puts the softmax over ALL experts UP FRONT — matching this
+# model's definition directly and giving a numerically-stable exp (vs exp-over-selected raw logits); the
+# linear renorm over the selected makes it mathematically ≡ softmax-over-selected, scale 1.0.
+n_group: 1
+score_func: softmax
+softmax_position: pre
+routed_scaling_factor: 1.0
+
+compute:
+  intermediate_size: 1024
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/qwen3_235b.yaml

@@ -0,0 +1,27 @@
+# qwen3_235b
+mesh_shape: [8, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 4096
+select_experts_k: 8
+num_routed_experts: 128
+num_shared_experts: 0
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Qwen3 MoE = softmax over all experts → top-k → renormalize (≡ softmax over selected, the global Z
+# cancels), no bias, no scaling. 128 experts → padded to the op's 256-face in TTMoEGate.forward.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+gate_notes: "Qwen3 235B: softmax kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top-8→renorm (global Z cancels, no bias). Set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 1536
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential

code/models/common/modules/moe/configs/qwen3_omni_talker.yaml

@@ -0,0 +1,29 @@
+# qwen3_omni_talker
+mesh_shape: [16, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 1024
+select_experts_k: 6
+num_routed_experts: 128
+num_shared_experts: 1
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Same as qwen3_omni_thinker but top-6: Qwen3 MoE = softmax over all → top-6 → renormalize (≡ softmax
+# over selected, Z cancels), no bias, no scaling. 128 experts → padded to the op's 256-face in
+# TTMoEGate.forward.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+gate_notes: "Qwen3 Omni Talker: softmax kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top-6→renorm (global Z cancels, no bias). Set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 384
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential
+  shared_expert_ids_to_devices: fully_replicated

code/models/common/modules/moe/configs/qwen3_omni_thinker.yaml

@@ -0,0 +1,27 @@
+# qwen3_omni_thinker
+mesh_shape: [16, 4]
+cluster_axis: 0
+batch_per_device: 32
+hidden_size: 2048
+select_experts_k: 8
+num_routed_experts: 128
+num_shared_experts: 0
+has_bias: false
+
+# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8).
+# Qwen3 MoE = softmax over all experts → top-8 → renormalize (≡ softmax over selected, the global Z
+# cancels), no bias, no scaling. 128 experts → padded to the op's 256-face in TTMoEGate.forward.
+n_group: 1
+score_func: softmax
+routed_scaling_factor: 1.0
+gate_notes: "Qwen3 Omni Thinker: softmax kept POST-topk (softmax-over-selected, the default softmax_position) — mathematically ≡ the model's softmax-over-all→top-8→renorm (global Z cancels, no bias). Set softmax_position: pre for the literal front softmax (numerically-stable exp over all experts)."
+
+compute:
+  intermediate_size: 768
+  activation_type: SILU
+
+reduce:
+  shared_expert_scale: 1.0
+
+experts:
+  expert_mapping: sequential

code/models/common/tests/demos/cleanup_utils.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+
+import ttnn
+from models.common.modules.lazy_weight import LazyWeight
+
+
+def cleanup_ttnn_value(value):
+    if value is None:
+        return
+
+    if isinstance(value, ttnn.Tensor):
+        ttnn.deallocate(value)
+        return
+
+    if isinstance(value, dict):
+        for nested_value in value.values():
+            cleanup_ttnn_value(nested_value)
+        return
+
+    if isinstance(value, (list, tuple, set)):
+        for nested_value in value:
+            cleanup_ttnn_value(nested_value)
+
+
+def cleanup_object_graph(obj, seen=None):
+    if obj is None:
+        return
+    if seen is None:
+        seen = set()
+
+    obj_id = id(obj)
+    if obj_id in seen:
+        return
+    seen.add(obj_id)
+
+    if isinstance(obj, ttnn.Tensor):
+        cleanup_ttnn_value(obj)
+        return
+
+    if isinstance(obj, LazyWeight):
+        if obj._value is not None:
+            cleanup_ttnn_value(obj._value)
+            obj._value = None
+        return
+
+    if isinstance(obj, dict):
+        for value in obj.values():
+            cleanup_object_graph(value, seen)
+        return
+
+    if isinstance(obj, (list, tuple, set)):
+        for value in obj:
+            cleanup_object_graph(value, seen)
+        return
+
+    state = getattr(obj, "__dict__", None)
+    if state is None:
+        return
+
+    for name, value in list(state.items()):
+        cleanup_object_graph(value, seen)
+        if isinstance(value, ttnn.Tensor):
+            setattr(obj, name, None)
+
+    if hasattr(obj, "_device_weights_loaded"):
+        obj._device_weights_loaded = False
+
+
+def cleanup_model_case(model, mesh_device):
+    ttnn.synchronize_device(mesh_device)
+    if model is not None:
+        cleanup_object_graph(model)
+    ttnn.synchronize_device(mesh_device)
+    gc.collect()
+
+
+def cleanup_dp_model_case(group, lanes, models, parent_mesh, submeshes):
+    """Release one carved DP case and restore command-queue ownership to its parent."""
+
+    failures = []
+
+    def run(action):
+        try:
+            action()
+        except Exception as error:
+            failures.append(error)
+
+    if group is not None:
+        run(group.cleanup)
+    else:
+        for lane in lanes:
+            run(lane.cleanup)
+
+    for model, submesh in models:
+        run(lambda model=model, submesh=submesh: cleanup_model_case(model, submesh))
+
+    # A carved child owns a distinct MeshCommandQueue over its parent's physical
+    # devices. Drain every live child through the parent before closing the child
+    # handles and returning command-queue ownership to the fixture-owned parent.
+    run(parent_mesh.quiesce_devices)
+    for submesh in submeshes:
+        run(lambda submesh=submesh: ttnn.close_mesh_device(submesh))
+
+    if failures:
+        primary = failures[0]
+        add_note = getattr(primary, "add_note", None)
+        if add_note is not None:
+            for failure in failures[1:]:
+                add_note(f"additional DP cleanup failure: {type(failure).__name__}: {failure}")
+        raise primary

code/models/common/tests/demos/run_helpers.py

@@ -0,0 +1,1071 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Reusable teacher-forcing and benchmark run helpers for demos and tests."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+from loguru import logger
+
+import ttnn
+
+_SAME_SAMPLING_PARAMS = object()
+
+
+def make_contiguous_page_table(batch_size: int, max_seq_len: int, block_size: int = 32) -> torch.Tensor:
+    """Create a contiguous demo page table with one disjoint block range per user."""
+    if min(batch_size, max_seq_len, block_size) <= 0:
+        raise ValueError("page-table dimensions must be positive")
+    blocks_per_user = (max_seq_len + block_size - 1) // block_size
+    return torch.arange(batch_size * blocks_per_user, dtype=torch.int32).reshape(batch_size, blocks_per_user)
+
+
+@dataclass
+class TeacherForceResult:
+    """Result from a teacher-forcing evaluation run."""
+
+    predicted_tokens: list[int]
+    predicted_tokens_per_user: list[list[int]]
+    reference_top5: torch.Tensor
+    prefill_time_s: float = 0.0
+    compile_decode_time_s: float = 0.0
+    decode_times_s: list[float] = field(default_factory=list)
+    batch_size: int = 1
+    prefill_len: int = 0
+
+    def top1_accuracy(self) -> float:
+        matches = sum(
+            1 for i, prediction in enumerate(self.predicted_tokens) if self.reference_top5[i, 0].item() == prediction
+        )
+        return matches / len(self.predicted_tokens)
+
+    def top5_accuracy(self) -> float:
+        matches = sum(
+            1 for i, prediction in enumerate(self.predicted_tokens) if prediction in self.reference_top5[i, :]
+        )
+        return matches / len(self.predicted_tokens)
+
+    @property
+    def ttft_ms(self) -> float:
+        return self.prefill_time_s / self.batch_size * 1000 if self.batch_size else 0.0
+
+    @property
+    def prefill_time_to_token_s(self) -> float:
+        return self.prefill_time_s / self.batch_size if self.batch_size else 0.0
+
+    @property
+    def prefill_tok_s(self) -> float:
+        return (self.batch_size * self.prefill_len) / self.prefill_time_s if self.prefill_time_s > 0 else 0.0
+
+    @property
+    def decode_tok_s_u(self) -> float:
+        return len(self.decode_times_s) / sum(self.decode_times_s) if self.decode_times_s else 0.0
+
+    @property
+    def decode_tok_s(self) -> float:
+        return self.decode_tok_s_u * self.batch_size
+
+
+@dataclass
+class PerfBenchmarkResult:
+    """Result from a performance benchmark run."""
+
+    prefill_time_s: float
+    compile_decode_time_s: float
+    decode_times_s: list[float]
+    batch_size: int
+    num_decode_tokens: int
+    generated_token_ids: list[list[int]]
+    decode_iteration_times_s: list[float] = field(default_factory=list)
+    argmax_top2_margins: list[list[float]] | None = None
+
+    @property
+    def ttft_ms(self) -> float:
+        """TTTv1-style average time to first token per user."""
+        return self.prefill_time_s / self.batch_size * 1000
+
+    @property
+    def tok_s_u(self) -> float:
+        """Tokens per second per user during steady-state decode."""
+        if not self.decode_times_s:
+            return 0.0
+        return len(self.decode_times_s) / sum(self.decode_times_s)
+
+    @property
+    def tok_s(self) -> float:
+        """Total decode throughput."""
+        return self.tok_s_u * self.batch_size
+
+    @property
+    def decode_latency_mean_ms(self) -> float:
+        if not self.decode_times_s:
+            return 0.0
+        return sum(self.decode_times_s) / len(self.decode_times_s) * 1000
+
+    def meets_target(self, expected: dict, tolerance: float = 0.05) -> dict[str, bool]:
+        """Check benchmark metrics against the expected thresholds."""
+        return {
+            "tok_s_u": self.tok_s_u >= expected["tok_s_u"] * (1 - tolerance),
+            "ttft_ms": self.ttft_ms <= expected["ttft_ms"] * (1 + tolerance),
+        }
+
+
+def _compile_prefill_and_decode(
+    execution_target,
+    *,
+    prefill_tokens: torch.Tensor,
+    prefill_page_table: torch.Tensor,
+    kv_cache=None,
+    prompt_lens: torch.Tensor | None = None,
+    empty_slots: list[int] | None = None,
+    start_pos: torch.Tensor | None = None,
+    sampling_params=None,
+    prefill_sampling_params=_SAME_SAMPLING_PARAMS,
+    decode_tokens: torch.Tensor | None = None,
+    decode_start_pos: torch.Tensor | None = None,
+    decode_page_table: torch.Tensor | None = None,
+) -> None:
+    """Compile the concrete prefill and decode cases through the public target surface."""
+    assert prefill_tokens.dim() == 2, f"prefill_tokens must be [batch_size, seq_len], got {prefill_tokens.dim()}D"
+    assert (
+        prefill_page_table.dim() == 2
+    ), f"prefill_page_table must be [batch_size, max_blocks], got {prefill_page_table.dim()}D"
+
+    batch_size = prefill_tokens.shape[0]
+    if decode_tokens is None:
+        decode_tokens = torch.zeros(batch_size, dtype=torch.long, device=prefill_tokens.device)
+    if decode_start_pos is None:
+        decode_start_pos = torch.full(
+            (decode_tokens.shape[0],),
+            prefill_tokens.shape[-1],
+            dtype=torch.long,
+            device=prefill_tokens.device,
+        )
+    if decode_page_table is None:
+        decode_page_table = prefill_page_table
+
+    if prefill_sampling_params is _SAME_SAMPLING_PARAMS:
+        prefill_sampling_params = sampling_params
+
+    if sampling_params is not None:
+        execution_target.compile_decode(
+            tokens=decode_tokens,
+            start_pos=decode_start_pos,
+            page_table=decode_page_table,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+        )
+        execution_target.compile_prefill(
+            tokens=prefill_tokens,
+            page_table=prefill_page_table,
+            kv_cache=kv_cache,
+            prompt_lens=prompt_lens,
+            empty_slots=empty_slots,
+            start_pos=start_pos,
+            sampling_params=prefill_sampling_params,
+        )
+        return
+
+    execution_target.compile_prefill(
+        tokens=prefill_tokens,
+        page_table=prefill_page_table,
+        kv_cache=kv_cache,
+        prompt_lens=prompt_lens,
+        empty_slots=empty_slots,
+        start_pos=start_pos,
+        sampling_params=None,
+    )
+    execution_target.compile_decode(
+        tokens=decode_tokens,
+        start_pos=decode_start_pos,
+        page_table=decode_page_table,
+        kv_cache=kv_cache,
+        sampling_params=None,
+    )
+
+
+def _profiler_start(profiler, name: str) -> None:
+    if profiler is not None:
+        profiler.start(name)
+
+
+def _profiler_end(profiler, name: str) -> None:
+    if profiler is not None:
+        profiler.end(name)
+
+
+def run_teacher_forcing(
+    executor,
+    *,
+    prompt_tokens: torch.Tensor,
+    reference_tokens: torch.Tensor,
+    top5_tokens: torch.Tensor,
+    kv_cache: list,
+    page_table: torch.Tensor,
+    max_batch_size: int = 1,
+    profiler=None,
+) -> TeacherForceResult:
+    """Run teacher-forcing accuracy measurement against an execution target."""
+    execution_target = executor
+    batch_size = prompt_tokens.shape[0]
+    assert (
+        batch_size == max_batch_size
+    ), f"Teacher forcing expects active batch to match max_batch_size, got {batch_size} vs {max_batch_size}"
+    prompt_len = prompt_tokens.shape[-1]
+    num_target = len(reference_tokens) - prompt_len
+    prompt_lens = torch.tensor([prompt_len] * batch_size)
+    empty_slots = list(range(batch_size))
+
+    _compile_prefill_and_decode(
+        execution_target,
+        prefill_tokens=prompt_tokens,
+        prefill_page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=prompt_lens,
+        empty_slots=empty_slots,
+    )
+
+    logger.info(f"Teacher forcing: prefilling {prompt_len} tokens with batch={batch_size}")
+    _profiler_start(profiler, "inference_prefill")
+    try:
+        start_time = time.perf_counter()
+        prefill_output = execution_target.prefill_forward(
+            prompt_tokens,
+            page_table=page_table,
+            kv_cache=kv_cache,
+            prompt_lens=prompt_lens,
+            empty_slots=empty_slots,
+        )
+        _synchronize_target(execution_target)
+        prefill_time_s = time.perf_counter() - start_time
+    finally:
+        _profiler_end(profiler, "inference_prefill")
+    first_tokens = torch.argmax(prefill_output, dim=-1).view(-1).tolist()
+    predicted_tokens_per_user = [[int(token)] for token in first_tokens]
+
+    logger.info(f"Teacher forcing: decoding {num_target - 1} tokens")
+    compile_decode_time_s = 0.0
+    decode_times_s = []
+    _profiler_start(profiler, "inference_decode")
+    try:
+        for step in range(1, num_target):
+            ground_truth_token = reference_tokens[prompt_len + step - 1]
+            decode_token = torch.full((batch_size,), ground_truth_token, dtype=torch.long)
+            current_pos = torch.full((batch_size,), prompt_len + step - 1, dtype=torch.long)
+            start_time = time.perf_counter()
+            logits, _ = execution_target.decode_forward(
+                decode_token,
+                current_pos,
+                page_table=page_table,
+                kv_cache=kv_cache,
+                read_from_device=True,
+            )
+            elapsed = time.perf_counter() - start_time
+            if step == 1:
+                compile_decode_time_s = elapsed
+            else:
+                decode_times_s.append(elapsed)
+
+            next_tokens = torch.argmax(logits[:, -1, :], dim=-1).view(-1).tolist()
+            for user_id, token in enumerate(next_tokens):
+                predicted_tokens_per_user[user_id].append(int(token))
+    finally:
+        _profiler_end(profiler, "inference_decode")
+
+    return TeacherForceResult(
+        predicted_tokens=predicted_tokens_per_user[0],
+        predicted_tokens_per_user=predicted_tokens_per_user,
+        reference_top5=top5_tokens[:num_target],
+        prefill_time_s=prefill_time_s,
+        compile_decode_time_s=compile_decode_time_s,
+        decode_times_s=decode_times_s,
+        batch_size=batch_size,
+        prefill_len=prompt_len,
+    )
+
+
+def _split_output(output):
+    return output if isinstance(output, tuple) else (output, None)
+
+
+def _target_mesh_device(execution_target):
+    return getattr(execution_target, "mesh_device", None)
+
+
+def _target_cluster_shape(execution_target):
+    cluster_shape = getattr(execution_target, "cluster_shape", None)
+    if cluster_shape is not None:
+        return list(cluster_shape)
+    mesh_device = _target_mesh_device(execution_target)
+    return list(mesh_device.shape) if mesh_device is not None else [1, 1]
+
+
+def _synchronize_target(execution_target):
+    mesh_devices = getattr(execution_target, "mesh_devices", None)
+    if mesh_devices is not None:
+        for mesh_device in mesh_devices:
+            if mesh_device is not None:
+                ttnn.synchronize_device(mesh_device)
+        return
+    mesh_device = _target_mesh_device(execution_target)
+    if mesh_device is not None:
+        ttnn.synchronize_device(mesh_device)
+
+
+def _concat_host_output(output, cluster_shape):
+    output_tensors = [ttnn.to_torch(tensor) for tensor in ttnn.get_device_tensors(output)]
+    _, columns = cluster_shape
+    mesh_rows = [output_tensors[i : i + columns] for i in range(0, len(output_tensors), columns)]
+    return torch.cat([torch.cat(row, dim=-1) for row in mesh_rows], dim=1)
+
+
+def _process_legacy_sampled_tokens(output, batch_size, cluster_shape):
+    torch_output = _concat_host_output(output, cluster_shape)
+    if torch_output.ndim >= 4:
+        if torch_output.shape[2] >= batch_size:
+            return torch_output[0, 0, :batch_size, 0]
+        if torch_output.shape[3] >= batch_size:
+            return torch_output[0, 0, 0, :batch_size]
+    return torch_output.reshape(-1)[:batch_size]
+
+
+def _to_host(value, *, blocking):
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        return value.cpu()
+    try:
+        return value.cpu(blocking=blocking)
+    except TypeError:
+        return value.cpu()
+
+
+def _submit_decode_read(execution_target, decode_output):
+    read_decode_output = getattr(execution_target, "read_decode_output", None)
+    if callable(read_decode_output):
+        return read_decode_output(decode_output, async_read=True)
+
+    output, log_probs = _split_output(decode_output)
+    host_output = (_to_host(output, blocking=False), _to_host(log_probs, blocking=False))
+    return host_output, [ttnn.record_event(_target_mesh_device(execution_target), 0)]
+
+
+def _synchronize_read_events(events):
+    if events is None:
+        return
+    if not isinstance(events, (list, tuple, set)):
+        events = [events]
+    for event in events:
+        ttnn.event_synchronize(event)
+
+
+def _consume_sampled_output(
+    execution_target,
+    host_output,
+    batch_size,
+    cluster_shape,
+    generated_token_ids,
+    *,
+    process_host_output,
+):
+    process_decode_output_host = (
+        getattr(execution_target, "process_decode_output_host", None) if process_host_output else None
+    )
+    if callable(process_decode_output_host):
+        tokens, _ = process_decode_output_host(host_output, is_tokens=True)
+    else:
+        tokens, _ = _split_output(host_output)
+        if isinstance(tokens, torch.Tensor):
+            tokens = tokens.view(-1)[:batch_size].detach().cpu()
+        else:
+            tokens = _process_legacy_sampled_tokens(tokens, batch_size, cluster_shape)
+            tokens = tokens.view(-1)[:batch_size].detach().cpu()
+
+    for user_id, token in enumerate(tokens.tolist()):
+        generated_token_ids[user_id].append(int(token))
+
+
+def _host_argmax_with_margins(logits: torch.Tensor, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return host argmax and top1-minus-top2 margins for eval diagnostics."""
+
+    rows = logits[:, -1, :]
+    top2 = torch.topk(rows.float(), k=2, dim=-1)
+    tokens = top2.indices[:, 0].view(-1)[:batch_size].detach().cpu()
+    margins = (top2.values[:, 0] - top2.values[:, 1]).view(-1)[:batch_size].detach().cpu()
+    return tokens, margins
+
+
+def run_perf_benchmark(
+    executor,
+    *,
+    tokens: torch.Tensor,
+    kv_cache: list,
+    page_table: torch.Tensor,
+    num_decode_tokens: int = 128,
+    max_batch_size: int = 1,
+    prompt_lens: torch.Tensor | None = None,
+    start_pos: list[int] | None = None,
+    sampling_params=None,
+    prefill_sampling_params=_SAME_SAMPLING_PARAMS,
+    pipeline_readback: bool = False,
+    profiler=None,
+    collect_argmax_diagnostics: bool = False,
+) -> PerfBenchmarkResult:
+    """Run the timed prefill and decode loop against a public execution target."""
+    execution_target = executor
+    mesh_device = _target_mesh_device(execution_target)
+    has_public_readback = callable(getattr(execution_target, "read_decode_output", None))
+    has_legacy_readback = mesh_device is not None and hasattr(ttnn, "record_event")
+    can_pipeline_readback = (
+        sampling_params is not None
+        and pipeline_readback
+        and hasattr(ttnn, "event_synchronize")
+        and (has_public_readback or has_legacy_readback)
+    )
+    if sampling_params is not None and pipeline_readback and not can_pipeline_readback:
+        logger.warning("PIPELINE_READBACK requested, but this execution target does not expose async readback")
+
+    batch_size, prompt_len = tokens.shape
+    max_batch_size = max(max_batch_size, batch_size)
+    cluster_shape = _target_cluster_shape(execution_target)
+    prompt_lens = prompt_lens if prompt_lens is not None else torch.tensor([prompt_len] * batch_size)
+    if prefill_sampling_params is _SAME_SAMPLING_PARAMS:
+        prefill_sampling_params = sampling_params
+    prefill_kwargs = dict(
+        # Decode always carries the full lane-capacity table. Prefill only owns
+        # the active request rows, and its public contract requires matching
+        # token/page-table batch dimensions.
+        page_table=page_table[:batch_size],
+        kv_cache=kv_cache,
+        prompt_lens=prompt_lens,
+        empty_slots=list(range(batch_size)),
+        start_pos=start_pos,
+    )
+
+    compile_tokens = torch.zeros(max_batch_size, prompt_len, dtype=tokens.dtype)
+    compile_tokens[:batch_size] = tokens
+    compile_prompt_lens = torch.zeros(max_batch_size, dtype=prompt_lens.dtype)
+    compile_prompt_lens[:batch_size] = prompt_lens
+    compile_page_table = page_table
+    compile_empty_slots = list(range(batch_size))
+    if batch_size < max_batch_size:
+        # A partial-cardinality diagnostic still decodes on the model's full
+        # lane-capacity program, but prefill owns only its active requests.
+        compile_tokens = tokens
+        compile_prompt_lens = prompt_lens
+        compile_page_table = page_table[:batch_size]
+    _compile_prefill_and_decode(
+        execution_target,
+        prefill_tokens=compile_tokens,
+        prefill_page_table=compile_page_table,
+        kv_cache=kv_cache,
+        prompt_lens=compile_prompt_lens,
+        empty_slots=compile_empty_slots,
+        start_pos=start_pos,
+        sampling_params=sampling_params,
+        prefill_sampling_params=prefill_sampling_params,
+        decode_tokens=torch.zeros(max_batch_size, dtype=torch.long),
+        decode_start_pos=torch.full((max_batch_size,), prompt_len, dtype=torch.long),
+        decode_page_table=page_table,
+    )
+
+    _profiler_start(profiler, "inference_prefill")
+    try:
+        start_time = time.perf_counter()
+        prefill_output = execution_target.prefill_forward(
+            tokens,
+            **prefill_kwargs,
+            sampling_params=prefill_sampling_params,
+        )
+        _synchronize_target(execution_target)
+        prefill_time = time.perf_counter() - start_time
+    finally:
+        _profiler_end(profiler, "inference_prefill")
+
+    first_margins = None
+    if collect_argmax_diagnostics and sampling_params is None:
+        first_token, first_margins = _host_argmax_with_margins(prefill_output, batch_size)
+    else:
+        first_token = prefill_output[0] if isinstance(prefill_output, tuple) else torch.argmax(prefill_output, dim=-1)
+        first_token = first_token.view(-1)[:batch_size].detach().cpu()
+    generated_token_ids = [[int(token)] for token in first_token.tolist()]
+    argmax_top2_margins = [[float(margin)] for margin in first_margins.tolist()] if first_margins is not None else None
+
+    current_tokens = torch.zeros(max_batch_size, dtype=torch.long)
+    current_tokens[:batch_size] = first_token
+    current_pos = torch.full((max_batch_size,), -1, dtype=torch.long)
+    current_pos[:batch_size] = prompt_lens[:batch_size]
+
+    compile_time = None
+    decode_times = []
+    decode_iteration_times = []
+    sampled_decode_start = None
+    pending_host_output = None
+    pending_read_events = None
+
+    _profiler_start(profiler, "inference_decode")
+    try:
+        for iteration in range(num_decode_tokens):
+            start_time = time.perf_counter()
+            read_from_device = sampling_params is None or not can_pipeline_readback
+            if sampling_params is not None and iteration == 1:
+                sampled_decode_start = start_time
+
+            decode_output = execution_target.decode_forward(
+                current_tokens,
+                current_pos,
+                page_table=page_table,
+                kv_cache=kv_cache,
+                read_from_device=read_from_device,
+                sampling_params=sampling_params,
+                reset_batch=iteration == 0,
+            )
+            output, _ = _split_output(decode_output)
+
+            completed_host_output = None
+            if can_pipeline_readback:
+                host_output, read_events = _submit_decode_read(execution_target, decode_output)
+                if pending_read_events is not None:
+                    _synchronize_read_events(pending_read_events)
+                    completed_host_output = pending_host_output
+                pending_host_output = host_output
+                pending_read_events = read_events
+
+            if sampling_params is not None and iteration == 0:
+                _synchronize_target(execution_target)
+            elapsed = time.perf_counter() - start_time
+
+            if iteration == 0:
+                compile_time = elapsed
+            else:
+                decode_iteration_times.append(elapsed)
+                if sampling_params is None or can_pipeline_readback:
+                    decode_times.append(elapsed)
+
+            if completed_host_output is not None:
+                _consume_sampled_output(
+                    execution_target,
+                    completed_host_output,
+                    batch_size,
+                    cluster_shape,
+                    generated_token_ids,
+                    process_host_output=True,
+                )
+
+            if sampling_params is None:
+                if isinstance(output, torch.Tensor) and output.dim() >= 2:
+                    if collect_argmax_diagnostics:
+                        next_token, next_margins = _host_argmax_with_margins(output, batch_size)
+                    else:
+                        next_token = torch.argmax(output[:, -1, :], dim=-1)
+                else:
+                    if collect_argmax_diagnostics:
+                        raise TypeError("argmax margin diagnostics require host decode logits")
+                    next_token = output
+                next_token = next_token.view(-1)[:batch_size].detach().cpu()
+                for user_id, token in enumerate(next_token.tolist()):
+                    generated_token_ids[user_id].append(int(token))
+                    if argmax_top2_margins is not None:
+                        argmax_top2_margins[user_id].append(float(next_margins[user_id]))
+                current_tokens[:batch_size] = next_token
+            elif not can_pipeline_readback:
+                _consume_sampled_output(
+                    execution_target,
+                    decode_output,
+                    batch_size,
+                    cluster_shape,
+                    generated_token_ids,
+                    process_host_output=False,
+                )
+            current_pos[:batch_size] += 1
+    finally:
+        _profiler_end(profiler, "inference_decode")
+
+    if sampling_params is not None:
+        if pending_read_events is not None:
+            _synchronize_read_events(pending_read_events)
+            _consume_sampled_output(
+                execution_target,
+                pending_host_output,
+                batch_size,
+                cluster_shape,
+                generated_token_ids,
+                process_host_output=True,
+            )
+        if sampled_decode_start is not None and not can_pipeline_readback:
+            _synchronize_target(execution_target)
+            sampled_decode_time = time.perf_counter() - sampled_decode_start
+            decode_times = [sampled_decode_time / (num_decode_tokens - 1)] * (num_decode_tokens - 1)
+
+    return PerfBenchmarkResult(
+        prefill_time_s=prefill_time,
+        compile_decode_time_s=compile_time or 0.0,
+        decode_times_s=decode_times,
+        batch_size=batch_size,
+        num_decode_tokens=num_decode_tokens,
+        generated_token_ids=generated_token_ids,
+        decode_iteration_times_s=decode_iteration_times,
+        argmax_top2_margins=argmax_top2_margins,
+    )
+
+
+def _add_token_ids(target: set[int], value) -> None:
+    if value is None or isinstance(value, bool):
+        return
+    if isinstance(value, int):
+        target.add(int(value))
+        return
+    if isinstance(value, (list, tuple, set)):
+        for item in value:
+            _add_token_ids(target, item)
+
+
+def _stop_token_ids(tokenizer) -> set[int]:
+    stop_ids: set[int] = set()
+    _add_token_ids(stop_ids, getattr(tokenizer, "eos_token_id", None))
+    _add_token_ids(stop_ids, getattr(tokenizer, "stop_tokens", None))
+    convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None)
+    if callable(convert_tokens_to_ids):
+        eot_id = convert_tokens_to_ids("<|eot_id|>")
+        if isinstance(eot_id, int) and eot_id >= 0:
+            stop_ids.add(eot_id)
+    return stop_ids
+
+
+def _truncate_at_stop(token_ids: list[int], stop_ids: set[int]) -> list[int]:
+    for index, token_id in enumerate(token_ids):
+        if token_id in stop_ids:
+            return token_ids[:index]
+    return token_ids
+
+
+def assert_no_special_tokens(
+    generated_token_ids: list[list[int]],
+    tokenizer,
+    *,
+    case_name: str = "",
+    is_ci_env: bool | None = None,
+) -> None:
+    """Warn locally; fail under CI or ``TT_DEMO_STRICT_SPECIAL_TOKENS=1``."""
+    if is_ci_env is None:
+        is_ci_env = os.environ.get("CI") == "true" or os.environ.get("TT_DEMO_STRICT_SPECIAL_TOKENS") == "1"
+
+    special_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+    stop_ids = _stop_token_ids(tokenizer)
+    offending_users = 0
+    for token_ids in generated_token_ids:
+        output_before_stop = _truncate_at_stop(list(token_ids), stop_ids)
+        if any(token_id in special_ids for token_id in output_before_stop):
+            offending_users += 1
+
+    if offending_users == 0:
+        return
+
+    prefix = f"[{case_name}] " if case_name else ""
+    message = f"{prefix}model produced special tokens ({offending_users}/{len(generated_token_ids)} users)"
+    logger.warning(message)
+    if is_ci_env:
+        raise AssertionError(message)
+
+
+def load_eval_repeat_prompts_batch32() -> list[str]:
+    """The 32 numeric sequence-continuation prompts TTTv1's ci-eval-32 uses (parity)."""
+    path = Path("models/tt_transformers/demo/sample_prompts/eval_repeat_prompts_batch32.json")
+    with open(path) as f:
+        data = json.load(f)
+    return [entry["prompt"] for entry in data]
+
+
+def rotate_prompts(all_prompts: list[str], repeat: int) -> list[str]:
+    """Rotate the prompt->slot assignment by ``repeat``: slot j holds prompt (j+repeat)%N."""
+    n = len(all_prompts)
+    return [all_prompts[(j + repeat) % n] for j in range(n)]
+
+
+def eval_page_table_for_repeat(page_table: torch.Tensor, repeat: int, *, mode: str) -> torch.Tensor:
+    """Select physical KV allocation for one rotated eval repeat.
+
+    ``slot-stable`` is the TTTv1 acceptance geometry: page-table row ``j``
+    remains at decode slot ``j`` while prompts rotate. ``prompt-stable`` is a
+    diagnostic A/B: rows rotate with prompts, so a prompt retains its original
+    physical KV blocks even as it moves to another decode slot.
+    """
+
+    if mode == "slot-stable":
+        return page_table
+    if mode == "prompt-stable":
+        return torch.roll(page_table, shifts=-repeat, dims=0)
+    raise ValueError(f"unsupported eval page-table mode {mode!r}; use 'slot-stable' or 'prompt-stable'")
+
+
+def eval_decode_trace_mode(mode: str) -> str:
+    """Resolve the eval decode execution A/B without changing its default gate."""
+
+    if mode == "traced":
+        return "decode_only"
+    if mode == "eager":
+        return "none"
+    raise ValueError(f"unsupported eval decode mode {mode!r}; use 'traced' or 'eager'")
+
+
+def require_canonical_eval_modes_in_ci(environ) -> None:
+    """Prevent diagnostic A/B knobs from replacing a canonical CI gate."""
+
+    if environ.get("CI") != "true":
+        return
+    noncanonical = []
+    if environ.get("EVAL_DECODE_MODE", "traced") != "traced":
+        noncanonical.append("EVAL_DECODE_MODE")
+    if environ.get("EVAL_PAGE_TABLE_MODE", "slot-stable") != "slot-stable":
+        noncanonical.append("EVAL_PAGE_TABLE_MODE")
+    noncanonical.extend(name for name in ("EVAL_IDENTICAL_PROMPT_INDEX", "EVAL_ACTIVE_BATCH_SIZE") if name in environ)
+    if noncanonical:
+        raise RuntimeError("diagnostic eval modes cannot replace the canonical CI gate: " + ", ".join(noncanonical))
+
+
+def truncate_at_stop(ids: list[int], stop_ids: set[int]) -> list[int]:
+    """Prefix of ``ids`` up to (excluding) the first id in ``stop_ids``."""
+    out: list[int] = []
+    for t in ids:
+        if t in stop_ids:
+            break
+        out.append(t)
+    return out
+
+
+def hf_stop_ids(tokenizer, hf_model_id: str | None = None) -> set[int]:
+    """Best-effort stop-token id set for an HF ``AutoTokenizer``.
+
+    Raw HF tokenizers have no ``.stop_tokens`` (that only exists on the TTTv1 wrapped
+    tokenizer). Build the set from ``eos_token_id`` (int|list|None), and — when an
+    ``hf_model_id`` is supplied — also fold in the model's ``generation_config`` eos ids,
+    since chat models (e.g. Llama-3 Instruct) often carry extra eot ids there rather than
+    on ``eos_token_id``. Missing/empty -> empty set (truncation simply runs full length).
+    """
+    stop: set[int] = set()
+
+    def _add(value) -> None:
+        if value is None:
+            return
+        if isinstance(value, bool):  # guard: bool is an int subclass
+            return
+        if isinstance(value, int):
+            stop.add(int(value))
+        elif isinstance(value, (list, tuple, set)):
+            for e in value:
+                _add(e)
+
+    _add(getattr(tokenizer, "eos_token_id", None))
+    # tt_transformers ModelArgs.tokenizer augments the HF tokenizer with ``stop_tokens``
+    # (eos + any extra eot ids); raw HF AutoTokenizers don't have it (getattr -> None).
+    _add(getattr(tokenizer, "stop_tokens", None))
+    if hf_model_id is not None:
+        try:
+            from transformers import GenerationConfig
+
+            gen_cfg = GenerationConfig.from_pretrained(hf_model_id)
+            _add(getattr(gen_cfg, "eos_token_id", None))
+        except Exception as e:  # generation_config absent / unreadable — eos_token_id is enough
+            logger.debug(f"ci-eval-32: could not read generation_config eos ids for {hf_model_id}: {e}")
+    return stop
+
+
+def decode_eval_output(tokenizer, token_ids: list[int], stop_ids: set[int]) -> str:
+    """Return the continuation text used by TTTv1's ci-eval-32 comparison.
+
+    Different valid BPE segmentations can decode to the same text.  The TTTv1
+    reference stores ``tokenizer.decode(...)`` results and compares those
+    strings, so comparing token-id lists here would make the port stricter than
+    its source workload and report false slot-rotation failures.
+    """
+    decoded = tokenizer.decode(truncate_at_stop(token_ids, stop_ids))
+    if not isinstance(decoded, str):
+        raise TypeError(f"tokenizer.decode must return str, got {type(decoded).__name__}")
+    return decoded
+
+
+def assert_cross_cardinality_consistency(
+    outputs_by_cardinality: dict[int, dict[str, str]],
+    *,
+    expected_cardinalities: tuple[int, ...] = (1, 2, 4, 32),
+) -> None:
+    """Require each fixed request's decoded output to be invariant as batch cardinality grows."""
+    if tuple(outputs_by_cardinality) != expected_cardinalities:
+        raise AssertionError(
+            f"cross-cardinality experiment expected {expected_cardinalities}, " f"got {tuple(outputs_by_cardinality)}"
+        )
+    reference: dict[str, tuple[int, str]] = {}
+    for cardinality, outputs in outputs_by_cardinality.items():
+        if len(outputs) != cardinality:
+            raise AssertionError(f"cardinality {cardinality} returned {len(outputs)} request outputs")
+        for request_id, output in outputs.items():
+            if request_id in reference:
+                reference_cardinality, reference_output = reference[request_id]
+                if output != reference_output:
+                    raise AssertionError(
+                        f"request {request_id!r} differs at cardinality {reference_cardinality}->{cardinality}: "
+                        f"{reference_output[:120]!r} != {output[:120]!r}"
+                    )
+            else:
+                reference[request_id] = (cardinality, output)
+
+
+def assert_cross_batch_consistency(
+    per_repeat_outputs: list[list[str]],
+    *,
+    per_repeat_token_ids: list[list[list[int]]] | None = None,
+    per_repeat_prompt_lens: list[list[int]] | None = None,
+    per_repeat_argmax_margins: list[list[list[float]] | None] | None = None,
+) -> None:
+    """Assert decoded prompt-position invariance across repeats.
+
+    ``per_repeat_outputs[b][u]`` = decoded continuation for slot ``u`` of repeat ``b``.
+    With slot j of repeat b holding prompt (j+b)%N (see ``rotate_prompts``), the same prompt
+    sits at slot (offset+1)%N of repeat b and slot offset of repeat b+1 — so those two
+    outputs must be identical if no per-user state leaks.
+    """
+    num_batches = len(per_repeat_outputs)
+    assert num_batches >= 2, "cross-batch consistency needs >=2 repeats"
+    n = len(per_repeat_outputs[0])
+    failed, total = 0, 0
+    first_failure = None
+    for b in range(num_batches - 1):
+        cur, nxt = per_repeat_outputs[b], per_repeat_outputs[b + 1]
+        for offset in range(n):
+            total += 1
+            if cur[(offset + 1) % n] != nxt[offset]:
+                failed += 1
+                if first_failure is None:
+                    current_slot = (offset + 1) % n
+                    prompt_index = (offset + b + 1) % n
+                    token_detail = ""
+                    if per_repeat_token_ids is not None:
+                        current_tokens = per_repeat_token_ids[b][current_slot]
+                        next_tokens = per_repeat_token_ids[b + 1][offset]
+                        common_tokens = 0
+                        for current_token, next_token in zip(current_tokens, next_tokens):
+                            if current_token != next_token:
+                                break
+                            common_tokens += 1
+                        current_token = current_tokens[common_tokens] if common_tokens < len(current_tokens) else None
+                        next_token = next_tokens[common_tokens] if common_tokens < len(next_tokens) else None
+                        token_detail = (
+                            f"; first token divergence at generation step {common_tokens} "
+                            f"({current_token!r} != {next_token!r})"
+                        )
+                        if per_repeat_argmax_margins is not None:
+                            current_repeat_margins = per_repeat_argmax_margins[b]
+                            next_repeat_margins = per_repeat_argmax_margins[b + 1]
+                            if current_repeat_margins is not None and next_repeat_margins is not None:
+                                current_margin = current_repeat_margins[current_slot][common_tokens]
+                                next_margin = next_repeat_margins[offset][common_tokens]
+                                token_detail += f"; top2 margins {current_margin:.6g} and {next_margin:.6g}"
+                    length_detail = ""
+                    if per_repeat_prompt_lens is not None:
+                        current_length = per_repeat_prompt_lens[b][current_slot]
+                        next_length = per_repeat_prompt_lens[b + 1][offset]
+                        length_detail = f"; prompt lengths {current_length} and {next_length}"
+                    first_failure = (
+                        b,
+                        offset,
+                        current_slot,
+                        prompt_index,
+                        cur[current_slot],
+                        nxt[offset],
+                        token_detail,
+                        length_detail,
+                    )
+    assert failed == 0, (
+        f"ci-eval-32: {failed}/{total} cross-batch consistency checks failed "
+        f"(first at repeat {first_failure[0]} slot {first_failure[2]} -> "
+        f"repeat {first_failure[0] + 1} slot {first_failure[1]}, prompt index {first_failure[3]}"
+        f"{first_failure[6]}{first_failure[7]}; "
+        f"decoded outputs {first_failure[4][:80]!r} != {first_failure[5][:80]!r})"
+    )
+
+
+def assert_within_batch_slot_consistency(
+    decoded_outputs: list[str],
+    *,
+    token_ids: list[list[int]],
+    argmax_margins: list[list[float]] | None,
+    prompt_index: int,
+) -> None:
+    """Prove one fixed request is invariant across logical slots in one run."""
+
+    reference = decoded_outputs[0]
+    for slot, decoded in enumerate(decoded_outputs[1:], start=1):
+        if decoded == reference:
+            continue
+        reference_tokens = token_ids[0]
+        slot_tokens = token_ids[slot]
+        common_tokens = 0
+        for reference_token, slot_token in zip(reference_tokens, slot_tokens):
+            if reference_token != slot_token:
+                break
+            common_tokens += 1
+        reference_token = reference_tokens[common_tokens] if common_tokens < len(reference_tokens) else None
+        slot_token = slot_tokens[common_tokens] if common_tokens < len(slot_tokens) else None
+        margin_detail = ""
+        if argmax_margins is not None:
+            margin_detail = (
+                f"; top2 margins {argmax_margins[0][common_tokens]:.6g} "
+                f"and {argmax_margins[slot][common_tokens]:.6g}"
+            )
+        raise AssertionError(
+            f"ci-eval-32 identical-request diagnostic: prompt index {prompt_index} differs between "
+            f"logical slots 0 and {slot} at generation step {common_tokens} "
+            f"({reference_token!r} != {slot_token!r}){margin_detail}; "
+            f"decoded outputs {reference[:80]!r} != {decoded[:80]!r}"
+        )
+
+
+def run_eval_repeat_batch32(
+    *,
+    make_executor,
+    allocate_kv_cache,
+    page_table: torch.Tensor,
+    prompts: list[str],
+    tokenizer,
+    tokenize_fn,
+    num_decode_tokens: int,
+    max_batch_size: int,
+    sampling_params=None,
+    repeat_batches: int = 3,
+    hf_model_id: str | None = None,
+    first_repeat_profiler=None,
+    page_table_mode: str = "slot-stable",
+    identical_prompt_index: int | None = None,
+    active_batch_size: int | None = None,
+) -> PerfBenchmarkResult:
+    """Drive the ci-eval-32 determinism case, building a fresh traced executor per repeat.
+
+    Each repeat builds its own traced executor (``make_executor()``) and its own zeroed KV
+    cache (``allocate_kv_cache(executor)``), so the rotated batches are fully independent —
+    no shared device or host state can leak across repeats. The executor is cleaned up after
+    each repeat. (The model is bit-deterministic across repeats either way; fresh-per-repeat
+    is simply the cleanest independence guarantee for a determinism test, and the trace
+    recapture cost is negligible at batch-32.)
+
+    Args:
+        make_executor: Zero-arg callable returning a fresh traced executor
+            (``run_perf_benchmark`` requires traced). Called once per repeat.
+        allocate_kv_cache: Callable(executor) -> fresh zeroed kv_cache bound on that executor.
+        page_table: Fixed contiguous page table (shared across repeats).
+        prompts: The N (=max_batch_size) prompts to rotate (TTTv1 ci-eval-32 numeric prompts;
+            see the module note above re: degenerate-output sensitivity on small models).
+        tokenizer: HF tokenizer (for stop / special ids).
+        tokenize_fn: Callable(list[str]) -> (tokens, prompt_lens).
+        num_decode_tokens: Decode steps per repeat.
+        max_batch_size: Padded batch (== len(prompts) for this fixed-32 case).
+        sampling_params: None -> host argmax (deterministic, mesh-agnostic default).
+        repeat_batches: Number of rotated repeats (TTTv1 uses 3).
+        hf_model_id: Optional, to enrich stop ids from generation_config.
+        first_repeat_profiler: Optional profiler passed only to the first repeat, allowing a caller
+            to emit perf telemetry without changing the three-repeat determinism geometry.
+        page_table_mode: ``slot-stable`` preserves TTTv1 acceptance geometry.
+            ``prompt-stable`` keeps each prompt on the same physical KV blocks
+            as a diagnostic A/B while it rotates through decode slots.
+        identical_prompt_index: Diagnostic-only fixed request replicated into
+            every logical slot. Slot consistency is checked after the first
+            batch, before any repeat-lifecycle comparison.
+        active_batch_size: Diagnostic-only number of leading logical slots to
+            activate while retaining the model's full lane capacity. This is
+            restricted to the identical-request probe so inactive lanes cannot
+            complicate prompt-rotation semantics.
+    """
+    assert (
+        len(prompts) == max_batch_size
+    ), f"ci-eval-32 expects len(prompts)==max_batch_size; got {len(prompts)} vs {max_batch_size}"
+    if active_batch_size is not None:
+        if identical_prompt_index is None:
+            raise ValueError("active_batch_size requires identical_prompt_index")
+        if not 1 <= active_batch_size <= max_batch_size:
+            raise ValueError(f"active_batch_size must be in [1, {max_batch_size}]")
+    if identical_prompt_index is not None:
+        if not 0 <= identical_prompt_index < len(prompts):
+            raise ValueError(f"identical_prompt_index must be in [0, {len(prompts) - 1}]")
+        prompts = [prompts[identical_prompt_index]] * (active_batch_size or len(prompts))
+    stop_ids = hf_stop_ids(tokenizer, hf_model_id)
+    special_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+    # Garbage guard targets only special tokens that are NOT recognized stops: a legitimate
+    # stop is removed by truncation, so anything special left in the body is degenerate output.
+    garbage_ids = special_ids - stop_ids
+    logger.info(
+        f"ci-eval-32: repeat_batches={repeat_batches}, N={len(prompts)}, "
+        f"stop_ids={sorted(stop_ids)}, |special_ids|={len(special_ids)}, sampling_params={sampling_params}, "
+        f"page_table_mode={page_table_mode}, identical_prompt_index={identical_prompt_index}, "
+        f"active_batch_size={len(prompts)}"
+    )
+
+    per_repeat: list[list[str]] = []
+    per_repeat_token_ids: list[list[list[int]]] = []
+    per_repeat_prompt_lens: list[list[int]] = []
+    per_repeat_argmax_margins: list[list[list[float]] | None] = []
+    first_result = None
+    for i in range(repeat_batches):
+        traced_executor = make_executor()
+        try:
+            kv_cache = allocate_kv_cache(traced_executor)
+            rotated = rotate_prompts(prompts, i)
+            tokens, prompt_lens = tokenize_fn(rotated)
+            repeat_page_table = eval_page_table_for_repeat(page_table, i, mode=page_table_mode)
+            result = run_perf_benchmark(
+                traced_executor,
+                tokens=tokens,
+                kv_cache=kv_cache,
+                page_table=repeat_page_table,
+                num_decode_tokens=num_decode_tokens,
+                max_batch_size=max_batch_size,
+                prompt_lens=prompt_lens,
+                sampling_params=sampling_params,
+                profiler=first_repeat_profiler if i == 0 else None,
+                collect_argmax_diagnostics=sampling_params is None,
+            )
+        finally:
+            traced_executor.cleanup()
+        truncated = [truncate_at_stop(ids, stop_ids) for ids in result.generated_token_ids]
+        for u, ids in enumerate(truncated):
+            bad = set(ids) & garbage_ids
+            assert not bad, f"ci-eval-32: user {u} produced special token(s) {sorted(bad)} mid-stream"
+        decoded = [decode_eval_output(tokenizer, ids, stop_ids) for ids in result.generated_token_ids]
+        per_repeat.append(decoded)
+        per_repeat_token_ids.append(result.generated_token_ids)
+        per_repeat_prompt_lens.append([int(length) for length in prompt_lens])
+        per_repeat_argmax_margins.append(getattr(result, "argmax_top2_margins", None))
+        if identical_prompt_index is not None:
+            assert_within_batch_slot_consistency(
+                decoded,
+                token_ids=result.generated_token_ids,
+                argmax_margins=getattr(result, "argmax_top2_margins", None),
+                prompt_index=identical_prompt_index,
+            )
+        if i == 0:
+            first_result = result
+        logger.info(
+            f"ci-eval-32 repeat {i}: truncated token lengths = {[len(t) for t in truncated]}, "
+            f"decoded character lengths = {[len(text) for text in decoded]}"
+        )
+
+    if identical_prompt_index is None:
+        assert_cross_batch_consistency(
+            per_repeat,
+            per_repeat_token_ids=per_repeat_token_ids,
+            per_repeat_prompt_lens=per_repeat_prompt_lens,
+            per_repeat_argmax_margins=per_repeat_argmax_margins,
+        )
+        logger.info(f"ci-eval-32: all {(repeat_batches - 1) * len(prompts)} cross-batch consistency checks passed")
+    else:
+        logger.info(
+            f"ci-eval-32 identical-request diagnostic: prompt index {identical_prompt_index} "
+            f"is invariant across all {len(prompts)} logical slots"
+        )
+    assert first_result is not None
+    return first_result

code/models/common/tests/demos/test_cleanup_utils.py

@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from models.common.tests.demos import cleanup_utils
+
+
+def test_cleanup_dp_model_case_orders_owners_before_parent_and_children(monkeypatch):
+    calls = []
+    parent = SimpleNamespace(quiesce_devices=lambda: calls.append("parent-quiesce"))
+    submeshes = [object(), object()]
+    group = SimpleNamespace(cleanup=lambda: calls.append("group-cleanup"))
+    models = [("model-0", submeshes[0]), ("model-1", submeshes[1])]
+
+    monkeypatch.setattr(
+        cleanup_utils,
+        "cleanup_model_case",
+        lambda model, submesh: calls.append(("model-cleanup", model, submesh)),
+    )
+    monkeypatch.setattr(
+        cleanup_utils.ttnn,
+        "close_mesh_device",
+        lambda submesh: calls.append(("child-close", submesh)),
+    )
+
+    cleanup_utils.cleanup_dp_model_case(group, [], models, parent, submeshes)
+
+    assert calls == [
+        "group-cleanup",
+        ("model-cleanup", "model-0", submeshes[0]),
+        ("model-cleanup", "model-1", submeshes[1]),
+        "parent-quiesce",
+        ("child-close", submeshes[0]),
+        ("child-close", submeshes[1]),
+    ]
+
+
+def test_cleanup_dp_model_case_closes_every_carved_child_after_partial_build(monkeypatch):
+    calls = []
+    submeshes = [object(), object(), object(), object()]
+    lanes = [SimpleNamespace(cleanup=lambda: calls.append("lane-cleanup"))]
+    parent = SimpleNamespace(quiesce_devices=lambda: calls.append("parent-quiesce"))
+
+    monkeypatch.setattr(
+        cleanup_utils,
+        "cleanup_model_case",
+        lambda model, submesh: calls.append(("model-cleanup", model, submesh)),
+    )
+    monkeypatch.setattr(
+        cleanup_utils.ttnn,
+        "close_mesh_device",
+        lambda submesh: calls.append(("child-close", submesh)),
+    )
+
+    cleanup_utils.cleanup_dp_model_case(None, lanes, [("model-0", submeshes[0])], parent, submeshes)
+
+    assert calls[:3] == ["lane-cleanup", ("model-cleanup", "model-0", submeshes[0]), "parent-quiesce"]
+    assert calls[3:] == [("child-close", submesh) for submesh in submeshes]
+
+
+def test_cleanup_dp_model_case_attempts_all_children_after_cleanup_failure(monkeypatch, expect_error):
+    closed = []
+    submeshes = [object(), object(), object()]
+
+    def fail_group_cleanup():
+        raise RuntimeError("group cleanup failed")
+
+    monkeypatch.setattr(cleanup_utils, "cleanup_model_case", lambda *_: None)
+    monkeypatch.setattr(cleanup_utils.ttnn, "close_mesh_device", closed.append)
+
+    with expect_error(RuntimeError, "group cleanup failed"):
+        cleanup_utils.cleanup_dp_model_case(
+            SimpleNamespace(cleanup=fail_group_cleanup),
+            [],
+            [],
+            SimpleNamespace(quiesce_devices=lambda: None),
+            submeshes,
+        )
+
+    assert closed == submeshes
+
+
+def test_two_sequential_dp_profiles_release_children_between_carves(monkeypatch):
+    class Child:
+        def __init__(self, generation):
+            self.generation = generation
+            self.closed = False
+
+    class Parent:
+        def __init__(self):
+            self.children = []
+            self.quiesce_count = 0
+
+        def quiesce_devices(self):
+            self.quiesce_count += 1
+
+        def carve(self, generation):
+            assert all(child.closed for child in self.children)
+            children = [Child(generation), Child(generation)]
+            self.children.extend(children)
+            return children
+
+    parent = Parent()
+    monkeypatch.setattr(cleanup_utils.ttnn, "close_mesh_device", lambda child: setattr(child, "closed", True))
+
+    for generation in ("performance", "accuracy"):
+        parent.quiesce_devices()
+        children = parent.carve(generation)
+        cleanup_utils.cleanup_dp_model_case(None, [], [], parent, children)
+
+    assert parent.quiesce_count == 4
+    assert all(child.closed for child in parent.children)
+
+
+@pytest.mark.parametrize(
+    ("demo_path", "create_name"),
+    [
+        ("models/common/tests/demos/qwen2_7b/demo.py", "_create_dp_submeshes"),
+        ("models/common/tests/demos/qwen25_7b/demo.py", "_create_dp_submeshes"),
+        ("models/common/tests/demos/deepseek_r1_distill_qwen_14b/demo.py", "_create_dp_submeshes"),
+        ("models/common/tests/demos/llama32_1b/demo.py", "create_dp_submeshes"),
+        ("models/common/tests/demos/llama32_3b/demo.py", "create_dp_submeshes"),
+    ],
+)
+def test_dp_demos_quiesce_before_carving_and_use_shared_teardown(demo_path, create_name):
+    tree = ast.parse(Path(demo_path).read_text(encoding="utf-8"), filename=demo_path)
+    function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_dp_smoke")
+    calls = [node for node in ast.walk(function) if isinstance(node, ast.Call)]
+    create_call = next(node for node in calls if isinstance(node.func, ast.Name) and node.func.id == create_name)
+    parent_quiesce = next(
+        node
+        for node in calls
+        if isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "mesh_device"
+        and node.func.attr == "quiesce_devices"
+    )
+    teardown_call = next(
+        node for node in calls if isinstance(node.func, ast.Name) and node.func.id == "cleanup_dp_model_case"
+    )
+
+    assert parent_quiesce.lineno < create_call.lineno < teardown_call.lineno
+    assert [ast.unparse(argument) for argument in teardown_call.args] == [
+        "group",
+        "lanes",
+        "models",
+        "mesh_device",
+        "submeshes",
+    ]

code/models/common/tests/demos/test_run_helpers.py

@@ -0,0 +1,687 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from models.common.tests.demos import run_helpers
+from models.common.tests.demos.run_helpers import run_perf_benchmark, run_teacher_forcing
+
+
+class FakeProfiler:
+    def __init__(self):
+        self.events = []
+
+    def start(self, name):
+        self.events.append(("start", name))
+
+    def end(self, name):
+        self.events.append(("end", name))
+
+
+def _logits(token_ids, vocab_size=8):
+    output = torch.zeros(len(token_ids), 1, vocab_size)
+    for row, token_id in enumerate(token_ids):
+        output[row, 0, token_id] = 1
+    return output
+
+
+class FakeExecutionTarget:
+    def __init__(self, *, compile_prefill_output, prefill_output, decode_outputs):
+        self.compile_prefill_output = compile_prefill_output
+        self.prefill_output = prefill_output
+        self.decode_outputs = list(decode_outputs)
+        self.calls = []
+
+    def _record(self, method_name, arguments):
+        self.calls.append(
+            (
+                method_name,
+                {name: value for name, value in arguments.items() if name != "self"},
+            )
+        )
+
+    @property
+    def _engine(self):
+        raise AssertionError("execution helpers must not inspect a private wrapped engine")
+
+    def compile_prefill(
+        self,
+        *,
+        tokens,  # ↓ Core request
+        page_table,
+        prompt_lens=None,  # ↓ Sequence metadata
+        start_pos=None,
+        empty_slots=None,  # ↓ Lane routing
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("compile_prefill", locals())
+        return self.compile_prefill_output
+
+    def compile_decode(
+        self,
+        *,
+        tokens,  # ↓ Core request
+        start_pos,
+        page_table,
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        reset_batch=False,  # ↓ State transition
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("compile_decode", locals())
+
+    def prefill_forward(
+        self,
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,  # ↓ Sequence metadata
+        start_pos=None,
+        empty_slots=None,  # ↓ Lane routing
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("prefill_forward", locals())
+        return self.prefill_output
+
+    def decode_forward(
+        self,
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        reset_batch=False,  # ↓ State transition
+        read_from_device=True,  # ↓ Output policy
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("decode_forward", locals())
+        return self.decode_outputs.pop(0)
+
+
+def test_compile_only_helper_ignores_compiled_prefill_programs():
+    target = FakeExecutionTarget(
+        compile_prefill_output=(object(),),
+        prefill_output=None,
+        decode_outputs=[],
+    )
+
+    run_helpers._compile_prefill_and_decode(
+        target,
+        prefill_tokens=torch.tensor([[1, 2], [3, 4]]),
+        prefill_page_table=torch.zeros(2, 1, dtype=torch.int32),
+    )
+
+    assert [name for name, _ in target.calls] == ["compile_prefill", "compile_decode"]
+    assert torch.equal(target.calls[1][1]["tokens"], torch.zeros(2, dtype=torch.long))
+
+
+def test_teacher_forcing_uses_public_target_surface_and_preserves_user_order():
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([3, 4]),
+        prefill_output=_logits([3, 4]),
+        decode_outputs=[(_logits([5, 1]), None)],
+    )
+    top5_tokens = torch.tensor([[3, 0, 1, 2, 4], [5, 0, 1, 2, 3]])
+
+    result = run_teacher_forcing(
+        executor=target,
+        prompt_tokens=torch.tensor([[1, 2], [1, 2]]),
+        reference_tokens=torch.tensor([1, 2, 6, 7]),
+        top5_tokens=top5_tokens,
+        kv_cache=[],
+        page_table=torch.zeros(2, 1, dtype=torch.int32),
+        max_batch_size=2,
+    )
+
+    assert result.predicted_tokens_per_user == [[3, 5], [4, 1]]
+    assert result.top1_accuracy() == 1.0
+    assert result.top5_accuracy() == 1.0
+    assert [name for name, _ in target.calls] == [
+        "compile_prefill",
+        "compile_decode",
+        "prefill_forward",
+        "decode_forward",
+    ]
+
+
+def test_teacher_forcing_times_prefill_excludes_first_decode_and_brackets_profiler(monkeypatch):
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([3, 4]),
+        prefill_output=_logits([3, 4]),
+        decode_outputs=[(_logits([5, 1]), None), (_logits([6, 2]), None)],
+    )
+    profiler = FakeProfiler()
+    times = iter([0.0, 0.2, 1.0, 1.1, 2.0, 2.25])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+
+    result = run_teacher_forcing(
+        executor=target,
+        prompt_tokens=torch.tensor([[1, 2], [1, 2]]),
+        reference_tokens=torch.tensor([1, 2, 8, 9, 10]),
+        top5_tokens=torch.tensor([[3, 0, 1, 2, 4], [5, 0, 1, 2, 3], [6, 0, 1, 2, 3]]),
+        kv_cache=[],
+        page_table=torch.zeros(2, 1, dtype=torch.int32),
+        max_batch_size=2,
+        profiler=profiler,
+    )
+
+    assert result.predicted_tokens_per_user == [[3, 5, 6], [4, 1, 2]]
+    assert result.prefill_time_s == pytest.approx(0.2)
+    assert result.compile_decode_time_s == pytest.approx(0.1)
+    assert result.decode_times_s == pytest.approx([0.25])
+    assert result.ttft_ms == pytest.approx(100.0)
+    assert result.prefill_tok_s == pytest.approx(20.0)
+    assert result.decode_tok_s_u == pytest.approx(4.0)
+    assert result.decode_tok_s == pytest.approx(8.0)
+    assert profiler.events == [
+        ("start", "inference_prefill"),
+        ("end", "inference_prefill"),
+        ("start", "inference_decode"),
+        ("end", "inference_decode"),
+    ]
+
+
+def test_perf_benchmark_host_argmax_path_preserves_timing_and_tokens(monkeypatch):
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([2]),
+        prefill_output=_logits([2]),
+        decode_outputs=[(_logits([3]), None), (_logits([4]), None), (_logits([5]), None)],
+    )
+    times = iter([0.0, 0.1, 1.0, 1.2, 2.0, 2.25, 3.0, 3.3])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+
+    result = run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2]]),
+        kv_cache=[],
+        page_table=torch.zeros(1, 1, dtype=torch.int32),
+        num_decode_tokens=3,
+    )
+
+    assert result.prefill_time_s == pytest.approx(0.1)
+    assert result.compile_decode_time_s == pytest.approx(0.2)
+    assert result.decode_times_s == pytest.approx([0.25, 0.3])
+    assert result.decode_iteration_times_s == pytest.approx([0.25, 0.3])
+    assert result.generated_token_ids == [[2, 3, 4, 5]]
+
+
+def test_perf_benchmark_slices_prefill_page_rows_but_keeps_full_decode_capacity(monkeypatch):
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([2, 2, 2, 2]),
+        prefill_output=_logits([2, 2]),
+        decode_outputs=[(_logits([3, 3, 0, 0]), None)],
+    )
+    times = iter([0.0, 0.1, 1.0, 1.1])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+    page_table = torch.arange(8, dtype=torch.int32).reshape(4, 2)
+
+    run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2], [1, 2]]),
+        kv_cache=[],
+        page_table=page_table,
+        num_decode_tokens=1,
+        max_batch_size=4,
+    )
+
+    calls = {name: arguments for name, arguments in target.calls}
+    assert calls["compile_prefill"]["page_table"].shape[0] == 2
+    torch.testing.assert_close(calls["prefill_forward"]["page_table"], page_table[:2])
+    torch.testing.assert_close(calls["decode_forward"]["page_table"], page_table)
+
+
+def test_perf_benchmark_brackets_profiler_without_changing_host_argmax(monkeypatch):
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([2]),
+        prefill_output=_logits([2]),
+        decode_outputs=[(_logits([3]), None), (_logits([4]), None)],
+    )
+    profiler = FakeProfiler()
+    times = iter([0.0, 0.1, 1.0, 1.2, 2.0, 2.25])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+
+    result = run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2]]),
+        kv_cache=[],
+        page_table=torch.zeros(1, 1, dtype=torch.int32),
+        num_decode_tokens=2,
+        profiler=profiler,
+    )
+
+    assert result.generated_token_ids == [[2, 3, 4]]
+    assert result.prefill_time_s == pytest.approx(0.1)
+    assert result.compile_decode_time_s == pytest.approx(0.2)
+    assert result.decode_times_s == pytest.approx([0.25])
+    assert profiler.events == [
+        ("start", "inference_prefill"),
+        ("end", "inference_prefill"),
+        ("start", "inference_decode"),
+        ("end", "inference_decode"),
+    ]
+
+
+class PublicReadbackTarget(FakeExecutionTarget):
+    def __init__(self):
+        super().__init__(
+            compile_prefill_output=(torch.tensor([2]), None),
+            prefill_output=(torch.tensor([2]), None),
+            decode_outputs=[
+                (torch.tensor([3]), None),
+                (torch.tensor([4]), None),
+                (torch.tensor([5]), None),
+            ],
+        )
+        self.mesh_device = SimpleNamespace(shape=(1, 1))
+        self.next_event = 0
+        self.pending_events = []
+
+    def read_decode_output(self, tt_out, *, async_read=False):
+        assert async_read
+        event = self.next_event
+        self.next_event += 1
+        self.pending_events.append(event)
+        self.calls.append(("read_decode_output", {"event": event}))
+        return tt_out, [event]
+
+    def process_decode_output_host(self, tt_out, *, is_tokens=False):
+        assert is_tokens
+        run_helpers.ttnn.event_synchronize(self.pending_events.pop(0))
+        self.calls.append(("process_decode_output_host", {}))
+        return tt_out
+
+
+def test_perf_benchmark_uses_public_async_readback_without_trace_introspection(monkeypatch):
+    target = PublicReadbackTarget()
+    profiler = FakeProfiler()
+    synchronized_events = []
+    times = iter([0.0, 0.1, 1.0, 1.1, 2.0, 2.1, 3.0, 3.1, 3.4])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+    monkeypatch.setattr(run_helpers.ttnn, "synchronize_device", lambda mesh: None, raising=False)
+    monkeypatch.setattr(run_helpers.ttnn, "event_synchronize", synchronized_events.append, raising=False)
+
+    result = run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2]]),
+        kv_cache=[],
+        page_table=torch.zeros(1, 1, dtype=torch.int32),
+        num_decode_tokens=3,
+        sampling_params=object(),
+        pipeline_readback=True,
+        profiler=profiler,
+    )
+
+    assert [name for name, _ in target.calls[:2]] == ["compile_decode", "compile_prefill"]
+    # The benchmark owns the device-paced wait; the public normalizer may
+    # defensively observe the already-completed event again while retiring it.
+    assert synchronized_events == [0, 0, 1, 1, 2, 2]
+    assert result.generated_token_ids == [[2, 3, 4, 5]]
+    assert len(result.decode_times_s) == 2
+    assert result.decode_iteration_times_s == pytest.approx([0.1, 0.1])
+    assert profiler.events == [
+        ("start", "inference_prefill"),
+        ("end", "inference_prefill"),
+        ("start", "inference_decode"),
+        ("end", "inference_decode"),
+    ]
+
+
+def test_perf_benchmark_does_not_reprocess_blocking_sampled_output(monkeypatch):
+    target = PublicReadbackTarget()
+    times = iter([0.0, 0.1, 1.0, 1.1])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+    monkeypatch.setattr(run_helpers.ttnn, "synchronize_device", lambda mesh: None, raising=False)
+
+    result = run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2]]),
+        kv_cache=[],
+        page_table=torch.zeros(1, 1, dtype=torch.int32),
+        num_decode_tokens=1,
+        sampling_params=object(),
+        pipeline_readback=False,
+    )
+
+    assert result.generated_token_ids == [[2, 3]]
+    assert "process_decode_output_host" not in [name for name, _ in target.calls]
+
+
+def test_perf_benchmark_can_use_host_prefill_with_sampled_decode(monkeypatch):
+    target = FakeExecutionTarget(
+        compile_prefill_output=_logits([2]),
+        prefill_output=_logits([2]),
+        decode_outputs=[(torch.tensor([3]), None)],
+    )
+    times = iter([0.0, 0.1, 1.0, 1.1])
+    monkeypatch.setattr(run_helpers.time, "perf_counter", lambda: next(times))
+
+    result = run_perf_benchmark(
+        target,
+        tokens=torch.tensor([[1, 2]]),
+        kv_cache=[],
+        page_table=torch.zeros(1, 1, dtype=torch.int32),
+        num_decode_tokens=1,
+        sampling_params=object(),
+        prefill_sampling_params=None,
+    )
+
+    calls = {name: arguments for name, arguments in target.calls}
+    assert calls["compile_prefill"]["sampling_params"] is None
+    assert calls["prefill_forward"]["sampling_params"] is None
+    assert calls["compile_decode"]["sampling_params"] is not None
+    assert calls["decode_forward"]["sampling_params"] is not None
+    assert result.generated_token_ids == [[2, 3]]
+
+
+def test_composite_target_synchronizes_each_lane_mesh(monkeypatch):
+    target = SimpleNamespace(mesh_device="parent", mesh_devices=("lane-0", "lane-1"))
+    synchronized = []
+    monkeypatch.setattr(
+        run_helpers.ttnn,
+        "synchronize_device",
+        lambda mesh_device: synchronized.append(mesh_device),
+        raising=False,
+    )
+
+    run_helpers._synchronize_target(target)
+
+    assert synchronized == ["lane-0", "lane-1"]
+
+
+def test_special_token_guard_truncates_at_stop_tokens_warns_locally_and_fails_in_ci(monkeypatch, expect_error):
+    tokenizer = SimpleNamespace(
+        all_special_ids=[0, 1, 2, 99],
+        eos_token_id=2,
+        stop_tokens=[3],
+        convert_tokens_to_ids=lambda token: 99 if token == "<|eot_id|>" else -1,
+    )
+    generated_token_ids = [
+        [5, 2, 0],
+        [5, 3, 0],
+        [6, 99, 1],
+        [7, 0, 8],
+        [9, 1, 2],
+    ]
+    warnings = []
+    monkeypatch.setattr(run_helpers.logger, "warning", warnings.append)
+
+    run_helpers.assert_no_special_tokens(generated_token_ids, tokenizer, case_name="batch-1", is_ci_env=False)
+
+    assert warnings == ["[batch-1] model produced special tokens (2/5 users)"]
+    with expect_error(AssertionError, "2/5 users"):
+        run_helpers.assert_no_special_tokens(generated_token_ids, tokenizer, case_name="batch-1", is_ci_env=True)
+
+
+def test_special_token_guard_without_stop_tokens_keeps_generated_tail_visible(expect_error):
+    tokenizer = SimpleNamespace(
+        all_special_ids=[0],
+        eos_token_id=2,
+        convert_tokens_to_ids=lambda _token: -1,
+    )
+
+    with expect_error(AssertionError, "1/1 users"):
+        run_helpers.assert_no_special_tokens([[5, 3, 0]], tokenizer, is_ci_env=True)
+
+
+def test_eval_repeat_compares_decoded_text_like_tttv1_despite_different_bpe_segmentations():
+    pieces = {10: "12", 11: "3", 12: "1", 13: "23", 99: "<eos>", 77: "ignored"}
+    tokenizer = SimpleNamespace(decode=lambda token_ids: "".join(pieces[token] for token in token_ids))
+
+    first_segmentation = run_helpers.decode_eval_output(tokenizer, [10, 11, 99, 77], {99})
+    second_segmentation = run_helpers.decode_eval_output(tokenizer, [12, 13], {99})
+
+    assert first_segmentation == second_segmentation == "123"
+    # Repeat 1 rotates prompt 1 into slot 0 and prompt 0 into slot 1.
+    run_helpers.assert_cross_batch_consistency(
+        [
+            [first_segmentation, "456"],
+            ["456", second_segmentation],
+        ]
+    )
+
+
+def test_eval_page_table_ab_preserves_slots_or_prompt_physical_blocks(expect_error):
+    page_table = torch.tensor([[0, 1], [10, 11], [20, 21]], dtype=torch.int32)
+
+    assert run_helpers.eval_page_table_for_repeat(page_table, 2, mode="slot-stable") is page_table
+    torch.testing.assert_close(
+        run_helpers.eval_page_table_for_repeat(page_table, 1, mode="prompt-stable"),
+        torch.tensor([[10, 11], [20, 21], [0, 1]], dtype=torch.int32),
+    )
+    with expect_error(ValueError, "slot-stable.*prompt-stable"):
+        run_helpers.eval_page_table_for_repeat(page_table, 0, mode="unsupported")
+
+
+def test_eval_decode_ab_defaults_to_trace_and_can_isolate_eager_execution(expect_error):
+    assert run_helpers.eval_decode_trace_mode("traced") == "decode_only"
+    assert run_helpers.eval_decode_trace_mode("eager") == "none"
+    with expect_error(ValueError, "traced.*eager"):
+        run_helpers.eval_decode_trace_mode("unsupported")
+
+
+@pytest.mark.parametrize(
+    "override",
+    [
+        {"EVAL_DECODE_MODE": "eager"},
+        {"EVAL_PAGE_TABLE_MODE": "prompt-stable"},
+        {"EVAL_IDENTICAL_PROMPT_INDEX": "8"},
+        {"EVAL_ACTIVE_BATCH_SIZE": "24"},
+    ],
+)
+def test_ci_rejects_diagnostic_eval_modes(override, expect_error):
+    with expect_error(RuntimeError, "diagnostic eval modes cannot replace the canonical CI gate"):
+        run_helpers.require_canonical_eval_modes_in_ci({"CI": "true", **override})
+
+
+def test_non_ci_diagnostics_and_canonical_ci_are_allowed():
+    run_helpers.require_canonical_eval_modes_in_ci({"EVAL_DECODE_MODE": "eager", "EVAL_IDENTICAL_PROMPT_INDEX": "8"})
+    run_helpers.require_canonical_eval_modes_in_ci(
+        {"CI": "true", "EVAL_DECODE_MODE": "traced", "EVAL_PAGE_TABLE_MODE": "slot-stable"}
+    )
+
+
+def test_host_argmax_diagnostic_reports_top1_minus_top2_margin():
+    logits = torch.tensor(
+        [
+            [[0.0, 4.0, 3.5]],
+            [[9.0, 1.0, 8.875]],
+        ]
+    )
+
+    tokens, margins = run_helpers._host_argmax_with_margins(logits, 2)
+
+    torch.testing.assert_close(tokens, torch.tensor([1, 0]))
+    torch.testing.assert_close(margins, torch.tensor([0.5, 0.125]))
+
+
+def test_eval_repeat_driver_canonicalizes_each_rotated_output_through_tokenizer(monkeypatch):
+    pieces = {10: "12", 11: "3", 12: "1", 13: "23", 20: "456"}
+    tokenizer = SimpleNamespace(
+        decode=lambda token_ids: "".join(pieces[token] for token in token_ids),
+        eos_token_id=None,
+        stop_tokens=[],
+        all_special_ids=[],
+    )
+    generated = iter(
+        [
+            [[10, 11], [20]],
+            [[20], [12, 13]],
+        ]
+    )
+    monkeypatch.setattr(
+        run_helpers,
+        "run_perf_benchmark",
+        lambda *_args, **_kwargs: SimpleNamespace(generated_token_ids=next(generated)),
+    )
+
+    class FakeEvalExecutor:
+        def cleanup(self):
+            pass
+
+    result = run_helpers.run_eval_repeat_batch32(
+        make_executor=FakeEvalExecutor,
+        allocate_kv_cache=lambda _executor: [],
+        page_table=torch.zeros(2, 1, dtype=torch.int32),
+        prompts=["prompt-0", "prompt-1"],
+        tokenizer=tokenizer,
+        tokenize_fn=lambda prompts: (torch.zeros(len(prompts), 1, dtype=torch.long), torch.ones(len(prompts))),
+        num_decode_tokens=2,
+        max_batch_size=2,
+        repeat_batches=2,
+    )
+
+    assert result.generated_token_ids == [[10, 11], [20]]
+
+
+def test_eval_repeat_still_rejects_genuine_decoded_text_divergence(expect_error):
+    with expect_error(AssertionError, "1/2 cross-batch consistency checks failed"):
+        run_helpers.assert_cross_batch_consistency(
+            [
+                ["123", "456"],
+                ["different", "123"],
+            ]
+        )
+
+
+def test_eval_repeat_failure_localizes_prompt_slots_and_first_token_divergence(expect_error):
+    with expect_error(
+        AssertionError,
+        r"repeat 0 slot 1 -> repeat 1 slot 0, prompt index 1; "
+        r"first token divergence at generation step 1 \(12 != 99\); "
+        r"top2 margins 0.125 and 0.25; prompt lengths 80 and 80",
+    ):
+        run_helpers.assert_cross_batch_consistency(
+            [
+                ["alpha", "beta-left"],
+                ["beta-right", "alpha"],
+            ],
+            per_repeat_token_ids=[
+                [[1], [11, 12, 13]],
+                [[11, 99, 13], [1]],
+            ],
+            per_repeat_prompt_lens=[
+                [64, 80],
+                [80, 64],
+            ],
+            per_repeat_argmax_margins=[
+                [[1.0], [0.5, 0.125, 0.75]],
+                [[0.5, 0.25, 0.75], [1.0]],
+            ],
+        )
+
+
+def test_identical_request_diagnostic_localizes_logical_slot_divergence(expect_error):
+    with expect_error(
+        AssertionError,
+        r"prompt index 8 differs between logical slots 0 and 1 at generation step 2 "
+        r"\(13 != 99\); top2 margins 0.125 and 0.25",
+    ):
+        run_helpers.assert_within_batch_slot_consistency(
+            ["same prefix left", "same prefix right"],
+            token_ids=[[11, 12, 13], [11, 12, 99]],
+            argmax_margins=[[1.0, 0.5, 0.125], [1.0, 0.5, 0.25]],
+            prompt_index=8,
+        )
+
+
+def test_identical_request_diagnostic_accepts_slot_invariant_decoded_text():
+    run_helpers.assert_within_batch_slot_consistency(
+        ["same", "same", "same"],
+        token_ids=[[1], [1], [1]],
+        argmax_margins=[[0.5], [0.5], [0.5]],
+        prompt_index=8,
+    )
+
+
+def test_identical_request_driver_uses_one_fixed_prompt_and_needs_no_repeat(monkeypatch):
+    seen_prompts = []
+    tokenizer = SimpleNamespace(
+        decode=lambda token_ids: "same",
+        eos_token_id=None,
+        stop_tokens=[],
+        all_special_ids=[],
+    )
+    monkeypatch.setattr(
+        run_helpers,
+        "run_perf_benchmark",
+        lambda *_args, **_kwargs: SimpleNamespace(
+            generated_token_ids=[[1]],
+            argmax_top2_margins=[[0.5]],
+        ),
+    )
+
+    class FakeEvalExecutor:
+        def cleanup(self):
+            pass
+
+    run_helpers.run_eval_repeat_batch32(
+        make_executor=FakeEvalExecutor,
+        allocate_kv_cache=lambda _executor: [],
+        page_table=torch.zeros(2, 1, dtype=torch.int32),
+        prompts=["prompt-0", "prompt-1"],
+        tokenizer=tokenizer,
+        tokenize_fn=lambda prompts: (
+            seen_prompts.append(prompts) or torch.zeros(len(prompts), 1, dtype=torch.long),
+            torch.ones(len(prompts)),
+        ),
+        num_decode_tokens=1,
+        max_batch_size=2,
+        repeat_batches=1,
+        identical_prompt_index=1,
+        active_batch_size=1,
+    )
+
+    assert seen_prompts == [["prompt-1"]]
+
+
+def test_active_batch_diagnostic_requires_identical_request(expect_error):
+    with expect_error(ValueError, "active_batch_size requires identical_prompt_index"):
+        run_helpers.run_eval_repeat_batch32(
+            make_executor=lambda: None,
+            allocate_kv_cache=lambda _executor: [],
+            page_table=torch.zeros(2, 1, dtype=torch.int32),
+            prompts=["prompt-0", "prompt-1"],
+            tokenizer=SimpleNamespace(eos_token_id=None, stop_tokens=[], all_special_ids=[]),
+            tokenize_fn=lambda prompts: (torch.zeros(len(prompts), 1), torch.ones(len(prompts))),
+            num_decode_tokens=1,
+            max_batch_size=2,
+            repeat_batches=1,
+            active_batch_size=1,
+        )
+
+
+def test_cross_cardinality_consistency_accepts_fixed_request_prefixes():
+    run_helpers.assert_cross_cardinality_consistency(
+        {
+            1: {"r0": "a"},
+            2: {"r0": "a", "r1": "b"},
+            4: {"r0": "a", "r1": "b", "r2": "c", "r3": "d"},
+            32: {f"r{i}": chr(97 + i) for i in range(32)},
+        }
+    )
+
+
+def test_cross_cardinality_consistency_reports_request_and_cardinalities(expect_error):
+    with expect_error(AssertionError, "request 'r0' differs at cardinality 1->2"):
+        run_helpers.assert_cross_cardinality_consistency(
+            {1: {"r0": "same"}, 2: {"r0": "different", "r1": "x"}},
+            expected_cardinalities=(1, 2),
+        )
+
+
+def test_loop_policy_is_not_exported_from_production_executor():
+    try:
+        from models.common.llm_runtime import execution as production_executor
+    except AttributeError as exc:
+        pytest.skip(f"production executor import requires full ttnn runtime: {exc}")
+
+    for name in ("TeacherForceResult", "PerfBenchmarkResult", "run_teacher_forcing", "run_perf_benchmark"):
+        assert not hasattr(production_executor, name)

code/models/common/tests/host/test_metrics_pytorch_only.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""PyTorch-only metric tests that do not require TTNN hardware."""
+
+import torch
+
+from models.common.metrics import comp_allclose, compute_max_abs_error, compute_mean_abs_error, compute_pcc
+
+
+def test_metrics_pytorch_only():
+    """Verify fallback metric implementations using pure PyTorch tensors."""
+
+    torch_a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.float32)
+    torch_b = torch.tensor([[1.0, 2.0, 3.5], [4.0, 5.0, 6.0]], dtype=torch.float32)
+
+    max_error = compute_max_abs_error(torch_a, torch_b)
+    mean_error = compute_mean_abs_error(torch_a, torch_b)
+    pcc = compute_pcc(torch_a, torch_b)
+
+    expected_max = 0.5
+    expected_mean = (torch_a - torch_b).abs().mean().item()
+
+    assert abs(max_error - expected_max) < 1e-6, f"Expected {expected_max}, got {max_error}"
+    assert abs(mean_error - expected_mean) < 1e-6, "Mean error mismatch"
+    assert 0.99 <= pcc <= 1.0, f"PCC should be high (~1.0) for similar tensors, got {pcc}"
+
+
+def test_comp_allclose_pytorch_only():
+    """PyTorch-only tests for comp_allclose covering pass/fail and edge cases."""
+
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    b = a.clone()
+    passed, msg = comp_allclose(a, b)
+    assert passed, f"Expected pass for exact equality. Got: {msg}"
+    assert "Max ATOL Delta" in msg and "Max RTOL Delta" in msg
+
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    b = torch.tensor([1.0 + 1e-7, 2.0 - 1e-7, 3.0], dtype=torch.float32)
+    passed, msg = comp_allclose(a, b)
+    assert passed, f"Expected pass within default tolerance. Got: {msg}"
+
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    b = torch.tensor([1.0, 2.1, 3.0], dtype=torch.float32)
+    passed, msg = comp_allclose(a, b)
+    assert not passed and "Allclose check failed" in msg
+
+    a = torch.tensor([float("nan"), 1.0, 2.0], dtype=torch.float32)
+    b = torch.tensor([float("nan"), 1.0, 2.0], dtype=torch.float32)
+    passed, _ = comp_allclose(a, b)
+    assert passed, "Both NaNs at same positions should pass"
+
+    a = torch.tensor([float("inf"), -float("inf"), 1.0])
+    b = torch.tensor([float("inf"), -float("inf"), 1.0])
+    passed, _ = comp_allclose(a, b)
+    assert passed, "Same sign infinities should pass"
+
+    a = torch.tensor([float("inf"), -float("inf"), 1.0])
+    b = torch.tensor([float("inf"), float("inf"), 1.0])
+    passed, msg = comp_allclose(a, b)
+    assert not passed and "Allclose check failed" in msg

code/models/common/tests/host/test_utility_functions_imports.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Import-dependency tests for models.common.utility_functions."""
+
+import ast
+from pathlib import Path
+
+
+def test_pytest_is_not_imported_at_module_scope():
+    utility_functions_path = Path(__file__).parents[2] / "utility_functions.py"
+    syntax_tree = ast.parse(utility_functions_path.read_text())
+
+    top_level_pytest_imports = [
+        node
+        for node in syntax_tree.body
+        if (isinstance(node, ast.Import) and any(alias.name == "pytest" for alias in node.names))
+        or (isinstance(node, ast.ImportFrom) and node.module == "pytest")
+    ]
+
+    assert not top_level_pytest_imports, "pytest must remain an optional test-only dependency"
+
+
+def test_ti_skip_imports_pytest_lazily():
+    utility_functions_path = Path(__file__).parents[2] / "utility_functions.py"
+    syntax_tree = ast.parse(utility_functions_path.read_text())
+    ti_skip = next(node for node in syntax_tree.body if isinstance(node, ast.FunctionDef) and node.name == "ti_skip")
+
+    assert any(
+        isinstance(node, ast.Import) and any(alias.name == "pytest" for alias in node.names) for node in ti_skip.body
+    ), "ti_skip must import pytest when the test helper is used"

code/models/common/tests/llm_runtime/test_config.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import FrozenInstanceError, fields
+
+import pytest
+
+import ttnn
+from models.common.llm_runtime import config as runtime_config
+from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig
+from models.common.models.llama3_8b.executor import Llama3ExecutorConfig
+
+
+class _TraceConfigSubclass(TraceConfig):
+    pass
+
+
+def _paged_config(**overrides):
+    kwargs = {
+        "block_size": 32,
+        "max_num_blocks": 1024,
+        "dtype": ttnn.bfloat8_b,
+    }
+    kwargs.update(overrides)
+    return PagedKVCacheConfig(**kwargs)
+
+
+def test_executor_config_has_exact_static_policy_owners_and_is_frozen(expect_error):
+    config = Llama3ExecutorConfig(
+        trace=TraceConfig(mode="all"),
+        warmup=WarmupConfig(),
+        paged_kv_cache=_paged_config(),
+        device_sampling_enabled=True,
+    )
+
+    assert [field.name for field in fields(config)] == [
+        "trace",
+        "warmup",
+        "paged_kv_cache",
+        "device_sampling_enabled",
+        "allow_batched_prefill_with_device_sampling_for_diagnostics",
+    ]
+    assert not config.allow_batched_prefill_with_device_sampling_for_diagnostics
+    forbidden = {
+        "model",
+        "mesh_device",
+        "hf_model",
+        "tokenizer",
+        "dtype",
+        "n_layers",
+        "sampling_config",
+        "sampling_output_dtype",
+    }
+    assert forbidden.isdisjoint(field.name for field in fields(config))
+    assert not hasattr(runtime_config, "LLMGraphCompilerConfig")
+    assert not hasattr(runtime_config, "LLMExecutorConfig")
+    assert not hasattr(runtime_config, "Sampling1DConfig")
+    with expect_error(FrozenInstanceError, ""):
+        config.device_sampling_enabled = False
+
+
+@pytest.mark.parametrize(
+    ("field_name", "invalid_value"),
+    [
+        ("trace", WarmupConfig()),
+        ("trace", _TraceConfigSubclass()),
+        ("warmup", TraceConfig()),
+        ("paged_kv_cache", WarmupConfig()),
+    ],
+)
+def test_executor_config_rejects_non_exact_nested_config_types(field_name, invalid_value, expect_error):
+    values = {
+        "trace": TraceConfig(),
+        "warmup": WarmupConfig(),
+        "paged_kv_cache": _paged_config(),
+        "device_sampling_enabled": False,
+    }
+    values[field_name] = invalid_value
+
+    with expect_error(TypeError, rf"{field_name} must be exactly"):
+        Llama3ExecutorConfig(**values)
+
+
+@pytest.mark.parametrize(
+    ("mode", "prefill", "decode"),
+    [("none", False, False), ("decode_only", False, True), ("all", True, True)],
+)
+def test_trace_config_selects_static_coverage(mode, prefill, decode, expect_error):
+    config = TraceConfig(mode=mode)
+
+    assert config.prefill_enabled is prefill
+    assert config.decode_enabled is decode
+    with expect_error(FrozenInstanceError, ""):
+        config.mode = "none"
+
+
+def test_trace_config_rejects_unknown_mode(expect_error):
+    with expect_error(ValueError, "Unsupported trace mode"):
+        TraceConfig(mode="prefill_only")
+
+
+def test_warmup_config_keeps_model_derived_defaults_and_is_deeply_immutable(expect_error):
+    config = WarmupConfig()
+
+    assert config.prefill_seq_lens is None
+    assert config.prefill_batch_sizes == (1, 2, 4, 8, 16, 32)
+    assert config.include_decode_top_k is False
+    with expect_error(TypeError, "must be a tuple"):
+        WarmupConfig(prefill_batch_sizes=[1, 2])
+
+
+def test_paged_kv_config_has_plan_fields_and_resolved_capacity(expect_error):
+    unresolved = _paged_config()
+    resolved = _paged_config(num_blocks=512)
+
+    assert [field.name for field in fields(unresolved)] == [
+        "block_size",
+        "max_num_blocks",
+        "dtype",
+        "memory_config",
+        "num_blocks",
+    ]
+    assert unresolved.memory_config == ttnn.DRAM_MEMORY_CONFIG
+    assert not unresolved.is_resolved()
+    assert resolved.is_resolved()
+    with expect_error(FrozenInstanceError, ""):
+        resolved.num_blocks = 256
+
+
+def test_paged_kv_config_rejects_invalid_capacity(expect_error):
+    with expect_error(ValueError, "exceeds max_num_blocks"):
+        _paged_config(num_blocks=1025)
+    with expect_error(ValueError, "block_size"):
+        _paged_config(block_size=0)
+
+
+def test_page_table_layout_is_resolved_without_warmup_policy():
+    layout = PageTableLayout.resolve(
+        block_size=32,
+        model_max_sequence_length=4096,
+        physical_num_blocks=100,
+        max_prefill_chunk_size=2048,
+    )
+
+    assert layout.raw_capacity_width == 100
+    assert layout.decode_width == 104
+    assert layout.prefill_width == 168

code/models/common/tests/llm_runtime/test_decode_runtime.py

@@ -0,0 +1,1337 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import dataclasses
+import inspect
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+import models.common.llm_runtime.decode as decode_module
+import ttnn
+from models.common.llm_runtime.config import PageTableLayout
+from models.common.llm_runtime.decode import (
+    DecodeDeviceInputs,
+    DecodePersistentInputs,
+    DecodeProgramSignature,
+    DecodeRuntime,
+    DecodeRuntimeConfig,
+    DecodeTraceSignature,
+    InvocationResult,
+)
+from models.common.llm_runtime.output_reader import OutputReader, PendingRead
+from models.common.modules.sampling.params import PreparedSamplingParams
+from models.common.sampling.sampling_params import SamplingParams
+
+
+class FakeMesh:
+    shape = (1, 1)
+
+
+class FakeSampling:
+    def __init__(self, seed_buffer=None):
+        self.config = SimpleNamespace(allow_force_argmax=True, max_batch_size=2, max_top_k=32, seeds=seed_buffer)
+
+    def decode_forward(
+        self,
+        logits,
+        *,
+        k=None,
+        p=None,
+        temp=None,
+        seeds=None,
+        tt_out_tok=None,
+        enable_log_probs=False,
+    ):
+        return logits, None
+
+
+class FakeRope:
+    def get_rot_idxs(self, positions, *, on_host):
+        assert on_host
+        return ("rotary", positions.clone())
+
+    def get_rot_mats(self, rotary_indices):
+        return ("cos", "sin")
+
+
+class FakeModel:
+    def __init__(self, seed_buffer=None):
+        self.config = SimpleNamespace(max_batch_size=2)
+        self.sampling = FakeSampling(seed_buffer)
+        self.rope_setup = FakeRope()
+        self.vocab_size = 8
+        self.num_devices = 1
+
+    def iter_executor_named_modules(self):
+        return iter(())
+
+    def increment_positions(self, positions, rotary_indices):
+        return None
+
+
+class FakeLazySeedBuffer:
+    def __init__(self):
+        self.source = torch.arange(2, dtype=torch.int64)
+        self._value = object()
+        self.updates = []
+
+    def update(self, source):
+        self.source = source
+        self.updates.append(source.clone())
+
+    def get_device_buffer(self):
+        return self._value
+
+
+def make_runtime(*, sampling=True, force_greedy_top_k=False, seed_buffer=None):
+    mesh = FakeMesh()
+    model = FakeModel(seed_buffer)
+    config = DecodeRuntimeConfig.resolve(
+        model=model,
+        output_reader=OutputReader(mesh),
+        lane_capacity=2,
+        page_table_layout=page_table_layout(),
+        device_sampling_enabled=sampling,
+        force_greedy_top_k=force_greedy_top_k,
+    )
+    return DecodeRuntime(config)
+
+
+def page_table_layout(*, raw_width=8, block_size=32):
+    return PageTableLayout(
+        block_size=block_size,
+        raw_capacity_width=raw_width,
+        prefill_width=((raw_width + 7) // 8) * 8,
+        decode_width=((raw_width + 7) // 8) * 8,
+    )
+
+
+def greedy_sampling():
+    return SamplingParams(temperature=[0.0, 0.0], top_k=[1, 1], top_p=[1.0, 1.0])
+
+
+def test_prepare_uses_resolved_sampler_capacity_and_neutral_inactive_rows():
+    prepared = prepare(
+        make_runtime(),
+        positions=(0, -1),
+        sampling_params=SamplingParams(temperature=1.0, top_k=32, top_p=0.08),
+    ).prepared_sampling
+
+    assert isinstance(prepared, PreparedSamplingParams)
+    assert prepared.batch_size == 2
+    assert prepared.active_rows == 1
+    assert prepared.active_mask == (True, False)
+    assert prepared.top_k == (32, 1)
+    assert prepared.top_p == pytest.approx((0.08, 0.0))
+    assert prepared.temperature == (1.0, 1.0)
+    assert prepared.row_paths == ("topk", "inactive")
+
+
+def test_prepare_accepts_vector_tensor_fields_for_full_lane():
+    prepared = prepare(
+        make_runtime(),
+        positions=(0, 0),
+        sampling_params=SamplingParams(
+            temperature=torch.ones(2),
+            top_k=torch.full((2,), 32, dtype=torch.int32),
+            top_p=torch.full((2,), 0.08),
+        ),
+    ).prepared_sampling
+
+    assert isinstance(prepared, PreparedSamplingParams)
+    assert prepared.active_mask == (True, True)
+    assert prepared.top_k == (32, 32)
+    assert prepared.top_p == pytest.approx((0.08, 0.08))
+    assert prepared.temperature == (1.0, 1.0)
+
+
+def prepare(
+    runtime,
+    *,
+    positions=(0, -1),
+    page_table=None,
+    sampling_params=None,
+    prompt_tokens=None,
+    output_tokens=None,
+    slot_remap=None,
+    reset=False,
+):
+    if page_table is None:
+        page_table = torch.tensor([[3, 4, 5], [6, 7, 8]], dtype=torch.int32)
+    return runtime.prepare(
+        torch.tensor([11, 0]),
+        torch.tensor(positions),
+        page_table,
+        sampling_params=sampling_params,
+        prompt_tokens=prompt_tokens,
+        output_tokens=output_tokens,
+        slot_remap=slot_remap,
+        reset_batch=reset,
+    )
+
+
+def test_prepare_places_only_start_pos_active_rows_and_neutralizes_gap_sentinels():
+    runtime = make_runtime()
+    sampling = SamplingParams(
+        temperature=[0.8, 0.8],
+        top_k=[999, 7],
+        top_p=[0.9, 0.8],
+        seed=[-1, -1],
+        enable_log_probs=[False, False],
+        num_logprobs=[-2, -2],
+    )
+
+    prepared = prepare(
+        runtime,
+        positions=(-1, 4),
+        sampling_params=sampling,
+        prompt_tokens=torch.tensor([[10, -1], [20, 21]]),
+        output_tokens=[[30, -1], [40, 41]],
+        slot_remap=torch.tensor([1, 0]),
+        reset=True,
+    ).prepared_sampling
+
+    assert prepared is not None
+    assert prepared.active_mask == (False, True)
+    assert prepared.row_paths == ("inactive", "topk")
+    assert prepared.top_k == (1, 7)
+    assert prepared.top_p == pytest.approx((0.0, 0.8))
+    assert prepared.seeds == (None, None)
+    assert prepared.enable_log_probs == (False, False)
+    assert prepared.num_logprobs == (0, 0)
+    assert prepared.prompt_tokens.tolist() == [[-1, -1], [20, 21]]
+    assert prepared.output_tokens == [[-1, -1], [40, 41]]
+    assert prepared.slot_remap.tolist() == [1, 0]
+
+
+def seeded_sampling(seed0, seed1=None):
+    return SamplingParams(
+        temperature=[0.8, 0.8],
+        top_k=[32, 32],
+        top_p=[0.95, 0.95],
+        seed=[seed0, seed1],
+    )
+
+
+def stochastic_sampling(seed=None):
+    return SamplingParams(
+        temperature=[0.8, 0.8],
+        top_k=[32, 32],
+        top_p=[0.95, 0.95],
+        seed=seed,
+    )
+
+
+def make_four_slot_seed_runtime():
+    seed_buffer = FakeLazySeedBuffer()
+    seed_buffer.source = torch.arange(4, dtype=torch.int64)
+    model = FakeModel(seed_buffer)
+    model.config.max_batch_size = 4
+    model.sampling.config.max_batch_size = 4
+    runtime = DecodeRuntime(
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(FakeMesh()),
+            lane_capacity=4,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+        )
+    )
+    return runtime, seed_buffer
+
+
+def four_slot_sampled_warmup(runtime):
+    return runtime.prepare(
+        torch.zeros(4, dtype=torch.long),
+        torch.zeros(4, dtype=torch.long),
+        torch.zeros((4, 8), dtype=torch.int32),
+        sampling_params=SamplingParams(
+            temperature=torch.ones(4),
+            top_k=torch.full((4,), 32, dtype=torch.int32),
+            top_p=torch.full((4,), 0.08),
+            seed=[11, 22, 33, 44],
+        ),
+    )
+
+
+def _stub_compile_only_decode(runtime, monkeypatch, run_body):
+    monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: "host")
+    monkeypatch.setattr(
+        runtime,
+        "_stage_inputs_and_kpt",
+        lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None),
+    )
+    monkeypatch.setattr(runtime, "_run_body", run_body)
+
+
+def test_compile_only_sampled_decode_temporarily_admits_and_resets_fallback_seed_slots(monkeypatch, expect_error):
+    runtime, seed_buffer = make_four_slot_seed_runtime()
+    defaults = seed_buffer.source.clone()
+    prepared = four_slot_sampled_warmup(runtime)
+    during_compile = []
+
+    def run_body(*args, **kwargs):
+        during_compile.append(runtime._seed_state.snapshot())
+        assert not kwargs["count_tokens"]
+        assert not kwargs["advance_seeds"]
+        return object()
+
+    _stub_compile_only_decode(runtime, monkeypatch, run_body)
+
+    runtime.invoke(prepared, count_tokens=False)
+
+    assert during_compile[0].active_slots == (0, 1, 2, 3)
+    assert not during_compile[0].buffer_is_default
+    reset = runtime._seed_state.snapshot()
+    assert reset.active_slots == ()
+    assert reset.buffer_is_default
+    assert torch.equal(seed_buffer.source, defaults)
+    with expect_error(RuntimeError, "reset_batch=True"):
+        runtime._refresh_sampling_seeds(prepared)
+
+
+def test_compile_only_sampled_decode_resets_fallback_seed_slots_after_failure(monkeypatch, expect_error):
+    runtime, seed_buffer = make_four_slot_seed_runtime()
+    defaults = seed_buffer.source.clone()
+    prepared = four_slot_sampled_warmup(runtime)
+
+    def run_body(*args, **kwargs):
+        assert runtime._seed_state.snapshot().active_slots == (0, 1, 2, 3)
+        raise RuntimeError("compile boom")
+
+    _stub_compile_only_decode(runtime, monkeypatch, run_body)
+    monkeypatch.setattr(runtime, "_release_or_retain_transient", lambda owned: [])
+
+    with expect_error(RuntimeError, "compile boom"):
+        runtime.invoke(prepared, count_tokens=False)
+
+    reset = runtime._seed_state.snapshot()
+    assert reset.active_slots == ()
+    assert reset.buffer_is_default
+    assert torch.equal(seed_buffer.source, defaults)
+
+
+def test_runtime_seed_same_request_and_absolute_position_are_cardinality_independent():
+    first_buffer = FakeLazySeedBuffer()
+    first = make_runtime(seed_buffer=first_buffer)
+    first_prepared = prepare(first, positions=(41, -1), sampling_params=seeded_sampling(1234), reset=True)
+    first._refresh_sampling_seeds(first_prepared)
+
+    remapped_buffer = FakeLazySeedBuffer()
+    remapped = make_runtime(seed_buffer=remapped_buffer)
+    remapped_prepared = prepare(
+        remapped,
+        positions=(-1, 41),
+        sampling_params=seeded_sampling(None, 1234),
+        reset=True,
+    )
+    remapped._refresh_sampling_seeds(remapped_prepared)
+
+    assert int(first_buffer.updates[-1][0]) == int(remapped_buffer.updates[-1][1])
+    assert first._seed_state.snapshot().active == (True, False)
+    assert remapped._seed_state.snapshot().active == (False, True)
+
+
+@pytest.mark.parametrize("seed", [1234, torch.tensor(1234)])
+def test_runtime_scalar_seed_belongs_to_one_request_and_is_not_broadcast(seed):
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    sampling_params = SamplingParams(
+        temperature=[0.8, 0.8],
+        top_k=[32, 32],
+        top_p=[0.95, 0.95],
+        seed=seed,
+    )
+
+    prepared = prepare(runtime, positions=(17, 17), sampling_params=sampling_params, reset=True)
+    runtime._refresh_sampling_seeds(prepared)
+
+    assert prepared.prepared_sampling.seeds == (1234, None)
+    snapshot = runtime._seed_state.snapshot()
+    assert snapshot.request_seeds == (1234, None)
+    assert snapshot.active == (True, True)
+    assert snapshot.current_device_seeds[0] is not None
+    assert snapshot.current_device_seeds[1] is not None
+
+
+@pytest.mark.parametrize("seed", [[111, 222], torch.tensor([111, 222])])
+def test_runtime_vector_seed_remains_slot_indexed(seed):
+    runtime = make_runtime(seed_buffer=FakeLazySeedBuffer())
+    sampling_params = SamplingParams(
+        temperature=[0.8, 0.8],
+        top_k=[32, 32],
+        top_p=[0.95, 0.95],
+        seed=seed,
+    )
+
+    prepared = prepare(runtime, positions=(5, 5), sampling_params=sampling_params, reset=True)
+
+    assert prepared.prepared_sampling.seeds == (111, 222)
+
+
+def test_runtime_simultaneous_equal_request_seeds_share_vllm_stream():
+    runtime = make_runtime(seed_buffer=FakeLazySeedBuffer())
+    prepared = prepare(
+        runtime,
+        positions=(5, 5),
+        sampling_params=seeded_sampling(77, 77),
+        reset=True,
+    )
+
+    runtime._refresh_sampling_seeds(prepared)
+
+    snapshot = runtime._seed_state.snapshot()
+    assert runtime._seed_manager.salt_duplicate_seeds is False
+    assert snapshot.request_seeds == (77, 77)
+    assert snapshot.salts == (0, 0)
+    assert snapshot.current_device_seeds[0] == snapshot.current_device_seeds[1]
+
+
+def test_runtime_vllm_uniform_seed_is_deterministic_across_concurrent_slots():
+    seed_buffer = FakeLazySeedBuffer()
+    seed_buffer.source = torch.arange(32, dtype=torch.int64)
+    model = FakeModel(seed_buffer)
+    model.config.max_batch_size = 32
+    model.sampling.config.max_batch_size = 32
+    runtime = DecodeRuntime(
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(FakeMesh()),
+            lane_capacity=32,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+        )
+    )
+    prepared = runtime.prepare(
+        torch.zeros(32, dtype=torch.long),
+        torch.zeros(32, dtype=torch.long),
+        torch.zeros((32, 8), dtype=torch.int32),
+        sampling_params=SamplingParams(
+            temperature=torch.ones(32),
+            top_k=torch.full((32,), 32, dtype=torch.int32),
+            top_p=torch.full((32,), 0.95),
+            seed=[1234] * 32,
+        ),
+        reset_batch=True,
+    )
+
+    runtime._refresh_sampling_seeds(prepared)
+
+    snapshot = runtime._seed_state.snapshot()
+    assert runtime._seed_manager.salt_duplicate_seeds is False
+    assert snapshot.salts == (0,) * 32
+    assert len(set(snapshot.current_device_seeds)) == 1
+
+
+def test_runtime_slot_remap_moves_complete_seed_stream_before_refresh():
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    initial = prepare(
+        runtime,
+        positions=(5, -1),
+        sampling_params=seeded_sampling(77),
+        reset=True,
+    )
+    runtime._refresh_sampling_seeds(initial)
+    original_device_seed = int(seed_buffer.updates[-1][0])
+    original_state = runtime._seed_state.snapshot()
+
+    moved = prepare(
+        runtime,
+        positions=(-1, 5),
+        sampling_params=seeded_sampling(None, 77),
+        slot_remap=torch.tensor([0, 0]),
+        reset=False,
+    )
+    runtime._refresh_sampling_seeds(moved)
+
+    state = runtime._seed_state.snapshot()
+    assert state.active == (False, True)
+    assert state.request_seeds == (None, 77)
+    assert state.token_counters[1] == original_state.token_counters[0]
+    assert int(seed_buffer.updates[-1][1]) == original_device_seed
+
+
+def test_runtime_seed_changes_with_request_seed_and_decode_position():
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+
+    runtime._refresh_sampling_seeds(
+        prepare(runtime, positions=(7, -1), sampling_params=seeded_sampling(101), reset=True)
+    )
+    position_7 = int(seed_buffer.updates[-1][0])
+    runtime._refresh_sampling_seeds(
+        prepare(runtime, positions=(8, -1), sampling_params=seeded_sampling(101), reset=False)
+    )
+    position_8 = int(seed_buffer.updates[-1][0])
+    runtime._refresh_sampling_seeds(
+        prepare(runtime, positions=(7, -1), sampling_params=seeded_sampling(202), reset=True)
+    )
+    different_request = int(seed_buffer.updates[-1][0])
+
+    assert len({position_7, position_8, different_request}) == 3
+
+
+def test_runtime_explicit_seed_absolute_position_is_stable_across_reset_boundaries():
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    request = seeded_sampling(909)
+
+    runtime._refresh_sampling_seeds(prepare(runtime, positions=(19, -1), sampling_params=request, reset=True))
+    original = seed_buffer.updates[-1].clone()
+    runtime._refresh_sampling_seeds(prepare(runtime, positions=(20, -1), sampling_params=request, reset=False))
+    continued = seed_buffer.updates[-1].clone()
+    runtime._refresh_sampling_seeds(prepare(runtime, positions=(19, -1), sampling_params=request, reset=True))
+    restarted = seed_buffer.updates[-1].clone()
+    runtime._refresh_sampling_seeds(prepare(runtime, positions=(20, -1), sampling_params=request, reset=True))
+    resumed = seed_buffer.updates[-1].clone()
+
+    assert torch.equal(original, restarted)
+    assert torch.equal(continued, resumed)
+
+
+def test_runtime_seed_refreshes_before_eager_model_invocation(monkeypatch):
+    events = []
+    seed_buffer = FakeLazySeedBuffer()
+    original_update = seed_buffer.update
+
+    def record_update(source):
+        events.append("seed")
+        original_update(source)
+
+    seed_buffer.update = record_update
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    prepared = prepare(runtime, positions=(3, -1), sampling_params=seeded_sampling(77), reset=True)
+    monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: object())
+    monkeypatch.setattr(
+        runtime,
+        "_stage_inputs_and_kpt",
+        lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None),
+    )
+
+    def run_body(*args, **kwargs):
+        events.append("invoke")
+        return object()
+
+    monkeypatch.setattr(runtime, "_run_body", run_body)
+
+    runtime.invoke(prepared)
+
+    assert events[-1] == "invoke"
+    assert events[:-1]
+    assert set(events[:-1]) == {"seed"}
+
+
+def test_runtime_trace_captures_stable_seed_handle_and_refreshes_before_replay(monkeypatch):
+    events = []
+    seed_buffer = FakeLazySeedBuffer()
+    original_update = seed_buffer.update
+
+    def record_update(source):
+        events.append("seed")
+        original_update(source)
+
+    seed_buffer.update = record_update
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    prepared = prepare(runtime, positions=(12, -1), sampling_params=seeded_sampling(88), reset=True)
+    monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: object())
+    monkeypatch.setattr(
+        runtime,
+        "_stage_inputs_and_kpt",
+        lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None),
+    )
+
+    persistent = runtime.capture_plan(prepared).prepare_inputs()
+    assert persistent.seed_buffer is seed_buffer.get_device_buffer()
+    sampling = prepared.prepared_sampling
+    assert sampling is not None
+    persistent = dataclasses.replace(
+        persistent,
+        kpt_signature=[(sampling.top_k, sampling.top_p, sampling.temperature)],
+    )
+    runtime.refresh_trace(persistent, prepared, SimpleNamespace(full=False, page_table=False))
+    events.append("replay")
+
+    assert events[-1] == "replay"
+    assert events[:-1]
+    assert set(events[:-1]) == {"seed"}
+
+
+def test_runtime_seed_handling_does_not_mutate_sampling_params():
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    sampling_params = seeded_sampling(11, 22)
+    before = dataclasses.asdict(sampling_params)
+
+    prepared = prepare(runtime, positions=(4, 9), sampling_params=sampling_params, reset=True)
+    runtime._refresh_sampling_seeds(prepared)
+
+    assert dataclasses.asdict(sampling_params) == before
+
+
+def test_runtime_unseeded_stream_varies_and_same_absolute_position_is_idempotent():
+    seed_buffer = FakeLazySeedBuffer()
+    defaults = seed_buffer.source.clone()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+
+    unseeded = prepare(runtime, positions=(5, -1), sampling_params=stochastic_sampling(), reset=True)
+    runtime._refresh_sampling_seeds(unseeded)
+    first = seed_buffer.updates[-1].clone()
+    first_state = runtime._seed_state.snapshot()
+    assert int(first[0]) != int(defaults[0])
+    assert int(first[1]) == int(defaults[1])
+
+    runtime._refresh_sampling_seeds(dataclasses.replace(unseeded, reset_batch=False))
+    repeated = seed_buffer.updates[-1].clone()
+    assert torch.equal(repeated, first)
+
+    advanced = prepare(
+        runtime,
+        positions=(6, -1),
+        sampling_params=stochastic_sampling(),
+        reset=False,
+    )
+    runtime._refresh_sampling_seeds(advanced)
+    snapshot = runtime._seed_state.snapshot()
+    assert snapshot.request_seeds == (None, None)
+    assert snapshot.active == (True, False)
+    assert snapshot.token_counters[0] == 2
+    assert snapshot.unseeded_rng_states[0] != first_state.unseeded_rng_states[0]
+
+
+def test_runtime_seed_change_requires_reset_and_preserves_unseeded_survivor(expect_error):
+    seed_buffer = FakeLazySeedBuffer()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    unseeded = prepare(
+        runtime,
+        positions=(10, 10),
+        sampling_params=stochastic_sampling(seed=[None, None]),
+        reset=True,
+    )
+    runtime._refresh_sampling_seeds(unseeded)
+    initial = runtime._seed_state.snapshot()
+
+    mixed = prepare(
+        runtime,
+        positions=(11, 11),
+        sampling_params=stochastic_sampling(seed=[None, 42]),
+        reset=False,
+    )
+    with expect_error(RuntimeError, "reset_batch=True"):
+        runtime._refresh_sampling_seeds(mixed)
+
+    runtime._refresh_sampling_seeds(dataclasses.replace(mixed, reset_batch=True))
+    admitted = runtime._seed_state.snapshot()
+    assert admitted.request_seeds == (None, 42)
+    assert admitted.active == (True, True)
+    assert admitted.token_counters[0] == initial.token_counters[0] + 1
+
+    continued = dataclasses.replace(mixed, start_pos=torch.tensor([12, 12]))
+    runtime._refresh_sampling_seeds(continued)
+    continued_state = runtime._seed_state.snapshot()
+    assert continued_state.token_counters[0] == admitted.token_counters[0] + 1
+    assert continued_state.request_seeds == (None, 42)
+
+
+def test_runtime_inactive_peer_is_cleaned_up_without_readmitting_survivor():
+    seed_buffer = FakeLazySeedBuffer()
+    defaults = seed_buffer.source.clone()
+    runtime = make_runtime(seed_buffer=seed_buffer)
+    initial = prepare(
+        runtime,
+        positions=(20, 20),
+        sampling_params=stochastic_sampling(seed=[None, 42]),
+        reset=True,
+    )
+    runtime._refresh_sampling_seeds(initial)
+    initial_state = runtime._seed_state.snapshot()
+
+    peer_left = prepare(
+        runtime,
+        positions=(21, -1),
+        sampling_params=stochastic_sampling(seed=[None, 42]),
+        reset=False,
+    )
+    runtime._refresh_sampling_seeds(peer_left)
+
+    state = runtime._seed_state.snapshot()
+    assert state.active == (True, False)
+    assert state.request_seeds == (None, None)
+    assert state.token_counters[0] == initial_state.token_counters[0] + 1
+    assert int(seed_buffer.updates[-1][1]) == int(defaults[1])
+
+
+@pytest.mark.parametrize(
+    ("method", "expected"),
+    (
+        (
+            DecodeRuntime.prepare,
+            (
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("start_pos", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("prompt_tokens", inspect.Parameter.KEYWORD_ONLY, None),
+                ("output_tokens", inspect.Parameter.KEYWORD_ONLY, None),
+                ("slot_remap", inspect.Parameter.KEYWORD_ONLY, None),
+                ("reset_batch", inspect.Parameter.KEYWORD_ONLY, False),
+            ),
+        ),
+        (
+            DecodeRuntime.read_decode_output,
+            (
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("async_read", inspect.Parameter.KEYWORD_ONLY, False),
+            ),
+        ),
+        (
+            DecodeRuntime.process_decode_output_host,
+            (
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("is_tokens", inspect.Parameter.KEYWORD_ONLY, False),
+            ),
+        ),
+    ),
+)
+def test_runtime_api_signatures_are_exact(method, expected):
+    parameters = inspect.signature(method).parameters.values()
+
+    assert tuple((parameter.name, parameter.kind, parameter.default) for parameter in parameters) == expected
+
+
+def test_runtime_api_preserves_positional_prefixes_and_rejects_extra_fields(expect_error):
+    cases = (
+        (DecodeRuntime.prepare, ("tokens", "start_pos", "page_table"), "sampling_params"),
+        (DecodeRuntime.read_decode_output, ("tt_out",), "async_read"),
+        (DecodeRuntime.process_decode_output_host, ("tt_out",), "is_tokens"),
+    )
+
+    for method, positional_prefix, keyword_name in cases:
+        signature = inspect.signature(method)
+        signature.bind(None, *positional_prefix, **{keyword_name: True})
+        with expect_error(TypeError, "too many positional arguments"):
+            signature.bind(None, *positional_prefix, True)
+        with expect_error(TypeError, "unexpected keyword argument 'unknown'"):
+            signature.bind(None, *positional_prefix, unknown=True)
+
+    assert inspect.get_annotations(DecodeRuntime.prepare, eval_str=True)["sampling_params"] is Any
+
+
+def test_config_resolves_canonical_static_capabilities_and_is_frozen(expect_error):
+    mesh = FakeMesh()
+    model = FakeModel()
+    config = DecodeRuntimeConfig.resolve(
+        model=model,
+        output_reader=OutputReader(mesh),
+        lane_capacity=2,
+        page_table_layout=page_table_layout(),
+        device_sampling_enabled=True,
+    )
+
+    assert config.cluster_shape == (1, 1)
+    assert config.num_devices == 1
+    assert config.vocab_size == 8
+    assert config.allow_force_argmax
+    assert config.max_device_top_k == 32
+    assert config.sampling_batch_size == 2
+    assert config.sampling_state_controller is None
+    assert config.sampling_state is None
+    assert config.position_feedback_capable
+    with expect_error(dataclasses.FrozenInstanceError, "cannot assign to field"):
+        config.lane_capacity = 1
+    with expect_error(TypeError, "DecodeRuntimeConfig"):
+        DecodeRuntime(config=None)
+
+
+def test_config_rejects_inconsistent_collaborators_and_dimensions(expect_error):
+    mesh = FakeMesh()
+    model = FakeModel()
+    model.config.mesh_device = mesh
+    with expect_error(ValueError, "same mesh_device"):
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(FakeMesh()),
+            lane_capacity=2,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+        )
+    with expect_error(ValueError, "positive integer"):
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(mesh),
+            lane_capacity=True,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+        )
+    with expect_error(TypeError, "PageTableLayout"):
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(mesh),
+            lane_capacity=2,
+            page_table_layout=SimpleNamespace(raw_capacity_width=8, decode_width=8, block_size=32),
+            device_sampling_enabled=True,
+        )
+
+
+def test_layout_replacement_is_immutable_and_bounded(expect_error):
+    runtime = make_runtime()
+    original = runtime.config
+    replacement = page_table_layout(raw_width=4)
+
+    runtime.configure_page_table_layout(replacement)
+
+    assert runtime.config is not original
+    assert original.page_table_layout.raw_capacity_width == 8
+    assert runtime.config.page_table_layout is replacement
+    with expect_error(ValueError, "block_size"):
+        runtime.configure_page_table_layout(page_table_layout(raw_width=4, block_size=16))
+    with expect_error(ValueError, "ceiling"):
+        runtime.configure_page_table_layout(page_table_layout(raw_width=9))
+    with expect_error(ValueError, "decode width"):
+        runtime.configure_page_table_layout(PageTableLayout(32, 4, 8, 1024))
+
+
+def test_sampling_admission_follows_resolved_configuration(expect_error):
+    runtime = make_runtime(sampling=False)
+    with expect_error(ValueError, "device sampling is disabled"):
+        prepare(runtime, sampling_params=greedy_sampling())
+
+
+def test_feedback_and_sampling_path_follow_resolved_capabilities():
+    argmax_runtime = make_runtime()
+    argmax_prepared = prepare(argmax_runtime, sampling_params=greedy_sampling())
+    assert argmax_prepared.device_feedback
+    assert argmax_prepared.sampling_path == "argmax"
+
+    model = FakeModel()
+    model.increment_positions = None
+    mesh = FakeMesh()
+    no_feedback = DecodeRuntime(
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(mesh),
+            lane_capacity=2,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+            force_greedy_top_k=True,
+        )
+    )
+    prepared = prepare(no_feedback, sampling_params=greedy_sampling())
+
+    assert not prepared.device_feedback
+    assert prepared.sampling_path == "topk"
+
+
+def test_single_and_multi_device_logits_conversion(monkeypatch):
+    single = make_runtime()
+    logits = torch.arange(16, dtype=torch.float32).reshape(1, 1, 2, 8)
+    monkeypatch.setattr(ttnn, "to_torch", lambda value: logits)
+    converted, _ = single._normalize_host_output("single-device", is_tokens=False)
+    assert converted.shape == (2, 1, 8)
+
+    mesh = SimpleNamespace(shape=(1, 2))
+    model = FakeModel()
+    model.num_devices = 2
+    multi = DecodeRuntime(
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(mesh),
+            lane_capacity=2,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+        )
+    )
+    calls = []
+    monkeypatch.setattr(
+        decode_module,
+        "_concat_host_output",
+        lambda value, shape: calls.append((value, shape)) or logits,
+    )
+    converted, _ = multi._normalize_host_output("multi-device", is_tokens=False)
+    assert converted.shape == (2, 1, 8)
+    assert calls == [("multi-device", (1, 2))]
+
+
+def test_signatures_expose_ordered_material_and_separate_types():
+    runtime = make_runtime()
+    prepared = prepare(runtime, sampling_params=greedy_sampling())
+
+    program = runtime.program_signature(prepared)
+    trace = runtime.trace_signature(prepared)
+
+    assert isinstance(program, DecodeProgramSignature)
+    assert isinstance(trace, DecodeTraceSignature)
+    assert program.key_material() == (
+        ("operation", "decode"),
+        ("batch_size", 2),
+        ("page_table_width", 8),
+        ("sampling_path", "argmax"),
+        ("device_feedback", True),
+    )
+    assert trace.key_material() == program.key_material()
+    assert runtime.program_signature(prepare(runtime)).sampling_path == "logits"
+
+
+def test_signature_tracks_native_penalty_and_sampled_logprob_program_modes():
+    runtime = make_runtime()
+    sampling = SamplingParams(
+        temperature=[0.0, 0.0],
+        top_k=[1, 1],
+        top_p=[1.0, 1.0],
+        presence_penalty=[0.5, 0.0],
+        enable_log_probs=[True, False],
+        num_logprobs=[0, -2],
+    )
+
+    prepared = prepare(runtime, sampling_params=sampling)
+    native = prepared.prepared_sampling
+    signature = runtime.program_signature(prepared)
+
+    assert native is not None
+    assert native.logprob_modes == ("sampled_token", "none")
+    assert native.penalties_enabled
+    assert native.log_probs_enabled
+    assert prepared.sampling_path == "topk"
+    assert signature.penalties_enabled
+    assert signature.logprobs_enabled
+    assert signature.key_material()[-2:] == (
+        ("penalties_enabled", True),
+        ("logprobs_enabled", True),
+    )
+
+
+def test_configured_topk_policy_is_not_collapsed_to_argmax_by_greedy_temperature():
+    runtime = make_runtime(force_greedy_top_k=True)
+    sampling = SamplingParams(temperature=[0.0, 0.0], top_k=[32, 32], top_p=[0.08, 0.08])
+
+    prepared = prepare(runtime, sampling_params=sampling)
+
+    assert prepared.sampling_path == "topk"
+
+
+def test_unconfigured_topk_values_use_argmax_for_greedy_temperature():
+    runtime = make_runtime()
+    sampling = SamplingParams(temperature=[0.0, 0.0], top_k=[32, 32], top_p=[0.08, 0.08])
+
+    assert prepare(runtime, sampling_params=sampling).sampling_path == "argmax"
+
+
+def test_sampling_params_are_prepared_once_and_reused_for_kpt(monkeypatch):
+    runtime = make_runtime(force_greedy_top_k=True)
+    calls = []
+    formatter = decode_module.prepare_sampling_params
+
+    def formatter_spy(*args, **kwargs):
+        calls.append((args, kwargs))
+        return formatter(*args, **kwargs)
+
+    monkeypatch.setattr(
+        decode_module,
+        "prepare_sampling_params",
+        formatter_spy,
+    )
+    prepared = prepare(runtime, sampling_params=greedy_sampling())
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: "mapper")
+    monkeypatch.setattr(ttnn, "from_torch", lambda value, **kwargs: value)
+
+    runtime._make_host_kpt(prepared)
+    assert len(calls) == 1
+
+
+def test_tile_padded_sampler_preserves_one_semantic_lane_and_neutralizes_inactive_rows(monkeypatch):
+    mesh = FakeMesh()
+    model = FakeModel()
+    model.config.max_batch_size = 1
+    model.sampling.config.max_batch_size = 32
+    runtime = DecodeRuntime(
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=OutputReader(mesh),
+            lane_capacity=1,
+            page_table_layout=page_table_layout(),
+            device_sampling_enabled=True,
+            force_greedy_top_k=True,
+        )
+    )
+    prepared = runtime.prepare(
+        torch.tensor([11]),
+        torch.tensor([0]),
+        torch.tensor([[3, 4, 5]], dtype=torch.int32),
+        sampling_params=SamplingParams(
+            temperature=0.7,
+            top_k=32,
+            top_p=0.08,
+            presence_penalty=0.25,
+            frequency_penalty=0.5,
+            repetition_penalty=1.2,
+            seed=17,
+            enable_log_probs=True,
+            num_logprobs=1,
+        ),
+    )
+    sampling = prepared.prepared_sampling
+
+    assert runtime.config.lane_capacity == 1
+    assert runtime.config.sampling_batch_size == 32
+    assert sampling.batch_size == 32
+    assert sampling.active_rows == 1
+    assert sampling.active_mask == (True,) + (False,) * 31
+    assert sampling.top_k == (32,) + (1,) * 31
+    assert sampling.top_p == pytest.approx((0.08,) + (0.0,) * 31)
+    assert sampling.temperature == pytest.approx((1.0 / 0.7,) + (1.0,) * 31)
+    assert sampling.seeds == (17,) + (None,) * 31
+    assert sampling.presence_penalty == pytest.approx((0.25,) + (0.0,) * 31)
+    assert sampling.frequency_penalty == pytest.approx((0.5,) + (0.0,) * 31)
+    assert sampling.repetition_penalty == pytest.approx((1.2,) + (1.0,) * 31)
+    assert sampling.enable_log_probs == (True,) + (False,) * 31
+    assert sampling.num_logprobs == (1,) + (0,) * 31
+    assert sampling.logprob_modes[1:] == ("none",) * 31
+
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh_device: "mapper")
+    monkeypatch.setattr(ttnn, "from_torch", lambda value, **kwargs: value)
+    k, p, temperature = runtime._make_host_kpt(prepared)
+
+    assert tuple(k.shape) == tuple(p.shape) == tuple(temperature.shape) == (32,)
+    assert k.tolist() == [32] + [1] * 31
+    assert p.tolist() == pytest.approx([0.08] + [0.0] * 31)
+    assert temperature.tolist() == pytest.approx([1.0 / 0.7] + [1.0] * 31)
+
+
+def test_normalization_preserves_feedback_lookahead_and_inactive_convention():
+    runtime = make_runtime()
+    page_table = torch.tensor([[10, 11, 99], [20, 21, 98]], dtype=torch.int64)
+
+    prepared = prepare(
+        runtime,
+        positions=(31, -1),
+        page_table=page_table,
+        sampling_params=greedy_sampling(),
+    )
+
+    assert prepared.page_table.dtype == torch.int32
+    assert prepared.page_table.shape == (2, 8)
+    assert prepared.page_table[0].tolist() == [10, 11, 0, 0, 0, 0, 0, 0]
+    assert prepared.page_table[1].tolist() == [0, 0, 0, 0, 0, 0, 0, 0]
+
+
+def test_normalization_reuses_equal_source_with_same_copy_counts():
+    runtime = make_runtime()
+    first = prepare(runtime, positions=(0, -1))
+    second = prepare(runtime, positions=(1, -1))
+
+    assert second.page_table is first.page_table
+
+
+def test_normalization_cache_detects_in_place_source_mutation():
+    runtime = make_runtime()
+    source = torch.tensor([[3, 4], [6, 7]], dtype=torch.int32)
+    first = prepare(runtime, positions=(0, -1), page_table=source)
+    source[0, 0] = 9
+    second = prepare(runtime, positions=(0, -1), page_table=source)
+
+    assert second.page_table is not first.page_table
+    assert second.page_table[0, 0].item() == 9
+
+
+def test_normalization_cache_misses_when_copy_counts_or_feedback_change():
+    runtime = make_runtime()
+    source = torch.tensor([[3, 4], [6, 7]], dtype=torch.int32)
+    one_block = prepare(runtime, positions=(0, -1), page_table=source)
+    two_blocks = prepare(runtime, positions=(32, -1), page_table=source)
+    no_feedback = prepare(runtime, positions=(31, -1), page_table=source)
+    with_feedback = prepare(
+        runtime,
+        positions=(31, -1),
+        page_table=source,
+        sampling_params=greedy_sampling(),
+    )
+
+    assert two_blocks.page_table is not one_block.page_table
+    assert no_feedback.page_table[0, 1].item() == 0
+    assert with_feedback.page_table is not no_feedback.page_table
+    assert with_feedback.page_table[0, 1].item() == 4
+
+
+def test_fixed_capacity_and_page_table_capacity_are_validated(expect_error):
+    runtime = make_runtime()
+    with expect_error(ValueError, "must equal lane capacity"):
+        runtime.prepare(torch.tensor([1]), torch.tensor([0]), torch.tensor([[1]]))
+    with expect_error(ValueError, "batches must match"):
+        runtime.prepare(
+            torch.tensor([1, 2]),
+            torch.tensor([0]),
+            torch.tensor([[1], [2]]),
+        )
+    with expect_error(ValueError, "paged-KV capacity"):
+        prepare(runtime, positions=(8 * 32, -1))
+    with expect_error(ValueError, "too narrow"):
+        prepare(
+            runtime,
+            positions=(64, -1),
+            page_table=torch.tensor([[1, 2], [0, 0]], dtype=torch.int32),
+        )
+
+
+def test_preparation_tracks_first_used_page_change_reset_and_ignores_unused_tail():
+    runtime = make_runtime()
+    first = prepare(runtime, positions=(0, -1), reset=True)
+    assert first.page_table_changed
+    assert first.reset_batch
+
+    runtime.note_submitted(first)
+    same_semantics = prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[3, 90, 91], [88, 87, 86]], dtype=torch.int32),
+    )
+    assert not same_semantics.page_table_changed
+    assert not same_semantics.reset_batch
+
+    changed = prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[4, 90, 91], [88, 87, 86]], dtype=torch.int32),
+    )
+    assert changed.page_table_changed
+
+
+def test_submission_state_tracks_last_table_despite_stale_prepare_change_hint():
+    runtime = make_runtime()
+    baseline = prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[3], [0]], dtype=torch.int32),
+    )
+    runtime.note_submitted(baseline)
+
+    changed = prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[4], [0]], dtype=torch.int32),
+    )
+    back_to_baseline = prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[3], [0]], dtype=torch.int32),
+    )
+    assert changed.page_table_changed
+    assert not back_to_baseline.page_table_changed
+
+    runtime.note_submitted(changed)
+    runtime.note_submitted(back_to_baseline)
+
+    assert not prepare(runtime, page_table=back_to_baseline.page_table).page_table_changed
+    assert prepare(
+        runtime,
+        positions=(0, -1),
+        page_table=torch.tensor([[4], [0]], dtype=torch.int32),
+    ).page_table_changed
+
+
+def test_capture_plan_describes_full_step_refresh_and_typed_persistent_inputs(monkeypatch):
+    runtime = make_runtime()
+    prepared = prepare(runtime, sampling_params=greedy_sampling())
+    device = DecodeDeviceInputs("tokens", "positions", "rotary", "page_table")
+    monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda request: "host")
+    monkeypatch.setattr(runtime, "_stage_inputs_and_kpt", lambda host, request: (device, "kpt"))
+
+    def run_body(
+        inputs,
+        prepared,
+        kpt,
+        *,
+        device_feedback,
+        count_tokens=True,
+        advance_seeds=True,
+    ):
+        return "captured"
+
+    monkeypatch.setattr(runtime, "_run_body", run_body)
+
+    plan = runtime.capture_plan(prepared)
+    persistent = plan.prepare_inputs()
+
+    assert persistent.device_inputs is device
+    assert persistent.kpt == "kpt"
+    sampling = prepared.prepared_sampling
+    assert sampling is not None
+    assert persistent.kpt_signature == [(sampling.top_k, sampling.top_p, sampling.temperature)]
+    assert plan.capture(persistent) == "captured"
+    assert plan.refresh_policy.every_replay == ("sampling",)
+    assert plan.refresh_policy.full_on_batch_reset
+    assert plan.refresh_policy.full_on_graph_switch
+    assert plan.refresh_policy.full_without_device_feedback
+    assert plan.refresh_policy.refresh_page_table_on_change
+
+
+def test_trace_refresh_skips_unchanged_sampling_values(monkeypatch):
+    runtime = make_runtime(force_greedy_top_k=True)
+    prepared = prepare(runtime, sampling_params=greedy_sampling())
+    sampling = prepared.prepared_sampling
+    assert sampling is not None
+    persistent = DecodePersistentInputs(
+        device_inputs=DecodeDeviceInputs("tokens", "positions", "rotary", "page_table"),
+        kpt="kpt",
+        kpt_signature=[(sampling.top_k, sampling.top_p, sampling.temperature)],
+    )
+
+    def fail_refresh_kpt(device_kpt, prepared):
+        pytest.fail("unchanged KPT was refreshed")
+
+    monkeypatch.setattr(runtime, "_refresh_kpt", fail_refresh_kpt)
+
+    runtime.refresh_trace(
+        persistent,
+        prepared,
+        SimpleNamespace(full=False, page_table=False),
+    )
+
+
+def test_eager_invoke_returns_owned_result_and_advances_submission_state(monkeypatch):
+    runtime = make_runtime()
+    prepared = prepare(runtime)
+    device = DecodeDeviceInputs("tokens", "positions", "rotary", "page_table")
+    calls = []
+    monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda request: "host")
+    monkeypatch.setattr(runtime, "_stage_inputs_and_kpt", lambda host, request: (device, None))
+    monkeypatch.setattr(
+        runtime,
+        "_run_body",
+        lambda inputs, prepared, kpt, *, device_feedback, **kwargs: calls.append(device_feedback) or ("raw", None),
+    )
+
+    result = runtime.invoke(prepared)
+
+    assert isinstance(result, InvocationResult)
+    assert result.value == ("raw", None)
+    assert result.owned == (("raw", None), (device, None))
+    assert not result.is_tokens
+    assert calls == [False]
+    assert not prepare(runtime, page_table=prepared.page_table).page_table_changed
+
+
+def test_blocking_consume_normalizes_logits_and_releases_owned_values(monkeypatch):
+    runtime = make_runtime()
+    host_logits = torch.arange(16, dtype=torch.float32).reshape(1, 1, 2, 8)
+    released = []
+    monkeypatch.setattr(runtime.config.output_reader, "read", lambda value, *, blocking: (host_logits, "probs"))
+    monkeypatch.setattr(runtime, "_release_or_retain_transient", lambda value: released.append(value) or [])
+    result = InvocationResult(value="raw", owned="owned", is_tokens=False)
+
+    logits, log_probs = runtime.consume(result)
+
+    assert logits.shape == (2, 1, 8)
+    assert log_probs == "probs"
+    assert released == ["owned"]
+
+
+def test_sampled_token_logprobs_are_flattened_to_lane_row_order():
+    runtime = make_runtime()
+    host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32)
+    host_log_probs = torch.tensor([[[[-0.25, -0.75]]]], dtype=torch.bfloat16)
+
+    tokens, log_probs = runtime._normalize_host_output(
+        (host_tokens, host_log_probs),
+        is_tokens=True,
+    )
+
+    assert tokens.tolist() == [7, 8]
+    assert tokens.dtype == torch.int64
+    assert log_probs.tolist() == pytest.approx([-0.25, -0.75])
+    assert log_probs.dtype == torch.float32
+
+
+def test_raw_blocking_and_async_leases_release_exact_records(monkeypatch):
+    runtime = make_runtime()
+    deallocated = []
+    monkeypatch.setattr(
+        decode_module,
+        "best_effort_deallocate_owned_tensors",
+        lambda values, completed: deallocated.append(values) or [],
+    )
+
+    first = InvocationResult(value=object(), owned="first-owned", is_tokens=False)
+    assert runtime.consume(first, read_from_device=False) is first.value
+    monkeypatch.setattr(runtime.config.output_reader, "read", lambda value, *, blocking: "first-host")
+    assert runtime.read_decode_output(first.value) == "first-host"
+
+    second = InvocationResult(value=object(), owned="second-owned", is_tokens=True)
+    runtime.consume(second, read_from_device=False)
+    host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32)
+    pending = PendingRead(value=(host_tokens, None), events=("event",), sequence=4, _owner=object())
+    monkeypatch.setattr(runtime.config.output_reader, "submit", lambda value: pending)
+    monkeypatch.setattr(runtime.config.output_reader, "complete", lambda value: pending.value)
+
+    host, events = runtime.read_decode_output(second.value, async_read=True)
+    assert host is pending.value
+    assert events == ["event"]
+    tokens, log_probs = runtime.process_decode_output_host(host, is_tokens=True)
+    assert tokens.tolist() == [7, 8]
+    assert tokens.dtype == torch.int64
+    assert log_probs is None
+    assert deallocated == [
+        (first.value, "first-owned"),
+        (second.value, "second-owned"),
+    ]
+
+
+def test_async_trace_lease_never_releases_borrowed_trace_output(monkeypatch):
+    runtime = make_runtime()
+    deallocated = []
+    raw = object()
+    host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32)
+    pending = PendingRead(value=(host_tokens, None), events=("event",), sequence=4, _owner=object())
+    monkeypatch.setattr(
+        decode_module,
+        "best_effort_deallocate_owned_tensors",
+        lambda values, completed: deallocated.append(values) or [],
+    )
+    monkeypatch.setattr(runtime.config.output_reader, "submit", lambda value: pending)
+    monkeypatch.setattr(runtime.config.output_reader, "complete", lambda value: pending.value)
+
+    result = InvocationResult(value=raw, owned=None, is_tokens=True)
+    assert runtime.consume(result, read_from_device=False) is raw
+    host, events = runtime.read_decode_output(raw, async_read=True)
+    assert events == ["event"]
+    tokens, log_probs = runtime.process_decode_output_host(host, is_tokens=True)
+
+    assert tokens.tolist() == [7, 8]
+    assert log_probs is None
+    assert deallocated == []
+
+
+def test_failed_transient_release_blocks_use_and_cleanup_retries(monkeypatch, expect_error):
+    runtime = make_runtime()
+
+    class FakeTensor:
+        pass
+
+    tensor = FakeTensor()
+    attempts = []
+    monkeypatch.setattr(decode_module.ttnn, "Tensor", FakeTensor)
+
+    def deallocate(value):
+        attempts.append(value)
+        if len(attempts) == 1:
+            raise RuntimeError("release failed")
+
+    monkeypatch.setattr(decode_module.ttnn, "deallocate", deallocate)
+
+    failures = runtime._release_or_retain_transient(tensor)
+    assert [str(error) for error in failures] == ["release failed"]
+    assert runtime.transient_orphan_count == 1
+    with expect_error(RuntimeError, "unreleased transient"):
+        prepare(runtime)
+
+    runtime.cleanup_transients()
+    assert attempts == [tensor, tensor]
+    assert runtime.transient_orphan_count == 0
+    assert prepare(runtime).sampling_path == "logits"

code/models/common/tests/llm_runtime/test_execution.py

@@ -0,0 +1,831 @@
+# 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),
+    ]

code/models/common/tests/llm_runtime/test_executor_integration.py

@@ -0,0 +1,1833 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Generic host-only executor/generator integration for migrated siblings."""
+
+import inspect
+from dataclasses import replace
+from importlib import import_module
+from types import SimpleNamespace
+from unittest.mock import MagicMock, create_autospec
+
+import pytest
+import torch
+
+import ttnn
+from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig
+from models.common.llm_runtime.decode import DecodeTraceSignature
+from models.common.llm_runtime.execution import EagerExecutor
+from models.common.llm_runtime.lane_group import LaneGroupExecutor
+from models.common.llm_runtime.prefill.signatures import PrefillProgramSignature
+from models.common.llm_runtime.program_compiler import ProgramKey
+from models.common.llm_runtime.warmup import _build_plan
+from models.common.models import executor as shared_model_executor
+from models.common.models import llama3_executor as llama3_family_executor
+from models.common.models import qwen2_executor as qwen2_family_executor
+from models.common.models.deepseek_r1_distill_qwen_14b import executor as deepseek_executor
+from models.common.models.deepseek_r1_distill_qwen_14b import generator as deepseek_generator
+from models.common.models.llama32_1b import executor as llama32_executor
+from models.common.models.llama32_1b import generator as llama32_generator
+from models.common.models.llama32_3b import executor as llama32_3b_executor
+from models.common.models.llama32_3b import generator as llama32_3b_generator
+from models.common.models.llama33_70b import executor as llama33_70b_executor
+from models.common.models.llama33_70b import generator as llama33_70b_generator
+from models.common.models.mistral_7b import executor as mistral_executor
+from models.common.models.mistral_7b import generator as mistral_generator
+from models.common.models.phi4 import executor as phi4_executor
+from models.common.models.phi4 import generator as phi4_generator
+from models.common.models.qwen2_7b import executor as qwen2_executor
+from models.common.models.qwen2_7b import generator as qwen2_generator
+from models.common.models.qwen3_32b import executor as qwen3_32b_executor
+from models.common.models.qwen3_32b import generator as qwen3_32b_generator
+from models.common.models.qwen25_7b import executor as qwen25_executor
+from models.common.models.qwen25_7b import generator as qwen25_generator
+from models.common.models.qwen25_72b import executor as qwen25_72b_executor
+from models.common.models.qwen25_72b import generator as qwen25_72b_generator
+from models.common.models.qwen25_coder_32b import executor as qwen25_coder_32b_executor
+from models.common.models.qwen25_coder_32b import generator as qwen25_coder_32b_generator
+
+EXECUTOR_BINDINGS = {
+    "llama32_1b": SimpleNamespace(
+        executor_module=llama32_executor,
+        executor_class=llama32_executor.Llama32_1BExecutor,
+        executor_config_class=llama32_executor.Llama32_1BExecutorConfig,
+        generator_module=llama32_generator,
+        generator_class=llama32_generator.Llama32_1BGenerator,
+        generator_config_class=llama32_generator.Llama32_1BGeneratorConfig,
+        build_generator_name="build_llama32_1b_generator",
+        build_executor_name="build_llama32_1b_executor",
+        make_model=lambda **kwargs: _make_llama32_model(**kwargs),
+        make_runtime_config=lambda: _make_llama32_runtime_config(),
+        make_executor_config=lambda mode="none": _make_llama32_executor_config(mode),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_llama32_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_llama32_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="meta-llama/Llama-3.2-1B-Instruct",
+    ),
+    "llama32_3b": SimpleNamespace(
+        executor_module=llama32_3b_executor,
+        executor_class=llama32_3b_executor.Llama32_3BExecutor,
+        executor_config_class=llama32_3b_executor.Llama32_3BExecutorConfig,
+        generator_module=llama32_3b_generator,
+        generator_class=llama32_3b_generator.Llama32_3BGenerator,
+        generator_config_class=llama32_3b_generator.Llama32_3BGeneratorConfig,
+        build_generator_name="build_llama32_3b_generator",
+        build_executor_name="build_llama32_3b_executor",
+        make_model=lambda **kwargs: _make_llama32_model(**kwargs),
+        make_runtime_config=lambda: _make_llama32_runtime_config(),
+        make_executor_config=lambda mode="none": _make_llama32_executor_config(mode, module=llama32_3b_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_llama32_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_llama32_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="meta-llama/Llama-3.2-3B-Instruct",
+    ),
+    "llama33_70b": SimpleNamespace(
+        executor_module=llama33_70b_executor,
+        executor_class=llama33_70b_executor.Llama33_70BExecutor,
+        executor_config_class=llama33_70b_executor.Llama33_70BExecutorConfig,
+        generator_module=llama33_70b_generator,
+        generator_class=llama33_70b_generator.Llama33_70BGenerator,
+        generator_config_class=llama33_70b_generator.Llama33_70BGeneratorConfig,
+        build_generator_name="build_llama33_70b_generator",
+        build_executor_name="build_llama33_70b_executor",
+        make_model=lambda **kwargs: _make_llama32_model(**kwargs),
+        make_runtime_config=lambda: _make_llama32_runtime_config(),
+        make_executor_config=lambda mode="none": _make_llama32_executor_config(mode, module=llama33_70b_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(
+            _make_llama32_model(),
+            request_state_fields=llama33_70b_executor.Llama33_70BExecutor.request_state_fields,
+            **kwargs,
+        ),
+        make_product=lambda mesh_device, max_batch_size: _make_llama32_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(
+            llm, config, request_state_fields=llama33_70b_executor.Llama33_70BExecutor.request_state_fields
+        ),
+        hf_model="meta-llama/Llama-3.3-70B-Instruct",
+    ),
+    "qwen2_7b": SimpleNamespace(
+        executor_module=qwen2_executor,
+        executor_class=qwen2_executor.Qwen2Executor,
+        executor_config_class=qwen2_executor.Qwen2ExecutorConfig,
+        generator_module=qwen2_generator,
+        generator_class=qwen2_generator.Qwen2Generator,
+        generator_config_class=qwen2_generator.Qwen2GeneratorConfig,
+        build_generator_name="build_qwen2_7b_generator",
+        build_executor_name="build_qwen2_7b_executor",
+        make_model=lambda **kwargs: _make_qwen2_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen2_runtime_config(),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen2_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen2_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="Qwen/Qwen2-7B-Instruct",
+    ),
+    "qwen25_7b": SimpleNamespace(
+        executor_module=qwen25_executor,
+        executor_class=qwen25_executor.Qwen25Executor,
+        executor_config_class=qwen25_executor.Qwen25ExecutorConfig,
+        generator_module=qwen25_generator,
+        generator_class=qwen25_generator.Qwen25Generator,
+        generator_config_class=qwen25_generator.Qwen25GeneratorConfig,
+        build_generator_name="build_qwen25_7b_generator",
+        build_executor_name="build_qwen25_7b_executor",
+        make_model=lambda **kwargs: _make_qwen2_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen2_runtime_config(max_prefill_batch_size=8),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=qwen25_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen2_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen2_product(
+            mesh_device, max_batch_size, max_prefill_batch_size=8
+        ),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="Qwen/Qwen2.5-7B-Instruct",
+    ),
+    "qwen25_72b": SimpleNamespace(
+        executor_module=qwen25_72b_executor,
+        executor_class=qwen25_72b_executor.Qwen25_72BExecutor,
+        executor_config_class=qwen25_72b_executor.Qwen25_72BExecutorConfig,
+        generator_module=qwen25_72b_generator,
+        generator_class=qwen25_72b_generator.Qwen25_72BGenerator,
+        generator_config_class=qwen25_72b_generator.Qwen25_72BGeneratorConfig,
+        build_generator_name="build_qwen25_72b_generator",
+        build_executor_name="build_qwen25_72b_executor",
+        make_model=lambda **kwargs: _make_qwen25_72b_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen25_72b_runtime_config(),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=qwen25_72b_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen25_72b_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen25_72b_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="Qwen/Qwen2.5-72B-Instruct",
+    ),
+    "qwen25_coder_32b": SimpleNamespace(
+        executor_module=qwen25_coder_32b_executor,
+        executor_class=qwen25_coder_32b_executor.Qwen25Coder32BExecutor,
+        executor_config_class=qwen25_coder_32b_executor.Qwen25Coder32BExecutorConfig,
+        generator_module=qwen25_coder_32b_generator,
+        generator_class=qwen25_coder_32b_generator.Qwen25Coder32BGenerator,
+        generator_config_class=qwen25_coder_32b_generator.Qwen25Coder32BGeneratorConfig,
+        build_generator_name="build_qwen25_coder_32b_generator",
+        build_executor_name="build_qwen25_coder_32b_executor",
+        make_model=lambda **kwargs: _make_qwen25_coder_32b_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen25_coder_32b_runtime_config(),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=qwen25_coder_32b_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen25_coder_32b_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen25_coder_32b_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="Qwen/Qwen2.5-Coder-32B-Instruct",
+    ),
+    "qwen3_32b": SimpleNamespace(
+        executor_module=qwen3_32b_executor,
+        executor_class=qwen3_32b_executor.Qwen3_32BExecutor,
+        executor_config_class=qwen3_32b_executor.Qwen3_32BExecutorConfig,
+        generator_module=qwen3_32b_generator,
+        generator_class=qwen3_32b_generator.Qwen3_32BGenerator,
+        generator_config_class=qwen3_32b_generator.Qwen3_32BGeneratorConfig,
+        build_generator_name="build_qwen3_32b_generator",
+        build_executor_name="build_qwen3_32b_executor",
+        make_model=lambda **kwargs: _make_qwen3_32b_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen3_32b_runtime_config(),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=qwen3_32b_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(
+            _make_qwen3_32b_model(),
+            request_state_fields=qwen3_32b_executor.Qwen3_32BExecutor.request_state_fields,
+            **kwargs,
+        ),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen3_32b_product(mesh_device, max_batch_size),
+        make_lane=lambda llm, config: _FakeLane(
+            llm, config, request_state_fields=qwen3_32b_executor.Qwen3_32BExecutor.request_state_fields
+        ),
+        hf_model="Qwen/Qwen3-32B",
+    ),
+    "deepseek_r1_distill_qwen_14b": SimpleNamespace(
+        executor_module=deepseek_executor,
+        executor_class=deepseek_executor.DeepSeekR1Qwen14BExecutor,
+        executor_config_class=deepseek_executor.DeepSeekR1Qwen14BExecutorConfig,
+        generator_module=deepseek_generator,
+        generator_class=deepseek_generator.DeepSeekR1Qwen14BGenerator,
+        generator_config_class=deepseek_generator.DeepSeekR1Qwen14BGeneratorConfig,
+        build_generator_name="build_deepseek_r1_distill_qwen_14b_generator",
+        build_executor_name="build_deepseek_r1_distill_qwen_14b_executor",
+        make_model=lambda **kwargs: _make_qwen2_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen2_runtime_config(max_prefill_batch_size=32),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=deepseek_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen2_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen2_product(
+            mesh_device, max_batch_size, max_prefill_batch_size=32
+        ),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="deepseek-ai/DeepSeek-R1-Distill-Qwen-14B",
+    ),
+    "mistral_7b": SimpleNamespace(
+        executor_module=mistral_executor,
+        executor_class=mistral_executor.Mistral7BExecutor,
+        executor_config_class=mistral_executor.Mistral7BExecutorConfig,
+        generator_module=mistral_generator,
+        generator_class=mistral_generator.Mistral7BGenerator,
+        generator_config_class=mistral_generator.Mistral7BGeneratorConfig,
+        build_generator_name="build_mistral_7b_generator",
+        build_executor_name="build_mistral_7b_executor",
+        make_model=lambda **kwargs: _make_qwen2_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen2_runtime_config(max_prefill_batch_size=8),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=mistral_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen2_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen2_product(
+            mesh_device, max_batch_size, max_prefill_batch_size=8
+        ),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="mistralai/Mistral-7B-Instruct-v0.3",
+    ),
+    "phi4": SimpleNamespace(
+        executor_module=phi4_executor,
+        executor_class=phi4_executor.Phi4Executor,
+        executor_config_class=phi4_executor.Phi4ExecutorConfig,
+        generator_module=phi4_generator,
+        generator_class=phi4_generator.Phi4Generator,
+        generator_config_class=phi4_generator.Phi4GeneratorConfig,
+        build_generator_name="build_phi4_generator",
+        build_executor_name="build_phi4_executor",
+        make_model=lambda **kwargs: _make_qwen2_model(**kwargs),
+        make_runtime_config=lambda: _make_qwen2_runtime_config(max_prefill_batch_size=8),
+        make_executor_config=lambda mode="none": _make_qwen2_executor_config(mode, module=phi4_executor),
+        make_recording_target=lambda **kwargs: _RecordingTarget(_make_qwen2_model(), **kwargs),
+        make_product=lambda mesh_device, max_batch_size: _make_qwen2_product(
+            mesh_device, max_batch_size, max_prefill_batch_size=8
+        ),
+        make_lane=lambda llm, config: _FakeLane(llm, config),
+        hf_model="microsoft/phi-4",
+    ),
+}
+
+GENERATOR_PATHS = {
+    "llama32_1b": "models.common.models.llama32_1b.generator:Llama32_1BGenerator",
+    "llama32_3b": "models.common.models.llama32_3b.generator:Llama32_3BGenerator",
+    "llama33_70b": "models.common.models.llama33_70b.generator:Llama33_70BGenerator",
+    "mistral_7b": "models.common.models.mistral_7b.generator:Mistral7BGenerator",
+    "phi4": "models.common.models.phi4.generator:Phi4Generator",
+    "qwen2_7b": "models.common.models.qwen2_7b.generator:Qwen2Generator",
+    "qwen25_7b": "models.common.models.qwen25_7b.generator:Qwen25Generator",
+    "qwen25_72b": "models.common.models.qwen25_72b.generator:Qwen25_72BGenerator",
+    "qwen25_coder_32b": "models.common.models.qwen25_coder_32b.generator:Qwen25Coder32BGenerator",
+    "qwen3_32b": "models.common.models.qwen3_32b.generator:Qwen3_32BGenerator",
+    "deepseek_r1_distill_qwen_14b": (
+        "models.common.models.deepseek_r1_distill_qwen_14b.generator:DeepSeekR1Qwen14BGenerator"
+    ),
+}
+
+
+@pytest.fixture(params=EXECUTOR_BINDINGS.items(), ids=lambda item: item[0])
+def binding(request):
+    return request.param[1]
+
+
+_LLAMA_FAMILY_EXECUTOR_MODULES = (llama32_executor, llama32_3b_executor, llama33_70b_executor)
+_QWEN2_FAMILY_EXECUTOR_MODULES = (
+    qwen2_executor,
+    qwen25_executor,
+    qwen25_72b_executor,
+    qwen25_coder_32b_executor,
+)
+_SHARED_MODEL_EXECUTOR_MODULES = (
+    *_LLAMA_FAMILY_EXECUTOR_MODULES,
+    *_QWEN2_FAMILY_EXECUTOR_MODULES,
+    qwen3_32b_executor,
+)
+
+
+def _composition_module(binding):
+    if binding.executor_module in _SHARED_MODEL_EXECUTOR_MODULES:
+        return shared_model_executor
+    return binding.executor_module
+
+
+def _sampling_policy_module(binding):
+    if binding.executor_module in _LLAMA_FAMILY_EXECUTOR_MODULES:
+        return llama3_family_executor
+    if binding.executor_module in _QWEN2_FAMILY_EXECUTOR_MODULES:
+        return qwen2_family_executor
+    return binding.executor_module
+
+
+class _Mesh:
+    shape = (1, 1)
+
+    @staticmethod
+    def get_num_devices():
+        return 1
+
+
+class _Mesh2:
+    shape = (1, 2)
+
+    @staticmethod
+    def get_num_devices():
+        return 2
+
+
+class _Mesh8:
+    shape = (1, 8)
+
+    @staticmethod
+    def get_num_devices():
+        return 8
+
+
+def _make_llama32_model(max_batch_size=4):
+    paged = SimpleNamespace(block_size=32, max_num_blocks=132)
+    attention = SimpleNamespace(
+        n_kv_heads=8,
+        head_dim=64,
+        kv_cache_dtype=ttnn.bfloat8_b,
+        paged_attention_config=paged,
+        use_vllm_paged_kv_cache=True,
+        kv_cache=None,
+    )
+    live = SimpleNamespace(config=attention, kv_cache=None)
+    model = SimpleNamespace(
+        config=SimpleNamespace(
+            mesh_device=_Mesh(),
+            max_batch_size=max_batch_size,
+            max_seq_len=4096,
+            n_layers=1,
+            num_devices=1,
+            block_configs=(SimpleNamespace(attention_config=attention),),
+        ),
+        layers=(SimpleNamespace(attention=live),),
+        iter_executor_named_modules=lambda: (),
+        vocab_size=128256,
+        num_devices=1,
+    )
+
+    def configure_paged_attention(*, block_size, max_num_blocks):
+        assert attention.kv_cache is None
+        assert live.kv_cache is None
+        attention.paged_attention_config = SimpleNamespace(
+            block_size=block_size,
+            max_num_blocks=max_num_blocks,
+        )
+
+    model.configure_paged_attention = configure_paged_attention
+    return model
+
+
+def _make_llama32_runtime_config():
+    return SimpleNamespace(
+        model_cache_path="cache",
+        max_prefill_chunk_size=2048,
+        trace_prefill_supported_seq_lens=(128,),
+        can_enable_trace=lambda length, num_cached_tokens=0: length == 128,
+        supports_batched_prefill=True,
+        disable_batched_prefill=False,
+        max_prefill_batch_size=32,
+        batched_prefill_batched_extract=True,
+    )
+
+
+def _make_llama32_executor_config(mode="none", *, module=llama32_executor):
+    config_class = next(
+        getattr(module, name)
+        for name in (
+            "Llama32_1BExecutorConfig",
+            "Llama32_3BExecutorConfig",
+            "Llama33_70BExecutorConfig",
+        )
+        if hasattr(module, name)
+    )
+    return config_class(
+        trace=TraceConfig(mode),
+        warmup=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        paged_kv_cache=PagedKVCacheConfig(block_size=32, max_num_blocks=132, dtype=ttnn.bfloat8_b),
+        device_sampling_enabled=False,
+    )
+
+
+def _make_llama32_product(mesh_device, max_batch_size):
+    model = _make_llama32_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = mesh_device
+    return SimpleNamespace(model=model, runtime_config=_make_llama32_runtime_config())
+
+
+def _make_qwen2_model(max_batch_size=4):
+    model = _make_llama32_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = _Mesh2()
+    model.config.num_devices = 2
+    model.num_devices = 2
+    attention = model.layers[0].attention.config
+    attention.n_kv_heads = 4
+    attention.head_dim = 128
+    return model
+
+
+def _make_qwen2_runtime_config(*, max_prefill_batch_size=32):
+    runtime = _make_llama32_runtime_config()
+    runtime.trace_prefill_supported_seq_lens = (128, 1024)
+    runtime.can_enable_trace = lambda length, num_cached_tokens=0: length in (128, 1024)
+    runtime.max_prefill_batch_size = max_prefill_batch_size
+    return runtime
+
+
+def _make_qwen2_executor_config(mode="none", *, module=qwen2_executor):
+    config_class = (
+        getattr(module, "Qwen2ExecutorConfig", None)
+        or getattr(module, "Qwen25ExecutorConfig", None)
+        or getattr(module, "Qwen25_72BExecutorConfig", None)
+        or getattr(module, "Qwen25Coder32BExecutorConfig", None)
+        or getattr(module, "Qwen3_32BExecutorConfig", None)
+        or getattr(module, "DeepSeekR1Qwen14BExecutorConfig", None)
+        or getattr(module, "Mistral7BExecutorConfig", None)
+        or module.Phi4ExecutorConfig
+    )
+    return config_class(
+        trace=TraceConfig(mode),
+        warmup=WarmupConfig(prefill_seq_lens=(128, 1024), prefill_batch_sizes=(1,)),
+        paged_kv_cache=PagedKVCacheConfig(block_size=32, max_num_blocks=132, dtype=ttnn.bfloat8_b),
+        device_sampling_enabled=False,
+    )
+
+
+def _make_qwen2_product(mesh_device, max_batch_size, *, max_prefill_batch_size=32):
+    model = _make_qwen2_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = mesh_device
+    return SimpleNamespace(
+        model=model,
+        runtime_config=_make_qwen2_runtime_config(max_prefill_batch_size=max_prefill_batch_size),
+    )
+
+
+def _make_qwen25_72b_model(max_batch_size=4):
+    model = _make_llama32_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = _Mesh8()
+    model.config.num_devices = 8
+    model.num_devices = 8
+    attention = model.layers[0].attention.config
+    attention.n_kv_heads = 8
+    attention.head_dim = 128
+    return model
+
+
+def _make_qwen25_72b_runtime_config():
+    return _make_qwen2_runtime_config(max_prefill_batch_size=32)
+
+
+def _make_qwen25_72b_product(mesh_device, max_batch_size):
+    model = _make_qwen25_72b_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = mesh_device
+    return SimpleNamespace(model=model, runtime_config=_make_qwen25_72b_runtime_config())
+
+
+def _make_qwen25_coder_32b_model(max_batch_size=4):
+    model = _make_qwen25_72b_model(max_batch_size=max_batch_size)
+    model.config.dim = 5120
+    model.config.n_heads = 40
+    model.config.hidden_dim = 27648
+    model.config.hf_model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
+    return model
+
+
+def _make_qwen25_coder_32b_runtime_config():
+    runtime = _make_qwen25_72b_runtime_config()
+    runtime.max_prefill_chunk_size = 4096
+    runtime.can_enable_trace = lambda length, num_cached_tokens=0: num_cached_tokens == 0 and length in (128, 1024)
+    return runtime
+
+
+def _make_qwen25_coder_32b_product(mesh_device, max_batch_size):
+    model = _make_qwen25_coder_32b_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = mesh_device
+    return SimpleNamespace(model=model, runtime_config=_make_qwen25_coder_32b_runtime_config())
+
+
+def _make_qwen3_32b_model(max_batch_size=4):
+    model = _make_qwen25_72b_model(max_batch_size=max_batch_size)
+    model.config.dim = 5120
+    model.config.n_heads = 64
+    model.config.hidden_dim = 27648
+    model.config.vocab_size = 151936
+    model.config.hf_model_id = "Qwen/Qwen3-32B"
+    model.padded_vocab_size = 152064
+    return model
+
+
+def _make_qwen3_32b_runtime_config():
+    runtime = _make_qwen25_coder_32b_runtime_config()
+    runtime.max_prefill_chunk_size = 4096
+    return runtime
+
+
+def _make_qwen3_32b_product(mesh_device, max_batch_size):
+    model = _make_qwen3_32b_model(max_batch_size=max_batch_size)
+    model.config.mesh_device = mesh_device
+    return SimpleNamespace(model=model, runtime_config=_make_qwen3_32b_runtime_config())
+
+
+def test_qwen2_binding_preserves_tp2_runtime_and_sampling_defaults():
+    model = _make_qwen2_model()
+    _, num_layers, kv_heads_per_device, head_dim = qwen2_generator._model_kv_metadata(model)
+    runtime = _make_qwen2_runtime_config()
+    config = qwen2_generator.Qwen2GeneratorConfig(
+        hf_model="Qwen/Qwen2-7B-Instruct",
+        hf_revision="test-revision",
+        mesh_device=model.config.mesh_device,
+        max_batch_size=32,
+        max_seq_len=4096,
+    )
+
+    assert model.config.mesh_device.shape == (1, 2)
+    assert (num_layers, kv_heads_per_device, head_dim) == (1, 2, 128)
+    assert runtime.trace_prefill_supported_seq_lens == (128, 1024)
+    assert runtime.max_prefill_chunk_size == 2048
+    assert runtime.max_prefill_batch_size == 32
+    assert config.hf_revision == "test-revision"
+    assert config.device_sampling_enabled is False
+
+
+def test_qwen25_72b_binding_preserves_tp8_runtime_and_sampling_defaults():
+    model = _make_qwen25_72b_model()
+    _, num_layers, kv_heads_per_device, head_dim = qwen25_72b_generator._model_kv_metadata(model)
+    runtime = _make_qwen25_72b_runtime_config()
+    config = qwen25_72b_generator.Qwen25_72BGeneratorConfig(
+        hf_model="Qwen/Qwen2.5-72B-Instruct",
+        mesh_device=model.config.mesh_device,
+        max_batch_size=32,
+        max_seq_len=4096,
+    )
+
+    assert model.config.mesh_device.shape == (1, 8)
+    assert (num_layers, kv_heads_per_device, head_dim) == (1, 1, 128)
+    assert runtime.trace_prefill_supported_seq_lens == (128, 1024)
+    assert runtime.max_prefill_chunk_size == 2048
+    assert runtime.max_prefill_batch_size == 32
+    assert config.hf_revision == qwen25_72b_generator.DEFAULT_HF_REVISION
+    assert config.device_sampling_enabled is False
+
+
+def test_qwen25_coder_32b_binding_preserves_tp8_runtime_and_sampling_defaults():
+    model = _make_qwen25_coder_32b_model()
+    _, num_layers, kv_heads_per_device, head_dim = qwen25_coder_32b_generator._model_kv_metadata(model)
+    runtime = _make_qwen25_coder_32b_runtime_config()
+    config = qwen25_coder_32b_generator.Qwen25Coder32BGeneratorConfig(
+        hf_model="Qwen/Qwen2.5-Coder-32B-Instruct",
+        mesh_device=model.config.mesh_device,
+        max_batch_size=32,
+        max_seq_len=4096,
+    )
+
+    assert model.config.mesh_device.shape == (1, 8)
+    assert (num_layers, kv_heads_per_device, head_dim) == (1, 1, 128)
+    assert runtime.trace_prefill_supported_seq_lens == (128, 1024)
+    assert runtime.max_prefill_chunk_size == 4096
+    assert runtime.can_enable_trace(128, 0) is True
+    assert runtime.can_enable_trace(128, 1) is False
+    assert runtime.max_prefill_batch_size == 32
+    assert config.hf_revision == qwen25_coder_32b_generator.DEFAULT_HF_REVISION
+    assert config.device_sampling_enabled is False
+
+
+@pytest.mark.parametrize(
+    "product_binding",
+    (EXECUTOR_BINDINGS["qwen25_72b"], EXECUTOR_BINDINGS["qwen25_coder_32b"]),
+    ids=("qwen25_72b", "qwen25_coder_32b"),
+)
+@pytest.mark.parametrize("device_sampling_enabled", (False, True), ids=("sampling-off", "sampling-on"))
+def test_large_qwen_builder_threads_exact_decode_sampling_coverage(
+    monkeypatch,
+    product_binding,
+    device_sampling_enabled,
+):
+    mesh_device = _Mesh()
+    product = product_binding.make_product(mesh_device, 4)
+    executor_configs = []
+
+    monkeypatch.setattr(product_binding.generator_module, "from_pretrained", lambda **kwargs: product)
+    monkeypatch.setattr(
+        product_binding.generator_module,
+        "_model_kv_metadata",
+        lambda model: ((ttnn.bfloat8_b,), 1, 1, 128),
+    )
+
+    def build_executor(llm, config):
+        executor_configs.append(config)
+        return product_binding.make_lane(llm, config)
+
+    monkeypatch.setattr(product_binding.generator_module, product_binding.build_executor_name, build_executor)
+    generator = getattr(product_binding.generator_module, product_binding.build_generator_name)(
+        product_binding.generator_config_class(
+            hf_model=product_binding.hf_model,
+            mesh_device=mesh_device,
+            max_batch_size=4,
+            max_seq_len=1024,
+            n_layers=1,
+            device_sampling_enabled=device_sampling_enabled,
+        )
+    )
+
+    try:
+        assert len(executor_configs) == 1
+        warmup = executor_configs[0].warmup
+        assert warmup.include_decode_top_k is device_sampling_enabled
+        plan = _build_plan(
+            warmup=warmup,
+            layout=PageTableLayout(block_size=32, raw_capacity_width=32, prefill_width=64, decode_width=32),
+            prefill_sequence_lengths=(128,),
+            lane_batch_size=4,
+            allow_force_argmax=True,
+            can_sample_on_device=device_sampling_enabled,
+        )
+        assert [case.sampling_path for case in plan.decode] == (
+            ["logits", "argmax", "topk"] if device_sampling_enabled else ["logits"]
+        )
+    finally:
+        generator.cleanup()
+
+
+def test_qwen3_32b_binding_preserves_tp8_runtime_and_padded_vocab_defaults():
+    model = _make_qwen3_32b_model()
+    _, num_layers, kv_heads_per_device, head_dim = qwen3_32b_generator._model_kv_metadata(model)
+    runtime = _make_qwen3_32b_runtime_config()
+    config = qwen3_32b_generator.Qwen3_32BGeneratorConfig(
+        hf_model="Qwen/Qwen3-32B",
+        mesh_device=model.config.mesh_device,
+        max_batch_size=32,
+        max_seq_len=4096,
+    )
+
+    assert model.config.mesh_device.shape == (1, 8)
+    assert (num_layers, kv_heads_per_device, head_dim) == (1, 1, 128)
+    assert model.config.vocab_size == 151936
+    assert model.padded_vocab_size == 152064
+    assert runtime.trace_prefill_supported_seq_lens == (128, 1024)
+    assert runtime.max_prefill_chunk_size == 4096
+    assert runtime.can_enable_trace(128, 0) is True
+    assert runtime.can_enable_trace(128, 1) is False
+    assert config.hf_revision == qwen3_32b_generator.DEFAULT_HF_REVISION
+    assert config.device_sampling_enabled is False
+
+
+@pytest.mark.parametrize(
+    ("cluster_shape", "disable_batched_prefill", "advertised_lengths", "expected"),
+    [
+        ([1, 8], False, (128, 1024), (128, 1024)),
+        ([1, 8], True, (128, 1024), (128, 1024)),
+        ([1, 4], False, (128, 1024), ()),
+        ([1, 8], False, (128,), (128,)),
+    ],
+    ids=("t3k-batched", "t3k-sequential", "bh-batched", "t3k-low-ceiling"),
+)
+def test_qwen3_prefill_capture_primes_are_t3k_product_owned(
+    cluster_shape,
+    disable_batched_prefill,
+    advertised_lengths,
+    expected,
+):
+    runtime = SimpleNamespace(
+        cluster_shape=cluster_shape,
+        disable_batched_prefill=disable_batched_prefill,
+        trace_prefill_supported_seq_lens=advertised_lengths,
+        can_enable_trace=lambda length, cached: length in advertised_lengths and cached == 0,
+    )
+
+    num_devices = int(cluster_shape[0]) * int(cluster_shape[1])
+    assert (
+        qwen3_32b_executor._resolve_trace_capture_prime_sequence_lengths(
+            runtime,
+            num_devices=num_devices,
+        )
+        == expected
+    )
+
+
+def test_phi4_binding_preserves_cap8_trace_buckets_and_pinned_revision():
+    runtime = _make_qwen2_runtime_config(max_prefill_batch_size=8)
+    config = phi4_generator.Phi4GeneratorConfig(
+        hf_model="microsoft/phi-4",
+        mesh_device=_Mesh2(),
+        max_batch_size=32,
+        max_seq_len=4096,
+    )
+
+    assert runtime.trace_prefill_supported_seq_lens == (128, 1024)
+    assert runtime.max_prefill_chunk_size == 2048
+    assert runtime.max_prefill_batch_size == 8
+    assert config.hf_revision == phi4_generator.DEFAULT_HF_REVISION
+    assert config.device_sampling_enabled is False
+
+
+@pytest.mark.parametrize("mode", ["none", "decode_only", "all"])
+def test_model_owned_executor_has_exact_composition_and_owner_counts(binding, mode, monkeypatch):
+    owner_names = (
+        "PagedKVCacheManager",
+        "OutputReader",
+        "PrefillRuntime",
+        "DecodeRuntime",
+        "ProgramCompiler",
+        "EagerExecutor",
+        "TraceCompiler",
+        "TracedExecutor",
+        "WarmupCoordinator",
+    )
+    composition_module = _composition_module(binding)
+    owner_factories = {}
+    for name in owner_names:
+        factory = MagicMock(wraps=getattr(composition_module, name))
+        monkeypatch.setattr(composition_module, name, factory)
+        owner_factories[name] = factory
+
+    executor = binding.executor_class(
+        binding.make_model(),
+        binding.make_runtime_config(),
+        binding.make_executor_config(mode),
+    )
+    expected_counts = {name: 1 for name in owner_names}
+    if mode == "none":
+        expected_counts["TraceCompiler"] = 0
+        expected_counts["TracedExecutor"] = 0
+    assert {name: factory.call_count for name, factory in owner_factories.items()} == expected_counts
+    assert executor.eager_executor.program_compiler is executor.program_compiler
+    assert executor.eager_executor.prefill is executor.prefill_runtime
+    assert executor.eager_executor.decode is executor.decode_runtime
+    assert executor.warmup.eager is executor.eager_executor
+    assert executor.warmup.trace_compiler is executor.trace_compiler
+    assert executor.eager_execution is executor.eager_executor
+    assert executor.prefill_runtime.config.trace_capture_prime_sequence_lengths == (
+        (128, 1024) if binding.executor_module is qwen3_32b_executor else ()
+    )
+    if mode == "none":
+        assert executor.warmup.execution is executor.eager_executor
+        assert executor.trace_compiler is None
+        assert executor.traced_executor is None
+        assert executor.traced_prefill_execution is None
+        assert executor.traced_decode_execution is None
+    else:
+        assert executor.warmup.execution is executor.traced_executor
+        assert executor.traced_executor.eager_executor is executor.eager_executor
+        assert executor.traced_executor.trace_compiler is executor.trace_compiler
+        assert executor.trace_compiler.program_compiler is executor.program_compiler
+        assert (executor.traced_prefill_execution is not None) is (mode == "all")
+        assert executor.traced_decode_execution is executor.traced_executor
+
+
+def test_llama3_create_sampling_state_disables_duplicate_seed_salting(monkeypatch):
+    captured = {}
+
+    class FakeSampling1D:
+        def __init__(self):
+            self.config = SimpleNamespace(is_resolved=lambda: True)
+
+    class FakeSamplingState1D:
+        def __init__(self, sampling, *, salt_duplicate_seeds=True, **_kwargs):
+            captured["salt_duplicate_seeds"] = salt_duplicate_seeds
+            self.sampling = sampling
+
+        def create_state(self):
+            return object()
+
+    monkeypatch.setattr(llama3_family_executor, "Sampling1D", FakeSampling1D)
+    monkeypatch.setattr(llama3_family_executor, "SamplingState1D", FakeSamplingState1D)
+
+    controller, state = llama3_family_executor._create_sampling_state(SimpleNamespace(sampling=FakeSampling1D()), True)
+
+    assert captured["salt_duplicate_seeds"] is False
+    assert controller is not None
+    assert state is not None
+
+
+def test_qwen3_create_sampling_state_disables_duplicate_seed_salting(monkeypatch):
+    captured = {}
+
+    class FakeSampling1D:
+        def __init__(self):
+            self.config = SimpleNamespace(is_resolved=lambda: True)
+
+    class FakeSamplingState1D:
+        def __init__(self, sampling, *, salt_duplicate_seeds=True, **_kwargs):
+            captured["salt_duplicate_seeds"] = salt_duplicate_seeds
+            self.sampling = sampling
+
+        def create_state(self):
+            return object()
+
+    monkeypatch.setattr(qwen3_32b_executor, "Sampling1D", FakeSampling1D)
+    monkeypatch.setattr(qwen3_32b_executor, "SamplingState1D", FakeSamplingState1D)
+
+    controller, state = qwen3_32b_executor._create_sampling_state(SimpleNamespace(sampling=FakeSampling1D()), True)
+
+    assert captured["salt_duplicate_seeds"] is False
+    assert controller is not None
+    assert state is not None
+
+
+def _device_sampling_executor(binding, monkeypatch, *, runtime_disable: bool):
+    class FakeSampling1D:
+        config = SimpleNamespace(
+            is_resolved=lambda: True,
+            allow_force_argmax=False,
+            max_batch_size=32,
+            max_top_k=32,
+        )
+
+        def decode_forward(self):
+            raise AssertionError("construction-policy test must not execute sampling")
+
+    sampling_policy_module = _sampling_policy_module(binding)
+    monkeypatch.setattr(sampling_policy_module, "Sampling1D", FakeSampling1D)
+    model = binding.make_model()
+    model.sampling = FakeSampling1D()
+    if binding.executor_module in (llama33_70b_executor, qwen3_32b_executor):
+
+        class FakeSamplingState1D:
+            def __init__(self, sampling, **_kwargs):
+                self.sampling = sampling
+                self.seed_manager = SimpleNamespace()
+
+            def create_state(self):
+                return SimpleNamespace(seed_state=SimpleNamespace(capacity=32))
+
+            def admit(self, *args, **kwargs):
+                return None
+
+            def decode_forward(self, *args, **kwargs):
+                return None
+
+            def release(self, *args, **kwargs):
+                return None
+
+        monkeypatch.setattr(sampling_policy_module, "SamplingState1D", FakeSamplingState1D)
+    runtime_config = binding.make_runtime_config()
+    runtime_config.disable_batched_prefill = runtime_disable
+    config = replace(
+        binding.make_executor_config("none"),
+        device_sampling_enabled=True,
+    )
+    return binding.executor_class(model, runtime_config, config)
+
+
+@pytest.mark.parametrize(
+    ("runtime_disable", "environment_disable", "expected_kinds"),
+    [
+        (False, False, ("batched",)),
+        (True, False, ("single", "single")),
+        (False, True, ("single", "single")),
+    ],
+    ids=("device-sampled-batched", "runtime-disabled", "environment-disabled"),
+)
+def test_device_sampling_prefill_batch_policy_is_model_owned(
+    binding,
+    monkeypatch,
+    runtime_disable,
+    environment_disable,
+    expected_kinds,
+):
+    if environment_disable:
+        monkeypatch.setenv("DISABLE_BATCHED_PREFILL", "1")
+    else:
+        monkeypatch.delenv("DISABLE_BATCHED_PREFILL", raising=False)
+    executor = _device_sampling_executor(
+        binding,
+        monkeypatch,
+        runtime_disable=runtime_disable,
+    )
+
+    prepared = executor.prefill_runtime.prepare(
+        tokens=torch.ones((2, 128), dtype=torch.long),
+        page_table=torch.arange(8, dtype=torch.int32).reshape(2, 4),
+        prompt_lens=torch.full((2,), 128, dtype=torch.long),
+        empty_slots=[0, 1],
+    )
+
+    if binding.executor_module in (llama33_70b_executor, qwen3_32b_executor):
+        expected_kinds = ("single", "single")
+    assert tuple(item.request.kind for item in prepared) == expected_kinds
+    if expected_kinds == ("batched",):
+        assert prepared[0].request.source_rows == (0, 1)
+        assert not prepared[0].request.uses_chunked_prefill
+
+
+def test_llama32_1b_warms_every_q128_topk_tile_start_once_per_execution_mode():
+    executor = object.__new__(llama32_executor.Llama32_1BExecutor)
+    executor._q128_topk_tile_ends_warmed = set()
+    executor.eager_executor = object()
+    executor.traced_executor = object()
+    executor.page_table_layout = SimpleNamespace(block_size=32)
+    executor.prefill_runtime = SimpleNamespace(config=SimpleNamespace(static_q128_topk_supported=True))
+    executor.warmup = SimpleNamespace(
+        config=SimpleNamespace(
+            prefill_sequence_lengths=(128,),
+            prime_q128_tile_ends=False,
+        )
+    )
+    executor.compile_prefill = MagicMock()
+    kv_cache = object()
+
+    for enable_trace in (False, False, True, True):
+        executor._warmup_q128_topk_tile_ends(
+            kv_cache=kv_cache,
+            can_sample_on_device=True,
+            enable_trace=enable_trace,
+        )
+
+    assert executor.compile_prefill.call_count == 6
+    calls = executor.compile_prefill.call_args_list
+    assert [call.kwargs["tokens"].shape[1] for call in calls] == [32, 64, 96, 32, 64, 96]
+    assert [call.kwargs["page_table"].shape[1] for call in calls] == [1, 2, 3, 1, 2, 3]
+    assert all(call.kwargs["kv_cache"] is kv_cache for call in calls)
+    assert all(call.kwargs["execution"] is executor.eager_executor for call in calls[:3])
+    assert all(call.kwargs["execution"] is executor.traced_executor for call in calls[3:])
+    assert executor._q128_topk_tile_ends_warmed == {False, True}
+
+
+@pytest.mark.parametrize(
+    ("enable_trace", "expected_order"),
+    [
+        (False, ("default", 96, 0, 32, 64)),
+        (True, (0, 32, 64, "default", 96)),
+    ],
+    ids=("eager", "traced"),
+)
+def test_qwen3_lane4_warms_every_runtime_q128_topk_signature_before_activation(enable_trace, expected_order):
+    executor = SimpleNamespace(
+        _q128_topk_tile_ends_warmed=set(),
+        eager_executor=object(),
+        traced_executor=object(),
+        page_table_layout=SimpleNamespace(block_size=32),
+        prefill_runtime=SimpleNamespace(config=SimpleNamespace(static_q128_topk_supported=True)),
+        warmup=SimpleNamespace(
+            config=SimpleNamespace(
+                prefill_sequence_lengths=(128, 1024),
+                prime_q128_tile_ends=False,
+            )
+        ),
+    )
+    compiled = []
+    order = []
+    activation = []
+
+    def record_signature(prompt_length):
+        assert not activation
+        order.append(((prompt_length - 1) // 32) * 32)
+        compiled.append(
+            PrefillProgramSignature(
+                operation_variant="regular-single",
+                padded_batch_size=1,
+                invocation_sequence_length=128,
+                page_table_width=64,
+                chunk_page_table_width=None,
+                sampling_path="topk",
+                penalties_enabled=False,
+                logprobs_enabled=False,
+                last_token_tile_start=((prompt_length - 1) // 32) * 32,
+            )
+        )
+
+    def compile_prefill(**kwargs):
+        expected_execution = executor.traced_executor if enable_trace else executor.eager_executor
+        assert kwargs["execution"] is expected_execution
+        assert kwargs["kv_cache"] == "cache"
+        assert kwargs["sampling_params"].top_k.tolist() == [32]
+        record_signature(int(kwargs["prompt_lens"][0]))
+
+    executor.compile_prefill = compile_prefill
+
+    default_warmed = False
+
+    def default_warmup():
+        nonlocal default_warmed
+        if default_warmed:
+            return
+        default_warmed = True
+        order.append("default")
+        # The coordinator's ordinary Q128 top-k case covers the final tile.
+        record_signature(128)
+
+    for _ in range(2):
+        qwen3_32b_executor._warmup_q128_around_prefill(
+            executor,
+            default_warmup,
+            kv_cache="cache",
+            can_sample_on_device=True,
+            enable_trace=enable_trace,
+        )
+    activation.append(True)
+
+    assert tuple(order) == expected_order
+    assert {signature.last_token_tile_start for signature in compiled} == {0, 32, 64, 96}
+    compiled_keys = {ProgramKey.from_signature(signature) for signature in compiled}
+    runtime_signatures = {replace(compiled[0], last_token_tile_start=tile_start) for tile_start in (0, 32, 64, 96)}
+    assert {ProgramKey.from_signature(signature) for signature in runtime_signatures} == compiled_keys
+    assert executor._q128_topk_tile_ends_warmed == {enable_trace}
+
+
+def test_qwen3_lane4_executor_installs_model_owned_q128_warmup(monkeypatch):
+    executor = _device_sampling_executor(EXECUTOR_BINDINGS["qwen3_32b"], monkeypatch, runtime_disable=False)
+
+    assert executor._prefill_warmup is qwen3_32b_executor._warmup_q128_around_prefill
+    assert executor._q128_topk_tile_ends_warmed == set()
+
+
+@pytest.mark.parametrize("device_sampling_enabled", (False, True), ids=("disabled", "enabled"))
+def test_qwen3_generator_sampling_policy_controls_decode_topk_warmup(monkeypatch, device_sampling_enabled):
+    mesh_device = _Mesh8()
+    product = _make_qwen3_32b_product(mesh_device, max_batch_size=4)
+    executor_configs = []
+
+    monkeypatch.setattr(qwen3_32b_generator, "from_pretrained", lambda **kwargs: product)
+    monkeypatch.setattr(
+        qwen3_32b_generator,
+        "_model_kv_metadata",
+        lambda model: ((ttnn.bfloat8_b,), 1, 8, 64),
+    )
+
+    def build_executor(llm, config):
+        executor_configs.append(config)
+        return _FakeLane(
+            llm,
+            config,
+            request_state_fields=qwen3_32b_executor.Qwen3_32BExecutor.request_state_fields,
+        )
+
+    monkeypatch.setattr(qwen3_32b_generator, "build_qwen3_32b_executor", build_executor)
+    generator = qwen3_32b_generator.build_qwen3_32b_generator(
+        qwen3_32b_generator.Qwen3_32BGeneratorConfig(
+            hf_model="Qwen/Qwen3-32B",
+            mesh_device=mesh_device,
+            max_batch_size=4,
+            max_seq_len=1024,
+            trace_mode="decode_only",
+            device_sampling_enabled=device_sampling_enabled,
+        )
+    )
+
+    assert len(executor_configs) == 1
+    assert executor_configs[0].warmup.include_decode_top_k is device_sampling_enabled
+    if device_sampling_enabled:
+        observed_missing_signature = DecodeTraceSignature(
+            batch_size=4,
+            page_table_width=32,
+            sampling_path="topk",
+            device_feedback=True,
+        )
+        assert (
+            ProgramKey.from_signature(observed_missing_signature).digest
+            == "6f8351f51a0c90eaea5fca6700b3887e380a015dad7ee6f6a9e8be971dfebbd5"
+        )
+    generator.cleanup()
+
+
+@pytest.mark.parametrize(
+    "method,positional,keyword_only",
+    [
+        (
+            "compile_prefill",
+            ["self"],
+            [
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "execution",
+            ],
+        ),
+        (
+            "compile_decode",
+            ["self"],
+            [
+                "tokens",
+                "start_pos",
+                "page_table",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+                "execution",
+            ],
+        ),
+        (
+            "prefill_forward",
+            ["self", "tokens", "page_table"],
+            [
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "execution",
+            ],
+        ),
+        (
+            "decode_forward",
+            ["self", "tokens", "start_pos", "page_table"],
+            [
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+                "read_from_device",
+                "execution",
+            ],
+        ),
+        ("read_decode_output", ["self", "tt_out"], ["async_read"]),
+        ("process_decode_output_host", ["self", "tt_out"], ["is_tokens"]),
+        ("can_trace_prefill", ["self"], ["tokens", "prompt_lens", "start_pos", "empty_slots"]),
+        ("warmup_model_prefill", ["self"], ["kv_cache", "can_sample_on_device", "enable_trace"]),
+        (
+            "warmup_model_decode",
+            ["self"],
+            ["kv_cache", "max_batch_size", "num_blocks", "can_sample_on_device", "enable_trace"],
+        ),
+    ],
+)
+def test_executor_call_contract(binding, method, positional, keyword_only):
+    if binding.executor_module not in (llama33_70b_executor, qwen3_32b_executor):
+        keyword_only = [name for name in keyword_only if name not in {"prompt_tokens", "output_tokens", "slot_remap"}]
+    signature = inspect.signature(getattr(binding.executor_class, method))
+    parameters = signature.parameters
+    required = {
+        "compile_prefill": {"tokens", "page_table"},
+        "compile_decode": {"tokens", "start_pos", "page_table"},
+        "prefill_forward": {"tokens", "page_table"},
+        "decode_forward": {"tokens", "start_pos", "page_table"},
+        "read_decode_output": {"tt_out"},
+        "process_decode_output_host": {"tt_out"},
+        "can_trace_prefill": {"tokens"},
+        "warmup_model_prefill": {"kv_cache", "can_sample_on_device", "enable_trace"},
+        "warmup_model_decode": {
+            "kv_cache",
+            "max_batch_size",
+            "num_blocks",
+            "can_sample_on_device",
+            "enable_trace",
+        },
+    }[method]
+    non_none_defaults = {"reset_batch": False, "read_from_device": True, "async_read": False, "is_tokens": False}
+
+    assert list(parameters) == positional + keyword_only
+    assert all(parameters[name].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD for name in positional)
+    assert all(parameters[name].kind is inspect.Parameter.KEYWORD_ONLY for name in keyword_only)
+    for name, parameter in tuple(parameters.items())[1:]:
+        expected_default = inspect.Parameter.empty if name in required else non_none_defaults.get(name)
+        assert parameter.default == expected_default
+        assert parameter.annotation is not inspect.Parameter.empty
+    assert signature.return_annotation is not inspect.Signature.empty
+
+
+@pytest.mark.parametrize(
+    "executor_class",
+    [
+        qwen3_32b_executor.Qwen3_32BExecutor,
+    ],
+)
+def test_qwen3_delegates_resolved_sampling_warmup_and_activation_to_coordinator(executor_class):
+    warmup = SimpleNamespace(warmup_prefill=MagicMock(), warmup_decode=MagicMock())
+    executor = SimpleNamespace(_ensure_active=MagicMock(), warmup=warmup)
+
+    executor_class.warmup_model_prefill(
+        executor,
+        kv_cache="cache",
+        can_sample_on_device=True,
+        enable_trace=True,
+    )
+    executor_class.warmup_model_decode(
+        executor,
+        kv_cache="cache",
+        max_batch_size=32,
+        num_blocks=64,
+        can_sample_on_device=True,
+        enable_trace=True,
+    )
+
+    warmup.warmup_prefill.assert_called_once_with(
+        kv_cache="cache",
+        can_sample_on_device=True,
+        enable_trace=True,
+    )
+    warmup.warmup_decode.assert_called_once_with(
+        kv_cache="cache",
+        max_batch_size=32,
+        num_blocks=64,
+        can_sample_on_device=True,
+        enable_trace=True,
+    )
+    assert "capture_all" not in inspect.getsource(executor_class.warmup_model_decode)
+
+
+class _RecordingTarget:
+    model_args = object()
+    mesh_device = object()
+    cache_path = "cache"
+    already_warmed_up_prefill = False
+    eager_execution = object()
+    traced_prefill_execution = object()
+    traced_decode_execution = object()
+
+    def __init__(self, model, traceable=True, request_state_fields=()):
+        self.model = model
+        self.traceable = traceable
+        self._request_state_fields = tuple(request_state_fields)
+        self.calls = []
+
+    def can_trace_prefill(self, **kwargs):
+        self.calls.append(("can_trace_prefill", kwargs))
+        return self.traceable
+
+    def prefill_forward(self, **kwargs):
+        self.calls.append(("prefill_forward", kwargs))
+        return kwargs["execution"]
+
+    def decode_forward(self, **kwargs):
+        self.calls.append(("decode_forward", kwargs))
+        return kwargs["execution"]
+
+    def cleanup(self):
+        self.calls.append(("cleanup", {}))
+
+
+def test_generator_preserves_required_trace_intent_for_ineligible_prefill(binding):
+    target = binding.make_recording_target(traceable=False)
+    target.config = binding.make_executor_config("all")
+    generator = binding.generator_class(target, binding.generator_module._build_vllm_adapter(target))
+    tokens = __import__("torch").tensor([[1]])
+    page_table = __import__("torch").tensor([[0]], dtype=__import__("torch").int32)
+    assert generator.prefill_forward(tokens, page_table, enable_trace=True) is target.traced_prefill_execution
+    assert [name for name, _ in target.calls] == ["prefill_forward"]
+
+
+def test_generator_routes_external_decode_only_policy_with_all_trace_targets(binding):
+    target = binding.make_recording_target()
+    target.config = binding.make_executor_config("all")
+    generator = binding.generator_class(target, binding.generator_module._build_vllm_adapter(target))
+    torch = __import__("torch")
+    tokens = torch.tensor([[1]])
+    page_table = torch.tensor([[0]], dtype=torch.int32)
+    start_pos = torch.tensor([0])
+
+    assert generator.prefill_forward(tokens, page_table, enable_trace=False) is target.eager_execution
+    assert (
+        generator.decode_forward(tokens[:, 0], start_pos, page_table, enable_trace=True)
+        is target.traced_decode_execution
+    )
+    assert [name for name, _ in target.calls] == ["prefill_forward", "decode_forward"]
+
+
+def test_executor_validates_borrowed_cache_then_omits_it_from_execution(binding):
+    execution = create_autospec(EagerExecutor, instance=True)
+    events = []
+    execution.compile_prefill.side_effect = lambda **kwargs: events.append("dispatch_compile_prefill")
+    execution.compile_decode.side_effect = lambda **kwargs: events.append("dispatch_compile_decode")
+    execution.prefill_forward.side_effect = lambda **kwargs: events.append("dispatch_prefill") or "prefill"
+    execution.decode_forward.side_effect = lambda **kwargs: events.append("dispatch_decode") or "decode"
+    executor = object.__new__(binding.executor_class)
+    executor._prefill_execution = execution
+    executor._decode_execution = execution
+    executor._ensure_active = lambda: None
+    executor._validate_bound_cache = lambda cache: events.append(("validate_cache", cache))
+    executor._ensure_sampling_for = lambda params: events.append(("validate_sampling", params))
+
+    tokens = torch.zeros((1, 4), dtype=torch.long)
+    start_pos = torch.zeros((1,), dtype=torch.long)
+    page_table = torch.zeros((1, 1), dtype=torch.int32)
+    prompt_lens = torch.full((1,), 4, dtype=torch.long)
+    empty_slots = [0]
+    kv_cache = object()
+    sampling_params = object()
+
+    executor.compile_prefill(
+        tokens=tokens,
+        page_table=page_table,
+        prompt_lens=prompt_lens,
+        start_pos=start_pos,
+        empty_slots=empty_slots,
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+    )
+    executor.compile_decode(
+        tokens=tokens,
+        start_pos=start_pos,
+        page_table=page_table,
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+        reset_batch=True,
+    )
+    assert (
+        executor.prefill_forward(
+            tokens,
+            page_table,
+            prompt_lens=prompt_lens,
+            start_pos=start_pos,
+            empty_slots=empty_slots,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+        )
+        == "prefill"
+    )
+    assert (
+        executor.decode_forward(
+            tokens,
+            start_pos,
+            page_table,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+            reset_batch=True,
+            read_from_device=False,
+        )
+        == "decode"
+    )
+
+    expected_validation = [("validate_cache", kv_cache), ("validate_sampling", sampling_params)]
+    assert events == [
+        *expected_validation,
+        "dispatch_compile_prefill",
+        *expected_validation,
+        "dispatch_compile_decode",
+        *expected_validation,
+        "dispatch_prefill",
+        *expected_validation,
+        "dispatch_decode",
+    ]
+    request_state_names = (
+        ("prompt_tokens", "output_tokens", "slot_remap")
+        if "prompt_tokens" in inspect.signature(binding.executor_class.compile_prefill).parameters
+        else ()
+    )
+    for target, expected_names in (
+        (
+            execution.compile_prefill,
+            (
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "sampling_params",
+                *request_state_names,
+            ),
+        ),
+        (
+            execution.compile_decode,
+            ("tokens", "start_pos", "page_table", "sampling_params", *request_state_names, "reset_batch"),
+        ),
+        (
+            execution.prefill_forward,
+            (
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "sampling_params",
+                *request_state_names,
+            ),
+        ),
+        (
+            execution.decode_forward,
+            (
+                "tokens",
+                "start_pos",
+                "page_table",
+                "sampling_params",
+                *request_state_names,
+                "reset_batch",
+                "read_from_device",
+            ),
+        ),
+    ):
+        assert target.call_count == 1
+        assert tuple(target.call_args.kwargs) == expected_names
+        assert "kv_cache" not in target.call_args.kwargs
+
+
+def test_late_capacity_reconfigures_existing_owners_before_allocation(binding, monkeypatch):
+    executor = binding.executor_class(
+        binding.make_model(), binding.make_runtime_config(), binding.make_executor_config()
+    )
+    owner_ids = tuple(
+        id(owner)
+        for owner in (executor.prefill_runtime, executor.decode_runtime, executor.warmup, executor.program_compiler)
+    )
+    assert executor.page_table_layout.raw_capacity_width == 128
+
+    executor.configure_paged_kv_cache(
+        PagedKVCacheConfig(
+            block_size=16,
+            max_num_blocks=200,
+            dtype=ttnn.bfloat8_b,
+            num_blocks=200,
+        )
+    )
+
+    assert (
+        tuple(
+            id(owner)
+            for owner in (executor.prefill_runtime, executor.decode_runtime, executor.warmup, executor.program_compiler)
+        )
+        == owner_ids
+    )
+    assert executor.config.paged_kv_cache is executor.kv_cache_manager.config
+    assert executor.kv_cache_manager.config.block_size == 16
+    assert executor.kv_cache_manager.config.max_num_blocks == executor.kv_cache_manager.config.num_blocks == 200
+    assert executor.model.layers[0].attention.config.paged_attention_config.block_size == 16
+    assert executor.model.layers[0].attention.config.paged_attention_config.max_num_blocks == 200
+    assert executor.page_table_layout.block_size == 16
+    assert executor.page_table_layout.raw_capacity_width == 200
+    assert executor.prefill_runtime.config.page_table_layout is executor.page_table_layout
+    assert executor.decode_runtime.config.page_table_layout is executor.page_table_layout
+    assert executor.warmup.config.page_table_layout is executor.page_table_layout
+    assert executor.prefill_runtime.config.trace_capture_prime_sequence_lengths == (
+        (128, 1024) if binding.executor_module is qwen3_32b_executor else ()
+    )
+
+    def fake_allocate():
+        assert executor._runtime_configuration_sealed
+        assert executor.warmup._configuration_sealed
+        return ["allocated"]
+
+    monkeypatch.setattr(executor.kv_cache_manager, "allocate", fake_allocate)
+    assert executor.allocate_kv_cache() == ["allocated"]
+
+
+def test_late_capacity_failure_is_atomic(binding, expect_error):
+    executor = binding.executor_class(
+        binding.make_model(), binding.make_runtime_config(), binding.make_executor_config()
+    )
+    executor._seal_runtime_configuration()
+    unresolved = executor.kv_cache_manager.config
+    original_layout = executor.page_table_layout
+    original_model_paged = executor.model.layers[0].attention.config.paged_attention_config
+
+    with expect_error(RuntimeError, "runtime configuration is sealed"):
+        executor.configure_paged_kv_cache(
+            PagedKVCacheConfig(
+                block_size=32,
+                max_num_blocks=132,
+                dtype=ttnn.bfloat8_b,
+                num_blocks=64,
+            )
+        )
+
+    assert executor.kv_cache_manager.config is unresolved
+    assert not unresolved.is_resolved()
+    assert executor.page_table_layout is original_layout
+    assert executor.model.layers[0].attention.config.paged_attention_config is original_model_paged
+
+
+def test_generator_resolves_configures_then_allocates_vllm_kv_shape(binding):
+    events = []
+    resolved = object()
+    cache = object()
+    shape = (129, 8, 64, 128)
+    dtype = object()
+    target = SimpleNamespace(
+        configure_paged_kv_cache=lambda config: events.append(("configure", config)),
+        allocate_kv_cache=lambda: events.append(("allocate",)) or cache,
+    )
+    adapter = SimpleNamespace(
+        resolve_legacy_kv_cache_config=lambda *args: events.append(("resolve", args)) or resolved,
+    )
+    generator = binding.generator_class(target, adapter)
+
+    assert generator.allocate_kv_cache(shape, dtype, 32) is cache
+    assert events == [
+        ("resolve", (shape, dtype, 32)),
+        ("configure", resolved),
+        ("allocate",),
+    ]
+
+
+def test_generator_allocates_model_owned_kv_without_reconfiguration(binding):
+    events = []
+    cache = object()
+    target = SimpleNamespace(
+        configure_paged_kv_cache=lambda config: events.append(("configure", config)),
+        allocate_kv_cache=lambda: events.append(("allocate",)) or cache,
+    )
+    adapter = SimpleNamespace(
+        resolve_legacy_kv_cache_config=lambda *args: events.append(("resolve", args)),
+    )
+    generator = binding.generator_class(target, adapter)
+
+    assert generator.allocate_kv_cache() is cache
+    assert events == [("allocate",)]
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    (
+        ((64, 8, 32, 128), None, None),
+        (None, object(), None),
+        (None, None, 32),
+        ((64, 8, 32, 128), object(), None),
+    ),
+)
+def test_generator_rejects_partial_vllm_kv_shape_atomically(binding, arguments, expect_error):
+    events = []
+    target = SimpleNamespace(
+        configure_paged_kv_cache=lambda config: events.append(("configure", config)),
+        allocate_kv_cache=lambda: events.append(("allocate",)),
+    )
+    adapter = SimpleNamespace(
+        resolve_legacy_kv_cache_config=lambda *args: events.append(("resolve", args)),
+    )
+    generator = binding.generator_class(target, adapter)
+
+    with expect_error(TypeError, "must be supplied together"):
+        generator.allocate_kv_cache(*arguments)
+
+    assert events == []
+
+
+def test_generator_does_not_configure_or_allocate_after_vllm_kv_resolution_failure(binding, expect_error):
+    events = []
+
+    def fail_resolution(*args):
+        events.append(("resolve", args))
+        raise ValueError("invalid vLLM KV geometry")
+
+    target = SimpleNamespace(
+        configure_paged_kv_cache=lambda config: events.append(("configure", config)),
+        allocate_kv_cache=lambda: events.append(("allocate",)),
+    )
+    generator = binding.generator_class(
+        target,
+        SimpleNamespace(resolve_legacy_kv_cache_config=fail_resolution),
+    )
+
+    with expect_error(ValueError, "invalid vLLM KV geometry"):
+        generator.allocate_kv_cache((129, 8, 64, 128), object(), 32)
+
+    assert tuple(name for name, *_ in events) == ("resolve",)
+
+
+def test_generator_reports_unmultiplied_per_submesh_token_capacity(binding):
+    assert (
+        binding.generator_class.get_max_tokens_all_users(
+            model_name="ignored",
+            num_devices=8,
+            tt_data_parallel=4,
+            max_model_len=32768,
+            max_num_seqs=64,
+        )
+        == 32768
+    )
+
+
+def test_generator_rejects_unavailable_traced_execution(binding, expect_error):
+    target = binding.make_recording_target()
+    target.traced_decode_execution = None
+    target.config = binding.make_executor_config("none")
+    generator = binding.generator_class(target, binding.generator_module._build_vllm_adapter(target))
+
+    with expect_error(RuntimeError, "unavailable traced decode execution"):
+        generator._select_execution("decode", True)
+
+
+def test_initialize_vllm_model_threads_policy(binding, monkeypatch):
+    captured = []
+    sentinel = object()
+    mesh_device = object()
+    monkeypatch.setattr(
+        binding.generator_module,
+        binding.build_generator_name,
+        lambda config: captured.append(config) or sentinel,
+    )
+
+    result = binding.generator_class.initialize_vllm_model(
+        SimpleNamespace(_name_or_path=binding.hf_model),
+        mesh_device,
+        8,
+        4096,
+        n_layers=3,
+        tt_data_parallel=2,
+        optimizations="accuracy",
+        trace_mode="decode_only",
+        device_sampling_enabled=True,
+    )
+
+    assert result is sentinel
+    config = captured[0]
+    assert isinstance(config, binding.generator_config_class)
+    assert config.hf_model == binding.hf_model
+    assert config.mesh_device is mesh_device
+    assert config.max_batch_size == 8
+    assert config.max_seq_len == 4096
+    assert config.n_layers == 3
+    assert config.tt_data_parallel == 2
+    assert config.optimizations == "accuracy"
+    assert config.trace_mode == "decode_only"
+    assert config.device_sampling_enabled is True
+
+
+@pytest.mark.parametrize("model_id,generator_path", GENERATOR_PATHS.items(), ids=GENERATOR_PATHS)
+def test_vllm_generator_path_and_construction_defaults(model_id, generator_path):
+    module_name, class_name = generator_path.split(":", maxsplit=1)
+    generator_class = getattr(import_module(module_name), class_name)
+
+    assert generator_class is EXECUTOR_BINDINGS[model_id].generator_class
+    assert callable(getattr(generator_class, "initialize_vllm_model", None))
+
+    parameters = inspect.signature(generator_class.initialize_vllm_model).parameters
+    assert parameters["trace_mode"].default == "all"
+    assert parameters["device_sampling_enabled"].default is True
+
+
+class _FakeLane:
+    requires_prefill_trace_warmup = True
+
+    def __init__(self, llm, config, request_state_fields=()):
+        self.model = llm.model
+        self.model_args = llm.runtime_config
+        self.mesh_device = llm.model.config.mesh_device
+        self.cache_path = llm.runtime_config.model_cache_path
+        self.config = config
+        self._request_state_fields = tuple(request_state_fields)
+        self.paged_kv_cache_config = config.paged_kv_cache
+        self.already_warmed_up_prefill = False
+        self.eager_execution = object()
+        self.traced_prefill_execution = object()
+        self.traced_decode_execution = object()
+        self.cleanup_calls = 0
+
+    def cleanup(self):
+        self.cleanup_calls += 1
+
+
+def test_generator_constructs_data_parallel_lane_group(binding, monkeypatch):
+    executor_calls = []
+    built_lanes = []
+    pretrained_calls = []
+    parent_mesh = object()
+    submeshes = [_Mesh(), _Mesh()]
+    create_submeshes = MagicMock(return_value=submeshes)
+    monkeypatch.setattr(binding.generator_module, "_create_submeshes", create_submeshes)
+
+    def fake_from_pretrained(mesh_device, **kwargs):
+        pretrained_calls.append((mesh_device, kwargs))
+        return binding.make_product(mesh_device, kwargs["max_batch_size"])
+
+    def fake_build_executor(llm, config):
+        executor_calls.append((llm, config))
+        lane = binding.make_lane(llm, config)
+        built_lanes.append(lane)
+        return lane
+
+    monkeypatch.setattr(binding.generator_module, "from_pretrained", fake_from_pretrained)
+    monkeypatch.setattr(binding.generator_module, binding.build_executor_name, fake_build_executor)
+    monkeypatch.setattr(
+        binding.generator_module,
+        "_model_kv_metadata",
+        lambda model: ((ttnn.bfloat8_b,), 1, 8, 64),
+    )
+
+    generator = getattr(binding.generator_module, binding.build_generator_name)(
+        binding.generator_config_class(
+            hf_model=binding.hf_model,
+            mesh_device=parent_mesh,
+            max_batch_size=4,
+            max_seq_len=4096,
+            n_layers=1,
+            tt_data_parallel=2,
+            trace_mode="all",
+            device_sampling_enabled=True,
+        )
+    )
+
+    try:
+        create_submeshes.assert_called_once_with(parent_mesh, 2)
+        assert [mesh for mesh, _ in pretrained_calls] == submeshes
+        assert all(call[1]["max_batch_size"] == 2 for call in pretrained_calls)
+        assert all(call[1]["max_seq_len"] == 4096 for call in pretrained_calls)
+        assert all(call[1]["n_layers"] == 1 for call in pretrained_calls)
+        assert isinstance(generator.target, LaneGroupExecutor)
+        assert generator.target.mesh_device is parent_mesh
+        assert generator.target.tt_data_parallel == 2
+        assert len(executor_calls) == 2
+        assert executor_calls[0][0] is not executor_calls[1][0]
+        assert generator.target.lanes == built_lanes
+        assert [lane.model for lane in generator.target.lanes] == [llm.model for llm, _ in executor_calls]
+        assert [lane.mesh_device for lane in generator.target.lanes] == submeshes
+        assert len({id(lane) for lane in generator.target.lanes}) == 2
+        assert all(isinstance(config, binding.executor_config_class) for _, config in executor_calls)
+        assert all(llm.model.config.max_batch_size == 2 for llm, _ in executor_calls)
+        assert generator._adapter.config.trace.mode == "all"
+        assert generator._adapter.config.expected_num_layers == 1
+        assert generator._adapter.config.expected_kv_heads_per_device == 8
+        assert generator._adapter.config.expected_head_dim == 64
+    finally:
+        generator.cleanup()
+
+
+def test_executor_cleanup_is_ordered_retryable_and_idempotent(binding, expect_error):
+    calls = []
+    failures = {"reader", "trace"}
+
+    class _Owner:
+        def __init__(self, name):
+            self.name = name
+
+        def cleanup(self, *args):
+            calls.append(self.name)
+            if self.name in failures:
+                raise RuntimeError(self.name)
+
+        drain = cleanup
+        drain_external_outputs = cleanup
+        cleanup_transients = cleanup
+        release = cleanup
+
+    executor = object.__new__(binding.executor_class)
+    executor._terminal = False
+    executor._cleaned_up = False
+    executor.decode_runtime = _Owner("decode-external")
+    executor.output_reader = _Owner("reader")
+    executor.prefill_runtime = _Owner("prefill")
+    executor.trace_compiler = _Owner("trace")
+    executor.program_compiler = _Owner("program")
+    executor.config = SimpleNamespace(device_sampling_enabled=True)
+    executor.model = SimpleNamespace(sampling=_Owner("sampling"))
+    if binding.executor_module in (llama33_70b_executor, qwen3_32b_executor):
+        executor.sampling_state_controller = _Owner("sampling-state")
+        executor.sampling_state = object()
+    else:
+        executor.sampling_state_controller = None
+        executor.sampling_state = None
+    executor.kv_cache_manager = _Owner("kv")
+
+    with expect_error(RuntimeError, "reader") as raised:
+        executor.cleanup()
+
+    expected_order = [
+        "decode-external",
+        "reader",
+        "prefill",
+        "decode-external",
+        "trace",
+        "program",
+    ]
+    if binding.executor_module in (llama33_70b_executor, qwen3_32b_executor):
+        expected_order.append("sampling-state")
+    expected_order.extend(["sampling", "kv"])
+    assert calls == expected_order
+    assert tuple(error.args[0] for error in raised.value.cleanup_failures) == ("trace",)
+    assert executor.terminal
+    assert not executor._cleaned_up
+
+    failures.clear()
+    executor.cleanup()
+    assert calls == expected_order * 2
+    assert executor._cleaned_up
+
+    executor.cleanup()
+    assert calls == expected_order * 2
+
+
+def test_llama33_generator_emits_runtime_summary_before_owned_cleanup():
+    events = []
+    traced = SimpleNamespace(log_runtime_summary=lambda **kwargs: events.append(("summary", kwargs)))
+    target = SimpleNamespace(
+        traced_executor=traced,
+        cleanup=lambda: events.append("cleanup"),
+    )
+    generator = llama33_70b_generator.Llama33_70BGenerator(target, SimpleNamespace())
+
+    generator.cleanup()
+
+    assert events == [("summary", {"phase": "shutdown"}), "cleanup"]
+
+
+def test_llama33_generator_emits_serving_ready_and_idempotent_shutdown_summaries():
+    phases = []
+    trace_compiler = SimpleNamespace(trace_active=False)
+    traced = SimpleNamespace(
+        trace_compiler=trace_compiler,
+        log_runtime_summary=lambda **kwargs: phases.append(kwargs["phase"]),
+    )
+    target = SimpleNamespace(
+        traced_executor=traced,
+        warmup_model_prefill=lambda **kwargs: None,
+        warmup_model_decode=lambda **kwargs: None,
+        cleanup=lambda: None,
+    )
+    generator = llama33_70b_generator.Llama33_70BGenerator(target, SimpleNamespace())
+
+    generator.warmup_model_prefill(kv_cache="cache", can_sample_on_device=True, enable_trace=True)
+    trace_compiler.trace_active = True
+    generator.warmup_model_decode(
+        kv_cache="cache",
+        max_batch_size=16,
+        num_blocks=128,
+        can_sample_on_device=True,
+        enable_trace=True,
+    )
+    generator.warmup_model_prefill(kv_cache="cache", can_sample_on_device=True, enable_trace=True)
+    generator._shutdown_summary_callback()
+    generator.cleanup()
+
+    assert phases == ["serving_ready", "shutdown"]

code/models/common/tests/llm_runtime/test_lane_group.py

@@ -0,0 +1,1373 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from contextlib import contextmanager
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from models.common.llm_runtime.lane_group import LaneGroupExecutor
+
+
+@dataclass(frozen=True)
+class _SamplingParams:
+    temperature: list[float]
+    top_k: torch.Tensor
+    top_p: float
+
+
+@dataclass(frozen=True)
+class _SingletonSamplingParams:
+    temperature: list[float]
+    top_k: torch.Tensor
+    top_p: tuple[float, ...]
+
+
+@dataclass(frozen=True)
+class _StatefulSamplingParams:
+    temperature: list[float]
+    top_k: torch.Tensor
+    top_p: list[float]
+    presence_penalty: list[float]
+    frequency_penalty: list[float]
+    repetition_penalty: list[float]
+    seed: list[int | None]
+    enable_log_probs: list[bool]
+    num_logprobs: list[int]
+
+
+class _Lane:
+    def __init__(self, lane_idx, *, capacity=2):
+        self.lane_idx = lane_idx
+        self.model = SimpleNamespace(config=SimpleNamespace(max_batch_size=capacity))
+        self.model_args = f"args-{lane_idx}"
+        self.mesh_device = f"mesh-{lane_idx}"
+        self.cache_path = f"cache-{lane_idx}"
+        self.already_warmed_up_prefill = False
+        self.calls = []
+        self.cleanup_calls = 0
+        self.fail_method = None
+        self.cleanup_error = None
+
+    def _call(self, method, kwargs):
+        self.calls.append((method, kwargs))
+        if self.fail_method == method:
+            raise RuntimeError(f"{method} boom {self.lane_idx}")
+
+    def configure_paged_kv_cache(self, config):
+        self._call("configure", {"config": config})
+        self.paged_kv_cache_config = config
+
+    def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None):
+        self._call(
+            "allocate",
+            {
+                "kv_cache_shape": kv_cache_shape,
+                "dtype": dtype,
+                "num_layers": num_layers,
+            },
+        )
+        return f"cache-handle-{self.lane_idx}"
+
+    def compile_prefill(
+        self,
+        *,
+        tokens,
+        page_table,
+        prompt_lens=None,
+        start_pos=None,
+        empty_slots=None,
+        kv_cache=None,
+        sampling_params=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        slot_remap=None,
+        execution=None,
+    ):
+        self._call(
+            "compile_prefill",
+            {
+                "tokens": tokens,
+                "page_table": page_table,
+                "prompt_lens": prompt_lens,
+                "start_pos": start_pos,
+                "empty_slots": empty_slots,
+                "kv_cache": kv_cache,
+                "sampling_params": sampling_params,
+                "prompt_tokens": prompt_tokens,
+                "output_tokens": output_tokens,
+                "slot_remap": slot_remap,
+                "execution": execution,
+            },
+        )
+
+    def compile_decode(
+        self,
+        *,
+        tokens,
+        start_pos,
+        page_table,
+        kv_cache=None,
+        sampling_params=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        slot_remap=None,
+        reset_batch=False,
+        execution=None,
+    ):
+        self._call(
+            "compile_decode",
+            {
+                "tokens": tokens,
+                "start_pos": start_pos,
+                "page_table": page_table,
+                "kv_cache": kv_cache,
+                "sampling_params": sampling_params,
+                "prompt_tokens": prompt_tokens,
+                "output_tokens": output_tokens,
+                "slot_remap": slot_remap,
+                "reset_batch": reset_batch,
+                "execution": execution,
+            },
+        )
+
+    def warmup_model_prefill(self, *, kv_cache, can_sample_on_device, enable_trace):
+        self._call(
+            "warmup_prefill",
+            {
+                "kv_cache": kv_cache,
+                "can_sample_on_device": can_sample_on_device,
+                "enable_trace": enable_trace,
+            },
+        )
+        self.already_warmed_up_prefill = True
+
+    def warmup_model_decode(
+        self,
+        *,
+        kv_cache,
+        max_batch_size,
+        num_blocks,
+        can_sample_on_device,
+        enable_trace,
+    ):
+        self._call(
+            "warmup_decode",
+            {
+                "kv_cache": kv_cache,
+                "max_batch_size": max_batch_size,
+                "num_blocks": num_blocks,
+                "can_sample_on_device": can_sample_on_device,
+                "enable_trace": enable_trace,
+            },
+        )
+
+    def prefill_forward(
+        self,
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,
+        start_pos=None,
+        empty_slots=None,
+        kv_cache=None,
+        sampling_params=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        slot_remap=None,
+        execution=None,
+    ):
+        kwargs = {
+            "tokens": tokens,
+            "page_table": page_table,
+            "prompt_lens": prompt_lens,
+            "start_pos": start_pos,
+            "empty_slots": empty_slots,
+            "kv_cache": kv_cache,
+            "sampling_params": sampling_params,
+            "prompt_tokens": prompt_tokens,
+            "output_tokens": output_tokens,
+            "slot_remap": slot_remap,
+            "execution": execution,
+        }
+        self._call("prefill", kwargs)
+        values = tokens[:, 0] + self.lane_idx * 100
+        if sampling_params is not None:
+            return values.to(torch.int64), None
+        return values.float().view(-1, 1, 1)
+
+    def decode_forward(
+        self,
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,
+        sampling_params=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        slot_remap=None,
+        reset_batch=False,
+        read_from_device=True,
+        execution=None,
+    ):
+        kwargs = {
+            "tokens": tokens,
+            "start_pos": start_pos,
+            "page_table": page_table,
+            "kv_cache": kv_cache,
+            "sampling_params": sampling_params,
+            "prompt_tokens": prompt_tokens,
+            "output_tokens": output_tokens,
+            "slot_remap": slot_remap,
+            "reset_batch": reset_batch,
+            "read_from_device": read_from_device,
+            "execution": execution,
+        }
+        self._call("decode", kwargs)
+        values = tokens + self.lane_idx * 100
+        if not read_from_device:
+            return f"raw-{self.lane_idx}", None
+        if sampling_params is not None:
+            return values.to(torch.int64), None
+        return values.float().view(-1, 1, 1), None
+
+    def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None, empty_slots=None):
+        self._call(
+            "can_trace_prefill",
+            {
+                "tokens": tokens,
+                "prompt_lens": prompt_lens,
+                "start_pos": start_pos,
+                "empty_slots": empty_slots,
+            },
+        )
+        return False
+
+    def read_decode_output(self, tt_out, *, async_read=False):
+        self._call("read", {"tt_out": tt_out, "async_read": async_read})
+        host = (torch.tensor([self.lane_idx * 2, self.lane_idx * 2 + 1]), None)
+        if async_read:
+            return host, [f"event-{self.lane_idx}"]
+        return host
+
+    def process_decode_output_host(self, tt_out, *, is_tokens=False):
+        self._call("process", {"tt_out": tt_out, "is_tokens": is_tokens})
+        if is_tokens:
+            return torch.tensor([self.lane_idx * 2, self.lane_idx * 2 + 1], dtype=torch.int64), None
+        return torch.full((2, 1, 3), float(self.lane_idx)), None
+
+    def cleanup(self):
+        self.cleanup_calls += 1
+        if self.cleanup_error is not None:
+            raise self.cleanup_error
+
+
+def _sampling():
+    return _SamplingParams(
+        temperature=[0.1, 0.2, 0.3, 0.4],
+        top_k=torch.tensor([1, 2, 3, 4]),
+        top_p=0.9,
+    )
+
+
+_POSITIONAL = inspect.Parameter.POSITIONAL_OR_KEYWORD
+_KEYWORD_ONLY = inspect.Parameter.KEYWORD_ONLY
+_REQUIRED = inspect.Parameter.empty
+
+_PUBLIC_SIGNATURES = {
+    "allocate_kv_cache": (
+        ("kv_cache_shape", _POSITIONAL, None),
+        ("dtype", _POSITIONAL, None),
+        ("num_layers", _POSITIONAL, None),
+    ),
+    "compile_prefill": (
+        ("tokens", _POSITIONAL, _REQUIRED),
+        ("page_table", _POSITIONAL, _REQUIRED),
+        ("prompt_lens", _KEYWORD_ONLY, None),
+        ("start_pos", _KEYWORD_ONLY, None),
+        ("empty_slots", _KEYWORD_ONLY, None),
+        ("kv_cache", _KEYWORD_ONLY, None),
+        ("sampling_params", _KEYWORD_ONLY, None),
+        ("prompt_tokens", _KEYWORD_ONLY, None),
+        ("output_tokens", _KEYWORD_ONLY, None),
+        ("slot_remap", _KEYWORD_ONLY, None),
+        ("execution", _KEYWORD_ONLY, None),
+    ),
+    "compile_decode": (
+        ("tokens", _POSITIONAL, _REQUIRED),
+        ("start_pos", _POSITIONAL, _REQUIRED),
+        ("page_table", _POSITIONAL, _REQUIRED),
+        ("kv_cache", _KEYWORD_ONLY, None),
+        ("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),
+        ("execution", _KEYWORD_ONLY, None),
+    ),
+    "warmup_model_prefill": (
+        ("kv_cache", _KEYWORD_ONLY, _REQUIRED),
+        ("can_sample_on_device", _KEYWORD_ONLY, _REQUIRED),
+        ("enable_trace", _KEYWORD_ONLY, _REQUIRED),
+    ),
+    "warmup_model_decode": (
+        ("kv_cache", _KEYWORD_ONLY, _REQUIRED),
+        ("max_batch_size", _KEYWORD_ONLY, _REQUIRED),
+        ("num_blocks", _KEYWORD_ONLY, _REQUIRED),
+        ("can_sample_on_device", _KEYWORD_ONLY, _REQUIRED),
+        ("enable_trace", _KEYWORD_ONLY, _REQUIRED),
+    ),
+    "prefill_forward": (
+        ("tokens", _POSITIONAL, _REQUIRED),
+        ("page_table", _POSITIONAL, _REQUIRED),
+        ("prompt_lens", _KEYWORD_ONLY, None),
+        ("start_pos", _KEYWORD_ONLY, None),
+        ("empty_slots", _KEYWORD_ONLY, None),
+        ("kv_cache", _KEYWORD_ONLY, None),
+        ("sampling_params", _KEYWORD_ONLY, None),
+        ("prompt_tokens", _KEYWORD_ONLY, None),
+        ("output_tokens", _KEYWORD_ONLY, None),
+        ("slot_remap", _KEYWORD_ONLY, None),
+        ("execution", _KEYWORD_ONLY, None),
+    ),
+    "can_trace_prefill": (
+        ("tokens", _KEYWORD_ONLY, _REQUIRED),
+        ("prompt_lens", _KEYWORD_ONLY, None),
+        ("start_pos", _KEYWORD_ONLY, None),
+        ("empty_slots", _KEYWORD_ONLY, None),
+    ),
+    "decode_forward": (
+        ("tokens", _POSITIONAL, _REQUIRED),
+        ("start_pos", _POSITIONAL, _REQUIRED),
+        ("page_table", _POSITIONAL, _REQUIRED),
+        ("kv_cache", _KEYWORD_ONLY, None),
+        ("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),
+        ("read_from_device", _KEYWORD_ONLY, True),
+        ("execution", _KEYWORD_ONLY, None),
+    ),
+    "read_decode_output": (
+        ("tt_out", _POSITIONAL, _REQUIRED),
+        ("async_read", _KEYWORD_ONLY, False),
+    ),
+    "process_decode_output_host": (
+        ("tt_out", _POSITIONAL, _REQUIRED),
+        ("is_tokens", _KEYWORD_ONLY, False),
+    ),
+}
+
+
+@pytest.mark.parametrize(("method_name", "expected"), _PUBLIC_SIGNATURES.items())
+def test_public_api_signatures_are_exact(method_name, expected):
+    signature = inspect.signature(getattr(LaneGroupExecutor, method_name))
+    parameters = tuple(signature.parameters.values())[1:]
+
+    assert tuple((parameter.name, parameter.kind, parameter.default) for parameter in parameters) == expected
+    assert all(parameter.annotation is not inspect.Parameter.empty for parameter in parameters)
+    assert signature.return_annotation is not inspect.Signature.empty
+
+
+def test_prefill_routes_global_slots_and_restores_source_row_order():
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+    tokens = torch.tensor([[10], [11], [12]])
+    page_table = torch.tensor([[0], [1], [2]], dtype=torch.int32)
+
+    output = group.prefill_forward(
+        tokens,
+        page_table,
+        empty_slots=[3, 0, 2],
+        prompt_lens=torch.tensor([1, 1, 1]),
+        sampling_params=_sampling(),
+        prompt_tokens=torch.tensor([[100], [101], [102], [103]]),
+        output_tokens=[[200], [201], [202], [203]],
+        slot_remap=torch.tensor([0, 1, 2, 3]),
+        kv_cache=["kv-0", "kv-1"],
+    )
+
+    assert isinstance(output, tuple)
+    assert output[0].tolist() == [110, 11, 112]
+    assert output[0].dtype == torch.int64
+    assert output[1] is None
+
+    lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "prefill")
+    lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "prefill")
+    assert lane0_kwargs["empty_slots"] == [0]
+    assert lane1_kwargs["empty_slots"] == [1, 0]
+    assert lane0_kwargs["tokens"].flatten().tolist() == [11]
+    assert lane1_kwargs["tokens"].flatten().tolist() == [10, 12]
+    assert lane0_kwargs["sampling_params"].temperature == [0.2]
+    assert lane1_kwargs["sampling_params"].temperature == [0.1, 0.3]
+    assert lane0_kwargs["sampling_params"].top_k.tolist() == [2]
+    assert lane1_kwargs["sampling_params"].top_k.tolist() == [1, 3]
+    assert lane0_kwargs["sampling_params"].top_p == 0.9
+    assert lane0_kwargs["prompt_tokens"].flatten().tolist() == [100, -1]
+    assert lane1_kwargs["prompt_tokens"].flatten().tolist() == [102, 103]
+    assert lane0_kwargs["output_tokens"] == [[200], [-1]]
+    assert lane1_kwargs["output_tokens"] == [[202], [203]]
+    assert lane0_kwargs["slot_remap"].tolist() == [0, 1]
+    assert lane1_kwargs["slot_remap"].tolist() == [0, 1]
+    assert lane0_kwargs["kv_cache"] == "kv-0"
+    assert lane1_kwargs["kv_cache"] == "kv-1"
+
+
+def test_prefill_log_probs_restore_source_rows_and_fill_nonrequesting_lane():
+    lanes = [_Lane(0), _Lane(1)]
+
+    def prefill_with_log_probs(
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,
+        start_pos=None,
+        empty_slots=None,
+        kv_cache=None,
+        sampling_params=None,
+        execution=None,
+    ):
+        assert tokens.shape[0] == 2
+        return torch.zeros(2, dtype=torch.int64), torch.tensor([0.25, 0.75])
+
+    lanes[1].prefill_forward = prefill_with_log_probs
+    group = LaneGroupExecutor(lanes)
+
+    output, log_probs = group.prefill_forward(
+        tokens=torch.tensor([[10], [11], [12]]),
+        page_table=torch.tensor([[0], [1], [2]], dtype=torch.int32),
+        empty_slots=[3, 0, 2],
+        sampling_params=_sampling(),
+    )
+
+    assert output.tolist() == [0, 11, 0]
+    assert isinstance(log_probs, torch.Tensor)
+    assert log_probs.tolist() == [0.25, 1.0, 0.75]
+
+
+def test_decode_slices_fixed_lane_capacity_and_normalizes_sampled_tokens():
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+
+    output = group.decode_forward(
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([1, 2, 3, 4]),
+        torch.arange(4, dtype=torch.int32).view(4, 1),
+        sampling_params=_sampling(),
+        prompt_tokens=torch.tensor([[100], [101], [102], [103]]),
+        output_tokens=[[200], [201], [202], [203]],
+        slot_remap=torch.tensor([1, 0, 3, 2]),
+        kv_cache=["kv-0", "kv-1"],
+        reset_batch=True,
+    )
+
+    assert output[0].tolist() == [10, 11, 112, 113]
+    assert output[0].dtype == torch.int64
+    assert output[1] is None
+    lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "decode")
+    lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "decode")
+    assert lane0_kwargs["tokens"].tolist() == [10, 11]
+    assert lane1_kwargs["tokens"].tolist() == [12, 13]
+    assert lane0_kwargs["sampling_params"].temperature == [0.1, 0.2]
+    assert lane1_kwargs["sampling_params"].temperature == [0.3, 0.4]
+    assert lane0_kwargs["prompt_tokens"].flatten().tolist() == [100, 101]
+    assert lane1_kwargs["prompt_tokens"].flatten().tolist() == [102, 103]
+    assert lane0_kwargs["output_tokens"] == [[200], [201]]
+    assert lane1_kwargs["output_tokens"] == [[202], [203]]
+    assert lane0_kwargs["slot_remap"].tolist() == [1, 0]
+    assert lane1_kwargs["slot_remap"].tolist() == [1, 0]
+    assert lane0_kwargs["reset_batch"] is lane1_kwargs["reset_batch"] is True
+
+
+@pytest.mark.parametrize(
+    (
+        "layout",
+        "prefill_tokens",
+        "empty_slots",
+        "prefill_sampling",
+        "decode_tokens",
+        "decode_positions",
+        "decode_sampling",
+        "slot_remap",
+        "expected_prefill_tokens",
+        "expected_decode_tokens",
+        "expected_lane_slots",
+        "expected_lane_seeds",
+    ),
+    [
+        pytest.param(
+            "front_packed_gathered",
+            torch.tensor([[10], [11], [12], [13]]),
+            [0, 1, 2, 3],
+            _StatefulSamplingParams(
+                temperature=[0.0, 0.8, 1.0, 1.2],
+                top_k=torch.tensor([1, 5, 7, 9]),
+                top_p=[0.0, 0.7, 0.8, 0.9],
+                presence_penalty=[0.0, 0.1, 0.2, 0.3],
+                frequency_penalty=[0.0, 0.4, 0.5, 0.6],
+                repetition_penalty=[1.0, 1.1, 1.2, 1.3],
+                seed=[100, 101, 102, 103],
+                enable_log_probs=[False, True, False, True],
+                num_logprobs=[0, 0, 0, 0],
+            ),
+            torch.tensor([10, 11, 12, 13]),
+            torch.tensor([1, 2, 3, 4]),
+            _StatefulSamplingParams(
+                temperature=[0.0, 0.8, 1.0, 1.2],
+                top_k=torch.tensor([1, 5, 7, 9]),
+                top_p=[0.0, 0.7, 0.8, 0.9],
+                presence_penalty=[0.0, 0.1, 0.2, 0.3],
+                frequency_penalty=[0.0, 0.4, 0.5, 0.6],
+                repetition_penalty=[1.0, 1.1, 1.2, 1.3],
+                seed=[100, 101, 102, 103],
+                enable_log_probs=[False, True, False, True],
+                num_logprobs=[0, 0, 0, 0],
+            ),
+            torch.tensor([1, 0, 3, 2]),
+            [10, 11, 112, 113],
+            [10, 11, 112, 113],
+            ([0, 1], [0, 1]),
+            ([100, 101], [102, 103]),
+            id="front_packed_gathered",
+        ),
+        pytest.param(
+            "stable_gap_lane",
+            torch.tensor([[11], [13]]),
+            [1, 3],
+            _StatefulSamplingParams(
+                temperature=[0.8, 1.2],
+                top_k=torch.tensor([5, 9]),
+                top_p=[0.7, 0.9],
+                presence_penalty=[0.1, 0.3],
+                frequency_penalty=[0.4, 0.6],
+                repetition_penalty=[1.1, 1.3],
+                seed=[101, 103],
+                enable_log_probs=[True, True],
+                num_logprobs=[0, 0],
+            ),
+            torch.tensor([0, 11, 0, 13]),
+            torch.tensor([-1, 2, -1, 4]),
+            _StatefulSamplingParams(
+                temperature=[0.0, 0.8, 0.0, 1.2],
+                top_k=torch.tensor([1, 5, 1, 9]),
+                top_p=[0.0, 0.7, 0.0, 0.9],
+                presence_penalty=[0.0, 0.1, 0.0, 0.3],
+                frequency_penalty=[0.0, 0.4, 0.0, 0.6],
+                repetition_penalty=[1.0, 1.1, 1.0, 1.3],
+                seed=[None, 101, None, 103],
+                enable_log_probs=[False, True, False, True],
+                num_logprobs=[0, 0, 0, 0],
+            ),
+            torch.tensor([0, 1, 2, 3]),
+            [11, 113],
+            [0, 11, 100, 113],
+            ([1], [1]),
+            ([None, 101], [None, 103]),
+            id="stable_gap_lane",
+        ),
+    ],
+)
+def test_sampling_state_boundary_contract_for_gathered_and_lane_layouts(
+    layout,
+    prefill_tokens,
+    empty_slots,
+    prefill_sampling,
+    decode_tokens,
+    decode_positions,
+    decode_sampling,
+    slot_remap,
+    expected_prefill_tokens,
+    expected_decode_tokens,
+    expected_lane_slots,
+    expected_lane_seeds,
+):
+    """Validate layouts already constructed at the tt-metal target boundary.
+
+    ``front_packed_gathered`` models the dense, rank-segmented payload supplied
+    after gathered-DP packing. ``stable_gap_lane`` models the fixed stable-slot
+    grid supplied by lane-DP. This deliberately does not claim to test the vLLM
+    code that constructs or gathers either payload.
+    """
+
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+    prompt_tokens = torch.tensor([[200, -1], [210, 211], [220, -1], [230, 231]])
+    output_tokens = [[300, -1], [310, 311], [320, -1], [330, 331]]
+    execution = [f"trace-{layout}-0", f"trace-{layout}-1"]
+
+    prefill_result = group.prefill_forward(
+        prefill_tokens,
+        torch.arange(len(prefill_tokens), dtype=torch.int32).view(-1, 1),
+        empty_slots=empty_slots,
+        sampling_params=prefill_sampling,
+        prompt_tokens=prompt_tokens,
+        output_tokens=output_tokens,
+        slot_remap=slot_remap,
+        execution=execution,
+    )
+    assert prefill_result[0].tolist() == expected_prefill_tokens
+
+    decode_result = group.decode_forward(
+        decode_tokens,
+        decode_positions,
+        torch.arange(4, dtype=torch.int32).view(4, 1),
+        sampling_params=decode_sampling,
+        prompt_tokens=prompt_tokens,
+        output_tokens=output_tokens,
+        slot_remap=slot_remap,
+        reset_batch=True,
+        execution=execution,
+    )
+    assert decode_result[0].tolist() == expected_decode_tokens
+
+    if layout == "front_packed_gathered":
+        prefill_rows = ([0, 1], [2, 3])
+        prefill_prompt = (
+            prompt_tokens[0:2].tolist(),
+            prompt_tokens[2:4].tolist(),
+        )
+        prefill_output = (output_tokens[0:2], output_tokens[2:4])
+        lane_remap = ([1, 0], [1, 0])
+    else:
+        prefill_rows = ([0], [1])
+        prefill_prompt = (
+            [[-1, -1], prompt_tokens[1].tolist()],
+            [[-1, -1], prompt_tokens[3].tolist()],
+        )
+        prefill_output = (
+            [[-1, -1], output_tokens[1]],
+            [[-1, -1], output_tokens[3]],
+        )
+        lane_remap = ([0, 1], [0, 1])
+
+    for lane_idx, lane in enumerate(lanes):
+        prefill = next(kwargs for method, kwargs in lane.calls if method == "prefill")
+        decode = next(kwargs for method, kwargs in lane.calls if method == "decode")
+        source_rows = prefill_rows[lane_idx]
+        prefill_params = prefill["sampling_params"]
+        decode_params = decode["sampling_params"]
+        assert prefill["empty_slots"] == expected_lane_slots[lane_idx]
+        assert prefill["execution"] == execution[lane_idx]
+        assert prefill_params.temperature == [prefill_sampling.temperature[row] for row in source_rows]
+        assert prefill_params.top_k.tolist() == [int(prefill_sampling.top_k[row]) for row in source_rows]
+        assert prefill_params.top_p == [prefill_sampling.top_p[row] for row in source_rows]
+        assert prefill_params.presence_penalty == [prefill_sampling.presence_penalty[row] for row in source_rows]
+        assert prefill_params.frequency_penalty == [prefill_sampling.frequency_penalty[row] for row in source_rows]
+        assert prefill_params.repetition_penalty == [prefill_sampling.repetition_penalty[row] for row in source_rows]
+        assert prefill_params.seed == [prefill_sampling.seed[row] for row in source_rows]
+        assert prefill_params.enable_log_probs == [prefill_sampling.enable_log_probs[row] for row in source_rows]
+        assert prefill_params.num_logprobs == [prefill_sampling.num_logprobs[row] for row in source_rows]
+        assert prefill["prompt_tokens"].tolist() == prefill_prompt[lane_idx]
+        assert prefill["output_tokens"] == prefill_output[lane_idx]
+        assert prefill["slot_remap"].tolist() == lane_remap[lane_idx]
+        assert decode["execution"] == execution[lane_idx]
+        lane_slice = slice(lane_idx * 2, lane_idx * 2 + 2)
+        assert decode_params.temperature == decode_sampling.temperature[lane_slice]
+        assert decode_params.top_k.tolist() == decode_sampling.top_k[lane_slice].tolist()
+        assert decode_params.top_p == decode_sampling.top_p[lane_slice]
+        assert decode_params.presence_penalty == decode_sampling.presence_penalty[lane_slice]
+        assert decode_params.frequency_penalty == decode_sampling.frequency_penalty[lane_slice]
+        assert decode_params.repetition_penalty == decode_sampling.repetition_penalty[lane_slice]
+        assert decode_params.seed == expected_lane_seeds[lane_idx]
+        assert decode_params.enable_log_probs == decode_sampling.enable_log_probs[lane_slice]
+        assert decode_params.num_logprobs == decode_sampling.num_logprobs[lane_slice]
+        assert decode["prompt_tokens"].tolist() == prompt_tokens[lane_idx * 2 : lane_idx * 2 + 2].tolist()
+        assert decode["output_tokens"] == output_tokens[lane_idx * 2 : lane_idx * 2 + 2]
+        assert decode["slot_remap"].tolist() == lane_remap[lane_idx]
+        assert decode["reset_batch"] is True
+
+    raw = group.decode_forward(
+        decode_tokens,
+        decode_positions,
+        torch.arange(4, dtype=torch.int32).view(4, 1),
+        sampling_params=decode_sampling,
+        read_from_device=False,
+    )
+    host_outputs, events = group.read_decode_output(raw, async_read=True)
+    completed = group.process_decode_output_host(host_outputs, is_tokens=True)
+    assert events == ["event-0", "event-1"]
+    assert completed[0].tolist() == [0, 1, 2, 3]
+    assert completed[1] is None
+
+
+def test_decode_broadcasts_singleton_sampling_values_to_later_dp_lanes():
+    lanes = [_Lane(0, capacity=1), _Lane(1, capacity=1)]
+    group = LaneGroupExecutor(lanes)
+    sampling_params = _SingletonSamplingParams(
+        temperature=[0.7],
+        top_k=torch.tensor([8]),
+        top_p=(0.4,),
+    )
+
+    output = group.decode_forward(
+        tokens=torch.tensor([10, 11]),
+        start_pos=torch.tensor([1, 2]),
+        page_table=torch.arange(2, dtype=torch.int32).view(2, 1),
+        sampling_params=sampling_params,
+    )
+
+    assert output[0].tolist() == [10, 111]
+    lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "decode")
+    lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "decode")
+    for lane_kwargs in (lane0_kwargs, lane1_kwargs):
+        sliced = lane_kwargs["sampling_params"]
+        assert sliced.temperature == [0.7]
+        assert sliced.top_k.tolist() == [8]
+        assert sliced.top_p == (0.4,)
+
+
+def test_decode_without_readback_returns_lane_local_outputs():
+    group = LaneGroupExecutor([_Lane(0), _Lane(1)])
+
+    output = group.decode_forward(
+        tokens=torch.tensor([0, 1, 2, 3]),
+        start_pos=torch.tensor([0, 0, 0, 0]),
+        page_table=torch.zeros((4, 1), dtype=torch.int32),
+        read_from_device=False,
+    )
+
+    assert output == [("raw-0", None), ("raw-1", None)]
+
+
+def test_async_read_and_host_processing_run_per_lane_and_preserve_lane_order():
+    group = LaneGroupExecutor([_Lane(0), _Lane(1)])
+
+    host_outputs, events = group.read_decode_output(
+        tt_out=[("raw-0", None), ("raw-1", None)],
+        async_read=True,
+    )
+    processed = group.process_decode_output_host(tt_out=host_outputs, is_tokens=True)
+
+    assert events == ["event-0", "event-1"]
+    assert processed[0].tolist() == [0, 1, 2, 3]
+    assert processed[0].dtype == torch.int64
+    assert processed[1] is None
+
+
+def test_warmup_replicates_lane_local_case_and_cache_to_every_lane():
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+
+    group.warmup_model_prefill(
+        kv_cache=["kv-0", "kv-1"],
+        can_sample_on_device=True,
+        enable_trace=False,
+    )
+    group.warmup_model_decode(
+        kv_cache=["kv-0", "kv-1"],
+        enable_trace=False,
+        max_batch_size=4,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    for lane_idx, lane in enumerate(lanes):
+        prefill = next(kwargs for method, kwargs in lane.calls if method == "warmup_prefill")
+        decode = next(kwargs for method, kwargs in lane.calls if method == "warmup_decode")
+        assert prefill["kv_cache"] == f"kv-{lane_idx}"
+        assert decode["kv_cache"] == f"kv-{lane_idx}"
+        assert decode["max_batch_size"] == 2
+    assert group.already_warmed_up_prefill
+
+
+def test_traced_warmup_without_coordinators_fails_before_invoking_a_lane(expect_error):
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+
+    with expect_error(RuntimeError, "Every DP lane must expose the trace activation barrier"):
+        group.warmup_model_decode(
+            kv_cache=["kv-0", "kv-1"],
+            enable_trace=True,
+            max_batch_size=4,
+            num_blocks=8,
+            can_sample_on_device=False,
+        )
+
+    assert [lane.calls for lane in lanes] == [[], []]
+
+
+class _DeferredCapture:
+    def __init__(self, lane_idx, events, *, ready=True):
+        self.lane_idx = lane_idx
+        self.events = events
+        self.ready = ready
+        self.capture_pending = False
+        self.trace_activated = False
+        self.capture_calls = 0
+
+    @contextmanager
+    def defer_capture(self):
+        self.events.append(("enter", self.lane_idx))
+        try:
+            yield self
+        finally:
+            self.capture_pending = False
+            self.events.append(("exit", self.lane_idx))
+
+    def register(self):
+        self.events.append(("register", self.lane_idx))
+        self.capture_pending = self.ready
+
+    def activate_pending_capture(self):
+        assert self.capture_pending
+        self.events.append(("capture", self.lane_idx))
+        self.capture_calls += 1
+        self.trace_activated = True
+        self.capture_pending = False
+
+
+class _BarrierLane(_Lane):
+    def __init__(self, lane_idx, events, *, ready=True, fail=False):
+        super().__init__(lane_idx)
+        self.warmup = _DeferredCapture(lane_idx, events, ready=ready)
+        self.events = events
+        self.fail = fail
+
+    def warmup_model_decode(self, **kwargs):
+        self.events.append(("warmup", self.lane_idx))
+        if kwargs["enable_trace"]:
+            self.warmup.register()
+        if self.fail:
+            raise RuntimeError(f"warmup boom {self.lane_idx}")
+        return super().warmup_model_decode(**kwargs)
+
+
+def test_traced_warmup_registers_every_lane_before_any_capture():
+    events = []
+    lanes = [_BarrierLane(0, events), _BarrierLane(1, events)]
+    group = LaneGroupExecutor(lanes)
+
+    group.warmup_model_decode(
+        kv_cache=["kv-0", "kv-1"],
+        enable_trace=True,
+        max_batch_size=4,
+        num_blocks=8,
+        can_sample_on_device=False,
+    )
+
+    assert events == [
+        ("enter", 0),
+        ("enter", 1),
+        ("warmup", 0),
+        ("register", 0),
+        ("warmup", 1),
+        ("register", 1),
+        ("capture", 0),
+        ("capture", 1),
+        ("exit", 1),
+        ("exit", 0),
+    ]
+
+
+@pytest.mark.parametrize(
+    ("mode", "message"),
+    [
+        ("exception", "warmup boom 1"),
+        ("asymmetry", "mixed trace activation readiness"),
+    ],
+)
+def test_traced_warmup_failure_or_asymmetry_captures_no_lane(mode, message, expect_error):
+    events = []
+    group_lanes = [
+        _BarrierLane(0, events, ready=True),
+        _BarrierLane(1, events, ready=mode != "asymmetry", fail=mode == "exception"),
+    ]
+    group = LaneGroupExecutor(group_lanes)
+
+    with expect_error(RuntimeError, message):
+        group.warmup_model_decode(
+            kv_cache=["kv-0", "kv-1"],
+            enable_trace=True,
+            max_batch_size=4,
+            num_blocks=8,
+            can_sample_on_device=False,
+        )
+
+    assert [lane.warmup.capture_calls for lane in group_lanes] == [0, 0]
+    assert not any(lane.warmup.capture_pending for lane in group_lanes)
+
+
+def test_eager_warmup_does_not_enter_trace_activation_deferral():
+    events = []
+    lanes = [_BarrierLane(0, events), _BarrierLane(1, events)]
+    group = LaneGroupExecutor(lanes)
+
+    group.warmup_model_decode(
+        kv_cache=["kv-0", "kv-1"],
+        enable_trace=False,
+        max_batch_size=4,
+        num_blocks=8,
+        can_sample_on_device=False,
+    )
+
+    assert events == [
+        ("warmup", 0),
+        ("warmup", 1),
+    ]
+    assert [lane.warmup.capture_calls for lane in lanes] == [0, 0]
+
+
+def test_compile_methods_slice_requests_to_lane_executors():
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+
+    group.compile_prefill(
+        torch.tensor([[10], [11], [12]]),
+        torch.tensor([[0], [1], [2]], dtype=torch.int32),
+        empty_slots=[3, 0, 2],
+        kv_cache=["kv-0", "kv-1"],
+    )
+    group.compile_decode(
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([1, 2, 3, 4]),
+        torch.arange(4, dtype=torch.int32).view(4, 1),
+        kv_cache=["kv-0", "kv-1"],
+    )
+
+    lane0_prefill = next(kwargs for method, kwargs in lanes[0].calls if method == "compile_prefill")
+    lane1_prefill = next(kwargs for method, kwargs in lanes[1].calls if method == "compile_prefill")
+    lane0_decode = next(kwargs for method, kwargs in lanes[0].calls if method == "compile_decode")
+    lane1_decode = next(kwargs for method, kwargs in lanes[1].calls if method == "compile_decode")
+    assert lane0_prefill["tokens"].flatten().tolist() == [11]
+    assert lane1_prefill["tokens"].flatten().tolist() == [10, 12]
+    assert lane0_decode["tokens"].tolist() == [10, 11]
+    assert lane1_decode["tokens"].tolist() == [12, 13]
+
+
+def test_concrete_execution_targets_preflight_every_traced_lane_before_any_execution(expect_error):
+    class DispatchLane(_Lane):
+        def __init__(self, lane_idx, *, traceable):
+            super().__init__(lane_idx)
+            self.traceable = traceable
+            self.eager_execution = object()
+            self.traced_prefill_execution = object()
+            self.traced_decode_execution = object()
+
+        def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None):
+            self._call(
+                "can_trace_prefill",
+                {
+                    "tokens": tokens,
+                    "prompt_lens": prompt_lens,
+                    "start_pos": start_pos,
+                },
+            )
+            return self.traceable
+
+        def prefill_forward(
+            self,
+            tokens,
+            page_table,
+            *,
+            prompt_lens=None,
+            start_pos=None,
+            empty_slots=None,
+            kv_cache=None,
+            sampling_params=None,
+            execution=None,
+        ):
+            kwargs = {
+                "tokens": tokens,
+                "page_table": page_table,
+                "prompt_lens": prompt_lens,
+                "start_pos": start_pos,
+                "empty_slots": empty_slots,
+                "kv_cache": kv_cache,
+                "sampling_params": sampling_params,
+                "execution": execution,
+            }
+            self._call("prefill", kwargs)
+            return tokens[:, :1].float().view(-1, 1, 1)
+
+    lanes = [DispatchLane(0, traceable=True), DispatchLane(1, traceable=False)]
+    group = LaneGroupExecutor(lanes)
+    tokens = torch.tensor([[10], [11]])
+    page_table = torch.tensor([[0], [1]], dtype=torch.int32)
+    prompt_lens = torch.tensor([7, 9])
+    start_pos = torch.tensor([3, 4])
+    empty_slots = [0, 2]
+
+    assert not group.can_trace_prefill(
+        tokens=tokens,
+        prompt_lens=prompt_lens,
+        start_pos=start_pos,
+        empty_slots=empty_slots,
+    )
+    kwargs = {
+        "tokens": tokens,
+        "page_table": page_table,
+        "prompt_lens": prompt_lens,
+        "start_pos": start_pos,
+        "empty_slots": empty_slots,
+    }
+    assert group.prefill_forward(execution=group.eager_execution, **kwargs).flatten().tolist() == [10.0, 11.0]
+    for lane_idx, lane in enumerate(lanes):
+        eager_call = next(call for call in lane.calls if call[0] == "prefill")
+        assert eager_call[1]["execution"] is group.eager_execution[lane_idx]
+        lane.calls.clear()
+
+    with expect_error(RuntimeError, "no DP lane executed"):
+        group.prefill_forward(execution=group.traced_prefill_execution, **kwargs)
+
+    assert [[method for method, _ in lane.calls] for lane in lanes] == [
+        ["can_trace_prefill"],
+        ["can_trace_prefill"],
+    ]
+    for lane_idx, lane in enumerate(lanes):
+        trace_kwargs = lane.calls[0][1]
+        assert set(trace_kwargs) == {"tokens", "prompt_lens", "start_pos"}
+        assert trace_kwargs["tokens"].flatten().tolist() == [10 + lane_idx]
+        assert trace_kwargs["prompt_lens"].tolist() == [7 + 2 * lane_idx]
+        assert trace_kwargs["start_pos"].tolist() == [3 + lane_idx]
+
+
+def test_dp1_trace_classification_slices_only_the_exact_lane_subset():
+    class TraceLane(_Lane):
+        def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None):
+            self._call(
+                "can_trace_prefill",
+                {
+                    "tokens": tokens,
+                    "prompt_lens": prompt_lens,
+                    "start_pos": start_pos,
+                },
+            )
+            return True
+
+    lane = TraceLane(0)
+    group = LaneGroupExecutor([lane])
+
+    assert group.can_trace_prefill(
+        tokens=torch.tensor([[10], [11]]),
+        prompt_lens=torch.tensor([7, 9]),
+        start_pos=torch.tensor([3, 4]),
+        empty_slots=[1, 0],
+    )
+
+    trace_kwargs = next(kwargs for method, kwargs in lane.calls if method == "can_trace_prefill")
+    assert set(trace_kwargs) == {"tokens", "prompt_lens", "start_pos"}
+    assert trace_kwargs["tokens"].flatten().tolist() == [10, 11]
+    assert trace_kwargs["prompt_lens"].tolist() == [7, 9]
+    assert trace_kwargs["start_pos"].tolist() == [3, 4]
+
+
+def test_successful_traced_prefill_preflights_all_lanes_before_first_execution():
+    events = []
+
+    class TraceLane(_Lane):
+        def __init__(self, lane_idx):
+            super().__init__(lane_idx)
+            self.traced_prefill_execution = object()
+
+        def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None):
+            events.append(("preflight", self.lane_idx))
+            return True
+
+        def prefill_forward(self, *args, **kwargs):
+            events.append(("execute", self.lane_idx))
+            return super().prefill_forward(*args, **kwargs)
+
+    group = LaneGroupExecutor([TraceLane(0), TraceLane(1)])
+
+    output = group.prefill_forward(
+        tokens=torch.tensor([[10], [11]]),
+        page_table=torch.tensor([[0], [1]], dtype=torch.int32),
+        empty_slots=[0, 2],
+        execution=group.traced_prefill_execution,
+    )
+
+    assert output.flatten().tolist() == [10.0, 111.0]
+    assert events == [("preflight", 0), ("preflight", 1), ("execute", 0), ("execute", 1)]
+
+
+@pytest.mark.parametrize(
+    "unrelated_kwargs",
+    (
+        {"page_table": torch.zeros((1, 1), dtype=torch.int32)},
+        {"kv_cache": ["kv-0"]},
+        {"sampling_params": None},
+        {"execution": (object(),)},
+    ),
+)
+def test_trace_classification_rejects_unrelated_request_fields(unrelated_kwargs, expect_error):
+    group = LaneGroupExecutor([_Lane(0)])
+
+    with expect_error(TypeError, "unexpected keyword argument"):
+        group.can_trace_prefill(tokens=torch.tensor([[10]]), **unrelated_kwargs)
+
+    group.cleanup()
+
+
+def test_configure_and_allocate_fan_out_with_distinct_lane_configs():
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+    config = SimpleNamespace(num_blocks=8)
+
+    group.configure_paged_kv_cache(config)
+    handles = group.allocate_kv_cache((8, 2, 32, 64), torch.bfloat16, 2)
+
+    assert handles == ["cache-handle-0", "cache-handle-1"]
+    lane_configs = [next(kwargs["config"] for method, kwargs in lane.calls if method == "configure") for lane in lanes]
+    assert lane_configs[0] is not lane_configs[1]
+    assert lane_configs[0].num_blocks == lane_configs[1].num_blocks == 8
+    for lane in lanes:
+        allocation = next(kwargs for method, kwargs in lane.calls if method == "allocate")
+        assert allocation == {
+            "kv_cache_shape": (8, 2, 32, 64),
+            "dtype": torch.bfloat16,
+            "num_layers": 2,
+        }
+
+
+def test_side_effect_execution_target_methods_return_none_even_if_lanes_return_values():
+    class ReturningLane(_Lane):
+        def configure_paged_kv_cache(self, config):
+            super().configure_paged_kv_cache(config)
+            return self.lane_idx
+
+        def compile_prefill(
+            self,
+            *,
+            tokens,
+            page_table,
+            prompt_lens=None,
+            start_pos=None,
+            empty_slots=None,
+            kv_cache=None,
+            sampling_params=None,
+            execution=None,
+        ):
+            super().compile_prefill(
+                tokens=tokens,
+                page_table=page_table,
+                prompt_lens=prompt_lens,
+                start_pos=start_pos,
+                empty_slots=empty_slots,
+                kv_cache=kv_cache,
+                sampling_params=sampling_params,
+                execution=execution,
+            )
+            return self.lane_idx
+
+        def compile_decode(
+            self,
+            *,
+            tokens,
+            start_pos,
+            page_table,
+            kv_cache=None,
+            sampling_params=None,
+            reset_batch=False,
+            execution=None,
+        ):
+            super().compile_decode(
+                tokens=tokens,
+                start_pos=start_pos,
+                page_table=page_table,
+                kv_cache=kv_cache,
+                sampling_params=sampling_params,
+                reset_batch=reset_batch,
+                execution=execution,
+            )
+            return self.lane_idx
+
+        def warmup_model_prefill(self, *, kv_cache, can_sample_on_device, enable_trace):
+            super().warmup_model_prefill(
+                kv_cache=kv_cache,
+                can_sample_on_device=can_sample_on_device,
+                enable_trace=enable_trace,
+            )
+            return self.lane_idx
+
+        def warmup_model_decode(
+            self,
+            *,
+            kv_cache,
+            max_batch_size,
+            num_blocks,
+            can_sample_on_device,
+            enable_trace,
+        ):
+            super().warmup_model_decode(
+                kv_cache=kv_cache,
+                max_batch_size=max_batch_size,
+                num_blocks=num_blocks,
+                can_sample_on_device=can_sample_on_device,
+                enable_trace=enable_trace,
+            )
+            return self.lane_idx
+
+    group = LaneGroupExecutor([ReturningLane(0), ReturningLane(1)])
+    config = SimpleNamespace(num_blocks=8)
+    results = (
+        group.configure_paged_kv_cache(config),
+        group.compile_prefill(
+            tokens=torch.tensor([[10], [11]]),
+            page_table=torch.tensor([[0], [1]], dtype=torch.int32),
+            empty_slots=[0, 2],
+        ),
+        group.compile_decode(
+            tokens=torch.tensor([10, 11, 12, 13]),
+            start_pos=torch.tensor([1, 2, 3, 4]),
+            page_table=torch.arange(4, dtype=torch.int32).view(4, 1),
+        ),
+        group.warmup_model_prefill(
+            kv_cache=None,
+            can_sample_on_device=False,
+            enable_trace=False,
+        ),
+        group.warmup_model_decode(
+            kv_cache=None,
+            max_batch_size=4,
+            num_blocks=8,
+            can_sample_on_device=False,
+            enable_trace=False,
+        ),
+    )
+
+    assert results == (None, None, None, None, None)
+
+
+def test_constructor_validation_failure_cleans_every_supplied_lane(expect_error):
+    lanes = [_Lane(0, capacity=2), _Lane(1, capacity=4)]
+
+    with expect_error(ValueError, "same fixed capacity"):
+        LaneGroupExecutor(lanes)
+
+    assert [lane.cleanup_calls for lane in lanes] == [1, 1]
+
+
+def test_operation_failure_is_primary_group_becomes_terminal_and_all_lanes_cleanup(expect_error):
+    lanes = [_Lane(0), _Lane(1)]
+    lanes[1].fail_method = "allocate"
+    lanes[0].cleanup_error = RuntimeError("cleanup boom")
+    group = LaneGroupExecutor(lanes)
+
+    with expect_error(RuntimeError, "allocate boom 1") as exc_info:
+        group.allocate_kv_cache()
+
+    assert lanes[0].cleanup_calls == 1
+    assert lanes[1].cleanup_calls == 1
+    assert [str(error) for error in exc_info.value.cleanup_failures] == ["cleanup boom"]
+    with expect_error(RuntimeError, "terminal"):
+        group.allocate_kv_cache()
+
+
+def test_non_null_lane_log_probs_are_aggregated_in_global_row_order():
+    lanes = [_Lane(0), _Lane(1)]
+
+    def decode_with_log_probs(
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,
+        sampling_params=None,
+        reset_batch=False,
+        read_from_device=True,
+        execution=None,
+    ):
+        return (tokens + 100).to(torch.int64), torch.tensor([0.25, 0.75])
+
+    lanes[1].decode_forward = decode_with_log_probs
+    group = LaneGroupExecutor(lanes)
+
+    output, log_probs = group.decode_forward(
+        tokens=torch.tensor([0, 1, 2, 3]),
+        start_pos=torch.tensor([0, 0, 0, 0]),
+        page_table=torch.zeros((4, 1), dtype=torch.int32),
+        sampling_params=_sampling(),
+    )
+
+    assert output.tolist() == [0, 1, 102, 103]
+    assert isinstance(log_probs, torch.Tensor)
+    assert log_probs.tolist() == [1.0, 1.0, 0.25, 0.75]
+
+
+def test_cleanup_is_idempotent_and_terminal(expect_error):
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+
+    group.cleanup()
+    group.cleanup()
+
+    assert [lane.cleanup_calls for lane in lanes] == [1, 1]
+    with expect_error(RuntimeError, "terminal"):
+        group.prefill_forward(
+            tokens=torch.zeros((1, 1), dtype=torch.long),
+            page_table=torch.zeros((1, 1), dtype=torch.int32),
+        )
+
+
+def test_cleanup_retries_only_lanes_that_failed_cleanup_and_stays_terminal(expect_error):
+    class FailOnceCleanupLane(_Lane):
+        def cleanup(self):
+            self.cleanup_calls += 1
+            if self.cleanup_calls == 1:
+                raise RuntimeError("lane cleanup failed once")
+
+    failing_lane = FailOnceCleanupLane(0)
+    successful_lane = _Lane(1)
+    group = LaneGroupExecutor([failing_lane, successful_lane])
+
+    with expect_error(RuntimeError, "lane cleanup failed once"):
+        group.cleanup()
+
+    assert group.terminal
+    assert [failing_lane.cleanup_calls, successful_lane.cleanup_calls] == [1, 1]
+    with expect_error(RuntimeError, "terminal"):
+        group.allocate_kv_cache()
+
+    group.cleanup()
+    group.cleanup()
+
+    assert group.terminal
+    assert [failing_lane.cleanup_calls, successful_lane.cleanup_calls] == [2, 1]
+
+
+def test_cleanup_retries_failed_pool_shutdown_without_recleaning_lanes(expect_error):
+    class FailOnceShutdownPool:
+        def __init__(self, pool):
+            self.pool = pool
+            self.shutdown_calls = 0
+
+        def submit(self, *args, **kwargs):
+            return self.pool.submit(*args, **kwargs)
+
+        def shutdown(self, *, wait):
+            self.shutdown_calls += 1
+            if self.shutdown_calls == 1:
+                raise RuntimeError("pool shutdown failed once")
+            self.pool.shutdown(wait=wait)
+
+    lanes = [_Lane(0), _Lane(1)]
+    group = LaneGroupExecutor(lanes)
+    assert group._output_pool is not None
+    pool = FailOnceShutdownPool(group._output_pool)
+    group._output_pool = pool
+
+    with expect_error(RuntimeError, "pool shutdown failed once"):
+        group.cleanup()
+
+    assert group.terminal
+    assert [lane.cleanup_calls for lane in lanes] == [1, 1]
+
+    group.cleanup()
+    group.cleanup()
+
+    assert pool.shutdown_calls == 2
+    assert [lane.cleanup_calls for lane in lanes] == [1, 1]
+
+
+def test_explicit_mesh_device_is_preserved_by_identity():
+    mesh_device = object()
+    group = LaneGroupExecutor([_Lane(0), _Lane(1)], mesh_device=mesh_device)
+
+    assert group.mesh_device is mesh_device
+
+    group.cleanup()

code/models/common/tests/llm_runtime/test_llama3_8b_integration.py

@@ -0,0 +1,1065 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from types import SimpleNamespace
+from typing import Any, Sequence
+from unittest.mock import create_autospec
+
+import pytest
+import torch
+
+import ttnn
+from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig
+from models.common.llm_runtime.execution import EagerExecutor, TracedExecutor
+from models.common.llm_runtime.lane_group import LaneGroupExecutor
+from models.common.llm_runtime.warmup import _build_plan
+from models.common.models import llama3_executor as llama3_family_executor
+from models.common.models.llama3_8b import executor as llama_executor
+from models.common.models.llama3_8b import generator as llama_generator
+from models.common.models.llama3_8b import model as llama_model
+from models.common.tests.demos.llama3_8b import demo as llama_demo
+
+
+class _Mesh:
+    shape = (1, 1)
+
+    @staticmethod
+    def get_num_devices():
+        return 1
+
+
+def _model(*, max_batch_size=4, max_seq_len=4096):
+    paged = SimpleNamespace(block_size=32, max_num_blocks=132)
+    attention = SimpleNamespace(
+        n_kv_heads=8,
+        head_dim=128,
+        kv_cache_dtype=ttnn.bfloat8_b,
+        paged_attention_config=paged,
+        use_vllm_paged_kv_cache=True,
+        kv_cache=None,
+    )
+    attention_module = SimpleNamespace(config=attention, kv_cache=None)
+    model = SimpleNamespace(
+        config=SimpleNamespace(
+            mesh_device=_Mesh(),
+            max_batch_size=max_batch_size,
+            max_seq_len=max_seq_len,
+            n_layers=1,
+            num_devices=1,
+            block_configs=(SimpleNamespace(attention_config=attention),),
+        ),
+        layers=(SimpleNamespace(attention=attention_module),),
+        iter_executor_named_modules=lambda: (),
+        vocab_size=128,
+        num_devices=1,
+    )
+
+    def configure_paged_attention(*, block_size, max_num_blocks):
+        assert attention.kv_cache is None
+        assert attention_module.kv_cache is None
+        attention.paged_attention_config = SimpleNamespace(
+            block_size=block_size,
+            max_num_blocks=max_num_blocks,
+        )
+
+    model.configure_paged_attention = configure_paged_attention
+    return model
+
+
+def _runtime_config():
+    return SimpleNamespace(
+        model_cache_path="cache",
+        max_prefill_chunk_size=2048,
+        trace_prefill_supported_seq_lens=(128, 1024),
+        supports_batched_prefill=True,
+        disable_batched_prefill=False,
+        max_prefill_batch_size=32,
+        batched_prefill_batched_extract=True,
+        can_enable_trace=lambda sequence_length, num_cached_tokens=0: (
+            num_cached_tokens == 0 and sequence_length in (128, 1024)
+        ),
+    )
+
+
+def _config(mode="none", *, num_blocks=None):
+    return llama_executor.Llama3ExecutorConfig(
+        trace=TraceConfig(mode),
+        warmup=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        paged_kv_cache=PagedKVCacheConfig(
+            block_size=32,
+            max_num_blocks=132,
+            dtype=ttnn.bfloat8_b,
+            num_blocks=num_blocks,
+        ),
+        device_sampling_enabled=False,
+    )
+
+
+@pytest.mark.parametrize(
+    ("runtime_disabled", "device_sampling_enabled", "diagnostic_override", "expected_disabled"),
+    [
+        (False, False, False, False),
+        (True, False, False, True),
+        (False, True, False, True),
+        (True, True, False, True),
+        (False, True, True, False),
+        (True, True, True, False),
+    ],
+)
+def test_executor_resolves_batched_prefill_policy(
+    monkeypatch, runtime_disabled, device_sampling_enabled, diagnostic_override, expected_disabled
+):
+    runtime_config = _runtime_config()
+    runtime_config.disable_batched_prefill = runtime_disabled
+    base_config = _config()
+    config = llama_executor.Llama3ExecutorConfig(
+        trace=base_config.trace,
+        warmup=base_config.warmup,
+        paged_kv_cache=base_config.paged_kv_cache,
+        device_sampling_enabled=device_sampling_enabled,
+        allow_batched_prefill_with_device_sampling_for_diagnostics=diagnostic_override,
+    )
+    model = _model()
+    if device_sampling_enabled:
+
+        class _Sampling:
+            def decode_forward(self):
+                raise AssertionError("construction-policy test must not execute sampling")
+
+        class _SamplingState:
+            def __init__(self, sampling, **_kwargs):
+                self.sampling = sampling
+                self.seed_manager = SimpleNamespace()
+
+            def create_state(self):
+                return SimpleNamespace(seed_state=SimpleNamespace(capacity=32))
+
+            def admit(self, *args, **kwargs):
+                return None
+
+            def decode_forward(self, *args, **kwargs):
+                return None
+
+            def release(self, *args, **kwargs):
+                return None
+
+        monkeypatch.setattr(llama3_family_executor, "Sampling1D", _Sampling)
+        monkeypatch.setattr(llama3_family_executor, "SamplingState1D", _SamplingState)
+        sampler = _Sampling()
+        sampler.config = SimpleNamespace(
+            is_resolved=lambda: True,
+            allow_force_argmax=True,
+            max_batch_size=32,
+            max_top_k=32,
+        )
+        model.sampling = sampler
+
+    executor = llama_executor.Llama3Executor(model, runtime_config, config)
+
+    assert executor.prefill_runtime.config.supports_batched_prefill is True
+    assert executor.prefill_runtime.config.disable_batched_prefill is expected_disabled
+    assert executor.prefill_runtime.config.max_prefill_batch_size == 32
+    assert executor.prefill_runtime.config.batched_prefill_batched_extract is True
+
+
+def test_executor_rejects_batched_prefill_diagnostic_override_without_device_sampling(expect_error):
+    base_config = _config()
+
+    with expect_error(ValueError, "requires device_sampling_enabled"):
+        llama_executor.Llama3ExecutorConfig(
+            trace=base_config.trace,
+            warmup=base_config.warmup,
+            paged_kv_cache=base_config.paged_kv_cache,
+            device_sampling_enabled=False,
+            allow_batched_prefill_with_device_sampling_for_diagnostics=True,
+        )
+
+
+@pytest.mark.parametrize("mode", ["none", "decode_only", "all"])
+def test_model_owned_executor_constructs_exact_composition(mode):
+    executor = llama_executor.Llama3Executor(_model(), _runtime_config(), _config(mode))
+
+    assert executor.eager_executor.program_compiler is executor.program_compiler
+    assert executor.eager_executor.prefill is executor.prefill_runtime
+    assert executor.eager_executor.decode is executor.decode_runtime
+    assert executor.warmup.eager is executor.eager_executor
+    assert executor.warmup.trace_compiler is executor.trace_compiler
+    if mode == "none":
+        assert executor.warmup.execution is executor.eager_executor
+        assert executor.eager_execution is executor.eager_executor
+        assert executor.traced_prefill_execution is None
+        assert executor.traced_decode_execution is None
+        assert executor.trace_compiler is None
+        assert executor.traced_executor is None
+    else:
+        assert executor.warmup.execution is executor.traced_executor
+        expected_prefill = executor.traced_executor if mode == "all" else None
+        assert executor.eager_execution is executor.eager_executor
+        assert executor.traced_prefill_execution is expected_prefill
+        assert executor.traced_decode_execution is executor.traced_executor
+        assert executor.traced_executor.eager_executor is executor.eager_executor
+        assert executor.traced_executor.trace_compiler is executor.trace_compiler
+        assert executor.trace_compiler.program_compiler is executor.program_compiler
+
+
+@pytest.mark.parametrize(
+    ("method_name", "positional_names", "keyword_only_names"),
+    [
+        (
+            "compile_prefill",
+            ("self",),
+            (
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "execution",
+            ),
+        ),
+        (
+            "compile_decode",
+            ("self",),
+            (
+                "tokens",
+                "start_pos",
+                "page_table",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+                "execution",
+            ),
+        ),
+        (
+            "prefill_forward",
+            ("self", "tokens", "page_table"),
+            (
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "execution",
+            ),
+        ),
+        (
+            "decode_forward",
+            ("self", "tokens", "start_pos", "page_table"),
+            (
+                "kv_cache",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+                "read_from_device",
+                "execution",
+            ),
+        ),
+        ("read_decode_output", ("self", "tt_out"), ("async_read",)),
+        ("process_decode_output_host", ("self", "tt_out"), ("is_tokens",)),
+        (
+            "can_trace_prefill",
+            ("self",),
+            ("tokens", "prompt_lens", "start_pos", "empty_slots"),
+        ),
+        (
+            "warmup_model_prefill",
+            ("self",),
+            ("kv_cache", "can_sample_on_device", "enable_trace"),
+        ),
+        (
+            "warmup_model_decode",
+            ("self",),
+            ("kv_cache", "max_batch_size", "num_blocks", "can_sample_on_device", "enable_trace"),
+        ),
+    ],
+)
+def test_model_owned_executor_has_exact_call_contract(method_name, positional_names, keyword_only_names):
+    signature = inspect.signature(getattr(llama_executor.Llama3Executor, method_name))
+    parameters = signature.parameters
+    required_names = {
+        "compile_prefill": {"tokens", "page_table"},
+        "compile_decode": {"tokens", "start_pos", "page_table"},
+        "prefill_forward": {"tokens", "page_table"},
+        "decode_forward": {"tokens", "start_pos", "page_table"},
+        "read_decode_output": {"tt_out"},
+        "process_decode_output_host": {"tt_out"},
+        "can_trace_prefill": {"tokens"},
+        "warmup_model_prefill": {"kv_cache", "can_sample_on_device", "enable_trace"},
+        "warmup_model_decode": {
+            "kv_cache",
+            "max_batch_size",
+            "num_blocks",
+            "can_sample_on_device",
+            "enable_trace",
+        },
+    }[method_name]
+    special_defaults = {
+        "reset_batch": False,
+        "read_from_device": True,
+        "async_read": False,
+        "is_tokens": False,
+    }
+
+    assert tuple(parameters) == positional_names + keyword_only_names
+    assert all(parameters[name].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD for name in positional_names)
+    assert all(parameters[name].kind is inspect.Parameter.KEYWORD_ONLY for name in keyword_only_names)
+    for name, parameter in tuple(parameters.items())[1:]:
+        expected_default = inspect.Parameter.empty if name in required_names else special_defaults.get(name)
+        assert parameter.default == expected_default
+        assert parameter.annotation is not inspect.Parameter.empty
+    assert signature.return_annotation is not inspect.Signature.empty
+    if "execution" in parameters:
+        assert parameters["execution"].annotation == "EagerExecutor | TracedExecutor | None"
+
+
+def test_model_owned_executor_validates_cache_then_omits_it_from_execution():
+    execution = create_autospec(EagerExecutor, instance=True)
+    execution.prefill_forward.return_value = "prefill"
+    execution.decode_forward.return_value = "decode"
+    executor = object.__new__(llama_executor.Llama3Executor)
+    executor._prefill_execution = execution
+    executor._decode_execution = execution
+    executor._ensure_active = lambda: None
+    validated_caches = []
+    sampling_values = []
+    executor._validate_bound_cache = validated_caches.append
+    executor._ensure_sampling_for = sampling_values.append
+
+    tokens = torch.zeros((1, 4), dtype=torch.long)
+    start_pos = torch.zeros((1,), dtype=torch.long)
+    page_table = torch.zeros((1, 1), dtype=torch.int32)
+    prompt_lens = torch.full((1,), 4, dtype=torch.long)
+    empty_slots = [0]
+    kv_cache = object()
+    sampling_params = object()
+
+    executor.compile_prefill(
+        tokens=tokens,
+        page_table=page_table,
+        prompt_lens=prompt_lens,
+        start_pos=start_pos,
+        empty_slots=empty_slots,
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+    )
+    executor.compile_decode(
+        tokens=tokens,
+        start_pos=start_pos,
+        page_table=page_table,
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+        reset_batch=True,
+    )
+    assert (
+        executor.prefill_forward(
+            tokens,
+            page_table,
+            prompt_lens=prompt_lens,
+            start_pos=start_pos,
+            empty_slots=empty_slots,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+        )
+        == "prefill"
+    )
+    assert (
+        executor.decode_forward(
+            tokens,
+            start_pos,
+            page_table,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+            reset_batch=True,
+            read_from_device=False,
+        )
+        == "decode"
+    )
+
+    assert validated_caches == [kv_cache] * 4
+    assert sampling_values == [sampling_params] * 4
+    for target, expected_names in (
+        (
+            execution.compile_prefill,
+            (
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+            ),
+        ),
+        (
+            execution.compile_decode,
+            (
+                "tokens",
+                "start_pos",
+                "page_table",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+            ),
+        ),
+        (
+            execution.prefill_forward,
+            (
+                "tokens",
+                "page_table",
+                "prompt_lens",
+                "start_pos",
+                "empty_slots",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+            ),
+        ),
+        (
+            execution.decode_forward,
+            (
+                "tokens",
+                "start_pos",
+                "page_table",
+                "sampling_params",
+                "prompt_tokens",
+                "output_tokens",
+                "slot_remap",
+                "reset_batch",
+                "read_from_device",
+            ),
+        ),
+    ):
+        assert target.call_count == 1
+        assert tuple(target.call_args.kwargs) == expected_names
+
+
+def test_model_owned_executor_trace_output_and_warmup_forwarding_is_named():
+    calls = []
+
+    class _PrefillRuntime:
+        def can_trace(
+            self,
+            *,
+            tokens: torch.Tensor,  # ↓ Core request
+            prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
+            start_pos: torch.Tensor | None = None,
+        ) -> bool:
+            calls.append(("can_trace", tokens, prompt_lens, start_pos))
+            return True
+
+    class _DecodeRuntime:
+        def read_decode_output(self, tt_out: Any, *, async_read: bool = False) -> Any:
+            calls.append(("read_decode_output", tt_out, async_read))
+            return "read"
+
+        def process_decode_output_host(self, tt_out: Any, *, is_tokens: bool = False) -> tuple[Any, Any]:
+            calls.append(("process_decode_output_host", tt_out, is_tokens))
+            return "processed", "event"
+
+    class _Warmup:
+        def warmup_prefill(
+            self,
+            *,
+            kv_cache: Any,  # ↓ Borrowed resources
+            can_sample_on_device: bool,  # ↓ Execution policy
+            enable_trace: bool,
+        ) -> None:
+            calls.append(("warmup_prefill", kv_cache, can_sample_on_device, enable_trace))
+
+        def warmup_decode(
+            self,
+            *,
+            kv_cache: Any,  # ↓ Borrowed resources
+            max_batch_size: int,  # ↓ Coverage dimensions
+            num_blocks: int,
+            can_sample_on_device: bool,  # ↓ Execution policy
+            enable_trace: bool,
+        ) -> None:
+            calls.append(("warmup_decode", kv_cache, max_batch_size, num_blocks, can_sample_on_device, enable_trace))
+
+    executor = object.__new__(llama_executor.Llama3Executor)
+    executor.traced_executor = object()
+    executor.config = SimpleNamespace(trace=SimpleNamespace(prefill_enabled=True))
+    executor.prefill_runtime = _PrefillRuntime()
+    executor.decode_runtime = _DecodeRuntime()
+    executor.warmup = _Warmup()
+    executor._ensure_active = lambda: None
+
+    tokens = torch.zeros((1, 4), dtype=torch.long)
+    prompt_lens = torch.full((1,), 4, dtype=torch.long)
+    start_pos = torch.zeros((1,), dtype=torch.long)
+    kv_cache = object()
+
+    assert executor.can_trace_prefill(
+        tokens=tokens,
+        prompt_lens=prompt_lens,
+        start_pos=start_pos,
+        empty_slots=[0],
+    )
+    assert executor.read_decode_output("device-output", async_read=True) == "read"
+    assert executor.process_decode_output_host("host-output", is_tokens=True) == ("processed", "event")
+    executor.warmup_model_prefill(
+        kv_cache=kv_cache,
+        can_sample_on_device=True,
+        enable_trace=False,
+    )
+    executor.warmup_model_decode(
+        kv_cache=kv_cache,
+        max_batch_size=4,
+        num_blocks=8,
+        can_sample_on_device=True,
+        enable_trace=False,
+    )
+
+    assert calls == [
+        ("can_trace", tokens, prompt_lens, start_pos),
+        ("read_decode_output", "device-output", True),
+        ("process_decode_output_host", "host-output", True),
+        ("warmup_prefill", kv_cache, True, False),
+        ("warmup_decode", kv_cache, 4, 8, True, False),
+    ]
+
+
+def test_vllm_capacity_resolution_reconfigures_existing_runtime_owners_before_allocation(monkeypatch):
+    executor = llama_executor.Llama3Executor(_model(), _runtime_config(), _config())
+    owner_ids = tuple(
+        id(owner)
+        for owner in (
+            executor.prefill_runtime,
+            executor.decode_runtime,
+            executor.warmup,
+            executor.program_compiler,
+        )
+    )
+    assert executor.page_table_layout.raw_capacity_width == 128
+
+    executor.configure_paged_kv_cache(
+        PagedKVCacheConfig(
+            block_size=16,
+            max_num_blocks=200,
+            dtype=ttnn.bfloat8_b,
+            num_blocks=200,
+        )
+    )
+
+    assert (
+        tuple(
+            id(owner)
+            for owner in (
+                executor.prefill_runtime,
+                executor.decode_runtime,
+                executor.warmup,
+                executor.program_compiler,
+            )
+        )
+        == owner_ids
+    )
+    assert executor.config.paged_kv_cache is executor.kv_cache_manager.config
+    assert executor.kv_cache_manager.config.block_size == 16
+    assert executor.kv_cache_manager.config.max_num_blocks == executor.kv_cache_manager.config.num_blocks == 200
+    assert executor.model.layers[0].attention.config.paged_attention_config.block_size == 16
+    assert executor.model.layers[0].attention.config.paged_attention_config.max_num_blocks == 200
+    assert executor.page_table_layout.block_size == 16
+    assert executor.page_table_layout.raw_capacity_width == 200
+    assert executor.prefill_runtime.config.page_table_layout is executor.page_table_layout
+    assert executor.decode_runtime.config.page_table_layout is executor.page_table_layout
+    assert executor.warmup.config.page_table_layout is executor.page_table_layout
+
+    def fake_allocate():
+        assert executor._runtime_configuration_sealed
+        assert executor.warmup._configuration_sealed
+        return ["allocated"]
+
+    monkeypatch.setattr(executor.kv_cache_manager, "allocate", fake_allocate)
+    assert executor.allocate_kv_cache() == ["allocated"]
+
+
+def test_model_reconfigures_construction_and_live_attention_without_allocating():
+    construction_paged = llama_model.Llama31_8BPagedAttentionConfig(block_size=32, max_num_blocks=132)
+    live_paged = llama_model.Llama31_8BPagedAttentionConfig(block_size=32, max_num_blocks=132)
+    construction_attention = SimpleNamespace(
+        use_vllm_paged_kv_cache=True,
+        paged_attention_config=construction_paged,
+    )
+    live_attention = SimpleNamespace(
+        use_vllm_paged_kv_cache=True,
+        paged_attention_config=live_paged,
+        kv_cache=None,
+    )
+    model = object.__new__(llama_model.Llama3Transformer1D)
+    model.config = SimpleNamespace(
+        block_configs=(SimpleNamespace(attention_config=construction_attention),),
+    )
+    model.layers = (SimpleNamespace(attention=SimpleNamespace(config=live_attention, kv_cache=None)),)
+
+    model.configure_paged_attention(block_size=16, max_num_blocks=200)
+
+    assert construction_attention.paged_attention_config.block_size == 16
+    assert construction_attention.paged_attention_config.max_num_blocks == 200
+    assert live_attention.paged_attention_config.block_size == 16
+    assert live_attention.paged_attention_config.max_num_blocks == 200
+    assert live_attention.kv_cache is None
+    assert model.layers[0].attention.kv_cache is None
+
+
+def test_late_vllm_capacity_resolution_fails_before_mutating_kv_configuration(expect_error):
+    executor = llama_executor.Llama3Executor(_model(), _runtime_config(), _config())
+    executor._seal_runtime_configuration()
+    unresolved = executor.kv_cache_manager.config
+
+    with expect_error(RuntimeError, "runtime configuration is sealed"):
+        executor.configure_paged_kv_cache(
+            PagedKVCacheConfig(
+                block_size=32,
+                max_num_blocks=132,
+                dtype=ttnn.bfloat8_b,
+                num_blocks=64,
+            )
+        )
+
+    assert executor.kv_cache_manager.config is unresolved
+    assert not executor.kv_cache_manager.config.is_resolved()
+
+
+def test_direct_demo_resolves_physical_capacity_to_configured_maximum(monkeypatch):
+    attention = SimpleNamespace(
+        paged_attention_config=SimpleNamespace(block_size=32, max_num_blocks=128),
+        kv_cache_dtype=ttnn.bfloat8_b,
+    )
+    llm = SimpleNamespace(
+        model=SimpleNamespace(config=SimpleNamespace(block_configs=(SimpleNamespace(attention_config=attention),)))
+    )
+    captured = []
+    monkeypatch.setattr(
+        llama_demo, "build_llama3_executor", lambda product, config: captured.append(config) or object()
+    )
+
+    llama_demo._build_demo_executor(llm, trace_mode="all", device_sampling_enabled=False)
+
+    assert captured[0].paged_kv_cache.num_blocks == captured[0].paged_kv_cache.max_num_blocks == 128
+
+
+def test_model_owned_cleanup_is_ordered_best_effort_retryable_and_idempotent(expect_error):
+    calls = []
+    failures = {"reader", "trace"}
+
+    class _Owner:
+        def __init__(self, name):
+            self.name = name
+
+        def cleanup(self, *args):
+            calls.append(self.name)
+            if self.name in failures:
+                raise RuntimeError(self.name)
+
+        drain = cleanup
+        drain_external_outputs = cleanup
+        cleanup_transients = cleanup
+        release = cleanup
+
+    executor = object.__new__(llama_executor.Llama3Executor)
+    executor._terminal = False
+    executor._cleaned_up = False
+    executor.decode_runtime = _Owner("decode-external")
+    executor.output_reader = _Owner("reader")
+    executor.prefill_runtime = _Owner("prefill")
+    executor.trace_compiler = _Owner("trace")
+    executor.program_compiler = _Owner("program")
+    executor.config = SimpleNamespace(device_sampling_enabled=True)
+    executor.model = SimpleNamespace(sampling=_Owner("sampling"))
+    executor.sampling_state_controller = _Owner("sampling-state")
+    executor.sampling_state = object()
+    executor.kv_cache_manager = _Owner("kv")
+
+    with expect_error(RuntimeError, "reader") as raised:
+        executor.cleanup()
+
+    expected_order = [
+        "decode-external",
+        "reader",
+        "prefill",
+        "decode-external",
+        "trace",
+        "program",
+        "sampling-state",
+        "sampling",
+        "kv",
+    ]
+    assert calls == expected_order
+    assert tuple(error.args[0] for error in raised.value.cleanup_failures) == ("trace",)
+    assert executor.terminal
+    assert not executor._cleaned_up
+
+    failures.clear()
+    executor.cleanup()
+    assert calls == expected_order * 2
+    assert executor._cleaned_up
+
+    executor.cleanup()
+    assert calls == expected_order * 2
+
+
+def test_build_llama3_executor_uses_prebuilt_product(monkeypatch):
+    sentinel = object()
+    calls = []
+    llm = SimpleNamespace(model=object(), runtime_config=object())
+    monkeypatch.setattr(
+        llama_executor,
+        "Llama3Executor",
+        lambda model, runtime_config, config: calls.append((model, runtime_config, config)) or sentinel,
+    )
+
+    config = object()
+    assert llama_executor.build_llama3_executor(llm, config) is sentinel
+    assert calls == [(llm.model, llm.runtime_config, config)]
+
+
+def test_configured_path_has_no_legacy_or_common_aggregate_surface():
+    source = inspect.getsource(llama_executor)
+    assert not hasattr(llama_executor, "EagerLlamaExecutor")
+    assert not hasattr(llama_executor, "TracedLlamaExecutor")
+    assert "models.common.models.executor" not in source
+    assert "llm_runtime.executor" not in source
+    assert "class LLMExecutor" not in source
+    assert EagerExecutor not in TracedExecutor.__mro__
+
+
+def test_generator_explicitly_accepts_static_trace_mode():
+    assert llama_generator.Llama3Generator.model_capabilities["accepts_trace_mode"] is True
+
+
+def test_initialize_vllm_model_threads_policy(monkeypatch):
+    captured = []
+    sentinel = object()
+    monkeypatch.setattr(
+        llama_generator,
+        "build_llama3_generator",
+        lambda config: captured.append(config) or sentinel,
+    )
+
+    result = llama_generator.Llama3Generator.initialize_vllm_model(
+        SimpleNamespace(_name_or_path="meta-llama/Llama-3.1-8B-Instruct"),
+        object(),
+        8,
+        4096,
+        n_layers=3,
+        tt_data_parallel=2,
+        optimizations="accuracy",
+        trace_mode="decode_only",
+        device_sampling_enabled=True,
+    )
+
+    assert result is sentinel
+    assert captured[0].tt_data_parallel == 2
+    assert captured[0].trace_mode == "decode_only"
+    assert captured[0].device_sampling_enabled is True
+
+
+class _FakeLane:
+    requires_prefill_trace_warmup = True
+
+    def __init__(self, llm, config):
+        self.model = llm.model
+        self.model_args = llm.runtime_config
+        self.mesh_device = llm.model.config.mesh_device
+        self.cache_path = llm.runtime_config.model_cache_path
+        self._request_state_fields = ("prompt_tokens", "output_tokens", "slot_remap")
+        self.config = config
+        self.paged_kv_cache_config = config.paged_kv_cache
+        self.already_warmed_up_prefill = False
+        self.cleanup_calls = 0
+
+    def cleanup(self):
+        self.cleanup_calls += 1
+
+
+@pytest.mark.parametrize("device_sampling_enabled", (False, True), ids=("sampling-off", "sampling-on"))
+def test_generator_constructs_model_owned_lane_configs_with_exact_decode_coverage(
+    monkeypatch,
+    device_sampling_enabled,
+):
+    executor_calls = []
+    monkeypatch.setattr(llama_generator, "_create_submeshes", lambda mesh, dp: [_Mesh(), _Mesh()])
+
+    def fake_from_pretrained(
+        mesh_device,
+        *,
+        hf_model: str | None = None,
+        instruct: bool | None = None,
+        max_batch_size: int,
+        max_seq_len: int,
+        optimizations="performance",
+        n_layers: int | None = None,
+        dtype=ttnn.bfloat8_b,
+        paged_attention_config=None,
+    ):
+        return SimpleNamespace(model=_model(max_batch_size=max_batch_size), runtime_config=_runtime_config())
+
+    def fake_build_executor(llm, config):
+        executor_calls.append((llm, config))
+        return _FakeLane(llm, config)
+
+    monkeypatch.setattr(llama_generator, "from_pretrained", fake_from_pretrained)
+    monkeypatch.setattr(llama_generator, "build_llama3_executor", fake_build_executor)
+    monkeypatch.setattr(llama_generator, "_model_kv_metadata", lambda model: ((ttnn.bfloat8_b,), 1, 8, 128))
+
+    generator = llama_generator.build_llama3_generator(
+        llama_generator.Llama3GeneratorConfig(
+            hf_model="meta-llama/Llama-3.1-8B-Instruct",
+            mesh_device=object(),
+            max_batch_size=4,
+            max_seq_len=4096,
+            n_layers=1,
+            tt_data_parallel=2,
+            trace_mode="all",
+            device_sampling_enabled=device_sampling_enabled,
+        )
+    )
+
+    assert isinstance(generator.target, LaneGroupExecutor)
+    assert len(executor_calls) == 2
+    assert all(isinstance(config, llama_executor.Llama3ExecutorConfig) for _, config in executor_calls)
+    assert all(config.warmup.include_decode_top_k is device_sampling_enabled for _, config in executor_calls)
+    for _, executor_config in executor_calls:
+        plan = _build_plan(
+            warmup=executor_config.warmup,
+            layout=PageTableLayout(block_size=32, raw_capacity_width=32, prefill_width=64, decode_width=32),
+            prefill_sequence_lengths=(128,),
+            lane_batch_size=2,
+            allow_force_argmax=True,
+            can_sample_on_device=device_sampling_enabled,
+        )
+        assert [case.sampling_path for case in plan.decode] == (
+            ["logits", "argmax", "topk"] if device_sampling_enabled else ["logits"]
+        )
+    assert isinstance(generator._adapter.config, llama_generator.VLLMAdapterConfig)
+    assert vars(generator._adapter) == {"config": generator._adapter.config}
+    assert generator._adapter.config.trace.mode == "all"
+    assert generator._adapter.config.expected_num_layers == 1
+    assert generator._adapter.config.expected_kv_heads_per_device == 8
+    assert generator._adapter.config.expected_head_dim == 128
+
+
+@pytest.mark.parametrize(
+    ("sampling_mode", "sampling_params", "num_devices", "expected"),
+    [
+        ("on_device_topk", object(), 1, False),
+        ("on_device_topk", object(), 2, False),
+        ("on_device_topk", object(), 8, True),
+        ("on_device", object(), 8, False),
+        ("on_device_topk", None, 8, False),
+    ],
+)
+def test_direct_demo_forces_decode_top_k_only_on_t3k(sampling_mode, sampling_params, num_devices, expected):
+    assert llama_demo._force_decode_top_k(sampling_mode, sampling_params, num_devices) is expected
+
+
+class _RecordingTarget:
+    model = SimpleNamespace(config=SimpleNamespace(max_batch_size=4))
+    model_args = object()
+    mesh_device = object()
+    cache_path = "cache"
+    already_warmed_up_prefill = False
+    eager_execution = object()
+    traced_prefill_execution = object()
+    traced_decode_execution = object()
+
+    def __init__(self, *, traceable_prefill=True):
+        self.calls = []
+        self.traceable_prefill = traceable_prefill
+        self._request_state_fields = ("prompt_tokens", "output_tokens", "slot_remap")
+
+    def _record(self, name, arguments):
+        arguments = {key: value for key, value in arguments.items() if key != "self"}
+        self.calls.append((name, (), arguments))
+        return name
+
+    def can_trace_prefill(
+        self,
+        *,
+        tokens: torch.Tensor,  # ↓ Core request
+        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
+        start_pos: torch.Tensor | None = None,
+        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
+    ) -> bool:
+        self._record("can_trace_prefill", locals())
+        return self.traceable_prefill
+
+    def prefill_forward(
+        self,
+        tokens: torch.Tensor,
+        page_table: torch.Tensor,
+        *,
+        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
+        start_pos: torch.Tensor | None = None,
+        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
+        kv_cache: Any = None,  # ↓ Borrowed resources
+        sampling_params: Any = None,  # ↓ Sampling
+        prompt_tokens: Any = None,
+        output_tokens: Any = None,
+        slot_remap: Any = None,
+        execution: EagerExecutor | TracedExecutor | None = None,  # ↓ Internal dispatch
+    ) -> str:
+        return self._record("prefill_forward", locals())
+
+    def decode_forward(
+        self,
+        tokens: torch.Tensor,
+        start_pos: torch.Tensor,
+        page_table: torch.Tensor,
+        *,
+        kv_cache: Any = None,  # ↓ Borrowed resources
+        sampling_params: Any = None,  # ↓ Sampling
+        prompt_tokens: Any = None,
+        output_tokens: Any = None,
+        slot_remap: Any = None,
+        reset_batch: bool = False,  # ↓ State transition
+        read_from_device: bool = True,  # ↓ Output policy
+        execution: EagerExecutor | TracedExecutor | None = None,  # ↓ Internal dispatch
+    ) -> str:
+        return self._record("decode_forward", locals())
+
+    def cleanup(self) -> str:
+        self.calls.append(("cleanup", (), {}))
+        return "cleanup"
+
+
+def _recording_generator(target):
+    target.model = _model()
+    target.config = _config("all")
+    return llama_generator.Llama3Generator(target, adapter=llama_generator._build_vllm_adapter(target))
+
+
+def test_generator_delegates_without_concrete_type_checks():
+    target = _RecordingTarget()
+    generator = _recording_generator(target)
+    tokens = torch.tensor([[1]], dtype=torch.long)
+    start_pos = torch.tensor([0], dtype=torch.long)
+    page_table = torch.tensor([[0]], dtype=torch.int32)
+
+    assert generator.prefill_forward(tokens, page_table, enable_trace=True) == "prefill_forward"
+    assert generator.decode_forward(tokens, start_pos, page_table, enable_trace=True) == "decode_forward"
+    assert generator.cleanup() == "cleanup"
+    assert [name for name, _, _ in target.calls] == [
+        "prefill_forward",
+        "decode_forward",
+        "cleanup",
+    ]
+    assert target.calls[0][2]["execution"] is target.traced_prefill_execution
+    assert target.calls[1][2]["execution"] is target.traced_decode_execution
+
+
+def test_generator_preserves_required_trace_intent_for_ineligible_prefill():
+    target = _RecordingTarget(traceable_prefill=False)
+    generator = _recording_generator(target)
+    tokens = torch.tensor([[1]], dtype=torch.long)
+    page_table = torch.tensor([[0]], dtype=torch.int32)
+
+    assert generator.prefill_forward(tokens, page_table, enable_trace=True) == "prefill_forward"
+    assert [name for name, _, _ in target.calls] == ["prefill_forward"]
+    assert target.calls[0][2]["execution"] is target.traced_prefill_execution
+
+
+def test_generator_rejects_unavailable_traced_execution(expect_error):
+    target = _RecordingTarget()
+    target.traced_decode_execution = None
+    generator = _recording_generator(target)
+    tokens = torch.tensor([1], dtype=torch.long)
+    start_pos = torch.tensor([0], dtype=torch.long)
+    page_table = torch.tensor([[0]], dtype=torch.int32)
+
+    with expect_error(RuntimeError, "unavailable traced decode execution"):
+        generator.decode_forward(tokens, start_pos, page_table, enable_trace=True)
+
+    assert target.calls == []
+
+
+def test_demo_uses_model_owned_config_and_order_independent_warmup(monkeypatch):
+    attention = SimpleNamespace(
+        paged_attention_config=SimpleNamespace(block_size=32, max_num_blocks=128),
+        kv_cache_dtype=ttnn.bfloat8_b,
+    )
+    llm = SimpleNamespace(
+        model=SimpleNamespace(config=SimpleNamespace(block_configs=(SimpleNamespace(attention_config=attention),)))
+    )
+    captured = []
+    monkeypatch.setattr(
+        llama_demo, "build_llama3_executor", lambda product, config: captured.append(config) or object()
+    )
+
+    llama_demo._build_demo_executor(llm, trace_mode="all", device_sampling_enabled=False)
+    assert isinstance(captured[0], llama_executor.Llama3ExecutorConfig)
+
+    calls = []
+
+    def fake_warmup_model_prefill(
+        *,
+        kv_cache: Any,  # ↓ Borrowed resources
+        can_sample_on_device: bool,  # ↓ Execution policy
+        enable_trace: bool,
+    ) -> None:
+        calls.append(("prefill", kv_cache, can_sample_on_device, enable_trace))
+
+    def fake_warmup_model_decode(
+        *,
+        kv_cache: Any,  # ↓ Borrowed resources
+        max_batch_size: int,  # ↓ Coverage dimensions
+        num_blocks: int,
+        can_sample_on_device: bool,  # ↓ Execution policy
+        enable_trace: bool,
+    ) -> None:
+        calls.append(("decode", kv_cache, max_batch_size, num_blocks, can_sample_on_device, enable_trace))
+
+    executor = SimpleNamespace(
+        config=SimpleNamespace(
+            trace=TraceConfig("all"),
+            device_sampling_enabled=False,
+        ),
+        model=SimpleNamespace(config=SimpleNamespace(max_batch_size=4)),
+        warmup_model_prefill=fake_warmup_model_prefill,
+        warmup_model_decode=fake_warmup_model_decode,
+    )
+    kv_cache = object()
+    llama_demo._warmup_demo_executor(executor, kv_cache=kv_cache, page_table=SimpleNamespace(shape=(4, 8)))
+    assert calls == [
+        ("prefill", kv_cache, False, False),
+        ("decode", kv_cache, 4, 8, False, False),
+        ("prefill", kv_cache, False, True),
+        ("decode", kv_cache, 4, 8, False, True),
+    ]
+
+    calls.clear()
+    executor.config.device_sampling_enabled = True
+    llama_demo._warmup_demo_executor(
+        executor,
+        kv_cache=kv_cache,
+        page_table=SimpleNamespace(shape=(4, 8)),
+        prefill_can_sample_on_device=False,
+    )
+    assert calls == [
+        ("prefill", kv_cache, False, False),
+        ("decode", kv_cache, 4, 8, True, False),
+        ("prefill", kv_cache, False, True),
+        ("decode", kv_cache, 4, 8, True, True),
+    ]

code/models/common/tests/llm_runtime/test_llama3_8b_model_contract.py

@@ -0,0 +1,349 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+import pytest
+
+from models.common.models.llama3_8b import model as llama_model
+from models.common.models.llama3_8b.model import Llama3Transformer1D, Llama3Transformer1DConfig, TransformerBlock1D
+
+
+def _layer():
+    return SimpleNamespace(
+        attention_norm=object(),
+        attention=SimpleNamespace(config=SimpleNamespace(kv_cache=None), kv_cache=None),
+        ff_norm=object(),
+        feed_forward=object(),
+    )
+
+
+def test_transformer_config_exposes_hidden_width_for_runtime_slices():
+    config = Llama3Transformer1DConfig(
+        n_layers=1,
+        vocab_size=128256,
+        max_batch_size=1,
+        max_seq_len=1024,
+        dim=4096,
+        num_devices=1,
+        mesh_device=object(),
+        embedding_config=object(),
+        rope_config=object(),
+        block_configs=[object()],
+        norm_config=object(),
+        lm_head_config=object(),
+    )
+
+    assert config.dim == 4096
+
+
+def test_iter_executor_named_modules_preserves_names_and_order():
+    layers = [_layer(), _layer()]
+    model = SimpleNamespace(layers=layers, norm=object(), lm_head=object())
+
+    named_modules = list(Llama3Transformer1D.iter_executor_named_modules(model))
+
+    assert named_modules == [
+        ("layer[0].attn_norm", layers[0].attention_norm),
+        ("layer[0].attention", layers[0].attention),
+        ("layer[0].ff_norm", layers[0].ff_norm),
+        ("layer[0].mlp", layers[0].feed_forward),
+        ("layer[1].attn_norm", layers[1].attention_norm),
+        ("layer[1].attention", layers[1].attention),
+        ("layer[1].ff_norm", layers[1].ff_norm),
+        ("layer[1].mlp", layers[1].feed_forward),
+        ("final_norm", model.norm),
+        ("lm_head", model.lm_head),
+    ]
+
+
+def test_iter_executor_named_modules_without_layers_yields_nothing():
+    model = SimpleNamespace(norm=object(), lm_head=object())
+
+    assert list(Llama3Transformer1D.iter_executor_named_modules(model)) == []
+
+
+def test_set_kv_cache_binds_and_unbinds_config_and_runtime_references():
+    layers = [_layer(), _layer()]
+    model = SimpleNamespace(layers=layers)
+    kv_cache = [[object(), object()], [object(), object()]]
+
+    Llama3Transformer1D.set_kv_cache(model, kv_cache)
+
+    for layer, expected in zip(layers, kv_cache):
+        bound = layer.attention.config.kv_cache
+        assert bound == tuple(expected)
+        assert bound[0] is expected[0]
+        assert bound[1] is expected[1]
+        assert layer.attention.kv_cache is bound
+
+    Llama3Transformer1D.set_kv_cache(model, None)
+    Llama3Transformer1D.set_kv_cache(model, None)
+
+    for layer in layers:
+        assert layer.attention.config.kv_cache is None
+        assert layer.attention.kv_cache is None
+
+
+def test_set_kv_cache_rejects_wrong_layer_count_before_binding(expect_error):
+    layers = [_layer(), _layer()]
+    model = SimpleNamespace(layers=layers)
+
+    with expect_error(ValueError, "model has 2 layers"):
+        Llama3Transformer1D.set_kv_cache(model, [[object(), object()]])
+
+    assert all(layer.attention.config.kv_cache is None for layer in layers)
+    assert all(layer.attention.kv_cache is None for layer in layers)
+
+
+@pytest.mark.parametrize("bad_pair", [[object()], object()])
+def test_set_kv_cache_validates_every_pair_before_binding(bad_pair, expect_error):
+    layers = [_layer(), _layer()]
+    model = SimpleNamespace(layers=layers)
+
+    with expect_error((TypeError, ValueError), "layer 1.*K/V tensor"):
+        Llama3Transformer1D.set_kv_cache(
+            model,
+            [[object(), object()], bad_pair],
+        )
+
+    assert all(layer.attention.config.kv_cache is None for layer in layers)
+    assert all(layer.attention.kv_cache is None for layer in layers)
+
+
+class _FakeTensor:
+    def __init__(self, shape=(1, 1, 96, 4096), dtype=None):
+        self.shape = shape
+        self.dtype = dtype
+        self.deallocate = MagicMock()
+
+
+def _identity_all_gather(_norm, tensor, *, memory_config=None):
+    return tensor
+
+
+def _model_ttnn(
+    *, slice_result=None, split_result=None, typecast_result=None, embedding_result=None, unsqueeze_result=None
+):
+    slice_mock = MagicMock(return_value=slice_result)
+    return SimpleNamespace(
+        DRAM_MEMORY_CONFIG=object(),
+        TILE_LAYOUT=object(),
+        bfloat16=object(),
+        add=MagicMock(side_effect=lambda lhs, _rhs, **_kwargs: lhs),
+        concat=MagicMock(return_value=slice_result),
+        deallocate=MagicMock(),
+        embedding=MagicMock(return_value=embedding_result),
+        interleaved_to_sharded=MagicMock(side_effect=lambda tensor, _memcfg: tensor),
+        reshape=MagicMock(side_effect=lambda tensor, *_args, **_kwargs: tensor),
+        slice=slice_mock,
+        split=MagicMock(return_value=split_result),
+        to_memory_config=MagicMock(side_effect=lambda tensor, *_args, **_kwargs: tensor),
+        typecast=MagicMock(return_value=typecast_result),
+        unsqueeze_to_4D=MagicMock(return_value=unsqueeze_result),
+    )
+
+
+@pytest.mark.parametrize("chunk_start_idx_tensor", [None, object()])
+def test_transformer_block_prefill_forwards_scalar_and_tensor_chunk_start(monkeypatch, chunk_start_idx_tensor):
+    x = _FakeTensor()
+    attn_output = _FakeTensor()
+    attention = SimpleNamespace(prefill_forward=MagicMock(return_value=attn_output))
+    block = SimpleNamespace(
+        activation_dtype=None,
+        attention=attention,
+        attention_norm=SimpleNamespace(prefill_forward=MagicMock(return_value=x)),
+        feed_forward=SimpleNamespace(prefill_forward=MagicMock(side_effect=lambda tensor: tensor)),
+        ff_norm=SimpleNamespace(prefill_forward=MagicMock(side_effect=lambda tensor: tensor)),
+        prefill_residual_memcfg=object(),
+    )
+    fake_ttnn = _model_ttnn()
+    monkeypatch.setattr(llama_model, "ttnn", fake_ttnn)
+    monkeypatch.setattr(llama_model, "_all_gather_rmsnorm_tensor", _identity_all_gather)
+
+    TransformerBlock1D.prefill_forward(
+        block,
+        x,
+        ("cos", "sin"),
+        user_id=3,
+        page_table="page-table",
+        chunk_page_table="chunk-page-table",
+        chunk_start_idx=96,
+        chunk_start_idx_tensor=chunk_start_idx_tensor,
+    )
+
+    attention.prefill_forward.assert_called_once_with(
+        x,
+        ("cos", "sin"),
+        user_id=3,
+        page_table="page-table",
+        chunk_page_table="chunk-page-table",
+        chunk_start_idx=96,
+        chunk_start_idx_tensor=chunk_start_idx_tensor,
+    )
+
+
+def test_prefill_uses_runtime_block_slice_then_embeds_exact_row(monkeypatch):
+    hidden_states = _FakeTensor(shape=(1, 1, 96, 4096))
+    sliced = _FakeTensor(shape=(1, 1, 32, 4096))
+    converted = _FakeTensor(shape=(1, 1, 32, 4096))
+    embedded = _FakeTensor(shape=(1, 1, 4096))
+    selected = _FakeTensor(shape=(1, 1, 1, 4096))
+    start_tensor = object()
+    end_tensor = object()
+    row_index = object()
+    fake_ttnn = _model_ttnn(
+        slice_result=sliced,
+        typecast_result=converted,
+        embedding_result=embedded,
+        unsqueeze_result=selected,
+    )
+    model = SimpleNamespace(
+        activation_dtypes=[],
+        layers=[],
+        lm_head=SimpleNamespace(
+            config=SimpleNamespace(input_memcfg=None),
+            forward=MagicMock(side_effect=lambda tensor: tensor),
+        ),
+        norm=SimpleNamespace(prefill_forward=MagicMock(side_effect=lambda tensor: tensor)),
+    )
+    monkeypatch.setattr(llama_model, "ttnn", fake_ttnn)
+    monkeypatch.setattr(llama_model, "_all_gather_rmsnorm_tensor", _identity_all_gather)
+
+    Llama3Transformer1D.prefill_forward(
+        model,
+        hidden_states,
+        ("cos", "sin"),
+        get_last_token=95,
+        last_token_slice=(start_tensor, end_tensor),
+        last_token_index=row_index,
+    )
+
+    fake_ttnn.slice.assert_called_once_with(
+        hidden_states,
+        start_tensor,
+        end_tensor,
+        slice_dim=2,
+        num_devices=3,
+    )
+    fake_ttnn.typecast.assert_called_once_with(sliced, fake_ttnn.bfloat16)
+    fake_ttnn.embedding.assert_called_once_with(row_index, converted, layout=fake_ttnn.TILE_LAYOUT)
+    fake_ttnn.unsqueeze_to_4D.assert_called_once_with(embedded)
+    assert fake_ttnn.deallocate.call_args_list == [call(hidden_states), call(sliced), call(converted)]
+    model.norm.prefill_forward.assert_called_once_with(selected)
+
+
+def test_post_process_prefill_uses_runtime_bounds_for_aligned_block_slice(monkeypatch):
+    hidden_states = _FakeTensor(shape=(1, 1, 96, 4096))
+    sliced = _FakeTensor(shape=(1, 1, 32, 4096))
+    start_tensor = object()
+    end_tensor = object()
+    fake_ttnn = _model_ttnn(slice_result=sliced)
+    model = SimpleNamespace(
+        lm_head=SimpleNamespace(
+            config=SimpleNamespace(input_memcfg=None),
+            forward=MagicMock(side_effect=lambda tensor: tensor),
+        ),
+        norm=SimpleNamespace(prefill_forward=MagicMock(side_effect=lambda tensor: tensor)),
+    )
+    monkeypatch.setattr(llama_model, "ttnn", fake_ttnn)
+    monkeypatch.setattr(llama_model, "_all_gather_rmsnorm_tensor", _identity_all_gather)
+
+    Llama3Transformer1D.post_process_prefill_output(
+        model,
+        hidden_states,
+        last_token_idx=95,
+        last_token_slice=(start_tensor, end_tensor),
+    )
+
+    fake_ttnn.slice.assert_called_once_with(
+        hidden_states,
+        start_tensor,
+        end_tensor,
+        slice_dim=2,
+        num_devices=3,
+    )
+
+
+def test_batched_post_process_uses_runtime_bounds_for_each_active_slot(monkeypatch):
+    hidden_states = _FakeTensor(shape=(1, 1, 256, 4096))
+    user_states = [_FakeTensor(shape=(1, 1, 128, 4096)) for _ in range(2)]
+    block = _FakeTensor(shape=(1, 1, 32, 4096))
+    embedded = _FakeTensor(shape=(1, 1, 4096))
+    selected = _FakeTensor(shape=(1, 1, 1, 4096))
+    start_tensor = object()
+    end_tensor = object()
+    row_index = object()
+    fake_ttnn = _model_ttnn(
+        slice_result=block,
+        split_result=user_states,
+        embedding_result=embedded,
+        unsqueeze_result=selected,
+    )
+    model = SimpleNamespace(
+        lm_head=SimpleNamespace(
+            config=SimpleNamespace(input_memcfg=None),
+            forward=MagicMock(side_effect=lambda tensor: tensor),
+        ),
+        norm=SimpleNamespace(prefill_forward=MagicMock(side_effect=lambda tensor: tensor)),
+    )
+    monkeypatch.setattr(llama_model, "ttnn", fake_ttnn)
+    monkeypatch.setattr(llama_model, "_all_gather_rmsnorm_tensor", _identity_all_gather)
+
+    Llama3Transformer1D.post_process_batched_prefill_output(
+        model,
+        hidden_states,
+        last_token_idx_list=[4, 17],
+        padded_batch=2,
+        prefill_seq_len=128,
+        last_token_slice=(start_tensor, end_tensor),
+        last_token_index=row_index,
+    )
+
+    assert fake_ttnn.slice.call_args_list == [
+        call(user_states[0], start_tensor, end_tensor, slice_dim=2, num_devices=4),
+        call(user_states[1], start_tensor, end_tensor, slice_dim=2, num_devices=4),
+    ]
+    assert fake_ttnn.embedding.call_args_list == [
+        call(row_index, block, layout=fake_ttnn.TILE_LAYOUT),
+        call(row_index, block, layout=fake_ttnn.TILE_LAYOUT),
+    ]
+    assert fake_ttnn.unsqueeze_to_4D.call_args_list == [call(embedded), call(embedded)]
+    assert fake_ttnn.deallocate.call_args_list == [call(block), call(block)]
+
+
+def test_prepare_prefill_rot_mats_gathers_runtime_positions_without_slicing(monkeypatch):
+    position_indices = object()
+    cos_matrix = object()
+    sin_matrix = object()
+    cos_rows = object()
+    sin_rows = object()
+    cos_4d = object()
+    sin_4d = object()
+    embedding = MagicMock(side_effect=[cos_rows, sin_rows])
+    unsqueeze_to_4d = MagicMock(side_effect=[cos_4d, sin_4d])
+    fake_ttnn = SimpleNamespace(
+        TILE_LAYOUT=object(),
+        embedding=embedding,
+        unsqueeze_to_4D=unsqueeze_to_4d,
+    )
+    rope_setup = SimpleNamespace(
+        cos_matrix=cos_matrix,
+        sin_matrix=sin_matrix,
+        load_device_weights=MagicMock(),
+    )
+    monkeypatch.setattr(llama_model, "ttnn", fake_ttnn)
+
+    result = Llama3Transformer1D.prepare_prefill_rot_mats(
+        SimpleNamespace(rope_setup=rope_setup),
+        position_indices,
+    )
+
+    rope_setup.load_device_weights.assert_called_once_with()
+    assert embedding.call_args_list == [
+        call(position_indices, cos_matrix, layout=fake_ttnn.TILE_LAYOUT),
+        call(position_indices, sin_matrix, layout=fake_ttnn.TILE_LAYOUT),
+    ]
+    assert unsqueeze_to_4d.call_args_list == [call(cos_rows), call(sin_rows)]
+    assert result == (cos_4d, sin_4d)

code/models/common/tests/llm_runtime/test_model_contract.py

@@ -0,0 +1,972 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Generic host-only tensor-model contract for migrated sibling models."""
+
+import inspect
+from dataclasses import dataclass, fields
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+import pytest
+
+from models.common.models.deepseek_r1_distill_qwen_14b import model as deepseek_model
+from models.common.models.llama32_1b import model as llama32_model
+from models.common.models.llama32_3b import model as llama32_3b_model
+from models.common.models.llama33_70b import model as llama33_70b_model
+from models.common.models.mistral_7b import model as mistral_model
+from models.common.models.phi4 import model as phi4_model
+from models.common.models.qwen2_7b import model as qwen2_model
+from models.common.models.qwen3_32b import model as qwen3_32b_model
+from models.common.models.qwen25_7b import model as qwen25_model
+from models.common.models.qwen25_72b import model as qwen25_72b_model
+from models.common.models.qwen25_coder_32b import model as qwen25_coder_32b_model
+
+MODEL_CONTRACTS = {
+    "llama32_1b": SimpleNamespace(
+        module=llama32_model,
+        model_class=llama32_model.Llama32_1BTransformer1D,
+        config_class=llama32_model.Llama32_1BTransformer1DConfig,
+        attention_config_class=llama32_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_llama32_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "llama32_3b": SimpleNamespace(
+        module=llama32_3b_model,
+        model_class=llama32_3b_model.Llama32_3BTransformer1D,
+        config_class=llama32_3b_model.Llama32_3BTransformer1DConfig,
+        attention_config_class=llama32_3b_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_llama32_config(module=llama32_3b_model, **kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch, module=llama32_3b_model),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "llama33_70b": SimpleNamespace(
+        module=llama33_70b_model,
+        model_class=llama33_70b_model.Llama33_70BTransformer1D,
+        config_class=llama33_70b_model.Llama33_70BTransformer1DConfig,
+        attention_config_class=llama33_70b_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_llama32_config(module=llama33_70b_model, **kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch, module=llama33_70b_model),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "qwen2_7b": SimpleNamespace(
+        module=qwen2_model,
+        model_class=qwen2_model.Qwen2_7B,
+        config_class=qwen2_model.Qwen2_7BTransformerConfig,
+        attention_config_class=qwen2_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen2_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "qwen25_7b": SimpleNamespace(
+        module=qwen25_model,
+        model_class=qwen25_model.Qwen25_7B,
+        config_class=qwen25_model.Qwen25_7BTransformerConfig,
+        attention_config_class=qwen25_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen2_config(module=qwen25_model, **kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch, module=qwen25_model),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "qwen25_72b": SimpleNamespace(
+        module=qwen25_72b_model,
+        model_class=qwen25_72b_model.Qwen25_72B,
+        config_class=qwen25_72b_model.Qwen25_72BConfig,
+        attention_config_class=qwen25_72b_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen25_72b_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen25_72b_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "qwen25_coder_32b": SimpleNamespace(
+        module=qwen25_coder_32b_model,
+        model_class=qwen25_coder_32b_model.Qwen25Coder32B,
+        config_class=qwen25_coder_32b_model.Qwen25Coder32BConfig,
+        attention_config_class=qwen25_coder_32b_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen25_coder_32b_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen25_coder_32b_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "qwen3_32b": SimpleNamespace(
+        module=qwen3_32b_model,
+        model_class=qwen3_32b_model.Qwen3_32B,
+        config_class=qwen3_32b_model.Qwen3_32BConfig,
+        attention_config_class=qwen3_32b_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen3_32b_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen3_32b_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "deepseek_r1_distill_qwen_14b": SimpleNamespace(
+        module=deepseek_model,
+        model_class=deepseek_model.DeepSeekR1Qwen14B,
+        config_class=deepseek_model.DeepSeekR1Qwen14BTransformerConfig,
+        attention_config_class=deepseek_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_qwen2_config(module=deepseek_model, **kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch, module=deepseek_model),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "mistral_7b": SimpleNamespace(
+        module=mistral_model,
+        model_class=mistral_model.Mistral7B,
+        config_class=mistral_model.Mistral7BTransformerConfig,
+        attention_config_class=mistral_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_mistral_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_mistral_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+    "phi4": SimpleNamespace(
+        module=phi4_model,
+        model_class=phi4_model.Phi4Transformer,
+        config_class=phi4_model.Phi4TransformerConfig,
+        attention_config_class=phi4_model.Attention1DConfig,
+        make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs),
+        make_config=lambda **kwargs: _make_phi4_config(**kwargs),
+        make_layer=lambda attention_config=None: _make_llama32_layer(attention_config),
+        construct_model=lambda monkeypatch: _construct_phi4_model(monkeypatch),
+        expected_module_names=(
+            "layer[0].attn_norm",
+            "layer[0].attention",
+            "layer[0].ff_norm",
+            "layer[0].mlp",
+            "layer[1].attn_norm",
+            "layer[1].attention",
+            "layer[1].ff_norm",
+            "layer[1].mlp",
+            "final_norm",
+            "lm_head",
+        ),
+    ),
+}
+
+
+@pytest.fixture(params=MODEL_CONTRACTS.items(), ids=lambda item: item[0])
+def contract(request):
+    return request.param[1]
+
+
+@dataclass(frozen=True)
+class _PagedAttentionConfig:
+    block_size: int
+    max_num_blocks: int
+
+
+def _make_llama32_attention_config(*, n_kv_heads=8, kv_cache=None):
+    return SimpleNamespace(
+        n_kv_heads=n_kv_heads,
+        use_vllm_paged_kv_cache=True,
+        paged_attention_config=_PagedAttentionConfig(block_size=32, max_num_blocks=128),
+        kv_cache=kv_cache,
+    )
+
+
+def _make_llama32_layer(attention_config=None):
+    attention_config = attention_config or _make_llama32_attention_config()
+    attention = SimpleNamespace(config=attention_config, kv_cache=attention_config.kv_cache)
+    return SimpleNamespace(
+        attention_norm=object(),
+        attention=attention,
+        self_attn=attention,
+        ff_norm=object(),
+        feed_forward=object(),
+    )
+
+
+def _make_llama32_config(*, module=llama32_model, n_layers=1, num_devices=2, n_kv_heads=8, sampling_config=None):
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    config_class = next(
+        getattr(module, name)
+        for name in (
+            "Llama32_1BTransformer1DConfig",
+            "Llama32_3BTransformer1DConfig",
+            "Llama33_70BTransformer1DConfig",
+        )
+        if hasattr(module, name)
+    )
+    return config_class(
+        n_layers=n_layers,
+        vocab_size=128256,
+        max_batch_size=4,
+        max_seq_len=4096,
+        dim=2048,
+        num_devices=num_devices,
+        mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices),
+        embedding_config=object(),
+        rope_config=object(),
+        block_configs=block_configs,
+        norm_config=object(),
+        lm_head_config=object(),
+        sampling_config=sampling_config,
+    )
+
+
+def _construct_llama32_model(monkeypatch, *, module=llama32_model):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    for owner_name, sentinel_name in (
+        ("Embedding1D", "embedding"),
+        ("RotarySetup1D", "rope_setup"),
+        ("TransformerBlock1D", "layer"),
+        ("RMSNorm1D", "norm"),
+        ("LMHead1D", "lm_head"),
+        ("Sampling1D", "sampling"),
+    ):
+        monkeypatch.setattr(
+            getattr(module, owner_name),
+            "from_config",
+            MagicMock(return_value=sentinels[sentinel_name]),
+        )
+    config = _make_llama32_config(module=module, num_devices=1, sampling_config=object())
+    model_class = next(
+        getattr(module, name)
+        for name in (
+            "Llama32_1BTransformer1D",
+            "Llama32_3BTransformer1D",
+            "Llama33_70BTransformer1D",
+        )
+        if hasattr(module, name)
+    )
+    return model_class(config), config, sentinels
+
+
+def _make_qwen2_config(*, module=qwen2_model, n_layers=1, num_devices=2, n_kv_heads=4, sampling_config=None):
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    config_class = (
+        getattr(module, "Qwen2_7BTransformerConfig", None)
+        or getattr(module, "Qwen25_7BTransformerConfig", None)
+        or module.DeepSeekR1Qwen14BTransformerConfig
+    )
+    return config_class(
+        n_layers=n_layers,
+        vocab_size=152064,
+        max_batch_size=4,
+        max_seq_len=4096,
+        dim=3584,
+        num_devices=num_devices,
+        mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices),
+        embedding_config=object(),
+        rope_config=object(),
+        block_configs=block_configs,
+        norm_config=object(),
+        lm_head_config=object(),
+        sampling_config=sampling_config,
+        tt_ccl=object(),
+    )
+
+
+def _construct_qwen2_model(monkeypatch, *, module=qwen2_model):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    layer_class = (
+        "Qwen2_7BDecoderLayer"
+        if module is qwen2_model
+        else ("Qwen25_7BDecoderLayer" if module is qwen25_model else "DeepSeekR1Qwen14BDecoderLayer")
+    )
+    for owner_name, sentinel_name in (
+        ("Embedding1D", "embedding"),
+        ("RotarySetup1D", "rope_setup"),
+        (layer_class, "layer"),
+        ("RMSNorm1D", "norm"),
+        ("LMHead1D", "lm_head"),
+        ("Sampling1D", "sampling"),
+    ):
+        monkeypatch.setattr(
+            getattr(module, owner_name),
+            "from_config",
+            MagicMock(return_value=sentinels[sentinel_name]),
+        )
+    config = _make_qwen2_config(module=module, sampling_config=object())
+    model_class = getattr(module, "Qwen2_7B", None) or getattr(module, "Qwen25_7B", None) or module.DeepSeekR1Qwen14B
+    return model_class(config), config, sentinels
+
+
+def _make_qwen25_72b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None):
+    del sampling_config
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    return qwen25_72b_model.Qwen25_72BConfig(
+        hf_model_id="Qwen/Qwen2.5-72B-Instruct",
+        dim=8192,
+        n_heads=64,
+        n_kv_heads=n_kv_heads,
+        head_dim=128,
+        hidden_dim=29568,
+        vocab_size=152064,
+        rms_norm_eps=1e-6,
+        rope_theta=1_000_000.0,
+        num_hidden_layers=n_layers,
+        max_batch_size=4,
+        max_seq_len=4096,
+        rope_table_len=8192,
+        num_devices=num_devices,
+        mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices),
+        block_configs=block_configs,
+    )
+
+
+def _construct_qwen25_72b_model(monkeypatch):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    monkeypatch.setattr(qwen25_72b_model, "get_tt_ccl", MagicMock(return_value=object()))
+    monkeypatch.setattr(qwen25_72b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"]))
+    config = _make_qwen25_72b_config()
+    return (
+        qwen25_72b_model.Qwen25_72B(
+            config,
+            sentinels["embedding"],
+            sentinels["rope_setup"],
+            [sentinels["layer"]],
+            sentinels["norm"],
+            sentinels["lm_head"],
+            config.mesh_device,
+        ),
+        config,
+        sentinels,
+    )
+
+
+def _make_qwen25_coder_32b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None):
+    del sampling_config
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    return qwen25_coder_32b_model.Qwen25Coder32BConfig(
+        hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+        dim=5120,
+        n_heads=40,
+        n_kv_heads=n_kv_heads,
+        head_dim=128,
+        hidden_dim=27648,
+        vocab_size=152064,
+        rms_norm_eps=1e-6,
+        rope_theta=1_000_000.0,
+        num_hidden_layers=n_layers,
+        max_batch_size=4,
+        max_seq_len=4096,
+        rope_table_len=8192,
+        num_devices=num_devices,
+        mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices),
+        block_configs=block_configs,
+    )
+
+
+def _construct_qwen25_coder_32b_model(monkeypatch):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    monkeypatch.setattr(qwen25_coder_32b_model, "get_tt_ccl", MagicMock(return_value=object()))
+    monkeypatch.setattr(qwen25_coder_32b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"]))
+    config = _make_qwen25_coder_32b_config()
+    return (
+        qwen25_coder_32b_model.Qwen25Coder32B(
+            config,
+            sentinels["embedding"],
+            sentinels["rope_setup"],
+            [sentinels["layer"]],
+            sentinels["norm"],
+            sentinels["lm_head"],
+            config.mesh_device,
+        ),
+        config,
+        sentinels,
+    )
+
+
+def _make_qwen3_32b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None):
+    del sampling_config
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    return qwen3_32b_model.Qwen3_32BConfig(
+        hf_model_id="Qwen/Qwen3-32B",
+        dim=5120,
+        n_heads=64,
+        n_kv_heads=n_kv_heads,
+        head_dim=128,
+        hidden_dim=27648,
+        vocab_size=151936,
+        rms_norm_eps=1e-6,
+        rope_theta=1_000_000.0,
+        num_hidden_layers=n_layers,
+        max_batch_size=4,
+        max_seq_len=4096,
+        rope_table_len=8192,
+        num_devices=num_devices,
+        mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices),
+        block_configs=block_configs,
+    )
+
+
+def _construct_qwen3_32b_model(monkeypatch):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    monkeypatch.setattr(qwen3_32b_model, "get_tt_ccl", MagicMock(return_value=object()))
+    monkeypatch.setattr(qwen3_32b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"]))
+    monkeypatch.setattr(qwen3_32b_model.weight_utils, "lm_head_padded_vocab_size", MagicMock(return_value=152064))
+    config = _make_qwen3_32b_config()
+    return (
+        qwen3_32b_model.Qwen3_32B(
+            config,
+            sentinels["embedding"],
+            sentinels["rope_setup"],
+            [sentinels["layer"]],
+            sentinels["norm"],
+            sentinels["lm_head"],
+            config.mesh_device,
+        ),
+        config,
+        sentinels,
+    )
+
+
+def _make_mistral_config(*, n_layers=1, num_devices=2, n_kv_heads=8, sampling_config=None):
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    return mistral_model.Mistral7BTransformerConfig(
+        n_layers=n_layers,
+        vocab_size=32768,
+        max_batch_size=4,
+        max_seq_len=4096,
+        dim=4096,
+        num_devices=num_devices,
+        mesh_device=object(),
+        embedding_config=object(),
+        rope_config=object(),
+        block_configs=block_configs,
+        norm_config=object(),
+        lm_head_config=object(),
+        sampling_config=sampling_config,
+        tt_ccl=object(),
+    )
+
+
+def _construct_mistral_model(monkeypatch):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    for owner_name, sentinel_name in (
+        ("Embedding1D", "embedding"),
+        ("RotarySetup1D", "rope_setup"),
+        ("Mistral7BDecoderLayer", "layer"),
+        ("RMSNorm1D", "norm"),
+        ("LMHead1D", "lm_head"),
+        ("Sampling1D", "sampling"),
+    ):
+        monkeypatch.setattr(
+            getattr(mistral_model, owner_name),
+            "from_config",
+            MagicMock(return_value=sentinels[sentinel_name]),
+        )
+    config = _make_mistral_config(sampling_config=object())
+    return mistral_model.Mistral7B(config), config, sentinels
+
+
+def _make_phi4_config(*, n_layers=1, num_devices=2, n_kv_heads=10, sampling_config=None):
+    block_configs = [
+        SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers)
+    ]
+    return phi4_model.Phi4TransformerConfig(
+        n_layers=n_layers,
+        vocab_size=100352,
+        max_batch_size=4,
+        max_seq_len=4096,
+        dim=5120,
+        num_devices=num_devices,
+        mesh_device=object(),
+        embedding_config=object(),
+        rope_config=object(),
+        block_configs=block_configs,
+        norm_config=object(),
+        lm_head_config=object(),
+        sampling_config=sampling_config,
+        tt_ccl=object(),
+    )
+
+
+def _construct_phi4_model(monkeypatch):
+    sentinels = {
+        "embedding": object(),
+        "rope_setup": object(),
+        "layer": _make_llama32_layer(),
+        "norm": object(),
+        "lm_head": object(),
+        "sampling": object(),
+    }
+    for owner_name, sentinel_name in (
+        ("Embedding1D", "embedding"),
+        ("RotarySetup1D", "rope_setup"),
+        ("Phi4DecoderLayer", "layer"),
+        ("RMSNorm1D", "norm"),
+        ("LMHead1D", "lm_head"),
+        ("Sampling1D", "sampling"),
+    ):
+        monkeypatch.setattr(
+            getattr(phi4_model, owner_name),
+            "from_config",
+            MagicMock(return_value=sentinels[sentinel_name]),
+        )
+    config = _make_phi4_config(sampling_config=object())
+    return phi4_model.Phi4Transformer(config), config, sentinels
+
+
+def test_config_exposes_complete_runtime_metadata(contract):
+    names = {field.name for field in fields(contract.config_class)}
+    assert {
+        "dim",
+        "mesh_device",
+        "num_devices",
+        "n_layers",
+        "max_batch_size",
+        "max_seq_len",
+        "block_configs",
+    } <= names
+    assert "n_kv_heads" in {field.name for field in fields(contract.attention_config_class)}
+
+    config = contract.make_config(n_layers=2, num_devices=2, n_kv_heads=8)
+    assert len(config.block_configs) == config.n_layers
+    for block in config.block_configs:
+        assert hasattr(block.attention_config, "n_kv_heads")
+        assert block.attention_config.n_kv_heads % config.num_devices == 0
+
+
+@pytest.mark.parametrize(
+    "method,names,keyword_only,defaults",
+    [
+        (
+            "prefill_forward",
+            (
+                "self",
+                "x_embed",
+                "rot_mats",
+                "user_id",
+                "page_table",
+                "chunk_page_table",
+                "chunk_start_idx",
+                "get_last_token",
+                "batch_size",
+                "chunk_start_idx_tensor",
+                "last_token_slice",
+                "last_token_index",
+            ),
+            (),
+            {
+                "user_id": 0,
+                "page_table": None,
+                "chunk_page_table": None,
+                "chunk_start_idx": None,
+                "get_last_token": -1,
+                "batch_size": 1,
+                "chunk_start_idx_tensor": None,
+                "last_token_slice": None,
+                "last_token_index": None,
+            },
+        ),
+        (
+            "post_process_prefill_output",
+            ("self", "hidden_states", "last_token_idx", "last_token_slice", "last_token_index"),
+            (),
+            {"last_token_slice": None, "last_token_index": None},
+        ),
+        (
+            "post_process_batched_prefill_output",
+            (
+                "self",
+                "hidden_states",
+                "last_token_idx_list",
+                "padded_batch",
+                "prefill_seq_len",
+                "last_token_slice",
+                "last_token_index",
+            ),
+            (),
+            {"last_token_slice": None, "last_token_index": None},
+        ),
+        ("set_kv_cache", ("self", "kv_cache"), (), {}),
+        (
+            "configure_paged_attention",
+            ("self", "block_size", "max_num_blocks"),
+            ("block_size", "max_num_blocks"),
+            {},
+        ),
+        ("prepare_prefill_rot_mats", ("self", "position_indices"), (), {}),
+        ("iter_executor_named_modules", ("self",), (), {}),
+    ],
+)
+def test_method_signatures_are_exact(contract, method, names, keyword_only, defaults):
+    signature = inspect.signature(getattr(contract.model_class, method))
+    parameters = signature.parameters
+    assert tuple(parameters) == names
+    for name, parameter in parameters.items():
+        expected_kind = (
+            inspect.Parameter.KEYWORD_ONLY if name in keyword_only else inspect.Parameter.POSITIONAL_OR_KEYWORD
+        )
+        assert parameter.kind is expected_kind
+        expected_default = defaults.get(name, inspect.Parameter.empty)
+        assert parameter.default == expected_default
+
+    annotations = {
+        "prefill_forward": (
+            (
+                inspect.Parameter.empty,
+                "ttnn.Tensor",
+                "tuple[ttnn.Tensor, ttnn.Tensor]",
+                "int",
+                "ttnn.Tensor | None",
+                "ttnn.Tensor | None",
+                "int | None",
+                "int",
+                "int",
+                "ttnn.Tensor | None",
+                "tuple[ttnn.Tensor, ttnn.Tensor] | None",
+                "ttnn.Tensor | None",
+            ),
+            "ttnn.Tensor",
+        ),
+        "post_process_prefill_output": (
+            (
+                inspect.Parameter.empty,
+                "ttnn.Tensor",
+                "int",
+                "tuple[ttnn.Tensor, ttnn.Tensor] | None",
+                "ttnn.Tensor | None",
+            ),
+            "ttnn.Tensor",
+        ),
+        "post_process_batched_prefill_output": (
+            (
+                inspect.Parameter.empty,
+                "ttnn.Tensor",
+                "list[int]",
+                "int",
+                "int",
+                "tuple[ttnn.Tensor, ttnn.Tensor] | None",
+                "ttnn.Tensor | None",
+            ),
+            "ttnn.Tensor",
+        ),
+        "set_kv_cache": ((inspect.Parameter.empty, "list | None"), "None"),
+        "configure_paged_attention": ((inspect.Parameter.empty, "int", "int"), "None"),
+        "prepare_prefill_rot_mats": (
+            (inspect.Parameter.empty, "ttnn.Tensor"),
+            "tuple[ttnn.Tensor, ttnn.Tensor]",
+        ),
+        "iter_executor_named_modules": (
+            (inspect.Parameter.empty,),
+            inspect.Signature.empty,
+        ),
+    }
+    expected_parameter_annotations, expected_return_annotation = annotations[method]
+    assert tuple(parameter.annotation for parameter in parameters.values()) == expected_parameter_annotations
+    assert signature.return_annotation == expected_return_annotation
+
+
+def test_required_methods_exist(contract):
+    for method in (
+        "set_kv_cache",
+        "configure_paged_attention",
+        "prepare_prefill_rot_mats",
+        "iter_executor_named_modules",
+        "embed_decode",
+        "embed_prefill",
+        "gather_and_untilize_logits",
+        "increment_positions",
+    ):
+        assert callable(getattr(contract.model_class, method))
+
+
+def test_constructed_model_resolves_runtime_surface(contract, monkeypatch):
+    model, config, sentinels = contract.construct_model(monkeypatch)
+
+    assert model.config is config
+    assert model.embedding is sentinels["embedding"]
+    assert model.rope_setup is sentinels["rope_setup"]
+    assert model.layers == [sentinels["layer"]]
+    assert model.norm is sentinels["norm"]
+    assert model.lm_head is sentinels["lm_head"]
+    assert model.sampling is sentinels["sampling"]
+    assert model.supports_on_device_sampling
+    assert model.mesh_device is config.mesh_device
+    assert model.vocab_size == config.vocab_size
+    assert model.n_layers == config.n_layers
+    assert model.num_devices == config.num_devices
+    assert model.model_args is None
+
+
+def test_named_modules_are_complete_unique_and_ordered(contract):
+    layers = [contract.make_layer(), contract.make_layer()]
+    model = SimpleNamespace(layers=layers, norm=object(), lm_head=object())
+    named = list(contract.model_class.iter_executor_named_modules(model))
+
+    assert tuple(name for name, _ in named) == contract.expected_module_names
+    assert len(named) == 4 * len(layers) + 2
+    assert len({name for name, _ in named}) == len(named)
+    assert tuple(module for _, module in named) == (
+        layers[0].attention_norm,
+        layers[0].attention,
+        layers[0].ff_norm,
+        layers[0].feed_forward,
+        layers[1].attention_norm,
+        layers[1].attention,
+        layers[1].ff_norm,
+        layers[1].feed_forward,
+        model.norm,
+        model.lm_head,
+    )
+
+
+def test_named_modules_without_layers_yields_nothing(contract):
+    assert list(contract.model_class.iter_executor_named_modules(SimpleNamespace())) == []
+
+
+def test_set_kv_cache_binds_identity_and_unbinds_idempotently(contract):
+    layers = [contract.make_layer(), contract.make_layer()]
+    model = SimpleNamespace(layers=layers)
+    cache = [[object(), object()], [object(), object()]]
+
+    contract.model_class.set_kv_cache(model, cache)
+    for layer, expected in zip(layers, cache):
+        bound = layer.attention.config.kv_cache
+        assert bound == tuple(expected)
+        assert bound[0] is expected[0]
+        assert bound[1] is expected[1]
+        assert layer.attention.kv_cache is bound
+
+    contract.model_class.set_kv_cache(model, None)
+    contract.model_class.set_kv_cache(model, None)
+    assert all(layer.attention.config.kv_cache is None for layer in layers)
+    assert all(layer.attention.kv_cache is None for layer in layers)
+
+
+def test_set_kv_cache_rejects_wrong_layer_count_before_binding(contract, expect_error):
+    layers = [contract.make_layer(), contract.make_layer()]
+    model = SimpleNamespace(layers=layers)
+
+    with expect_error(ValueError, "model has 2 layers"):
+        contract.model_class.set_kv_cache(model, [[object(), object()]])
+
+    assert all(layer.attention.config.kv_cache is None for layer in layers)
+    assert all(layer.attention.kv_cache is None for layer in layers)
+
+
+@pytest.mark.parametrize("bad_pair", [[object()], object()])
+def test_set_kv_cache_validates_all_pairs_before_binding(contract, bad_pair, expect_error):
+    layers = [contract.make_layer(), contract.make_layer()]
+    model = SimpleNamespace(layers=layers)
+
+    with expect_error((TypeError, ValueError), "layer 1.*K/V tensor"):
+        contract.model_class.set_kv_cache(model, [[object(), object()], bad_pair])
+
+    assert all(layer.attention.config.kv_cache is None for layer in layers)
+    assert all(layer.attention.kv_cache is None for layer in layers)
+
+
+def test_configure_paged_attention_updates_construction_and_live_configs_and_rejects_bound_cache(
+    contract, expect_error
+):
+    construction = contract.make_attention_config()
+    live = contract.make_attention_config()
+    model = SimpleNamespace(
+        config=SimpleNamespace(block_configs=(SimpleNamespace(attention_config=construction),)),
+        layers=(contract.make_layer(live),),
+    )
+
+    contract.model_class.configure_paged_attention(model, block_size=16, max_num_blocks=200)
+    assert construction.paged_attention_config.block_size == live.paged_attention_config.block_size == 16
+    assert construction.paged_attention_config.max_num_blocks == live.paged_attention_config.max_num_blocks == 200
+
+    bound = object()
+    live.kv_cache = (bound, bound)
+    with expect_error(RuntimeError, "already has a bound KV cache"):
+        contract.model_class.configure_paged_attention(model, block_size=32, max_num_blocks=128)
+    assert construction.paged_attention_config.block_size == live.paged_attention_config.block_size == 16
+
+
+def test_prepare_prefill_rot_mats_gathers_device_rows(contract, monkeypatch):
+    position_indices = object()
+    cos_matrix, sin_matrix = object(), object()
+    cos_rows, sin_rows = object(), object()
+    cos_4d, sin_4d = object(), object()
+    fake_ttnn = SimpleNamespace(
+        TILE_LAYOUT=object(),
+        embedding=MagicMock(side_effect=[cos_rows, sin_rows]),
+        unsqueeze_to_4D=MagicMock(side_effect=[cos_4d, sin_4d]),
+    )
+    rope = SimpleNamespace(cos_matrix=cos_matrix, sin_matrix=sin_matrix, load_device_weights=MagicMock())
+    monkeypatch.setattr(contract.module, "ttnn", fake_ttnn)
+
+    result = contract.model_class.prepare_prefill_rot_mats(SimpleNamespace(rope_setup=rope), position_indices)
+
+    rope.load_device_weights.assert_called_once_with()
+    assert fake_ttnn.embedding.call_args_list == [
+        call(position_indices, cos_matrix, layout=fake_ttnn.TILE_LAYOUT),
+        call(position_indices, sin_matrix, layout=fake_ttnn.TILE_LAYOUT),
+    ]
+    assert result == (cos_4d, sin_4d)

code/models/common/tests/llm_runtime/test_model_executor.py

@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Focused contracts for the family-neutral model composition root."""
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+import ttnn
+from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, WarmupConfig
+from models.common.models import executor as executor_module
+from models.common.models.executor import ModelExecutor, ModelExecutorConfig
+
+_EXECUTOR_PATH = Path(__file__).parents[2] / "models" / "executor.py"
+_MODELS_ROOT = _EXECUTOR_PATH.parent
+
+
+def _config(*, device_sampling_enabled: bool = False) -> ModelExecutorConfig:
+    return ModelExecutorConfig(
+        trace=TraceConfig(mode="none"),
+        warmup=WarmupConfig(),
+        paged_kv_cache=PagedKVCacheConfig(
+            block_size=32,
+            max_num_blocks=128,
+            num_blocks=128,
+            dtype=ttnn.bfloat8_b,
+        ),
+        device_sampling_enabled=device_sampling_enabled,
+    )
+
+
+def test_common_executor_has_no_concrete_model_dependencies_or_dispatch() -> None:
+    tree = ast.parse(_EXECUTOR_PATH.read_text())
+    imports = {node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) and node.module is not None}
+    imported_names = {alias.name for node in ast.walk(tree) if isinstance(node, ast.Import) for alias in node.names}
+    assert not any(module.startswith("models.common.models.") for module in imports)
+    assert not any(name.startswith("models.common.models.") for name in imported_names)
+
+    control_flow = [
+        node.test for node in ast.walk(tree) if isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert))
+    ]
+    dispatch_names = {"model_name", "model_id", "provider_id", "checkpoint_path", "model_version"}
+    assert not any(
+        isinstance(node, ast.Name) and node.id in dispatch_names
+        for expression in control_flow
+        for node in ast.walk(expression)
+    )
+
+    torch_calls = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "torch"
+    ]
+    assert torch_calls == []
+
+
+def test_model_layer_has_only_the_approved_family_modules_and_readmes() -> None:
+    assert (_MODELS_ROOT / "llama3_executor.py").is_file()
+    assert (_MODELS_ROOT / "qwen2_executor.py").is_file()
+    assert not (_MODELS_ROOT / "qwen3_executor.py").exists()
+
+    model_directories = sorted(
+        path for path in _MODELS_ROOT.iterdir() if path.is_dir() and (path / "model.py").is_file()
+    )
+    assert len(model_directories) == 12
+    assert all((path / "README.md").is_file() for path in model_directories)
+
+
+@pytest.mark.parametrize(
+    "relative_path",
+    (
+        "deepseek_r1_distill_qwen_14b/executor.py",
+        "mistral_7b/executor.py",
+        "phi4/executor.py",
+    ),
+)
+def test_direct_composition_examples_do_not_depend_on_shared_family_executors(relative_path: str) -> None:
+    tree = ast.parse((_MODELS_ROOT / relative_path).read_text())
+    imports = {node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)}
+    assert "models.common.models.executor" not in imports
+    assert "models.common.models.llama3_executor" not in imports
+    assert "models.common.models.qwen2_executor" not in imports
+
+
+def test_common_config_is_frozen_and_rejects_non_exact_nested_config_types(expect_error) -> None:
+    config = _config()
+    with expect_error(AttributeError, "cannot assign to field"):
+        config.device_sampling_enabled = True
+
+    class TraceConfigSubclass(TraceConfig):
+        pass
+
+    with expect_error(TypeError, "trace must be exactly TraceConfig"):
+        ModelExecutorConfig(
+            trace=TraceConfigSubclass(mode="none"),
+            warmup=config.warmup,
+            paged_kv_cache=config.paged_kv_cache,
+            device_sampling_enabled=False,
+        )
+
+
+def test_sampling_state_inputs_are_an_optional_owned_pair(expect_error) -> None:
+    with expect_error(ValueError, "must be supplied together"):
+        ModelExecutor(
+            None,
+            None,
+            _config(device_sampling_enabled=True),
+            sampling_state_controller=object(),
+        )
+    with expect_error(ValueError, "requires device sampling"):
+        ModelExecutor(
+            None,
+            None,
+            _config(),
+            sampling_state_controller=object(),
+            sampling_state=object(),
+        )
+
+
+@pytest.mark.parametrize(
+    ("enable_trace", "expected"),
+    [(True, ["prime", "coordinator"]), (False, ["coordinator", "prime"])],
+)
+def test_prefill_warmup_policy_controls_order_through_one_continuation(enable_trace, expected) -> None:
+    events = []
+
+    def policy(executor, default_warmup, *, kv_cache, can_sample_on_device, enable_trace):
+        assert executor is target
+        assert kv_cache is cache
+        assert can_sample_on_device
+        if enable_trace:
+            events.append("prime")
+        default_warmup()
+        if not enable_trace:
+            events.append("prime")
+
+    cache = object()
+    target = object.__new__(ModelExecutor)
+    target._terminal = False
+    target.prefill_runtime = SimpleNamespace(transient_orphan_count=0)
+    target.decode_runtime = SimpleNamespace(transient_orphan_count=0)
+    target._prefill_warmup = policy
+    target.warmup = SimpleNamespace(
+        warmup_prefill=lambda **kwargs: events.append("coordinator"),
+    )
+
+    target.warmup_model_prefill(
+        kv_cache=cache,
+        can_sample_on_device=True,
+        enable_trace=enable_trace,
+    )
+
+    assert events == expected
+
+
+def test_request_state_and_execution_target_are_forwarded_by_identity() -> None:
+    target = object.__new__(ModelExecutor)
+    target._terminal = False
+    target.prefill_runtime = SimpleNamespace(transient_orphan_count=0)
+    target.decode_runtime = SimpleNamespace(transient_orphan_count=0)
+    target._validate_bound_cache = MagicMock()
+    target._ensure_sampling_for = MagicMock()
+    target._prefill_execution = MagicMock()
+    target._request_state_fields = ("prompt_tokens", "output_tokens", "slot_remap")
+
+    values = {name: object() for name in ("tokens", "page_table", "prompt_tokens", "output_tokens", "slot_remap")}
+    target.compile_prefill(
+        tokens=values["tokens"],
+        page_table=values["page_table"],
+        prompt_tokens=values["prompt_tokens"],
+        output_tokens=values["output_tokens"],
+        slot_remap=values["slot_remap"],
+    )
+
+    forwarded = target._prefill_execution.compile_prefill.call_args.kwargs
+    for name, value in values.items():
+        assert forwarded[name] is value
+
+
+def test_layout_refresh_preserves_owner_and_sampling_state_identity(monkeypatch) -> None:
+    state = object()
+    layout = object()
+    prefill_config = SimpleNamespace(
+        max_batch_size=4,
+        max_prefill_chunk_size=2048,
+        device_sampling_enabled=True,
+        can_enable_trace=lambda *_: True,
+        supports_batched_prefill=True,
+        disable_batched_prefill=True,
+        max_prefill_batch_size=4,
+        batched_prefill_batched_extract=True,
+        trace_capture_prime_sequence_lengths=(128,),
+        sampling_state_controller=object(),
+        sampling_state=state,
+    )
+    decode_config = SimpleNamespace(
+        lane_capacity=4,
+        device_sampling_enabled=True,
+        force_greedy_top_k=True,
+        sampling_state_controller=prefill_config.sampling_state_controller,
+        sampling_state=state,
+    )
+    warmup_config = SimpleNamespace(warmup=object(), prefill_sequence_lengths=(128,))
+    resolved_prefill = SimpleNamespace(page_table_layout=layout, sampling_state=state)
+    resolved_decode = SimpleNamespace(page_table_layout=layout, sampling_state=state)
+    resolved_warmup = SimpleNamespace(page_table_layout=layout)
+    prefill_resolve = MagicMock(return_value=resolved_prefill)
+    decode_resolve = MagicMock(return_value=resolved_decode)
+    warmup_resolve = MagicMock(return_value=resolved_warmup)
+    monkeypatch.setattr(executor_module.PrefillRuntimeConfig, "resolve", prefill_resolve)
+    monkeypatch.setattr(executor_module.DecodeRuntimeConfig, "resolve", decode_resolve)
+    monkeypatch.setattr(executor_module.WarmupCoordinatorConfig, "resolve", warmup_resolve)
+
+    target = object.__new__(ModelExecutor)
+    target.model = object()
+    target.output_reader = object()
+    target.config = SimpleNamespace(trace=object())
+    target.prefill_runtime = SimpleNamespace(config=prefill_config)
+    target.decode_runtime = SimpleNamespace(config=decode_config)
+    target.warmup = SimpleNamespace(config=warmup_config)
+    target._resolve_page_table_layout = lambda: layout
+    owners = (target.prefill_runtime, target.decode_runtime, target.warmup)
+
+    target._refresh_page_table_layout()
+
+    assert (target.prefill_runtime, target.decode_runtime, target.warmup) == owners
+    assert target.page_table_layout is layout
+    assert all(owner.config.page_table_layout is layout for owner in owners)
+    assert target.prefill_runtime.config.sampling_state is state
+    assert target.decode_runtime.config.sampling_state is state
+    assert prefill_resolve.call_args.kwargs["trace_capture_prime_sequence_lengths"] == (128,)
+    assert prefill_resolve.call_args.kwargs["sampling_state_controller"] is prefill_config.sampling_state_controller
+    assert prefill_resolve.call_args.kwargs["sampling_state"] is state
+    assert decode_resolve.call_args.kwargs["sampling_state_controller"] is prefill_config.sampling_state_controller
+    assert decode_resolve.call_args.kwargs["sampling_state"] is state
+
+
+def test_cleanup_is_ordered_retryable_idempotent_and_terminal(expect_error) -> None:
+    events = []
+    failing = {"reader", "trace"}
+
+    class _Owner:
+        def __init__(self, name):
+            self.name = name
+
+        def action(self, *args):
+            events.append(self.name)
+            if self.name in failing:
+                raise RuntimeError(self.name)
+
+        cleanup = action
+        drain = action
+        drain_external_outputs = action
+        cleanup_transients = action
+        release = action
+
+    target = object.__new__(ModelExecutor)
+    target._terminal = False
+    target._cleaned_up = False
+    target._owner_name = "TestExecutor"
+    target.decode_runtime = _Owner("decode")
+    target.output_reader = _Owner("reader")
+    target.prefill_runtime = _Owner("prefill")
+    target.trace_compiler = _Owner("trace")
+    target.program_compiler = _Owner("program")
+    target.config = SimpleNamespace(device_sampling_enabled=True)
+    target.sampling_state_controller = _Owner("sampling-state")
+    target.sampling_state = object()
+    target.model = SimpleNamespace(sampling=_Owner("sampling"))
+    target.kv_cache_manager = _Owner("kv")
+
+    expected = ["decode", "reader", "prefill", "decode", "trace", "program", "sampling-state", "sampling", "kv"]
+    with expect_error(RuntimeError, "reader") as raised:
+        target.cleanup()
+    assert events == expected
+    assert [str(error) for error in raised.value.cleanup_failures] == ["trace"]
+    assert target.terminal
+    assert not target._cleaned_up
+
+    failing.clear()
+    target.cleanup()
+    target.cleanup()
+    assert events == expected * 2
+    assert target._cleaned_up

code/models/common/tests/llm_runtime/test_output_reader.py

@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import weakref
+
+import torch
+
+import ttnn
+from models.common.llm_runtime.output_reader import OutputReader, PendingRead
+
+
+class FakeDeviceTensor:
+    def __init__(self, name: str):
+        self.name = name
+        self.cpu_calls = []
+        self.host_value = object()
+
+    def cpu(self, *, blocking: bool):
+        self.cpu_calls.append(blocking)
+        return self.host_value
+
+
+class FailingDeviceTensor:
+    def __init__(self):
+        self.cpu_calls = []
+
+    def cpu(self, *, blocking: bool):
+        self.cpu_calls.append(blocking)
+        raise RuntimeError("copy failed")
+
+
+def _install_events(monkeypatch):
+    events = []
+    synchronized = []
+
+    def record_event(device, queue_id):
+        event = object()
+        events.append((device, queue_id, event))
+        return event
+
+    monkeypatch.setattr(ttnn, "record_event", record_event)
+    monkeypatch.setattr(ttnn, "event_synchronize", synchronized.append)
+    return events, synchronized
+
+
+def test_blocking_read_returns_nested_host_payload_without_retention(monkeypatch):
+    events, synchronized = _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    output = FakeDeviceTensor("output")
+    log_probs = FakeDeviceTensor("log_probs")
+
+    host = reader.read((output, log_probs), blocking=True)
+
+    assert host == (output.host_value, log_probs.host_value)
+    assert output.cpu_calls == [True]
+    assert log_probs.cpu_calls == [True]
+    assert events == []
+    assert synchronized == []
+
+
+def test_synchronized_read_submits_every_copy_then_synchronizes_device_once(monkeypatch):
+    synchronized = []
+    monkeypatch.setattr(ttnn, "synchronize_device", synchronized.append)
+    reader = OutputReader("mesh")
+    output = FakeDeviceTensor("output")
+    log_probs = FakeDeviceTensor("log_probs")
+
+    host = reader.read_synchronized((output, log_probs))
+
+    assert host == (output.host_value, log_probs.host_value)
+    assert output.cpu_calls == [False]
+    assert log_probs.cpu_calls == [False]
+    assert synchronized == ["mesh"]
+
+
+def test_async_read_retains_destination_and_event_until_completion(monkeypatch):
+    events, synchronized = _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    output = FakeDeviceTensor("output")
+
+    pending = reader.submit(output)
+
+    assert isinstance(pending, PendingRead)
+    assert pending.value is output.host_value
+    assert pending.events == (events[0][2],)
+    assert events[0][:2] == ("mesh", 0)
+    assert output.cpu_calls == [False]
+
+    assert reader.complete(pending) is output.host_value
+    assert synchronized == [events[0][2]]
+
+    assert reader.complete(pending) is output.host_value
+    assert synchronized == [events[0][2]]
+
+
+def test_async_read_can_be_completed_by_exact_unwrapped_value(monkeypatch):
+    events, synchronized = _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    output = FakeDeviceTensor("output")
+    pending = reader.submit(output)
+
+    assert reader.complete(pending.value) is pending.value
+    assert synchronized == [events[0][2]]
+
+
+def test_reader_rejects_another_readers_pending_handle(monkeypatch, expect_error):
+    _install_events(monkeypatch)
+    first_reader = OutputReader("first_mesh")
+    second_reader = OutputReader("second_mesh")
+    foreign_pending = first_reader.submit(FakeDeviceTensor("first"))
+    second_reader.submit(FakeDeviceTensor("second"))
+
+    with expect_error(ValueError, "not owned"):
+        second_reader.complete(foreign_pending)
+
+    first_reader.drain()
+    second_reader.drain()
+
+
+def test_reader_rejects_another_readers_completed_pending_handle(monkeypatch, expect_error):
+    _install_events(monkeypatch)
+    first_reader = OutputReader("first_mesh")
+    second_reader = OutputReader("second_mesh")
+    foreign_pending = first_reader.submit(torch.tensor([1]))
+
+    with expect_error(ValueError, "not owned"):
+        second_reader.complete(foreign_pending)
+
+
+def test_drain_completes_every_pending_read_and_is_idempotent(monkeypatch, expect_error):
+    events, synchronized = _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    first = reader.submit(FakeDeviceTensor("first"))
+    second = reader.submit(FakeDeviceTensor("second"))
+
+    reader.drain()
+
+    assert synchronized == [events[0][2], events[1][2]]
+    assert reader.complete(first) is first.value
+    assert reader.complete(second) is second.value
+
+    reader.drain()
+    assert synchronized == [events[0][2], events[1][2]]
+
+
+def test_async_host_only_payload_is_already_complete(monkeypatch):
+    events, synchronized = _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    host_tensor = torch.tensor([1, 2])
+
+    pending = reader.submit((host_tensor, None))
+
+    assert torch.equal(pending.value[0], host_tensor)
+    assert pending.value[1] is None
+    assert pending.events == ()
+    assert reader.complete(pending) is pending.value
+    assert events == []
+    assert synchronized == []
+
+
+def test_nested_dict_and_list_preserve_shape(monkeypatch):
+    _install_events(monkeypatch)
+    reader = OutputReader("mesh")
+    output = FakeDeviceTensor("output")
+    log_probs = FakeDeviceTensor("log_probs")
+
+    pending = reader.submit({"outputs": [output], "log_probs": log_probs})
+
+    assert pending.value == {"outputs": [output.host_value], "log_probs": log_probs.host_value}
+    reader.complete(pending)
+
+
+def test_record_event_failure_synchronizes_device(monkeypatch):
+    device_synchronizations = []
+
+    # record_event is an external nanobind API with optional backend arguments.
+    def fail_record_event(*_args, **_kwargs):
+        raise RuntimeError("record failed")
+
+    monkeypatch.setattr(ttnn, "record_event", fail_record_event)
+    monkeypatch.setattr(ttnn, "synchronize_device", device_synchronizations.append)
+    reader = OutputReader("mesh")
+
+    try:
+        reader.submit(FakeDeviceTensor("output"))
+    except RuntimeError as error:
+        assert str(error) == "record failed"
+    else:
+        raise AssertionError("record_event failure was not propagated")
+
+    assert device_synchronizations == ["mesh"]
+
+
+def test_partial_nested_async_copy_failure_synchronizes_while_destinations_are_retained(monkeypatch, expect_error):
+    destination_refs = []
+    device_synchronizations = []
+
+    class HostDestination:
+        pass
+
+    class EphemeralDeviceTensor:
+        def cpu(self, *, blocking: bool):
+            destination = HostDestination()
+            destination_refs.append(weakref.ref(destination))
+            return destination
+
+    def synchronize_device(mesh_device):
+        assert destination_refs[0]() is not None
+        device_synchronizations.append(mesh_device)
+
+    monkeypatch.setattr(ttnn, "synchronize_device", synchronize_device)
+    reader = OutputReader("mesh")
+    failing = FailingDeviceTensor()
+
+    with expect_error(RuntimeError, "copy failed"):
+        reader.submit((EphemeralDeviceTensor(), failing))
+
+    assert failing.cpu_calls == [False]
+    assert device_synchronizations == ["mesh"]

code/models/common/tests/llm_runtime/test_paged_kv_cache.py

@@ -0,0 +1,415 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import ttnn
+from models.common.llm_runtime.config import PagedKVCacheConfig
+from models.common.llm_runtime.paged_kv_cache import PagedKVCacheManager, torch_dtype_for_ttnn
+
+
+class FakeMesh:
+    def __init__(self, num_devices=2):
+        self._num_devices = num_devices
+
+    def get_num_devices(self):
+        return self._num_devices
+
+
+class FakeTensor:
+    def __init__(self, shape, dtype):
+        self.shape = tuple(shape)
+        self.dtype = dtype
+
+    def volume(self):
+        result = 1
+        for dimension in self.shape:
+            result *= dimension
+        return result
+
+    def element_size(self):
+        return 1 if self.dtype in (ttnn.bfloat8_b, ttnn.bfloat4_b) else 2
+
+
+class FakeModel:
+    def __init__(self, dtypes=(ttnn.bfloat8_b, ttnn.bfloat8_b), *, bind_fails=False):
+        mesh = FakeMesh(2)
+        blocks = []
+        for dtype in dtypes:
+            attention = SimpleNamespace(
+                n_kv_heads=8,
+                head_dim=16,
+                kv_cache_dtype=dtype,
+                paged_attention_config=SimpleNamespace(block_size=32, max_num_blocks=8),
+            )
+            blocks.append(SimpleNamespace(attention_config=attention))
+        self.config = SimpleNamespace(
+            block_configs=blocks,
+            n_layers=len(blocks),
+            num_devices=2,
+            mesh_device=mesh,
+        )
+        self.set_calls = []
+        self.bound_cache = None
+        self.bind_fails = bind_fails
+
+    def set_kv_cache(self, cache):
+        self.set_calls.append(cache)
+        if cache is not None and self.bind_fails:
+            raise RuntimeError("bind failed")
+        self.bound_cache = cache
+
+
+def cache_config(**overrides):
+    values = {
+        "block_size": 32,
+        "max_num_blocks": 8,
+        "dtype": ttnn.bfloat8_b,
+    }
+    values.update(overrides)
+    return PagedKVCacheConfig(**values)
+
+
+@pytest.fixture
+def fake_allocator(monkeypatch):
+    allocated = []
+    deallocated = []
+
+    # Non-failure TTNN fakes retain overloaded backend keyword options for assertions.
+    def as_tensor(host_tensor, **kwargs):
+        tensor = FakeTensor(host_tensor.shape, kwargs["dtype"])
+        allocated.append((tensor, host_tensor, kwargs))
+        return tensor
+
+    monkeypatch.setattr(ttnn, "as_tensor", as_tensor)
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: ("replicate", mesh))
+    monkeypatch.setattr(ttnn, "deallocate", lambda tensor: deallocated.append(tensor))
+    return allocated, deallocated
+
+
+def test_unresolved_config_accepts_one_resolved_replacement(expect_error):
+    manager = PagedKVCacheManager(FakeModel(), cache_config())
+
+    resolved = replace(manager.config, num_blocks=4)
+    manager.configure(resolved)
+
+    assert manager.config is resolved
+    assert manager.config.num_blocks == 4
+    with expect_error(RuntimeError, "only once"):
+        manager.configure(replace(resolved, num_blocks=5))
+
+
+def test_resolved_replacement_accepts_model_aligned_external_geometry(expect_error):
+    model = FakeModel()
+    manager = PagedKVCacheManager(model, cache_config())
+
+    for block in model.config.block_configs:
+        block.attention_config.paged_attention_config.block_size = 16
+        block.attention_config.paged_attention_config.max_num_blocks = 12
+    manager.configure(cache_config(block_size=16, max_num_blocks=12, num_blocks=12))
+
+    assert manager.config.block_size == 16
+    assert manager.config.max_num_blocks == manager.config.num_blocks == 12
+
+    other = PagedKVCacheManager(FakeModel(), cache_config())
+    with expect_error(ValueError, "dtype"):
+        other.configure(cache_config(dtype=ttnn.bfloat16, num_blocks=4))
+
+
+def test_pre_resolved_config_starts_configured_and_cannot_be_replaced(expect_error):
+    manager = PagedKVCacheManager(FakeModel(), cache_config(num_blocks=4))
+
+    with expect_error(RuntimeError, "only once"):
+        manager.configure(cache_config(num_blocks=5))
+
+
+def test_model_paged_attention_policy_and_uniform_dtype_are_validated(expect_error):
+    model = FakeModel()
+    model.config.block_configs[1].attention_config.paged_attention_config.block_size = 16
+    with expect_error(ValueError, "block_size"):
+        PagedKVCacheManager(model, cache_config())
+
+    with expect_error(ValueError, "model-owned dtype"):
+        PagedKVCacheManager(FakeModel(), cache_config(dtype=ttnn.bfloat16))
+
+
+def test_vllm_torch_dtype_mapping_includes_quantized_surrogate(expect_error):
+    manager = PagedKVCacheManager(FakeModel(), cache_config())
+
+    assert torch_dtype_for_ttnn(ttnn.bfloat8_b) == torch.bfloat16
+    manager.validate_vllm_cache_spec(block_size=32, dtype=torch.bfloat16, num_blocks=8)
+    with expect_error(ValueError, "incompatible"):
+        manager.validate_vllm_cache_spec(block_size=32, dtype=torch.float32)
+    with expect_error(ValueError, "block_size"):
+        manager.validate_vllm_cache_spec(block_size=16, dtype=torch.bfloat16)
+    with expect_error(ValueError, "exceeds configured maximum"):
+        manager.validate_vllm_cache_spec(block_size=32, dtype=torch.bfloat16, num_blocks=9)
+
+
+def test_nonuniform_model_dtypes_remain_per_layer():
+    manager = PagedKVCacheManager(
+        FakeModel(dtypes=(ttnn.bfloat8_b, ttnn.bfloat16)),
+        cache_config(dtype=ttnn.bfloat8_b),
+    )
+
+    assert manager.per_layer_dtypes == (ttnn.bfloat8_b, ttnn.bfloat16)
+    manager.validate_vllm_cache_spec(block_size=32, dtype=torch.bfloat16)
+
+
+def test_allocate_derives_shapes_and_dtypes_binds_exact_borrowed_handle(fake_allocator, expect_error):
+    allocated, _ = fake_allocator
+    model = FakeModel(dtypes=(ttnn.bfloat8_b, ttnn.bfloat16))
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+
+    cache = manager.allocate()
+
+    assert cache is model.bound_cache
+    assert manager.cache_shapes == ((4, 4, 32, 16), (4, 4, 32, 16))
+    assert [entry[0].dtype for entry in allocated] == [
+        ttnn.bfloat8_b,
+        ttnn.bfloat8_b,
+        ttnn.bfloat16,
+        ttnn.bfloat16,
+    ]
+    assert all(tuple(entry[1].shape) == (4, 4, 32, 16) for entry in allocated)
+    assert len({id(entry[1]) for entry in allocated}) == 1
+    assert all(entry[2]["memory_config"] == ttnn.DRAM_MEMORY_CONFIG for entry in allocated)
+    assert all(entry[2]["cache_file_name"] is None for entry in allocated)
+    manager.validate_borrowed_handle(cache)
+    with expect_error(ValueError, "exact manager-owned"):
+        manager.validate_borrowed_handle([pair[:] for pair in cache])
+    assert manager.bound_context.config is manager.config
+    assert manager.bound_context.tensors[0][0] is cache[0][0]
+
+
+def test_allocate_reuses_legacy_cache_files_when_dtype_is_unambiguous(fake_allocator, tmp_path):
+    allocated, _ = fake_allocator
+    model = FakeModel()
+    model.model_args = SimpleNamespace(model_cache_path=tmp_path)
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+
+    manager.allocate()
+
+    shape = (4, 4, 32, 16)
+    expected = [
+        tmp_path / f"empty_kcache_paged_attention{shape}",
+        tmp_path / f"empty_vcache_paged_attention{shape}",
+    ]
+    assert [entry[2]["cache_file_name"] for entry in allocated] == expected * 2
+    assert len({id(entry[1]) for entry in allocated}) == 1
+
+
+def test_allocate_avoids_cache_file_collision_for_nonuniform_device_dtypes(fake_allocator, tmp_path):
+    allocated, _ = fake_allocator
+    model = FakeModel(dtypes=(ttnn.bfloat8_b, ttnn.bfloat16))
+    model.model_args = SimpleNamespace(model_cache_path=tmp_path)
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+
+    manager.allocate()
+
+    assert all(entry[2]["cache_file_name"] is None for entry in allocated)
+
+
+def test_allocation_requires_resolved_capacity_and_happens_once(fake_allocator, expect_error):
+    model = FakeModel()
+    manager = PagedKVCacheManager(model, cache_config())
+    with expect_error(RuntimeError, "must be resolved"):
+        manager.allocate()
+
+    manager.configure(replace(manager.config, num_blocks=4))
+    cache = manager.allocate()
+    with expect_error(RuntimeError, "already been allocated"):
+        manager.allocate()
+    assert cache is model.bound_cache
+
+
+def test_partial_allocation_failure_deallocates_created_tensors(monkeypatch, expect_error):
+    first = FakeTensor((4, 4, 32, 16), ttnn.bfloat8_b)
+    calls = 0
+    deallocated = []
+
+    def fail_second_allocation(
+        host_tensor,
+        *,
+        device,
+        mesh_mapper,
+        layout,
+        memory_config,
+        dtype,
+        cache_file_name,
+    ):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("allocation failed")
+        return first
+
+    monkeypatch.setattr(ttnn, "as_tensor", fail_second_allocation)
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: None)
+    monkeypatch.setattr(ttnn, "deallocate", lambda tensor: deallocated.append(tensor))
+    model = FakeModel()
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+
+    with expect_error(RuntimeError, "allocation failed"):
+        manager.allocate()
+
+    assert deallocated == [first]
+    assert model.set_calls == []
+
+
+def test_bind_failure_unbinds_then_deallocates_all_tensors(fake_allocator, expect_error):
+    allocated, deallocated = fake_allocator
+    model = FakeModel(bind_fails=True)
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+
+    with expect_error(RuntimeError, "bind failed"):
+        manager.allocate()
+
+    assert model.set_calls[-1] is None
+    assert deallocated == [entry[0] for entry in reversed(allocated)]
+    assert model.bound_cache is None
+
+
+def test_release_unbinds_before_deallocating_and_is_idempotent(monkeypatch, expect_error):
+    operations = []
+    model = FakeModel()
+    original_set = model.set_kv_cache
+
+    def set_kv_cache(cache):
+        operations.append(("bind", cache))
+        original_set(cache)
+
+    model.set_kv_cache = set_kv_cache
+    monkeypatch.setattr(
+        ttnn,
+        "as_tensor",
+        lambda host_tensor, **kwargs: FakeTensor(host_tensor.shape, kwargs["dtype"]),
+    )
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: None)
+    monkeypatch.setattr(ttnn, "deallocate", lambda tensor: operations.append(("deallocate", tensor)))
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+    cache = manager.allocate()
+    operations.clear()
+
+    manager.release()
+
+    assert operations[0] == ("bind", None)
+    assert [operation[1] for operation in operations[1:]] == [tensor for pair in cache for tensor in pair]
+    assert manager.bound_context is None
+
+    manager.release()
+    assert len(operations) == 5
+    with expect_error(RuntimeError, "terminal"):
+        manager.allocate()
+
+
+def test_borrowed_handle_mutation_cannot_redirect_owned_tensor_release(fake_allocator, expect_error):
+    allocated, deallocated = fake_allocator
+    model = FakeModel()
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+    cache = manager.allocate()
+    owned_tensors = [entry[0] for entry in allocated]
+    replacement = FakeTensor(cache[0][0].shape, cache[0][0].dtype)
+
+    cache[0][0] = replacement
+
+    with expect_error(ValueError, "exact manager-owned K/V tensors"):
+        manager.validate_borrowed_handle(cache)
+
+    manager.release()
+
+    assert deallocated == owned_tensors
+    assert replacement not in deallocated
+
+
+def test_release_failure_retains_only_failed_tensor_and_retries(monkeypatch, expect_error):
+    tensors = []
+
+    def as_tensor(host_tensor, **kwargs):
+        tensor = FakeTensor(host_tensor.shape, kwargs["dtype"])
+        tensors.append(tensor)
+        return tensor
+
+    monkeypatch.setattr(ttnn, "as_tensor", as_tensor)
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: None)
+    model = FakeModel()
+    manager = PagedKVCacheManager(model, cache_config(num_blocks=4))
+    manager.allocate()
+    failed_tensor = tensors[0]
+    attempts = []
+
+    def fail_once(tensor):
+        attempts.append(tensor)
+        if tensor is failed_tensor and attempts.count(tensor) == 1:
+            raise RuntimeError("deallocate failed")
+
+    monkeypatch.setattr(ttnn, "deallocate", fail_once)
+
+    with expect_error(RuntimeError, "Failed to deallocate 1"):
+        manager.release()
+
+    assert manager.bound_context is None
+    assert model.bound_cache is None
+    assert attempts == tensors
+
+    manager.release()
+
+    assert attempts.count(failed_tensor) == 2
+    assert all(attempts.count(tensor) == 1 for tensor in tensors[1:])
+
+
+def test_partial_allocation_cleanup_failure_preserves_tensor_for_release_retry(monkeypatch, expect_error):
+    tensor = FakeTensor((4, 4, 32, 16), ttnn.bfloat8_b)
+    allocation_calls = 0
+    deallocation_calls = 0
+
+    def fail_second_allocation(
+        host_tensor,
+        *,
+        device,
+        mesh_mapper,
+        layout,
+        memory_config,
+        dtype,
+        cache_file_name,
+    ):
+        nonlocal allocation_calls
+        allocation_calls += 1
+        if allocation_calls == 2:
+            raise RuntimeError("allocation failed")
+        return tensor
+
+    def fail_first_deallocation(value):
+        nonlocal deallocation_calls
+        deallocation_calls += 1
+        if deallocation_calls == 1:
+            raise RuntimeError("cleanup failed")
+
+    monkeypatch.setattr(ttnn, "as_tensor", fail_second_allocation)
+    monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: None)
+    monkeypatch.setattr(ttnn, "deallocate", fail_first_deallocation)
+    manager = PagedKVCacheManager(FakeModel(), cache_config(num_blocks=4))
+
+    with expect_error(RuntimeError, "allocation failed") as exc_info:
+        manager.allocate()
+
+    assert [str(error) for error in exc_info.value.cleanup_failures] == ["cleanup failed"]
+    manager.release()
+
+    assert deallocation_calls == 2
+    with expect_error(RuntimeError, "terminal"):
+        manager.allocate()
+
+
+def test_release_before_allocation_is_terminal_and_idempotent(expect_error):
+    manager = PagedKVCacheManager(FakeModel(), cache_config())
+
+    manager.release()
+    manager.release()
+
+    with expect_error(RuntimeError, "terminal"):
+        manager.allocate()

code/models/common/tests/llm_runtime/test_prefill_inputs.py

@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import models.common.llm_runtime.prefill.inputs as inputs_module
+from models.common.llm_runtime.prefill.inputs import (
+    PrefillDeviceInputs,
+    PrefillHostInputs,
+    PrefillInputStager,
+    allocate_device_tensors,
+    copy_into_device_tensors,
+)
+
+
+class _Model:
+    def __init__(self, *, rotary_capacity=8, rotary_outputs=("cos", "sin")):
+        self.config = SimpleNamespace(dim=64)
+        self.rope_setup = SimpleNamespace(
+            cos_matrix=torch.zeros(1, 1, rotary_capacity),
+            load_device_weights=lambda: None,
+        )
+        self.rotary_outputs = rotary_outputs
+
+    def prepare_prefill_rot_mats(self, position_indices):
+        del position_indices
+        return self.rotary_outputs
+
+
+def _stager(*, model=None, released=None):
+    released = [] if released is None else released
+    return PrefillInputStager(
+        model=_Model() if model is None else model,
+        mesh_device="mesh",
+        release_transient=lambda values: released.append(values) or [],
+    )
+
+
+def _patch_host_conversion(monkeypatch):
+    converted = []
+    monkeypatch.setattr(inputs_module.ttnn, "ReplicateTensorToMesh", lambda mesh: ("mapper", mesh))
+    monkeypatch.setattr(
+        inputs_module.ttnn,
+        "from_torch",
+        lambda value, **kwargs: converted.append((value.clone(), kwargs)) or value.clone(),
+    )
+    return converted
+
+
+@pytest.mark.parametrize("shape", [(8,), (1, 2, 8)])
+def test_prepare_host_inputs_rejects_non_matrix_tokens_before_conversion(monkeypatch, expect_error, shape):
+    monkeypatch.setattr(inputs_module.ttnn, "from_torch", lambda *args, **kwargs: pytest.fail("converted"))
+
+    with expect_error(ValueError, "rank 2"):
+        _stager().prepare_host_inputs(torch.zeros(shape), torch.zeros(1, 1, dtype=torch.int32))
+
+
+def test_prepare_host_inputs_rejects_negative_start_and_last_token_beyond_rotary_capacity(
+    monkeypatch,
+    expect_error,
+):
+    _patch_host_conversion(monkeypatch)
+    stager = _stager(model=_Model(rotary_capacity=8))
+    tokens = torch.zeros(1, 4, dtype=torch.long)
+    page_table = torch.zeros(1, 1, dtype=torch.int32)
+
+    with expect_error(ValueError, "start position must be nonnegative"):
+        stager.prepare_host_inputs(tokens, page_table, start_pos=-1)
+    with expect_error(ValueError, "exceeds rotary capacity 8"):
+        stager.prepare_host_inputs(tokens, page_table, last_token_idx=8)
+
+
+def test_prepare_host_inputs_clamps_padded_positions_to_last_rotary_entry(monkeypatch):
+    converted = _patch_host_conversion(monkeypatch)
+
+    _stager(model=_Model(rotary_capacity=8)).prepare_host_inputs(
+        torch.zeros(1, 4, dtype=torch.long),
+        torch.zeros(1, 1, dtype=torch.int32),
+        start_pos=6,
+    )
+
+    position_indices = converted[1][0]
+    assert position_indices.tolist() == [[6, 7, 7, 7]]
+
+
+@pytest.mark.parametrize("relative_last,sequence_length", [(-1, 32), (32, 32), (0, 0)])
+def test_prepare_position_inputs_rejects_positions_outside_padded_sequence(
+    monkeypatch,
+    expect_error,
+    relative_last,
+    sequence_length,
+):
+    monkeypatch.setattr(inputs_module.ttnn, "from_torch", lambda *args, **kwargs: pytest.fail("converted"))
+
+    with expect_error(ValueError, "last-token position"):
+        _stager().prepare_position_inputs_host(relative_last, sequence_length)
+
+
+def test_allocate_device_tensors_releases_partial_allocation_on_failure(monkeypatch, expect_error):
+    first_device = object()
+    calls = []
+
+    def to_device(host_tensor, *, device):
+        calls.append((host_tensor, device))
+        if len(calls) == 2:
+            raise RuntimeError("allocation failed")
+        return first_device
+
+    released = []
+    monkeypatch.setattr(inputs_module.ttnn, "to_device", to_device)
+    monkeypatch.setattr(
+        inputs_module,
+        "best_effort_deallocate_owned_tensors",
+        lambda values: released.append(tuple(values)) or [],
+    )
+
+    with expect_error(RuntimeError, "allocation failed"):
+        allocate_device_tensors(("host-0", "host-1"), mesh_device="mesh")
+
+    assert released == [(first_device,)]
+
+
+def test_stage_device_inputs_releases_raw_and_malformed_rotary_outputs(monkeypatch, expect_error):
+    raw = ["tokens", "positions", "page", None, None]
+    released = []
+    model = _Model(rotary_outputs=("cos-only",))
+    monkeypatch.setattr(inputs_module, "allocate_device_tensors", lambda values, *, mesh_device: raw)
+
+    host = PrefillHostInputs("host-tokens", "host-positions", "host-page", None, None)
+    with expect_error(ValueError, "cosine and sine"):
+        _stager(model=model, released=released).stage_device_inputs(host)
+
+    assert released == [(model.rotary_outputs, raw)]
+
+
+def test_copy_rotary_inputs_rejects_malformed_output_count_and_releases_it(monkeypatch, expect_error):
+    released = []
+    model = _Model(rotary_outputs=("cos-only",))
+    device = PrefillDeviceInputs("tokens", "cos", "sin", "page", None, "positions", None)
+    monkeypatch.setattr(inputs_module.ttnn, "copy", lambda **kwargs: pytest.fail("copied"))
+
+    with expect_error(ValueError, "cosine and sine"):
+        _stager(model=model, released=released).copy_rotary_inputs(device)
+
+    assert released == [model.rotary_outputs]
+
+
+@pytest.mark.parametrize(
+    ("host", "device"),
+    [
+        ((None,), ("device",)),
+        (("host",), (None,)),
+        (("host", None), ("device", "unexpected-device")),
+        (("host", None), ("device",)),
+    ],
+)
+def test_copy_into_device_tensors_rejects_structure_changes_before_copy(
+    monkeypatch,
+    expect_error,
+    host,
+    device,
+):
+    monkeypatch.setattr(
+        inputs_module.ttnn,
+        "copy_host_to_device_tensor",
+        lambda *args: pytest.fail("copied"),
+    )
+
+    with expect_error(ValueError, "host/device"):
+        copy_into_device_tensors(host, device)

code/models/common/tests/llm_runtime/test_program_compiler.py

@@ -0,0 +1,268 @@
+# 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)

code/models/common/tests/llm_runtime/test_tensor_resources.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+
+from models.common.llm_runtime import tensor_resources
+from models.common.llm_runtime.decode import DecodeDeviceInputs, DecodePersistentInputs
+from models.common.llm_runtime.prefill.inputs import PrefillDeviceInputs, PrefillPositionInputs
+from models.common.llm_runtime.prefill.trace import PrefillHiddenPersistentInputs, PrefillReplayState
+
+
+def test_owned_runtime_containers_release_aliased_tensors_once(monkeypatch):
+    class FakeTensor:
+        pass
+
+    shared = FakeTensor()
+    other = FakeTensor()
+    released = []
+    monkeypatch.setattr(tensor_resources.ttnn, "Tensor", FakeTensor)
+    monkeypatch.setattr(tensor_resources.ttnn, "deallocate", released.append)
+
+    decode_inputs = DecodeDeviceInputs(shared, other, shared, None)
+    prefill_inputs = PrefillDeviceInputs(shared, other, None, shared, None, other, None)
+    positions = PrefillPositionInputs(other, shared, other)
+    values = (
+        decode_inputs,
+        DecodePersistentInputs(decode_inputs, (shared, other, shared)),
+        prefill_inputs,
+        positions,
+        PrefillHiddenPersistentInputs(prefill_inputs),
+        PrefillReplayState(positions, (shared, other, shared), shared),
+    )
+
+    assert tensor_resources.best_effort_deallocate_owned_tensors(values) == []
+    assert released == [shared, other]
+
+
+def test_arbitrary_dataclass_is_not_treated_as_an_ownership_projection(monkeypatch):
+    class FakeTensor:
+        pass
+
+    @dataclass
+    class BorrowedValue:
+        tensor: FakeTensor
+
+    released = []
+    monkeypatch.setattr(tensor_resources.ttnn, "Tensor", FakeTensor)
+    monkeypatch.setattr(tensor_resources.ttnn, "deallocate", released.append)
+
+    assert tensor_resources.best_effort_deallocate_owned_tensors(BorrowedValue(FakeTensor())) == []
+    assert released == []
+
+
+def test_raise_cleanup_failures_preserves_primary_and_attaches_rest(expect_error):
+    primary = RuntimeError("primary")
+    secondary = RuntimeError("secondary")
+
+    with expect_error(RuntimeError, "primary") as raised:
+        tensor_resources.raise_cleanup_failures((primary, secondary))
+
+    assert raised.value is primary
+    assert primary.cleanup_failures == (secondary,)

code/models/common/tests/llm_runtime/test_trace_compiler.py

@@ -0,0 +1,576 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+from itertools import permutations
+
+import pytest
+import torch
+
+import models.common.llm_runtime.trace_compiler as trace_compiler_module
+import ttnn
+from models.common.llm_runtime.decode import DecodeDeviceInputs, DecodePersistentInputs
+from models.common.llm_runtime.prefill.inputs import PrefillDeviceInputs, PrefillPositionInputs
+from models.common.llm_runtime.prefill.trace import PrefillHiddenPersistentInputs, PrefillReplayState
+from models.common.llm_runtime.program_compiler import ProgramCompiler
+from models.common.llm_runtime.trace_compiler import (
+    InputRefreshPolicy,
+    PersistentInputs,
+    TraceCapturePlan,
+    TraceCompiler,
+)
+
+
+@dataclass(frozen=True)
+class _Signature:
+    kind: str
+    variant: int
+
+    @property
+    def key_material(self):
+        return (("kind", self.kind), ("variant", self.variant))
+
+
+def _patch_backend(monkeypatch, events):
+    next_trace_id = iter(range(100, 200))
+    monkeypatch.setattr(ttnn, "synchronize_device", lambda mesh: events.append(("sync", mesh)))
+    monkeypatch.setattr(
+        ttnn,
+        "begin_trace_capture",
+        lambda mesh, cq_id: events.append(("begin", mesh, cq_id)) or next(next_trace_id),
+    )
+    monkeypatch.setattr(
+        ttnn,
+        "end_trace_capture",
+        lambda mesh, trace_id, cq_id: events.append(("end", trace_id, cq_id)),
+    )
+    monkeypatch.setattr(
+        ttnn,
+        "execute_trace",
+        lambda mesh, trace_id, cq_id, blocking: events.append(("execute", trace_id, cq_id, blocking)),
+    )
+    monkeypatch.setattr(ttnn, "release_trace", lambda mesh, trace_id: events.append(("release", trace_id)))
+    monkeypatch.setattr(trace_compiler_module, "_trim_host_allocator", lambda: events.append(("trim",)))
+
+
+def _compiled_program(program_compiler, monkeypatch, variant):
+    monkeypatch.setattr(ttnn, "synchronize_device", lambda mesh: None)
+    return program_compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1))
+
+
+def _plan(program, variant, events, *, operation="decode", policy=InputRefreshPolicy()):
+    return TraceCapturePlan(
+        program_key=program.key,
+        trace_signature=_Signature("trace", variant),
+        operation=operation,
+        prepare_inputs=lambda: events.append(("prepare", variant)) or (),
+        capture=lambda persistent: events.append(("capture", variant)) or torch.zeros(1),
+        refresh_policy=policy,
+    )
+
+
+def test_trace_compiler_retains_exact_program_compiler_and_separate_registries(monkeypatch):
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = _compiled_program(compiler, monkeypatch, 1)
+    trace = TraceCompiler(compiler)
+    trace_key = trace.register_capture_plan(_plan(program, 1, []))
+
+    assert trace.program_compiler is compiler
+    assert trace.mesh_device is compiler.mesh_device
+    assert trace.trace_key_for_program(program.key) == trace_key
+    assert trace.get(trace_key) is not None
+    assert compiler.require_compiled(program.key) is program
+    assert not hasattr(program, "artifact")
+    assert not hasattr(trace, "programs")
+
+
+def test_trace_aliases_share_one_artifact_without_copying_program_records(monkeypatch):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    first = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    second = compiler.compile(_Signature("program", 2), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+    shared_signature = _Signature("trace", 1)
+    first_key = trace.register_capture_plan(
+        TraceCapturePlan(first.key, shared_signature, "decode", lambda: (), lambda persistent: torch.zeros(1))
+    )
+    second_key = trace.register_capture_plan(
+        TraceCapturePlan(second.key, shared_signature, "decode", lambda: (), lambda persistent: torch.zeros(1))
+    )
+
+    trace.capture_all()
+
+    assert first_key == second_key
+    assert trace.get(first_key) is not None
+    assert trace.trace_key_for_program(first.key) == first_key
+    assert trace.trace_key_for_program(second.key) == first_key
+    assert compiler.require_compiled(first.key) is first
+    assert compiler.require_compiled(second.key) is second
+    assert [event[0] for event in events].count("begin") == 1
+    assert trace.trace_count == 1
+    assert trace.trace_association_count == 2
+    assert trace.registered_coverage("decode") == ((first_key, shared_signature),)
+    assert trace.registered_coverage("prefill") == ()
+
+
+@pytest.mark.parametrize("order", tuple(permutations(("logits", "argmax", "topk"))))
+def test_hidden_trace_alias_workspaces_are_registration_order_independent(monkeypatch, order):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = {
+        path: compiler.compile(_Signature("program", index), lambda context: torch.zeros(1))
+        for index, path in enumerate(("logits", "argmax", "topk"), 1)
+    }
+    trace = TraceCompiler(compiler)
+    signature = _Signature("shared-hidden", 1)
+
+    for path in order:
+        trace.register_capture_plan(
+            TraceCapturePlan(
+                programs[path].key,
+                signature,
+                "prefill",
+                lambda: "hidden-inputs",
+                lambda persistent: "hidden-output",
+                schema_fingerprint=("hidden-v2",),
+                prepare_workspace=lambda path=path: f"{path}-workspace",
+                workspace_fingerprint=("postprocess-v1", path),
+            )
+        )
+
+    trace.capture_all()
+
+    assert [event[0] for event in events].count("begin") == 1
+    assert {path: trace.workspace_for_program(program.key) for path, program in programs.items()} == {
+        "logits": "logits-workspace",
+        "argmax": "argmax-workspace",
+        "topk": "topk-workspace",
+    }
+
+
+def test_trace_alias_rejects_a_mismatched_persistent_schema(monkeypatch, expect_error):
+    compiler = ProgramCompiler("mesh", lambda: object())
+    first = _compiled_program(compiler, monkeypatch, 1)
+    second = _compiled_program(compiler, monkeypatch, 2)
+    trace = TraceCompiler(compiler)
+    signature = _Signature("trace", 1)
+    trace.register_capture_plan(
+        TraceCapturePlan(first.key, signature, "prefill", lambda: (), lambda _: (), schema_fingerprint=("v1",))
+    )
+
+    with expect_error(ValueError, "different schema fingerprint"):
+        trace.register_capture_plan(
+            TraceCapturePlan(second.key, signature, "prefill", lambda: (), lambda _: (), schema_fingerprint=("v2",))
+        )
+
+
+def test_capture_allocates_every_input_before_capture_and_coordinates_gates(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    trace = TraceCompiler(compiler)
+
+    def capture_with_gate(persistent):
+        events.append(("capture", 1))
+        with expect_error(RuntimeError, "capture is in progress"):
+            compiler.compile(_Signature("program", 3), lambda context: torch.zeros(1))
+        return torch.zeros(1)
+
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            programs[0].key,
+            _Signature("trace", 1),
+            "decode",
+            lambda: events.append(("prepare", 1)) or (),
+            capture_with_gate,
+        )
+    )
+    trace.register_capture_plan(_plan(programs[1], 2, events))
+    events.clear()
+
+    trace.capture_all()
+
+    first_begin = next(index for index, event in enumerate(events) if event[0] == "begin")
+    assert events[:first_begin] == [("prepare", 1), ("prepare", 2)]
+    assert trace.trace_active and compiler.trace_active
+    assert not compiler.trace_capture_in_progress
+    with expect_error(RuntimeError, "after trace activation"):
+        compiler.compile(_Signature("program", 3), lambda context: torch.zeros(1))
+
+
+def test_opt_in_capture_prime_runs_after_allocations_and_before_capture_gate(monkeypatch):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+
+    def prime(persistent):
+        assert not compiler.trace_capture_in_progress
+        events.append(("prime", persistent))
+        return "prime-output"
+
+    def release_prime_output(output):
+        events.append(("release-prime", output))
+        return []
+
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            program.key,
+            _Signature("trace", 1),
+            "prefill",
+            lambda: events.append(("prepare", 1)) or "persistent",
+            lambda persistent: events.append(("capture", persistent)) or torch.zeros(1),
+            prepare_workspace=lambda: events.append(("workspace", 1)) or "workspace",
+            workspace_fingerprint=("workspace",),
+            prime=prime,
+            release_prime_output=release_prime_output,
+        )
+    )
+    events.clear()
+
+    trace.capture_all()
+
+    first_begin = next(index for index, event in enumerate(events) if event[0] == "begin")
+    assert events[:first_begin] == [
+        ("prepare", 1),
+        ("workspace", 1),
+        ("prime", PersistentInputs("persistent")),
+        ("sync", "mesh"),
+        ("release-prime", "prime-output"),
+        ("sync", "mesh"),
+    ]
+    assert ("capture", PersistentInputs("persistent")) in events[first_begin:]
+
+
+def test_capture_prime_failure_rolls_back_without_beginning_trace(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+
+    class OwnedTensor:
+        pass
+
+    released = []
+    monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
+    monkeypatch.setattr(ttnn, "deallocate", released.append)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+    persistent = OwnedTensor()
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            program.key,
+            _Signature("trace", 1),
+            "prefill",
+            lambda: persistent,
+            lambda _: pytest.fail("capture must not begin after prime failure"),
+            prime=lambda _: (_ for _ in ()).throw(RuntimeError("prime failed")),
+            release_prime_output=lambda _: [],
+        )
+    )
+
+    with expect_error(RuntimeError, "prime failed"):
+        trace.capture_all()
+
+    assert not any(event[0] == "begin" for event in events)
+    assert released == [persistent]
+    assert not trace.trace_active and not compiler.trace_active
+
+
+def test_capture_prime_release_failure_still_synchronizes_before_rollback(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+    release_error = RuntimeError("prime release failed")
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            program.key,
+            _Signature("trace", 1),
+            "prefill",
+            lambda: (),
+            lambda _: pytest.fail("capture must not begin after prime release failure"),
+            prime=lambda _: events.append(("prime",)) or "prime-output",
+            release_prime_output=lambda output: events.append(("release-prime", output)) or [release_error],
+        )
+    )
+    events.clear()
+
+    with expect_error(RuntimeError, "prime release failed") as caught:
+        trace.capture_all()
+
+    assert caught.value is release_error
+    assert events[:4] == [
+        ("prime",),
+        ("sync", "mesh"),
+        ("release-prime", "prime-output"),
+        ("sync", "mesh"),
+    ]
+    assert not any(event[0] == "begin" for event in events)
+
+
+def test_capture_orders_decode_before_prefill_after_allocating_every_input(monkeypatch):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    trace = TraceCompiler(compiler)
+    trace.register_capture_plan(_plan(programs[0], 1, events, operation="prefill"))
+    trace.register_capture_plan(_plan(programs[1], 2, events, operation="decode"))
+    events.clear()
+
+    trace.capture_all()
+
+    first_begin = next(index for index, event in enumerate(events) if event[0] == "begin")
+    assert events[:first_begin] == [("prepare", 1), ("prepare", 2)]
+    assert [event for event in events if event[0] == "capture"] == [("capture", 2), ("capture", 1)]
+
+
+def test_capture_failure_rolls_back_traces_and_uncaptured_inputs(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+
+    class OwnedTensor:
+        pass
+
+    released = []
+    monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
+    monkeypatch.setattr(ttnn, "deallocate", released.append)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    trace = TraceCompiler(compiler)
+    first_input, first_output, second_input = OwnedTensor(), OwnedTensor(), OwnedTensor()
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            programs[0].key,
+            _Signature("trace", 1),
+            "decode",
+            lambda: first_input,
+            lambda persistent: first_output,
+        )
+    )
+    primary = RuntimeError("second capture failed")
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            programs[1].key,
+            _Signature("trace", 2),
+            "decode",
+            lambda: second_input,
+            lambda persistent: (_ for _ in ()).throw(primary),
+        )
+    )
+
+    with expect_error(RuntimeError, "second capture failed") as caught:
+        trace.capture_all()
+
+    assert caught.value is primary
+    assert released.count(first_input) == 1
+    assert released.count(first_output) == 1
+    assert released.count(second_input) == 1
+    assert not trace.trace_active and not compiler.trace_active
+    for program in programs:
+        trace_key = trace.trace_key_for_program(program.key)
+        assert trace_key is not None
+        assert trace.get(trace_key).artifact is None
+
+
+def test_incomplete_capture_rollback_keeps_program_gate_closed_until_cleanup(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+
+    class OwnedTensor:
+        pass
+
+    retry = OwnedTensor()
+    attempts = []
+
+    def deallocate(value):
+        attempts.append(value)
+        if value is retry and attempts.count(retry) == 1:
+            raise RuntimeError("release once")
+
+    monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
+    monkeypatch.setattr(ttnn, "deallocate", deallocate)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    trace = TraceCompiler(compiler)
+    trace.register_capture_plan(
+        TraceCapturePlan(programs[0].key, _Signature("trace", 1), "decode", lambda: retry, lambda _: torch.zeros(1))
+    )
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            programs[1].key,
+            _Signature("trace", 2),
+            "decode",
+            lambda: (_ for _ in ()).throw(RuntimeError("prepare failed")),
+            lambda _: torch.zeros(1),
+        )
+    )
+
+    with expect_error(RuntimeError, "prepare failed"):
+        trace.capture_all()
+
+    assert trace.trace_active and compiler.trace_active
+    with expect_error(RuntimeError, "after trace activation"):
+        compiler.compile(_Signature("program", 3), lambda context: torch.zeros(1))
+
+    trace.cleanup()
+    assert attempts.count(retry) == 2
+    assert not compiler.trace_active
+
+
+def test_replay_refresh_decisions_cover_first_replay_page_change_feedback_and_switch(monkeypatch):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    policy = InputRefreshPolicy(every_replay=("position", "sampling"))
+    trace = TraceCompiler(compiler)
+    for variant, program in enumerate(programs, 1):
+        trace.register_capture_plan(_plan(program, variant, events, policy=policy))
+    trace.capture_all()
+    decisions = []
+
+    trace.replay(
+        programs[0].key,
+        lambda artifact, decision: decisions.append(decision),
+        device_feedback_enabled=True,
+        feedback_compatible=True,
+    )
+    trace.replay(
+        programs[0].key,
+        lambda artifact, decision: decisions.append(decision),
+        device_feedback_enabled=True,
+        feedback_compatible=True,
+        page_table_changed=True,
+    )
+    trace.replay(
+        programs[1].key,
+        lambda artifact, decision: decisions.append(decision),
+        device_feedback_enabled=True,
+        feedback_compatible=True,
+    )
+    trace.replay(
+        programs[1].key,
+        lambda artifact, decision: decisions.append(decision),
+        device_feedback_enabled=False,
+    )
+
+    assert [decision.full for decision in decisions] == [True, False, True, True]
+    assert [decision.page_table for decision in decisions] == [False, True, False, False]
+    assert all(decision.fields == ("position", "sampling") for decision in decisions)
+    assert [event[0] for event in events].count("execute") == 4
+    assert trace.replay_count == 4
+    assert trace.replay_counts == {"prefill": 0, "decode": 4}
+
+
+def test_replay_counters_increment_only_after_successful_submission(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+    trace.register_capture_plan(_plan(program, 1, events, operation="prefill"))
+    trace.capture_all()
+
+    monkeypatch.setattr(ttnn, "execute_trace", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("submit")))
+    with expect_error(RuntimeError, "submit"):
+        trace.replay(program.key, lambda artifact, decision: None)
+    assert trace.replay_count == 0
+    assert trace.replay_counts == {"prefill": 0, "decode": 0}
+
+
+def test_cleanup_retries_trace_release_before_deallocating_and_does_not_own_programs(monkeypatch, expect_error):
+    events = []
+    _patch_backend(monkeypatch, events)
+
+    class OwnedTensor:
+        pass
+
+    persistent, output = OwnedTensor(), OwnedTensor()
+    deallocated = []
+    release_attempts = []
+
+    def release(mesh, trace_id):
+        release_attempts.append(trace_id)
+        if len(release_attempts) == 1:
+            raise RuntimeError("trace release once")
+
+    monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
+    monkeypatch.setattr(ttnn, "deallocate", deallocated.append)
+    monkeypatch.setattr(ttnn, "release_trace", release)
+    compiler = ProgramCompiler("mesh", lambda: object())
+    program = compiler.compile(_Signature("program", 1), lambda context: torch.zeros(1))
+    trace = TraceCompiler(compiler)
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            program.key,
+            _Signature("trace", 1),
+            "decode",
+            lambda: persistent,
+            lambda values: output,
+        )
+    )
+    trace.capture_all()
+
+    with expect_error(RuntimeError, "Failed to release"):
+        trace.cleanup()
+    assert deallocated == []
+
+    trace.cleanup()
+    trace.cleanup()
+    assert release_attempts == [100, 100]
+    assert deallocated.count(persistent) == 1
+    assert deallocated.count(output) == 1
+    compiler.cleanup()
+
+
+def test_cleanup_releases_operation_owned_persistent_dataclasses_once(monkeypatch):
+    events = []
+    _patch_backend(monkeypatch, events)
+
+    class OwnedTensor:
+        pass
+
+    values = [OwnedTensor() for _ in range(20)]
+    decode = DecodePersistentInputs(
+        DecodeDeviceInputs(*values[:4]),
+        tuple(values[4:7]),
+    )
+    prefill = PrefillHiddenPersistentInputs(PrefillDeviceInputs(*values[7:14]))
+    prefill_workspace = PrefillReplayState(
+        PrefillPositionInputs(*values[14:17]),
+        (values[17], values[17], values[17]),
+        values[18],
+    )
+    deallocated = []
+    monkeypatch.setattr(ttnn, "Tensor", OwnedTensor)
+    monkeypatch.setattr(ttnn, "deallocate", deallocated.append)
+
+    compiler = ProgramCompiler("mesh", lambda: object())
+    programs = [compiler.compile(_Signature("program", variant), lambda context: torch.zeros(1)) for variant in (1, 2)]
+    trace = TraceCompiler(compiler)
+    trace.register_capture_plan(
+        TraceCapturePlan(programs[0].key, _Signature("trace", 1), "decode", lambda: decode, lambda _: values[0])
+    )
+    trace.register_capture_plan(
+        TraceCapturePlan(
+            programs[1].key,
+            _Signature("trace", 2),
+            "prefill",
+            lambda: prefill,
+            lambda _: values[19],
+            prepare_workspace=lambda: prefill_workspace,
+            workspace_fingerprint=("prefill-workspace",),
+        )
+    )
+
+    trace.capture_all()
+    trace.cleanup()
+
+    assert len(deallocated) == len(values)
+    assert {id(value) for value in deallocated} == {id(value) for value in values}

code/models/common/tests/llm_runtime/test_vllm_adapter.py

@@ -0,0 +1,1025 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import inspect
+from types import SimpleNamespace
+from unittest.mock import create_autospec
+
+import pytest
+import torch
+
+import ttnn
+from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig
+from models.common.llm_runtime.vllm_adapter import (
+    NormalizedDecodeKwargs,
+    NormalizedPrefillKwargs,
+    VLLMAdapter,
+    VLLMAdapterConfig,
+)
+from models.common.models.llama3_8b import generator as llama3_generator_module
+from models.common.models.llama3_8b.executor import Llama3Executor
+from models.common.models.llama3_8b.generator import Llama3Generator
+from models.common.models.llama33_70b import generator as llama70_generator_module
+from models.common.models.qwen2_7b.generator import Qwen2Generator
+from models.common.models.qwen3_32b import generator as qwen3_generator_module
+from models.common.models.qwen25_7b.generator import Qwen25Generator
+from models.common.models.qwen25_72b.generator import Qwen25_72BGenerator
+from models.common.models.qwen25_coder_32b.generator import Qwen25Coder32BGenerator
+
+
+def _adapter(
+    *,
+    trace=None,
+    paged_config=None,
+    model_dtype=ttnn.bfloat8_b,
+    request_state_fields=("prompt_tokens", "output_tokens", "slot_remap"),
+):
+    return VLLMAdapter(
+        VLLMAdapterConfig.resolve(
+            trace=trace or TraceConfig(mode="all"),
+            paged_kv_cache=paged_config or PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b),
+            expected_num_layers=32,
+            expected_kv_heads_per_device=8,
+            expected_head_dim=128,
+            model_kv_cache_dtype=model_dtype,
+            request_state_fields=request_state_fields,
+        )
+    )
+
+
+def test_config_resolves_canonical_static_policy_and_is_frozen(expect_error):
+    trace = TraceConfig(mode="all")
+    paged_kv_cache = PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b)
+
+    config = VLLMAdapterConfig.resolve(
+        trace=trace,
+        paged_kv_cache=paged_kv_cache,
+        expected_num_layers=32.0,
+        expected_kv_heads_per_device=8.0,
+        expected_head_dim=128.0,
+        model_kv_cache_dtype=[ttnn.bfloat8_b] * 32,
+    )
+
+    assert config.trace is trace
+    assert config.paged_kv_cache is paged_kv_cache
+    assert config.expected_num_layers == 32
+    assert isinstance(config.expected_num_layers, int)
+    assert config.expected_kv_heads_per_device == 8
+    assert isinstance(config.expected_kv_heads_per_device, int)
+    assert config.expected_head_dim == 128
+    assert isinstance(config.expected_head_dim, int)
+    assert config.model_kv_cache_dtypes == (ttnn.bfloat8_b,) * 32
+    assert config.request_state_fields == ()
+    with expect_error(dataclasses.FrozenInstanceError, "cannot assign to field"):
+        config.expected_num_layers = 1
+    with expect_error(ValueError, "expected_num_layers"):
+        VLLMAdapterConfig(
+            trace=trace,
+            paged_kv_cache=paged_kv_cache,
+            expected_num_layers=0,
+            expected_kv_heads_per_device=8,
+            expected_head_dim=128,
+            model_kv_cache_dtypes=(ttnn.bfloat8_b,),
+        )
+    with expect_error(TypeError, "must be a tuple"):
+        VLLMAdapterConfig(
+            trace=trace,
+            paged_kv_cache=paged_kv_cache,
+            expected_num_layers=32,
+            expected_kv_heads_per_device=8,
+            expected_head_dim=128,
+            model_kv_cache_dtypes=[ttnn.bfloat8_b],
+        )
+
+
+@pytest.mark.parametrize(
+    ("overrides", "error_type", "message"),
+    [
+        ({"trace": object()}, TypeError, "TraceConfig"),
+        ({"paged_kv_cache": object()}, TypeError, "PagedKVCacheConfig"),
+        ({"expected_num_layers": 0}, ValueError, "positive integer"),
+        ({"expected_num_layers": True}, ValueError, "positive integer"),
+        ({"expected_kv_heads_per_device": 0}, ValueError, "positive integer"),
+        ({"expected_kv_heads_per_device": -1}, ValueError, "positive integer"),
+        ({"expected_kv_heads_per_device": True}, ValueError, "positive integer"),
+        ({"expected_kv_heads_per_device": 8.5}, ValueError, "positive integer"),
+        ({"expected_head_dim": -1}, ValueError, "positive integer"),
+        ({"expected_head_dim": 0}, ValueError, "positive integer"),
+        ({"expected_head_dim": True}, ValueError, "positive integer"),
+        ({"expected_head_dim": 128.5}, ValueError, "positive integer"),
+        ({"model_kv_cache_dtype": ()}, ValueError, "cannot be empty"),
+        ({"model_kv_cache_dtype": (ttnn.bfloat8_b,) * 2}, ValueError, "one dtype per model layer"),
+        ({"model_kv_cache_dtype": None}, TypeError, "model metadata"),
+    ],
+)
+def test_config_rejects_inconsistent_static_inputs(overrides, error_type, message, expect_error):
+    arguments = {
+        "trace": TraceConfig(mode="all"),
+        "paged_kv_cache": PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b),
+        "expected_num_layers": 32,
+        "expected_kv_heads_per_device": 8,
+        "expected_head_dim": 128,
+        "model_kv_cache_dtype": ttnn.bfloat8_b,
+    }
+
+    with expect_error(error_type, message):
+        VLLMAdapterConfig.resolve(**(arguments | overrides))
+
+
+def test_config_requires_exact_static_policy_types(expect_error):
+    class TraceConfigSubclass(TraceConfig):
+        pass
+
+    class PagedKVCacheConfigSubclass(PagedKVCacheConfig):
+        pass
+
+    arguments = {
+        "trace": TraceConfig(mode="all"),
+        "paged_kv_cache": PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b),
+        "expected_num_layers": 32,
+        "model_kv_cache_dtype": ttnn.bfloat8_b,
+    }
+
+    with expect_error(TypeError, "TraceConfig"):
+        VLLMAdapterConfig.resolve(**(arguments | {"trace": TraceConfigSubclass(mode="all")}))
+    with expect_error(TypeError, "PagedKVCacheConfig"):
+        VLLMAdapterConfig.resolve(
+            **(
+                arguments
+                | {
+                    "paged_kv_cache": PagedKVCacheConfigSubclass(
+                        block_size=32,
+                        max_num_blocks=128,
+                        dtype=ttnn.bfloat8_b,
+                    )
+                }
+            )
+        )
+
+
+def test_adapter_is_plain_orchestration_with_one_config_surface(expect_error):
+    adapter = _adapter()
+
+    assert tuple(inspect.signature(VLLMAdapter).parameters) == ("config",)
+    assert vars(adapter) == {"config": adapter.config}
+    with expect_error(TypeError, "VLLMAdapterConfig"):
+        VLLMAdapter(config=None)
+
+
+@pytest.mark.parametrize(
+    ("method_name", "expected"),
+    [
+        (
+            "normalize_prefill",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("prompt_lens", inspect.Parameter.KEYWORD_ONLY, None),
+                ("start_pos", inspect.Parameter.KEYWORD_ONLY, None),
+                ("empty_slots", inspect.Parameter.KEYWORD_ONLY, None),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("compatibility_kwargs", inspect.Parameter.KEYWORD_ONLY, None),
+            ],
+        ),
+        (
+            "normalize_decode",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("start_pos", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("reset_batch", inspect.Parameter.KEYWORD_ONLY, False),
+                ("compatibility_kwargs", inspect.Parameter.KEYWORD_ONLY, None),
+            ],
+        ),
+    ],
+)
+def test_normalizer_signatures_are_explicit(method_name, expected):
+    parameters = inspect.signature(getattr(VLLMAdapter, method_name)).parameters
+
+    assert [(name, parameter.kind, parameter.default) for name, parameter in parameters.items()] == expected
+
+
+def test_normalized_typed_dicts_have_stable_key_order_and_optional_request_state():
+    assert tuple(NormalizedPrefillKwargs.__annotations__) == (
+        "tokens",
+        "page_table",
+        "prompt_lens",
+        "start_pos",
+        "empty_slots",
+        "kv_cache",
+        "sampling_params",
+        "prompt_tokens",
+        "output_tokens",
+        "slot_remap",
+    )
+    assert NormalizedPrefillKwargs.__required_keys__ == frozenset(
+        {"tokens", "page_table", "prompt_lens", "start_pos", "empty_slots", "kv_cache", "sampling_params"}
+    )
+    assert NormalizedPrefillKwargs.__optional_keys__ == frozenset({"prompt_tokens", "output_tokens", "slot_remap"})
+    assert tuple(NormalizedDecodeKwargs.__annotations__) == (
+        "tokens",
+        "start_pos",
+        "page_table",
+        "kv_cache",
+        "sampling_params",
+        "reset_batch",
+        "prompt_tokens",
+        "output_tokens",
+        "slot_remap",
+    )
+    assert NormalizedDecodeKwargs.__required_keys__ == frozenset(
+        {"tokens", "start_pos", "page_table", "kv_cache", "sampling_params", "reset_batch"}
+    )
+    assert NormalizedDecodeKwargs.__optional_keys__ == frozenset({"prompt_tokens", "output_tokens", "slot_remap"})
+
+
+def test_normalize_prefill_positional_call_without_mutating_caller_kwargs():
+    adapter = _adapter()
+    compatibility_kwargs = {
+        "page_tables_per_layer": object(),
+        "prompt_tokens": object(),
+        "output_tokens": object(),
+        "slot_remap": object(),
+        "rope_deltas_all_users": object(),
+    }
+
+    normalized, enable_trace = adapter.normalize_prefill(
+        [[1, 2, 3, 4], [5, 6, 0, 0]],
+        [[0, 1], [2, 3]],
+        enable_trace=True,
+        prompt_lens=[4, 3],
+        start_pos=[0, 1],
+        sampling_params="sampling",
+        compatibility_kwargs=compatibility_kwargs,
+    )
+
+    assert tuple(normalized) == tuple(NormalizedPrefillKwargs.__annotations__)
+    assert normalized["tokens"].dtype == torch.long
+    assert normalized["page_table"].dtype == torch.int32
+    assert normalized["prompt_lens"].dtype == torch.long
+    assert normalized["start_pos"].dtype == torch.long
+    assert normalized["empty_slots"] is None
+    assert normalized["kv_cache"] is None
+    assert normalized["sampling_params"] == "sampling"
+    assert normalized["prompt_tokens"] is compatibility_kwargs["prompt_tokens"]
+    assert normalized["output_tokens"] is compatibility_kwargs["output_tokens"]
+    assert normalized["slot_remap"] is compatibility_kwargs["slot_remap"]
+    assert enable_trace is True
+    assert tuple(compatibility_kwargs) == (
+        "page_tables_per_layer",
+        "prompt_tokens",
+        "output_tokens",
+        "slot_remap",
+        "rope_deltas_all_users",
+    )
+
+
+def test_normalize_decode_converts_existing_tensors_and_flattens_column_tokens():
+    adapter = _adapter(trace=TraceConfig(mode="decode_only"))
+
+    normalized, enable_trace = adapter.normalize_decode(
+        torch.tensor([[1], [2]], dtype=torch.int32),
+        torch.tensor([3, 4], dtype=torch.int32),
+        torch.tensor([[0], [1]], dtype=torch.int64),
+        enable_trace=True,
+        compatibility_kwargs={"slot_remap": [0, 1]},
+    )
+
+    assert tuple(normalized) == (
+        "tokens",
+        "start_pos",
+        "page_table",
+        "kv_cache",
+        "sampling_params",
+        "reset_batch",
+        "slot_remap",
+    )
+    assert normalized["tokens"].shape == (2,)
+    assert normalized["tokens"].dtype == torch.long
+    assert normalized["start_pos"].dtype == torch.long
+    assert normalized["page_table"].dtype == torch.int32
+    assert normalized["kv_cache"] is None
+    assert normalized["sampling_params"] is None
+    assert "prompt_tokens" not in normalized
+    assert "output_tokens" not in normalized
+    assert normalized["slot_remap"] == [0, 1]
+    assert normalized["reset_batch"] is False
+    assert enable_trace is True
+
+
+@pytest.mark.parametrize("method_name", ["normalize_prefill", "normalize_decode"])
+@pytest.mark.parametrize(
+    "compatibility_kwargs",
+    (None, {"prompt_tokens": None, "output_tokens": None, "slot_remap": None}),
+)
+def test_normalize_omits_unsupplied_or_none_request_state(method_name, compatibility_kwargs):
+    adapter = _adapter(trace=TraceConfig(mode="all"), request_state_fields=())
+    args = (
+        (torch.zeros((1, 1)), torch.zeros((1, 1)))
+        if method_name == "normalize_prefill"
+        else (torch.zeros(1), torch.zeros(1), torch.zeros((1, 1)))
+    )
+
+    normalized, _ = getattr(adapter, method_name)(
+        *args,
+        enable_trace=True,
+        compatibility_kwargs=compatibility_kwargs,
+    )
+
+    assert not ({"prompt_tokens", "output_tokens", "slot_remap"} & normalized.keys())
+
+
+class _NarrowQwenRequestTarget:
+    def __init__(self):
+        self.calls = []
+
+    def prefill_forward(
+        self,
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,
+        start_pos=None,
+        empty_slots=None,
+        kv_cache=None,
+        sampling_params=None,
+        execution=None,
+    ):
+        self.calls.append(
+            (
+                "prefill",
+                {
+                    "tokens": tokens,
+                    "page_table": page_table,
+                    "prompt_lens": prompt_lens,
+                    "start_pos": start_pos,
+                    "empty_slots": empty_slots,
+                    "kv_cache": kv_cache,
+                    "sampling_params": sampling_params,
+                    "execution": execution,
+                },
+            )
+        )
+        return "prefill"
+
+    def decode_forward(
+        self,
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,
+        sampling_params=None,
+        reset_batch=False,
+        read_from_device=True,
+        execution=None,
+    ):
+        self.calls.append(
+            (
+                "decode",
+                {
+                    "tokens": tokens,
+                    "start_pos": start_pos,
+                    "page_table": page_table,
+                    "kv_cache": kv_cache,
+                    "sampling_params": sampling_params,
+                    "reset_batch": reset_batch,
+                    "read_from_device": read_from_device,
+                    "execution": execution,
+                },
+            )
+        )
+        return "decode"
+
+
+@pytest.mark.parametrize(
+    "generator_class",
+    (Qwen2Generator, Qwen25Generator, Qwen25_72BGenerator, Qwen25Coder32BGenerator),
+)
+def test_qwen_generator_dispatch_omits_absent_state_for_narrow_request_surface(generator_class):
+    generator = object.__new__(generator_class)
+    generator._adapter = _adapter(request_state_fields=())
+    generator.target = _NarrowQwenRequestTarget()
+    prefill_execution = object()
+    decode_execution = object()
+    generator._select_prefill_execution = lambda normalized, requested: prefill_execution
+    generator._select_execution = lambda operation, requested: decode_execution
+
+    request_state = {"prompt_tokens": object(), "output_tokens": object(), "slot_remap": [0]}
+    assert generator.prefill_forward([[1]], [[0]], enable_trace=False, **request_state) == "prefill"
+    assert generator.decode_forward([1], [0], [[0]], enable_trace=False, **request_state) == "decode"
+
+    assert [name for name, _ in generator.target.calls] == ["prefill", "decode"]
+    assert generator.target.calls[0][1]["execution"] is prefill_execution
+    assert generator.target.calls[1][1]["execution"] is decode_execution
+
+
+@pytest.mark.parametrize(
+    "generator_module",
+    (llama3_generator_module, llama70_generator_module, qwen3_generator_module),
+)
+def test_stateful_generator_adapter_reuses_executor_request_state_allowlist(monkeypatch, generator_module):
+    request_state_fields = ("prompt_tokens", "output_tokens", "slot_remap")
+    lane = SimpleNamespace(
+        model=object(),
+        config=SimpleNamespace(
+            trace=TraceConfig(mode="decode_only"),
+            paged_kv_cache=PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b),
+        ),
+        _request_state_fields=request_state_fields,
+    )
+    monkeypatch.setattr(
+        generator_module,
+        "_model_kv_metadata",
+        lambda model: ((ttnn.bfloat8_b,), 1, 1, 128),
+    )
+
+    adapter = generator_module._build_vllm_adapter(lane)
+
+    assert adapter.config.request_state_fields == request_state_fields
+
+
+@pytest.mark.parametrize(
+    ("method_name", "args", "trace", "hint"),
+    [
+        (
+            "normalize_prefill",
+            (torch.zeros((1, 1)), torch.zeros((1, 1))),
+            TraceConfig(mode="decode_only"),
+            True,
+        ),
+        (
+            "normalize_decode",
+            (torch.zeros(1), torch.zeros(1), torch.zeros((1, 1))),
+            TraceConfig(mode="none"),
+            True,
+        ),
+    ],
+)
+def test_normalize_rejects_trace_hint_that_disagrees_with_static_policy(method_name, args, trace, hint, expect_error):
+    adapter = _adapter(trace=trace)
+
+    with expect_error(ValueError, "enable_trace"):
+        getattr(adapter, method_name)(*args, enable_trace=hint)
+
+
+def test_normalize_accepts_unselected_operation_with_available_trace_target():
+    adapter = _adapter(trace=TraceConfig(mode="all"))
+
+    _, enable_trace = adapter.normalize_prefill(
+        torch.zeros((1, 1)),
+        torch.zeros((1, 1)),
+        enable_trace=False,
+    )
+
+    assert enable_trace is False
+
+
+@pytest.mark.parametrize(
+    ("mode", "prefill_enabled", "decode_enabled"),
+    [
+        ("none", False, False),
+        ("decode_only", False, True),
+        ("all", True, True),
+    ],
+)
+def test_normalize_accepts_operation_specific_static_trace_policy(mode, prefill_enabled, decode_enabled):
+    adapter = _adapter(trace=TraceConfig(mode=mode))
+
+    _, normalized_prefill_trace = adapter.normalize_prefill(
+        torch.zeros((1, 1)),
+        torch.zeros((1, 1)),
+        enable_trace=prefill_enabled,
+    )
+    _, normalized_decode_trace = adapter.normalize_decode(
+        torch.zeros(1),
+        torch.zeros(1),
+        torch.zeros((1, 1)),
+        enable_trace=decode_enabled,
+    )
+
+    assert normalized_prefill_trace is prefill_enabled
+    assert normalized_decode_trace is decode_enabled
+
+
+@pytest.mark.parametrize("hint", [None, "true", 1])
+def test_normalize_requires_an_explicit_boolean_trace_selection(hint, expect_error):
+    adapter = _adapter()
+
+    with expect_error(TypeError, "enable_trace"):
+        if hint is None:
+            adapter.normalize_prefill(torch.zeros((1, 1)), torch.zeros((1, 1)))
+        else:
+            adapter.normalize_prefill(torch.zeros((1, 1)), torch.zeros((1, 1)), enable_trace=hint)
+
+
+def test_normalize_rejects_duplicate_positional_and_keyword_argument(expect_error):
+    adapter = _adapter()
+
+    with expect_error(TypeError, "tokens"):
+        adapter.normalize_prefill(
+            torch.zeros((1, 1)),
+            torch.zeros((1, 1)),
+            tokens=torch.zeros((1, 1)),
+            enable_trace=True,
+        )
+
+
+@pytest.mark.parametrize("method_name", ["normalize_prefill", "normalize_decode"])
+def test_normalize_rejects_unknown_compatibility_keys(method_name, expect_error):
+    adapter = _adapter()
+    args = (
+        (torch.zeros((1, 1)), torch.zeros((1, 1)))
+        if method_name == "normalize_prefill"
+        else (torch.zeros(1), torch.zeros(1), torch.zeros((1, 1)))
+    )
+
+    with expect_error(TypeError, "unexpected keyword argument 'unknown_plugin_field'"):
+        getattr(adapter, method_name)(
+            *args,
+            enable_trace=True,
+            compatibility_kwargs={"unknown_plugin_field": object()},
+        )
+
+
+def _signature_entries(method):
+    return [
+        (name, parameter.kind, parameter.default) for name, parameter in inspect.signature(method).parameters.items()
+    ]
+
+
+@pytest.mark.parametrize(
+    ("method_name", "expected"),
+    [
+        (
+            "compile_prefill",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("prompt_lens", inspect.Parameter.KEYWORD_ONLY, None),
+                ("start_pos", inspect.Parameter.KEYWORD_ONLY, None),
+                ("empty_slots", inspect.Parameter.KEYWORD_ONLY, None),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+            ],
+        ),
+        (
+            "compile_decode",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("start_pos", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("reset_batch", inspect.Parameter.KEYWORD_ONLY, False),
+            ],
+        ),
+        (
+            "prefill_forward",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("prompt_lens", inspect.Parameter.KEYWORD_ONLY, None),
+                ("start_pos", inspect.Parameter.KEYWORD_ONLY, None),
+                ("empty_slots", inspect.Parameter.KEYWORD_ONLY, None),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("compatibility_kwargs", inspect.Parameter.VAR_KEYWORD, inspect.Parameter.empty),
+            ],
+        ),
+        (
+            "decode_forward",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("start_pos", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, None),
+                ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None),
+                ("reset_batch", inspect.Parameter.KEYWORD_ONLY, False),
+                ("read_from_device", inspect.Parameter.KEYWORD_ONLY, True),
+                ("compatibility_kwargs", inspect.Parameter.VAR_KEYWORD, inspect.Parameter.empty),
+            ],
+        ),
+        (
+            "read_decode_output",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("async_read", inspect.Parameter.KEYWORD_ONLY, False),
+            ],
+        ),
+        (
+            "process_decode_output_host",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("is_tokens", inspect.Parameter.KEYWORD_ONLY, False),
+            ],
+        ),
+        (
+            "warmup_model_prefill",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("can_sample_on_device", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+            ],
+        ),
+        (
+            "warmup_model_decode",
+            [
+                ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty),
+                ("kv_cache", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("max_batch_size", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("num_blocks", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("can_sample_on_device", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+                ("enable_trace", inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.empty),
+            ],
+        ),
+    ],
+)
+def test_registered_generator_signatures_are_exact(method_name, expected):
+    assert _signature_entries(getattr(Llama3Generator, method_name)) == expected
+
+
+def test_registered_generator_sizing_signature_has_no_compatibility_bag():
+    assert _signature_entries(Llama3Generator.get_max_tokens_all_users) == [
+        ("model_name", inspect.Parameter.POSITIONAL_OR_KEYWORD, ""),
+        ("num_devices", inspect.Parameter.POSITIONAL_OR_KEYWORD, 1),
+        ("tt_data_parallel", inspect.Parameter.POSITIONAL_OR_KEYWORD, 1),
+        ("max_model_len", inspect.Parameter.POSITIONAL_OR_KEYWORD, 0),
+        ("max_num_seqs", inspect.Parameter.POSITIONAL_OR_KEYWORD, 1),
+    ]
+
+
+def test_registered_generator_compatibility_bags_exist_only_on_forward_methods():
+    variadic_keyword_methods = {
+        method_name
+        for method_name, method in vars(Llama3Generator).items()
+        if inspect.isfunction(method)
+        and any(
+            parameter.kind is inspect.Parameter.VAR_KEYWORD
+            for parameter in inspect.signature(method).parameters.values()
+        )
+    }
+
+    assert variadic_keyword_methods == {"prefill_forward", "decode_forward"}
+
+
+class _ExplicitGeneratorTarget:
+    model = SimpleNamespace()
+    model_args = object()
+    mesh_device = object()
+    cache_path = "cache"
+    already_warmed_up_prefill = False
+    eager_execution = object()
+    traced_prefill_execution = object()
+    traced_decode_execution = object()
+
+    def __init__(self):
+        self.calls = []
+
+    def _record(self, method_name, arguments):
+        self.calls.append(
+            (
+                method_name,
+                {name: value for name, value in arguments.items() if name != "self"},
+            )
+        )
+
+    def can_trace_prefill(
+        self,
+        *,
+        tokens,  # ↓ Core request
+        prompt_lens=None,  # ↓ Sequence metadata
+        start_pos=None,
+        empty_slots=None,  # ↓ Lane routing
+    ):
+        self._record("can_trace_prefill", locals())
+        return True
+
+    def compile_prefill(
+        self,
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,  # ↓ Sequence metadata
+        start_pos=None,
+        empty_slots=None,  # ↓ Lane routing
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        prompt_tokens=None,  # ↓ Request-owned sampling state
+        output_tokens=None,
+        slot_remap=None,
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("compile_prefill", locals())
+
+    def compile_decode(
+        self,
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        prompt_tokens=None,  # ↓ Request-owned sampling state
+        output_tokens=None,
+        slot_remap=None,
+        reset_batch=False,  # ↓ State transition
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("compile_decode", locals())
+
+    def prefill_forward(
+        self,
+        tokens,
+        page_table,
+        *,
+        prompt_lens=None,  # ↓ Sequence metadata
+        start_pos=None,
+        empty_slots=None,  # ↓ Lane routing
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        prompt_tokens=None,  # ↓ Request-owned sampling state
+        output_tokens=None,
+        slot_remap=None,
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("prefill_forward", locals())
+        return "prefill"
+
+    def decode_forward(
+        self,
+        tokens,
+        start_pos,
+        page_table,
+        *,
+        kv_cache=None,  # ↓ Borrowed resources
+        sampling_params=None,  # ↓ Sampling
+        prompt_tokens=None,  # ↓ Request-owned sampling state
+        output_tokens=None,
+        slot_remap=None,
+        reset_batch=False,  # ↓ State transition
+        read_from_device=True,  # ↓ Output policy
+        execution=None,  # ↓ Internal dispatch
+    ):
+        self._record("decode_forward", locals())
+        return "decode"
+
+
+def test_registered_generator_compile_methods_normalize_and_select_execution():
+    target = _ExplicitGeneratorTarget()
+    generator = Llama3Generator(target, _adapter())
+    tokens = torch.tensor([[1, 2]])
+    page_table = torch.tensor([[0]])
+    prompt_lens = torch.tensor([2])
+    start_pos = torch.tensor([0])
+    kv_cache = object()
+    sampling_params = object()
+
+    generator.compile_prefill(
+        tokens,
+        page_table,
+        enable_trace=True,
+        prompt_lens=prompt_lens,
+        start_pos=start_pos,
+        empty_slots=[0],
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+    )
+    assert target.calls[0][0] == "compile_prefill"
+    assert target.calls[0][1]["execution"] is target.traced_prefill_execution
+    assert set(target.calls[0][1]) == set(NormalizedPrefillKwargs.__annotations__) | {"execution"}
+
+    generator.compile_decode(
+        tokens[:, 0],
+        start_pos,
+        page_table,
+        enable_trace=True,
+        kv_cache=kv_cache,
+        sampling_params=sampling_params,
+        reset_batch=True,
+    )
+    assert target.calls[1][0] == "compile_decode"
+    assert target.calls[1][1]["execution"] is target.traced_decode_execution
+    assert set(target.calls[1][1]) == set(NormalizedDecodeKwargs.__annotations__) | {"execution"}
+
+
+def test_registered_generator_discards_allowlisted_compatibility_and_limits_trace_classification():
+    target = _ExplicitGeneratorTarget()
+    generator = Llama3Generator(target, _adapter())
+    tokens = torch.tensor([[1, 2]])
+    page_table = torch.tensor([[0]])
+    prompt_lens = torch.tensor([2])
+    start_pos = torch.tensor([0])
+    kv_cache = object()
+    sampling_params = object()
+
+    assert (
+        generator.prefill_forward(
+            tokens,
+            page_table,
+            enable_trace=True,
+            prompt_lens=prompt_lens,
+            start_pos=start_pos,
+            empty_slots=[0],
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+            page_tables_per_layer=object(),
+        )
+        == "prefill"
+    )
+    assert target.calls[0][0] == "prefill_forward"
+    assert target.calls[0][1]["execution"] is target.traced_prefill_execution
+    assert set(target.calls[0][1]) == set(NormalizedPrefillKwargs.__annotations__) | {"execution"}
+
+    assert (
+        generator.decode_forward(
+            tokens[:, 0],
+            start_pos,
+            page_table,
+            enable_trace=True,
+            kv_cache=kv_cache,
+            sampling_params=sampling_params,
+            reset_batch=True,
+            read_from_device=False,
+            slot_remap=[0],
+        )
+        == "decode"
+    )
+    assert target.calls[1][0] == "decode_forward"
+    assert target.calls[1][1]["execution"] is target.traced_decode_execution
+    assert target.calls[1][1]["read_from_device"] is False
+    assert set(target.calls[1][1]) == set(NormalizedDecodeKwargs.__annotations__) | {
+        "read_from_device",
+        "execution",
+    }
+
+
+@pytest.mark.parametrize(
+    ("method_name", "args"),
+    [
+        ("prefill_forward", (torch.zeros((1, 1)), torch.zeros((1, 1)))),
+        ("decode_forward", (torch.zeros(1), torch.zeros(1), torch.zeros((1, 1)))),
+    ],
+)
+def test_registered_generator_rejects_unknown_compatibility_before_target_selection(
+    method_name,
+    args,
+    expect_error,
+):
+    target = _ExplicitGeneratorTarget()
+    generator = Llama3Generator(target, _adapter())
+
+    with expect_error(TypeError, "unexpected keyword argument 'unknown_plugin_field'"):
+        getattr(generator, method_name)(
+            *args,
+            enable_trace=True,
+            unknown_plugin_field=object(),
+        )
+
+    assert target.calls == []
+
+
+def test_registered_generator_rejects_unknown_nonforward_keywords(expect_error):
+    with expect_error(TypeError, "unexpected keyword argument"):
+        Llama3Generator.get_max_tokens_all_users(unknown_plugin_field=True)
+
+
+def test_registered_generator_forwards_output_and_warmup_arguments_by_name():
+    pending_output = object()
+    read_events = [object()]
+    target = create_autospec(Llama3Executor, instance=True)
+    target.read_decode_output.return_value = pending_output, read_events
+    target.process_decode_output_host.return_value = "tokens", "log-probs"
+    generator = Llama3Generator(target, adapter=object())
+    tt_out = object()
+    kv_cache = object()
+
+    assert generator.read_decode_output(tt_out, async_read=True) == (pending_output, read_events)
+    target.read_decode_output.assert_called_once_with(tt_out=tt_out, async_read=True)
+    assert generator.process_decode_output_host(pending_output, is_tokens=True) == ("tokens", "log-probs")
+    target.process_decode_output_host.assert_called_once_with(tt_out=pending_output, is_tokens=True)
+
+    generator.warmup_model_prefill(
+        kv_cache=kv_cache,
+        can_sample_on_device=True,
+        enable_trace=False,
+    )
+    target.warmup_model_prefill.assert_called_once_with(
+        kv_cache=kv_cache,
+        can_sample_on_device=True,
+        enable_trace=False,
+    )
+    generator.warmup_model_decode(
+        kv_cache=kv_cache,
+        max_batch_size=8,
+        num_blocks=128,
+        can_sample_on_device=False,
+        enable_trace=True,
+    )
+    target.warmup_model_decode.assert_called_once_with(
+        kv_cache=kv_cache,
+        max_batch_size=8,
+        num_blocks=128,
+        can_sample_on_device=False,
+        enable_trace=True,
+    )
+
+
+def test_resolve_legacy_kv_cache_returns_new_immutable_config():
+    base = PagedKVCacheConfig(block_size=32, max_num_blocks=128, dtype=ttnn.bfloat8_b)
+    adapter = _adapter(paged_config=base)
+
+    resolved = adapter.resolve_legacy_kv_cache_config(
+        (129, 8, 64, 128),
+        torch.bfloat16,
+        32,
+    )
+
+    assert resolved is not base
+    assert base.num_blocks is None
+    assert base.block_size == 32
+    assert base.max_num_blocks == 128
+    assert resolved.num_blocks == 129
+    assert resolved.block_size == 64
+    assert resolved.max_num_blocks == 129
+    assert resolved.dtype == base.dtype
+    assert resolved.memory_config == base.memory_config
+
+
+@pytest.mark.parametrize(
+    ("shape", "dtype", "num_layers", "message"),
+    [
+        ((64, 8, 0, 128), torch.bfloat16, 32, "block_size"),
+        ((64, 4, 32, 128), torch.bfloat16, 32, "KV heads"),
+        ((64, 8, 32, 64), torch.bfloat16, 32, "head dimension"),
+        ((0, 8, 32, 128), torch.bfloat16, 32, "num_blocks"),
+        ((64, 8, 32, 128), torch.float32, 32, "dtype"),
+        ((64, 8, 32, 128), torch.bfloat16, 31, "layer count"),
+    ],
+)
+def test_resolve_legacy_kv_cache_rejects_mismatched_vllm_spec(shape, dtype, num_layers, message, expect_error):
+    adapter = _adapter()
+
+    with expect_error((TypeError, ValueError), message):
+        adapter.resolve_legacy_kv_cache_config(shape, dtype, num_layers)
+
+
+def test_adapter_rejects_static_dtype_that_disagrees_with_model_owned_dtype(expect_error):
+    with expect_error(ValueError, "model-owned"):
+        _adapter(model_dtype=ttnn.bfloat16)
+
+
+def test_adapter_requires_explicit_model_owned_dtype_metadata(expect_error):
+    with expect_error(TypeError, "must be supplied from model metadata"):
+        _adapter(model_dtype=None)
+
+
+def test_bfloat4_model_dtype_uses_shared_bfloat16_torch_surrogate():
+    config = PagedKVCacheConfig(
+        block_size=32,
+        max_num_blocks=128,
+        dtype=ttnn.bfloat4_b,
+    )
+    adapter = _adapter(paged_config=config, model_dtype=ttnn.bfloat4_b)
+
+    resolved = adapter.resolve_legacy_kv_cache_config(
+        (64, 8, 32, 128),
+        torch.bfloat16,
+        32,
+    )
+
+    assert resolved.dtype == ttnn.bfloat4_b
+    assert resolved.num_blocks == 64
+
+
+def test_resolve_legacy_kv_cache_rejects_replacing_resolved_capacity(expect_error):
+    adapter = _adapter(
+        paged_config=PagedKVCacheConfig(
+            block_size=32,
+            max_num_blocks=128,
+            dtype=ttnn.bfloat8_b,
+            num_blocks=32,
+        )
+    )
+
+    with expect_error(ValueError, "already resolved"):
+        adapter.resolve_legacy_kv_cache_config((64, 8, 32, 128), torch.bfloat16, 32)

code/models/common/tests/llm_runtime/test_warmup.py

@@ -0,0 +1,1012 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import Any, Sequence
+
+import pytest
+import torch
+
+from models.common.llm_runtime.config import PageTableLayout, TraceConfig, WarmupConfig
+from models.common.llm_runtime.decode import DecodeRuntimeConfig
+from models.common.llm_runtime.output_reader import OutputReader
+from models.common.llm_runtime.prefill.config import PrefillRuntimeConfig
+from models.common.llm_runtime.program_compiler import CompiledProgram, OutputSpec, ProgramKey
+from models.common.llm_runtime.warmup import (
+    CoverageAlias,
+    WarmupCoordinator,
+    WarmupCoordinatorConfig,
+    _resolve_coverage_manifest,
+)
+
+
+class RecordingExecution:
+    def __init__(self, events=None):
+        self.prefill_calls = []
+        self.decode_calls = []
+        self.events = events if events is not None else []
+        self.fail_decode_call = None
+        self.prefill_replays = []
+
+    def compile_prefill(
+        self,
+        *,
+        tokens: torch.Tensor,  # ↓ Core request
+        page_table: torch.Tensor,
+        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
+        start_pos: torch.Tensor | None = None,
+        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
+        sampling_params: Any = None,  # ↓ Sampling
+    ) -> None:
+        self.events.append("compile_prefill")
+        self.prefill_calls.append(
+            {
+                "tokens": tokens,
+                "page_table": page_table,
+                "prompt_lens": prompt_lens,
+                "start_pos": start_pos,
+                "empty_slots": empty_slots,
+                "sampling_params": sampling_params,
+            }
+        )
+
+    def compile_decode(
+        self,
+        *,
+        tokens: torch.Tensor,  # ↓ Core request
+        start_pos: torch.Tensor,
+        page_table: torch.Tensor,
+        sampling_params: Any = None,  # ↓ Sampling
+        reset_batch: bool = False,  # ↓ State transition
+    ) -> None:
+        call = len(self.decode_calls) + 1
+        self.events.append("compile_decode")
+        if call == self.fail_decode_call:
+            self.fail_decode_call = None
+            raise RuntimeError("decode compile failed")
+        self.decode_calls.append(
+            {
+                "tokens": tokens,
+                "start_pos": start_pos,
+                "page_table": page_table,
+                "sampling_params": sampling_params,
+                "reset_batch": reset_batch,
+            }
+        )
+
+    def prefill_forward(
+        self,
+        *,
+        tokens: torch.Tensor,  # ↓ Core request
+        page_table: torch.Tensor,
+        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
+        start_pos: torch.Tensor | None = None,
+        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
+        sampling_params: Any = None,  # ↓ Sampling
+    ) -> None:
+        self.events.append("prefill_replay")
+        self.prefill_replays.append(
+            {
+                "tokens": tokens,
+                "page_table": page_table,
+                "prompt_lens": prompt_lens,
+                "start_pos": start_pos,
+                "empty_slots": empty_slots,
+                "sampling_params": sampling_params,
+            }
+        )
+
+
+class RecordingTraceCompiler:
+    def __init__(self, events=None):
+        self.calls = 0
+        self.events = events if events is not None else []
+
+    def capture_all(self):
+        self.events.append("capture")
+        self.calls += 1
+
+
+class Mesh:
+    shape = (1, 1)
+
+
+def make_runtime_configs(
+    *,
+    sampling=True,
+    lane_capacity=4,
+    allow_force_argmax=True,
+    page_table_layout=None,
+    sampling_config=None,
+    model=None,
+):
+    mesh = Mesh()
+    sampling_config = sampling_config or SimpleNamespace(
+        allow_force_argmax=allow_force_argmax,
+        max_top_k=32,
+    )
+    sampling_config.max_batch_size = lane_capacity
+    model = model or SimpleNamespace(
+        config=SimpleNamespace(max_batch_size=lane_capacity, mesh_device=mesh, num_devices=1),
+        sampling=SimpleNamespace(
+            config=sampling_config,
+            decode_forward=lambda logits, *, k=None, p=None, temp=None, seeds=None, tt_out_tok=None, enable_log_probs=False: None,
+        ),
+        vocab_size=128,
+    )
+    mesh = model.config.mesh_device
+    layout = page_table_layout or PageTableLayout(
+        block_size=32,
+        raw_capacity_width=128,
+        prefill_width=192,
+        decode_width=128,
+    )
+    output_reader = OutputReader(mesh)
+    return (
+        PrefillRuntimeConfig.resolve(
+            model=model,
+            output_reader=output_reader,
+            page_table_layout=layout,
+            max_batch_size=lane_capacity,
+            max_prefill_chunk_size=128,
+            device_sampling_enabled=sampling,
+            can_enable_trace=lambda _sequence_length, _batch_size: True,
+        ),
+        DecodeRuntimeConfig.resolve(
+            model=model,
+            output_reader=output_reader,
+            lane_capacity=lane_capacity,
+            page_table_layout=layout,
+            device_sampling_enabled=sampling,
+        ),
+    )
+
+
+def make_coordinator(
+    *,
+    trace_mode="all",
+    sampling=True,
+    warmup_config=None,
+    sequence_lengths=(128, 1024),
+    lane_capacity=4,
+    execution=None,
+    trace_compiler=None,
+    events=None,
+    allow_force_argmax=True,
+    page_table_layout=None,
+    sampling_config=None,
+):
+    events = events if events is not None else []
+    execution = execution or RecordingExecution(events)
+    if trace_compiler is None and trace_mode != "none":
+        trace_compiler = RecordingTraceCompiler(events)
+    sampling_calls = []
+    bound_calls = []
+
+    def ensure_sampling():
+        events.append("sampling")
+        sampling_calls.append(True)
+
+    def validate_bound(value):
+        bound_calls.append(value)
+
+    layout = page_table_layout or PageTableLayout(
+        block_size=32,
+        raw_capacity_width=128,
+        prefill_width=192,
+        decode_width=128,
+    )
+    prefill_config, decode_config = make_runtime_configs(
+        sampling=sampling,
+        lane_capacity=lane_capacity,
+        allow_force_argmax=allow_force_argmax,
+        page_table_layout=layout,
+        sampling_config=sampling_config,
+    )
+    execution.prefill = SimpleNamespace(config=prefill_config)
+    execution.decode = SimpleNamespace(config=decode_config)
+    execution.eager_executor = execution
+    execution.trace_compiler = trace_compiler
+
+    coordinator = WarmupCoordinator(
+        config=WarmupCoordinatorConfig.resolve(
+            warmup=warmup_config or WarmupConfig(),
+            trace=TraceConfig(trace_mode),
+            prefill=prefill_config,
+            decode=decode_config,
+            prefill_sequence_lengths=sequence_lengths,
+        ),
+        execution=execution,
+        ensure_sampling_buffers=ensure_sampling,
+        validate_bound_cache=validate_bound,
+    )
+    return coordinator, execution, trace_compiler, sampling_calls, bound_calls, events
+
+
+@pytest.mark.parametrize(
+    ("method_name", "parameter_names"),
+    [
+        ("warmup_prefill", ("self", "kv_cache", "can_sample_on_device", "enable_trace")),
+        (
+            "warmup_decode",
+            ("self", "kv_cache", "max_batch_size", "num_blocks", "can_sample_on_device", "enable_trace"),
+        ),
+    ],
+)
+def test_warmup_signatures_match_registered_plugin_contract(method_name, parameter_names):
+    parameters = inspect.signature(getattr(WarmupCoordinator, method_name)).parameters
+
+    assert tuple(parameters) == parameter_names
+    assert parameters["self"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD
+    for name in parameter_names[1:]:
+        assert parameters[name].kind is inspect.Parameter.KEYWORD_ONLY
+        assert parameters[name].default is inspect.Parameter.empty
+
+
+def test_registered_plugin_warmup_calls_validate_cache_without_forwarding_it():
+    cache = object()
+    coordinator, execution, _, _, bound_calls, _ = make_coordinator(
+        trace_mode="none",
+        sampling=False,
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        lane_capacity=1,
+    )
+    prefill_kwargs = {
+        "kv_cache": cache,
+        "can_sample_on_device": False,
+    }
+    decode_kwargs = {
+        "kv_cache": cache,
+        "max_batch_size": 1,
+        "num_blocks": 8,
+        "can_sample_on_device": False,
+    }
+
+    coordinator.warmup_prefill(enable_trace=False, **prefill_kwargs)
+    coordinator.warmup_decode(enable_trace=False, **decode_kwargs)
+
+    assert coordinator.coverage_manifest is None
+    assert bound_calls == [cache, cache]
+    assert all(
+        tuple(call) == ("tokens", "page_table", "prompt_lens", "start_pos", "empty_slots", "sampling_params")
+        for call in execution.prefill_calls
+    )
+    assert all(
+        tuple(call) == ("tokens", "start_pos", "page_table", "sampling_params", "reset_batch")
+        for call in execution.decode_calls
+    )
+
+
+@pytest.mark.parametrize(
+    ("method_name", "plugin_kwargs", "unexpected_name"),
+    [
+        (
+            "warmup_prefill",
+            {"kv_cache": "cache", "can_sample_on_device": False, "enable_trace": False},
+            "greedy_only",
+        ),
+        (
+            "warmup_decode",
+            {
+                "kv_cache": "cache",
+                "max_batch_size": 1,
+                "num_blocks": 8,
+                "can_sample_on_device": False,
+                "enable_trace": False,
+            },
+            "read_from_device",
+        ),
+    ],
+)
+def test_warmup_contract_rejects_unregistered_plugin_keywords(
+    method_name,
+    plugin_kwargs,
+    unexpected_name,
+    expect_error,
+):
+    coordinator, *_ = make_coordinator(trace_mode="none", sampling=False, lane_capacity=1)
+    plugin_kwargs[unexpected_name] = False
+
+    with expect_error(TypeError, unexpected_name):
+        getattr(coordinator, method_name)(**plugin_kwargs)
+
+
+def test_configured_prefill_lengths_override_model_supported_defaults():
+    coordinator, execution, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(1024,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        sampling=False,
+    )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+
+    regular_lengths = [int(call["tokens"].shape[-1]) for call in execution.prefill_calls if call["start_pos"] is None]
+    assert regular_lengths == [1024]
+
+
+@pytest.mark.parametrize(
+    ("sequence_lengths", "message"),
+    [
+        ((), "non-empty tuple"),
+        ((True,), "positive integers"),
+        ((128, 128), "unique"),
+    ],
+)
+def test_model_supported_prefill_lengths_are_validated_once(sequence_lengths, message, expect_error):
+    with expect_error(ValueError, message):
+        make_coordinator(sequence_lengths=sequence_lengths)
+
+
+def test_sampler_argmax_capability_is_resolved_once():
+    class SamplingConfig:
+        reads = 0
+        max_top_k = 32
+
+        @property
+        def allow_force_argmax(self):
+            self.reads += 1
+            return True
+
+    sampling_config = SamplingConfig()
+    coordinator, *_ = make_coordinator(sampling_config=sampling_config)
+    resolved_reads = sampling_config.reads
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=4,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    assert sampling_config.reads == resolved_reads
+
+
+def test_page_table_layout_can_be_reconfigured_only_before_use(expect_error):
+    coordinator, execution, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sampling=False,
+    )
+    final_layout = PageTableLayout(
+        block_size=32,
+        raw_capacity_width=4,
+        prefill_width=64,
+        decode_width=8,
+    )
+
+    coordinator.configure_page_table_layout(final_layout)
+    assert coordinator.config.page_table_layout is final_layout
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+
+    assert all(call["start_pos"] is None for call in execution.prefill_calls)
+    with expect_error(RuntimeError, "configuration is sealed"):
+        coordinator.configure_page_table_layout(final_layout)
+
+
+def test_explicit_configuration_seal_precedes_physical_kv_allocation(expect_error):
+    coordinator, *_ = make_coordinator()
+    coordinator.seal_configuration()
+
+    with expect_error(RuntimeError, "configuration is sealed"):
+        coordinator.configure_page_table_layout(
+            PageTableLayout(
+                block_size=32,
+                raw_capacity_width=64,
+                prefill_width=128,
+                decode_width=64,
+            )
+        )
+
+
+def test_page_table_layout_reconfiguration_requires_immutable_layout(expect_error):
+    coordinator, *_ = make_coordinator()
+
+    with expect_error(TypeError, "PageTableLayout"):
+        coordinator.configure_page_table_layout(SimpleNamespace(block_size=32))
+
+
+def test_resolved_config_is_frozen_and_owns_both_coverage_plans(expect_error):
+    coordinator, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        lane_capacity=2,
+    )
+
+    assert coordinator.config.eager_plan.decode == (coordinator.config.eager_plan.decode[0],)
+    assert [case.sampling_path for case in coordinator.config.sampled_plan.decode] == ["logits", "argmax"]
+    with expect_error(AttributeError, "cannot assign"):
+        coordinator.config.lane_batch_size = 4
+
+
+def test_direct_config_construction_rejects_inconsistent_derived_plan(expect_error):
+    coordinator, *_ = make_coordinator()
+
+    with expect_error(ValueError, "plans must match"):
+        replace(coordinator.config, eager_plan=coordinator.config.sampled_plan)
+
+
+@pytest.mark.parametrize(
+    ("mismatch", "message"),
+    [
+        ("model", "share one model"),
+        ("layout", "share one page-table layout"),
+        ("lane", "share one lane capacity"),
+        ("sampling", "share device-sampling policy"),
+        ("argmax", "share force-argmax capability"),
+        ("raw_ceiling", "share one page-table layout ceiling"),
+        ("decode_ceiling", "share one page-table layout ceiling"),
+    ],
+)
+def test_resolution_rejects_inconsistent_runtime_configs(mismatch, message, expect_error):
+    prefill, decode = make_runtime_configs()
+    if mismatch == "model":
+        _, decode = make_runtime_configs()
+    elif mismatch == "layout":
+        decode = decode.with_page_table_layout(PageTableLayout(32, 64, 128, 64))
+    elif mismatch == "lane":
+        decode = DecodeRuntimeConfig.resolve(
+            model=prefill.model,
+            output_reader=prefill.output_reader,
+            lane_capacity=2,
+            page_table_layout=prefill.page_table_layout,
+            device_sampling_enabled=True,
+        )
+    elif mismatch == "sampling":
+        decode = DecodeRuntimeConfig.resolve(
+            model=prefill.model,
+            output_reader=prefill.output_reader,
+            lane_capacity=prefill.max_batch_size,
+            page_table_layout=prefill.page_table_layout,
+            device_sampling_enabled=False,
+        )
+    elif mismatch == "argmax":
+        prefill.model.sampling.config.allow_force_argmax = False
+        decode = DecodeRuntimeConfig.resolve(
+            model=prefill.model,
+            output_reader=prefill.output_reader,
+            lane_capacity=prefill.max_batch_size,
+            page_table_layout=prefill.page_table_layout,
+            device_sampling_enabled=True,
+        )
+    elif mismatch == "raw_ceiling":
+        ceiling = prefill.page_table_layout_ceiling
+        prefill = replace(
+            prefill,
+            page_table_layout_ceiling=replace(
+                ceiling,
+                raw_capacity_width=ceiling.raw_capacity_width + 1,
+                decode_width=ceiling.decode_width + 8,
+            ),
+        )
+    else:
+        ceiling = prefill.page_table_layout_ceiling
+        prefill = replace(
+            prefill,
+            page_table_layout_ceiling=replace(ceiling, decode_width=ceiling.decode_width + 8),
+        )
+
+    with expect_error(ValueError, message):
+        WarmupCoordinatorConfig.resolve(
+            warmup=WarmupConfig(),
+            trace=TraceConfig("all"),
+            prefill=prefill,
+            decode=decode,
+            prefill_sequence_lengths=(128,),
+        )
+
+
+def test_constructor_rejects_execution_disagreement_with_resolved_config(expect_error):
+    coordinator, execution, *_ = make_coordinator()
+    config = coordinator.config
+    execution.prefill.config = replace(
+        execution.prefill.config,
+        page_table_layout=PageTableLayout(32, 64, 128, 64),
+    )
+
+    with expect_error(ValueError, "warmup config page-table layout"):
+        WarmupCoordinator(
+            config=config,
+            execution=execution,
+            ensure_sampling_buffers=lambda: None,
+            validate_bound_cache=lambda _: None,
+        )
+
+
+def test_runtime_does_not_copy_static_config_fields():
+    coordinator, *_ = make_coordinator()
+
+    assert {
+        "page_table_layout",
+        "prefill_sequence_lengths",
+        "lane_batch_size",
+        "device_sampling_enabled",
+        "allow_force_argmax",
+        "prime_q128_tile_ends",
+        "prefill_trace_enabled",
+        "decode_trace_enabled",
+        "eager_plan",
+        "sampled_plan",
+    }.isdisjoint(vars(coordinator))
+
+
+def test_layout_replacement_is_immutable_bounded_and_rebuilds_coverage(expect_error):
+    coordinator, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        page_table_layout=PageTableLayout(32, 128, 192, 128),
+    )
+    original = coordinator.config
+    replacement = PageTableLayout(32, 4, 64, 8)
+
+    coordinator.configure_page_table_layout(replacement)
+
+    assert coordinator.config is not original
+    assert coordinator.config.page_table_layout_ceiling is original.page_table_layout
+    assert original.page_table_layout.raw_capacity_width == 128
+    assert not any(case.cached_tokens for case in coordinator.config.eager_plan.prefill)
+    with expect_error(ValueError, "cannot change block_size"):
+        original.with_page_table_layout(PageTableLayout(16, 4, 64, 8))
+    with expect_error(ValueError, "capacity ceiling"):
+        original.with_page_table_layout(PageTableLayout(32, 129, 192, 136))
+    with expect_error(ValueError, "canonical geometry"):
+        original.with_page_table_layout(PageTableLayout(32, 128, 200, 128))
+
+
+def test_q128_batches_are_capped_by_lane_and_non128_is_batch_one():
+    config = WarmupConfig(prefill_batch_sizes=(1, 2, 4, 8, 16, 32))
+    coordinator, execution, *_ = make_coordinator(warmup_config=config, lane_capacity=8)
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+
+    regular_q128 = [
+        int(call["tokens"].shape[0])
+        for call in execution.prefill_calls
+        if int(call["tokens"].shape[-1]) == 128 and call["start_pos"] is None
+    ]
+    regular_q1024 = [
+        int(call["tokens"].shape[0])
+        for call in execution.prefill_calls
+        if int(call["tokens"].shape[-1]) == 1024 and call["start_pos"] is None
+    ]
+    assert regular_q128 == [1, 2, 4, 8]
+    assert regular_q1024 == [1]
+
+
+def test_sampling_paths_include_forced_prefill_topk_and_opt_in_true_topk_decode():
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,), include_decode_top_k=True)
+    coordinator, execution, *_ = make_coordinator(warmup_config=config, sequence_lengths=(128,), lane_capacity=2)
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=2,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    assert execution.prefill_calls[0]["sampling_params"] is None
+    assert execution.prefill_calls[1]["sampling_params"].top_k.tolist() == [32]
+    assert execution.decode_calls[0]["sampling_params"] is None
+    assert execution.decode_calls[1]["sampling_params"].top_k.tolist() == [1, 1]
+    assert execution.decode_calls[2]["sampling_params"].top_k.tolist() == [32, 32]
+    # Preserve the established true-top-k recipe, not merely a top-k label.
+    assert execution.decode_calls[2]["sampling_params"].top_p.tolist() == pytest.approx([0.08, 0.08])
+
+
+def test_q128_single_topk_primes_all_tile_ends():
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, *_ = make_coordinator(
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=32,
+    )
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
+
+    topk_calls = [
+        call
+        for call in execution.prefill_calls
+        if call["sampling_params"] is not None
+        and float(call["sampling_params"].temperature[0]) == 1.0
+        and call["start_pos"] is None
+    ]
+    assert [int(call["prompt_lens"][0]) for call in topk_calls] == [32, 64, 96, 128]
+    assert [int(call["tokens"].shape[-1]) for call in topk_calls] == [32, 64, 96, 128]
+    argmax_calls = [
+        call
+        for call in execution.prefill_calls
+        if call["sampling_params"] is not None
+        and float(call["sampling_params"].temperature[0]) == 0.0
+        and call["start_pos"] is None
+    ]
+    assert [int(call["prompt_lens"][0]) for call in argmax_calls] == [32, 64, 96, 128]
+
+
+def test_decode_warmup_uses_topk_as_the_platform_greedy_path_when_argmax_is_disabled():
+    coordinator, execution, *_ = make_coordinator(
+        sequence_lengths=(128,),
+        lane_capacity=2,
+        allow_force_argmax=False,
+    )
+
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=2,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    assert execution.decode_calls[0]["sampling_params"] is None
+    assert execution.decode_calls[1]["sampling_params"].top_k.tolist() == [32, 32]
+
+
+def test_eager_and_trace_coverage_are_separately_idempotent():
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, trace_compiler, *_ = make_coordinator(
+        warmup_config=config, sequence_lengths=(128,), lane_capacity=1
+    )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+    eager_calls = len(execution.prefill_calls)
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+    assert len(execution.prefill_calls) == eager_calls
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
+    trace_calls = len(execution.prefill_calls)
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
+    assert len(execution.prefill_calls) == trace_calls
+    assert trace_compiler.calls == 0
+
+
+def test_trace_warmup_routes_cached_prefill_through_traced_execution_target():
+    coordinator, eager, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        sampling=False,
+    )
+    traced = RecordingExecution()
+    coordinator.execution = traced
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+
+    assert not eager.prefill_calls
+    assert any(call["start_pos"] is None for call in traced.prefill_calls)
+    assert any(call["start_pos"] is not None for call in traced.prefill_calls)
+
+
+def test_coverage_manifest_uses_compiler_registries_and_deduplicates_trace_identities():
+    eager_program = CompiledProgram(ProgramKey("0" * 64), "eager", OutputSpec((1,), torch.float32))
+    first_traced = CompiledProgram(ProgramKey("1" * 64), "traced-a", OutputSpec((1,), torch.float32))
+    second_traced = CompiledProgram(ProgramKey("2" * 64), "traced-b", OutputSpec((1,), torch.float32))
+    shared_trace_key = ProgramKey("a" * 64)
+    program_compiler = SimpleNamespace(compiled_programs=(eager_program, first_traced, second_traced))
+    eager = SimpleNamespace(program_compiler=program_compiler)
+    trace_compiler = SimpleNamespace(
+        trace_key_for_program=lambda key: None if key == eager_program.key else shared_trace_key,
+        get=lambda key: SimpleNamespace(signature="shared-trace") if key == shared_trace_key else None,
+    )
+
+    manifest = _resolve_coverage_manifest(eager, trace_compiler)
+
+    assert manifest.eager_program_signatures == ("eager",)
+    assert manifest.traced_source_program_signatures == ("traced-a", "traced-b")
+    assert manifest.trace_signatures == ("shared-trace",)
+    assert manifest.aliases == (
+        CoverageAlias("traced-a", "shared-trace"),
+        CoverageAlias("traced-b", "shared-trace"),
+    )
+
+
+def test_coverage_manifest_rejects_any_missing_required_trace_alias(expect_error):
+    first_traced = CompiledProgram(ProgramKey("1" * 64), "traced-a", OutputSpec((1,), torch.float32))
+    second_traced = CompiledProgram(ProgramKey("2" * 64), "traced-b", OutputSpec((1,), torch.float32))
+    trace_key = ProgramKey("a" * 64)
+    eager = SimpleNamespace(program_compiler=SimpleNamespace(compiled_programs=(first_traced, second_traced)))
+    trace_compiler = SimpleNamespace(
+        trace_key_for_program=lambda key: trace_key if key == first_traced.key else None,
+        get=lambda key: SimpleNamespace(signature="trace-a") if key == trace_key else None,
+    )
+
+    with expect_error(RuntimeError, "required trace alias"):
+        _resolve_coverage_manifest(
+            eager,
+            trace_compiler,
+            required_program_keys={first_traced.key, second_traced.key},
+            required_trace_program_keys={first_traced.key, second_traced.key},
+        )
+
+
+def test_activation_validates_every_program_returned_by_trace_warmup(expect_error):
+    coordinator, execution, trace_compiler, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        sampling=False,
+    )
+    programs = tuple(
+        CompiledProgram(ProgramKey(str(index) * 64), f"program-{index}", OutputSpec((1,), torch.float32))
+        for index in range(1, 4)
+    )
+    execution.program_compiler = SimpleNamespace(compiled_programs=programs)
+    prefill_programs = iter(programs[:2])
+    execution.compile_prefill = lambda **_kwargs: (next(prefill_programs),)
+    execution.compile_decode = lambda **_kwargs: programs[2]
+    trace_keys = {
+        programs[0].key: ProgramKey("a" * 64),
+        programs[2].key: ProgramKey("c" * 64),
+    }
+    trace_compiler.trace_key_for_program = trace_keys.get
+    trace_compiler.get = lambda key: SimpleNamespace(signature=f"trace-{key.digest[0]}")
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+    with expect_error(RuntimeError, programs[1].key.digest):
+        coordinator.warmup_decode(
+            kv_cache="cache",
+            enable_trace=True,
+            max_batch_size=1,
+            num_blocks=8,
+            can_sample_on_device=False,
+        )
+
+    assert trace_compiler.calls == 0
+
+
+@pytest.mark.parametrize("order", [("prefill", "decode"), ("decode", "prefill")])
+def test_prefill_decode_order_is_independent_and_capture_waits_for_both(order):
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, trace_compiler, *_ = make_coordinator(
+        warmup_config=config, sequence_lengths=(128,), lane_capacity=1
+    )
+
+    def run(operation):
+        if operation == "prefill":
+            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
+        else:
+            coordinator.warmup_decode(
+                kv_cache="cache",
+                enable_trace=True,
+                max_batch_size=1,
+                num_blocks=8,
+                can_sample_on_device=True,
+            )
+
+    run(order[0])
+    assert trace_compiler.calls == 0
+    run(order[1])
+    assert trace_compiler.calls == 1
+    run(order[0])
+    run(order[1])
+    assert trace_compiler.calls == 1
+
+
+@pytest.mark.parametrize("order", [("prefill", "decode"), ("decode", "prefill")])
+def test_capture_uses_phase_specific_sampling_decisions(order):
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, trace_compiler, *_ = make_coordinator(
+        trace_mode="all",
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=1,
+    )
+
+    def run(operation):
+        if operation == "prefill":
+            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+        else:
+            coordinator.warmup_decode(
+                kv_cache="cache",
+                enable_trace=True,
+                max_batch_size=1,
+                num_blocks=8,
+                can_sample_on_device=True,
+            )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+    with coordinator.defer_capture():
+        run(order[0])
+        assert trace_compiler.calls == 0
+        run(order[1])
+        assert coordinator.capture_pending
+        coordinator.activate_pending_capture()
+
+    assert trace_compiler.calls == 1
+    assert coordinator.already_warmed_up_prefill
+    assert all(call["sampling_params"] is None for call in execution.prefill_calls)
+    assert any(call["sampling_params"] is not None for call in execution.decode_calls)
+    assert not execution.prefill_replays
+
+
+def test_capture_deferral_stages_complete_registration_until_explicit_activation():
+    coordinator, _, trace_compiler, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        sampling=False,
+    )
+
+    with coordinator.defer_capture():
+        coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+        coordinator.warmup_decode(
+            kv_cache="cache",
+            enable_trace=True,
+            max_batch_size=1,
+            num_blocks=8,
+            can_sample_on_device=False,
+        )
+        assert coordinator.capture_pending
+        assert not coordinator.trace_activated
+        assert trace_compiler.calls == 0
+        coordinator.activate_pending_capture()
+        assert coordinator.trace_activated
+        assert trace_compiler.calls == 1
+
+    assert not coordinator.capture_pending
+
+
+def test_capture_deferral_exception_discards_pending_activation(expect_error):
+    coordinator, _, trace_compiler, *_ = make_coordinator(
+        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        sampling=False,
+    )
+
+    with expect_error(RuntimeError, "staging failed"):
+        with coordinator.defer_capture():
+            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+            coordinator.warmup_decode(
+                kv_cache="cache",
+                enable_trace=True,
+                max_batch_size=1,
+                num_blocks=8,
+                can_sample_on_device=False,
+            )
+            assert coordinator.capture_pending
+            raise RuntimeError("staging failed")
+
+    assert not coordinator.capture_pending
+    assert not coordinator.trace_activated
+    assert trace_compiler.calls == 0
+
+
+def test_static_all_can_capture_decode_only_runtime_trace():
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, trace_compiler, *_ = make_coordinator(
+        trace_mode="all",
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=1,
+    )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
+    assert coordinator.already_warmed_up_prefill
+    assert trace_compiler.calls == 0
+
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=True,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    assert trace_compiler.calls == 1
+    assert not execution.prefill_replays
+
+
+def test_two_phase_static_all_waits_for_phase_two_decode_before_capture():
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, _, trace_compiler, *_ = make_coordinator(
+        trace_mode="all",
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=1,
+    )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
+    assert trace_compiler.calls == 0
+
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=True,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+
+    assert trace_compiler.calls == 1
+
+
+def test_sampling_buffers_are_materialized_before_first_compile_and_capture():
+    events = []
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, _, _, _, _, events = make_coordinator(
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        events=events,
+    )
+
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=True,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
+
+    assert events.index("sampling") < events.index("compile_decode")
+    assert events.index("sampling") < events.index("compile_prefill")
+    assert max(index for index, event in enumerate(events) if event.startswith("compile_")) < events.index("capture")
+
+
+def test_failed_case_is_not_marked_complete_and_retry_skips_completed_case(expect_error):
+    execution = RecordingExecution()
+    execution.fail_decode_call = 2
+    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
+    coordinator, execution, *_ = make_coordinator(
+        trace_mode="none",
+        sampling=True,
+        warmup_config=config,
+        sequence_lengths=(128,),
+        lane_capacity=1,
+        execution=execution,
+    )
+
+    with expect_error(RuntimeError, "decode compile failed"):
+        coordinator.warmup_decode(
+            kv_cache="cache",
+            enable_trace=False,
+            max_batch_size=1,
+            num_blocks=8,
+            can_sample_on_device=True,
+        )
+    assert len(execution.decode_calls) == 1
+
+    coordinator.warmup_decode(
+        kv_cache="cache",
+        enable_trace=False,
+        max_batch_size=1,
+        num_blocks=8,
+        can_sample_on_device=True,
+    )
+    assert len(execution.decode_calls) == 2
+
+
+def test_dynamic_hints_cannot_expand_static_trace_or_sampling_ceilings(expect_error):
+    coordinator, *_ = make_coordinator(trace_mode="decode_only", sampling=False)
+
+    with expect_error(ValueError, "prefill trace warmup exceeds"):
+        coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
+    with expect_error(ValueError, "statically disabled"):
+        coordinator.warmup_decode(
+            kv_cache="cache",
+            enable_trace=False,
+            max_batch_size=4,
+            num_blocks=8,
+            can_sample_on_device=True,
+        )
+
+    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)

code/models/common/tests/models/deepseek_r1_distill_qwen_14b/test_demo_contract.py

@@ -0,0 +1,452 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from models.common.llm_runtime.config import TraceConfig
+from models.common.llm_runtime.prefill.plan import _plan_prefill_requests
+
+_DEMO_PATH = "models/common/tests/demos/deepseek_r1_distill_qwen_14b/demo.py"
+_DEMO_TREE = ast.parse(Path(_DEMO_PATH).read_text(encoding="utf-8"), filename=_DEMO_PATH)
+
+
+def _demo_function(name, namespace=None):
+    function = next(node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == name)
+    namespace = {} if namespace is None else namespace
+    exec(compile(ast.Module(body=[function], type_ignores=[]), _DEMO_PATH, "exec"), namespace)
+    return namespace[name]
+
+
+def _called_names(function_name):
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == function_name
+    )
+    return [
+        node.func.id for node in ast.walk(function) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    ]
+
+
+def test_demo_case_manifest_is_preserved():
+    test_function = next(
+        node
+        for node in _DEMO_TREE.body
+        if isinstance(node, ast.FunctionDef) and node.name == "test_deepseek_r1_qwen_14b"
+    )
+    decorators = [node for node in test_function.decorator_list if isinstance(node, ast.Call)]
+    test_config = next(node for node in decorators if ast.literal_eval(node.args[0]) == "test_config")
+    optimizations = next(node for node in decorators if ast.literal_eval(node.args[0]) == "optimizations")
+    case_ids = [ast.literal_eval(element.keywords[0].value) for element in test_config.args[1].elts]
+    assert case_ids == [
+        "token-accuracy",
+        "batch-1",
+        "batch-32",
+        "batch-32-ci",
+        "eval-32",
+        "ci-b1-DP-2",
+        "ci-b1-DP-4",
+        "ci-b1-DP-8",
+        "ci-b1-DP-16",
+        "ci-b1-DP-32",
+    ]
+    assert ast.literal_eval(optimizations.args[1]) == ["performance", "accuracy"]
+
+
+def test_demo_reserves_trace_space_by_mesh(monkeypatch):
+    for mesh_name, mesh_shape, trace_region_size in (
+        ("N300", (1, 2), 50_000_000),
+        ("T3K", (1, 8), 100_000_000),
+    ):
+        monkeypatch.setenv("MESH_DEVICE", mesh_name)
+        device_params = _demo_function(
+            "_ttnn_mesh_device_param_from_env",
+            {
+                "os": os,
+                "pytest": pytest,
+                "_MESH_DEVICE_TO_SHAPE": {mesh_name: mesh_shape},
+                "ttnn": SimpleNamespace(FabricConfig=SimpleNamespace(FABRIC_1D=object())),
+            },
+        )()
+
+        assert device_params["mesh_shape"] == mesh_shape
+        assert device_params["trace_region_size"] == trace_region_size
+
+
+def test_demo_warmup_compiles_eager_programs_before_trace_capture():
+    calls = []
+    config = SimpleNamespace(trace=TraceConfig("all"), device_sampling_enabled=True)
+    executor = SimpleNamespace(
+        config=config,
+        model=SimpleNamespace(config=SimpleNamespace(max_batch_size=4)),
+        warmup_model_prefill=lambda **kwargs: calls.append(("prefill", kwargs)),
+        warmup_model_decode=lambda **kwargs: calls.append(("decode", kwargs)),
+    )
+    warmup = _demo_function("_warmup_demo_executor")
+    kv_cache = object()
+    warmup(executor, kv_cache=kv_cache, page_table=SimpleNamespace(shape=(4, 8)))
+
+    assert [(kind, kwargs["enable_trace"]) for kind, kwargs in calls] == [
+        ("decode", False),
+        ("prefill", False),
+        ("prefill", True),
+        ("decode", True),
+    ]
+    assert all(kwargs["kv_cache"] is kv_cache for _, kwargs in calls)
+
+
+def test_demo_warmup_registers_concrete_prefill_before_trace_capture():
+    calls = []
+    config = SimpleNamespace(trace=TraceConfig("all"), device_sampling_enabled=False)
+    eager_execution = object()
+    executor = SimpleNamespace(
+        config=config,
+        eager_execution=eager_execution,
+        model=SimpleNamespace(config=SimpleNamespace(max_batch_size=32)),
+        warmup_model_prefill=lambda **kwargs: calls.append(("prefill", kwargs)),
+        warmup_model_decode=lambda **kwargs: calls.append(("decode", kwargs)),
+        compile_prefill=lambda **kwargs: calls.append(("compile_prefill", kwargs)),
+    )
+    warmup = _demo_function("_warmup_demo_executor")
+    tokens = torch.zeros((32, 700), dtype=torch.long)
+    prompt_lens = torch.tensor([64] * 30 + [400, 700])
+    page_table = torch.zeros((32, 64), dtype=torch.int32)
+    kv_cache = object()
+
+    warmup(
+        executor,
+        kv_cache=kv_cache,
+        page_table=page_table,
+        prefill_compile_case=(tokens, prompt_lens),
+    )
+
+    assert [(kind, kwargs.get("enable_trace")) for kind, kwargs in calls] == [
+        ("decode", False),
+        ("prefill", False),
+        ("compile_prefill", None),
+        ("prefill", True),
+        ("decode", True),
+    ]
+    compile_kwargs = calls[2][1]
+    assert compile_kwargs["tokens"] is tokens
+    assert compile_kwargs["prompt_lens"] is prompt_lens
+    assert compile_kwargs["page_table"] is page_table
+    assert compile_kwargs["kv_cache"] is kv_cache
+    assert compile_kwargs["empty_slots"] == list(range(32))
+    assert compile_kwargs["execution"] is eager_execution
+
+
+def test_demo_warmup_uses_lane_group_capacity_and_lane_trace_policy():
+    calls = []
+    lane_config = SimpleNamespace(trace=TraceConfig("all"), device_sampling_enabled=True)
+    group = SimpleNamespace(
+        lanes=[SimpleNamespace(config=lane_config) for _ in range(4)],
+        max_batch_size=4,
+        warmup_model_prefill=lambda **kwargs: calls.append(("prefill", kwargs)),
+        warmup_model_decode=lambda **kwargs: calls.append(("decode", kwargs)),
+    )
+    warmup = _demo_function("_warmup_demo_executor")
+    kv_cache = [object() for _ in range(4)]
+    warmup(group, kv_cache=kv_cache, page_table=SimpleNamespace(shape=(4, 128)))
+
+    decode_calls = [kwargs for kind, kwargs in calls if kind == "decode"]
+    assert len(decode_calls) == 2
+    assert all(kwargs["max_batch_size"] == 4 for kwargs in decode_calls)
+    assert all(kwargs["num_blocks"] == 128 for kwargs in decode_calls)
+    assert all(kwargs["kv_cache"] is kv_cache for _, kwargs in calls)
+
+
+def test_eval_prefill_signature_multiset_is_rotation_invariant_and_not_static_warmup_shaped():
+    tokens = torch.zeros((32, 700), dtype=torch.long)
+    prompt_lens = torch.tensor([64] * 30 + [400, 700])
+    page_table = torch.zeros((32, 64), dtype=torch.int32)
+
+    def planned_shapes(offset):
+        rotated_tokens = torch.roll(tokens, shifts=-offset, dims=0)
+        rotated_lens = torch.roll(prompt_lens, shifts=-offset, dims=0)
+        requests = _plan_prefill_requests(
+            tokens=rotated_tokens,
+            page_table=page_table,
+            prompt_lens=rotated_lens,
+            empty_slots=list(range(32)),
+            start_pos=None,
+            block_size=32,
+            max_batch_size=32,
+            max_prefill_chunk_size=1024,
+            supports_batched_prefill=True,
+            max_prefill_batch_size=8,
+            max_actual_page_table_width=32,
+            canonical_page_table_width=64,
+        )
+        return sorted(
+            (request.padded_sequence_length, request.padded_batch_size, len(request.source_rows))
+            for request in requests
+        )
+
+    expected = [(128, 8, 6), (128, 8, 8), (128, 8, 8), (128, 8, 8), (1024, 2, 2)]
+    assert planned_shapes(0) == expected
+    assert planned_shapes(1) == expected
+    assert planned_shapes(2) == expected
+
+
+@pytest.mark.parametrize("function_name", ["_run_perf_benchmark", "_run_eval_repeat_batch32"])
+def test_traced_demo_paths_warm_up_fresh_executor(function_name):
+    assert "_warmup_demo_executor" in _called_names(function_name)
+
+
+def test_create_executor_uses_model_owned_executor_and_resolved_cache():
+    captured = {}
+
+    def executor_config(**kwargs):
+        captured.update(kwargs)
+        return SimpleNamespace(**kwargs)
+
+    namespace = {
+        "DeepSeekR1Qwen14B": object,
+        "DeepSeekR1Qwen14BExecutor": lambda model, runtime_config, config: config,
+        "DeepSeekR1Qwen14BExecutorConfig": executor_config,
+        "PagedKVCacheConfig": lambda **kwargs: SimpleNamespace(**kwargs),
+        "TraceConfig": TraceConfig,
+        "WarmupConfig": lambda: object(),
+    }
+    create_executor = _demo_function("create_executor", namespace)
+    model = SimpleNamespace(
+        model_args=object(),
+        config=SimpleNamespace(
+            max_seq_len=2048,
+            max_batch_size=32,
+            block_configs=[SimpleNamespace(attention_config=SimpleNamespace(kv_cache_dtype=object()))],
+        ),
+    )
+
+    result = create_executor(model, traced=True, device_sampling_enabled=True)
+
+    assert result.trace.mode == "all"
+    assert result.device_sampling_enabled is True
+    assert captured["paged_kv_cache"].num_blocks == 2048
+
+
+def test_eval_uses_decode_only_trace_while_ordinary_traced_executor_uses_all():
+    create_executor = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "create_executor"
+    )
+    trace_config = next(
+        node
+        for node in ast.walk(create_executor)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "TraceConfig"
+    )
+    assert isinstance(trace_config.keywords[0].value, ast.Name)
+    assert trace_config.keywords[0].value.id == "trace_mode"
+    derived_mode = next(
+        node
+        for node in ast.walk(create_executor)
+        if isinstance(node, ast.Assign)
+        and any(isinstance(target, ast.Name) and target.id == "trace_mode" for target in node.targets)
+    )
+    assert ast.unparse(derived_mode.value) == "'all' if traced else 'none'"
+
+    eval_function = next(
+        node
+        for node in _DEMO_TREE.body
+        if isinstance(node, ast.FunctionDef) and node.name == "_run_eval_repeat_batch32"
+    )
+    eval_create = next(
+        node
+        for node in ast.walk(eval_function)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "create_executor"
+    )
+    keywords = {keyword.arg: keyword.value for keyword in eval_create.keywords}
+    assert ast.literal_eval(keywords["traced"]) is True
+    assert ast.literal_eval(keywords["trace_mode"]) == "decode_only"
+
+
+def test_deepseek_stop_guard_truncates_eos_but_not_ordinary_reasoning_tokens(expect_error, monkeypatch):
+    shared_calls = []
+
+    def shared_guard(generated_token_ids, tokenizer, **kwargs):
+        shared_calls.append((generated_token_ids, kwargs))
+        if kwargs["is_ci_env"] is None and os.environ.get("TT_DEMO_STRICT_SPECIAL_TOKENS") == "1":
+            outputs_before_eos = [
+                output[: output.index(tokenizer.eos_token_id)] if tokenizer.eos_token_id in output else output
+                for output in generated_token_ids
+            ]
+            if any(99 in output for output in outputs_before_eos):
+                raise AssertionError("model produced special tokens")
+
+    guard = _demo_function("assert_no_special_tokens", {"assert_no_special_tokens_shared": shared_guard})
+    tokenizer = SimpleNamespace(
+        all_special_ids=[10, 99],
+        eos_token_id=10,
+    )
+
+    monkeypatch.setenv("TT_DEMO_STRICT_SPECIAL_TOKENS", "1")
+    guard([[1, 10, 99], [2, 3, 4]], tokenizer)
+    assert shared_calls[-1][0] == [[1], [2, 3, 4]]
+    with expect_error(AssertionError, "model produced special tokens"):
+        guard([[1, 99]], tokenizer)
+
+
+def test_dp_smoke_uses_model_owned_lane_group_execution():
+    calls = _called_names("_run_dp_smoke")
+    assert "_dp_lane_tp_or_skip" in calls
+    assert "_create_dp_submeshes" in calls
+    assert "create_executor" in calls
+    assert "LaneGroupExecutor" in calls
+    assert "run_perf_benchmark" in calls
+    assert "cleanup_dp_model_case" in calls
+    assert "_skip_below_min_tp_devices" not in calls
+
+
+def test_runnable_dp_lane_build_errors_are_not_converted_to_topology_skips():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_dp_smoke"
+    )
+    pytest_skip_calls = [
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "pytest"
+        and node.func.attr == "skip"
+    ]
+    assert pytest_skip_calls == []
+
+
+def test_deepseek_dp_topology_accepts_t3k_dp2_tp4_and_dp4_tp2(expect_error):
+    topology = _demo_function(
+        "_dp_lane_tp_or_skip",
+        {"ttnn": SimpleNamespace(MeshDevice=object), "pytest": pytest, "_MIN_TP_DEVICES": 2},
+    )
+    t3k = SimpleNamespace(get_num_devices=lambda: 8)
+
+    assert topology(t3k, 2) == 4
+    assert topology(t3k, 4) == 2
+    with expect_error(pytest.skip.Exception, "DP-8 on 8 devices creates TP1 lanes"):
+        topology(t3k, 8)
+    with expect_error(pytest.skip.Exception, "DP-16 cannot partition 8 devices"):
+        topology(t3k, 16)
+
+
+def test_deepseek_dp4_partitions_four_tp2_submeshes():
+    calls = []
+    submeshes = [object() for _ in range(4)]
+    parent = SimpleNamespace(
+        create_submeshes=lambda shape: calls.append(shape) or submeshes,
+    )
+    fake_ttnn = SimpleNamespace(MeshDevice=object, MeshShape=lambda rows, columns: (rows, columns))
+    create_submeshes = _demo_function("_create_dp_submeshes", {"ttnn": fake_ttnn})
+
+    assert create_submeshes(parent, 4, 2) == submeshes
+    assert calls == [(1, 2)]
+
+
+def test_deepseek_dp_lane_cache_reuses_lane_topology(tmp_path):
+    cache_dir = tmp_path / "DeepSeek-R1-Distill-Qwen-14B" / "T3K"
+    cache_dir.mkdir(parents=True)
+    lane_cache_dir = _demo_function("_dp_lane_cache_dir", {"Path": Path})(cache_dir, 2)
+
+    assert lane_cache_dir == cache_dir.parent / "N300"
+    assert lane_cache_dir.is_dir()
+    assert _demo_function("_dp_lane_cache_dir", {"Path": Path})(cache_dir, 4) == cache_dir.parent / "N150x4"
+
+
+def test_deepseek_dp_lane_contract_checks_heads_capacity_and_cache(expect_error):
+    validate = _demo_function(
+        "_validate_dp_lane",
+        {
+            "DeepSeekR1Qwen14B": object,
+            "DeepSeekR1Qwen14BExecutor": object,
+            "math": __import__("math"),
+        },
+    )
+    attention = SimpleNamespace(n_heads=40, n_kv_heads=8)
+    model = SimpleNamespace(
+        config=SimpleNamespace(
+            num_devices=2,
+            max_batch_size=1,
+            block_configs=[SimpleNamespace(attention_config=attention)],
+        )
+    )
+    cache = SimpleNamespace(max_num_blocks=128, num_blocks=128)
+    lane = SimpleNamespace(config=SimpleNamespace(paged_kv_cache=cache))
+
+    validate(model, lane, 2, 4096)
+    model.config.num_devices = 4
+    with expect_error(ValueError, "expected TP2, model uses TP4"):
+        validate(model, lane, 2, 4096)
+    model.config.num_devices = 2
+    model.config.max_batch_size = 2
+    with expect_error(ValueError, "capacity 1"):
+        validate(model, lane, 2, 4096)
+    model.config.max_batch_size = 1
+    cache.num_blocks = None
+    with expect_error(ValueError, "cache must contain 128 blocks"):
+        validate(model, lane, 2, 4096)
+
+
+def test_token_accuracy_cleans_up_executor_in_finally():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_token_accuracy"
+    )
+    cleanup_calls = [
+        statement
+        for node in ast.walk(function)
+        if isinstance(node, ast.Try)
+        for statement in node.finalbody
+        if isinstance(statement, ast.Expr)
+        and isinstance(statement.value, ast.Call)
+        and isinstance(statement.value.func, ast.Attribute)
+        and statement.value.func.attr == "cleanup"
+    ]
+    assert len(cleanup_calls) == 1
+
+
+def test_main_demo_does_not_synchronize_parent_mesh_after_prebuild_skip():
+    function = next(
+        node
+        for node in _DEMO_TREE.body
+        if isinstance(node, ast.FunctionDef) and node.name == "test_deepseek_r1_qwen_14b"
+    )
+    try_node = next(node for node in function.body if isinstance(node, ast.Try))
+
+    assert len(try_node.finalbody) == 1
+    guard = try_node.finalbody[0]
+    assert isinstance(guard, ast.If)
+    assert ast.unparse(guard.test) == "model is not None"
+    assert any(
+        isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "cleanup_model_case"
+        for node in ast.walk(guard)
+    )
+
+
+@pytest.mark.parametrize("function_name", ["_run_token_accuracy", "_run_perf_benchmark", "_run_eval_repeat_batch32"])
+def test_demo_reads_model_geometry_from_model_config(function_name):
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == function_name
+    )
+    model_args_aliases = [
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.Assign)
+        and isinstance(node.value, ast.Attribute)
+        and isinstance(node.value.value, ast.Name)
+        and node.value.value.id == "model"
+        and node.value.attr == "model_args"
+    ]
+    config_fields = {
+        node.attr
+        for node in ast.walk(function)
+        if isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Attribute)
+        and isinstance(node.value.value, ast.Name)
+        and node.value.value.id == "model"
+        and node.value.attr == "config"
+    }
+
+    assert model_args_aliases == []
+    assert {"max_batch_size", "max_seq_len"} <= config_fields

code/models/common/tests/models/deepseek_r1_distill_qwen_14b/test_hf_adaptor.py

@@ -0,0 +1,290 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from types import SimpleNamespace
+
+import torch
+from transformers import Qwen2Config, Qwen2ForCausalLM
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RotaryEmbedding
+
+from models.common.models.deepseek_r1_distill_qwen_14b import generator, hf_adaptor
+from models.common.models.deepseek_r1_distill_qwen_14b import model as qwen_model
+from models.common.models.deepseek_r1_distill_qwen_14b import weight_utils
+from models.common.models.deepseek_r1_distill_qwen_14b.hf_adaptor import DeepSeekR1Qwen14BForCausalLM as DeepSeekProduct
+from models.common.models.deepseek_r1_distill_qwen_14b.hf_adaptor import (
+    DeepSeekR1Qwen14BRuntimeConfig,
+    _trace_seq_lens,
+    convert_hf_model_weights,
+)
+
+
+def test_runtime_config_preserves_tp2_trace_and_batched_prefill_policy():
+    runtime = DeepSeekR1Qwen14BRuntimeConfig(
+        model_name="DeepSeek-R1-Distill-Qwen-14B",
+        model_cache_path=None,
+        max_prefill_chunk_size=2048,
+        max_context_len=32768,
+        max_seq_len=4096,
+        trace_prefill_supported_seq_lens=(128, 1024),
+    )
+    assert runtime.can_enable_trace(128, num_cached_tokens=32)
+    assert runtime.can_enable_trace(1024)
+    assert not runtime.can_enable_trace(2048)
+    assert runtime.supports_batched_prefill
+    assert runtime.max_prefill_batch_size == 32
+    assert runtime.batched_prefill_batched_extract
+    assert _trace_seq_lens(2, 2048, 4096) == (128, 1024)
+    assert _trace_seq_lens(4, 2048, 4096) == (128,)
+    assert _trace_seq_lens(8, 2048, 4096) == (128, 1024)
+
+
+def test_pinned_revision_is_the_provider_and_generator_default():
+    expected = "1df8507178afcc1bef68cd8c393f61a886323761"
+    assert hf_adaptor.DEFAULT_HF_REVISION == expected
+    assert generator.DeepSeekR1Qwen14BGeneratorConfig.__dataclass_fields__["hf_revision"].default == expected
+
+
+def test_generator_keeps_deepseek_chat_template_enabled():
+    source = inspect.getsource(generator.build_deepseek_r1_distill_qwen_14b_generator)
+    assert "instruct=True" in source
+
+
+def test_provider_rejects_below_capacity_before_loading_hf(expect_error):
+    mesh = SimpleNamespace(get_num_devices=lambda: 1)
+    with expect_error(ValueError, "supports logical TP2/TP4/TP8"):
+        hf_adaptor.from_pretrained(mesh)
+
+
+def test_product_binds_runtime_config_and_stop_tokens():
+    model = SimpleNamespace(config=SimpleNamespace(max_seq_len=4096), model_args=None)
+    tokenizer = SimpleNamespace(stop_tokens=[151643, 151644])
+    runtime = DeepSeekR1Qwen14BRuntimeConfig(
+        model_name="model",
+        model_cache_path=None,
+        max_prefill_chunk_size=2048,
+        max_context_len=32768,
+        max_seq_len=4096,
+        trace_prefill_supported_seq_lens=(128, 1024),
+    )
+    product = DeepSeekProduct(model=model, tokenizer=tokenizer, runtime_config=runtime)
+    assert model.model_args is runtime
+    assert product.generation_config.stop_token_ids == (151643, 151644)
+    assert product.max_seq_len == 4096
+    assert product.max_context_len == 32768
+
+
+def test_tokenizer_adds_eos_and_threads_revision(monkeypatch):
+    tokenizer = SimpleNamespace(
+        eos_token_id=151643,
+        convert_tokens_to_ids=lambda token: -1,
+    )
+    seen = {}
+
+    def fake_from_pretrained(model, **kwargs):
+        seen.update(model=model, **kwargs)
+        return tokenizer
+
+    monkeypatch.setattr(hf_adaptor.AutoTokenizer, "from_pretrained", fake_from_pretrained)
+    assert hf_adaptor.load_tokenizer("deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", "revision") is tokenizer
+    assert tokenizer.stop_tokens == [151643]
+    assert seen["revision"] == "revision"
+
+
+def test_qkv_weights_and_bias_use_reverse_permutation_and_device_major_packing():
+    hidden_size = 16
+    n_heads = 4
+    n_kv_heads = 2
+    head_dim = 4
+    num_devices = 2
+    kv_width = n_kv_heads * head_dim
+    q = torch.arange(hidden_size * hidden_size, dtype=torch.float32).reshape(hidden_size, hidden_size)
+    k = torch.arange(kv_width * hidden_size, dtype=torch.float32).reshape(kv_width, hidden_size) + 10_000
+    v = k + 10_000
+    o = q + 30_000
+    bq = torch.arange(hidden_size, dtype=torch.float32)
+    bk = torch.arange(kv_width, dtype=torch.float32) + 100
+    bv = torch.arange(kv_width, dtype=torch.float32) + 200
+    attention = SimpleNamespace(
+        config=SimpleNamespace(
+            hidden_size=hidden_size,
+            num_attention_heads=n_heads,
+            num_key_value_heads=n_kv_heads,
+        ),
+        q_proj=SimpleNamespace(weight=q, bias=bq),
+        k_proj=SimpleNamespace(weight=k, bias=bk),
+        v_proj=SimpleNamespace(weight=v, bias=bv),
+        o_proj=SimpleNamespace(weight=o),
+    )
+
+    wqkv, wo, q_norm, k_norm, bias = weight_utils.attention_wqkv_wo_from_hf_layer(attention, num_devices)
+    q_meta = weight_utils.reverse_permute(q, n_heads, hidden_size, hidden_size).T
+    k_meta = weight_utils.reverse_permute(k, n_kv_heads, kv_width, hidden_size).T
+    bq_meta = weight_utils.reverse_permute_1d(bq.view(n_heads, head_dim)).view(-1)
+    bk_meta = weight_utils.reverse_permute_1d(bk.view(n_kv_heads, head_dim)).view(-1)
+    expected_weights = (
+        torch.cat(
+            [
+                torch.cat(parts, dim=-1)
+                for parts in zip(
+                    torch.chunk(q_meta, num_devices, dim=1),
+                    torch.chunk(k_meta, num_devices, dim=1),
+                    torch.chunk(v.T, num_devices, dim=1),
+                )
+            ],
+            dim=-1,
+        )
+        .unsqueeze(0)
+        .unsqueeze(0)
+    )
+    expected_bias = torch.cat(
+        [
+            torch.cat(parts, dim=-1)
+            for parts in zip(
+                torch.chunk(bq_meta, num_devices),
+                torch.chunk(bk_meta, num_devices),
+                torch.chunk(bv, num_devices),
+            )
+        ]
+    )
+
+    torch.testing.assert_close(wqkv, expected_weights)
+    torch.testing.assert_close(wo, o.T.unsqueeze(0).unsqueeze(0))
+    torch.testing.assert_close(bias, expected_bias)
+    assert q_norm is None and k_norm is None
+
+
+def test_hf_rope_tables_preserve_plain_theta_one_million():
+    head_dim = 16
+    table_len = 128
+    config = Qwen2Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=32768,
+        rope_parameters={"rope_type": "default", "rope_theta": 1_000_000.0},
+    )
+    rotary = Qwen2RotaryEmbedding(config)
+    cos, sin = weight_utils.build_rope_cos_sin_torch(rotary, table_len, head_dim, torch.bfloat16)
+    x = torch.zeros(1, 1, table_len, head_dim, dtype=torch.bfloat16)
+    positions = torch.arange(table_len).unsqueeze(0)
+    with torch.no_grad():
+        hf_cos, hf_sin = rotary(x, positions)
+    expected_cos, expected_sin = weight_utils.permute_hf_rope_to_meta_tables(hf_cos.float(), hf_sin.float())
+    assert config.rope_parameters["rope_theta"] == 1_000_000.0
+    torch.testing.assert_close(cos, expected_cos.to(torch.bfloat16))
+    torch.testing.assert_close(sin, expected_sin.to(torch.bfloat16))
+
+
+def test_conversion_covers_qkv_bias_and_untied_lm_head():
+    config = Qwen2Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=128,
+        max_position_embeddings=32768,
+        rope_parameters={"rope_type": "default", "rope_theta": 1_000_000.0},
+        tie_word_embeddings=False,
+    )
+    hf = Qwen2ForCausalLM(config).eval()
+    weights = convert_hf_model_weights(
+        hf,
+        config,
+        n_layers=1,
+        num_devices=2,
+        rope_table_len=128,
+        head_dim=16,
+    )
+    layer = weights.layers[0]
+    assert layer.wqkv.shape == (1, 1, 64, 128)
+    assert layer.wqkv_bias.shape == (128,)
+    assert layer.wo.shape == (1, 1, 64, 64)
+    assert layer.w1.shape == layer.w3.shape == (64, 2048)
+    assert layer.w2.shape == (2048, 64)
+    assert torch.count_nonzero(layer.w1[:, 128:]) == 0
+    assert torch.count_nonzero(layer.w3[:, 128:]) == 0
+    assert torch.count_nonzero(layer.w2[128:, :]) == 0
+    torch.testing.assert_close(weights.lm_head, hf.lm_head.weight.detach().to(torch.bfloat16))
+    assert weights.lm_head.data_ptr() != weights.embedding.data_ptr()
+
+
+def test_config_builder_is_owned_by_model_module():
+    assert (
+        hf_adaptor.build_deepseek_r1_distill_qwen_14b_transformer_config
+        is qwen_model.build_deepseek_r1_distill_qwen_14b_transformer_config
+    )
+    assert qwen_model.build_deepseek_r1_distill_qwen_14b_transformer_config.__module__ == qwen_model.__name__
+
+
+def test_post_attention_norm_program_and_memory_use_same_mlp_grid(monkeypatch):
+    grid = SimpleNamespace(num_cores=28)
+    program = object()
+    memory = object()
+    captured = {}
+
+    monkeypatch.setattr(qwen_model, "get_padded_hidden_dim", lambda *_: 18944)
+    monkeypatch.setattr(qwen_model, "_dram_shard_core_grid_k_n", lambda *_: grid)
+    monkeypatch.setattr(
+        qwen_model,
+        "_create_sharded_norm_program_config",
+        lambda dim, selected_grid, rows, tile: captured.update(program=(dim, selected_grid, rows, tile)) or program,
+    )
+    monkeypatch.setattr(
+        qwen_model.ttnn,
+        "create_sharded_memory_config",
+        lambda shape, selected_grid, *args, **kwargs: captured.update(memory=(shape, selected_grid)) or memory,
+    )
+
+    assert qwen_model._post_attn_norm_decode_configs(
+        dim=3584,
+        hidden_dim=18944,
+        num_devices=2,
+        max_batch_size=32,
+    ) == (program, memory)
+    assert captured["program"] == (3584, grid, 32, 32)
+    assert captured["memory"] == ((32, 128), grid)
+
+
+def test_decoder_layer_prefill_calls_chunk_capable_attention_entrypoint(monkeypatch):
+    captured = {}
+    attention_output = object()
+    final_output = object()
+    attention = SimpleNamespace(
+        prefill_forward=lambda x, rot_mats, **kwargs: captured.update(attention=(x, rot_mats, kwargs))
+        or attention_output
+    )
+    layer = qwen_model.DeepSeekR1Qwen14BDecoderLayer(
+        input_layernorm=SimpleNamespace(prefill_forward=lambda x: x),
+        self_attn=attention,
+        post_attention_layernorm=SimpleNamespace(prefill_forward=lambda x: x),
+        mlp=SimpleNamespace(prefill_forward=lambda x: x),
+    )
+    monkeypatch.setattr(qwen_model, "_all_gather_rmsnorm_tensor", lambda _norm, x: x)
+    monkeypatch.setattr(
+        qwen_model.ttnn,
+        "add",
+        lambda *_args, **_kwargs: final_output,
+    )
+
+    chunk_start_idx_tensor = object()
+    rot_mats = (object(), object())
+    assert (
+        layer.prefill_forward(
+            object(),
+            rot_mats,
+            user_id=[0, 1],
+            page_table=object(),
+            chunk_page_table=object(),
+            chunk_start_idx=128,
+            batch_size=2,
+            chunk_start_idx_tensor=chunk_start_idx_tensor,
+        )
+        is final_output
+    )
+    assert captured["attention"][1] is rot_mats
+    assert captured["attention"][2]["chunk_start_idx_tensor"] is chunk_start_idx_tensor
+    assert captured["attention"][2]["batch_size"] == 2

code/models/common/tests/models/deepseek_r1_distill_qwen_14b/test_prefill_last_token_contract.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+from models.common.models.deepseek_r1_distill_qwen_14b import model as qwen_model
+
+
+def test_prefill_runtime_slice_and_index_override_full_hidden_state_return(monkeypatch):
+    calls = []
+    hidden = SimpleNamespace(shape=(1, 1, 128, 640), dtype=qwen_model.ttnn.bfloat16)
+    sliced = SimpleNamespace(dtype=qwen_model.ttnn.bfloat16)
+    selected = object()
+    selected_4d = object()
+    logits = object()
+    slice_start = object()
+    slice_end = object()
+    last_token_index = object()
+    model = SimpleNamespace(
+        layers=[],
+        _last_tile_logits=lambda value: calls.append(("last_tile_logits", value)) or logits,
+    )
+
+    monkeypatch.setattr(
+        qwen_model.ttnn,
+        "slice",
+        lambda value, start, end, **kwargs: calls.append(("slice", value, start, end, kwargs)) or sliced,
+    )
+    monkeypatch.setattr(
+        qwen_model.ttnn,
+        "embedding",
+        lambda index, value, **kwargs: calls.append(("embedding", index, value, kwargs)) or selected,
+    )
+    monkeypatch.setattr(
+        qwen_model.ttnn,
+        "unsqueeze_to_4D",
+        lambda value: calls.append(("unsqueeze_to_4D", value)) or selected_4d,
+    )
+    monkeypatch.setattr(qwen_model.ttnn, "deallocate", lambda value: calls.append(("deallocate", value)))
+
+    result = qwen_model.DeepSeekR1Qwen14B.prefill_forward(
+        model,
+        hidden,
+        rot_mats=(object(), object()),
+        get_last_token=-1,
+        last_token_slice=(slice_start, slice_end),
+        last_token_index=last_token_index,
+    )
+
+    assert result is logits
+    assert calls == [
+        ("slice", hidden, slice_start, slice_end, {"slice_dim": 2, "num_devices": 4}),
+        ("deallocate", hidden),
+        ("embedding", last_token_index, sliced, {"layout": qwen_model.ttnn.TILE_LAYOUT}),
+        ("unsqueeze_to_4D", selected),
+        ("deallocate", sliced),
+        ("last_tile_logits", selected_4d),
+    ]
+
+
+def test_prefill_runtime_index_requires_runtime_slice(expect_error):
+    model = SimpleNamespace(layers=[])
+
+    with expect_error(ValueError, "last_token_index is required with a runtime last_token_slice"):
+        qwen_model.DeepSeekR1Qwen14B.prefill_forward(
+            model,
+            object(),
+            rot_mats=(object(), object()),
+            get_last_token=-1,
+            last_token_index=object(),
+        )

code/models/common/tests/models/llama32_1b/test_batched_prefill_postprocess.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+
+import models.common.models.llama32_1b.model as model_module
+from models.common.models.llama32_1b.model import Llama32_1BTransformer1D
+
+
+def test_batched_prefill_postprocess_gathers_each_slots_last_token_before_norm(monkeypatch):
+    calls = []
+    hidden = object()
+    selector_tt = object()
+    gathered = object()
+    normalized = object()
+    all_gathered = object()
+    logits = object()
+    output = object()
+    mesh = SimpleNamespace(arch=lambda: "wormhole")
+
+    class FakeTTNN:
+        bfloat16 = "bfloat16"
+        TILE_LAYOUT = "tile"
+        DRAM_MEMORY_CONFIG = "dram"
+        MathFidelity = SimpleNamespace(HiFi4="hifi4")
+
+        @staticmethod
+        def ReplicateTensorToMesh(device):
+            assert device is mesh
+            return "replicate"
+
+        @staticmethod
+        def from_torch(selector, **kwargs):
+            calls.append(("from_torch", selector.clone(), kwargs))
+            return selector_tt
+
+        @staticmethod
+        def init_device_compute_kernel_config(arch, **kwargs):
+            assert arch == "wormhole"
+            return kwargs
+
+        @staticmethod
+        def matmul(lhs, rhs, **kwargs):
+            calls.append(("matmul", lhs, rhs, kwargs))
+            return gathered
+
+        @staticmethod
+        def deallocate(tensor):
+            calls.append(("deallocate", tensor))
+
+        @staticmethod
+        def to_memory_config(tensor, memory_config):
+            calls.append(("to_memory_config", tensor, memory_config))
+            return output
+
+    fake_norm = SimpleNamespace(prefill_forward=lambda tensor: calls.append(("norm", tensor)) or normalized)
+    fake_lm_head = SimpleNamespace(
+        config=SimpleNamespace(input_memcfg=None),
+        forward=lambda tensor: calls.append(("lm_head", tensor)) or logits,
+    )
+    model = SimpleNamespace(mesh_device=mesh, norm=fake_norm, lm_head=fake_lm_head)
+    monkeypatch.setattr(model_module, "ttnn", FakeTTNN)
+    monkeypatch.setattr(
+        model_module,
+        "_all_gather_rmsnorm_tensor",
+        lambda norm, tensor: calls.append(("all_gather", norm, tensor)) or all_gathered,
+    )
+
+    result = Llama32_1BTransformer1D.post_process_batched_prefill_output(
+        model,
+        hidden,
+        last_token_idx_list=[3, 7, 11, 0],
+        padded_batch=4,
+        prefill_seq_len=32,
+    )
+
+    assert result is output
+    selector = calls[0][1]
+    assert selector.shape == (1, 1, 32, 128)
+    assert selector.dtype == torch.bfloat16
+    assert torch.count_nonzero(selector).item() == 4
+    assert [selector[0, 0, row].nonzero().item() for row in range(4)] == [3, 39, 75, 96]
+    assert calls[0][2] == {
+        "device": mesh,
+        "dtype": "bfloat16",
+        "layout": "tile",
+        "mesh_mapper": "replicate",
+    }
+    assert [call[0] for call in calls] == [
+        "from_torch",
+        "matmul",
+        "deallocate",
+        "norm",
+        "all_gather",
+        "lm_head",
+        "to_memory_config",
+    ]
+    assert calls[1][1:3] == (selector_tt, hidden)
+    assert calls[2] == ("deallocate", selector_tt)
+    assert calls[3] == ("norm", gathered)
+    assert calls[4] == ("all_gather", fake_norm, normalized)
+    assert calls[5] == ("lm_head", all_gathered)

code/models/common/tests/models/llama32_1b/test_demo_warmup.py

@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from models.common.llm_runtime.config import TraceConfig
+
+_DEMO_PATH = "models/common/tests/demos/llama32_1b/demo.py"
+_DEMO_TREE = ast.parse(Path(_DEMO_PATH).read_text(encoding="utf-8"), filename=_DEMO_PATH)
+
+
+def _demo_function(name, namespace=None):
+    function = next(node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == name)
+    namespace = {} if namespace is None else namespace
+    exec(compile(ast.Module(body=[function], type_ignores=[]), _DEMO_PATH, "exec"), namespace)
+    return namespace[name]
+
+
+_warmup_demo_executor = _demo_function("_warmup_demo_executor")
+
+
+@pytest.mark.parametrize("lane_group", [False, True])
+def test_demo_warmup_compiles_eager_programs_before_trace_capture(lane_group):
+    calls = []
+    config = SimpleNamespace(trace=TraceConfig("all"), device_sampling_enabled=True)
+
+    def warmup_prefill(**kwargs):
+        calls.append(("prefill", kwargs))
+
+    def warmup_decode(**kwargs):
+        calls.append(("decode", kwargs))
+
+    executor = SimpleNamespace(
+        warmup_model_prefill=warmup_prefill,
+        warmup_model_decode=warmup_decode,
+        max_batch_size=4,
+    )
+    if lane_group:
+        executor.lanes = [SimpleNamespace(config=config)]
+    else:
+        executor.config = config
+        executor.model = SimpleNamespace(config=SimpleNamespace(max_batch_size=4))
+
+    kv_cache = object()
+    page_table = SimpleNamespace(shape=(4, 8))
+    _warmup_demo_executor(executor, kv_cache=kv_cache, page_table=page_table)
+
+    assert [(kind, kwargs["enable_trace"]) for kind, kwargs in calls] == [
+        ("decode", False),
+        ("prefill", False),
+        ("prefill", True),
+        ("decode", True),
+    ]
+    for _, kwargs in calls:
+        assert kwargs["kv_cache"] is kv_cache
+        assert kwargs["can_sample_on_device"] is True
+    for kind, kwargs in calls:
+        if kind == "decode":
+            assert kwargs["max_batch_size"] == 4
+            assert kwargs["num_blocks"] == 8
+
+
+def _called_names(function_name):
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == function_name
+    )
+    return [
+        node.func.id for node in ast.walk(function) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    ]
+
+
+@pytest.mark.parametrize(("data_parallel", "expected_tp_devices"), [(4, 2), (8, 1)])
+def test_t3k_dp_topology_preserves_supported_tp_lanes(data_parallel, expected_tp_devices):
+    helper = _demo_function("_dp_tp_devices_or_skip", {"pytest": pytest, "ttnn": SimpleNamespace(MeshDevice=object)})
+    mesh = SimpleNamespace(get_num_devices=lambda: 8)
+
+    assert helper(mesh, data_parallel) == expected_tp_devices
+
+
+def test_t3k_dp2_skips_unsupported_tp4_lanes(expect_error):
+    helper = _demo_function("_dp_tp_devices_or_skip", {"pytest": pytest, "ttnn": SimpleNamespace(MeshDevice=object)})
+    mesh = SimpleNamespace(get_num_devices=lambda: 8)
+
+    with expect_error(pytest.skip.Exception, "creates TP4 lanes"):
+        helper(mesh, 2)
+
+
+def test_dp_build_validates_and_resolves_cache_from_each_lane_submesh():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_dp_smoke"
+    )
+    lane_loop = next(
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.For) and isinstance(node.target, ast.Name) and node.target.id == "sm"
+    )
+    calls = [node for node in ast.walk(lane_loop) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)]
+    call_names = [node.func.id for node in calls]
+    assert "_skip_unless_heads_divide_mesh" in call_names
+    assert "lazy_weight_cache_dir_for_demo" in call_names
+
+    from_pretrained_call = next(node for node in calls if node.func.id == "from_pretrained")
+    cache_dir = next(keyword.value for keyword in from_pretrained_call.keywords if keyword.arg == "cache_dir")
+    assert isinstance(cache_dir, ast.Name)
+    assert cache_dir.id == "lane_cache_dir"
+
+
+@pytest.mark.parametrize("function_name", ["_run_perf_benchmark", "_run_dp_smoke"])
+def test_traced_demo_paths_warm_up_before_benchmark(function_name):
+    calls = _called_names(function_name)
+    assert calls.index("_warmup_demo_executor") < calls.index("run_perf_benchmark")
+
+
+def test_eval_repeat_warms_each_fresh_executor():
+    calls = _called_names("_run_eval_repeat_batch32")
+    assert "_warmup_demo_executor" in calls
+
+
+def test_perf_path_enables_pipeline_readback_by_default():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_perf_benchmark"
+    )
+    benchmark_call = next(
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "run_perf_benchmark"
+    )
+    keywords = {keyword.arg: keyword.value for keyword in benchmark_call.keywords}
+    assert isinstance(keywords["pipeline_readback"], ast.Name)
+    assert keywords["pipeline_readback"].id == "pipeline_readback"

code/models/common/tests/models/llama32_1b/test_hf_adaptor.py

@@ -0,0 +1,264 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+from transformers import LlamaConfig, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+
+from models.common.models.llama32_1b import hf_adaptor
+from models.common.models.llama32_1b import model as llama_model
+from models.common.models.llama32_1b import weight_utils
+from models.common.models.llama32_1b.hf_adaptor import (
+    Llama32_1BForCausalLM,
+    Llama32_1BRuntimeConfig,
+    _trace_seq_lens,
+    convert_hf_model_weights,
+)
+
+LLAMA32_ROPE_PARAMETERS = {
+    "rope_type": "llama3",
+    "factor": 32.0,
+    "low_freq_factor": 1.0,
+    "high_freq_factor": 4.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+}
+
+
+def test_runtime_config_preserves_trace_and_batched_prefill_policy():
+    runtime = Llama32_1BRuntimeConfig(
+        model_name="Llama-3.2-1B-Instruct",
+        model_cache_path=None,
+        max_prefill_chunk_size=2048,
+        max_context_len=131072,
+        max_seq_len=4096,
+        trace_prefill_supported_seq_lens=(128, 1024),
+    )
+    assert runtime.can_enable_trace(128, num_cached_tokens=32)
+    assert runtime.can_enable_trace(1024)
+    assert not runtime.can_enable_trace(2048)
+    assert runtime.supports_batched_prefill
+    assert runtime.max_prefill_batch_size == 32
+    assert runtime.batched_prefill_batched_extract
+
+
+def test_trace_matrix_is_device_specific_and_bounded():
+    assert _trace_seq_lens(1, 2048, 4096) == (128,)
+    assert _trace_seq_lens(2, 2048, 4096) == (128, 1024)
+    assert _trace_seq_lens(8, 2048, 4096) == (128, 1024)
+
+
+def test_product_binds_runtime_config_unconditionally():
+    model = SimpleNamespace(config=SimpleNamespace(max_seq_len=4096), model_args=None)
+    tokenizer = SimpleNamespace(stop_tokens=[128001])
+    runtime = Llama32_1BRuntimeConfig(
+        model_name="model",
+        model_cache_path=None,
+        max_prefill_chunk_size=2048,
+        max_context_len=131072,
+        max_seq_len=4096,
+        trace_prefill_supported_seq_lens=(128,),
+    )
+    product = Llama32_1BForCausalLM(model=model, tokenizer=tokenizer, runtime_config=runtime)
+    assert model.model_args is runtime
+    assert product.generation_config.stop_token_ids == (128001,)
+    assert product.model_name == "model"
+    assert product.model_cache_path is None
+    assert product.max_seq_len == 4096
+    assert product.max_context_len == 131072
+
+
+def test_hf_attention_and_mlp_weights_match_reference_layouts():
+    hidden_size = 128
+    num_attention_heads = 32
+    num_key_value_heads = 8
+    num_devices = 8
+    head_dim = hidden_size // num_attention_heads
+    kv_width = num_key_value_heads * head_dim
+    config = SimpleNamespace(
+        num_attention_heads=num_attention_heads,
+        num_key_value_heads=num_key_value_heads,
+        hidden_size=hidden_size,
+    )
+    q = torch.arange(hidden_size * hidden_size, dtype=torch.float32).reshape(hidden_size, hidden_size)
+    k = torch.arange(kv_width * hidden_size, dtype=torch.float32).reshape(kv_width, hidden_size) + 100_000
+    v = k + 100_000
+    o = q + 300_000
+    attention = SimpleNamespace(
+        config=config,
+        q_proj=SimpleNamespace(weight=q),
+        k_proj=SimpleNamespace(weight=k),
+        v_proj=SimpleNamespace(weight=v),
+        o_proj=SimpleNamespace(weight=o),
+    )
+
+    wqkv, wo = weight_utils.attention_wqkv_wo_from_hf_layer(attention, num_devices=num_devices)
+    q_meta = q.view(num_attention_heads, 2, head_dim // 2, hidden_size).transpose(1, 2).reshape(q.shape).T
+    k_meta = k.view(num_key_value_heads, 2, head_dim // 2, hidden_size).transpose(1, 2).reshape(k.shape).T
+    expected_qkv = (
+        torch.cat(
+            [
+                torch.cat(parts, dim=-1)
+                for parts in zip(
+                    torch.chunk(q_meta, num_devices, dim=1),
+                    torch.chunk(k_meta, num_devices, dim=1),
+                    torch.chunk(v.T, num_devices, dim=1),
+                )
+            ],
+            dim=-1,
+        )
+        .unsqueeze(0)
+        .unsqueeze(0)
+    )
+    assert wqkv.shape == (1, 1, hidden_size, hidden_size + 2 * kv_width)
+    torch.testing.assert_close(wqkv, expected_qkv)
+    torch.testing.assert_close(wo, o.T.unsqueeze(0).unsqueeze(0))
+
+    gate = torch.arange(48, dtype=torch.float32).reshape(6, 8)
+    down = torch.arange(48, dtype=torch.float32).reshape(8, 6)
+    up = gate + 100
+    mlp = SimpleNamespace(
+        gate_proj=SimpleNamespace(weight=gate),
+        down_proj=SimpleNamespace(weight=down),
+        up_proj=SimpleNamespace(weight=up),
+    )
+    w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(mlp)
+    torch.testing.assert_close(w1, gate.T)
+    torch.testing.assert_close(w2, down.T)
+    torch.testing.assert_close(w3, up.T)
+
+
+def test_hf_rope_tables_match_real_llama32_scaled_rotary_reference():
+    head_dim = 64
+    table_len = LLAMA32_ROPE_PARAMETERS["original_max_position_embeddings"] + 128
+    config = LlamaConfig(
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=head_dim,
+        max_position_embeddings=131072,
+        rope_parameters=LLAMA32_ROPE_PARAMETERS,
+    )
+    rotary = LlamaRotaryEmbedding(config)
+
+    cos, sin = weight_utils.build_rope_cos_sin_torch(
+        rotary, table_len=table_len, head_dim=head_dim, dtype=torch.bfloat16
+    )
+    x = torch.zeros(1, 1, table_len, head_dim, dtype=torch.bfloat16)
+    position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0)
+    with torch.no_grad():
+        hf_cos, hf_sin = rotary(x, position_ids)
+    expected_cos = hf_cos.float().squeeze(0)[:, : head_dim // 2].repeat_interleave(2, dim=-1).unsqueeze(0).unsqueeze(0)
+    expected_sin = hf_sin.float().squeeze(0)[:, : head_dim // 2].repeat_interleave(2, dim=-1).unsqueeze(0).unsqueeze(0)
+
+    assert config.rope_parameters == LLAMA32_ROPE_PARAMETERS
+    assert cos.shape == sin.shape == (1, 1, table_len, head_dim)
+    assert cos.dtype == sin.dtype == torch.bfloat16
+    torch.testing.assert_close(cos, expected_cos.to(torch.bfloat16))
+    torch.testing.assert_close(sin, expected_sin.to(torch.bfloat16))
+
+
+def test_convert_hf_model_weights_covers_real_nonempty_llama_layer():
+    config = LlamaConfig(
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=1,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=4,
+        vocab_size=128,
+        max_position_embeddings=131072,
+        rope_parameters=LLAMA32_ROPE_PARAMETERS,
+        tie_word_embeddings=True,
+    )
+    hf = LlamaForCausalLM(config).eval()
+    weights = convert_hf_model_weights(
+        hf,
+        config,
+        n_layers=1,
+        num_devices=8,
+        rope_table_len=128,
+        head_dim=4,
+    )
+
+    assert len(weights.layers) == 1
+    layer_weights = weights.layers[0]
+    assert layer_weights.wqkv.shape == (1, 1, 128, 192)
+    assert layer_weights.wo.shape == (1, 1, 128, 128)
+    assert layer_weights.w1.shape == (128, 256)
+    assert layer_weights.w2.shape == (256, 128)
+    assert layer_weights.w3.shape == (128, 256)
+    assert layer_weights.attention_norm.shape == layer_weights.ff_norm.shape == (128,)
+    assert weights.embedding.shape == (1, 1, 128, 128)
+    assert weights.rope_cos.shape == weights.rope_sin.shape == (1, 1, 128, 4)
+    assert weights.final_norm.shape == (128,)
+    torch.testing.assert_close(weights.lm_head, hf.model.embed_tokens.weight.detach().to(torch.bfloat16))
+
+
+def test_tied_embedding_is_explicit_lm_head_construction_source():
+    class Rotary:
+        def __call__(self, x, position_ids):
+            return torch.ones(1, position_ids.shape[-1], x.shape[-1]), torch.zeros(
+                1, position_ids.shape[-1], x.shape[-1]
+            )
+
+    tied_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
+    decoy_lm_head = torch.full((6, 4), -99.0)
+    base = SimpleNamespace(
+        embed_tokens=SimpleNamespace(weight=tied_weight),
+        rotary_emb=Rotary(),
+        layers=[],
+        norm=SimpleNamespace(weight=torch.ones(4)),
+    )
+    hf = SimpleNamespace(model=base, lm_head=SimpleNamespace(weight=decoy_lm_head))
+    config = SimpleNamespace(tie_word_embeddings=True)
+    weights = convert_hf_model_weights(
+        hf,
+        config,
+        n_layers=0,
+        num_devices=1,
+        rope_table_len=8,
+        head_dim=4,
+    )
+
+    torch.testing.assert_close(weights.lm_head, tied_weight.to(torch.bfloat16))
+    assert not torch.equal(weights.lm_head, decoy_lm_head.to(torch.bfloat16))
+    assert weights.embedding.shape == (1, 1, 6, 4)
+
+
+def test_config_builder_is_owned_by_model_module():
+    assert hf_adaptor.build_llama32_1b_transformer_1d_config is llama_model.build_llama32_1b_transformer_1d_config
+    assert llama_model.build_llama32_1b_transformer_1d_config.__module__ == llama_model.__name__
+
+
+def test_post_attention_norm_program_and_memory_use_same_mlp_grid(monkeypatch):
+    grid = SimpleNamespace(num_cores=64)
+    program = object()
+    memory = object()
+    captured = {}
+
+    monkeypatch.setattr(llama_model, "get_padded_hidden_dim", lambda *_: 8192)
+    monkeypatch.setattr(llama_model, "_dram_shard_core_grid_k_n", lambda *_: grid)
+    monkeypatch.setattr(
+        llama_model,
+        "_create_sharded_norm_program_config",
+        lambda dim, selected_grid, rows, tile: captured.update(program=(dim, selected_grid, rows, tile)) or program,
+    )
+    monkeypatch.setattr(
+        llama_model.ttnn,
+        "create_sharded_memory_config",
+        lambda shape, selected_grid, *args, **kwargs: captured.update(memory=(shape, selected_grid)) or memory,
+    )
+
+    assert llama_model._post_attn_norm_decode_configs(
+        dim=2048,
+        hidden_dim=8192,
+        num_devices=1,
+        max_batch_size=1,
+    ) == (program, memory)
+    assert captured["program"] == (2048, grid, 32, 32)
+    assert captured["memory"] == ((32, 32), grid)

code/models/common/tests/models/llama32_3b/test_batched_prefill_postprocess.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+
+import models.common.models.llama32_3b.model as model_module
+from models.common.models.llama32_3b.model import Llama32_3BTransformer1D
+
+
+def test_batched_prefill_postprocess_gathers_each_slots_last_token_before_norm(monkeypatch):
+    calls = []
+    hidden = object()
+    selector_tt = object()
+    gathered = object()
+    normalized = object()
+    all_gathered = object()
+    logits = object()
+    output = object()
+    mesh = SimpleNamespace(arch=lambda: "wormhole")
+
+    class FakeTTNN:
+        bfloat16 = "bfloat16"
+        TILE_LAYOUT = "tile"
+        DRAM_MEMORY_CONFIG = "dram"
+        MathFidelity = SimpleNamespace(HiFi4="hifi4")
+
+        @staticmethod
+        def ReplicateTensorToMesh(device):
+            assert device is mesh
+            return "replicate"
+
+        @staticmethod
+        def from_torch(selector, **kwargs):
+            calls.append(("from_torch", selector.clone(), kwargs))
+            return selector_tt
+
+        @staticmethod
+        def init_device_compute_kernel_config(arch, **kwargs):
+            assert arch == "wormhole"
+            return kwargs
+
+        @staticmethod
+        def matmul(lhs, rhs, **kwargs):
+            calls.append(("matmul", lhs, rhs, kwargs))
+            return gathered
+
+        @staticmethod
+        def deallocate(tensor):
+            calls.append(("deallocate", tensor))
+
+        @staticmethod
+        def to_memory_config(tensor, memory_config):
+            calls.append(("to_memory_config", tensor, memory_config))
+            return output
+
+    fake_norm = SimpleNamespace(prefill_forward=lambda tensor: calls.append(("norm", tensor)) or normalized)
+    fake_lm_head = SimpleNamespace(
+        config=SimpleNamespace(input_memcfg=None),
+        forward=lambda tensor: calls.append(("lm_head", tensor)) or logits,
+    )
+    model = SimpleNamespace(mesh_device=mesh, norm=fake_norm, lm_head=fake_lm_head)
+    monkeypatch.setattr(model_module, "ttnn", FakeTTNN)
+    monkeypatch.setattr(
+        model_module,
+        "_all_gather_rmsnorm_tensor",
+        lambda norm, tensor: calls.append(("all_gather", norm, tensor)) or all_gathered,
+    )
+
+    result = Llama32_3BTransformer1D.post_process_batched_prefill_output(
+        model,
+        hidden,
+        last_token_idx_list=[3, 7, 11, 0],
+        padded_batch=4,
+        prefill_seq_len=32,
+    )
+
+    assert result is output
+    selector = calls[0][1]
+    assert selector.shape == (1, 1, 32, 128)
+    assert selector.dtype == torch.bfloat16
+    assert torch.count_nonzero(selector).item() == 4
+    assert [selector[0, 0, row].nonzero().item() for row in range(4)] == [3, 39, 75, 96]
+    assert calls[0][2] == {
+        "device": mesh,
+        "dtype": "bfloat16",
+        "layout": "tile",
+        "mesh_mapper": "replicate",
+    }
+    assert [call[0] for call in calls] == [
+        "from_torch",
+        "matmul",
+        "deallocate",
+        "norm",
+        "all_gather",
+        "lm_head",
+        "to_memory_config",
+    ]
+    assert calls[1][1:3] == (selector_tt, hidden)
+    assert calls[2] == ("deallocate", selector_tt)
+    assert calls[3] == ("norm", gathered)
+    assert calls[4] == ("all_gather", fake_norm, normalized)
+    assert calls[5] == ("lm_head", all_gathered)

code/models/common/tests/models/llama32_3b/test_demo_warmup.py

@@ -0,0 +1,218 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from models.common.llm_runtime.config import TraceConfig
+
+_DEMO_PATH = "models/common/tests/demos/llama32_3b/demo.py"
+_DEMO_TREE = ast.parse(Path(_DEMO_PATH).read_text(encoding="utf-8"), filename=_DEMO_PATH)
+
+
+def _demo_function(name, namespace=None):
+    function = next(node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == name)
+    namespace = {} if namespace is None else namespace
+    exec(compile(ast.Module(body=[function], type_ignores=[]), _DEMO_PATH, "exec"), namespace)
+    return namespace[name]
+
+
+_warmup_demo_executor = _demo_function("_warmup_demo_executor")
+
+
+@pytest.mark.parametrize("lane_group", [False, True])
+@pytest.mark.parametrize(
+    ("trace_mode", "expected_trace_calls"),
+    [
+        ("all", [("prefill", True), ("decode", True)]),
+        ("decode_only", [("decode", True)]),
+    ],
+)
+def test_demo_warmup_compiles_eager_programs_before_enabled_trace_capture(lane_group, trace_mode, expected_trace_calls):
+    calls = []
+    config = SimpleNamespace(trace=TraceConfig(trace_mode), device_sampling_enabled=True)
+
+    def warmup_prefill(**kwargs):
+        calls.append(("prefill", kwargs))
+
+    def warmup_decode(**kwargs):
+        calls.append(("decode", kwargs))
+
+    executor = SimpleNamespace(
+        warmup_model_prefill=warmup_prefill,
+        warmup_model_decode=warmup_decode,
+        max_batch_size=4,
+    )
+    if lane_group:
+        executor.lanes = [SimpleNamespace(config=config)]
+    else:
+        executor.config = config
+        executor.model = SimpleNamespace(config=SimpleNamespace(max_batch_size=4))
+
+    kv_cache = object()
+    page_table = SimpleNamespace(shape=(4, 8))
+    _warmup_demo_executor(executor, kv_cache=kv_cache, page_table=page_table)
+
+    eager_calls = [("decode", False), ("prefill", False)]
+    assert [(kind, kwargs["enable_trace"]) for kind, kwargs in calls] == eager_calls + expected_trace_calls
+    for _, kwargs in calls:
+        assert kwargs["kv_cache"] is kv_cache
+        assert kwargs["can_sample_on_device"] is True
+    for kind, kwargs in calls:
+        if kind == "decode":
+            assert kwargs["max_batch_size"] == 4
+            assert kwargs["num_blocks"] == 8
+
+
+@pytest.mark.parametrize(
+    ("num_devices", "traced", "expected_mode"),
+    [(1, True, "decode_only"), (2, True, "all"), (8, True, "all"), (1, False, "none")],
+)
+def test_create_executor_preserves_3b_trace_device_matrix(num_devices, traced, expected_mode):
+    captured = {}
+
+    def executor_config(**kwargs):
+        captured.update(kwargs)
+        return SimpleNamespace(**kwargs)
+
+    namespace = {
+        "Llama32_3BTransformer1D": object,
+        "Llama32_3BExecutor": lambda model, model_args, config: config,
+        "Llama32_3BExecutorConfig": executor_config,
+        "PagedKVCacheConfig": lambda **kwargs: SimpleNamespace(**kwargs),
+        "TraceConfig": TraceConfig,
+        "WarmupConfig": lambda: object(),
+    }
+    create_executor = _demo_function("create_executor", namespace)
+    model = SimpleNamespace(
+        model_args=object(),
+        config=SimpleNamespace(
+            max_seq_len=4096,
+            max_batch_size=32,
+            num_devices=num_devices,
+            block_configs=[SimpleNamespace(attention_config=SimpleNamespace(kv_cache_dtype=object()))],
+        ),
+    )
+
+    result = create_executor(model, traced=traced, device_sampling_enabled=True)
+
+    assert result.trace.mode == expected_mode
+    assert captured["device_sampling_enabled"] is True
+    assert captured["paged_kv_cache"].num_blocks == 4096
+
+
+def _called_names(function_name):
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == function_name
+    )
+    return [
+        node.func.id for node in ast.walk(function) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    ]
+
+
+@pytest.mark.parametrize(("data_parallel", "expected_tp_devices"), [(4, 2), (8, 1)])
+def test_t3k_dp_topology_preserves_supported_tp_lanes(data_parallel, expected_tp_devices):
+    helper = _demo_function("_dp_tp_devices_or_skip", {"pytest": pytest, "ttnn": SimpleNamespace(MeshDevice=object)})
+    mesh = SimpleNamespace(get_num_devices=lambda: 8)
+
+    assert helper(mesh, data_parallel) == expected_tp_devices
+
+
+def test_t3k_dp2_skips_unsupported_tp4_lanes(expect_error):
+    helper = _demo_function("_dp_tp_devices_or_skip", {"pytest": pytest, "ttnn": SimpleNamespace(MeshDevice=object)})
+    mesh = SimpleNamespace(get_num_devices=lambda: 8)
+
+    with expect_error(pytest.skip.Exception, "creates TP4 lanes"):
+        helper(mesh, 2)
+
+
+def test_dp_build_validates_and_resolves_cache_from_each_lane_submesh():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_dp_smoke"
+    )
+    lane_loop = next(
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.For) and isinstance(node.target, ast.Name) and node.target.id == "sm"
+    )
+    calls = [node for node in ast.walk(lane_loop) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)]
+    call_names = [node.func.id for node in calls]
+    assert "_skip_unless_heads_divide_mesh" in call_names
+    assert "lazy_weight_cache_dir_for_demo" in call_names
+
+    from_pretrained_call = next(node for node in calls if node.func.id == "from_pretrained")
+    cache_dir = next(keyword.value for keyword in from_pretrained_call.keywords if keyword.arg == "cache_dir")
+    assert isinstance(cache_dir, ast.Name)
+    assert cache_dir.id == "lane_cache_dir"
+
+
+@pytest.mark.parametrize("function_name", ["_run_perf_benchmark", "_run_dp_smoke"])
+def test_traced_demo_paths_warm_up_before_benchmark(function_name):
+    calls = _called_names(function_name)
+    assert calls.index("_warmup_demo_executor") < calls.index("run_perf_benchmark")
+
+
+def test_eval_repeat_warms_each_fresh_executor():
+    calls = _called_names("_run_eval_repeat_batch32")
+    assert "_warmup_demo_executor" in calls
+
+
+def test_perf_path_enables_pipeline_readback_by_default():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_perf_benchmark"
+    )
+    benchmark_call = next(
+        node
+        for node in ast.walk(function)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "run_perf_benchmark"
+    )
+    keywords = {keyword.arg: keyword.value for keyword in benchmark_call.keywords}
+    assert isinstance(keywords["pipeline_readback"], ast.Name)
+    assert keywords["pipeline_readback"].id == "pipeline_readback"
+
+
+def test_create_model_preserves_reduced_layer_diagnostic_override(monkeypatch):
+    captured = {}
+    model = SimpleNamespace()
+
+    def from_pretrained(*args, **kwargs):
+        captured.update(kwargs)
+        return SimpleNamespace(model=model, tokenizer=object())
+
+    namespace = {
+        "Path": Path,
+        "Llama32_3BTransformer1D": object,
+        "LLAMA32_3B_ACCURACY": object(),
+        "LLAMA32_3B_PERFORMANCE": object(),
+        "_skip_unless_heads_divide_mesh": lambda *_: None,
+        "from_pretrained": from_pretrained,
+        "os": os,
+        "pytest": pytest,
+        "ttnn": SimpleNamespace(MeshDevice=object),
+    }
+    create_model = _demo_function("create_model", namespace)
+    monkeypatch.setenv("LLAMA32_3B_DEMO_NUM_LAYERS", "3")
+
+    assert create_model(object(), "performance", Path("cache")) is model
+    assert captured["n_layers"] == 3
+
+
+def test_token_accuracy_cleans_up_executor_in_finally():
+    function = next(
+        node for node in _DEMO_TREE.body if isinstance(node, ast.FunctionDef) and node.name == "_run_token_accuracy"
+    )
+    cleanup_finally = [
+        statement
+        for node in ast.walk(function)
+        if isinstance(node, ast.Try)
+        for statement in node.finalbody
+        if isinstance(statement, ast.Expr)
+        and isinstance(statement.value, ast.Call)
+        and isinstance(statement.value.func, ast.Attribute)
+        and statement.value.func.attr == "cleanup"
+    ]
+    assert len(cleanup_finally) == 1